第 14 章 · Capstone: 用自己写的 kernel 跑 GPT-2 small

学习目标

完成这个 capstone 后，你应能解释 GPT-2 small 的端到端数据流、权重布局、kernel 串联方式与生成循环，并能为后续优化建立可复现基线。

前置知识

到第 13 章，你已经手写过：

• kernel 启动模板（Ch2-3）
• shared memory tile（Ch6）
• warp shuffle reduce（Ch7）
• tiled GEMM（Ch9）
• 数值稳定 / online softmax（Ch10）
• attention 三阶段 + FlashAttention（Ch11-12）
• RoPE / SwiGLU / KV cache / sampling（Ch13）

本章把它们串成一个能跑的 GPT-2 small (124M)。

核心概念

14.1 GPT-2 small 架构速览

graph TD
    Tok["token ids (T,)"] --> Emb["embedding: wte + wpe
(T, 768)"]
    Emb --> Loop["× 12 layers"]
    Loop --> LN1["LayerNorm"]
    LN1 --> Attn["fused QKV → causal MHA → out_proj"]
    Attn --> Res1["+ residual"]
    Res1 --> LN2["LayerNorm"]
    LN2 --> MLP["fc (768→3072) → GeLU → fc_proj (3072→768)"]
    MLP --> Res2["+ residual"]
    Res2 --> Loop
    Res2 --> LNF["final LayerNorm"]
    LNF --> Head["logits = h @ wte^T  (T, 50257)"]
    Head --> Smp["argmax / top-k / top-p → next token"]
    style Loop fill:#f3f1e8,stroke:#8b1538
    style Attn fill:#f3f1e8,stroke:#a86420

关键代码

14.2 三步上手

Step 1 — 下载权重

pip install transformers torch
python data/download_gpt2.py --out data/gpt2-small.bin
# 输出: data/gpt2-small.bin (~250 MB, fp16)

脚本会把 HuggingFace 的 gpt2 权重按固定二进制 layout 序列化，方便 C++ 一次 fread。 文件头是 16 个 int32（含 n_layer / n_head / d_model 等），后接 fp16 张量流。

Step 2 — 编译与运行

cd code/ch14_mini_llm
make ARCH=sm_80
./mini_llm --weights=../../data/gpt2-small.bin \
           --tokens=15496,11,612,318 \
           --max_new=10

Step 3 — token ↔ 文本

本仓库为简化没在 C++ 里写 BPE tokenizer。 用 Python 编码 / 解码：

from transformers import GPT2Tokenizer
tok = GPT2Tokenizer.from_pretrained("gpt2")
ids = tok.encode("Hello, there is")
print(ids)                  # → [15496, 11, 612, 318]
# 把上面 ids 传给 mini_llm --tokens=...
# 然后把它打印出的 final tokens 复制回 Python:
print(tok.decode([15496, 11, 612, 318, 1043, 257, 1810, 2950, 287]))

运行结果

正确运行应完成权重加载、逐层 forward、logits 计算和 token 选择，并输出可由同一 GPT-2 tokenizer 解码的 token 序列。GPU 数值结果与生成文本需在目标环境实测：TODO(on GPU)。

没有 GPT-2 权重时，可先运行确定性的 toy transformer 对拍；它覆盖 embedding、LayerNorm、causal attention、MLP、residual、final norm 与 final logits，不替代真实权重的逐层回归：

./mini_llm --self-test --seed=1234
# [mini_llm_self_test_final_logits] PASS ...

14.3 forward() 的结构

源码：mini_llm.cu。 核心 forward 函数把所有 kernel 按顺序串起来，每步：

void forward(const Weights& W, std::vector<int>& tokens, int* next) {
    embed_kernel(wte, wpe, tokens, x);            // (T, D)

    for (int l = 0; l < n_layer; ++l) {
        layernorm_row(x, ln1_w, ln1_b, normed);
        matmul_kernel(normed, qkv_w, qkv);        // Ch6 tile
        bias_add(qkv, qkv_b);
        split_qkv(qkv, Q, K, V);
        mha_naive(Q, K, V, attn_out);             // Ch11 naive
        matmul_kernel(attn_out, proj_w, proj_out);
        bias_add(proj_out, proj_b);
        residual_add(x, proj_out);                // x += proj_out

        layernorm_row(x, ln2_w, ln2_b, normed);
        matmul_kernel(normed, fc_w, ff_mid);      // (T, d_ff)
        bias_add(ff_mid, fc_b);
        gelu_kernel(ff_mid);                       // Ch13 风格激活
        matmul_kernel(ff_mid, fcp_w, ff_out);      // (T, D)
        bias_add(ff_out, fcp_b);
        residual_add(x, ff_out);
    }

    layernorm_row(x, lnf_w, lnf_b, normed);
    gemv_logits(normed[T-1], wte, logits);         // (V,)
    greedy_argmax(logits, next);
}

14.4 简化与代价

即便如此，本章 GPT-2 small 已足以验证端到端流程。吞吐请在目标 GPU 上实测； 优化前后记录 tokens/s、TTFT、TPOT、batch、prompt length 和编译架构，不在教程里写死未经验证的数值。

性能数据

TODO(on GPU)：固定 GPU、driver、CUDA、commit、prompt length 与生成长度，报告 TTFT、TPOT、tokens/s 和峰值显存；同时分别记录 warmup 与 steady-state。

自检清单

• 能画出 GPT-2 一层中 LayerNorm、attention、MLP 与 residual 的数据流。
• 能解释为什么本章每步重算全序列，以及 KV cache 会改变哪部分复杂度。
• 能说明 fp32、batch=1 和朴素 attention 为什么只适合作为教学基线。
• 能用固定 token 输入比较 CPU/框架参考与 CUDA 输出，而不是只看生成文本“像不像”。

练习题

以下扩展任务按难度递进：

★ 入门

1. 支持 prompt 长度 > 16：测试 --max_new=64
2. 把 greedy 换成 temperature=0.8 + top-k=40
3. 用 nsys 看一次 forward 的 timeline，找最大 kernel

★★ 进阶

1. 加 KV cache：维护 K_cache / V_cache (n_layer, n_head, T_max, D_head)，每步只 forward 1 token。预期方向是把单步复杂度从重算全序列降到只读历史 KV；具体延迟写入 benchmark 表。
2. fp16 路径：fc/proj/qkv 用 WMMA。Ch9 模板可移植。
3. 用 Ch12 FlashAttention 替换 mha_naive。

★★★ 挑战

1. continuous batching：多请求并发，按 token 调度而非按请求。这是 vLLM 的核心创新。
2. W4A16 量化：把权重压成 INT4，runtime 解 → fp16 GEMM。看 llama.cpp / AWQ 实现。
3. speculative decoding：用 distilgpt2 当 draft，gpt2 当 verifier，记录 accept rate、draft cost、verify batch 和端到端 TPOT。
4. 移植到 Llama 架构：把 LayerNorm 换 RMSNorm、GeLU 换 SwiGLU、加 RoPE、加 GQA。

14.6 工业实战：从能跑到能上线

capstone 能在单卡跑通 GPT-2 small。距离给真实用户服务还差 6 个工程组件，本节给概念地图。

14.6.1 连续批处理 (Continuous Batching)

朴素 batch：N 个请求攒齐再 forward，慢的拖快的（"head-of-line blocking"）。Continuous batching 是 iteration-level scheduling——每个 token 步骤独立组 batch，新请求随时插入，完成的随时移出。

sequenceDiagram
    participant Q as RequestQueue
    participant S as Scheduler
    participant E as Engine
    Q->>S: req A, prompt T=100
    Q->>S: req B, prompt T=50
    S->>E: prefill A, B together
    Note over E: 同 batch 处理 prefill
    E->>S: gen token A=t1, B=t1
    Q->>S: req C, prompt T=20
    S->>E: prefill C plus decode A, B
    E->>S: A=t2, B=t2, C=t1
    Note over Q,E: B 提前完成, 出 batch, 槽位让给新请求

核心数据结构：request slot 数组，每个 slot 持有 (state=prefill/decode, KV pages, current token, sampling params)。每步根据 slot 状态组 batch。

vLLM、TGI、TensorRT-LLM 都内置 continuous batching。收益取决于请求长度分布、arrival rate 和 KV pool；上线前要用真实 trace 压测，而不是引用固定倍数。

14.6.2 PagedAttention KV 管理（集成版）

12.9.4 给了算法，在 mini_llm 里集成步骤：

1. K_cache / V_cache 改成 page pool：K_blocks (n_blocks, block_size, n_head, D_head)
2. 每请求一个 block_table：logical_block_id → physical_block_id
3. attention kernel 通过 block_table 间接寻址
4. Scheduler 维护 free block pool；新请求 alloc，结束归还

陷阱：block 大小要平衡。太大（256+）→ 内部碎片；太小（4-）→ block_table 大、kernel indirection 开销大。生产典型 16 token / block。

14.6.3 多 GPU：TP / PP / EP / SP

单卡装不下（70B fp16 = 140 GB）时分卡跑：

典型 70B 推理：TP=8, PP=1（单节点 8×A100/H100）。kernel 改动：QKV proj 输出维度 ÷TP，attention head ÷TP，FFN 中间维 ÷TP，每层最后 1 次 all-reduce (NCCL)。

14.6.4 CUDA Graph capture — 杀掉 launch overhead

单 token decode 会触发数十个小 kernel，launch gap 很容易和实际计算同量级。capture 成一个 Graph 后重放只保留一次 graph launch，具体收益用第 15 章的 nsys 流程采集：

// warmup 时同时 capture
cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal);
decode_one_step(...);
cudaStreamEndCapture(stream, &graph);
cudaGraphInstantiate(&exec, graph, ...);

