| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """The ROCm fused KDA decode kernel must match the Triton chain it replaces. |
| 4 | |
| 5 | The fused kernel folds the packed causal conv1d update, the gated delta-rule |
| 6 | recurrence and the gated output RMSNorm into one launch, and updates both the |
| 7 | conv state and the recurrent state in place. Every one of those outputs is |
| 8 | compared against the three-kernel path the AMD layer falls back to. |
| 9 | """ |
| 10 | |
| 11 | import pytest |
| 12 | import torch |
| 13 | |
| 14 | from vllm.platforms import current_platform |
| 15 | |
| 16 | |
| 17 | def _on_supported_arch() -> bool: |
| 18 | if not current_platform.is_rocm(): |
| 19 | return False |
| 20 | from vllm.platforms.rocm import on_gfx942, on_gfx950 |
| 21 | |
| 22 | return on_gfx950() or on_gfx942() |
| 23 | |
| 24 | |
| 25 | pytestmark = pytest.mark.skipif( |
| 26 | not _on_supported_arch(), |
| 27 | reason="The fused KDA decode kernel is only built for gfx942 / gfx950", |
| 28 | ) |
| 29 | |
| 30 | # Kimi-K3 KDA: 96 heads x 128, conv width 4, gate_lower_bound -5.0. |
| 31 | HEAD_DIM = 128 |
| 32 | CONV_WIDTH = 4 |
| 33 | GATE_LOWER_BOUND = -5.0 |
| 34 | NORM_EPS = 1e-5 |
| 35 | DTYPE = torch.bfloat16 |
| 36 | |
| 37 | |
| 38 | def _requires_kernel() -> None: |
| 39 | if not hasattr(torch.ops._C, "fused_kda_decode"): |
| 40 | pytest.skip("vLLM was built without the fused KDA decode kernel") |
| 41 | |
| 42 | |
| 43 | class KdaDecodeInputs: |
| 44 | """One decode step of a Kimi-K3 KDA layer, in the layer's own layout.""" |
| 45 | |
| 46 | def __init__( |
| 47 | self, |
| 48 | num_tokens: int, |
| 49 | num_heads: int, |
| 50 | num_slots: int, |
| 51 | seed: int = 0, |
| 52 | ) -> None: |
| 53 | torch.manual_seed(seed) |
| 54 | device = "cuda" |
| 55 | dim = num_heads * HEAD_DIM |
| 56 | self.num_heads = num_heads |
| 57 | self.dim = dim |
| 58 | |
| 59 | self.mixed_qkv = ( |
| 60 | torch.randn(num_tokens, 3 * dim, device=device, dtype=DTYPE) * 0.5 |
| 61 | ) |
| 62 | # conv1d weight as the layer holds it: [3 * dim, width] fp32. |
| 63 | self.conv_weights = ( |
| 64 | torch.randn(3 * dim, CONV_WIDTH, device=device, dtype=torch.float32) * 0.3 |
| 65 | ) |
| 66 | # Width-major fp32 copy the fused kernel indexes: [3, width, dim]. |
| 67 | self.decode_conv1d_weight = torch.stack( |
| 68 | [ |
| 69 | self.conv_weights[i * dim : (i + 1) * dim].transpose(0, 1).contiguous() |
| 70 | for i in range(3) |
| 71 | ] |
| 72 | ) |
| 73 | # SD cache layout: [slots, width - 1, 3 * dim]; the layer transposes it. |
| 74 | self.conv_state = ( |
| 75 | torch.randn(num_slots, CONV_WIDTH - 1, 3 * dim, device=device, dtype=DTYPE) |
| 76 | * 0.5 |
| 77 | ) |
| 78 | self.recurrent_state = ( |
| 79 | torch.randn( |
| 80 | num_slots, |
| 81 | num_heads, |
| 82 | HEAD_DIM, |
| 83 | HEAD_DIM, |
| 84 | device=device, |
| 85 | dtype=torch.float32, |
| 86 | ) |
| 87 | * 0.1 |
| 88 | ) |
| 89 | self.g1 = ( |
| 90 | torch.randn(1, num_tokens, num_heads, HEAD_DIM, device=device, dtype=DTYPE) |
| 91 | * 0.5 |
| 92 | ) |
| 93 | self.g2 = ( |
| 94 | torch.randn(num_tokens, num_heads, HEAD_DIM, device=device, dtype=DTYPE) |
| 95 | * 0.5 |
| 96 | ) |
| 97 | self.beta = torch.randn(1, num_tokens, num_heads, device=device, dtype=DTYPE) |
| 98 | self.A_log = torch.randn(num_heads, device=device, dtype=torch.float32) * 0.5 |
| 99 | self.dt_bias = torch.randn(dim, device=device, dtype=torch.float32) * 0.1 |
| 100 | self.norm_weight_bf16 = 1 + 0.1 * torch.randn( |
| 101 | HEAD_DIM, device=device, dtype=DTYPE |
| 102 | ) |
| 103 | self.decode_norm_weight = self.norm_weight_bf16.float() |
| 104 | # Distinct, shuffled slots: the kernel must honour the indirection. |
| 105 | # Slot 0 is NULL_BLOCK_ID and never allocated for live state, and the |
| 106 | # Triton reference short-circuits it, so draw from [1, num_slots). |
| 107 | assert num_slots > num_tokens, "need a spare slot to detect stray writes" |
| 108 | self.state_indices = ( |
| 109 | torch.randperm(num_slots - 1, device=device)[:num_tokens] + 1 |
| 110 | ).to(torch.int32) |
| 111 | |
| 112 | def conv_state_view(self, state: torch.Tensor) -> torch.Tensor: |
| 113 | return state.transpose(-1, -2) |
| 114 | |
| 115 | |
| 116 | def _gated_rmsnorm( |
| 117 | x: torch.Tensor, gate: torch.Tensor, weight: torch.Tensor, eps: float |
| 118 | ) -> torch.Tensor: |
| 119 | x_float = x.float() |
| 120 | variance = x_float.pow(2).mean(dim=-1, keepdim=True) |
| 121 | normed = x_float * torch.rsqrt(variance + eps) * weight.float() |
| 122 | return (normed * torch.sigmoid(gate.float())).to(x.dtype) |
| 123 | |
| 124 | |
| 125 | def _run_triton_chain( |
| 126 | inp: KdaDecodeInputs, |
| 127 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 128 | """conv1d update -> recurrent decode -> gated norm, as the layer runs it.""" |
| 129 | from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( |
| 130 | causal_conv1d_update, |
| 131 | ) |
| 132 | from vllm.models.kimi_k3.amd.ops.third_party.kda import ( |
| 133 | fused_recurrent_kda_packed_decode, |
| 134 | ) |
| 135 | |
| 136 | conv_state = inp.conv_state.clone() |
| 137 | recurrent_state = inp.recurrent_state.clone() |
| 138 | conv_out = torch.empty_like(inp.mixed_qkv) |
| 139 | causal_conv1d_update( |
| 140 | inp.mixed_qkv, |
| 141 | inp.conv_state_view(conv_state), |
| 142 | inp.conv_weights, |
| 143 | None, |
| 144 | activation="silu", |
| 145 | conv_state_indices=inp.state_indices, |
| 146 | validate_data=True, |
| 147 | out=conv_out, |
| 148 | ) |
| 149 | core_attn_out, _ = fused_recurrent_kda_packed_decode( |
| 150 | mixed_qkv=conv_out, |
| 151 | raw_g=inp.g1, |
| 152 | raw_beta=inp.beta, |
| 153 | A_log=inp.A_log, |
| 154 | dt_bias=inp.dt_bias, |
| 155 | lower_bound=GATE_LOWER_BOUND, |
| 156 | initial_state=recurrent_state, |
| 157 | state_indices=inp.state_indices, |
| 158 | ) |
| 159 | out = _gated_rmsnorm(core_attn_out, inp.g2, inp.norm_weight_bf16, NORM_EPS) |
| 160 | return out, conv_state, recurrent_state |
| 161 | |
| 162 | |
| 163 | def _run_fused( |
| 164 | inp: KdaDecodeInputs, |
| 165 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 166 | from vllm import _custom_ops as ops |
| 167 | |
| 168 | conv_state = inp.conv_state.clone() |
| 169 | recurrent_state = inp.recurrent_state.clone() |
| 170 | num_tokens = inp.mixed_qkv.shape[0] |
| 171 | out = torch.empty( |
| 172 | 1, |
| 173 | num_tokens, |
| 174 | inp.num_heads, |
| 175 | HEAD_DIM, |
| 176 | device=inp.mixed_qkv.device, |
| 177 | dtype=DTYPE, |
| 178 | ) |
| 179 | ops.fused_kda_decode( |
| 180 | x=inp.mixed_qkv, |
| 181 | weight=inp.decode_conv1d_weight, |
| 182 | bias=None, |
| 183 | conv_state=inp.conv_state_view(conv_state), |
| 184 | raw_g=inp.g1, |
| 185 | raw_beta=inp.beta, |
| 186 | A_log=inp.A_log, |
| 187 | dt_bias=inp.dt_bias, |
| 188 | state_indices=inp.state_indices, |
| 189 | state=recurrent_state, |
| 190 | out=out, |
| 191 | lower_bound=GATE_LOWER_BOUND, |
| 192 | output_gate=inp.g2, |
| 193 | norm_weight=inp.decode_norm_weight, |
| 194 | norm_eps=NORM_EPS, |
| 195 | ) |
| 196 | return out, conv_state, recurrent_state |
| 197 | |
| 198 | |
| 199 | @pytest.mark.parametrize("num_heads", [12, 24, 96]) |
| 200 | @pytest.mark.parametrize("num_tokens", [1, 7, 128]) |
| 201 | @torch.inference_mode() |
| 202 | def test_fused_kda_decode_matches_triton_chain(num_heads: int, num_tokens: int) -> None: |
| 203 | _requires_kernel() |
| 204 | inp = KdaDecodeInputs(num_tokens, num_heads, num_slots=max(num_tokens, 4) + 3) |
| 205 | |
| 206 | expected_out, expected_conv, expected_state = _run_triton_chain(inp) |
| 207 | actual_out, actual_conv, actual_state = _run_fused(inp) |
| 208 | |
| 209 | # The fused kernel keeps the recurrent output in fp32 through the norm, |
| 210 | # while the Triton chain rounds it to BF16 in between. |
| 211 | torch.testing.assert_close(actual_out, expected_out, atol=3e-2, rtol=3e-2) |
| 212 | # The conv state is a pure shift-and-append: it must be bit-exact. |
| 213 | torch.testing.assert_close(actual_conv, expected_conv, atol=0, rtol=0) |
| 214 | torch.testing.assert_close(actual_state, expected_state, atol=2e-3, rtol=2e-3) |
| 215 | |
| 216 | |
| 217 | @torch.inference_mode() |
| 218 | def test_fused_kda_decode_leaves_untouched_slots_alone() -> None: |
| 219 | """Only the slots named by state_indices may be written.""" |
| 220 | _requires_kernel() |
| 221 | num_tokens, num_heads, num_slots = 4, 12, 9 |
| 222 | inp = KdaDecodeInputs(num_tokens, num_heads, num_slots, seed=3) |
| 223 | |
| 224 | _, actual_conv, actual_state = _run_fused(inp) |
| 225 | |
| 226 | touched = set(inp.state_indices.tolist()) |
| 227 | untouched = [slot for slot in range(num_slots) if slot not in touched] |
| 228 | assert untouched, "test needs at least one unused slot" |
| 229 | torch.testing.assert_close( |
| 230 | actual_conv[untouched], inp.conv_state[untouched], atol=0, rtol=0 |
| 231 | ) |
| 232 | torch.testing.assert_close( |
| 233 | actual_state[untouched], inp.recurrent_state[untouched], atol=0, rtol=0 |
| 234 | ) |
| 235 | |
| 236 | |
| 237 | @torch.inference_mode() |
| 238 | def test_fused_kda_decode_without_output_norm() -> None: |
| 239 | """Omitting the gate/norm pair returns the raw recurrent output.""" |
| 240 | _requires_kernel() |
| 241 | from vllm import _custom_ops as ops |
| 242 | |
| 243 | inp = KdaDecodeInputs(5, 12, num_slots=8, seed=7) |
| 244 | _, _, expected_state = _run_triton_chain(inp) |
| 245 | |
| 246 | conv_state = inp.conv_state.clone() |
| 247 | recurrent_state = inp.recurrent_state.clone() |
| 248 | out = torch.empty(1, 5, inp.num_heads, HEAD_DIM, device="cuda", dtype=DTYPE) |
| 249 | ops.fused_kda_decode( |
| 250 | x=inp.mixed_qkv, |
| 251 | weight=inp.decode_conv1d_weight, |
| 252 | bias=None, |
| 253 | conv_state=inp.conv_state_view(conv_state), |
| 254 | raw_g=inp.g1, |
| 255 | raw_beta=inp.beta, |
| 256 | A_log=inp.A_log, |
| 257 | dt_bias=inp.dt_bias, |
| 258 | state_indices=inp.state_indices, |
| 259 | state=recurrent_state, |
| 260 | out=out, |
| 261 | lower_bound=GATE_LOWER_BOUND, |
| 262 | ) |
| 263 | |
| 264 | torch.testing.assert_close(recurrent_state, expected_state, atol=2e-3, rtol=2e-3) |
| 265 | assert torch.isfinite(out).all() |
| 266 | |
| 267 | |
| 268 | @torch.inference_mode() |
| 269 | def test_fused_kda_decode_skips_null_block_padding() -> None: |
| 270 | """A CUDA-graph decode batch is padded with NULL_BLOCK_ID (0). |
| 271 | |
| 272 | gdn_attn.py fills the tail of ``non_spec_state_indices_tensor`` with 0, and |
| 273 | the Triton chain zeroes those rows' output while leaving slot 0 alone. The |
| 274 | fused kernel must do the same rather than read-modify-write slot 0's state |
| 275 | once per padded row. |
| 276 | """ |
| 277 | _requires_kernel() |
| 278 | num_real, num_padded, num_heads = 3, 5, 12 |
| 279 | inp = KdaDecodeInputs(num_real + num_padded, num_heads, num_slots=12, seed=11) |
| 280 | inp.state_indices[num_real:] = 0 |
| 281 | |
| 282 | expected_out, expected_conv, expected_state = _run_triton_chain(inp) |
| 283 | actual_out, actual_conv, actual_state = _run_fused(inp) |
| 284 | |
| 285 | torch.testing.assert_close(actual_out, expected_out, atol=3e-2, rtol=3e-2) |
| 286 | assert not actual_out[0, num_real:].any(), "padded rows must produce zeros" |
| 287 | # Slot 0 must be untouched by both paths. |
| 288 | torch.testing.assert_close(actual_conv[0], inp.conv_state[0], atol=0, rtol=0) |
| 289 | torch.testing.assert_close(actual_state[0], inp.recurrent_state[0], atol=0, rtol=0) |
| 290 | torch.testing.assert_close(actual_conv, expected_conv, atol=0, rtol=0) |
| 291 | torch.testing.assert_close(actual_state, expected_state, atol=2e-3, rtol=2e-3) |