[Infra-13] vLLM 如何修复 Mamba Prefix Cache Poisoning

一个 Hybrid GDN/Mamba 模型在单请求测试里连续运行都正常，到了并发环境却偶发重复输出、生成循环或结果漂移；关闭 prefix caching 或 MTP 后，问题又随之消失。

这种现象很容易让人先怀疑 sampling、数值精度、speculative verification，甚至底层 kernel。但在 vLLM PR #51113 中，真正的根因来自一个更隐蔽的跨模块错误：scheduler 切出的 prefill chunk 破坏了 Mamba recurrent state 的缓存边界，导致 prefix cache 的元数据与物理状态不一致。

PR 的标题是：

[Bugfix] Keep mamba align prefill chunks block-aligned past last_cache_position

它修复的是 mamba_cache_mode="align" 下，MTP/EAGLE 与 prefix caching 共同启用时可能发生的 prefix-cache poisoning。代码改动只有几行，但它重新建立了整个系统最重要的 correctness invariant：

physical recurrent state offset
==
cache hash advertised offset

本文从这个 invariant 出发，完整还原 bug 如何发生、为什么并发和 EAGLE 会放大它，以及两处 scheduler 修改为什么足以修复问题。

1. Mamba Cache 与 KV Cache 有什么不同#

对于普通 Transformer attention，一段 token 的 KV Cache 可以按 token 或 block 保存：

token 0 ... token 1599
        ↓
KV cache block 0

GDN/Mamba 则维护随序列不断推进的 recurrent state：

S_0 --token 1--> S_1 --token 2--> S_2 ... --token 1600--> S_1600

在 vLLM 的 mamba_cache_mode="align" 中，这些状态按 MAMBA_BLOCK_SIZE 对齐保存。假设：

MAMBA_BLOCK_SIZE = 1600

那么每个物理 slot 的语义是：

slot 0 <=> state@1600
slot 1 <=> state@3200
slot 2 <=> state@4800
...

一般地，必须满足：

slot p holds state@((p + 1) * block_size)

这也是 PR 新增测试明确写下的 invariant。它不是一种性能偏好，而是 prefix cache 能否安全复用 Mamba state 的前提。

如果 cache metadata 声称 slot 0 保存的是 state@1600，物理 slot 中就必须真的是处理完前 1600 个 token 后的 recurrent state。任何偏差都会让后续请求从错误的历史状态继续计算。

2. Cache Poisoning 是如何产生的#

考虑一组贯穿全文的参数：

model              = Qwen3.6-27B（Hybrid GDN）
mamba block size   = 1600
prompt length      = 2002
MTP/EAGLE          = enabled
prefix caching     = enabled

高并发时，scheduler 的 token budget 由多个 request 共享。Request A 在第一个 scheduling step 中可能只分到 364 个 token，于是第一次 prefill 变成：

token 0 -> token 364

GDN forward 结束后，物理 slot 中暂存的是：

slot 0 = state@364

这本身未必错误。它可以只是一个尚未完成的 running state。危险在于，旧 scheduler 允许下一个 chunk 从 364 直接运行到 prompt 末尾：

364 ------------------------------> 2002
              ↑
             1600

这个 chunk 在内部跨过了第一个 Mamba block boundary。cache manager 看到 block 0 已经完成，便为它计算代表前 1600 个 token 的 block hash：

cache metadata: slot 0 = state@1600
physical data:  slot 0 = state@364

从这一刻起，prefix cache 已被污染。

后续 Request B 带着相同 prefix 到来。cache lookup 告诉它系统已经拥有 state@1600，但实际恢复出来的却是 state@364。Request B 随后直接从 token 1601 继续运行：

expected: S_1600 -> token 1601 -> token 1602 -> ...
actual:   S_364  -> token 1601 -> token 1602 -> ...

token 365 到 1600 的信息从 recurrent state 中完全缺失。其结果可能是 logits 轻微漂移，也可能从首个不同 token 开始改变整条 autoregressive trajectory，最终表现为：

• 同一个 prompt 偶发得到不同结果；
• 输出出现重复、循环或明显精度异常；
• 单请求难以复现，并发压力下出现；
• 关闭 prefix caching 或 MTP 后消失。

3. 为什么单请求通常“意外安全”#

单请求下，scheduler 往往拥有充足的 token budget。例如 max_num_batched_tokens = 8192 时，一个 2002-token prompt 可以在一个 step 内完成：

0 --------------------------------> 2002

它不会被切成：

