Hybrid 模型的 prefix cache
Qwen3-Next、Qwen3.5、Kimi Linear 这类 hybrid 模型同时有 attention 层和 recurrent 层(GDN、Mamba)。两种层记东西的方式完全不同,prefix cache 必须同时照顾两边。下面每一节都有一个可以动手的小实验,讲 vLLM、SGLang 和 Marconi 各自怎么做。
两种记忆
attention 层给每个 token 留一块 KV。想从第 p 个 token 接着算,把 KV 截到 p 就行。
recurrent 层只有一个固定大小的 state,每读一个 token 就原地覆盖一次。它回不到过去:想从 p 接着算,必须事先在 p 处存过一份 state 的拷贝,也就是 checkpoint。
拷贝很贵。以一个 GDN 模型的配置为例,一个 GDN state 是 3.12 MiB,相当于 799 个 token 的 KV。
所以 hybrid 模型的命中长度,等于 attention 和 recurrent 两边都能恢复的公共位置。剩下的问题只有三个。
每 N 个 token 存一个,只在 chunk 边界存,还是只在分支点存?
存太密会占满显存。vLLM 观察到密集快照占掉 80% 的 pool。
attention 匹配到 2432,最近的 checkpoint 在 1792。
只能从 1792 恢复,多匹配的 640 个 token 要重算。
KV 的大小正比于它省下的算力,state 不是。
代表 token 越多的 state 越值钱,LRU 会误删它们。
内存布局
vLLM 让所有层共用一个 block pool。为了让同一个 block 既能装 attention KV,也能装一个 state,它把 attention 的 block 撑大:3.12 MiB ÷ 4 KiB ≈ 799,向上对齐到 128 的倍数就是 896。启动日志里的 Setting attention block size to 896 tokens 就是这么来的。
SGLang 分两个池:state 按请求分 slot,KV 按 token 分 page。attention 的粒度不受影响;代价是 state 池按比例静态切出去,每个 running request 还要预留 3 到 5 个 slot。
vLLM · 匹配、写入、896
匹配时,每种层一个 manager,各找各的命中长度,然后取所有层都能恢复的最长位置。候选长度只缩不涨,直到不再变化。
写入更受限。GDN 类模型只能用 align 模式:只有当一个 prefill chunk 恰好结束在 896 的倍数上,那里的 state 才会被冻结成能命中的 checkpoint。所以 scheduler 会主动把 chunk 切在 block 边界上。
问题出在它用哪个网格切。scheduler 读的是 cache_config.block_size,也就是所有 cache group 里最小的 block size。只要有一个更小 block 的 group,切分网格就和 GDN 网格对不上。例如 128 对 896 的配置就是这样。
简化模型:三个请求依次到达,前 S 个 token 相同,之后各不相同。按 vLLM align 模式的切分规则推演,假设密集保留,prefill 预算 2048,attention block 896。用来解释机制,不是 benchmark。
用 Mamba group 自己的 block size,不再用所有 group 的最小值。
今天 main 的 scheduler.py:443 仍是 cache_config.block_size。
每跨过一个 block 边界都结束 chunk,从对齐位置开始也一样。
原来一个 chunk 能跨 k 个 block,中间 k − 1 个 slot 一直为空。
最后一个 block 的回退只对 eagle、eagle3、mtp 生效。
DFlash 和 DSpark 用自己的 KV,不再回退。对 MTP 没有影响。
| 情况 | 例子 | 不修会怎样 | 修后(社区部署数据) |
|---|---|---|---|
| 网格一致 | padding 之后 block size 相等的常规配置 | 没有问题 | 无变化,两组 A/B 零差异 |
| 有更细的 group | DFlash drafter 816 对 Mamba 1648;GLM-5.3 drafter 1152 对 2304;128 对 896 | 要么 state 从不落地,复用为 0;要么错位发布,命中过时的 state,输出悄悄变错 | 16.7K prompt 复用 0 → 100%,TTFT 8.18 s → 0.41 s;cold 与 warm 的 |Δlogprob| 0.25–0.36 → 0.02–0.04 |
状态:9-30 仍是 open,没有 maintainer review,pre-run-check 失败。数据来自 PR 描述和评论区的部署报告,不是 CI 结果;一位评论者撤回了自己的数字。
attention 命中而 Mamba 没命中,就在共享前缀交点补存 state
Mamba group 支持稀疏保留,每 k 个 token 留一个
稀疏保留时仍然保住交点 checkpoint
稀疏保留成为默认,间隔 0 只留语义 checkpoint
交点 checkpoint 可以不在 block 网格上
同期的默认值变化:所有 Mamba 模型默认 align(#50991);all 模式 deprecated(#55041),GDN 从来没支持过 all;KDA 能在一次 prefill 内导出中间 state(#52789),GDN 还没有。淘汰仍是一条全局 LRU。
Marconi 与 SGLang · radix tree
SGLang 和 Marconi 都用 radix tree。边上存 KV,节点上可以挂一份 state。匹配时沿树往下走,只能停在挂着 state 的位置;attention 多匹配的部分要重算。
三种 admission 策略的区别,在于什么时候往树上挂 state:
淘汰:LRU 只看最近用没用过。Marconi 的 FLOP-aware 还看这个 state 代表多少 token,越长越值钱。
教学模型:token 长度为示意,chunked prefill 每块 1024 token,只限制 state 数量,KV 视为足够。试试:选 Marconi,把上限调到 5,再在 LRU 和 FLOP-aware 之间切换。LRU 会丢掉文档 A 分支点上的 state,R5 只能从 1,024 恢复;FLOP-aware 保住它,R5 从 2,560 恢复。
intermediate_ssm,verify 后把被接受的那一步拷回。放在一起看
三者都收敛到同一件事:attention 决定上界,recurrent state 决定能不能落地。差别在 checkpoint 的密度、位置策略,以及显存怎么分。
| 维度 | vLLM | SGLang | Marconi |
|---|---|---|---|
| state 放哪 | 和 KV 共用一个 pool,block 撑大到 896 | 独立的请求级 slot 池 | radix tree 节点 |
| checkpoint 密度 | chunk 边界 ∩ block 网格,默认稀疏保留 | chunk 边界、请求结束、每 256 token | 每个序列最多 2 个 |
| 分支点 | attention 命中而 Mamba 未命中时推断交点 | branching seqlen(extra_buffer),SWA 也有 | 投机插入 radix tree |
| CoW | 只在 partial tail 命中时 | 每次命中都拷一份 | 无 |
| 淘汰 | 一条全局 LRU | 每组件 LRU 或 TLRU,内部节点 tombstone | recency + α · FLOPs/byte |
| 显存预算 | 一个 pool,自动分配 | 按比例静态切分,或 unified memory | 无 |
这是读代码和公开数据得出的判断,没有在同一个 workload 上做过对比实测。