From 01549796563f2e8bd0d6b08243cfb954a7b42ec8 Mon Sep 17 00:00:00 2001 From: kliuae <17350011+kliuae@users.noreply.github.com> Date: Fri, 2 Oct 2026 14:58:04 +0800 Subject: [PATCH] [ROCm][Perf] Kimi-K3 enable prefill checkpoints on ROCm (#58344) Signed-off-by: kliuae --- .../models/kimi_k3/test_amd_kda_checkpoint.py | 521 ++++++++++++++++++ tests/models/kimi_k3/test_amd_kda_chunk.py | 62 ++- vllm/models/kimi_k3/amd/kda.py | 62 ++- vllm/models/kimi_k3/amd/kda_metadata.py | 56 ++ vllm/models/kimi_k3/amd/ops/kda_checkpoint.py | 41 ++ vllm/models/kimi_k3/amd/ops/kda_chunk.py | 12 +- vllm/models/kimi_k3/amd/ops/kda_prefill.py | 8 +- 7 files changed, 748 insertions(+), 14 deletions(-) create mode 100644 tests/models/kimi_k3/test_amd_kda_checkpoint.py create mode 100644 vllm/models/kimi_k3/amd/ops/kda_checkpoint.py diff --git a/tests/models/kimi_k3/test_amd_kda_checkpoint.py b/tests/models/kimi_k3/test_amd_kda_checkpoint.py new file mode 100644 index 0000000000..0a622bfa04 --- /dev/null +++ b/tests/models/kimi_k3/test_amd_kda_checkpoint.py @@ -0,0 +1,521 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The ROCm KDA prefill checkpoint: where it is taken and what it stores. + +Two things have to line up for a checkpoint to be safe: the metadata must name +a position the chunk kernel will actually write, and the conv window saved +there must be the one a resumed prefill would have left. +""" + +import pytest +import torch + +from tests.v1.attention.utils import ( + BatchSpec, + create_common_attn_metadata, + create_vllm_config, +) +from vllm.model_executor.layers.mamba.checkpoint import ( + MambaPrefillCheckpointExporter, + MambaPrefillCheckpointMetadata, +) +from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention +from vllm.model_executor.layers.mamba.ops.causal_conv1d import causal_conv1d_fn +from vllm.models.kimi_k3.amd.kda import KimiK3DeltaAttention +from vllm.models.kimi_k3.amd.kda_metadata import KimiK3ROCmKDAMetadataBuilder +from vllm.models.kimi_k3.amd.ops.kda_checkpoint import ( + KimiK3ROCmKDAPrefillCheckpointExporter, +) +from vllm.models.kimi_k3.amd.ops.kda_chunk import ( + KDA_CHECKPOINT_ALIGNMENT, + is_fused_kda_chunk_supported, +) +from vllm.platforms import current_platform +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID +from vllm.v1.kv_cache_interface import MambaSpec + +pytestmark = pytest.mark.skipif( + not ( + current_platform.is_rocm() + and torch.cuda.is_available() + and is_fused_kda_chunk_supported() + ), + reason="The ROCm KDA prefill checkpoint is exported by the fused chunk " + "kernel, which needs a gfx950 build that includes it", +) + +# A mamba block size the chunk walk lands on. The real K3 value is 768; any +# multiple of the kernel's 64-token chunk behaves the same and keeps the +# fixtures small. +BLOCK_SIZE = 64 + + +class _StubSpeculativeConfig: + method = "dspark" + parallel_drafting = False + + def __init__(self, num_speculative_tokens: int) -> None: + self.num_speculative_tokens = num_speculative_tokens + + def use_eagle_block_drop(self) -> bool: + return False + + +def _builder( + num_prefill_checkpoint_blocks: int = 1, + mamba_cache_mode: str = "align", + block_size: int = BLOCK_SIZE, + prefix_match_unit: int | None = None, + num_spec: int = 0, +) -> KimiK3ROCmKDAMetadataBuilder: + vllm_config = create_vllm_config( + model_name="Qwen/Qwen3.5-0.8B", block_size=block_size + ) + vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode + vllm_config.cache_config.prefix_match_unit = prefix_match_unit + if num_spec: + vllm_config.speculative_config = _StubSpeculativeConfig(num_spec) + return KimiK3ROCmKDAMetadataBuilder( + kv_cache_spec=MambaSpec( + block_size=block_size, + shapes=((16, 64),), + dtypes=(torch.float16,), + mamba_cache_mode=mamba_cache_mode, + num_prefill_checkpoint_blocks=num_prefill_checkpoint_blocks, + prefill_checkpoint_alignment=( + KDA_CHECKPOINT_ALIGNMENT if num_prefill_checkpoint_blocks else None + ), + num_speculative_blocks=num_spec, + ), + layer_names=["layer.0"], + vllm_config=vllm_config, + device=torch.device("cuda"), + ) + + +def _build(batch: BatchSpec, block_size: int = BLOCK_SIZE, **kw): + common = create_common_attn_metadata( + batch, block_size, torch.device("cuda"), arange_block_indices=True + ) + return _builder(block_size=block_size, **kw).build(0, common) + + +def _i32(values: list[int]) -> torch.Tensor: + return torch.tensor(values, dtype=torch.int32, device="cuda") + + +def test_checkpoint_metadata_targets_last_aligned_boundary() -> None: + # Slot p holds the state after (p + 1) * block_size tokens, so a 200-token + # prompt checkpoints at 192 into slot 2 and a 128-token one at 64 into slot + # 0. The block table is an arange, so request 1 owns blocks 4..7. + md = _build(BatchSpec(seq_lens=[200, 128], query_lens=[200, 128])) + + assert md.checkpoint is not None + torch.testing.assert_close(md.checkpoint.checkpoint_offsets, _i32([192, 64])) + torch.testing.assert_close(md.checkpoint.state_indices, _i32([2, 4])) + + +def test_checkpoint_offset_is_relative_to_query_start() -> None: + # 256 tokens computed, 144 more this step: the boundary at 384 lands 128 + # tokens into the chunk, and cdiv(400, 64) - 2 = 5 is its slot. + md = _build(BatchSpec(seq_lens=[400, 64], query_lens=[144, 64])) + + assert md.checkpoint is not None + torch.testing.assert_close(md.checkpoint.checkpoint_offsets, _i32([128, 0])) + torch.testing.assert_close(md.checkpoint.state_indices, _i32([5, NULL_BLOCK_ID])) + + +def test_checkpoint_metadata_avoids_live_state_row() -> None: + # Writing the row the walk itself owns would race it. + md = _build(BatchSpec(seq_lens=[200, 130], query_lens=[200, 130])) + + assert md.checkpoint is not None + # `mamba_get_block_table_tensor` hands the walk block (seq_len - 1) // bs. + live = torch.tensor([(200 - 1) // BLOCK_SIZE, (130 - 1) // BLOCK_SIZE]) + table = torch.arange(2 * 4).view(2, 4) + live_rows = table[torch.arange(2), live].cuda() + assert not bool((md.checkpoint.state_indices == live_rows).any()) + + +def test_checkpoint_metadata_skips_chunk_without_boundary() -> None: + # 128 computed, 40 more: the only boundary at 128 is the state this chunk + # started from, so there is nothing new to snapshot. + assert _build(BatchSpec(seq_lens=[168], query_lens=[40])).checkpoint is None + + +def test_checkpoint_metadata_skips_null_reservation() -> None: + # Only the chunk that finishes a prompt gets a block reserved; every + # earlier one leaves the shared null block, which nothing hashes and + # writing it would corrupt every other request's scratch row. + batch = BatchSpec(seq_lens=[200], query_lens=[200]) + common = create_common_attn_metadata( + batch, BLOCK_SIZE, torch.device("cuda"), arange_block_indices=True + ) + # cdiv(200, 64) - 2 = 2 is the checkpoint slot. + common.block_table_tensor[0, 2] = NULL_BLOCK_ID + + md = _builder().build(0, common) + + assert md.checkpoint is not None + torch.testing.assert_close(md.checkpoint.state_indices, _i32([NULL_BLOCK_ID])) + + +def test_checkpoint_metadata_absent_without_reserved_block() -> None: + md = _build( + BatchSpec(seq_lens=[200], query_lens=[200]), + num_prefill_checkpoint_blocks=0, + ) + assert getattr(md, "checkpoint", None) is None + + +def test_spec_checkpoint_metadata_gathers_non_spec_rows_by_index() -> None: + # gdn_attn selects the non-spec rows by mask and reclassifies their + # one-token decodes as prefills, so the chunk kernel is handed rows 1, 3 + # and 4 here while a num_decodes: slice would hand it rows 0, 1 and 2. + num_spec = 2 + batch = BatchSpec( + seq_lens=[300, 200, 400, 512, 128], + query_lens=[num_spec + 1, 200, num_spec + 1, 1, 128], + ) + common = create_common_attn_metadata( + batch, BLOCK_SIZE, torch.device("cuda"), arange_block_indices=True + ) + # Align mode reads 1 + num_speculative_blocks columns, so widen the table. + rows, cols = common.block_table_tensor.shape + width = cols + num_spec + common.block_table_tensor = torch.arange( + rows * width, dtype=torch.int32, device="cuda" + ).view(rows, width) + + md = _builder(num_spec=num_spec).build( + 0, + common, + num_accepted_tokens=torch.ones(rows, dtype=torch.int32, device="cuda"), + num_decode_draft_tokens_cpu=torch.tensor( + [num_spec, -1, num_spec, -1, -1], dtype=torch.int32 + ), + ) + + assert md.spec_sequence_masks is not None + assert md.num_decodes == 0 and md.num_prefills == 3 + assert md.checkpoint is not None + # Row 3 resumes mid-block, so its chunk holds no reusable boundary. + torch.testing.assert_close(md.checkpoint.checkpoint_offsets, _i32([192, 0, 64])) + # Request r owns blocks r*width.., so slot cdiv(200, 64) - 2 = 2 for + # request 1 and cdiv(128, 64) - 2 = 0 for request 4. + torch.testing.assert_close( + md.checkpoint.state_indices, _i32([width + 2, NULL_BLOCK_ID, 4 * width]) + ) + + +def test_checkpoint_metadata_absent_outside_align_mode() -> None: + md = _build( + BatchSpec(seq_lens=[200], query_lens=[200]), + mamba_cache_mode="none", + ) + assert getattr(md, "checkpoint", None) is None + + +# --------------------------------------------------------------------------- +# Conv half +# --------------------------------------------------------------------------- + +DIM = 384 +WIDTH = 4 +STATE_LEN = WIDTH - 1 + + +def store_conv_checkpoints(x, conv_state, cu, offsets, rows, state_len) -> None: + KimiK3ROCmKDAPrefillCheckpointExporter(state_len=state_len).export( + MambaPrefillCheckpointMetadata(offsets, rows), + raw_qkv=x, + conv_state=conv_state, + cu_seqlens=cu, + ) + + +def _conv_fixture( + seqlens: list[int], slots: int = 16, seed: int = 0, num_spec: int = 0 +): + # kda_state_shape widens the conv row to conv_kernel_size - 1 + num_spec, + # but the window every reader starts from stays the leading STATE_LEN. + torch.manual_seed(seed) + total = sum(seqlens) + x = torch.randn(total, DIM, device="cuda", dtype=torch.bfloat16) + cu = torch.tensor( + [0] + torch.tensor(seqlens).cumsum(0).tolist(), + device="cuda", + dtype=torch.int32, + ) + conv_state = torch.full( + (slots, DIM, STATE_LEN + num_spec), + float("nan"), + device="cuda", + dtype=torch.bfloat16, + ) + return x, cu, conv_state + + +def test_conv_checkpoint_saves_window_before_offset() -> None: + seqlens = [200, 130] + x, cu, conv_state = _conv_fixture(seqlens) + offsets = _i32([192, 128]) + rows = _i32([5, 9]) + + store_conv_checkpoints(x, conv_state, cu, offsets, rows, STATE_LEN) + + for seq, (offset, row) in enumerate(zip([192, 128], [5, 9])): + end = int(cu[seq]) + offset + expected = x[end - STATE_LEN : end].transpose(0, 1) + torch.testing.assert_close(conv_state[row], expected) + + +def test_conv_checkpoint_skips_opted_out_sequences() -> None: + x, cu, conv_state = _conv_fixture([200, 200]) + before = conv_state.clone() + + store_conv_checkpoints( + x, conv_state, cu, _i32([0, 192]), _i32([5, NULL_BLOCK_ID]), STATE_LEN + ) + + assert torch.equal(conv_state.isnan(), before.isnan()) + + +def test_exporter_stores_conv_window_through_shared_interface() -> None: + x, cu, conv_state = _conv_fixture([200, 200]) + exporter = KimiK3ROCmKDAPrefillCheckpointExporter(state_len=STATE_LEN) + assert isinstance(exporter, MambaPrefillCheckpointExporter) + checkpoint = MambaPrefillCheckpointMetadata( + _i32([192, 0]), _i32([5, NULL_BLOCK_ID]) + ) + + exporter.export(checkpoint, raw_qkv=x, conv_state=conv_state, cu_seqlens=cu) + + torch.testing.assert_close(conv_state[5], x[192 - STATE_LEN : 192].transpose(0, 1)) + untouched = torch.ones(16, dtype=torch.bool, device="cuda") + untouched[5] = False + assert bool(conv_state[untouched].isnan().all()) + + +@pytest.mark.parametrize("num_spec", [0, 2, 8]) +def test_conv_checkpoint_matches_prefill_truncated_at_offset(num_spec: int) -> None: + # causal_conv1d_fn's own state write is the reference: a prefill that + # stopped at the offset leaves exactly the row a later prefix-cache hit + # reads back. Speculative decoding widens the row without moving that + # window, so a checkpoint sized off the row would sit num_spec tokens in + # the past. + seqlens = [200, 130] + offsets = [192, 64] + x, cu, conv_state = _conv_fixture(seqlens, seed=3, num_spec=num_spec) + weight = torch.randn(DIM, WIDTH, device="cuda", dtype=torch.bfloat16) * 0.3 + bias = torch.randn(DIM, device="cuda", dtype=torch.bfloat16) + rows = [5, 9] + + store_conv_checkpoints(x, conv_state, cu, _i32(offsets), _i32(rows), STATE_LEN) + + # A prefill of just the prefix, into its own cache. Row 0 is the null block + # the conv kernel skips by design, so the reference rows start at 1. + prefix = torch.cat( + [x[int(cu[n]) : int(cu[n]) + off] for n, off in enumerate(offsets)] + ) + prefix_cu = torch.tensor( + [0] + torch.tensor(offsets).cumsum(0).tolist(), + device="cuda", + dtype=torch.int32, + ) + reference = torch.zeros_like(conv_state) + causal_conv1d_fn( + prefix.transpose(0, 1), + weight, + bias, + activation="silu", + conv_states=reference, + has_initial_state=torch.zeros(2, dtype=torch.bool, device="cuda"), + cache_indices=_i32([1, 2]), + query_start_loc=prefix_cu, + ) + + # causal_conv1d_fn writes only the leading STATE_LEN columns. + window = slice(None), slice(None, STATE_LEN) + torch.testing.assert_close(conv_state[rows[0]][window], reference[1][window]) + torch.testing.assert_close(conv_state[rows[1]][window], reference[2][window]) + + +def test_conv_checkpoint_leaves_spec_tail_of_row_alone() -> None: + # The draft-token columns belong to causal_conv1d_update. Filling the whole + # row would hand every reader, which indexes from column 0, a window + # shifted num_spec tokens into the past. + num_spec = 4 + x, cu, conv_state = _conv_fixture([200], num_spec=num_spec, seed=11) + + store_conv_checkpoints(x, conv_state, cu, _i32([192]), _i32([5]), STATE_LEN) + + torch.testing.assert_close( + conv_state[5, :, :STATE_LEN], x[192 - STATE_LEN : 192].transpose(0, 1) + ) + assert bool(conv_state[5, :, STATE_LEN:].isnan().all()), ( + "the checkpoint wrote into the speculative-decode tail of the row" + ) + + +def test_conv_checkpoint_honours_transposed_cache_view() -> None: + # The SD conv layout reaches the layer transposed, not contiguous. + x, cu, _ = _conv_fixture([200]) + raw = torch.full( + (16, STATE_LEN, DIM), float("nan"), device="cuda", dtype=torch.bfloat16 + ) + conv_state = raw.transpose(-1, -2) + assert not conv_state.is_contiguous() + + store_conv_checkpoints(x, conv_state, cu, _i32([192]), _i32([7]), STATE_LEN) + + torch.testing.assert_close(conv_state[7], x[192 - STATE_LEN : 192].transpose(0, 1)) + untouched = torch.ones(16, dtype=torch.bool, device="cuda") + untouched[7] = False + assert bool(raw[untouched].isnan().all()) + + +def test_checkpoint_metadata_indexes_prefill_rows_after_decodes() -> None: + # The chunk kernel receives the prefill tail alone, so the metadata has to + # be indexed the same way or it snapshots the wrong request's boundary. + md = _build(BatchSpec(seq_lens=[64, 200], query_lens=[1, 200])) + + assert md.num_decodes == 1 and md.num_prefills == 1 + assert md.checkpoint is not None + torch.testing.assert_close(md.checkpoint.checkpoint_offsets, _i32([192])) + # Request 1 owns blocks 4..7, so its slot-2 state block is 6. + torch.testing.assert_close(md.checkpoint.state_indices, _i32([6])) + + +def test_conv_checkpoint_addresses_row_past_int32_limit() -> None: + # state_stride_0 is the whole mamba page, padded up to the attention page, + # so row * stride passes 2**31 a few thousand blocks into a large cache and + # an int32 wrap stores outside the tensor. The margin on the plain prefill + # path is 1.6x, which is luck: a bigger cache, a higher utilization or the + # wider row speculative decoding allocates all close it. + stride_row = DIM * (STATE_LEN + 8) + row = 2**31 // stride_row + 1 + need = (row + 1) * stride_row * 2 + free, _ = torch.accelerator.get_memory_info() + if free < need + (4 << 30): + pytest.skip(f"needs {need >> 30} GiB of free VRAM to reach the wrap") + + x, cu, _ = _conv_fixture([200]) + conv_state = torch.zeros( + row + 1, DIM, STATE_LEN + 8, device="cuda", dtype=torch.bfloat16 + ) + assert conv_state.stride(0) * row > 2**31, "fixture does not reach the wrap" + + store_conv_checkpoints(x, conv_state, cu, _i32([192]), _i32([row]), STATE_LEN) + + torch.testing.assert_close( + conv_state[row, :, :STATE_LEN], x[192 - STATE_LEN : 192].transpose(0, 1) + ) + assert not bool(conv_state[:row].any()), "the store wrote outside its row" + + +def test_conv_checkpoint_reads_row_strided_projection() -> None: + # mixed_qkv is a band of the fused QKVGFAB projection, so its row stride is + # the whole projection width, not DIM. + torch.manual_seed(5) + total, pad = 200, 177 + projection = torch.randn(total, DIM + pad, device="cuda", dtype=torch.bfloat16) + x = projection[:, :DIM] + assert x.stride(0) == DIM + pad + + cu = _i32([0, total]) + conv_state = torch.full( + (16, DIM, STATE_LEN), float("nan"), device="cuda", dtype=torch.bfloat16 + ) + store_conv_checkpoints(x, conv_state, cu, _i32([192]), _i32([4]), STATE_LEN) + + torch.testing.assert_close(conv_state[4], x[192 - STATE_LEN : 192].transpose(0, 1)) + + +def test_finer_prefix_match_unit_moves_checkpoint_to_hash_grid() -> None: + # A partial prefix hit matches at prefix_match_unit, so a state captured + # only at block boundaries is too coarse to resume from. Block 256 with + # unit 128 moves a 700-token prompt's checkpoint from 512 to 640. + md = _build( + BatchSpec(seq_lens=[700], query_lens=[700]), + block_size=256, + prefix_match_unit=128, + ) + + assert md.checkpoint is not None + torch.testing.assert_close(md.checkpoint.checkpoint_offsets, _i32([640])) + # Slot is still block-indexed: cdiv(700, 256) - 2 = 1, request 0 owns 0..2. + torch.testing.assert_close(md.checkpoint.state_indices, _i32([1])) + + +def test_checkpoint_metadata_skips_unaligned_offset() -> None: + # The scheduler declines a boundary off the published alignment and + # reserves nothing, so the builder has to reach the same verdict or it + # writes a block nobody hashed. + # Block 192, unit 96: a 700-token prompt's last hash boundary is 672, which + # is 32 tokens off the chunk grid. + md = _build( + BatchSpec(seq_lens=[700], query_lens=[700]), + block_size=192, + prefix_match_unit=96, + ) + + assert md.checkpoint is None + + +# --------------------------------------------------------------------------- +# Opting in +# --------------------------------------------------------------------------- + + +def _spec_for(monkeypatch, **attrs) -> MambaSpec: + """Run the layer's `get_kv_cache_spec` over a stubbed base spec.""" + base = MambaSpec(block_size=BLOCK_SIZE, shapes=((16, 64),), dtypes=(torch.float16,)) + monkeypatch.setattr( + GatedDeltaNetAttention, "get_kv_cache_spec", lambda self, cfg: base + ) + layer = KimiK3DeltaAttention.__new__(KimiK3DeltaAttention) + for name, value in dict( + use_prefill_checkpoint=True, + use_fused_chunk=True, + num_spec=0, + use_safe_gate=True, + ).items(): + setattr(layer, name, attrs.get(name, value)) + vllm_config = create_vllm_config(model_name="Qwen/Qwen3.5-0.8B") + vllm_config.cache_config.mamba_cache_mode = "align" + return layer.get_kv_cache_spec(vllm_config) + + +def test_opt_in_publishes_chunk_walk_alignment(monkeypatch) -> None: + # The scheduler learns the kernel's granularity only from the spec, so + # reserving a block without publishing the alignment would hash a state + # that is never written. + spec = _spec_for(monkeypatch) + + assert spec.num_prefill_checkpoint_blocks == 1 + assert spec.prefill_checkpoint_alignment == KDA_CHECKPOINT_ALIGNMENT + + +@pytest.mark.parametrize( + "attr,value", + [ + ("use_prefill_checkpoint", False), + ("use_fused_chunk", False), + ("use_safe_gate", False), + ], +) +def test_paths_that_cannot_export_stay_opted_out(monkeypatch, attr, value) -> None: + spec = _spec_for(monkeypatch, **{attr: value}) + + assert spec.num_prefill_checkpoint_blocks == 0 + assert spec.prefill_checkpoint_alignment is None + + +def test_speculative_decoding_does_not_opt_out(monkeypatch) -> None: + # num_spec reaches the kernels as a wider conv row and an interleaved + # batch, both of which the layer handles. + spec = _spec_for(monkeypatch, num_spec=8) + + assert spec.num_prefill_checkpoint_blocks == 1 + assert spec.prefill_checkpoint_alignment == KDA_CHECKPOINT_ALIGNMENT diff --git a/tests/models/kimi_k3/test_amd_kda_chunk.py b/tests/models/kimi_k3/test_amd_kda_chunk.py index 3e99018edc..c53e9bacdb 100644 --- a/tests/models/kimi_k3/test_amd_kda_chunk.py +++ b/tests/models/kimi_k3/test_amd_kda_chunk.py @@ -13,17 +13,17 @@ import torch from vllm.platforms import current_platform -def _on_gfx950() -> bool: - if not current_platform.is_rocm(): +def _fused_chunk_available() -> bool: + if not (current_platform.is_rocm() and torch.cuda.is_available()): return False - from vllm.platforms.rocm import on_gfx950 + from vllm.models.kimi_k3.amd.ops.kda_chunk import is_fused_kda_chunk_supported - return on_gfx950() + return is_fused_kda_chunk_supported() pytestmark = pytest.mark.skipif( - not _on_gfx950(), - reason="The fused KDA chunk kernel is only built for gfx950", + not _fused_chunk_available(), + reason="The fused KDA chunk kernel needs a gfx950 build that includes it", ) HEAD_DIM = 128 @@ -701,6 +701,56 @@ def test_cold_rows_start_from_zero_not_from_cache_junk(use_fused: bool) -> None: torch.testing.assert_close(cache[idx.long()], ht_ref, rtol=5e-3, atol=5e-3) +@pytest.mark.parametrize("use_fused", [False, True]) +def test_state_indices_accept_a_block_table_column(use_fused: bool) -> None: + # The layer passes block_table[:, 0], which is contiguous only while the + # table has one column. Speculative decoding reserves a state page per + # draft token, so the table widens and that column becomes strided. + _requires_kernel() + seqlens = [512, 320, 64] + inp = _inputs(seqlens, seed=43) + n = len(seqlens) + rows = [3 * i + 5 for i in range(n)] + cache, idx, warm = _paged_cache(inp, [True] * n, rows, slots=3 * n + 12) + ref_cache = cache.clone() + + table = torch.zeros(n, 9, dtype=torch.int32, device="cuda") + table[:, 0] = idx + column = table[:, 0] + assert not column.is_contiguous() + + o_ref, _ = _run_paged(inp, ref_cache, idx, warm, use_fused=use_fused) + o_got, _ = _run_paged(inp, cache, column, warm, use_fused=use_fused) + + torch.testing.assert_close(o_got, o_ref) + torch.testing.assert_close(cache, ref_cache) + + +@pytest.mark.parametrize("use_fused", [False, True]) +def test_padded_zero_length_sequence_is_inert(use_fused: bool) -> None: + # A batch carrying speculative decodes pads its non-spec group with + # zero-length sequences whose cache row still belongs to a request, so + # consuming no tokens has to mean writing nothing and perturbing nobody. + _requires_kernel() + live = [200, 1, 128] + # _inputs draws h0 last, so the padded batch's live rows are the same draw + # as the unpadded one's and the two runs compare elementwise. + padded, unpadded = _inputs([*live, 0], seed=41), _inputs(live, seed=41) + rows = [3 * i + 5 for i in range(4)] + + cache, idx, warm = _paged_cache(padded, [True] * 4, rows, slots=24) + before = cache.clone() + _run_paged(padded, cache, idx, warm, use_fused=use_fused) + + ref_cache, ref_idx, ref_warm = _paged_cache(unpadded, [True] * 3, rows[:3], 24) + _run_paged(unpadded, ref_cache, ref_idx, ref_warm, use_fused=use_fused) + + assert torch.equal(cache[rows[3]], before[rows[3]]), ( + "the padded sequence's cache row was written" + ) + torch.testing.assert_close(cache[ref_idx.long()], ref_cache[ref_idx.long()]) + + @pytest.mark.parametrize("groups", [2, 4]) def test_paged_state_survives_the_group_split(groups: int, monkeypatch) -> None: """At G>1 the read and the write are different blocks of one launch.""" diff --git a/vllm/models/kimi_k3/amd/kda.py b/vllm/models/kimi_k3/amd/kda.py index d98d16521f..e2dea74b6b 100644 --- a/vllm/models/kimi_k3/amd/kda.py +++ b/vllm/models/kimi_k3/amd/kda.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from dataclasses import replace + import torch from einops import rearrange from torch import nn @@ -38,7 +40,11 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader from vllm.model_executor.utils import set_weight_attrs from vllm.models.kimi_k3.amd.kda_metadata import KimiK3ROCmKDABackend +from vllm.models.kimi_k3.amd.ops.kda_checkpoint import ( + KimiK3ROCmKDAPrefillCheckpointExporter, +) from vllm.models.kimi_k3.amd.ops.kda_chunk import ( + KDA_CHECKPOINT_ALIGNMENT, is_fused_kda_chunk_supported, ) 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 from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig from vllm.v1.attention.backend import AttentionBackend from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata +from vllm.v1.kv_cache_interface import MambaSpec logger = init_logger(__name__) @@ -79,6 +86,23 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): num_spec=self.num_spec, ) + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> MambaSpec: + spec = super().get_kv_cache_spec(vllm_config) + assert isinstance(spec, MambaSpec) + # Only the fused path exports the checkpoint snapshots + can_checkpoint = ( + self.use_prefill_checkpoint and self.use_fused_chunk and self.use_safe_gate + ) + if can_checkpoint and vllm_config.cache_config.mamba_cache_mode == "align": + logger.info_once("Kimi-K3 KDA prefill checkpoint enabled.") + return replace( + spec, + num_prefill_checkpoint_blocks=int(can_checkpoint), + prefill_checkpoint_alignment=( + KDA_CHECKPOINT_ALIGNMENT if can_checkpoint else None + ), + ) + def __init__( self, config: KimiLinearConfig, @@ -226,6 +250,15 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): "Kimi-K3 KDA prefill backend: %s", "fused" if self.use_fused_chunk else "triton", ) + # TODO: Remove the extra config arg once tested + self.use_prefill_checkpoint = ( + additional_config.get("kda_prefill_checkpoint", True) + if isinstance(additional_config, dict) + else True + ) + self._checkpoint_exporter = KimiK3ROCmKDAPrefillCheckpointExporter( + state_len=self.conv_size - 1 + ) self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid") decode_norm_weight = None @@ -511,10 +544,9 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): assert non_spec_state_indices_tensor is not None assert has_initial_state is not None - # A mixed non-spec batch is decode-first: the decodes are - # length-1 sequences, and the chunk kernel returns NaN for those. - # Send them to the recurrent kernel and give the chunk kernel the - # prefill tail only. + # Decode-first, so peel the length-1 decodes to the recurrent + # kernel. A spec batch interleaves its non-spec rows, so the + # peel cannot apply and the chunk kernel takes them. core_attn_out_decode = None split_non_spec = spec_sequence_masks is None and m.num_decodes > 0 # Without a spec split the non-spec tokens are exactly @@ -522,7 +554,6 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): # can write their slice straight into it. With one, the tokens # are interleaved and have to be scattered by index afterwards. direct = spec_sequence_masks is None - use_fused_chunk = self.use_fused_chunk and spec_sequence_masks is None if split_non_spec: assert non_spec_query_start_loc is not None nd_tok = m.num_decode_tokens @@ -559,6 +590,22 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): prefill_has_initial_state = has_initial_state nd_tok = 0 + checkpoint = getattr(m, "checkpoint", None) + checkpoint_state = None + checkpoint_offsets = None + checkpoint_state_indices = None + if checkpoint is not None: + assert prefill_query_start_loc is not None + checkpoint_state = recurrent_state + checkpoint_offsets = checkpoint.checkpoint_offsets + checkpoint_state_indices = checkpoint.state_indices + self._checkpoint_exporter.export( + checkpoint, + raw_qkv=mixed_qkv_ns[nd_tok:], + conv_state=conv_state, + cu_seqlens=prefill_query_start_loc, + ) + ( core_attn_out_non_spec, _, @@ -578,8 +625,11 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): cu_seqlens=prefill_query_start_loc, chunk_indices=m.chunk_indices, chunk_offsets=m.chunk_offsets, - use_fused_chunk=use_fused_chunk, + use_fused_chunk=self.use_fused_chunk, out=core_attn_out[:, nd_tok:num_actual_tokens] if direct else None, + checkpoint_state=checkpoint_state, + checkpoint_offsets=checkpoint_offsets, + checkpoint_state_indices=checkpoint_state_indices, ) # chunk_kda_prefill updates `recurrent_state` in place, so # there is no final state to write back here. diff --git a/vllm/models/kimi_k3/amd/kda_metadata.py b/vllm/models/kimi_k3/amd/kda_metadata.py index 9681aca1d8..dcc03d2814 100644 --- a/vllm/models/kimi_k3/amd/kda_metadata.py +++ b/vllm/models/kimi_k3/amd/kda_metadata.py @@ -7,16 +7,23 @@ The request classification and cudagraph staging intentionally mirror differently on device rather than on the host. """ +from dataclasses import replace + import torch from vllm.logger import init_logger +from vllm.model_executor.layers.mamba.checkpoint import ( + MambaPrefillCheckpointMetadata, +) from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE from vllm.triton_utils import tl, triton from vllm.utils.math_utils import next_power_of_2 from vllm.v1.attention.backends.gdn_attn import ( GDNAttentionBackend, + GDNAttentionMetadata, GDNAttentionMetadataBuilder, ) +from vllm.v1.attention.backends.utils import CommonAttentionMetadata logger = init_logger(__name__) @@ -86,6 +93,55 @@ def prepare_chunk_metadata_device( class KimiK3ROCmKDAMetadataBuilder(GDNAttentionMetadataBuilder): + def build( # type: ignore[override] + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + num_accepted_tokens: torch.Tensor | None = None, + num_decode_draft_tokens_cpu: torch.Tensor | None = None, + fast_build: bool = False, + ) -> GDNAttentionMetadata: + attn_metadata = super().build( + common_prefix_len, + common_attn_metadata, + num_accepted_tokens, + num_decode_draft_tokens_cpu, + fast_build, + ) + if attn_metadata.checkpoint is None: + return attn_metadata + # The shared builder keys checkpoint rows on every request; the ROCm + # prefill path indexes them by prefill_query_start_loc, so rebuild with + # those rows. + return replace( + attn_metadata, + checkpoint=self._build_checkpoint_metadata( + common_attn_metadata, attn_metadata, num_decode_draft_tokens_cpu + ), + ) + + def _build_checkpoint_metadata( + self, + m: CommonAttentionMetadata, + attn_metadata: GDNAttentionMetadata, + num_decode_draft_tokens_cpu: torch.Tensor | None, + ) -> MambaPrefillCheckpointMetadata | None: + if attn_metadata.num_prefills == 0: + return None + # request_rows must line up with prefill_query_start_loc, either the + # prefill tail of a decode-first batch, or every non-spec row + # including padding + num_decodes = attn_metadata.num_decodes + request_rows = list( + range(num_decodes, num_decodes + attn_metadata.num_prefills) + ) + if attn_metadata.spec_sequence_masks is not None: + assert num_decode_draft_tokens_cpu is not None + request_rows = ( + (num_decode_draft_tokens_cpu < 0).nonzero().flatten().tolist() + ) + return self.checkpoint_builder.build(m, request_rows) + def _build_chunk_metadata( self, prefill_query_start_loc: torch.Tensor, diff --git a/vllm/models/kimi_k3/amd/ops/kda_checkpoint.py b/vllm/models/kimi_k3/amd/ops/kda_checkpoint.py new file mode 100644 index 0000000000..5560adc626 --- /dev/null +++ b/vllm/models/kimi_k3/amd/ops/kda_checkpoint.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Conv half of the ROCm Kimi-K3 prefill checkpoint.""" + +from dataclasses import dataclass + +import torch + +from vllm.model_executor.layers.mamba.checkpoint import ( + MambaPrefillCheckpointMetadata, +) +from vllm.model_executor.layers.mamba.kda_checkpoint import ( + FlashKDAPrefillCheckpointExporter, +) + + +@dataclass(frozen=True) +class KimiK3ROCmKDAPrefillCheckpointExporter(FlashKDAPrefillCheckpointExporter): + """Conv window only; `fused_kda_chunk` writes the recurrent half itself. + + `state_len` is `conv_kernel_size - 1`, not the spec-widened conv row. + """ + + def export( # type: ignore[override] + self, + checkpoint: MambaPrefillCheckpointMetadata, + *, + raw_qkv: torch.Tensor, + conv_state: torch.Tensor, + cu_seqlens: torch.Tensor, + ) -> None: + assert self.state_len is not None + # An empty recurrent checkpoint masks the kernel's recurrent half off. + super().export( + checkpoint, + raw_qkv=raw_qkv, + conv_state=conv_state, + recurrent_checkpoint=conv_state[:, :0], + recurrent_state=conv_state, + cu_seqlens=cu_seqlens, + ) diff --git a/vllm/models/kimi_k3/amd/ops/kda_chunk.py b/vllm/models/kimi_k3/amd/ops/kda_chunk.py index 556373c084..035a49f852 100644 --- a/vllm/models/kimi_k3/amd/ops/kda_chunk.py +++ b/vllm/models/kimi_k3/amd/ops/kda_chunk.py @@ -16,6 +16,7 @@ from vllm.third_party.flash_linear_attention.ops.index import ( prepare_chunk_indices, prepare_chunk_offsets, ) +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID CHUNK_SIZE = 64 HEAD_DIM = 128 @@ -28,6 +29,15 @@ _DEEP_LEN = 12 _MAX_GROUPS = 32 +def _hip_checkpoint_state_indices(state_indices: torch.Tensor) -> torch.Tensor: + """NULL_BLOCK_ID marks no checkpoint; the HIP kernel skips negative rows.""" + return ( + torch.where(state_indices != NULL_BLOCK_ID, state_indices, -1) + .to(torch.int32) + .contiguous() + ) + + @cache def is_fused_kda_chunk_supported() -> bool: from vllm.platforms.rocm import on_gfx950 @@ -295,7 +305,7 @@ def fused_kda_chunk( None if checkpoint_offsets is None else checkpoint_offsets.to(torch.int32), None if checkpoint_state_indices is None - else checkpoint_state_indices.to(torch.int32).contiguous(), + else _hip_checkpoint_state_indices(checkpoint_state_indices), state_cache, None if state_indices is None else state_indices.to(torch.int32).contiguous(), has_initial_state, diff --git a/vllm/models/kimi_k3/amd/ops/kda_prefill.py b/vllm/models/kimi_k3/amd/ops/kda_prefill.py index 56cb3e1482..5ba2fd6379 100644 --- a/vllm/models/kimi_k3/amd/ops/kda_prefill.py +++ b/vllm/models/kimi_k3/amd/ops/kda_prefill.py @@ -155,7 +155,13 @@ def chunk_kda_prefill( aqk=ws["aqk"], decay=ws["decay"], # v is dead by this point, so it doubles as the output buffer. - out=out if out is not None else v, + out=out + if out is not None + else ( + v + if v.is_contiguous() + else torch.empty_like(v, memory_format=torch.contiguous_format) + ), scale=scale, cu_seqlens=cu_seqlens, initial_state=initial_state,