[Kimi K3][Pref] Reuse internal checkpoint blocks for partial prefix caching

tests/kernels/mamba/test_causal_conv1d.py

@@ -392,3 +392,70 @@ def test_causal_conv1d_varlen(
     )
     unpadded_out = out[:, : out_ref_tensor.shape[-1]]
     assert torch.allclose(unpadded_out, out_ref_tensor, rtol=rtol, atol=atol)
+
+
+def test_causal_conv1d_exports_checkpoint_state_in_fwd_kernel():
+    device = DEVICE
+    set_random_seed(0)
+
+    dim = 64
+    width = 4
+    seqlens = [24, 24, 32]
+    query_start_loc = torch.tensor(
+        [0, *torch.tensor(seqlens).cumsum(0).tolist()],
+        dtype=torch.int32,
+        device=device,
+    )
+    total_tokens = sum(seqlens)
+    x = torch.randn(total_tokens, dim, device=device).transpose(0, 1)
+    weight = torch.randn(dim, width, device=device)
+    conv_states = torch.randn(8, dim, width - 1, device=device)
+    conv_states_before = conv_states.clone()
+    state_indices = torch.tensor([1, 2, 3], dtype=torch.int32, device=device)
+    has_initial_state = torch.tensor([True, False, True], device=device)
+
+    # Export rows 0 and 2 at chunk boundaries; the null block disables row 1.
+    checkpoint_offsets = torch.tensor([16, 16, 24], dtype=torch.int32, device=device)
+    checkpoint_state_indices = torch.tensor(
+        [5, NULL_BLOCK_ID, 6], dtype=torch.int64, device=device
+    )
+
+    out = causal_conv1d_fn(
+        x,
+        weight,
+        bias=None,
+        conv_states=conv_states,
+        query_start_loc=query_start_loc,
+        cache_indices=state_indices,
+        has_initial_state=has_initial_state,
+        checkpoint_offsets=checkpoint_offsets,
+        checkpoint_state_indices=checkpoint_state_indices,
+        validate_data=True,
+    )
+
+    expected_out = []
+    start = 0
+    for row, seqlen in enumerate(seqlens):
+        x_seq = x[:, start : start + seqlen]
+        initial_state = (
+            conv_states_before[state_indices[row]].unsqueeze(0)
+            if has_initial_state[row]
+            else None
+        )
+        ref_out, _ = causal_conv1d_ref(
+            x_seq.unsqueeze(0),
+            weight,
+            initial_states=initial_state,
+        )
+        expected_out.append(ref_out.squeeze(0))
+        start += seqlen
+    torch.testing.assert_close(out, torch.cat(expected_out, dim=1))
+
+    for row, destination in ((0, 5), (2, 6)):
+        start = int(query_start_loc[row].item())
+        offset = int(checkpoint_offsets[row].item())
+        expected_checkpoint = x[:, start + offset - (width - 1) : start + offset]
+        torch.testing.assert_close(conv_states[destination], expected_checkpoint)
+
+    # A disabled row must not write to any extra checkpoint cache line.
+    torch.testing.assert_close(conv_states[7], conv_states_before[7])

tests/models/kimi_k3/test_kda_metadata.py

