@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 54146a9bf9
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 4, 11:15 UTC

[Bugfix] correct h matrix layout in chunk_kda output kernel (#40956)

Signed-off-by: ChenxiQian <chenxi.qian.cq@outlook.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Chenxi Qian committed Apr 30, 2026 at 16:22 UTC 54146a9bf951b8c70ad85fb1a1bee241964209e0
3 files changed +173 -7
.buildkite/test_areas/kernels.yaml
+10
@@ -89,6 +89,16 @@ steps:
89 commands:
90 - pytest -v -s kernels/mamba
91
92 +- label: Kernels KDA Test
93 + timeout_in_minutes: 20
94 + source_file_dependencies:
95 + - vllm/model_executor/layers/fla/ops/kda.py
96 + - vllm/model_executor/layers/fla/ops/chunk_delta_h.py
97 + - vllm/model_executor/layers/fla/ops/l2norm.py
98 + - tests/kernels/test_kda.py
99 + commands:
100 + - pytest -v -s kernels/test_kda.py
101 +
102 - label: Kernels DeepGEMM Test (H100)
103 key: kernels-deepgemm-test-h100
104 timeout_in_minutes: 45
tests/kernels/test_kda.py new
+157
@@ -0,0 +1,157 @@
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 +import pytest
10 +import torch
11 +import torch.nn.functional as F
12 +
13 +from vllm.model_executor.layers.fla.ops.kda import chunk_kda
14 +from vllm.model_executor.layers.fla.ops.l2norm import l2norm_fwd
15 +
16 +DEVICE = "cuda"
17 +
18 +
19 +def naive_recurrent_kda(
20 + q: torch.Tensor,
21 + k: torch.Tensor,
22 + v: torch.Tensor,
23 + g: torch.Tensor,
24 + beta: torch.Tensor,
25 + scale: float | None = None,
26 + initial_state: torch.Tensor | None = None,
27 + output_final_state: bool = False,
28 +) -> tuple[torch.Tensor, torch.Tensor | None]:
29 + """Naive recurrent KDA reference, ported from FLA's naive.py."""
30 + dtype = v.dtype
31 + B, T, H, K = q.shape
32 + V = v.shape[-1]
33 + if scale is None:
34 + scale = K**-0.5
35 +
36 + q, k, v, g, beta = (x.to(torch.float) for x in [q, k, v, g, beta])
37 + q = q * scale
38 +
39 + S = k.new_zeros(B, H, K, V).to(q)
40 + if initial_state is not None:
41 + S += initial_state
42 + o = torch.zeros_like(v)
43 + for i in range(T):
44 + q_i, k_i, v_i, g_i, b_i = q[:, i], k[:, i], v[:, i], g[:, i], beta[:, i]
45 + S = S * g_i[..., None].exp()
46 + S = S + torch.einsum(
47 + "bhk,bhv->bhkv",
48 + b_i[..., None] * k_i,
49 + v_i - (k_i[..., None] * S).sum(-2),
50 + )
51 + o[:, i] = torch.einsum("bhk,bhkv->bhv", q_i, S)
52 + if not output_final_state:
53 + S = None
54 + return o.to(dtype), S
55 +
56 +
57 +def assert_close(
58 + name: str,
59 + ref: torch.Tensor,
60 + tri: torch.Tensor,
61 + ratio: float,
62 + err_atol: float = 1e-6,
63 +):
64 + """RMSE-based relative error comparison."""
65 + abs_err = (ref.detach() - tri.detach()).flatten().abs().max().item()
66 + rmse_diff = (ref.detach() - tri.detach()).flatten().square().mean().sqrt().item()
67 + rmse_base = ref.detach().flatten().square().mean().sqrt().item()
68 + rel_err = rmse_diff / (rmse_base + 1e-8)
69 + print(f"{name:>4} | abs={abs_err:.6f} | rmse={rel_err:.6f} | thr={ratio}")
70 + if abs_err <= err_atol:
71 + return
72 + assert not torch.isnan(ref).any(), f"{name}: NaN detected in ref"
73 + assert not torch.isnan(tri).any(), f"{name}: NaN detected in tri"
74 + assert rel_err < ratio, (
75 + f"{name}: max abs err {abs_err:.6f}, rmse ratio {rel_err:.6f} >= {ratio}"
76 + )
77 +
78 +
79 +@pytest.mark.parametrize(
80 + ("H", "D", "cu_seqlens", "dtype"),
81 + [
82 + pytest.param(
83 + *test,
84 + id="H{}-D{}-cu{}-{}".format(*test),
85 + )
86 + for test in [
87 + (32, 128, [0, 64], torch.float16),
88 + (32, 128, [0, 1024], torch.float16),
89 + (32, 128, [0, 15], torch.float16),
90 + (32, 128, [0, 256, 512, 768, 1024], torch.float16),
91 + (32, 128, [0, 15, 100, 300, 1200], torch.float16),
92 + (64, 128, [0, 256, 500, 1000], torch.float16),
93 + (32, 128, [0, 8192], torch.float16),
94 + (32, 128, [0, 256, 500, 1000], torch.bfloat16),
95 + ]
96 + ],
97 +)
98 +@torch.inference_mode()
99 +def test_chunk_kda(
100 + H: int,
101 + D: int,
102 + cu_seqlens: list[int],
103 + dtype: torch.dtype,
104 +):
105 + T = cu_seqlens[-1]
106 + torch.manual_seed(42)
107 + B = 1
108 + cu_seqlens_t = torch.LongTensor(cu_seqlens).to(DEVICE)
109 + N = len(cu_seqlens) - 1
110 +
111 + q = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE)
112 + k = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE)
113 + v = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE)
114 + g = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float32, device=DEVICE)).to(
115 + dtype
116 + )
117 + beta = torch.rand(B, T, H, dtype=dtype, device=DEVICE).sigmoid()
118 + h0 = torch.randn(N, H, D, D, dtype=torch.float32, device=DEVICE)
119 +
120 + # Naive reference with l2norm_fwd (same kernel as chunk_kda)
121 + ref_outputs = []
122 + ref_states = []
123 + for i in range(N):
124 + s, e = cu_seqlens[i], cu_seqlens[i + 1]
125 + q_i = l2norm_fwd(q[:, s:e].contiguous())
126 + k_i = l2norm_fwd(k[:, s:e].contiguous())
127 + o_i, ht_i = naive_recurrent_kda(
128 + q_i,
129 + k_i,
130 + v[:, s:e],
131 + g[:, s:e],
132 + beta[:, s:e],
133 + initial_state=h0[i],
134 + output_final_state=True,
135 + )
136 + ref_outputs.append(o_i)
137 + ref_states.append(ht_i)
138 + ref_o = torch.cat(ref_outputs, dim=1)
139 + ref_ht = torch.cat(ref_states, dim=0)
140 +
141 + # h0 transposed to (V, K) layout for the kernel; naive uses (K, V)
142 + tri_o, tri_ht = chunk_kda(
143 + q=q.clone(),
144 + k=k.clone(),
145 + v=v.clone(),
146 + g=g.clone(),
147 + beta=beta.clone(),
148 + initial_state=h0.transpose(-1, -2).contiguous().clone(),
149 + output_final_state=True,
150 + cu_seqlens=cu_seqlens_t,
151 + use_qk_l2norm_in_kernel=True,
152 + )
153 +
154 + assert not torch.isnan(tri_o).any(), "Triton output o contains NaN"
155 + assert not torch.isnan(tri_ht).any(), "Triton output ht contains NaN"
156 + assert_close("o", ref_o, tri_o, 0.005)
157 + assert_close("ht", ref_ht, tri_ht.transpose(-1, -2).contiguous(), 0.005)
vllm/model_executor/layers/fla/ops/kda.py
+6 -7
@@ -1076,10 +1076,10 @@ def chunk_gla_fwd_kernel_o(
1076 )
1077 p_h = tl.make_block_ptr(
1078 h + (i_tg * H + i_h) * K * V,
1079 - (K, V),
1080 - (V, 1),
1081 - (i_k * BK, i_v * BV),
1082 - (BK, BV),
1079 + (V, K),
1080 + (K, 1),
1081 + (i_v * BV, i_k * BK),
1082 + (BV, BK),
1083 (1, 0),
1084 )
1085
@@ -1090,12 +1090,11 @@ def chunk_gla_fwd_kernel_o(
1090 b_g = tl.load(p_g, boundary_check=(0, 1))
1091 # [BT, BK]
1092 b_qg = (b_q * exp(b_g)).to(b_q.dtype)
1093 - # [BK, BV]
1093 + # [BV, BK]
1094 b_h = tl.load(p_h, boundary_check=(0, 1))
1095 - # works but dkw, owing to divine benevolence
1095 # [BT, BV]
1096 if i_k >= 0:
1098 - b_o += tl.dot(b_qg, b_h.to(b_qg.dtype))
1097 + b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype))
1098 p_v = tl.make_block_ptr(
1099 v + (bos * H + i_h) * V,
1100 (T, V),