Building an LLM Inference Engine from ScratchPart V / Module 17

Part V · Performance

Module 17：下一步高级优化

在 correctness 与基础优化稳定后，本章讨论尚未实现或仍需完善的方向：量化、tiled GEMM/GEMV、fused epilogue、GPU sampling、speculative decoding 与 paged KV。

学习目标

理解 weight-only quantization 与 activation quantization 的差异。
理解 int8 / int4 / GGUF 对显存、带宽、精度和 kernel 复杂度的影响。
理解 tiled GEMM、tensor-core / WMMA 类路径为何需要 DXC SM6.x。
理解 fused epilogue 如何减少 dispatch 和中间 buffer。
理解 GPU top-k/top-p、speculative decoding、batching、多序列、paged KV cache 的基本思想。

17.1Quantization：int8 / int4 / GGUF

当前 xinfer 使用 f16 weights 与 f32 activations。量化是后续方向之一：用更少 bit 表示权重，以继续降低显存占用和权重读取带宽。

格式	字节/权重	优点	代价
f16	2	简单，硬件支持好	仍较大
int8	1	显存/带宽再减半	需要 scale / zero-point
int4	0.5	极省显存，适合大模型	解包复杂，精度风险更高

weight-only quantization 只量化权重，activation 仍用 f16/f32。典型反量化公式：

w \approx s \cdot (q - z)

其中 $q$ 是整数权重， $s$ 是 scale， $z$ 是 zero-point。许多 LLM 推理格式（例如 GGUF 的若干量化类型）按 block 存储 scale，使不同权重块使用不同缩放因子。xinfer README 将 int8/int4 quantization（GGUF）列为 future work。

图 17-1：量化减少权重字节数，但引入解包和数值误差。

17.2Tiled GEMM 与 tensor-core / WMMA 路径

当前 coalesced GEMV 已经消除了一部分低效访问，但矩阵乘法仍有进一步优化空间。后续路径会使用 tile：一个 threadgroup 处理 $M\times N$ 的小块，把 A 和 W 的 tile 缓存在 groupshared，然后重复利用。

C_{tile} = A_{tile} W_{tile}^{\top}

对大 $M$ 的 prefill，tiled GEMM 可以提高 arithmetic intensity。对 decode $M=1$ ，README 中列出的方向是 column-tiled GEMV：在 LDS 中缓存 activation vector，并同时计算多个输出列，减少对 $x$ 的重复读取。

Tensor-core / WMMA 类路径需要更现代的 shader 能力（例如 DXC / SM6.x 的 wave/tensor 操作）。这类路径可能提高吞吐，但也会增加平台、编译和兼容性约束。

图 17-2：Tiled GEMM 的核心是提高数据复用，提升 arithmetic intensity。

17.3Fused epilogues：把小算子合进大算子

当前 pipeline 中，若干小算子会产生中间 buffer，例如 linear → bias_add → RoPE，或 gate/up → SwiGLU → down。将 bias、activation、residual 等操作融合到 matmul epilogue 中，可以减少：

dispatch 次数；
中间 buffer 写入/读取；
UAV barrier；
显存带宽。

例如 Q projection 可以从：

Q = XW_Q^\top,\qquad Q'=Q+b_Q

融合成：

Q'_j = \sum_p X_p W_{Q,j,p}+b_{Q,j}

SwiGLU 也可作为 gate/up matmul 后的 fused epilogue；由于它依赖两个 matmul 的结果，调度和寄存器管理更复杂。README 将 bias / SwiGLU / residual fused epilogue 列为 future work。

17.4GPU sampling 与 speculative decoding

当前 xinfer 对 greedy 使用 GPU argmax；temperature/top-k/top-p 仍读回完整 logits 并在 CPU 上采样。后续可以实现 partial GPU top-k：

GPU partial top-k：在 GPU 上保留少量候选 token；
只读回 top-k token ids 和 logits，而不是完整 vocab；
CPU 或 GPU 完成最终随机采样。

speculative decoding 是另一类 future direction：小模型先草拟多个 token，大模型批量验证候选。理想情况下，一次大模型 forward 可接受多个候选 token，从而减少自回归循环次数。

图 17-3：GPU top-k 的目标：readback 大小从完整 vocab 变成少量候选。

17.5多序列 batching 与 paged KV cache

单用户本地 demo 通常只处理一个 sequence；服务端推理需要同时服务多个请求。Batching 将多个 sequence 合并执行，可提高 GEMM 利用率，但也带来复杂度：

不同请求 prompt 长度不同；
不同请求生成长度不同；
KV cache 需要动态分配和回收；
attention mask 变得更复杂。

