第 9 课：Blocked Matmul 与 FP32 累加

本课目标与完成标准

学完后你应能：实现带 K 尾块的 blocked matmul，解释 accumulator、mask 和 autotune 维度。

通过必须同时满足：

独立补齐本课唯一的代码填空题，并通过给定检查；
三个问答题均说明因果链，而不是只报术语；
能指出至少一个正确性边界和一个性能取舍；
总分不低于 8/10，且没有一票否决级概念错误。

前置关系

课程前置：Python、PyTorch 张量、CUDA 基本线程/内存概念
本课在路线中的作用：Triton GEMM 用二维 program grid 负责 C tile，在 K 维循环加载 A/B tiles 并用 tl.dot 累加。

核心心智模型

1. 它是什么，解决什么问题

通用矩阵乘法（GEMM: $C = A \times B$ ， $A \in \mathbb{R}^{M \times K}, B \in \mathbb{R}^{K \times N}$ ）是现代人工智能模型（Transformer、CNN、MoE）中计算量的主要来源。在 GPU 体系结构中，朴素的三重循环算法算术强度极低，会导致致命的显存带宽瓶颈。因此，高性能 GEMM 必须采用二维输出分块（Blocked Tiling）：网格划分为 $M \times N$ 的宏观瓦片，每个 Program 块负责计算并写回输出矩阵的一小块 $C[\text{BM}, \text{BN}]$ ；在内层沿 $K$ 维以跨步 $BK$ 循环切片，并在片上高速寄存器中不断执行外积累加。本课解决的核心工程问题是：在 Triton 中如何使用 tl.dot 驱动硬件底层 Tensor Core 核心，如何通过 FP32 累加器维持混合精度数值稳定性，以及如何为 $M, N, K$ 三维独立施加谓词掩码，优雅处理任意非对齐的尾块尺寸。

2. 它如何工作

分块矩阵乘瓦片流水线如图 9-1 至图 9-3 所示：

Triton Matmul：块指针与分组调度

2D 块指针

offs_m = pid_m*BM + tl.arange(0, BM)，offs_k = tl.arange(0, BK)

a_ptrs = a + offs_m[:, None]*stride_am + offs_k[None, :]*stride_ak

每轮 K 循环：a_ptrs += BK*stride_ak，b_ptrs += BK*stride_bk

Grouped PID（GROUP_M = G）

group_id = pid // (G*num_pid_n)，first_pid_m = group_id*G

gsize = min(num_pid_m - first_pid_m, G)

pid_m = first_pid_m + pid % gsize，pid_n = (pid % (G*num_pid_n)) // gsize

组内相邻 program 共享 B 的同一列块和 A 的同一组行块，提高 L2 命中率；收益取决于形状，Triton 官方教程在 A100 上的例子约提升 10%

图 9-1：块指针步进与 Program ID 分组

分块矩阵乘 (BLOCK_M × BLOCK_N)

A 瓦片 (BLOCK_M × BLOCK_K)

B 瓦片 (BLOCK_K × BLOCK_N)

累加器 accumulator (FP32 精度)

tl.dot(a, b, acc)tl.dot(a, b, acc)

图 9-2：Blocked Matmul 瓦片乘法与 FP32 累加器保持精度架构

初始化 acc = zeros((BLOCK_M, BLOCK_N), dtype=fp32)

K 维步进: 加载 A_tile 与 B_tile

acc = tl.dot(A_tile, B_tile, acc)

K 轴遍历完毕?

转回 FP16/BF16 并写回 C[i, j]

否是

图 9-3：Blocked Matmul K 维分块遍历累加流水线

累加器初始化：每个 Program 处理输出的一个 $(BM, BN)$ 子矩阵。在寄存器中初始化 FP32 全零矩阵： acc = tl.zeros((BM, BN), tl.float32)
K 维分块迭代与独立掩码保护：内层循环以 $BK$ 为步长推进：for k0 in range(0, K, BK):
- 加载 $A$ 的当前瓦片 $A_k \in \mathbb{R}^{BM \times BK}$ ，地址为 rm[:, None] * sam + rk[None, :] * sak，掩码必须同时约束 $M$ 轴与 $K$ 轴： mask = (rm[:, None] < M) & (rk[None, :] < K)，越界处填充加法单位元 other=0.0；
- 加载 $B$ 的当前瓦片 $B_k \in \mathbb{R}^{BK \times BN}$ ，掩码同时约束 $K$ 轴与 $N$ 轴： mask = (rk[:, None] < K) & (rn[None, :] < N)，越界处同样填充 other=0.0。
Tensor Core 驱动的片上累加（tl.dot）：执行 acc += tl.dot(av, bv)。在底层，Triton 编译器将 tl.dot 编译为 NVIDIA 原生 Tensor Core 硬件指令（如 Hopper 上的 wgmma 或 Ampere 上的 mma.sync.m16n8k16）。输入矩阵以 FP16/BF16 载入并相乘，但累加在 FP32 寄存器中进行，从硬件层面保证了深层网络训练中成千上万次乘加的累积精度。
截断与写回（Epilogue）： $K$ 维循环结束后，累加器 acc 携带完整的瓦片点积和。通过输出掩码 (rm[:, None] < M) & (rn[None, :] < N) 保护，将结果安全转换为目标数据类型（如 FP16）写回全局显存 $C$ 。

