[gcanlin]

Overview

With the extensible KV cache (on by default since #56492), MooncakeStoreConnector fails at startup. Its worker splits each layer's view by cache_config.num_gpu_blocks, but extend_kv_cache creates the connector against the committed views before it updates that value, so group_kernel_blocks asserts. The worker now takes the block count from its own kv_cache_config, as the rest of the worker already does.

mooncake/store/worker.py, in register_kv_caches
    group_kernel_blocks(view, self.num_blocks)
kv_cache_interface.py, in group_kernel_blocks
    assert cache.shape[0] % num_blocks == 0
AssertionError

No open PR covers this (searched for Mooncake + extensible KV cache). Related: #56492, #65225, #65137.

Validation

GLM-5.3-Flash, colocated DEP8 on 8x B200, MooncakeStoreConnector with a standalone store (1.3 TB DRAM plus a 6 TiB SSD tier), real MTP 5. Image nightly e0d00e4b8e with this change applied:

	Before	After
Startup	AssertionError in group_kernel_blocks	Serves; KV cache extended to 1920 blocks (105.38 GiB, 15,728,640 tokens per rank)
GSM8K, 5-shot	–	0.9727 (0.975 with --disable-extensible-kv-cache)
Needle read-back	–	16/16 on a cold store; 32/32 after 15 minutes of AgentX at concurrency 384, by which time 2.61 TB had spilled to SSD

Needle read-back: 16 contexts of about 20K tokens with 800 locker codes each. Each context is asked one code on a cold store, then the same code and a new one after the load.

The new unit test fails without the fix and passes with it. The rest of test_mooncake_store_worker.py, test_mooncake_store_connector.py and test_mooncake_store_hma_e2e.py give the same results with and without it.

pytest tests/v1/kv_connector/unit/test_mooncake_store_worker.py -k register_kv_caches

AI assistance (Claude Code) was used to find the cause, write the change and run the tests above.

🤖 Generated with Claude Code

[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.

[gcanlin]

@njhill Hi, there is another small bug from #56492 when I'm running mooncake with extensible KV cache. Could you help take a look? Thanks! #56492 is great and help reduce a lot memory.

[ivanium]

/ci run

[github-actions]

✅ @gcanlin, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• Run commands test your branch with every commit currently on its upstream target branch. If it is behind, the bot merges the target branch into it first (pull before pushing again). Append --allow-stale to a run command to test an outdated branch as is, at your own risk.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[github-actions]

✅ Triggered Buildkite CI #99764 for commit 69950c6e31c0.

ℹ️ This PR is 89 commits behind upstream main. Merged upstream main into this branch as 69950c6e31c0 before starting CI. Pull the branch before pushing more changes.

[njhill]

Thanks @gcanlin

I have these suggested additional changes but they can be handled in a follow-on: njhill@73db856