[Bugfix][KV Connector] MooncakeStore: exclude non-prefix-cacheable (QSA ring) groups; fix align-mode check by zhewenl · Pull Request #55027 · vllm-project/vllm

zhewenl · 2026-09-02T22:12:22Z

Purpose

Fix #54726: Qwen3.8-Flash-Next with prefix caching and MooncakeStoreConnector fails because its QSA ring scratch cache is transferable, but is not valid hash-addressed prefix data.

The QSA cache returns a CircularBufferSpec; CircularBufferSpec.prefix_cacheable is explicitly False. Its capacity is 4 without speculative tokens and 8 with MTP=3, while aligned attention/Mamba groups use block size 800. MooncakeStore previously built its hash/store/lookup layout from every enable_kv_transfer=True group, so the ring reached a divisibility check intended only for storable prefix groups.

A second issue was hidden behind that failure: MooncakeStore inferred Mamba align mode by comparing each Mamba block size with the globally rewritten cache_config.block_size. With heterogeneous groups, that global value can differ from the resolved Mamba group size even when mamba_cache_mode is correctly "align".

Fix

Add KVCacheConfig.prefix_cacheable_group_ids and prefix_cacheable_groups beside the existing transfer projection.
Extend KVCacheBlocks.get_block_ids(group_ids=...) so scheduler consumers can request group projections without adding a generic selector to KVCacheConfig.
Use only the prefix-cacheable transfer projection in MooncakeStore validation, coordinator, scheduler, and worker tensor-payload registration.
Preserve scheduler-to-store group remapping so non-storable groups such as the QSA ring are skipped without shifting store group IDs.
Validate Mamba support directly with spec.mamba_cache_mode == "align".

This deliberately does not change transfer_group_ids, transfer_groups, or select_transfer_block_ids. NIXL and other P/D connectors continue to transfer all enable_kv_transfer=True groups, including the QSA ring scratch state.

Duplicate-work check

Checked issue #54726 and its comments, plus open PR searches for 54726 in:body and Mooncake QSA prefix cacheable. The only matching open PR is this PR (#55027).

Tests

.venv/bin/python -m pytest \
  tests/v1/core/test_kv_cache_utils.py \
  tests/v1/kv_connector/unit/test_mooncake_store_connector.py \
  tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py \
  tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py \
  tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py \
  tests/v1/kv_connector/unit/test_mooncake_store_worker.py -q

Result on the current head: 412 passed, 14 warnings in 146.17s.

.venv/bin/python -m pytest \
  tests/v1/kv_connector/unit/test_nixl_connector_hma.py -q

Result: 77 passed. The remaining GPU integration case failed during connector initialization because NIXL is unavailable in the test environment; it did not reach block selection.

pre-commit run --files <all changed Python files>

Result: all applicable hooks passed, including Ruff and mypy.

Model evaluation

GSM8K 5-shot, thinking disabled, temperature 0, concurrency 512

Results:

Pass	flexible-extract	strict-match
First	0.9682	0.9689
After local prefix-cache reset	0.9697	0.9704

All 2638 requests returned a non-empty answer, all finished with stop, and there were no API errors. After the reset, Mooncake performed 1416 Get calls covering 7092 keys / 72.59 GB, with zero failed keys and zero transfer errors. The original block-size assertion did not occur.

AI assistance

AI coding agents (OpenAI Codex and Claude) assisted with analysis, implementation, testing, and review. The human submitter reviewed the changes. The commit includes DCO sign-off and AI co-author trailers.

coderabbitai · 2026-09-02T22:12:31Z

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

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

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-03T00:01:45Z

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-03T00:05:40.294673Z	`f6db171`	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.

mergify · 2026-09-07T08:16:41Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @zhewenl.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

ivanium

Nice fix. Only leaving some simplification suggestions. connector.py and scheduler.py can perhaps be simplified a bit too

ivanium · 2026-09-09T01:17:34Z

+        store_layer_names = {
+            layer_name
+            for group in self._kv_cache_groups
+            for layer_name in group.layer_names
+        }
        seen_storage_ptrs: set[int] = set()
        cache_tensors: list[torch.Tensor] = []

-        for cache in kv_caches.values():
+        for layer_name, cache in kv_caches.items():
+            if layer_name not in store_layer_names:
+                continue


Not sure if we really need this.

We still need this, otherwise non-cacheable groups like QSA ring will still be registered

mergify · 2026-09-12T05:59:21Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @zhewenl.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

ivanium

Overall good but I have an ask on interface design. I'll still approve to unblock.

ivanium · 2026-09-13T04:29:20Z

+            local_block_ids = self.kv_cache_config.select_block_ids(
+                blocks.get_block_ids(),
+                self._store_group_ids,
            )


after seeing many of this pattern, I am thinking if we should do

Suggested change

local_block_ids = self.kv_cache_config.select_block_ids(

blocks.get_block_ids(),

self._store_group_ids,

)

local_block_ids = \

blocks.get_block_ids(group_ids=self._store_group_ids)

Where we extend KVCacheBlocks.get_block_ids() with group_ids: Optional[list[int]]
There might be some other places scheduler has already passed full block_ids in, in which case we can do a manual for loop. This should help remove kv_cache_config.select_block_ids() interface

ivanium · 2026-09-13T04:35:47Z

