昇腾CANN ops-blas Batched GEMM：多头注意力的小矩阵乘批处理实战

雨季666

312人浏览 · 2026-05-23 22:55:00

雨季666 · 2026-05-23 22:55:00 发布

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__ void BatchedGEMMInterleaved(
    GlobalTensor<float16>& A_batched,  // [batch, M, K]
    GlobalTensor<float16>& B_batched,  // [batch, K, N]
    GlobalTensor<float16>& C_batched,  // [batch, M, N]
    int batch, int M, int N, int K
) {
    // 每个 block 处理一个 batch 的 GEMM
    for (int b = 0; b < batch; b++) {
        int block_id = b % gridDim.x;  // 轮询分配 block

        // 在 L1 中交错存储 32 个 batch 的 tile
        // 单个 tile 大小 = tile_M × tile_K = 16 × 16 = 256 elements
        LocalTensor<float16> A_tile(tile_M * tile_K);
        LocalTensor<float16> B_tile(tile_K * tile_N);
        LocalTensor<float16> C_tile(tile_M * tile_N);

        int A_offset = b * M * K;
        int B_offset = b * K * N;
        int C_offset = b * M * N;

        // 分块矩阵乘
        for (int m = 0; m < M; m += tile_M) {
            for (int n = 0; n < N; n += tile_N) {
                // 初始化累加器
                C_tile = 0.0f;

                for (int k = 0; k < K; k += tile_K) {
                    // 加载 A 和 B 的 tile 到 L1
                    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);

                    // 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__ void BatchedGEMMParallel(
    GlobalTensor<float16>& A_batched,  // [batch, M, K]
    GlobalTensor<float16>& B_batched,  // [batch, K, N]
    GlobalTensor<float16>& C_batched,  // [batch, M, N]
    int batch, int M, int N, int K
) {
    // 每个 block 处理一个独立的 batch（不是所有 block 处理同一 batch）
    // block 分配：block_id = b % num_batch_blocks
    // num_batch_blocks = gridDim.x / batch

    int num_batch_blocks = gridDim.x / batch;
    if (num_batch_blocks < 1) num_batch_blocks = 1;

    // 每个 batch 有 num_batch_blocks 个 block 在并行处理
    int batch_id = blockIdx.x / num_batch_blocks;
    int batch_block = blockIdx.x % num_batch_blocks;

    int A_offset = batch_id * M * K;
    int B_offset = batch_id * K * N;
    int C_offset = batch_id * M * N;

    // batch_block 决定此 block 处理矩阵的哪一部分
    // 把 M 维度均分给 num_batch_blocks 个 block
    int m_start = batch_block * (M / num_batch_blocks);
    int m_end = (batch_block + 1) * (M / num_batch_blocks);

    for (int m = m_start; m < m_end; m += tile_M) {
        for (int n = 0; n < N; n += tile_N) {
            LocalTensor<float16> C_tile(tile_M * tile_N);
            C_tile = 0.0f;

            for (int k = 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__ void BatchedGEMMHybrid(
    GlobalTensor<float16>& A_batched,  // [batch, M, K]
    GlobalTensor<float16>& B_batched,  // [batch, K, N]
    GlobalTensor<float16>& C_batched,  // [batch, M, N]
    int batch, int M, int N, int K
) {
    // 分组：每 group_size 个 batch 为一组
    // 组内用 Interleaved（充分利用 L1），组间用 Parallel
    int group_size = 4;  // 每组 4 个 batch

    int num_groups = (batch + group_size - 1) / group_size;
    int group_id = blockIdx.x % num_groups;  // 每个 block 处理一个 group

    // 组间并行处理
    int batch_start = group_id * group_size;
    int batch_end = min(batch_start + group_size, batch);

    // 组内 Interleaved
    for (int b = batch_start; b < batch_end; b++) {
        int A_offset = b * M * K;
        int B_offset = b * K * N;
        int C_offset = b * M * N;

        // 分块矩阵乘（同 Interleaved 策略）
        for (int m = 0; m < M; m += tile_M) {
            for (int n = 0; n < N; n += tile_N) {
                LocalTensor<float16> C_tile(tile_M * tile_N);
                C_tile = 0.0f;

                for (int k = 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 GEMM
import torch
import torch_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__ void BatchedGEMMStrided(
    GlobalTensor<float16>& A_batched,
    GlobalTensor<float16>& B_batched,
    GlobalTensor<float16>& C_batched,
    int batch, int M, int N, int K,
    int stride_A,   // A 的 batch stride（不连续时 > M*K）
    int stride_B,   // B 的 batch stride
    int stride_C    // C 的 batch stride
) {
    for (int b = 0; b < batch; b++) {
        // 使用 stride 替代 M*K
        int A_offset = b * stride_A;  // 不是 b * M * K
        int B_offset = b * stride_B;
        int C_offset = b * stride_C;
        // ... 其余同 Interleaved
    }
}

PyTorch 侧：

# 非连续 batch → 指定 stride
output = 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__ void BatchedGEMMMasked(
    GlobalTensor<float16>& A_batched,
    GlobalTensor<float16>& B_batched,
    GlobalTensor<float16>& C_batched,
    GlobalTensor<uint8>& batch_mask,  // [batch] 1=有效, 0=跳过
    int batch, int M, int N, int K
) {
    for (int b = 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 自动判断
if M > 4096 or K > 4096:
    # 大矩阵 → Merge 成一次大 GEMM
    # 好处：Cube 利用率高（tile 填满）
    return merged_GEMM(x, W_merged)
elif batch_count > 32:
    # 很多小 GEMM → Batched GEMM
    # 好处：一次 launch，减少开销
    return batched_GEMM(x, W_batched)
else:
    # 中等规模 → 混合策略
    return hybrid_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 调用。