[Bugfix][MooncakeStore] Fix finish-time save crash on hybrid models by zhewenl · Pull Request #54643 · vllm-project/vllm

zhewenl · 2026-08-31T22:22:35Z

What & why

Fixes a crash in MooncakeStoreConnector's save path on hybrid (multi KV-group) models such as Kimi-K3 (MLA + GDN):

worker.py: KVCacheStoreSendingThread → _handle_request
  group_blocks = block_ids_per_group[g_idx]
IndexError: tuple index out of range

Root cause

Under MultiConnector, update_state_after_alloc forwards the real num_external_tokens only to the connector that supplied them; every other connector is called with 0 (multi_connector.py). The store connector used that as a condition to record an empty block-table tuple in _unfinished_requests:

local_block_ids: tuple[list[int], ...] = ()
if num_external_tokens > 0:   # never true for the non-chosen connector
    local_block_ids = ...

The finish-time save path (scheduler.py, requests that left the schedule with a load spec) rebuilds the request tracker from that stored entry, so the worker's per-group save loop indexed into an empty tuple — on hybrid models with N>1 KV groups this raises IndexError per affected request. Each crash silently drops that request's store save (store pool stays under-populated); single-group models never hit it.

Observed live on Kimi-K3 PD + store offload (GB300): 997 crashes in one eval run while ~7 GB of saves still succeeded via the in-schedule path.

Fix

store/scheduler.py::update_state_after_alloc: always record the full-width per-group block table (drop the num_external_tokens > 0 condition). The stored tuple's only consumer is the finish-time save, which needs it regardless of who supplied external tokens.
store/worker.py::_handle_request: skip (with a warning) any group the request metadata does not cover, so a malformed entry can never kill the sending thread again.

Tests

pytest tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py \
       tests/v1/kv_connector/unit/test_mooncake_store_worker.py \
       tests/v1/kv_connector/unit/test_mooncake_store_connector.py -q
# 183 passed (GB300, aarch64, CUDA 13.3), incl. new regression test
# test_update_state_records_blocks_without_external_tokens

Live-verified on Kimi-K3 1P1D (TP8) + MultiConnector(NixlConnector, MooncakeStoreConnector) with a deliberately tiny GPU KV pool (--num-gpu-blocks-override 24) to force store read/write traffic: gsm8k accuracy holds at 0.967 (0 request errors) and the IndexError is gone.

AI assistance

Developed with AI assistance (Kimi Code); a human reviewed the diff and ran the tests above.

MultiConnector queries every child but selects only the first cache hit. A losing MooncakeStore lookup therefore leaves a LoadSpec with can_load=False. The pending-load metadata path treated any LoadSpec as actionable and converted the rejected load into a save, using an empty block table on hybrid models and crashing the send thread.\n\nOnly build pending metadata for an accepted load. Normal stores continue to get their block tables from SchedulerOutput after the request is scheduled. Signed-off-by: zhewenl <zhewenl@users.noreply.github.com>

Backport vllm-project#54643 at 8ad6f35. A speculative load rejected by MultiConnector must not create a save outside the scheduled block-table snapshot. Reproduced the exact missing-table assertion before the fix. 427 scheduler, offload, and metric regressions plus the focused override suite pass. Co-authored-by: Codex <noreply@openai.com> Signed-off-by: Cam Quilici <cjquilici@gmail.com>

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.

chatgpt-codex-connector · 2026-09-03T21:06:55Z

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-09-03T21:08:35.397033Z	`8ad6f35`	Draft marked ready

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

coderabbitai · 2026-09-03T21:08:26Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Team

Run ID: ca2e2913-12f7-40ea-a126-3c25407e4152

📥 Commits

Reviewing files that changed from the base of the PR and between 9cc7793 and a4907a3.

📒 Files selected for processing (2)

tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py

🚧 Files skipped from review as they are similar to previous changes (2)

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py
tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py

Included review availability: Your plan provides up to 10 included reviews per hour; 9 remain after this review.

📝 Summary

Summary by CodeRabbit

Bug Fixes
- Prevented rejected speculative loads from being issued by an unselected connector.
- Cleared stale pending-load state and allocated blocks when another connector is chosen.
- Ensured only valid, accepted load requests produce connector metadata.

Walkthrough

The Mooncake store scheduler now skips pending loads rejected by MultiConnector. A regression test verifies that the losing connector clears allocated blocks, emits no metadata, and removes stale request state.

Changes

Pending load rejection handling

Layer / File(s)	Summary
Scheduler guard and regression test `vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py`, `tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py`	`build_connector_meta` now requires `load_spec.can_load` before issuing a pending load. The test verifies cleanup for a rejected `LoadSpec`.

Estimated code review effort: 2 (Simple) | ~10 minutes

Merge Risk: ⚪ Minimal · up to a4907

