[Kimi-K3] support DCP partial prefix cache hit (#50493)
Signed-off-by: Summer Yang <girasoleyang@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>
Summer Yang committed
Aug 17, 2026 at 17:05 UTC
0db502c8d8a620f29dd3560e69a22f0a7ed3477c
8 files changed
+523
-37
tests/distributed/test_kimi_linear_context_parallel.py
+111
@@ -2,12 +2,15 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import math
5
+from pathlib import Path
6
7
import pytest
8
9
from tests.conftest import VllmRunner
10
from tests.utils import create_new_process_for_each_test
11
from vllm import SamplingParams, TokensPrompt
12
+from vllm.config import CUDAGraphMode
13
+from vllm.transformers_utils.configs.kimi_k3 import KimiK3Config
14
15
MODEL = "moonshotai/Kimi-Linear-48B-A3B-Instruct"
16
@@ -126,3 +129,111 @@ def test_kimi_linear_dcp_tiny(
129
logprob_drifts.append(abs(baseline_logprob - dcp_logprob))
130
131
assert max(logprob_drifts) <= 1e-2
132
+
133
+
134
+def _make_tiny_k3_config(model_dir: Path) -> str:
135
+ config = KimiK3Config(
136
+ text_config=_make_tiny_overrides(),
137
+ vision_config={
138
+ "vt_num_attention_heads": 2,
139
+ "vt_num_hidden_layers": 1,
140
+ "vt_hidden_size": 32,
141
+ "vt_intermediate_size": 64,
142
+ "qkv_hidden_size": 48,
143
+ },
144
+ architectures=["KimiK3ForConditionalGeneration"],
145
+ )
146
+ config.save_pretrained(model_dir)
147
+ return str(model_dir)
148
+
149
+
150
+def _run_k3_partial_prefix_reuse(
151
+ vllm_runner: type[VllmRunner], model_name: str, dcp_size: int
152
+) -> tuple[list[int], list[float], int]:
153
+ block_size = 256
154
+ common_prefix = [3 + (token_idx % 251) for token_idx in range(block_size)]
155
+ prime_prompt = TokensPrompt(prompt_token_ids=common_prefix + [17])
156
+ replay_prompt = TokensPrompt(prompt_token_ids=common_prefix + [18, 19])
157
+ sampling_params = SamplingParams(
158
+ temperature=0,
159
+ max_tokens=8,
160
+ seed=0,
161
+ ignore_eos=True,
162
+ logprobs=20,
163
+ )
164
+
165
+ with vllm_runner(
166
+ model_name=model_name,
167
+ skip_tokenizer_init=True,
168
+ load_format="dummy",
169
+ language_model_only=True,
170
+ tensor_parallel_size=2,
171
+ decode_context_parallel_size=dcp_size,
172
+ cp_kv_cache_interleave_size=1,
173
+ distributed_executor_backend="mp",
174
+ dtype="bfloat16",
175
+ seed=0,
176
+ max_model_len=512,
177
+ max_num_seqs=4,
178
+ max_num_batched_tokens=256,
179
+ gpu_memory_utilization=0.85,
180
+ block_size=block_size,
181
+ enable_chunked_prefill=True,
182
+ enable_prefix_caching=True,
183
+ mamba_cache_mode="align",
184
+ ) as runner:
185
+ assert (
186
+ runner.llm.llm_engine.vllm_config.compilation_config.cudagraph_mode
187
+ == CUDAGraphMode.FULL_AND_PIECEWISE
188
+ )
189
+ runner.llm.generate(
190
+ [prime_prompt],
191
+ sampling_params,
192
+ use_tqdm=False,
193
+ )
194
+ replay_output = runner.llm.generate(
195
+ [replay_prompt],
196
+ sampling_params,
197
+ use_tqdm=False,
198
+ )[0]
199
+
200
+ completion = replay_output.outputs[0]
201
+ token_ids = list(completion.token_ids)
202
+ assert len(token_ids) == sampling_params.max_tokens
203
+ assert completion.logprobs is not None
204
+ assert len(completion.logprobs) == sampling_params.max_tokens
205
+ selected_logprobs = [
206
+ step_logprobs[token_id].logprob
207
+ for token_id, step_logprobs in zip(token_ids, completion.logprobs)
208
+ ]
209
+ return token_ids, selected_logprobs, replay_output.num_cached_tokens
210
+
211
+
212
+@create_new_process_for_each_test()
213
+@pytest.mark.distributed(num_gpus=2)
214
+def test_kimi_k3_dcp_partial_prefix_reuse(
215
+ vllm_runner: type[VllmRunner],
216
+ num_gpus_available: int,
217
+ tmp_path: Path,
218
+) -> None:
219
+ if num_gpus_available < 2:
220
+ pytest.skip("Need at least 2 GPUs")
221
+
222
+ model_name = _make_tiny_k3_config(tmp_path / "tiny-kimi-k3")
223
+ baseline_tokens, baseline_logprobs, baseline_cached = _run_k3_partial_prefix_reuse(
224
+ vllm_runner, model_name, dcp_size=1
225
+ )
226
+ dcp_tokens, dcp_logprobs, dcp_cached = _run_k3_partial_prefix_reuse(
227
+ vllm_runner, model_name, dcp_size=2
228
+ )
229
+
230
+ assert baseline_cached == 256
231
+ assert dcp_cached == 256
232
+ assert dcp_tokens == baseline_tokens
233
+ assert len(dcp_logprobs) == len(baseline_logprobs) == 8
234
+ logprob_drifts = []
235
+ for baseline_logprob, dcp_logprob in zip(baseline_logprobs, dcp_logprobs):
236
+ assert math.isfinite(baseline_logprob)
237
+ assert math.isfinite(dcp_logprob)
238
+ logprob_drifts.append(abs(baseline_logprob - dcp_logprob))
239
+ assert max(logprob_drifts) <= 1e-2
tests/v1/attention/test_mla_backends.py
+43
@@ -27,6 +27,7 @@ from vllm.model_executor.layers.attention.mla_attention import (
27
MLAAttention,
28
QueryLenSupport,
29
_DecodeConcatQuantFP8,
30
+ build_mla_chunked_context_metadata,
31
)
32
from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
33
from vllm.platforms import current_platform
@@ -93,6 +94,48 @@ def test_mla_kv_cache_spec_uses_layer_cache_dtype(
94
assert spec.page_size_bytes == 64 * 656
95
96
97
+@pytest.mark.cpu_test
98
+def test_dcp_chunked_context_accepts_non_virtual_block_aligned_prefix():
99
+ context_lens_cpu = torch.tensor([65, 96], dtype=torch.int32)
100
+ prefill_query_start_loc_cpu = torch.tensor([0, 1, 2], dtype=torch.int32)
101
+
102
+ metadata = build_mla_chunked_context_metadata(
103
+ context_lens_cpu=context_lens_cpu,
104
+ prefill_query_start_loc_cpu=prefill_query_start_loc_cpu,
105
+ chunked_prefill_workspace=torch.empty(0),
106
+ chunked_prefill_workspace_size=128,
107
+ block_size=64,
108
+ align_chunk_to_block=True,
109
+ device=torch.device("cpu"),
110
+ dcp_world_size=4,
111
+ dcp_local_block_size=1,
112
+ dcp_virtual_block_size=4,
113
+ )
114
+
115
+ assert metadata is not None
116
+ assert metadata.context_lens.tolist() == [65, 96]
117
+ assert [chunk.seq_lens.tolist() for chunk in metadata.chunks] == [[65], [96]]
118
+ assert [chunk.starts.tolist() for chunk in metadata.chunks] == [[0], [0]]
119
+ assert [chunk.local_context_lens_allranks for chunk in metadata.chunks] == [
120
+ [[17, 16, 16, 16]],
121
+ [[24, 24, 24, 24]],
122
+ ]
123
+ assert [chunk.padded_local_seq_lens for chunk in metadata.chunks] == [
124
+ [17],
125
+ [24],
126
+ ]
127
+ assert [chunk.padded_local_cu_seq_lens.tolist() for chunk in metadata.chunks] == [
128
+ [0, 17],
129
+ [0, 24],
130
+ ]
131
+ assert [chunk.cu_seq_lens.tolist() for chunk in metadata.chunks] == [
132
+ [0, 65],
133
+ [0, 96],
134
+ ]
135
+ assert [chunk.num_context_tokens for chunk in metadata.chunks] == [65, 96]
136
+ assert [chunk.num_local_context_tokens for chunk in metadata.chunks] == [17, 24]
137
+
138
+
139
# Remove sm100 backends from the list if not using sm100
140
if not torch.cuda.is_available() or torch.cuda.get_device_properties(0).major < 10:
141
BACKENDS_TO_TEST.remove(AttentionBackendEnum.CUTLASS_MLA)
tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py
+203
-1
@@ -4,6 +4,7 @@
4
"align") models: scheduler chunk splitting, partial tail registration, CoW
5
on partial hits, and same-step deferral."""
6
7
+from math import lcm
8
from types import SimpleNamespace
9
from unittest.mock import MagicMock
10
@@ -68,12 +69,62 @@ def test_capable_connector_uses_divergent_partial_hit_lookup():
69
manager.get_computed_blocks.assert_not_called()
70
71
71
-def test_mamba_align_split_partial_tail_schedule():
72
+def make_full_mamba_manager(
73
+ *,
74
+ dcp_world_size: int,
75
+ hash_block_size: int = 2,
76
+ full_block_size: int = 4,
77
+ mamba_block_size: int = 4,
78
+ num_blocks: int = 32,
79
+ use_eagle: bool = False,
80
+):
81
+ kv_cache_config = KVCacheConfig(
82
+ num_blocks=num_blocks,
83
+ kv_cache_tensors=[],
84
+ kv_cache_groups=[
85
+ KVCacheGroupSpec(
86
+ ["full"],
87
+ FullAttentionSpec(
88
+ block_size=full_block_size,
89
+ num_kv_heads=1,
90
+ head_size=1,
91
+ dtype=torch.float32,
92
+ ),
93
+ ),
94
+ KVCacheGroupSpec(
95
+ ["mamba"],
96
+ MambaSpec(
97
+ block_size=mamba_block_size,
98
+ shapes=(1, 1),
99
+ dtypes=(torch.float32,),
100
+ mamba_cache_mode="align",
101
+ ),
102
+ ),
103
+ ],
104
+ )
105
+ scheduler_block_size = lcm(
106
+ full_block_size * dcp_world_size,
107
+ mamba_block_size,
108
+ )
109
+ return make_kv_cache_manager(
110
+ kv_cache_config=kv_cache_config,
111
+ max_model_len=8192,
112
+ enable_caching=True,
113
+ dcp_world_size=dcp_world_size,
114
+ scheduler_block_size=scheduler_block_size,
115
+ hash_block_size=hash_block_size,
116
+ use_eagle=use_eagle,
117
+ )
118
+
119
+
120
+@pytest.mark.parametrize("dcp_world_size", [1, 4])
121
+def test_mamba_align_split_partial_tail_schedule(dcp_world_size: int):
122
"""Chunk ends with partial hits on: block-aligned chunks, one extra stop
123
at the prompt's last hash boundary (registering the partial tail), then
124
the remaining tokens. block=512, hash=32, prompt=10000, budget=8192:
125
0 -> 8192 -> 9728 -> 9984 -> 10000."""
126
block_size = 512
127
+ scheduler_block_size = block_size * dcp_world_size
128
hash_block_size = 32
129
mock = SimpleNamespace(
130
cache_config=SimpleNamespace(block_size=block_size),
@@ -81,6 +132,8 @@ def test_mamba_align_split_partial_tail_schedule():
132
scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
133
use_eagle=False,
134
hash_block_size=hash_block_size,
135
+ dcp_world_size=dcp_world_size,
136
+ scheduler_block_size=scheduler_block_size,
137
mamba_partial_cache_hit=True,
138
)
139
split = Scheduler._mamba_block_aligned_split
@@ -109,6 +162,7 @@ def test_mamba_align_split_partial_tail_schedule():
162
# stops at the next block boundary (10240), later chunk ends re-align.
163
req2 = make_request("1", [0] * 12000, hash_block_size, sha256)
164
req2.num_computed_tokens = 9984
165
+ assert req2.num_computed_tokens % scheduler_block_size != 0
166
assert split(self=mock, request=req2, num_new_tokens=2016) == 256
167
req2.num_computed_tokens = 10240
168
assert split(self=mock, request=req2, num_new_tokens=1000) == 512
@@ -1499,3 +1553,151 @@ def test_hybrid_sliding_window_group_disables_partial_hash_hits():
1553
1554
assert num_computed == mamba_block_size
1555
assert len(computed_blocks.blocks[0]) * hash_block_size == num_computed
1556
+
1557
+
1558
+@pytest.mark.parametrize("dcp_world_size", [1, 2, 4])
1559
+def test_hybrid_partial_hash_hit_uses_cow_under_dcp(dcp_world_size: int):
1560
+ hash_block_size = 2
1561
+ physical_block_size = 4
1562
+ manager = make_full_mamba_manager(
1563
+ dcp_world_size=dcp_world_size,
1564
+ hash_block_size=hash_block_size,
1565
+ full_block_size=physical_block_size,
1566
+ mamba_block_size=physical_block_size,
1567
+ )
1568
+ assert manager.coordinator.enable_partial_hash_hits
1569
+
1570
+ req0 = make_request("dcp-owner", [0, 0, 1, 1, 2, 2], 2, sha256)
1571
+ computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
1572
+ assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
1573
+ manager.free(req0)
1574
+ manager.new_step_starts()
1575
+
1576
+ partial_hash = req0.block_hashes[2]
1577
+ partial_full_block = manager.block_pool.get_cached_block(partial_hash, [0])
1578
+ partial_mamba_block = manager.block_pool.get_cached_block(partial_hash, [1])
1579
+ assert partial_full_block is not None
1580
+ assert partial_mamba_block is not None
1581
+
1582
+ req1 = make_request("dcp-replay", [0, 0, 1, 1, 2, 2, 3, 3], 2, sha256)
1583
+ computed_blocks, num_computed, _ = manager.get_computed_blocks(req1)
1584
+ assert num_computed == 6
1585
+ full_block_size = physical_block_size * dcp_world_size
1586
+ assert [len(group) for group in computed_blocks.blocks] == [
1587
+ (6 + full_block_size - 1) // full_block_size,
1588
+ 2,
1589
+ ]
1590
+
1591
+ new_blocks = manager.allocate_slots(req1, 2, num_computed, computed_blocks)
1592
+ assert new_blocks is not None
1593
+ full_new_block_id = new_blocks.get_block_ids()[0][0]
1594
+ mamba_new_block_id = new_blocks.get_block_ids()[1][0]
1595
+ copies, retained = manager.take_kv_cache_block_copies()
1596
+ assert KVCacheBlockCopy(partial_full_block[0].block_id, full_new_block_id) in copies
1597
+ assert (
1598
+ KVCacheBlockCopy(partial_mamba_block[0].block_id, mamba_new_block_id) in copies
1599
+ )
1600
+ manager.block_pool.free_blocks(retained)
1601
+
1602
+
1603
+@pytest.mark.parametrize("dcp_world_size", [2, 4])
1604
+def test_dcp_partial_hit_resumes_on_replicated_mamba_snapshot(
1605
+ dcp_world_size: int,
1606
+):
1607
+ block_size = 4
1608
+ manager = make_full_mamba_manager(
1609
+ dcp_world_size=dcp_world_size,
1610
+ hash_block_size=block_size,
1611
+ full_block_size=block_size,
1612
+ mamba_block_size=block_size,
1613
+ )
1614
+ assert manager.coordinator.enable_partial_hash_hits
1615
+ assert manager.coordinator._cache_hit_alignment_tokens == block_size
1616
+ assert manager.coordinator.single_type_managers[0].block_size == (
1617
+ block_size * dcp_world_size
1618
+ )
1619
+ assert manager.coordinator.single_type_managers[1].block_size == block_size
1620
+
1621
+ prefix = list(range(12))
1622
+ req0 = make_request("snapshot-owner", prefix, block_size, sha256)
1623
+ computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
1624
+ assert manager.allocate_slots(req0, 12, num_computed, computed_blocks) is not None
1625
+ manager.free(req0)
1626
+ manager.new_step_starts()
1627
+
1628
+ req1 = make_request(
1629
+ "snapshot-replay", prefix + list(range(12, 16)), block_size, sha256
1630
+ )
1631
+
1632
+ computed_blocks, num_computed, _ = manager.get_computed_blocks(req1)
1633
+ assert num_computed == 12
1634
+ assert [len(group) for group in computed_blocks.blocks] == [
1635
+ (12 + block_size * dcp_world_size - 1) // (block_size * dcp_world_size),
1636
+ 3,
1637
+ ]
1638
+ partial_full_block = computed_blocks.blocks[0][-1]
1639
+ new_blocks = manager.allocate_slots(req1, 4, num_computed, computed_blocks)
1640
+ assert new_blocks is not None
1641
+ full_new_block_id = new_blocks.get_block_ids()[0][0]
1642
+ copies, retained = manager.take_kv_cache_block_copies()
1643
+ assert KVCacheBlockCopy(partial_full_block.block_id, full_new_block_id) in copies
1644
+ assert all(
1645
+ copy.src_block_id != computed_blocks.blocks[1][-1].block_id for copy in copies
1646
+ )
1647
+ manager.block_pool.free_blocks(retained)
1648
+
1649
+
1650
+def test_dcp_joint_hit_is_bounded_by_replicated_mamba_snapshots():
1651
+ block_size = 4
1652
+ manager = make_full_mamba_manager(
1653
+ dcp_world_size=2,
1654
+ hash_block_size=block_size,
1655
+ full_block_size=block_size,
1656
+ mamba_block_size=block_size,
1657
+ )
1658
+ prefix = list(range(12))
1659
+ req0 = make_request("joint-owner", prefix, block_size, sha256)
1660
+ computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
1661
+ assert manager.allocate_slots(req0, 8, num_computed, computed_blocks) is not None
1662
+ manager.new_step_starts()
1663
+
1664
+ replay_tokens = prefix + list(range(12, 16))
1665
+ req1 = make_request("joint-before", replay_tokens, block_size, sha256)
1666
+ computed_blocks, num_computed, _ = manager.get_computed_blocks(req1)
1667
+ assert num_computed == 8
1668
+ assert [len(group) for group in computed_blocks.blocks] == [1, 2]
1669
+
1670
+ req0.num_computed_tokens = 8
1671
+ assert manager.allocate_slots(req0, 4) is not None
1672
+ manager.new_step_starts()
1673
+
1674
+ req2 = make_request("joint-after", replay_tokens, block_size, sha256)
1675
+ computed_blocks, num_computed, _ = manager.get_computed_blocks(req2)
1676
+ assert num_computed == 12
1677
+ assert [len(group) for group in computed_blocks.blocks] == [2, 3]
1678
+
1679
+
1680
+def test_dcp_partial_hit_with_eagle_rewinds_one_hash_unit():
1681
+ hash_block_size = 2
1682
+ manager = make_full_mamba_manager(
1683
+ dcp_world_size=2,
1684
+ hash_block_size=hash_block_size,
1685
+ full_block_size=4,
1686
+ mamba_block_size=4,
1687
+ use_eagle=True,
1688
+ )
1689
+
1690
+ req0 = make_request("eagle-owner", [7] * 6, hash_block_size, sha256)
1691
+ computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
1692
+ assert manager.allocate_slots(req0, 4, num_computed, computed_blocks) is not None
1693
+ req0.num_computed_tokens = 4
1694
+ manager.new_step_starts()
1695
+ assert manager.allocate_slots(req0, 2) is not None
1696
+ req0.num_computed_tokens = 6
1697
+ manager.new_step_starts()
1698
+
1699
+ req1 = make_request("eagle-replay", [7] * 6 + [9] * 2, 2, sha256)
1700
+ computed_blocks, num_computed, _ = manager.get_computed_blocks(req1)
1701
+ assert num_computed == 4
1702
+ assert [len(group) for group in computed_blocks.blocks] == [1, 1]
1703
+ assert manager.allocate_slots(req1, 4, num_computed, computed_blocks) is not None
tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.py
+23
-21
@@ -55,9 +55,9 @@ def cache_full_block_and_partial_tail(
55
token_ids: list[int],
56
*,
57
enable_kv_cache_events: bool = False,
58
+ block_size: int = 6,
59
) -> tuple[BlockPool, Request, list[KVCacheBlock], BlockHash]:
60
hash_block_size = 2
60
- block_size = 6
61
kv_cache_group_id = 0
62
req = make_request("0", token_ids, hash_block_size, sha256)
63
pool = BlockPool(
@@ -389,27 +389,34 @@ def test_reset_prefix_cache_clears_partial_entry_metadata():
389
assert pool.cached_block_hashes_by_block == {}
390
391
392
-def test_evict_cached_block_removes_full_hash_and_partial_entry():
393
- pool, req, blocks, partial_hash_10 = cache_full_block_and_partial_tail(
394
- [0, 0, 1, 1, 2, 2, 3, 3, 4, 4]
392
+@pytest.mark.parametrize("dcp_world_size", [1, 2, 4])
393
+def test_evict_cached_block_removes_full_hash_and_partial_entry(
394
+ dcp_world_size: int,
395
+):
396
+ block_size = 6 * dcp_world_size
397
+ partial_num_tokens = 2 * block_size - 2
398
+ pool, req, blocks, partial_hash = cache_full_block_and_partial_tail(
399
+ list(range(partial_num_tokens)), block_size=block_size
400
)
396
- full_hash = BlockHashListWithBlockSize(req.block_hashes, 2, 6)[0]
401
+ full_hash = BlockHashListWithBlockSize(req.block_hashes, 2, block_size)[0]
402
403
assert pool.get_cached_block(full_hash, [0]) == [blocks[0]]
399
- assert pool.get_cached_block(partial_hash_10, [0]) == [blocks[1]]
404
+ assert pool.get_cached_block(partial_hash, [0]) == [blocks[1]]
405
406
pool.evict_blocks({blocks[0].block_id, blocks[1].block_id})
407
408
assert pool.get_cached_block(full_hash, [0]) is None
404
- assert pool.get_cached_block(partial_hash_10, [0]) is None
409
+ assert pool.get_cached_block(partial_hash, [0]) is None
410
assert pool.cached_block_hashes_by_block == {}
411
412
408
-def test_partial_block_promotes_to_direct_full_block_hash():
413
+@pytest.mark.parametrize("dcp_world_size", [1, 2, 4])
414
+def test_partial_block_promotes_to_direct_full_block_hash(dcp_world_size: int):
415
hash_block_size = 2
410
- block_size = 6
416
+ block_size = 6 * dcp_world_size
417
kv_cache_group_id = 0
412
- token_ids = [0, 0, 1, 1, 2, 2, 3, 3, 4, 4]
418
+ partial_num_tokens = 2 * block_size - hash_block_size
419
+ token_ids = list(range(partial_num_tokens))
420
req = make_request("0", token_ids, hash_block_size, sha256)
421
pool = BlockPool(
422
num_gpu_blocks=3,
@@ -426,27 +433,22 @@ def test_partial_block_promotes_to_direct_full_block_hash():
433
block_size=block_size,
434
kv_cache_group_id=kv_cache_group_id,
435
)
429
- partial_hash_10 = boundary_hash(req, hash_block_size, 10)
436
+ partial_hash = boundary_hash(req, hash_block_size, partial_num_tokens)
437
assert pool.cache_partial_block(
438
request=req,
439
block=blocks[1],
433
- num_tokens=10,
440
+ num_tokens=partial_num_tokens,
441
kv_cache_group_id=kv_cache_group_id,
442
block_size=block_size,
443
)
437
- assert pool.get_cached_block(partial_hash_10, [kv_cache_group_id]) == [blocks[1]]
444
+ assert pool.get_cached_block(partial_hash, [kv_cache_group_id]) == [blocks[1]]
445
439
- req.append_output_token_ids([5, 5])
446
+ req.append_output_token_ids(list(range(partial_num_tokens, 2 * block_size)))
447
full_hashes = BlockHashListWithBlockSize(
448
req.block_hashes, hash_block_size, block_size
449
)
450
promoted_full_hash = full_hashes[1]
444
- # The promoted full-block hash is the fine hash at the 12-token boundary,
445
- # not a concatenation of the fine hashes inside the block.
446
- assert promoted_full_hash == req.block_hashes[12 // hash_block_size - 1]
447
- assert promoted_full_hash != BlockHash(
448
- req.block_hashes[3] + req.block_hashes[4] + req.block_hashes[5]
449
- )
451
+ assert promoted_full_hash == req.block_hashes[2 * block_size // hash_block_size - 1]
452
453
pool.cache_full_blocks(
454
request=req,
@@ -457,4 +459,4 @@ def test_partial_block_promotes_to_direct_full_block_hash():
459
kv_cache_group_id=kv_cache_group_id,
460
)
461
assert pool.get_cached_block(promoted_full_hash, [kv_cache_group_id]) == [blocks[1]]
460
- assert pool.get_cached_block(partial_hash_10, [kv_cache_group_id]) is None
462
+ assert pool.get_cached_block(partial_hash, [kv_cache_group_id]) is None
tests/v1/worker/test_gpu_model_runner_v2.py
new
+118
@@ -0,0 +1,118 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from types import SimpleNamespace
5
+
6
+import pytest
7
+import torch
8
+
9
+import vllm.v1.worker.gpu.model_runner as model_runner_module
10
+from vllm.v1.kv_cache_interface import (
11
+ FullAttentionSpec,
12
+ KVCacheConfig,
13
+ KVCacheGroupSpec,
14
+ MambaSpec,
15
+)
16
+from vllm.v1.worker.gpu.block_table import BlockTables
17
+from vllm.v1.worker.gpu.model_runner import GPUModelRunner
18
+
19
+
20
+@pytest.mark.parametrize(
21
+ ("mamba_cache_mode", "num_speculative_blocks", "expected"),
22
+ [
23
+ pytest.param("align", 0, 65_536, id="align-prefix-cache"),
24
+ pytest.param("none", 7, 8, id="no-prefix-cache-with-speculation"),
25
+ ],
26
+)
27
+def test_initialize_kv_cache_does_not_dcp_shard_mamba_block_table(
28
+ monkeypatch,
29
+ mamba_cache_mode: str,
30
+ num_speculative_blocks: int,
31
+ expected: int,
32
+):
33
+ """Mamba/GDN block-table rows index global positions, unlike DCP KV."""
34
+
35
+ max_model_len = 1_048_576
36
+ attention_block_size = 1_536
37
+ mamba_block_size = 16
38
+ dcp_size = 8
39
+ full_attention_spec = FullAttentionSpec(
40
+ block_size=attention_block_size,
41
+ num_kv_heads=1,
42
+ head_size=1,
43
+ dtype=torch.bfloat16,
44
+ )
45
+ mamba_spec = MambaSpec(
46
+ shapes=((1,),),
47
+ dtypes=(torch.bfloat16,),
48
+ block_size=mamba_block_size,
49
+ mamba_cache_mode=mamba_cache_mode,
50
+ num_speculative_blocks=num_speculative_blocks,
51
+ )
52
+ kv_cache_config = KVCacheConfig(
53
+ num_blocks=1,
54
+ kv_cache_tensors=[],
55
+ kv_cache_groups=[
56
+ KVCacheGroupSpec(["attention"], full_attention_spec),
57
+ KVCacheGroupSpec(["kda"], mamba_spec),
58
+ ],
59
+ )
60
+ vllm_config = SimpleNamespace(
61
+ parallel_config=SimpleNamespace(decode_context_parallel_size=dcp_size),
62
+ cache_config=SimpleNamespace(mamba_cache_mode=mamba_cache_mode),
63
+ )
64
+ runner = SimpleNamespace(
65
+ max_model_len=max_model_len,
66
+ is_encoder_decoder=False,
67
+ vllm_config=vllm_config,
68
+ )
69
+
70
+ class _CapturedWidths(Exception):
71
+ pass
72
+
73
+ captured: list[int] = []
74
+
75
+ def capture_width(max_num_blocks: int, *_args, **_kwargs) -> int:
76
+ captured.append(max_num_blocks)
77
+ if len(captured) == 2:
78
+ raise _CapturedWidths
79
+ return max_num_blocks
80
+
81
+ monkeypatch.setattr(model_runner_module, "get_block_table_width", capture_width)
82
+
83
+ with pytest.raises(_CapturedWidths):
84
+ GPUModelRunner.initialize_kv_cache(runner, kv_cache_config)
85
+
86
+ # Attention KV is local to one of eight DCP ranks; KDA state is replicated
87
+ # and therefore needs one table entry for every global 16-token page.
88
+ assert captured == [86, expected]
89
+
90
+
91
+def test_append_block_ids_rejects_write_past_row_capacity():
92
+ """Reject an oversized staged write before it can corrupt the next row."""
93
+
94
+ class _BlockTable:
95
+ gpu = torch.empty((2, 4), dtype=torch.int32)
96
+
97
+ def stage_write(self, *_args):
98
+ pytest.fail("an oversized write must not be staged")
99
+
100
+ block_tables = BlockTables.__new__(BlockTables)
101
+ block_tables.num_kv_cache_groups = 1
102
+ block_tables.blocks_per_kv_block = [1]
103
+ block_tables.block_tables = [_BlockTable()]
104
+ block_tables.num_blocks = SimpleNamespace(
105
+ np=torch.tensor([[0, 3]], dtype=torch.int32)
106
+ )
107
+
108
+ with pytest.raises(
109
+ RuntimeError,
110
+ match=r"request 1, group 0 exceeds row capacity \(5 > 4\)",
111
+ ):
112
+ block_tables.append_block_ids(
113
+ req_index=1,
114
+ new_block_ids=([4, 5],),
115
+ overwrite=False,
116
+ )
117
+
118
+ assert block_tables.num_blocks.np[0, 1] == 3
vllm/v1/core/kv_cache_coordinator.py
+11
-4
@@ -619,15 +619,22 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
619
"full-attention and Mamba groups, got: "
620
f"{type(g.kv_cache_spec).__name__}."
621
)
622
- # Fine-grained hash hits require Mamba "align", no context
623
- # parallelism, and compatible cache managers in every group.
622
+ # Fine-grained hash hits require Mamba "align" and compatible cache
623
+ # managers in every group. TP needs hashing finer than the Mamba block;
624
+ # DCP accepts equality because it scales the effective full-attention
625
+ # block instead.
626
has_partial_mamba_group = any(
627
isinstance(g.kv_cache_spec, MambaSpec)
628
and g.kv_cache_spec.mamba_cache_mode == "align"
627
- and g.kv_cache_spec.block_size > hash_block_size
629
+ and (
630
+ (dcp_world_size == 1 and g.kv_cache_spec.block_size > hash_block_size)
631
+ or (
632
+ dcp_world_size > 1 and g.kv_cache_spec.block_size >= hash_block_size
633
+ )
634
+ )
635
for g in kv_cache_config.kv_cache_groups
636
)
630
- self.enable_partial_hash_hits = dcp_world_size == 1 and has_partial_mamba_group
637
+ self.enable_partial_hash_hits = has_partial_mamba_group
638
if self.enable_partial_hash_hits:
639
unsupported_partial_hit_managers = {
640
type(manager).__name__
vllm/v1/worker/gpu/block_table.py
+8
-1
@@ -116,8 +116,15 @@ class BlockTables:
116
bpk = self.blocks_per_kv_block[i]
117
if bpk > 1:
118
block_ids = [b * bpk + k for b in block_ids for k in range(bpk)]
119
+ end = start + len(block_ids)
120
+ row_capacity = self.block_tables[i].gpu.shape[1]
121
+ if end > row_capacity:
122
+ raise RuntimeError(
123
+ f"Block table write for request {req_index}, group {i} exceeds "
124
+ f"row capacity ({end} > {row_capacity})"
125
+ )
126
self.block_tables[i].stage_write(req_index, start, block_ids)
120
- self.num_blocks.np[i, req_index] = start + len(block_ids)
127
+ self.num_blocks.np[i, req_index] = end
128
129
def apply_staged_writes(self) -> None:
130
if self.num_kv_cache_groups == 0:
vllm/v1/worker/gpu/model_runner.py
+6
-10
@@ -60,7 +60,6 @@ from vllm.multimodal.encoder_budget import (
60
)
61
from vllm.sequence import IntermediateTensors
62
from vllm.tasks import SupportedTask
63
-from vllm.utils.math_utils import cdiv
63
from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
64
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
65
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
@@ -517,17 +516,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
516
for kv_cache_group in kv_cache_config.kv_cache_groups:
517
spec = kv_cache_group.kv_cache_spec
518
block_sizes.append(spec.block_size)
520
- # When using DCP, each request's KV cache is sharded among different ranks.
521
- # As a result, one block on the current rank covers `block_size * cp_size`
522
- # tokens in the full, global (unsharded) sequence.
523
- max_num_blocks = cdiv(
524
- block_table_max_model_len, spec.block_size * self.dcp_size
519
+ # Let each cache type account for CP. Attention KV is DCP-sharded,
520
+ # while Mamba/GDN recurrent state is replicated across DCP ranks.
521
+ max_num_blocks = spec.max_num_blocks_per_req(
522
+ self.vllm_config, block_table_max_model_len
523
)
526
- # For Mamba/Hybrid Model, KVCaches need extra blocks for speculative tokens
524
+ # Preserve each cache type's alignment requirements after applying
525
+ # its topology-aware block-table width.
526
if isinstance(spec, MambaSpec):
528
- max_num_blocks = (
529
- max_num_blocks if self.cache_config.enable_prefix_caching else 1
530
- ) + spec.num_speculative_blocks
527
max_num_blocks = get_block_table_width(
528
max_num_blocks, spec.block_size, token_alignment=None
529
)