// 之后每步只 1 次 launch
cudaGraphLaunch(exec, stream);

本章 mini_llm.cu 已经内置 NVTX range： decode_loop、layer_N、attention/*、mlp/*、logits/*、sampler/*。 运行时加 --profile 会用 cudaProfilerStart/Stop 把采集范围限制在 decode loop，避免把权重加载也录进去。

cd code/ch14_mini_llm
make ARCH=sm_80                  # 缺 libnvToolsExt 时: make NVTX=0

nsys profile --stats=true \
  --trace=cuda,nvtx,osrt \
  --capture-range=cudaProfilerApi \
  --capture-range-end=stop \
  --force-overwrite=true \
  -o mini_llm_decode \
  ./mini_llm --weights=../../data/gpt2-small.bin \
             --tokens=15496,11,612,318 \
             --max_new=4 \
             --profile

看 trace 的顺序：先找 decode_loop/step_* 之间是否有 CPU gap，再展开某个 layer_N， 比较 attention/qkv_matmul、attention/mha_naive_causal、mlp/fc_gelu、 logits/lm_head_gemv 和 sampler/greedy_argmax 的相对耗时。

坑：

• kernel 参数（KV 位置等）必须编译期常量。生产用 cudaGraphExecKernelNodeSetParams 动态更新
• batch size 变化要重新 capture；预 capture 几组 (1, 2, 4, 8, 16, 32) 的 graph
• 动态 KV 长度需要 padding 到固定 bucket，浪费一些算力换 graph 复用

14.6.5 服务化工程要点

• 异步 IO：HTTP queue → 内部 scheduler；asyncio / tokio。不能阻塞 GPU stream
• 流式输出：每生成一个 token 立即 SSE / gRPC streaming 推回 client
• 请求超时与取消：client 关连接，对应 slot 立即释放 KV blocks
• 动态 max_tokens：根据队列长度调整，繁忙时拒绝过长生成
• warmup：服务起后跑虚拟请求预热 cuBLAS / CUDA Graph，否则前几个真请求会混入初始化和 lazy loading 成本
• OOM 防护：admission control 估算每请求需要多少 KV blocks，不够就排队
• 多副本 + 自动扩缩容：K8s + HPA
• 监控：p50/p95/p99 latency、tokens/s、queue length、KV 占用率、GPU util

14.6.6 生产 LLM 推理栈对比（2025）

14.6.7 推荐工程化路径

1. 原型：vLLM docker 一行命令
2. 生产：vLLM 或 TRT-LLM；按 NVIDIA 平台粘度选
3. 极致优化：基于 CUTLASS / Triton 自写关键 kernel（attention、量化 GEMM），其他算子复用
4. 跨平台：llama.cpp (Mac / CPU)、MLX (Apple Silicon)、ROCm (AMD)

14.6.8 学完本教程能做什么

• 读懂 vLLM / TRT-LLM / llama.cpp 任一个的 CUDA kernel
• 给自己模型在 vLLM 里加自定义 fused 算子
• 实现 attention 变体（local / ALiBi / custom mask）
• 调优 LLM 服务并用 profile 证明 p99 latency 改善来自哪里
• 面试 GPU/CUDA 岗位，能答 FlashAttention 推导和 PagedAttention 原理

这是起点不是终点。顶级 CUDA 工程师还需 3-5 年实际项目经验，但你已经过了"看不懂代码"的门槛。

14.7 研究前沿（2025-2026）：Disaggregated Serving、MoE、Reasoning

14.7.1 Disaggregated Prefill / Decode（Mooncake、DistServe）

14.6 讲了 prefill 和 decode 是两个不同世界（compute-bound vs memory-bound）。共享 GPU 跑两者意味着两边都不能最优——prefill 抢占 SM 会拖慢 decode 的尾延迟。

2024-2025 的革命：把 prefill 和 decode 分到不同 GPU pool，KV cache 通过高速网络传递。

graph LR
    A["用户请求"] --> B["调度器"]
    B --> P1["Prefill GPU pool
(compute-bound, 大 batch)"]
    P1 -->|"KV 传输
NVLink/IB"| D1["Decode GPU pool
(memory-bound, 小 batch)"]
    D1 --> O["流式输出"]
    style P1 fill:#f3f1e8,stroke:#8b1538
    style D1 fill:#f3f1e8,stroke:#2f5d3a

主要工程实现：

• Mooncake（Moonshot AI / Kimi 2024）：业内首个生产 disaggregated 系统，KV 通过 RDMA 传输，全局 KV pool 跨节点共享
• DistServe（OSDI 2024 论文）：理论分析 + 早期开源实现
• vLLM V1：unified token-budget scheduler；disaggregated serving 能力按当前 feature guide 核验
• SGLang：同样支持
• llm-d（Red Hat / IBM 2025）：Kubernetes 原生 disaggregated 框架

典型目标是提高 tokens/$ 并降低 decode 尾延迟，但收益强依赖 prefill/decode 比例、网络、KV 传输和调度器复杂度。小规模团队（< 10 GPU）通常先优化单卡与单节点。

14.7.2 MoE 推理 — Expert Parallel 是新刚需

2024-2025 大模型几乎全转 MoE：DeepSeek-V3 671B（37B 激活）、Llama 4 Maverick 400B、Qwen 3 235B、Mixtral 8×22B 等。共同点：

• 总参数 100B+，但每 token 只激活 ~10-15% expert
• 显存 / HBM 流量瓶颈极重
• 路由不均衡（hot expert 卡死）

Expert Parallel (EP)

每张 GPU 装一部分 expert (例如 64 expert / 8 GPU = 8 expert/GPU)
每 token 路由到 top-k expert (典型 k=2 或 8)
跨 GPU 通信: all-to-all (一次发, 一次收)

挑战：

• all-to-all 时网络带宽是瓶颈 — 必须 NVLink 或 IB
• 路由不均衡 — 用 DeepSeek-V3 的 aux-loss-free balancing 或 expert capacity factor
• Expert 跨 GPU 调度 — vLLM fused MoE kernel、SGLang 都有自己实现

实战：DeepSeek-V3 推理架构

• 单实例 = 32 GPU H800 (4 节点 × 8 卡)
• TP=8（每节点）+ EP=32 + PP=1
• 每节点 NVLink 内 all-to-all，节点间 IB
• FlashMLA + paged KV + fp8 GEMM
• 吞吐需在目标 GPU 与真实并发上实测，记录 batch、accept rate、context length 和 tokens/s

14.7.3 Reasoning 模型的 serving

o1 / R1 / Gemini Thinking 等 reasoning 模型对推理服务的颠覆性影响：

服务系统的调整：

• BIG KV pool：单请求几 GB，PagedAttention 必须
• Speculative decoding 需要 trace 验证：统计 n-gram lookahead 命中率、accept rate、draft cost 和 verifier batch 效率
• Streaming 必须：用户看 thinking 流
• 取消机制：用户可能中途打断 thinking，要立即释放 KV

14.7.4 Prefix Caching 的演进

2024-2025 prefix caching 从 "memory cache" 升级到 "持久化共享池"：

• SGLang RadixAttention：进程内 trie 缓存
• LMCache / vLLM Cache（2024）：跨进程 KV 共享，RDMA 传输
• Mooncake Conductor：全集群 KV pool, 用 CRUSH 算法分布
• 磁盘 / SSD KV cache：冷 prefix 卸载，需要时拉回

对 chat 工作负载（系统 prompt + 多轮对话），prefix cache 的价值取决于系统 prompt 复用率、多轮上下文共享率和 KV 迁移成本。上线前要用真实 trace 统计命中率、节省 token 数、KV pool 压力和端到端吞吐。

14.7.5 多模态推理：vision + audio + text

Gemini 2.0 / GPT-4o / Llama 4 都是原生多模态。推理服务变化：

• vision encoder：ViT 或 SigLIP，前置阶段 GPU 跑（compute-bound）
• token 化：image / audio 转 visual tokens 喂给 LLM
• 异构 batch：text-only 请求和 multimodal 请求形状不同, scheduler 复杂
• encoder 缓存：相同图片不重复 encode

典型部署：vision encoder 作为独立微服务（甚至独立 GPU），LLM serving 拉它的输出 token。跟 disaggregated prefill 思路相通。

14.7.6 Agentic / Computer Use serving

2024-2025 Claude Computer Use、OpenAI Operator、AutoGPT 等"agent"：模型自动循环调用工具。serving 挑战：

• 长 turn-by-turn 对话：每 turn 都附上完整历史 + 工具调用结果
• 高度 prefix cache 命中：tool call 模式重复
• 不确定的 token 预算：可能 1 个 turn 也可能 50 个 turn
• 需要 KV cache 持久化：一个 session 可能跑数分钟，断开重连也要恢复

14.7.7 2026 LLM 推理栈对比（更新版）

14.7.8 推理硬件软件全栈 (2026 全景)

用户                ┌────────────────────────────────────┐
请求 ──────────────►│   API 网关 (FastAPI, gRPC, OpenAI 兼容)  │
                    └──────────────┬─────────────────────┘
                                   │
                    ┌──────────────▼─────────────────────┐
                    │   调度器 (continuous batching,      │
                    │   prefix caching, 投机, multi-LoRA) │
                    └──────────┬───┬─────────────────────┘
                               │   │
                    ┌──────────▼─┐ ▼──────────────────────┐
                    │ Prefill    │ │ Decode pool          │
                    │ pool       │ │ (B200 fp4)           │
                    │ (B200 fp8) │ │                      │
                    └─────┬──────┘ └──────┬───────────────┘
                          │  KV 传输      │
                          ▼  (NVLink-C2C  ▼
                          │   或 IB)
                  ┌───────┴──────────────┴────────────┐
                  │  全局 KV cache pool                │
                  │  (LMCache / Mooncake Conductor)    │
                  └────────────────────────────────────┘

14.7.9 学完本教程在 2026 LLM 行业能做什么

恭喜！你已经具备：

• 读懂 vLLM / SGLang / TRT-LLM / FlashAttention / FlashMLA / Marlin 任何一个 CUDA kernel
• 给生产推理引擎贡献自定义 fused 算子（norm + GEMM + activation）
• 实现新 attention 变体（local / sliding / linear / MLA）
• 调优 LLM 服务并用线上 trace 证明 p99 改善
• 面试 GPU/CUDA、LLM Inference Engineer 岗位
• 评估新硬件（B200 / MI325X / TPU v7）适合不适合自己工作负载
• 跟上 2025-2026 论文（FA v3/v4、MLA、QuaRot、QServe、Mooncake、EAGLE-3）

下一步建议：

1. 挑一个真实 OSS 项目（vLLM / SGLang）做一个 PR：加 fused 算子 / 修 bug / 加 model 支持
2. 读 DeepSeek-V3 / FlashAttention v3 / Mooncake 三篇论文 + 对应源码
3. 关注 github.com/NVIDIA/cutlass、HazyResearch/ThunderKittens、triton-lang/triton 的 release 动态
4. 关注 Lecture 系列：GPU MODE 是 2024-2026 最活跃的 GPU 编程社区

14.8 后续学习路径

• 📖 FlashAttention 官方 — 看 v2/v3 的 CUTLASS 实现
• 📖 vLLM — 学 PagedAttention、continuous batching
• 📖 TensorRT-LLM — 看工业级 builder + plugin
• 📖 llama.cpp — 看极致 CPU + GGUF 量化
• 📖 CUTLASS + CuTe — 看 GEMM 工厂
• 📖 Triton — Python-like GPU DSL
• 📖 SGLang — RadixAttention、结构化生成
• 📖 Megatron-LM — 多卡训练 TP/PP/SP 经典实现

14.9 总结

你从 __global__ void hello_kernel() 走到了用自己写的 kernel 跑出语言模型生成。 中间所有让你停下来调试的 bug，对将来读 vLLM 源码、写自定义算子都是直接经验。GPU 编程的"硬功夫"就这样炼出来的。

下一步：选一个真实仓库（推荐 llama.cpp 或 vLLM）读一个 kernel 实现，对照本教程章节复盘。 恭喜！🎉

14.10 CUDA 官方手册精讲（CUDA Programming Guide 13.2（核验：2026-07-20））

Prefill vs Decode、CUDA Graph 模板、Lazy Loading、NVTX、上线 checklist

Prompt 阶段 vs Decode 阶段 — 两套 kernel mix

同一个 forward()，处理 prompt（首次输入 T 个 token） 跟 decode（每步只来 1 个新 token） 的算子组合差别巨大。 生产推理引擎为它俩准备两套 kernel：

关键观察：M=1 的 GEMM 没法把 Tensor Core 喂饱（一次 mma.sync 至少要 M=8）。 所以 decode 阶段权重 GEMM 退化成 权重读 + dot-product reduce， 计算密度极低，带宽就是天花板。这也是为什么 13.7.1 W4A16 量化对 decode 收益最大。

Decode 阶段的 split-K GEMV 模板

// y[N] = W[N, K] * x[K], 单 M=1 GEMV
    // 朴素: 一个 block 一行, 一线程一列 -> 一个 SM 24 行/cycle 太少
    // 正确: split-K + atomic add or two-pass reduction
    template <int BK, int BN>
    __global__ void gemv_splitk(const __half* W, const __half* x, float* y,
                                 int N, int K, int splits) {
        int n = blockIdx.x * BN + threadIdx.y;     // 输出行
        int kstart = blockIdx.y * (K / splits);    // K 维分片
        int kend   = kstart + K / splits;
        if (n >= N) return;

        float acc = 0.f;
        for (int k = kstart + threadIdx.x; k < kend; k += BK) {
            acc += __half2float(W[n * K + k]) * __half2float(x[k]);
        }
        // warp reduce
        for (int off = 16; off > 0; off /= 2)
            acc += __shfl_xor_sync(0xFFFFFFFFu, acc, off);
        if ((threadIdx.x & 31) == 0) atomicAdd(&y[n], acc);
    }

性能数据采集表（Llama-7B 一层 hidden=4096 GEMV；需在目标 GPU 上实测）：

所以 capstone 第二步优化方向：把 decode 路径的 5 个 GEMM 全换成 split-K GEMV。 prompt 路径保留 tile GEMM 即可。

CUDA Graph 完整模板：把 30 次 launch 合成 1 次

GPT-2 small forward 一次 token 大约 12 层 × (LN + QKV + Attn + Proj + LN + FC + GeLU + FC) ≈ 60 个 kernel。每个 launch 都有固定调度开销；当 kernel 变短，launch 占比会明显上升。CUDA Graph capture 让所有 launch 合并成单次 GPU dispatch，开销几乎归零。

三步走

cudaGraph_t       graph;
    cudaGraphExec_t   exec;
    cudaStream_t      stream;
    cudaStreamCreate(&stream);

    // (1) warmup: 先跑一次普通 forward, 让 cuBLAS / cuRAND 完成 lazy init
    decode_one_step(W, ctx, stream);
    cudaStreamSynchronize(stream);

    // (2) capture: 关闭 cuBLAS / cuRAND 内部同步, 用 ThreadLocal 模式
    cudaStreamBeginCapture(stream, cudaStreamCaptureModeThreadLocal);
    decode_one_step(W, ctx, stream);
    cudaStreamEndCapture(stream, &graph);

    cudaGraphInstantiate(&exec, graph, nullptr, nullptr, 0);

    // (3) 之后每步:
    for (int step = 0; step < max_new; ++step) {
        update_runtime_params(ctx, step);     // 仅改 device-resident 计数器
        cudaGraphLaunch(exec, stream);
    }
    cudaStreamSynchronize(stream);

    cudaGraphExecDestroy(exec);
    cudaGraphDestroy(graph);
    cudaStreamDestroy(stream);

常见坑（每个都掉过）

• capture 模式：用 cudaStreamCaptureModeRelaxed 容易把外部 stream 操作 意外卷进 graph；ThreadLocal 最安全。
• cuBLAS handle：必须 capture 前调一次 cublasSetStream + dummy gemm， 否则 handle 内部还会建子 stream，capture 时 hang。
• 动态参数：t_pos / kv_seqlen 这类每步变化的参数要么放 __device__ 全局变量并在 graph 外 cudaMemcpyAsync 更新， 要么用 cudaGraphExecKernelNodeSetParams 在 launch 前 patch。
• printf / device assert 必须删：会破坏 capture。Debug 时改 host log。
• 动态 batch size：每个 batch 大小 capture 一份 graph，运行时按 batch 路由 （vLLM 预 capture 1/2/4/8/16/32 六份）。
• 看 CUDA_GRAPHS_USE_NODE_PRIORITY 环境变量（Programming Guide §5.2.3.5） 控制 graph 节点优先级，多模型协同时有用。

收益采集表（GPT-2 small, 单 token decode；需在目标 GPU 上实测）

服务冷启动：CUDA_MODULE_LOADING=LAZY 与 JIT 缓存

LLM 推理 binary 经常嵌入 50-200 个 cubin（FlashAttn、CUTLASS GEMM、采样、量化各一堆）， 每个 cubin 加载到 GPU 要消耗显存 + 启动时间。CUDA 11.7 引入 lazy module loading，CUDA 12.0 起默认开启， 可观察到的好处：

这张表应在你的服务镜像里采集：记录 CUDA/driver 版本、fatbin 数量、是否挂载 JIT cache、是否 warmup hot kernel。

环境变量（PDF §5.2.4.1）：

export CUDA_MODULE_LOADING=LAZY     # 默认, 显式声明
    # 仅 kernel lazy; 全局变量等 module data 也想 lazy:
    export CUDA_MODULE_DATA_LOADING=LAZY

    # 多线程 binary loader 加速 (CUDA 12.4+)
    export CUDA_BINARY_LOADER_THREAD_COUNT=4

配套优化 — warmup 阶段把所有 hot kernel 各跑 1 次，让 lazy load 提前完成：

void warmup(const Weights& W) {
        // 跑 prefill batch=1, T=16; 触发 prefill 全部 cubin 加载
        forward_prefill(W, dummy_ids_16, 1, /*max_new=*/0);
        // 跑 decode 各 batch size, 让 CUDA Graph 顺便 capture
        for (int B : {1, 2, 4, 8, 16, 32}) {
            forward_decode(W, dummy_kv, B);
        }
        cudaDeviceSynchronize();
    }

