@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 144e79c810

[Kimi K3] Support internal prefix checkpoints with partial prefix caching and spec-decoding (#53614)

Signed-off-by: Ziming Huang <zelda.huanghuang@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>

Ziming Huang committed Sep 6, 2026 at 13:18 UTC 144e79c8106da23141ac010394b782f730cc7fe8

11 files changed +481 -106

tests/models/kimi_k3/test_kda_metadata.py

+71 -1

@@ -92,6 +92,10 @@ def _make_builder(

92 mamba_cache_mode: str = "none",

93 use_recoverssm: bool = False,

94 num_prefill_checkpoint_blocks: int = 0,

95 + mamba_block_size: int = BLOCK_SIZE,

96 + prefix_match_unit: int | None = None,

97 + use_eagle: bool = False,

98 + disable_eagle_block_drop: bool = False,

99 ) -> AttentionMetadataBuilder:

100 vllm_config = create_vllm_config(

101 model_name="Qwen/Qwen3.5-0.8B",

@@ -102,20 +106,29 @@ def _make_builder(

106 method="ngram",

107 num_speculative_tokens=num_speculative_tokens,

108 )

109 + if use_eagle:

110 + vllm_config.speculative_config.method = "eagle3"

111 + vllm_config.speculative_config.disable_eagle_block_drop = (

112 + disable_eagle_block_drop

113 + )

114 vllm_config.compilation_config.cudagraph_mode = (

115 CUDAGraphMode.FULL_AND_PIECEWISE if full_cuda_graph else CUDAGraphMode.NONE

116 )

117 vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode

118 vllm_config.cache_config.use_replayssm = use_recoverssm

119 vllm_config.cache_config.use_kda_recoverssm = use_recoverssm

120 + vllm_config.cache_config.prefix_match_unit = prefix_match_unit

121 builder = builder_cls(

122 kv_cache_spec=MambaSpec(

113 - block_size=BLOCK_SIZE,

123 + block_size=mamba_block_size,

124 shapes=((16, 64),),

125 dtypes=(torch.float16,),

126 mamba_cache_mode=mamba_cache_mode,

127 num_speculative_blocks=(0 if use_recoverssm else num_speculative_tokens),

128 num_prefill_checkpoint_blocks=num_prefill_checkpoint_blocks,

129 + prefill_checkpoint_alignment=(

130 + 16 if num_prefill_checkpoint_blocks > 0 else None

131 + ),

132 ),

133 layer_names=["layer.0"],

134 vllm_config=vllm_config,

@@ -156,6 +169,7 @@ def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch):

169 mamba_type=MambaAttentionBackendEnum.GDN_ATTN,

170 mamba_cache_mode="align",

171 num_prefill_checkpoint_blocks=1,

172 + prefill_checkpoint_alignment=16,

173 )

174

175 # This is the same compatibility check performed while initializing the

@@ -177,11 +191,13 @@ def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch):

191 cache_config=SimpleNamespace(

192 mamba_cache_mode="align",

193 use_kda_recoverssm=True,

194 + prefix_match_unit=None,

195 ),

196 parallel_config=SimpleNamespace(decode_context_parallel_size=1),

197 speculative_config=SimpleNamespace(

198 num_speculative_tokens=2,

199 parallel_drafting=False,

200 + use_eagle_block_drop=Mock(return_value=False),

201 ),

202 compilation_config=SimpleNamespace(

203 cudagraph_mode=CUDAGraphMode.NONE,

@@ -300,6 +316,60 @@ def test_internal_checkpoint_metadata_targets_last_aligned_boundary():

316 )

317

318

319 +@pytest.mark.parametrize(

320 + ("disable_eagle_block_drop", "prefix_match_unit", "expected_offset"),

321 + [(False, 16, 80), (True, 16, 96), (False, 8, None)],

322 +)

323 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")

324 +def test_spec_internal_checkpoint_metadata_targets_replay_boundary(

325 + disable_eagle_block_drop: bool,

326 + prefix_match_unit: int,

327 + expected_offset: int | None,

328 +) -> None:

329 + device = torch.device("cuda")

330 + batch = BatchSpec(seq_lens=[100], query_lens=[100])

331 + common_attn_metadata = create_common_attn_metadata(

332 + batch, BLOCK_SIZE, device, arange_block_indices=True

333 + )

334 + common_attn_metadata = common_attn_metadata.replace(

335 + is_prefilling=torch.tensor([True]),

336 + block_table_tensor=common_attn_metadata.block_table_tensor + 1,

337 + )

338 + builder = _make_builder(

339 + KimiK3KDAMetadataBuilder,

340 + num_speculative_tokens=3,

341 + full_cuda_graph=False,

342 + mamba_cache_mode="align",

343 + num_prefill_checkpoint_blocks=1,

344 + mamba_block_size=64,

345 + prefix_match_unit=prefix_match_unit,

346 + use_eagle=True,

347 + disable_eagle_block_drop=disable_eagle_block_drop,

348 + device=device,

349 + )

350 + assert isinstance(builder, KimiK3KDAMetadataBuilder)

351 + builder.mamba_aligned_state_indices = mamba_get_block_table_tensor(

352 + common_attn_metadata.block_table_tensor,

353 + common_attn_metadata.seq_lens,

354 + builder.kv_cache_spec,

355 + "align",

356 + )

357 + actual = builder.build(0, common_attn_metadata)

358 +

359 + if expected_offset is None:

360 + assert actual.checkpoint is None

361 + return

362 + assert actual.checkpoint is not None

363 + torch.testing.assert_close(

364 + actual.checkpoint.state_indices,

365 + torch.tensor([1], dtype=torch.int32, device=device),

366 + )

367 + torch.testing.assert_close(

368 + actual.checkpoint.checkpoint_offsets,

369 + torch.tensor([expected_offset], dtype=torch.int32, device=device),

370 + )

371 +

372 +

373 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")

374 def test_internal_checkpoint_metadata_skips_unaligned_offset():

375 device = torch.device("cuda")

tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py

+104 -37

