#19 · Track B · Stage 进程与调度 · 书 p.84

P09 · Paged KV 写入 + decode attention（参考实现）

先用 PyTorch 写出正确的 paged KV 写入与 decode attention 参考实现——它就是之后 Triton kernel 的对拍标准。

⏱ 45 min 难度中实现 mine/p09_paged_attention.py 验收 ./check.sh p09 提示 A-19 · 参考解 C-19 solutions/p09_paged_attention.py 标记完成

题面速览

在系统中的位置：P06 决定每个 token 的 K/V 存在哪个物理 slot。model runner 每步：先按 slot_mapping 把新 token 的 K/V 写进 cache，再让 attention 按 block_tables 去 cache 里读历史。

cache 布局:   [num_blocks, block_size, num_kv_heads, head_dim]
写:           cache[slot // bs, slot % bs] = k[i]          （slot = -1 是 padding，不写）
decode 读:    序列 b 的前 context_lens[b] 个 token：
              k_b = cache[block_tables[b, :ceil(n/bs)]].reshape(-1, KVH, D)[:n]
              out_b = softmax(q_b · k_bᵀ · scale) · v_b    （GQA：head h 用 kv head h // (H/KVH)）

易错点（题面）：物理块不连续（测试会打乱分配）；block table 尾部未用的位置是垃圾值，不能读；CUDA graph 模式下 batch 会被 pad，pad 出来的 token 的 slot 是 -1；softmax 用 fp32。

验收要点：与 dense attention 数值一致（GQA、非连续物理块、block table 尾部是垃圾值）· slot -1 的 padding token 不写入 · 随机形状与 block_size。

直觉与类比

储物柜墙：slot 就是柜号

KV cache 是一面储物柜墙：num_blocks 组、每组 block_size 格。slot = 组号 × 每组格数 + 格号，所以“展平后的第 slot 格”和 cache[slot // bs, slot % bs] 是同一个位置。一个序列拿到的是一张取件码清单（block table），组号乱序、不连号都没问题。最危险的是 slot = -1：Python/PyTorch 里 -1 表示“最后一格”，于是本该丢掉的 padding token 会安静地覆盖掉最后一格里别人的包裹——不报错、不告警。

图书馆的索书单

block table 是印好格子的索书单（固定 max_blocks 行），你只填了前两行，后面是上一个读者留下的乱码——照着去取书，轻则取到别人的书（结果悄悄错），重则架号不存在（IndexError）。最后一卷只看前几页：“反正那两页是空白”这种侥幸不成立——空白页的 score = 0，在 softmax 里仍占一份权重（e⁰ = 1），只有截断（或置 -inf）才真正排除它。

框架视角 · shape trace

write_kv（可以整批向量化）：

步骤	张量 / 形状	做什么	易错点
1	slot_mapping `[T]`	`keep = slot_mapping >= 0`	没过滤 -1：padding 覆盖最后一个 slot，不报错
2	`s = slot_mapping[keep]` `[T']`	`blk = s // bs`、`off = s % bs`	int32 索引在某些后端不被接受，转 long
3	`k[keep]` `[T', KVH, D]`	`k_cache[blk, off] = k[keep]`（v 同理）	写成 `k_cache[blk][off]` 会先复制一份再写，改不到原 cache
4	—	返回 None（原地）	返回新张量 → 调用方拿不到写入结果

paged_decode_attention（逐序列循环即可，先写对）：

步骤	张量 / 形状	做什么	易错点
1	`n = int(context_lens[b])`	`nb = ceil(n / bs)`	用 `n // bs` → 最后一个不满的块被整块丢掉
2	`tbl = block_tables[b, :nb].long()`	只取有效项	取整行 → 读垃圾值：IndexError 或别人的 KV
3	`k_cache[tbl]` `[nb, bs, KVH, D]`	按 table 的顺序 gather，逻辑顺序自动正确	以为要先排序 table——不行，table 的顺序就是逻辑顺序
4	`.reshape(nb*bs, KVH, D)[:n]`	展平成 token 维，截掉最后一块多出来的位置	不截断 → 多看几个位置，结果悄悄错
5	`.repeat_interleave(H // KVH, dim=1)` → `[n, H, D]`	GQA	用 repeat → `h % KVH`；KVH=1 时两者相同，很难发现
6	scores `[H, n]` fp32	`einsum("hd,nhd->hn", q[b].float(), kb) * scale`，在 n 维 softmax	scale 乘在 softmax 之后；在 bf16 里做 softmax
7	`out[b]` `[H, D]`	`einsum("hn,nhd->hd", probs, vb).to(q.dtype)`	忘了转回 q.dtype；head 维和 token 维搞混

不变量	怎么验
I1 -1 的 slot 对应的值不出现在 cache 的任何位置	写一个所有元素都是哨兵值（如 9.0）的 padding token，写完 `assert (cache == 9).sum() == 0`
I2 只访问 `block_tables[b, :ceil(n/bs)]`	把 table 尾部填成 99999，正确实现不报错
I3 参与 attention 的 K/V 行数恰好 == context_lens[b]	`assert kb.shape[0] == n`，并用 n 不是 bs 倍数的长度专门测
I4 与 dense attention 数值一致	随机长度（含 1、bs 的倍数、非倍数）、打乱物理块、尾部填垃圾值对拍
I5 GQA：head h 用 kv head h // group	H=4, KVH=2：repeat_interleave 得 [0,0,1,1]，repeat 得 [0,1,0,1]

易错点

bug	症状
没过滤 slot -1	最后一个物理块的最后一个 slot 被 padding 覆盖，长序列结果偶发错误
读了 table 的整行	IndexError（垃圾值越界），或读到别人的 KV
没截掉最后一块的多余位置	上下文长度不是 bs 整数倍的序列结果错
GQA 用 repeat 而不是 repeat_interleave	num_kv_heads=1 时正确、>1 时错——很隐蔽
block_tables 是 int32 直接索引	新版 PyTorch CPU 上能用，旧版本 / 某些后端不支持；转 long 最稳
softmax 在 bf16 中做	与 fp32 参考有 1e-2 量级偏差

逐步图：例 2 序列 0 的 gather

▶ block_size=2, num_blocks=4, KVH=1, H=2, D=2, scale=1；序列 0 有 3 个 token，物理块 [2, 0]

1. 写入

K0 = [[1,0], [0,1], [1,1]]、V0 = [[1,2], [3,4], [5,6]] 写入 slot [4, 5, 0]：

物理块：块0 = [1,1]· 块1 = ·· 块2 = [1,0][0,1] 块3 = [2,2]·（块3 是序列 1）

block_tables = [[2, 0], [3, 77]]——77 是垃圾值（甚至超出 num_blocks）

2. n 与 nb

n = context_lens[0] = 3

nb = ? table = ?

nb = ceil(3/2) = 2；table = block_tables[0, :2] = [2, 0]

3. gather + reshape

k_cache[table] 形状 [2, 2, 1, 2] → reshape 后（逻辑位置 0..3）：

[1,0][0,1][1,1][0,0]

第 4 个位置是块 0 offset 1——未写入的 0 向量，不属于该序列。不截断的话它的 score 是 0 而不是 -inf，会被分到注意力。

4. 截断到前 n 个

[1,0][0,1][1,1]

5. head0：q = [1, 0]

scores = q·Kᵀ × 1.0 = [1, 0, 1] → softmax = [0.4223, 0.1554, 0.4223] → out = [3.0, 4.0]

head1：q = [0, 2] → scores [0, 2, 2] → [0.0634, 0.4683, 0.4683] → [3.8099, 4.8099]

6. 序列 1：只有 1 个 token

nb = 1，table = [3]（不读 77）。无论 q 是什么输出都是 V1 = [7, 8]。

结果：[[[3.0, 4.0], [3.8099, 4.8099]], [[7.0, 8.0], [7.0, 8.0]]]

三级提示

L1I/O 示例与关键问题I/O 示例 + 关键问题 + 细节（轻度剧透）

例 1：write_kv 与 padding（cache 形状 [3, 2, 1, 1]，初始全 0）：

k = [[[10.]], [[20.]], [[30.]], [[40.]]]          # [4, 1, 1]
v = k * 0.1
slot_mapping = tensor([4, 5, -1, 1])
write_kv(k_cache, v_cache, k, v, slot_mapping)
k_cache[:, :, 0, 0]  # -> [[ 0., 40.],    块0: offset1 <- token3 (slot 1)
                     #     [ 0.,  0.],    块1: 没人写
                     #     [10., 20.]]    块2: slot 4,5 <- token0, token1
v_cache[:, :, 0, 0]  # -> [[0., 4.], [0., 0.], [1., 2.]]

token2 的 slot 是 -1，30 不能出现在 cache 的任何位置——特别是最后一个 slot（块2 offset1）。

例 2：见逐步图。例 3：GQA 映射（H=8, KVH=2, group=4）：正确 h // group = 0 0 0 0 1 1 1 1；错误 h % KVH = 0 1 0 1 0 1 0 1。

关键问题

slot 7、block_size 4 对应哪个物理块、哪个 offset？cache 用 view(-1, KVH, D) 展平后 slot 和第一维下标什么关系？
直接 k_cache.view(-1, KVH, D)[slot_mapping] = k，slot -1 会写到哪里？会报错吗？
context_len = 5、block_size = 4：table 前几项有效？取出后一共几个 token 位置、要截掉几个？
block_tables 的 dtype 是什么？能直接用作高级索引吗？
GQA 下 repeat_interleave 和 repeat 哪个与 h // (H // KVH) 匹配？
“Softmax in fp32”：哪些张量转 fp32？输出最终什么 dtype？
scale 乘在哪一步？乘在 softmax 之后会怎样？

细节

前两维合起来就是“slot”维：cache[slot // bs, slot % bs] 与 cache_flat[slot] 同一位置。
block_tables 是 padding 成 [B, max_blocks] 的矩形，序列 b 只有前 ceil(context_lens[b] / bs) 项有意义。
decode 阶段每个序列只有最新的 1 个 query，它可以看见全部 context（包括自己，已先 write_kv 写进去），不需要 causal mask。

L2思路与常见坑思路拆分与常见坑（中度剧透）

write_kv（向量化）：valid = slot_mapping >= 0 → 取有效 slot 与 k/v 行 → blocks = slots // bs, offs = slots % bs → k_cache[blocks, offs] = k[valid]。原地写，返回 None。
decode（先写循环版）：对每个 b：n、nb = ceil(n / bs)、table = block_tables[b, :nb].long()。
gather + reshape + 截断：k_cache[table] [nb, bs, KVH, D] → reshape [nb*bs, KVH, D] → 取前 n 行。
GQA：k.repeat_interleave(H // KVH, dim=1)；也可以把 q reshape 成 [KVH, group, D] 不复制 K。
注意力：fp32 scores，softmax 在 j 维，加权和后 .to(q.dtype) 写入输出 [B, H, D]。einsum "hd,nhd->hn" / "hn,nhd->hd" 很方便。

L3详解：答案详解：答案、逐步追踪、骨架提示（重度剧透）

关键问题答案

块 1、offset 3；展平后第一维下标就是 slot。
写到展平后的最后一个位置（最后一块的最后一个 offset），不报错——这正是危险之处。
前 2 项；2·4 = 8 个位置，截掉后 3 个，只留前 5 个。
通常是 int32（vLLM 里为省显存/喂给 kernel）；新版 PyTorch 也接受，但转 long 再做高级索引最稳妥。
repeat_interleave（[a,b] → [a,a,a,a,b,b,b,b]）；repeat 得 [a,b,a,b,…]，对应 h % KVH。
q、k、v 都可以转 fp32 算 scores / softmax / 加权和；输出转回 q.dtype。
softmax 之前乘在 scores 上；乘在之后只是把输出整体缩放，注意力分布不对。

自检：自己写一个 dense 版本（直接用连续的 K/V 张量算注意力），随机生成若干序列长度（含 1、恰好 bs 的倍数、非倍数），打乱物理块分配、把 block table 尾部填成大垃圾值，与 paged 版本 torch.testing.assert_close 对拍。

面试官追问 & 对照源码

Triton 版：一个 program 处理 (seq, head)，按 block 循环做 online softmax。
prefill 的 varlen 版本怎么写？

对照源码：vllm/v1/attention/backends/triton_attn.py