[Bugfix][Mooncake] Fix mixed-TP GQA and indexer cache transfers for MiniMax-M3

[zhewenl]

Purpose

Fix Mooncake heterogeneous-TP transfers for MiniMax-M3's mixed GQA and lightning-indexer caches. With prefill TP4 and decode TP1/DP8, the main GQA block is TP-sharded, while the indexer registers an MLAAttentionSpec and holds the full block on every TP rank. The engine-wide TP-ratio validation rejects that indexer region:

region 60: local=32768, remote=32768, tp_ratio=4

Classify replication per transfer region from the layer's MLA spec type. Validate replicated regions with equal block lengths and reuse the existing whole-block replicated sender plan; retain TP-ratio slicing for the main GQA regions. The producer derives the classification from its local layer specs for aligned local/remote regions, so the wire format does not change. Per-rank KV head counts are not used to infer replication.

The diff contains only transfer-region classification, validation, sender planning, and corresponding tests. Scheduler and receive-failure handling are unchanged.

Related work

• [Bugfix] Fix Mooncake heterogeneous TP with replicated GQA heads #52516 fixes replicated GQA-head mappings across the KV-head replication boundary. This case instead combines sharded GQA and an independently TP-invariant MLA-typed side cache in one model; P=TP4/D=TP1 with four GQA KV heads does not cross that boundary.
• [Bugfix][KV Connector] Mooncake: heterogeneous-TP support for hybrid GDN/Mamba models #53078 covers hybrid GDN/Mamba state layouts. This change covers mixed attention regions without Mamba/GDN state.
• Failure reporting is tracked separately in [Bugfix][Mooncake] Report failed remote KV loads to the scheduler #50984. This PR does not include that work or the lifecycle follow-up proposed for it. It is independent of that PR.
• Fits the mixed-attention transfer gaps discussed in [RFC] [Feature]: PD Disaggregation with MooncakeConnector Roadmap #36948. Open PR searches for Mooncake/MiniMax/heterogeneous TP and the roadmap issue found no duplicate of this region-specific fix.

Tests

On OSS base f4eccdade, 74 passed:

PYTHONPATH=$PWD HF_HUB_OFFLINE=1 /home/zhewen/repos/vllm/.venv/bin/python -m pytest \
  tests/v1/kv_connector/unit/test_mooncake_connector.py \
  tests/v1/kv_connector/unit/test_mooncake_connector_hma.py \
  tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py \
  tests/v1/kv_connector/unit/test_mooncake_stats.py -q
/home/zhewen/repos/vigil/.venv/bin/pre-commit run --files \
  vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py \
  tests/v1/kv_connector/unit/test_mooncake_connector.py

All applicable pre-commit hooks passed, including ruff and mypy. New tests validate mixed GQA/MLA lengths in both TP directions, reject genuinely mismatched regions, check per-layer classification, and exercise sender offsets and the single sending rank for the replicated indexer. The older serving baseline's corresponding suites also pass: 72 passed.

Model evaluation

PASS: TP-only patch, 2026-09-06, vigil Slurm job 12216 on 12 GB300 GPUs. MiniMax-M3, prefill TP4 on one node and decode TP1/DP8 on two nodes; GSM8K 5-shot, greedy, concurrency 512.

Check	Result
GSM8K samples / request errors	1319/1319 / 0
Flexible-extract / strict-match	0.9143 / 0.9143
Prefill / decode HTTP statuses	1319 × 200 / 1319 × 200, no other statuses
Sampled producer transfers	2400 successful, 0 failed
TP-ratio rejection, scheduler assert, rendezvous timeout, local fallback, EngineCore fatal	0 each
Slurm terminal state	COMPLETED

cd /home/zhewen/repos/vigil
.venv/bin/vigil -c recipes/minimax_m3/1p1d_tp4_dep8_mooncake_gsm8k_tp_only.yaml --dry-run </dev/null
.venv/bin/vigil -c recipes/minimax_m3/1p1d_tp4_dep8_mooncake_gsm8k_tp_only.yaml </dev/null

For compatibility with the installed compiled artifacts, serving used 78edd1f70e plus the existing local baseline changes and only this TP patch, in /home/zhewen/repos/vllm-wt-m3-tp-only via PYTHONPATH. The runtime scheduler matched the original local baseline byte-for-byte, and the connector's receive-failure methods were unchanged. The separate router fixes (valid prefill DP rank and P-leg status checking) were present. This is a fresh run of the isolated patch, not the earlier combined TP/failure patch's result.

All-successful P/D access logs, positive transfer metrics, and zero fallback signatures establish real transfer coverage. Transfer counts are sampled metric-interval totals rather than unique requests. Local run evidence is under vigil/logs/1p1d_tp4_dep8_mooncake_gsm8k_tp_only/2026-09-06/20260906_014958/: results, source hashes, actual runtime diff, and unit/lint logs.

AI assistance

AI assistance (Codex) was used to prepare the implementation, regression tests, and validation report. The transfer-region design was isolated from the earlier combined patch developed with Kimi Code/Codex assistance. The reported tests were executed in the shared development environment.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}