Building an LLM Inference Engine from ScratchPart V / Module 16

Part V · Performance

Module 16：优化案例研究 — 六层优化

本章按工程记录分析六项优化：decode 从 3.6 tok/s 提升到 30.9 tok/s，权重从 f32 改为 f16，并将 greedy argmax 放到 GPU。

学习目标

理解每个优化 tier 对应的瓶颈、修改和收益。
区分同步优化、带宽优化、kernel 访问模式优化、readback 优化。
理解为什么某些优化在短上下文下收益不明显，但仍然有长期价值。
学习如何写优化报告：baseline、假设、改动、验证、收益、风险。
完成 Lab 16：实现两个优化 tier，并报告自己的 before/after 数据。

16.1总览：从 3.6 到 30.9 tok/s

本项目的性能提升来自连续的瓶颈拆解。每一层优化对应一个明确问题，并以 parity tests 保证输出保持一致。

阶段	优化	decode tok/s	主要瓶颈
Baseline	per-op submit, f32 weights	3.6	同步 + 带宽 + readback
Tier 1	每层一个 command list	8.5	GPU submission overhead
Tier 2	f16 weights	10.3	显存带宽 / VRAM
Tier 3	coalesced GEMV	12.3	非合并权重读取
Tier 4	single-pass attention	~12	attention K/V 读带宽
Tier 5	cooperative RMSNorm	—	串行 row reduction
Tier 6	GPU greedy argmax	30.9	完整 logits readback

图 16-1：优化是逐层消除瓶颈；最大两次收益来自减少 GPU submission 和消除 logits readback。

16.2Tier 1：把 dispatch 批到一个 command list

baseline decode 在每个小算子后调用submit_and_wait。一个 token 约包含：

22\;\text{ops/layer}\times24\;\text{layers}\approx528

次 GPU submission。每次 submission 都会引入 CPU/GPU 同步成本。Tier 1 将一个 layer 内的 dispatch 连续记录到同一 command list，用 UAV barrier 保证依赖顺序，最后提交一次。

528\;\text{submits/token}\quad\longrightarrow\quad 24\;\text{submits/token}

该修改不改变数学计算，只改变命令提交方式；在相同输入下输出保持一致。decode 吞吐从 3.6 提升到 8.5 tok/s。

16.3Tier 2：f16 weights，带宽与显存减半

decode 阶段的 matmul/GEMV 多数受显存带宽限制。将权重从 f32 改为 f16 后，每个权重从 4 bytes 变为 2 bytes：

\operatorname{bytes}_{f16}(W) = \frac{1}{2}\operatorname{bytes}_{f32}(W)

xinfer 使用ByteAddressBuffer与f16tof32解包权重，activation 仍保持 f32。权重显存从 1885 MiB 降到 942 MiB，decode 从 8.5 提升到 10.3 tok/s。

数值注意

f16 权重会引入舍入误差，因此 parity test 需要设置合理容差，并检查 argmax 是否稳定。

16.4Tier 3：Coalesced GEMV + groupshared reduction

naive GEMV 由一个线程计算一个输出元素，相邻线程读取的权重地址相隔 $K$ ，不利于 coalescing。Tier 3 改为一个 threadgroup 计算一个输出元素：组内线程沿 $K$ 方向读取相邻权重，并在 groupshared 中归约。

y_j=\sum_{p=0}^{K-1}x_pW_{j,p}

每个线程负责 $p=t,t+TG,t+2TG,\ldots$ 的部分和。同一 wave 中的线程因此更可能读取连续权重。LM head 输出维度超过 65535，还需要 2D dispatch grid。

图 16-2：Coalesced GEMV 改善权重读取模式，并用 groupshared 做部分和归约。

16.5Tier 4：Flash-style single-pass attention

传统稳定 softmax 通常先扫描 key 以获得最大值，再扫描一次计算权重与加权 V。Tier 4 使用 online softmax，在一次遍历中维护 running max、denominator 和 accumulator。

m'=\max(m,s),\qquad d'=d\cdot e^{m-m'}+e^{s-m'}

accumulator 同样乘以 $e^{m-m'}$ 做 rescale，然后加上当前 $e^{s-m'}v$ 。短上下文下 attention 占比较低，因此收益有限；上下文变长后，减少 K/V 重复读取的价值会提高。

16.6Tier 5：Cooperative RMSNorm

若 RMSNorm 由一个线程处理一整行，decode 时每层会出现两个串行的 896 元素 reduction。Tier 5 改为一个 threadgroup 处理一行：线程分摊平方和，再用 groupshared reduction 合并。

\sum_{j=1}^{H}x_j^2 = \sum_{\text{threads }t}\left(\sum_{j\in \mathcal{J}_t}x_j^2\right)

这类优化的短期收益小于 Tier 1 与 Tier 6，但它消除了串行热点，也为后续 fused norm kernel 提供基础。

