@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 162 lines 5.49 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import pytest
5 import torch
6
7 from vllm.platforms import current_platform
8 from vllm.third_party.flash_linear_attention.ops import (
9 fused_recurrent_gated_delta_rule,
10 fused_recurrent_gated_delta_rule_packed_decode,
11 )
12
13 DEVICE = current_platform.device_type
14
15 pytestmark = pytest.mark.skipif(
16 not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
17 reason="Gated delta rule Triton kernels require a CUDA-alike or XPU device.",
18 )
19
20
21 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
22 @pytest.mark.parametrize("strided_mixed_qkv", [False, True])
23 def test_fused_recurrent_packed_decode_matches_reference(
24 dtype: torch.dtype, strided_mixed_qkv: bool
25 ):
26 torch.manual_seed(0)
27
28 # Small but representative GDN config (Qwen3Next defaults are K=128, V=128).
29 B = 32
30 H = 4
31 HV = 8 # grouped value attention: HV must be divisible by H
32 K = 128
33 V = 128
34 qkv_dim = 2 * (H * K) + (HV * V)
35
36 device = torch.device(DEVICE)
37
38 if strided_mixed_qkv:
39 # Simulate a packed view into a larger projection buffer:
40 # mixed_qkv.stride(0) > mixed_qkv.shape[1]
41 proj = torch.randn((B, qkv_dim + 64), device=device, dtype=dtype)
42 mixed_qkv = proj[:, :qkv_dim]
43 else:
44 mixed_qkv = torch.randn((B, qkv_dim), device=device, dtype=dtype)
45
46 a = torch.randn((B, HV), device=device, dtype=dtype)
47 b = torch.randn((B, HV), device=device, dtype=dtype)
48 A_log = torch.randn((HV,), device=device, dtype=dtype)
49 dt_bias = torch.randn((HV,), device=device, dtype=dtype)
50
51 # Continuous batching indices (include PAD_SLOT_ID=-1 cases). Index 0 is
52 # reserved as NULL_BLOCK_ID (CUDA graph padding), so valid slots start at 1.
53 ssm_state_indices = torch.arange(1, B + 1, device=device, dtype=torch.int32)
54 ssm_state_indices[-3:] = -1
55
56 state0 = torch.randn((B + 1, HV, V, K), device=device, dtype=dtype)
57 state_ref = state0.clone()
58 state_packed = state0.clone()
59
60 out_packed = torch.empty((B, 1, HV, V), device=device, dtype=dtype)
61
62 # Reference path: materialize contiguous Q/K/V + explicit gating.
63 q, k, v = torch.split(mixed_qkv, [H * K, H * K, HV * V], dim=-1)
64 q = q.view(B, H, K).unsqueeze(1).contiguous()
65 k = k.view(B, H, K).unsqueeze(1).contiguous()
66 v = v.view(B, HV, V).unsqueeze(1).contiguous()
67
68 x = a.float() + dt_bias.float()
69 softplus_x = torch.where(
70 x <= 20.0, torch.log1p(torch.exp(torch.clamp(x, max=20.0))), x
71 )
72 g = (-torch.exp(A_log.float()) * softplus_x).unsqueeze(1)
73 beta = torch.sigmoid(b.float()).unsqueeze(1)
74
75 out_ref, state_ref = fused_recurrent_gated_delta_rule(
76 q=q,
77 k=k,
78 v=v,
79 g=g,
80 beta=beta,
81 scale=K**-0.5,
82 initial_state=state_ref,
83 inplace_final_state=True,
84 cu_seqlens=None,
85 ssm_state_indices=ssm_state_indices,
86 use_qk_l2norm_in_kernel=True,
87 )
88
89 # Packed path: fused gating + recurrent directly from packed mixed_qkv.
90 fused_recurrent_gated_delta_rule_packed_decode(
91 mixed_qkv=mixed_qkv,
92 a=a,
93 b=b,
94 A_log=A_log,
95 dt_bias=dt_bias,
96 scale=K**-0.5,
97 initial_state=state_packed,
98 out=out_packed,
99 ssm_state_indices=ssm_state_indices,
100 use_qk_l2norm_in_kernel=True,
101 )
102
103 atol = 2e-2 if dtype != torch.float32 else 1e-4
104 rtol = 1e-2 if dtype != torch.float32 else 1e-4
105 # Output rows for PAD_SLOT_ID entries are never written (uninitialized in
106 # both paths), so compare only the valid rows.
107 valid = ssm_state_indices > 0
108 torch.testing.assert_close(out_packed[valid], out_ref[valid], rtol=rtol, atol=atol)
109 torch.testing.assert_close(state_packed, state_ref, rtol=rtol, atol=atol)
110
111
112 def test_packed_decode_keeps_beta_in_fp32():
113 device = torch.device(DEVICE)
114 dtype = torch.bfloat16
115
116 mixed_qkv = torch.ones((1, 3), device=device, dtype=dtype)
117 a = torch.zeros((1, 1), device=device, dtype=dtype)
118 b = torch.full((1, 1), 0.5, device=device, dtype=dtype)
119 params = torch.zeros((1,), device=device, dtype=dtype)
120 ssm_state_indices = torch.ones((1,), device=device, dtype=torch.int32)
121
122 state_packed = torch.zeros((2, 1, 1, 1), device=device, dtype=torch.float32)
123 out_packed = torch.empty((1, 1, 1, 1), device=device, dtype=dtype)
124
125 fused_recurrent_gated_delta_rule_packed_decode(
126 mixed_qkv=mixed_qkv,
127 a=a,
128 b=b,
129 A_log=params,
130 dt_bias=params,
131 scale=1.0,
132 initial_state=state_packed,
133 out=out_packed,
134 ssm_state_indices=ssm_state_indices,
135 )
136
137 expected_beta = torch.sigmoid(b.float()).squeeze()
138 torch.testing.assert_close(
139 state_packed[1, 0, 0, 0], expected_beta, rtol=1e-6, atol=1e-6
140 )
141
142
143 def test_packed_decode_supports_large_batch_head_grid():
144 B, H, HV, K, V = 1024, 8, 64, 1, 1
145 device = torch.device(DEVICE)
146 gates = torch.empty((B, HV), device=device)
147 params = torch.empty((HV,), device=device)
148 out = torch.empty((B, 1, HV, V), device=device)
149
150 fused_recurrent_gated_delta_rule_packed_decode(
151 mixed_qkv=torch.empty((B, 2 * H * K + HV * V), device=device),
152 a=gates,
153 b=gates,
154 A_log=params,
155 dt_bias=params,
156 scale=1.0,
157 initial_state=torch.empty((1, HV, V, K), device=device),
158 out=out,
159 ssm_state_indices=torch.zeros((B,), device=device, dtype=torch.int32),
160 )
161
162 assert torch.count_nonzero(out).item() == 0