L2 讲解：paged decode attention 与 GQA

① 先看全貌：decode 就是读 KV

decode 的算术强度约为 1 FLOP/字节，远低于 GPU 的拐点。所以 decode attention 的时间 ≈ KV cache 的字节数 ÷ 带宽，所有优化都围绕“少读、读快”。

题目页的测试：B = 64、ctx = 1024、Hkv = 8、D = 128、fp16。KV 字节数 = 2 × 64 × 1024 × 8 × 128 × 2 = 268 MB，实测 176 us，即 1520 GB/s（峰值的 85%）。

② GQA：几个 query head 共享一个 KV head

LLaMA-3、Qwen 等模型用 GQA：Hq 个 query head 分成 Hkv 组，每组共享一个 KV head。题目的配置是 Hq = 32、Hkv = 8，GROUP = 4。

③ 循环外：program、q 的打包

drills30–45 行 · host 端 drills85–96 行

b   = tl.program_id(0).to(tl.int64)      # 第几个序列
kvh = tl.program_id(1).to(tl.int64)      # 第几个 KV head
ctx = tl.load(context_lens + b)          # 这个序列有多少个 KV token

offs_g  = tl.arange(0, GROUP_PAD)        # 组内的 query head（补齐到 16）
g_mask  = offs_g < GROUP
q_heads = kvh * GROUP + offs_g
q = tl.load(Q + b * stride_qb + q_heads[:, None] * stride_qh + offs_d[None, :],
            mask=g_mask[:, None], other=0.0)          # [GROUP_PAD, D]

• grid 是 (B, Hkv)，一个 program 负责一个序列的一个 KV head。
• 为什么补齐到 16 行：tl.dot 用 Tensor Core，要求矩阵的行数至少 16。GROUP = 4 时只有 4 行，所以 GROUP_PAD = max(16, next_power_of_2(GROUP)) = 16，多出的 12 行填 0。

KV head 3 对应的 q 块：

q_heads = 3 × 4 + offs_g。前 4 行是真实的 query head 12–15，后 12 行 mask 掉、值为 0，最后也不写回。Tensor Core 有 75% 的算力花在了 0 上，但 decode 本来就不是算力瓶颈，这点浪费无所谓；关键是 KV 只读了一次。

• m_i、l_i、acc 和 T5 一样，每一行一个状态，只是“行”从 query 位置变成了 query head。

④ page 循环：间接寻址

drills47–59 行

num_pages = tl.cdiv(ctx, PAGE)
for p in range(0, num_pages):
    phys  = tl.load(block_tables + b * stride_btb + p).to(tl.int64)    # 逻辑页 p → 物理块
    tok   = p * PAGE + offs_t
    t_mask = tok < ctx                                                 # 最后一页往往不满
    kv_off = phys * stride_cb + offs_t[:, None] * stride_ct + kvh * stride_ch + offs_d[None, :]
    k = tl.load(K_cache + kv_off, mask=t_mask[:, None], other=0.0)     # [PAGE, D]
    v = tl.load(V_cache + kv_off, mask=t_mask[:, None], other=0.0)
    s = tl.dot(q, tl.trans(k)) * scale                                 # [GROUP_PAD, PAGE]
    s = tl.where(t_mask[None, :], s, float("-inf"))

这里的 BLOCK_N 就是一页（PAGE = 16 个 token）。和 T5 唯一的区别是：下一块 K、V 在哪里，要先查 block table 才知道。

手算一个例子：ctx = 37，PAGE = 16，block_tables[b] = [9, 3, 7, …]。

num_pages = cdiv(37, 16) = 3。block_tables 第 3 项之后的内容不会被读到。

tl.trans(k)：k 读进来是 [PAGE, D]，转置成 [D, PAGE] 再和 [GROUP_PAD, D] 的 q 相乘。T5 是用 stride 直接读成转置的，这里先读再转置，两种写法都可以。

⑤ online softmax：和 T5 一样

drills63–73 行：m_new、p_、alpha、l_i、acc 的更新和 T5 逐字相同。第 0 页一定包含 token 0（context_lens ≥ 1），所以第一轮之后 m_new 就是有限值，不会出现 −inf − (−inf)。

循环结束 out = acc / l_i，只写回 g_mask 为 True 的那 GROUP 行。

⑥ 为什么 B = 1 时只有 73 GB/s

题目页的实测：

RTX 5090 有 170 个 SM。B = 1 时只有 8 个 program，162 个 SM 闲着；而这 8 个 program 每个都要串行走 32768 / 16 = 2048 页。总的读取量没变，但只用了不到 5% 的硬件。

⑦ 对照真实实现

• vLLM 的 Triton decode attention（两阶段，第一阶段就是这里的结构再加上分片）：vllm_fwd_kernel_stage1 里 vllmcur_kv_head = cur_head // kv_group_num 是 query head 到 KV head 的映射，vllmkv_page_number 是 page 的间接寻址；GQA 打包版是 vllm_fwd_grouped_kernel_stage1（BLOCK_H 就是这里的 GROUP_PAD）。
• vLLM 的 PagedAttention 设计文档：vllmpaged_attention.md，讲的是最早的 CUDA 版本（每个 warp 负责一部分 token），配了很多图。
• FlashInfer：flashinferBatchDecodeWithPagedKVCacheKernel；Python 端是 flashinferBatchDecodeWithPagedKVCacheWrapper，它的 plan() 会根据 batch 和长度决定要不要分片（F1 会用到）。

⑧ 常见坑与自测

• q_heads 算错：用了 kvh + offs_g * Hkv（交错分组），而模型实际是连续分组（head h 对应 KV head h // GROUP）。两种约定都存在，要和模型一致。
• 最后一页没 mask：读到上一个序列留下的旧 KV，结果悄悄出错。
• phys 没转 int64：phys * stride_cb 在大 cache 上溢出。
• 补齐行没 mask 写回：把 padding 行的 0/0 = NaN 写进了别的 head 的输出。

源码与资料索引

本仓库：drillssolutions/l2_paged_decode.py · drillsmine/l2_paged_decode.py · drillstests/test_l2_paged_decode.py · drillscontent/l2.md

vLLM：vllmtriton_decode_attention.py · vllmPagedAttention 设计文档

FlashInfer：flashinferBatchDecodeWithPagedKVCacheKernel · flashinferBatchDecodeWithPagedKVCacheWrapper

论文：GQA · PagedAttention