16.7Tier 6：GPU-side argmax

greedy 路径的 readback 是 decode 中的重要瓶颈。Qwen2.5 的词表大小为 151936；若 logits 为 f32，每步读回：

151936\times4\approx608\text{ KB}

greedy decoding 只需要最大 logit 的 index。Tier 6 增加 GPU argmax kernel，将最后一行 logits 在 GPU 上归约成一个u32token id，只读回 4 bytes。

608\text{ KB/token}\quad\longrightarrow\quad4\text{ bytes/token}

该优化使 decode 从 12.3 提升到 30.9 tok/s。它不减少模型计算量，但显著降低 CPU/GPU 传输和 CPU 扫描成本。

图 16-3：GPU argmax 把每 token readback 从 608 KB 降到 4 bytes。

16.8如何阅读结果：速度、显存、正确性一起看

综合结果如下：

指标	优化前	优化后	变化
decode tok/s	3.6	30.9	8.6×
prefill tok/s	~16	~378	~24×
weight VRAM	1885 MiB	942 MiB	-50%
输出	Paris	Paris	不变

优化报告应同时回答三个问题：

快了多少？
资源占用变了吗？
正确性是否保持？

Lab 16实现两个 tier，并写优化报告

从六个 tier 中选择两个实现。每个 tier 的报告必须包含：

baseline 数字；
你认为的瓶颈；
代码改动摘要；
正确性测试结果；
before/after 性能表；
如果没有提升，解释为什么。

参考格式见：

generated/optimization_report.md

小结

本章把 xinfer 的 decode 优化拆成六个案例：per-layer command list 将 submission 从约 530 次/token 降到约 24 次/token；f16 weights 将权重显存与读取带宽约减半；coalesced GEMV、single-pass attention 与 cooperative RMSNorm 改善 kernel 内部访问和归约；GPU greedy argmax 将 greedy 路径 readback 从约 608 KB/token 降到 4 bytes/token。Qwen2.5-0.5B 的 decode 吞吐由 3.6 tok/s 提升到 30.9 tok/s，所有阶段均以 parity tests 作为 regression gate。

思考与练习

基础哪一个 tier 主要减少 GPU submission？哪一个主要减少 readback？

T1（把整层前向录进一个命令列表、每层提交一次）主要减少 GPU submission——从约 530 submits/token 降到约 24。T6（GPU 端 greedy argmax）主要减少 readback——从 608 KB/token 降到 4 字节。

基础为什么 f16 权重能同时降低显存和带宽？

f16 每个权重 2 字节，是 f32 的一半。存储减半 → 显存从 1885 降到 942 MiB；而 decode 是带宽受限的，每次都要把整组权重从显存读进计算单元，读取字节减半 → 有效带宽翻倍，matmul 更快。代价是少量精度损失，需要用 fp32 累加、合理容差和 argmax 稳定性检查控制风险。

进阶为什么 attention 优化在短上下文下收益不明显？

因为 attention 的成本随上下文长度 $T$ 增长，而短上下文时 K/V 很少，attention 在总耗时中占比很小——瓶颈是各 matmul（QKV/MLP/LM head）。flash-style 单遍 attention 省的是对 K/V 的重复扫描带宽， $T$ 小时这点带宽微不足道。只有上下文变长（成百上千 token）后，attention 占比上升，单遍 + online softmax 的收益才显著。所以 xinfer 在短 ctx 下测得 T4≈中性。

进阶如果使用 temperature sampling，GPU argmax fast path 还能用吗？为什么？

不能直接用。temperature/top-k/top-p 需要完整（或至少 top-k 的）概率分布来做带随机性的采样，而 GPU argmax 只返回单个最大下标，丢掉了分布信息。所以带温度采样时要么读回完整 logits 在 CPU 采样，要么实现 GPU 端 top-k 只回传少量候选。greedy（等价 temperature=0）才只需 argmax，能走 4 字节 fast path。

挑战设计一个 GPU top-k 方案，估计它能减少多少 readback。

方案：在 GPU 上对 151936 个 logits 做并行 top-k。①每个 threadgroup 处理一段 logits，用局部堆/排序网络维护本段 top-k；②跨组归约合并各段的 top-k 得到全局 top-k（k 很小，可用一轮 LDS 归约或第二个 kernel）；③只把 k 对 (logit, index) 拷回 CPU，在 CPU 做温度缩放 + 归一化 + 采样。

readback 估计：原来回传全部 logits = $151936\times4\approx608$ KB。top-k 只回传 $k$ 个 (f32 logit + u32 index) = $k\times8$ 字节。取 $k=50$ ：约 400 字节，相比 608 KB 减少约 1500×（与 argmax 的 4 字节同量级）。代价是 top-k kernel 本身的复杂度与一点额外计算，但对受带宽或 PCIe readback 限制的 decode 通常有较高收益。