@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 327 lines 10.7 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.utils.math_utils import cdiv
9 from vllm.v1.attention.ops.triton_decode_attention import decode_attention_fwd
10
11 pytestmark = pytest.mark.skip_global_cleanup
12
13 DEVICE_TYPE = current_platform.device_type
14
15
16 @pytest.mark.parametrize("B", [3, 5])
17 @pytest.mark.parametrize("L", [1027, 1025])
18 @pytest.mark.parametrize("H_Q", [32])
19 @pytest.mark.parametrize("H_KV", [32, 8])
20 @pytest.mark.parametrize("D_QK", [128, 192, 576])
21 @pytest.mark.parametrize("D_V", [128, 512])
22 @pytest.mark.parametrize("CACHE_SIZE", [16384])
23 @pytest.mark.parametrize("PAGE_SIZE", [1, 16])
24 def test_decode_attention(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE):
25 assert CACHE_SIZE % PAGE_SIZE == 0
26 dtype = torch.bfloat16
27 seq_len = L # This represents the number of tokens already in the sequence
28 sm_scale = 1.0 / (D_QK**0.5)
29 num_kv_splits = 8
30
31 num_pages_per_batch = cdiv(seq_len, PAGE_SIZE)
32 req_to_page = torch.randint(
33 0, CACHE_SIZE // PAGE_SIZE, (B, num_pages_per_batch, 1), device=DEVICE_TYPE
34 )
35 req_to_token = req_to_page * PAGE_SIZE
36 req_to_token = req_to_token.expand(B, num_pages_per_batch, PAGE_SIZE)
37 req_to_token = req_to_token + torch.arange(PAGE_SIZE, device=DEVICE_TYPE).view(
38 1, 1, -1
39 )
40 req_to_token = req_to_token.view(B, -1)
41 req_to_token = req_to_token[:, :seq_len].contiguous()
42
43 # q represents the new token being generated, one per batch
44 q = torch.randn(B, H_Q, D_QK, dtype=dtype, device=DEVICE_TYPE)
45
46 # k_buffer and v_buffer represent all previous tokens
47 # Page size is 1.
48 k_buffer = torch.randn(CACHE_SIZE, H_KV, D_QK, dtype=dtype, device=DEVICE_TYPE)
49 v_buffer = torch.randn(CACHE_SIZE, H_KV, D_V, dtype=dtype, device=DEVICE_TYPE)
50
51 # o will have the same shape as q
52 o = torch.zeros(B, H_Q, D_V, dtype=dtype, device=DEVICE_TYPE)
53
54 lse = torch.zeros(B, H_Q, dtype=dtype, device=DEVICE_TYPE)
55
56 b_seq_len = torch.full((B,), seq_len, device=DEVICE_TYPE)
57
58 attn_logits = torch.empty(
59 (B, H_Q, num_kv_splits, D_V + 1),
60 dtype=torch.float32,
61 device=DEVICE_TYPE,
62 )
63
64 # Call the original implementation.
65 decode_attention_fwd(
66 q,
67 k_buffer,
68 v_buffer,
69 o,
70 lse,
71 req_to_token,
72 b_seq_len,
73 attn_logits,
74 num_kv_splits,
75 sm_scale,
76 )
77
78 # Page size can be larger than 1.
79 k_buffer = k_buffer.view(CACHE_SIZE // PAGE_SIZE, PAGE_SIZE, H_KV, D_QK)
80 v_buffer = v_buffer.view(CACHE_SIZE // PAGE_SIZE, PAGE_SIZE, H_KV, D_V)
81
82 o1 = torch.zeros_like(o)
83 lse1 = torch.zeros_like(lse)
84
85 decode_attention_fwd(
86 q,
87 k_buffer,
88 v_buffer,
89 o1,
90 lse1,
91 req_to_page,
92 b_seq_len,
93 attn_logits,
94 num_kv_splits,
95 sm_scale,
96 PAGE_SIZE,
97 )
98
99 assert torch.allclose(o, o1)
100
101
102 def _quantize_to_fp8(tensor: torch.Tensor):
103 """Quantize a BF16 tensor to FP8 e4m3fn with per-tensor scale.
104
105 Returns (fp8_tensor, scale) where:
106 fp8_tensor ≈ tensor / scale (stored as float8_e4m3fn)
107 tensor ≈ fp8_tensor.to(float32) * scale (dequantized)
108 """
109 amax = tensor.abs().amax()
110 # float8_e4m3fn max representable value is 448.0
111 scale = (amax / 448.0).clamp(min=1e-12).to(torch.float32)
112 fp8_tensor = (
113 (tensor.to(torch.float32) / scale).clamp(-448.0, 448.0).to(torch.float8_e4m3fn)
114 )
115 return fp8_tensor, scale
116
117
118 @pytest.mark.parametrize("B", [3])
119 @pytest.mark.parametrize("L", [1025])
120 @pytest.mark.parametrize("H_Q", [32])
121 @pytest.mark.parametrize("H_KV", [32, 8])
122 @pytest.mark.parametrize("D_QK", [128, 576])
123 @pytest.mark.parametrize("D_V", [128, 512])
124 @pytest.mark.parametrize("CACHE_SIZE", [16384])
125 @pytest.mark.parametrize("PAGE_SIZE", [1, 16])
126 def test_decode_attention_fp8(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE):
127 """Test FP8 KV cache path: quantize K/V to FP8, run kernel with scales,
128 and compare against BF16 reference output."""
129 assert CACHE_SIZE % PAGE_SIZE == 0
130 dtype = torch.bfloat16
131 seq_len = L
132 sm_scale = 1.0 / (D_QK**0.5)
133 num_kv_splits = 8
134
135 num_pages_per_batch = cdiv(seq_len, PAGE_SIZE)
136 req_to_page = torch.randint(
137 0, CACHE_SIZE // PAGE_SIZE, (B, num_pages_per_batch, 1), device=DEVICE_TYPE
138 )
139 req_to_token = req_to_page * PAGE_SIZE
140 req_to_token = req_to_token.expand(B, num_pages_per_batch, PAGE_SIZE)
141 req_to_token = req_to_token + torch.arange(PAGE_SIZE, device=DEVICE_TYPE).view(
142 1, 1, -1
143 )
144 req_to_token = req_to_token.view(B, -1)
145 req_to_token = req_to_token[:, :seq_len].contiguous()
146
147 q = torch.randn(B, H_Q, D_QK, dtype=dtype, device=DEVICE_TYPE)
148
149 # Create BF16 K/V as reference
150 k_bf16 = torch.randn(CACHE_SIZE, H_KV, D_QK, dtype=dtype, device=DEVICE_TYPE)
151 v_bf16 = torch.randn(CACHE_SIZE, H_KV, D_V, dtype=dtype, device=DEVICE_TYPE)
152
153 # --- BF16 reference ---
154 o_ref = torch.zeros(B, H_Q, D_V, dtype=dtype, device=DEVICE_TYPE)
155 lse_ref = torch.zeros(B, H_Q, dtype=dtype, device=DEVICE_TYPE)
156 attn_logits = torch.empty(
157 (B, H_Q, num_kv_splits, D_V + 1), dtype=torch.float32, device=DEVICE_TYPE
158 )
159
160 if PAGE_SIZE == 1:
161 decode_attention_fwd(
162 q,
163 k_bf16,
164 v_bf16,
165 o_ref,
166 lse_ref,
167 req_to_token,
168 b_seq_len=torch.full((B,), seq_len, device=DEVICE_TYPE),
169 attn_logits=attn_logits,
170 num_kv_splits=num_kv_splits,
171 sm_scale=sm_scale,
172 )
173 else:
174 k_paged = k_bf16.view(CACHE_SIZE // PAGE_SIZE, PAGE_SIZE, H_KV, D_QK)
175 v_paged = v_bf16.view(CACHE_SIZE // PAGE_SIZE, PAGE_SIZE, H_KV, D_V)
176 decode_attention_fwd(
177 q,
178 k_paged,
179 v_paged,
180 o_ref,
181 lse_ref,
182 req_to_page,
183 b_seq_len=torch.full((B,), seq_len, device=DEVICE_TYPE),
184 attn_logits=attn_logits,
185 num_kv_splits=num_kv_splits,
186 sm_scale=sm_scale,
187 page_size=PAGE_SIZE,
188 )
189
190 # --- FP8 path ---
191 k_fp8, k_scale = _quantize_to_fp8(k_bf16)
192 v_fp8, v_scale = _quantize_to_fp8(v_bf16)
193
194 o_fp8 = torch.zeros(B, H_Q, D_V, dtype=dtype, device=DEVICE_TYPE)
195 lse_fp8 = torch.zeros(B, H_Q, dtype=dtype, device=DEVICE_TYPE)
196 attn_logits_fp8 = torch.empty(
197 (B, H_Q, num_kv_splits, D_V + 1), dtype=torch.float32, device=DEVICE_TYPE
198 )
199
200 if PAGE_SIZE == 1:
201 decode_attention_fwd(
202 q,
203 k_fp8,
204 v_fp8,
205 o_fp8,
206 lse_fp8,
207 req_to_token,
208 b_seq_len=torch.full((B,), seq_len, device=DEVICE_TYPE),
209 attn_logits=attn_logits_fp8,
210 num_kv_splits=num_kv_splits,
211 sm_scale=sm_scale,
212 k_scale=k_scale,
213 v_scale=v_scale,
214 )
215 else:
216 k_fp8_paged = k_fp8.view(CACHE_SIZE // PAGE_SIZE, PAGE_SIZE, H_KV, D_QK)
217 v_fp8_paged = v_fp8.view(CACHE_SIZE // PAGE_SIZE, PAGE_SIZE, H_KV, D_V)
218 decode_attention_fwd(
219 q,
220 k_fp8_paged,
221 v_fp8_paged,
222 o_fp8,
223 lse_fp8,
224 req_to_page,
225 b_seq_len=torch.full((B,), seq_len, device=DEVICE_TYPE),
226 attn_logits=attn_logits_fp8,
227 num_kv_splits=num_kv_splits,
228 sm_scale=sm_scale,
229 page_size=PAGE_SIZE,
230 k_scale=k_scale,
231 v_scale=v_scale,
232 )
233
234 # FP8 tolerances match test_mla_backends.py test_backend_correctness.
235 torch.testing.assert_close(o_ref, o_fp8, atol=5e-1, rtol=1e-2)
236
237
238 @pytest.mark.parametrize(
239 "H_Q,H_KV,D_QK,D_V,is_mla",
240 [
241 (16, 1, 576, 512, True), # MLA path (grouped kernel, v = trans(k))
242 (32, 8, 128, 128, False), # GQA path (grouped kernel)
243 (32, 32, 128, 128, False), # MHA path (normal kernel)
244 ],
245 )
246 @pytest.mark.parametrize("PAGE_SIZE", [16])
247 def test_decode_attention_cross_layer_view(H_Q, H_KV, D_QK, D_V, is_mla, PAGE_SIZE):
248 """The kernel must honor the cache's page-dim stride, not assume pages are
249 packed back-to-back. A per-layer view into a cross-layer (block-major)
250 cache has stride(0) inflated by num_layers; outputs must match a
251 contiguous cache holding the same data exactly."""
252 B = 3
253 seq_len = 1027
254 CACHE_SIZE = 16384
255 NUM_LAYERS = 3
256 LAYER_IDX = 1
257 dtype = torch.bfloat16
258 sm_scale = 1.0 / (D_QK**0.5)
259 num_kv_splits = 8
260 num_pages = CACHE_SIZE // PAGE_SIZE
261
262 num_pages_per_batch = cdiv(seq_len, PAGE_SIZE)
263 req_to_page = torch.randint(
264 0, num_pages, (B, num_pages_per_batch), device=DEVICE_TYPE
265 )
266
267 q = torch.randn(B, H_Q, D_QK, dtype=dtype, device=DEVICE_TYPE)
268 b_seq_len = torch.full((B,), seq_len, device=DEVICE_TYPE)
269
270 # Reference: contiguous paged cache.
271 k_ref = torch.randn(
272 num_pages, PAGE_SIZE, H_KV, D_QK, dtype=dtype, device=DEVICE_TYPE
273 )
274 if is_mla:
275 v_ref = k_ref[..., :D_V]
276 else:
277 v_ref = torch.randn(
278 num_pages, PAGE_SIZE, H_KV, D_V, dtype=dtype, device=DEVICE_TYPE
279 )
280
281 # Cross-layer cache: all layers' pages for a block are adjacent. The
282 # per-layer view has the same shape as the contiguous cache but
283 # stride(0) is NUM_LAYERS x larger. Neighbor layers hold random data so
284 # any packed-pages addressing reads garbage rather than zeros.
285 k_xl = torch.randn(
286 num_pages, NUM_LAYERS, PAGE_SIZE, H_KV, D_QK, dtype=dtype, device=DEVICE_TYPE
287 )
288 k_view = k_xl[:, LAYER_IDX]
289 k_view.copy_(k_ref)
290 assert k_view.stride(0) == NUM_LAYERS * PAGE_SIZE * H_KV * D_QK
291 if is_mla:
292 v_view = k_view[..., :D_V]
293 else:
294 v_xl = torch.randn(
295 num_pages, NUM_LAYERS, PAGE_SIZE, H_KV, D_V, dtype=dtype, device=DEVICE_TYPE
296 )
297 v_view = v_xl[:, LAYER_IDX]
298 v_view.copy_(v_ref)
299
300 def run(k_buffer, v_buffer):
301 o = torch.zeros(B, H_Q, D_V, dtype=dtype, device=DEVICE_TYPE)
302 lse = torch.zeros(B, H_Q, dtype=dtype, device=DEVICE_TYPE)
303 attn_logits = torch.empty(
304 (B, H_Q, num_kv_splits, D_V + 1), dtype=torch.float32, device=DEVICE_TYPE
305 )
306 decode_attention_fwd(
307 q,
308 k_buffer,
309 v_buffer,
310 o,
311 lse,
312 req_to_page,
313 b_seq_len,
314 attn_logits,
315 num_kv_splits,
316 sm_scale,
317 PAGE_SIZE,
318 is_mla=is_mla,
319 )
320 return o, lse
321
322 o_ref, lse_ref = run(k_ref, v_ref)
323 o_xl, lse_xl = run(k_view, v_view)
324
325 # Same data and same compute order; only addressing differs.
326 assert torch.equal(o_ref, o_xl)
327 assert torch.equal(lse_ref, lse_xl)