[ROCm] Skip fresh Kimi-K3 KDA state copies

tests/models/kimi_k3/test_kda.py

@@ -15,6 +15,9 @@
 from vllm.model_executor.layers.mamba.ops.gather_initial_states import (
     gather_initial_states,
 )
+from vllm.models.kimi_k3.amd.ops.third_party.kda import (
+    chunk_kda_with_fused_gate as amd_chunk_kda_with_fused_gate,
+)
 from vllm.models.kimi_k3.nvidia.kda import (
     is_flashkda_supported,
     is_fused_kda_decode_supported,
@@ -270,6 +273,84 @@ def test_chunk_kda_fused_gate_cumsum_matches_unfused(
     assert_close("ht", old_ht, new_ht, 1e-3, err_atol=1e-3)
 
 
+@torch.inference_mode()
+def test_chunk_kda_none_matches_zero_initial_state():
+    H, D = 2, 128
+    cu_seqlens = torch.tensor([0, 17, 49], dtype=torch.int32, device=DEVICE)
+    T = 49
+    N = cu_seqlens.numel() - 1
+    torch.manual_seed(123)
+
+    kwargs = {
+        "q": torch.randn(1, T, H, D, dtype=torch.bfloat16, device=DEVICE),
+        "k": torch.randn(1, T, H, D, dtype=torch.bfloat16, device=DEVICE),
+        "v": torch.randn(1, T, H, D, dtype=torch.bfloat16, device=DEVICE),
+        "raw_g": torch.randn(1, T, H, D, dtype=torch.bfloat16, device=DEVICE),
+        "raw_beta": torch.randn(1, T, H, dtype=torch.bfloat16, device=DEVICE),
+        "A_log": torch.randn(H, dtype=torch.float32, device=DEVICE),
+        "g_bias": torch.randn(H * D, dtype=torch.float32, device=DEVICE),
+        "output_final_state": True,
+        "cu_seqlens": cu_seqlens,
+        "use_qk_l2norm_in_kernel": True,
+    }
+    zero_state = torch.zeros(
+        N,
+        H,
+        D,
+        D,
+        dtype=torch.float32,
+        device=DEVICE,
+    )
+
+    def run(initial_state: torch.Tensor | None, **extra_kwargs):
+        return amd_chunk_kda_with_fused_gate(
+            **{
+                key: value.clone() if isinstance(value, torch.Tensor) else value
+                for key, value in kwargs.items()
+            },
+            initial_state=initial_state,
+            **extra_kwargs,
+        )
+
+    output_with_zero, state_with_zero = run(zero_state)
+    output_without_state, state_without_state = run(None)
+
+    torch.testing.assert_close(output_without_state, output_with_zero)
+    torch.testing.assert_close(state_without_state, state_with_zero)
+
+    num_cache_rows = 5
+    row_stride = H * D * D + 17
+    cache_storage = torch.full(
+        (num_cache_rows * row_stride,),
+        torch.nan,
+        dtype=torch.float32,
+        device=DEVICE,
+    )
+    final_state_cache = torch.as_strided(
+        cache_storage,
+        (num_cache_rows, H, D, D),
+        (row_stride, D * D, D, 1),
+    )
+    final_state_indices = torch.tensor(
+        [3, 1],
+        dtype=torch.int32,
+        device=DEVICE,
+    )
+    output_direct, returned_state = run(
+        None,
+        final_state_cache=final_state_cache,
+        final_state_indices=final_state_indices,
+    )
+
+    assert returned_state is None
+    torch.testing.assert_close(output_direct, output_with_zero)
+    torch.testing.assert_close(
+        final_state_cache[final_state_indices.long()],
+        state_with_zero,
+    )
+    assert torch.isnan(final_state_cache[[0, 2, 4]]).all()
+
+
 @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

@@ -188,6 +188,45 @@ def test_kimi_k3_kda_metadata_matches_shared_gdn(
     _assert_matches_shared_gdn(reference, actual)
 
 
+@pytest.mark.parametrize(
+    ("batch", "expected"),
+    [
+        pytest.param(
+            BatchSpec(seq_lens=[16, 32], query_lens=[16, 32]),
+            True,
+            id="all-fresh",
+        ),
+        pytest.param(
+            BatchSpec(seq_lens=[20, 32], query_lens=[4, 32]),
+            False,
+            id="one-resumed",
+        ),
+        pytest.param(
+            BatchSpec(seq_lens=[16, 0], query_lens=[16, 0]),
+            False,
+            id="zero-length-padding",
+        ),
+    ],
+)
+@pytest.mark.parametrize(
+    "builder_cls",
+    [GDNAttentionMetadataBuilder, KimiK3KDAMetadataBuilder],
+)
+def test_kda_metadata_marks_all_fresh_prefills(
+    batch: BatchSpec,
+    expected: bool,
+    builder_cls: type[AttentionMetadataBuilder],
+):
+    common_attn_metadata = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE)
+    metadata = _make_builder(
+        builder_cls,
+        num_speculative_tokens=0,
+        full_cuda_graph=False,
+    ).build(0, common_attn_metadata)
+
+    assert metadata.all_initial_states_fresh is expected
+
+
 def test_mixed_regular_and_spec_decode_uses_packed_decode_metadata():
     batch = BatchSpec(seq_lens=[100, 65, 20], query_lens=[1, 1, 3])
     common_attn_metadata = create_common_attn_metadata(

vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py

@@ -556,10 +556,20 @@ def _prefill_conv(
 
                 assert non_spec_state_indices_tensor is not None
                 assert has_initial_state is not None
-                initial_state = gather_initial_states(
-                    recurrent_state,
-                    non_spec_state_indices_tensor,
-                    has_initial_state,
+                # Chunk KDA initializes its recurrence to zero when no initial
+                # state is supplied. On ROCm it can also store final states
+                # directly into their cache rows.
+                direct_final_state = (
+                    current_platform.is_rocm() and m.all_initial_states_fresh
+                )
+                initial_state = (
+                    None
+                    if direct_final_state
+                    else gather_initial_states(
+                        recurrent_state,
+                        non_spec_state_indices_tensor,
+                        has_initial_state,
+                    )
                 )
                 (
                     core_attn_out_non_spec,
@@ -577,9 +587,21 @@ def _prefill_conv(
                     output_final_state=True,
                     use_qk_l2norm_in_kernel=True,
                     cu_seqlens=non_spec_query_start_loc,
+                    final_state_cache=(
+                        recurrent_state if direct_final_state else None
+                    ),
+                    final_state_indices=(
+                        non_spec_state_indices_tensor if direct_final_state else None
+                    ),
                 )
-                # Init cache
-                recurrent_state[non_spec_state_indices_tensor] = last_recurrent_state
+                if direct_final_state:
+                    assert last_recurrent_state is None
+                else:
+                    # Init cache
+                    assert last_recurrent_state is not None
+                    recurrent_state[non_spec_state_indices_tensor] = (
+                        last_recurrent_state
+                    )
 
             else:
                 # pure-decode non-spec batch

vllm/models/kimi_k3/amd/ops/third_party/kda/__init__.py

@@ -16,9 +16,9 @@
 #   - OOB-mask correctness fix: present (all tl.load use mask=..., other=0).
 # Validated on gfx950: no core-dump, gsm8k 94.1%.
 #
-# AMD-specific deltas vs the NVIDIA copy: NONE yet (byte-identical). Keep in sync
-# with the NVIDIA copy on FLA updates; any divergence should be an intentional,
-# documented gfx950-specific change (a #869-style AMD-only fix).
+# AMD-specific deltas vs the NVIDIA copy: chunk KDA can write final recurrent
+# states directly to indexed cache rows for all-fresh ROCm prefills. Keep other
+# FLA updates in sync; divergence should remain intentional and documented.
 
 from .chunk import (
     chunk_kda,

vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py

@@ -596,6 +596,8 @@ def _chunk_kda_fwd_with_cumulative_g(
     chunk_indices: torch.Tensor | None = None,
     chunk_size: int = FLA_CHUNK_SIZE,
     safe_gate: bool = False,
+    final_state_cache: torch.Tensor | None = None,
+    final_state_indices: torch.Tensor | None = None,
 ):
     # `g` must already be chunk-local cumulatively-summed AND scaled by
     # RCP_LN2 (so the downstream exp2-based kernels reproduce exp(g)).
@@ -632,6 +634,8 @@ def _chunk_kda_fwd_with_cumulative_g(
         cu_seqlens=cu_seqlens,
         chunk_indices=chunk_indices,
         use_exp2=True,
+        final_state_cache=final_state_cache,
+        final_state_indices=final_state_indices,
     )
     del w, u, kg
     o = chunk_gla_fwd_o_gk(
@@ -704,6 +708,8 @@ def chunk_kda_with_fused_gate_fwd(
     output_final_state: bool,
     lower_bound: float | None = None,
     cu_seqlens: torch.Tensor | None = None,
+    final_state_cache: torch.Tensor | None = None,
+    final_state_indices: torch.Tensor | None = None,
 ):
     chunk_size = FLA_CHUNK_SIZE
     chunk_indices = (
@@ -734,6 +740,8 @@ def chunk_kda_with_fused_gate_fwd(
         chunk_indices=chunk_indices,
         chunk_size=chunk_size,
         safe_gate=lower_bound is not None,
+        final_state_cache=final_state_cache,
+        final_state_indices=final_state_indices,
     )
 
 
@@ -785,6 +793,8 @@ def chunk_kda_with_fused_gate(
     lower_bound: float | None = None,
     use_qk_l2norm_in_kernel: bool = False,
     cu_seqlens: torch.Tensor | None = None,
+    final_state_cache: torch.Tensor | None = None,
+    final_state_indices: torch.Tensor | None = None,
     **kwargs,
 ):
     """Run chunk KDA from raw gate and beta projections."""
@@ -808,6 +818,8 @@ def chunk_kda_with_fused_gate(
         output_final_state=output_final_state,
         lower_bound=lower_bound,
         cu_seqlens=cu_seqlens,
+        final_state_cache=final_state_cache,
+        final_state_indices=final_state_indices,
     )
     return o, final_state
 

vllm/models/kimi_k3/nvidia/kda_metadata.py

@@ -403,11 +403,27 @@ def build(  # type: ignore[override]
         # Unlike the shared GDN layer, Kimi-K3's prefill KDA wrapper prepares
         # its own chunk indices. Only causal-convolution metadata is needed here.
         nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None
+        all_initial_states_fresh = False
         if num_prefills > 0:
             has_initial_state = m.compute_num_computed_tokens() > 0
+            has_initial_state_cpu = m._num_computed_tokens_cpu
+            non_spec_query_lens_for_fresh_cpu = query_start_loc_cpu.diff()
             if spec_sequence_masks_cpu is not None:
                 has_initial_state = has_initial_state[active_non_spec_mask_cpu]
+                if has_initial_state_cpu is not None:
+                    has_initial_state_cpu = has_initial_state_cpu[
+                        active_non_spec_mask_cpu
+                    ]
+                non_spec_query_lens_for_fresh_cpu = (
+                    non_spec_query_lens_for_fresh_cpu[~spec_sequence_masks_cpu]
+                )
                 assert non_spec_query_start_loc_cpu is not None
+            if has_initial_state_cpu is not None:
+                all_initial_states_fresh = bool(
+                    non_spec_query_lens_for_fresh_cpu.numel() > 0
+                    and (non_spec_query_lens_for_fresh_cpu > 0).all().item()
+                    and not has_initial_state_cpu.any().item()
+                )
             nums_dict, batch_ptr, token_chunk_offset_ptr = (
                 compute_causal_conv1d_metadata(
                     non_spec_query_start_loc_cpu,
@@ -472,6 +488,7 @@ def build(  # type: ignore[override]
             num_spec_decode_tokens=num_spec_decode_tokens,
             num_actual_tokens=m.num_actual_tokens,
             has_initial_state=has_initial_state,
+            all_initial_states_fresh=all_initial_states_fresh,
             spec_query_start_loc=spec_query_start_loc,
             non_spec_query_start_loc=non_spec_query_start_loc,
             spec_state_indices_tensor=spec_state_indices_tensor,

vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py

@@ -27,6 +27,7 @@
         "USE_GK": lambda args: args["gk"] is not None,
         "USE_INITIAL_STATE": lambda args: args["h0"] is not None,
         "STORE_FINAL_STATE": lambda args: args["ht"] is not None,
+        "USE_FINAL_STATE_INDICES": lambda args: args["ht_indices"] is not None,
         "SAVE_NEW_VALUE": lambda args: args["v_new"] is not None,
         "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
     }
@@ -52,8 +53,10 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
     h,
     h0,
     ht,
+    ht_indices,
     cu_seqlens,
     chunk_offsets,
+    ht_stride_n,
     T,
     H: tl.constexpr,
     Hg: tl.constexpr,
@@ -65,6 +68,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
     USE_GK: tl.constexpr,
     USE_INITIAL_STATE: tl.constexpr,
     STORE_FINAL_STATE: tl.constexpr,
+    USE_FINAL_STATE_INDICES: tl.constexpr,
     SAVE_NEW_VALUE: tl.constexpr,
     IS_VARLEN: tl.constexpr,
     USE_EXP2: tl.constexpr,
@@ -107,7 +111,11 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
     if USE_INITIAL_STATE:
         h0 = h0 + i_nh * V * K
     if STORE_FINAL_STATE:
-        ht = ht + i_nh * V * K
+        if USE_FINAL_STATE_INDICES:
+            i_state = tl.load(ht_indices + i_n).to(tl.int64)
+            ht = ht + i_state * ht_stride_n + i_h * V * K
+        else:
+            ht = ht + i_nh * V * K
 
     # load initial state
     if USE_INITIAL_STATE:
@@ -331,7 +339,9 @@ def chunk_gated_delta_rule_fwd_h(
     chunk_indices: torch.Tensor | None = None,
     chunk_offsets: torch.Tensor | None = None,
     use_exp2: bool = False,
-) -> tuple[torch.Tensor, torch.Tensor]:
+    final_state_cache: torch.Tensor | None = None,
+    final_state_indices: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
     # This kernel is slightly different from fla to support Q/K with different head numbers.
     # In fla, Q/K always have the same head number, so Hg is always equal to H.
     B, T, Hg, K, V = *k.shape, u.shape[-1]
@@ -350,9 +360,24 @@ def chunk_gated_delta_rule_fwd_h(
     assert K <= 256, "current kernel does not support head dimension larger than 256."
 
     h = k.new_empty(B, NT, H, V, K)
-    final_state = (
-        k.new_empty(N, H, V, K, dtype=torch.float32) if output_final_state else None
-    )
+    if final_state_cache is not None:
+        assert output_final_state
+        assert final_state_indices is not None
+        assert final_state_indices.numel() == N
+        assert final_state_cache.shape[1:] == (H, V, K)
+        assert final_state_cache.dtype == torch.float32
+        assert final_state_cache.device == k.device
+        assert final_state_cache.stride()[1:] == (V * K, K, 1)
+        final_state = final_state_cache
+        returned_final_state = None
+    else:
+        assert final_state_indices is None
+        final_state = (
+            k.new_empty(N, H, V, K, dtype=torch.float32)
+            if output_final_state
+            else None
+        )
+        returned_final_state = final_state
 
     v_new = torch.empty_like(u) if save_new_value else None
 
@@ -369,8 +394,10 @@ def grid(meta):
         h=h,
         h0=initial_state,
         ht=final_state,
+        ht_indices=final_state_indices,
         cu_seqlens=cu_seqlens,
         chunk_offsets=chunk_offsets,
+        ht_stride_n=final_state.stride(0) if final_state is not None else 0,
         T=T,
         H=H,
         Hg=Hg,
@@ -379,4 +406,4 @@ def grid(meta):
         BT=BT,
         USE_EXP2=use_exp2,
     )
-    return h, v_new, final_state
+    return h, v_new, returned_final_state