[Bugfix][Kimi-K3] Give the AMD packed KDA decode kernel the state-index stride (#51682)
Signed-off-by: Lyu, Xudong <xudong.lyu@amd.com>
Lyu, Xudong committed
Aug 11, 2026 at 00:05 UTC
0e2d78028c473bb04fc74bf945eb32da29239b0c
2 files changed
+18
-4
tests/models/kimi_k3/test_kda.py
+13
-1
@@ -15,6 +15,9 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import causal_conv1d_upd
15
from vllm.model_executor.layers.mamba.ops.gather_initial_states import (
16
gather_initial_states,
17
)
18
+from vllm.models.kimi_k3.amd.ops.third_party.kda import (
19
+ fused_recurrent_kda_packed_decode as fused_recurrent_kda_packed_decode_amd,
20
+)
21
from vllm.models.kimi_k3.nvidia.kda import (
22
is_flashkda_supported,
23
is_fused_kda_decode_supported,
@@ -31,6 +34,13 @@ from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd
34
35
DEVICE = "cuda"
36
37
+# The AMD and NVIDIA copies of the KDA kernels are vendored separately and are
38
+# free to diverge, so the shared-semantics tests below run against both.
39
+PACKED_DECODE_IMPLS = {
40
+ "nvidia": fused_recurrent_kda_packed_decode,
41
+ "amd": fused_recurrent_kda_packed_decode_amd,
42
+}
43
+
44
45
@torch.inference_mode()
46
def test_gather_initial_states_correctness():
@@ -273,11 +283,13 @@ def test_chunk_kda_fused_gate_cumsum_matches_unfused(
283
@pytest.mark.parametrize("num_seqs", [1, 8, 32])
284
@pytest.mark.parametrize("lower_bound", [-5.0, None])
285
@pytest.mark.parametrize("state_indices_stride", [1, 8])
286
+@pytest.mark.parametrize("impl", PACKED_DECODE_IMPLS.keys())
287
@torch.inference_mode()
288
def test_packed_kda_decode_correctness(
289
num_seqs: int,
290
lower_bound: float | None,
291
state_indices_stride: int,
292
+ impl: str,
293
):
294
H, D = 8, 128
295
torch.manual_seed(321)
@@ -361,7 +373,7 @@ def test_packed_kda_decode_correctness(
373
use_qk_l2norm_in_kernel=True,
374
)
375
packed_state = state
364
- packed_out, _ = fused_recurrent_kda_packed_decode(
376
+ packed_out, _ = PACKED_DECODE_IMPLS[impl](
377
mixed_qkv=mixed_qkv,
378
raw_g=raw_g,
379
raw_beta=raw_beta,
vllm/models/kimi_k3/amd/ops/third_party/kda/fused_recurrent.py
+5
-3
@@ -459,6 +459,7 @@ def fused_recurrent_kda_packed_decode_kernel(
459
stride_g_token: tl.constexpr,
460
stride_beta_token: tl.constexpr,
461
stride_state_token: tl.constexpr,
462
+ stride_state_indices,
463
H: tl.constexpr,
464
K: tl.constexpr,
465
V: tl.constexpr,
@@ -476,7 +477,7 @@ def fused_recurrent_kda_packed_decode_kernel(
477
mask_v = o_v < V
478
mask_state = mask_v[:, None] & mask_k[None, :]
479
479
- state_idx = tl.load(state_indices + i_n).to(tl.int64)
480
+ state_idx = tl.load(state_indices + i_n * stride_state_indices).to(tl.int64)
481
p_out = out + (i_n * H + i_h) * V + o_v
482
if state_idx <= 0:
483
tl.store(p_out, tl.zeros([BV], dtype=tl.float32), mask=mask_v)
@@ -560,8 +561,8 @@ def fused_recurrent_kda_packed_decode(
561
raise ValueError("`raw_beta` heads must be contiguous.")
562
if initial_state.stride()[1:] != (V * K, K, 1):
563
raise ValueError("`initial_state` must be contiguous within each cache slot.")
563
- if state_indices.ndim != 1 or state_indices.stride(0) != 1:
564
- raise ValueError("`state_indices` must be contiguous and one-dimensional.")
564
+ if state_indices.ndim != 1:
565
+ raise ValueError("`state_indices` must be one-dimensional.")
566
if A_log.ndim != 1 or not A_log.is_contiguous():
567
raise ValueError("`A_log` must be contiguous and one-dimensional.")
568
if not dt_bias.is_contiguous():
@@ -608,6 +609,7 @@ def fused_recurrent_kda_packed_decode(
609
stride_g_token=raw_g.stride(1),
610
stride_beta_token=raw_beta.stride(1),
611
stride_state_token=initial_state.stride(0),
612
+ stride_state_indices=state_indices.stride(0),
613
H=H,
614
K=K,
615
V=V,