@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 939 lines 30.9 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5 import pytest
6 import torch
7
8 from vllm.platforms import current_platform
9 from vllm.triton_utils import tl, triton
10 from vllm.utils.math_utils import next_power_of_2
11 from vllm.utils.torch_utils import set_random_seed
12 from vllm.v1.attention.ops.triton_attention_helpers import (
13 apply_softcap,
14 compute_tile_loop_bounds,
15 )
16 from vllm.v1.attention.ops.triton_unified_attention import unified_attention
17 from vllm.v1.kv_cache_interface import KVQuantMode
18
19 pytestmark = pytest.mark.skip_global_cleanup
20
21 DEVICE_TYPE = current_platform.device_type
22
23 NUM_HEADS = [(4, 4), (8, 2), (5, 1)]
24 HEAD_SIZES = [128, 256]
25 BLOCK_SIZES = [16]
26
27 DTYPES = [torch.bfloat16]
28 QDTYPES = [None, current_platform.fp8_dtype()]
29 FP8_DTYPE = current_platform.fp8_dtype()
30
31 # one value large enough to test overflow in index calculation.
32 # one value small enough to test the schema op check
33 NUM_BLOCKS = [32768, 2048]
34
35 # 0: use 2D kernel for decode
36 # 8: use 3D kernel for decode
37 SEQ_THRESHOLD_3D_VALUES = [0, 8]
38
39
40 @triton.jit
41 def _compute_clamped_mm_tile_bounds(
42 output_ptr,
43 mm_prefix_range_ptr,
44 USE_CAUSAL: tl.constexpr,
45 USE_PER_SEQ_CAUSAL: tl.constexpr,
46 CHUNK_LOOKBACK: tl.constexpr,
47 CHUNK_SIZE: tl.constexpr,
48 ):
49 loop_lo, loop_hi, max_seq_prefix_len = compute_tile_loop_bounds(
50 0, # context_len
51 4096, # seq_len
52 4096, # cur_batch_query_len
53 140, # q_block_local_idx: query positions [1120, 1127]
54 0, # segm_idx_or_0
55 0, # tiles_per_segment_or_0
56 32, # TILE_SIZE
57 16, # BLOCK_M
58 8, # BLOCK_Q
59 2, # num_queries_per_kv
60 1024, # SLIDING_WINDOW
61 True, # USE_MM_PREFIX
62 False, # IS_3D
63 USE_CAUSAL,
64 USE_PER_SEQ_CAUSAL,
65 CHUNK_LOOKBACK,
66 CHUNK_SIZE,
67 False, # USE_R_SWA
68 True, # MM_PREFIX_CLAMP_SW
69 2, # MAX_MM_RANGES
70 mm_prefix_range_ptr,
71 0, # seq_idx
72 )
73 tl.store(output_ptr, loop_lo)
74 tl.store(output_ptr + 1, loop_hi)
75 tl.store(output_ptr + 2, max_seq_prefix_len)
76
77
78 def test_clamped_mm_prefix_prunes_sliding_window_tiles() -> None:
79 """Retain an intersecting image range without scanning the full sequence."""
80 mm_prefix_ranges = torch.tensor(
81 [[[1024, 2303], [0, 0]]], dtype=torch.int32, device=DEVICE_TYPE
82 )
83 bounds = torch.empty(3, dtype=torch.int32, device=DEVICE_TYPE)
84
85 _compute_clamped_mm_tile_bounds[(1,)](
86 bounds,
87 mm_prefix_ranges,
88 USE_CAUSAL=True,
89 USE_PER_SEQ_CAUSAL=False,
90 CHUNK_LOOKBACK=-1,
91 CHUNK_SIZE=-1,
92 )
93
94 # The range is wider than the sliding window, so the upper bound must use
95 # its inclusive endpoint rather than query_pos + sliding_window.
96 assert bounds.tolist() == [3, 72, 4096]
97
98
99 @pytest.mark.parametrize(
100 ("use_causal", "use_per_seq_causal"), [(False, False), (True, True)]
101 )
102 def test_clamped_mm_prefix_preserves_noncausal_right_window(
103 use_causal: bool, use_per_seq_causal: bool
104 ) -> None:
105 """Union the non-causal right window with intersecting image ranges."""
106 mm_prefix_ranges = torch.tensor(
107 [[[1024, 1500], [0, 0]]], dtype=torch.int32, device=DEVICE_TYPE
108 )
109 bounds = torch.empty(3, dtype=torch.int32, device=DEVICE_TYPE)
110
111 _compute_clamped_mm_tile_bounds[(1,)](
112 bounds,
113 mm_prefix_ranges,
114 USE_CAUSAL=use_causal,
115 USE_PER_SEQ_CAUSAL=use_per_seq_causal,
116 CHUNK_LOOKBACK=-1,
117 CHUNK_SIZE=-1,
118 )
119
120 # q_hi=1127 and window=1024 admit keys through 2150, beyond the image
121 # range endpoint at 1500. Tile 67 must therefore remain in the loop.
122 assert bounds.tolist() == [3, 68, 4096]
123
124
125 def ref_paged_attn(
126 query: torch.Tensor,
127 key_cache: torch.Tensor,
128 value_cache: torch.Tensor,
129 query_lens: list[int],
130 kv_lens: list[int],
131 block_tables: torch.Tensor,
132 scale: float,
133 sliding_window: int | None = None,
134 soft_cap: float | None = None,
135 ) -> torch.Tensor:
136 num_seqs = len(query_lens)
137 block_tables = block_tables.cpu().numpy()
138 _, block_size, num_kv_heads, head_size = key_cache.shape
139 head_size_v = value_cache.shape[-1]
140
141 outputs: list[torch.Tensor] = []
142 start_idx = 0
143 for i in range(num_seqs):
144 query_len = query_lens[i]
145 kv_len = kv_lens[i]
146 q = query[start_idx : start_idx + query_len]
147 q *= scale
148
149 num_kv_blocks = (kv_len + block_size - 1) // block_size
150 block_indices = block_tables[i, :num_kv_blocks]
151
152 k = key_cache[block_indices].view(-1, num_kv_heads, head_size)
153 k = k[:kv_len]
154 v = value_cache[block_indices].view(-1, num_kv_heads, head_size_v)
155 v = v[:kv_len]
156
157 if q.shape[1] != k.shape[1]:
158 k = torch.repeat_interleave(k, q.shape[1] // k.shape[1], dim=1)
159 v = torch.repeat_interleave(v, q.shape[1] // v.shape[1], dim=1)
160 attn = torch.einsum("qhd,khd->hqk", q, k).float()
161 empty_mask = torch.ones(query_len, kv_len)
162 mask = torch.triu(empty_mask, diagonal=kv_len - query_len + 1).bool()
163 if sliding_window is not None:
164 sliding_window_mask = (
165 torch.triu(
166 empty_mask, diagonal=kv_len - (query_len + sliding_window) + 1
167 )
168 .bool()
169 .logical_not()
170 )
171 mask |= sliding_window_mask
172 if soft_cap is not None and soft_cap > 0:
173 attn = soft_cap * torch.tanh(attn / soft_cap)
174 attn.masked_fill_(mask, float("-inf"))
175 attn = torch.softmax(attn, dim=-1).to(v.dtype)
176 out = torch.einsum("hqk,khd->qhd", attn, v)
177
178 outputs.append(out)
179 start_idx += query_len
180
181 return torch.cat(outputs, dim=0)
182
183
184 def ref_paged_clamped_mm_attn(
185 query: torch.Tensor,
186 key_cache: torch.Tensor,
187 value_cache: torch.Tensor,
188 query_lens: list[int],
189 kv_lens: list[int],
190 block_tables: torch.Tensor,
191 mm_ranges: list[list[tuple[int, int]]],
192 scale: float,
193 sliding_window: int,
194 chunk_lookback: int,
195 ) -> torch.Tensor:
196 block_tables_cpu = block_tables.cpu().numpy()
197 _, block_size, num_kv_heads, head_size = key_cache.shape
198 chunk_size = sliding_window // (chunk_lookback + 1)
199
200 outputs: list[torch.Tensor] = []
201 query_start = 0
202 for req_idx, (query_len, kv_len) in enumerate(zip(query_lens, kv_lens)):
203 query_i = query[query_start : query_start + query_len].float()
204 context_len = kv_len - query_len
205 query_pos = torch.arange(query_len, device=query.device) + context_len
206 key_pos = torch.arange(kv_len, device=query.device)
207
208 num_kv_blocks = (kv_len + block_size - 1) // block_size
209 block_indices = block_tables_cpu[req_idx, :num_kv_blocks]
210 key_i = key_cache[block_indices].view(-1, num_kv_heads, head_size)[:kv_len]
211 value_i = value_cache[block_indices].view(-1, num_kv_heads, head_size)[:kv_len]
212 if query_i.shape[1] != key_i.shape[1]:
213 repeats = query_i.shape[1] // key_i.shape[1]
214 key_i = torch.repeat_interleave(key_i, repeats, dim=1)
215 value_i = torch.repeat_interleave(value_i, repeats, dim=1)
216
217 delta = query_pos[:, None] - key_pos[None, :]
218 keep = (delta >= 0) & (
219 query_pos[:, None] // chunk_size - key_pos[None, :] // chunk_size
220 <= chunk_lookback
221 )
222 for range_start, range_end in mm_ranges[req_idx]:
223 q_in_range = (query_pos >= range_start) & (query_pos <= range_end)
224 k_in_range = (key_pos >= range_start) & (key_pos <= range_end)
225 mm_mask = q_in_range[:, None] & k_in_range[None, :]
226 keep |= mm_mask & (delta < sliding_window)
227
228 scores = torch.einsum("qhd,khd->hqk", query_i, key_i.float()) * scale
229 scores.masked_fill_(~keep[None], float("-inf"))
230 probs = scores.softmax(-1).to(value_i.dtype)
231 outputs.append(torch.einsum("hqk,khd->qhd", probs, value_i))
232 query_start += query_len
233
234 return torch.cat(outputs)
235
236
237 @torch.inference_mode()
238 def test_triton_unified_attn_clamped_mm_matches_dense_reference() -> None:
239 set_random_seed(0)
240 query_lens = [384, 96]
241 kv_lens_list = [384, 320]
242 mm_ranges = [[(32, 351)], [(160, 287)]]
243 sliding_window = 128
244 chunk_lookback = 0
245 block_size = 16
246 num_query_heads = 4
247 num_kv_heads = 2
248 head_size = 128
249 scale = head_size**-0.5
250
251 query = torch.randn(
252 sum(query_lens),
253 num_query_heads,
254 head_size,
255 dtype=torch.bfloat16,
256 device=DEVICE_TYPE,
257 )
258 num_blocks_per_req = [
259 (kv_len + block_size - 1) // block_size for kv_len in kv_lens_list
260 ]
261 total_blocks = sum(num_blocks_per_req)
262 key_cache = torch.randn(
263 total_blocks,
264 block_size,
265 num_kv_heads,
266 head_size,
267 dtype=torch.bfloat16,
268 device=DEVICE_TYPE,
269 )
270 value_cache = torch.randn_like(key_cache)
271
272 block_tables = torch.zeros(
273 len(query_lens),
274 max(num_blocks_per_req),
275 dtype=torch.int32,
276 device=DEVICE_TYPE,
277 )
278 block_start = 0
279 for req_idx, num_blocks in enumerate(num_blocks_per_req):
280 block_tables[req_idx, :num_blocks] = torch.arange(
281 block_start,
282 block_start + num_blocks,
283 dtype=torch.int32,
284 device=DEVICE_TYPE,
285 )
286 block_start += num_blocks
287
288 cu_seqlens_q = torch.tensor(
289 [0] + query_lens, dtype=torch.int32, device=DEVICE_TYPE
290 ).cumsum(0, dtype=torch.int32)
291 kv_lens = torch.tensor(kv_lens_list, dtype=torch.int32, device=DEVICE_TYPE)
292 mm_prefix_range = torch.tensor(
293 [[[32, 351]], [[160, 287]]], dtype=torch.int32, device=DEVICE_TYPE
294 )
295 actual = torch.empty_like(query)
296
297 unified_attention(
298 q=query,
299 k=key_cache,
300 v=value_cache,
301 out=actual,
302 cu_seqlens_q=cu_seqlens_q,
303 max_seqlen_q=max(query_lens),
304 seqused_k=kv_lens,
305 max_seqlen_k=max(kv_lens_list),
306 softmax_scale=scale,
307 causal=True,
308 window_size=(sliding_window - 1, 0),
309 block_table=block_tables,
310 softcap=0,
311 q_descale=None,
312 k_descale=None,
313 v_descale=None,
314 mm_prefix_range=mm_prefix_range,
315 chunk_lookback=chunk_lookback,
316 mm_prefix_clamp_sliding_window=True,
317 )
318
319 expected = ref_paged_clamped_mm_attn(
320 query,
321 key_cache,
322 value_cache,
323 query_lens,
324 kv_lens_list,
325 block_tables,
326 mm_ranges,
327 scale,
328 sliding_window,
329 chunk_lookback,
330 )
331 chunk_only = ref_paged_clamped_mm_attn(
332 query,
333 key_cache,
334 value_cache,
335 query_lens,
336 kv_lens_list,
337 block_tables,
338 [[], []],
339 scale,
340 sliding_window,
341 chunk_lookback,
342 )
343
344 torch.testing.assert_close(actual.float(), expected.float(), atol=2e-2, rtol=2e-2)
345 assert not torch.allclose(expected, chunk_only, atol=2e-2, rtol=2e-2)
346
347
348 @pytest.mark.parametrize(
349 "seq_lens", [[(1, 1328), (5, 18), (129, 463)], [(1, 523), (1, 37), (1, 2011)]]
350 )
351 @pytest.mark.parametrize("num_heads", NUM_HEADS)
352 @pytest.mark.parametrize("head_size", HEAD_SIZES)
353 @pytest.mark.parametrize("block_size", BLOCK_SIZES)
354 @pytest.mark.parametrize("sliding_window", [None, 64, 128, 256])
355 @pytest.mark.parametrize("dtype", DTYPES)
356 @pytest.mark.parametrize("soft_cap", [None, 50.0])
357 @pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
358 @pytest.mark.parametrize("q_dtype", QDTYPES)
359 @pytest.mark.parametrize("seq_threshold_3D", SEQ_THRESHOLD_3D_VALUES)
360 @torch.inference_mode()
361 def test_triton_unified_attn(
362 seq_lens: list[tuple[int, int]],
363 num_heads: tuple[int, int],
364 head_size: int,
365 sliding_window: int | None,
366 dtype: torch.dtype,
367 block_size: int,
368 soft_cap: float | None,
369 num_blocks: int,
370 q_dtype: torch.dtype | None,
371 seq_threshold_3D: int,
372 ) -> None:
373 torch.set_default_device(DEVICE_TYPE)
374
375 set_random_seed(0)
376 num_seqs = len(seq_lens)
377 query_lens = [x[0] for x in seq_lens]
378 kv_lens = [x[1] for x in seq_lens]
379 num_query_heads = num_heads[0]
380 num_kv_heads = num_heads[1]
381 assert num_query_heads % num_kv_heads == 0
382 max_query_len = max(query_lens)
383 max_kv_len = max(kv_lens)
384 window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
385 scale = head_size**-0.5
386
387 query = torch.randn(sum(query_lens), num_query_heads, head_size, dtype=dtype)
388 key_cache = torch.randn(
389 num_blocks, block_size, num_kv_heads, head_size, dtype=dtype
390 )
391 value_cache = torch.randn_like(key_cache)
392 cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
393 dim=0, dtype=torch.int32
394 )
395 kv_lens = torch.tensor(kv_lens, dtype=torch.int32)
396
397 max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
398 block_tables = torch.randint(
399 0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
400 )
401
402 output = torch.empty_like(query)
403
404 maybe_quantized_query = query
405 maybe_quantized_key_cache = key_cache
406 maybe_quantized_value_cache = value_cache
407 q_descale = None
408 k_descale = None
409 v_descale = None
410 kv_quant_mode = KVQuantMode.NONE
411 if q_dtype is not None:
412 # Use non-1 scales so FP8 Q/K/V descale handling is tested explicitly.
413 q_scale = torch.tensor(0.75, dtype=torch.float32)
414 k_scale = torch.tensor(0.5, dtype=torch.float32)
415 v_scale = torch.tensor(0.25, dtype=torch.float32)
416 q_descale = q_scale
417 scale_shape = (num_seqs, num_kv_heads)
418 k_descale = torch.full(scale_shape, k_scale.item(), dtype=torch.float32)
419 v_descale = torch.full(scale_shape, v_scale.item(), dtype=torch.float32)
420 maybe_quantized_query = (query / q_scale).to(q_dtype)
421 maybe_quantized_key_cache = (key_cache / k_scale).to(q_dtype)
422 maybe_quantized_value_cache = (value_cache / v_scale).to(q_dtype)
423 kv_quant_mode = KVQuantMode.FP8_PER_TENSOR
424
425 num_par_softmax_segments = 16
426 head_size_padded = next_power_of_2(head_size)
427 softmax_segm_output = torch.empty(
428 (seq_threshold_3D, num_query_heads, num_par_softmax_segments, head_size_padded),
429 dtype=torch.float32,
430 )
431 softmax_segm_max = torch.empty(
432 (seq_threshold_3D, num_query_heads, num_par_softmax_segments),
433 dtype=torch.float32,
434 )
435 softmax_segm_expsum = torch.empty(
436 (seq_threshold_3D, num_query_heads, num_par_softmax_segments),
437 dtype=torch.float32,
438 )
439
440 unified_attention(
441 q=maybe_quantized_query,
442 k=maybe_quantized_key_cache,
443 v=maybe_quantized_value_cache,
444 out=output,
445 cu_seqlens_q=cu_query_lens,
446 seqused_k=kv_lens,
447 max_seqlen_q=max_query_len,
448 max_seqlen_k=max_kv_len,
449 softmax_scale=scale,
450 causal=True,
451 window_size=window_size,
452 block_table=block_tables,
453 softcap=soft_cap if soft_cap is not None else 0,
454 q_descale=q_descale,
455 k_descale=k_descale,
456 v_descale=v_descale,
457 seq_threshold_3D=seq_threshold_3D,
458 num_par_softmax_segments=num_par_softmax_segments,
459 softmax_segm_output=softmax_segm_output,
460 softmax_segm_max=softmax_segm_max,
461 softmax_segm_expsum=softmax_segm_expsum,
462 kv_quant_mode=kv_quant_mode,
463 )
464
465 ref_output = ref_paged_attn(
466 query=query,
467 key_cache=key_cache,
468 value_cache=value_cache,
469 query_lens=query_lens,
470 kv_lens=kv_lens,
471 block_tables=block_tables,
472 scale=scale,
473 sliding_window=sliding_window,
474 soft_cap=soft_cap,
475 )
476 atol, rtol = 1.5e-2, 1e-2
477 if q_dtype is not None:
478 atol, rtol = 1.5e-1, 1.5e-1
479 (
480 torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol),
481 f"{torch.max(torch.abs(output - ref_output))}",
482 )
483
484
485 @pytest.mark.parametrize(
486 "seq_lens", [[(1, 1328), (5, 18), (129, 463)], [(1, 523), (1, 37), (1, 2011)]]
487 )
488 @pytest.mark.parametrize("num_heads", NUM_HEADS)
489 @pytest.mark.parametrize("head_size", HEAD_SIZES)
490 @pytest.mark.parametrize("block_size", BLOCK_SIZES)
491 @pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
492 @pytest.mark.parametrize("seq_threshold_3D", SEQ_THRESHOLD_3D_VALUES)
493 @torch.inference_mode()
494 def test_triton_unified_attn_bf16_query_fp8_kv(
495 seq_lens: list[tuple[int, int]],
496 num_heads: tuple[int, int],
497 head_size: int,
498 block_size: int,
499 num_blocks: int,
500 seq_threshold_3D: int,
501 ) -> None:
502 """Test bf16 Q with FP8 per-tensor KV cache (dequant via _cast_kv_tile)."""
503 torch.set_default_device(DEVICE_TYPE)
504 set_random_seed(0)
505
506 num_seqs = len(seq_lens)
507 query_lens = [x[0] for x in seq_lens]
508 kv_lens = [x[1] for x in seq_lens]
509 num_query_heads = num_heads[0]
510 num_kv_heads = num_heads[1]
511 assert num_query_heads % num_kv_heads == 0
512 max_query_len = max(query_lens)
513 max_kv_len = max(kv_lens)
514 window_size = (-1, -1)
515 scale = head_size**-0.5
516
517 dtype = torch.bfloat16
518 query = torch.randn(sum(query_lens), num_query_heads, head_size, dtype=dtype)
519 key_cache = torch.randn(
520 num_blocks, block_size, num_kv_heads, head_size, dtype=dtype
521 )
522 value_cache = torch.randn_like(key_cache)
523
524 k_scale = torch.tensor(0.5, dtype=torch.float32)
525 v_scale = torch.tensor(0.25, dtype=torch.float32)
526 fp8_key_cache = (key_cache / k_scale).to(FP8_DTYPE)
527 fp8_value_cache = (value_cache / v_scale).to(FP8_DTYPE)
528
529 scale_shape = (num_seqs, num_kv_heads)
530 k_descale = torch.full(scale_shape, k_scale.item(), dtype=torch.float32)
531 v_descale = torch.full(scale_shape, v_scale.item(), dtype=torch.float32)
532
533 cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
534 dim=0, dtype=torch.int32
535 )
536 kv_lens_t = torch.tensor(kv_lens, dtype=torch.int32)
537
538 max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
539 block_tables = torch.randint(
540 0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
541 )
542
543 output = torch.empty_like(query)
544
545 num_par_softmax_segments = 16
546 head_size_padded = next_power_of_2(head_size)
547 softmax_segm_output = torch.empty(
548 (seq_threshold_3D, num_query_heads, num_par_softmax_segments, head_size_padded),
549 dtype=torch.float32,
550 )
551 softmax_segm_max = torch.empty(
552 (seq_threshold_3D, num_query_heads, num_par_softmax_segments),
553 dtype=torch.float32,
554 )
555 softmax_segm_expsum = torch.empty(
556 (seq_threshold_3D, num_query_heads, num_par_softmax_segments),
557 dtype=torch.float32,
558 )
559
560 unified_attention(
561 q=query,
562 k=fp8_key_cache,
563 v=fp8_value_cache,
564 out=output,
565 cu_seqlens_q=cu_query_lens,
566 seqused_k=kv_lens_t,
567 max_seqlen_q=max_query_len,
568 max_seqlen_k=max_kv_len,
569 softmax_scale=scale,
570 causal=True,
571 window_size=window_size,
572 block_table=block_tables,
573 softcap=0,
574 q_descale=None,
575 k_descale=k_descale,
576 v_descale=v_descale,
577 seq_threshold_3D=seq_threshold_3D,
578 num_par_softmax_segments=num_par_softmax_segments,
579 softmax_segm_output=softmax_segm_output,
580 softmax_segm_max=softmax_segm_max,
581 softmax_segm_expsum=softmax_segm_expsum,
582 kv_quant_mode=KVQuantMode.FP8_PER_TENSOR,
583 )
584
585 ref_output = ref_paged_attn(
586 query=query,
587 key_cache=key_cache,
588 value_cache=value_cache,
589 query_lens=query_lens,
590 kv_lens=kv_lens,
591 block_tables=block_tables,
592 scale=scale,
593 )
594
595 atol, rtol = 1.5e-1, 1.5e-1
596 (
597 torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol),
598 f"{torch.max(torch.abs(output - ref_output))}",
599 )
600
601
602 @pytest.mark.parametrize(
603 "seq_lens",
604 [
605 [(1, 1328), (5, 18), (129, 463)],
606 [(1, 523), (1, 37), (1, 2011)],
607 [(1, 1)] * 533,
608 [(533, 533)] * 533,
609 ],
610 )
611 @pytest.mark.parametrize("num_heads", NUM_HEADS)
612 @pytest.mark.parametrize("head_size", HEAD_SIZES)
613 @pytest.mark.parametrize("block_size", BLOCK_SIZES)
614 @pytest.mark.parametrize("sliding_window", [None, 64, 128, 256])
615 @pytest.mark.parametrize("soft_cap", [None, 50.0])
616 @pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
617 @pytest.mark.parametrize("seq_threshold_3D", SEQ_THRESHOLD_3D_VALUES)
618 @torch.inference_mode()
619 def test_triton_unified_attn_fp16_input_fp8_output(
620 seq_lens: list[tuple[int, int]],
621 num_heads: tuple[int, int],
622 head_size: int,
623 sliding_window: int | None,
624 block_size: int,
625 soft_cap: float | None,
626 num_blocks: int,
627 seq_threshold_3D: int,
628 ) -> None:
629 """Test with fp16 input and fp8 output using output_scale."""
630 torch.set_default_device(DEVICE_TYPE)
631
632 set_random_seed(0)
633 num_seqs = len(seq_lens)
634 query_lens = [x[0] for x in seq_lens]
635 kv_lens = [x[1] for x in seq_lens]
636 num_query_heads = num_heads[0]
637 num_kv_heads = num_heads[1]
638 assert num_query_heads % num_kv_heads == 0
639 max_query_len = max(query_lens)
640 max_kv_len = max(kv_lens)
641 window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
642 scale = head_size**-0.5
643
644 dtype = torch.float16
645 query = torch.randn(sum(query_lens), num_query_heads, head_size, dtype=dtype)
646 key_cache = torch.randn(
647 num_blocks, block_size, num_kv_heads, head_size, dtype=dtype
648 )
649 value_cache = torch.randn_like(key_cache)
650 cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
651 dim=0, dtype=torch.int32
652 )
653 kv_lens_tensor = torch.tensor(kv_lens, dtype=torch.int32)
654
655 max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
656 block_tables = torch.randint(
657 0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
658 )
659
660 output = torch.empty(sum(query_lens), num_query_heads, head_size, dtype=FP8_DTYPE)
661
662 output_scale = torch.tensor(0.5, dtype=torch.float32)
663
664 num_par_softmax_segments = 16
665 head_size_padded = next_power_of_2(head_size)
666 softmax_segm_output = torch.empty(
667 (seq_threshold_3D, num_query_heads, num_par_softmax_segments, head_size_padded),
668 dtype=torch.float32,
669 )
670 softmax_segm_max = torch.empty(
671 (seq_threshold_3D, num_query_heads, num_par_softmax_segments),
672 dtype=torch.float32,
673 )
674 softmax_segm_expsum = torch.empty(
675 (seq_threshold_3D, num_query_heads, num_par_softmax_segments),
676 dtype=torch.float32,
677 )
678
679 unified_attention(
680 q=query,
681 k=key_cache,
682 v=value_cache,
683 out=output,
684 cu_seqlens_q=cu_query_lens,
685 seqused_k=kv_lens_tensor,
686 max_seqlen_q=max_query_len,
687 max_seqlen_k=max_kv_len,
688 softmax_scale=scale,
689 causal=True,
690 window_size=window_size,
691 block_table=block_tables,
692 softcap=soft_cap if soft_cap is not None else 0,
693 q_descale=None,
694 k_descale=None,
695 v_descale=None,
696 output_scale=output_scale,
697 seq_threshold_3D=seq_threshold_3D,
698 num_par_softmax_segments=num_par_softmax_segments,
699 softmax_segm_output=softmax_segm_output,
700 softmax_segm_max=softmax_segm_max,
701 softmax_segm_expsum=softmax_segm_expsum,
702 )
703
704 ref_output = ref_paged_attn(
705 query=query,
706 key_cache=key_cache,
707 value_cache=value_cache,
708 query_lens=query_lens,
709 kv_lens=kv_lens,
710 block_tables=block_tables,
711 scale=scale,
712 sliding_window=sliding_window,
713 soft_cap=soft_cap,
714 )
715
716 output_fp16 = output.to(torch.float32) * output_scale.item()
717 output_fp16 = output_fp16.to(torch.float16)
718
719 atol, rtol = 2e-1, 2e-1
720 (
721 torch.testing.assert_close(output_fp16, ref_output, atol=atol, rtol=rtol),
722 f"{torch.max(torch.abs(output_fp16 - ref_output))}",
723 )
724
725
726 # USE_TD path covers two head-size regimes:
727 # - pow2 (HEAD_SIZE == HEAD_SIZE_PADDED): full TD path including Q/O.
728 # - non-pow2 (96, HEAD_SIZE_PADDED=128): gates USE_TD_QO off — Q load
729 # and output store fall back to pointer path, KV tile TD load remains.
730 # The non-pow2 case mirrors real models like Phi-3-mini (head_size=96).
731 HEAD_SIZES_USE_TD = [128, 256, 96]
732
733
734 def _run_use_td_case(
735 seq_lens: list[tuple[int, int]],
736 num_heads: tuple[int, int],
737 head_size: int,
738 block_size: int,
739 sliding_window: int | None,
740 soft_cap: float | None,
741 seq_threshold_3D: int,
742 dtype: torch.dtype = torch.bfloat16,
743 num_blocks: int = 2048,
744 ) -> None:
745 """Shared driver for the USE_TD test cases.
746
747 Runs ``unified_attention(..., use_td=True)`` and compares against the
748 reference paged-attention implementation that the sibling non-TD
749 tests use.
750 """
751 torch.set_default_device(DEVICE_TYPE)
752 set_random_seed(0)
753
754 num_seqs = len(seq_lens)
755 query_lens = [x[0] for x in seq_lens]
756 kv_lens = [x[1] for x in seq_lens]
757 num_query_heads, num_kv_heads = num_heads
758 assert num_query_heads % num_kv_heads == 0
759 max_query_len = max(query_lens)
760 max_kv_len = max(kv_lens)
761 window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
762 scale = head_size**-0.5
763
764 query = torch.randn(sum(query_lens), num_query_heads, head_size, dtype=dtype)
765 key_cache = torch.randn(
766 num_blocks, block_size, num_kv_heads, head_size, dtype=dtype
767 )
768 value_cache = torch.randn_like(key_cache)
769 cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
770 dim=0, dtype=torch.int32
771 )
772 kv_lens_tensor = torch.tensor(kv_lens, dtype=torch.int32)
773
774 max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
775 block_tables = torch.randint(
776 0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
777 )
778
779 output = torch.empty_like(query)
780
781 num_par_softmax_segments = 16
782 head_size_padded = next_power_of_2(head_size)
783 softmax_segm_output = torch.empty(
784 (seq_threshold_3D, num_query_heads, num_par_softmax_segments, head_size_padded),
785 dtype=torch.float32,
786 )
787 softmax_segm_max = torch.empty(
788 (seq_threshold_3D, num_query_heads, num_par_softmax_segments),
789 dtype=torch.float32,
790 )
791 softmax_segm_expsum = torch.empty(
792 (seq_threshold_3D, num_query_heads, num_par_softmax_segments),
793 dtype=torch.float32,
794 )
795
796 unified_attention(
797 q=query,
798 k=key_cache,
799 v=value_cache,
800 out=output,
801 cu_seqlens_q=cu_query_lens,
802 seqused_k=kv_lens_tensor,
803 max_seqlen_q=max_query_len,
804 max_seqlen_k=max_kv_len,
805 softmax_scale=scale,
806 causal=True,
807 window_size=window_size,
808 block_table=block_tables,
809 softcap=soft_cap if soft_cap is not None else 0,
810 q_descale=None,
811 k_descale=None,
812 v_descale=None,
813 seq_threshold_3D=seq_threshold_3D,
814 num_par_softmax_segments=num_par_softmax_segments,
815 softmax_segm_output=softmax_segm_output,
816 softmax_segm_max=softmax_segm_max,
817 softmax_segm_expsum=softmax_segm_expsum,
818 use_td=True,
819 )
820
821 ref_output = ref_paged_attn(
822 query=query,
823 key_cache=key_cache,
824 value_cache=value_cache,
825 query_lens=query_lens,
826 kv_lens=kv_lens,
827 block_tables=block_tables,
828 scale=scale,
829 sliding_window=sliding_window,
830 soft_cap=soft_cap,
831 )
832 torch.testing.assert_close(output, ref_output, atol=1.5e-2, rtol=1e-2)
833
834
835 @pytest.mark.parametrize(
836 "seq_lens", [[(1, 1328), (5, 18), (129, 463)], [(1, 523), (1, 37), (1, 2011)]]
837 )
838 @pytest.mark.parametrize("num_heads", NUM_HEADS)
839 @pytest.mark.parametrize("head_size", HEAD_SIZES_USE_TD)
840 @pytest.mark.parametrize("block_size", BLOCK_SIZES)
841 @pytest.mark.parametrize("sliding_window", [None, 128])
842 @pytest.mark.parametrize("soft_cap", [None, 50.0])
843 @pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
844 @pytest.mark.parametrize("seq_threshold_3D", SEQ_THRESHOLD_3D_VALUES)
845 @torch.inference_mode()
846 def test_triton_unified_attn_use_td(
847 seq_lens: list[tuple[int, int]],
848 num_heads: tuple[int, int],
849 head_size: int,
850 sliding_window: int | None,
851 block_size: int,
852 soft_cap: float | None,
853 num_blocks: int,
854 seq_threshold_3D: int,
855 ) -> None:
856 """Exercise the USE_TD (tensor-descriptor) Q/K/V load/store path.
857
858 Covers both 2D and 3D kernels via ``seq_threshold_3D``. Two routes
859 to the USE_TD_QO=False fallback (pointer path for Q/O with TD still
860 active for KV tile loads):
861
862 - non-pow2 ``num_queries_per_kv`` via ``NUM_HEADS`` entry ``(5, 1)``,
863 - non-pow2 ``head_size`` via ``HEAD_SIZES_USE_TD`` entry ``96``.
864 """
865 _run_use_td_case(
866 seq_lens=seq_lens,
867 num_heads=num_heads,
868 head_size=head_size,
869 block_size=block_size,
870 sliding_window=sliding_window,
871 soft_cap=soft_cap,
872 seq_threshold_3D=seq_threshold_3D,
873 num_blocks=num_blocks,
874 )
875
876
877 # Prefill-heavy shape: long query drives the prefill kernel path where
878 # ``_get_tile_size`` returns 32, which exceeds block_size=16 and must be
879 # clamped by the fix in 'clamp TILE_SIZE to block_size when USE_TD'.
880 # Only the prefill launch exercises the clamp, so parameterize only over
881 # the (num_heads, seq_threshold_3D=0) combinations needed to cover it.
882 @pytest.mark.parametrize("num_heads", [(4, 4), (5, 1)])
883 @torch.inference_mode()
884 def test_triton_unified_attn_use_td_tile_clamp(
885 num_heads: tuple[int, int],
886 ) -> None:
887 """Regression guard: ``USE_TD`` needs ``BLOCK_SIZE % TILE_SIZE == 0``.
888
889 With ``block_size=16`` and ``head_size=128`` (non-Gemma3),
890 ``_get_tile_size`` returns 32 for prefill, which violates the
891 ``USE_TD`` constraint unless clamped to ``block_size``. Without
892 the clamp the triton kernel ``static_assert`` fires at compile time.
893 """
894 _run_use_td_case(
895 seq_lens=[(256, 256), (128, 128)],
896 num_heads=num_heads,
897 head_size=128,
898 block_size=16,
899 sliding_window=None,
900 soft_cap=None,
901 seq_threshold_3D=0,
902 )
903
904
905 @triton.jit
906 def _softcap_probe_kernel(
907 scores_ptr,
908 out_ptr,
909 numel,
910 soft_cap,
911 BLOCK: tl.constexpr,
912 ):
913 offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
914 mask = offs < numel
915 scores = tl.load(scores_ptr + offs, mask=mask, other=0.0)
916 tl.store(out_ptr + offs, apply_softcap(scores, soft_cap), mask=mask)
917
918
919 @torch.inference_mode()
920 def test_softcap_does_not_overflow_on_large_scores() -> None:
921 """Scores above ~88 * soft_cap must not overflow to inf/NaN.
922
923 The exp-based softcap computes ``(exp(y) - exp(-y)) / (exp(y) + exp(-y))``
924 with ``y = S / soft_cap``. For ``|y| > ~88`` the exponentials overflow to
925 ``inf`` and the ratio becomes ``inf / inf = NaN``, poisoning the whole
926 attention row. Gemma-2 style models use ``attn_logit_softcapping = 50``,
927 so scores above 4400 are in range for their large attention logits.
928 """
929 soft_cap = 50.0
930 scores = torch.tensor(
931 [1.0, 100.0, 1000.0, 3000.0, 5000.0, 10000.0, -1.0, -10000.0, 4400.0],
932 device=DEVICE_TYPE,
933 dtype=torch.float32,
934 )
935 out = torch.empty_like(scores)
936 _softcap_probe_kernel[(1,)](scores, out, scores.numel(), soft_cap, BLOCK=16)
937 ref = soft_cap * torch.tanh(scores / soft_cap)
938 assert torch.isfinite(out).all(), out
939 torch.testing.assert_close(out, ref, atol=1e-3, rtol=1e-3)