[Bugfix][Mamba] Zero mamba state blocks on (re)allocation to stop stale-state poisoning on hybrid models - #56524
datacrystals wants to merge 1 commit into
Conversation
Mamba/KDA state blocks were handed out by BlockPool.get_new_blocks without zeroing, and the worker-side zeroing queue recorded allocations only for AttentionSpec layers and skipped tuple-valued kv_cache bindings, so KDA state tensors (conv/recurrent/RecoverSSM records) could be read back with a previous tenant's bytes. On hybrid models (Kimi-K3 with KDA layers) this surfaced as cross-request state poisoning after aborted prefills. - _record_new_block_ids now covers MambaSpec managers - MambaManager.allocate_new_blocks (align mode) records blocks from its direct get_new_blocks() call - KVBlockZeroer walks MambaSpec groups and unpacks tuple/list kv_cache bindings into per-tensor segments Zeroing runs before CoW copies and the forward pass each step, so a stale slot read before its first legitimate write returns zeros (a fresh recurrent state) instead of a previous tenant's bytes.
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment Once the PR is approved or has the If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
|
This pull request has merge conflicts that must be resolved before it can be |
1 participant
[Bugfix] Zero mamba state blocks on (re)allocation to stop stale-state poisoning on hybrid models
Status: draft pending CI + unit tests. Validated in production on an 8×B300 node serving
moonshotai/Kimi-K3(hybrid MLA + KDA/mamba,--kv-cache-dtype fp8, prefix caching, dspark spec-decode).Problem
Mamba/KDA state blocks are handed out by
BlockPool.get_new_blockswithout zeroing, and the attention-only zeroing queue never covers them:KVCacheConfig.needs_kv_cache_zeroingis set toTruefor mamba models (kv_cache_interface.py, "Required for Mamba layers, whose state is read before it is fully written ([BUGFIX][Mamba][Qwen3.5] Zero freed SSM cache blocks on GPU #35219)"),SingleTypeKVCacheManager.__init__only sets_record_new_block_idsforAttentionSpec(vllm/v1/core/single_type_kv_cache_manager.py), so mamba block IDs are never recorded for the zeroing queue,KVBlockZeroer(vllm/v1/worker/utils.py) skips non-AttentionSpeclayers and non-Tensorkv_cachevalues — KDA layers bind a tuple of state tensors (conv_state, recurrent_state, *recoverssm_records = self.kv_cache,vllm/models/kimi_k3/nvidia/kda.py), so both guards exclude them.Any read of a mamba state slot that was not written by the current tenant therefore inherits arbitrary bytes from a previous tenant. We observed this in production as cross-request poisoning: after an aborted prefill, a retry resumed from a state block previously freed mid-write and emitted unbounded multilingual garbage until engine restart.
Fix
_record_new_block_ids: record allocations forMambaSpecmanagers (in addition to attention).MambaManager.allocate_new_blocks(align cache mode): record the IDs returned by its directblock_pool.get_new_blocks()call, which bypasses the base-class recording path.KVBlockZeroer.__init__: acceptMambaSpecgroups and unpack tuple/listkv_cachebindings, emitting one segment per state tensor. The existing per-block segment math applies unchanged (state tensors are[num_blocks, ...]block-strided carves of the shared page).Per-step ordering is already
zero → block copies → forward, so a slot read before its legitimate write returns zeros (i.e. a fresh recurrent state) instead of a previous tenant's bytes — degrading quality gracefully instead of corrupting the session.Cost
~0.6–1.0 MiB zeroed per mamba state block per layer group (K3: 69 KDA layers / 14 groups ⇒ a few MiB per block ID in practice), only at allocation time; mamba blocks are allocated at prefill-chunk boundaries. Measured overhead in production: none observable (throughput and TTFT unchanged within noise).
Validation
storm4.py, attached below): N workers sharing a 40k-token prefix; each cycle aborts mid-prefill (~1.2 s in) and immediately retries. Unpatched nightly: 1–2 in 36 retries poisoned with garbage (verified by full-text inspection). Patched: 0/72.idlegap.py: cache-hit-after-idle regression test, passes.verify.pyfunctional suite (text/reasoning/vision/prefix-cache) passes;stress_multiturn.py(3×6-round growing sessions) passes; throughput unchanged (243–246 tok/s single-conn decode).Related
AI-assisted (Claude Code / Kimi Code agent) root-cause and patch authoring, with human operator running all production experiments.
(POC repro script and full production logs available; happy to attach.)