给开源教程修了个 demo

前一篇 KV-Cache 的显存账本 里,我用公式算清了显存的去向。教程本身写得很好,但我在读它的第三篇配套实验(Tutorial 03: KV-Cache)时发现:§5 的 Paged Attention demo 对三种配置输出的结果完全一样。于是我开了个 issue,写了个 PR 重写它。这篇记录整个过程。一、demo 怎么了:一个跑不坏的实验教程 §5 想演示 PagedAttention 的价值:把 KV-Cache 切成固定大小的块,按需分配,比静态预留整段显存省得多。但 demo 跑出来是这样:SystemMax UsersThroughputFragSpeedupStatic (baseline)323225 t/s0.0%1.3×Paged (16 tok)323225 t/s0.0%1.3×Paged (64 tok)323225 t/s0.0%1.3×三种页大小,结果实验结果完全相同。它没有坏没有报错但它什么都演示不了。这是最危险的一种 bug:看起来在跑和真的在演示原理是两回事。二、原因:四个巧合叠在一起排查后,问题不是单点故障,而是四件事撞在一起:#原因结果1max_batch_size32被原样传给三种配置,且内存推导出的上限(~130)比它大min(上限, 32)恒等于 32,批量上限把结果钉死2seq_len4096能被 2048/16/64 全部整除三种页大小算出相同的填充后大小,末页碎片都是 0%3引擎只建模末页内部碎片,静态分配的真实浪费(按max_seq预留、churn 下的外部碎片)不在模型里碎片栏对三行都一样4speedup_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()内部会拒绝放不下的请求,于是后续的短请求跳过被拒的长请求继续进入,分配器挑食,数字偏乐观。维护者改成:deffill(alloc,lengths):Admit requests in arrival order; stop at the first one that does not fit.fortinlengths:ifnotalloc.admit(t):breakreturnalloc到达顺序截断。它把实验变得更加诚实:真实系统里请求不会因为你内存满了就不来,分配器也不该在拒了长请求之后偷偷收下短请求。2. 数字下修,拥抱不好看的真相。我的版本:314 users、1.78× 吞吐加速。修改合并后:259 users、1.47×。少的那部分正是挑食喂出来的。五、复盘:三条收获demo 代码是看起来在跑的重灾区。报错、崩溃会立刻被发现;而每行都输出数字的实验,可能什么都没演示。写 demo 前先问:这个输出会随自变量变化吗?能变多少?诚实的实验设置要按到达顺序截断。“喂完整条流会赋予分配器不该有的选择权。任何排队/分配类仿真,截断点都要放在第一个失败处,否则你测的是挑食的分配器”,不是分配器。建模粒度决定结论的适用边界。引擎把加速比固定成 0.6 因子,意味着分配策略收益这个维度在引擎里是失真的——结论只能在这个粒度内成立。看懂一个工具把哪些东西解耦,比背它的输出更重要。六、参考资料Issue:#2134 Paged Attention demo returns identical results for every page sizePR:#2135 fix(mlsysim): redesign §5 Paged Attention demo in KV-Cache tutorial(merged)相关笔记:KV-Cache 的显存账本——教程 03 的学习记录,就是它带我发现了这个 demo 的问题。如果这篇复盘对你有帮助,欢迎交流。