[Perf][Kimi K3] Support internal prefill checkpoints in Triton KDA

tests/models/kimi_k3/test_kda.py

@@ -352,6 +352,81 @@ def test_chunk_kda_fused_gate_cumsum_matches_unfused(
     assert_close("ht", old_ht, new_ht, 1e-3, err_atol=1e-3)
 
 
+@pytest.mark.parametrize("head_dim", [64, 128, 256])
+@pytest.mark.parametrize("lower_bound", [-5.0, None])
+@torch.inference_mode()
+def test_triton_prefill_checkpoint_resume(head_dim, lower_bound):
+    """Export FP32 chunk boundaries and resume without changing outputs."""
+    torch.manual_seed(123)
+    lengths = [193, 129, 65, 97, 1]
+    offsets = [128, 64, 65, 0, -1]
+    H, D = 2, head_dim
+    T = sum(lengths)
+    q, k, v, raw_g = [
+        torch.randn(1, T, H, D, dtype=torch.bfloat16, device=DEVICE) for _ in range(4)
+    ]
+    raw_beta = torch.randn(1, T, H, device=DEVICE)
+    A_log = torch.randn(H, device=DEVICE) * 0.1
+    bias = torch.randn(H * D, device=DEVICE) * 0.1
+    initial = torch.randn(len(lengths), H, D, D, device=DEVICE)
+    cu = torch.tensor([0, *lengths], dtype=torch.int32, device=DEVICE).cumsum(
+        0, dtype=torch.int32
+    )
+    checkpoints = torch.full_like(initial, float("nan"))
+
+    def run(start, end, state, cu_seqlens, **kwargs):
+        return chunk_kda_with_fused_gate(
+            q=q[:, start:end],
+            k=k[:, start:end],
+            v=v[:, start:end].clone(),
+            raw_g=raw_g[:, start:end],
+            raw_beta=raw_beta[:, start:end],
+            A_log=A_log,
+            g_bias=bias,
+            lower_bound=lower_bound,
+            initial_state=state,
+            output_final_state=True,
+            use_qk_l2norm_in_kernel=True,
+            cu_seqlens=cu_seqlens,
+            **kwargs,
+        )
+
+    expected_out, expected_final = run(0, T, initial, cu)
+    out, final = run(
+        0,
+        T,
+        initial,
+        cu,
+        checkpoint_state=checkpoints,
+        checkpoint_offsets=torch.tensor(offsets, dtype=torch.int32, device=DEVICE),
+    )
+    torch.testing.assert_close(out, expected_out, rtol=0, atol=0)
+    torch.testing.assert_close(final, expected_final, rtol=0, atol=0)
+    start = 0
+    for seq, (length, offset) in enumerate(zip(lengths, offsets)):
+        if offset > 0:
+            _, prefix_final = run(start, start + offset, initial[seq : seq + 1], None)
+            torch.testing.assert_close(
+                checkpoints[seq : seq + 1], prefix_final, rtol=1e-5, atol=1e-5
+            )
+            if offset < length:
+                resumed, resumed_final = run(
+                    start + offset, start + length, checkpoints[seq : seq + 1], None
+                )
+                torch.testing.assert_close(
+                    resumed,
+                    out[:, start + offset : start + length],
+                    rtol=1e-2,
+                    atol=1e-2,
+                )
+                torch.testing.assert_close(
+                    resumed_final, final[seq : seq + 1], rtol=1e-4, atol=1e-4
+                )
+        else:
+            assert checkpoints[seq].isnan().all()
+        start += length
+
+
 @pytest.mark.parametrize("num_seqs", [1, 8, 32])
 @pytest.mark.parametrize("lower_bound", [-5.0, None])
 @pytest.mark.parametrize("state_indices_stride", [1, 8])

tests/models/kimi_k3/test_kda_metadata.py

@@ -98,6 +98,7 @@ def _make_builder(
     prefix_match_unit: int | None = None,
     use_eagle: bool = False,
     disable_eagle_block_drop: bool = False,
+    checkpoint_alignment: int = 16,
 ) -> AttentionMetadataBuilder:
     vllm_config = create_vllm_config(
         model_name="Qwen/Qwen3.5-0.8B",
@@ -129,7 +130,7 @@ def _make_builder(
             num_speculative_blocks=(0 if use_recoverssm else num_speculative_tokens),
             num_prefill_checkpoint_blocks=num_prefill_checkpoint_blocks,
             prefill_checkpoint_alignment=(
-                16 if num_prefill_checkpoint_blocks > 0 else None
+                checkpoint_alignment if num_prefill_checkpoint_blocks > 0 else None
             ),
         ),
         layer_names=["layer.0"],
@@ -320,13 +321,26 @@ def test_internal_checkpoint_metadata_targets_last_aligned_boundary():
 
 
 @pytest.mark.parametrize(
-    ("disable_eagle_block_drop", "prefix_match_unit", "expected_offset"),
-    [(False, 16, 80), (True, 16, 96), (False, 8, None)],
+    (
+        "disable_eagle_block_drop",
+        "prefix_match_unit",
+        "checkpoint_alignment",
+        "expected_offset",
+    ),
+    [
+        (False, 16, 16, 80),
+        (True, 16, 16, 96),
+        (False, 8, 16, None),
+        (False, 16, 64, None),
+        (True, 16, 64, None),
+        (True, 64, 64, 64),
+    ],
 )
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
 def test_spec_internal_checkpoint_metadata_targets_replay_boundary(
     disable_eagle_block_drop: bool,
     prefix_match_unit: int,
+    checkpoint_alignment: int,
     expected_offset: int | None,
 ) -> None:
     device = torch.device("cuda")
@@ -348,6 +362,7 @@ def test_spec_internal_checkpoint_metadata_targets_replay_boundary(
         prefix_match_unit=prefix_match_unit,
         use_eagle=True,
         disable_eagle_block_drop=disable_eagle_block_drop,
+        checkpoint_alignment=checkpoint_alignment,
         device=device,
     )
     assert isinstance(builder, KimiK3KDAMetadataBuilder)

tests/v1/core/test_mamba_align_chunk_split.py

@@ -132,6 +132,7 @@ def _split(
     partial_hit: bool = False,
     num_prefill_checkpoint_blocks: int = 0,
     max_num_scheduled_tokens: int = 16384,
+    checkpoint_alignment: int = 16,
 ) -> int:
     """Call the real `Scheduler._mamba_block_aligned_split` on a stub self."""
     if use_eagle_block_drop is None:
@@ -148,36 +149,40 @@ def _split(
         hash_block_size=ATTN_BLOCK_SIZE,
         mamba_has_prefill_checkpoint_blocks=(num_prefill_checkpoint_blocks > 0),
         mamba_prefill_checkpoint_alignment=(
-            16 if num_prefill_checkpoint_blocks > 0 else None
+            checkpoint_alignment if num_prefill_checkpoint_blocks > 0 else None
         ),
     )
     return Scheduler._mamba_block_aligned_split(stub, request, num_new_tokens)
 
 
 @pytest.mark.parametrize(
-    ("prompt_len", "num_new_tokens", "use_eagle", "expected"),
+    ("prompt_len", "num_new_tokens", "use_eagle", "expected", "expected_triton"),
     [
-        (2002, 2002, False, 2002),
-        (3602, 2000, False, MAMBA_BLOCK_SIZE),
-        (3602, 3602, True, 3602),
-        (2002, 2002, True, 2002),
-        (3602, 512, True, 0),
-        (3602, 1700, True, MAMBA_BLOCK_SIZE),
+        (2002, 2002, False, 2002, MAMBA_BLOCK_SIZE),
+        (3602, 2000, False, MAMBA_BLOCK_SIZE, MAMBA_BLOCK_SIZE),
+        (3602, 3602, True, 3602, 3602),
+        (2002, 2002, True, 2002, 2002),
+        (3602, 512, True, 0, 0),
+        (3602, 1700, True, MAMBA_BLOCK_SIZE, MAMBA_BLOCK_SIZE),
     ],
 )
+@pytest.mark.parametrize("checkpoint_alignment", [16, 64])
 def test_internal_checkpoint_split(
-    prompt_len: int, num_new_tokens: int, use_eagle: bool, expected: int
+    prompt_len: int,
+    num_new_tokens: int,
+    use_eagle: bool,
+    expected: int,
+    expected_triton: int,
+    checkpoint_alignment: int,
 ) -> None:
     (request,) = create_requests(1, num_tokens=prompt_len, block_size=ATTN_BLOCK_SIZE)
-    assert (
-        _split(
-            request,
-            num_new_tokens,
-            use_eagle=use_eagle,
-            num_prefill_checkpoint_blocks=1,
-        )
-        == expected
-    )
+    assert _split(
+        request,
+        num_new_tokens,
+        use_eagle=use_eagle,
+        num_prefill_checkpoint_blocks=1,
+        checkpoint_alignment=checkpoint_alignment,
+    ) == (expected_triton if checkpoint_alignment == 64 else expected)
 
 
 def test_partial_checkpoint_resume_stops_at_mamba_block_boundary() -> None:

vllm/models/kimi_k3/nvidia/kda.py

@@ -48,6 +48,7 @@
 )
 from vllm.platforms import current_platform
 from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated
+from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE
 from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
 from vllm.triton_utils import tl, triton
 from vllm.utils.flashinfer import (
@@ -802,9 +803,15 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> MambaSpec:
         assert isinstance(spec, MambaSpec)
         return replace(
             spec,
-            num_prefill_checkpoint_blocks=int(self.kda_prefill_backend == "flashkda"),
+            num_prefill_checkpoint_blocks=int(
+                self.kda_prefill_backend in ("flashkda", "triton")
+            ),
             prefill_checkpoint_alignment=(
-                16 if self.kda_prefill_backend == "flashkda" else None
+                FLA_CHUNK_SIZE
+                if self.kda_prefill_backend == "triton"
+                else 16
+                if self.kda_prefill_backend == "flashkda"
+                else None
             ),
         )
 
@@ -1134,6 +1141,7 @@ def _prefill_conv(
                     non_spec_state_indices_tensor,
                     has_initial_state,
                 )
+                checkpoint_state = None
                 if self.kda_prefill_backend == "flashkda":
                     assert self.gate_lower_bound is not None
                     assert self._flashkda_buffer_specs is not None
@@ -1173,40 +1181,6 @@ def _prefill_conv(
                         )
                         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,
-                        )
                     else:
                         (
                             core_attn_out_non_spec,
@@ -1258,6 +1232,11 @@ def _prefill_conv(
                         seq_order=m.flashinfer_prefill_seq_order,
                     )
                 else:
+                    checkpoint_state = (
+                        torch.empty_like(initial_state, dtype=torch.float32)
+                        if checkpoint is not None
+                        else None
+                    )
                     (
                         core_attn_out_non_spec,
                         last_recurrent_state,
@@ -1275,6 +1254,49 @@ def _prefill_conv(
                         use_qk_l2norm_in_kernel=True,
                         cu_seqlens=non_spec_query_start_loc,
                         out=non_spec_out,
+                        checkpoint_state=checkpoint_state,
+                        checkpoint_offsets=(
+                            checkpoint.checkpoint_offsets
+                            if checkpoint is not None
+                            else None
+                        ),
+                    )
+                if checkpoint is not None:
+                    assert checkpoint_state is not None
+                    checkpoint_offsets = checkpoint.checkpoint_offsets
+                    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,
                     )
                 recurrent_state[non_spec_state_indices_tensor] = (
                     last_recurrent_state.to(recurrent_state.dtype)