@@ -88,6 +88,8 @@ def _make_builder(
     mamba_cache_mode: str = "none",
     use_recoverssm: bool = False,
     num_prefill_checkpoint_blocks: int = 0,
+    mamba_block_size: int = BLOCK_SIZE,
+    prefix_match_unit: int | None = None,
 ) -> AttentionMetadataBuilder:
     vllm_config = create_vllm_config(
         model_name="Qwen/Qwen3.5-0.8B",
@@ -104,9 +106,10 @@ def _make_builder(
     vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode
     vllm_config.cache_config.use_replayssm = use_recoverssm
     vllm_config.cache_config.use_kda_recoverssm = use_recoverssm
+    vllm_config.cache_config.prefix_match_unit = prefix_match_unit
     builder = builder_cls(
         kv_cache_spec=MambaSpec(
-            block_size=BLOCK_SIZE,
+            block_size=mamba_block_size,
             shapes=((16, 64),),
             dtypes=(torch.float16,),
             mamba_cache_mode=mamba_cache_mode,
@@ -150,6 +153,39 @@ def test_internal_checkpoint_metadata_targets_last_aligned_boundary():
     )
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_internal_checkpoint_metadata_targets_partial_hash_boundary():
+    device = torch.device("cuda")
+    batch = BatchSpec(seq_lens=[100], query_lens=[100])
+    common_attn_metadata = create_common_attn_metadata(
+        batch, BLOCK_SIZE, device, arange_block_indices=True
+    )
+    common_attn_metadata = common_attn_metadata.replace(
+        is_prefilling=torch.tensor([True]),
+        block_table_tensor=common_attn_metadata.block_table_tensor + 1,
+    )
+    actual = _make_builder(
+        KimiK3KDAMetadataBuilder,
+        num_speculative_tokens=0,
+        full_cuda_graph=False,
+        mamba_cache_mode="align",
+        num_prefill_checkpoint_blocks=1,
+        mamba_block_size=64,
+        prefix_match_unit=16,
+        device=device,
+    ).build(0, common_attn_metadata)
+
+    assert actual.checkpoint is not None
+    torch.testing.assert_close(
+        actual.checkpoint.state_indices,
+        torch.tensor([1], dtype=torch.int32, device=device),
+    )
+    torch.testing.assert_close(
+        actual.checkpoint.checkpoint_offsets,
+        torch.tensor([96], dtype=torch.int32, device=device),
+    )
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
 def test_internal_checkpoint_metadata_skips_unaligned_offset():
     device = torch.device("cuda")

tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py

@@ -78,6 +78,7 @@ def make_full_mamba_manager(
     num_blocks: int = 32,
     use_eagle: bool = False,
     num_prefill_checkpoint_blocks: int = 0,
+    prefill_checkpoint_alignment: int = 1,
 ):
     kv_cache_config = KVCacheConfig(
         num_blocks=num_blocks,
@@ -100,6 +101,7 @@ def make_full_mamba_manager(
                     dtypes=(torch.float32,),
                     mamba_cache_mode="align",
                     num_prefill_checkpoint_blocks=num_prefill_checkpoint_blocks,
+                    prefill_checkpoint_alignment=prefill_checkpoint_alignment,
                 ),
             ),
         ],
@@ -517,6 +519,45 @@ def test_partial_hit_then_internal_checkpoint_uses_distinct_mamba_blocks():
     assert mamba_blocks[3].block_id == running_block_id
 
 
+def test_internal_checkpoint_uses_partial_hash_lifecycle():
+    hash_block_size = 2
+    mamba_block_size = 4
+    manager = make_full_mamba_manager(
+        dcp_world_size=1,
+        hash_block_size=hash_block_size,
+        full_block_size=hash_block_size,
+        mamba_block_size=mamba_block_size,
+        num_prefill_checkpoint_blocks=1,
+    )
+    request = make_request("producer", list(range(15)), hash_block_size, sha256)
+
+    new_blocks = manager.allocate_slots(request, request.num_tokens)
+
+    assert new_blocks is not None
+    mamba_blocks = manager.get_blocks(request.request_id).blocks[1]
+    checkpoint_idx = 2
+    checkpoint_block = mamba_blocks[checkpoint_idx]
+    running_block = mamba_blocks[checkpoint_idx + 1]
+    assert not checkpoint_block.is_null
+    assert not running_block.is_null
+    assert checkpoint_block is not running_block
+
+    # The internal block is exported at state@14, replacing its temporary
+    # full-block state@12 key while retaining #52789's req_to_blocks ownership.
+    partial_hash = request.block_hashes[14 // hash_block_size - 1]
+    partial_hit = manager.block_pool.get_cached_block(partial_hash, [1])
+    assert partial_hit is not None
+    assert partial_hit[0] is checkpoint_block
+    assert checkpoint_block.block_hash_num_tokens == 14
+    full_hash = request.block_hashes[12 // hash_block_size - 1]
+    assert manager.block_pool.get_cached_block(full_hash, [1]) is None
+
+    manager.free(request)
+    replay = make_request("replay", list(range(15)), hash_block_size, sha256)
+    _, num_computed, _ = manager.get_computed_blocks(replay)
+    assert num_computed == 14
+
+
 def test_internal_checkpoint_requires_block_aligned_start():
     hash_block_size = 2
     mamba_block_size = 16
@@ -546,6 +587,27 @@ def test_internal_checkpoint_requires_block_aligned_start():
     assert manager.block_pool.get_cached_block(checkpoint_hash, [1]) is None
 
 
+def test_internal_checkpoint_requires_backend_aligned_offset():
+    hash_block_size = 2
+    mamba_block_size = 4
+    manager = make_full_mamba_manager(
+        dcp_world_size=1,
+        hash_block_size=hash_block_size,
+        full_block_size=hash_block_size,
+        mamba_block_size=mamba_block_size,
+        num_prefill_checkpoint_blocks=1,
+        prefill_checkpoint_alignment=4,
+    )
+    request = make_request("producer", list(range(15)), hash_block_size, sha256)
+
+    assert manager.allocate_slots(request, request.num_tokens) is not None
+
+    mamba_blocks = manager.get_blocks(request.request_id).blocks[1]
+    assert mamba_blocks[2].is_null
+    checkpoint_hash = request.block_hashes[14 // hash_block_size - 1]
+    assert manager.block_pool.get_cached_block(checkpoint_hash, [1]) is None
+
+
 def test_external_mamba_hit_same_block_uses_running_cow_on_continue():
     """An external mid-block hit must become a running request even when its
     first continuation does not need another Mamba block."""

vllm/model_executor/layers/mamba/ops/causal_conv1d.py

@@ -31,6 +31,8 @@ def _causal_conv1d_fwd_kernel(  # continuous batching
     block_idx_last_scheduled_token,  # (batch,)
     initial_state_idx,  # (batch,)
     num_computed_tokens,  # (batch,)
+    checkpoint_offsets_ptr,  # (batch,)
+    checkpoint_state_indices_ptr,  # (batch,)
     o_ptr,  # (dim, seqlen) - actually pointing to x_ptr
     # Matrix dimensions
     dim: tl.constexpr,
@@ -56,6 +58,7 @@ def _causal_conv1d_fwd_kernel(  # continuous batching
     SILU_ACTIVATION: tl.constexpr,
     IS_APC_ENABLED: tl.constexpr,
     HAS_NULL_BLOCK: tl.constexpr,
+    HAS_CHECKPOINT: tl.constexpr,
     NP2_STATELEN: tl.constexpr,
     BLOCK_M: tl.constexpr,
     BLOCK_N: tl.constexpr,
@@ -392,6 +395,44 @@ def _causal_conv1d_fwd_kernel(  # continuous batching
             tl.debug_barrier()  #  NOTE: use this due to bug in Triton compiler
             tl.store(conv_states_ptrs_target, loaded_x, mask)
 
+    # At a chunk-aligned checkpoint boundary, col0..colN already contain the
+    # exact causal-convolution history immediately before the checkpoint token.
+    # Export it here to avoid reconstructing the state in a separate kernel.
+    if HAS_CHECKPOINT:
+        checkpoint_offset = tl.load(checkpoint_offsets_ptr + idx_seq)
+        checkpoint_state_coord = tl.load(checkpoint_state_indices_ptr + idx_seq).to(
+            tl.int64
+        )
+        store_checkpoint = (checkpoint_state_coord != null_block_id) & (
+            token_offset == checkpoint_offset
+        )
+        checkpoint_state_base = (
+            conv_states_ptr
+            + checkpoint_state_coord * stride_conv_state_seq
+            + idx_feats * stride_conv_state_dim
+        )
+        checkpoint_mask = store_checkpoint & (idx_feats < dim)
+        if KERNEL_WIDTH >= 2:
+            tl.store(checkpoint_state_base, col0, mask=checkpoint_mask)
+        if KERNEL_WIDTH >= 3:
+            tl.store(
+                checkpoint_state_base + stride_conv_state_tok,
+                col1,
+                mask=checkpoint_mask,
+            )
+        if KERNEL_WIDTH >= 4:
+            tl.store(
+                checkpoint_state_base + 2 * stride_conv_state_tok,
+                col2,
+                mask=checkpoint_mask,
+            )
+        if KERNEL_WIDTH >= 5:
+            tl.store(
+                checkpoint_state_base + 3 * stride_conv_state_tok,
+                col3,
+                mask=checkpoint_mask,
+            )
+
     if HAS_BIAS:
         bias = bias_ptr + idx_feats
         mask_bias = idx_feats < dim
@@ -496,6 +537,8 @@ def causal_conv1d_fn(
     block_size_to_align=0,
     metadata=None,
     validate_data=False,
+    checkpoint_offsets: torch.Tensor | None = None,
+    checkpoint_state_indices: torch.Tensor | None = None,
 ):
     """support varlen + continuous batching when x is 2D tensor
 
@@ -546,6 +589,12 @@ def causal_conv1d_fn(
         The number of tokens already completed for each sequence
     block_size_to_align: int
         The block size to align the cached states to
+    checkpoint_offsets: (batch,) int32
+        Optional per-sequence checkpoint offsets. Each valid offset must be
+        aligned to the kernel's BLOCK_M chunk size.
+    checkpoint_state_indices: (batch,) int64
+        Cache line receiving each sequence's checkpoint state. ``null_block_id``
+        disables checkpoint export for that sequence.
     out: same shape as `x`
     """
     if isinstance(activation, bool) and activation:
@@ -580,6 +629,12 @@ def causal_conv1d_fn(
     np2_statelen = triton.next_power_of_2(state_len)
 
     padded_batch = query_start_loc.size(0) - 1
+    has_checkpoint = checkpoint_offsets is not None
+    assert has_checkpoint == (checkpoint_state_indices is not None), (
+        "checkpoint offsets and state indices must be provided together"
+    )
+    if has_checkpoint:
+        assert conv_states is not None
     stride_x_dim = x.stride(0)
     stride_x_token = x.stride(1)
     stride_w_dim = weight.stride(0)
@@ -628,6 +683,15 @@ def causal_conv1d_fn(
             assert conv_states is not None, (
                 "ERROR: `has_initial_state` is used, which needs also `conv_states`"
             )
+        if has_checkpoint:
+            assert checkpoint_offsets is not None
+            assert checkpoint_state_indices is not None
+            assert checkpoint_offsets.size() == (padded_batch,)
+            assert checkpoint_state_indices.size() == (padded_batch,)
+            valid_checkpoint_offsets = checkpoint_offsets[
+                checkpoint_state_indices != null_block_id
+            ]
+            assert torch.all(valid_checkpoint_offsets % BLOCK_M == 0)
         assert weight.stride(1) == 1
         assert (dim, width) == weight.shape
         assert is_channel_last, "Need to run in channel-last layout"
@@ -724,6 +788,8 @@ def grid(META):
         block_idx_last_scheduled_token,
         initial_state_idx,
         num_computed_tokens,
+        checkpoint_offsets,
+        checkpoint_state_indices,
         out,
         # Matrix dimensions
         dim,
@@ -749,6 +815,7 @@ def grid(META):
         SILU_ACTIVATION=activation in ["silu", "swish"],
         IS_APC_ENABLED=block_idx_last_scheduled_token is not None,
         HAS_NULL_BLOCK=null_block_id is not None,
+        HAS_CHECKPOINT=has_checkpoint,
         NP2_STATELEN=np2_statelen,
         # launch_cooperative_grid=True
         BLOCK_M=BLOCK_M,