[zyongye]

Overview

Preserve KV cache zeroing when mixed-precision UniformTypeKVCacheSpecs groups are reduced for the scheduler. Resolve the zeroing decision across workers so pipeline stages agree even when only one stage has mixed precision. Split out of #63232 (MiniMax-M3.1-Flash); this change remains independent of the other split PRs.

Claims

• Scheduler projection preserves the decision to zero newly allocated attention blocks.
• All PP workers initialize zeroing when any worker requires it. This prevents MRV2 from receiving zeroing work on a stage without a zeroer.
• Uniform-precision attention-only configurations continue to skip zeroing.

Validation

The same regression tests were run on original PR head 998f165a2ab0ce47f63ff280b96a600901c7d914 and this revision.

Check	Original PR head	This revision
Config matrix: uniform, mixed first stage, mixed second stage	1 passed, 2 failed (False is True on a uniform worker)	3 passed
Related scheduler, Mamba, and PP config tests	Not rerun	7 passed
Qwen3-0.6B, MRV2, PP=2, FP8 KV, BF16 layer 0	Worker_PP1 asserted that kv_block_zeroer is not None; first completion returned HTTP 500	1 passed; PP=2 and PP=1 first-token top-5 logprobs match at the helper's precision
Pre-commit on the seven changed files	Not rerun	Passed

CPU commands (with HF_HUB_OFFLINE=1 PYTHONPATH=$PWD CUDA_VISIBLE_DEVICES=''):

.venv/bin/python -m pytest tests/v1/core/test_kv_cache_utils.py \
  -k test_generate_scheduler_kv_cache_config_keeps_zeroing -q
.venv/bin/python -m pytest tests/v1/core/test_kv_cache_utils.py \
  tests/v1/core/test_kv_cache_planning.py \
  -k 'test_generate_scheduler_kv_cache_config or test_wrapped_mamba_group_requires_block_zeroing or test_get_kv_cache_configs_pp_sharding' -q

GPU command, run through Slurm with two visible GB200 GPUs and cached real weights (before: job 40854, 1 failed; after: job 40859, 1 passed in 171.17 s):

HF_HUB_OFFLINE=1 PYTHONPATH=$PWD CUDA_VISIBLE_DEVICES=0,1 \
  .venv/bin/python -m pytest tests/distributed/test_pipeline_parallel.py \
  -k test_pp_mixed_precision_kv_cache -v -s --tb=short

The regression compares PP=2 against PP=1 using the existing generate_close helper: greedy first-token top-5 logprobs rounded to two decimal places. Both revisions use the same existing native build. This validates the Python configuration/zeroing path with Qwen, not MiniMax native operators or model accuracy on a benchmark. No performance claim is made.

The E2E test is included in the existing pipeline/context parallelism CI steps, with dependencies added for the changed core files. NVIDIA coverage is in the existing optional nightly step; AMD coverage uses its existing PR step. AMD was not run locally.

Details

generate_scheduler_kv_cache_config reduces each UniformTypeKVCacheSpecs to one layer spec. Recomputing mixed precision after that reduction loses information. Saving only worker 0's mixed-precision result also leaves PP workers inconsistent: a uniform stage skips zeroer initialization although it receives the scheduler's new_block_ids_to_zero.

has_mixed_precision_kv_cache remains a local fact. needs_kv_cache_zeroing becomes a stored execution decision, initialized from each worker's full configuration. The planner ORs that decision across workers after final replanning, and scheduler projection preserves it through deepcopy. The tradeoff is that uniform stages also zero new attention blocks when another stage requires zeroing.

This extends the metadata-preservation fix in #49623 to intra-group precision and PP worker agreement. It is not a duplicate of the open fixes for transfer destinations/tails (#62313, #55085), backend-specific block reuse (#58716), or initial allocation zero fill (#59432).

AI assistance was used: Claude for the original split, and OpenAI Codex for this follow-up implementation, regression tests, and validation.

---

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.

[claude]

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

_{Tip: disable this comment in your organization's Code Review settings.}

[mergify]

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

[ivanium]

/ci run

[zyongye]

/ci retry

[github-actions]

✅ Queued 1 failed job(s) for retry in Buildkite CI #98505.

[zyongye]

/ci retry

[github-actions]

✅ Queued 1 failed job(s) for retry in Buildkite CI #98505.

[zyongye]

/ci run

[github-actions]

✅ Triggered Buildkite CI #98807 for commit aa59e0d3b48b.

[zyongye]

/ci retry

[github-actions]

✅ Queued 1 failed job(s) for retry in Buildkite CI #98807.

[zyongye]

/ci run

[github-actions]

✅ Triggered Buildkite CI #98993 for commit 76feeaafa709.

[zyongye]

/ci retry

[github-actions]

✅ Queued 1 failed job(s) for retry in Buildkite CI #98993.

[zyongye]

/ci retry

[github-actions]

✅ Queued 1 failed job(s) for retry in Buildkite CI #98993.

[zyongye]

/ci run

[github-actions]

✅ Triggered Buildkite CI #99089 for commit 377cb33faa46.

[mergify]

Hi @zyongye, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[zyongye]

/ci run

[github-actions]

✅ Triggered Buildkite CI #99284 for commit c17d8af8724e.

[zyongye]

/ci retry

[github-actions]

✅ No failed, timed-out, or expired jobs need retrying: https://buildkite.com/vllm/ci/builds/99284

[khluu]

Kevin's review agent: approve for CI changes only

[khluu]

Kevin's review agent: approve for CI changes only