给开源教程修了个 demo

前一篇 KV-Cache 的显存账本里,我用公式算清了显存的去向。教程本身写得很好,但我在读它的第三篇配套实验(Tutorial 03: KV-Cache)时发现:§5 的 Paged Attention demo 对三种配置输出的结果完全一样。于是我开了个 issue,写了个 PR 重写它。这篇记录整个过程。

一、demo 怎么了:一个”跑不坏”的实验

教程 §5 想演示 PagedAttention 的价值:把 KV-Cache 切成固定大小的块,按需分配,比静态预留整段显存省得多。但 demo 跑出来是这样:

System	Max Users	Throughput	Frag	Speedup
Static (baseline)	32	3225 t/s	0.0%	1.3×
Paged (16 tok)	32	3225 t/s	0.0%	1.3×
Paged (64 tok)	32	3225 t/s	0.0%	1.3×

三种页大小,结果实验结果完全相同。它没有坏（没有报错）但它什么都演示不了。这是最危险的一种 bug:“看起来在跑”和”真的在演示原理”是两回事。

二、原因:四个巧合叠在一起

排查后,问题不是单点故障,而是四件事撞在一起:

#	原因	结果
1	`max_batch_size=32` 被原样传给三种配置,且内存推导出的上限(~130)比它大	`min(上限, 32)` 恒等于 32,批量上限把结果钉死
2	`seq_len=4096` 能被 2048/16/64 全部整除	三种页大小算出相同的填充后大小,末页碎片都是 0%
3	引擎只建模”末页”内部碎片,静态分配的真实浪费(按 `max_seq` 预留、churn 下的外部碎片)不在模型里	碎片栏对三行都一样
4	`speedup_vs_static` 是固定的 0.6 因子(来自 Kwon et al. 2023),与分配策略无关;吞吐又只由内存受限的步时间决定	加速比、吞吐对分配策略零响应

第 4 条最致命:就算去掉批量上限和整除问题,这个 demo 的输出也不会随页大小变化——因为被评测的引擎本身不区分分配策略。结论:这个 demo 需要重写,而不是修补。

三、重写:让实验自己讲出理论

PR 的思路是放弃调用引擎模型,自己写解析仿真:两个小分配器共享同一个 KV 预算——

静态连续分配:每个请求预留一整段 max_seq,简单但浪费;
分页块池:固定大小的块按需分配,块表记账;

喂给它们同一条长尾分布的请求流(固定随机种子,可复现),于是上面的四个问题各自有了实验:

内部碎片:静态预留 77.8% 浪费 vs 16-token 分页 0.4%,并发用户 65 → 314;
外部碎片:churn 之后,静态 first-fit 碎片率 97.1%、10 个新 8K 请求只进得去 0 个;分页 10/10 全进;
页大小扫描:8→1024 token 的页,碎片与块表开销的权衡;
前缀共享:8 个请求共享 1K system prompt,峰值 KV 减 58%。

每个结论都能在表格里直接看到,数字全部可复现。PR 提交后,维护者 Vijay Janapa Reddi 很快回复:“diagnosis is spot on and the new demo actually shows the idea”,然后他 push 了一个小修,合了进去。

四、合并:更诚实的代码

合并前,维护者在我提交之上又做了一轮修改。逐条看,每一条都值得学:

1. 截断点:在第一个放不下的请求处停止,而不是喂完整条流。

我原来的写法是把 10 万条请求全部喂给分配器——包括内存耗尽之后到达的那些。PagedAllocator.admit() 内部会拒绝放不下的请求,于是后续的短请求跳过被拒的长请求继续进入,分配器”挑食”,数字偏乐观。维护者改成:

def fill(alloc, lengths):
    """Admit requests in arrival order; stop at the first one that does not fit."""
    for t in lengths:
        if not alloc.admit(t):
            break
    return alloc

到达顺序截断。它把实验变得更加”诚实”:真实系统里请求不会因为你内存满了就不来,分配器也不该在拒了长请求之后偷偷收下短请求。

2. 数字下修,拥抱不好看的真相。

我的版本:314 users、1.78× 吞吐加速。修改合并后:259 users、1.47×。少的那部分正是”挑食”喂出来的。

五、复盘:这次被改写的三条收获

demo 代码是”看起来在跑”的重灾区。 报错、崩溃会立刻被发现;而每行都输出数字的实验,可能什么都没演示。写 demo 前先问:这个输出会随自变量变化吗?能变多少?
诚实的实验要按到达顺序截断。 “喂完整条流”会赋予分配器不该有的选择权。任何排队/分配类仿真,截断点都要放在第一个失败处,否则你测的是”挑食的分配器”,不是”分配器”。
建模粒度决定结论的适用边界。 引擎把加速比固定成 0.6 因子,意味着”分配策略收益”这个维度在引擎里是失真的——结论只能在这个粒度内成立。看懂一个工具把哪些东西解耦,比背它的输出更重要。

六、参考资料

Issue:#2134 Paged Attention demo returns identical results for every page size
PR:#2135 fix(mlsysim): redesign §5 Paged Attention demo in KV-Cache tutorial(merged)
相关笔记:KV-Cache 的显存账本——教程 03 的学习记录,就是它带我发现了这个 demo 的问题。

如果这篇复盘对你有帮助,欢迎交流。

给开源教程修了个 demoFixing a Tutorial Demo: An Open-Source PR Retrospective

给 mlsysim 的 KV-Cache 教程重写了 §5 的 Paged Attention demo 并提了 PRI rewrote the §5 Paged Attention demo in mlsysim's KV-Cache tutorial and opened PR #2135; before merging, the maintainer revised it and shrank my numbers from 314 to 259. A retrospective of that experience.