@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 9eb9d9d395

[Perf] Support internal prefill checkpoints for Mamba prefix caching, 9%~25% TTFT improvement (#52789)

Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Ziming Huang <zelda.huanghuang@gmail.com> Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: Ziming Huang <zelda.huanghuang@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>

Wentao Ye committed Aug 22, 2026 at 13:32 UTC 9eb9d9d3953959695108600c8ed33d36bc6a1e5f

12 files changed +562 -48

cmake/external_projects/flashkda.cmake

+1 -1

@@ -13,7 +13,7 @@ else()

13 FetchContent_Declare(

14 flashkda

15 GIT_REPOSITORY https://github.com/vllm-project/FlashKDA.git

16 - GIT_TAG 053de1b716ef3255873e02d2d28f4adf09951978

16 + GIT_TAG ee0be888cd0e972f9409bf53756f8c38c6652173

17 GIT_PROGRESS TRUE

18 GIT_SUBMODULES cutlass

19 )

csrc/flashkda_registration.cpp

+2 -1

@@ -10,7 +10,8 @@ STABLE_TORCH_LIBRARY(_flashkda_C, m) {

10 "Tensor(a!) out, Tensor(c!) workspace, Tensor A_log, Tensor dt_bias, "

11 "float lower_bound, "

12 "Tensor? initial_state=None, Tensor(b!)? final_state=None, "

13 - "Tensor? cu_seqlens=None) -> ()");

13 + "Tensor? cu_seqlens=None, Tensor(d!)? checkpoint_state=None, "

14 + "Tensor? checkpoint_offsets=None) -> ()");

15 }

