iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Core][KDA] Generalize Mamba prefill checkpoint builder and exporter (#57783)

Browse Source 
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Co-authored-by: zjy0516 <riverclouds.zhu@qq.com>
This commit is contained in:
Chauncey
2026-09-21 04:46:11 +00:00
committed by GitHub
co-authored by zjy0516
parent ff0c885494
commit 3df4ae153e
5 changed files with 314 additions and 220 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/models/kimi_k3/test_kda.py
+11 -31
View File
@@ -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])
vllm/model_executor/layers/mamba/checkpoint.py
+137
View File
@@ -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
vllm/model_executor/layers/mamba/kda_checkpoint.py
+130
View File
@@ -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,
)
vllm/models/kimi_k3/nvidia/kda.py
+17 -102
View File
@@ -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:
(
vllm/models/kimi_k3/nvidia/kda_metadata.py
+19 -87
View File
@@ -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,