3. 正确性条件与常见误区

K 维尾块遗漏掩码的累积污染：若开发者仅在写回 $C$ 时施加了 $M$ 和 $N$ 的掩码，却在内层循环加载 $A$ 和 $B$ 时遗漏了 rk < K 的检查：对于尾块（ $K$ 不是 $BK$ 整数倍），tl.load 会将 $K$ 维度之后的非法显存脏数据读入寄存器。经过 Tensor Core 乘积累加后，这些脏数据产生的外积结果会被永久熔铸进累加器 acc 中。即使在写回时做了输出掩码，算出的数值也已经是错误的。
混合精度累加溢出陷阱：若累加器初始化为 FP16（tl.zeros((BM, BN), tl.float16)），在 $K$ 较大（如 $K=4096$ ）时，连续累加会导致 FP16 范围溢出（最大值仅 65504）或极严重的下溢与截断误差。必须严格遵循“低精度输入计算、高精度 FP32 累加”黄金法则。
tl.dot 的形状约束：tl.dot 的输入块各维必须 $\ge 16$ ，且块形状为 2 的幂；因此 BM、BN、BK 最小取 16。 $M$ 很小时（如 $M=1$ ）也只能用 BM=16 加 mask，或改为不用 tl.dot 的逐元素乘加 + tl.sum 写法。
浮点结合律与容差预期：由于 GPU Tensor Core 的并行点积顺序与 CPU 上的串行双精度循环截然不同，浮点加法不满足严格结合律，因此断言比较必须留有适当的相对容差（如 atol=2e-2, rtol=2e-2），不能按位要求 bitwise-exact。

4. 性能与工程取舍

Tile 形状与流水线深度（BM, BN, BK, num_warps, num_stages）： Triton GEMM 的性能极度依赖超参数选择：
- (BM, BN, BK) 决定了每个 Program 的计算强度与片上寄存器消耗；
- num_stages 决定了使用共享内存软件流水线（Async Pipeline，如 cp.async）预取的循环级数，通常 3~4 级最佳；
- num_warps（常用 4 或 8）是每个 Program（即一个 CTA）使用的 warp 数，线程数为 $32 \times$ num_warps；它影响每个线程分到的元素数、寄存器压力和 SM 上能同时驻留的 CTA 数。
训练 GEMM（大 M）与推理 Decode GEMV（M=1）的极化差异：在大模型预训练或 Prefill 阶段， $M$ 较大（如 $M = 4096$ ），属于典型的计算密集型（Compute-bound），需要选择大 Tile（如 $128 \times 128 \times 32$ ）打满 Tensor Core；但在自回归推理解码（Decode）阶段，每次仅输入 1 个 Token（ $M=1$ ），算子退化为矩阵-向量乘法（GEMV），极度受限于显存带宽（Memory-bound）。此时若强行使用大 Tile 会导致 99% 的 Program 空转，必须采用专门针对 $M=1$ 优化的 Split-K GEMV 内核。

具体演示

以 $M = 17, K = 33, N = 19$ 、分块大小 $BM = 32, BN = 32, BK = 32$ 为例：

网格发射 $\lceil 17/32 \rceil \times \lceil 19/32 \rceil = 1 \times 1 = 1$ 个 Program。
K 维循环迭代：⌈33/32⌉=2\lceil 33/32 \rceil = 2 轮。
- 第 1 轮（ $k0 = 0$ ）： $rk \in [0..31]$ 。 $A$ 掩码有效范围为 $[0..16, 0..31]$ ， $B$ 掩码有效范围为 $[0..31, 0..18]$ 。tl.dot 累加前 32 个通道的有效点积。
- 第 2 轮（ $k0 = 32$ ）： $rk \in [32..63]$ 。 $rk < 33$ 仅有第 32 这一列有效，后 31 列越界。由于指定了 other=0.0， $A$ 和 $B$ 的后 31 个切片全被填充为 0.0。tl.dot 仅计算第 32 通道的有效乘加，多余维度的 $0.0 \times 0.0 = 0.0$ 累加不改变精度。
最终写回仅输出前 $[17, 19]$ 区域；数值与 a @ b 在容差内一致（累加顺序不同，不要求逐位相同）。

请在阅读后先合上这一节，用自己的语言复述“输入状态 → 中间状态 → 输出状态”，再做练习。

实践任务：唯一代码填空题

补齐 K-tile 累加。

规则：只能修改 TODO/______ 所在位置；不要删除断言或放宽误差。代码注释说明了每个边界条件。

Code

import torch
import triton
import triton.language as tl

