第 13 章 · LLM 必备零件

本章拼齐第 14 章 Capstone 还差的算子。每个独立、简单，但少一个就跑不起来一个完整 LLM。

学习目标

• 理解并实现 RoPE、SwiGLU、KV cache 与基础采样算子。
• 把算子形状、显存字节和 prefill/decode 阶段联系起来。
• 能用 CPU reference 与 allclose 验证每个 CUDA 实现，再进入端到端组装。

前置知识

需要掌握第 7 章的 reduction/scan、第 9 章的 GEMM，以及第 10–12 章的 softmax 与 attention 数据流。

核心概念

本章关注 LLM 推理中 GEMM 和 attention 之外的必要部件，以及它们对显存容量、带宽和生成质量的影响。

关键代码

13.1 RoPE — Rotary Position Embedding

Llama / Mistral / Qwen 都用 RoPE 替代了 GPT-2 的绝对位置编码。 它把 Q、K 向量看成 D/2 个复数，第 i 对复数乘以 exp(j · θ_{t,i})，其中 θ_{t,i} = t / base^{2i/D}。

实现（in-place）

__global__ void rope_inplace(float* x, int T, int D, float base) {
    int t = blockIdx.x;
    int i = threadIdx.x;
    int half = D / 2;
    float theta = float(t) / powf(base, float(2*i) / float(D));
    float c = cosf(theta), s = sinf(theta);
    float x0 = x[t*D + i];
    float x1 = x[t*D + i + half];
    x[t*D + i       ] = x0 * c - x1 * s;
    x[t*D + i + half] = x0 * s + x1 * c;
}

性质：

• 不需要新增可学习参数
• "相对位置"可以通过 q^T_t · k_s 的内积自动浮现
• 可通过加大 base 或 NTK-aware 缩放外推到更长 context（Llama 2 → Llama 3 = 8K → 128K 就靠这个）

13.2 SwiGLU / SiLU — Llama FFN

GPT-2 用 FFN(x) = GELU(x @ W_1) @ W_2。 Llama 用双投影 + 门控的 SwiGLU：

FFN(x) = ( SiLU(x @ W_gate) ⊙ (x @ W_up) ) @ W_down

SiLU 定义：silu(x) = x · sigmoid(x) = x / (1 + e^{-x})

__device__ float silu(float x) {
    return x * (1.f / (1.f + __expf(-x)));
}

__global__ void swiglu(const float* G, const float* U, float* O, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) O[i] = silu(G[i]) * U[i];
}

"门控"让网络学会"对某些位置让信息通过、对另一些抑制"。代价：FFN 多一个 W_gate，参数量 ~1.5×；但 hidden = 2.66·D（GPT-2 是 4·D），总参数差不多。

13.3 KV Cache

自回归生成有个关键观察：第 t 步生成的 token i 时，第 t-1 步算出来的 K、V 仍然有效，不必重算。 于是推理代码维护一个 cache：

K_cache shape: (n_layers, n_heads, T_max, D_head)
V_cache shape: same

每步只算新 token 的 (K_new, V_new) 形状 (1, D_head)，append 到 t_pos 位置：

__global__ void append_kv(const float* K_new, const float* V_new,
                          float* K_cache, float* V_cache, int t_pos, int D) {
    int d = blockIdx.x * blockDim.x + threadIdx.x;
    if (d < D) {
        K_cache[t_pos * D + d] = K_new[d];
        V_cache[t_pos * D + d] = V_new[d];
    }
}

KV cache 的显存开销

Llama-7B (n_layers=32, n_heads=32, D_head=128, fp16) 一个 token 的 KV 占 32 × 32 × 128 × 2 × 2 B = 0.5 MB。 T=2048 时 1 GB；batch=8、T=2048 → 8 GB。这就是为什么 KV cache 管理（PagedAttention 等）这么重要。

13.4 Sampling 策略

greedy (argmax)

// 简化: V <= 65536 用单 block reduce 即可
template <int BLOCK>
__global__ void greedy_argmax(const float* logits, int* out, int V) {
    __shared__ float vals[BLOCK]; __shared__ int idxs[BLOCK];
    float bv = -INF; int bi = -1;
    for (int i = tid; i < V; i += BLOCK)
        if (logits[i] > bv) { bv = logits[i]; bi = i; }
    /* block-wide reduce on (vals, idxs) keeping max */
    if (tid == 0) *out = idxs[0];
}

