[0z5a]

Overview

In Mamba prefix-caching (mamba_cache_mode="align") mode, intermediate prefill chunks end on the block grid. With the large blocks hybrid models get after page-size alignment, a sizeable part of --max-num-batched-tokens can silently stay idle on every prefill step. This PR warns at startup when that happens, suggests a block-aligned value, and documents the interaction. It does not change scheduling behavior.

Claims

• VllmConfig.validate_block_size() now logs a warning when align-mode chunking would leave more than a quarter of the prefill budget idle. That method runs once the block size is final and already holds the align-mode constraints. The warning suggests the next block multiple plus decode headroom.
• The hybrid section of the automatic prefix caching guide explains how the token budget interacts with the block grid.
• For GLM-5.3-Flash on H100 64GB, setting the budget as the warning suggests improves 1/P90 (E2E seconds per output token) by 1.96× on a full agentic-coding trace replay (data below).

Validation

Unit tests (CPU-only environment, upstream main a9e90b9d + this change):

pytest tests/test_config.py -k mamba_align
# 3 passed: the 2 new cases + existing test_validate_mamba_align_subblock_prefill

The full tests/test_config.py was run with and without this change in the same offline CPU environment. Both runs had the same 106 failures (tests that need the HF hub or a GPU) and 7 skips. Without the change, 339 tests passed. With it, 341 passed: the same tests plus the 2 new ones. No new failures.

ruff check and ruff format are clean.

Why it matters: serving data (GLM-5.3-Flash, vLLM 0.31.0, 2 × 4 H100 64GB, TP8 + EP, MTP k=3, --max-model-len 1048576, fp8 KV cache). Metric: 1/P90 of per-request (E2E s / output tokens), 16 concurrent sessions replaying two-turn agentic-coding traces with ~230K-token prompts. GLM-5.3-Flash serves with 1152-token blocks here, and the API-server default budget on GPUs below 70 GiB is 2048.

--max-num-batched-tokens	prefill tokens / step	1/P90	TTFT P90	vs 2048
2048 (default)	1152	3.47 (n=8)	522 s	—
3072	2304	5.52	318 s	+59%
4096	3456	5.47 (n=5)	318 s	+58%
6144	5184	5.56	310 s	+60%
8192	6912	5.23 (n=4)	335 s	+51%

• The 3072–8192 rows also set --max-cudagraph-capture-size 128 --language-model-only. Those two flags alone gave +2% (3.53, n=2).
• 12288 and 16384 no longer leave enough KV cache for a 1M context on 64 GB GPUs.
• In a full trace-replay run (C16, 3600 s + 1800 s grace, seed 42; 238 / 353 profiled requests), 4096 vs 2048 measured 1/P90 0.642 vs 0.327 (1.96×) and TTFT P90 233.6 s vs 442.6 s. Both arms had 0 errors and passed the same output-quality checks.

This serving data was collected on vLLM 0.31.0, setting the flag by hand. The warning itself is verified by the unit test; it was not exercised in a GPU server run on main.

Details

• Root cause: Scheduler._mamba_block_aligned_split floors intermediate chunk ends to block_size so the cached SSM state stays block-aligned. A long prompt therefore prefills floor(budget / block) * block tokens per step, and decode tokens in the same step can lower this further. With 1152-token blocks and a 2048 budget, 44% of the budget goes unused, and nothing in the logs says so.
• Why warn instead of auto-adjusting: the final hybrid block size is only known after the workers align it with the Mamba page size. By then workers have already sized buffers from max_num_batched_tokens (profiling run, CUDA graph and compile ranges, backend workspaces), so raising the budget at that point would be unsafe.
• Scope:
  • The check is skipped when --long-prefill-token-threshold is set, because the user is then deliberately tuning chunk sizes.
  • It is also skipped when the block is larger than the budget. In that case the existing sub-block path keeps chunks running, so the budget is not idled.
  • The 25% threshold keeps it quiet for the 80 GB defaults: 8192 with 1152-token blocks keeps 6912 tokens, 84% of the budget.
• Related open PRs: [Bugfix][V1] Use the Mamba cache group's block size for align-mode chunk splitting #54076, [Bugfix][Core] Keep Mamba-aligned prefill chunks out of the MTP lookahead window #61845, [Bugfix] Stop mamba align prefill at the replay boundary #50409 and kv-cache: honor scheduler lookahead in mamba align-mode allocation #57605 change where align-mode chunks end. None of them surfaces or documents the idle budget, so this is not a duplicate. If [Bugfix][V1] Use the Mamba cache group's block size for align-mode chunk splitting #54076 lands, this check should read the same block size as the split. Today both read cache_config.block_size.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

Draft: opened for early feedback; the submitter will review every line before marking it ready.

🤖 Generated with Claude Code

[mergify]

Documentation preview: https://vllm--61878.org.readthedocs.build/en/61878/

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.