@triton.jit
def matmul_kernel(a, b, c, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr,
                  sam: tl.constexpr, sak: tl.constexpr,
                  sbk: tl.constexpr, sbn: tl.constexpr,
                  scm: tl.constexpr, scn: tl.constexpr,
                  BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr):
    rm = tl.program_id(0) * BM + tl.arange(0, BM)
    rn = tl.program_id(1) * BN + tl.arange(0, BN)
    acc = tl.zeros((BM, BN), tl.float32)
    for k0 in range(0, K, BK):
        rk = k0 + tl.arange(0, BK)
        av = tl.load(a + rm[:, None] * sam + rk[None, :] * sak,
                     mask=(rm[:, None] < M) & (rk[None, :] < K), other=0.0)
        bv = tl.load(b + rk[:, None] * sbk + rn[None, :] * sbn,
                     mask=(rk[:, None] < K) & (rn[None, :] < N), other=0.0)
        acc += ______  # TODO: tile 矩阵乘
    tl.store(c + rm[:, None] * scm + rn[None, :] * scn, acc,
             mask=(rm[:, None] < M) & (rn[None, :] < N))

def matmul(a, b):
    assert a.ndim == b.ndim == 2 and a.shape[1] == b.shape[0]
    M, K = a.shape; N = b.shape[1]
    c = torch.empty((M, N), device=a.device, dtype=a.dtype)
    grid = (triton.cdiv(M, 32), triton.cdiv(N, 32))
    matmul_kernel[grid](a, b, c, M, N, K, *a.stride(), *b.stride(), *c.stride(),
                        BM=32, BN=32, BK=32)
    return c

for shape in ((17, 33, 19), (64, 64, 64)):
    M, K, N = shape
    a = torch.randn((M,K), device="cuda", dtype=torch.float16)
    b = torch.randn((K,N), device="cuda", dtype=torch.float16)
    torch.testing.assert_close(matmul(a,b), a@b, atol=2e-2, rtol=2e-2)

检查方法

在 CUDA/Triton 环境运行本单元格；断言覆盖规则尺寸和非规则尾块。首次 JIT 不计入性能。

提交时请给出：补齐后的代码、实际运行输出（环境不可用时注明“仅静态审查”）以及对失败用例的解释。

Q1

不要背定义：请从输入、状态变化和输出三个阶段解释“Blocked Matmul 与 FP32 累加”的工作机制。

你的答案：

Q2

只 mask 输出而不 mask K 尾块会怎样？

你的答案：

Q3

为什么训练 GEMM 与 decode 的 M=1 GEMM 不应共用同一配置？

你的答案：

评分与通过规则

代码 4 分：正常输入 2 分，边界输入 1 分，解释实现 1 分；
Q1～Q3 各 2 分；
一票否决：结果碰巧正确但核心因果链错误、删除边界检查、把未运行结果说成实测。

需要提示时按四级机制请求：概念区域 → 具体方向 → 关键局部 → 完整答案。

💡 点击展开参考答案与详细代码实现

参考答案（仅 answer 分支）

先完成题目再核对。即使代码一致，也要能解释关键步骤，并尝试更换一个输入规模。

Code

import torch
import triton
import triton.language as tl

@triton.jit
def matmul_kernel(a, b, c, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr,
                  sam: tl.constexpr, sak: tl.constexpr,
                  sbk: tl.constexpr, sbn: tl.constexpr,
                  scm: tl.constexpr, scn: tl.constexpr,
                  BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr):
    rm = tl.program_id(0) * BM + tl.arange(0, BM)
    rn = tl.program_id(1) * BN + tl.arange(0, BN)
    acc = tl.zeros((BM, BN), tl.float32)
    for k0 in range(0, K, BK):
        rk = k0 + tl.arange(0, BK)
        av = tl.load(a + rm[:, None] * sam + rk[None, :] * sak,
                     mask=(rm[:, None] < M) & (rk[None, :] < K), other=0.0)
        bv = tl.load(b + rk[:, None] * sbk + rn[None, :] * sbn,
                     mask=(rk[:, None] < K) & (rn[None, :] < N), other=0.0)
        acc += tl.dot(av, bv)
    tl.store(c + rm[:, None] * scm + rn[None, :] * scn, acc,
             mask=(rm[:, None] < M) & (rn[None, :] < N))

def matmul(a, b):
    assert a.ndim == b.ndim == 2 and a.shape[1] == b.shape[0]
    M, K = a.shape; N = b.shape[1]
    c = torch.empty((M, N), device=a.device, dtype=a.dtype)
    grid = (triton.cdiv(M, 32), triton.cdiv(N, 32))
    matmul_kernel[grid](a, b, c, M, N, K, *a.stride(), *b.stride(), *c.stride(),
                        BM=32, BN=32, BK=32)
    return c

for shape in ((17, 33, 19), (64, 64, 64)):
    M, K, N = shape
    a = torch.randn((M,K), device="cuda", dtype=torch.float16)
    b = torch.randn((K,N), device="cuda", dtype=torch.float16)
    torch.testing.assert_close(matmul(a,b), a@b, atol=2e-2, rtol=2e-2)

Q1 参考答案

从三个阶段拆解其运行机制：

