[Bugfix][KV Offload] Skip non-prefix-cacheable groups in SimpleCPUOffload (GLM-5.3-Flash kpool tail) - #56810
Open
JaredforReal wants to merge 3 commits into
Open
[Bugfix][KV Offload] Skip non-prefix-cacheable groups in SimpleCPUOffload (GLM-5.3-Flash kpool tail)#56810JaredforReal wants to merge 3 commits into
JaredforReal wants to merge 3 commits into
Conversation
…load GLM-5.3-Flash carries a kpool-tail scratch group (KpoolTailSpec, prefix_cacheable=False) whose block_size is index_kpool=4 while the prefix-cacheable groups hash at 1152 tokens. Core, HybridKVCacheCoordinator and the MooncakeStore connector already skip such groups, but SimpleCPUOffloadScheduler treated every KV cache group as offloadable: - _select_eager_blocks_to_store viewed the request hashes at the scratch group's block size for BlockStored metadata, which trips resolve_block_hashes' divisibility assert as soon as kv cache events are enabled (engine dies on the first prefill store). - update_state_after_alloc mapped external tokens onto the scratch group and asserted alignment against its block size. Record which groups take part in offloading and skip the rest in store selection and load pairing. Their blocks never carry a hash, so nothing is stored or loaded for them either way; the scratch slot in the per-group hit tuple stays empty, as the coordinator already leaves it. Signed-off-by: Jared Wen <jaredwen@inferact.ai> Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Lazy target sizing still counts non-prefix-cacheable scratch groups.
Get a fresh assessment by requesting another Copilot review.
Pull request overview
Updates SimpleCPUOffloadScheduler to exclude non-prefix-cacheable scratch KV groups from CPU offload handling.
Changes:
- Filters scratch groups from eager store/load operations and KV events.
- Adds scheduler and KV-event regression tests.
- Lazy target sizing still requires adjustment to exclude scratch groups.
File summaries
| File | Description |
|---|---|
vllm/v1/simple_kv_offload/manager.py |
Filters non-cacheable groups from offload paths. |
tests/v1/simple_kv_offload/test_scheduler.py |
Tests scratch-group store/load exclusion. |
tests/v1/simple_kv_offload/test_kv_events.py |
Verifies events contain cacheable groups only. |
Review details
- Files reviewed: 3/3 changed files
- Comments generated: 1
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+172
to
+175
| self._group_offloadable = [ | ||
| g.kv_cache_spec.prefix_cacheable | ||
| for g in self.cpu_kv_cache_config.kv_cache_groups | ||
| ] |
…llm-project#55033 Fold the SimpleCPUOffloadScheduler changes of vllm-project#55033 into this branch so one PR covers both non-prefix-cacheable scratch layouts (GLM-5.3-Flash kpool tail and Qwen3.8-Flash-Next QSA ring): - Use the core projection KVCacheConfig.prefix_cacheable_group_ids instead of a private per-group list. - Skip non-prefix-cacheable groups in _estimate_lazy_target_blocks, so the ring's small block size no longer inflates the lazy-mode free-block target. - Add the QSA-shaped regression tests (store/load exclusion, lazy target sizing) alongside the kpool-tail ones. Co-authored-by: Zhewen Li <zhewenli@inferact.ai> Signed-off-by: Jared Wen <jaredwen@inferact.ai> Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
GLM-5.3-Flash carries a kpool-tail scratch group (
KpoolTailSpec,prefix_cacheable=False,block_size=index_kpool=4) next to prefix-cacheable groups that hash at 1152 tokens on this box (MLA 1152, 4× Mamba 1152 inalignmode). Core (resolve_kv_cache_block_sizes,HybridKVCacheCoordinator,KVCacheManager) and theMooncakeStore connectoralready skip such groups.SimpleCPUOffloadSchedulerdid not, and treated every KV cache group as offloadable:_select_eager_blocks_to_storeviews the request hashes at each group's block size to buildBlockStoredmetadata. For the tail group that isresolve_block_hashes(hashes, hash_block_size=1152, block_size=4), which trips the divisibility assert as soon as--kv-events-configenables KV cache events. The engine dies on the first prefill store.update_state_after_allocmaps external (CPU-hit) tokens onto the scratch group and asserts alignment against its block size, although the coordinator never returns hit blocks for it.This PR records which groups take part in offloading (
prefix_cacheable) and skips the rest in eager store selection and in load pairing. Scratch blocks never carry a hash, so nothing is stored or loaded for them either way; the scratch slot in the per-group hit tuple stays empty, exactly as the coordinator leaves it.Duplication check
Two open PRs touch the same code path: #55033 (draft since 2026-09-02, Qwen3.8-Flash-Next QSA ring, also adjusts
_estimate_lazy_target_blocks) and #56404 (DeepSeek-V4.1-Flash, gatesresolve_block_hashesin the store path only, currently conflicting). This PR is scoped to the GLM-5.3-Flash kpool-tail layout and adds what neither of them covers: a regression test for the actual crash trigger (enable_kv_cache_events=true,BlockStoredemitted for cacheable groups only) and an end-to-end store/load + accuracy check on GLM-5.3-Flash. Happy to fold the test into #55033 instead if the maintainers prefer to land that one. #54743 applies the same scoping to the OffloadingConnector (build_offloading_config) and does not touchvllm/v1/simple_kv_offload; the two are complementary.Test Plan
Unit:
E2E: GLM-5.3-Flash, 4× GB300, TP4 + EP,
--enforce-eager,--max-model-len 32768,SimpleCPUOffloadConnectorwith 16 GB pinned host memory per rank,--kv-events-config '{"enable_kv_cache_events":true,"publisher":"zmq"}', dev endpoint/reset_prefix_cacheused to evict the GPU prefix cache between runs (?reset_external=truealso clears the CPU pool).Test Result
E2E store/load (10,277-token prompt, greedy, 32 new tokens):
/reset_prefix_cache/reset_prefix_cacheagainHits are 1152-aligned (the Mamba/MLA block); the kpool tail block is a fresh per-request scratch block at every resume, as with a GPU hit.
Accuracy (4 natural-text prompts built from
docs/, 3.4k–24k tokens,max_tokens=16, greedy,logprobs=20at the first generated position; every case compared against the first fresh run of the same prompt):The CPU-hit deltas sit inside the band that two fresh recomputations of the same prompt already show on this model/backend (the fresh-vs-fresh spread is not caused by this PR; it is present with the connector disabled as well, see the
fresh againrow), so the offload round trip adds no measurable error over a plain GPU prefix hit.AI assistance: this change was developed with Claude Code (analysis, patch, tests and the reproduction scripts); the human submitter reviewed every changed line and ran the tests and evaluations above.