17 STABLE_TORCH_LIBRARY_IMPL(_flashkda_C, CompositeExplicitAutograd, m) {

tests/models/kimi_k3/test_kda.py

+81

@@ -21,6 +21,8 @@ from vllm.models.kimi_k3.amd.ops.third_party.kda import (

21 fused_recurrent_kda_packed_decode as fused_recurrent_kda_packed_decode_amd,

22 )

23 from vllm.models.kimi_k3.nvidia.kda import (

24 + _flashkda_prefill,

25 + _store_cache_checkpoints_kernel,

26 is_flashkda_supported,

27 is_fused_kda_decode_supported,

28 )

@@ -40,6 +42,7 @@ from vllm.models.kimi_k3.nvidia.ops.third_party.kda import (

42 )

43 from vllm.platforms import current_platform

44 from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd

45 +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID

47 DEVICE = current_platform.device_type

@@ -1088,6 +1091,17 @@ def test_flashkda_correctness():

1091 expected_states.append(final_state)

1092 expected_out = torch.cat(expected_outputs, dim=1)

1093 expected_state = torch.cat(expected_states).transpose(-1, -2).contiguous()

1094 + _, expected_checkpoint = naive_recurrent_kda(

1095 + q_norm[:, :16],

1096 + k_norm[:, :16],

1097 + v[:, :16],

1098 + gate[:, :16],

1099 + beta[:, :16],

1100 + initial_state=initial_state[0:1].transpose(-1, -2),

1101 + output_final_state=True,

1102 + )

1103 + assert expected_checkpoint is not None

1104 + expected_checkpoint = expected_checkpoint.transpose(-1, -2).contiguous()

1105

1106 actual_out = torch.empty_like(v)

1107 actual_state = torch.empty_like(initial_state)

@@ -1115,3 +1129,70 @@ def test_flashkda_correctness():

1129

1130 assert_close("o", expected_out, actual_out, 0.01)

1131 assert_close("ht", expected_state, actual_state, 0.01)

1132 +

1133 + checkpoint_out = torch.empty_like(v)

1134 + checkpoint_final_state = torch.empty_like(initial_state)

1135 + checkpoint_state = torch.empty_like(initial_state)

1136 + checkpoint_offsets = torch.tensor([16, 31], dtype=torch.int32, device=DEVICE)

1137 + _flashkda_prefill(

1138 + q=q,

1139 + k=k,

1140 + v=v,

1141 + g=raw_g,

1142 + beta=beta_logits,

1143 + A_log=A_log,

1144 + dt_bias=dt_bias,

1145 + lower_bound=lower_bound,

1146 + initial_state=initial_state,

1147 + cu_seqlens=cu_seqlens,

1148 + out=checkpoint_out,

1149 + final_state=checkpoint_final_state,

1150 + workspace=workspace,

1151 + checkpoint_state=checkpoint_state,

1152 + checkpoint_offsets=checkpoint_offsets,

1153 + )

1154 +

1155 + assert_close("checkpoint_o", expected_out, checkpoint_out, 0.01)

1156 + assert_close("checkpoint_ht", expected_state, checkpoint_final_state, 0.01)

1157 + assert_close("checkpoint", expected_checkpoint, checkpoint_state[:1], 0.01)

1158 +

1159 + conv_state = torch.zeros(2, H * D, 3, dtype=q.dtype, device=DEVICE)

1160 + recurrent_storage = torch.zeros(2, H * D * D + 8, device=DEVICE)

1161 + recurrent_state = recurrent_storage[:, : H * D * D].view(2, H, D, D)

1162 + conv_input = q[0].flatten(1)

1163 + checkpoint_state_indices = torch.tensor(

1164 + [1, NULL_BLOCK_ID], dtype=torch.int32, device=DEVICE

1165 + )

1166 + state_len = conv_state.shape[-1]

1167 + width = H * D

1168 + recurrent_row_size = checkpoint_state[0].numel()

1169 + block_size = 256

1170 + _store_cache_checkpoints_kernel[

1171 + (

1172 + checkpoint_state_indices.numel(),

1173 + (max(width * state_len, recurrent_row_size) + block_size - 1) // block_size,

1174 + )

1175 + ](

1176 + conv_input,

1177 + conv_state,

1178 + checkpoint_state,

1179 + recurrent_state,

1180 + cu_seqlens,

1181 + checkpoint_offsets,

1182 + checkpoint_state_indices,

1183 + conv_input.stride(0),

1184 + conv_input.stride(1),

1185 + conv_state.stride(0),

1186 + conv_state.stride(1),

1187 + conv_state.stride(2),

1188 + checkpoint_state.stride(0),

1189 + recurrent_state.stride(0),

1190 + checkpoint_offsets.stride(0),

1191 + state_len,

1192 + width,

1193 + recurrent_row_size,

1194 + NULL_BLOCK_ID,

1195 + block_size,

1196 + )

1197 + torch.testing.assert_close(conv_state[1], q[0, 13:16].flatten(1).transpose(0, 1))

1198 + torch.testing.assert_close(recurrent_state[1], checkpoint_state[0])

tests/models/kimi_k3/test_kda_metadata.py

+50

@@ -87,6 +87,7 @@ def _make_builder(

87 device: torch.device = DEVICE,

88 mamba_cache_mode: str = "none",

89 use_recoverssm: bool = False,

90 + num_prefill_checkpoint_blocks: int = 0,

91 ) -> AttentionMetadataBuilder:

92 vllm_config = create_vllm_config(

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

@@ -110,6 +111,7 @@ def _make_builder(

111 dtypes=(torch.float16,),

112 mamba_cache_mode=mamba_cache_mode,

113 num_speculative_blocks=(0 if use_recoverssm else num_speculative_tokens),

114 + num_prefill_checkpoint_blocks=num_prefill_checkpoint_blocks,

115 ),

116 layer_names=["layer.0"],

117 vllm_config=vllm_config,

@@ -121,6 +123,54 @@ def _make_builder(

123 return builder

124

125

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

127 +def test_internal_checkpoint_metadata_targets_last_aligned_boundary():

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

129 + batch = BatchSpec(seq_lens=[50, 32], query_lens=[50, 16])

130 + common_attn_metadata = create_common_attn_metadata(

131 + batch, BLOCK_SIZE, device, arange_block_indices=True

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

133 + actual = _make_builder(

134 + KimiK3KDAMetadataBuilder,

135 + num_speculative_tokens=0,

136 + full_cuda_graph=False,

137 + mamba_cache_mode="align",

138 + num_prefill_checkpoint_blocks=1,

139 + device=device,

140 + ).build(0, common_attn_metadata)

141 +

142 + assert actual.checkpoint is not None

143 + torch.testing.assert_close(

144 + actual.checkpoint.state_indices,

145 + torch.tensor([2, NULL_BLOCK_ID], dtype=torch.int32, device=device),

146 + )

147 + torch.testing.assert_close(

148 + actual.checkpoint.checkpoint_offsets,

149 + torch.tensor([48, 0], dtype=torch.int32, device=device),

150 + )

151 +

152 +

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

154 +def test_internal_checkpoint_metadata_skips_unaligned_offset():

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

156 + # The checkpoint block boundary is 48, but this query starts at token 1,

157 + # making the real checkpoint offset 47, which is not FlashKDA-aligned.

158 + batch = BatchSpec(seq_lens=[50], query_lens=[49])

159 + common_attn_metadata = create_common_attn_metadata(

160 + batch, BLOCK_SIZE, device, arange_block_indices=True

161 + ).replace(is_prefilling=torch.tensor([True]))

162 + actual = _make_builder(

163 + KimiK3KDAMetadataBuilder,

164 + num_speculative_tokens=0,

165 + full_cuda_graph=False,

166 + mamba_cache_mode="align",

167 + num_prefill_checkpoint_blocks=1,

168 + device=device,

169 + ).build(0, common_attn_metadata)

170 +

171 + assert actual.checkpoint is None

172 +

173 +

174 @pytest.mark.parametrize(

175 (

176 "batch",

tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py

+85

@@ -77,6 +77,7 @@ def make_full_mamba_manager(

77 mamba_block_size: int = 4,

78 num_blocks: int = 32,

79 use_eagle: bool = False,

80 + num_prefill_checkpoint_blocks: int = 0,

81 ):

82 kv_cache_config = KVCacheConfig(

83 num_blocks=num_blocks,

@@ -98,6 +99,7 @@ def make_full_mamba_manager(

99 shapes=(1, 1),

100 dtypes=(torch.float32,),

101 mamba_cache_mode="align",

102 + num_prefill_checkpoint_blocks=num_prefill_checkpoint_blocks,

103 ),

104 ),

105 ],

@@ -135,6 +137,7 @@ def test_mamba_align_split_partial_tail_schedule(dcp_world_size: int):

137 dcp_world_size=dcp_world_size,

138 scheduler_block_size=scheduler_block_size,

139 mamba_partial_cache_hit=True,

140 + mamba_has_prefill_checkpoint_blocks=False,

141 )

142 split = Scheduler._mamba_block_aligned_split

143

@@ -180,6 +183,7 @@ def test_mamba_align_split_when_block_exceeds_scheduling_budget():

183 use_eagle=False,

184 hash_block_size=32,

185 mamba_partial_cache_hit=False,

186 + mamba_has_prefill_checkpoint_blocks=False,

187 )

188 req = make_request("0", [0] * prompt_length, 32, sha256)

189 split = Scheduler._mamba_block_aligned_split

@@ -218,6 +222,7 @@ def test_mamba_align_split_when_block_exceeds_long_prefill_threshold():

222 use_eagle=False,

223 hash_block_size=32,

224 mamba_partial_cache_hit=False,

225 + mamba_has_prefill_checkpoint_blocks=False,

226 )

227 req = make_request("0", [0] * prompt_length, 32, sha256)

228 split = Scheduler._mamba_block_aligned_split

@@ -462,6 +467,85 @@ def test_hybrid_mamba_partial_tail_owner_uses_cow_on_continue():

467 assert moved[0].block_hash_num_tokens == 6

468

469

470 +def test_partial_hit_then_internal_checkpoint_uses_distinct_mamba_blocks():

471 + hash_block_size = 2

472 + mamba_block_size = 4

473 + manager = make_full_mamba_manager(

474 + dcp_world_size=1,

475 + hash_block_size=hash_block_size,

476 + full_block_size=hash_block_size,

477 + mamba_block_size=mamba_block_size,

478 + num_prefill_checkpoint_blocks=1,

479 + )

480 +

481 + owner = make_request("owner", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)

482 + computed_blocks, num_computed, _ = manager.get_computed_blocks(owner)

483 + assert manager.allocate_slots(owner, 6, num_computed, computed_blocks) is not None

484 + manager.free(owner)

485 + manager.new_step_starts()

486 +

487 + partial_hash = owner.block_hashes[2]

488 + partial_block = manager.block_pool.get_cached_block(partial_hash, [1])

489 + assert partial_block is not None

490 + partial_block_id = partial_block[0].block_id

491 +

492 + replay = make_request(

493 + "replay",

494 + [0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6],

495 + hash_block_size,

496 + sha256,

497 + )

498 + computed_blocks, num_computed, _ = manager.get_computed_blocks(replay)

499 + assert num_computed == 6

500 +

501 + realigned_blocks = manager.allocate_slots(replay, 2, num_computed, computed_blocks)

502 + assert realigned_blocks is not None

503 + mamba_cow_block_id = realigned_blocks.get_block_ids()[1][0]

504 + copies, retained = manager.take_kv_cache_block_copies()

505 + assert KVCacheBlockCopy(partial_block_id, mamba_cow_block_id) in copies

506 + manager.block_pool.free_blocks(retained)

507 +

508 + replay.num_computed_tokens = 8

509 + manager.new_step_starts()

510 + final_blocks = manager.allocate_slots(replay, 6)

511 + assert final_blocks is not None

512 +

513 + checkpoint_block_id, running_block_id = final_blocks.get_block_ids()[1]

514 + assert len({mamba_cow_block_id, checkpoint_block_id, running_block_id}) == 3

515 + mamba_blocks = manager.get_blocks(replay.request_id).blocks[1]

516 + assert mamba_blocks[2].block_id == checkpoint_block_id

517 + assert mamba_blocks[3].block_id == running_block_id

518 +

519 +

520 +def test_internal_checkpoint_requires_block_aligned_start():

521 + hash_block_size = 2

522 + mamba_block_size = 16

523 + manager = make_full_mamba_manager(

524 + dcp_world_size=1,

525 + hash_block_size=hash_block_size,

526 + full_block_size=hash_block_size,

527 + mamba_block_size=mamba_block_size,

528 + num_prefill_checkpoint_blocks=1,

529 + )

530 + request = make_request("producer", list(range(50)), hash_block_size, sha256)

531 +

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

533 + # aligned to the Mamba block size.

534 + assert manager.allocate_slots(request, 1) is not None

535 + request.num_computed_tokens = 1

536 + manager.new_step_starts()

537 +

538 + new_blocks = manager.allocate_slots(request, 49)

539 +

540 + assert new_blocks is not None

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

542 + checkpoint_block_idx = 48 // mamba_block_size - 1

543 + assert mamba_blocks[checkpoint_block_idx].is_null

544 + assert not mamba_blocks[-1].is_null

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

546 + assert manager.block_pool.get_cached_block(checkpoint_hash, [1]) is None

547 +

548 +

549 def test_external_mamba_hit_same_block_uses_running_cow_on_continue():

550 """An external mid-block hit must become a running request even when its

551 first continuation does not need another Mamba block."""

@@ -560,6 +644,7 @@ def test_take_partial_tail_offloads_returns_cow_target():

644 shapes=(1, 1),

645 dtypes=(torch.float32,),

646 mamba_cache_mode="align",

647 + num_prefill_checkpoint_blocks=1,

648 ),

649 ),

