PR #40384: [Bugfix] Exclude O(1) Mamba groups from hybrid KV cache token capacity

• Repository: vllm-project/vllm
• Author: jhsmith409
• State: open | merged: False
• Link: https://github.com/vllm-project/vllm/pull/40384

Description (problem / solution / changelog)

Summary

On hybrid attention + Mamba models (Qwen3-Next, Qwen3.5/3.6 MoE hybrids, RecurrentGemma, Jamba, Zamba2, Nemotron-H, …), the reported GPU KV cache token capacity and the scheduler's max_num_kv_tokens are deflated by the number of Mamba groups, which in the default mamba_cache_mode='none' (and 'align') pre-reserve a fixed number of blocks and do not scale with sequence length.

Both _report_kv_cache_config() (vllm/v1/core/kv_cache_utils.py) and Scheduler.__init__ (vllm/v1/core/sched/scheduler.py) currently compute per-token capacity as:

num_tokens = num_blocks // len(kv_cache_config.kv_cache_groups) * min_block_size

For a typical hybrid with one attention group and N Mamba groups, that's off by a factor of (1 + N) / 1 — 2× understatement for the common case, 4× for Nemotron-H-style 1 attn + 3 mamba groups. The max_num_kv_tokens number is what sizes the routed_experts buffer for MoE and what the scheduler believes is its budget; getting this wrong shows up as (a) misleading boot-time logs and (b) over-conservative scheduling of concurrent requests on the very models (hybrid MoE) where extra concurrency is the whole point.

Fix

• Factor the filter into a tiny helper token_capacity_kv_cache_groups(vllm_config, kv_cache_config) in kv_cache_utils.py that returns only the groups that scale with sequence length (attention always, Mamba only when mamba_cache_mode == 'all').
• Use that helper in both _report_kv_cache_config and Scheduler.__init__.
• Fall back to all groups if the filter would produce an empty list (preserves dense-model and Mamba-only paths).

The helper is exported (no leading underscore) because scheduler.py imports it; if the maintainers would rather keep it scheduler-local or inline it, happy to rewrite.

Why this is not duplicating an existing PR

Checked on 2026-04-20:

• gh pr list --repo vllm-project/vllm --state open --search \"max_num_kv_tokens\" → none.
• gh pr list --repo vllm-project/vllm --state open --search \"mamba kv_cache_groups token\" → none.
• Referenced in https://github.com/vllm-project/vllm/issues/40124 (patch 9) as still not upstream.

Test plan + results

python -m pytest tests/v1/core/test_kv_cache_utils.py -v

No existing test exercises the filter directly; I'll follow up with a small unit test in a separate commit once PR feedback lands (or now, if reviewers prefer). Syntax check (python -m py_compile) is clean; ruff check/ruff format were not available in my local sandbox but the edits follow the surrounding style.

End-to-end verification on our runtime stack (cu130-nightly + TurboQuant hybrid overlay + RedHatAI/Qwen3.6-35B-A3B-NVFP4, turboquant_k8v4, max_model_len=8192, max_num_seqs=8, --gpu-memory-utilization=0.85, torch.compile + cudagraph):

• Before: INFO kv_cache_utils.py:1363] GPU KV cache size: 143,936 tokens
• After: I'll reply in a follow-up comment with the log delta from a re-run on this branch; expect a clean 2× jump on 1 attn + 1 mamba group.

AI-assist disclosure (per AGENTS.md)

Change was drafted with help from Claude (Anthropic); human submitter reviewed every line end-to-end and understands the hybrid KV cache group semantics. Original bug identification and filter design credit to @Sandermage — ref his issue #40124 tracking table (patch 9) and the ai-jz/vllm#1 approach he references. Co-authored-by: trailers included.

Changed files

• tests/v1/core/test_kv_cache_utils.py (modified, +80/-0)
• vllm/v1/core/kv_cache_utils.py (modified, +30/-6)
• vllm/v1/core/sched/scheduler.py (modified, +5/-7)

PR #37118: [Bugfix] out-of-bounds error for routed experts capture

• Repository: vllm-project/vllm
• Author: HollowMan6
• State: open | merged: False
• Link: https://github.com/vllm-project/vllm/pull/37118

Description (problem / solution / changelog)

Purpose

This PR fixes an out-of-bounds error in routed expert capture when enable_return_routed_experts=True is used with hybrid KV cache groups.

Traceback (most recent call last):
  File "/usr/local/lib/python3.12/dist-packages/vllm/v1/executor/multiproc_executor.py", line 927, in worker_busy_loop
    output = func(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_contextlib.py", line 124, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/vllm/v1/worker/gpu_worker.py", line 756, in sample_tokens
    return self.model_runner.sample_tokens(grammar_output)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_contextlib.py", line 124, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/vllm/v1/worker/gpu_model_runner.py", line 4060, in sample_tokens
    capturer.save_captured_experts(indices=self.slot_mapping)  # noqa
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py", line 217, in save_captured_experts
    self._host_buffer_view[indices, :, :] = data
    ~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^
IndexError: index 699312 is out of bounds for axis 0 with size 699312

The routed-experts side buffer was sized with:

(num_blocks // num_groups) * min_block_size

on both the worker and scheduler sides.

That formula is only a coarse aggregate/token-capacity estimate for hybrid KV cache layouts. However, routed expert capture/readback indexes the buffer with the selected attention group's actual slot_mapping, whose address space is based on that attention KV group directly.

As a result, in hybrid/padded KV-cache configurations, the routed-experts buffer can be smaller than the valid range of slot_mapping, which leads to out-of-bounds writes/reads.

Use the routed-experts attention group's full KV address space to size the buffer consistently on both sides:

kv_cache_config.num_blocks * attn_group.kv_cache_spec.block_size

Routed expert capture is indexed by the attention group's slot_mapping, so the auxiliary buffer must match the full addressable range of that mapping. Sizing it from the specific attention group keeps the writer and reader aligned

• fixes crashes when enable_return_routed_experts=True
• no behavior change when routed expert return is disabled
• no change to model weights, routing logic, or sampling semantics

Test Plan

End to end tests

Test Result

Now error is gone

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

Changed files

• vllm/v1/core/sched/scheduler.py (modified, +5/-9)
• vllm/v1/worker/gpu_model_runner.py (modified, +5/-9)