[Bugfix][V1] Support heterogeneous KV page sizes in KVBlockZeroer

[Fangzhou-Ai]

Purpose

KVBlockZeroer assumed every attention layer shares one physical page size:

assert page_size_el == cur_page_el, (
    f"Non-uniform page sizes: {page_size_el} vs {cur_page_el}"
)

Any model that mixes page sizes across its attention specs therefore fails at startup,
inside initialize_from_config, before serving a single request.

Reproducer

MiniMax-M3 on ROCm, TP4, --kv-cache-dtype fp8:

AssertionError: Non-uniform page sizes: 8192 vs 4096
  vllm/v1/worker/utils.py:155 in KVBlockZeroer.__init__
  <- gpu_model_runner.py:1136 _init_kv_zero_meta
  <- gpu_worker.py:660 initialize_from_config
RuntimeError: Engine core initialization failed.

Two independent properties of the model combine to make this unavoidable:

1. The page sizes differ by construction. MiniMaxM3Indexer.get_kv_cache_spec
   registers a key-only MLAAttentionSpec for the indexer side cache (1 head,
   index_head_dim), alongside the main K+V attention cache. 8192 vs 4096 elements.
2. KV zeroing is unavoidably on. needs_kv_cache_zeroing is
   has_mamba_layers or has_mixed_precision_kv_cache. M3 has no Mamba layers, but with
   --kv-cache-dtype fp8 the main cache is fp8 while the indexer side cache is bf16 —
   mixed precision, so zeroing engages and KVBlockZeroer is constructed.

Neither is opt-in, so MiniMax-M3 with an fp8 KV cache cannot start on current main.
This is independent of VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT / the AITER sparse-PA path —
get_kv_cache_spec has no dependence on either.

EAGLE3 draft layers with a different GQA width under TP are a second instance of the
same shape.

Fix

Group segment addresses by page size (dict[int, list[int]]) instead of asserting a
single one, store one metadata entry per distinct page size, and launch the zeroing
kernel once per entry. A single-page-size model keeps a one-element list and identical
behavior — same grid, same kernel arguments.

Test Plan and Result

tests/v1/worker/test_kv_block_zeroer.py updated for the new _meta shape (list of
tuples rather than a bare tuple); the in-flight-copy behavior it asserts is unchanged.

End-to-end on MI355X, MiniMax-M3-MXFP4, TP4, --kv-cache-dtype fp8:

• Before: engine dies in initialize_from_config, as above.
• After: init engine (profile, create kv cache, warmup model) took 28.56 s, server
  serves normally.

Lint: ruff check and ruff format --check (v0.14.0, as pinned in
.pre-commit-config.yaml) clean. pre-commit itself could not run — no pre_commit
module in the venv — so the hooks were run manually.

Duplicate check

gh pr list --repo vllm-project/vllm --state open --search over "KVBlockZeroer",
"Non-uniform page sizes", "needs_kv_cache_zeroing", "kv cache zeroing page size",
"mixed precision kv cache zero", plus an issue search. Nothing covers this assertion.
The closest, vllm-project#36617, is prefix caching for hybrid models with non-uniform page sizes — a
different code path (scheduler-side allocation, not the worker-side zeroer).

Notes

Draft pending the full benchmark sweep this unblocks. AI assistance (Claude Code) was
used for this change.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

[Fangzhou-Ai]

Closing: this fixes the symptom in shared vLLM infrastructure, and the preference is to keep MiniMax-M3 changes inside the model.

The cause is a MiniMax-M3 configuration that should not be built in the first place: a bf16 indexer side cache beside an fp8 main KV cache makes the KV cache mixed-precision, which turns on block zeroing, and the zeroer then rejects the two page sizes (8192 vs 4096).

NVIDIA already avoids this by construction — minimaxm3_fp8_b300.sh sets --attention-config.indexer_kv_dtype "fp8" alongside --kv-cache-dtype fp8, so precision is uniform and KVBlockZeroer is never constructed. The same pairing works on ROCm once #14 makes indexer_kv_dtype=fp8 reachable there, so no base-file change is needed for the target configuration.

Known gap left open: setting --kv-cache-dtype fp8 without an fp8 indexer still fails on the opaque AssertionError: Non-uniform page sizes rather than a clear message. Recorded here rather than fixed.