@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 258 lines 7.61 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 import vllm.models.inkling.nvidia.ops.sconv as sconv
8 from vllm.models.inkling.nvidia.ops import qkvr_prep
9 from vllm.platforms import current_platform
10
11 _cap = current_platform.get_device_capability() if current_platform.is_cuda() else None
12
13
14 def _make_inputs(*, is_local: bool, tokens: int = 33, tp_size: int = 4):
15 torch.manual_seed(0)
16 heads = 64 // tp_size
17 kv_heads = (16 if is_local else 8) // tp_size
18 head_dim = 128
19 d_rel = 16
20 rel_extent = 512 if is_local else 1024
21 page_size = 16
22 num_blocks = (tokens + page_size - 1) // page_size
23 q_width = heads * head_dim
24 kv_width = kv_heads * head_dim
25 r_width = heads * d_rel
26 device = "cuda"
27
28 qkvr = torch.randn(
29 tokens,
30 q_width + 2 * kv_width + r_width,
31 device=device,
32 dtype=torch.bfloat16,
33 )
34 k_weight = torch.randn(kv_width, 4, device=device, dtype=torch.bfloat16)
35 v_weight = torch.randn_like(k_weight)
36 q_norm_weight = torch.randn(head_dim, device=device, dtype=torch.bfloat16)
37 k_norm_weight = torch.randn_like(q_norm_weight)
38 rel_proj = torch.randn(d_rel, rel_extent, device=device, dtype=torch.bfloat16)
39 conv_cache = torch.zeros(
40 num_blocks,
41 kv_heads,
42 page_size,
43 2 * head_dim,
44 device=device,
45 dtype=torch.bfloat16,
46 )
47 key_cache = torch.empty(
48 num_blocks,
49 page_size,
50 kv_heads,
51 head_dim,
52 device=device,
53 dtype=torch.bfloat16,
54 )
55 value_cache = torch.empty_like(key_cache)
56 positions = torch.arange(tokens, device=device, dtype=torch.int64)
57 block_table = torch.arange(num_blocks, device=device, dtype=torch.int32)[None]
58 seq_idx = torch.zeros(tokens, device=device, dtype=torch.int32)
59 slots = torch.arange(tokens, device=device, dtype=torch.int64)
60 query_start = torch.zeros(tokens, device=device, dtype=torch.int32)
61 log_scaling_n_floor = None if is_local else 128000
62 log_scaling = None
63 if log_scaling_n_floor is not None:
64 log_scaling = torch.linspace(
65 1.0,
66 1.1,
67 tokens,
68 device=device,
69 dtype=torch.float32,
70 )
71 return (
72 qkvr,
73 k_weight,
74 v_weight,
75 q_norm_weight,
76 k_norm_weight,
77 rel_proj,
78 1e-6,
79 heads,
80 kv_heads,
81 head_dim,
82 d_rel,
83 conv_cache,
84 key_cache,
85 value_cache,
86 positions,
87 block_table,
88 seq_idx,
89 slots,
90 query_start,
91 slots,
92 0,
93 head_dim,
94 page_size,
95 log_scaling,
96 )
97
98
99 def _reference(args):
100 qkvr = args[0]
101 heads, kv_heads, head_dim, d_rel = args[7:11]
102 q_width = heads * head_dim
103 kv_width = kv_heads * head_dim
104 q, k, v, r = qkvr.split((q_width, kv_width, kv_width, heads * d_rel), dim=1)
105 conv_cache = args[11].clone()
106 k = sconv.fused_sconv(
107 k.contiguous(),
108 args[1],
109 conv_cache,
110 args[14],
111 args[15],
112 args[16],
113 args[17],
114 args[18],
115 args[20],
116 head_dim,
117 args[22],
118 )
119 v = sconv.fused_sconv(
120 v.contiguous(),
121 args[2],
122 conv_cache,
123 args[14],
124 args[15],
125 args[16],
126 args[17],
127 args[18],
128 args[21],
129 head_dim,
130 args[22],
131 )
132
133 def rms_norm(x, weight):
134 x = x.reshape(-1, head_dim).float()
135 rstd = torch.rsqrt(x.square().mean(1, keepdim=True) + args[6])
136 return (x * rstd * weight.float()).to(qkvr.dtype)
137
138 q = rms_norm(q, args[3]).view(qkvr.shape[0], heads, head_dim)
139 k = rms_norm(k, args[4]).view(qkvr.shape[0], kv_heads, head_dim)
140 v = v.view(qkvr.shape[0], kv_heads, head_dim)
141 rel = torch.mm(r.reshape(-1, d_rel), args[5]).view(qkvr.shape[0], heads, -1)
142 if args[23] is not None:
143 q = (q.float() * args[23][:, None, None]).to(q.dtype)
144 rel = (rel.float() * args[23][:, None, None]).to(rel.dtype)
145
146 key_cache = args[12].clone()
147 value_cache = args[13].clone()
148 slots = args[19]
149 valid = slots >= 0
150 key_cache.view(-1, kv_heads, head_dim)[slots[valid]] = k[valid]
151 value_cache.view(-1, kv_heads, head_dim)[slots[valid]] = v[valid]
152 return q.flatten(1), rel, key_cache, value_cache
153
154
155 @pytest.mark.skipif(
156 _cap is None,
157 reason="Inkling QKVR prep kernels require CUDA",
158 )
159 @pytest.mark.parametrize(
160 ("is_local", "tokens", "tp_size"),
161 [
162 (True, 9, 4),
163 (False, 33, 4),
164 (True, 33, 8),
165 (True, 128, 4),
166 (False, 128, 8),
167 (True, 512, 4),
168 (False, 640, 8),
169 ],
170 )
171 @torch.inference_mode()
172 def test_qkvr_prep_matches_reference(is_local, tokens, tp_size):
173 args = _make_inputs(is_local=is_local, tokens=tokens, tp_size=tp_size)
174 ref_q, ref_rel, ref_key, ref_value = _reference(args)
175
176 q, rel = qkvr_prep.fused_qkvr_prep(*args)
177
178 torch.testing.assert_close(q, ref_q, rtol=0.01, atol=0.02)
179 torch.testing.assert_close(rel, ref_rel, rtol=0.02, atol=0.125)
180 torch.testing.assert_close(args[12], ref_key, rtol=0.01, atol=0.01)
181 torch.testing.assert_close(args[13], ref_value, rtol=0, atol=0)
182
183
184 @pytest.mark.skipif(
185 _cap is None,
186 reason="Inkling QKVR prep kernels require CUDA",
187 )
188 @pytest.mark.parametrize("tokens", [9, 128])
189 @torch.inference_mode()
190 def test_qkvr_log_scaling_preserves_bf16_norm_output(tokens):
191 args = list(_make_inputs(is_local=False, tokens=tokens, tp_size=8))
192 tau = args[23]
193 assert tau is not None
194
195 scaled_q, scaled_rel = qkvr_prep.fused_qkvr_prep(*args)
196 args[23] = None
197 unscaled_q, unscaled_rel = qkvr_prep.fused_qkvr_prep(*args)
198
199 expected_q = (unscaled_q.float() * tau[:, None]).to(unscaled_q.dtype)
200 expected_rel = (unscaled_rel.float() * tau[:, None, None]).to(unscaled_rel.dtype)
201 torch.testing.assert_close(scaled_q, expected_q, rtol=0, atol=0)
202 torch.testing.assert_close(scaled_rel, expected_rel, rtol=0, atol=0)
203
204
205 @pytest.mark.skipif(
206 _cap is None,
207 reason="Inkling QKVR prep kernels require CUDA",
208 )
209 @pytest.mark.parametrize("tokens", [9, 128])
210 @torch.inference_mode()
211 def test_qkvr_prep_negative_slots_skip_cache_writes(tokens):
212 args = list(_make_inputs(is_local=True, tokens=tokens))
213 args[17].fill_(-1)
214 args[19].fill_(-1)
215 conv_cache = args[11].clone()
216 key_cache = args[12].clone()
217 value_cache = args[13].clone()
218
219 qkvr_prep.fused_qkvr_prep(*args)
220
221 torch.testing.assert_close(args[11], conv_cache, rtol=0, atol=0)
222 torch.testing.assert_close(args[12], key_cache, rtol=0, atol=0)
223 torch.testing.assert_close(args[13], value_cache, rtol=0, atol=0)
224
225
226 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
227 @torch.inference_mode()
228 def test_fused_sconv_negative_slots_skip_cache_writes():
229 args = list(_make_inputs(is_local=True, tokens=9))
230 args[17].fill_(-1)
231 cache = args[11].clone()
232
233 sconv.fused_sconv(
234 args[0][:, 16 * 128 : 16 * 128 + 4 * 128].contiguous(),
235 args[1],
236 args[11],
237 args[14],
238 args[15],
239 args[16],
240 args[17],
241 args[18],
242 0,
243 128,
244 args[22],
245 )
246
247 torch.testing.assert_close(args[11], cache, rtol=0, atol=0)
248
249
250 @pytest.mark.parametrize(
251 ("rel_extent", "last_latency_rows", "first_throughput_rows"),
252 [(512, 8191, 8192), (1024, 2047, 2048)],
253 )
254 def test_rel_projection_schedule_crossover(
255 rel_extent, last_latency_rows, first_throughput_rows
256 ):
257 assert not qkvr_prep.use_rel_proj_throughput(last_latency_rows, rel_extent)
258 assert qkvr_prep.use_rel_proj_throughput(first_throughput_rows, rel_extent)