[ROCm][Perf] Kimi-K3 enable prefill checkpoints on ROCm (#58344)
Signed-off-by: kliuae <kuanfu.liu@embeddedllm.com>
kliuae committed
Oct 2, 2026 at 14:58 UTC
01549796563f2e8bd0d6b08243cfb954a7b42ec8
7 files changed
+748
-14
tests/models/kimi_k3/test_amd_kda_checkpoint.py
new
+521
@@ -0,0 +1,521 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""The ROCm KDA prefill checkpoint: where it is taken and what it stores.
4
+
5
+Two things have to line up for a checkpoint to be safe: the metadata must name
6
+a position the chunk kernel will actually write, and the conv window saved
7
+there must be the one a resumed prefill would have left.
8
+"""
9
+
10
+import pytest
11
+import torch
12
+
13
+from tests.v1.attention.utils import (
14
+ BatchSpec,
15
+ create_common_attn_metadata,
16
+ create_vllm_config,
17
+)
18
+from vllm.model_executor.layers.mamba.checkpoint import (
19
+ MambaPrefillCheckpointExporter,
20
+ MambaPrefillCheckpointMetadata,
21
+)
22
+from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention
23
+from vllm.model_executor.layers.mamba.ops.causal_conv1d import causal_conv1d_fn
24
+from vllm.models.kimi_k3.amd.kda import KimiK3DeltaAttention
25
+from vllm.models.kimi_k3.amd.kda_metadata import KimiK3ROCmKDAMetadataBuilder
26
+from vllm.models.kimi_k3.amd.ops.kda_checkpoint import (
27
+ KimiK3ROCmKDAPrefillCheckpointExporter,
28
+)
29
+from vllm.models.kimi_k3.amd.ops.kda_chunk import (
30
+ KDA_CHECKPOINT_ALIGNMENT,
31
+ is_fused_kda_chunk_supported,
32
+)
33
+from vllm.platforms import current_platform
34
+from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
35
+from vllm.v1.kv_cache_interface import MambaSpec
36
+
37
+pytestmark = pytest.mark.skipif(
38
+ not (
39
+ current_platform.is_rocm()
40
+ and torch.cuda.is_available()
41
+ and is_fused_kda_chunk_supported()
42
+ ),
43
+ reason="The ROCm KDA prefill checkpoint is exported by the fused chunk "
44
+ "kernel, which needs a gfx950 build that includes it",
45
+)
46
+
47
+# A mamba block size the chunk walk lands on. The real K3 value is 768; any
48
+# multiple of the kernel's 64-token chunk behaves the same and keeps the
49
+# fixtures small.
50
+BLOCK_SIZE = 64
51
+
52
+
53
+class _StubSpeculativeConfig:
54
+ method = "dspark"
55
+ parallel_drafting = False
56
+
57
+ def __init__(self, num_speculative_tokens: int) -> None:
58
+ self.num_speculative_tokens = num_speculative_tokens
59
+
60
+ def use_eagle_block_drop(self) -> bool:
61
+ return False
62
+
63
+
64
+def _builder(
65
+ num_prefill_checkpoint_blocks: int = 1,
66
+ mamba_cache_mode: str = "align",
67
+ block_size: int = BLOCK_SIZE,
68
+ prefix_match_unit: int | None = None,
69
+ num_spec: int = 0,
70
+) -> KimiK3ROCmKDAMetadataBuilder:
71
+ vllm_config = create_vllm_config(
72
+ model_name="Qwen/Qwen3.5-0.8B", block_size=block_size
73
+ )
74
+ vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode
75
+ vllm_config.cache_config.prefix_match_unit = prefix_match_unit
76
+ if num_spec:
77
+ vllm_config.speculative_config = _StubSpeculativeConfig(num_spec)
78
+ return KimiK3ROCmKDAMetadataBuilder(
79
+ kv_cache_spec=MambaSpec(
80
+ block_size=block_size,
81
+ shapes=((16, 64),),
82
+ dtypes=(torch.float16,),
83
+ mamba_cache_mode=mamba_cache_mode,
84
+ num_prefill_checkpoint_blocks=num_prefill_checkpoint_blocks,
85
+ prefill_checkpoint_alignment=(
86
+ KDA_CHECKPOINT_ALIGNMENT if num_prefill_checkpoint_blocks else None
87
+ ),
88
+ num_speculative_blocks=num_spec,
89
+ ),
90
+ layer_names=["layer.0"],
91
+ vllm_config=vllm_config,
92
+ device=torch.device("cuda"),
93
+ )
94
+
95
+
96
+def _build(batch: BatchSpec, block_size: int = BLOCK_SIZE, **kw):
97
+ common = create_common_attn_metadata(
98
+ batch, block_size, torch.device("cuda"), arange_block_indices=True
99
+ )
100
+ return _builder(block_size=block_size, **kw).build(0, common)
101
+
102
+
103
+def _i32(values: list[int]) -> torch.Tensor:
104
+ return torch.tensor(values, dtype=torch.int32, device="cuda")
105
+
106
+
107
+def test_checkpoint_metadata_targets_last_aligned_boundary() -> None:
108
+ # Slot p holds the state after (p + 1) * block_size tokens, so a 200-token
109
+ # prompt checkpoints at 192 into slot 2 and a 128-token one at 64 into slot
110
+ # 0. The block table is an arange, so request 1 owns blocks 4..7.
111
+ md = _build(BatchSpec(seq_lens=[200, 128], query_lens=[200, 128]))
112
+
113
+ assert md.checkpoint is not None
114
+ torch.testing.assert_close(md.checkpoint.checkpoint_offsets, _i32([192, 64]))
115
+ torch.testing.assert_close(md.checkpoint.state_indices, _i32([2, 4]))
116
+
117
+
118
+def test_checkpoint_offset_is_relative_to_query_start() -> None:
119
+ # 256 tokens computed, 144 more this step: the boundary at 384 lands 128
120
+ # tokens into the chunk, and cdiv(400, 64) - 2 = 5 is its slot.
121
+ md = _build(BatchSpec(seq_lens=[400, 64], query_lens=[144, 64]))
122
+
123
+ assert md.checkpoint is not None
124
+ torch.testing.assert_close(md.checkpoint.checkpoint_offsets, _i32([128, 0]))
125
+ torch.testing.assert_close(md.checkpoint.state_indices, _i32([5, NULL_BLOCK_ID]))
126
+
127
+
128
+def test_checkpoint_metadata_avoids_live_state_row() -> None:
129
+ # Writing the row the walk itself owns would race it.
130
+ md = _build(BatchSpec(seq_lens=[200, 130], query_lens=[200, 130]))
131
+
132
+ assert md.checkpoint is not None
133
+ # `mamba_get_block_table_tensor` hands the walk block (seq_len - 1) // bs.
134
+ live = torch.tensor([(200 - 1) // BLOCK_SIZE, (130 - 1) // BLOCK_SIZE])
135
+ table = torch.arange(2 * 4).view(2, 4)
136
+ live_rows = table[torch.arange(2), live].cuda()
137
+ assert not bool((md.checkpoint.state_indices == live_rows).any())
138
+
139
+
140
+def test_checkpoint_metadata_skips_chunk_without_boundary() -> None:
141
+ # 128 computed, 40 more: the only boundary at 128 is the state this chunk
142
+ # started from, so there is nothing new to snapshot.
143
+ assert _build(BatchSpec(seq_lens=[168], query_lens=[40])).checkpoint is None
144
+
145
+
146
+def test_checkpoint_metadata_skips_null_reservation() -> None:
147
+ # Only the chunk that finishes a prompt gets a block reserved; every
148
+ # earlier one leaves the shared null block, which nothing hashes and
149
+ # writing it would corrupt every other request's scratch row.
150
+ batch = BatchSpec(seq_lens=[200], query_lens=[200])
151
+ common = create_common_attn_metadata(
152
+ batch, BLOCK_SIZE, torch.device("cuda"), arange_block_indices=True
153
+ )
154
+ # cdiv(200, 64) - 2 = 2 is the checkpoint slot.
155
+ common.block_table_tensor[0, 2] = NULL_BLOCK_ID
156
+
157
+ md = _builder().build(0, common)
158
+
159
+ assert md.checkpoint is not None
160
+ torch.testing.assert_close(md.checkpoint.state_indices, _i32([NULL_BLOCK_ID]))
161
+
162
+
163
+def test_checkpoint_metadata_absent_without_reserved_block() -> None:
164
+ md = _build(
165
+ BatchSpec(seq_lens=[200], query_lens=[200]),
166
+ num_prefill_checkpoint_blocks=0,
167
+ )
168
+ assert getattr(md, "checkpoint", None) is None
169
+
170
+
171
+def test_spec_checkpoint_metadata_gathers_non_spec_rows_by_index() -> None:
172
+ # gdn_attn selects the non-spec rows by mask and reclassifies their
173
+ # one-token decodes as prefills, so the chunk kernel is handed rows 1, 3
174
+ # and 4 here while a num_decodes: slice would hand it rows 0, 1 and 2.
175
+ num_spec = 2
176
+ batch = BatchSpec(
177
+ seq_lens=[300, 200, 400, 512, 128],
178
+ query_lens=[num_spec + 1, 200, num_spec + 1, 1, 128],
179
+ )
180
+ common = create_common_attn_metadata(
181
+ batch, BLOCK_SIZE, torch.device("cuda"), arange_block_indices=True
182
+ )
183
+ # Align mode reads 1 + num_speculative_blocks columns, so widen the table.
184
+ rows, cols = common.block_table_tensor.shape
185
+ width = cols + num_spec
186
+ common.block_table_tensor = torch.arange(
187
+ rows * width, dtype=torch.int32, device="cuda"
188
+ ).view(rows, width)
189
+
190
+ md = _builder(num_spec=num_spec).build(
191
+ 0,
192
+ common,
193
+ num_accepted_tokens=torch.ones(rows, dtype=torch.int32, device="cuda"),
194
+ num_decode_draft_tokens_cpu=torch.tensor(
195
+ [num_spec, -1, num_spec, -1, -1], dtype=torch.int32
196
+ ),
197
+ )
198
+
199
+ assert md.spec_sequence_masks is not None
200
+ assert md.num_decodes == 0 and md.num_prefills == 3
201
+ assert md.checkpoint is not None
202
+ # Row 3 resumes mid-block, so its chunk holds no reusable boundary.
203
+ torch.testing.assert_close(md.checkpoint.checkpoint_offsets, _i32([192, 0, 64]))
204
+ # Request r owns blocks r*width.., so slot cdiv(200, 64) - 2 = 2 for
205
+ # request 1 and cdiv(128, 64) - 2 = 0 for request 4.
206
+ torch.testing.assert_close(
207
+ md.checkpoint.state_indices, _i32([width + 2, NULL_BLOCK_ID, 4 * width])
208
+ )
209
+
210
+
211
+def test_checkpoint_metadata_absent_outside_align_mode() -> None:
212
+ md = _build(
213
+ BatchSpec(seq_lens=[200], query_lens=[200]),
214
+ mamba_cache_mode="none",
215
+ )
216
+ assert getattr(md, "checkpoint", None) is None
217
+
218
+
219
+# ---------------------------------------------------------------------------
220
+# Conv half
221
+# ---------------------------------------------------------------------------
222
+
223
+DIM = 384
224
+WIDTH = 4
225
+STATE_LEN = WIDTH - 1
226
+
227
+
228
+def store_conv_checkpoints(x, conv_state, cu, offsets, rows, state_len) -> None:
229
+ KimiK3ROCmKDAPrefillCheckpointExporter(state_len=state_len).export(
230
+ MambaPrefillCheckpointMetadata(offsets, rows),
231
+ raw_qkv=x,
232
+ conv_state=conv_state,
233
+ cu_seqlens=cu,
234
+ )
235
+
236
+
237
+def _conv_fixture(
238
+ seqlens: list[int], slots: int = 16, seed: int = 0, num_spec: int = 0
239
+):
240
+ # kda_state_shape widens the conv row to conv_kernel_size - 1 + num_spec,
241
+ # but the window every reader starts from stays the leading STATE_LEN.
242
+ torch.manual_seed(seed)
243
+ total = sum(seqlens)
244
+ x = torch.randn(total, DIM, device="cuda", dtype=torch.bfloat16)
245
+ cu = torch.tensor(
246
+ [0] + torch.tensor(seqlens).cumsum(0).tolist(),
247
+ device="cuda",
248
+ dtype=torch.int32,
249
+ )
250
+ conv_state = torch.full(
251
+ (slots, DIM, STATE_LEN + num_spec),
252
+ float("nan"),
253
+ device="cuda",
254
+ dtype=torch.bfloat16,
255
+ )
256
+ return x, cu, conv_state
257
+
258
+
259
+def test_conv_checkpoint_saves_window_before_offset() -> None:
260
+ seqlens = [200, 130]
261
+ x, cu, conv_state = _conv_fixture(seqlens)
262
+ offsets = _i32([192, 128])
263
+ rows = _i32([5, 9])
264
+
265
+ store_conv_checkpoints(x, conv_state, cu, offsets, rows, STATE_LEN)
266
+
267
+ for seq, (offset, row) in enumerate(zip([192, 128], [5, 9])):
268
+ end = int(cu[seq]) + offset
269
+ expected = x[end - STATE_LEN : end].transpose(0, 1)
270
+ torch.testing.assert_close(conv_state[row], expected)
271
+
272
+
273
+def test_conv_checkpoint_skips_opted_out_sequences() -> None:
274
+ x, cu, conv_state = _conv_fixture([200, 200])
275
+ before = conv_state.clone()
276
+
277
+ store_conv_checkpoints(
278
+ x, conv_state, cu, _i32([0, 192]), _i32([5, NULL_BLOCK_ID]), STATE_LEN
279
+ )
280
+
281
+ assert torch.equal(conv_state.isnan(), before.isnan())
282
+
283
+
284
+def test_exporter_stores_conv_window_through_shared_interface() -> None:
285
+ x, cu, conv_state = _conv_fixture([200, 200])
286
+ exporter = KimiK3ROCmKDAPrefillCheckpointExporter(state_len=STATE_LEN)
287
+ assert isinstance(exporter, MambaPrefillCheckpointExporter)
288
+ checkpoint = MambaPrefillCheckpointMetadata(
289
+ _i32([192, 0]), _i32([5, NULL_BLOCK_ID])
290
+ )
291
+
292
+ exporter.export(checkpoint, raw_qkv=x, conv_state=conv_state, cu_seqlens=cu)
293
+
294
+ torch.testing.assert_close(conv_state[5], x[192 - STATE_LEN : 192].transpose(0, 1))
295
+ untouched = torch.ones(16, dtype=torch.bool, device="cuda")
296
+ untouched[5] = False
297
+ assert bool(conv_state[untouched].isnan().all())
298
+
299
+
300
+@pytest.mark.parametrize("num_spec", [0, 2, 8])
301
+def test_conv_checkpoint_matches_prefill_truncated_at_offset(num_spec: int) -> None:
302
+ # causal_conv1d_fn's own state write is the reference: a prefill that
303
+ # stopped at the offset leaves exactly the row a later prefix-cache hit
304
+ # reads back. Speculative decoding widens the row without moving that
305
+ # window, so a checkpoint sized off the row would sit num_spec tokens in
306
+ # the past.
307
+ seqlens = [200, 130]
308
+ offsets = [192, 64]
309
+ x, cu, conv_state = _conv_fixture(seqlens, seed=3, num_spec=num_spec)
310
+ weight = torch.randn(DIM, WIDTH, device="cuda", dtype=torch.bfloat16) * 0.3
311
+ bias = torch.randn(DIM, device="cuda", dtype=torch.bfloat16)
312
+ rows = [5, 9]
313
+
314
+ store_conv_checkpoints(x, conv_state, cu, _i32(offsets), _i32(rows), STATE_LEN)
315
+
316
+ # A prefill of just the prefix, into its own cache. Row 0 is the null block
317
+ # the conv kernel skips by design, so the reference rows start at 1.
318
+ prefix = torch.cat(
319
+ [x[int(cu[n]) : int(cu[n]) + off] for n, off in enumerate(offsets)]
320
+ )
321
+ prefix_cu = torch.tensor(
322
+ [0] + torch.tensor(offsets).cumsum(0).tolist(),
323
+ device="cuda",
324
+ dtype=torch.int32,
325
+ )
326
+ reference = torch.zeros_like(conv_state)
327
+ causal_conv1d_fn(
328
+ prefix.transpose(0, 1),
329
+ weight,
330
+ bias,
331
+ activation="silu",
332
+ conv_states=reference,
333
+ has_initial_state=torch.zeros(2, dtype=torch.bool, device="cuda"),
334
+ cache_indices=_i32([1, 2]),
335
+ query_start_loc=prefix_cu,
336
+ )
337
+
338
+ # causal_conv1d_fn writes only the leading STATE_LEN columns.
339
+ window = slice(None), slice(None, STATE_LEN)
340
+ torch.testing.assert_close(conv_state[rows[0]][window], reference[1][window])
341
+ torch.testing.assert_close(conv_state[rows[1]][window], reference[2][window])
342
+
343
+
344
+def test_conv_checkpoint_leaves_spec_tail_of_row_alone() -> None:
345
+ # The draft-token columns belong to causal_conv1d_update. Filling the whole
346
+ # row would hand every reader, which indexes from column 0, a window
347
+ # shifted num_spec tokens into the past.
348
+ num_spec = 4
349
+ x, cu, conv_state = _conv_fixture([200], num_spec=num_spec, seed=11)
350
+
351
+ store_conv_checkpoints(x, conv_state, cu, _i32([192]), _i32([5]), STATE_LEN)
352
+
353
+ torch.testing.assert_close(
354
+ conv_state[5, :, :STATE_LEN], x[192 - STATE_LEN : 192].transpose(0, 1)
355
+ )
356
+ assert bool(conv_state[5, :, STATE_LEN:].isnan().all()), (
357
+ "the checkpoint wrote into the speculative-decode tail of the row"
358
+ )
359
+
360
+
361
+def test_conv_checkpoint_honours_transposed_cache_view() -> None:
362
+ # The SD conv layout reaches the layer transposed, not contiguous.
363
+ x, cu, _ = _conv_fixture([200])
364
+ raw = torch.full(
365
+ (16, STATE_LEN, DIM), float("nan"), device="cuda", dtype=torch.bfloat16
366
+ )
367
+ conv_state = raw.transpose(-1, -2)
368
+ assert not conv_state.is_contiguous()
369
+
370
+ store_conv_checkpoints(x, conv_state, cu, _i32([192]), _i32([7]), STATE_LEN)
371
+
372
+ torch.testing.assert_close(conv_state[7], x[192 - STATE_LEN : 192].transpose(0, 1))
373
+ untouched = torch.ones(16, dtype=torch.bool, device="cuda")
374
+ untouched[7] = False
375
+ assert bool(raw[untouched].isnan().all())
376
+
377
+
378
+def test_checkpoint_metadata_indexes_prefill_rows_after_decodes() -> None:
379
+ # The chunk kernel receives the prefill tail alone, so the metadata has to
380
+ # be indexed the same way or it snapshots the wrong request's boundary.
381
+ md = _build(BatchSpec(seq_lens=[64, 200], query_lens=[1, 200]))
382
+
383
+ assert md.num_decodes == 1 and md.num_prefills == 1
384
+ assert md.checkpoint is not None
385
+ torch.testing.assert_close(md.checkpoint.checkpoint_offsets, _i32([192]))
386
+ # Request 1 owns blocks 4..7, so its slot-2 state block is 6.
387
+ torch.testing.assert_close(md.checkpoint.state_indices, _i32([6]))
388
+
389
+
390
+def test_conv_checkpoint_addresses_row_past_int32_limit() -> None:
391
+ # state_stride_0 is the whole mamba page, padded up to the attention page,
392
+ # so row * stride passes 2**31 a few thousand blocks into a large cache and
393
+ # an int32 wrap stores outside the tensor. The margin on the plain prefill
394
+ # path is 1.6x, which is luck: a bigger cache, a higher utilization or the
395
+ # wider row speculative decoding allocates all close it.
396
+ stride_row = DIM * (STATE_LEN + 8)
397
+ row = 2**31 // stride_row + 1
398
+ need = (row + 1) * stride_row * 2
399
+ free, _ = torch.accelerator.get_memory_info()
400
+ if free < need + (4 << 30):
401
+ pytest.skip(f"needs {need >> 30} GiB of free VRAM to reach the wrap")
402
+
403
+ x, cu, _ = _conv_fixture([200])
404
+ conv_state = torch.zeros(
405
+ row + 1, DIM, STATE_LEN + 8, device="cuda", dtype=torch.bfloat16
406
+ )
407
+ assert conv_state.stride(0) * row > 2**31, "fixture does not reach the wrap"
408
+
409
+ store_conv_checkpoints(x, conv_state, cu, _i32([192]), _i32([row]), STATE_LEN)
410
+
411
+ torch.testing.assert_close(
412
+ conv_state[row, :, :STATE_LEN], x[192 - STATE_LEN : 192].transpose(0, 1)
413
+ )
414
+ assert not bool(conv_state[:row].any()), "the store wrote outside its row"
415
+
416
+
417
+def test_conv_checkpoint_reads_row_strided_projection() -> None:
418
+ # mixed_qkv is a band of the fused QKVGFAB projection, so its row stride is
419
+ # the whole projection width, not DIM.
420
+ torch.manual_seed(5)
421
+ total, pad = 200, 177
422
+ projection = torch.randn(total, DIM + pad, device="cuda", dtype=torch.bfloat16)
423
+ x = projection[:, :DIM]
424
+ assert x.stride(0) == DIM + pad
425
+
426
+ cu = _i32([0, total])
427
+ conv_state = torch.full(
428
+ (16, DIM, STATE_LEN), float("nan"), device="cuda", dtype=torch.bfloat16
429
+ )
430
+ store_conv_checkpoints(x, conv_state, cu, _i32([192]), _i32([4]), STATE_LEN)
431
+
432
+ torch.testing.assert_close(conv_state[4], x[192 - STATE_LEN : 192].transpose(0, 1))
433
+
434
+
435
+def test_finer_prefix_match_unit_moves_checkpoint_to_hash_grid() -> None:
436
+ # A partial prefix hit matches at prefix_match_unit, so a state captured
437
+ # only at block boundaries is too coarse to resume from. Block 256 with
438
+ # unit 128 moves a 700-token prompt's checkpoint from 512 to 640.
439
+ md = _build(
440
+ BatchSpec(seq_lens=[700], query_lens=[700]),
441
+ block_size=256,
442
+ prefix_match_unit=128,
443
+ )
444
+
445
+ assert md.checkpoint is not None
446
+ torch.testing.assert_close(md.checkpoint.checkpoint_offsets, _i32([640]))
447
+ # Slot is still block-indexed: cdiv(700, 256) - 2 = 1, request 0 owns 0..2.
448
+ torch.testing.assert_close(md.checkpoint.state_indices, _i32([1]))
449
+
450
+
451
+def test_checkpoint_metadata_skips_unaligned_offset() -> None:
452
+ # The scheduler declines a boundary off the published alignment and
453
+ # reserves nothing, so the builder has to reach the same verdict or it
454
+ # writes a block nobody hashed.
455
+ # Block 192, unit 96: a 700-token prompt's last hash boundary is 672, which
456
+ # is 32 tokens off the chunk grid.
457
+ md = _build(
458
+ BatchSpec(seq_lens=[700], query_lens=[700]),
459
+ block_size=192,
460
+ prefix_match_unit=96,
461
+ )
462
+
463
+ assert md.checkpoint is None
464
+
465
+
466
+# ---------------------------------------------------------------------------
467
+# Opting in
468
+# ---------------------------------------------------------------------------
469
+
470
+
471
+def _spec_for(monkeypatch, **attrs) -> MambaSpec:
472
+ """Run the layer's `get_kv_cache_spec` over a stubbed base spec."""
473
+ base = MambaSpec(block_size=BLOCK_SIZE, shapes=((16, 64),), dtypes=(torch.float16,))
474
+ monkeypatch.setattr(
475
+ GatedDeltaNetAttention, "get_kv_cache_spec", lambda self, cfg: base
476
+ )
477
+ layer = KimiK3DeltaAttention.__new__(KimiK3DeltaAttention)
478
+ for name, value in dict(
479
+ use_prefill_checkpoint=True,
480
+ use_fused_chunk=True,
481
+ num_spec=0,
482
+ use_safe_gate=True,
483
+ ).items():
484
+ setattr(layer, name, attrs.get(name, value))
485
+ vllm_config = create_vllm_config(model_name="Qwen/Qwen3.5-0.8B")
486
+ vllm_config.cache_config.mamba_cache_mode = "align"
487
+ return layer.get_kv_cache_spec(vllm_config)
488
+
489
+
490
+def test_opt_in_publishes_chunk_walk_alignment(monkeypatch) -> None:
491
+ # The scheduler learns the kernel's granularity only from the spec, so
492
+ # reserving a block without publishing the alignment would hash a state
493
+ # that is never written.
494
+ spec = _spec_for(monkeypatch)
495
+
496
+ assert spec.num_prefill_checkpoint_blocks == 1
497
+ assert spec.prefill_checkpoint_alignment == KDA_CHECKPOINT_ALIGNMENT
498
+
499
+
500
+@pytest.mark.parametrize(
501
+ "attr,value",
502
+ [
503
+ ("use_prefill_checkpoint", False),
504
+ ("use_fused_chunk", False),
505
+ ("use_safe_gate", False),
506
+ ],
507
+)
508
+def test_paths_that_cannot_export_stay_opted_out(monkeypatch, attr, value) -> None:
509
+ spec = _spec_for(monkeypatch, **{attr: value})
510
+
511
+ assert spec.num_prefill_checkpoint_blocks == 0
512
+ assert spec.prefill_checkpoint_alignment is None
513
+
514
+
515
+def test_speculative_decoding_does_not_opt_out(monkeypatch) -> None:
516
+ # num_spec reaches the kernels as a wider conv row and an interleaved
517
+ # batch, both of which the layer handles.
518
+ spec = _spec_for(monkeypatch, num_spec=8)
519
+
520
+ assert spec.num_prefill_checkpoint_blocks == 1
521
+ assert spec.prefill_checkpoint_alignment == KDA_CHECKPOINT_ALIGNMENT
tests/models/kimi_k3/test_amd_kda_chunk.py
+56
-6
@@ -13,17 +13,17 @@ import torch
13
from vllm.platforms import current_platform
14
15
16
-def _on_gfx950() -> bool:
17
- if not current_platform.is_rocm():
16
+def _fused_chunk_available() -> bool:
17
+ if not (current_platform.is_rocm() and torch.cuda.is_available()):
18
return False
19
- from vllm.platforms.rocm import on_gfx950
19
+ from vllm.models.kimi_k3.amd.ops.kda_chunk import is_fused_kda_chunk_supported
20
21
- return on_gfx950()
21
+ return is_fused_kda_chunk_supported()
22
23
24
pytestmark = pytest.mark.skipif(
25
- not _on_gfx950(),
26
- reason="The fused KDA chunk kernel is only built for gfx950",
25
+ not _fused_chunk_available(),
26
+ reason="The fused KDA chunk kernel needs a gfx950 build that includes it",
27
)
28
29
HEAD_DIM = 128
@@ -701,6 +701,56 @@ def test_cold_rows_start_from_zero_not_from_cache_junk(use_fused: bool) -> None:
701
torch.testing.assert_close(cache[idx.long()], ht_ref, rtol=5e-3, atol=5e-3)
702
703
704
+@pytest.mark.parametrize("use_fused", [False, True])
705
+def test_state_indices_accept_a_block_table_column(use_fused: bool) -> None:
706
+ # The layer passes block_table[:, 0], which is contiguous only while the
707
+ # table has one column. Speculative decoding reserves a state page per
708
+ # draft token, so the table widens and that column becomes strided.
709
+ _requires_kernel()
710
+ seqlens = [512, 320, 64]
711
+ inp = _inputs(seqlens, seed=43)
712
+ n = len(seqlens)
713
+ rows = [3 * i + 5 for i in range(n)]
714
+ cache, idx, warm = _paged_cache(inp, [True] * n, rows, slots=3 * n + 12)
715
+ ref_cache = cache.clone()
716
+
717
+ table = torch.zeros(n, 9, dtype=torch.int32, device="cuda")
718
+ table[:, 0] = idx
719
+ column = table[:, 0]
720
+ assert not column.is_contiguous()
721
+
722
+ o_ref, _ = _run_paged(inp, ref_cache, idx, warm, use_fused=use_fused)
723
+ o_got, _ = _run_paged(inp, cache, column, warm, use_fused=use_fused)
724
+
725
+ torch.testing.assert_close(o_got, o_ref)
726
+ torch.testing.assert_close(cache, ref_cache)
727
+
728
+
729
+@pytest.mark.parametrize("use_fused", [False, True])
730
+def test_padded_zero_length_sequence_is_inert(use_fused: bool) -> None:
731
+ # A batch carrying speculative decodes pads its non-spec group with
732
+ # zero-length sequences whose cache row still belongs to a request, so
733
+ # consuming no tokens has to mean writing nothing and perturbing nobody.
734
+ _requires_kernel()
735
+ live = [200, 1, 128]
736
+ # _inputs draws h0 last, so the padded batch's live rows are the same draw
737
+ # as the unpadded one's and the two runs compare elementwise.
738
+ padded, unpadded = _inputs([*live, 0], seed=41), _inputs(live, seed=41)
739
+ rows = [3 * i + 5 for i in range(4)]
740
+
741
+ cache, idx, warm = _paged_cache(padded, [True] * 4, rows, slots=24)
742
+ before = cache.clone()
743
+ _run_paged(padded, cache, idx, warm, use_fused=use_fused)
744
+
745
+ ref_cache, ref_idx, ref_warm = _paged_cache(unpadded, [True] * 3, rows[:3], 24)
746
+ _run_paged(unpadded, ref_cache, ref_idx, ref_warm, use_fused=use_fused)
747
+
748
+ assert torch.equal(cache[rows[3]], before[rows[3]]), (
749
+ "the padded sequence's cache row was written"
750
+ )
751
+ torch.testing.assert_close(cache[ref_idx.long()], ref_cache[ref_idx.long()])
752
+
753
+
754
@pytest.mark.parametrize("groups", [2, 4])
755
def test_paged_state_survives_the_group_split(groups: int, monkeypatch) -> None:
756
"""At G>1 the read and the write are different blocks of one launch."""
vllm/models/kimi_k3/amd/kda.py
+56
-6
@@ -1,6 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from dataclasses import replace
5
+
6
import torch
7
from einops import rearrange
8
from torch import nn
@@ -38,7 +40,11 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
40
from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader
41
from vllm.model_executor.utils import set_weight_attrs
42
from vllm.models.kimi_k3.amd.kda_metadata import KimiK3ROCmKDABackend
43
+from vllm.models.kimi_k3.amd.ops.kda_checkpoint import (
44
+ KimiK3ROCmKDAPrefillCheckpointExporter,
45
+)
46
from vllm.models.kimi_k3.amd.ops.kda_chunk import (
47
+ KDA_CHECKPOINT_ALIGNMENT,
48
is_fused_kda_chunk_supported,
49
)
50
from vllm.models.kimi_k3.amd.ops.kda_decode import (
@@ -55,6 +61,7 @@ from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated
61
from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
62
from vllm.v1.attention.backend import AttentionBackend
63
from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
64
+from vllm.v1.kv_cache_interface import MambaSpec
65
66
logger = init_logger(__name__)
67
@@ -79,6 +86,23 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
86
num_spec=self.num_spec,
87
)
88
89
+ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> MambaSpec:
90
+ spec = super().get_kv_cache_spec(vllm_config)
91
+ assert isinstance(spec, MambaSpec)
92
+ # Only the fused path exports the checkpoint snapshots
93
+ can_checkpoint = (
94
+ self.use_prefill_checkpoint and self.use_fused_chunk and self.use_safe_gate
95
+ )
96
+ if can_checkpoint and vllm_config.cache_config.mamba_cache_mode == "align":
97
+ logger.info_once("Kimi-K3 KDA prefill checkpoint enabled.")
98
+ return replace(
99
+ spec,
100
+ num_prefill_checkpoint_blocks=int(can_checkpoint),
101
+ prefill_checkpoint_alignment=(
102
+ KDA_CHECKPOINT_ALIGNMENT if can_checkpoint else None
103
+ ),
104
+ )
105
+
106
def __init__(
107
self,
108
config: KimiLinearConfig,
@@ -226,6 +250,15 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
250
"Kimi-K3 KDA prefill backend: %s",
251
"fused" if self.use_fused_chunk else "triton",
252
)
253
+ # TODO: Remove the extra config arg once tested
254
+ self.use_prefill_checkpoint = (
255
+ additional_config.get("kda_prefill_checkpoint", True)
256
+ if isinstance(additional_config, dict)
257
+ else True
258
+ )
259
+ self._checkpoint_exporter = KimiK3ROCmKDAPrefillCheckpointExporter(
260
+ state_len=self.conv_size - 1
261
+ )
262
263
self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid")
264
decode_norm_weight = None
@@ -511,10 +544,9 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
544
assert non_spec_state_indices_tensor is not None
545
assert has_initial_state is not None
546
514
- # A mixed non-spec batch is decode-first: the decodes are
515
- # length-1 sequences, and the chunk kernel returns NaN for those.
516
- # Send them to the recurrent kernel and give the chunk kernel the
517
- # prefill tail only.
547
+ # Decode-first, so peel the length-1 decodes to the recurrent
548
+ # kernel. A spec batch interleaves its non-spec rows, so the
549
+ # peel cannot apply and the chunk kernel takes them.
550
core_attn_out_decode = None
551
split_non_spec = spec_sequence_masks is None and m.num_decodes > 0
552
# Without a spec split the non-spec tokens are exactly
@@ -522,7 +554,6 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
554
# can write their slice straight into it. With one, the tokens
555
# are interleaved and have to be scattered by index afterwards.
556
direct = spec_sequence_masks is None
525
- use_fused_chunk = self.use_fused_chunk and spec_sequence_masks is None
557
if split_non_spec:
558
assert non_spec_query_start_loc is not None
559
nd_tok = m.num_decode_tokens
@@ -559,6 +590,22 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
590
prefill_has_initial_state = has_initial_state
591
nd_tok = 0
592
593
+ checkpoint = getattr(m, "checkpoint", None)
594
+ checkpoint_state = None
595
+ checkpoint_offsets = None
596
+ checkpoint_state_indices = None
597
+ if checkpoint is not None:
598
+ assert prefill_query_start_loc is not None
599
+ checkpoint_state = recurrent_state
600
+ checkpoint_offsets = checkpoint.checkpoint_offsets
601
+ checkpoint_state_indices = checkpoint.state_indices
602
+ self._checkpoint_exporter.export(
603
+ checkpoint,
604
+ raw_qkv=mixed_qkv_ns[nd_tok:],
605
+ conv_state=conv_state,
606
+ cu_seqlens=prefill_query_start_loc,
607
+ )
608
+
609
(
610
core_attn_out_non_spec,
611
_,
@@ -578,8 +625,11 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
625
cu_seqlens=prefill_query_start_loc,
626
chunk_indices=m.chunk_indices,
627
chunk_offsets=m.chunk_offsets,
581
- use_fused_chunk=use_fused_chunk,
628
+ use_fused_chunk=self.use_fused_chunk,
629
out=core_attn_out[:, nd_tok:num_actual_tokens] if direct else None,
630
+ checkpoint_state=checkpoint_state,
631
+ checkpoint_offsets=checkpoint_offsets,
632
+ checkpoint_state_indices=checkpoint_state_indices,
633
)
634
# chunk_kda_prefill updates `recurrent_state` in place, so
635
# there is no final state to write back here.
vllm/models/kimi_k3/amd/kda_metadata.py
+56
@@ -7,16 +7,23 @@ The request classification and cudagraph staging intentionally mirror
7
differently on device rather than on the host.
8
"""
9
10
+from dataclasses import replace
11
+
12
import torch
13
14
from vllm.logger import init_logger
15
+from vllm.model_executor.layers.mamba.checkpoint import (
16
+ MambaPrefillCheckpointMetadata,
17
+)
18
from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE
19
from vllm.triton_utils import tl, triton
20
from vllm.utils.math_utils import next_power_of_2
21
from vllm.v1.attention.backends.gdn_attn import (
22
GDNAttentionBackend,
23
+ GDNAttentionMetadata,
24
GDNAttentionMetadataBuilder,
25
)
26
+from vllm.v1.attention.backends.utils import CommonAttentionMetadata
27
28
logger = init_logger(__name__)
29
@@ -86,6 +93,55 @@ def prepare_chunk_metadata_device(
93
94
95
class KimiK3ROCmKDAMetadataBuilder(GDNAttentionMetadataBuilder):
96
+ def build( # type: ignore[override]
97
+ self,
98
+ common_prefix_len: int,
99
+ common_attn_metadata: CommonAttentionMetadata,
100
+ num_accepted_tokens: torch.Tensor | None = None,
101
+ num_decode_draft_tokens_cpu: torch.Tensor | None = None,
102
+ fast_build: bool = False,
103
+ ) -> GDNAttentionMetadata:
104
+ attn_metadata = super().build(
105
+ common_prefix_len,
106
+ common_attn_metadata,
107
+ num_accepted_tokens,
108
+ num_decode_draft_tokens_cpu,
109
+ fast_build,
110
+ )
111
+ if attn_metadata.checkpoint is None:
112
+ return attn_metadata
113
+ # The shared builder keys checkpoint rows on every request; the ROCm
114
+ # prefill path indexes them by prefill_query_start_loc, so rebuild with
115
+ # those rows.
116
+ return replace(
117
+ attn_metadata,
118
+ checkpoint=self._build_checkpoint_metadata(
119
+ common_attn_metadata, attn_metadata, num_decode_draft_tokens_cpu
120
+ ),
121
+ )
122
+
123
+ def _build_checkpoint_metadata(
124
+ self,
125
+ m: CommonAttentionMetadata,
126
+ attn_metadata: GDNAttentionMetadata,
127
+ num_decode_draft_tokens_cpu: torch.Tensor | None,
128
+ ) -> MambaPrefillCheckpointMetadata | None:
129
+ if attn_metadata.num_prefills == 0:
130
+ return None
131
+ # request_rows must line up with prefill_query_start_loc, either the
132
+ # prefill tail of a decode-first batch, or every non-spec row
133
+ # including padding
134
+ num_decodes = attn_metadata.num_decodes
135
+ request_rows = list(
136
+ range(num_decodes, num_decodes + attn_metadata.num_prefills)
137
+ )
138
+ if attn_metadata.spec_sequence_masks is not None:
139
+ assert num_decode_draft_tokens_cpu is not None
140
+ request_rows = (
141
+ (num_decode_draft_tokens_cpu < 0).nonzero().flatten().tolist()
142
+ )
143
+ return self.checkpoint_builder.build(m, request_rows)
144
+
145
def _build_chunk_metadata(
146
self,
147
prefill_query_start_loc: torch.Tensor,
vllm/models/kimi_k3/amd/ops/kda_checkpoint.py
new
+41
@@ -0,0 +1,41 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Conv half of the ROCm Kimi-K3 prefill checkpoint."""
4
+
5
+from dataclasses import dataclass
6
+
7
+import torch
8
+
9
+from vllm.model_executor.layers.mamba.checkpoint import (
10
+ MambaPrefillCheckpointMetadata,
11
+)
12
+from vllm.model_executor.layers.mamba.kda_checkpoint import (
13
+ FlashKDAPrefillCheckpointExporter,
14
+)
15
+
16
+
17
+@dataclass(frozen=True)
18
+class KimiK3ROCmKDAPrefillCheckpointExporter(FlashKDAPrefillCheckpointExporter):
19
+ """Conv window only; `fused_kda_chunk` writes the recurrent half itself.
20
+
21
+ `state_len` is `conv_kernel_size - 1`, not the spec-widened conv row.
22
+ """
23
+
24
+ def export( # type: ignore[override]
25
+ self,
26
+ checkpoint: MambaPrefillCheckpointMetadata,
27
+ *,
28
+ raw_qkv: torch.Tensor,
29
+ conv_state: torch.Tensor,
30
+ cu_seqlens: torch.Tensor,
31
+ ) -> None:
32
+ assert self.state_len is not None
33
+ # An empty recurrent checkpoint masks the kernel's recurrent half off.
34
+ super().export(
35
+ checkpoint,
36
+ raw_qkv=raw_qkv,
37
+ conv_state=conv_state,
38
+ recurrent_checkpoint=conv_state[:, :0],
39
+ recurrent_state=conv_state,
40
+ cu_seqlens=cu_seqlens,
41
+ )
vllm/models/kimi_k3/amd/ops/kda_chunk.py
+11
-1
@@ -16,6 +16,7 @@ from vllm.third_party.flash_linear_attention.ops.index import (
16
prepare_chunk_indices,
17
prepare_chunk_offsets,
18
)
19
+from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
20
21
CHUNK_SIZE = 64
22
HEAD_DIM = 128
@@ -28,6 +29,15 @@ _DEEP_LEN = 12
29
_MAX_GROUPS = 32
30
31
32
+def _hip_checkpoint_state_indices(state_indices: torch.Tensor) -> torch.Tensor:
33
+ """NULL_BLOCK_ID marks no checkpoint; the HIP kernel skips negative rows."""
34
+ return (
35
+ torch.where(state_indices != NULL_BLOCK_ID, state_indices, -1)
36
+ .to(torch.int32)
37
+ .contiguous()
38
+ )
39
+
40
+
41
@cache
42
def is_fused_kda_chunk_supported() -> bool:
43
from vllm.platforms.rocm import on_gfx950
@@ -295,7 +305,7 @@ def fused_kda_chunk(
305
None if checkpoint_offsets is None else checkpoint_offsets.to(torch.int32),
306
None
307
if checkpoint_state_indices is None
298
- else checkpoint_state_indices.to(torch.int32).contiguous(),
308
+ else _hip_checkpoint_state_indices(checkpoint_state_indices),
309
state_cache,
310
None if state_indices is None else state_indices.to(torch.int32).contiguous(),
311
has_initial_state,
vllm/models/kimi_k3/amd/ops/kda_prefill.py
+7
-1
@@ -155,7 +155,13 @@ def chunk_kda_prefill(
155
aqk=ws["aqk"],
156
decay=ws["decay"],
157
# v is dead by this point, so it doubles as the output buffer.
158
- out=out if out is not None else v,
158
+ out=out
159
+ if out is not None
160
+ else (
161
+ v
162
+ if v.is_contiguous()
163
+ else torch.empty_like(v, memory_format=torch.contiguous_format)
164
+ ),
165
scale=scale,
166
cu_seqlens=cu_seqlens,
167
initial_state=initial_state,