| 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 |