Restore the per-group KV pool in the v0.31.0 seam

[arbi-dev]

Restores the per-group KV pool that #30 dropped. That drop was a mistake: #30 called the pools "conflicting with upstream's single-pool design". But the shared pool is exactly what wasted the KV budget on hybrid models, and why the seam had its own pools. turbo-attn ATTRIBUTION §7.3 measured Qwen3.5-0.8B at a 10 GB budget: 5,728 tokens with a shared pool, 355,696 with per-group pools.

Shape

The branch is v0.31.0 + one carry commit (e54368700: #30's carry with the pool restored), then git merge -s ours origin/main so it merges without force-pushing main. The merge leaves the tree untouched (verified equal). Diff vs main: 16 files, +949/−141. The fork now changes 20 vllm/ files vs v0.31.0. #30 changed 12.

How the pool is re-expressed on v0.31

v0.31 allocates one backing buffer and places a [B, H, N, C] view per layer (vllm#51718). On v0.28 the pools were separate tensors; that code can't be replayed. The capability is rebuilt on upstream's structures instead.

• Eligibility (_mamba_pool_eligible): Mamba cache mode none/align; no HiSparse, KV connector, DCP/PCP or hidden-state layers. Those features address blocks across groups by one id space, so they keep upstream's shared pool.
• Planning:
  • The Mamba/GDN state pages are no longer unified with attention pages. Attention layers are still unified among themselves.
  • The state pool holds max_num_seqs × Σ per-request blocks per group + 1 blocks.
  • That pool plus the sampler reserve (VLLM_SAMPLER_RESERVE_MIB) comes off the budget before attention is planned. So auto-fit, num_gpu_blocks_override, the admission check and the per-rank re-plan all plan the attention pool only.
  • Both pools are regions of the one buffer, each laid out by upstream's own placement (_layout_kv_cache_tensors, factored out unchanged). The regions never alias.
  • New fields: KVCacheConfig.mamba_pool_num_blocks / mamba_pool_offset / group_num_blocks(), and KVCacheTensor.num_blocks.
• Scheduler:
  • The coordinator gets a second BlockPool for the state groups.
  • Admission checks each pool against its own free blocks (get_num_blocks_to_allocate_per_pool, KVCacheManager._fits). Watermark and reservations count main-pool blocks.
  • Usage reports the fuller pool, and events come from both pools.
  • Deferred frees already route by KVCacheBlock.pool upstream; a test confirms it.
• Worker:
  • The two pools' block ids overlap. CoW copies of state blocks carry ids shifted past num_blocks, and copy_kv_cache_blocks_inplace applies them only to views inside the state region.
  • MRV2 profiling dummy contexts (set_dummy_context) and both warmups number each pool's blocks separately.

Per-hunk rows (changes vs #30)

file	what	verdict
v1/core/kv_cache_utils.py	eligibility, split planner, region helper, grouping without state/attention unification, capacity excludes the state pool, get_kv_cache_configs fixed-cost accounting	still needed, re-expressed on v0.31 placement
v1/kv_cache_interface.py	pool fields on KVCacheConfig / KVCacheTensor	still needed
v1/core/kv_cache_coordinator.py	second BlockPool, per-pool block needs	still needed
v1/core/kv_cache_manager.py	per-pool admission, usage, events, CoW id shift	still needed
v1/worker/utils.py, gpu_model_runner.py, gpu/model_runner.py	CoW copies per pool	still needed (CoW is new upstream since v0.28)
v1/worker/gpu/input_batch.py, gpu/warmup.py	per-pool dummy/warmup block ids	still needed (MRV2 is the default runner in v0.31)
envs.py	VLLM_SAMPLER_RESERVE_MIB	still needed (restored)
PROTOCOL.md	new section	doc

Not restored, with the reason:

• Padded-page block fill (_fill_padded_page_block_size): it only worked around state/attention page unification, which the split removes.
• Drain hook / on_kv_manager_created and the zeroer clamp: no turbo-attn consumer.

Composite (smart-mix) layout

aggregated_layer_count is unchanged from #30: every fused layer's view is the same dense page region of the attention pool. test_mamba_pool.py::test_composite_attention_shares_one_dense_page_beside_the_state_pool checks it.

The plugin side adds one fix (turbo-attn PR, linked below). It puts each composite layer's region view into runner.kv_caches in place of the composite view. Otherwise a prefix-cache CoW copy, which works by page index, would move other layers' bytes in the layer-major composite.

Verification

CPU only, in vllm/vllm-openai:v0.31.0 with the 20 files overlaid, on 10.1.0.202.

• tests/v1/core/test_mamba_pool.py (new): 12 passed. It covers:
  • pool sizing, and 16.9× the attention blocks vs the shared pool on a bf16 Qwen3.5-like layout;
  • no aliasing, and the allocation bound;
  • the eligibility exclusions;
  • each pool allocating from itself, and admission refusing when the state pool is full even with attention blocks free;
  • deferred frees returning to their own pools, and usage;
  • CoW copy routing;
  • composite on the split layout.
• Seam tests: all pass. test_kv_seam_invariants, test_aggregated_layer_count, test_mla_wrapper_selection, and the hook tests in test_gpu_worker.
• Upstream suites: tests/v1/core, tests/v1/worker and tests/engine/test_arg_utils.py, stock v0.31.0 vs overlay. The overlay shows no new failures and 18 more passes (the seam's own tests). The remaining failures are CPU-container device-inference errors, identical on stock. The one per-run difference (test_prefix_cache_default) fails identically on both when run alone.
• turbo-attn vLLM tests (with the plugin PR): 2099 passed. The 2 failures are already stale on turbo-attn main, and the 7 errors are GPU fixtures.

Not yet run

GPU validation of the engine image. Hybrid capacity needs measuring on a real boot (Qwen3.5-0.8B tkv: blocks reported vs v0.28 image).

🤖 Generated with Claude Code