算子与运行时

全章的路线：从一次调用出发，由近到远，最后回到用户等待的完整请求。

调用一次 torch.matmul，CPU 先跑框架代码：读形状和类型、挑实现、准备地址参数，再经运行时和驱动把任务交给 GPU。先分清三个经常混用的词：

CPU 提交、GPU 执行，两者构成一条流水。大矩阵算得久，CPU 来得及准备下一项；小算子转瞬结束，GPU 就停下来等 CPU。所以加速有两个方向：缩短 kernel，或者减少加速器等主机的时间。

数据在哪：H2D、D2H、D2D

普通内存的数据要先被 CPU 复制进锁页缓冲，复制引擎才能搬，等于多走一遍。让 CPU 直接在可复用的锁页缓冲里准备输入，能省掉这次中转和反复分配。

stream 与 event：谁先谁后

stream 是一串按顺序执行的 GPU 任务：H2D 和读它的 kernel 放进同一个 stream，计算自然排在复制之后。复制和计算放在两个 stream 时，用 event 连起来：复制流在 H2D 后记录事件，计算流等它。这个等待发生在 GPU 的队列里，CPU 照样往下提交。

一个 kernel 在 SM 上怎么跑

图 5-3 里那段 20 μs 的 kernel 能不能缩短，要看 SM 里发生了什么。固定一个线程块：输出块 128×128、沿 K 每步 64；A 块、W 块各 16 KiB，FP32 累加器 64 KiB，都放共享内存，工作集共 96 KiB，256 个线程（8 个 warp）。硬件取 H100 的一个 SM。

把投影写成 C = A·W，M=1024，K=4096，N=12288，A、W、C 分别 8、96、24 MiB。算 C[i,j] 要 A 的第 i 行和 W 的第 j 列；接着算右边的 C[i,j+1]，A 还是同一行。这行 A 若还在片上，就能直接复用；若已被换走，就要重新搬。

分块就是让搬来的数据多用几次。固定一块 m×n 的输出，沿 K 依次搬入 A 块（m×k）和 W 块（k×n），一直累加到同一份 FP32 部分和里，最后才写回。A 块每个元素服务 n 列输出，W 块每个元素服务 m 行输出。

片上需求 S = 2mk (A块) + 2kn (W块) + 4mn (累加器) bytes
接口读写 V = 2MK·(N/n) + 2KN·(M/m) + 2MN

数据布局：让 32 个请求落到 32 个 bank

共享内存分 32 个 bank，每轮每个 bank 给一个 32-bit 字，bank 号 = 字地址 mod 32。一个 warp 的 32 个 lane 同时读 32×32 FP32 块的一列，行跨度 32 时全部撞进同一个 bank，要排 32 轮。

沿哪个维度切：并行任务与部分和

RMSNorm 对每行先求平方和（一次归约）再缩放。1024 行时一组一行就有 1024 个任务；decode 只有 1 行，同样的分法只产生 1 个任务，GPU 大面积空闲。把一行拆成几段，各自求局部平方和，再合并，再重读输入做归一化。

切 M 或 N，各执行者得到不同位置的完整结果，需要时拼起来；切 K，各执行者得到同一输出的部分和，必须相加。5.2.2 的输出块就是同时切 M、N，再在一个线程块内沿 K 累加，不跨任何边界。到第 6 章，这三种切法会变成多卡并行。

拆分还改变了数值

浮点加法不满足结合律。在 FP32 里算 1 + 2⁻²⁴ + 2⁻²⁴：

FFN 里 G = XW_g、U = XW_u，然后 Z = SiLU(G) ⊙ U，再交给 down 投影。G、U、Z 都是 [1024, 12288]，BF16 各 24 MiB。SiLU 和 ⊙ 都是逐元素运算：Z[i,j] 只依赖 G[i,j]、U[i,j]，读进来就能直接算出最终结果。

融合省掉的是中间结果的一次写回和一次重读，两步运算本身一步也没少。每少保存一个 24 MiB 的中间量，就少 48 MiB 读写。

双缓冲：搬下一块的同时算这一块

布局转换也有同样的取舍：另建一份 24 MiB 的重排结果要多 48 MiB 读写，不如让前一个算子直接按后一个需要的布局写出。缓冲则决定两个阶段能否重叠。槽 A 的“搬完了”和“用完了”是两个时刻，只有最后一次读结束，才能覆盖它。

