triton-ascend-matmul

MatMul 算子优化

适用于矩阵乘法及相关运算

核心数选择（硬约束）

• 涉及 tl.dot / 矩阵乘法运算 → 必须使用 CUBE_CORE_NUM
• 混合运算（先 matmul 再 elementwise 后处理）→ CUBE_CORE_NUM
• 纯 elementwise / 标量运算 → VEC_CORE_NUM

使用 VEC_CORE_NUM 启动 matmul kernel 会导致数值结果错误。

Tile Size 限制（硬件约束）

MatMul 数据走 L0A/L0B/L0C，tile 大小受硬件存储容量限制，超出会导致 ub overflow / cbuf overflow 编译错误。

约束公式（具体容量参考目标硬件信息文档）：

• L0A：BLOCK_M × BLOCK_K × sizeof(dtype) ≤ L0A容量
• L0B：BLOCK_K × BLOCK_N × sizeof(dtype) ≤ L0B容量
• L0C：BLOCK_M × BLOCK_N × sizeof(acc_dtype) ≤ L0C容量

遇 ub overflow / cbuf overflow → 缩小 BLOCK_M, BLOCK_N 或 BLOCK_K

Ascend 后端切分优化

关键原则: 充分发挥带宽，算子行宽为 512B 的整数倍。

以 fp16/bf16 为例（每个元素 2 字节）：

切分配置（根据转置情况）

1. A、B 都不转置

   • 分块行宽分别为 K0 和 N0
   • 推荐: M0=128, K0=256, N0=256

2. A 不转置，B 转置

   • 分块行宽都是 K0
   • 推荐: K0=256, M0 和 N0 影响较小

3. A、B 都转置

   • 分块行宽分别为 M0 和 K0
   • 推荐: M0=256, K0=256, N0=128

4. A 转置，B 不转置

   • 分块行宽分别为 M0 和 N0
   • 注意: 左右矩阵均无法同时满足 512B 的整数倍，需根据实际情况调整

为什么是 512B？

• 512B = 256 个 fp16/bf16 元素（256 × 2 字节）
• NPU 的最佳带宽对齐
• 确保每次内存访问充分利用带宽

固定核心数启动

MatMul 算子使用 CUBE核心数（矩阵计算核心）。

关键: 使用 grid=(num_cores,) 而非 (NUM_BLOCKS,)

@triton.jit
def matmul_kernel(
    a_ptr, b_ptr, c_ptr,
    M, N, K,
    num_cores: tl.constexpr,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
):
    # 关键:使用固定核心数启动,每个核心处理多个块
    pid = tl.program_id(0)  # 核心ID: 0~num_cores-1
    NUM_BLOCKS_M = triton.cdiv(M, BLOCK_M)
    NUM_BLOCKS_N = triton.cdiv(N, BLOCK_N)
    NUM_BLOCKS = NUM_BLOCKS_M * NUM_BLOCKS_N

    # 每个核心循环处理多个块
    for block_idx in range(pid, NUM_BLOCKS, num_cores):
        # 计算当前块的2D索引
        block_m = block_idx // NUM_BLOCKS_N
        block_n = block_idx % NUM_BLOCKS_N

        # 初始化累加器
        accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

        # K维度循环
        for k in range(0, K, BLOCK_K):
            # 加载A块
            a_offset = (block_m * BLOCK_M + tl.arange(0, BLOCK_M))[:, None] * K + \
                       (k + tl.arange(0, BLOCK_K))[None, :]
            a_mask = (block_m * BLOCK_M + tl.arange(0, BLOCK_M))[:, None] < M
            a = tl.load(a_ptr + a_offset, mask=a_mask, other=0.0)

            # 加载B块
            b_offset = (k + tl.arange(0, BLOCK_K))[:, None] * N + \
                       (block_n * BLOCK_N + tl.arange(0, BLOCK_N))[None, :]
            b_mask = (block_n * BLOCK_N + tl.arange(0, BLOCK_N))[None, :] < N
            b = tl.load(b_ptr + b_offset, mask=b_mask, other=0.0)

            # 矩阵乘累加
            accumulator += tl.dot(a, b)

        # 存储结果
        c_offset = (block_m * BLOCK_M + tl.arange(0, BLOCK_M))[:, None] * N + \
                   (block_n * BLOCK_N + tl.arange(0, BLOCK_N))[None, :]
        c_mask = ((block_m * BLOCK_M + tl.arange(0, BLOCK_M))[:, None] < M) & \
                 ((block_n * BLOCK_N + tl.arange(0, BLOCK_N))[None, :] < N)
        tl.store(c_ptr + c_offset, accumulator, mask=c_mask)

class ModelNew(torch.nn.Module):
    def __init__(self):
        super().__init__()
        try:
            self.CUBE_CORE_NUM = torch_npu.npu.npu_config.get_device_limit(0).get("cube_core_num", 20)
        except:
            self.CUBE_CORE_NUM = 20

    def forward(self, a, b):
        M, K = a.shape
        K2, N = b.shape
        assert K == K2
        c = torch.empty((M, N), device=a.device, dtype=a.dtype)

        num_cores = self.CUBE_CORE_NUM
        BLOCK_M, BLOCK_N, BLOCK_K = 128, 256, 256

        matmul_kernel[(num_cores,)](
            a, b, c, M, N, K, num_cores,
            BLOCK_M, BLOCK_N, BLOCK_K
        )
        return c