本文详解如何将多个独立矩阵乘法任务在cpu端通过内存视图（`as_strided`）实现零开销并行化，并扩展至cupy gpu环境，支持同步批处理与异步提交，兼顾性能、可读性与通用性。
在科学计算与机器学习中，常需对一组结构相似的矩阵对执行独立的乘法运算（如滑动窗口、多头注意力子块、批量参数投影等）。虽然NumPy底层BLAS已自动并行化单次np.dot，但外层Python循环仍是串行瓶颈。真正的加速关键在于避免显式循环，转而利用内存布局优化与硬件级并行能力。
✅ CPU端：零拷贝向量化——as_strided + np.matmul
核心思想是将原始矩阵 b（形状 (20, 1000)）按列切片（每片宽20列）构造成一个“虚拟三维张量”，其形状为 (981, 20, 20)，其中每个 (20, 20) 切片对应 b[:, i:i+20]。借助 numpy.lib.stride_tricks.as_strided，我们仅修改内存步长（strides）和形状（shape），不复制数据，实现零开销视图构造：
import numpy as np
from numpy.lib.stride_tricks import as_strided
def batch_matmul_vectorized(a: np.ndarray, b: np.ndarray, window_size: int = 20) -> np.ndarray:
"""
向量化批量矩阵乘法：a @ b[:, i:i+window_size] for all valid i
Parameters:
-----------
a : (M, K) ndarray
b : (K, N) ndarray
window_size : int, 每次取 b 的 window_size 列
Returns:
--------
(num_windows, M, window_size) ndarray
"""
assert b.shape[0] == a.shape[1], "Inner dimensions must match"
num_windows = b.shape[1] - window_size + 1
# 构造 b 的滑动窗口视图：(num_windows, K, window_size)
strides = (b.strides[1],) + b.strides # 沿列方向步进
b_windows = as_strided(
b,
shape=(num_windows, b.shape[0], window_size),
strides=strides
)
# 扩展 a 维度以支持广播：(1, M, K) → (num_windows, M, K)
a_expanded = a[np.newaxis, ...]
# 批量矩阵乘法：(num_windows, M, K) @ (num_windows, K, window_size) → (num_windows, M, window_size)
return np.matmul(a_expanded, b_windows)
# 示例调用
a = np.random.rand(10, 20)
b = np.random.rand(20, 1000)
result = batch_matmul_vectorized(a, b) # shape: (981, 10, 20)⚠️ 重要注意事项 as_strided 是“危险但强大”的工具：错误的 strides 或 shape 可能导致内存越界或静默错误。务必确保 num_windows > 0 且 window_size
? GPU端：CuPy批处理与异步流控制
迁移到GPU时，CuPy 提供与NumPy几乎一致的API，且原生支持批量矩阵乘（cp.matmul）及CUDA流（stream）异步调度：
import cupy as cp
def batch_matmul_cupy(a: cp.ndarray, b: cp.ndarray, window_size: int = 20,
stream: cp.cuda.Stream = None) -> cp.ndarray:
"""CuPy版批量矩阵乘，支持可选CUDA流异步执行"""
num_windows = b.shape[1] - window_size + 1
strides = (b.strides[1],) + b.strides
b_windows = cp.lib.stride_tricks.as_strided(
b,
shape=(num_windows, b.shape[0], window_size),
strides=strides
)
a_expanded = a[np.newaxis, ...]
# 指定流（若提供），实现异步计算
if stream is not None:
with stream:
result = cp.matmul(a_expanded, b_windows)
stream.synchronize() # 可选：等待完成
else:
result = cp.matmul(a_expanded, b_windows)
return result
# 示例：使用默认流（同步）
a_gpu = cp.asarray(a)
b_gpu = cp.asarray(b)
result_gpu = batch_matmul_cupy(a_gpu, b_gpu)
# 示例：使用自定义流实现异步重叠计算与数据传输
with cp.cuda.Stream() as stream:
# 异步拷贝 + 计算（需配合 cp.asarray(..., stream=stream)）
a_async = cp.asarray(a, dtype=cp.float64, stream=stream)
b_async = cp.asarray(b, dtype=cp.float64, stream=stream)
out_async = batch_matmul_cupy(a_async, b_async, stream=stream)
# 此处可插入其他CPU任务，GPU计算在后台运行
stream.synchronize()? GPU最佳实践提示
- 批量尺寸建议 ≥ 32，以充分占用GPU SM单元；过小的 window_size 或 num_windows 会导致内核启动开销占比过高。
- 使用 cp.cuda.MemoryPool 管理显存，避免频繁分配释放。
- 对于“任意矩阵对”（非滑动窗口），可预分配 cp.stack([a1, a2, ..., an]) 和 cp.stack([b1, b2, ..., bn])，再调用 cp.matmul(A_batch, B_batch) 实现真·批量乘法。
? 总结
- 首选方案（推荐）：用 as_strided + np.matmul / cp.matmul 实现零拷贝向量化批处理，兼顾极致性能与简洁代码；
- 扩展场景：当矩阵对无共享维度（即非滑动窗口）时，统一堆叠为 (N, M, K) 和 (N, K, P) 后调用批矩阵乘；
- GPU进阶：结合CUDA流与显存池，实现计算-传输重叠，最大化吞吐；
- 慎用替代：multiprocessing 或 threading 在NumPy/CuPy场景下通常画蛇添足——前者因进程间数据拷贝昂贵，后者受GIL限制无法加速计算。
向量化不是魔法，而是对内存与计算范式的深度理解。掌握 as_strided 与 matmul 的组合，你已握有解锁高性能线性代数的关键钥匙。