JIT 编译缓存（PTX → SASS）

新一代 GPU 上没有对应 cubin 时，PTX 会被驱动 JIT 编译。第一次跑会卡 1-5 s （取决于 kernel 复杂度）。开启磁盘缓存（PDF §5.2.2.1-5.2.2.3）：

export CUDA_CACHE_DISABLE=0                       # 默认开启
    export CUDA_CACHE_PATH=/var/cache/cuda-jit       # 自定义路径, 容器化时挂卷
    export CUDA_CACHE_MAXSIZE=$((4*1024*1024*1024))  # 4 GB 上限

容器化部署关键点：把 ~/.nv/ComputeCache 或自定义路径 挂出来到持久卷，否则每次 pod 重启 JIT 都要重跑一遍，冷启 +30 s。

nvJitLink — 跨编译单元 LTO

CUDA 12 引入 nvJitLinkCompile / nvJitLinkAddFile（CUDA Toolkit 头文件 nvJitLink.h），可以运行时把多个 LTO-IR 文件链接成单 cubin， 做 device-side link-time optimization。LLM 用法：

• fused kernel 的 device 函数库（norm / activation / dequant）以 LTO-IR 发布
• 启动时根据当前模型架构（fp8 / fp4 / GQA / MoE）即时链接出最优 cubin
• 避免编译时穷举所有变种导致 binary 膨胀