top-k (truncated sampling)

从 logits 取最大的 k 个，softmax，按概率抽样。常见 k = 40~50。GPU 实现：用 block-wide top-k（heap 或 partial sort）。

top-p (nucleus sampling)

排序 logits → 累计概率 → 截断到累积 ≥ p 的前缀（典型 p = 0.9~0.95）→ 重新归一 → 抽样。需要 sort + prefix scan，第 7 章的 scan 模板在这里能复用。

temperature

logits /= T，T > 1 让分布更平（更随机），T < 1 让分布更尖（更确定）。在 softmax 前应用。

运行结果

在目标 NVIDIA GPU 上分别运行本章示例；每个程序应打印 CPU/GPU 对拍结果，性能数字保留为本机实测值，不跨 GPU 照搬。

自检清单

练习题

1. 实现 top-k sampling：先 block-wide top-k（用堆），再 softmax 抽样。
2. 实现 top-p (nucleus)：sort + scan + 截断。
3. 给 rope.cu 加 theta_base = 500000（Llama 3 用），看长 T 时角度变化。
4. 把 swiglu 改成 fused：把 silu(G) * U 合并到 W_down 的 GEMM 里（pre-multiplier trick）。

13.7 工业实战：量化、投机解码、采样工程

13.7.1 W4A16 量化 — LLM 推理头号加速武器

LLM decode 阶段瓶颈是权重的 HBM 读带宽（M=1 时 FLOPs 极少）。 最低延迟可以先用 weight_bytes / measured_HBM_bandwidth 估算。 把权重从 fp16 压成 INT4 会减少编码后的权重字节；端到端结果还要计入 dequant、scale 读取、cache reuse 和 batch 复用。这就是 W4A16（W=int4，A=fp16）。

W4A16 kernel 的关键：dequant + GEMM 融合

朴素做法：先把 INT4 dequant 成 fp16 写回 HBM 再调常规 GEMM —— 这样消除不了带宽瓶颈。
正确做法：kernel 内即时 dequant，权重始终以 INT4 从 HBM 读：

for (int kt = 0; kt < K; kt += BK) {
    // 1) HBM 只读 INT4 weight tile (BM*BK/2 字节, 2 个 INT4 / byte)
    load_int4_to_shared(W_int4_tile, gmem_ptr);
    // 2) shared 内 dequant: int4 -> fp16, 乘 per-group scale
    dequant_inplace(W_int4_tile, W_fp16_tile, scales[group_id]);
    // 3) WMMA / mma.sync 用 fp16
    wmma::mma_sync(acc, A_frag, B_frag, acc);
}

典型实现：TensorRT-LLM weight_only_gemm，vLLM 的 Marlin kernel。实际 tokens/s 与 GPU、batch、上下文长度、group size 和调度器强相关；本教程不写死未经本仓 benchmark 验证的吞吐数字。

13.7.2 投机解码 (Speculative Decoding)

用小模型快速生成 N 个候选 token，大模型一次性验证。接受多少多少出，拒绝就退回。

