diff --git a/tests/models/kimi_k3/test_kda.py b/tests/models/kimi_k3/test_kda.py index be916f6c0d..9fae1c2bcd 100644 --- a/tests/models/kimi_k3/test_kda.py +++ b/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, diff --git a/vllm/models/kimi_k3/amd/ops/third_party/kda/fused_recurrent.py b/vllm/models/kimi_k3/amd/ops/third_party/kda/fused_recurrent.py index 2f512df626..db519fb6f0 100644 --- a/vllm/models/kimi_k3/amd/ops/third_party/kda/fused_recurrent.py +++ b/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,