TensorRT-LLM 9+ 的 "JIT engine" 模式就是这套。开源端 vLLM 仍是 AOT cubin。

NVTX 端到端 profiling + Tensor Core 吞吐速查表

14.5 推荐用 nsys 看 kernel timeline。当前 mini_llm.cu 已经把 decode path 标成可读的 NVTX range： decode_loop、layer_N、attention/qkv_matmul、 attention/mha_naive_causal、mlp/fc_gelu、logits/lm_head_gemv 和 sampler/greedy_argmax。

编译时默认链接 -lnvToolsExt；如果环境缺少 NVTX 库，用 make NVTX=0 仍可运行，只是 timeline 上没有 range 名称。

cd code/ch14_mini_llm
make ARCH=sm_80

../../scripts/bench.sh nsys-range ./mini_llm \
  --weights=../../data/gpt2-small.bin \
  --tokens=15496,11,612,318 \
  --max_new=8 \
  --profile

# 输出 mini_llm.range.nsys-rep；Nsight Systems GUI 打开看 timeline

各代 Tensor Core 吞吐速查（理论 peak, 单 SM）

PDF §5.1 Table 33（lines 751-766）只列了"哪些 dtype 支持"。这里加上吞吐数字（来自 NVIDIA 白皮书）：

用法：选模型权重精度时先反查硬件支持，再用第 15 章的 roofline 方法验证 workload 到底吃带宽还是吃算力。例如 7B 模型在 4090 上：

