[Bugfix][Kimi-K3] Give the AMD packed KDA decode kernel the state-index stride (#51682)

Signed-off-by: Lyu, Xudong <xudong.lyu@amd.com>

tests/models/kimi_k3/test_kda.py

@@ -15,6 +15,9 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import causal_conv1d_upd
 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 (
+    fused_recurrent_kda_packed_decode as fused_recurrent_kda_packed_decode_amd,
+)
 from vllm.models.kimi_k3.nvidia.kda import (
     is_flashkda_supported,
     is_fused_kda_decode_supported,
@@ -31,6 +34,13 @@ from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd
 
 DEVICE = "cuda"
 
+# The AMD and NVIDIA copies of the KDA kernels are vendored separately and are
+# free to diverge, so the shared-semantics tests below run against both.
+PACKED_DECODE_IMPLS = {
+    "nvidia": fused_recurrent_kda_packed_decode,
+    "amd": fused_recurrent_kda_packed_decode_amd,
+}
+
 
 @torch.inference_mode()
 def test_gather_initial_states_correctness():
@@ -273,11 +283,13 @@ def test_chunk_kda_fused_gate_cumsum_matches_unfused(
 @pytest.mark.parametrize("num_seqs", [1, 8, 32])
 @pytest.mark.parametrize("lower_bound", [-5.0, None])
 @pytest.mark.parametrize("state_indices_stride", [1, 8])
+@pytest.mark.parametrize("impl", PACKED_DECODE_IMPLS.keys())
 @torch.inference_mode()
 def test_packed_kda_decode_correctness(
     num_seqs: int,
     lower_bound: float | None,
     state_indices_stride: int,
+    impl: str,
 ):
     H, D = 8, 128
     torch.manual_seed(321)
@@ -361,7 +373,7 @@ def test_packed_kda_decode_correctness(
         use_qk_l2norm_in_kernel=True,
     )
     packed_state = state
-    packed_out, _ = fused_recurrent_kda_packed_decode(
+    packed_out, _ = PACKED_DECODE_IMPLS[impl](
         mixed_qkv=mixed_qkv,
         raw_g=raw_g,
         raw_beta=raw_beta,

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

@@ -459,6 +459,7 @@ def fused_recurrent_kda_packed_decode_kernel(
     stride_g_token: tl.constexpr,
     stride_beta_token: tl.constexpr,
     stride_state_token: tl.constexpr,
+    stride_state_indices,
     H: tl.constexpr,
     K: tl.constexpr,
     V: tl.constexpr,
@@ -476,7 +477,7 @@ def fused_recurrent_kda_packed_decode_kernel(
     mask_v = o_v < V
     mask_state = mask_v[:, None] & mask_k[None, :]
 
-    state_idx = tl.load(state_indices + i_n).to(tl.int64)
+    state_idx = tl.load(state_indices + i_n * stride_state_indices).to(tl.int64)
     p_out = out + (i_n * H + i_h) * V + o_v
     if state_idx <= 0:
         tl.store(p_out, tl.zeros([BV], dtype=tl.float32), mask=mask_v)
@@ -560,8 +561,8 @@ def fused_recurrent_kda_packed_decode(
         raise ValueError("`raw_beta` heads must be contiguous.")
     if initial_state.stride()[1:] != (V * K, K, 1):
         raise ValueError("`initial_state` must be contiguous within each cache slot.")
-    if state_indices.ndim != 1 or state_indices.stride(0) != 1:
-        raise ValueError("`state_indices` must be contiguous and one-dimensional.")
+    if state_indices.ndim != 1:
+        raise ValueError("`state_indices` must be one-dimensional.")
     if A_log.ndim != 1 or not A_log.is_contiguous():
         raise ValueError("`A_log` must be contiguous and one-dimensional.")
     if not dt_bias.is_contiguous():
@@ -608,6 +609,7 @@ def fused_recurrent_kda_packed_decode(
         stride_g_token=raw_g.stride(1),
         stride_beta_token=raw_beta.stride(1),
         stride_state_token=initial_state.stride(0),
+        stride_state_indices=state_indices.stride(0),
         H=H,
         K=K,
         V=V,