0 -> 364
364 -> 2002

因此也不容易在第一个 slot 中留下 state@364，再由后续 chunk 跨过 1600 boundary。这里并不是单请求调度在语义上更正确，而是足够大的 budget 恰好避开了危险切分。

并发 prefill 则会争用全局 budget。请求可能在任意 token offset 被截断，因此更容易暴露 scheduler chunk boundary 与 Mamba cache boundary 之间的不一致。这解释了一个典型的线上模式：

单请求回归 100 次：正常
并发 workload：偶发错误

4. EAGLE 为什么会放大问题#

旧逻辑使用 last_cache_position 决定是否继续强制 Mamba block alignment。启用 EAGLE 时，相关代码还会将它向前移动一个 Mamba block：

if self.use_eagle:
    last_cache_position = max(
        last_cache_position - block_size,
        0,
    )

对于 prompt 长度 2002、block size 1600 的请求，last_cache_position 可能因此变成 0。旧 guard：

if end < last_cache_position:

随即退化为：

if end < 0:

它永远为 False，等于整个 prefill 都失去了 block-alignment 保护。scheduler 可以自由地切出 0 -> 364，再让下一段跨过 1600。

问题的本质在于：

prefill_end
    = computation correctness boundary
 
last_cache_position
    = cache policy / hit-rate boundary

last_cache_position 仍然可以作为提升 cache hit rate 的 mandatory stop，但它不能决定 recurrent computation 在哪里可以不再遵守对齐约束。旧代码混淆了缓存策略边界与计算正确性边界。

5. 第一处修复：直到 Prefill 结束都保持对齐#

PR 先显式提取了 prefill_end：

prefill_end = max(
    request.num_prompt_tokens,
    request.num_tokens - 1,
)
 
if start >= prefill_end:
    return num_new_tokens

变量提取本身不改变行为，真正关键的是将：

if end < last_cache_position:

改为：

if end < prefill_end:

新规则可以翻译成一句话：

只要当前 chunk 还不是整个 prefill 的最后一个 chunk，它就必须遵循 Mamba block alignment。

回到 364 / 1600 / 2002 的例子。scheduler 本轮只剩 364 token budget 时，发现 364 小于 prefill_end=2002，说明这不是 final prefill chunk；一个安全的非最终 chunk 必须运行到下一个 boundary，也就是 1600。

但本轮 budget 不足以覆盖 1600 个 token，于是 scheduler 返回 0，将请求 defer 到后续 step：

step 1: budget = 364
        0 -> 1600 cannot fit
        request deferred
 
step 2: budget >= 1600
        0 -> 1600
        slot 0 = state@1600

这样一来，slot 0 在成为可哈希、可复用的完整 block 时，物理数据就与它声明的 offset 完全一致。

6. 为什么最后一个 Prefill Chunk 可以不对齐#

prompt 长度是 2002，因此处理完 0 -> 1600 后仍剩下 402 个 token：

1600 -> 2002

PR 特意豁免了最后一个 prefill chunk：

# Exempt: the prompt's last chunk,
# whose slot decode advances
# to the boundary.

这是安全的，因为 state@2002 仍然只是 slot 1 中的 running state，系统没有把它声明为完整的 state@3200。随后的 decode 会继续推进：

state@2002 -> state@2003 -> ... -> state@3200

直到真正到达 boundary，它才具备完整 block 的语义。因此正确规则不是“所有 chunk end 都必须对齐”，而是：

intermediate prefill chunk: must be block-aligned
final prefill chunk:        may end at prompt_end

7. 第二处修复：Mid-block Resume 必须先补齐当前 Block#

另一类危险来自不对齐的 chunk start。假设请求从 token 2531 恢复：

1600 -------- 2531 ---------------- 3200 -------- 3602
               ↑                     ↑
             start             next boundary

正确行为必须先运行到 3200：

2531 -> 3200

旧代码只有在 next boundary 不超过 last_cache_position 时，才把它加入 mandatory stops：

next_block_boundary
if (
    start % block_size != 0
    and next_block_boundary <= last_cache_position
)
else 0

如果 last_cache_position=1600，那么 3200 <= 1600 为假，3200 不会成为 stop。scheduler 可能直接执行 2531 -> 3602，再次跨过本应完整落盘的 boundary。

PR 删除了与缓存策略有关的第二个条件：

next_block_boundary
if start % block_size != 0
else 0

现在，只要请求从 block 中间恢复，就必须先补齐当前 block，无论 last_cache_position 在哪里。