@@ -92,8 +92,23 @@ def make_full_mamba_manager(

92 mamba_block_size: int = 4,

93 num_blocks: int = 32,

94 use_eagle: bool = False,

95 + num_speculative_blocks: int = 0,

96 num_prefill_checkpoint_blocks: int = 0,

97 ):

98 + mamba_group = KVCacheGroupSpec(

99 + ["mamba"],

100 + MambaSpec(

101 + block_size=mamba_block_size,

102 + shapes=(1, 1),

103 + dtypes=(torch.float32,),

104 + mamba_cache_mode="align",

105 + num_speculative_blocks=num_speculative_blocks,

106 + num_prefill_checkpoint_blocks=num_prefill_checkpoint_blocks,

107 + prefill_checkpoint_alignment=(

108 + 16 if num_prefill_checkpoint_blocks > 0 else None

109 + ),

110 + ),

111 + )

112 kv_cache_config = KVCacheConfig(

113 num_blocks=num_blocks,

114 kv_cache_tensors=[],

@@ -107,16 +122,7 @@ def make_full_mamba_manager(

122 dtype=torch.float32,

123 ),

124 ),

110 - KVCacheGroupSpec(

111 - ["mamba"],

112 - MambaSpec(

113 - block_size=mamba_block_size,

114 - shapes=(1, 1),

115 - dtypes=(torch.float32,),

116 - mamba_cache_mode="align",

117 - num_prefill_checkpoint_blocks=num_prefill_checkpoint_blocks,

118 - ),

119 - ),

125 + mamba_group,

126 ],

127 )

