[Core][KDA] Generalize Mamba prefill checkpoint builder and exporter (#57783)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: zjy0516 <riverclouds.zhu@qq.com>
Chauncey committed
Sep 21, 2026 at 12:46 UTC
3df4ae153eb385e27b52f26c81f8edb9e20b9984
5 files changed
+314
-220
tests/models/kimi_k3/test_kda.py
+11
-31
@@ -14,6 +14,10 @@ import torch.nn.functional as F
14
15
from vllm import _custom_ops as ops
16
from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
17
+from vllm.model_executor.layers.mamba.checkpoint import MambaPrefillCheckpointMetadata
18
+from vllm.model_executor.layers.mamba.kda_checkpoint import (
19
+ FlashKDAPrefillCheckpointExporter,
20
+)
21
from vllm.model_executor.layers.mamba.ops.causal_conv1d import causal_conv1d_update
22
from vllm.model_executor.layers.mamba.ops.gather_initial_states import (
23
gather_initial_states,
@@ -25,7 +29,6 @@ from vllm.models.kimi_k3.nvidia import kda as nvidia_kda
29
from vllm.models.kimi_k3.nvidia.kda import (
30
_flashinfer_kda_prefill,
31
_flashkda_prefill,
28
- _store_cache_checkpoints_kernel,
32
is_flashinfer_fused_kda_decode_supported,
33
is_flashinfer_recurrent_kda_prefill_supported,
34
is_flashkda_supported,
@@ -1498,36 +1501,13 @@ def test_flashkda_checkpoint_correctness(state_dtype: torch.dtype, tolerance: fl
1501
checkpoint_state_indices = torch.tensor(
1502
[1, NULL_BLOCK_ID], dtype=torch.int32, device=DEVICE
1503
)
1501
- state_len = conv_state.shape[-1]
1502
- width = H * D
1503
- recurrent_row_size = checkpoint_state[0].numel()
1504
- block_size = 256
1505
- _store_cache_checkpoints_kernel[
1506
- (
1507
- checkpoint_state_indices.numel(),
1508
- (max(width * state_len, recurrent_row_size) + block_size - 1) // block_size,
1509
- )
1510
- ](
1511
- conv_input,
1512
- conv_state,
1513
- checkpoint_state,
1514
- recurrent_state,
1515
- cu_seqlens,
1516
- checkpoint_offsets,
1517
- checkpoint_state_indices,
1518
- conv_input.stride(0),
1519
- conv_input.stride(1),
1520
- conv_state.stride(0),
1521
- conv_state.stride(1),
1522
- conv_state.stride(2),
1523
- checkpoint_state.stride(0),
1524
- recurrent_state.stride(0),
1525
- checkpoint_offsets.stride(0),
1526
- state_len,
1527
- width,
1528
- recurrent_row_size,
1529
- NULL_BLOCK_ID,
1530
- block_size,
1504
+ FlashKDAPrefillCheckpointExporter().export(
1505
+ MambaPrefillCheckpointMetadata(checkpoint_offsets, checkpoint_state_indices),
1506
+ raw_qkv=conv_input,
1507
+ conv_state=conv_state,
1508
+ recurrent_checkpoint=checkpoint_state,
1509
+ recurrent_state=recurrent_state,
1510
+ cu_seqlens=cu_seqlens,
1511
)
1512
torch.testing.assert_close(conv_state[1], q[0, 13:16].flatten(1).transpose(0, 1))
1513
torch.testing.assert_close(recurrent_state[1], checkpoint_state[0])
vllm/model_executor/layers/mamba/checkpoint.py
new
+137
@@ -0,0 +1,137 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from abc import ABC, abstractmethod
5
+from dataclasses import dataclass
6
+
7
+import torch
8
+
9
+from vllm.config import VllmConfig
10
+from vllm.utils.math_utils import cdiv
11
+from vllm.utils.torch_utils import async_tensor_h2d
12
+from vllm.v1.attention.backend import CommonAttentionMetadata
13
+from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
14
+from vllm.v1.kv_cache_interface import (
15
+ MambaSpec,
16
+ get_mamba_prefill_checkpoint_position,
17
+ is_mamba_prefill_checkpoint_valid,
18
+)
19
+
20
+
21
+def compute_mamba_prefill_checkpoints(
22
+ seq_lens: list[int],
23
+ query_lens: list[int],
24
+ hash_block_size: int,
25
+ mamba_block_size: int,
26
+ checkpoint_alignment: int | None,
27
+ drop_eagle_block: bool,
28
+) -> tuple[list[int], list[int]]:
29
+ """Per-row internal prefill checkpoint offsets and cache block columns.
30
+
31
+ Backends call this instead of re-deriving the rules, so they decline in
32
+ lockstep with the scheduler and ``MambaManager``: allocating a checkpoint
33
+ block without writing it leaves the prefix cache serving uninitialized
34
+ state.
35
+
36
+ Returns:
37
+ ``(offsets, cols)``: the checkpoint's offset into each row's query and
38
+ its block-table column. ``0`` and ``-1`` mean the row has none.
39
+
40
+ """
41
+ offsets: list[int] = []
42
+ cols: list[int] = []
43
+ for seq_len, query_len in zip(seq_lens, query_lens):
44
+ query_start = seq_len - query_len
45
+ position = get_mamba_prefill_checkpoint_position(
46
+ seq_len, hash_block_size, drop_eagle_block=drop_eagle_block
47
+ )
48
+ valid = is_mamba_prefill_checkpoint_valid(
49
+ query_start=query_start,
50
+ query_end=seq_len,
51
+ checkpoint_position=position,
52
+ hash_block_size=hash_block_size,
53
+ mamba_block_size=mamba_block_size,
54
+ checkpoint_alignment=checkpoint_alignment,
55
+ )
56
+ offsets.append(position - query_start if valid else 0)
57
+ cols.append(cdiv(seq_len, mamba_block_size) - 2 if valid else -1)
58
+ return offsets, cols
59
+
60
+
61
+@dataclass
62
+class MambaPrefillCheckpointMetadata:
63
+ checkpoint_offsets: torch.Tensor
64
+ state_indices: torch.Tensor
65
+
66
+
67
+class MambaPrefillCheckpointBuilder:
68
+ """Build backend-neutral prefill checkpoint locations."""
69
+
70
+ def __init__(self, vllm_config: VllmConfig, kv_cache_spec: MambaSpec) -> None:
71
+ self.vllm_config = vllm_config
72
+ self.kv_cache_spec = kv_cache_spec
73
+
74
+ def build(
75
+ self,
76
+ m: CommonAttentionMetadata,
77
+ request_rows: list[int],
78
+ ) -> MambaPrefillCheckpointMetadata | None:
79
+ if self.kv_cache_spec.num_prefill_checkpoint_blocks == 0:
80
+ return None
81
+ assert m.seq_lens_cpu_upper_bound is not None
82
+ all_query_lens = m.query_start_loc_cpu.diff().tolist()
83
+ query_lens = [all_query_lens[row] for row in request_rows]
84
+ seq_lens = m.seq_lens_cpu_upper_bound.tolist()
85
+ block_size = self.kv_cache_spec.block_size
86
+ hash_block_size = self.vllm_config.cache_config.prefix_match_unit or block_size
87
+ speculative_config = self.vllm_config.speculative_config
88
+ drop_eagle_block = (
89
+ speculative_config is not None and speculative_config.use_eagle_block_drop()
90
+ )
91
+ checkpoint_offsets, checkpoint_cols = compute_mamba_prefill_checkpoints(
92
+ [seq_lens[row] for row in request_rows],
93
+ query_lens,
94
+ hash_block_size=hash_block_size,
95
+ mamba_block_size=block_size,
96
+ checkpoint_alignment=self.kv_cache_spec.prefill_checkpoint_alignment,
97
+ drop_eagle_block=drop_eagle_block,
98
+ )
99
+ if not any(checkpoint_offsets):
100
+ return None
101
+ checkpoint_offsets_tensor = async_tensor_h2d(
102
+ checkpoint_offsets,
103
+ dtype=torch.int32,
104
+ device=m.query_start_loc.device,
105
+ )
106
+ request_rows_tensor = async_tensor_h2d(
107
+ request_rows, dtype=torch.int64, device=m.query_start_loc.device
108
+ )
109
+ checkpoint_cols_tensor = async_tensor_h2d(
110
+ checkpoint_cols, dtype=torch.int64, device=m.query_start_loc.device
111
+ )
112
+ checkpoint_state_indices = m.block_table_tensor[
113
+ request_rows_tensor, checkpoint_cols_tensor
114
+ ]
115
+ checkpoint_state_indices = torch.where(
116
+ checkpoint_cols_tensor >= 0,
117
+ checkpoint_state_indices,
118
+ NULL_BLOCK_ID,
119
+ )
120
+ return MambaPrefillCheckpointMetadata(
121
+ checkpoint_offsets_tensor,
122
+ checkpoint_state_indices,
123
+ )
124
+
125
+
126
+class MambaPrefillCheckpointExporter(ABC):
127
+ """Export a mid-prefill checkpoint into backend-specific paged states."""
128
+
129
+ @abstractmethod
130
+ def export(
131
+ self,
132
+ checkpoint: MambaPrefillCheckpointMetadata,
133
+ *args,
134
+ **kwargs,
135
+ ) -> None:
136
+ """Write checkpoint state into the paged cache."""
137
+ raise NotImplementedError
vllm/model_executor/layers/mamba/kda_checkpoint.py
new
+130
@@ -0,0 +1,130 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from dataclasses import dataclass
5
+
6
+import torch
7
+
8
+from vllm.model_executor.layers.mamba.checkpoint import (
9
+ MambaPrefillCheckpointExporter,
10
+ MambaPrefillCheckpointMetadata,
11
+)
12
+from vllm.triton_utils import tl, triton
13
+from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
14
+
15
+
16
+def kda_prefill_checkpoint_alignment(backend: str) -> int | None:
17
+ return 16 if backend == "flashkda" else None
18
+
19
+
20
+@dataclass(frozen=True)
21
+class FlashKDAPrefillCheckpointExporter(MambaPrefillCheckpointExporter):
22
+ """Store FlashKDA recurrent and convolution checkpoint states."""
23
+
24
+ state_len: int | None = None
25
+
26
+ def export(
27
+ self,
28
+ checkpoint: MambaPrefillCheckpointMetadata,
29
+ *,
30
+ raw_qkv: torch.Tensor,
31
+ conv_state: torch.Tensor,
32
+ recurrent_checkpoint: torch.Tensor,
33
+ recurrent_state: torch.Tensor,
34
+ cu_seqlens: torch.Tensor,
35
+ ) -> None:
36
+ state_len = (
37
+ self.state_len if self.state_len is not None else conv_state.shape[-1]
38
+ )
39
+ width = raw_qkv.shape[-1]
40
+ recurrent_row_size = recurrent_checkpoint[0].numel()
41
+ block_size = 256
42
+ store_cache_checkpoints_kernel[
43
+ (
44
+ checkpoint.checkpoint_offsets.numel(),
45
+ triton.cdiv(max(width * state_len, recurrent_row_size), block_size),
46
+ )
47
+ ](
48
+ raw_qkv,
49
+ conv_state,
50
+ recurrent_checkpoint,
51
+ recurrent_state,
52
+ cu_seqlens,
53
+ checkpoint.checkpoint_offsets,
54
+ checkpoint.state_indices,
55
+ raw_qkv.stride(0),
56
+ raw_qkv.stride(1),
57
+ conv_state.stride(0),
58
+ conv_state.stride(1),
59
+ conv_state.stride(2),
60
+ recurrent_checkpoint.stride(0),
61
+ recurrent_state.stride(0),
62
+ checkpoint.checkpoint_offsets.stride(0),
63
+ state_len,
64
+ width,
65
+ recurrent_row_size,
66
+ NULL_BLOCK_ID,
67
+ block_size,
68
+ )
69
+
70
+
71
+@triton.jit
72
+def store_cache_checkpoints_kernel(
73
+ x_ptr,
74
+ conv_state_ptr,
75
+ recurrent_checkpoint_ptr,
76
+ recurrent_state_ptr,
77
+ query_start_loc_ptr,
78
+ checkpoint_offsets_ptr,
79
+ checkpoint_state_indices_ptr,
80
+ x_stride_0: tl.constexpr,
81
+ x_stride_1: tl.constexpr,
82
+ state_stride_0: tl.constexpr,
83
+ state_stride_1: tl.constexpr,
84
+ state_stride_2: tl.constexpr,
85
+ checkpoint_stride_0: tl.constexpr,
86
+ recurrent_state_stride_0: tl.constexpr,
87
+ checkpoint_offset_stride: tl.constexpr,
88
+ STATE_LEN: tl.constexpr,
89
+ WIDTH: tl.constexpr,
90
+ RECURRENT_ROW_SIZE: tl.constexpr,
91
+ NULL_STATE_IDX: tl.constexpr,
92
+ BLOCK_SIZE: tl.constexpr,
93
+):
94
+ seq_idx = tl.program_id(0)
95
+ cols = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
96
+ state_idx = tl.load(checkpoint_state_indices_ptr + seq_idx).to(tl.int64)
97
+ checkpoint_offset = tl.load(
98
+ checkpoint_offsets_ptr + seq_idx * checkpoint_offset_stride
99
+ )
100
+ valid_checkpoint = (state_idx != NULL_STATE_IDX) & (checkpoint_offset > 0)
101
+ valid_conv = (
102
+ (cols < WIDTH * STATE_LEN) & valid_checkpoint & (checkpoint_offset >= STATE_LEN)
103
+ )
104
+ width_idx = cols // STATE_LEN
105
+ history_idx = cols % STATE_LEN
106
+ checkpoint_end = tl.load(query_start_loc_ptr + seq_idx) + checkpoint_offset
107
+ token_idx = checkpoint_end - STATE_LEN + history_idx
108
+ values = tl.load(
109
+ x_ptr + token_idx * x_stride_0 + width_idx * x_stride_1,
110
+ mask=valid_conv,
111
+ )
112
+ tl.store(
113
+ conv_state_ptr
114
+ + state_idx * state_stride_0
115
+ + width_idx * state_stride_1
116
+ + history_idx * state_stride_2,
117
+ values,
118
+ mask=valid_conv,
119
+ )
120
+
121
+ valid_recurrent = (cols < RECURRENT_ROW_SIZE) & valid_checkpoint
122
+ recurrent = tl.load(
123
+ recurrent_checkpoint_ptr + seq_idx * checkpoint_stride_0 + cols,
124
+ mask=valid_recurrent,
125
+ )
126
+ tl.store(
127
+ recurrent_state_ptr + state_idx * recurrent_state_stride_0 + cols,
128
+ recurrent,
129
+ mask=valid_recurrent,
130
+ )
vllm/models/kimi_k3/nvidia/kda.py
+17
-102
@@ -21,6 +21,10 @@ from vllm.model_executor.layers.linear import (
21
RowParallelLinear,
22
)
23
from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention
24
+from vllm.model_executor.layers.mamba.kda_checkpoint import (
25
+ FlashKDAPrefillCheckpointExporter,
26
+ kda_prefill_checkpoint_alignment,
27
+)
28
from vllm.model_executor.layers.mamba.mamba_utils import (
29
MambaStateDtypeCalculator,
30
MambaStateShapeCalculator,
@@ -49,7 +53,6 @@ from vllm.models.kimi_k3.nvidia.kda_metadata import (
53
from vllm.platforms import current_platform
54
from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated
55
from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
52
-from vllm.triton_utils import tl, triton
56
from vllm.utils.flashinfer import (
57
flashinfer_fused_kda_decode,
58
flashinfer_recurrent_kda,
@@ -57,7 +60,6 @@ from vllm.utils.flashinfer import (
60
has_flashinfer_recurrent_kda,
61
)
62
from vllm.v1.attention.backend import AttentionBackend
60
-from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
63
from vllm.v1.kv_cache_interface import MambaSpec
64
from vllm.v1.worker.workspace import current_workspace_manager
65
@@ -389,69 +391,6 @@ def _flashinfer_kda_prefill(
391
return output, initial_state
392
393
392
-@triton.jit
393
-def _store_cache_checkpoints_kernel(
394
- x_ptr,
395
- conv_state_ptr,
396
- recurrent_checkpoint_ptr,
397
- recurrent_state_ptr,
398
- query_start_loc_ptr,
399
- checkpoint_offsets_ptr,
400
- checkpoint_state_indices_ptr,
401
- x_stride_0: tl.constexpr,
402
- x_stride_1: tl.constexpr,
403
- state_stride_0: tl.constexpr,
404
- state_stride_1: tl.constexpr,
405
- state_stride_2: tl.constexpr,
406
- checkpoint_stride_0: tl.constexpr,
407
- recurrent_state_stride_0: tl.constexpr,
408
- checkpoint_offset_stride: tl.constexpr,
409
- STATE_LEN: tl.constexpr,
410
- WIDTH: tl.constexpr,
411
- RECURRENT_ROW_SIZE: tl.constexpr,
412
- NULL_STATE_IDX: tl.constexpr,
413
- BLOCK_SIZE: tl.constexpr,
414
-):
415
- # store checkpoints to cache
416
- seq_idx = tl.program_id(0)
417
- cols = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
418
- state_idx = tl.load(checkpoint_state_indices_ptr + seq_idx).to(tl.int64)
419
- checkpoint_offset = tl.load(
420
- checkpoint_offsets_ptr + seq_idx * checkpoint_offset_stride
421
- )
422
- valid_checkpoint = (state_idx != NULL_STATE_IDX) & (checkpoint_offset > 0)
423
- valid_conv = (
424
- (cols < WIDTH * STATE_LEN) & valid_checkpoint & (checkpoint_offset >= STATE_LEN)
425
- )
426
- width_idx = cols // STATE_LEN
427
- history_idx = cols % STATE_LEN
428
- checkpoint_end = tl.load(query_start_loc_ptr + seq_idx) + checkpoint_offset
429
- token_idx = checkpoint_end - STATE_LEN + history_idx
430
- values = tl.load(
431
- x_ptr + token_idx * x_stride_0 + width_idx * x_stride_1,
432
- mask=valid_conv,
433
- )
434
- tl.store(
435
- conv_state_ptr
436
- + state_idx * state_stride_0
437
- + width_idx * state_stride_1
438
- + history_idx * state_stride_2,
439
- values,
440
- mask=valid_conv,
441
- )
442
-
443
- valid_recurrent = (cols < RECURRENT_ROW_SIZE) & valid_checkpoint
444
- recurrent = tl.load(
445
- recurrent_checkpoint_ptr + seq_idx * checkpoint_stride_0 + cols,
446
- mask=valid_recurrent,
447
- )
448
- tl.store(
449
- recurrent_state_ptr + state_idx * recurrent_state_stride_0 + cols,
450
- recurrent,
451
- mask=valid_recurrent,
452
- )
453
-
454
-
394
def resolve_kda_prefill_backend(
395
backend: str,
396
head_dim: int,
@@ -735,6 +674,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
674
self._flashinfer_kda_output_spec: tuple[tuple[int, ...], torch.dtype] | None = (
675
None
676
)
677
+ self._checkpoint_exporter: FlashKDAPrefillCheckpointExporter | None = None
678
if self.kda_prefill_backend == "flashkda":
679
T = vllm_config.scheduler_config.max_num_batched_tokens
680
N = vllm_config.scheduler_config.max_num_seqs
@@ -748,6 +688,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
688
((N, H, D, D), self.get_state_dtype()[1]),
689
((workspace_size,), torch.uint8),
690
)
691
+ self._checkpoint_exporter = FlashKDAPrefillCheckpointExporter()
692
elif self.kda_prefill_backend == "flashinfer":
693
T = vllm_config.scheduler_config.max_num_batched_tokens
694
H, D = self.local_num_heads, self.head_dim
@@ -800,12 +741,11 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
741
def get_kv_cache_spec(self, vllm_config: VllmConfig) -> MambaSpec:
742
spec = super().get_kv_cache_spec(vllm_config)
743
assert isinstance(spec, MambaSpec)
744
+ alignment = kda_prefill_checkpoint_alignment(self.kda_prefill_backend)
745
return replace(
746
spec,
805
- num_prefill_checkpoint_blocks=int(self.kda_prefill_backend == "flashkda"),
806
- prefill_checkpoint_alignment=(
807
- 16 if self.kda_prefill_backend == "flashkda" else None
808
- ),
747
+ num_prefill_checkpoint_blocks=int(alignment is not None),
748
+ prefill_checkpoint_alignment=alignment,
749
)
750
751
def forward(
@@ -1173,39 +1113,14 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
1113
)
1114
core_attn_out_non_spec = flashkda_out
1115
last_recurrent_state = final_state
1176
- state_len = conv_state.shape[-1]
1177
- width = mixed_qkv_ns.shape[-1]
1178
- recurrent_row_size = checkpoint_state[0].numel()
1179
- block_size = 256
1180
- _store_cache_checkpoints_kernel[
1181
- (
1182
- checkpoint_offsets.numel(),
1183
- triton.cdiv(
1184
- max(width * state_len, recurrent_row_size),
1185
- block_size,
1186
- ),
1187
- )
1188
- ](
1189
- mixed_qkv_ns,
1190
- conv_state,
1191
- checkpoint_state,
1192
- recurrent_state,
1193
- non_spec_query_start_loc,
1194
- checkpoint_offsets,
1195
- checkpoint.state_indices,
1196
- mixed_qkv_ns.stride(0),
1197
- mixed_qkv_ns.stride(1),
1198
- conv_state.stride(0),
1199
- conv_state.stride(1),
1200
- conv_state.stride(2),
1201
- checkpoint_state.stride(0),
1202
- recurrent_state.stride(0),
1203
- checkpoint_offsets.stride(0),
1204
- state_len,
1205
- width,
1206
- recurrent_row_size,
1207
- NULL_BLOCK_ID,
1208
- block_size,
1116
+ assert self._checkpoint_exporter is not None
1117
+ self._checkpoint_exporter.export(
1118
+ checkpoint,
1119
+ raw_qkv=mixed_qkv_ns,
1120
+ conv_state=conv_state,
1121
+ recurrent_checkpoint=checkpoint_state,
1122
+ recurrent_state=recurrent_state,
1123
+ cu_seqlens=non_spec_query_start_loc,
1124
)
1125
else:
1126
(
vllm/models/kimi_k3/nvidia/kda_metadata.py
+19
-87
@@ -18,9 +18,12 @@ from typing import TYPE_CHECKING
18
import torch
19
20
from vllm.config import VllmConfig
21
+from vllm.model_executor.layers.mamba.checkpoint import (
22
+ MambaPrefillCheckpointBuilder,
23
+ MambaPrefillCheckpointMetadata,
24
+)
25
from vllm.platforms import current_platform
26
from vllm.triton_utils import tl, triton
23
-from vllm.utils.math_utils import cdiv
27
from vllm.utils.torch_utils import async_tensor_h2d
28
from vllm.v1.attention.backend import CommonAttentionMetadata
29
from vllm.v1.attention.backends.gdn_attn import (
@@ -37,11 +40,7 @@ from vllm.v1.attention.backends.utils import (
40
compute_causal_conv1d_metadata,
41
split_decodes_and_prefills,
42
)
40
-from vllm.v1.kv_cache_interface import (
41
- MambaSpec,
42
- get_mamba_prefill_checkpoint_position,
43
- is_mamba_prefill_checkpoint_valid,
44
-)
43
+from vllm.v1.kv_cache_interface import MambaSpec
44
45
if TYPE_CHECKING:
46
from vllm.models.kimi_k3.nvidia.ops.recoverssm import (
@@ -262,12 +261,6 @@ class KDARecoverSSMCommitMetadata:
261
align: KDARecoverSSMAlignMetadata | None
262
263
265
-@dataclass
266
-class KDACheckpointMetadata:
267
- checkpoint_offsets: torch.Tensor
268
- state_indices: torch.Tensor
269
-
270
-
264
@dataclass
265
class KimiK3KDAMetadata(GDNAttentionMetadata, RecoverSSMMetadata):
266
spec_token_start: int | None = None
@@ -278,7 +271,7 @@ class KimiK3KDAMetadata(GDNAttentionMetadata, RecoverSSMMetadata):
271
recoverssm_context: "KDARecoverSSMCommitContext | None" = field(
272
default=None, repr=False, compare=False
273
)
281
- checkpoint: KDACheckpointMetadata | None = None
274
+ checkpoint: MambaPrefillCheckpointMetadata | None = None
275
276
def commit_recoverssm_state(
277
self, num_accepted_tokens: torch.Tensor
@@ -322,6 +315,9 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
315
device: torch.device,
316
) -> None:
317
super().__init__(kv_cache_spec, layer_names, vllm_config, device)
318
+ self.checkpoint_builder = MambaPrefillCheckpointBuilder(
319
+ vllm_config, kv_cache_spec
320
+ )
321
additional_config = vllm_config.additional_config
322
self.use_flashinfer_prefill = (
323
isinstance(additional_config, dict)
@@ -400,8 +396,10 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
396
if not self.use_spec_decode or num_decode_draft_tokens_cpu is None:
397
spec_sequence_masks_cpu = None
398
num_spec_decodes = 0
399
+ active_non_spec_mask_cpu = None
400
else:
401
spec_sequence_masks_cpu = num_decode_draft_tokens_cpu >= 0
402
+ active_non_spec_mask_cpu = None
403
if self.use_recoverssm:
404
assert m.is_prefilling is not None
405
assert m.is_prefilling.device.type == "cpu"
@@ -418,6 +416,7 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
416
):
417
spec_sequence_masks_cpu = None
418
num_spec_decodes = 0
419
+ active_non_spec_mask_cpu = None
420
else:
421
num_spec_decodes = spec_sequence_masks_cpu.sum().item()
422
if num_spec_decodes == 0:
@@ -624,80 +623,6 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
623
else:
624
has_initial_state = None
625
627
- checkpoint = None
628
- checkpoint_enabled = (
629
- self.vllm_config.cache_config.mamba_cache_mode == "align"
630
- and self.kv_cache_spec.num_prefill_checkpoint_blocks > 0
631
- )
632
- if num_prefills > 0 and checkpoint_enabled:
633
- # prepare checkpoint metadata
634
- assert m.seq_lens_cpu_upper_bound is not None
635
- request_rows = list(range(m.num_reqs))
636
- if spec_sequence_masks_cpu is not None:
637
- # get non spec request rows
638
- request_rows = active_non_spec_mask_cpu.nonzero().flatten().tolist()
639
- all_query_lens = query_start_loc_cpu.diff().tolist()
640
- query_lens = [all_query_lens[row] for row in request_rows]
641
- seq_lens = m.seq_lens_cpu_upper_bound.tolist()
642
- block_size = self.kv_cache_spec.block_size
643
- hash_block_size = (
644
- self.vllm_config.cache_config.prefix_match_unit or block_size
645
- )
646
- speculative_config = self.vllm_config.speculative_config
647
- drop_eagle_block = (
648
- speculative_config is not None
649
- and speculative_config.use_eagle_block_drop()
650
- )
651
- checkpoint_splits = []
652
- checkpoint_cols = []
653
- for row, query_len in zip(request_rows, query_lens):
654
- seq_len = seq_lens[row]
655
- query_start = seq_len - query_len
656
- checkpoint_position = get_mamba_prefill_checkpoint_position(
657
- seq_len,
658
- hash_block_size,
659
- drop_eagle_block=drop_eagle_block,
660
- )
661
- offset = checkpoint_position - query_start
662
- valid = is_mamba_prefill_checkpoint_valid(
663
- query_start=query_start,
664
- query_end=seq_len,
665
- checkpoint_position=checkpoint_position,
666
- hash_block_size=hash_block_size,
667
- mamba_block_size=block_size,
668
- checkpoint_alignment=(
669
- self.kv_cache_spec.prefill_checkpoint_alignment
670
- ),
671
- )
672
- offset = offset if valid else 0
673
- first_len = offset or query_len
674
- checkpoint_splits.append((first_len, query_len - first_len))
675
- checkpoint_cols.append(cdiv(seq_len, block_size) - 2 if valid else -1)
676
- if any(tail for _, tail in checkpoint_splits):
677
- checkpoint_offsets_tensor = async_tensor_h2d(
678
- [first if tail else 0 for first, tail in checkpoint_splits],
679
- dtype=torch.int32,
680
- device=query_start_loc.device,
681
- )
682
- request_rows_tensor = async_tensor_h2d(
683
- request_rows, dtype=torch.int64, device=query_start_loc.device
684
- )
685
- checkpoint_cols_tensor = async_tensor_h2d(
686
- checkpoint_cols, dtype=torch.int64, device=query_start_loc.device
687
- )
688
- checkpoint_state_indices = m.block_table_tensor[
689
- request_rows_tensor, checkpoint_cols_tensor
690
- ]
691
- checkpoint_state_indices = torch.where(
692
- checkpoint_cols_tensor >= 0,
693
- checkpoint_state_indices,
694
- NULL_BLOCK_ID,
695
- )
696
- checkpoint = KDACheckpointMetadata(
697
- checkpoint_offsets_tensor,
698
- checkpoint_state_indices,
699
- )
700
-
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.
@@ -776,6 +701,13 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
701
flashinfer_prefill_query_start_loc.diff(), descending=True
702
).to(torch.int32)
703
704
+ checkpoint = None
705
+ if num_prefills > 0:
706
+ request_rows = list(range(m.num_reqs))
707
+ if active_non_spec_mask_cpu is not None:
708
+ request_rows = active_non_spec_mask_cpu.nonzero().flatten().tolist()
709
+ checkpoint = self.checkpoint_builder.build(m, request_rows)
710
+
711
return KimiK3KDAMetadata(
712
num_prefills=num_prefills,
713
num_prefill_tokens=num_prefill_tokens,