两处修改共同维护了同一个 invariant：

end alignment:
  非最终 prefill chunk 不能停在 block 中间
 
start re-alignment:
  从 block 中间恢复时，不能跨过下一个 boundary

换句话说，任何 Mamba block boundary 都不能被一个不安全的 prefill chunk 从内部跨越。

8. Regression Test 测的不是输出，而是系统不变量#

PR 新增了 tests/v1/core/test_mamba_align_chunk_split.py。测试没有只检查某次生成是否碰巧正确，而是构造真实的：

FullAttentionSpec
    +
MambaSpec
    +
KVCacheManager

它用以下参数模拟对应 workload：

ATTN_BLOCK_SIZE = 16
MAMBA_BLOCK_SIZE = 1600
NUM_SPEC = 3
PROMPT_LEN = 2002

测试同时记录两个事实：

block.block_hash_num_tokens
    = cache metadata 声称该 block 对应哪个 token offset
 
state_at[block.block_id]
    = 物理 Mamba slot 实际保存哪个 token offset 的 state

然后直接断言二者一致。旧版本暴露出的 failure 正是：

mamba slot 0 is hashed as state@1600
but holds state@364

这种测试比只看最终 accuracy 更接近根因。生成结果具有随机性，cache poisoning 也依赖调度时序；而系统 invariant 一旦被破坏，就可以在错误发生的最早位置被稳定捕获。

测试还覆盖了 1536 和 12288 等不同 Mamba block size，说明这不是 1600 或某个 Qwen 模型特有的问题，而是所有按 block boundary 缓存 recurrent state 的 hybrid linear-attention serving 系统都需要遵守的约束。

PR 给出的结果是：

before fix: 14 failed, 6 passed
after fix:  20 passed

真实 workload audit 也使用 hybrid GDN、3 个 speculative tokens、prefix caching 和并发请求捕获到过：

slot 6 hashed as state@3920
but holds state@3662

修复前出现 1 个 poisoned entry，修复后降为 0。这证明问题不仅存在于构造出来的边界测试中，也能在正常 serving workload 里发生。

9. 从重复输出到 Root Cause#

现在可以把完整因果链串起来：

MTP / EAGLE
    ↓
last_cache_position 向前移动
    ↓
旧 alignment guard 提前失效
    ↓
并发 prefill 被切成任意长度
    ↓
GDN/Mamba 在 mid-block chunk end 写入 recurrent state
    ↓
下一 chunk 跨过 Mamba block boundary
    ↓
cache manager 为 slot 注册 boundary hash
    ↓
state@364 被声明成 state@1600
    ↓
prefix cache poisoning
    ↓
其他 request 命中并恢复错误 recurrent state
    ↓
logits 与 token trajectory 漂移
    ↓
偶发重复、循环或精度异常

因此，当 Hybrid GDN/Mamba 模型出现“单请求正常、并发异常，关闭 APC 或 speculative decoding 后恢复”的现象时，排查范围不应只局限在 sampling、kernel 或 verify 阶段。还需要检查：

• scheduler 实际切出了哪些 prefill chunk；
• chunk start/end 是否跨过 recurrent-state boundary；
• cache metadata 给每个 slot 声明了哪个 token offset；
• 物理 slot 实际保存的是哪个 offset 的 state；
• EAGLE/MTP 是否改变了用于 alignment 的 cache position。

10. 这个 PR 真正修复了什么#

表面上，PR #51113 只是让 Mamba prefill chunk “更对齐”。更准确地说，它在 scheduler 层重新建立了一个跨模块的全局一致性约束：

Prefix cache 对外声明的 recurrent state offset，必须与物理 slot 中实际完成计算的 token offset 完全相同。

这个 bug 之所以隐蔽，是因为涉及的每个模块单独看都似乎合理：

Scheduler
    决定本轮执行多少 token
 
GDN/Mamba kernel
    在 chunk end 写 recurrent state
 
Mamba cache manager
    按 block boundary 管理和哈希 state
 
Prefix cache
    跨 request 复用已完成 block

真正出错的是它们对同一个问题没有共享一致答案：这个 slot 到底代表处理完多少个 token 后的状态？

这也是 PR 最值得借鉴的地方。很多严重的 AI Infra correctness bug，并不来自某一段显然错误的计算，而来自调度、kernel、缓存和元数据之间被悄悄破坏的语义契约。修复这种问题，关键不是在最终输出端增加更多特例，而是找到那个跨模块 invariant，并让测试直接验证它。