@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 9a161307f5
Fork 0

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 14, 05:30 UTC

[torch.compile][ROCm][V1] Enable attention output FP8 fusion for V1 attention backends (#19767)

Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com> Signed-off-by: Luka Govedič <lgovedic@redhat.com> Co-authored-by: Luka Govedič <lgovedic@redhat.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>

Gregory Shtrasberg committed Sep 10, 2025 at 16:59 UTC 9a161307f5f096c63ae4134c5055d87a36d224a8
8 files changed +249 -135
tests/compile/test_fusion_attn.py
+87 -42
@@ -40,13 +40,12 @@ backend_unfused: Optional[TestBackend] = None
40 @pytest.mark.parametrize(
41 "model, quant_key",
42 [("amd/Llama-3.1-8B-Instruct-FP8-KV", kFp8StaticTensorSym)])
43 -@pytest.mark.parametrize(
44 - "use_triton_fa", [True, False] if current_platform.is_rocm() else [False])
43 +@pytest.mark.parametrize("use_triton_fa", [True, False])
44 @pytest.mark.skipif(not current_platform.supports_fp8(), reason="Need FP8")
46 -@pytest.mark.skipif(not current_platform.is_cuda_alike(),
47 - reason="Only test CUDA and ROCm")
48 -def test_attention_fusion(example_prompts, monkeypatch, model: str,
49 - quant_key: QuantKey, use_triton_fa: bool):
45 +@pytest.mark.skipif(not current_platform.is_rocm(),
46 + reason="V0 attn quant fusion only on ROCm")
47 +def test_attention_fusion_v0(example_prompts, monkeypatch, model: str,
48 + quant_key: QuantKey, use_triton_fa: bool):
49 # Clean Dynamo cache to avoid reusing other test cases
50 # (for some reason the reset at the end is not enough)
51 torch._dynamo.reset()
@@ -69,13 +68,17 @@ def test_attention_fusion(example_prompts, monkeypatch, model: str,
68 backend="tests.compile.test_fusion_attn.backend_unfused",
69 custom_ops=["+quant_fp8"],
70 )
72 - vllm_config = VllmConfig(compilation_config=compile_config)
71 + vllm_config = VllmConfig(compilation_config=compile_config,
72 + model_config=ModelConfig(
73 + model=model,
74 + dtype=torch.bfloat16,
75 + ))
76 backend_unfused = TestBackend(NoOpEliminationPass(vllm_config))
77
78 llm = LLM(model,
79 enforce_eager=True,
80 compilation_config=compile_config,
78 - gpu_memory_utilization=0.9,
81 + gpu_memory_utilization=0.5,
82 max_model_len=2048)
83
84 sampling_params = SamplingParams(temperature=0.0,
@@ -93,7 +96,11 @@ def test_attention_fusion(example_prompts, monkeypatch, model: str,
96 backend="tests.compile.test_fusion_attn.backend",
97 custom_ops=["+quant_fp8"],
98 )
96 - vllm_config = VllmConfig(compilation_config=compile_config)
99 + vllm_config = VllmConfig(compilation_config=compile_config,
100 + model_config=ModelConfig(
101 + model=model,
102 + dtype=torch.bfloat16,
103 + ))
104
105 # AttnFusionPass needs attention layers to be registered in config upon init
106 # so we initialize it during compilation.
@@ -102,7 +109,7 @@ def test_attention_fusion(example_prompts, monkeypatch, model: str,
109 llm2 = LLM(model,
110 enforce_eager=True,
111 compilation_config=compile_config,
105 - gpu_memory_utilization=0.9,
112 + gpu_memory_utilization=0.5,
113 max_model_len=2048)
114
115 # check support
@@ -171,6 +178,8 @@ class AttentionQuantPatternModel(torch.nn.Module):
178 cache_config=vllm_config.cache_config,
179 prefix="model.layers.0.self_attn.attn",
180 )
181 + self.attn._k_scale = self.attn._k_scale.to(device)
182 + self.attn._v_scale = self.attn._v_scale.to(device)
183
184 self.block_size = 16
185
@@ -188,7 +197,7 @@ class AttentionQuantPatternModel(torch.nn.Module):
197 device=self.device,
198 )
199
191 - def build_attn_metadata(self, batch_size: int):
200 + def build_attn_metadata(self, batch_size: int, use_hnd: bool):
201 """Initialize attention metadata."""
202
203 # Create common attn metadata
@@ -205,10 +214,8 @@ class AttentionQuantPatternModel(torch.nn.Module):
214 num_blocks = batch_size * max_blocks
215
216 # Create dummy KV cache for FlashInfer TRTLLM
208 - # - NHD: [num_blocks, 2, block_size, num_kv_heads, head_size]
209 - # - HND: [num_blocks, 2, num_kv_heads, block_size, head_size]
210 - # Create kv_cache in HND layout and permute to NHD layout
211 - # (later will be permuted back to HND layout in forward pass)
217 + # - NHD: [num_blocks, block_size, num_kv_heads, head_size]
218 + # - HND: [num_blocks, num_kv_heads, block_size, head_size]
219 kv_cache = torch.zeros(num_blocks,
220 2,
221 self.num_kv_heads,
@@ -216,7 +223,17 @@ class AttentionQuantPatternModel(torch.nn.Module):
223 self.head_size,
224 dtype=self.kv_cache_dtype,
225 device=self.device)
219 - kv_cache = kv_cache.permute(0, 1, 3, 2, 4)
226 + if current_platform.is_rocm():
227 + # k/v as 1st dimention
228 + if use_hnd:
229 + kv_cache = kv_cache.permute(1, 0, 2, 3, 4)
230 + else:
231 + kv_cache = kv_cache.permute(1, 0, 3, 2, 4)
232 + else:
233 + # k/v as 2nd dimention
234 + # Create kv_cache in HND layout and permute to NHD layout
235 + # (later will be permuted back to HND layout in forward pass)
236 + kv_cache = kv_cache.permute(0, 1, 3, 2, 4)
237 self.attn.kv_cache = [kv_cache]
238
239 # Build attn metadata
@@ -296,28 +313,51 @@ class TestAttentionNvfp4QuantPatternModel(AttentionQuantPatternModel):
313 out_dtype=attn_output.dtype)
314
315
299 -@pytest.mark.parametrize("num_qo_heads, num_kv_heads", [(64, 8), (40, 8)])
316 +if current_platform.is_cuda():
317 + MODELS = [("nvidia/Llama-4-Scout-17B-16E-Instruct-FP8",
318 + TestAttentionFp8StaticQuantPatternModel),
319 + ("nvidia/Llama-4-Scout-17B-16E-Instruct-FP4",
320 + TestAttentionNvfp4QuantPatternModel)]
321 + HEADS = [(64, 8), (40, 8)]
322 +elif current_platform.is_rocm():
323 + MODELS = [("amd/Llama-3.1-8B-Instruct-FP8-KV",
324 + TestAttentionFp8StaticQuantPatternModel)]
325 + HEADS = [(32, 8), (40, 8)]
326 +else:
327 + MODELS = []
328 + HEADS = []
329 +
330 +
331 +@pytest.mark.parametrize("num_qo_heads, num_kv_heads", HEADS)
332 @pytest.mark.parametrize("head_size", [128])
301 -@pytest.mark.parametrize("batch_size", [7, 256, 533])
302 -@pytest.mark.parametrize("dtype", [torch.bfloat16])
303 -@pytest.mark.parametrize("model_name, model_class",
304 - [("nvidia/Llama-4-Scout-17B-16E-Instruct-FP8",
305 - TestAttentionFp8StaticQuantPatternModel),
306 - ("nvidia/Llama-4-Scout-17B-16E-Instruct-FP4",
307 - TestAttentionNvfp4QuantPatternModel)])
308 -@pytest.mark.parametrize("backend", [_Backend.FLASHINFER])
309 -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA")
333 +@pytest.mark.parametrize("batch_size",
334 + [7, 256, 533] if current_platform.is_cuda() else [8])
335 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
336 +@pytest.mark.parametrize("model_name, model_class", MODELS)
337 +@pytest.mark.parametrize("backend", [_Backend.FLASHINFER] if
338 + current_platform.is_cuda() else [_Backend.ROCM_FLASH])
339 +@pytest.mark.parametrize(
340 + "split_attention",
341 + [False, True] if current_platform.is_rocm() else [False])
342 +@pytest.mark.skipif(not current_platform.is_cuda_alike(),
343 + reason="Only test ROCm or CUDA")
344 @pytest.mark.skipif(not current_platform.supports_fp8(), reason="Need FP8")
311 -@pytest.mark.skipif(not current_platform.is_device_capability((10, 0)),
312 - reason="Only test on SM100(Blackwell)")
345 +@pytest.mark.skipif(current_platform.is_cuda()
346 + and not current_platform.is_device_capability((10, 0)),
347 + reason="On CUDA only test on SM100(Blackwell)")
348 +@pytest.mark.skipif(not current_platform.is_cuda_alike(),
349 + reason="Only test ROCm or CUDA")
350 def test_attention_quant_pattern(num_qo_heads: int, num_kv_heads: int,
351 head_size: int, batch_size: int,
352 dtype: torch.dtype, model_name: str,
353 model_class: type[AttentionQuantPatternModel],
317 - backend: _Backend, monkeypatch, dist_init):
354 + backend: _Backend, split_attention: bool,
355 + monkeypatch, dist_init):
356 """Test AttentionStaticQuantPattern fusion pass"""
357
358 monkeypatch.setenv("VLLM_USE_V1", "1")
359 + if split_attention:
360 + monkeypatch.setenv("VLLM_V1_USE_PREFILL_DECODE_ATTENTION", "1")
361
362 device = torch.device("cuda:0")
363 torch.manual_seed(42)
@@ -326,6 +366,7 @@ def test_attention_quant_pattern(num_qo_heads: int, num_kv_heads: int,
366 model_config=ModelConfig(
367 model=model_name,
368 max_model_len=2048,
369 + dtype=dtype,
370 ),
371 scheduler_config=SchedulerConfig(max_num_seqs=1024),
372 compilation_config=CompilationConfig(
@@ -368,7 +409,7 @@ def test_attention_quant_pattern(num_qo_heads: int, num_kv_heads: int,
409
410 forward_ctx = get_forward_context()
411 forward_ctx.attn_metadata = model_unfused.build_attn_metadata(
371 - batch_size)
412 + batch_size, use_hnd=split_attention)
413
414 # Run model directly without compilation and fusion
415 result_unfused = model_unfused(q, k, v)
@@ -389,7 +430,8 @@ def test_attention_quant_pattern(num_qo_heads: int, num_kv_heads: int,
430 model_fused = model_fused.to(device)
431
432 forward_ctx = get_forward_context()
392 - forward_ctx.attn_metadata = model_fused.build_attn_metadata(batch_size)
433 + forward_ctx.attn_metadata = model_fused.build_attn_metadata(
434 + batch_size, use_hnd=split_attention)
435
436 # Create test backend with fusion passes enabled
437 noop_pass = NoOpEliminationPass(vllm_config)
@@ -404,12 +446,19 @@ def test_attention_quant_pattern(num_qo_heads: int, num_kv_heads: int,
446 assert model_compiled.attn._o_scale_float is None
447 result_fused_1 = model_compiled(q, k, v)
448
407 - # After the 1st round of the forward pass, output quant scale should be
408 - # loaded into the attn layer's _o_scale_float, the 2nd round should
409 - # reuse the loaded _o_scale_float
410 - assert model_compiled.attn._o_scale_float is not None
411 - result_fused_2 = model_compiled(q, k, v)
412 - assert model_compiled.attn._o_scale_float is not None
449 + if backend == _Backend.FLASHINFER:
450 + # With the Flashinfer backend after the 1st round of the forward
451 + # pass, output quant scale should be loaded into the attn layer's
452 + # _o_scale_float, the 2nd round should reuse the loaded
453 + # _o_scale_float
454 + assert model_compiled.attn._o_scale_float is not None
455 + result_fused_2 = model_compiled(q, k, v)
456 + assert model_compiled.attn._o_scale_float is not None
457 +
458 + torch.testing.assert_close(result_unfused,
459 + result_fused_2,
460 + atol=1e-2,
461 + rtol=1e-2)
462
463 # Check attn fusion support
464 quant_key = model_class.quant_key
@@ -444,12 +493,8 @@ def test_attention_quant_pattern(num_qo_heads: int, num_kv_heads: int,
493 assert attn_nodes_post[0].kwargs.get("output_block_scale") is not None, \
494 "Attention should have output_block_scale after FP4 fusion" # noqa: E501
495
447 - # Check that results are closed
496 + # Check that results are close
497 torch.testing.assert_close(result_unfused,
498 result_fused_1,
499 atol=1e-2,
500 rtol=1e-2)
452 - torch.testing.assert_close(result_unfused,
453 - result_fused_2,
454 - atol=1e-2,
455 - rtol=1e-2)
vllm/attention/ops/chunked_prefill_paged_decode.py
+16 -2
@@ -15,6 +15,8 @@ from vllm.triton_utils import tl, triton
15
16 from .prefix_prefill import context_attention_fwd
17
18 +float8_info = torch.finfo(current_platform.fp8_dtype())
19 +
20
21 @triton.jit
22 def cdiv_fn(x, y):
@@ -34,6 +36,7 @@ def kernel_paged_attention_2d(
36 scale, # float32
37 k_scale, # float32
38 v_scale, # float32
39 + out_scale_inv,
40 num_query_heads: tl.constexpr, # int
41 num_queries_per_kv: tl.constexpr, # int
42 num_queries_per_kv_padded: tl.constexpr, # int
@@ -60,7 +63,9 @@ def kernel_paged_attention_2d(
63 filter_by_query_len: tl.constexpr, # bool
64 query_start_len_ptr, # [num_seqs+1]
65 USE_SINKS: tl.constexpr, # bool
63 -):
66 + USE_FP8: tl.constexpr,
67 + FP8_MIN: tl.constexpr = float8_info.min,
68 + FP8_MAX: tl.constexpr = float8_info.max):
69 seq_idx = tl.program_id(0)
70 kv_head_idx = tl.program_id(1)
71
@@ -204,6 +209,9 @@ def kernel_paged_attention_2d(
209
210 # epilogue
211 acc = acc / L[:, None]
212 + if USE_FP8:
213 + acc = acc * tl.load(out_scale_inv)
214 + acc = tl.clamp(acc, FP8_MIN, FP8_MAX)
215
216 output_offset = (cur_batch_in_all_start_index * output_stride_0 +
217 query_head_idx * output_stride_1)
@@ -234,6 +242,7 @@ def chunked_prefill_paged_decode(
242 alibi_slopes=None,
243 sliding_window=None,
244 sm_scale=None,
245 + output_scale=None,
246 # Optional tensor for sinks
247 sinks=None,
248 ):
@@ -266,6 +275,7 @@ def chunked_prefill_paged_decode(
275 sliding_window=sliding_window,
276 sm_scale=sm_scale,
277 skip_decode=True,
278 + fp8_out_scale=output_scale,
279 sinks=sinks,
280 )
281
@@ -316,7 +326,7 @@ def chunked_prefill_paged_decode(
326 tmp_output = torch.empty(
327 size=(total_num_seq, num_query_heads, max_num_partitions,
328 head_size),
319 - dtype=output.dtype,
329 + dtype=query.dtype,
330 device=output.device,
331 )
332 exp_sums = torch.empty(
@@ -345,6 +355,7 @@ def chunked_prefill_paged_decode(
355 kv_cache_dtype=kv_cache_dtype,
356 k_scale=k_scale,
357 v_scale=v_scale,
358 + fp8_out_scale=output_scale,
359 )
360 else:
361 kernel_paged_attention_2d[(
@@ -362,6 +373,8 @@ def chunked_prefill_paged_decode(
373 scale=sm_scale,
374 k_scale=k_scale,
375 v_scale=v_scale,
376 + out_scale_inv=1.0 /
377 + output_scale if output_scale is not None else 1.0,
378 num_query_heads=num_query_heads,
379 num_queries_per_kv=num_queries_per_kv,
380 num_queries_per_kv_padded=num_queries_per_kv_padded,
@@ -388,4 +401,5 @@ def chunked_prefill_paged_decode(
401 filter_by_query_len=True,
402 query_start_len_ptr=query_start_loc,
403 USE_SINKS=sinks is not None,
404 + USE_FP8=output_scale is not None,
405 )
vllm/attention/ops/prefix_prefill.py
+13 -1
@@ -15,6 +15,7 @@ NUM_WARPS = 4 if current_platform.is_rocm() else 8
15
16 # To check compatibility
17 IS_TURING = current_platform.get_device_capability() == (7, 5)
18 +float8_info = torch.finfo(current_platform.fp8_dtype())
19
20
21 # Here's an example autotuner config for this kernel. This config does provide
@@ -43,6 +44,7 @@ def _fwd_kernel(Q,
44 sm_scale,
45 k_scale,
46 v_scale,
47 + out_scale_inv,
48 B_Start_Loc,
49 B_Seqlen,
50 x: tl.constexpr,
@@ -82,8 +84,11 @@ def _fwd_kernel(Q,
84 num_unroll_request: tl.constexpr,
85 SKIP_DECODE: tl.constexpr,
86 USE_SINKS: tl.constexpr,
87 + USE_FP8: tl.constexpr,
88 MAX_Q_LEN: tl.constexpr = 0,
86 - MAX_CTX_LEN: tl.constexpr = 0):
89 + MAX_CTX_LEN: tl.constexpr = 0,
90 + FP8_MIN: tl.constexpr = float8_info.min,
91 + FP8_MAX: tl.constexpr = float8_info.max):
92
93 cur_batch = tl.program_id(0)
94 cur_head = tl.program_id(1)
@@ -284,6 +289,9 @@ def _fwd_kernel(Q,
289 off_o = ((cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs +
290 cur_head * stride_oh + offs_d[None, :] * stride_od)
291 out_ptrs = Out + off_o
292 + if USE_FP8:
293 + acc = acc * tl.load(out_scale_inv)
294 + acc = tl.clamp(acc, FP8_MIN, FP8_MAX)
295 tl.store(out_ptrs,
296 acc,
297 mask=dim_mask[None, :] & (offs_m[:, None] < cur_batch_query_len))
@@ -743,6 +751,7 @@ def context_attention_fwd(q,
751 sliding_window=None,
752 sm_scale=None,
753 skip_decode=False,
754 + fp8_out_scale=None,
755 sinks=None):
756
757 q_dtype_is_f32 = q.dtype is torch.float32
@@ -793,6 +802,7 @@ def context_attention_fwd(q,
802
803 if alibi_slopes is not None:
804 assert sinks is None, "Sinks arg is not supported with alibi"
805 + assert fp8_out_scale is None, "FP8 output not supported with alibi"
806 # need to reduce num. blocks when using fp32
807 # due to increased use of GPU shared memory
808 # if q.dtype is torch.float32:
@@ -870,6 +880,7 @@ def context_attention_fwd(q,
880 sm_scale,
881 k_scale,
882 v_scale,
883 + 1.0 / fp8_out_scale if fp8_out_scale is not None else 1.0,
884 b_start_loc,
885 b_seq_len,
886 k_cache.shape[4],
@@ -905,6 +916,7 @@ def context_attention_fwd(q,
916 BLOCK_DMODEL_PADDED=Lk_padded,
917 SLIDING_WINDOW=sliding_window,
918 SKIP_DECODE=skip_decode,
919 + USE_FP8=fp8_out_scale is not None,
920 BLOCK_M=128,
921 BLOCK_N=64,
922 num_unroll_cache=4,
vllm/attention/ops/triton_unified_attention.py
+81 -58
@@ -10,9 +10,11 @@
10 import torch
11
12 from vllm.logger import init_logger
13 +from vllm.platforms import current_platform
14 from vllm.triton_utils import tl, triton
15
16 logger = init_logger(__name__)
17 +float8_info = torch.finfo(current_platform.fp8_dtype())
18
19
20 @triton.jit
@@ -48,47 +50,51 @@ def find_seq_idx(query_start_len_ptr, target_idx, num_seqs,
50
51 @triton.jit
52 def kernel_unified_attention_2d(
51 - output_ptr, # [num_tokens, num_query_heads, head_size]
52 - query_ptr, # [num_tokens, num_query_heads, head_size]
53 - key_cache_ptr, # [num_blks, blk_size, num_kv_heads, head_size]
54 - value_cache_ptr, # [num_blks, blk_size, num_kv_heads, head_size]
55 - sink_ptr, # [num_query_heads]
56 - block_tables_ptr, # [num_seqs, max_num_blocks_per_seq]
57 - seq_lens_ptr, # [num_seqs]
58 - alibi_slopes_ptr, # [num_query_heads]
59 - qq_bias_ptr, # [num_query_tokens, num_query_tokens]
60 - scale, # float32
61 - k_scale, # float32
62 - v_scale, # float32
63 - softcap, # float32
64 - num_query_heads: tl.constexpr, # int
65 - num_queries_per_kv: tl.constexpr, # int
66 - block_table_stride: tl.int64, # int
67 - query_stride_0: tl.int64, # int
68 - query_stride_1: tl.int64, # int, should be equal to head_size
69 - output_stride_0: tl.int64, # int
70 - output_stride_1: tl.int64, # int, should be equal to head_size
71 - qq_bias_stride_0: tl.int64, # int
72 - BLOCK_SIZE: tl.constexpr, # int
73 - HEAD_SIZE: tl.constexpr, # int
74 - HEAD_SIZE_PADDED: tl.constexpr, # int, must be power of 2
75 - USE_ALIBI_SLOPES: tl.constexpr, # bool
76 - USE_QQ_BIAS: tl.constexpr, # bool
77 - USE_SOFTCAP: tl.constexpr, # bool
78 - USE_SINKS: tl.constexpr, # bool
79 - SLIDING_WINDOW: tl.constexpr, # int
80 - stride_k_cache_0: tl.int64, # int
81 - stride_k_cache_1: tl.int64, # int
82 - stride_k_cache_2: tl.int64, # int
83 - stride_k_cache_3: tl.constexpr, # int
84 - stride_v_cache_0: tl.int64, # int
85 - stride_v_cache_1: tl.int64, # int
86 - stride_v_cache_2: tl.int64, # int
87 - stride_v_cache_3: tl.constexpr, # int
88 - query_start_len_ptr, # [num_seqs+1]
89 - BLOCK_Q: tl.constexpr, # int
90 - num_seqs: tl.int32,
91 - BLOCK_M: tl.constexpr, # int
53 + output_ptr, # [num_tokens, num_query_heads, head_size]
54 + query_ptr, # [num_tokens, num_query_heads, head_size]
55 + key_cache_ptr, # [num_blks, blk_size, num_kv_heads, head_size]
56 + value_cache_ptr, # [num_blks, blk_size, num_kv_heads, head_size]
57 + sink_ptr, # [num_query_heads]
58 + block_tables_ptr, # [num_seqs, max_num_blocks_per_seq]
59 + seq_lens_ptr, # [num_seqs]
60 + alibi_slopes_ptr, # [num_query_heads]
61 + qq_bias_ptr, # [num_query_tokens, num_query_tokens]
62 + scale, # float32
63 + k_scale, # float32
64 + v_scale, # float32
65 + out_scale, # float32
66 + softcap, # float32
67 + num_query_heads: tl.constexpr, # int
68 + num_queries_per_kv: tl.constexpr, # int
69 + block_table_stride: tl.int64, # int
70 + query_stride_0: tl.int64, # int
71 + query_stride_1: tl.int64, # int, should be equal to head_size
72 + output_stride_0: tl.int64, # int
73 + output_stride_1: tl.int64, # int, should be equal to head_size
74 + qq_bias_stride_0: tl.int64, # int
75 + BLOCK_SIZE: tl.constexpr, # int
76 + HEAD_SIZE: tl.constexpr, # int
77 + HEAD_SIZE_PADDED: tl.constexpr, # int, must be power of 2
78 + USE_ALIBI_SLOPES: tl.constexpr, # bool
79 + USE_QQ_BIAS: tl.constexpr, # bool
80 + USE_SOFTCAP: tl.constexpr, # bool
81 + USE_SINKS: tl.constexpr, # bool
82 + SLIDING_WINDOW: tl.constexpr, # int
83 + stride_k_cache_0: tl.int64, # int
84 + stride_k_cache_1: tl.int64, # int
85 + stride_k_cache_2: tl.int64, # int
86 + stride_k_cache_3: tl.constexpr, # int
87 + stride_v_cache_0: tl.int64, # int
88 + stride_v_cache_1: tl.int64, # int
89 + stride_v_cache_2: tl.int64, # int
90 + stride_v_cache_3: tl.constexpr, # int
91 + query_start_len_ptr, # [num_seqs+1]
92 + BLOCK_Q: tl.constexpr, # int
93 + num_seqs: tl.int32,
94 + BLOCK_M: tl.constexpr, # int
95 + USE_FP8: tl.constexpr, # bool
96 + FP8_MIN: tl.constexpr = float8_info.min,
97 + FP8_MAX: tl.constexpr = float8_info.max,
98 ):
99 q_block_global_idx = tl.program_id(0)
100 kv_head_idx = tl.program_id(1)
@@ -281,6 +287,9 @@ def kernel_unified_attention_2d(
287
288 # epilogue
289 acc = acc / L[:, None]
290 + if USE_FP8:
291 + acc = acc * tl.load(out_scale)
292 + acc = tl.clamp(acc, FP8_MIN, FP8_MAX)
293
294 output_offset = (query_offset_0[:, None] * output_stride_0 +
295 query_offset_1[:, None] * output_stride_1 +
@@ -552,23 +561,27 @@ def kernel_unified_attention_3d(
561
562 @triton.jit
563 def reduce_segments(
555 - output_ptr, # [num_tokens, num_query_heads, head_size]
556 - segm_output_ptr,
557 - #[num_tokens, num_query_heads, max_num_segments, head_size]
558 - segm_max_ptr, # [num_tokens, num_query_heads, max_num_segments]
559 - segm_expsum_ptr, # [num_tokens, num_query_heads, max_num_segments]
560 - seq_lens_ptr, # [num_seqs]
561 - num_seqs, # int
562 - num_query_heads: tl.constexpr, # int
563 - output_stride_0: tl.int64, # int
564 - output_stride_1: tl.int64, # int, should be equal to head_size
565 - block_table_stride: tl.int64, # int
566 - BLOCK_SIZE: tl.constexpr, # int
567 - HEAD_SIZE: tl.constexpr, # int, must be power of 2
568 - HEAD_SIZE_PADDED: tl.constexpr, # int, must be power of 2
569 - query_start_len_ptr, # [num_seqs+1]
570 - BLOCK_Q: tl.constexpr, # int
571 - NUM_SEGMENTS_PER_SEQ: tl.constexpr, # int
564 + output_ptr, # [num_tokens, num_query_heads, head_size]
565 + segm_output_ptr,
566 + #[num_tokens, num_query_heads, max_num_segments, head_size]
567 + segm_max_ptr, # [num_tokens, num_query_heads, max_num_segments]
568 + segm_expsum_ptr, # [num_tokens, num_query_heads, max_num_segments]
569 + seq_lens_ptr, # [num_seqs]
570 + num_seqs, # int
571 + num_query_heads: tl.constexpr, # int
572 + out_scale_inv, # float32
573 + output_stride_0: tl.int64, # int
574 + output_stride_1: tl.int64, # int, should be equal to head_size
575 + block_table_stride: tl.int64, # int
576 + BLOCK_SIZE: tl.constexpr, # int
577 + HEAD_SIZE: tl.constexpr, # int, must be power of 2
578 + HEAD_SIZE_PADDED: tl.constexpr, # int, must be power of 2
579 + query_start_len_ptr, # [num_seqs+1]
580 + BLOCK_Q: tl.constexpr, # int
581 + NUM_SEGMENTS_PER_SEQ: tl.constexpr, # int
582 + USE_FP8: tl.constexpr, # bool
583 + FP8_MIN: tl.constexpr = float8_info.min,
584 + FP8_MAX: tl.constexpr = float8_info.max,
585 ):
586 query_token_idx = tl.program_id(0)
587 query_head_idx = tl.program_id(1)
@@ -624,6 +637,10 @@ def reduce_segments(
637 # safely divide by overall_expsum, returning 0.0 if overall_expsum is 0
638 acc = tl.where(overall_expsum == 0.0, 0.0, acc_sum / overall_expsum)
639
640 + if USE_FP8:
641 + acc = acc * tl.load(out_scale_inv)
642 + acc = tl.clamp(acc, FP8_MIN, FP8_MAX)
643 +
644 # write result
645 output_offset = (query_token_idx * output_stride_0 +
646 query_head_idx * output_stride_1 +
@@ -649,6 +666,7 @@ def unified_attention(
666 k_descale,
667 v_descale,
668 alibi_slopes=None,
669 + output_scale=None,
670 qq_bias=None,
671 # Optional tensor for sinks
672 sinks=None,
@@ -707,6 +725,7 @@ def unified_attention(
725 scale=softmax_scale,
726 k_scale=k_descale,
727 v_scale=v_descale,
728 + out_scale=1 / output_scale if output_scale is not None else 1.0,
729 softcap=softcap,
730 num_query_heads=num_query_heads,
731 num_queries_per_kv=num_queries_per_kv,
@@ -736,6 +755,7 @@ def unified_attention(
755 BLOCK_Q=BLOCK_Q,
756 num_seqs=num_seqs,
757 BLOCK_M=BLOCK_M,
758 + USE_FP8=output_scale is not None,
759 )
760 else:
761 # for initial version, NUM_SEGMENTS = 16 is chosen as a default
@@ -819,6 +839,8 @@ def unified_attention(
839 seq_lens_ptr=seqused_k,
840 num_seqs=num_seqs,
841 num_query_heads=num_query_heads,
842 + out_scale_inv=1 /
843 + output_scale if output_scale is not None else 1.0,
844 output_stride_0=out.stride(0),
845 output_stride_1=out.stride(1),
846 block_table_stride=block_table.stride(0),
@@ -828,4 +850,5 @@ def unified_attention(
850 query_start_len_ptr=cu_seqlens_q,
851 BLOCK_Q=BLOCK_Q,
852 NUM_SEGMENTS_PER_SEQ=NUM_SEGMENTS,
853 + USE_FP8=output_scale is not None,
854 )
vllm/compilation/backends.py
+2 -1
@@ -454,11 +454,12 @@ class VllmBackend:
454 inductor_config = config.inductor_compile_config
455 PASS_KEY = "post_grad_custom_post_pass"
456 if PASS_KEY in inductor_config:
457 - # Config should automatically wrap all inductor passes
457 if isinstance(inductor_config[PASS_KEY], PostGradPassManager):
458 + # PassManager already added to config, make sure it's correct
459 assert (inductor_config[PASS_KEY].uuid() ==
460 self.post_grad_pass_manager.uuid())
461 else:
462 + # Config should automatically wrap all inductor passes
463 assert isinstance(inductor_config[PASS_KEY], InductorPass)
464 self.post_grad_pass_manager.add(inductor_config[PASS_KEY])
465 inductor_config[PASS_KEY] = self.post_grad_pass_manager
vllm/compilation/fusion_attn.py
+22 -9
@@ -39,6 +39,7 @@ class AttentionQuantPattern(ABC):
39 self,
40 layer: Attention,
41 quant_key: QuantKey,
42 + dtype: torch.dtype,
43 ):
44 self.layer = layer
45 self.layer_name = layer.layer_name
@@ -46,11 +47,16 @@ class AttentionQuantPattern(ABC):
47 self.head_size = layer.head_size
48 self.quant_key = quant_key
49 self.quant_dtype = quant_key.dtype
50 + self.dtype = dtype
51
52 assert self.quant_key in QUANT_OPS, \
53 f"unsupported quantization scheme {self.quant_key}"
54 self.QUANT_OP = QUANT_OPS[self.quant_key]
55
56 + def empty(self, *args, **kwargs):
57 + kwargs = {'dtype': self.dtype, 'device': "cuda", **kwargs}
58 + return torch.empty(*args, **kwargs)
59 +
60 def empty_quant(self, *args, **kwargs):
61 kwargs = {'dtype': self.quant_dtype, 'device': "cuda", **kwargs}
62 return torch.empty(*args, **kwargs)
@@ -91,12 +97,13 @@ class AttentionFp8StaticQuantPattern(AttentionQuantPattern):
97 def __init__(
98 self,
99 layer: Attention,
100 + dtype: torch.dtype,
101 symmetric: bool = True,
102 ):
103 quant_key = QuantKey(dtype=FP8_DTYPE,
104 scale=kStaticTensorScale,
105 symmetric=symmetric)
99 - super().__init__(layer, quant_key)
106 + super().__init__(layer, quant_key, dtype)
107
108 def _register(self, pm_pass: PatternMatcherPass):
109
@@ -139,10 +146,14 @@ class AttentionFp8StaticQuantPattern(AttentionQuantPattern):
146 return RESHAPE_OP(at1[1], [-1, self.num_heads * self.head_size])
147
148 inputs = [
142 - empty_bf16(5, self.num_heads, self.head_size), # q
143 - empty_bf16(5, self.num_heads, self.head_size), # k
144 - empty_bf16(5, self.num_heads, self.head_size), # v
145 - empty_bf16(5, self.num_heads, self.head_size), # attn_output
149 + self.empty(5, self.num_heads, self.head_size,
150 + dtype=self.dtype), # q
151 + self.empty(5, self.num_heads, self.head_size,
152 + dtype=self.dtype), # k
153 + self.empty(5, self.num_heads, self.head_size,
154 + dtype=self.dtype), # v
155 + self.empty(5, self.num_heads, self.head_size,
156 + dtype=self.dtype), # attn_output
157 self.empty_quant(5,
158 self.num_heads * self.head_size), # quant_output
159 empty_fp32(1, 1) # scale
@@ -165,8 +176,8 @@ class AttentionNvfp4QuantPattern(AttentionQuantPattern):
176 will be passed into Attention op as the `output_scale` argument.
177 """
178
168 - def __init__(self, layer: Attention):
169 - super().__init__(layer, kNvfp4Quant)
179 + def __init__(self, layer: Attention, dtype: torch.dtype):
180 + super().__init__(layer, kNvfp4Quant, dtype)
181
182 def _register(self, pm_pass: PatternMatcherPass):
183
@@ -255,12 +266,14 @@ class AttnFusionPass(VllmInductorPass):
266
267 attn_layers = get_layers_from_vllm_config(config, Attention)
268 for layer_name, layer in attn_layers.items():
258 - pattern_fp8 = AttentionFp8StaticQuantPattern(layer)
269 + pattern_fp8 = AttentionFp8StaticQuantPattern(
270 + layer, config.model_config.dtype)
271 pattern_fp8.register_if_supported(self.patterns)
272
273 if current_platform.is_cuda() and hasattr(torch.ops._C,
274 "scaled_fp4_quant"):
263 - pattern_nvfp4 = AttentionNvfp4QuantPattern(layer)
275 + pattern_nvfp4 = AttentionNvfp4QuantPattern(
276 + layer, config.model_config.dtype)
277 pattern_nvfp4.register_if_supported(self.patterns)
278
279 if len(attn_layers) == 0:
vllm/model_executor/layers/quantization/utils/w8a8_utils.py
+19 -19
@@ -171,10 +171,12 @@ def flashinfer_w8a8_scaled_mm(*, qinput: torch.Tensor, weight: torch.Tensor,
171 bias=bias)
172
173
174 -def rocm_per_tensor_w8a8_scaled_mm_impl(
175 - qinput: torch.Tensor, weight: torch.Tensor, out_dtype: torch.dtype,
176 - scale_a: torch.Tensor, scale_b: torch.Tensor, bias: torch.Tensor,
177 - input_2d: torch.Tensor) -> torch.Tensor:
174 +def rocm_per_tensor_w8a8_scaled_mm_impl(qinput: torch.Tensor,
175 + weight: torch.Tensor,
176 + out_dtype: torch.dtype,
177 + scale_a: torch.Tensor,
178 + scale_b: torch.Tensor,
179 + bias: torch.Tensor) -> torch.Tensor:
180 from vllm.platforms.rocm import on_mi3xx
181 if envs.VLLM_ROCM_USE_SKINNY_GEMM and on_mi3xx(
182 ) and qinput.shape[0] == 1 and qinput.shape[1] % 16 == 0:
@@ -190,10 +192,12 @@ def rocm_per_tensor_w8a8_scaled_mm_impl(
192 return output
193
194
193 -def rocm_per_tensor_w8a8_scaled_mm_fake(
194 - qinput: torch.Tensor, weight: torch.Tensor, out_dtype: torch.dtype,
195 - scale_a: torch.Tensor, scale_b: torch.Tensor, bias: torch.Tensor,
196 - input_2d: torch.Tensor) -> torch.Tensor:
195 +def rocm_per_tensor_w8a8_scaled_mm_fake(qinput: torch.Tensor,
196 + weight: torch.Tensor,
197 + out_dtype: torch.dtype,
198 + scale_a: torch.Tensor,
199 + scale_b: torch.Tensor,
200 + bias: torch.Tensor) -> torch.Tensor:
201 return qinput.new_empty((*qinput.shape[:-1], weight.shape[1]),
202 dtype=out_dtype)
203
@@ -203,11 +207,10 @@ def rocm_per_tensor_w8a8_scaled_mm(*, qinput: torch.Tensor,
207 out_dtype: torch.dtype,
208 scale_a: torch.Tensor,
209 scale_b: torch.Tensor, bias: torch.Tensor,
206 - input_2d: torch.Tensor,
210 output_shape: list) -> torch.Tensor:
211 output = torch.ops.vllm.rocm_per_tensor_w8a8_scaled_mm_impl(
209 - qinput, weight, out_dtype, scale_a, scale_b, bias, input_2d)
210 - return torch.narrow(output, 0, 0, input_2d.shape[0]).view(*output_shape)
212 + qinput, weight, out_dtype, scale_a, scale_b, bias)
213 + return torch.narrow(output, 0, 0, qinput.shape[0]).view(*output_shape)
214
215
216 direct_register_custom_op(
@@ -224,7 +227,6 @@ def torch_per_tensor_w8a8_scaled_mm(*, qinput: torch.Tensor,
227 out_dtype: torch.dtype,
228 scale_a: torch.Tensor,
229 scale_b: torch.Tensor, bias: torch.Tensor,
227 - input_2d: torch.Tensor,
230 output_shape: list) -> torch.Tensor:
231 output = torch._scaled_mm(qinput,
232 weight,
@@ -237,7 +239,7 @@ def torch_per_tensor_w8a8_scaled_mm(*, qinput: torch.Tensor,
239 if type(output) is tuple and len(output) == 2:
240 output = output[0]
241
240 - return torch.narrow(output, 0, 0, input_2d.shape[0]).view(*output_shape)
242 + return torch.narrow(output, 0, 0, qinput.shape[0]).view(*output_shape)
243
244
245 def torch_per_token_w8a8_scaled_mm(*, qinput: torch.Tensor,
@@ -245,7 +247,7 @@ def torch_per_token_w8a8_scaled_mm(*, qinput: torch.Tensor,
247 out_dtype: torch.dtype,
248 scale_a: torch.Tensor,
249 scale_b: torch.Tensor, bias: torch.Tensor,
248 - input_2d: torch.Tensor, output_shape: list,
250 + output_shape: list,
251 **kwargs) -> torch.Tensor:
252 # Note: Callers of this function should check USE_ROWWISE_TORCH_SCALED_MM
253 # when using it.
@@ -265,7 +267,7 @@ def torch_per_token_w8a8_scaled_mm(*, qinput: torch.Tensor,
267 scale_b=scale_b.t(),
268 bias=bias)
269
268 - output = torch.narrow(output, 0, 0, input_2d.shape[0])
270 + output = torch.narrow(output, 0, 0, qinput.shape[0])
271 output = output.view(*output_shape)
272 return output
273
@@ -275,7 +277,6 @@ def torch_channelwise_w8a8_scaled_mm(*, qinput: torch.Tensor,
277 out_dtype: torch.dtype,
278 scale_a: torch.Tensor,
279 scale_b: torch.Tensor, bias: torch.Tensor,
278 - input_2d: torch.Tensor,
280 output_shape: list,
281 **kwargs) -> torch.Tensor:
282 # Use unfused DQ due to limitations with scaled_mm
@@ -305,8 +306,8 @@ def torch_channelwise_w8a8_scaled_mm(*, qinput: torch.Tensor,
306 if type(output) is tuple and len(output) == 2:
307 output = output[0]
308 # Unpad (undo num_token_padding)
308 - output = torch.narrow(output, 0, 0, input_2d.shape[0])
309 - x_scale = torch.narrow(scale_a, 0, 0, input_2d.shape[0])
309 + output = torch.narrow(output, 0, 0, qinput.shape[0])
310 + x_scale = torch.narrow(scale_a, 0, 0, qinput.shape[0])
311
312 # DQ
313 # C = sw * sx * (X * W) + bias
@@ -430,7 +431,6 @@ class Fp8LinearOp:
431 scale_a=x_scale,
432 scale_b=weight_scale,
433 bias=bias,
433 - input_2d=input_2d,
434 output_shape=output_shape)
435
436
vllm/v1/attention/backends/triton_attn.py
+9 -3
@@ -15,6 +15,8 @@ from vllm.attention.ops.chunked_prefill_paged_decode import (
15 from vllm.attention.ops.paged_attn import PagedAttention
16 from vllm.config import VllmConfig
17 from vllm.logger import init_logger
18 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
19 + QuantKey, kFp8StaticTensorSym)
20 from vllm.platforms import current_platform
21 from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
22 from vllm.v1.attention.backends.utils import (AttentionCGSupport,
@@ -202,6 +204,9 @@ def use_aiter_unified_attention() -> bool:
204
205 class TritonAttentionImpl(AttentionImpl):
206
207 + def fused_output_quant_supported(self, quant_key: QuantKey):
208 + return quant_key == kFp8StaticTensorSym
209 +
210 def __init__(
211 self,
212 num_heads: int,
@@ -297,9 +302,9 @@ class TritonAttentionImpl(AttentionImpl):
302 """
303 assert output is not None, "Output tensor must be provided."
304
300 - if output_scale is not None or output_block_scale is not None:
305 + if output_block_scale is not None:
306 raise NotImplementedError(
302 - "fused output quantization is not yet supported"
307 + "fused block_scale output quantization is not yet supported"
308 " for TritonAttentionImpl")
309
310 if attn_metadata is None:
@@ -394,6 +399,7 @@ class TritonAttentionImpl(AttentionImpl):
399 alibi_slopes=self.alibi_slopes,
400 sliding_window=self.sliding_window[0],
401 sm_scale=self.scale,
402 + output_scale=output_scale,
403 sinks=self.sinks,
404 )
405
@@ -419,6 +425,6 @@ class TritonAttentionImpl(AttentionImpl):
425 k_descale=layer._k_scale.expand(descale_shape),
426 v_descale=layer._v_scale.expand(descale_shape),
427 sinks=self.sinks,
422 - )
428 + output_scale=output_scale)
429
430 return output