输入阶段（2D 宏观网格与 K 维切片加载）：以二维网格映射输出矩阵瓦片，Program 块初始化 FP32 寄存器累加器 acc。外层循环沿 $K$ 维以 $BK$ 为步长推进。每一轮迭代中，Program 分别构造 $A$ 与 $B$ 瓦片的二维指针，并在双向掩码（ $A$ 受控于 $(M, K)$ ， $B$ 受控于 $(K, N)$ ）保护下，结合 other=0.0 将当前切片合并载入片上寄存器。
状态变化阶段（Tensor Core 驱动的片上点积累加）：调用 tl.dot(av, bv) 将当前 $A$ 瓦片与 $B$ 瓦片送入 GPU 硬件底层的 Tensor Core 矩阵乘加流水线。计算采用 FP16 乘法并在 FP32 累加器中就地累加，经过多轮 $K$ 维切片迭代后，累加器逐步收敛出该瓦片的最终真实点积和，整个循环计算完全在片上寄存器展开。
输出阶段（输出掩码截断与类型转换写回）：将高精度累加器 acc 转换为目标张量数据类型（FP16/BF16），结合输出边界掩码 (rm < M) & (rn < N)，通过 tl.store 写回全局目标矩阵 $C$ 。越界位置完全被硬件屏蔽，保证零越界写入。

Q2 参考答案

若仅在输出写回 $C$ 时施加掩码，而在内层循环加载 $A$ 和 $B$ 时遗漏了针对 $K$ 维度的掩码保护（即没有 rk < K 约束）：

尾块越界脏数据强行读入：当 $K$ 不是分块尺寸 $BK$ 的整数倍时（例如 $K=33, BK=32$ ，在第二轮迭代中有效 $K$ 仅有 1 项，后 31 项为越界内存），内核会绕过边界检查，从显存越界地址中强行读取邻近内存的未初始化脏数据。
脏数据永久熔铸进累加器：在执行 tl.dot(av, bv) 时，这 31 个非法的脏数据向量会被送进 Tensor Core 进行点乘求和，产生不可预测的伪数值并直接累加进片上累加器 acc 中。
输出毁坏：即使在最后写回 $C$ 时施加了 $M$ 和 $N$ 维度的输出掩码，累加器内存储的数值本身已经被脏数据污染，最终写回目标张量的所有有效数值都会变成不可控的随机荒谬值。

Q3 参考答案

大模型训练 GEMM 与自回归推理解码（Decode）阶段的计算特征处于完全不同的物理极端，绝对不能共用同一套分块配置：

训练 GEMM（ $M$ 极大，如 $M=4096, N=4096, K=4096$ ）：
- 体系特征：属于典型的计算密集型（Compute-bound）算子，算术强度极高。
- 最优配置：需要采用大分块尺寸（如 $BM=128, BN=128, BK=32$ 或 64），以最大化寄存器与片上数据复用率，释放 Tensor Core 的峰值 TFLOPs 吞吐，并使用 4~8 个 Warp 和深度流水线。
推理 Decode GEMM（ $M=1$ ）：
- 体系特征：退化为矩阵-向量乘法（GEMV），算术强度极低（每个权重只被使用 1 次），属于绝对的显存带宽密集型（Memory-bound）算子。
- 配置问题：若在 $M=1$ 时沿用训练的大分块（如 $BM=128$ ），tile 中 128 行有 127 行被 mask 掉，Tensor Core 算力基本空转；同时 grid 只有 $\lceil N/BN \rceil$ 个 Program（如 $N=4096, BN=128$ 时仅 32 个），远少于 H100 的 132 个 SM。 $M=1$ 时应使用小 $BM$ （用 tl.dot 时最小为 16，或改为不用 tl.dot 的乘加 + tl.sum GEMV 写法）、较小的 $BN$ 以增加 Program 数，并结合 Split-K（沿 $K$ 维切分到多个 Program，最后用原子加或第二个 kernel 归约）。

参考资料

资料用于建立事实基线；面试回答仍需用自己的语言组织。

← 上一课
第 8 课：二维 Tile 与转置 下一课 →
第 10 课：无状态随机数与 Dropout

