[v1][kvcache] Honor prefix-cache retention interval for Mamba/linear attention

[Dao007forever]

Purpose

Wire VLLM_PREFIX_CACHE_RETENTION_INTERVAL to Mamba/linear-attention KV-cache groups, completing the in-code # TODO: Support Mamba/linear attention left by #43447 (which added the mechanism for sliding-window attention only).

Background. Mamba/KDA prefix caching retains a full recurrent-state snapshot once per block_size-token boundary. At small attention block sizes (e.g. 128) each snapshot spans several base blocks, and dense per-boundary retention saturates the KV pool — at block_size 128 the Mamba snapshots occupy ~80% of the blocks — leaving no uncached headroom, so the allocator is forced to evict live attention prefixes. The prefix-cache hit rate then collapses late in long multi-turn runs, dragging down throughput and tail latency, while larger block sizes (256/512) are unaffected (this is in a test setting which allows shorter block-size)

Change. MambaManager.reachable_block_mask now sparsifies state-snapshot retention the same way SlidingWindowManager does: keep one cached state per retention_interval-sized segment (plus the latest replay boundary) instead of one per block. A hit resumes from the nearest retained boundary (at most retention_interval tokens coarser), costing negligible extra prefill while freeing the intermediate snapshots for reuse. Also:

• MambaManager.cache_blocks now tolerates sparse (unhashed) blocks in the cached range.
• _validate_prefix_cache_retention_interval now accepts models with a Mamba group.

Default behavior is unchanged: with the interval unset, Mamba caches densely (every boundary), exactly as before.

Why this is not a duplicate

• [Prefix Caching] DeepSeekv4 - Support selective prefix-cache retention for sliding-window KV cache #43447 (merged) added VLLM_PREFIX_CACHE_RETENTION_INTERVAL for sliding-window KV cache only, and explicitly left # TODO: Support Mamba/linear attention. This PR completes that TODO.
• [Feature] Context-Aware KV-Cache Retention API (#37003) #38514 is a separate feature (a request-level retention-directive API + priority eviction queue); it does not touch the env-var interval path or reachable_block_mask.
• No other open PR wires retention-interval sparsification into Mamba/linear-attention groups.

Test Plan

.venv/bin/python -m pytest tests/v1/core/test_prefix_caching.py \
  -k "retention or reachable" -v

.venv/bin/python -m pytest tests/v1/core/ \
  -k "retention or mamba or prefix_cach" -q

E2E: Kimi-Linear-48B-A3B-Instruct (block_size 512, TP4, multi-turn prefix-on, 50 prompts × 60 turns), comparing default vs VLLM_PREFIX_CACHE_RETENTION_INTERVAL=2048.

Test Result

Unit tests (CPU; this change is pure-Python KV-cache scheduling logic):

• test_prefix_caching.py -k "retention or reachable" — including the new test_mamba_reachable_block_mask_sparsifies_retention — 10 passed.
• tests/v1/core/ -k "retention or mamba or prefix_cach" — all change-relevant tests pass in a CPU-only env; the remaining e2e/flash-attn tests in this selection require a full CUDA build + GPU and were not run here.

E2E: validated on Kimi-Linear-48B-A3B-Instruct with the config above; with VLLM_PREFIX_CACHE_RETENTION_INTERVAL=2048

---

AI assistance disclosure: AI assistance (Claude) was used for this change. The human submitter has reviewed every changed line and run the tests above.

🤖 Generated with Claude Code

[wzhao18]

Looks good to me. Thanks!

Just need to update the description for VLLM_PREFIX_CACHE_RETENTION_INTERVAL:

vllm/vllm/envs.py

Line 1055 in 475a6ad

# Applies to sliding-window attention for now but not yet Mamba/linear attention.

to include mamba.

[ivanium]

Make sense to me at a high level. Left one comment on mamba+spec decoding. Also cc @ZJY0516

[QilaiZhang]

Thanks for the PR. I have a small clarification question. My understanding is that Kimi-Linear-48B-A3B-Instruct prefix caching currently only supports align mode, and in align mode it does not save a recurrent-state snapshot at every block_size boundary. I may be missing something here, though. Could you explain how this applies to the Kimi-Linear benchmark case?

[Dao007forever]

Hi @QilaiZhang, in align-mode, it still deposits a snapshot every block_size tokens. Those linger as unreferenced-but-cached blocks and accumulate as context grows. This causes block reused and a drop in prefix cache hit rate. Live footprint stays ~2 blocks; cached footprint grows ~1 per block_size tokens, which is what the retention interval thins.

(I was running in a test setting which work around the uniform block-size of HMA.)

[QilaiZhang]

@Dao007forever Thanks, that makes sense regarding freed state blocks remaining as unreferenced cached blocks.

One follow-up: in upstream align mode, _mamba_block_aligned_split() seems to align scheduled chunks to a multiple of block_size, but not necessarily force one scheduled chunk per block. If a scheduled chunk spans multiple blocks, intermediate entries are null_blocks and only the chunk-end state is cached.

Was your Kimi-Linear benchmark configured so that aligned prefill chunks are effectively one block_size each, perhaps due to the HMA workaround you mentioned? That would explain the “one snapshot per block_size tokens” behavior.

[Dao007forever]

You are right that in prefill, we snapshot per chunk end, but in decode, we snapshot every block size still.

[QilaiZhang]

Thanks, that makes sense. I was only thinking about the prefill path and missed the decode-side behavior across block boundaries. The distinction between live footprint and unreferenced-but-cached footprint is helpful. Thanks for clarifying.

[ivanium]

One quick comment

[ivanium]

not sure if this part can work because for mamba, we need scheduler side changes to cache the end of the prompt. Maybe we can raise NotImplementedError when num_prompt_tokens is given?

[Dao007forever]

Good catch to raise it! I think we're safe here though — the mask is purely subtractive. A True never forces a block to be cached, it just declines to skip it; the real gate is blk.is_null in cache_full_blocks:

if blk.is_null or (block_mask is not None and not block_mask[i]):
    continue

So the layers stay separate: the scheduler decides where a snapshot lands (which blocks are non-null), and the mask just picks which existing ones to keep. If there's no state at the boundary, it's a null_block and gets skipped regardless of the mask — so we can never fabricate a cache entry over a stateless block. Worst case without scheduler changes is that the branch is inert, not incorrect (and in decode we snapshot every block_size, so it usually does fire).

[ivanium]

Make sense!

[ivanium]

Make sense!

[ivanium]

LGTM!