[ROCm][DSpark][DCP] Serve prefix cache hits under DCP for Kimi-K3 - #53598
Conversation
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment Once the PR is approved or has the If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| has_dcp_partial_full_attention_group = dcp_world_size > 1 and any( | ||
| isinstance(g.kv_cache_spec, FullAttentionSpec) | ||
| and manager.block_size > hash_block_size | ||
| for g, manager in zip( | ||
| kv_cache_config.kv_cache_groups, self.single_type_managers | ||
| ) | ||
| ) | ||
| self.enable_partial_hash_hits = ( | ||
| has_partial_mamba_group or has_dcp_partial_full_attention_group | ||
| ) |
There was a problem hiding this comment.
This is enabled with has_partial_mamba_group. Can you enble --mamba-cache-mode align?
The logic here enables partial prefix hit in a full attention case. Which we do not intend to enable at the moment.
|
Thank you for breaking dow the change to individual commits. changes in b968057 looks good. For the rest please see comments. |
Construct cache managers with the effective DCP size of each KV group and use that manager-owned geometry for local prefix-cache lookup. This keeps DCP-sharded full-attention groups and replicated state groups from sharing one incorrect block-size assumption. Assisted-by: OpenAI Codex Signed-off-by: andyluo7 <andy.luo@amd.com> # Conflicts: # vllm/v1/core/kv_cache_utils.py
3f010a4 to
6d4f041
Compare
#53598 will remain scoped to per-group DCP cache geometry. The replay-boundary issue is real but separable and will be handled independently. @GirasoleY #50897 is too big, and we might also check if successor-aware hashing is enabled to do legacy drop or not? @ZJY0516 |
Hybrid DCP lookup can use a finer hash alignment than the scheduler LCM. Propagate that validated alignment to cache managers so sparse Mamba retention does not discard checkpoints that local lookup can reuse. Keep scheduler alignment as the default and whenever fine-grained lookup is disabled, including CPU offload coordinators. Cover the behavior without EAGLE to keep this geometry fix independent from replay-boundary handling. Co-authored-by: andyluo7 <andy.luo@amd.com> Signed-off-by: Yichao Zhu <Yichao.Zhu@amd.com>
|
/amd-ci run |
|
❌ @YukioZzz, A reviewer with write access must run |
|
/amd-ci run |
|
✅ Triggered Buildkite AMD CI #12457 for commit |
|
/ci run |
|
❌ @YukioZzz, A reviewer with write access must run |
|
✅ @YukioZzz, CI is now available for this PR.
|
…lm-project#53598) Cherry-picked from upstream 2ba984a onto build/kimi-k3-v2-dcp16-20260819 (base 99e62b8). Root-cause fix for the Kimi-K3 TP16/DCP16 long-context needle failures: pre-fix, every single-type KV cache manager was constructed with the process DCP world size, so the Mamba/KDA manager applied x16 block geometry to per-rank-replicated state - prefix-cache hits then resumed at positions where no KDA checkpoint exists, positionally desyncing the KDA and full-attention layer families on every cache-hit request. Adaptations to this tree (vs the upstream commit): - kept this branch's kv_cache_utils imports (_use_lockstep_mla_allocation, is_deepseek_v4_hybrid_kv_cache_config) alongside the new helper - unified has_partial_mamba_group with this branch's DCP1-only expression (upstream relaxes partial hash hits to DCP>1 when mamba block >= hash block) - ported the small KVCacheSpec.prefix_cacheable property (base True + UniformTypeKVCacheSpecs all()) which upstream added separately; this tree has no CircularBufferSpec so no False override exists - test fixture: dropped MambaSpec.num_prefill_checkpoint_blocks (upstream prefill-checkpoint feature absent here), added the lcm import, and kept this tree's unparametrized test_mamba_align_split_partial_tail_schedule (the dcp mock attrs it threads are upstream-scheduler-only) Original-author: Yichao Zhu <Yichao.Zhu@amd.com> (and co-authors, vllm-project#53598) Claude-Session: https://claude.ai/code/session_01SnjnheY7rHpk73rGeNS5Sk
…lm-project#53598) Signed-off-by: Yichao Zhu <Yichao.Zhu@amd.com> Co-authored-by: andyluo7 <andy.luo@amd.com>
…lm-project#53598) Signed-off-by: Yichao Zhu <Yichao.Zhu@amd.com> Co-authored-by: andyluo7 <andy.luo@amd.com>
…lm-project#53598) Signed-off-by: Yichao Zhu <Yichao.Zhu@amd.com> Co-authored-by: andyluo7 <andy.luo@amd.com>
5 participants
Summary
Use each KV-cache group's effective DCP geometry when constructing cache managers and performing hybrid prefix-cache lookup.
Kimi-K3 mixes DCP-sharded MLA/full-attention groups with replicated Mamba groups. Applying the process-wide DCP size to every group makes the scheduler, block hashing, and cache managers disagree about block geometry.
This PR is intentionally limited to per-group DCP cache geometry. It does not include:
Implementation
Validation
Validation used an integration stack containing #51705 and this PR with DCP8, DSpark, FP8 KV cache, prefix caching, eager execution, and block rejection sampling.
Server reproduction args:
vllm serve /data/model/Kimi-K3 \ --served-model-name moonshotai/Kimi-K3 \ --host 0.0.0.0 \ --port 8892 \ --trust-remote-code \ --moe-backend auto \ --tensor-parallel-size 8 \ --decode-context-parallel-size 8 \ --load-format fastsafetensors \ --gpu-memory-utilization 0.82 \ --language-model-only \ --max-num-seqs 20 \ --enable-auto-tool-choice \ --tool-call-parser kimi_k3 \ --reasoning-parser kimi_k3 \ --max-model-len 1048576 \ --enable-prefix-caching \ --kv-cache-dtype fp8 \ --enforce-eager \ --speculative-config '{"model":"/data/model/Kimi-K3-DSpark","num_speculative_tokens":2,"method":"dspark","attention_backend":"TRITON_MLA","kv_cache_dtype":"auto","draft_sample_method":"probabilistic","rejection_sample_method":"block"}'Accuracy
GSM8K full-set, 5-shot,
num_concurrent=64:0.9682 +/- 0.0048.0.9682 +/- 0.0048.After adding hash-aligned manager retention in
8d28f83fcc, a full non-DSpark max-throughput run produced:0.9644 +/- 0.0051.0.9629 +/- 0.0052.8,775,168 / 11,093,014 = 79.11%.84.39%.10,627.90 / 139.31 tokens/s.186 / 926.13s, with TTFT and ITL coverage both100%.Reproduction command:
AIPerf fast
The sustained agentic replay run used:
This PR fixes the DCP geometry layer but does not by itself retain the Mamba state at the EAGLE-adjusted replay boundary. The distinction is visible in server-counter deltas:
1,155,072 / 7,404,11815.60%80.37%7,649,280 / 10,638,81571.90%83.92%The remaining DSpark replay issue is separate from this alignment fix. In the reproduced DCP layout, the scheduler block is
24,576tokens while the hash and Mamba blocks are1,536tokens. The dense groups reconcile to the EAGLE-adjusted boundary at12,288, but the ordinary prompt boundary is13,824; without explicitly materializing and retaining the EAGLE boundary, the Mamba groups still reduce that hybrid hit to zero.A short latest-head smoke test combined #51705
cc504d70a522, this PR6d4f04123a, and the separate replay-boundary fix. With dummy loading and eager execution, cold/hot/reuse returned0,12,288, and12,288cached tokens. A 180-second AIPerf run completed with1,181,184 / 2,506,238 = 47.13%actual local prefix hits versus61.63%theoretical; these short dummy-load numbers are cache-path validation only, not performance results.Related
Development note
AI assistance was used for code investigation, test construction, and validation orchestration. The resulting changes and validation results were reviewed by the human submitter.