The scheduler now drops rejected pending-load work and clears its stale state, preventing non-selected connectors from emitting metadata or retaining allocations. The targeted regression coverage and reported test results indicate no remaining merge-readiness risk.

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Title check	⚠️ Warning	The title describes fixing a finish-time save crash on hybrid models, but the summarized changes only tighten pending-load handling for rejected non-chosen connectors and add a regression test. The ti…	Update the title to describe the actual change, such as "[Bugfix][MooncakeStore] Skip rejected pending loads from non-chosen connectors".
Docstring Coverage	⚠️ Warning	Docstring coverage is 75.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 4 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description check	✅ Passed	The description addresses the MooncakeStore crash, MultiConnector behavior, scheduler state, worker handling, and regression testing. It is related to the stated objectives, although parts of it do no…
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

Full details: Title check

Explanation

The title describes fixing a finish-time save crash on hybrid models, but the summarized changes only tighten pending-load handling for rejected non-chosen connectors and add a regression test. The title does not accurately identify the primary change.

Fix all pre-merge checks with AI

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

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

ivanium

Thanks for the fix!

coderabbitai · 2026-09-07T04:31:56Z

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

zhewenl · 2026-09-07T04:36:08Z

/ci run

github-actions · 2026-09-07T04:36:25Z

✅ Triggered Buildkite CI #87487 for commit a4907a3ee071.

…llm-project#54643) Signed-off-by: zhewenl <zhewenl@users.noreply.github.com> Co-authored-by: zhewenl <zhewenl@users.noreply.github.com> Signed-off-by: Jyotirmoy Roy <jyotirmoyroy649@gmail.com>

mergify Bot added bug Something isn't working kv-connector labels Aug 31, 2026

zhewenl force-pushed the mooncake-store-finish-save-fix branch from 1fe0f87 to 9132c71 Compare August 31, 2026 22:32

zhewenl force-pushed the mooncake-store-finish-save-fix branch from 9132c71 to 8ad6f35 Compare August 31, 2026 23:47

zhewenl mentioned this pull request Sep 2, 2026

[Core][KV Connector] Resolve connector block tables for every scheduled request #54853

Merged

4 tasks

cquil11 mentioned this pull request Sep 3, 2026

Validate vLLM P/D cache-source metrics on GB300 SemiAnalysisAI/InferenceX#2797

Open

zhewenl marked this pull request as ready for review September 3, 2026 21:06

zhewenl requested review from ApostaC, NickLucche, ivanium, orozery and xuechendi as code owners September 3, 2026 21:06

claude Bot reviewed Sep 3, 2026

View reviewed changes

ivanium approved these changes Sep 7, 2026

View reviewed changes

Merge branch 'main' into mooncake-store-finish-save-fix

a4907a3

ivanium enabled auto-merge (squash) September 7, 2026 04:31

github-actions Bot added the ready ONLY add when PR is ready to merge/full CI is needed label Sep 7, 2026

ivanium merged commit 34b1e9f into vllm-project:main Sep 7, 2026
28 checks passed

Woolgathererer mentioned this pull request Sep 9, 2026

[Bugfix] Drop Mooncake store saves missing from the block-state snapshot #55066

Closed

Navigation Menu

Uh oh!

Uh oh!

[Bugfix][MooncakeStore] Fix finish-time save crash on hybrid models - #54643

[Bugfix][MooncakeStore] Fix finish-time save crash on hybrid models#54643
ivanium merged 2 commits into
vllm-project:mainfrom
zhewenl:mooncake-store-finish-save-fix

zhewenl commented Aug 31, 2026

Uh oh!

claude Bot left a comment

Uh oh!

chatgpt-codex-connector Bot commented Sep 3, 2026 •

edited

Loading

Uh oh!

coderabbitai Bot commented Sep 3, 2026 •

edited

Loading

Summary by CodeRabbit

Walkthrough

Changes

❌ Failed checks (2 warnings)

Uh oh!

ivanium left a comment

Uh oh!

coderabbitai Bot commented Sep 7, 2026

Uh oh!

zhewenl commented Sep 7, 2026

Uh oh!

github-actions Bot commented Sep 7, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Uh oh!

Conversation

zhewenl commented Aug 31, 2026

What & why

Root cause

Fix

Tests

AI assistance

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

chatgpt-codex-connector Bot commented Sep 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codex Review Summary

Uh oh!

coderabbitai Bot commented Sep 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Walkthrough

Changes

❌ Failed checks (2 warnings)

Uh oh!

ivanium left a comment

Choose a reason for hiding this comment

Uh oh!

coderabbitai Bot commented Sep 7, 2026

Uh oh!

zhewenl commented Sep 7, 2026

Uh oh!

github-actions Bot commented Sep 7, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

chatgpt-codex-connector Bot commented Sep 3, 2026 •

edited

Loading

coderabbitai Bot commented Sep 3, 2026 •

edited

Loading