# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Precision tests for vllm's chunk_kda Triton operator. Compares chunk_kda against a naive recurrent reference (float32). Uses torch.rand for q/k/v to match FLA's test pattern. """ import pytest import torch import torch.nn.functional as F from vllm.model_executor.layers.fla.ops.kda import ( chunk_kda, chunk_kda_with_fused_gate, fused_kda_gate, ) from vllm.model_executor.layers.fla.ops.l2norm import l2norm_fwd DEVICE = "cuda" def naive_recurrent_kda( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, g: torch.Tensor, beta: torch.Tensor, scale: float | None = None, initial_state: torch.Tensor | None = None, output_final_state: bool = False, ) -> tuple[torch.Tensor, torch.Tensor | None]: """Naive recurrent KDA reference, ported from FLA's naive.py.""" dtype = v.dtype B, T, H, K = q.shape V = v.shape[-1] if scale is None: scale = K**-0.5 q, k, v, g, beta = (x.to(torch.float) for x in [q, k, v, g, beta]) q = q * scale S = k.new_zeros(B, H, K, V).to(q) if initial_state is not None: S += initial_state o = torch.zeros_like(v) for i in range(T): q_i, k_i, v_i, g_i, b_i = q[:, i], k[:, i], v[:, i], g[:, i], beta[:, i] S = S * g_i[..., None].exp() S = S + torch.einsum( "bhk,bhv->bhkv", b_i[..., None] * k_i, v_i - (k_i[..., None] * S).sum(-2), ) o[:, i] = torch.einsum("bhk,bhkv->bhv", q_i, S) if not output_final_state: S = None return o.to(dtype), S def assert_close( name: str, ref: torch.Tensor, tri: torch.Tensor, ratio: float, err_atol: float = 1e-6, ): """RMSE-based relative error comparison.""" abs_err = (ref.detach() - tri.detach()).flatten().abs().max().item() rmse_diff = (ref.detach() - tri.detach()).flatten().square().mean().sqrt().item() rmse_base = ref.detach().flatten().square().mean().sqrt().item() rel_err = rmse_diff / (rmse_base + 1e-8) print(f"{name:>4} | abs={abs_err:.6f} | rmse={rel_err:.6f} | thr={ratio}") if abs_err <= err_atol: return assert not torch.isnan(ref).any(), f"{name}: NaN detected in ref" assert not torch.isnan(tri).any(), f"{name}: NaN detected in tri" assert rel_err < ratio, ( f"{name}: max abs err {abs_err:.6f}, rmse ratio {rel_err:.6f} >= {ratio}" ) @pytest.mark.parametrize( ("H", "D", "cu_seqlens", "dtype"), [ pytest.param( *test, id="H{}-D{}-cu{}-{}".format(*test), ) for test in [ (32, 128, [0, 64], torch.float16), (32, 128, [0, 1024], torch.float16), (32, 128, [0, 15], torch.float16), (32, 128, [0, 256, 512, 768, 1024], torch.float16), (32, 128, [0, 15, 100, 300, 1200], torch.float16), (64, 128, [0, 256, 500, 1000], torch.float16), (32, 128, [0, 8192], torch.float16), (32, 128, [0, 256, 500, 1000], torch.bfloat16), ] ], ) @torch.inference_mode() def test_chunk_kda( H: int, D: int, cu_seqlens: list[int], dtype: torch.dtype, ): T = cu_seqlens[-1] torch.manual_seed(42) B = 1 cu_seqlens_t = torch.LongTensor(cu_seqlens).to(DEVICE) N = len(cu_seqlens) - 1 q = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE) k = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE) v = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE) g = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float32, device=DEVICE)).to( dtype ) beta = torch.rand(B, T, H, dtype=dtype, device=DEVICE).sigmoid() h0 = torch.randn(N, H, D, D, dtype=torch.float32, device=DEVICE) # Naive reference with l2norm_fwd (same kernel as chunk_kda) ref_outputs = [] ref_states = [] for i in range(N): s, e = cu_seqlens[i], cu_seqlens[i + 1] q_i = l2norm_fwd(q[:, s:e].contiguous()) k_i = l2norm_fwd(k[:, s:e].contiguous()) o_i, ht_i = naive_recurrent_kda( q_i, k_i, v[:, s:e], g[:, s:e], beta[:, s:e], initial_state=h0[i], output_final_state=True, ) ref_outputs.append(o_i) ref_states.append(ht_i) ref_o = torch.cat(ref_outputs, dim=1) ref_ht = torch.cat(ref_states, dim=0) # h0 transposed to (V, K) layout for the kernel; naive uses (K, V) tri_o, tri_ht = chunk_kda( q=q.clone(), k=k.clone(), v=v.clone(), g=g.clone(), beta=beta.clone(), initial_state=h0.transpose(-1, -2).contiguous().clone(), output_final_state=True, cu_seqlens=cu_seqlens_t, use_qk_l2norm_in_kernel=True, ) assert not torch.isnan(tri_o).any(), "Triton output o contains NaN" assert not torch.isnan(tri_ht).any(), "Triton output ht contains NaN" assert_close("o", ref_o, tri_o, 0.005) assert_close("ht", ref_ht, tri_ht.transpose(-1, -2).contiguous(), 0.005) @pytest.mark.parametrize( ("cu_seqlens", "dtype"), [ ([0, 64], torch.float16), ([0, 15, 100, 300], torch.bfloat16), ], ) @torch.inference_mode() def test_chunk_kda_fused_gate_cumsum_matches_unfused( cu_seqlens: list[int], dtype: torch.dtype, ): H, D = 8, 64 T = cu_seqlens[-1] N = len(cu_seqlens) - 1 torch.manual_seed(123) cu_seqlens_t = torch.tensor(cu_seqlens, dtype=torch.int32, device=DEVICE) q = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE) k = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE) v = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE) raw_g = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE) beta = torch.rand(1, T, H, dtype=dtype, device=DEVICE).sigmoid() A_log = (torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5).contiguous() dt_bias = ( torch.randn(H * D, dtype=torch.float32, device=DEVICE) * 0.1 ).contiguous() h0 = torch.randn(N, H, D, D, dtype=torch.float32, device=DEVICE) initial_state = h0.transpose(-1, -2).contiguous() gate = fused_kda_gate( raw_g.reshape(T, H * D), A_log, D, g_bias=dt_bias, ).unsqueeze(0) old_o, old_ht = chunk_kda( q=q.clone(), k=k.clone(), v=v.clone(), g=gate, beta=beta.clone(), initial_state=initial_state.clone(), output_final_state=True, cu_seqlens=cu_seqlens_t, use_qk_l2norm_in_kernel=True, ) new_o, new_ht = chunk_kda_with_fused_gate( q=q.clone(), k=k.clone(), v=v.clone(), raw_g=raw_g, beta=beta.clone(), A_log=A_log, g_bias=dt_bias, initial_state=initial_state.clone(), output_final_state=True, cu_seqlens=cu_seqlens_t, use_qk_l2norm_in_kernel=True, ) assert_close("o", old_o, new_o, 1e-3, err_atol=1e-3) assert_close("ht", old_ht, new_ht, 1e-3, err_atol=1e-3)