[KV Connector] Support MooncakeStore with hybrid DCP prefix caching - #53324
[KV Connector] Support MooncakeStore with hybrid DCP prefix caching#53324wzhao18 wants to merge 1 commit into
Conversation
Scale attention cache blocks for DCP while keeping recurrent state replicated. Preserve the actual stored hash boundary through hybrid prefix-hit reconciliation so loads do not derive missing Mooncake keys. Align coordinator behavior for multiple attention groups and add regression coverage for DCP, Mamba, EAGLE, and partial-prefix loads. Assisted-by: OpenAI Codex Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
cc @GirasoleY for a look too |
ivanium
left a comment
ivanium
left a comment
There was a problem hiding this comment.
Thanks for the feature and the fix! Left one nit and another PCP question for @GirasoleY to confirm.
| # to final hit_length. | ||
| for group in self.attention_groups: | ||
| if not isinstance(group.spec, FullAttentionSpec): | ||
| continue |
There was a problem hiding this comment.
I think we have sorted the groups to always put FullAttn groups at the beginning. Can you verify if this is still the case? If so, we can break here
| continue | |
| break |
| def resolve_dcp_kv_block_size(spec: KVCacheSpec, dcp_world_size: int) -> int: | ||
| """Return the token span of a cache block under DCP.""" | ||
| if isinstance(spec, AttentionSpec): | ||
| return spec.block_size * dcp_world_size | ||
| return spec.block_size |
There was a problem hiding this comment.
This is good; but I heard from @GirasoleY that now we have revised the PCP kv cache layout design, and technically we should exclude PCP world size from all this kind of block size calculation. In other words, we should make resolve_dcp_kv_block_size the right way for resolve_dcp_kv_block_size. @GirasoleY is my understanding correct?
There was a problem hiding this comment.
I feel this is non-blocking though. If confirmed, I am okay if we just leave a TODO note here for a future fix PR
There was a problem hiding this comment.
Right PCP should be excluded from block calculation now. The current implementation is sound here.
This comment was marked as duplicate.
This comment was marked as duplicate.
Sorry, something went wrong.
3 participants
Purpose
This PR enables
MooncakeStoreConnectorfor mamba hybrid models (e.g., kimi k3) using decode context parallelism (DCP) and fixes a bug where incorrect load keys are constructed under fine-grained prefix hits, causing KV load failures.Failure Example
Consider Kimi-K3 with DCP 8 + DSpark and
--prefix-match-unit 128:Reproduction steps:
Cache blocks stored and keys:
H_12288H_14848H_13824H_14848Lookup
H_14848.H_13824.hit_length = 13,824.The existing load path derives the key for a group's partially reused physical block from
hit_length. As a result, it requestedH_13824for the attention block, even though that physical block was stored underH_14848. Mooncake therefore returned-704because the derived key did not exist.This PR fixes this load failure by recording the precise token boundaries for partial cache store so that subsequent load requests can use the correct load key.
Test Plan
Test Result
-704load failure now passes with the fix.Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.