FlashAttention：Softmax 也能逐块算

逐元素融合到 Softmax 这里卡住了：分母是整行指数和，只读到前一块时，既不知道最终分母，也不知道后面会不会出现更大的分数。标准做法 S = QKᵀ/√d、P = softmax(S)、O = PV，中间矩阵随序列长度平方增长。

前两节是手算的。实际模型有各种形状、各种算子组合、各种硬件，编译器做的是三件事：把运算写成循环与数组访问，按依赖和数值规则排除非法变换，结合容量与实测比较成本。

用循环变换表达分块与融合

以 Y = SiLU(AW) 为例。split 把 i 拆成块号 io 和块内坐标 ii（i = 64·io + ii）；reorder 把块循环提到外层；把 K 按 32 拆成 128 个归约块；最后把 SiLU 挪到每个输出块的 ko 循环之后（Halide/TVM 叫 compute_at）。

for i in range(M):
  for j in range(N):
    acc = 0.0
    for k in range(K):
      acc += A[i,k] * W[k,j]
    C[i,j] = round_to_bf16(acc)
for i in range(M):           # 再读一遍 C
  for j in range(N):
    Y[i,j] = silu(C[i,j])

for io in range(ceil_div(M, BM)):
  for jo in range(ceil_div(N, BN)):
    acc = zeros_fp32(BM, BN)    # 16 KiB
    for ko in range(ceil_div(K, BK)):
      a = load_tile(A, io, ko)   # 4 KiB
      w = load_tile(W, ko, jo)   # 4 KiB
      matmul_accumulate(acc, a, w)
    c = round_to_bf16(acc)
    store_tile(Y, io, jo, silu(c))

缩进层级就是保存期限。计算位置还决定复用还是重算：放在外层，先算完整结果供大家读；放进消费者的分块循环，只算当前需要的部分，多个块要用同一份时就会重复计算。

哪些变换不合法

AKG：多面体编译

AKG 把规则循环表示成三类信息：要执行哪些迭代、每次读写哪些数组元素、哪些读写必须保持先后。循环边界和下标能用线性约束描述（“多面体”由此得名），编译器据此推导依赖、调整顺序，并由数组访问反推每个 tile 需要哪块数据（64×64×32 的 tile 需要 A 的 64×32、W 的 32×64）。外层安排算子顺序，内层安排块的连续计算和局部缓存。手工做的“拆循环—定读取区域—留累加结果—归约后激活”，在这里成了一组联动的编译步骤。

成本估计与实测选择

实测往往推翻直觉：配套的矩阵转置程序在 RTX PRO 6000 上，16×16 块 7.9 μs，32×32 块 2.9 μs，2.7×。TVM 自动调优或 Agent 改 kernel，评价对象都是编译后的程序，因为编译器可能把源码里分开的算子融合成一个 kernel。

编译器选好程序后，运行时每一步还要准备输入、挑程序、提交任务、管理缓冲。kernel 代码一字不改，换一种提交方式或复用已有的准备结果，总时间也会变。

CUDA Graph：重放提交序列

先把一串 kernel 及其依赖录成图，之后主机每次只发一次 graph launch。GPU 还是逐个执行图里的节点。图重放减少的是主机提交次数，减少 kernel 数量靠的是融合。

动态形状：通用、分桶，还是特化

Persistent kernel：按数据块交接

后一个 kernel 通常要等前一个全部算完。如果它只需要前一个的一块结果，就把等待条件缩小到数据块：kernel 常驻 GPU，从队列取任务，用事件确认数据就绪（Mirage Persistent Kernel 即此类）。例子：512 个 token 的 up 投影接 SiLU，每 64 个 token 一块，共 8 块。

局部加速比必须放回整条请求里看。串行阶段按 Amdahl 定律相加；并行分支要沿依赖图找最长路径；共享资源还会让一个分支的优化拖慢另一个。

综合案例：一次 Qwen3-8B 请求

RTX PRO 6000，输入 7239 token，输出 32 token，并发 1，BF16、eager。36 层 × 32 步 = 1152 次 SwiGLU 激活调用：prefill 36 次（每次 7239 行），decode 1116 次（每次 1 行）。替换成一个新调度的 kernel（block256 / 4 warps）。

下一章把这些关系推广到多张加速器：一片结果何时就绪、谁要读、保存多久，将决定通信何时开始，以及两端要多少缓冲。