Dedicated mamba block pools for GDN prefix caching (align mode), no vLLM change

[wenxindongwork]

Summary

With --enable-prefix-caching, hybrid GDN models (Qwen3.5) run vLLM's mamba align mode, which addresses recurrent state by block id. #3422 therefore skips the compact-mamba sizing and every block id gets a GDN state slot in all 45 GDN layers. Because vLLM hands out block ids from one pool to the attention group and to each of the 3 GDN groups, a block id used for attention still reserves (and never touches) 183 MiB of GDN state, and a GDN id of one group reserves slots in the other two groups' layers.

Measured on Qwen3.5-397B-A17B-FP8, 8x v7x, TP8 / attn-DP 4, bf16 KV, block 1024, 286 GiB KV budget:

	per 1024-token block id	pool	conversations (30k ctx) per attn-DP rank before LRU thrash
today (align, uniform)	30 MiB attn + 183 MiB GDN = 213 MiB	1372 ids (343 / rank)	~9
this PR (MAMBA_ALIGN_STATE_MEM_FRACTION=0.3)	attn 30 MiB; GDN 61 MiB per group slot	~6.6k attn ids + ~130 GDN slots per group per rank	~32 (attention-bound past ~55)

A running request pins only 2 GDN slots per group and a retained prefix snapshot costs 1, so today's layout leaves the GDN arrays ~95% idle while they hold ~85% of the budget. An LRU replay of the agentic multi-turn trace (16 streams/group, ~30 turns, ctx 7.6k -> 30k) goes from 0% turn-level prefix hits at 4+ concurrent groups to ~88% through 8 groups.

Change (no vLLM change needed)

• tpu_inference/core/mamba_block_pool.py (new): TPUMambaSpec is a MambaSpec carrying num_blocks; TPUMambaManager is a MambaManager that owns a private BlockPool of that size instead of the shared pool the coordinator passes in. Registered through vLLM's platform hook TpuPlatform.register_custom_kv_cache_specs, so every process that builds a coordinator maps the spec to the manager. The manager covers the places where vLLM assumes a single pool:
  • admission: get_num_blocks_to_allocate returns "shared pool size + 1" (the sentinel MambaManager already uses) when the private pool cannot serve, so allocate_slots returns None and the scheduler waits/preempts;
  • cache hits: the classmethod find_longest_cache_hit receives the shared pool and a batch of group ids sharing a spec; each group is looked up in its own private pool and the hit length is reconciled;
  • deferred frees: the scheduler returns popped blocks to the shared pool, so popped mamba blocks are held and released max_concurrent_batches + 1 scheduler steps later (clocked by new_step_starts);
  • prefix-cache reset: detected at the next step (the shared pool's hash map object is replaced) and applied to the private pool;
  • partial-tail offload hand-offs and KV events are disabled for private pools.
• runner/kv_cache_manager.py: in align mode, _maybe_set_align_mamba_num_blocks_override splits the KV budget: MAMBA_ALIGN_STATE_MEM_FRACTION (default 0.3) buys S state slots per GDN layer, attention gets the rest via num_gpu_blocks_override. Both counts are rounded to the sharding divisor and the DP size. get_kv_cache_spec hands vLLM TPUMambaSpec(num_blocks=S) and initialize_kv_cache allocates mamba arrays with that many slots. The HBM probe + divisor logic is factored into _probe_kv_cache_budget, shared with the compact path (unchanged behavior there).
• core/sched/dp_scheduler.py: shard TPUMambaSpec.num_blocks per rank like num_blocks.
• envs.py: MAMBA_ALIGN_STATE_MEM_FRACTION.
• GDN block tables stay group-local, so gdn_attention_op is unchanged.

Known limits: vLLM's "GPU KV cache size" line still counts 2 mamba blocks per group against the shared pool (slightly pessimistic); the deferred-free window is derived from max_concurrent_batches rather than the scheduler's own fence.

Validation (CPU only, no TPU run yet)

Against stock vLLM (58302b4591):

• tests/core/test_mamba_block_pool.py (new, drives vLLM's KVCacheManager as the scheduler does: private pools, batched-group hits, admission gating, step-clocked deferred free, reset propagation), tests/core/test_dp_scheduler.py, and the sizing / spec-wrapping tests in tests/runner/test_kv_cache_manager.py: 63 passed on this branch.
• The initialize_kv_cache allocation test passes on branch wxd-mamba-split-pool-pr3422base (same change on the Enable Prefix Caching for Hybrid Linear-Attention (GDN) Models with DP Support #3422 merge commit, matching the pre-KVCacheTensor.layers vLLM available locally); on main those tests need the newer vLLM API.

TODO before undraft (needs the TPU): tests/e2e/test_mamba_prefix_caching.py for accuracy on cache hits; serve Qwen3.5-397B with attn-DP 4 and check the Mamba-align KV cache: log line and the agentic benchmark at 4 and 8 concurrent groups; sweep MAMBA_ALIGN_STATE_MEM_FRACTION (0.2-0.4); confirm the deferred-free window under --async-scheduling.

🤖 Generated with Claude Code

https://claude.ai/code/session_01Sw2sUivv1hHHcGUY8U4fyS

[github-actions]

Description

Start with a short description of what the PR does and how this is a change from
the past.

The rest of the description includes relevant details and context, examples:

• why is this change being made,
• the problem being solved and any relevant context,
• why this is a good solution,
• some information about the specific implementation,
• shortcomings of the solution and possible future improvements.

If the change fixes a Github issue, please include a link, e.g.,:
FIXES: #123456

Tests

Please describe how you tested this change, and include any instructions and/or
commands to reproduce.

Checklist

Before submitting this PR, please make sure:

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have made or will make corresponding changes to any relevant documentation.