diff --git a/tests/models/kimi_k3/test_kda.py b/tests/models/kimi_k3/test_kda.py index 84a5e7eed4..f03d2cdcae 100644 --- a/tests/models/kimi_k3/test_kda.py +++ b/tests/models/kimi_k3/test_kda.py @@ -14,6 +14,10 @@ import torch.nn.functional as F from vllm import _custom_ops as ops from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture +from vllm.model_executor.layers.mamba.checkpoint import MambaPrefillCheckpointMetadata +from vllm.model_executor.layers.mamba.kda_checkpoint import ( + FlashKDAPrefillCheckpointExporter, +) from vllm.model_executor.layers.mamba.ops.causal_conv1d import causal_conv1d_update from vllm.model_executor.layers.mamba.ops.gather_initial_states import ( gather_initial_states, @@ -25,7 +29,6 @@ from vllm.models.kimi_k3.nvidia import kda as nvidia_kda from vllm.models.kimi_k3.nvidia.kda import ( _flashinfer_kda_prefill, _flashkda_prefill, - _store_cache_checkpoints_kernel, is_flashinfer_fused_kda_decode_supported, is_flashinfer_recurrent_kda_prefill_supported, is_flashkda_supported, @@ -1498,36 +1501,13 @@ def test_flashkda_checkpoint_correctness(state_dtype: torch.dtype, tolerance: fl checkpoint_state_indices = torch.tensor( [1, NULL_BLOCK_ID], dtype=torch.int32, device=DEVICE ) - state_len = conv_state.shape[-1] - width = H * D - recurrent_row_size = checkpoint_state[0].numel() - block_size = 256 - _store_cache_checkpoints_kernel[ - ( - checkpoint_state_indices.numel(), - (max(width * state_len, recurrent_row_size) + block_size - 1) // block_size, - ) - ]( - conv_input, - conv_state, - checkpoint_state, - recurrent_state, - cu_seqlens, - checkpoint_offsets, - checkpoint_state_indices, - conv_input.stride(0), - conv_input.stride(1), - conv_state.stride(0), - conv_state.stride(1), - conv_state.stride(2), - checkpoint_state.stride(0), - recurrent_state.stride(0), - checkpoint_offsets.stride(0), - state_len, - width, - recurrent_row_size, - NULL_BLOCK_ID, - block_size, + FlashKDAPrefillCheckpointExporter().export( + MambaPrefillCheckpointMetadata(checkpoint_offsets, checkpoint_state_indices), + raw_qkv=conv_input, + conv_state=conv_state, + recurrent_checkpoint=checkpoint_state, + recurrent_state=recurrent_state, + cu_seqlens=cu_seqlens, ) torch.testing.assert_close(conv_state[1], q[0, 13:16].flatten(1).transpose(0, 1)) torch.testing.assert_close(recurrent_state[1], checkpoint_state[0]) diff --git a/vllm/model_executor/layers/mamba/checkpoint.py b/vllm/model_executor/layers/mamba/checkpoint.py new file mode 100644 index 0000000000..3ba4062079 --- /dev/null +++ b/vllm/model_executor/layers/mamba/checkpoint.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from abc import ABC, abstractmethod +from dataclasses import dataclass + +import torch + +from vllm.config import VllmConfig +from vllm.utils.math_utils import cdiv +from vllm.utils.torch_utils import async_tensor_h2d +from vllm.v1.attention.backend import CommonAttentionMetadata +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID +from vllm.v1.kv_cache_interface import ( + MambaSpec, + get_mamba_prefill_checkpoint_position, + is_mamba_prefill_checkpoint_valid, +) + + +def compute_mamba_prefill_checkpoints( + seq_lens: list[int], + query_lens: list[int], + hash_block_size: int, + mamba_block_size: int, + checkpoint_alignment: int | None, + drop_eagle_block: bool, +) -> tuple[list[int], list[int]]: + """Per-row internal prefill checkpoint offsets and cache block columns. + + Backends call this instead of re-deriving the rules, so they decline in + lockstep with the scheduler and ``MambaManager``: allocating a checkpoint + block without writing it leaves the prefix cache serving uninitialized + state. + + Returns: + ``(offsets, cols)``: the checkpoint's offset into each row's query and + its block-table column. ``0`` and ``-1`` mean the row has none. + + """ + offsets: list[int] = [] + cols: list[int] = [] + for seq_len, query_len in zip(seq_lens, query_lens): + query_start = seq_len - query_len + position = get_mamba_prefill_checkpoint_position( + seq_len, hash_block_size, drop_eagle_block=drop_eagle_block + ) + valid = is_mamba_prefill_checkpoint_valid( + query_start=query_start, + query_end=seq_len, + checkpoint_position=position, + hash_block_size=hash_block_size, + mamba_block_size=mamba_block_size, + checkpoint_alignment=checkpoint_alignment, + ) + offsets.append(position - query_start if valid else 0) + cols.append(cdiv(seq_len, mamba_block_size) - 2 if valid else -1) + return offsets, cols + + +@dataclass +class MambaPrefillCheckpointMetadata: + checkpoint_offsets: torch.Tensor + state_indices: torch.Tensor + + +class MambaPrefillCheckpointBuilder: + """Build backend-neutral prefill checkpoint locations.""" + + def __init__(self, vllm_config: VllmConfig, kv_cache_spec: MambaSpec) -> None: + self.vllm_config = vllm_config + self.kv_cache_spec = kv_cache_spec + + def build( + self, + m: CommonAttentionMetadata, + request_rows: list[int], + ) -> MambaPrefillCheckpointMetadata | None: + if self.kv_cache_spec.num_prefill_checkpoint_blocks == 0: + return None + assert m.seq_lens_cpu_upper_bound is not None + all_query_lens = m.query_start_loc_cpu.diff().tolist() + query_lens = [all_query_lens[row] for row in request_rows] + seq_lens = m.seq_lens_cpu_upper_bound.tolist() + block_size = self.kv_cache_spec.block_size + hash_block_size = self.vllm_config.cache_config.prefix_match_unit or block_size + speculative_config = self.vllm_config.speculative_config + drop_eagle_block = ( + speculative_config is not None and speculative_config.use_eagle_block_drop() + ) + checkpoint_offsets, checkpoint_cols = compute_mamba_prefill_checkpoints( + [seq_lens[row] for row in request_rows], + query_lens, + hash_block_size=hash_block_size, + mamba_block_size=block_size, + checkpoint_alignment=self.kv_cache_spec.prefill_checkpoint_alignment, + drop_eagle_block=drop_eagle_block, + ) + if not any(checkpoint_offsets): + return None + checkpoint_offsets_tensor = async_tensor_h2d( + checkpoint_offsets, + dtype=torch.int32, + device=m.query_start_loc.device, + ) + request_rows_tensor = async_tensor_h2d( + request_rows, dtype=torch.int64, device=m.query_start_loc.device + ) + checkpoint_cols_tensor = async_tensor_h2d( + checkpoint_cols, dtype=torch.int64, device=m.query_start_loc.device + ) + checkpoint_state_indices = m.block_table_tensor[ + request_rows_tensor, checkpoint_cols_tensor + ] + checkpoint_state_indices = torch.where( + checkpoint_cols_tensor >= 0, + checkpoint_state_indices, + NULL_BLOCK_ID, + ) + return MambaPrefillCheckpointMetadata( + checkpoint_offsets_tensor, + checkpoint_state_indices, + ) + + +class MambaPrefillCheckpointExporter(ABC): + """Export a mid-prefill checkpoint into backend-specific paged states.""" + + @abstractmethod + def export( + self, + checkpoint: MambaPrefillCheckpointMetadata, + *args, + **kwargs, + ) -> None: + """Write checkpoint state into the paged cache.""" + raise NotImplementedError diff --git a/vllm/model_executor/layers/mamba/kda_checkpoint.py b/vllm/model_executor/layers/mamba/kda_checkpoint.py new file mode 100644 index 0000000000..20741623d1 --- /dev/null +++ b/vllm/model_executor/layers/mamba/kda_checkpoint.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import dataclass + +import torch + +from vllm.model_executor.layers.mamba.checkpoint import ( + MambaPrefillCheckpointExporter, + MambaPrefillCheckpointMetadata, +) +from vllm.triton_utils import tl, triton +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID + + +def kda_prefill_checkpoint_alignment(backend: str) -> int | None: + return 16 if backend == "flashkda" else None + + +@dataclass(frozen=True) +class FlashKDAPrefillCheckpointExporter(MambaPrefillCheckpointExporter): + """Store FlashKDA recurrent and convolution checkpoint states.""" + + state_len: int | None = None + + def export( + self, + checkpoint: MambaPrefillCheckpointMetadata, + *, + raw_qkv: torch.Tensor, + conv_state: torch.Tensor, + recurrent_checkpoint: torch.Tensor, + recurrent_state: torch.Tensor, + cu_seqlens: torch.Tensor, + ) -> None: + state_len = ( + self.state_len if self.state_len is not None else conv_state.shape[-1] + ) + width = raw_qkv.shape[-1] + recurrent_row_size = recurrent_checkpoint[0].numel() + block_size = 256 + store_cache_checkpoints_kernel[ + ( + checkpoint.checkpoint_offsets.numel(), + triton.cdiv(max(width * state_len, recurrent_row_size), block_size), + ) + ]( + raw_qkv, + conv_state, + recurrent_checkpoint, + recurrent_state, + cu_seqlens, + checkpoint.checkpoint_offsets, + checkpoint.state_indices, + raw_qkv.stride(0), + raw_qkv.stride(1), + conv_state.stride(0), + conv_state.stride(1), + conv_state.stride(2), + recurrent_checkpoint.stride(0), + recurrent_state.stride(0), + checkpoint.checkpoint_offsets.stride(0), + state_len, + width, + recurrent_row_size, + NULL_BLOCK_ID, + block_size, + ) + + +@triton.jit +def store_cache_checkpoints_kernel( + x_ptr, + conv_state_ptr, + recurrent_checkpoint_ptr, + recurrent_state_ptr, + query_start_loc_ptr, + checkpoint_offsets_ptr, + checkpoint_state_indices_ptr, + x_stride_0: tl.constexpr, + x_stride_1: tl.constexpr, + state_stride_0: tl.constexpr, + state_stride_1: tl.constexpr, + state_stride_2: tl.constexpr, + checkpoint_stride_0: tl.constexpr, + recurrent_state_stride_0: tl.constexpr, + checkpoint_offset_stride: tl.constexpr, + STATE_LEN: tl.constexpr, + WIDTH: tl.constexpr, + RECURRENT_ROW_SIZE: tl.constexpr, + NULL_STATE_IDX: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + seq_idx = tl.program_id(0) + cols = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + state_idx = tl.load(checkpoint_state_indices_ptr + seq_idx).to(tl.int64) + checkpoint_offset = tl.load( + checkpoint_offsets_ptr + seq_idx * checkpoint_offset_stride + ) + valid_checkpoint = (state_idx != NULL_STATE_IDX) & (checkpoint_offset > 0) + valid_conv = ( + (cols < WIDTH * STATE_LEN) & valid_checkpoint & (checkpoint_offset >= STATE_LEN) + ) + width_idx = cols // STATE_LEN + history_idx = cols % STATE_LEN + checkpoint_end = tl.load(query_start_loc_ptr + seq_idx) + checkpoint_offset + token_idx = checkpoint_end - STATE_LEN + history_idx + values = tl.load( + x_ptr + token_idx * x_stride_0 + width_idx * x_stride_1, + mask=valid_conv, + ) + tl.store( + conv_state_ptr + + state_idx * state_stride_0 + + width_idx * state_stride_1 + + history_idx * state_stride_2, + values, + mask=valid_conv, + ) + + valid_recurrent = (cols < RECURRENT_ROW_SIZE) & valid_checkpoint + recurrent = tl.load( + recurrent_checkpoint_ptr + seq_idx * checkpoint_stride_0 + cols, + mask=valid_recurrent, + ) + tl.store( + recurrent_state_ptr + state_idx * recurrent_state_stride_0 + cols, + recurrent, + mask=valid_recurrent, + ) diff --git a/vllm/models/kimi_k3/nvidia/kda.py b/vllm/models/kimi_k3/nvidia/kda.py index 92326944f5..c920cdbbdf 100644 --- a/vllm/models/kimi_k3/nvidia/kda.py +++ b/vllm/models/kimi_k3/nvidia/kda.py @@ -21,6 +21,10 @@ from vllm.model_executor.layers.linear import ( RowParallelLinear, ) from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention +from vllm.model_executor.layers.mamba.kda_checkpoint import ( + FlashKDAPrefillCheckpointExporter, + kda_prefill_checkpoint_alignment, +) from vllm.model_executor.layers.mamba.mamba_utils import ( MambaStateDtypeCalculator, MambaStateShapeCalculator, @@ -49,7 +53,6 @@ from vllm.models.kimi_k3.nvidia.kda_metadata import ( from vllm.platforms import current_platform from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig -from vllm.triton_utils import tl, triton from vllm.utils.flashinfer import ( flashinfer_fused_kda_decode, flashinfer_recurrent_kda, @@ -57,7 +60,6 @@ from vllm.utils.flashinfer import ( has_flashinfer_recurrent_kda, ) from vllm.v1.attention.backend import AttentionBackend -from vllm.v1.attention.backends.utils import NULL_BLOCK_ID from vllm.v1.kv_cache_interface import MambaSpec from vllm.v1.worker.workspace import current_workspace_manager @@ -389,69 +391,6 @@ def _flashinfer_kda_prefill( return output, initial_state -@triton.jit -def _store_cache_checkpoints_kernel( - x_ptr, - conv_state_ptr, - recurrent_checkpoint_ptr, - recurrent_state_ptr, - query_start_loc_ptr, - checkpoint_offsets_ptr, - checkpoint_state_indices_ptr, - x_stride_0: tl.constexpr, - x_stride_1: tl.constexpr, - state_stride_0: tl.constexpr, - state_stride_1: tl.constexpr, - state_stride_2: tl.constexpr, - checkpoint_stride_0: tl.constexpr, - recurrent_state_stride_0: tl.constexpr, - checkpoint_offset_stride: tl.constexpr, - STATE_LEN: tl.constexpr, - WIDTH: tl.constexpr, - RECURRENT_ROW_SIZE: tl.constexpr, - NULL_STATE_IDX: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - # store checkpoints to cache - seq_idx = tl.program_id(0) - cols = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - state_idx = tl.load(checkpoint_state_indices_ptr + seq_idx).to(tl.int64) - checkpoint_offset = tl.load( - checkpoint_offsets_ptr + seq_idx * checkpoint_offset_stride - ) - valid_checkpoint = (state_idx != NULL_STATE_IDX) & (checkpoint_offset > 0) - valid_conv = ( - (cols < WIDTH * STATE_LEN) & valid_checkpoint & (checkpoint_offset >= STATE_LEN) - ) - width_idx = cols // STATE_LEN - history_idx = cols % STATE_LEN - checkpoint_end = tl.load(query_start_loc_ptr + seq_idx) + checkpoint_offset - token_idx = checkpoint_end - STATE_LEN + history_idx - values = tl.load( - x_ptr + token_idx * x_stride_0 + width_idx * x_stride_1, - mask=valid_conv, - ) - tl.store( - conv_state_ptr - + state_idx * state_stride_0 - + width_idx * state_stride_1 - + history_idx * state_stride_2, - values, - mask=valid_conv, - ) - - valid_recurrent = (cols < RECURRENT_ROW_SIZE) & valid_checkpoint - recurrent = tl.load( - recurrent_checkpoint_ptr + seq_idx * checkpoint_stride_0 + cols, - mask=valid_recurrent, - ) - tl.store( - recurrent_state_ptr + state_idx * recurrent_state_stride_0 + cols, - recurrent, - mask=valid_recurrent, - ) - - def resolve_kda_prefill_backend( backend: str, head_dim: int, @@ -735,6 +674,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): self._flashinfer_kda_output_spec: tuple[tuple[int, ...], torch.dtype] | None = ( None ) + self._checkpoint_exporter: FlashKDAPrefillCheckpointExporter | None = None if self.kda_prefill_backend == "flashkda": T = vllm_config.scheduler_config.max_num_batched_tokens N = vllm_config.scheduler_config.max_num_seqs @@ -748,6 +688,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): ((N, H, D, D), self.get_state_dtype()[1]), ((workspace_size,), torch.uint8), ) + self._checkpoint_exporter = FlashKDAPrefillCheckpointExporter() elif self.kda_prefill_backend == "flashinfer": T = vllm_config.scheduler_config.max_num_batched_tokens H, D = self.local_num_heads, self.head_dim @@ -800,12 +741,11 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): def get_kv_cache_spec(self, vllm_config: VllmConfig) -> MambaSpec: spec = super().get_kv_cache_spec(vllm_config) assert isinstance(spec, MambaSpec) + alignment = kda_prefill_checkpoint_alignment(self.kda_prefill_backend) return replace( spec, - num_prefill_checkpoint_blocks=int(self.kda_prefill_backend == "flashkda"), - prefill_checkpoint_alignment=( - 16 if self.kda_prefill_backend == "flashkda" else None - ), + num_prefill_checkpoint_blocks=int(alignment is not None), + prefill_checkpoint_alignment=alignment, ) def forward( @@ -1173,39 +1113,14 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): ) core_attn_out_non_spec = flashkda_out last_recurrent_state = final_state - state_len = conv_state.shape[-1] - width = mixed_qkv_ns.shape[-1] - recurrent_row_size = checkpoint_state[0].numel() - block_size = 256 - _store_cache_checkpoints_kernel[ - ( - checkpoint_offsets.numel(), - triton.cdiv( - max(width * state_len, recurrent_row_size), - block_size, - ), - ) - ]( - mixed_qkv_ns, - conv_state, - checkpoint_state, - recurrent_state, - non_spec_query_start_loc, - checkpoint_offsets, - checkpoint.state_indices, - mixed_qkv_ns.stride(0), - mixed_qkv_ns.stride(1), - conv_state.stride(0), - conv_state.stride(1), - conv_state.stride(2), - checkpoint_state.stride(0), - recurrent_state.stride(0), - checkpoint_offsets.stride(0), - state_len, - width, - recurrent_row_size, - NULL_BLOCK_ID, - block_size, + assert self._checkpoint_exporter is not None + self._checkpoint_exporter.export( + checkpoint, + raw_qkv=mixed_qkv_ns, + conv_state=conv_state, + recurrent_checkpoint=checkpoint_state, + recurrent_state=recurrent_state, + cu_seqlens=non_spec_query_start_loc, ) else: ( diff --git a/vllm/models/kimi_k3/nvidia/kda_metadata.py b/vllm/models/kimi_k3/nvidia/kda_metadata.py index d342adb6a6..d1c2b31921 100644 --- a/vllm/models/kimi_k3/nvidia/kda_metadata.py +++ b/vllm/models/kimi_k3/nvidia/kda_metadata.py @@ -18,9 +18,12 @@ from typing import TYPE_CHECKING import torch from vllm.config import VllmConfig +from vllm.model_executor.layers.mamba.checkpoint import ( + MambaPrefillCheckpointBuilder, + MambaPrefillCheckpointMetadata, +) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton -from vllm.utils.math_utils import cdiv from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backend import CommonAttentionMetadata from vllm.v1.attention.backends.gdn_attn import ( @@ -37,11 +40,7 @@ from vllm.v1.attention.backends.utils import ( compute_causal_conv1d_metadata, split_decodes_and_prefills, ) -from vllm.v1.kv_cache_interface import ( - MambaSpec, - get_mamba_prefill_checkpoint_position, - is_mamba_prefill_checkpoint_valid, -) +from vllm.v1.kv_cache_interface import MambaSpec if TYPE_CHECKING: from vllm.models.kimi_k3.nvidia.ops.recoverssm import ( @@ -262,12 +261,6 @@ class KDARecoverSSMCommitMetadata: align: KDARecoverSSMAlignMetadata | None -@dataclass -class KDACheckpointMetadata: - checkpoint_offsets: torch.Tensor - state_indices: torch.Tensor - - @dataclass class KimiK3KDAMetadata(GDNAttentionMetadata, RecoverSSMMetadata): spec_token_start: int | None = None @@ -278,7 +271,7 @@ class KimiK3KDAMetadata(GDNAttentionMetadata, RecoverSSMMetadata): recoverssm_context: "KDARecoverSSMCommitContext | None" = field( default=None, repr=False, compare=False ) - checkpoint: KDACheckpointMetadata | None = None + checkpoint: MambaPrefillCheckpointMetadata | None = None def commit_recoverssm_state( self, num_accepted_tokens: torch.Tensor @@ -322,6 +315,9 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder): device: torch.device, ) -> None: super().__init__(kv_cache_spec, layer_names, vllm_config, device) + self.checkpoint_builder = MambaPrefillCheckpointBuilder( + vllm_config, kv_cache_spec + ) additional_config = vllm_config.additional_config self.use_flashinfer_prefill = ( isinstance(additional_config, dict) @@ -400,8 +396,10 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder): if not self.use_spec_decode or num_decode_draft_tokens_cpu is None: spec_sequence_masks_cpu = None num_spec_decodes = 0 + active_non_spec_mask_cpu = None else: spec_sequence_masks_cpu = num_decode_draft_tokens_cpu >= 0 + active_non_spec_mask_cpu = None if self.use_recoverssm: assert m.is_prefilling is not None assert m.is_prefilling.device.type == "cpu" @@ -418,6 +416,7 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder): ): spec_sequence_masks_cpu = None num_spec_decodes = 0 + active_non_spec_mask_cpu = None else: num_spec_decodes = spec_sequence_masks_cpu.sum().item() if num_spec_decodes == 0: @@ -624,80 +623,6 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder): else: has_initial_state = None - checkpoint = None - checkpoint_enabled = ( - self.vllm_config.cache_config.mamba_cache_mode == "align" - and self.kv_cache_spec.num_prefill_checkpoint_blocks > 0 - ) - if num_prefills > 0 and checkpoint_enabled: - # prepare checkpoint metadata - assert m.seq_lens_cpu_upper_bound is not None - request_rows = list(range(m.num_reqs)) - if spec_sequence_masks_cpu is not None: - # get non spec request rows - request_rows = active_non_spec_mask_cpu.nonzero().flatten().tolist() - all_query_lens = query_start_loc_cpu.diff().tolist() - query_lens = [all_query_lens[row] for row in request_rows] - seq_lens = m.seq_lens_cpu_upper_bound.tolist() - block_size = self.kv_cache_spec.block_size - hash_block_size = ( - self.vllm_config.cache_config.prefix_match_unit or block_size - ) - speculative_config = self.vllm_config.speculative_config - drop_eagle_block = ( - speculative_config is not None - and speculative_config.use_eagle_block_drop() - ) - checkpoint_splits = [] - checkpoint_cols = [] - for row, query_len in zip(request_rows, query_lens): - seq_len = seq_lens[row] - query_start = seq_len - query_len - checkpoint_position = get_mamba_prefill_checkpoint_position( - seq_len, - hash_block_size, - drop_eagle_block=drop_eagle_block, - ) - offset = checkpoint_position - query_start - valid = is_mamba_prefill_checkpoint_valid( - query_start=query_start, - query_end=seq_len, - checkpoint_position=checkpoint_position, - hash_block_size=hash_block_size, - mamba_block_size=block_size, - checkpoint_alignment=( - self.kv_cache_spec.prefill_checkpoint_alignment - ), - ) - offset = offset if valid else 0 - first_len = offset or query_len - checkpoint_splits.append((first_len, query_len - first_len)) - checkpoint_cols.append(cdiv(seq_len, block_size) - 2 if valid else -1) - if any(tail for _, tail in checkpoint_splits): - checkpoint_offsets_tensor = async_tensor_h2d( - [first if tail else 0 for first, tail in checkpoint_splits], - dtype=torch.int32, - device=query_start_loc.device, - ) - request_rows_tensor = async_tensor_h2d( - request_rows, dtype=torch.int64, device=query_start_loc.device - ) - checkpoint_cols_tensor = async_tensor_h2d( - checkpoint_cols, dtype=torch.int64, device=query_start_loc.device - ) - checkpoint_state_indices = m.block_table_tensor[ - request_rows_tensor, checkpoint_cols_tensor - ] - checkpoint_state_indices = torch.where( - checkpoint_cols_tensor >= 0, - checkpoint_state_indices, - NULL_BLOCK_ID, - ) - checkpoint = KDACheckpointMetadata( - checkpoint_offsets_tensor, - checkpoint_state_indices, - ) - # Prepare per-request tensors for cudagraph replay. num_actual_tokens # may be token-padded, while state/query/acceptance metadata is indexed # by request. @@ -776,6 +701,13 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder): flashinfer_prefill_query_start_loc.diff(), descending=True ).to(torch.int32) + checkpoint = None + if num_prefills > 0: + request_rows = list(range(m.num_reqs)) + if active_non_spec_mask_cpu is not None: + request_rows = active_non_spec_mask_cpu.nonzero().flatten().tolist() + checkpoint = self.checkpoint_builder.build(m, request_rows) + return KimiK3KDAMetadata( num_prefills=num_prefills, num_prefill_tokens=num_prefill_tokens,