第 17 章 高性能 Kernel、图编译与执行后端¶
开场问题：所有该开的开关都开了，离硬件理论上限还差 25%。这 25% 在哪几个 kernel 里？
17.1 先 profile，再动手¶
模型代码里的一行 x = x + attn(norm(x))，在 GPU 上可能是七八个 kernel，每个都要把数据从显存读进来、算完写回去。性能工程关心的从来不是代码行数，而是最终发射了多少 kernel、每个 kernel 读写了多少字节。
看 profiler 时间线要回答四个问题：GPU 在算的时间占多少、kernel 之间有没有空隙（CPU 提交跟不上）、通信有没有和计算重叠、有没有某个意外的 kernel 占了大头。这四个问题对应四种完全不同的修法，猜错方向做的优化全是白工。
17.2 算子融合：主要为了少读写¶
大量逐元素算子（norm、激活、残差加、RoPE、量化和反量化）的计算量微不足道，成本几乎全在"从显存读进来、写回去"。把它们融合成一个 kernel，中间结果留在寄存器或共享内存里，一次读写解决。
常见的融合包括：RMSNorm + 残差、SwiGLU 的门控乘法、QKV 投影后的 RoPE、以及采样链路上的一串小操作。融合也不是越大越好——寄存器用量超限会导致占用率下降，反而变慢。
17.3 FlashAttention 作为方法论¶
第一册 3.4 节讲过它：分块加在线 softmax，让 n×n 的注意力矩阵从不写进显存。
这里要强调的是方法而不是这个具体算子：先识别出瓶颈是访存而非计算，再重排计算顺序让数据留在片上，数学结果完全不变。第 11 章的分页 KV 要求 attention kernel 支持间接寻址，第 14 章的树状验证要求特制 mask，都是在这个框架下的进一步改造。
17.4 Decode 阶段需要专门的 kernel¶
Prefill 和 Decode 的 attention 形状差别极大：Prefill 是 [n, n] 的方阵，Decode 是 [1, n] 的细长条。同一个 kernel 不可能对两者都最优。
Decode kernel 的核心困难是并行度不足：只有一个 query，如果按常规方式并行，大部分 SM 都闲着。解法是沿 KV 长度方向切分（split-K），让多个线程块各算一段，最后归并。GQA 还有额外的优化空间（多个 Q 头共享同一份 KV，可以合并读取）。
MoE 则需要 grouped GEMM：把路由到不同专家的 token 分组，一次 kernel 调用完成多个形状不同的矩阵乘。
17.5 CUDA Graph：消掉 CPU 开销¶
每个 kernel 的启动都要 CPU 参与，一层十几个 kernel、几十层下来就是几百次启动。Decode 阶段单步 GPU 计算可能只有几毫秒，而 CPU 提交这几百个 kernel 也要几毫秒——CPU 成了瓶颈。
CUDA Graph 的做法是把一整段 GPU 操作序列捕获下来，之后每次一条命令回放整张图，CPU 开销几乎归零。小模型上收益可以到 30% 以上。
代价是图要求形状固定，而在线推理的 batch 和序列长度一直在变。工程上的做法是预先捕获若干个 batch 档位（比如 1/2/4/8/16/32）的图，实际运行时向上取整到最近的档位并 padding。这也是为什么分页 KV 的设计要尽量让形状变化可枚举。
17.6 编译栈¶
torch.compile 做图捕获、算子融合和内存规划，对训练和小规模推理都有效；Triton 让你用 Python 风格写 kernel，性能接近手写 CUDA 而开发成本低得多，现在大量新算子首发在 Triton；TensorRT-LLM 这类 AOT 编译引擎做得更彻底，离线编译出高度特化的执行计划，代价是编译耗时长、动态性差。
共同的敌人是动态形状：形状一变就可能触发重新编译或退回慢路径。冷启动时间（加载权重 + 编译 + warmup）在弹性扩缩容场景里是实打实的成本，第 19 章会再提到它。
17.7 什么时候该自己写 kernel¶
这是个成本判断，不是技术判断。三个条件同时满足才值得：profiler 证明这个算子长期占据关键路径的显著比例；现有后端确实没有高效实现（先去上游仓库确认一遍，很多时候已经有人在做了）；你有能力持续维护它（硬件换代、模型结构变化都要跟着改）。
大多数团队更划算的做法是：把形状和配置调到能命中现成的高效 kernel，而不是自己写一个。
核心问答¶
Q：两个引擎跑同一个模型，FLOPs 一样，为什么速度差几倍？ FLOPs 不包含访存量、kernel 启动开销、融合程度和调度效率，而这些才是推理性能的主要构成。
Q：CUDA Graph 为什么和动态 batch 冲突？ 图捕获时形状就固定了。解法是准备若干形状档位并 padding，用一点浪费换 CPU 开销归零。
Q：Triton 会取代 CUDA 吗？ 大部分常规热点会被 Triton 覆盖，但极致优化和新硬件特性的首发支持仍需要 CUDA。两者是分工不是替代。
Q：优化的顺序应该是什么？ 先前缀缓存和调度（消掉工作量），再量化（减少访存），最后才是 kernel 和编译（抠常数项）。很多团队顺序搞反，在最难的地方花最多时间拿最少收益。