graph LR
    Start["prompt"] --> Draft["draft model
(TinyLlama 1B 等)
快速 4-8 token"]
    Draft --> Verify["target model
(Llama 70B)
一次前向算 5 个位置 logits"]
    Verify --> Accept{"p_target / p_draft
是否接受?"}
    Accept -->|"接受"| Continue["输出, 继续 draft"]
    Accept -->|"拒绝"| Resample["从拒绝位重采样, 退回 draft"]
    Continue --> Draft
    Resample --> Draft
    style Draft fill:#f3f1e8,stroke:#8b1538
    style Verify fill:#f3f1e8,stroke:#2f5d3a

关键洞察：

• target model 一次 forward 同时算 N 个位置 logits（沿 batch 维并行），verify 的额外成本要按 batch、KV 读和 draft 长度测量
• 端到端收益由接受率、draft 成本、verify batch 形状和调度器共同决定，不能只看论文平均值

draft model 候选：

• TinyLlama 1B（同架构小模型）
• EAGLE（专门训的 draft head）
• Medusa（多 head 并行预测）
• n-gram cache（重复 token 直接命中）

vLLM、TensorRT-LLM 都内置。EAGLE / Medusa 需要额外训练，n-gram 0 训练成本。

13.7.3 生产 sampler — 不只是 argmax

// 完整 pipeline: logits -> /temp -> top-k -> top-p -> softmax -> cuRAND 抽样
__global__ void sample_kernel(float* logits, int V, float temp, int top_k, float top_p,
                              unsigned int seed, int* out_token) {
    // 1) 除 temperature
    for (int i = tid; i < V; i += BLOCK) logits[i] /= temp;
    // 2) top-k: 保留最大 k 个, 其他 -inf  (block-wide heap 或 partial sort)
    block_topk(logits, V, top_k);
    // 3) softmax + inclusive scan + 找 top-p 截断点 + 重新归一
    block_softmax(logits, V);
    block_inclusive_scan(logits, V);
    int cutoff = block_find_first_ge(logits, top_p);
    // 4) 用 cuRAND 抽样
    curandState s; curand_init(seed, tid, 0, &s);
    *out_token = sample_from_dist(logits, cutoff, curand_uniform(&s));
}

性能：请用第 15 章的 nsys/ncu 流程在目标 GPU 上采集。记录项包括 vocab size、top-k/top-p、batch、sampler kernel latency 与 decode 总耗时占比。

13.7.4 KV cache 量化（int8 / fp8）

权重量化外，KV cache 也可以量化。编码字节减少后可能容纳更大 batch 或更长 context，但需计入 scale、metadata 与 allocator 对齐。

实现：per-group scale，attention kernel 内 load 时 dequant：

// attn 读 K 时:
int8_t k_i8 = K_cache_int8[t * D + d];
half scale  = K_scales[t / GROUP_SIZE];          // 每 64 token 一个 scale
half k      = __int2half_rn(k_i8) * scale;

TensorRT-LLM 与 vLLM V1 都提供量化 KV cache 路径，但支持矩阵、scale 粒度和质量约束必须按当前文档核验。 陷阱：per-tensor、per-channel、per-group 各有精度、元数据和调度成本，不能把固定 group_size 当作所有模型的生产默认值。

13.7.5 RoPE 工程注意点

• fuse 到 QKV proj：QKV GEMM 出来立即 RoPE，不要单起 kernel
• 预计算 cos/sin 表：每步 forward 重算 sin(t/base^(2i/D)) 浪费；服务启动时预计算 (T_max, D/2) 表存显存
• NTK-aware / YaRN scaling：base 从 10000 改 500000+，让模型外推到更长 context
• 正确性测试：用 PyTorch rotary_emb 对拍，pair (偶位置, 奇位置) 必须严格对齐

13.7.6 推理优化优先级

顺序建议：先用 vLLM/TRT-LLM 这类成熟引擎建立可用 baseline，再针对自己模型做 fp8 量化 / spec decoding 增量调优，并用同一套 trace 比较。

13.8 研究前沿（2025-2026）：EAGLE-3、Lookahead、KV 压缩、W4A4 / FP4

13.8.1 投机解码 2024-2026 演进

2026 工业现实：

• vLLM / TRT-LLM 默认开 EAGLE-2 / Medusa-2
• 不愿意训 draft 的接 Lookahead（n-gram 命中率随 prompt 类型变化大）
• Reasoning 模型的 thinking 是否有高重复率，要用真实 prompt trace 统计 n-gram 命中，而不是假设固定收益

13.8.2 KV cache 压缩前沿（2024-2026）

除了 13.7.4 的量化路线，2024-2026 出现了大量"丢 token"型压缩：

组合：MLA + fp8 KV + SnapKV / Quest 是 DeepSeek-V3 / Kimi 等方向的工业实践。长 context 的 KV 显存节省要按模型层数、head_dim、保留 token 数和精度策略逐项测算。

13.8.3 W4A4 / NVFP4 — 极致量化（2024-2026）

9.10.4 给了对比表，这里详述算法：

QuaRot（ICML 2024）— Hadamard 旋转解 outlier

激活的 outlier 是 W4A4 最大障碍（少数 channel 数值范围比平均大 100×）。QuaRot 用 Hadamard 矩阵 H 旋转：x' = xH, w' = wH^T，数学等价 + 旋转后分布更均匀，outlier 被打散。

SpinQuant（Meta 2024）— 可学习旋转

不用固定 Hadamard，训练学习旋转矩阵。精度比 QuaRot 再好 0.2-0.5 PPL。

Atom（MLSys 2024）— 混合 W4A4 + Heuristic outlier

正常 channel 用 INT4，outlier channel 单独留 INT8。fp16 baseline 之上无损。

NVFP4（Blackwell 2025）— 硬件原生

fp4 (E2M1) + per-16-block fp8 scale，Tensor Core 直接吞。无需 dequant 到 fp16/fp8 中转，吞吐拉满。是 2025+ B200 推理的事实标准。

13.8.4 BitDelta / 1-bit weight diff

BitDelta（Anthropic 2024）：把 fine-tuned 模型相对 base 的 diff 量化到 1 bit。

• base model (full fp16) 一份
• 每个 fine-tune 的 delta 只占 1 bit/param + scale
• Llama-70B fine-tune 模型从 140 GB 缩到 ~9 GB + 1 GB
• 精度损失极小（< 0.5 LR-bench）

意义：多租户服务能存几百个 fine-tune 版本同时运行，每个只多花一点显存。Anthropic / Together AI / Fireworks 都在用。

13.8.5 RoPE 长 context 外推：YaRN、PI、NTK

原始 RoPE base=10000 训练长度内表现好，超出训练长度急剧退化。2023-2025 主要技术：

Llama 3 直接训练到 8K，推理用 YaRN 外推到 128K。Llama 4 据传训练到 32K，外推到 10M。不重训能让模型多用 10×+ context，性价比极高。

13.8.6 Reasoning 模型的采样新需求

o1 / R1 类 reasoning 模型，采样策略影响最终能力：

• 多次采样 + 投票 (best-of-N, BoN)：同一 prompt 跑 N 次, 选 reward 最高的。N=64 是常见 ceiling
• Self-Consistency：跑多次, 答案投票
• Process Reward Model (PRM)：用专门的 reward model 给中间 thinking 打分, 引导更好的路径
• Beam Search 复兴：reasoning 模型的 thinking 路径有"枝", beam 重新有意义
• MCTS（Monte Carlo Tree Search）：DeepSeek-R1 等用类似思路在 thinking 期间探索

对 sampler kernel 的影响：不再是单 next-token argmax，而是 N 个并行候选 + reward model 评估 + 选择，采样跟 attention 计算量同阶。

13.8.7 2026 LLM 算子优化总览

不要把单项收益直接相乘：很多技术优化的是同一段瓶颈，叠加后会互相稀释。工业评估应以固定模型、固定 SLO、固定硬件的 tokens/$、TTFT、TPOT 和 p99 为准。

常见坑

• RoPE pair 索引写错（i + half 和 i 颠倒）→ Q/K 一致地坏，模型仍能跑但精度差
• KV cache 越界写（t_pos > T_max）→ illegal memory access
• greedy argmax 在 V 不是 BLOCK 倍数时漏元素 → 偶发概率错
• temperature = 0 → 除零，应改为 greedy 分支
• W4A16 用错 scale layout（per-tensor vs per-group）→ 数值崩，模型胡言乱语
• 投机解码 draft 模型没用同 tokenizer → 投机 token id 跟 target 对不上
• top-p 排序对 V=128K 词表过慢 → 用 radix-select 或先 top-k 截断到 1000 再 top-p
• fp8 KV cache 用 per-tensor scale → 长 T 时精度损失明显，必须 per-group
• EAGLE / Medusa 跟原模型不同 tokenizer 或 vocab → 投机一直被拒, 反而拖慢
• Lookahead n-gram cache 不清理 → 跨请求污染，输出错乱
• YaRN 配置忘了应用到 model 加载时, 只在 inference 阶段调 base → RoPE 实际没变
• BitDelta 用错 base model 版本（如 Llama-2 vs Llama-2-chat） → 解出来全是噪声

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

__sincosf / __half2 / KV layout / cuRAND / cudaMallocAsync

RoPE 工程提速：__sincosf 取代 sinf/cosf

原版 cosf / sinf 走的是 CUDA 数学库的 IEEE 精度路径，单次调用展开为 ~20+ SASS 指令；而 device intrinsic __sincosf(x, &s, &c) 一次调用同时算出 sin/cos，只用 SFU（Special Function Unit）路径。 对 RoPE 这种每元素一次三角函数的算子，通常值得替换后用 SASS 和 benchmark 复核。

CUDA Programming Guide §5.5.9.2 给出的 ULP 误差：

改写后的 RoPE：

__global__ void rope_inplace_fast(float* x, int T, int D, float base) {
        int t = blockIdx.x;
        int i = threadIdx.x;
        int half = D / 2;
        // __powf uses the fast-math path; verify accuracy before replacing powf.
        float theta = float(t) * __powf(base, -float(2*i) / float(D));
        float s, c;
        __sincosf(theta, &s, &c);          // 一次 SFU 调用
        float x0 = x[t*D + i];
        float x1 = x[t*D + i + half];
        x[t*D + i       ] = x0 * c - x1 * s;
        x[t*D + i + half] = x0 * s + x1 * c;
    }

注意点：

• 编译时也可以加 --use_fast_math，nvcc 会自动把 sinf/cosf/expf/powf/tanhf/log*f 等替换成 intrinsic（见 Programming Guide §5.5.9.3 Table 59）。但全工程开启可能影响 LayerNorm 等数值敏感算子， 生产 LLM kernel 普遍只在 RoPE / SwiGLU / softmax 等场景手写 intrinsic。
• RoPE 的 base 是常数（如 10000 或 500000），__powf(base, -2i/D) 可以预计算到 cos/sin LUT（典型 shape (T_max, D/2)）一次性放显存， decode 每步直接 lookup。这是 vLLM / TRT-LLM 的实际做法。

SwiGLU 用 __half2 打包加速 (2 元素 / 指令)

Llama / Qwen 等模型 FFN 走 fp16 / bf16 路径。Hopper 之前的硬件没有原生 fp16-vector ALU，但 __half2（CUDA Math API 提供）把两个 fp16 打包到 32-bit 寄存器， 一条 PTX 指令一次算两个元素：HBM 读带宽利用率 +100%、寄存器压力 -50%。

#include <cuda_fp16.h>

    __device__ __forceinline__ __half2 silu2(__half2 v) {
        // SiLU(x) = x / (1 + exp(-x))
        // hexp2 + h2div + h2add  (HW 加速; sm_53+)
        __half2 one  = __float2half2_rn(1.0f);
        __half2 neg  = __hneg2(v);
        __half2 e    = h2exp(neg);                  // 包含 __hexp on each lane
        __half2 den  = __hadd2(one, e);
        return __h2div(v, den);
    }

    __global__ void swiglu_fp16(const __half2* G, const __half2* U,
                                __half2* O, int n_pairs) {
        int i = blockIdx.x * blockDim.x + threadIdx.x;
        if (i >= n_pairs) return;
        __half2 g = G[i], u = U[i];
        O[i] = __hmul2(silu2(g), u);                // 一条 HMMA-style mul, 2 元素
    }

调用端 launch 时 n_pairs = N / 2，原始 N 个 fp16 元素以 __half2* 重新解释即可。

更激进的写法（compute capability 8.x+）：

• 用 __nv_bfloat162（cuda_bf16.h）做 bf16 路径——Llama 训练用 bf16 保留 fp32 动态范围
• 用 __hexp2 / __hcos2 等 native fp16 intrinsic（精度低于 fp32 intrinsic 但单指令）
• Hopper（sm_90）+ 提供 FP8 (E4M3/E5M2) 类型 __nv_fp8x2_e4m3；吞吐取决于量化、shape 与 Tensor Core 路径

陷阱：

• fp16 表达范围只到 65504；SiLU 在 |x| > 16 时 exp 溢出。Llama 训练通常先 LayerNorm 把激活拉回 [-5, 5]，但 fine-tune 模型如果 LN 学坏会让 SwiGLU 溢出，输出 NaN。 Production kernel 会加 __hmax2(v, __float2half2_rn(-80.f)) 这种 clamp。
• __half2 要求张量基地址按 4-byte（2 个 fp16）对齐——Llama hidden_size 通常是 256 / 512，天然对齐。

KV Cache 三种内存布局 (BSHD / BHSD / Paged) 对比

同样的 4-D 张量 (Batch, Heads, Seq, Dim)，存储 stride 不同会导致 attention kernel 的访存模式完全不一样。三个主流 layout：

(1) 朴素 BHSD append

// K_cache shape (B, H, T_max, D), row-major BHSD
    // append 第 t_pos 步的 K_new shape (B, H, D)
    __global__ void append_kv_bhsd(const __half* K_new, __half* K_cache,
                                    int B, int H, int T_max, int D, int t_pos) {
        int b = blockIdx.z;
        int h = blockIdx.y;
        int d = blockIdx.x * blockDim.x + threadIdx.x;
        if (d >= D) return;
        int dst = ((b * H + h) * T_max + t_pos) * D + d;
        int src = (b * H + h) * D + d;
        K_cache[dst] = K_new[src];
    }

K_cache 显存预分配 B*H*T_max*D*sizeof(half)。如果 batch 中真实序列长度差异 大（比如 chat 场景的 prompt 长度跨度很大），按 T_max 为每个请求预分配会造成随 trace 变化的显存浪费。 这就是 vLLM 提出 PagedAttention 的动机。

(2) Paged 布局 — 仿 OS 虚拟内存

// K_pool: (n_pages, page_size, H, D), 全局物理页池
    // block_table: (n_req, max_pages_per_req)  逻辑页 -> 物理页
    __device__ __half* k_at(int req_id, int t_pos,
                             const __half* K_pool, const int* block_table,
                             int page_size, int H, int D) {
        int logical_page = t_pos / page_size;
        int offset_in_page = t_pos % page_size;
        int phys_page = block_table[req_id * MAX_PAGES + logical_page];
        return (__half*)K_pool + (((size_t)phys_page * page_size + offset_in_page) * H * D);
    }

典型 page_size = 16 token（vLLM 默认）。优势：

• 每请求只 alloc 需要的 page 数，显著降低按 T_max 预分配造成的内部碎片
• 多请求可以共享 prefix page（system prompt 等），即 prefix caching
• 请求结束页可立即回收，supporting continuous batching

代价：attention kernel 每访问一个 (b, t, h, d) 元素要先查 block_table，多一层间接寻址； FlashAttention v2.4+ 的 paged kernel 把 block_table 缓存到 shared memory 摊平开销。

采样工程：cuRAND Philox + warp-level reduce

生产 sampler 必须满足三个性质：(1) 可复现（同 seed 同输出），(2) 一次启动多请求并行采样不串扰， (3) latency 需要足够低，不能盖过 sampler 本身。cuRAND 的 curandStatePhilox4_32_10_t 是常用选择： 无状态、counter-based、单线程一次产 4 个 uniform 数，特别适合 LLM 大 batch 解码。

#include <curand_kernel.h>

    // 一个请求一个 state, 跨步骤增 counter
    __device__ float draw_uniform(curandStatePhilox4_32_10_t* st) {
        return curand_uniform(st);   // 单次调用 1 个 [0, 1) float
    }

    __global__ void sample_topp_kernel(const float* sorted_probs,   // 已 softmax 排过序
                                        const int*   sorted_idx,
                                        int V, float top_p,
                                        uint64_t seed, uint64_t step,
                                        int req_id, int* out_token) {
        extern __shared__ float scan[];

        // 1) 初始化此请求此步骤的 Philox state
        curandStatePhilox4_32_10_t state;
        curand_init(/*seed=*/seed, /*subseq=*/req_id, /*offset=*/step, &state);

        // 2) block-wide inclusive scan 找 top-p 截断点 (Ch7 模板)
        //    cum_p[V] 累计概率
        block_inclusive_scan(sorted_probs, scan, V);

        // 3) lane 0 取一个 [0, 1) 样本, broadcast
        float r = 0.f;
        if (threadIdx.x == 0) r = curand_uniform(&state);
        r = __shfl_sync(0xFFFFFFFFu, r, 0);

        // 4) 二分找 r 落在 scan[] 哪个 bucket; 注意截断到 cum_p <= top_p
        //    简化版: 每线程 stride 扫一遍, atomicMin 记录命中 index
        __shared__ int hit;
        if (threadIdx.x == 0) hit = V - 1;
        __syncthreads();
        float r_scaled = r * min(scan[V - 1], top_p);
        for (int i = threadIdx.x; i < V; i += blockDim.x) {
            if (scan[i] >= r_scaled) atomicMin(&hit, i);
        }
        __syncthreads();
        if (threadIdx.x == 0) *out_token = sorted_idx[hit];
    }

三个关键工程点：

• Philox state 不要每步重 init：用 (seed, req_id, step) 三元组， 同 seed 同 req_id 重跑 严格相同；不同 req_id 子序列独立。
• scan 用 warp shuffle 加速：Programming Guide §5.4.6.5 warp shuffle (__shfl_up_sync, __shfl_xor_sync) 用于 prefix scan 可避免 shared memory bank conflict；收益用 sampler benchmark 记录。
• compute_80+ 用 warp reduce：找 argmax 时 __reduce_max_sync(0xFFFFFFFFu, val)（§5.4.6.4）单指令完成 warp 内 max， 原来要 5 步 shuffle。

// Hopper / Ampere 上的 fast argmax
    __device__ int warp_argmax(float val, int idx) {
        // 把 (val, idx) 编码到 uint64 高低位
        uint64_t pack = (uint64_t(__float_as_uint(val)) << 32) | unsigned(idx);
        // 注意: __reduce_max_sync 仅支持 unsigned/int, 这里走手写 shuffle
        for (int off = 16; off > 0; off /= 2) {
            uint64_t other = __shfl_xor_sync(0xFFFFFFFFu, pack, off);
            if (other > pack) pack = other;
        }
        return int(pack & 0xFFFFFFFFu);
    }

性能数据采集表（V=50257, 单请求；需在目标 GPU 上实测）：

cudaMallocAsync — LLM 推理的零开销 allocator

朴素 LLM 服务在请求到来时 cudaMalloc 一块 KV cache，请求结束 cudaFree。 这俩调用是同步、跨进程串行化的，在高 QPS 服务里会放大尾延迟；应该离线固定或启动时设置。

CUDA 11.2+ 引入 stream-ordered allocator 解决这个问题：

// 主线程一次性创建 pool, 调好上限
    cudaMemPool_t pool;
    cudaDeviceGetDefaultMemPool(&pool, 0);
    size_t threshold = size_t(40) << 30;  // 40 GB
    cudaMemPoolSetAttribute(pool, cudaMemPoolAttrReleaseThreshold, &threshold);

    // 每请求所在 stream 上 allocate, 不会阻塞
    void* kv;
    cudaMallocAsync(&kv, kv_bytes, stream);
    // ... decode_step_kernel<<<..., stream>>>(kv, ...);
    cudaFreeAsync(kv, stream);   // 还是非阻塞, GPU 跑完才回收

关键好处：

• 0 同步：alloc/free 走 stream，跟 kernel 同队列，不戳 CPU↔GPU
• 页缓存复用：pool 内部按 size class 缓存；典型 LLM 解码 KV 都是同一尺寸，alloc 几乎 0 开销
• 跨 stream 共享：用 cudaMemPoolSetAccess 给其他 GPU/stream 授权即可远程释放
• 配合 cudaMemPoolAttrReleaseThreshold 控制上限，pool 一旦超阈值才真还给驱动

vLLM / SGLang 这类推理框架会避免在 hot path 使用同步分配。实际 p99 改善幅度必须用目标 workload 压测确认，记录 QPS、batch 分布、KV 分配大小和 allocator 配置。

Per-request 显存与 cudaStreamAttachMemAsync

多租户场景下，多个请求共享 GPU。要保证某请求结束后 它的 KV 立即被回收给同 stream， 而不是等 pool flush，可用：

// 把 KV 标记成仅 stream 内可见
    cudaStreamAttachMemAsync(req_stream, kv, 0, cudaMemAttachSingle);
    // ... decode ...
    cudaFreeAsync(kv, req_stream);  // stream 完成立即回收

NVIDIA 官方 deep learning examples 也用同一招做 batch 异构调度，可参考 §5.2.3 的 CUDA_DEVICE_MAX_CONNECTIONS 环境变量来调整并发 stream 数。

下一章导览

下一章会把这些零件与前面手写的 CUDA kernel 串成 GPT-2 small 教学推理程序，并明确它与生产引擎之间的差距。