[ROCm][Perf] Kimi-K3 enable prefill checkpoints on ROCm (#58344)
Signed-off-by: kliuae <kuanfu.liu@embeddedllm.com>
This commit is contained in:
kliuae
@@ -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
|
||||
@@ -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."""
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user