昇腾CANN ops-blas Batched GEMM：多头注意力的小矩阵乘批处理实战
Transformer 的 Multi-Head Attention 有 H 个注意力头——每个头独立做矩阵乘（Qh×Kh^T、Attn×Vh）。H=32 时，一个 BatchNorm 后面紧跟着 32 个小矩阵乘（每个头独立）。单独启动 32 次 GEMM 会有 32 次 launch 开销（~50μs/次 → 1.6ms 总开销），加上 32 次 kernel 启动带来的流水线 flush。
ops-blas 的 Batched GEMM 把 32 个小矩阵乘合并成一个 kernel——一次 launch 处理全部 32 个头。
Batched GEMM 的三种策略
ops-blas 根据 batched GEMM 的形状自动选择策略：
策略选择逻辑 if (batch_count >= 32 && M * N * K < 4096): → 策略 1：Interleaved Batching（交错批处理） 把 32 个小 GEMM 交织在一个 block 内执行 elif (batch_count < 16 && M * N * K >= 4096): → 策略 2：Parallel Batching（并行批处理） 给每个小 GEMM 分配独立 block else: → 策略 3：Hybrid Batching（混合批处理） 分组内交错的组外并行策略 1：Interleaved Batching
// ops-blas/kernels/batched_gemm_interleaved.cpp__aicore__voidBatchedGEMMInterleaved(GlobalTensor<float16>&A_batched,// [batch, M, K]GlobalTensor<float16>&B_batched,// [batch, K, N]GlobalTensor<float16>&C_batched,// [batch, M, N]intbatch,intM,intN,intK){// 每个 block 处理一个 batch 的 GEMMfor(intb=0;b<batch;b++){intblock_id=b%gridDim.x;// 轮询分配 block// 在 L1 中交错存储 32 个 batch 的 tile// 单个 tile 大小 = tile_M × tile_K = 16 × 16 = 256 elementsLocalTensor<float16>A_tile(tile_M*tile_K);LocalTensor<float16>B_tile(tile_K*tile_N);LocalTensor<float16>C_tile(tile_M*tile_N);intA_offset=b*M*K;intB_offset=b*K*N;intC_offset=b*M*N;// 分块矩阵乘for(intm=0;m<M;m+=tile_M){for(intn=0;n<N;n+=tile_N){// 初始化累加器C_tile=0.0f;for(intk=0;k<K;k+=tile_K){// 加载 A 和 B 的 tile 到 L1DataCopy(A_tile,A_batched+A_offset+m*K+k,tile_M*tile_K);DataCopy(B_tile,B_batched+B_offset+k*N+n,tile_K*tile_N);// Cube 单元矩阵乘累加MMA(C_tile,A_tile,B_tile,tile_M,tile_N,tile_K);}// 写回结果DataCopy(C_batched+C_offset+m*N+n,C_tile,tile_M*tile_N);}}}}策略 2：Parallel Batching
// ops-blas/kernels/batched_gemm_parallel.cpp__aicore__voidBatchedGEMMParallel(GlobalTensor<float16>&A_batched,// [batch, M, K]GlobalTensor<float16>&B_batched,// [batch, K, N]GlobalTensor<float16>&C_batched,// [batch, M, N]intbatch,intM,intN,intK){// 每个 block 处理一个独立的 batch（不是所有 block 处理同一 batch）// block 分配：block_id = b % num_batch_blocks// num_batch_blocks = gridDim.x / batchintnum_batch_blocks=gridDim.x/batch;if(num_batch_blocks<1)num_batch_blocks=1;// 每个 batch 有 num_batch_blocks 个 block 在并行处理intbatch_id=blockIdx.x/num_batch_blocks;intbatch_block=blockIdx.x%num_batch_blocks;intA_offset=batch_id*M*K;intB_offset=batch_id*K*N;intC_offset=batch_id*M*N;// batch_block 决定此 block 处理矩阵的哪一部分// 把 M 维度均分给 num_batch_blocks 个 blockintm_start=batch_block*(M/num_batch_blocks);intm_end=(batch_block+1)*(M/num_batch_blocks);for(intm=m_start;m<m_end;m+=tile_M){for(intn=0;n<N;n+=tile_N){LocalTensor<float16>C_tile(tile_M*tile_N);C_tile=0.0f;for(intk=0;k<K;k+=tile_K){LocalTensor<float16>A_tile(tile_M*tile_K);LocalTensor<float16>B_tile(tile_K*tile_N);DataCopy(A_tile,A_batched+A_offset+m*K+k,tile_M*tile_K);DataCopy(B_tile,B_batched+B_offset+k*N+n,tile_K*tile_N);MMA(C_tile,A_tile,B_tile,tile_M,tile_N,tile_K);}DataCopy(C_batched+C_offset+m*N+n,C_tile,tile_M*tile_N);}}}策略 3：Hybrid Batching
// ops-blas/kernels/batched_gemm_hybrid.cpp__aicore__voidBatchedGEMMHybrid(GlobalTensor<float16>&A_batched,// [batch, M, K]GlobalTensor<float16>&B_batched,// [batch, K, N]GlobalTensor<float16>&C_batched,// [batch, M, N]intbatch,intM,intN,intK){// 分组：每 group_size 个 batch 为一组// 组内用 Interleaved（充分利用 L1），组间用 Parallelintgroup_size=4;// 每组 4 个 batchintnum_groups=(batch+group_size-1)/group_size;intgroup_id=blockIdx.x%num_groups;// 每个 block 处理一个 group// 组间并行处理intbatch_start=group_id*group_size;intbatch_end=min(batch_start+group_size,batch);// 组内 Interleavedfor(intb=batch_start;b<batch_end;b++){intA_offset=b*M*K;intB_offset=b*K*N;intC_offset=b*M*N;// 分块矩阵乘（同 Interleaved 策略）for(intm=0;m<M;m+=tile_M){for(intn=0;n<N;n+=tile_N){LocalTensor<float16>C_tile(tile_M*tile_N);C_tile=0.0f;for(intk=0;k<K;k+=tile_K){LocalTensor<float16>A_tile(tile_M*tile_K);LocalTensor<float16>B_tile(tile_K*tile_N);DataCopy(A_tile,A_batched+A_offset+m*K+k,tile_M*tile_K);DataCopy(B_tile,B_batched+B_offset+k*N+n,tile_K*tile_N);MMA(C_tile,A_tile,B_tile,tile_M,tile_N,tile_K);}DataCopy(C_batched+C_offset+m*N+n,C_tile,tile_M*tile_N);}}}}Multi-Head Attention 的 Batched GEMM 应用
Transformer 中 Multi-Head Attention 的三种 GEMM 都可以用 Batched GEMM 加速：
# PyTorch 自动路由到 ops-blas 的 Batched GEMMimporttorchimporttorch_npu# MHA 的三个 GEMM 步骤# 输入：x [batch, seq, d_model] (如 [1, 2048, 4096])# H=32 heads, d_head = d_model // H = 128# 1. QKV projection（每个头独立，共 3H 个小 GEMM）# x @ W_q[head] → Q[head] [batch, seq, d_head]# 转成 batched form: [batch*seq, d_model] @ [3, head, d_model, d_head]qkv=torch.nn.functional.linear(x,W_qkv)# 底层用 Batched GEMM# 2. Attention score（每个头独立，H 个小 GEMM）# Q[head] @ K[head]^T → scores[head] [batch, seq, seq]# batched form: [head, batch*seq, d_head] @ [head, d_head, batch*seq]attn_scores=torch.bmm(Q.reshape(-1,seq,d_head).transpose(0,1),K.reshape(-1,seq,d_head).transpose(0,1).transpose(1,2))# 底层用 Batched GEMM，一次 launch 处理 H 个头# 3. Output projection（每个头独立，H 个小 GEMM）# attn[head] @ V[head] → output[head] [batch, seq, d_head]# batched form 同理output=torch.bmm(attn_weights,V.reshape(-1,seq,d_head).transpose(0,1))关键：python 侧看到的torch.bmm(batched matrix multiplication)——底层自动映射到 ops-blas 的 Batched GEMM。
踩坑一：batch 维度的 stride 不连续
标准 Batched GEMM 假设 A 和 B 的 batch 维度是连续存储的 ([batch, M, K])。但 MHA 中 QKV projection 的 weight 是 [num_heads, d_model, d_head]——head 维度的 stride = d_model * d_head，不是 K * d_head。
修复：ops-blas 的 Batched GEMM 支持 stride 参数：
// 支持 stride 参数__aicore__voidBatchedGEMMStrided(GlobalTensor<float16>&A_batched,GlobalTensor<float16>&B_batched,GlobalTensor<float16>&C_batched,intbatch,intM,intN,intK,intstride_A,// A 的 batch stride（不连续时 > M*K）intstride_B,// B 的 batch strideintstride_C// C 的 batch stride){for(intb=0;b<batch;b++){// 使用 stride 替代 M*KintA_offset=b*stride_A;// 不是 b * M * KintB_offset=b*stride_B;intC_offset=b*stride_C;// ... 其余同 Interleaved}}PyTorch 侧：
# 非连续 batch → 指定 strideoutput=torch_npu.batched_gemm(A_strided,B_strided,stride_A=d_model*d_head,stride_B=d_head*seq)踩坑二：batch 中 GEMM 形状不一致
MHA 的 32 个头可能形状不同（某些头是 padding 头，不需要计算）。Batched GEMM 默认假设所有 batch 的 shape 相同——形状不一致时，padding 头浪费计算。
修复：使用 mask 跳过不需要的 batch：
__aicore__voidBatchedGEMMMasked(GlobalTensor<float16>&A_batched,GlobalTensor<float16>&B_batched,GlobalTensor<float16>&C_batched,GlobalTensor<uint8>&batch_mask,// [batch] 1=有效, 0=跳过intbatch,intM,intN,intK){for(intb=0;b<batch;b++){if(!batch_mask[b]){continue;// 跳过这个 batch — 节省 Cube 和时间}// ... 正常计算}}Mask 由上层（ATB）传入——对于 padding 头，batch_mask = 0。
踩坑三：Batched GEMM 和单次大 GEMM 的取舍
Merge QKV projection：把 3H 个小 GEMM 合并成 1 次大 GEMM——x @ [W_q, W_k, W_v]。形状是[batch*seq, d_model] @ [d_model, 3*head*d_head]——一次 GEMM 代替 3H 次小 GEMM。
选择逻辑：
# ops-blas 自动判断ifM>4096orK>4096:# 大矩阵 → Merge 成一次大 GEMM# 好处：Cube 利用率高（tile 填满）returnmerged_GEMM(x,W_merged)elifbatch_count>32:# 很多小 GEMM → Batched GEMM# 好处：一次 launch，减少开销returnbatched_GEMM(x,W_batched)else:# 中等规模 → 混合策略returnhybrid_GEMM(x,W_batched)经验规则：
- MHA 推理（batch=1, seq=128, head=32）→ Batched GEMM（32 个小矩阵）
- MHA 训练（batch=8, seq=2048, head=32）→ Merged GEMM（1 次大矩阵大 GEMM）
- 形状阈值：M×K > 4096×4096 → Merge；否则 → Batched
Batched GEMM 解决的不只是计算效率——而是 launch 开销和流水线中断。32 次 HEAD MM 各 launch 一次（32×50μs=1.6ms 开销）vs 一次 Batched GEMM launch（50μs）。在推理管线的 2ms 总时间中，launch 开销占比从 80% 降到 2.5%。ops-blas 的 Batched GEMM 自动选择策略（Interleaved/Parallel/Hybrid）、支持 stride 和 mask——让 MHA 的 H 个小矩阵乘变成一次 kernel 调用。