手写 Triton PagedAttention decode kernel，与 PyTorch baseline 做可复现的 正确性与性能对比。目标是把 vLLM 的核心机制（分页 KV cache + online softmax） 实现到能讲透每一行的程度。
项目定位：学习与原型验证项目。只做 decode attention 算子层， 没有 scheduler、没有 serving、没有量化。不是生产级推理引擎。
| 项目 | 作者 | 关系 |
|---|---|---|
| vLLM | UC Berkeley 等 | PagedAttention 的原始论文与工业实现，本项目的思想来源 |
| nano-vllm | 俞星凯（DeepSeek） | 名字相近但无代码关系。它是 ~1200 行 Python/PyTorch 的完整引擎复刻（含 scheduler/engine） |
| 本项目 | — | 独立实现，聚焦点不同：只做 decode attention 的手写 Triton kernel，配三层对比 baseline + profiler |
本项目没有 fork 任何仓库，代码从零编写。与 nano-vllm 的区别在于：
它复刻的是引擎架构，本项目做的是单算子的 kernel 级优化与量化对比。
decode 阶段每步只生成 1 个 token，却要对全部历史 KV 做 attention。 朴素做法有三笔随 context 放大的开销：
| 开销 | 朴素做法 | 本项目做法 |
|---|---|---|
| 显存 | materialize [seq_len] 分数矩阵，O(N) |
online softmax 递推，中间状态只有 BLOCK_N，O(1) |
| 访存 | 每步重新 gather 并拷贝 KV | block table 间接寻址，直接读物理显存，零拷贝 |
| Kernel launch | gather / matmul / softmax / matmul 多次 launch | 全部融合进一个 Triton kernel，一次 launch |
online softmax 与标准 softmax 数学上完全等价（不是近似）， 但把「必须看到全部分数才能归一化」变成了流式递推。
| 实现 | 做法 | 定位 |
|---|---|---|
paged_attention_ref |
Python 三重循环 | 正确性金标准 |
paged_attention_vectorized |
高级索引 gather + einsum |
公平 baseline：已向量化，但仍 materialize 分数矩阵 |
paged_attention_decode_triton |
手写 Triton kernel，online softmax | 本项目产出 |
关键对比是 vectorized vs triton。只和 Python 循环比是没有说服力的 ——
两边都向量化之后的差距，才真实反映 kernel 优化（online softmax + 单次 launch
- 分块 gather）的收益。
pip install torch
python tests/test_kernel_logic_cpu.py用 PyTorch 在 CPU 上逐行复刻 Triton kernel 的指针算术和 online softmax 递推，与标准 softmax 对拍。7 个 case 覆盖 GQA / MHA / MQA 和 4 种 context 边界。
这一步把「算法 / 索引错误」和「Triton 语法错误」解耦，排查范围小一个数量级。
环境：CUDA GPU、Python 3.10+、PyTorch ≥ 2.1（Linux 版自带 Triton）
# 1) 正确性验收
python tests/test_paged_attention.py
# 2) 单点三层对比
python bench/bench_decode.py
# 3) 扫描 context_len，出曲线数据
python bench/bench_decode.py --scan --csv results.csv
# 4) profiler trace（chrome://tracing 或 ui.perfetto.dev 打开）
python bench/bench_decode.py --profile --ctx 2048 --trace-out trace.json
# 5) 出图 + 打印 markdown 结果表
python bench/plot_results.py --csv results.csv --out-dir figs --gpu "RTX 3090"完整的一步步操作流程见 RUNBOOK.md。
nanoVLLM/
├── nanoVLLM/ops/
│ ├── attention.py # Triton PagedAttention decode kernel（核心）
│ └── paged_attention_ref.py # PyTorch 参考实现 + 向量化 baseline
├── tests/
│ ├── test_kernel_logic_cpu.py # CPU 逻辑验证（不需要 GPU）
│ └── test_paged_attention.py # GPU 正确性验收（GQA/MHA/MQA + 边界长度）
├── bench/
│ ├── bench_decode.py # 时间 / 显存 / profiler
│ └── plot_results.py # CSV -> 曲线图 + markdown 表
└── RUNBOOK.md # 一步步操作手册
query: [num_seqs, num_heads, head_dim] fp16
key_cache: [num_blocks, block_size, num_kv_heads, head_dim] fp16
value_cache: 同上
block_tables: [num_seqs, max_blocks_per_seq] int32
context_lens: [num_seqs] int32
key_cache 的第 0 维是物理 block 编号，与序列无关 —— 它是全局显存池。
序列第 b 个逻辑 block 存在哪，由 block_tables[seq, b] 决定。
这和操作系统的虚拟内存页表同构，这就是 "paged" 的全部含义。
grid = (num_heads, num_seqs)，每个 program 负责一个 (序列, head) 组合，
串行遍历该序列的所有 block。
已知局限：program 总数 = num_heads × num_seqs。小 batch 时（如 8×8=64）
跑不满 RTX 3090 的 82 个 SM，occupancy 不足。生产级实现（vLLM 的 v2 kernel）
会在 context 维度再切 partition 做 split-K：先并行算局部 softmax，再二次归约。
本项目 v1 不做，但 benchmark 里提供了 --num-seqs 64 对照来量化这个影响。
- GQA 映射 —
kv_head = head // (num_heads // num_kv_heads)。num_kv_heads == 1（MQA）时公式自动退化成立，无需分支。 - 尾部 mask —
context_len通常不是block_size整数倍，最后一个 block 尾部是其他序列释放后的残留数据。除了把分数压成-inf，tl.load上也挂mask=valid, other=0.0，防止残留的inf/nan污染整行乘加。 -inf相减产生nan—m_i初值-inf，第一轮alpha = exp(-inf - m_new) = 0是正确的。但若某 block 整块被 mask，m_new仍为-inf，exp(-inf - (-inf)) = nan会污染全部结果。ctx_len >= 1时不会发生；ctx_len == 0时循环不执行，归一化前做了保护除。HEAD_DIM必须是 2 的幂 —tl.arange的编译期约束。- 精度 — fp16 存储、fp32 累加。fp16 只有 10 位尾数，4096 token 累加
误差按 √N 放大，且
exp极易溢出。m_i / l_i / acc全用 fp32，只在写回时转 fp16。
| 项目 | 值 |
|---|---|
| GPU | |
| Driver / CUDA / PyTorch / Triton | |
| dtype | fp16 存储，fp32 累加 |
| num_seqs (batch) | 8（另有 64 的对照组） |
| num_heads / num_kv_heads | 8 / 4（GQA，与 Qwen3 同构） |
| head_dim | 128 |
| block_size | 16 |
| context_len | 128 / 256 / 512 / 1024 / 2048 / 4096 |
| 预热 | 10 次 |
| 计时 | 重复 50 次取平均，前后 torch.cuda.synchronize() |
| 显存 | torch.cuda.max_memory_allocated() 峰值增量 |
| 数值容差 | max abs diff < 2e-2（fp16 合理范围） |
| context_len | vectorized (ms) | triton (ms) | speedup | vectorized 显存 (MB) | triton 显存 (MB) |
|---|---|---|---|---|---|
| 128 | |||||
| 256 | |||||
| 512 | |||||
| 1024 | |||||
| 2048 | |||||
| 4096 |
流式 ASR 和 LLM decode 是同一个问题：变长序列 + 增量推理 + 历史状态复用。
| LLM 推理 | 流式 ASR |
|---|---|
| prefill 处理 prompt | encoder 处理一个音频 chunk |
| decode 逐 token 生成 | decoder 逐 token 输出文本 |
| PagedAttention block 管理 | chunk 级 KV cache 池 |
| prefix caching | 热词 / 上下文历史复用 |
| continuous batching | 多路音频流实时调度 |
context_lens 变长 mask |
各路音频流长度不一 |
所以这套「block 管理 + online softmax decode kernel」可以直接迁移到流式 ASR 的 decoder 侧。差异点在于 ASR 对尾延迟（tail latency）更敏感， 所以适合固定 chunk 的确定性调度，而不适合引入时间抖动的投机解码。
- Triton PagedAttention decode kernel
- CPU 逻辑验证（不依赖 GPU）
- 三层对比 baseline + 正确性验收（GQA/MHA/MQA + 边界长度）
- benchmark / profiler / 出图脚本
- 在 RTX 3090 上跑出结果并填入上表
- split-K partition 版 kernel，解决小 batch occupancy 不足
- KV cache block manager + prefill/decode 引擎循环
- Kwon et al., Efficient Memory Management for Large Language Model Serving with PagedAttention, SOSP 2023
- Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, NeurIPS 2022
- Milakov & Gimelshein, Online normalizer calculation for softmax, 2018
MIT