paged KV cache 将 KV cache 切成固定大小的 page，并通过 page table 管理逻辑位置到物理 page 的映射，从而避免为每个请求预留最大长度的连续 buffer。

\text{logical token position} \rightarrow (\text{page id},\text{offset within page})

从教学引擎到生产引擎

xinfer 当前是以单序列和可读性为主要目标的教学引擎。面向生产系统时，batching、paged KV、调度器和量化会成为核心工程问题。

Lab 17提出一个高级优化设计

本实验不要求完整实现。请选择一个高级方向，写一份设计说明，明确它属于 future work 还是已实现路径的扩展：

你要解决什么瓶颈？
需要改哪些模块？
正确性如何验证？
性能如何测量？
可能引入哪些数值或平台风险？

可选主题：int4 weight-only quantization、GPU top-k、fused bias epilogue、tiled GEMV、paged KV cache、speculative decoding。

小结

本章讨论的方向大多属于 future work。int8/int4 quantization（GGUF）继续降低权重显存和带宽；column-tiled GEMV 通过在 LDS 缓存 activation vector 提高 decode 的数据复用；fused epilogue 可减少 dispatch 与中间 buffer；partial GPU top-k 可降低 temperature/top-k/top-p sampling 的 readback；speculative decoding、batching 与 paged KV 则面向多请求和长上下文场景。这些方向都应先定义 correctness gate，再用稳定 benchmark 验证收益。

思考与练习

基础weight-only quantization 与 activation quantization 有什么区别？

weight-only 只把权重量化到低位（如 int4/int8），激活仍保持 f16/f32，计算时把权重解量化回浮点再算——主要节省显存与权重读取带宽，对 decode 这种带宽受限场景有效，精度风险相对较低。activation quantization 还把激活也量化，能用整数矩阵乘进一步加速计算，但激活分布动态范围大、更难量化，精度下降风险较高，常需 per-tensor/per-token scale 与校准。LLM 推理常先采用 weight-only。

基础为什么 int4 权重能显著减少带宽？

int4 每个权重 0.5 字节，是 f16 的 1/4、f32 的 1/8。decode 受限于把权重从显存读入计算单元的带宽，读取字节数随位宽下降；实际收益还要扣除解量化、scale 读取和访存对齐开销。显存占用也随之降低，可容纳更大模型。代价是需要 group-wise scale/zero-point 和解量化 kernel，并接受一定精度损失。

进阶举例说明 fused epilogue 如何减少一次中间 buffer 写读。

例：MLP 的 SwiGLU。未融合时 gate 投影输出写回显存、up 投影输出写回显存，再启动一个 elementwise kernel 读回两者算silu(gate)*up又写回。融合后，在 matmul kernel 的尾段（epilogue）算出 gate/up 结果时，直接在寄存器里做silu(gate)*up再写一次结果，省掉中间两次写、一次读和一个 dispatch。类似地可把 bias、residual add 融进 matmul epilogue。每减少一次中间写读，就减少一部分带宽和 dispatch 开销；在带宽受限的 decode 中，这类收益可以累积。

进阶GPU top-k 相比 GPU argmax 多了哪些难点？

argmax 只需一个 max 归约（结合律好、单值、一轮WaveActiveMax/LDS 即可）。top-k 要维护“前 k 大”的有序集合：①每个线程/组要保留 k 个候选而非 1 个（堆或排序网络，寄存器/LDS 压力大）；②跨组合并多个局部 top-k 成全局 top-k 需要第二阶段归约，且合并不是简单求 max；③要同时跟踪 value 和原始 index；④处理并列、k 与 group 大小不整除等边界。复杂度和实现成本都明显高于 argmax。

挑战设计 paged KV cache 的数据结构，并说明 attention kernel 如何根据 page table 读取 K/V。

数据结构：把 KV cache 切成固定大小的 page（如每 page 存 16 个 token 的 K/V）。维护一个 page pool（一堆等大的 GPU block）和每个序列的page_table：一个数组，第 $b$ 项给出该序列第 $b$ 个逻辑块对应的物理 page id。token 位置 $p$ → 逻辑块 $b=\lfloor p/\text{page\_size}\rfloor$ ，块内偏移 $o=p\bmod\text{page\_size}$ 。

kernel 读取：attention 遍历 key 位置 $p$ 时，不再用连续地址，而是：查page_table[b]得到物理 page，基址 = page_base + $o\times$ （每 token 的 K/V 字节）。即一次额外的间接寻址。好处：①显存按需分配、几乎零碎片（vLLM 思路），能服务更多并发序列；②序列间可共享 page（如相同前缀）。代价是每次 K/V 访问多一层 page_table 查找，且 kernel 要传入 page_table 与 page_size。