• fp16/fp8 的 Tensor Core 峰值需要按具体显卡官方规格确认，不能只按架构推断。
• decode 常先受权重读取带宽限制，量化收益要看 dequant、scale 读取和 batch reuse。
• int4 (W4A16) 往往瓶颈在 HBM 读和 dequant pipeline，FLOPS 不是唯一指标。

"你已经实现的微型版本对应业界哪个特性"

读到这里你已经掌握了上述每个文件的 核心 kernel 写法，剩下的差距是工程量（调度器、KV pool、多 GPU 通信），不是算法新概念。

上线前 12 个 CUDA 检查点

把 capstone 改造成"能撑住生产流量"的服务，下面 12 条来自 NVIDIA Programming Guide §5.2-5.4 的最佳实践 + 工业事故复盘：

1. 每个 CUDA API 调用都包 CUDA_CHECK；kernel launch 后立刻 KERNEL_CHECK （cudaPeekAtLastError 或 cudaGetLastError + cudaDeviceSynchronize 配对）。 Debug 时设 CUDA_LAUNCH_BLOCKING=1（§5.2.3.1）让错误马上暴露在源头。
2. 错误日志：开 CUDA_LOG_FILE=/var/log/cuda.log（§5.2.5.1）， kernel launch 配置错误会写人类可读消息进去（默认只有 generic error code）。
3. 合理设置 CUDA_DEVICE_MAX_CONNECTIONS（§5.2.3.2）。 默认 8，超过 8 个并发 stream 后会出现 false dependency 拖慢。 LLM 服务一般设到 32（最大值）：export CUDA_DEVICE_MAX_CONNECTIONS=32。
4. 不要在 hot path 调 cudaMalloc；用 stream-ordered allocator 或预分配 KV pool（见 13.7.7）。
5. kernel 用 __launch_bounds__(MAX_THREADS, MIN_BLOCKS)（§5.4.3.2）， 避免编译器为单 block 优化导致 occupancy 低。LLM kernel 通常用 __launch_bounds__(256, 4)。
6. warp intrinsic 严格遵守 mask 约束（§5.4.6.6 Warp __sync Intrinsic Constraints）： 所有参与线程的 mask 必须相同；不参与线程的 bit 必须为 0。 违反 → 偶发 hang，常温 1 万 QPS 跑 24 小时才出。
7. fp16/bf16 kernel 加 NaN/Inf 检查 hook。debug 模式下：

   if (!isfinite(__half2float(val))) {
           printf("NaN at layer %d head %d\n", layer, head);
           __trap();   // 立即 abort, GDB 可以接住
       }

   生产模式去掉 hook；偶发 NaN 用 cuda-memcheck 复现。