--- title: "第 9 课：Blocked Matmul 与 FP32 累加" jupyter: python3 --- ## 本课目标与完成标准学完后你应能：实现带 K 尾块的 blocked matmul，解释 accumulator、mask 和 autotune 维度。通过必须同时满足： - 独立补齐本课唯一的代码填空题，并通过给定检查； - 三个问答题均说明因果链，而不是只报术语； - 能指出至少一个正确性边界和一个性能取舍； - 总分不低于 8/10，且没有一票否决级概念错误。 ## 前置关系 - 课程前置：Python、PyTorch 张量、CUDA 基本线程/内存概念 - 本课在路线中的作用：Triton GEMM 用二维 program grid 负责 C tile，在 K 维循环加载 A/B tiles 并用 `tl.dot` 累加。 ## 核心心智模型 ### 1. 它是什么，解决什么问题通用矩阵乘法（GEMM: $C = A \times B$，$A \in \mathbb{R}^{M \times K}, B \in \mathbb{R}^{K \times N}$）是现代人工智能模型（Transformer、CNN、MoE）中计算量的主要来源。在 GPU 体系结构中，朴素的三重循环算法算术强度极低，会导致致命的显存带宽瓶颈。因此，高性能 GEMM 必须采用**二维输出分块（Blocked Tiling）**：网格划分为 $M \times N$ 的宏观瓦片，每个 Program 块负责计算并写回输出矩阵的一小块 $C[\text{BM}, \text{BN}]$；在内层沿 $K$ 维以跨步 $BK$ 循环切片，并在片上高速寄存器中不断执行外积累加。本课解决的核心工程问题是：**在 Triton 中如何使用 `tl.dot` 驱动硬件底层 Tensor Core 核心，如何通过 FP32 累加器维持混合精度数值稳定性，以及如何为 $M, N, K$ 三维独立施加谓词掩码，优雅处理任意非对齐的尾块尺寸**。 ### 2. 它如何工作分块矩阵乘瓦片流水线如图 9-1 至图 9-3 所示： :::: {.vfig} ::: {.vfig-title} Triton Matmul：块指针与分组调度 ::: ::::: {.vfig-grid style="--cols:2"} ::: {.vfig-card .c-blue} **2D 块指针** `offs_m = pid_m*BM + tl.arange(0, BM)`，`offs_k = tl.arange(0, BK)` `a_ptrs = a + offs_m[:, None]*stride_am + offs_k[None, :]*stride_ak` 每轮 K 循环：`a_ptrs += BK*stride_ak`，`b_ptrs += BK*stride_bk` ::: ::: {.vfig-card .c-green} **Grouped PID（GROUP_M = G）** `group_id = pid // (G*num_pid_n)`，`first_pid_m = group_id*G` `gsize = min(num_pid_m - first_pid_m, G)` `pid_m = first_pid_m + pid % gsize`，`pid_n = (pid % (G*num_pid_n)) // gsize` ::: ::::: ::: {.vfig-key} 组内相邻 program 共享 B 的同一列块和 A 的同一组行块，提高 L2 命中率；收益取决于形状，Triton 官方教程在 A100 上的例子约提升 10% ::: :::: <p class="caption">图 9-1：块指针步进与 Program ID 分组</p> ```{=html} <div class="fd fd-flow" role="img"><div class="fd-stack"><div class="fd-comp fd-td"><div class="fd-sg" data-nid="BlockGEMM" data-depth="0"><div class="fd-sg-title">分块矩阵乘 (BLOCK_M × BLOCK_N)</div><div class="fd-stack"><div class="fd-comp fd-td"><div class="fd-rank fd-few"><div class="fd-node fd-box" data-nid="A_k">A 瓦片 (BLOCK_M × BLOCK_K)</div><div class="fd-node fd-box" data-nid="B_k">B 瓦片 (BLOCK_K × BLOCK_N)</div></div><div class="fd-gap" style="--gl:157px;--gt:53px"></div><div class="fd-rank"><div class="fd-node fd-box" data-nid="Acc">累加器 accumulator (FP32 精度)</div></div></div></div></div></div></div><i class="fd-e" data-s="A_k" data-t="Acc" data-k="solid" data-h="1" data-b="0">tl.dot(a, b, acc)</i><i class="fd-e" data-s="B_k" data-t="Acc" data-k="solid" data-h="1" data-b="0">tl.dot(a, b, acc)</i></div> ``` <p class="caption" align="center"><em>图 9-2：Blocked Matmul 瓦片乘法与 FP32 累加器保持精度架构</em></p> ```{=html} <div class="fd fd-flow" role="img"><div class="fd-stack"><div class="fd-comp fd-lr" data-minw="1115"><div class="fd-rank"><div class="fd-node fd-box" data-nid="Init">初始化 acc = zeros((BLOCK_M, BLOCK_N), dtype=fp32)</div></div><div class="fd-gap" style="--gl:44px;--gt:38px"></div><div class="fd-rank"><div class="fd-node fd-box" data-nid="LoopK">K 维步进: 加载 A_tile 与 B_tile</div></div><div class="fd-gap" style="--gl:58px;--gt:53px"></div><div class="fd-rank"><div class="fd-node fd-box" data-nid="Mma">acc = tl.dot(A_tile, B_tile, acc)</div></div><div class="fd-gap" style="--gl:44px;--gt:38px"></div><div class="fd-rank"><div class="fd-node fd-dec" data-nid="Check">K 轴遍历完毕?</div></div><div class="fd-gap" style="--gl:58px;--gt:53px"></div><div class="fd-rank"><div class="fd-node fd-box" data-nid="CastStore">转回 FP16/BF16 并写回 C[i, j]</div></div></div></div><i class="fd-e" data-s="Init" data-t="LoopK" data-k="solid" data-h="1" data-b="0"></i><i class="fd-e" data-s="LoopK" data-t="Mma" data-k="solid" data-h="1" data-b="0"></i><i class="fd-e" data-s="Mma" data-t="Check" data-k="solid" data-h="1" data-b="0"></i><i class="fd-e" data-s="Check" data-t="LoopK" data-k="solid" data-h="1" data-b="0">否</i><i class="fd-e" data-s="Check" data-t="CastStore" data-k="solid" data-h="1" data-b="0">是</i></div> ``` <p class="caption" align="center"><em>图 9-3：Blocked Matmul K 维分块遍历累加流水线</em></p> 1. **累加器初始化**：每个 Program 处理输出的一个 $(BM, BN)$ 子矩阵。在寄存器中初始化 FP32 全零矩阵： `acc = tl.zeros((BM, BN), tl.float32)` 2. **K 维分块迭代与独立掩码保护**：内层循环以 $BK$ 为步长推进：`for k0 in range(0, K, BK):` - 加载 $A$ 的当前瓦片 $A_k \in \mathbb{R}^{BM \times BK}$，地址为 `rm[:, None] * sam + rk[None, :] * sak`，掩码必须同时约束 $M$ 轴与 $K$ 轴： `mask = (rm[:, None] < M) & (rk[None, :] < K)`，越界处填充加法单位元 `other=0.0`； - 加载 $B$ 的当前瓦片 $B_k \in \mathbb{R}^{BK \times BN}$，掩码同时约束 $K$ 轴与 $N$ 轴： `mask = (rk[:, None] < K) & (rn[None, :] < N)`，越界处同样填充 `other=0.0`。 3. **Tensor Core 驱动的片上累加（`tl.dot`）**：执行 `acc += tl.dot(av, bv)`。在底层，Triton 编译器将 `tl.dot` 编译为 NVIDIA 原生 Tensor Core 硬件指令（如 Hopper 上的 `wgmma` 或 Ampere 上的 `mma.sync.m16n8k16`）。输入矩阵以 FP16/BF16 载入并相乘，但累加在 FP32 寄存器中进行，从硬件层面保证了深层网络训练中成千上万次乘加的累积精度。 4. **截断与写回（Epilogue）**： $K$ 维循环结束后，累加器 `acc` 携带完整的瓦片点积和。通过输出掩码 `(rm[:, None] < M) & (rn[None, :] < N)` 保护，将结果安全转换为目标数据类型（如 FP16）写回全局显存 $C$。 ### 3. 正确性条件与常见误区 - **K 维尾块遗漏掩码的累积污染**：若开发者仅在写回 $C$ 时施加了 $M$ 和 $N$ 的掩码，却在内层循环加载 $A$ 和 $B$ 时遗漏了 `rk < K` 的检查：对于尾块（$K$ 不是 $BK$ 整数倍），`tl.load` 会将 $K$ 维度之后的非法显存脏数据读入寄存器。经过 Tensor Core 乘积累加后，这些脏数据产生的外积结果会被永久熔铸进累加器 `acc` 中。即使在写回时做了输出掩码，算出的数值也已经是错误的。 - **混合精度累加溢出陷阱**：若累加器初始化为 FP16（`tl.zeros((BM, BN), tl.float16)`），在 $K$ 较大（如 $K=4096$）时，连续累加会导致 FP16 范围溢出（最大值仅 65504）或极严重的下溢与截断误差。必须严格遵循“低精度输入计算、高精度 FP32 累加”黄金法则。 - **`tl.dot` 的形状约束**：`tl.dot` 的输入块各维必须 $\ge 16$，且块形状为 2 的幂；因此 `BM`、`BN`、`BK` 最小取 16。$M$ 很小时（如 $M=1$）也只能用 `BM=16` 加 mask，或改为不用 `tl.dot` 的逐元素乘加 + `tl.sum` 写法。 - **浮点结合律与容差预期**：由于 GPU Tensor Core 的并行点积顺序与 CPU 上的串行双精度循环截然不同，浮点加法不满足严格结合律，因此断言比较必须留有适当的相对容差（如 `atol=2e-2, rtol=2e-2`），不能按位要求 bitwise-exact。 ### 4. 性能与工程取舍 - **Tile 形状与流水线深度（`BM, BN, BK`, `num_warps`, `num_stages`）**： Triton GEMM 的性能极度依赖超参数选择： - `(BM, BN, BK)` 决定了每个 Program 的计算强度与片上寄存器消耗； - `num_stages` 决定了使用共享内存软件流水线（Async Pipeline，如 `cp.async`）预取的循环级数，通常 3~4 级最佳； - `num_warps`（常用 4 或 8）是每个 Program（即一个 CTA）使用的 warp 数，线程数为 $32 \times$ `num_warps`；它影响每个线程分到的元素数、寄存器压力和 SM 上能同时驻留的 CTA 数。 - **训练 GEMM（大 M）与推理 Decode GEMV（M=1）的极化差异**：在大模型预训练或 Prefill 阶段，$M$ 较大（如 $M = 4096$），属于典型的计算密集型（Compute-bound），需要选择大 Tile（如 $128 \times 128 \times 32$）打满 Tensor Core；但在自回归推理解码（Decode）阶段，每次仅输入 1 个 Token（$M=1$），算子退化为矩阵-向量乘法（GEMV），极度受限于显存带宽（Memory-bound）。此时若强行使用大 Tile 会导致 99% 的 Program 空转，必须采用专门针对 $M=1$ 优化的 Split-K GEMV 内核。 ## 具体演示以 $M = 17, K = 33, N = 19$、分块大小 $BM = 32, BN = 32, BK = 32$ 为例： - 网格发射 $\lceil 17/32 \rceil \times \lceil 19/32 \rceil = 1 \times 1 = 1$ 个 Program。 - **K 维循环迭代**：$\lceil 33/32 \rceil = 2$ 轮。 - **第 1 轮**（$k0 = 0$）：$rk \in [0..31]$。$A$ 掩码有效范围为 $[0..16, 0..31]$，$B$ 掩码有效范围为 $[0..31, 0..18]$。`tl.dot` 累加前 32 个通道的有效点积。 - **第 2 轮**（$k0 = 32$）：$rk \in [32..63]$。$rk < 33$ 仅有第 32 这一列有效，后 31 列越界。由于指定了 `other=0.0`，$A$ 和 $B$ 的后 31 个切片全被填充为 0.0。`tl.dot` 仅计算第 32 通道的有效乘加，多余维度的 $0.0 \times 0.0 = 0.0$ 累加不改变精度。 - 最终写回仅输出前 $[17, 19]$ 区域；数值与 `a @ b` 在容差内一致（累加顺序不同，不要求逐位相同）。请在阅读后先合上这一节，用自己的语言复述“输入状态 → 中间状态 → 输出状态”，再做练习。 ## 实践任务：唯一代码填空题补齐 K-tile 累加。规则：只能修改 `TODO`/`______` 所在位置；不要删除断言或放宽误差。代码注释说明了每个边界条件。 ```{python} #| course_role: exercise import torch import triton import triton.language as tl @triton.jit def matmul_kernel(a, b, c, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr, sam: tl.constexpr, sak: tl.constexpr, sbk: tl.constexpr, sbn: tl.constexpr, scm: tl.constexpr, scn: tl.constexpr, BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr): rm = tl.program_id(0) * BM + tl.arange(0, BM) rn = tl.program_id(1) * BN + tl.arange(0, BN) acc = tl.zeros((BM, BN), tl.float32) for k0 in range(0, K, BK): rk = k0 + tl.arange(0, BK) av = tl.load(a + rm[:, None] * sam + rk[None, :] * sak, mask=(rm[:, None] < M) & (rk[None, :] < K), other=0.0) bv = tl.load(b + rk[:, None] * sbk + rn[None, :] * sbn, mask=(rk[:, None] < K) & (rn[None, :] < N), other=0.0) acc += ______ # TODO: tile 矩阵乘 tl.store(c + rm[:, None] * scm + rn[None, :] * scn, acc, mask=(rm[:, None] < M) & (rn[None, :] < N)) def matmul(a, b): assert a.ndim == b.ndim == 2 and a.shape[1] == b.shape[0] M, K = a.shape; N = b.shape[1] c = torch.empty((M, N), device=a.device, dtype=a.dtype) grid = (triton.cdiv(M, 32), triton.cdiv(N, 32)) matmul_kernel[grid](a, b, c, M, N, K, *a.stride(), *b.stride(), *c.stride(), BM=32, BN=32, BK=32) return c for shape in ((17, 33, 19), (64, 64, 64)): M, K, N = shape a = torch.randn((M,K), device="cuda", dtype=torch.float16) b = torch.randn((K,N), device="cuda", dtype=torch.float16) torch.testing.assert_close(matmul(a,b), a@b, atol=2e-2, rtol=2e-2) ``` ### 检查方法在 CUDA/Triton 环境运行本单元格；断言覆盖规则尺寸和非规则尾块。首次 JIT 不计入性能。提交时请给出：补齐后的代码、实际运行输出（环境不可用时注明“仅静态审查”）以及对失败用例的解释。 ### Q1 不要背定义：请从输入、状态变化和输出三个阶段解释“Blocked Matmul 与 FP32 累加”的工作机制。 **你的答案：** ### Q2 只 mask 输出而不 mask K 尾块会怎样？ **你的答案：** ### Q3 为什么训练 GEMM 与 decode 的 M=1 GEMM 不应共用同一配置？ **你的答案：** ## 评分与通过规则 - 代码 4 分：正常输入 2 分，边界输入 1 分，解释实现 1 分； - Q1～Q3 各 2 分； - 一票否决：结果碰巧正确但核心因果链错误、删除边界检查、把未运行结果说成实测。需要提示时按四级机制请求：概念区域 → 具体方向 → 关键局部 → 完整答案。 ::: {.callout-tip collapse="true" title="💡 点击展开参考答案与详细代码实现" icon="false"} ## 参考答案（仅 answer 分支）先完成题目再核对。即使代码一致，也要能解释关键步骤，并尝试更换一个输入规模。 ```{python} #| course_role: solution import torch import triton import triton.language as tl @triton.jit def matmul_kernel(a, b, c, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr, sam: tl.constexpr, sak: tl.constexpr, sbk: tl.constexpr, sbn: tl.constexpr, scm: tl.constexpr, scn: tl.constexpr, BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr): rm = tl.program_id(0) * BM + tl.arange(0, BM) rn = tl.program_id(1) * BN + tl.arange(0, BN) acc = tl.zeros((BM, BN), tl.float32) for k0 in range(0, K, BK): rk = k0 + tl.arange(0, BK) av = tl.load(a + rm[:, None] * sam + rk[None, :] * sak, mask=(rm[:, None] < M) & (rk[None, :] < K), other=0.0) bv = tl.load(b + rk[:, None] * sbk + rn[None, :] * sbn, mask=(rk[:, None] < K) & (rn[None, :] < N), other=0.0) acc += tl.dot(av, bv) tl.store(c + rm[:, None] * scm + rn[None, :] * scn, acc, mask=(rm[:, None] < M) & (rn[None, :] < N)) def matmul(a, b): assert a.ndim == b.ndim == 2 and a.shape[1] == b.shape[0] M, K = a.shape; N = b.shape[1] c = torch.empty((M, N), device=a.device, dtype=a.dtype) grid = (triton.cdiv(M, 32), triton.cdiv(N, 32)) matmul_kernel[grid](a, b, c, M, N, K, *a.stride(), *b.stride(), *c.stride(), BM=32, BN=32, BK=32) return c for shape in ((17, 33, 19), (64, 64, 64)): M, K, N = shape a = torch.randn((M,K), device="cuda", dtype=torch.float16) b = torch.randn((K,N), device="cuda", dtype=torch.float16) torch.testing.assert_close(matmul(a,b), a@b, atol=2e-2, rtol=2e-2) ``` ### Q1 参考答案从三个阶段拆解其运行机制： 1. **输入阶段（2D 宏观网格与 K 维切片加载）**：以二维网格映射输出矩阵瓦片，Program 块初始化 FP32 寄存器累加器 `acc`。外层循环沿 $K$ 维以 $BK$ 为步长推进。每一轮迭代中，Program 分别构造 $A$ 与 $B$ 瓦片的二维指针，并在双向掩码（$A$ 受控于 $(M, K)$，$B$ 受控于 $(K, N)$）保护下，结合 `other=0.0` 将当前切片合并载入片上寄存器。 2. **状态变化阶段（Tensor Core 驱动的片上点积累加）**：调用 `tl.dot(av, bv)` 将当前 $A$ 瓦片与 $B$ 瓦片送入 GPU 硬件底层的 Tensor Core 矩阵乘加流水线。计算采用 FP16 乘法并在 FP32 累加器中就地累加，经过多轮 $K$ 维切片迭代后，累加器逐步收敛出该瓦片的最终真实点积和，整个循环计算完全在片上寄存器展开。 3. **输出阶段（输出掩码截断与类型转换写回）**：将高精度累加器 `acc` 转换为目标张量数据类型（FP16/BF16），结合输出边界掩码 `(rm < M) & (rn < N)`，通过 `tl.store` 写回全局目标矩阵 $C$。越界位置完全被硬件屏蔽，保证零越界写入。 ### Q2 参考答案若仅在输出写回 $C$ 时施加掩码，而在内层循环加载 $A$ 和 $B$ 时遗漏了针对 $K$ 维度的掩码保护（即没有 `rk < K` 约束）： 1. **尾块越界脏数据强行读入**：当 $K$ 不是分块尺寸 $BK$ 的整数倍时（例如 $K=33, BK=32$，在第二轮迭代中有效 $K$ 仅有 1 项，后 31 项为越界内存），内核会绕过边界检查，从显存越界地址中强行读取邻近内存的未初始化脏数据。 2. **脏数据永久熔铸进累加器**：在执行 `tl.dot(av, bv)` 时，这 31 个非法的脏数据向量会被送进 Tensor Core 进行点乘求和，产生不可预测的伪数值并直接累加进片上累加器 `acc` 中。 3. **输出毁坏**：即使在最后写回 $C$ 时施加了 $M$ 和 $N$ 维度的输出掩码，累加器内存储的数值本身已经被脏数据污染，最终写回目标张量的所有有效数值都会变成不可控的随机荒谬值。 ### Q3 参考答案大模型训练 GEMM 与自回归推理解码（Decode）阶段的计算特征处于完全不同的物理极端，绝对不能共用同一套分块配置： 1. **训练 GEMM（$M$ 极大，如 $M=4096, N=4096, K=4096$）**： - **体系特征**：属于典型的**计算密集型（Compute-bound）**算子，算术强度极高。 - **最优配置**：需要采用大分块尺寸（如 $BM=128, BN=128, BK=32$ 或 64），以最大化寄存器与片上数据复用率，释放 Tensor Core 的峰值 TFLOPs 吞吐，并使用 4~8 个 Warp 和深度流水线。 2. **推理 Decode GEMM（$M=1$）**： - **体系特征**：退化为矩阵-向量乘法（GEMV），算术强度极低（每个权重只被使用 1 次），属于绝对的**显存带宽密集型（Memory-bound）**算子。 - **配置问题**：若在 $M=1$ 时沿用训练的大分块（如 $BM=128$），tile 中 128 行有 127 行被 mask 掉，Tensor Core 算力基本空转；同时 grid 只有 $\lceil N/BN \rceil$ 个 Program（如 $N=4096, BN=128$ 时仅 32 个），远少于 H100 的 132 个 SM。$M=1$ 时应使用小 $BM$（用 `tl.dot` 时最小为 16，或改为不用 `tl.dot` 的乘加 + `tl.sum` GEMV 写法）、较小的 $BN$ 以增加 Program 数，并结合 Split-K（沿 $K$ 维切分到多个 Program，最后用原子加或第二个 kernel 归约）。 ## 参考资料 - [Triton Tutorials](https://triton-lang.org/main/getting-started/tutorials/) - [Triton language API](https://triton-lang.org/main/python-api/triton.language.html) 资料用于建立事实基线；面试回答仍需用自己的语言组织。 ::: ::: {.course-pagination} [**← 上一课**<br>第 8 课：二维 Tile 与转置](lesson08_transpose.qmd){.pagination-link .pagination-prev} [**下一课 →**<br>第 10 课：无状态随机数与 Dropout](lesson10_dropout.qmd){.pagination-link .pagination-next} :::