| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Precision tests for vllm's chunk_kda Triton operator. |
| 4 | |
| 5 | Compares chunk_kda against a naive recurrent reference (float32). |
| 6 | Uses torch.rand for q/k/v to match FLA's test pattern. |
| 7 | """ |
| 8 | |
| 9 | from types import SimpleNamespace |
| 10 | |
| 11 | import pytest |
| 12 | import torch |
| 13 | import torch.nn.functional as F |
| 14 | |
| 15 | from vllm import _custom_ops as ops |
| 16 | from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture |
| 17 | from vllm.model_executor.layers.mamba.checkpoint import MambaPrefillCheckpointMetadata |
| 18 | from vllm.model_executor.layers.mamba.kda_checkpoint import ( |
| 19 | FlashKDAPrefillCheckpointExporter, |
| 20 | ) |
| 21 | from vllm.model_executor.layers.mamba.ops.causal_conv1d import causal_conv1d_update |
| 22 | from vllm.model_executor.layers.mamba.ops.gather_initial_states import ( |
| 23 | gather_initial_states, |
| 24 | ) |
| 25 | from vllm.models.kimi_k3.amd.ops.third_party.kda import ( |
| 26 | fused_recurrent_kda as fused_recurrent_kda_amd, |
| 27 | ) |
| 28 | from vllm.models.kimi_k3.amd.ops.third_party.kda import ( |
| 29 | fused_recurrent_kda_packed_decode as fused_recurrent_kda_packed_decode_amd, |
| 30 | ) |
| 31 | from vllm.models.kimi_k3.nvidia import kda as nvidia_kda |
| 32 | from vllm.models.kimi_k3.nvidia.kda import ( |
| 33 | KimiK3DeltaAttention, |
| 34 | _flashinfer_kda_prefill, |
| 35 | _flashkda_prefill, |
| 36 | is_flashinfer_fused_kda_decode_supported, |
| 37 | is_flashinfer_fused_kda_spec_decode_supported, |
| 38 | is_flashinfer_recurrent_kda_prefill_supported, |
| 39 | is_flashkda_supported, |
| 40 | is_fused_kda_decode_supported, |
| 41 | resolve_kda_spec_decode_backend, |
| 42 | ) |
| 43 | from vllm.models.kimi_k3.nvidia.kda_metadata import KimiK3KDAMetadata |
| 44 | from vllm.models.kimi_k3.nvidia.model import KimiLinearForCausalLM |
| 45 | from vllm.models.kimi_k3.nvidia.ops import recoverssm as recoverssm_ops |
| 46 | from vllm.models.kimi_k3.nvidia.ops.recoverssm import ( |
| 47 | KDARecoverSSMCommitContext, |
| 48 | kda_recoverssm_verify, |
| 49 | ) |
| 50 | from vllm.models.kimi_k3.nvidia.ops.third_party.kda import ( |
| 51 | chunk_kda, |
| 52 | chunk_kda_with_fused_gate, |
| 53 | fused_kda_gate, |
| 54 | fused_recurrent_kda_fwd, |
| 55 | fused_recurrent_kda_packed_decode, |
| 56 | ) |
| 57 | from vllm.models.kimi_k3.nvidia.ops.third_party.kda import ( |
| 58 | fused_recurrent_kda as fused_recurrent_kda_nvidia, |
| 59 | ) |
| 60 | from vllm.platforms import current_platform |
| 61 | from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd |
| 62 | from vllm.utils.flashinfer import flashinfer_fused_kda_decode |
| 63 | from vllm.v1.attention.backends.utils import NULL_BLOCK_ID |
| 64 | |
| 65 | DEVICE = current_platform.device_type |
| 66 | |
| 67 | pytestmark = pytest.mark.skipif( |
| 68 | not (current_platform.is_cuda_alike() or current_platform.is_xpu()), |
| 69 | reason="The KDA kernels require a CUDA-alike or XPU device.", |
| 70 | ) |
| 71 | |
| 72 | # The AMD and NVIDIA copies of the KDA kernels are vendored separately and are |
| 73 | # free to diverge, so the shared-semantics tests below run against both. |
| 74 | PACKED_DECODE_IMPLS = { |
| 75 | "nvidia": fused_recurrent_kda_packed_decode, |
| 76 | "amd": fused_recurrent_kda_packed_decode_amd, |
| 77 | } |
| 78 | SPEC_DECODE_IMPLS = { |
| 79 | "nvidia": fused_recurrent_kda_nvidia, |
| 80 | "amd": fused_recurrent_kda_amd, |
| 81 | } |
| 82 | |
| 83 | |
| 84 | def test_kda_warmup_skips_missing_metadata(monkeypatch): |
| 85 | monkeypatch.setattr( |
| 86 | nvidia_kda, |
| 87 | "get_forward_context", |
| 88 | lambda: SimpleNamespace(attn_metadata={}), |
| 89 | ) |
| 90 | layer = object.__new__(nvidia_kda.KimiK3DeltaAttention) |
| 91 | object.__setattr__(layer, "prefix", "language_model.model.layers.0.self_attn") |
| 92 | empty = torch.empty(0, device=DEVICE) |
| 93 | |
| 94 | assert layer._forward(empty, empty, empty, empty, empty) is None |
| 95 | |
| 96 | |
| 97 | def test_resolve_kda_spec_decode_backend(monkeypatch: pytest.MonkeyPatch): |
| 98 | monkeypatch.setattr( |
| 99 | "vllm.models.kimi_k3.nvidia.kda.is_flashinfer_fused_kda_spec_decode_supported", |
| 100 | lambda *args: True, |
| 101 | ) |
| 102 | args = (12, 128, 4, 6, torch.bfloat16, torch.bfloat16, torch.float32, -5.0) |
| 103 | assert resolve_kda_spec_decode_backend("auto", *args, False) == "flashinfer" |
| 104 | assert resolve_kda_spec_decode_backend("native", *args, False) == "native" |
| 105 | assert resolve_kda_spec_decode_backend("flashinfer", *args, False) == "flashinfer" |
| 106 | |
| 107 | monkeypatch.setattr( |
| 108 | "vllm.models.kimi_k3.nvidia.kda.is_flashinfer_fused_kda_spec_decode_supported", |
| 109 | lambda *args: False, |
| 110 | ) |
| 111 | assert resolve_kda_spec_decode_backend("auto", *args, False) == "native" |
| 112 | with pytest.raises(RuntimeError, match="packed_fused_kda_decode"): |
| 113 | resolve_kda_spec_decode_backend("flashinfer", *args, False) |
| 114 | |
| 115 | |
| 116 | def test_kda_recoverssm_config_state_layout(): |
| 117 | vllm_config = SimpleNamespace( |
| 118 | model_config=SimpleNamespace( |
| 119 | dtype=torch.bfloat16, |
| 120 | hf_config=SimpleNamespace( |
| 121 | linear_attn_config={ |
| 122 | "num_heads": 4, |
| 123 | "head_dim": 32, |
| 124 | "short_conv_kernel_size": 4, |
| 125 | } |
| 126 | ), |
| 127 | ), |
| 128 | cache_config=SimpleNamespace( |
| 129 | mamba_cache_dtype="auto", |
| 130 | mamba_ssm_cache_dtype="auto", |
| 131 | use_kda_recoverssm=True, |
| 132 | ), |
| 133 | parallel_config=SimpleNamespace(tensor_parallel_size=1), |
| 134 | speculative_config=SimpleNamespace(num_speculative_tokens=2), |
| 135 | ) |
| 136 | |
| 137 | assert KimiLinearForCausalLM.get_mamba_state_dtype_from_config(vllm_config) == ( |
| 138 | torch.bfloat16, |
| 139 | torch.float32, |
| 140 | torch.float32, |
| 141 | torch.bfloat16, |
| 142 | ) |
| 143 | assert KimiLinearForCausalLM.get_mamba_state_shape_from_config(vllm_config)[2:] == ( |
| 144 | (4, 3, 32), |
| 145 | (4, 3, 64), |
| 146 | ) |
| 147 | |
| 148 | |
| 149 | @torch.inference_mode() |
| 150 | def test_gather_initial_states_correctness(): |
| 151 | row_size = 8 * 128 * 128 |
| 152 | storage = torch.randn(5, row_size + 256, dtype=torch.float32, device=DEVICE) |
| 153 | state = storage[:, :row_size].view(5, 8, 128, 128) |
| 154 | assert not state.is_contiguous() |
| 155 | assert state[0].is_contiguous() |
| 156 | indices = torch.tensor([4, 1, 3], dtype=torch.int32, device=DEVICE) |
| 157 | has_initial_state = torch.tensor([True, False, True], device=DEVICE) |
| 158 | |
| 159 | expected = state[indices].clone() |
| 160 | expected[~has_initial_state] = 0 |
| 161 | |
| 162 | torch.testing.assert_close( |
| 163 | gather_initial_states(state, indices, has_initial_state), |
| 164 | expected, |
| 165 | ) |
| 166 | |
| 167 | |
| 168 | def naive_recurrent_kda( |
| 169 | q: torch.Tensor, |
| 170 | k: torch.Tensor, |
| 171 | v: torch.Tensor, |
| 172 | g: torch.Tensor, |
| 173 | beta: torch.Tensor, |
| 174 | scale: float | None = None, |
| 175 | initial_state: torch.Tensor | None = None, |
| 176 | output_final_state: bool = False, |
| 177 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 178 | """Naive recurrent KDA reference, ported from FLA's naive.py.""" |
| 179 | dtype = v.dtype |
| 180 | B, T, H, K = q.shape |
| 181 | V = v.shape[-1] |
| 182 | if scale is None: |
| 183 | scale = K**-0.5 |
| 184 | |
| 185 | q, k, v, g, beta = (x.to(torch.float) for x in [q, k, v, g, beta]) |
| 186 | q = q * scale |
| 187 | |
| 188 | S = k.new_zeros(B, H, K, V).to(q) |
| 189 | if initial_state is not None: |
| 190 | S += initial_state |
| 191 | o = torch.zeros_like(v) |
| 192 | for i in range(T): |
| 193 | q_i, k_i, v_i, g_i, b_i = q[:, i], k[:, i], v[:, i], g[:, i], beta[:, i] |
| 194 | S = S * g_i[..., None].exp() |
| 195 | S = S + torch.einsum( |
| 196 | "bhk,bhv->bhkv", |
| 197 | b_i[..., None] * k_i, |
| 198 | v_i - (k_i[..., None] * S).sum(-2), |
| 199 | ) |
| 200 | o[:, i] = torch.einsum("bhk,bhkv->bhv", q_i, S) |
| 201 | if not output_final_state: |
| 202 | S = None |
| 203 | return o.to(dtype), S |
| 204 | |
| 205 | |
| 206 | def assert_close( |
| 207 | name: str, |
| 208 | ref: torch.Tensor, |
| 209 | tri: torch.Tensor, |
| 210 | ratio: float, |
| 211 | err_atol: float = 1e-6, |
| 212 | ): |
| 213 | """RMSE-based relative error comparison.""" |
| 214 | abs_err = (ref.detach() - tri.detach()).flatten().abs().max().item() |
| 215 | rmse_diff = (ref.detach() - tri.detach()).flatten().square().mean().sqrt().item() |
| 216 | rmse_base = ref.detach().flatten().square().mean().sqrt().item() |
| 217 | rel_err = rmse_diff / (rmse_base + 1e-8) |
| 218 | print(f"{name:>4} | abs={abs_err:.6f} | rmse={rel_err:.6f} | thr={ratio}") |
| 219 | if abs_err <= err_atol: |
| 220 | return |
| 221 | assert not torch.isnan(ref).any(), f"{name}: NaN detected in ref" |
| 222 | assert not torch.isnan(tri).any(), f"{name}: NaN detected in tri" |
| 223 | assert rel_err < ratio, ( |
| 224 | f"{name}: max abs err {abs_err:.6f}, rmse ratio {rel_err:.6f} >= {ratio}" |
| 225 | ) |
| 226 | |
| 227 | |
| 228 | @pytest.mark.parametrize( |
| 229 | ("H", "D", "cu_seqlens", "dtype"), |
| 230 | [ |
| 231 | pytest.param( |
| 232 | *test, |
| 233 | id="H{}-D{}-cu{}-{}".format(*test), |
| 234 | ) |
| 235 | for test in [ |
| 236 | (32, 128, [0, 64], torch.float16), |
| 237 | (32, 128, [0, 1024], torch.float16), |
| 238 | (32, 128, [0, 15], torch.float16), |
| 239 | (32, 128, [0, 256, 512, 768, 1024], torch.float16), |
| 240 | (32, 128, [0, 15, 100, 300, 1200], torch.float16), |
| 241 | (64, 128, [0, 256, 500, 1000], torch.float16), |
| 242 | (32, 128, [0, 8192], torch.float16), |
| 243 | (32, 128, [0, 256, 500, 1000], torch.bfloat16), |
| 244 | ] |
| 245 | ], |
| 246 | ) |
| 247 | @torch.inference_mode() |
| 248 | def test_chunk_kda( |
| 249 | H: int, |
| 250 | D: int, |
| 251 | cu_seqlens: list[int], |
| 252 | dtype: torch.dtype, |
| 253 | ): |
| 254 | T = cu_seqlens[-1] |
| 255 | torch.manual_seed(42) |
| 256 | B = 1 |
| 257 | cu_seqlens_t = torch.LongTensor(cu_seqlens).to(DEVICE) |
| 258 | N = len(cu_seqlens) - 1 |
| 259 | |
| 260 | q = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE) |
| 261 | k = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE) |
| 262 | v = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE) |
| 263 | g = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float32, device=DEVICE)).to( |
| 264 | dtype |
| 265 | ) |
| 266 | beta = torch.rand(B, T, H, dtype=dtype, device=DEVICE).sigmoid() |
| 267 | h0 = torch.randn(N, H, D, D, dtype=torch.float32, device=DEVICE) |
| 268 | |
| 269 | # Naive reference with l2norm_fwd (same kernel as chunk_kda) |
| 270 | ref_outputs = [] |
| 271 | ref_states = [] |
| 272 | for i in range(N): |
| 273 | s, e = cu_seqlens[i], cu_seqlens[i + 1] |
| 274 | q_i = l2norm_fwd(q[:, s:e].contiguous()) |
| 275 | k_i = l2norm_fwd(k[:, s:e].contiguous()) |
| 276 | o_i, ht_i = naive_recurrent_kda( |
| 277 | q_i, |
| 278 | k_i, |
| 279 | v[:, s:e], |
| 280 | g[:, s:e], |
| 281 | beta[:, s:e], |
| 282 | initial_state=h0[i], |
| 283 | output_final_state=True, |
| 284 | ) |
| 285 | ref_outputs.append(o_i) |
| 286 | ref_states.append(ht_i) |
| 287 | ref_o = torch.cat(ref_outputs, dim=1) |
| 288 | ref_ht = torch.cat(ref_states, dim=0) |
| 289 | |
| 290 | # h0 transposed to (V, K) layout for the kernel; naive uses (K, V) |
| 291 | tri_o, tri_ht = chunk_kda( |
| 292 | q=q.clone(), |
| 293 | k=k.clone(), |
| 294 | v=v.clone(), |
| 295 | g=g.clone(), |
| 296 | beta=beta.clone(), |
| 297 | initial_state=h0.transpose(-1, -2).contiguous().clone(), |
| 298 | output_final_state=True, |
| 299 | cu_seqlens=cu_seqlens_t, |
| 300 | use_qk_l2norm_in_kernel=True, |
| 301 | ) |
| 302 | |
| 303 | assert not torch.isnan(tri_o).any(), "Triton output o contains NaN" |
| 304 | assert not torch.isnan(tri_ht).any(), "Triton output ht contains NaN" |
| 305 | assert_close("o", ref_o, tri_o, 0.005) |
| 306 | assert_close("ht", ref_ht, tri_ht.transpose(-1, -2).contiguous(), 0.005) |
| 307 | |
| 308 | |
| 309 | @pytest.mark.parametrize( |
| 310 | ("cu_seqlens", "dtype", "lower_bound"), |
| 311 | [ |
| 312 | ([0, 64], torch.float16, None), |
| 313 | ([0, 15, 100, 300], torch.bfloat16, None), |
| 314 | ([0, 15, 100, 300], torch.bfloat16, -3.0), |
| 315 | ], |
| 316 | ) |
| 317 | @torch.inference_mode() |
| 318 | def test_chunk_kda_fused_gate_cumsum_matches_unfused( |
| 319 | cu_seqlens: list[int], |
| 320 | dtype: torch.dtype, |
| 321 | lower_bound: float | None, |
| 322 | ): |
| 323 | H, D = 8, 64 |
| 324 | T = cu_seqlens[-1] |
| 325 | N = len(cu_seqlens) - 1 |
| 326 | torch.manual_seed(123) |
| 327 | |
| 328 | cu_seqlens_t = torch.tensor(cu_seqlens, dtype=torch.int32, device=DEVICE) |
| 329 | q = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE) |
| 330 | k = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE) |
| 331 | v = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE) |
| 332 | raw_g = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE) |
| 333 | beta_storage = torch.randn(1, T, 2 * H + 3, dtype=dtype, device=DEVICE) |
| 334 | raw_beta = beta_storage[..., 1 : 2 * H + 1 : 2] |
| 335 | beta = raw_beta.float().sigmoid() |
| 336 | A_log = (torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5).contiguous() |
| 337 | dt_bias = ( |
| 338 | torch.randn(H * D, dtype=torch.float32, device=DEVICE) * 0.1 |
| 339 | ).contiguous() |
| 340 | h0 = torch.randn(N, H, D, D, dtype=torch.float32, device=DEVICE) |
| 341 | initial_state = h0.transpose(-1, -2).contiguous() |
| 342 | |
| 343 | gate = fused_kda_gate( |
| 344 | raw_g.reshape(T, H * D), |
| 345 | A_log, |
| 346 | D, |
| 347 | g_bias=dt_bias, |
| 348 | lower_bound=lower_bound, |
| 349 | ) |
| 350 | if lower_bound is not None: |
| 351 | expected_gate = lower_bound * torch.sigmoid( |
| 352 | A_log.exp()[None, :, None] |
| 353 | * (raw_g.float().view(T, H, D) + dt_bias.view(H, D)) |
| 354 | ) |
| 355 | torch.testing.assert_close(gate, expected_gate) |
| 356 | gate = gate.unsqueeze(0) |
| 357 | old_o, old_ht = chunk_kda( |
| 358 | q=q.clone(), |
| 359 | k=k.clone(), |
| 360 | v=v.clone(), |
| 361 | g=gate, |
| 362 | beta=beta, |
| 363 | initial_state=initial_state.clone(), |
| 364 | output_final_state=True, |
| 365 | cu_seqlens=cu_seqlens_t, |
| 366 | use_qk_l2norm_in_kernel=True, |
| 367 | ) |
| 368 | output = torch.empty_like(v) |
| 369 | new_o, new_ht = chunk_kda_with_fused_gate( |
| 370 | q=q.clone(), |
| 371 | k=k.clone(), |
| 372 | v=v.clone(), |
| 373 | raw_g=raw_g, |
| 374 | raw_beta=raw_beta, |
| 375 | A_log=A_log, |
| 376 | g_bias=dt_bias, |
| 377 | lower_bound=lower_bound, |
| 378 | initial_state=initial_state.clone(), |
| 379 | output_final_state=True, |
| 380 | cu_seqlens=cu_seqlens_t, |
| 381 | use_qk_l2norm_in_kernel=True, |
| 382 | out=output, |
| 383 | ) |
| 384 | |
| 385 | assert new_o.data_ptr() == output.data_ptr() |
| 386 | assert_close("o", old_o, new_o, 1e-3, err_atol=1e-3) |
| 387 | assert_close("ht", old_ht, new_ht, 1e-3, err_atol=1e-3) |
| 388 | |
| 389 | |
| 390 | @pytest.mark.parametrize("num_seqs", [1, 8, 32]) |
| 391 | @pytest.mark.parametrize("lower_bound", [-5.0, None]) |
| 392 | @pytest.mark.parametrize("state_indices_stride", [1, 8]) |
| 393 | @pytest.mark.parametrize("impl", PACKED_DECODE_IMPLS.keys()) |
| 394 | @torch.inference_mode() |
| 395 | def test_packed_kda_decode_correctness( |
| 396 | num_seqs: int, |
| 397 | lower_bound: float | None, |
| 398 | state_indices_stride: int, |
| 399 | impl: str, |
| 400 | ): |
| 401 | H, D = 8, 128 |
| 402 | torch.manual_seed(321) |
| 403 | |
| 404 | packed_storage = torch.randn( |
| 405 | num_seqs, |
| 406 | 3 * H * D + 1, |
| 407 | dtype=torch.bfloat16, |
| 408 | device=DEVICE, |
| 409 | ) |
| 410 | mixed_qkv = packed_storage[:, : 3 * H * D] |
| 411 | assert mixed_qkv.stride(0) == 3 * H * D + 1 |
| 412 | q, k, v = ( |
| 413 | x.contiguous().view(1, num_seqs, H, D) for x in mixed_qkv.split(H * D, dim=-1) |
| 414 | ) |
| 415 | raw_g = torch.randn( |
| 416 | 1, |
| 417 | num_seqs, |
| 418 | H, |
| 419 | D, |
| 420 | dtype=torch.bfloat16, |
| 421 | device=DEVICE, |
| 422 | ) |
| 423 | raw_beta = torch.randn( |
| 424 | 1, |
| 425 | num_seqs, |
| 426 | H, |
| 427 | dtype=torch.bfloat16, |
| 428 | device=DEVICE, |
| 429 | ) |
| 430 | beta = raw_beta.float().sigmoid() |
| 431 | A_log = torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5 |
| 432 | dt_bias = torch.randn(H, D, dtype=torch.float32, device=DEVICE) * 0.1 |
| 433 | state_storage = torch.randn( |
| 434 | num_seqs + 1, |
| 435 | H * D * D + 17, |
| 436 | dtype=torch.float32, |
| 437 | device=DEVICE, |
| 438 | ) |
| 439 | state = state_storage[:, : H * D * D].view(num_seqs + 1, H, D, D) |
| 440 | assert not state.is_contiguous() |
| 441 | assert state.stride()[1:] == (D * D, D, 1) |
| 442 | state_indices_storage = torch.zeros( |
| 443 | num_seqs, |
| 444 | state_indices_stride, |
| 445 | dtype=torch.int32, |
| 446 | device=DEVICE, |
| 447 | ) |
| 448 | state_indices = state_indices_storage[:, 0] |
| 449 | state_indices.copy_( |
| 450 | torch.arange( |
| 451 | 1, |
| 452 | num_seqs + 1, |
| 453 | dtype=torch.int32, |
| 454 | device=DEVICE, |
| 455 | ) |
| 456 | ) |
| 457 | gate = fused_kda_gate( |
| 458 | raw_g.reshape(num_seqs, H * D), |
| 459 | A_log, |
| 460 | D, |
| 461 | g_bias=dt_bias, |
| 462 | lower_bound=lower_bound, |
| 463 | ).unsqueeze(0) |
| 464 | dense_state = state.clone() |
| 465 | dense_out, _ = fused_recurrent_kda_fwd( |
| 466 | q=q, |
| 467 | k=k, |
| 468 | v=v, |
| 469 | g=gate, |
| 470 | beta=beta, |
| 471 | scale=D**-0.5, |
| 472 | initial_state=dense_state, |
| 473 | inplace_final_state=True, |
| 474 | cu_seqlens=torch.arange( |
| 475 | num_seqs + 1, |
| 476 | dtype=torch.int32, |
| 477 | device=DEVICE, |
| 478 | ), |
| 479 | ssm_state_indices=state_indices, |
| 480 | use_qk_l2norm_in_kernel=True, |
| 481 | ) |
| 482 | packed_state = state |
| 483 | packed_output = torch.empty_like(dense_out) |
| 484 | packed_out, _ = PACKED_DECODE_IMPLS[impl]( |
| 485 | mixed_qkv=mixed_qkv, |
| 486 | raw_g=raw_g, |
| 487 | raw_beta=raw_beta, |
| 488 | A_log=A_log, |
| 489 | dt_bias=dt_bias, |
| 490 | lower_bound=lower_bound, |
| 491 | initial_state=packed_state, |
| 492 | state_indices=state_indices, |
| 493 | **({"out": packed_output} if impl == "nvidia" else {}), |
| 494 | ) |
| 495 | |
| 496 | if impl == "nvidia": |
| 497 | assert packed_out.data_ptr() == packed_output.data_ptr() |
| 498 | assert_close("o", dense_out, packed_out, 1e-3, err_atol=1e-3) |
| 499 | assert_close("ht", dense_state, packed_state, 1e-3, err_atol=1e-3) |
| 500 | |
| 501 | |
| 502 | @pytest.mark.parametrize( |
| 503 | ("impl", "H", "fuse_gate", "num_seqs", "query_len"), |
| 504 | [ |
| 505 | *[ |
| 506 | pytest.param( |
| 507 | impl, |
| 508 | H, |
| 509 | fuse_gate, |
| 510 | 3, |
| 511 | 3, |
| 512 | id=f"{impl}-H{H}-fuse-{fuse_gate}", |
| 513 | ) |
| 514 | for impl in SPEC_DECODE_IMPLS |
| 515 | for H, fuse_gate in [ |
| 516 | (12, True), |
| 517 | (12, False), |
| 518 | (12, None), |
| 519 | (96, None), |
| 520 | ] |
| 521 | ], |
| 522 | pytest.param("amd", 12, True, 1, 1, id="amd-single-token"), |
| 523 | pytest.param("amd", 12, True, 1, 8, id="amd-single-sequence-short"), |
| 524 | pytest.param( |
| 525 | "amd", |
| 526 | 96, |
| 527 | None, |
| 528 | 1, |
| 529 | 8, |
| 530 | id="amd-single-sequence-many-heads", |
| 531 | ), |
| 532 | pytest.param("amd", 12, True, 2, 4, id="amd-uniform-two-sequences"), |
| 533 | pytest.param("amd", 12, True, 4, 4, id="amd-uniform-four-sequences"), |
| 534 | pytest.param("amd", 12, True, 8, 7, id="amd-uniform-eight-sequences"), |
| 535 | ], |
| 536 | ) |
| 537 | @pytest.mark.parametrize("lower_bound", [-5.0, None]) |
| 538 | @torch.inference_mode() |
| 539 | def test_kda_spec_decode_correctness( |
| 540 | H: int, |
| 541 | fuse_gate: bool | None, |
| 542 | num_seqs: int, |
| 543 | query_len: int, |
| 544 | lower_bound: float | None, |
| 545 | impl: str, |
| 546 | ): |
| 547 | D = 128 |
| 548 | T = num_seqs * query_len |
| 549 | torch.manual_seed(1234) |
| 550 | |
| 551 | qkv_storage = torch.randn( |
| 552 | 1, |
| 553 | T, |
| 554 | 3 * H * D + 7, |
| 555 | dtype=torch.bfloat16, |
| 556 | device=DEVICE, |
| 557 | ) |
| 558 | packed_qkv = qkv_storage[..., : 3 * H * D] |
| 559 | q, k, v = (x.view(1, T, H, D) for x in packed_qkv.split(H * D, dim=-1)) |
| 560 | gate_storage = torch.randn( |
| 561 | 1, |
| 562 | T, |
| 563 | H * D + 5, |
| 564 | dtype=torch.bfloat16, |
| 565 | device=DEVICE, |
| 566 | ) |
| 567 | raw_g = gate_storage[..., : H * D].view(1, T, H, D) |
| 568 | beta_storage = torch.randn( |
| 569 | 1, |
| 570 | T, |
| 571 | H + 1, |
| 572 | dtype=torch.bfloat16, |
| 573 | device=DEVICE, |
| 574 | ) |
| 575 | raw_beta = beta_storage[..., :H] |
| 576 | A_log = 0.5 * torch.randn(H, dtype=torch.float32, device=DEVICE) |
| 577 | dt_bias = 0.1 * torch.randn(H, D, dtype=torch.float32, device=DEVICE) |
| 578 | cu_seqlens = torch.arange( |
| 579 | 0, |
| 580 | T + 1, |
| 581 | query_len, |
| 582 | dtype=torch.int32, |
| 583 | device=DEVICE, |
| 584 | ) |
| 585 | state_indices = torch.arange( |
| 586 | 1, |
| 587 | T + 1, |
| 588 | dtype=torch.int32, |
| 589 | device=DEVICE, |
| 590 | ).view(num_seqs, query_len) |
| 591 | num_accepted_tokens = ( |
| 592 | torch.arange(num_seqs, dtype=torch.int32, device=DEVICE) % query_len + 1 |
| 593 | ) |
| 594 | state_storage = 0.01 * torch.randn( |
| 595 | T + 1, |
| 596 | H * D * D + 17, |
| 597 | dtype=torch.float32, |
| 598 | device=DEVICE, |
| 599 | ) |
| 600 | state = state_storage[:, : H * D * D].view(T + 1, H, D, D) |
| 601 | output_storage = torch.full( |
| 602 | (1, T, H * D + 11), |
| 603 | torch.nan, |
| 604 | dtype=torch.bfloat16, |
| 605 | device=DEVICE, |
| 606 | ) |
| 607 | output = output_storage[..., : H * D].view(1, T, H, D) |
| 608 | |
| 609 | gate = fused_kda_gate( |
| 610 | raw_g.contiguous().view(T, H * D), |
| 611 | A_log, |
| 612 | D, |
| 613 | g_bias=dt_bias, |
| 614 | lower_bound=lower_bound, |
| 615 | ).unsqueeze(0) |
| 616 | beta = raw_beta.float().sigmoid() |
| 617 | q_norm = l2norm_fwd(q.contiguous()) |
| 618 | k_norm = l2norm_fwd(k.contiguous()) |
| 619 | expected_state = state.clone() |
| 620 | expected_outputs = [] |
| 621 | for seq, accepted in enumerate(num_accepted_tokens.tolist()): |
| 622 | recurrent_state = expected_state[state_indices[seq, accepted - 1]].transpose( |
| 623 | -1, -2 |
| 624 | ) |
| 625 | start = seq * query_len |
| 626 | for token in range(query_len): |
| 627 | token_slice = slice(start + token, start + token + 1) |
| 628 | token_output, recurrent_state = naive_recurrent_kda( |
| 629 | q_norm[:, token_slice], |
| 630 | k_norm[:, token_slice], |
| 631 | v[:, token_slice], |
| 632 | gate[:, token_slice], |
| 633 | beta[:, token_slice], |
| 634 | initial_state=recurrent_state, |
| 635 | output_final_state=True, |
| 636 | ) |
| 637 | assert recurrent_state is not None |
| 638 | expected_outputs.append(token_output) |
| 639 | expected_state[state_indices[seq, token]] = recurrent_state.transpose( |
| 640 | -1, -2 |
| 641 | ) |
| 642 | expected = torch.cat(expected_outputs, dim=1) |
| 643 | |
| 644 | actual_state = state.clone() |
| 645 | extra_args = ( |
| 646 | {"uniform_sequence_length": query_len} |
| 647 | if impl == "amd" and num_seqs in (2, 4, 8) |
| 648 | else {} |
| 649 | ) |
| 650 | actual, _ = SPEC_DECODE_IMPLS[impl]( |
| 651 | q=q, |
| 652 | k=k, |
| 653 | v=v, |
| 654 | raw_g=raw_g, |
| 655 | raw_beta=raw_beta, |
| 656 | A_log=A_log, |
| 657 | dt_bias=dt_bias, |
| 658 | lower_bound=lower_bound, |
| 659 | initial_state=actual_state, |
| 660 | cu_seqlens=cu_seqlens, |
| 661 | ssm_state_indices=state_indices, |
| 662 | num_accepted_tokens=num_accepted_tokens, |
| 663 | out=output, |
| 664 | fuse_gate=fuse_gate, |
| 665 | **extra_args, |
| 666 | ) |
| 667 | |
| 668 | assert actual.data_ptr() == output.data_ptr() |
| 669 | assert_close("o", expected, actual, 1e-3, err_atol=1e-3) |
| 670 | used_states = state_indices.flatten().long() |
| 671 | assert_close( |
| 672 | "ht", |
| 673 | expected_state[used_states], |
| 674 | actual_state[used_states], |
| 675 | 3e-3, |
| 676 | err_atol=3e-3, |
| 677 | ) |
| 678 | assert torch.isnan(output_storage[..., H * D :]).all() |
| 679 | |
| 680 | |
| 681 | @pytest.mark.parametrize( |
| 682 | ( |
| 683 | "conv_state_dim_first", |
| 684 | "use_request_indices", |
| 685 | "lower_bound", |
| 686 | "align_mode", |
| 687 | ), |
| 688 | [ |
| 689 | pytest.param(False, False, None, False, id="baseline"), |
| 690 | pytest.param(True, True, -5.0, True, id="all-features"), |
| 691 | pytest.param(False, True, -5.0, False, id="request-indexed"), |
| 692 | pytest.param(True, False, None, True, id="aligned"), |
| 693 | ], |
| 694 | ) |
| 695 | @torch.inference_mode() |
| 696 | def test_kda_recoverssm_verify_and_group_commit( |
| 697 | monkeypatch: pytest.MonkeyPatch, |
| 698 | lower_bound: float | None, |
| 699 | use_request_indices: bool, |
| 700 | conv_state_dim_first: bool, |
| 701 | align_mode: bool, |
| 702 | ): |
| 703 | monkeypatch.setattr( |
| 704 | recoverssm_ops, |
| 705 | "is_conv_state_dim_first", |
| 706 | lambda: conv_state_dim_first, |
| 707 | ) |
| 708 | num_layers, num_seqs, query_len = 2, 2, 8 |
| 709 | num_blocks, num_heads, dim = (7 if align_mode else 3), 4, 128 |
| 710 | total_tokens = num_seqs * query_len |
| 711 | torch.manual_seed(20260808) |
| 712 | |
| 713 | q, k, v, raw_g = [ |
| 714 | torch.randn( |
| 715 | 1, |
| 716 | total_tokens, |
| 717 | num_heads, |
| 718 | dim, |
| 719 | dtype=torch.bfloat16, |
| 720 | device=DEVICE, |
| 721 | ) |
| 722 | for _ in range(4) |
| 723 | ] |
| 724 | raw_beta = torch.randn( |
| 725 | 1, |
| 726 | total_tokens, |
| 727 | num_heads, |
| 728 | dtype=torch.bfloat16, |
| 729 | device=DEVICE, |
| 730 | ) |
| 731 | query_start_loc = torch.arange( |
| 732 | 0, |
| 733 | total_tokens + 1, |
| 734 | query_len, |
| 735 | dtype=torch.int32, |
| 736 | device=DEVICE, |
| 737 | ) |
| 738 | state_indices = torch.tensor( |
| 739 | [5, 6] if align_mode else [1, 2], dtype=torch.int32, device=DEVICE |
| 740 | ) |
| 741 | accepted = [2, 8] |
| 742 | if use_request_indices: |
| 743 | global_num_accepted = torch.tensor( |
| 744 | [0, accepted[0], 0, accepted[1]], |
| 745 | dtype=torch.int32, |
| 746 | device=DEVICE, |
| 747 | ) |
| 748 | request_indices = torch.tensor([1, 3], dtype=torch.int32, device=DEVICE) |
| 749 | else: |
| 750 | global_num_accepted = torch.tensor(accepted, dtype=torch.int32, device=DEVICE) |
| 751 | request_indices = None |
| 752 | |
| 753 | block_table = None |
| 754 | num_computed_tokens = None |
| 755 | mamba_block_size = None |
| 756 | if align_mode: |
| 757 | batch_size = 4 if use_request_indices else num_seqs |
| 758 | block_table = torch.full((batch_size, 2), -1, dtype=torch.int32, device=DEVICE) |
| 759 | rows = ( |
| 760 | request_indices |
| 761 | if request_indices is not None |
| 762 | else torch.arange(num_seqs, device=DEVICE) |
| 763 | ) |
| 764 | block_table[rows] = torch.tensor( |
| 765 | [[1, 5], [2, 6]], |
| 766 | dtype=torch.int32, |
| 767 | device=DEVICE, |
| 768 | ) |
| 769 | num_computed_tokens = torch.zeros(batch_size, dtype=torch.int32, device=DEVICE) |
| 770 | num_computed_tokens[rows] = 4 |
| 771 | mamba_block_size = 8 |
| 772 | |
| 773 | layers = [] |
| 774 | expected_outputs = [] |
| 775 | expected_states = [] |
| 776 | initial_states = [] |
| 777 | initial_conv_states = [] |
| 778 | history_len, conv_dim = 3, 12 |
| 779 | for layer_idx in range(num_layers): |
| 780 | A_log = ( |
| 781 | 0.2 * torch.randn(num_heads, dtype=torch.float32, device=DEVICE) |
| 782 | + layer_idx * 0.03 |
| 783 | ).contiguous() |
| 784 | dt_bias = ( |
| 785 | 0.1 * torch.randn(num_heads, dim, dtype=torch.float32, device=DEVICE) |
| 786 | ).contiguous() |
| 787 | checkpoint = 0.01 * torch.randn( |
| 788 | num_blocks, |
| 789 | num_heads, |
| 790 | dim, |
| 791 | dim, |
| 792 | dtype=torch.float32, |
| 793 | device=DEVICE, |
| 794 | ) |
| 795 | conv_shape = ( |
| 796 | (num_blocks, conv_dim, history_len + query_len - 1) |
| 797 | if conv_state_dim_first |
| 798 | else (num_blocks, history_len + query_len - 1, conv_dim) |
| 799 | ) |
| 800 | conv_state = torch.randn(conv_shape, dtype=torch.bfloat16, device=DEVICE) |
| 801 | correction_cache = torch.empty( |
| 802 | num_blocks, |
| 803 | num_heads, |
| 804 | query_len, |
| 805 | dim, |
| 806 | dtype=torch.float32, |
| 807 | device=DEVICE, |
| 808 | ) |
| 809 | kg_cache = torch.empty( |
| 810 | num_blocks, |
| 811 | num_heads, |
| 812 | query_len, |
| 813 | 2 * dim, |
| 814 | dtype=torch.bfloat16, |
| 815 | device=DEVICE, |
| 816 | ) |
| 817 | layer = SimpleNamespace( |
| 818 | kv_cache=( |
| 819 | conv_state, |
| 820 | checkpoint, |
| 821 | correction_cache, |
| 822 | kg_cache, |
| 823 | ), |
| 824 | A_log=A_log, |
| 825 | dt_bias=dt_bias, |
| 826 | local_num_heads=num_heads, |
| 827 | head_dim=dim, |
| 828 | gate_lower_bound=lower_bound, |
| 829 | ) |
| 830 | layers.append(layer) |
| 831 | initial_states.append(checkpoint.clone()) |
| 832 | initial_conv_states.append(conv_state.clone()) |
| 833 | |
| 834 | actual_output = kda_recoverssm_verify( |
| 835 | q=q, |
| 836 | k=k, |
| 837 | v=v, |
| 838 | raw_g=raw_g, |
| 839 | raw_beta=raw_beta, |
| 840 | A_log=A_log, |
| 841 | dt_bias=dt_bias, |
| 842 | lower_bound=lower_bound, |
| 843 | checkpoint_state=checkpoint, |
| 844 | correction_cache=correction_cache, |
| 845 | kg_cache=kg_cache, |
| 846 | query_start_loc=query_start_loc, |
| 847 | state_indices=state_indices, |
| 848 | spec_query_len=query_len, |
| 849 | ) |
| 850 | |
| 851 | normalized_q = q.float() * torch.rsqrt( |
| 852 | q.float().square().sum(dim=-1, keepdim=True) + 1e-6 |
| 853 | ) |
| 854 | normalized_k = k.float() * torch.rsqrt( |
| 855 | k.float().square().sum(dim=-1, keepdim=True) + 1e-6 |
| 856 | ) |
| 857 | gate_input = raw_g.float() + dt_bias.view(1, 1, num_heads, dim) |
| 858 | if lower_bound is None: |
| 859 | gate = -A_log.exp().view(1, 1, num_heads, 1) * F.softplus(gate_input) |
| 860 | else: |
| 861 | gate = lower_bound * torch.sigmoid( |
| 862 | A_log.exp().view(1, 1, num_heads, 1) * gate_input |
| 863 | ) |
| 864 | beta = raw_beta.float().sigmoid() |
| 865 | |
| 866 | reference_output = [] |
| 867 | committed_states = checkpoint.clone() |
| 868 | for seq_idx, commit_len in enumerate(accepted): |
| 869 | start = seq_idx * query_len |
| 870 | end = start + query_len |
| 871 | output, _ = naive_recurrent_kda( |
| 872 | normalized_q[:, start:end], |
| 873 | normalized_k[:, start:end], |
| 874 | v[:, start:end], |
| 875 | gate[:, start:end], |
| 876 | beta[:, start:end], |
| 877 | initial_state=checkpoint[state_indices[seq_idx]].transpose(-1, -2), |
| 878 | ) |
| 879 | reference_output.append(output) |
| 880 | _, committed_state = naive_recurrent_kda( |
| 881 | normalized_q[:, start : start + commit_len], |
| 882 | normalized_k[:, start : start + commit_len], |
| 883 | v[:, start : start + commit_len], |
| 884 | gate[:, start : start + commit_len], |
| 885 | beta[:, start : start + commit_len], |
| 886 | initial_state=checkpoint[state_indices[seq_idx]].transpose(-1, -2), |
| 887 | output_final_state=True, |
| 888 | ) |
| 889 | assert committed_state is not None |
| 890 | final_block = state_indices[seq_idx] |
| 891 | if align_mode: |
| 892 | assert block_table is not None |
| 893 | row = request_indices[seq_idx] if use_request_indices else seq_idx |
| 894 | final_block = block_table[row, (4 + commit_len) // 8] |
| 895 | committed_states[final_block] = committed_state.transpose(-1, -2) |
| 896 | if align_mode and 4 + commit_len >= 8: |
| 897 | _, boundary_state = naive_recurrent_kda( |
| 898 | normalized_q[:, start : start + 4], |
| 899 | normalized_k[:, start : start + 4], |
| 900 | v[:, start : start + 4], |
| 901 | gate[:, start : start + 4], |
| 902 | beta[:, start : start + 4], |
| 903 | initial_state=checkpoint[state_indices[seq_idx]].transpose(-1, -2), |
| 904 | output_final_state=True, |
| 905 | ) |
| 906 | assert boundary_state is not None |
| 907 | assert block_table is not None |
| 908 | row = request_indices[seq_idx] if use_request_indices else seq_idx |
| 909 | committed_states[block_table[row, 0]] = boundary_state.transpose(-1, -2) |
| 910 | expected_outputs.append(torch.cat(reference_output, dim=1)) |
| 911 | expected_states.append(committed_states) |
| 912 | torch.testing.assert_close(checkpoint, initial_states[-1]) |
| 913 | torch.testing.assert_close( |
| 914 | actual_output, |
| 915 | expected_outputs[-1], |
| 916 | atol=3e-2, |
| 917 | rtol=3e-2, |
| 918 | ) |
| 919 | |
| 920 | context = KDARecoverSSMCommitContext.create( |
| 921 | layers, |
| 922 | spec_query_len=query_len, |
| 923 | max_num_reqs=global_num_accepted.shape[0], |
| 924 | ) |
| 925 | context.commit( |
| 926 | global_num_accepted, |
| 927 | state_indices, |
| 928 | query_start_loc, |
| 929 | request_indices=request_indices, |
| 930 | block_table=block_table, |
| 931 | num_computed_tokens=num_computed_tokens, |
| 932 | mamba_block_size=mamba_block_size, |
| 933 | ) |
| 934 | |
| 935 | for layer_idx, layer in enumerate(layers): |
| 936 | torch.testing.assert_close( |
| 937 | layer.kv_cache[1], |
| 938 | expected_states[layer_idx], |
| 939 | atol=3e-3, |
| 940 | rtol=3e-3, |
| 941 | ) |
| 942 | for seq_idx, commit_len in enumerate(accepted): |
| 943 | block = state_indices[seq_idx] |
| 944 | if align_mode: |
| 945 | assert block_table is not None |
| 946 | row = request_indices[seq_idx] if use_request_indices else seq_idx |
| 947 | block = block_table[row, (4 + commit_len) // 8] |
| 948 | source_block = state_indices[seq_idx] if align_mode else block |
| 949 | if conv_state_dim_first: |
| 950 | actual_conv = layer.kv_cache[0][block, :, :history_len] |
| 951 | expected_conv = initial_conv_states[layer_idx][ |
| 952 | source_block, |
| 953 | :, |
| 954 | commit_len - 1 : commit_len - 1 + history_len, |
| 955 | ] |
| 956 | else: |
| 957 | actual_conv = layer.kv_cache[0][block, :history_len] |
| 958 | expected_conv = initial_conv_states[layer_idx][ |
| 959 | source_block, |
| 960 | commit_len - 1 : commit_len - 1 + history_len, |
| 961 | ] |
| 962 | torch.testing.assert_close(actual_conv, expected_conv) |
| 963 | if align_mode and 4 + commit_len >= 8: |
| 964 | assert block_table is not None |
| 965 | boundary_block = block_table[row, 0] |
| 966 | if conv_state_dim_first: |
| 967 | actual_boundary_conv = layer.kv_cache[0][ |
| 968 | boundary_block, :, :history_len |
| 969 | ] |
| 970 | expected_boundary_conv = initial_conv_states[layer_idx][ |
| 971 | state_indices[seq_idx], :, 3 : 3 + history_len |
| 972 | ] |
| 973 | else: |
| 974 | actual_boundary_conv = layer.kv_cache[0][ |
| 975 | boundary_block, :history_len |
| 976 | ] |
| 977 | expected_boundary_conv = initial_conv_states[layer_idx][ |
| 978 | state_indices[seq_idx], 3 : 3 + history_len |
| 979 | ] |
| 980 | torch.testing.assert_close(actual_boundary_conv, expected_boundary_conv) |
| 981 | |
| 982 | |
| 983 | @pytest.mark.parametrize( |
| 984 | ("num_heads", "num_seqs", "lower_bound", "fuse_output_norm", "conv_layout"), |
| 985 | [ |
| 986 | (12, 1, -5.0, True, "SD"), |
| 987 | (12, 4, None, False, "SD"), |
| 988 | (24, 4, None, False, "SD"), |
| 989 | (48, 1, -5.0, True, "SD"), |
| 990 | (96, 1, -5.0, True, "SD"), |
| 991 | (12, 1, -5.0, True, "DS"), |
| 992 | (12, 4, None, False, "DS"), |
| 993 | (24, 4, None, False, "DS"), |
| 994 | (48, 1, -5.0, True, "DS"), |
| 995 | (96, 1, -5.0, True, "DS"), |
| 996 | ], |
| 997 | ) |
| 998 | @pytest.mark.parametrize("decode_backend", ["native", "flashinfer"]) |
| 999 | @torch.inference_mode() |
| 1000 | def test_fused_kda_decode_correctness( |
| 1001 | num_heads: int, |
| 1002 | num_seqs: int, |
| 1003 | lower_bound: float | None, |
| 1004 | fuse_output_norm: bool, |
| 1005 | conv_layout: str, |
| 1006 | decode_backend: str, |
| 1007 | ): |
| 1008 | D, W = 128, 4 |
| 1009 | state_dtype = torch.bfloat16 if decode_backend == "flashinfer" else torch.float32 |
| 1010 | if decode_backend == "flashinfer": |
| 1011 | if conv_layout == "DS": |
| 1012 | pytest.skip("FlashInfer fused decode requires SD conv-state layout") |
| 1013 | if not fuse_output_norm: |
| 1014 | pytest.skip("FlashInfer's fused decode always applies output norm") |
| 1015 | supported = is_flashinfer_fused_kda_decode_supported( |
| 1016 | num_heads, |
| 1017 | D, |
| 1018 | W, |
| 1019 | num_spec=0, |
| 1020 | input_dtype=torch.bfloat16, |
| 1021 | conv_state_dtype=torch.bfloat16, |
| 1022 | recurrent_state_dtype=state_dtype, |
| 1023 | ) |
| 1024 | else: |
| 1025 | supported = is_fused_kda_decode_supported( |
| 1026 | num_heads, |
| 1027 | D, |
| 1028 | W, |
| 1029 | num_spec=0, |
| 1030 | input_dtype=torch.bfloat16, |
| 1031 | conv_state_dtype=torch.bfloat16, |
| 1032 | recurrent_state_dtype=state_dtype, |
| 1033 | ) |
| 1034 | if not supported: |
| 1035 | pytest.skip(f"{decode_backend} fused KDA decode is not supported") |
| 1036 | torch.manual_seed(967 + num_heads + num_seqs + (conv_layout == "DS")) |
| 1037 | dim = num_heads * D |
| 1038 | slots = num_seqs + 2 |
| 1039 | packed_x_storage = torch.randn( |
| 1040 | num_seqs, 3 * dim + 17, dtype=torch.bfloat16, device=DEVICE |
| 1041 | ) |
| 1042 | packed_x = packed_x_storage[:, : 3 * dim] |
| 1043 | weight = 0.1 * torch.randn(3 * dim, W, dtype=torch.float32, device=DEVICE) |
| 1044 | if conv_layout == "DS": |
| 1045 | # DS cache layout: per slot the taps are innermost |
| 1046 | # (stride (W-1, 1)), matching VLLM_SSM_CONV_STATE_LAYOUT=DS. |
| 1047 | conv_seed = 0.1 * torch.randn( |
| 1048 | slots, |
| 1049 | 3 * dim, |
| 1050 | W - 1, |
| 1051 | dtype=torch.bfloat16, |
| 1052 | device=DEVICE, |
| 1053 | ) |
| 1054 | else: |
| 1055 | # SD cache layout: per slot the channels are innermost. |
| 1056 | conv_seed = 0.1 * torch.randn( |
| 1057 | slots, |
| 1058 | W - 1, |
| 1059 | 3 * dim, |
| 1060 | dtype=torch.bfloat16, |
| 1061 | device=DEVICE, |
| 1062 | ).transpose(1, 2) |
| 1063 | raw_g = torch.randn( |
| 1064 | 1, |
| 1065 | num_seqs, |
| 1066 | num_heads, |
| 1067 | D, |
| 1068 | dtype=torch.bfloat16, |
| 1069 | device=DEVICE, |
| 1070 | ) |
| 1071 | raw_beta_storage = torch.randn( |
| 1072 | 1, |
| 1073 | num_seqs, |
| 1074 | num_heads + 1, |
| 1075 | dtype=torch.bfloat16, |
| 1076 | device=DEVICE, |
| 1077 | ) |
| 1078 | raw_beta = raw_beta_storage[:, :, :num_heads] |
| 1079 | output_gate_storage = torch.randn( |
| 1080 | num_seqs, |
| 1081 | dim + 7, |
| 1082 | dtype=torch.bfloat16, |
| 1083 | device=DEVICE, |
| 1084 | ) |
| 1085 | output_gate = output_gate_storage[:, :dim].view(num_seqs, num_heads, D) |
| 1086 | norm_weight = torch.randn(D, dtype=torch.float32, device=DEVICE) |
| 1087 | norm_eps = 1e-5 |
| 1088 | A_log = 0.5 * torch.randn(num_heads, dtype=torch.float32, device=DEVICE) |
| 1089 | dt_bias = 0.1 * torch.randn(dim, dtype=torch.float32, device=DEVICE) |
| 1090 | state_indices = torch.arange( |
| 1091 | num_seqs, |
| 1092 | 0, |
| 1093 | -1, |
| 1094 | dtype=torch.int32, |
| 1095 | device=DEVICE, |
| 1096 | ) |
| 1097 | state_seed = 0.01 * torch.randn( |
| 1098 | slots, |
| 1099 | num_heads, |
| 1100 | D, |
| 1101 | D, |
| 1102 | dtype=torch.float32, |
| 1103 | device=DEVICE, |
| 1104 | ).to(state_dtype) |
| 1105 | |
| 1106 | conv_ref = conv_seed.clone() |
| 1107 | state_ref = state_seed.clone() |
| 1108 | mixed_qkv = causal_conv1d_update( |
| 1109 | packed_x, |
| 1110 | conv_ref, |
| 1111 | weight, |
| 1112 | activation="silu", |
| 1113 | conv_state_indices=state_indices, |
| 1114 | validate_data=True, |
| 1115 | out=torch.empty_like(packed_x), |
| 1116 | ) |
| 1117 | expected, _ = fused_recurrent_kda_packed_decode( |
| 1118 | mixed_qkv=mixed_qkv, |
| 1119 | raw_g=raw_g, |
| 1120 | raw_beta=raw_beta, |
| 1121 | A_log=A_log, |
| 1122 | dt_bias=dt_bias, |
| 1123 | lower_bound=lower_bound, |
| 1124 | initial_state=state_ref, |
| 1125 | state_indices=state_indices, |
| 1126 | ) |
| 1127 | if fuse_output_norm: |
| 1128 | expected_float = expected.float() |
| 1129 | expected = ( |
| 1130 | expected_float |
| 1131 | * torch.rsqrt(expected_float.square().mean(dim=-1, keepdim=True) + norm_eps) |
| 1132 | * norm_weight |
| 1133 | * output_gate.float().sigmoid().unsqueeze(0) |
| 1134 | ).to(expected.dtype) |
| 1135 | |
| 1136 | conv_slot_elements = 3 * dim * (W - 1) |
| 1137 | state_slot_elements = num_heads * D * D |
| 1138 | conv_slot_bytes = conv_slot_elements * torch.bfloat16.itemsize |
| 1139 | page_bytes = conv_slot_bytes + state_slot_elements * state_dtype.itemsize |
| 1140 | cache_storage = torch.empty(slots * page_bytes, dtype=torch.uint8, device=DEVICE) |
| 1141 | conv_actual = torch.as_strided( |
| 1142 | cache_storage.view(torch.bfloat16), |
| 1143 | size=(slots, 3 * dim, W - 1), |
| 1144 | stride=( |
| 1145 | page_bytes // torch.bfloat16.itemsize, |
| 1146 | (W - 1) if conv_layout == "DS" else 1, |
| 1147 | 1 if conv_layout == "DS" else 3 * dim, |
| 1148 | ), |
| 1149 | ) |
| 1150 | state_actual = torch.as_strided( |
| 1151 | cache_storage.view(state_dtype), |
| 1152 | size=(slots, num_heads, D, D), |
| 1153 | stride=(page_bytes // state_dtype.itemsize, D * D, D, 1), |
| 1154 | storage_offset=conv_slot_bytes // state_dtype.itemsize, |
| 1155 | ) |
| 1156 | conv_actual.copy_(conv_seed) |
| 1157 | state_actual.copy_(state_seed) |
| 1158 | fused_weight = weight.reshape(3, dim, W).transpose(1, 2).contiguous() |
| 1159 | if decode_backend == "flashinfer": |
| 1160 | output = torch.empty( |
| 1161 | 1, |
| 1162 | num_seqs, |
| 1163 | num_heads, |
| 1164 | D, |
| 1165 | dtype=torch.bfloat16, |
| 1166 | device=DEVICE, |
| 1167 | ) |
| 1168 | |
| 1169 | actual = flashinfer_fused_kda_decode( |
| 1170 | x=packed_x, |
| 1171 | weight=fused_weight, |
| 1172 | conv_state=conv_actual, |
| 1173 | raw_gate=raw_g, |
| 1174 | raw_beta=raw_beta, |
| 1175 | A_log=A_log, |
| 1176 | dt_bias=dt_bias, |
| 1177 | state_indices=state_indices, |
| 1178 | state=state_actual, |
| 1179 | output_gate=output_gate, |
| 1180 | norm_weight=norm_weight, |
| 1181 | lower_bound=lower_bound, |
| 1182 | norm_eps=norm_eps, |
| 1183 | output=output, |
| 1184 | ) |
| 1185 | else: |
| 1186 | actual = ops.fused_kda_decode( |
| 1187 | x=packed_x, |
| 1188 | weight=fused_weight, |
| 1189 | bias=None, |
| 1190 | conv_state=conv_actual, |
| 1191 | raw_g=raw_g, |
| 1192 | raw_beta=raw_beta, |
| 1193 | A_log=A_log, |
| 1194 | dt_bias=dt_bias, |
| 1195 | state_indices=state_indices, |
| 1196 | state=state_actual, |
| 1197 | lower_bound=lower_bound, |
| 1198 | output_gate=output_gate if fuse_output_norm else None, |
| 1199 | norm_weight=norm_weight if fuse_output_norm else None, |
| 1200 | norm_eps=norm_eps, |
| 1201 | ) |
| 1202 | |
| 1203 | torch.testing.assert_close(actual, expected, atol=3e-2, rtol=3e-2) |
| 1204 | torch.testing.assert_close(conv_actual, conv_ref, atol=0, rtol=0) |
| 1205 | torch.testing.assert_close(state_actual, state_ref, atol=3e-2, rtol=3e-2) |
| 1206 | |
| 1207 | |
| 1208 | @torch.inference_mode() |
| 1209 | @pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16]) |
| 1210 | def test_flashinfer_fused_kda_spec_decode_integration( |
| 1211 | monkeypatch: pytest.MonkeyPatch, |
| 1212 | state_dtype: torch.dtype, |
| 1213 | ): |
| 1214 | if not is_flashinfer_fused_kda_spec_decode_supported( |
| 1215 | num_heads=12, |
| 1216 | head_dim=128, |
| 1217 | conv_width=4, |
| 1218 | num_spec=6, |
| 1219 | input_dtype=torch.bfloat16, |
| 1220 | conv_state_dtype=torch.bfloat16, |
| 1221 | recurrent_state_dtype=state_dtype, |
| 1222 | lower_bound=-5.0, |
| 1223 | use_recoverssm=False, |
| 1224 | ): |
| 1225 | pytest.skip("FlashInfer packed fused KDA decode is not supported") |
| 1226 | |
| 1227 | generator = torch.Generator(device=DEVICE).manual_seed(20260820) |
| 1228 | H, D, W, T = 12, 128, 4, 7 |
| 1229 | hidden_size = H * D |
| 1230 | num_rows = 10 |
| 1231 | num_slots = 16 |
| 1232 | |
| 1233 | x = torch.randn( |
| 1234 | num_rows, |
| 1235 | 3 * hidden_size, |
| 1236 | dtype=torch.bfloat16, |
| 1237 | device=DEVICE, |
| 1238 | generator=generator, |
| 1239 | ) |
| 1240 | conv_weight = 0.1 * torch.randn( |
| 1241 | 3 * hidden_size, |
| 1242 | W, |
| 1243 | dtype=torch.float32, |
| 1244 | device=DEVICE, |
| 1245 | generator=generator, |
| 1246 | ) |
| 1247 | fused_weight = conv_weight.view(3, hidden_size, W).transpose(1, 2).contiguous() |
| 1248 | raw_gate = torch.randn( |
| 1249 | 1, num_rows, H, D, dtype=torch.bfloat16, device=DEVICE, generator=generator |
| 1250 | ) |
| 1251 | raw_beta = torch.randn( |
| 1252 | 1, num_rows, H, dtype=torch.bfloat16, device=DEVICE, generator=generator |
| 1253 | ) |
| 1254 | output_gate = torch.randn( |
| 1255 | num_rows, H, D, dtype=torch.bfloat16, device=DEVICE, generator=generator |
| 1256 | ) |
| 1257 | A_log = torch.log(torch.arange(1, H + 1, dtype=torch.float32, device=DEVICE)) |
| 1258 | dt_bias = torch.empty(hidden_size, dtype=torch.float32, device=DEVICE).uniform_( |
| 1259 | -7.0, -4.0, generator=generator |
| 1260 | ) |
| 1261 | norm_weight = torch.randn( |
| 1262 | D, dtype=torch.float32, device=DEVICE, generator=generator |
| 1263 | ) |
| 1264 | norm_eps = 1e-5 |
| 1265 | state_indices = torch.tensor( |
| 1266 | [ |
| 1267 | list(range(1, T + 1)), |
| 1268 | list(range(T + 1, 2 * T + 1)), |
| 1269 | [0] * T, |
| 1270 | ], |
| 1271 | dtype=torch.int32, |
| 1272 | device=DEVICE, |
| 1273 | ) |
| 1274 | query_start_loc = torch.tensor([0, 7, 10, 10], dtype=torch.int32, device=DEVICE) |
| 1275 | num_accepted_tokens = torch.tensor([1, 4, 1], dtype=torch.int32, device=DEVICE) |
| 1276 | conv_seed = ( |
| 1277 | 0.5 |
| 1278 | * torch.randn( |
| 1279 | num_slots, |
| 1280 | T + W - 2, |
| 1281 | 3 * hidden_size, |
| 1282 | dtype=torch.bfloat16, |
| 1283 | device=DEVICE, |
| 1284 | generator=generator, |
| 1285 | ) |
| 1286 | ).transpose(1, 2) |
| 1287 | state_seed = ( |
| 1288 | 0.5 |
| 1289 | * torch.randn( |
| 1290 | num_slots, |
| 1291 | H, |
| 1292 | D, |
| 1293 | D, |
| 1294 | dtype=torch.float32, |
| 1295 | device=DEVICE, |
| 1296 | generator=generator, |
| 1297 | ) |
| 1298 | ).to(state_dtype) |
| 1299 | metadata = KimiK3KDAMetadata( |
| 1300 | num_prefills=0, |
| 1301 | num_prefill_tokens=0, |
| 1302 | num_decodes=0, |
| 1303 | num_decode_tokens=0, |
| 1304 | num_spec_decodes=2, |
| 1305 | num_spec_decode_tokens=num_rows, |
| 1306 | num_actual_tokens=num_rows, |
| 1307 | spec_query_start_loc=query_start_loc, |
| 1308 | spec_state_indices_tensor=state_indices, |
| 1309 | num_accepted_tokens=num_accepted_tokens, |
| 1310 | ) |
| 1311 | monkeypatch.setattr( |
| 1312 | "vllm.models.kimi_k3.nvidia.kda.get_forward_context", |
| 1313 | lambda: SimpleNamespace(attn_metadata={"test.layer": metadata}), |
| 1314 | ) |
| 1315 | |
| 1316 | def clone_strided(tensor: torch.Tensor) -> torch.Tensor: |
| 1317 | clone = torch.empty_strided( |
| 1318 | tensor.shape, tensor.stride(), dtype=tensor.dtype, device=tensor.device |
| 1319 | ) |
| 1320 | clone.copy_(tensor) |
| 1321 | return clone |
| 1322 | |
| 1323 | def composed_reference( |
| 1324 | conv_state: torch.Tensor, |
| 1325 | recurrent_state: torch.Tensor, |
| 1326 | ) -> torch.Tensor: |
| 1327 | mixed_qkv = causal_conv1d_update( |
| 1328 | x, |
| 1329 | conv_state, |
| 1330 | conv_weight, |
| 1331 | bias=None, |
| 1332 | activation="silu", |
| 1333 | conv_state_indices=state_indices[:, 0], |
| 1334 | num_accepted_tokens=num_accepted_tokens, |
| 1335 | query_start_loc=query_start_loc, |
| 1336 | max_query_len=T, |
| 1337 | validate_data=False, |
| 1338 | out=torch.empty_like(x), |
| 1339 | ) |
| 1340 | q, k, v = mixed_qkv.view(num_rows, 3, H, D).unbind(1) |
| 1341 | recurrent_out, _ = fused_recurrent_kda_nvidia( |
| 1342 | q=q.unsqueeze(0), |
| 1343 | k=k.unsqueeze(0), |
| 1344 | v=v.unsqueeze(0), |
| 1345 | raw_g=raw_gate, |
| 1346 | raw_beta=raw_beta, |
| 1347 | A_log=A_log, |
| 1348 | dt_bias=dt_bias, |
| 1349 | lower_bound=-5.0, |
| 1350 | initial_state=recurrent_state, |
| 1351 | cu_seqlens=query_start_loc, |
| 1352 | ssm_state_indices=state_indices, |
| 1353 | num_accepted_tokens=num_accepted_tokens, |
| 1354 | ) |
| 1355 | recurrent_float = recurrent_out.float() |
| 1356 | return ( |
| 1357 | recurrent_float |
| 1358 | * torch.rsqrt( |
| 1359 | recurrent_float.square().mean(dim=-1, keepdim=True) + norm_eps |
| 1360 | ) |
| 1361 | * norm_weight |
| 1362 | * output_gate.float().sigmoid().unsqueeze(0) |
| 1363 | ).to(torch.bfloat16) |
| 1364 | |
| 1365 | def run_flashinfer( |
| 1366 | conv_state: torch.Tensor, |
| 1367 | recurrent_state: torch.Tensor, |
| 1368 | output: torch.Tensor, |
| 1369 | ) -> torch.Tensor: |
| 1370 | layer = SimpleNamespace( |
| 1371 | prefix="test.layer", |
| 1372 | kv_cache=(conv_state.transpose(-1, -2), recurrent_state), |
| 1373 | kda_spec_decode_backend="flashinfer", |
| 1374 | decode_conv1d_weight=fused_weight, |
| 1375 | decode_norm_weight=norm_weight, |
| 1376 | gate_lower_bound=-5.0, |
| 1377 | A_log=A_log, |
| 1378 | dt_bias=dt_bias, |
| 1379 | o_norm=SimpleNamespace(eps=norm_eps), |
| 1380 | ) |
| 1381 | KimiK3DeltaAttention._forward( |
| 1382 | layer, |
| 1383 | mixed_qkv=x, |
| 1384 | g1=raw_gate, |
| 1385 | g2=output_gate, |
| 1386 | beta=raw_beta, |
| 1387 | core_attn_out=output, |
| 1388 | ) |
| 1389 | return output |
| 1390 | |
| 1391 | expected_conv = clone_strided(conv_seed) |
| 1392 | expected_state = state_seed.clone() |
| 1393 | expected = composed_reference(expected_conv, expected_state) |
| 1394 | actual_conv = clone_strided(conv_seed) |
| 1395 | actual_state = state_seed.clone() |
| 1396 | output = torch.empty_like(expected) |
| 1397 | actual = run_flashinfer(actual_conv, actual_state, output) |
| 1398 | |
| 1399 | assert actual.data_ptr() == output.data_ptr() |
| 1400 | torch.testing.assert_close(actual_conv, expected_conv, atol=0, rtol=0) |
| 1401 | torch.testing.assert_close(actual_state, expected_state, atol=2e-3, rtol=3e-2) |
| 1402 | torch.testing.assert_close(actual, expected, atol=2e-2, rtol=3e-2) |
| 1403 | assert torch.count_nonzero(actual[:, :7]) > 0 |
| 1404 | torch.testing.assert_close(actual_state[0], state_seed[0], atol=0, rtol=0) |
| 1405 | |
| 1406 | graph = torch.cuda.CUDAGraph() |
| 1407 | with torch.cuda.graph(graph): |
| 1408 | captured = run_flashinfer(actual_conv, actual_state, output) |
| 1409 | |
| 1410 | actual_conv.copy_(conv_seed) |
| 1411 | actual_state.copy_(state_seed) |
| 1412 | query_start_loc.copy_(torch.tensor([0, 6, 10, 10], device=DEVICE)) |
| 1413 | num_accepted_tokens.copy_(torch.tensor([3, 2, 1], device=DEVICE)) |
| 1414 | expected_conv.copy_(conv_seed) |
| 1415 | expected_state.copy_(state_seed) |
| 1416 | expected = composed_reference(expected_conv, expected_state) |
| 1417 | graph.replay() |
| 1418 | torch.accelerator.synchronize() |
| 1419 | |
| 1420 | assert captured.data_ptr() == output.data_ptr() |
| 1421 | torch.testing.assert_close(output, expected, atol=2e-2, rtol=3e-2) |
| 1422 | torch.testing.assert_close(actual_conv, expected_conv, atol=0, rtol=0) |
| 1423 | torch.testing.assert_close(actual_state, expected_state, atol=2e-3, rtol=3e-2) |
| 1424 | |
| 1425 | |
| 1426 | def test_fused_kda_decode_rejects_speculative_conv_state(): |
| 1427 | assert not is_fused_kda_decode_supported( |
| 1428 | num_heads=12, |
| 1429 | head_dim=128, |
| 1430 | conv_width=4, |
| 1431 | num_spec=2, |
| 1432 | input_dtype=torch.bfloat16, |
| 1433 | conv_state_dtype=torch.bfloat16, |
| 1434 | recurrent_state_dtype=torch.float32, |
| 1435 | ) |
| 1436 | |
| 1437 | |
| 1438 | def _make_kda_prefill_inputs( |
| 1439 | state_dtype: torch.dtype, |
| 1440 | *, |
| 1441 | lower_bound: float, |
| 1442 | ) -> SimpleNamespace: |
| 1443 | B, T, H, D = 1, 48, 2, 128 |
| 1444 | torch.manual_seed(11) |
| 1445 | q, k, v, raw_g = [ |
| 1446 | torch.randn(B, T, H, D, dtype=torch.bfloat16, device=DEVICE) for _ in range(4) |
| 1447 | ] |
| 1448 | raw_beta = torch.randn(B, T, H, dtype=torch.bfloat16, device=DEVICE) |
| 1449 | A_log = torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5 |
| 1450 | dt_bias = torch.randn(H, D, dtype=torch.float32, device=DEVICE) * 0.1 |
| 1451 | initial_state = torch.randn(2, H, D, D, dtype=torch.float32, device=DEVICE).to( |
| 1452 | state_dtype |
| 1453 | ) |
| 1454 | cu_seqlens = torch.tensor([0, 17, T], dtype=torch.int32, device=DEVICE) |
| 1455 | return SimpleNamespace( |
| 1456 | q=q, |
| 1457 | k=k, |
| 1458 | v=v, |
| 1459 | raw_g=raw_g, |
| 1460 | raw_beta=raw_beta, |
| 1461 | A_log=A_log, |
| 1462 | dt_bias=dt_bias, |
| 1463 | initial_state=initial_state, |
| 1464 | cu_seqlens=cu_seqlens, |
| 1465 | lower_bound=lower_bound, |
| 1466 | ) |
| 1467 | |
| 1468 | |
| 1469 | def _require_kda_prefill_backend( |
| 1470 | backend: str, |
| 1471 | state_dtype: torch.dtype, |
| 1472 | lower_bound: float, |
| 1473 | ) -> None: |
| 1474 | if backend == "flashinfer": |
| 1475 | supported = is_flashinfer_recurrent_kda_prefill_supported( |
| 1476 | 128, torch.bfloat16, state_dtype, lower_bound |
| 1477 | ) |
| 1478 | else: |
| 1479 | assert backend == "flashkda" |
| 1480 | supported = is_flashkda_supported(128, torch.bfloat16, state_dtype, lower_bound) |
| 1481 | if not supported: |
| 1482 | pytest.skip(f"{backend} KDA prefill is not supported on this platform") |
| 1483 | |
| 1484 | |
| 1485 | def _run_kda_prefill_backend( |
| 1486 | backend: str, |
| 1487 | *, |
| 1488 | q: torch.Tensor, |
| 1489 | k: torch.Tensor, |
| 1490 | v: torch.Tensor, |
| 1491 | raw_g: torch.Tensor, |
| 1492 | raw_beta: torch.Tensor, |
| 1493 | A_log: torch.Tensor, |
| 1494 | dt_bias: torch.Tensor, |
| 1495 | initial_state: torch.Tensor, |
| 1496 | cu_seqlens: torch.Tensor, |
| 1497 | lower_bound: float, |
| 1498 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 1499 | initial_state = initial_state.clone() |
| 1500 | output = torch.empty_like(v) |
| 1501 | if backend == "flashinfer": |
| 1502 | flashinfer_query_start_loc = cu_seqlens.to(torch.int64) |
| 1503 | seq_order = None |
| 1504 | if q.shape[1] > initial_state.shape[0]: |
| 1505 | seq_order = torch.argsort( |
| 1506 | flashinfer_query_start_loc.diff(), descending=True |
| 1507 | ).to(torch.int32) |
| 1508 | return _flashinfer_kda_prefill( |
| 1509 | q=q, |
| 1510 | k=k, |
| 1511 | v=v, |
| 1512 | raw_g=raw_g, |
| 1513 | raw_beta=raw_beta, |
| 1514 | A_log=A_log, |
| 1515 | dt_bias=dt_bias, |
| 1516 | lower_bound=lower_bound, |
| 1517 | initial_state=initial_state, |
| 1518 | cu_seqlens=flashinfer_query_start_loc, |
| 1519 | out=output, |
| 1520 | seq_order=seq_order, |
| 1521 | ) |
| 1522 | |
| 1523 | assert backend == "flashkda" |
| 1524 | import vllm._flashkda_C # noqa: F401 |
| 1525 | |
| 1526 | final_state = torch.empty_like(initial_state) |
| 1527 | workspace = torch.empty( |
| 1528 | torch.ops._flashkda_C.get_workspace_size( |
| 1529 | q.shape[1], q.shape[2], cu_seqlens.numel() - 1 |
| 1530 | ), |
| 1531 | dtype=torch.uint8, |
| 1532 | device=q.device, |
| 1533 | ) |
| 1534 | return _flashkda_prefill( |
| 1535 | q=q, |
| 1536 | k=k, |
| 1537 | v=v, |
| 1538 | g=raw_g, |
| 1539 | beta=raw_beta, |
| 1540 | A_log=A_log, |
| 1541 | dt_bias=dt_bias, |
| 1542 | lower_bound=lower_bound, |
| 1543 | initial_state=initial_state, |
| 1544 | cu_seqlens=cu_seqlens, |
| 1545 | out=output, |
| 1546 | final_state=final_state, |
| 1547 | workspace=workspace, |
| 1548 | ) |
| 1549 | |
| 1550 | |
| 1551 | def _kda_prefill_reference( |
| 1552 | inputs: SimpleNamespace, |
| 1553 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 1554 | gate = inputs.lower_bound * torch.sigmoid( |
| 1555 | inputs.A_log.exp()[None, None, :, None] |
| 1556 | * (inputs.raw_g.float() + inputs.dt_bias[None, None]) |
| 1557 | ) |
| 1558 | beta = inputs.raw_beta.float().sigmoid() |
| 1559 | q_norm = l2norm_fwd(inputs.q.contiguous()) |
| 1560 | k_norm = l2norm_fwd(inputs.k.contiguous()) |
| 1561 | expected_outputs = [] |
| 1562 | expected_states = [] |
| 1563 | for i, (start, end) in enumerate( |
| 1564 | zip(inputs.cu_seqlens[:-1].tolist(), inputs.cu_seqlens[1:].tolist()) |
| 1565 | ): |
| 1566 | output, final_state = naive_recurrent_kda( |
| 1567 | q_norm[:, start:end], |
| 1568 | k_norm[:, start:end], |
| 1569 | inputs.v[:, start:end], |
| 1570 | gate[:, start:end], |
| 1571 | beta[:, start:end], |
| 1572 | initial_state=inputs.initial_state[i : i + 1].transpose(-1, -2), |
| 1573 | output_final_state=True, |
| 1574 | ) |
| 1575 | assert final_state is not None |
| 1576 | expected_outputs.append(output) |
| 1577 | expected_states.append(final_state) |
| 1578 | return ( |
| 1579 | torch.cat(expected_outputs, dim=1), |
| 1580 | torch.cat(expected_states).transpose(-1, -2).contiguous(), |
| 1581 | ) |
| 1582 | |
| 1583 | |
| 1584 | @pytest.mark.parametrize( |
| 1585 | ("backend", "state_dtype"), |
| 1586 | [ |
| 1587 | pytest.param("flashkda", torch.float32, id="flashkda"), |
| 1588 | pytest.param("flashinfer", torch.bfloat16, id="flashinfer"), |
| 1589 | ], |
| 1590 | ) |
| 1591 | @torch.inference_mode() |
| 1592 | def test_kda_prefill_near_collinear_keys_remain_finite( |
| 1593 | backend: str, |
| 1594 | state_dtype: torch.dtype, |
| 1595 | ): |
| 1596 | """Guard against unstable inversion of near-collinear key blocks.""" |
| 1597 | lower_bound = -5.0 |
| 1598 | _require_kda_prefill_backend(backend, state_dtype, lower_bound) |
| 1599 | |
| 1600 | T, H, D = 16384, 1, 128 |
| 1601 | torch.manual_seed(0) |
| 1602 | key = torch.randn(1, 1, H, D, dtype=torch.bfloat16, device=DEVICE) |
| 1603 | qk = key.expand(1, T, H, D).contiguous() |
| 1604 | value_block = torch.randn(1, 16, H, D, dtype=torch.bfloat16, device=DEVICE) |
| 1605 | value = value_block.repeat(1, T // 16, 1, 1) |
| 1606 | raw_gate = torch.full_like(qk, -12.0) |
| 1607 | raw_beta = torch.full((1, T, H), 8.0, dtype=qk.dtype, device=DEVICE) |
| 1608 | A_log = torch.zeros(H, dtype=torch.float32, device=DEVICE) |
| 1609 | dt_bias = torch.zeros(H, D, dtype=torch.float32, device=DEVICE) |
| 1610 | initial_state = torch.zeros(1, H, D, D, dtype=state_dtype, device=DEVICE) |
| 1611 | cu_seqlens = torch.tensor([0, T], dtype=torch.int32, device=DEVICE) |
| 1612 | output, final_state = _run_kda_prefill_backend( |
| 1613 | backend, |
| 1614 | q=qk, |
| 1615 | k=qk, |
| 1616 | v=value, |
| 1617 | raw_g=raw_gate, |
| 1618 | raw_beta=raw_beta, |
| 1619 | A_log=A_log, |
| 1620 | dt_bias=dt_bias, |
| 1621 | initial_state=initial_state, |
| 1622 | cu_seqlens=cu_seqlens, |
| 1623 | lower_bound=lower_bound, |
| 1624 | ) |
| 1625 | |
| 1626 | assert torch.isfinite(output).all() |
| 1627 | assert torch.isfinite(final_state).all() |
| 1628 | |
| 1629 | |
| 1630 | @pytest.mark.parametrize( |
| 1631 | ("backend", "state_dtype", "tolerance"), |
| 1632 | [ |
| 1633 | pytest.param("flashinfer", torch.bfloat16, 0.03, id="flashinfer-bf16"), |
| 1634 | pytest.param("flashkda", torch.bfloat16, 0.03, id="flashkda-bf16"), |
| 1635 | pytest.param("flashkda", torch.float32, 0.01, id="flashkda-fp32"), |
| 1636 | ], |
| 1637 | ) |
| 1638 | @torch.inference_mode() |
| 1639 | def test_kda_prefill_correctness( |
| 1640 | backend: str, |
| 1641 | state_dtype: torch.dtype, |
| 1642 | tolerance: float, |
| 1643 | ): |
| 1644 | lower_bound = -5.0 |
| 1645 | _require_kda_prefill_backend(backend, state_dtype, lower_bound) |
| 1646 | inputs = _make_kda_prefill_inputs(state_dtype, lower_bound=lower_bound) |
| 1647 | expected_out, expected_state = _kda_prefill_reference(inputs) |
| 1648 | actual_out, actual_state = _run_kda_prefill_backend( |
| 1649 | backend, |
| 1650 | **vars(inputs), |
| 1651 | ) |
| 1652 | |
| 1653 | assert_close("o", expected_out, actual_out, tolerance) |
| 1654 | assert_close("ht", expected_state, actual_state, tolerance) |
| 1655 | |
| 1656 | |
| 1657 | @torch.inference_mode() |
| 1658 | def test_flashinfer_kda_prefill_breakable_graph_cross_stream(): |
| 1659 | if not is_flashinfer_recurrent_kda_prefill_supported( |
| 1660 | 128, |
| 1661 | torch.bfloat16, |
| 1662 | torch.bfloat16, |
| 1663 | -5.0, |
| 1664 | ): |
| 1665 | pytest.skip("FlashInfer KDA prefill is not supported on this platform") |
| 1666 | |
| 1667 | B, T, H, D = 1, 8, 12, 128 |
| 1668 | q, k, v, raw_g = [ |
| 1669 | torch.randn(B, T, H, D, dtype=torch.bfloat16, device=DEVICE) for _ in range(4) |
| 1670 | ] |
| 1671 | raw_beta = torch.randn(B, T, H, dtype=torch.bfloat16, device=DEVICE) |
| 1672 | initial_state = torch.randn( |
| 1673 | B, |
| 1674 | H, |
| 1675 | D, |
| 1676 | D, |
| 1677 | dtype=torch.bfloat16, |
| 1678 | device=DEVICE, |
| 1679 | ) |
| 1680 | kwargs = { |
| 1681 | "q": q, |
| 1682 | "k": k, |
| 1683 | "v": v, |
| 1684 | "raw_g": raw_g, |
| 1685 | "raw_beta": raw_beta, |
| 1686 | "A_log": torch.randn(H, dtype=torch.float32, device=DEVICE), |
| 1687 | "dt_bias": torch.randn(H, D, dtype=torch.float32, device=DEVICE), |
| 1688 | "lower_bound": -5.0, |
| 1689 | "initial_state": initial_state, |
| 1690 | "cu_seqlens": torch.tensor([0, T], dtype=torch.int64, device=DEVICE), |
| 1691 | "out": torch.empty_like(q), |
| 1692 | "seq_order": torch.zeros(B, dtype=torch.int32, device=DEVICE), |
| 1693 | } |
| 1694 | |
| 1695 | capture_stream = torch.Stream(device=DEVICE) |
| 1696 | original_stream = torch.accelerator.current_stream() |
| 1697 | torch.accelerator.set_stream(capture_stream) |
| 1698 | try: |
| 1699 | graph_value = torch.zeros(1, device=DEVICE) |
| 1700 | capture = BreakableCUDAGraphCapture() |
| 1701 | with capture: |
| 1702 | graph_value.add_(1) |
| 1703 | capture.add_eager(lambda: _flashinfer_kda_prefill(**kwargs)) |
| 1704 | graph_value.add_(1) |
| 1705 | capture_stream.synchronize() |
| 1706 | finally: |
| 1707 | torch.accelerator.set_stream(original_stream) |
| 1708 | |
| 1709 | torch.testing.assert_close(graph_value, torch.zeros_like(graph_value)) |
| 1710 | capture.replay() |
| 1711 | torch.accelerator.synchronize() |
| 1712 | torch.testing.assert_close(graph_value, torch.full_like(graph_value, 2)) |
| 1713 | |
| 1714 | |
| 1715 | @pytest.mark.parametrize( |
| 1716 | ("state_dtype", "tolerance"), |
| 1717 | [ |
| 1718 | pytest.param(torch.bfloat16, 0.03, id="bf16"), |
| 1719 | pytest.param(torch.float32, 0.01, id="fp32"), |
| 1720 | ], |
| 1721 | ) |
| 1722 | @torch.inference_mode() |
| 1723 | def test_flashkda_checkpoint_correctness(state_dtype: torch.dtype, tolerance: float): |
| 1724 | lower_bound = -3.0 |
| 1725 | _require_kda_prefill_backend("flashkda", state_dtype, lower_bound) |
| 1726 | |
| 1727 | import vllm._flashkda_C # noqa: F401 |
| 1728 | |
| 1729 | inputs = _make_kda_prefill_inputs(state_dtype, lower_bound=lower_bound) |
| 1730 | q, k, v = inputs.q, inputs.k, inputs.v |
| 1731 | raw_g, raw_beta = inputs.raw_g, inputs.raw_beta |
| 1732 | A_log, dt_bias = inputs.A_log, inputs.dt_bias |
| 1733 | initial_state, cu_seqlens = inputs.initial_state, inputs.cu_seqlens |
| 1734 | _, T, H, D = q.shape |
| 1735 | expected_out, expected_state = _kda_prefill_reference(inputs) |
| 1736 | gate = lower_bound * torch.sigmoid( |
| 1737 | A_log.exp()[None, None, :, None] * (raw_g.float() + dt_bias[None, None, :, :]) |
| 1738 | ) |
| 1739 | beta = raw_beta.float().sigmoid() |
| 1740 | q_norm = l2norm_fwd(q.contiguous()) |
| 1741 | k_norm = l2norm_fwd(k.contiguous()) |
| 1742 | _, expected_checkpoint = naive_recurrent_kda( |
| 1743 | q_norm[:, :16], |
| 1744 | k_norm[:, :16], |
| 1745 | v[:, :16], |
| 1746 | gate[:, :16], |
| 1747 | beta[:, :16], |
| 1748 | initial_state=initial_state[0:1].transpose(-1, -2), |
| 1749 | output_final_state=True, |
| 1750 | ) |
| 1751 | assert expected_checkpoint is not None |
| 1752 | expected_checkpoint = expected_checkpoint.transpose(-1, -2).contiguous() |
| 1753 | |
| 1754 | workspace = torch.empty( |
| 1755 | torch.ops._flashkda_C.get_workspace_size(T, H, cu_seqlens.numel() - 1), |
| 1756 | dtype=torch.uint8, |
| 1757 | device=DEVICE, |
| 1758 | ) |
| 1759 | checkpoint_out = torch.empty_like(v) |
| 1760 | checkpoint_final_state = torch.empty_like(initial_state) |
| 1761 | checkpoint_state = torch.empty_like(initial_state) |
| 1762 | checkpoint_offsets = torch.tensor([16, 31], dtype=torch.int32, device=DEVICE) |
| 1763 | _flashkda_prefill( |
| 1764 | q=q, |
| 1765 | k=k, |
| 1766 | v=v, |
| 1767 | g=raw_g, |
| 1768 | beta=raw_beta, |
| 1769 | A_log=A_log, |
| 1770 | dt_bias=dt_bias, |
| 1771 | lower_bound=lower_bound, |
| 1772 | initial_state=initial_state, |
| 1773 | cu_seqlens=cu_seqlens, |
| 1774 | out=checkpoint_out, |
| 1775 | final_state=checkpoint_final_state, |
| 1776 | workspace=workspace, |
| 1777 | checkpoint_state=checkpoint_state, |
| 1778 | checkpoint_offsets=checkpoint_offsets, |
| 1779 | ) |
| 1780 | |
| 1781 | assert_close("checkpoint_o", expected_out, checkpoint_out, tolerance) |
| 1782 | assert_close("checkpoint_ht", expected_state, checkpoint_final_state, tolerance) |
| 1783 | assert_close("checkpoint", expected_checkpoint, checkpoint_state[:1], tolerance) |
| 1784 | |
| 1785 | conv_state = torch.zeros(2, H * D, 3, dtype=q.dtype, device=DEVICE) |
| 1786 | recurrent_storage = torch.zeros(2, H * D * D + 8, dtype=state_dtype, device=DEVICE) |
| 1787 | recurrent_state = recurrent_storage[:, : H * D * D].view(2, H, D, D) |
| 1788 | conv_input = q[0].flatten(1) |
| 1789 | checkpoint_state_indices = torch.tensor( |
| 1790 | [1, NULL_BLOCK_ID], dtype=torch.int32, device=DEVICE |
| 1791 | ) |
| 1792 | FlashKDAPrefillCheckpointExporter().export( |
| 1793 | MambaPrefillCheckpointMetadata(checkpoint_offsets, checkpoint_state_indices), |
| 1794 | raw_qkv=conv_input, |
| 1795 | conv_state=conv_state, |
| 1796 | recurrent_checkpoint=checkpoint_state, |
| 1797 | recurrent_state=recurrent_state, |
| 1798 | cu_seqlens=cu_seqlens, |
| 1799 | ) |
| 1800 | torch.testing.assert_close(conv_state[1], q[0, 13:16].flatten(1).transpose(0, 1)) |
| 1801 | torch.testing.assert_close(recurrent_state[1], checkpoint_state[0]) |