8. memcpy_async / pipeline（§5.6.2）必须保证 size_and_align ∈ {4, 8, 16} 且 dst/src 地址对齐 — 不对齐会静默退化到 sync copy。
9. atomic 操作选最低 scope：block-only 计数器用 atomicAdd_block 而非 atomicAdd（§5.4.5.1），后者会引入跨 device fence，是否成为瓶颈要用多 GPU trace 验证。 cross-process 才需要 _system 后缀。
10. JIT 缓存挂卷：容器化部署设 CUDA_CACHE_PATH=/cache/jit 并挂出来（§5.2.2.2）， pod 重启不重 JIT。
11. compute_XYa 架构后缀谨慎用（§5.1.2 Architecture-Specific Features）： compute_90a 编出来的 cubin 只 在 sm_90 跑， 升级到 B200/GB200 不能直接复用。生产应为目标架构分别生成 fatbin，并保留合适的 PTX fallback。
12. CUDA Graph 捕获前 warmup cuBLAS / cuRAND，handle 内部 lazy state 不能在 capture 阶段触发（见 14.5.1）。

这 12 条直接对应 NVIDIA 内部 "CUDA Code Review Checklist"。 贡献到 vLLM / TRT-LLM 的 PR 几乎所有 review 意见都落在这上面。

常见坑

• 权重二进制布局与模型配置不一致，会产生能运行但完全错误的输出。
• 只检查最终 token 会掩盖中间层数值漂移；应逐算子或逐层与参考实现对拍。
• 首次运行包含初始化和 lazy loading，不能直接当 steady-state 性能。
• 把本章无 KV cache 的全序列重算速度外推到生产 decode，会得出错误结论。

下一章导览

第 15 章把这个“能跑”的基线放进 Nsight Systems、Nsight Compute 与 roofline 工作流，学习如何用证据定位推理瓶颈。