| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for GDN post-convolution kernels. |
| 4 | |
| 5 | The prefill tests cover preparation and the MTP tests cover recurrent state |
| 6 | updates plus output normalization and gating. |
| 7 | """ |
| 8 | |
| 9 | import pytest |
| 10 | import torch |
| 11 | import torch.nn.functional as F |
| 12 | |
| 13 | from vllm import _custom_ops as ops |
| 14 | from vllm.third_party.flash_linear_attention.ops import ( |
| 15 | fused_sigmoid_gating_delta_rule_update, |
| 16 | ) |
| 17 | from vllm.third_party.flash_linear_attention.ops.fused_gdn_prefill_post_conv import ( |
| 18 | fused_post_conv_prep, |
| 19 | ) |
| 20 | from vllm.third_party.flash_linear_attention.ops.layernorm_guard import ( |
| 21 | rmsnorm_fn, |
| 22 | ) |
| 23 | |
| 24 | |
| 25 | def reference_post_conv( |
| 26 | conv_output: torch.Tensor, |
| 27 | a: torch.Tensor, |
| 28 | b: torch.Tensor, |
| 29 | A_log: torch.Tensor, |
| 30 | dt_bias: torch.Tensor, |
| 31 | H: int, |
| 32 | K: int, |
| 33 | V: int, |
| 34 | apply_l2norm: bool = True, |
| 35 | output_g_exp: bool = False, |
| 36 | ): |
| 37 | """Reference implementation using individual ops.""" |
| 38 | L = conv_output.shape[0] |
| 39 | HV = A_log.shape[0] |
| 40 | |
| 41 | # Split |
| 42 | q_flat, k_flat, v_flat = torch.split(conv_output, [H * K, H * K, HV * V], dim=-1) |
| 43 | |
| 44 | # Rearrange + contiguous |
| 45 | q = q_flat.view(L, H, K).contiguous() |
| 46 | k = k_flat.view(L, H, K).contiguous() |
| 47 | v = v_flat.view(L, HV, V).contiguous() |
| 48 | |
| 49 | # L2 norm |
| 50 | if apply_l2norm: |
| 51 | q = F.normalize(q.float(), p=2, dim=-1, eps=1e-6).to(conv_output.dtype) |
| 52 | k = F.normalize(k.float(), p=2, dim=-1, eps=1e-6).to(conv_output.dtype) |
| 53 | |
| 54 | # Gating |
| 55 | x = a.float() + dt_bias.float() |
| 56 | sp = F.softplus(x, beta=1.0, threshold=20.0) |
| 57 | g = -torch.exp(A_log.float()) * sp |
| 58 | |
| 59 | if output_g_exp: |
| 60 | g = torch.exp(g) |
| 61 | |
| 62 | beta_out = torch.sigmoid(b.float()) |
| 63 | |
| 64 | return q, k, v, g, beta_out |
| 65 | |
| 66 | |
| 67 | # Qwen3.5-35B config: H=16, HV=32, K=128, V=128 |
| 68 | # Qwen3.5-397B config: H=16, HV=64, K=128, V=128 |
| 69 | @pytest.mark.parametrize( |
| 70 | "H, HV, K, V", |
| 71 | [ |
| 72 | (16, 32, 128, 128), # 35B |
| 73 | (16, 64, 128, 128), # 397B |
| 74 | (4, 8, 64, 64), # small |
| 75 | ], |
| 76 | ) |
| 77 | @pytest.mark.parametrize("L", [1, 16, 128, 512, 2048]) |
| 78 | @pytest.mark.parametrize("apply_l2norm", [True, False]) |
| 79 | @pytest.mark.parametrize("output_g_exp", [True, False]) |
| 80 | @pytest.mark.parametrize("dtype", [torch.bfloat16]) |
| 81 | def test_fused_post_conv_correctness(H, HV, K, V, L, apply_l2norm, output_g_exp, dtype): |
| 82 | """Test fused kernel matches reference for all configs.""" |
| 83 | torch.manual_seed(42) |
| 84 | device = "cuda" |
| 85 | qkv_dim = 2 * H * K + HV * V |
| 86 | |
| 87 | conv_output = torch.randn(L, qkv_dim, dtype=dtype, device=device) |
| 88 | a = torch.randn(L, HV, dtype=dtype, device=device) |
| 89 | b = torch.randn(L, HV, dtype=dtype, device=device) |
| 90 | A_log = torch.randn(HV, dtype=torch.float32, device=device) - 2.0 |
| 91 | dt_bias = torch.randn(HV, dtype=torch.float32, device=device) * 0.1 |
| 92 | |
| 93 | # Reference |
| 94 | ref_q, ref_k, ref_v, ref_g, ref_beta = reference_post_conv( |
| 95 | conv_output, |
| 96 | a, |
| 97 | b, |
| 98 | A_log, |
| 99 | dt_bias, |
| 100 | H, |
| 101 | K, |
| 102 | V, |
| 103 | apply_l2norm, |
| 104 | output_g_exp, |
| 105 | ) |
| 106 | |
| 107 | # Fused kernel |
| 108 | fused_q, fused_k, fused_v, fused_g, fused_beta = fused_post_conv_prep( |
| 109 | conv_output, |
| 110 | a, |
| 111 | b, |
| 112 | A_log, |
| 113 | dt_bias, |
| 114 | num_k_heads=H, |
| 115 | head_k_dim=K, |
| 116 | head_v_dim=V, |
| 117 | apply_l2norm=apply_l2norm, |
| 118 | output_g_exp=output_g_exp, |
| 119 | ) |
| 120 | |
| 121 | # Check shapes |
| 122 | assert fused_q.shape == (L, H, K), f"q shape: {fused_q.shape}" |
| 123 | assert fused_k.shape == (L, H, K), f"k shape: {fused_k.shape}" |
| 124 | assert fused_v.shape == (L, HV, V), f"v shape: {fused_v.shape}" |
| 125 | assert fused_g.shape == (L, HV), f"g shape: {fused_g.shape}" |
| 126 | assert fused_beta.shape == (L, HV), f"beta shape: {fused_beta.shape}" |
| 127 | |
| 128 | # Check dtypes |
| 129 | assert fused_q.dtype == dtype |
| 130 | assert fused_k.dtype == dtype |
| 131 | assert fused_v.dtype == dtype |
| 132 | assert fused_g.dtype == torch.float32 |
| 133 | assert fused_beta.dtype == torch.float32 |
| 134 | |
| 135 | # Check contiguity |
| 136 | assert fused_q.is_contiguous() |
| 137 | assert fused_k.is_contiguous() |
| 138 | assert fused_v.is_contiguous() |
| 139 | |
| 140 | # Check values |
| 141 | atol_qkv = 1e-2 if apply_l2norm else 1e-3 |
| 142 | rtol_qkv = 1e-2 if apply_l2norm else 1e-3 |
| 143 | |
| 144 | torch.testing.assert_close(fused_q, ref_q, atol=atol_qkv, rtol=rtol_qkv) |
| 145 | torch.testing.assert_close(fused_k, ref_k, atol=atol_qkv, rtol=rtol_qkv) |
| 146 | torch.testing.assert_close(fused_v, ref_v, atol=1e-3, rtol=1e-3) |
| 147 | torch.testing.assert_close(fused_g, ref_g, atol=1e-4, rtol=1e-4) |
| 148 | torch.testing.assert_close(fused_beta, ref_beta, atol=1e-4, rtol=1e-4) |
| 149 | |
| 150 | |
| 151 | @pytest.mark.parametrize("L", [1, 64, 256]) |
| 152 | def test_fused_post_conv_sanity(L): |
| 153 | """Sanity checks: no NaN, unit-norm q/k, beta in (0,1).""" |
| 154 | torch.manual_seed(0) |
| 155 | device = "cuda" |
| 156 | H, HV, K, V = 16, 32, 128, 128 |
| 157 | qkv_dim = 2 * H * K + HV * V |
| 158 | |
| 159 | conv_output = torch.randn(L, qkv_dim, dtype=torch.bfloat16, device=device) |
| 160 | a = torch.randn(L, HV, dtype=torch.bfloat16, device=device) |
| 161 | b = torch.randn(L, HV, dtype=torch.bfloat16, device=device) |
| 162 | A_log = torch.randn(HV, dtype=torch.float32, device=device) - 2.0 |
| 163 | dt_bias = torch.randn(HV, dtype=torch.float32, device=device) |
| 164 | |
| 165 | q, k, v, g, beta = fused_post_conv_prep( |
| 166 | conv_output, |
| 167 | a, |
| 168 | b, |
| 169 | A_log, |
| 170 | dt_bias, |
| 171 | num_k_heads=H, |
| 172 | head_k_dim=K, |
| 173 | head_v_dim=V, |
| 174 | ) |
| 175 | |
| 176 | # Basic sanity |
| 177 | assert not torch.isnan(q).any(), "NaN in q" |
| 178 | assert not torch.isnan(k).any(), "NaN in k" |
| 179 | assert not torch.isnan(v).any(), "NaN in v" |
| 180 | assert not torch.isnan(g).any(), "NaN in g" |
| 181 | assert not torch.isnan(beta).any(), "NaN in beta" |
| 182 | |
| 183 | # L2 norm check: each head vector should have unit norm |
| 184 | q_norms = torch.norm(q.float(), dim=-1) |
| 185 | k_norms = torch.norm(k.float(), dim=-1) |
| 186 | torch.testing.assert_close(q_norms, torch.ones_like(q_norms), atol=1e-3, rtol=1e-3) |
| 187 | torch.testing.assert_close(k_norms, torch.ones_like(k_norms), atol=1e-3, rtol=1e-3) |
| 188 | |
| 189 | # Beta should be in (0, 1) |
| 190 | assert (beta >= 0).all() and (beta <= 1).all(), "beta out of range" |
| 191 | |
| 192 | |
| 193 | def test_fused_post_conv_l0(): |
| 194 | """Test L=0 edge case.""" |
| 195 | device = "cuda" |
| 196 | H, HV, K, V = 16, 32, 128, 128 |
| 197 | qkv_dim = 2 * H * K + HV * V |
| 198 | |
| 199 | conv_output = torch.empty(0, qkv_dim, dtype=torch.bfloat16, device=device) |
| 200 | a = torch.empty(0, HV, dtype=torch.bfloat16, device=device) |
| 201 | b = torch.empty(0, HV, dtype=torch.bfloat16, device=device) |
| 202 | A_log = torch.randn(HV, dtype=torch.float32, device=device) |
| 203 | dt_bias = torch.randn(HV, dtype=torch.float32, device=device) |
| 204 | |
| 205 | q, k, v, g, beta = fused_post_conv_prep( |
| 206 | conv_output, |
| 207 | a, |
| 208 | b, |
| 209 | A_log, |
| 210 | dt_bias, |
| 211 | num_k_heads=H, |
| 212 | head_k_dim=K, |
| 213 | head_v_dim=V, |
| 214 | ) |
| 215 | assert q.shape == (0, H, K) |
| 216 | assert g.shape == (0, HV) |
| 217 | |
| 218 | |
| 219 | @pytest.mark.parametrize( |
| 220 | "head_ratio,tp_size,query_lengths,state_dtype,norm_dtype", |
| 221 | [ |
| 222 | pytest.param(8, 16, (4, 4), torch.bfloat16, torch.bfloat16, id="tp16-bf16"), |
| 223 | pytest.param(8, 4, (4, 4), torch.float32, torch.float32, id="tp4-fp32"), |
| 224 | pytest.param(8, 16, (4, 2, 0), torch.bfloat16, torch.float32, id="tp16-ragged"), |
| 225 | pytest.param(8, 4, (4, 2, 0), torch.float32, torch.bfloat16, id="tp4-ragged"), |
| 226 | pytest.param(8, 16, (8,), torch.float32, torch.bfloat16, id="tp16-max"), |
| 227 | pytest.param(8, 4, (8,), torch.bfloat16, torch.float32, id="tp4-max"), |
| 228 | pytest.param( |
| 229 | 1, |
| 230 | 1, |
| 231 | (4, 4), |
| 232 | torch.float32, |
| 233 | torch.bfloat16, |
| 234 | id="ratio1-tp1-fp32", |
| 235 | ), |
| 236 | pytest.param( |
| 237 | 2, |
| 238 | 1, |
| 239 | (4, 2, 0), |
| 240 | torch.bfloat16, |
| 241 | torch.bfloat16, |
| 242 | id="ratio2-tp1-ragged-bf16", |
| 243 | ), |
| 244 | pytest.param( |
| 245 | 2, |
| 246 | 4, |
| 247 | (8,), |
| 248 | torch.float32, |
| 249 | torch.float32, |
| 250 | id="ratio2-tp4-max-fp32", |
| 251 | ), |
| 252 | pytest.param( |
| 253 | 3, |
| 254 | 2, |
| 255 | (4, 2, 0), |
| 256 | torch.float32, |
| 257 | torch.float32, |
| 258 | id="ratio3-tp2-ragged-fp32", |
| 259 | ), |
| 260 | pytest.param( |
| 261 | 4, |
| 262 | 4, |
| 263 | (8,), |
| 264 | torch.float32, |
| 265 | torch.bfloat16, |
| 266 | id="ratio4-tp4-max-fp32", |
| 267 | ), |
| 268 | ], |
| 269 | ) |
| 270 | @pytest.mark.parametrize("output_gate_activation", ["silu", "sigmoid"]) |
| 271 | @torch.inference_mode() |
| 272 | def test_fused_gdn_decode_post_conv_mtp_head_ratios( |
| 273 | head_ratio: int, |
| 274 | tp_size: int, |
| 275 | query_lengths: tuple[int, ...], |
| 276 | state_dtype: torch.dtype, |
| 277 | norm_dtype: torch.dtype, |
| 278 | output_gate_activation: str, |
| 279 | ) -> None: |
| 280 | if torch.cuda.get_device_capability() < (8, 0): |
| 281 | pytest.skip("fused GDN decode MTP requires compute capability 8.0+") |
| 282 | if not hasattr(torch.ops._C, "fused_gdn_decode_post_conv_mtp"): |
| 283 | pytest.skip("fused GDN decode MTP op is not built") |
| 284 | |
| 285 | torch.manual_seed(0) |
| 286 | device = "cuda" |
| 287 | H = 16 // tp_size |
| 288 | HV = head_ratio * H |
| 289 | K = V = 128 |
| 290 | num_reqs = len(query_lengths) |
| 291 | state_width = max(query_lengths) |
| 292 | num_tokens = sum(query_lengths) |
| 293 | num_slots = num_reqs * state_width + 1 |
| 294 | scale = K**-0.5 |
| 295 | eps = 1e-6 |
| 296 | |
| 297 | mixed_qkv = torch.randn( |
| 298 | num_tokens, |
| 299 | 2 * H * K + HV * V, |
| 300 | dtype=torch.bfloat16, |
| 301 | device=device, |
| 302 | ) |
| 303 | query, key, value = torch.split( |
| 304 | mixed_qkv, |
| 305 | [H * K, H * K, HV * V], |
| 306 | dim=-1, |
| 307 | ) |
| 308 | query = query.view(1, num_tokens, H, K) |
| 309 | key = key.view(1, num_tokens, H, K) |
| 310 | value = value.view(1, num_tokens, HV, V) |
| 311 | ba = torch.randn(num_tokens, 2 * HV, dtype=torch.bfloat16, device=device) |
| 312 | b, a = ba.chunk(2, dim=-1) |
| 313 | assert not a.is_contiguous() |
| 314 | assert not b.is_contiguous() |
| 315 | A_log = 0.5 * torch.randn(HV, dtype=torch.float32, device=device) |
| 316 | dt_bias = 0.1 * torch.randn(HV, dtype=torch.float32, device=device) |
| 317 | output_gate = torch.randn(num_tokens, HV, V, dtype=torch.bfloat16, device=device) |
| 318 | norm_weight = torch.randn(V, dtype=norm_dtype, device=device) |
| 319 | state_ref = ( |
| 320 | 0.01 * torch.randn(num_slots, HV, V, K, dtype=torch.float32, device=device) |
| 321 | ).to(state_dtype) |
| 322 | state_actual = state_ref.clone() |
| 323 | state_indices = torch.arange(1, num_slots, dtype=torch.int32, device=device).view( |
| 324 | num_reqs, state_width |
| 325 | ) |
| 326 | cu_seqlens = torch.tensor( |
| 327 | [0, *torch.tensor(query_lengths).cumsum(0).tolist()], |
| 328 | dtype=torch.int32, |
| 329 | device=device, |
| 330 | ) |
| 331 | num_accepted_tokens = torch.ones(num_reqs, dtype=torch.int32, device=device) |
| 332 | if query_lengths[-1] == 0: |
| 333 | state_indices[-1].zero_() |
| 334 | |
| 335 | for step, accepted_tokens in enumerate((1, min(2, state_width), state_width)): |
| 336 | num_accepted_tokens.fill_(accepted_tokens) |
| 337 | if query_lengths[-1] == 0: |
| 338 | num_accepted_tokens[-1] = 1 |
| 339 | raw_ref, _ = fused_sigmoid_gating_delta_rule_update( |
| 340 | A_log=A_log, |
| 341 | a=a, |
| 342 | b=b, |
| 343 | dt_bias=dt_bias, |
| 344 | q=query, |
| 345 | k=key, |
| 346 | v=value, |
| 347 | initial_state=state_ref, |
| 348 | inplace_final_state=True, |
| 349 | cu_seqlens=cu_seqlens, |
| 350 | ssm_state_indices=state_indices, |
| 351 | num_accepted_tokens=num_accepted_tokens, |
| 352 | scale=scale, |
| 353 | use_qk_l2norm_in_kernel=True, |
| 354 | ) |
| 355 | expected = rmsnorm_fn( |
| 356 | raw_ref.squeeze(0), |
| 357 | norm_weight, |
| 358 | None, |
| 359 | z=output_gate, |
| 360 | eps=eps, |
| 361 | norm_before_gate=True, |
| 362 | activation=output_gate_activation, |
| 363 | ) |
| 364 | actual = ops.fused_gdn_decode_post_conv_mtp( |
| 365 | mixed_qkv=mixed_qkv, |
| 366 | a=a, |
| 367 | b=b, |
| 368 | A_log=A_log, |
| 369 | dt_bias=dt_bias, |
| 370 | state_indices=state_indices, |
| 371 | cu_seqlens=cu_seqlens, |
| 372 | num_accepted_tokens=num_accepted_tokens, |
| 373 | state=state_actual, |
| 374 | output_gate=output_gate, |
| 375 | norm_weight=norm_weight, |
| 376 | out=torch.empty_like(output_gate), |
| 377 | scale=scale, |
| 378 | norm_eps=eps, |
| 379 | output_gate_activation=output_gate_activation, |
| 380 | ) |
| 381 | |
| 382 | output_error = (actual.float() - expected.float()).norm() |
| 383 | output_relative_l2 = output_error / expected.float().norm().clamp_min(1e-20) |
| 384 | assert output_relative_l2 < 5e-4, ( |
| 385 | f"MTP output relative L2 mismatch at step {step}: " |
| 386 | f"{output_relative_l2.item():.6g}" |
| 387 | ) |
| 388 | |
| 389 | torch.testing.assert_close(state_actual, state_ref, atol=3e-2, rtol=3e-2) |