A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 03:45 UTC
Add unit tests for fp8 output fusion of triton_attn (#34228)
Signed-off-by: Burkhard Ringlein <ngl@zurich.ibm.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Burkhard Ringlein committed
Feb 18, 2026 at 12:22 UTC
e24663c5a958f9ad2cf787ba2e9b1da0ba558768
1 file changed
+127
-1
tests/kernels/attention/test_triton_unified_attention.py
+127
-1
@@ -10,7 +10,7 @@ from vllm.utils.math_utils import next_power_of_2
10
from vllm.utils.torch_utils import set_random_seed
11
from vllm.v1.attention.ops.triton_unified_attention import unified_attention
12
13
-NUM_HEADS = [(4, 4), (8, 2)]
13
+NUM_HEADS = [(4, 4), (8, 2), (5, 1)]
14
HEAD_SIZES = [128, 256]
15
BLOCK_SIZES = [16]
16
@@ -20,6 +20,8 @@ QDTYPES = (
20
if not current_platform.is_rocm()
21
else [None, torch.float8_e4m3fnuz]
22
)
23
+FP8_DTYPE = current_platform.fp8_dtype()
24
+
25
# one value large enough to test overflow in index calculation.
26
# one value small enough to test the schema op check
27
NUM_BLOCKS = [32768, 2048]
@@ -217,3 +219,127 @@ def test_triton_unified_attn(
219
torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol),
220
f"{torch.max(torch.abs(output - ref_output))}",
221
)
222
+
223
+
224
+@pytest.mark.parametrize(
225
+ "seq_lens",
226
+ [
227
+ [(1, 1328), (5, 18), (129, 463)],
228
+ [(1, 523), (1, 37), (1, 2011)],
229
+ [(1, 1)] * 533,
230
+ [(533, 533)] * 533,
231
+ ],
232
+)
233
+@pytest.mark.parametrize("num_heads", NUM_HEADS)
234
+@pytest.mark.parametrize("head_size", HEAD_SIZES)
235
+@pytest.mark.parametrize("block_size", BLOCK_SIZES)
236
+@pytest.mark.parametrize("sliding_window", [None, 64, 128, 256])
237
+@pytest.mark.parametrize("soft_cap", [None, 50.0])
238
+@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
239
+@pytest.mark.parametrize("seq_threshold_3D", SEQ_THRESHOLD_3D_VALUES)
240
+@torch.inference_mode()
241
+def test_triton_unified_attn_fp16_input_fp8_output(
242
+ seq_lens: list[tuple[int, int]],
243
+ num_heads: tuple[int, int],
244
+ head_size: int,
245
+ sliding_window: int | None,
246
+ block_size: int,
247
+ soft_cap: float | None,
248
+ num_blocks: int,
249
+ seq_threshold_3D: int,
250
+) -> None:
251
+ """Test with fp16 input and fp8 output using output_scale."""
252
+ torch.set_default_device("cuda")
253
+
254
+ set_random_seed(0)
255
+ num_seqs = len(seq_lens)
256
+ query_lens = [x[0] for x in seq_lens]
257
+ kv_lens = [x[1] for x in seq_lens]
258
+ num_query_heads = num_heads[0]
259
+ num_kv_heads = num_heads[1]
260
+ assert num_query_heads % num_kv_heads == 0
261
+ max_query_len = max(query_lens)
262
+ max_kv_len = max(kv_lens)
263
+ window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
264
+ scale = head_size**-0.5
265
+
266
+ dtype = torch.float16
267
+ query = torch.randn(sum(query_lens), num_query_heads, head_size, dtype=dtype)
268
+ key_cache = torch.randn(
269
+ num_blocks, block_size, num_kv_heads, head_size, dtype=dtype
270
+ )
271
+ value_cache = torch.randn_like(key_cache)
272
+ cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
273
+ dim=0, dtype=torch.int32
274
+ )
275
+ kv_lens_tensor = torch.tensor(kv_lens, dtype=torch.int32)
276
+
277
+ max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
278
+ block_tables = torch.randint(
279
+ 0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
280
+ )
281
+
282
+ output = torch.empty(sum(query_lens), num_query_heads, head_size, dtype=FP8_DTYPE)
283
+
284
+ output_scale = torch.tensor(0.5, dtype=torch.float32)
285
+
286
+ num_par_softmax_segments = 16
287
+ head_size_padded = next_power_of_2(head_size)
288
+ softmax_segm_output = torch.empty(
289
+ (seq_threshold_3D, num_query_heads, num_par_softmax_segments, head_size_padded),
290
+ dtype=torch.float32,
291
+ )
292
+ softmax_segm_max = torch.empty(
293
+ (seq_threshold_3D, num_query_heads, num_par_softmax_segments),
294
+ dtype=torch.float32,
295
+ )
296
+ softmax_segm_expsum = torch.empty(
297
+ (seq_threshold_3D, num_query_heads, num_par_softmax_segments),
298
+ dtype=torch.float32,
299
+ )
300
+
301
+ unified_attention(
302
+ q=query,
303
+ k=key_cache,
304
+ v=value_cache,
305
+ out=output,
306
+ cu_seqlens_q=cu_query_lens,
307
+ seqused_k=kv_lens_tensor,
308
+ max_seqlen_q=max_query_len,
309
+ max_seqlen_k=max_kv_len,
310
+ softmax_scale=scale,
311
+ causal=True,
312
+ window_size=window_size,
313
+ block_table=block_tables,
314
+ softcap=soft_cap if soft_cap is not None else 0,
315
+ q_descale=None,
316
+ k_descale=None,
317
+ v_descale=None,
318
+ output_scale=output_scale,
319
+ seq_threshold_3D=seq_threshold_3D,
320
+ num_par_softmax_segments=num_par_softmax_segments,
321
+ softmax_segm_output=softmax_segm_output,
322
+ softmax_segm_max=softmax_segm_max,
323
+ softmax_segm_expsum=softmax_segm_expsum,
324
+ )
325
+
326
+ ref_output = ref_paged_attn(
327
+ query=query,
328
+ key_cache=key_cache,
329
+ value_cache=value_cache,
330
+ query_lens=query_lens,
331
+ kv_lens=kv_lens,
332
+ block_tables=block_tables,
333
+ scale=scale,
334
+ sliding_window=sliding_window,
335
+ soft_cap=soft_cap,
336
+ )
337
+
338
+ output_fp16 = output.to(torch.float32) * output_scale.item()
339
+ output_fp16 = output_fp16.to(torch.float16)
340
+
341
+ atol, rtol = 2e-1, 2e-1
342
+ (
343
+ torch.testing.assert_close(output_fp16, ref_output, atol=atol, rtol=rtol),
344
+ f"{torch.max(torch.abs(output_fp16 - ref_output))}",
345
+ )