128 scheduler_block_size = lcm(

@@ -169,6 +175,7 @@ def test_mamba_align_split_partial_tail_schedule(dcp_world_size: int):

175 scheduler_block_size = block_size * dcp_world_size

176 hash_block_size = 32

177 mock = SimpleNamespace(

178 + block_size=scheduler_block_size,

179 cache_config=SimpleNamespace(block_size=block_size),

180 max_num_scheduled_tokens=8192,

181 scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),

@@ -217,6 +224,7 @@ def test_mamba_align_split_when_block_exceeds_scheduling_budget():

224 token_budget = 8192

225 prompt_length = 30000

226 mock = SimpleNamespace(

227 + block_size=block_size,

228 cache_config=SimpleNamespace(block_size=block_size),

229 max_num_scheduled_tokens=token_budget,

230 scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),

@@ -254,6 +262,7 @@ def test_mamba_align_split_when_block_exceeds_long_prefill_threshold():

262 long_prefill_threshold = 384

263 prompt_length = 1300

264 mock = SimpleNamespace(

265 + block_size=block_size,

266 cache_config=SimpleNamespace(block_size=block_size),

267 max_num_scheduled_tokens=token_budget,

268 scheduler_config=SimpleNamespace(

@@ -557,9 +566,9 @@ def test_partial_hit_then_internal_checkpoint_uses_distinct_mamba_blocks():

566 assert mamba_blocks[3].block_id == running_block_id

567

568

560 -def test_internal_checkpoint_requires_block_aligned_start():

561 - hash_block_size = 2

562 - mamba_block_size = 16

569 +def test_internal_checkpoint_uses_partial_hash_lifecycle():

570 + hash_block_size = 16

571 + mamba_block_size = 32

572 manager = make_full_mamba_manager(

573 dcp_world_size=1,

574 hash_block_size=hash_block_size,

@@ -567,23 +576,90 @@ def test_internal_checkpoint_requires_block_aligned_start():

576 mamba_block_size=mamba_block_size,

577 num_prefill_checkpoint_blocks=1,

578 )

570 - request = make_request("producer", list(range(50)), hash_block_size, sha256)

579 + request = make_request("producer", list(range(120)), hash_block_size, sha256)

580

572 - # Compute token 0 first, so the next query starts at token 1, which is not

573 - # aligned to the Mamba block size.

574 - assert manager.allocate_slots(request, 1) is not None

575 - request.num_computed_tokens = 1

576 - manager.new_step_starts()

581 + new_blocks = manager.allocate_slots(request, request.num_tokens)

582 +

583 + assert new_blocks is not None

584 + mamba_blocks = manager.get_blocks(request.request_id).blocks[1]

585 + checkpoint_idx = 2

586 + checkpoint_block = mamba_blocks[checkpoint_idx]

587 + running_block = mamba_blocks[checkpoint_idx + 1]

588 + assert not checkpoint_block.is_null

589 + assert not running_block.is_null

590 + assert checkpoint_block is not running_block

591 +

592 + # The internal block is exported at state@112, replacing its temporary

593 + # full-block state@96 key while retaining #52789's req_to_blocks ownership.

594 + partial_hash = request.block_hashes[112 // hash_block_size - 1]

595 + partial_hit = manager.block_pool.get_cached_block(partial_hash, [1])

596 + assert partial_hit is not None

597 + assert partial_hit[0] is checkpoint_block

598 + assert checkpoint_block.block_hash_num_tokens == 112

599 + full_hash = request.block_hashes[96 // hash_block_size - 1]

600 + assert manager.block_pool.get_cached_block(full_hash, [1]) is None

601 + assert drain_boundary_state_offloads(manager) == {

602 + request.request_id: [(1, checkpoint_block.block_id, 112)]

603 + }

604 +

605 + manager.free(request)

606 + replay = make_request("replay", list(range(120)), hash_block_size, sha256)

607 + _, num_computed, _ = manager.get_computed_blocks(replay)

608 + assert num_computed == 112

609 +

610 +

611 +def test_eagle_block_aligned_checkpoint_replaces_newer_hash():

612 + hash_block_size = mamba_block_size = 32

613 + manager = make_full_mamba_manager(

614 + dcp_world_size=1,

615 + hash_block_size=hash_block_size,

616 + full_block_size=hash_block_size,

617 + mamba_block_size=mamba_block_size,

618 + use_eagle=True,

619 + num_speculative_blocks=3,

620 + num_prefill_checkpoint_blocks=1,

621 + )

622 + request = make_request("producer", list(range(104)), hash_block_size, sha256)

623 +

624 + assert manager.allocate_slots(request, request.num_tokens) is not None

625 +

626 + mamba_blocks = manager.get_blocks(request.request_id).blocks[1]

627 + checkpoint_block = mamba_blocks[2]

628 + checkpoint_hash = request.block_hashes[64 // hash_block_size - 1]

629 + checkpoint_hit = manager.block_pool.get_cached_block(checkpoint_hash, [1])

630 + assert checkpoint_hit is not None

631 + assert checkpoint_hit[0] is checkpoint_block

632 + assert checkpoint_block.block_hash_num_tokens == 64

633 +

634 + overwritten_full_hash = request.block_hashes[96 // hash_block_size - 1]

635 + assert manager.block_pool.get_cached_block(overwritten_full_hash, [1]) is None

636 +

637 +

638 +def test_hash_aligned_query_end_uses_regular_partial_tail():

639 + hash_block_size = 2

640 + mamba_block_size = 4

641 + manager = make_full_mamba_manager(

642 + dcp_world_size=1,

643 + hash_block_size=hash_block_size,

644 + full_block_size=hash_block_size,

645 + mamba_block_size=mamba_block_size,

646 + num_prefill_checkpoint_blocks=1,

647 + )

648 + request = make_request("producer", list(range(14)), hash_block_size, sha256)

649

578 - new_blocks = manager.allocate_slots(request, 49)

650 + new_blocks = manager.allocate_slots(request, request.num_tokens)

651

652 assert new_blocks is not None

653 + mamba_manager = manager.coordinator.single_type_managers[1]

654 + assert mamba_manager._checkpoint_positions[request.request_id] == 0

655 mamba_blocks = manager.get_blocks(request.request_id).blocks[1]

582 - checkpoint_block_idx = 48 // mamba_block_size - 1

583 - assert mamba_blocks[checkpoint_block_idx].is_null

584 - assert not mamba_blocks[-1].is_null

585 - checkpoint_hash = request.block_hashes[48 // hash_block_size - 1]

586 - assert manager.block_pool.get_cached_block(checkpoint_hash, [1]) is None

656 + assert sum(not block.is_null for block in mamba_blocks) == 1

657 +

658 + partial_hash = request.block_hashes[14 // hash_block_size - 1]

659 + cached = manager.block_pool.get_cached_block(partial_hash, [1])

660 + assert cached is not None

661 + assert cached[0] is mamba_blocks[14 // mamba_block_size]

662 + assert cached[0].block_hash_num_tokens == 14

663

664

665 def test_external_mamba_hit_same_block_uses_running_cow_on_continue():

@@ -699,18 +775,9 @@ def test_boundary_state_offloads_returns_cow_target():

775 computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)

776 assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None

777

702 - # Step A offers the materialized aligned checkpoint immediately.

703 - ((group_id, block_id, boundary_tokens),) = drain_boundary_state_offloads(manager)[

704 - "0"

705 - ]

706 - assert group_id == 1

707 - assert boundary_tokens == 4

708 - aligned_hash = req0.block_hashes[4 // hash_block_size - 1]

709 - aligned_block = manager.block_pool.get_cached_block(

710 - aligned_hash, kv_cache_group_ids=[1]

711 - )

712 - assert aligned_block is not None

713 - assert block_id == aligned_block[0].block_id

778 + # The final hash-aligned boundary uses the regular partial-tail path rather

779 + # than allocating an internal checkpoint at the query end.

780 + assert drain_boundary_state_offloads(manager) == {}

781

782 partial_mamba_hash = req0.block_hashes[6 // hash_block_size - 1]

783 source_block = manager.block_pool.get_cached_block(

tests/v1/core/test_mamba_align_chunk_split.py

+92 -13

@@ -8,6 +8,7 @@ the wrong offset, and a later chunk crossing that boundary publishes it anyway.

8 Requests resuming from it then restore a truncated state (#43559).

9 """

11 +import sys

12 from types import SimpleNamespace

14 import pytest

@@ -21,6 +22,7 @@ from vllm.v1.kv_cache_interface import (

22 KVCacheConfig,

23 KVCacheGroupSpec,

24 MambaSpec,

25 + is_mamba_prefill_checkpoint_valid,

26 )

27 from vllm.v1.request import Request

@@ -37,6 +39,49 @@ PROMPT_LEN = 2002

39 MAMBA_GROUP_ID = 1

42 +@pytest.mark.parametrize(

43 + (

44 + "query_start",

45 + "query_end",

46 + "checkpoint_position",

47 + "hash_size",

48 + "block_size",

49 + "alignment",

50 + "expected",

51 + ),

52 + [

53 + (0, 100, 88, 8, 64, 16, False),

54 + (0, 100, 96, 8, 64, 16, True),

55 + (0, 100, 96, 8, 64, None, False),

56 + (8, 100, 96, 8, 64, 16, False),

57 + (15360, 16896, 16768, 128, 1536, 16, False),

58 + (15360, 16960, 16896, 128, 1536, 16, True),

59 + (15360, 16960, 16768, 128, 1536, 16, True),

60 + (16064, 18433, 18432, 128, 1536, 16, False),

61 + ],

62 +)

63 +def test_mamba_prefill_checkpoint_valid(

64 + query_start: int,

65 + query_end: int,

66 + checkpoint_position: int,

67 + hash_size: int,

68 + block_size: int,

69 + alignment: int | None,

70 + expected: bool,

71 +) -> None:

72 + assert (

73 + is_mamba_prefill_checkpoint_valid(

74 + query_start=query_start,

75 + query_end=query_end,

76 + checkpoint_position=checkpoint_position,

77 + hash_block_size=hash_size,

78 + mamba_block_size=block_size,

79 + checkpoint_alignment=alignment,

80 + )

81 + is expected

82 + )

83 +

84 +

85 def _make_hybrid_kv_cache_manager(

86 num_prefill_checkpoint_blocks: int = 0,

87 ) -> KVCacheManager:

@@ -62,6 +107,9 @@ def _make_hybrid_kv_cache_manager(

107 mamba_cache_mode="align",

108 num_speculative_blocks=NUM_SPEC,

109 num_prefill_checkpoint_blocks=num_prefill_checkpoint_blocks,

110 + prefill_checkpoint_alignment=(

111 + 16 if num_prefill_checkpoint_blocks > 0 else None

112 + ),

113 ),

114 ),

115 ],

@@ -83,20 +131,23 @@ def _split(

131 use_eagle_block_drop: bool | None = None,

132 partial_hit: bool = False,

133 num_prefill_checkpoint_blocks: int = 0,

134 + max_num_scheduled_tokens: int = 16384,

135 ) -> int:

136 """Call the real `Scheduler._mamba_block_aligned_split` on a stub self."""

137 if use_eagle_block_drop is None:

138 use_eagle_block_drop = use_eagle

139 stub = SimpleNamespace(

140 + block_size=MAMBA_BLOCK_SIZE,

141 cache_config=SimpleNamespace(block_size=MAMBA_BLOCK_SIZE),

142 use_eagle_block_drop=use_eagle_block_drop,

93 - max_num_scheduled_tokens=16384,

143 + max_num_scheduled_tokens=max_num_scheduled_tokens,

144 scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),

145 # `prefix_match_unit` finer than the block size (#46384).

146 mamba_partial_cache_hit=partial_hit,

147 hash_block_size=ATTN_BLOCK_SIZE,

98 - mamba_has_prefill_checkpoint_blocks=(

99 - num_prefill_checkpoint_blocks > 0 and not use_eagle

148 + mamba_has_prefill_checkpoint_blocks=(num_prefill_checkpoint_blocks > 0),

149 + mamba_prefill_checkpoint_alignment=(

150 + 16 if num_prefill_checkpoint_blocks > 0 else None

151 ),

152 )

153 return Scheduler._mamba_block_aligned_split(stub, request, num_new_tokens)

@@ -106,8 +157,11 @@ def _split(

157 ("prompt_len", "num_new_tokens", "use_eagle", "expected"),

158 [

159 (2002, 2002, False, 2002),

109 - (3602, 2000, False, 2000),

110 - (3602, 3602, True, MAMBA_BLOCK_SIZE),

160 + (3602, 2000, False, MAMBA_BLOCK_SIZE),

161 + (3602, 3602, True, 3602),

162 + (2002, 2002, True, 2002),

163 + (3602, 512, True, 0),

164 + (3602, 1700, True, MAMBA_BLOCK_SIZE),

165 ],

166 )

167 def test_internal_checkpoint_split(

@@ -123,12 +177,39 @@ def test_internal_checkpoint_split(

177 )

178 == expected

179 )

126 - if not use_eagle:

127 - manager = _make_hybrid_kv_cache_manager(num_prefill_checkpoint_blocks=1)

128 - assert manager.allocate_slots(request, expected) is not None

129 - mamba_manager = manager.coordinator.single_type_managers[MAMBA_GROUP_ID]

130 - blocks = mamba_manager.req_to_blocks[request.request_id]

131 - assert all(not block.is_null for block in blocks) # checkpoint + running state

180 +

181 +

182 +def test_partial_checkpoint_resume_stops_at_mamba_block_boundary() -> None:

183 + prompt_len = 3602

184 + resume_at = 1984

185 + (request,) = create_requests(1, num_tokens=prompt_len, block_size=ATTN_BLOCK_SIZE)

186 + request.num_computed_tokens = resume_at

187 +

188 + # Without an internal checkpoint export, the forward must stop at the

189 + # physical Mamba boundary so it does not publish a stale running state.

190 + assert (

191 + _split(

192 + request,

193 + prompt_len - resume_at,

194 + use_eagle=False,

195 + partial_hit=True,

196 + )

197 + == MAMBA_BLOCK_SIZE - resume_at % MAMBA_BLOCK_SIZE

198 + )

199 +

200 + # A hash-aligned resume inside a physical Mamba block cannot use the fixed

201 + # internal-checkpoint column: it aliases the initial-state column. Enabling

202 + # checkpoints must therefore preserve the normal aligned split.

203 + assert (

204 + _split(

205 + request,

206 + prompt_len - resume_at,

207 + use_eagle=False,

208 + partial_hit=True,

209 + num_prefill_checkpoint_blocks=1,

210 + )

211 + == MAMBA_BLOCK_SIZE - resume_at % MAMBA_BLOCK_SIZE

212 + )

213

214

215 def test_disabling_eagle_block_drop_keeps_the_trailing_cache_boundary() -> None:

@@ -252,8 +333,6 @@ def test_poisoning_is_block_size_independent(

333 budgets: list[int],

334 ) -> None:

335 """The invariant is per-block, so large mamba blocks are not safer."""

255 - import sys

256 -

336 monkeypatch.setattr(sys.modules[__name__], "MAMBA_BLOCK_SIZE", block_size)

337 assert _prefill(prompt_len, budgets=budgets) > 0

338

tests/v1/core/test_prefix_caching.py

@@ -1162,6 +1162,7 @@ def test_hybrid_cache_mamba_align_shared_prefix_detection():

1162 # junction (absolute) and stops the chunk there (req_2 has num_computed 0).

1163 # Create minimal mock with just the needed attributes

1164 mock = SimpleNamespace(

1165 + block_size=block_size,

1166 cache_config=SimpleNamespace(block_size=block_size),

1167 max_num_scheduled_tokens=3 * block_size,

1168 scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),

vllm/models/kimi_k3/nvidia/kda.py

@@ -609,6 +609,9 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):

609 return replace(

610 spec,

611 num_prefill_checkpoint_blocks=int(self.kda_prefill_backend == "flashkda"),

612 + prefill_checkpoint_alignment=(

613 + 16 if self.kda_prefill_backend == "flashkda" else None

614 + ),

615 )

616

617 def forward(

vllm/models/kimi_k3/nvidia/kda_metadata.py

+35 -15

@@ -20,6 +20,7 @@ import torch

20 from vllm.config import VllmConfig

21 from vllm.platforms import current_platform

22 from vllm.triton_utils import tl, triton

23 +from vllm.utils.math_utils import cdiv

24 from vllm.utils.torch_utils import async_tensor_h2d

25 from vllm.v1.attention.backend import CommonAttentionMetadata

26 from vllm.v1.attention.backends.gdn_attn import (

@@ -36,7 +37,11 @@ from vllm.v1.attention.backends.utils import (

37 compute_causal_conv1d_metadata,

38 split_decodes_and_prefills,

39 )

39 -from vllm.v1.kv_cache_interface import MambaSpec

40 +from vllm.v1.kv_cache_interface import (

41 + MambaSpec,

42 + get_mamba_prefill_checkpoint_position,

43 + is_mamba_prefill_checkpoint_valid,

44 +)

46 if TYPE_CHECKING:

47 from vllm.models.kimi_k3.nvidia.ops.recoverssm import (

@@ -44,9 +49,6 @@ if TYPE_CHECKING:

49 )

47 -FLASHKDA_CHUNK_SIZE = 16

48 -

49 -

52 @cache

53 def _metadata_launch_pdl() -> bool:

54 return current_platform.is_arch_support_pdl()

@@ -582,11 +584,11 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):

584 has_initial_state = None

585

586 checkpoint = None

585 - if (

586 - num_prefills > 0

587 + checkpoint_enabled = (

588 + self.vllm_config.cache_config.mamba_cache_mode == "align"

589 and self.kv_cache_spec.num_prefill_checkpoint_blocks > 0

588 - and self.vllm_config.cache_config.mamba_cache_mode == "align"

589 - ):

590 + )

591 + if num_prefills > 0 and checkpoint_enabled:

592 # prepare checkpoint metadata

593 assert m.seq_lens_cpu_upper_bound is not None

594 request_rows = list(range(m.num_reqs))

@@ -597,21 +599,39 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):

599 query_lens = [all_query_lens[row] for row in request_rows]

600 seq_lens = m.seq_lens_cpu_upper_bound.tolist()

601 block_size = self.kv_cache_spec.block_size

602 + hash_block_size = (

603 + self.vllm_config.cache_config.prefix_match_unit or block_size

604 + )

605 + speculative_config = self.vllm_config.speculative_config

606 + drop_eagle_block = (

607 + speculative_config is not None

608 + and speculative_config.use_eagle_block_drop()

609 + )

610 checkpoint_splits = []

611 checkpoint_cols = []

612 for row, query_len in zip(request_rows, query_lens):

613 seq_len = seq_lens[row]

604 - offset = seq_len // block_size * block_size - (seq_len - query_len)

605 - # offset should be less than query_len

606 - valid = (

607 - seq_len % block_size != 0

608 - and 0 < offset < query_len

609 - and offset % FLASHKDA_CHUNK_SIZE == 0

614 + query_start = seq_len - query_len

615 + checkpoint_position = get_mamba_prefill_checkpoint_position(

616 + seq_len,

617 + hash_block_size,

618 + drop_eagle_block=drop_eagle_block,

619 + )

620 + offset = checkpoint_position - query_start

621 + valid = is_mamba_prefill_checkpoint_valid(

622 + query_start=query_start,

623 + query_end=seq_len,

624 + checkpoint_position=checkpoint_position,

625 + hash_block_size=hash_block_size,

626 + mamba_block_size=block_size,

627 + checkpoint_alignment=(

628 + self.kv_cache_spec.prefill_checkpoint_alignment

629 + ),

630 )

631 offset = offset if valid else 0

632 first_len = offset or query_len

633 checkpoint_splits.append((first_len, query_len - first_len))

614 - checkpoint_cols.append(seq_len // block_size - 1 if valid else -1)

634 + checkpoint_cols.append(cdiv(seq_len, block_size) - 2 if valid else -1)

635 if any(tail for _, tail in checkpoint_splits):

636 checkpoint_offsets_tensor = async_tensor_h2d(

637 [first if tail else 0 for first, tail in checkpoint_splits],

vllm/v1/core/block_pool.py

+12 -3

@@ -450,6 +450,7 @@ class BlockPool:

450 num_tokens: int,

451 kv_cache_group_id: int,

452 block_size: int,

453 + replace_existing_hashes: bool = False,

454 ) -> BlockHashWithGroupId | None:

455 """Register a partial prefix-cache entry for an existing block.

456

@@ -477,6 +478,9 @@ class BlockPool:

478 entry hash itself is always the prefix-chain hash at

479 ``num_tokens``; ``block_size`` is used to assert that the

480 entry is partial within the owning cache block.

481 + replace_existing_hashes: Whether the block contents were replaced

482 + and all existing cache entries must be removed before the new

483 + entry is registered.

484

485 Returns:

486 The hash key with group ID if a partial entry can be registered;

@@ -485,9 +489,10 @@ class BlockPool:

489 if block.is_null:

490 return None

491

488 - assert block_size > self.hash_block_size

492 assert block_size % self.hash_block_size == 0

490 - assert num_tokens % block_size != 0

493 + assert replace_existing_hashes or (

494 + block_size > self.hash_block_size and num_tokens % block_size != 0

495 + )

496 block_hash = self._get_partial_block_hash(request, num_tokens)

497 num_hash_blocks = num_tokens // self.hash_block_size

498 block_hash_with_group_id = make_block_hash_with_group_id(

@@ -498,7 +503,11 @@ class BlockPool:

503 block_hash_with_group_id, block.block_id

504 )

505 )

501 - if (

506 + if replace_existing_hashes:

507 + removed_hashes = self._remove_cached_block_hashes(block)

508 + self._emit_block_removed_events(removed_hashes)

509 + already_cached = False

510 + elif (

511 not already_cached

512 and block.block_hash is not None

513 and block.block_hash_num_tokens is not None

vllm/v1/core/kv_cache_coordinator.py

@@ -15,6 +15,7 @@ from vllm.v1.core.kv_cache_utils import (

15 )

16 from vllm.v1.core.single_type_kv_cache_manager import (

17 CrossAttentionManager,

18 + MambaManager,

19 SingleTypeKVCacheManager,

20 get_manager_for_kv_cache_spec,

21 )

@@ -151,6 +152,11 @@ class KVCacheCoordinator(ABC):

152 )

153 for i, kv_cache_group in enumerate(self.kv_cache_config.kv_cache_groups)

154 )

155 + # Match Mamba checkpoints to Eagle's attention replay boundary.

156 + if use_eagle:

157 + for manager in self.single_type_managers:

158 + if isinstance(manager, MambaManager):

159 + manager.drop_eagle_checkpoint_block = True

160

161 # A positive retention interval must be a multiple of the base hit granularity

162 # (``scheduler_block_size``) to land on real cache-hit boundaries.

vllm/v1/core/sched/scheduler.py

+39 -13

@@ -54,7 +54,12 @@ from vllm.v1.core.sched.request_queue import (

54 )

55 from vllm.v1.core.sched.utils import check_stop, remove_all

56 from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs

57 -from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec

57 +from vllm.v1.kv_cache_interface import (

58 + KVCacheConfig,

59 + MambaSpec,

60 + get_mamba_prefill_checkpoint_position,

61 + is_mamba_prefill_checkpoint_valid,

62 +)

63 from vllm.v1.metrics.perf import ModelMetrics, PerfStats

64 from vllm.v1.metrics.stats import (

65 PrefixCacheStats,

@@ -334,15 +339,20 @@ class Scheduler(SchedulerInterface):

339 self.need_mamba_block_aligned_split = (

340 self.has_mamba_layers and self.cache_config.mamba_cache_mode == "align"

341 )

337 - self.mamba_has_prefill_checkpoint_blocks = (

338 - self.has_mamba_layers

339 - # TODO: support spec decoding

340 - and not self.use_eagle

341 - and all(

342 - not isinstance(group.kv_cache_spec, MambaSpec)

343 - or group.kv_cache_spec.num_prefill_checkpoint_blocks > 0

342 + # TODO: Support models with multiple Mamba specs that require different

343 + # prefill checkpoint alignments instead of selecting the first one.

344 + self.mamba_prefill_checkpoint_alignment = next(

345 + (

346 + group.kv_cache_spec.prefill_checkpoint_alignment

347 for group in kv_cache_config.kv_cache_groups

345 - )

348 + if isinstance(group.kv_cache_spec, MambaSpec)

349 + ),

350 + None,

351 + )

352 + self.mamba_has_prefill_checkpoint_blocks = self.has_mamba_layers and all(

353 + not isinstance(group.kv_cache_spec, MambaSpec)

354 + or group.kv_cache_spec.num_prefill_checkpoint_blocks > 0

355 + for group in kv_cache_config.kv_cache_groups

356 )

357 # A finer prefix_match_unit is configured: a mamba partial tail entry

358 # can only be registered by a step ending exactly at the prompt's last

@@ -426,8 +436,22 @@ class Scheduler(SchedulerInterface):

436 last_cache_position = max(last_cache_position - block_size, 0)

437

438 end = start + num_new_tokens

439 + checkpoint_position = get_mamba_prefill_checkpoint_position(

440 + prefill_end,

441 + self.hash_block_size,

442 + drop_eagle_block=self.use_eagle_block_drop,

443 + )

444 use_internal_checkpoint = (

430 - self.mamba_has_prefill_checkpoint_blocks and start % block_size == 0

445 + self.mamba_has_prefill_checkpoint_blocks

446 + and end >= prefill_end

447 + and is_mamba_prefill_checkpoint_valid(

448 + query_start=start,

449 + query_end=end,

450 + checkpoint_position=checkpoint_position,

451 + hash_block_size=self.hash_block_size,

452 + mamba_block_size=block_size,

453 + checkpoint_alignment=self.mamba_prefill_checkpoint_alignment,

454 + )

455 )

456 if use_internal_checkpoint:

457 last_cache_position = 0

@@ -436,7 +460,7 @@ class Scheduler(SchedulerInterface):

460 # aligned. Exempt: the prompt's last chunk, whose slot decode advances

461 # to the boundary. A block too wide for one chunk advances sub-block

462 # and re-aligns at the next boundary.

439 - if end < prefill_end and not use_internal_checkpoint:

463 + if end < prefill_end:

464 max_prefill_tokens = self.max_num_scheduled_tokens

465 long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold

466 if long_prefill_threshold > 0:

@@ -448,13 +472,15 @@ class Scheduler(SchedulerInterface):

472 next_block_boundary = (start // block_size + 1) * block_size

473 tail_boundary = (

474 request.num_prompt_tokens // self.hash_block_size * self.hash_block_size

451 - if self.mamba_partial_cache_hit

475 + if self.mamba_partial_cache_hit and not use_internal_checkpoint

476 else 0

477 )

478 stops = (

479 # Same invariant: a chunk starting mid-block stops at the boundary

480 # rather than running past it.

457 - next_block_boundary if start % block_size != 0 else 0,

481 + next_block_boundary

482 + if start % block_size != 0 and not use_internal_checkpoint

483 + else 0,

484 # Never run past the last cacheable block boundary mid-chunk.

485 last_cache_position,

486 # Fine-grained hits: the prompt's partial-tail entry can only be

vllm/v1/core/single_type_kv_cache_manager.py

+80 -24

@@ -31,6 +31,8 @@ from vllm.v1.kv_cache_interface import (

31 SinkFullAttentionSpec,

32 SlidingWindowMLASpec,

33 SlidingWindowSpec,

34 + get_mamba_prefill_checkpoint_position,

35 + is_mamba_prefill_checkpoint_valid,

36 )

37 from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry

38 from vllm.v1.request import Request

@@ -1399,6 +1401,12 @@ class MambaManager(SingleTypeKVCacheManager):

1401 self.block_size = kv_cache_spec.block_size

1402 self.mamba_cache_mode = kv_cache_spec.mamba_cache_mode

1403 self.num_speculative_blocks: int = kv_cache_spec.num_speculative_blocks

1404 + self.has_prefill_checkpoint_blocks = (

1405 + self.mamba_cache_mode == "align"

1406 + and kv_cache_spec.num_prefill_checkpoint_blocks > 0

1407 + )

1408 + # Mamba checkpoints follow Eagle's global replay boundary.

1409 + self.drop_eagle_checkpoint_block = False

1410 self.cached_blocks_this_step: set[BlockHashWithGroupId] = set()

1411 if self.mamba_cache_mode == "align":

1412 # Mapping from request ID to the index of the block

@@ -1406,9 +1414,9 @@ class MambaManager(SingleTypeKVCacheManager):

1414 self.last_state_block_idx: dict[str, int] = {}

1415 # The set of the requests that have been allocated blocks

1416 self._allocated_block_reqs: set[str] = set()

1409 - # Number of internal checkpoint blocks required by each request's

1417 + # Absolute checkpoint position selected by the scheduler for the

1418 # current allocation.

1411 - self._num_checkpoint_blocks: dict[str, int] = {}

1419 + self._checkpoint_positions: dict[str, int] = {}

1420 # Requests that registered their own last-prompt-boundary partial

1421 # tail (producers). A later CoW hands its private copy to the

1422 # connector; a request that finishes first hands off this table

@@ -1592,18 +1600,32 @@ class MambaManager(SingleTypeKVCacheManager):

1600 def _needs_internal_checkpoint(

1601 self,

1602 request_id: str,

1595 - num_tokens: int,

1596 - num_computed_tokens: int,

1603 + query_start: int,

1604 + query_end: int,

1605 + checkpoint_position: int,

1606 ) -> bool:

1607 assert isinstance(self.kv_cache_spec, MambaSpec)

1599 - checkpoint_idx = cdiv(num_tokens, self.block_size) - 2

1608 + checkpoint_idx = cdiv(query_end, self.block_size) - 2

1609 blocks = self.req_to_blocks[request_id]

1610 return (

1602 - self.kv_cache_spec.num_prefill_checkpoint_blocks > 0

1603 - and num_tokens % self.block_size != 0

1604 - and num_computed_tokens % self.block_size == 0

1611 + self.has_prefill_checkpoint_blocks

1612 + and is_mamba_prefill_checkpoint_valid(

1613 + query_start=query_start,

1614 + query_end=query_end,

1615 + checkpoint_position=checkpoint_position,

1616 + hash_block_size=self.block_pool.hash_block_size,

1617 + mamba_block_size=self.block_size,

1618 + checkpoint_alignment=(self.kv_cache_spec.prefill_checkpoint_alignment),

1619 + )

1620 and checkpoint_idx >= 0

1606 - and (checkpoint_idx >= len(blocks) or blocks[checkpoint_idx].is_null)

1621 + and (

1622 + checkpoint_idx >= len(blocks)

1623 + or blocks[checkpoint_idx].is_null

1624 + or (

1625 + request_id in self._allocated_block_reqs

1626 + and checkpoint_idx >= len(blocks) - self.num_speculative_blocks

1627 + )

1628 + )

1629 )

1630

1631 def get_num_blocks_to_allocate(

@@ -1668,17 +1690,27 @@ class MambaManager(SingleTypeKVCacheManager):

1690 )

1691 if has_partial_hit:

1692 num_new_blocks = max(num_new_blocks, 0) + 1

1671 - checkpoint_block = int(

1672 - self._needs_internal_checkpoint(

1673 - request_id, num_tokens, total_computed_tokens

1674 - )

1693 + checkpoint_position = get_mamba_prefill_checkpoint_position(

1694 + num_tokens,

1695 + self.block_pool.hash_block_size,

1696 + self.drop_eagle_checkpoint_block,

1697 )

1698 + if not self._needs_internal_checkpoint(

1699 + request_id,

1700 + total_computed_tokens,

1701 + num_tokens,

1702 + checkpoint_position,

1703 + ):

1704 + checkpoint_position = 0

1705 + checkpoint_block = int(checkpoint_position > 0)

1706 if not apply_admission_cap:

1677 - self._num_checkpoint_blocks[request_id] = checkpoint_block

1707 + self._checkpoint_positions[request_id] = checkpoint_position

1708 if num_new_blocks > 0:

1679 - num_new_blocks = 1 + int(has_partial_hit) + checkpoint_block

1680 - if request_id not in self._allocated_block_reqs:

1681 - num_new_blocks += self.num_speculative_blocks

1709 + blocks_allocated = request_id in self._allocated_block_reqs

1710 + if not (checkpoint_block and blocks_allocated):

1711 + num_new_blocks = 1 + int(has_partial_hit) + checkpoint_block

1712 + if not blocks_allocated:

1713 + num_new_blocks += self.num_speculative_blocks

1714

1715 num_evictable_computed_blocks = self._get_num_evictable_blocks(

1716 new_computed_blocks

@@ -1710,12 +1742,17 @@ class MambaManager(SingleTypeKVCacheManager):

1742 num_required_blocks = (

1743 cdiv(num_tokens, self.block_size) + self.num_speculative_blocks

1744 )

1713 - checkpoint_block = self._num_checkpoint_blocks.get(request_id, 0)

1745 + checkpoint_position = self._checkpoint_positions.get(request_id, 0)

1746 + checkpoint_block = int(checkpoint_position > 0)

1747 partial_hit = self._partial_hit_reqs.get(request_id)

1748 has_partial_hit = partial_hit is not None

1749 # `num_required_blocks` might be less than `len(req_blocks)` if blocks are

1750 # over-allocated at last round.

1718 - if num_required_blocks <= len(req_blocks) and not has_partial_hit:

1751 + if (

1752 + num_required_blocks <= len(req_blocks)

1753 + and not has_partial_hit

1754 + and not checkpoint_block

1755 + ):

1756 self._allocated_block_reqs.add(request_id)

1757 return []

1758 else:

@@ -1745,7 +1782,7 @@ class MambaManager(SingleTypeKVCacheManager):

1782 [self._null_block for _ in range(prev_block_len, null_end)]

1783 )

1784

1748 - if blocks_allocated:

1785 + if blocks_allocated and not checkpoint_block:

1786 # Relocate exclusively owned speculative scratch blocks.

1787 for block_idx in range(

1788 prev_block_len - self.num_speculative_blocks, prev_block_len

@@ -1754,11 +1791,11 @@ class MambaManager(SingleTypeKVCacheManager):

1791 self._relocate_speculative_block(req_blocks, block_idx)

1792 else:

1793 break

1757 - num_new_blocks = num_required_blocks - len(req_blocks)

1794 + num_new_blocks = max(num_required_blocks - len(req_blocks), 0)

1795 if has_partial_hit:

1796 num_new_blocks = max(num_new_blocks, 0) + 1

1797 max_new_blocks = 1 + int(has_partial_hit) + checkpoint_block

1761 - if not blocks_allocated:

1798 + if not blocks_allocated or checkpoint_block:

1799 max_new_blocks += self.num_speculative_blocks

1800 assert num_new_blocks <= max_new_blocks

1801 new_blocks = self.block_pool.get_new_blocks(num_new_blocks)

@@ -1838,7 +1875,7 @@ class MambaManager(SingleTypeKVCacheManager):

1875 if self.mamba_cache_mode == "align":

1876 self._allocated_block_reqs.discard(request_id)

1877 self.last_state_block_idx.pop(request_id, None)

1841 - self._num_checkpoint_blocks.pop(request_id, None)

1878 + self._checkpoint_positions.pop(request_id, None)

1879 self._producer_partial_tail_reqs.pop(request_id, None)

1880 # An offer is only guaranteed to hold committed bytes until the end

1881 # of the pass that made it. This request's blocks are going back to

@@ -1885,6 +1922,7 @@ class MambaManager(SingleTypeKVCacheManager):

1922 continue

1923 self.cached_blocks_this_step.add(block.block_hash)

1924 if self.mamba_cache_mode == "align":

1925 + assert block.block_hash_num_tokens is not None

1926 # Offer every retained boundary with its exact block.

1927 # The connector filters against its save window, which may

1928 # extend past the original prompt during resumed prefill.

@@ -1893,7 +1931,7 @@ class MambaManager(SingleTypeKVCacheManager):

1931 request.request_id,

1932 self.kv_cache_group_id,

1933 block,

1896 - (idx + 1) * self.block_size,

1934 + block.block_hash_num_tokens,

1935 )

1936 )

1937

@@ -1906,6 +1944,24 @@ class MambaManager(SingleTypeKVCacheManager):

1944 num_tokens: int,

1945 ) -> BlockHashWithGroupId | None:

1946 hash_block_size = self.block_pool.hash_block_size

1947 + # Re-key the reserved block at its exported checkpoint boundary.

1948 + checkpoint_position = self._checkpoint_positions.get(request.request_id, 0)

1949 + if checkpoint_position > 0:

1950 + # TODO: Store the reserved slot explicitly for multi-module MTP.

1951 + checkpoint_idx = cdiv(num_tokens, self.block_size) - 2

1952 + blocks = self.req_to_blocks[request.request_id]

1953 + assert 0 <= checkpoint_idx < len(blocks)

1954 + checkpoint_block = blocks[checkpoint_idx]

1955 + if checkpoint_block.block_hash_num_tokens == checkpoint_position:

1956 + return None

1957 + return self.block_pool.cache_partial_block(

1958 + request=request,

1959 + block=checkpoint_block,

1960 + num_tokens=checkpoint_position,

1961 + kv_cache_group_id=self.kv_cache_group_id,

1962 + block_size=self.block_size,

1963 + replace_existing_hashes=True,

1964 + )

1965 if self.block_size == hash_block_size:

1966 return None

1967 if num_tokens % self.block_size == 0:

vllm/v1/kv_cache_interface.py

+38

@@ -890,6 +890,7 @@ class MambaSpec(KVCacheSpec):

890 mamba_cache_mode: str = "none"

891 num_speculative_blocks: int = 0

892 num_prefill_checkpoint_blocks: int = 0

893 + prefill_checkpoint_alignment: int | None = None

894 num_heads: int = 1

895 tokens_per_state: int = -1

896 # False: the state is sharded across TP ranks (e.g. GDN). True: every TP

@@ -950,12 +951,49 @@ class MambaSpec(KVCacheSpec):

951 isinstance(spec, MambaSpec)

952 and spec.num_speculative_blocks == self.num_speculative_blocks

953 and spec.num_prefill_checkpoint_blocks == self.num_prefill_checkpoint_blocks

954 + and spec.prefill_checkpoint_alignment == self.prefill_checkpoint_alignment

955 and spec.page_size_bytes == self.page_size_bytes

956 and spec.tp_replicated == self.tp_replicated

957 for spec in kv_cache_specs.values()

958 )

959

960

961 +def get_mamba_prefill_checkpoint_position(

962 + num_tokens: int,

963 + hash_block_size: int,

964 + drop_eagle_block: bool,

965 +) -> int:

966 + """Return the reusable Mamba checkpoint boundary for a prefill."""

967 + checkpoint_position = (num_tokens - 1) // hash_block_size * hash_block_size

968 + if drop_eagle_block:

969 + checkpoint_position -= hash_block_size

970 + return max(checkpoint_position, 0)

971 +

972 +

973 +def is_mamba_prefill_checkpoint_valid(

974 + query_start: int,

975 + query_end: int,

976 + checkpoint_position: int,

977 + hash_block_size: int,

978 + mamba_block_size: int,

979 + checkpoint_alignment: int | None,

980 +) -> bool:

981 + """Whether a backend can export the checkpoint in this query."""

982 + if checkpoint_alignment is None:

983 + return False

984 + assert checkpoint_alignment > 0

985 +

986 + initial_state_col = (query_start - 1) // mamba_block_size

987 + checkpoint_col = cdiv(query_end, mamba_block_size) - 2

988 + return (

989 + query_start % hash_block_size == 0

990 + and checkpoint_col > initial_state_col

991 + and query_start + hash_block_size <= checkpoint_position

992 + and query_start < checkpoint_position < query_end

993 + and (checkpoint_position - query_start) % checkpoint_alignment == 0

994 + )

995 +

996 +

997 @dataclass(frozen=True)

998 class EncoderOnlyAttentionSpec(AttentionSpec):

999 def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: