[ROCm][DSpark][DCP] Serve prefix cache hits under DCP for Kimi-K3 (#53598)
Signed-off-by: Yichao Zhu <Yichao.Zhu@amd.com> Co-authored-by: andyluo7 <andy.luo@amd.com>
YukioZzz committed
Sep 1, 2026 at 02:18 UTC
2ba984a5d06db414f3b2474fe9338faf6cd80a1c
6 files changed
+161
-11
tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py
+25
@@ -132,6 +132,31 @@ def make_full_mamba_manager(
132
)
133
134
135
+def test_dcp_fine_hit_retention_uses_hash_alignment_without_eagle():
136
+ """DCP must not discard a reusable Mamba state at hash alignment."""
137
+ hash_block_size = 2
138
+ manager = make_full_mamba_manager(
139
+ dcp_world_size=4,
140
+ hash_block_size=hash_block_size,
141
+ full_block_size=hash_block_size,
142
+ mamba_block_size=hash_block_size,
143
+ num_blocks=64,
144
+ )
145
+ manager.coordinator.retention_interval = 0
146
+
147
+ token_ids = list(range(7))
148
+ producer = make_request("producer", token_ids, hash_block_size, sha256)
149
+ computed_blocks, num_computed, _ = manager.get_computed_blocks(producer)
150
+ assert num_computed == 0
151
+ assert manager.allocate_slots(producer, 6, 0, computed_blocks) is not None
152
+ manager.free(producer)
153
+ manager.new_step_starts()
154
+
155
+ consumer = make_request("consumer", token_ids, hash_block_size, sha256)
156
+ _, num_computed, _ = manager.get_computed_blocks(consumer)
157
+ assert num_computed == 6
158
+
159
+
160
@pytest.mark.parametrize("dcp_world_size", [1, 4])
161
def test_mamba_align_split_partial_tail_schedule(dcp_world_size: int):
162
"""Chunk ends with partial hits on: block-aligned chunks, one extra stop
tests/v1/core/test_kv_cache_utils.py
+15
@@ -1318,6 +1318,21 @@ def test_project_kv_cache_groups_to_worker():
1318
assert set(proj_spec.kv_cache_specs.keys()) == {"layer1", "layer3"}
1319
1320
1321
+def test_dcp_world_size_for_kv_cache_spec_shards_full_attention_only():
1322
+ dcp = 8
1323
+ full = FullAttentionSpec(
1324
+ block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float32
1325
+ )
1326
+ mla = new_mla_spec()
1327
+ mamba = new_mamba_spec()
1328
+ uniform_mla = UniformTypeKVCacheSpecs(block_size=16, kv_cache_specs={"layer": mla})
1329
+ assert kv_cache_utils.dcp_world_size_for_kv_cache_spec(full, dcp) == dcp
1330
+ assert kv_cache_utils.dcp_world_size_for_kv_cache_spec(mla, dcp) == dcp
1331
+ assert kv_cache_utils.dcp_world_size_for_kv_cache_spec(uniform_mla, dcp) == dcp
1332
+ assert kv_cache_utils.dcp_world_size_for_kv_cache_spec(mamba, dcp) == 1
1333
+ assert kv_cache_utils.dcp_world_size_for_kv_cache_spec(full, 1) == 1
1334
+
1335
+
1336
@pytest.mark.parametrize(
1337
"layer_type,dcp_size,expected_width",
1338
[
tests/v1/core/test_prefix_caching.py
+77
@@ -232,6 +232,83 @@ def make_kv_cache_config_three_types(
232
)
233
234
235
+def test_prefix_cache_hit_uses_per_group_dcp_geometry():
236
+ """Prefix lookup must use each group's DCP size, not the process-wide one.
237
+
238
+ Target and draft MLA are both sharded (DCP=8); Mamba stays replicated
239
+ (DCP=1). Hits then align to the sharded full-attention block, not the
240
+ unsharded page size.
241
+ """
242
+ block_size = 16
243
+ dcp = 8
244
+ sharded_block = block_size * dcp
245
+ kv_cache_config = KVCacheConfig(
246
+ num_blocks=64,
247
+ kv_cache_tensors=[],
248
+ kv_cache_groups=[
249
+ KVCacheGroupSpec(
250
+ ["target_mla"],
251
+ FullAttentionSpec(
252
+ block_size=block_size,
253
+ num_kv_heads=1,
254
+ head_size=1,
255
+ dtype=torch.float32,
256
+ ),
257
+ ),
258
+ KVCacheGroupSpec(
259
+ ["draft_mla"],
260
+ MLAAttentionSpec(
261
+ block_size=block_size,
262
+ num_kv_heads=1,
263
+ head_size=1,
264
+ dtype=torch.float32,
265
+ ),
266
+ ),
267
+ KVCacheGroupSpec(
268
+ ["mamba"],
269
+ MambaSpec(
270
+ block_size=block_size,
271
+ shapes=(1, 1),
272
+ dtypes=(torch.float32,),
273
+ ),
274
+ ),
275
+ ],
276
+ )
277
+ manager = KVCacheManager(
278
+ kv_cache_config,
279
+ max_model_len=8192,
280
+ enable_caching=True,
281
+ hash_block_size=block_size,
282
+ scheduler_block_size=sharded_block,
283
+ dcp_world_size=dcp,
284
+ )
285
+ target_mgr, draft_mgr, mamba_mgr = manager.coordinator.single_type_managers
286
+ assert target_mgr.dcp_world_size == dcp
287
+ assert draft_mgr.dcp_world_size == dcp
288
+ assert mamba_mgr.dcp_world_size == 1
289
+ assert target_mgr.block_size == sharded_block
290
+ assert draft_mgr.block_size == sharded_block
291
+ assert mamba_mgr.block_size == block_size
292
+
293
+ hash_fn = sha256
294
+ common_token_ids = [i for i in range(2) for _ in range(sharded_block)]
295
+ req0 = make_request("0", common_token_ids + [99] * 7, block_size, hash_fn)
296
+ computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)
297
+ assert num_computed_tokens == 0
298
+ blocks = manager.allocate_slots(
299
+ req0, len(req0.prompt_token_ids), 0, computed_blocks
300
+ )
301
+ assert blocks is not None
302
+
303
+ req1 = make_request("1", common_token_ids + [100] * 5, block_size, hash_fn)
304
+ computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)
305
+ assert num_computed_tokens == 2 * sharded_block
306
+ assert [len(group) for group in computed_blocks.blocks] == [2, 2, 16]
307
+
308
+ manager.free(req0)
309
+ manager.free(req1)
310
+
311
+
312
@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
313
def test_prefill(hash_fn):
314
block_size = 16
vllm/v1/core/kv_cache_coordinator.py
+18
-10
@@ -11,6 +11,7 @@ from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
11
from vllm.v1.core.kv_cache_utils import (
12
BlockHash,
13
KVCacheBlock,
14
+ dcp_world_size_for_kv_cache_spec,
15
)
16
from vllm.v1.core.single_type_kv_cache_manager import (
17
CrossAttentionManager,
@@ -141,7 +142,9 @@ class KVCacheCoordinator(ABC):
142
block_pool=self.block_pool,
143
enable_caching=enable_caching,
144
kv_cache_group_id=i,
144
- dcp_world_size=dcp_world_size,
145
+ dcp_world_size=dcp_world_size_for_kv_cache_spec(
146
+ kv_cache_group.kv_cache_spec, dcp_world_size
147
+ ),
148
pcp_world_size=pcp_world_size,
149
scheduler_block_size=self.scheduler_block_size,
150
needs_kv_cache_zeroing=self.kv_cache_config.needs_kv_cache_zeroing,
@@ -506,11 +509,9 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):
509
num_prefill_lookahead=num_prefill_lookahead,
510
)
511
self.kv_cache_spec = self.kv_cache_config.kv_cache_groups[0].kv_cache_spec
509
- self.block_size = self.kv_cache_spec.block_size
510
- self.dcp_world_size = dcp_world_size
512
+ self.dcp_world_size = self.single_type_managers[0].dcp_world_size
513
self.pcp_world_size = pcp_world_size
512
- if dcp_world_size > 1:
513
- self.block_size *= dcp_world_size
514
+ self.block_size = self.single_type_managers[0].block_size
515
# For models using only Mamba, block_size is set to max_model_len when
516
# prefix caching is disabled, and hash_block_size validation is skipped.
517
assert not enable_caching or (hash_block_size == self.block_size), (
@@ -656,6 +657,9 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
657
"cache managers require block-aligned lookups: %s.",
658
", ".join(sorted(unsupported_partial_hit_managers)),
659
)
660
+ cache_hit_alignment_tokens = self._cache_hit_alignment_tokens
661
+ for manager in self.single_type_managers:
662
+ manager.cache_hit_alignment_tokens = cache_hit_alignment_tokens
663
self.verify_and_split_kv_cache_groups()
664
665
@property
@@ -852,11 +856,12 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
856
kv_cache_spec=spec,
857
drop_eagle_block=drop_eagle_block,
858
alignment_tokens=self._cache_hit_alignment_tokens,
855
- dcp_world_size=(
856
- self.dcp_world_size
857
- if isinstance(spec, FullAttentionSpec)
858
- else 1
859
- ),
859
+ dcp_world_size=self.single_type_managers[
860
+ first_group_id
861
+ ].dcp_world_size,
862
+ pcp_world_size=self.single_type_managers[
863
+ first_group_id
864
+ ].pcp_world_size,
865
)
866
if drop_eagle_block:
867
eagle_verified.add(idx)
@@ -913,6 +918,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
918
hit_lengths: list[int] = [0] * num_groups
919
920
for spec, group_ids, manager_cls, use_eagle in self.attention_groups:
921
+ manager = self.single_type_managers[group_ids[0]]
922
blocks, group_hit = manager_cls.find_longest_cache_hit(
923
block_hashes=block_hashes,
924
max_length=max_cache_hit_length,
@@ -921,6 +927,8 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
927
kv_cache_spec=spec,
928
drop_eagle_block=use_eagle,
929
alignment_tokens=self._cache_hit_alignment_tokens,
930
+ dcp_world_size=manager.dcp_world_size,
931
+ pcp_world_size=manager.pcp_world_size,
932
)
933
for gid, blks in zip(group_ids, blocks):
934
hit_blocks[gid] = blks
vllm/v1/core/kv_cache_utils.py
+22
@@ -675,6 +675,28 @@ def resolve_dcp_kv_cache_spec(spec: KVCacheSpec, dcp_world_size: int) -> KVCache
675
return replace(spec, block_size=block_size)
676
677
678
+def dcp_world_size_for_kv_cache_spec(spec: KVCacheSpec, dcp_world_size: int) -> int:
679
+ """Return the DCP size that owns this group's block geometry.
680
+
681
+ Full-attention KV (including MLA) is sharded across DCP ranks, so prefix
682
+ hashing and manager ``block_size`` use the process DCP size. Other specs
683
+ keep replicated per-rank state (Mamba, sliding window, chunked-local) and
684
+ must keep ``dcp_world_size=1`` even when the process runs with DCP > 1.
685
+
686
+ Draft MLA groups on the sharded DSpark path are ``FullAttentionSpec`` /
687
+ ``MLAAttentionSpec`` and therefore keep the process DCP size. A replicated
688
+ draft group would need a different spec, not this helper.
689
+ """
690
+ if dcp_world_size <= 1:
691
+ return 1
692
+ inner = spec
693
+ if isinstance(spec, UniformTypeKVCacheSpecs):
694
+ inner = next(iter(spec.kv_cache_specs.values()))
695
+ if isinstance(inner, FullAttentionSpec):
696
+ return dcp_world_size
697
+ return 1
698
+
699
+
700
def resolve_kv_cache_block_sizes(
701
kv_cache_config: KVCacheConfig,
702
vllm_config: VllmConfig,
vllm/v1/core/single_type_kv_cache_manager.py
+4
-1
@@ -72,6 +72,9 @@ class SingleTypeKVCacheManager(ABC):
72
block until the request finishes.
73
"""
74
self.scheduler_block_size = scheduler_block_size
75
+ # Hybrid fine-grained lookup may lower this after all participating
76
+ # managers have been validated by the coordinator.
77
+ self.cache_hit_alignment_tokens = scheduler_block_size
78
# The block size for this manager; used for actual block allocation.
79
self.block_size = kv_cache_spec.block_size
80
self.dcp_world_size = dcp_world_size
@@ -460,7 +463,7 @@ class SingleTypeKVCacheManager(ABC):
463
block_mask = self.reachable_block_mask(
464
start_block=num_cached_blocks,
465
end_block=num_full_blocks,
463
- alignment_tokens=self.scheduler_block_size,
466
+ alignment_tokens=self.cache_hit_alignment_tokens,
467
kv_cache_spec=self.kv_cache_spec,
468
use_eagle=self.use_eagle,
469
retention_interval=retention_interval,