+    def select_block_ids(
+        self,
+        block_ids: Sequence[list[int]],


see comments in mooncake store scheduler.py. Perhaps we can remove this API

zhewenl · 2026-09-13T05:51:05Z

/ci run

github-actions · 2026-09-13T05:51:23Z

✅ Triggered Buildkite CI #88610 for commit a6916cc7728d.

Keep transferable KV cache groups distinct from the prefix-cacheable subset used by hash-addressed stores. Preserve the full transfer projection for connectors such as NIXL while making MooncakeStore scheduler, worker, and coordinator paths operate only on storable groups. Validate Mamba align mode from the spec directly and add regression coverage for QSA ring groups, block-id projection, and worker tensor registration. Fixes vllm-project#54726 Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Zhewen Li <zhewenli@inferact.ai>

Use KVCacheConfig's prefix-cacheable group projection directly in MooncakeStore and remove the connector-local helper. Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Zhewen Li <zhewenli@inferact.ai>

zhewenl · 2026-09-13T21:19:30Z

/ci run

github-actions · 2026-09-13T21:19:48Z

✅ Triggered Buildkite CI #88686 for commit 59f605639a69.

…SA ring) groups; fix align-mode check (vllm-project#55027) Signed-off-by: Zhewen Li <zhewenli@inferact.ai> Co-authored-by: Zhewen Li <zhewenli@inferact.ai> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

mergify Bot added bug Something isn't working kv-connector kv-cache-manager labels Sep 2, 2026

zhewenl mentioned this pull request Sep 2, 2026

[Bugfix][KV Offload] SimpleCPUOffload: skip non-prefix-cacheable (QSA ring) groups #55033

Draft

zhewenl force-pushed the fix-mooncake-qsa-prefix-store branch from fa6456d to f6db171 Compare September 2, 2026 23:53

zhewenl marked this pull request as ready for review September 3, 2026 00:01

zhewenl requested review from ApostaC, NickLucche, WoosukKwon, alexm-redhat, heheda12345, ivanium, njhill, orozery, robertgshaw2-redhat, xuechendi and ywang96 as code owners September 3, 2026 00:01

claude Bot reviewed Sep 3, 2026

View reviewed changes

zhewenl mentioned this pull request Sep 3, 2026

[Bug]: Qwen3.8-Flash-Next-FP8 + MooncakeStoreConnector fails at startup #54726

Closed

1 task

drakosha mentioned this pull request Sep 3, 2026

[KV Offload] Scope offload group configs to prefix-cacheable KV cache groups #54743

Open

mergify Bot added the needs-rebase label Sep 7, 2026

zhewenl force-pushed the fix-mooncake-qsa-prefix-store branch from f6db171 to 928ba63 Compare September 8, 2026 23:55

mergify Bot removed the needs-rebase label Sep 8, 2026

ivanium reviewed Sep 9, 2026

View reviewed changes

zhewenl force-pushed the fix-mooncake-qsa-prefix-store branch 3 times, most recently from 3d290b8 to ef9c42f Compare September 9, 2026 06:33

zhewenl force-pushed the fix-mooncake-qsa-prefix-store branch from ef9c42f to 2b123ab Compare September 9, 2026 09:42

mergify Bot added the needs-rebase label Sep 12, 2026

zhewenl force-pushed the fix-mooncake-qsa-prefix-store branch from 2b123ab to 45bac07 Compare September 12, 2026 17:56

mergify Bot removed the needs-rebase label Sep 12, 2026

zhewenl force-pushed the fix-mooncake-qsa-prefix-store branch from 45bac07 to d962d27 Compare September 12, 2026 18:21

ivanium approved these changes Sep 13, 2026

View reviewed changes

zhewenl force-pushed the fix-mooncake-qsa-prefix-store branch from d962d27 to 96f18ff Compare September 13, 2026 05:18

Zhewen Li and others added 2 commits September 13, 2026 21:09

refactor: simplify prefix-cacheable group selection

59f6056

Use KVCacheConfig's prefix-cacheable group projection directly in MooncakeStore and remove the connector-local helper. Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Zhewen Li <zhewenli@inferact.ai>

zhewenl force-pushed the fix-mooncake-qsa-prefix-store branch from a6916cc to 59f6056 Compare September 13, 2026 21:13

ivanium approved these changes Sep 14, 2026

View reviewed changes

ivanium merged commit 58d45fd into vllm-project:main Sep 14, 2026
114 checks passed

ivanium mentioned this pull request Sep 16, 2026

[Bugfix][KV Offload] Skip non-prefix-cacheable groups in SimpleCPUOffload (GLM-5.3-Flash kpool tail and QSA) #56810

Open

Uh oh!

Uh oh!

Conversation

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

Uh oh!

Purpose

Fix

Duplicate-work check

Tests

Model evaluation

AI assistance

Uh oh!

coderabbitai Bot commented Sep 2, 2026

Draft PR not reviewed

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!

mergify Bot commented Sep 7, 2026

Uh oh!

ivanium left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ivanium Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

zhewenl Sep 12, 2026

Choose a reason for hiding this comment

Uh oh!

mergify Bot commented Sep 12, 2026

Uh oh!

ivanium left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ivanium Sep 13, 2026

Choose a reason for hiding this comment

Uh oh!

ivanium Sep 13, 2026

Choose a reason for hiding this comment

Uh oh!

zhewenl commented Sep 13, 2026

Uh oh!

github-actions Bot commented Sep 13, 2026

Uh oh!

zhewenl commented Sep 13, 2026

Uh oh!

github-actions Bot commented Sep 13, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

zhewenl commented Sep 2, 2026 •

edited

Loading

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

edited

Loading

ivanium left a comment •

edited

Loading