650 ],

tests/v1/core/test_mamba_align_chunk_split.py

+52 -3

@@ -37,7 +37,9 @@ PROMPT_LEN = 2002

37 MAMBA_GROUP_ID = 1

40 -def _make_hybrid_kv_cache_manager() -> KVCacheManager:

40 +def _make_hybrid_kv_cache_manager(

41 + num_prefill_checkpoint_blocks: int = 0,

42 +) -> KVCacheManager:

43 config = KVCacheConfig(

44 num_blocks=10000,

45 kv_cache_tensors=[],

@@ -59,6 +61,7 @@ def _make_hybrid_kv_cache_manager() -> KVCacheManager:

61 dtypes=(torch.float32,),

62 mamba_cache_mode="align",

63 num_speculative_blocks=NUM_SPEC,

64 + num_prefill_checkpoint_blocks=num_prefill_checkpoint_blocks,

65 ),

66 ),

67 ],

@@ -78,6 +81,7 @@ def _split(

81 num_new_tokens: int,

82 use_eagle: bool = True,

83 partial_hit: bool = False,

84 + num_prefill_checkpoint_blocks: int = 0,

85 ) -> int:

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

87 stub = SimpleNamespace(

@@ -88,10 +92,42 @@ def _split(

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

93 mamba_partial_cache_hit=partial_hit,

94 hash_block_size=ATTN_BLOCK_SIZE,

95 + mamba_has_prefill_checkpoint_blocks=(

96 + num_prefill_checkpoint_blocks > 0 and not use_eagle

97 + ),

98 )

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

100

101

102 +@pytest.mark.parametrize(

103 + ("prompt_len", "num_new_tokens", "use_eagle", "expected"),

104 + [

105 + (2002, 2002, False, 2002),

106 + (3602, 2000, False, 2000),

107 + (3602, 3602, True, MAMBA_BLOCK_SIZE),

108 + ],

109 +)

110 +def test_internal_checkpoint_split(

111 + prompt_len: int, num_new_tokens: int, use_eagle: bool, expected: int

112 +) -> None:

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

114 + assert (

115 + _split(

116 + request,

117 + num_new_tokens,

118 + use_eagle=use_eagle,

119 + num_prefill_checkpoint_blocks=1,

120 + )

121 + == expected

122 + )

123 + if not use_eagle:

124 + manager = _make_hybrid_kv_cache_manager(num_prefill_checkpoint_blocks=1)

125 + assert manager.allocate_slots(request, expected) is not None

126 + mamba_manager = manager.coordinator.single_type_managers[MAMBA_GROUP_ID]

127 + blocks = mamba_manager.req_to_blocks[request.request_id]

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

129 +

130 +

131 def _run_chunked_prefill(

132 manager: KVCacheManager, request: Request, budgets: list[int]

133 ) -> dict[int, int]:

@@ -211,8 +247,15 @@ def test_poisoning_is_block_size_independent(

247

248 @pytest.mark.parametrize("partial_hit", [False, True])

249 @pytest.mark.parametrize("resume_at", [331, 1599, 1601, 2531, 3011])

250 +@pytest.mark.parametrize(

251 + ("num_prefill_checkpoint_blocks", "use_eagle"),

252 + [(0, True), (1, False)],

253 +)

254 def test_unaligned_resume_never_runs_past_its_block(

215 - partial_hit: bool, resume_at: int

255 + partial_hit: bool,

256 + resume_at: int,

257 + num_prefill_checkpoint_blocks: int,

258 + use_eagle: bool,

259 ) -> None:

260 """A prefill resuming mid-block must re-align before crossing a boundary.

261

@@ -227,7 +270,13 @@ def test_unaligned_resume_never_runs_past_its_block(

270 pos, ends = resume_at, []

271 while pos < prompt_len:

272 request.num_computed_tokens = pos

230 - num_new = _split(request, prompt_len - pos, partial_hit=partial_hit)

273 + num_new = _split(

274 + request,

275 + prompt_len - pos,

276 + use_eagle=use_eagle,

277 + partial_hit=partial_hit,

278 + num_prefill_checkpoint_blocks=num_prefill_checkpoint_blocks,

279 + )

280 assert num_new > 0, f"no progress at {pos}"

281 if pos % MAMBA_BLOCK_SIZE != 0:

282 block_end = (pos // MAMBA_BLOCK_SIZE + 1) * MAMBA_BLOCK_SIZE

tests/v1/core/test_prefix_caching.py

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

1089 use_eagle=False,

1090 hash_block_size=block_size,

1091 mamba_partial_cache_hit=False,

1092 + mamba_has_prefill_checkpoint_blocks=False,

1093 )

1094 req_2.shared_prefix_boundary = shared_prefix_boundary

1095 num_new_tokens_adjusted = Scheduler._mamba_block_aligned_split(

vllm/models/kimi_k3/nvidia/kda.py

+161 -19

@@ -2,6 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 from collections.abc import Callable

5 +from dataclasses import replace

7 import torch

8 from einops import rearrange

@@ -45,7 +46,10 @@ from vllm.models.kimi_k3.nvidia.kda_metadata import (

46 from vllm.platforms import current_platform

47 from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated

48 from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig

49 +from vllm.triton_utils import tl, triton

50 from vllm.v1.attention.backend import AttentionBackend

51 +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID

52 +from vllm.v1.kv_cache_interface import MambaSpec

53 from vllm.v1.worker.workspace import current_workspace_manager

55 logger = init_logger(__name__)

@@ -199,6 +203,8 @@ def _flashkda_prefill(

203 out: torch.Tensor,

204 final_state: torch.Tensor,

205 workspace: torch.Tensor,

206 + checkpoint_state: torch.Tensor | None = None,

207 + checkpoint_offsets: torch.Tensor | None = None,

208 ) -> tuple[torch.Tensor, torch.Tensor]:

209 import vllm._flashkda_C # noqa: F401

210

@@ -221,10 +227,75 @@ def _flashkda_prefill(

227 initial_state.contiguous(),

228 final_state,

229 cu_seqlens.contiguous(),

230 + checkpoint_state,

231 + checkpoint_offsets.contiguous() if checkpoint_offsets is not None else None,

232 )

233 return out, final_state

234

235

236 +@triton.jit

237 +def _store_cache_checkpoints_kernel(

238 + x_ptr,

239 + conv_state_ptr,

240 + recurrent_checkpoint_ptr,

241 + recurrent_state_ptr,

242 + query_start_loc_ptr,

243 + checkpoint_offsets_ptr,

244 + checkpoint_state_indices_ptr,

245 + x_stride_0: tl.constexpr,

246 + x_stride_1: tl.constexpr,

247 + state_stride_0: tl.constexpr,

248 + state_stride_1: tl.constexpr,

249 + state_stride_2: tl.constexpr,

250 + checkpoint_stride_0: tl.constexpr,

251 + recurrent_state_stride_0: tl.constexpr,

252 + checkpoint_offset_stride: tl.constexpr,

253 + STATE_LEN: tl.constexpr,

254 + WIDTH: tl.constexpr,

255 + RECURRENT_ROW_SIZE: tl.constexpr,

256 + NULL_STATE_IDX: tl.constexpr,

257 + BLOCK_SIZE: tl.constexpr,

258 +):

259 + # store checkpoints to cache

260 + seq_idx = tl.program_id(0)

261 + cols = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)

262 + state_idx = tl.load(checkpoint_state_indices_ptr + seq_idx)

263 + checkpoint_offset = tl.load(

264 + checkpoint_offsets_ptr + seq_idx * checkpoint_offset_stride

265 + )

266 + valid_checkpoint = (state_idx != NULL_STATE_IDX) & (checkpoint_offset > 0)

267 + valid_conv = (

268 + (cols < WIDTH * STATE_LEN) & valid_checkpoint & (checkpoint_offset >= STATE_LEN)

269 + )

270 + width_idx = cols // STATE_LEN

271 + history_idx = cols % STATE_LEN

272 + checkpoint_end = tl.load(query_start_loc_ptr + seq_idx) + checkpoint_offset

273 + token_idx = checkpoint_end - STATE_LEN + history_idx

274 + values = tl.load(

275 + x_ptr + token_idx * x_stride_0 + width_idx * x_stride_1,

276 + mask=valid_conv,

277 + )

278 + tl.store(

279 + conv_state_ptr

280 + + state_idx * state_stride_0

281 + + width_idx * state_stride_1

282 + + history_idx * state_stride_2,

283 + values,

284 + mask=valid_conv,

285 + )

286 +

287 + valid_recurrent = (cols < RECURRENT_ROW_SIZE) & valid_checkpoint

288 + recurrent = tl.load(

289 + recurrent_checkpoint_ptr + seq_idx * checkpoint_stride_0 + cols,

290 + mask=valid_recurrent,

291 + )

292 + tl.store(

293 + recurrent_state_ptr + state_idx * recurrent_state_stride_0 + cols,

294 + recurrent,

295 + mask=valid_recurrent,

296 + )

297 +

298 +

299 def resolve_kda_prefill_backend(

300 backend: str,

301 head_dim: int,

@@ -471,6 +542,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):

542 self._flashkda_buffer_specs = (

543 ((1, T, H, D), self.model_config.dtype),

544 ((N, H, D, D), self.get_state_dtype()[1]),

545 + ((N, H, D, D), self.get_state_dtype()[1]),

546 ((workspace_size,), torch.uint8),

547 )

548

@@ -518,6 +590,14 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):

590 raise ValueError(f"Duplicate layer name: {prefix}")

591 compilation_config.static_forward_context[prefix] = self

592

593 + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> MambaSpec:

594 + spec = super().get_kv_cache_spec(vllm_config)

595 + assert isinstance(spec, MambaSpec)

596 + return replace(

597 + spec,

598 + num_prefill_checkpoint_blocks=int(self.kda_prefill_backend == "flashkda"),

599 + )

600 +

601 def forward(

602 self,

603 hidden_states: torch.Tensor,

@@ -590,6 +670,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):

670 spec_query_start_loc = m.spec_query_start_loc

671 num_accepted_tokens = m.num_accepted_tokens

672 num_actual_tokens = m.num_actual_tokens

673 + checkpoint = m.checkpoint

674 has_spec_decode = m.num_spec_decodes > 0

675 mixed_qkv = mixed_qkv[:num_actual_tokens]

676 g1 = g1[:, :num_actual_tokens]

@@ -780,7 +861,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):

861 if self.kda_prefill_backend == "flashkda":

862 assert self.gate_lower_bound is not None

863 assert self._flashkda_buffer_specs is not None

783 - workspace_out, final_state, workspace = (

864 + workspace_out, final_state, checkpoint_state, workspace = (

865 current_workspace_manager().get_simultaneous(

866 *self._flashkda_buffer_specs

867 )

@@ -788,24 +869,85 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):

869 flashkda_out = (

870 workspace_out if has_spec_decode else core_attn_out

871 )[:, : q_ns.shape[1]]

791 - (

792 - core_attn_out_non_spec,

793 - last_recurrent_state,

794 - ) = _flashkda_prefill(

795 - q=q_ns,

796 - k=k_ns,

797 - v=v_ns,

798 - g=g1_ns,

799 - beta=beta_ns,

800 - A_log=self.A_log,

801 - dt_bias=self.dt_bias,

802 - lower_bound=self.gate_lower_bound,

803 - initial_state=initial_state,

804 - cu_seqlens=non_spec_query_start_loc,

805 - out=flashkda_out,

806 - final_state=final_state[: initial_state.shape[0]],

807 - workspace=workspace,

808 - )

872 + if checkpoint is not None:

873 + assert non_spec_query_start_loc is not None

874 + num_sequences = initial_state.shape[0]

875 + assert checkpoint.checkpoint_offsets.shape == (num_sequences,)

876 + final_state = final_state[:num_sequences]

877 + checkpoint_state = checkpoint_state[:num_sequences]

878 + checkpoint_offsets = checkpoint.checkpoint_offsets

879 + _flashkda_prefill(

880 + q=q_ns,

881 + k=k_ns,

882 + v=v_ns,

883 + g=g1_ns,

884 + beta=beta_ns,

885 + A_log=self.A_log,

886 + dt_bias=self.dt_bias,

887 + lower_bound=self.gate_lower_bound,

888 + initial_state=initial_state,

889 + cu_seqlens=non_spec_query_start_loc,

890 + out=flashkda_out,

891 + final_state=final_state,

892 + workspace=workspace,

893 + checkpoint_state=checkpoint_state,

894 + checkpoint_offsets=checkpoint_offsets,

895 + )

896 + core_attn_out_non_spec = flashkda_out

897 + last_recurrent_state = final_state

898 + state_len = conv_state.shape[-1]

899 + width = mixed_qkv_ns.shape[-1]

900 + recurrent_row_size = checkpoint_state[0].numel()

901 + block_size = 256

902 + _store_cache_checkpoints_kernel[

903 + (

904 + checkpoint_offsets.numel(),

905 + triton.cdiv(

906 + max(width * state_len, recurrent_row_size),

907 + block_size,

908 + ),

909 + )

910 + ](

911 + mixed_qkv_ns,

912 + conv_state,

913 + checkpoint_state,

914 + recurrent_state,

915 + non_spec_query_start_loc,

916 + checkpoint_offsets,

917 + checkpoint.state_indices,

918 + mixed_qkv_ns.stride(0),

919 + mixed_qkv_ns.stride(1),

920 + conv_state.stride(0),

921 + conv_state.stride(1),

922 + conv_state.stride(2),

923 + checkpoint_state.stride(0),

924 + recurrent_state.stride(0),

925 + checkpoint_offsets.stride(0),

926 + state_len,

927 + width,

928 + recurrent_row_size,

929 + NULL_BLOCK_ID,

930 + block_size,

931 + )

932 + else:

933 + (

934 + core_attn_out_non_spec,

935 + last_recurrent_state,

936 + ) = _flashkda_prefill(

937 + q=q_ns,

938 + k=k_ns,

939 + v=v_ns,

940 + g=g1_ns,

941 + beta=beta_ns,

942 + A_log=self.A_log,

943 + dt_bias=self.dt_bias,

944 + lower_bound=self.gate_lower_bound,

945 + initial_state=initial_state,

946 + cu_seqlens=non_spec_query_start_loc,

947 + out=flashkda_out,

948 + final_state=final_state[: initial_state.shape[0]],

949 + workspace=workspace,

950 + )

951 else:

952 (

953 core_attn_out_non_spec,

vllm/models/kimi_k3/nvidia/kda_metadata.py

+67

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

44 )

47 +FLASHKDA_CHUNK_SIZE = 16

48 +

49 +

50 @cache

51 def _metadata_launch_pdl() -> bool:

52 return current_platform.is_arch_support_pdl()

@@ -257,12 +260,19 @@ class KDARecoverSSMCommitMetadata:

260 align: KDARecoverSSMAlignMetadata | None

261

262

263 +@dataclass

264 +class KDACheckpointMetadata:

265 + checkpoint_offsets: torch.Tensor

266 + state_indices: torch.Tensor

267 +

268 +

269 @dataclass

270 class KimiK3KDAMetadata(GDNAttentionMetadata, RecoverSSMMetadata):

271 recoverssm_commit: KDARecoverSSMCommitMetadata | None = None

272 recoverssm_context: "KDARecoverSSMCommitContext | None" = field(

273 default=None, repr=False, compare=False

274 )

275 + checkpoint: KDACheckpointMetadata | None = None

276

277 def commit_recoverssm_state(

278 self, num_accepted_tokens: torch.Tensor

@@ -557,6 +567,62 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):

567 else:

568 has_initial_state = None

569

570 + checkpoint = None

571 + if (

572 + num_prefills > 0

573 + and self.kv_cache_spec.num_prefill_checkpoint_blocks > 0

574 + and self.vllm_config.cache_config.mamba_cache_mode == "align"

575 + ):

576 + # prepare checkpoint metadata

577 + assert m.seq_lens_cpu_upper_bound is not None

578 + request_rows = list(range(m.num_reqs))

579 + if spec_sequence_masks_cpu is not None:

580 + # get non spec request rows

581 + request_rows = active_non_spec_mask_cpu.nonzero().flatten().tolist()

582 + all_query_lens = query_start_loc_cpu.diff().tolist()

583 + query_lens = [all_query_lens[row] for row in request_rows]

584 + seq_lens = m.seq_lens_cpu_upper_bound.tolist()

585 + block_size = self.kv_cache_spec.block_size

586 + checkpoint_splits = []

587 + checkpoint_cols = []

588 + for row, query_len in zip(request_rows, query_lens):

589 + seq_len = seq_lens[row]

590 + offset = seq_len // block_size * block_size - (seq_len - query_len)

591 + # offset should be less than query_len

592 + valid = (

593 + seq_len % block_size != 0

594 + and 0 < offset < query_len

595 + and offset % FLASHKDA_CHUNK_SIZE == 0

596 + )

597 + offset = offset if valid else 0

598 + first_len = offset or query_len

599 + checkpoint_splits.append((first_len, query_len - first_len))

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

601 + if any(tail for _, tail in checkpoint_splits):

602 + checkpoint_offsets_tensor = async_tensor_h2d(

603 + [first if tail else 0 for first, tail in checkpoint_splits],

604 + dtype=torch.int32,

605 + device=query_start_loc.device,

606 + )

607 + request_rows_tensor = async_tensor_h2d(

608 + request_rows, dtype=torch.int64, device=query_start_loc.device

609 + )

610 + checkpoint_cols_tensor = async_tensor_h2d(

611 + checkpoint_cols, dtype=torch.int64, device=query_start_loc.device

612 + )

613 + checkpoint_state_indices = m.block_table_tensor[

614 + request_rows_tensor, checkpoint_cols_tensor

615 + ]

616 + checkpoint_state_indices = torch.where(

617 + checkpoint_cols_tensor >= 0,

618 + checkpoint_state_indices,

619 + NULL_BLOCK_ID,

620 + )

621 + checkpoint = KDACheckpointMetadata(

622 + checkpoint_offsets_tensor,

623 + checkpoint_state_indices,

624 + )

625 +

626 # Prepare per-request tensors for cudagraph replay. num_actual_tokens

627 # may be token-padded, while state/query/acceptance metadata is indexed

628 # by request.

@@ -647,6 +713,7 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):

713 nums_dict=nums_dict,

714 batch_ptr=batch_ptr,

715 token_chunk_offset_ptr=token_chunk_offset_ptr,

716 + checkpoint=checkpoint,

717 )

718

719

vllm/v1/core/sched/scheduler.py

+17 -2

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

53 )

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

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

56 -from vllm.v1.kv_cache_interface import KVCacheConfig

56 +from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec

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

58 from vllm.v1.metrics.stats import (

59 PrefixCacheStats,

@@ -324,6 +324,16 @@ class Scheduler(SchedulerInterface):

324 self.need_mamba_block_aligned_split = (

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

326 )

327 + self.mamba_has_prefill_checkpoint_blocks = (

328 + self.has_mamba_layers

329 + # TODO: support spec decoding

330 + and not self.use_eagle

331 + and all(

332 + not isinstance(group.kv_cache_spec, MambaSpec)

333 + or group.kv_cache_spec.num_prefill_checkpoint_blocks > 0

334 + for group in kv_cache_config.kv_cache_groups

335 + )

336 + )

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

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

339 # hash boundary, so the split adds that stop.

@@ -406,12 +416,17 @@ class Scheduler(SchedulerInterface):

416 last_cache_position = max(last_cache_position - block_size, 0)

417

418 end = start + num_new_tokens

419 + use_internal_checkpoint = (

420 + self.mamba_has_prefill_checkpoint_blocks and start % block_size == 0

421 + )

422 + if use_internal_checkpoint:

423 + last_cache_position = 0

424 # Invariant: slot p holds the state after exactly (p + 1) * block_size

425 # tokens. State is written at chunk ends, so chunk ends must be block

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

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

428 # and re-aligns at the next boundary.

414 - if end < prefill_end:

429 + if end < prefill_end and not use_internal_checkpoint:

430 max_prefill_tokens = self.max_num_scheduled_tokens

431 long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold

432 if long_prefill_threshold > 0:

vllm/v1/core/single_type_kv_cache_manager.py

+40 -21

@@ -1285,6 +1285,9 @@ class MambaManager(SingleTypeKVCacheManager):

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

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

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

1288 + # Number of internal checkpoint blocks required by each request's

1289 + # current allocation.

1290 + self._num_checkpoint_blocks: dict[str, int] = {}

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

1292 # tail (producers). On the next step's CoW the boundary state moves

1293 # into a private cow_block; we record that block for connector

@@ -1464,6 +1467,23 @@ class MambaManager(SingleTypeKVCacheManager):

1467 """

1468 return 0

1469

1470 + def _needs_internal_checkpoint(

1471 + self,

1472 + request_id: str,

1473 + num_tokens: int,

1474 + num_computed_tokens: int,

1475 + ) -> bool:

1476 + assert isinstance(self.kv_cache_spec, MambaSpec)

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

1478 + blocks = self.req_to_blocks[request_id]

1479 + return (

1480 + self.kv_cache_spec.num_prefill_checkpoint_blocks > 0

1481 + and num_tokens % self.block_size != 0

1482 + and num_computed_tokens % self.block_size == 0

1483 + and checkpoint_idx >= 0

1484 + and (checkpoint_idx >= len(blocks) or blocks[checkpoint_idx].is_null)

1485 + )

1486 +

1487 def get_num_blocks_to_allocate(

1488 self,

1489 request_id: str,

@@ -1526,18 +1546,17 @@ class MambaManager(SingleTypeKVCacheManager):

1546 )

1547 if has_partial_hit:

1548 num_new_blocks = max(num_new_blocks, 0) + 1

1549 + checkpoint_block = int(

1550 + self._needs_internal_checkpoint(

1551 + request_id, num_tokens, total_computed_tokens

1552 + )

1553 + )

1554 + if not apply_admission_cap:

1555 + self._num_checkpoint_blocks[request_id] = checkpoint_block

1556 if num_new_blocks > 0:

1530 - if request_id in self._allocated_block_reqs:

1531 - # Old request. Needs at most 1 more blocks as we can reuse the

1532 - # speculative blocks in previous step.

1533 - num_new_blocks = 1 + int(has_partial_hit)

1534 - else:

1535 - # First prefill. Allocate 1 block for running state, the

1536 - # speculative blocks, and one extra block if a partial cache

1537 - # hit must be copy-on-written before the new tokens run.

1538 - num_new_blocks = (

1539 - 1 + self.num_speculative_blocks + int(has_partial_hit)

1540 - )

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

1558 + if request_id not in self._allocated_block_reqs:

1559 + num_new_blocks += self.num_speculative_blocks

1560

1561 num_evictable_computed_blocks = self._get_num_evictable_blocks(

1562 new_computed_blocks

@@ -1569,6 +1588,7 @@ class MambaManager(SingleTypeKVCacheManager):

1588 num_required_blocks = (

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

1590 )

1591 + checkpoint_block = self._num_checkpoint_blocks.get(request_id, 0)

1592 partial_hit = self._partial_hit_reqs.get(request_id)

1593 has_partial_hit = partial_hit is not None

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

@@ -1596,11 +1616,11 @@ class MambaManager(SingleTypeKVCacheManager):

1616 )

1617 # null blocks

1618 if prev_block_len < num_skipped_blocks:

1619 + # minus the internal checkpoint block

1620 + # so we don't set null for that block

1621 + null_end = num_skipped_blocks - checkpoint_block

1622 req_blocks.extend(

1600 - [

1601 - self._null_block

1602 - for _ in range(prev_block_len, num_skipped_blocks)

1603 - ]

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

1624 )

1625

1626 if blocks_allocated:

@@ -1616,12 +1636,10 @@ class MambaManager(SingleTypeKVCacheManager):

1636 num_new_blocks = num_required_blocks - len(req_blocks)

1637 if has_partial_hit:

1638 num_new_blocks = max(num_new_blocks, 0) + 1

1619 - if blocks_allocated:

1620 - assert num_new_blocks <= 1 + int(has_partial_hit)

1621 - else:

1622 - assert num_new_blocks <= self.num_speculative_blocks + 1 + int(

1623 - has_partial_hit

1624 - )

1639 + max_new_blocks = 1 + int(has_partial_hit) + checkpoint_block

1640 + if not blocks_allocated:

1641 + max_new_blocks += self.num_speculative_blocks

1642 + assert num_new_blocks <= max_new_blocks

1643 new_blocks = self.block_pool.get_new_blocks(num_new_blocks)

1644 returned_blocks = req_blocks[prev_block_len:]

1645 if partial_hit is not None:

@@ -1670,6 +1688,7 @@ class MambaManager(SingleTypeKVCacheManager):

1688 if self.mamba_cache_mode == "align":

1689 self._allocated_block_reqs.discard(request_id)

1690 self.last_state_block_idx.pop(request_id, None)

1691 + self._num_checkpoint_blocks.pop(request_id, None)

1692 self._producer_partial_tail_reqs.pop(request_id, None)

1693 # A hand-off whose request died in this same scheduling pass must

1694 # not reach the connector: its unpin hook (free) has already run.

vllm/v1/kv_cache_interface.py

+5 -1

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

815 mamba_type: MambaAttentionBackendEnum = MambaAttentionBackendEnum.MAMBA2

816 mamba_cache_mode: str = "none"

817 num_speculative_blocks: int = 0

818 + num_prefill_checkpoint_blocks: int = 0

819 num_heads: int = 1

820 tokens_per_state: int = -1

821

@@ -843,7 +844,9 @@ class MambaSpec(KVCacheSpec):

844 cdiv(max_model_len, self.block_size) + self.num_speculative_blocks

845 ) * self.page_size_bytes

846 elif vllm_config.cache_config.mamba_cache_mode == "align":

846 - return self.page_size_bytes * (2 + self.num_speculative_blocks)

847 + return self.page_size_bytes * (

848 + 2 + self.num_speculative_blocks + self.num_prefill_checkpoint_blocks

849 + )

850 else:

851 return self.page_size_bytes * (1 + self.num_speculative_blocks)

852

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

868 return all(

869 isinstance(spec, MambaSpec)

870 and spec.num_speculative_blocks == self.num_speculative_blocks

871 + and spec.num_prefill_checkpoint_blocks == self.num_prefill_checkpoint_blocks

872 for spec in kv_cache_specs.values()

873 )

874