@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 199 lines 6.64 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 import torch.nn.functional as F
7
8 from vllm.platforms import current_platform
9 from vllm.third_party.flash_linear_attention.ops import (
10 fused_recurrent_gated_delta_rule,
11 fused_sigmoid_gating_delta_rule_update,
12 )
13 from vllm.utils.torch_utils import set_random_seed
14
15 DEVICE = current_platform.device_type
16
17
18 @pytest.mark.parametrize("tp_size", [1])
19 @pytest.mark.parametrize("num_reqs", [1, 2, 4])
20 @pytest.mark.parametrize("num_k_heads", [16])
21 @pytest.mark.parametrize("num_v_heads", [32])
22 @pytest.mark.parametrize("head_k_dim", [128])
23 @pytest.mark.parametrize("head_v_dim", [128])
24 @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
25 def test_fused_sigmoid_gating_delta_rule_update_non_spec(
26 tp_size: int,
27 num_reqs: int,
28 num_k_heads: int,
29 num_v_heads: int,
30 head_k_dim: int,
31 head_v_dim: int,
32 dtype: torch.dtype,
33 ) -> None:
34 torch.set_default_device(DEVICE)
35 set_random_seed(0)
36 key_dim = head_k_dim * num_k_heads
37 value_dim = head_v_dim * num_v_heads
38 mixed_qkv_dim = (key_dim * 2 + value_dim) // tp_size
39 seq_len = 1 # seq_len is 1 for decode
40 num_tokens = num_reqs * seq_len
41 total_entries = num_tokens * 2
42
43 mixed_qkv = torch.rand(num_tokens, mixed_qkv_dim, dtype=dtype)
44 query, key, value = torch.split(
45 mixed_qkv,
46 [
47 key_dim // tp_size,
48 key_dim // tp_size,
49 value_dim // tp_size,
50 ],
51 dim=-1,
52 )
53 query = query.view(1, num_tokens, num_k_heads, head_k_dim)
54 key = key.view(1, num_tokens, num_k_heads, head_k_dim)
55 value = value.view(1, num_tokens, num_v_heads, head_v_dim)
56
57 A_log = torch.rand(num_v_heads // tp_size, dtype=dtype)
58 dt_bias = torch.rand(num_v_heads // tp_size, dtype=dtype)
59 a = torch.rand(num_tokens, num_v_heads, dtype=dtype)
60 b = torch.rand(num_tokens, num_v_heads, dtype=dtype)
61 # Entry 0 is reserved as NULL_BLOCK_ID (CUDA graph padding), so valid
62 # state indices start at 1.
63 ssm_state = torch.rand(
64 total_entries + 1, num_v_heads, head_k_dim, head_v_dim, dtype=dtype
65 )
66 state_indices = (torch.randperm(total_entries, dtype=torch.int32) + 1)[:num_tokens]
67 cu_seqlens = torch.arange(0, num_tokens + 1, dtype=torch.int32)
68
69 beta = b.sigmoid()
70 g = -A_log.float().exp() * F.softplus(a.float() + dt_bias)
71 core_attn_out_ref, last_recurrent_state_ref = fused_recurrent_gated_delta_rule(
72 q=query,
73 k=key,
74 v=value,
75 g=g.unsqueeze(0),
76 beta=beta.unsqueeze(0),
77 initial_state=ssm_state.clone(),
78 inplace_final_state=True,
79 ssm_state_indices=state_indices,
80 cu_seqlens=cu_seqlens,
81 use_qk_l2norm_in_kernel=True,
82 )
83
84 core_attn_out, last_recurrent_state = fused_sigmoid_gating_delta_rule_update(
85 A_log=A_log,
86 a=a,
87 b=b,
88 dt_bias=dt_bias,
89 q=query,
90 k=key,
91 v=value,
92 initial_state=ssm_state,
93 inplace_final_state=True,
94 ssm_state_indices=state_indices,
95 cu_seqlens=cu_seqlens,
96 use_qk_l2norm_in_kernel=True,
97 )
98
99 torch.testing.assert_close(core_attn_out, core_attn_out_ref, atol=1e-2, rtol=1e-2)
100 torch.testing.assert_close(
101 last_recurrent_state, last_recurrent_state_ref, atol=1e-2, rtol=1e-2
102 )
103
104
105 @pytest.mark.parametrize("tp_size", [1])
106 @pytest.mark.parametrize("num_reqs", [1, 2, 4])
107 @pytest.mark.parametrize("num_k_heads", [16])
108 @pytest.mark.parametrize("num_v_heads", [32])
109 @pytest.mark.parametrize("head_k_dim", [128])
110 @pytest.mark.parametrize("head_v_dim", [128])
111 @pytest.mark.parametrize("num_speculative_tokens", [1, 3])
112 @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
113 def test_fused_sigmoid_gating_delta_rule_update_spec(
114 tp_size: int,
115 num_reqs: int,
116 num_k_heads: int,
117 num_v_heads: int,
118 head_k_dim: int,
119 head_v_dim: int,
120 num_speculative_tokens: int,
121 dtype: torch.dtype,
122 ) -> None:
123 torch.set_default_device(DEVICE)
124 set_random_seed(0)
125 key_dim = head_k_dim * num_k_heads
126 value_dim = head_v_dim * num_v_heads
127 mixed_qkv_dim = (key_dim * 2 + value_dim) // tp_size
128 num_tokens = num_reqs * (num_speculative_tokens + 1)
129 total_entries = num_tokens * 2
130
131 mixed_qkv = torch.rand(num_tokens, mixed_qkv_dim, dtype=dtype)
132 query, key, value = torch.split(
133 mixed_qkv,
134 [
135 key_dim // tp_size,
136 key_dim // tp_size,
137 value_dim // tp_size,
138 ],
139 dim=-1,
140 )
141 query = query.view(1, num_tokens, num_k_heads, head_k_dim)
142 key = key.view(1, num_tokens, num_k_heads, head_k_dim)
143 value = value.view(1, num_tokens, num_v_heads, head_v_dim)
144
145 A_log = torch.rand(num_v_heads // tp_size, dtype=dtype)
146 dt_bias = torch.rand(num_v_heads // tp_size, dtype=dtype)
147 a = torch.rand(num_tokens, num_v_heads, dtype=dtype)
148 b = torch.rand(num_tokens, num_v_heads, dtype=dtype)
149 # Entry 0 is reserved as NULL_BLOCK_ID (CUDA graph padding), so valid
150 # state indices start at 1.
151 ssm_state = torch.rand(
152 total_entries + 1, num_v_heads, head_k_dim, head_v_dim, dtype=dtype
153 )
154 state_indices = (torch.randperm(total_entries, dtype=torch.int32) + 1)[
155 :num_tokens
156 ].view(num_reqs, num_speculative_tokens + 1)
157 num_accepted_tokens = torch.randint(
158 1, num_speculative_tokens + 1, (num_reqs,), dtype=torch.int32
159 )
160 cu_seqlens = torch.arange(
161 0, num_tokens + 1, num_speculative_tokens + 1, dtype=torch.int32
162 )
163
164 beta = b.sigmoid()
165 g = -A_log.float().exp() * F.softplus(a.float() + dt_bias)
166 core_attn_out_ref, last_recurrent_state_ref = fused_recurrent_gated_delta_rule(
167 q=query,
168 k=key,
169 v=value,
170 g=g.unsqueeze(0),
171 beta=beta.unsqueeze(0),
172 initial_state=ssm_state.clone(),
173 inplace_final_state=True,
174 ssm_state_indices=state_indices,
175 cu_seqlens=cu_seqlens,
176 num_accepted_tokens=num_accepted_tokens,
177 use_qk_l2norm_in_kernel=True,
178 )
179
180 core_attn_out, last_recurrent_state = fused_sigmoid_gating_delta_rule_update(
181 A_log=A_log,
182 a=a,
183 b=b,
184 dt_bias=dt_bias,
185 q=query,
186 k=key,
187 v=value,
188 initial_state=ssm_state,
189 inplace_final_state=True,
190 ssm_state_indices=state_indices,
191 cu_seqlens=cu_seqlens,
192 num_accepted_tokens=num_accepted_tokens,
193 use_qk_l2norm_in_kernel=True,
194 )
195
196 torch.testing.assert_close(core_attn_out, core_attn_out_ref, atol=1e-2, rtol=1e-2)
197 torch.testing.assert_close(
198 last_recurrent_state, last_recurrent_state_ref, atol=1e-2, rtol=1e-2
199 )