@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 24d0c9e6ed
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

[NVIDIA][torch.compile] Support Flashinfer TRTLLM FP8-q/kv NVFP4-out Attention Kernel (#22703)

Signed-off-by: elvischenv <219235043+elvischenv@users.noreply.github.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>

elvischenv committed Aug 23, 2025 at 06:09 UTC 24d0c9e6edc4299f62053ee0cb0154ce86b08cb8
27 files changed +598 -202
benchmarks/kernels/benchmark_trtllm_decode_attention.py
+37 -13
@@ -9,8 +9,11 @@ from typing import Optional
9 import flashinfer
10 import torch
11
12 +from vllm.utils import round_up
13 +
14 FLOAT32_BYTES = torch.finfo(torch.float).bits // 8
15 FP8_DTYPE = torch.float8_e4m3fn
16 +FP4_DTYPE = torch.uint8
17
18
19 def to_float8(x, dtype=torch.float8_e4m3fn):
@@ -61,13 +64,13 @@ def benchmark_decode(
64 else:
65 raise ValueError(f"Invalid kv_layout: {kv_layout}")
66
64 - query = torch.randn(batch_size, num_qo_heads, head_size, dtype=dtype)
67 + # Always using 1.0 scale to reflect the real perf in benchmarking
68 + q_scale = 1.0
69 + ref_query = torch.randn(batch_size, num_qo_heads, head_size, dtype=dtype)
70 if q_quant_dtype == FP8_DTYPE:
66 - query, q_scale = to_float8(query)
67 - ref_query = query.to(dtype) * q_scale
71 + query, _ = to_float8(ref_query)
72 else:
69 - q_scale = 1.0
70 - ref_query = query
73 + query = ref_query
74
75 kv_lens = torch.randint(1, max_seq_len, (batch_size,), dtype=torch.int32)
76 kv_lens[-1] = max_seq_len
@@ -75,14 +78,13 @@ def benchmark_decode(
78 seq_lens = kv_lens
79 max_seq_len = torch.max(seq_lens).item()
80
78 - kv_cache = torch.randn(kv_cache_shape, dtype=dtype)
81 + # Always using 1.0 scale to reflect the real perf in benchmarking
82 + k_scale = v_scale = 1.0
83 + ref_kv_cache = torch.randn(kv_cache_shape, dtype=dtype)
84 if kv_quant_dtype == FP8_DTYPE:
80 - kv_cache, kv_scale = to_float8(kv_cache)
81 - ref_kv_cache = kv_cache.to(dtype) * kv_scale
85 + kv_cache, _ = to_float8(ref_kv_cache)
86 else:
83 - kv_scale = 1.0
84 - ref_kv_cache = kv_cache
85 - k_scale = v_scale = kv_scale
87 + kv_cache = ref_kv_cache
88
89 max_num_blocks_per_seq = (max_seq_len + block_size - 1) // block_size
90 block_tables = torch.randint(
@@ -142,11 +144,31 @@ def benchmark_decode(
144 return sum(times) / len(times), torch.std(torch.tensor(times))
145
146 o_scale = 1.0
147 + o_sf_scale = None
148 output_baseline = torch.empty(ref_query.shape, dtype=dtype)
146 - output_trtllm = torch.empty(query.shape, dtype=o_quant_dtype)
149 + if o_quant_dtype == FP4_DTYPE:
150 + o_sf_scale = 500.0
151 + output_trtllm = flashinfer.utils.FP4Tensor(
152 + torch.empty(query.shape[:-1] + (query.shape[-1] // 2,), dtype=torch.uint8),
153 + torch.empty(
154 + (
155 + round_up(query.shape[0], 128),
156 + round_up(query.shape[1] * query.shape[2] // 16, 4),
157 + ),
158 + dtype=torch.float8_e4m3fn,
159 + ),
160 + )
161 + else:
162 + output_trtllm = torch.empty(query.shape, dtype=o_quant_dtype)
163
164 def baseline_decode():
149 - return wrapper.run(ref_query, ref_kv_cache, out=output_baseline)
165 + return wrapper.run(
166 + ref_query,
167 + ref_kv_cache,
168 + k_scale=k_scale,
169 + v_scale=v_scale,
170 + out=output_baseline,
171 + )
172
173 def trtllm_decode():
174 return flashinfer.decode.trtllm_batch_decode_with_kv_cache(
@@ -158,6 +180,7 @@ def benchmark_decode(
180 max_seq_len=max_seq_len,
181 bmm1_scale=q_scale * k_scale * sm_scale,
182 bmm2_scale=v_scale / o_scale,
183 + o_sf_scale=o_sf_scale,
184 out=output_trtllm,
185 )
186
@@ -237,6 +260,7 @@ if __name__ == "__main__":
260 (None, None, None),
261 (None, FP8_DTYPE, None),
262 (FP8_DTYPE, FP8_DTYPE, FP8_DTYPE),
263 + (FP8_DTYPE, FP8_DTYPE, FP4_DTYPE),
264 ]
265
266 for quant_dtype in quant_dtypes:
benchmarks/kernels/benchmark_trtllm_prefill_attention.py
+39 -13
@@ -9,8 +9,11 @@ from typing import Optional
9 import flashinfer
10 import torch
11
12 +from vllm.utils import round_up
13 +
14 FLOAT32_BYTES = torch.finfo(torch.float).bits // 8
15 FP8_DTYPE = torch.float8_e4m3fn
16 +FP4_DTYPE = torch.uint8
17
18
19 def to_float8(x, dtype=torch.float8_e4m3fn):
@@ -72,13 +75,15 @@ def benchmark_prefill(
75 ]
76 )
77
75 - query = torch.randn(torch.sum(q_lens).item(), num_qo_heads, head_size, dtype=dtype)
78 + # Always using 1.0 scale to reflect the real perf in benchmarking
79 + q_scale = 1.0
80 + ref_query = torch.randn(
81 + torch.sum(q_lens).item(), num_qo_heads, head_size, dtype=dtype
82 + )
83 if q_quant_dtype == FP8_DTYPE:
77 - query, q_scale = to_float8(query)
78 - ref_query = query.to(dtype) * q_scale
84 + query, _ = to_float8(ref_query)
85 else:
80 - q_scale = 1.0
81 - ref_query = query
86 + query = ref_query
87
88 kv_lens = torch.randint(0, max_kv_len, (batch_size,), dtype=torch.int32)
89 kv_lens[-1] = max_kv_len
@@ -86,14 +91,13 @@ def benchmark_prefill(
91 seq_lens = kv_lens + q_lens
92 max_seq_len = torch.max(seq_lens).item()
93
89 - kv_cache = torch.randn(kv_cache_shape, dtype=dtype)
94 + # Always using 1.0 scale to reflect the real perf in benchmarking
95 + k_scale = v_scale = 1.0
96 + ref_kv_cache = torch.randn(kv_cache_shape, dtype=dtype)
97 if kv_quant_dtype == FP8_DTYPE:
91 - kv_cache, kv_scale = to_float8(kv_cache)
92 - ref_kv_cache = kv_cache.to(dtype) * kv_scale
98 + kv_cache, _ = to_float8(ref_kv_cache)
99 else:
94 - kv_scale = 1.0
95 - ref_kv_cache = kv_cache
96 - k_scale = v_scale = kv_scale
100 + kv_cache = ref_kv_cache
101
102 max_num_blocks_per_seq = (max_seq_len + block_size - 1) // block_size
103 block_tables = torch.randint(
@@ -152,11 +156,31 @@ def benchmark_prefill(
156 return sum(times) / len(times), torch.std(torch.tensor(times))
157
158 o_scale = 1.0
159 + o_sf_scale = None
160 output_baseline = torch.empty(ref_query.shape, dtype=dtype)
156 - output_trtllm = torch.empty(query.shape, dtype=o_quant_dtype)
161 + if o_quant_dtype == FP4_DTYPE:
162 + o_sf_scale = 500.0
163 + output_trtllm = flashinfer.utils.FP4Tensor(
164 + torch.empty(query.shape[:-1] + (query.shape[-1] // 2,), dtype=torch.uint8),
165 + torch.empty(
166 + (
167 + round_up(query.shape[0], 128),
168 + round_up(query.shape[1] * query.shape[2] // 16, 4),
169 + ),
170 + dtype=torch.float8_e4m3fn,
171 + ),
172 + )
173 + else:
174 + output_trtllm = torch.empty(query.shape, dtype=o_quant_dtype)
175
176 def baseline_prefill():
159 - return wrapper.run(ref_query, ref_kv_cache, out=output_baseline)
177 + return wrapper.run(
178 + ref_query,
179 + ref_kv_cache,
180 + k_scale=k_scale,
181 + v_scale=v_scale,
182 + out=output_baseline,
183 + )
184
185 def trtllm_prefill():
186 return flashinfer.prefill.trtllm_batch_context_with_kv_cache(
@@ -172,6 +196,7 @@ def benchmark_prefill(
196 batch_size=batch_size,
197 cum_seq_lens_q=q_indptr,
198 cum_seq_lens_kv=kv_indptr,
199 + o_sf_scale=o_sf_scale,
200 out=output_trtllm,
201 )
202
@@ -250,6 +275,7 @@ if __name__ == "__main__":
275 # (q_quant_dtype, kv_quant_dtype, o_quant_dtype)
276 (None, None, None),
277 (FP8_DTYPE, FP8_DTYPE, FP8_DTYPE),
278 + (FP8_DTYPE, FP8_DTYPE, FP4_DTYPE),
279 ]
280
281 for quant_dtype in quant_dtypes:
tests/compile/test_functionalization.py
+3 -2
@@ -8,11 +8,12 @@ import vllm.envs as envs
8 from vllm import LLM, SamplingParams
9 from vllm.compilation.activation_quant_fusion import ActivationQuantFusionPass
10 from vllm.compilation.fix_functionalization import FixFunctionalizationPass
11 -from vllm.compilation.fusion import (FUSED_OPS, FusionPass, QuantKey,
12 - kFp8DynamicTokenSym, kFp8StaticTensorSym)
11 +from vllm.compilation.fusion import FUSED_OPS, FusionPass
12 from vllm.compilation.fx_utils import find_auto_fn, find_auto_fn_maybe, is_func
13 from vllm.compilation.noop_elimination import NoOpEliminationPass
14 from vllm.config import CompilationConfig, PassConfig, VllmConfig
15 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
16 + QuantKey, kFp8DynamicTokenSym, kFp8StaticTensorSym)
17
18 from .backend import TestBackend
19
tests/compile/test_fusion.py
+5 -5
@@ -7,11 +7,13 @@ import torch
7 import vllm.envs as envs
8 import vllm.plugins
9 from vllm.compilation.fusion import (FUSED_OPS, QUANT_OPS, FusedRMSQuantKey,
10 - FusionPass, GroupShape, QuantKey)
10 + FusionPass)
11 from vllm.compilation.noop_elimination import NoOpEliminationPass
12 from vllm.config import (CompilationConfig, CompilationLevel, PassConfig,
13 VllmConfig)
14 from vllm.model_executor.layers.layernorm import RMSNorm
15 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
16 + GroupShape, QuantKey, ScaleDesc)
17 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
18 CUTLASS_FP8_SUPPORTED, Fp8LinearOp, maybe_create_device_identity)
19 from vllm.platforms import current_platform
@@ -30,10 +32,8 @@ class TestModel(torch.nn.Module):
32 self.norm = [RMSNorm(hidden_size, eps) for _ in range(3)]
33 self.wscale = [torch.rand(1, dtype=torch.float32) for _ in range(2)]
34 group_shape = GroupShape.PER_TENSOR if static else GroupShape.PER_TOKEN
33 - self.key = QuantKey(dtype=FP8_DTYPE,
34 - static=static,
35 - group_shape=group_shape,
36 - symmetric=True)
35 + quant_scale = ScaleDesc(torch.float32, static, group_shape)
36 + self.key = QuantKey(dtype=FP8_DTYPE, scale=quant_scale, symmetric=True)
37 if static:
38 self.scale = [torch.rand(1, dtype=torch.float32) for _ in range(2)]
39 else:
tests/compile/test_fusion_attn.py
+113 -39
@@ -11,9 +11,10 @@ from tests.models.utils import check_outputs_equal
11 from tests.v1.attention.utils import (BatchSpec, _Backend,
12 create_common_attn_metadata)
13 from vllm import LLM, SamplingParams
14 +from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
15 from vllm.attention import Attention
16 from vllm.attention.selector import global_force_attn_backend_context_manager
16 -from vllm.compilation.fusion import QUANT_OPS, QuantKey, kFp8StaticTensorSym
17 +from vllm.compilation.fusion import QUANT_OPS
18 from vllm.compilation.fusion_attn import ATTN_OP, AttnFusionPass
19 from vllm.compilation.fx_utils import find_op_nodes
20 from vllm.compilation.noop_elimination import NoOpEliminationPass
@@ -22,13 +23,14 @@ from vllm.config import (CacheConfig, CompilationConfig, CompilationLevel,
23 set_current_vllm_config)
24 from vllm.forward_context import get_forward_context, set_forward_context
25 from vllm.model_executor.layers.quantization.utils.quant_utils import (
25 - GroupShape)
26 + QuantKey, kFp8StaticTensorSym, kNvfp4Quant)
27 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
28 Fp8LinearOp)
29 from vllm.platforms import current_platform
30 from vllm.v1.kv_cache_interface import AttentionSpec
31
32 FP8_DTYPE = current_platform.fp8_dtype()
33 +FP4_DTYPE = torch.uint8
34
35 # globals needed for string-import custom Dynamo backend field
36 backend: Optional[TestBackend] = None
@@ -105,9 +107,7 @@ def test_attention_fusion(example_prompts, monkeypatch, model: str,
107
108 # check support
109 attn_fusion_supported = [
108 - layer.impl.fused_output_quant_supported(quant_key.dtype,
109 - quant_key.static,
110 - quant_key.group_shape)
110 + layer.impl.fused_output_quant_supported(quant_key)
111 for key, layer in compile_config.static_forward_context.items()
112 ]
113
@@ -149,12 +149,12 @@ def test_attention_fusion(example_prompts, monkeypatch, model: str,
149 backend = None
150
151
152 -class TestAttentionStaticQuantPatternModel(torch.nn.Module):
153 - """Test model for AttentionStaticQuantPattern fusion."""
152 +class AttentionQuantPatternModel(torch.nn.Module):
153 + """Base model for AttentionQuantPattern fusion."""
154
155 def __init__(self, num_qo_heads: int, num_kv_heads: int, head_size: int,
156 kv_cache_dtype: torch.dtype, device: torch.device,
157 - vllm_config: VllmConfig):
157 + vllm_config: VllmConfig, **kwargs):
158 super().__init__()
159 self.num_qo_heads = num_qo_heads
160 self.num_kv_heads = num_kv_heads
@@ -172,11 +172,6 @@ class TestAttentionStaticQuantPatternModel(torch.nn.Module):
172 prefix="model.layers.0.self_attn.attn",
173 )
174
175 - self.fp8_linear = Fp8LinearOp(
176 - act_quant_static=True, act_quant_group_shape=GroupShape.PER_TENSOR)
177 - self.wscale = torch.tensor([1.0], dtype=torch.float32)
178 - self.scale = torch.tensor([1.0], dtype=torch.float32)
179 -
175 self.block_size = 16
176
177 # Initialize attn MetadataBuilder
@@ -230,23 +225,86 @@ class TestAttentionStaticQuantPatternModel(torch.nn.Module):
225
226 return self.attn_metadata
227
233 - def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
234 - w: torch.Tensor):
228 +
229 +class TestAttentionFp8StaticQuantPatternModel(AttentionQuantPatternModel):
230 + """Test model for AttentionFp8StaticQuantPattern fusion."""
231 +
232 + quant_key = kFp8StaticTensorSym
233 +
234 + def __init__(self, *args, **kwargs):
235 + super().__init__(*args, **kwargs)
236 +
237 + self.fp8_linear = Fp8LinearOp(
238 + act_quant_static=self.quant_key.scale.static,
239 + act_quant_group_shape=self.quant_key.scale.group_shape)
240 +
241 + hidden_size = self.num_qo_heads * self.head_size
242 + self.w = kwargs.get(
243 + "w", {
244 + "weight":
245 + torch.randn(hidden_size, hidden_size).to(
246 + dtype=FP8_DTYPE, device=self.device).t(),
247 + "wscale":
248 + torch.tensor([1.0], dtype=torch.float32, device=self.device),
249 + "scale":
250 + torch.tensor([1.0], dtype=torch.float32, device=self.device),
251 + })
252 +
253 + def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
254 """Forward pass that creates the pattern to be fused."""
255 attn_output = self.attn(q, k, v)
256 return self.fp8_linear.apply(input=attn_output,
238 - weight=w,
239 - weight_scale=self.wscale,
240 - input_scale=self.scale)
257 + weight=self.w["weight"],
258 + weight_scale=self.w["wscale"],
259 + input_scale=self.w["scale"])
260 +
261 +
262 +class TestAttentionNvfp4QuantPatternModel(AttentionQuantPatternModel):
263 + """Test model for AttentionNvfp4QuantPattern fusion."""
264 +
265 + quant_key = kNvfp4Quant
266 +
267 + def __init__(self, *args, **kwargs):
268 + super().__init__(*args, **kwargs)
269 +
270 + hidden_size = self.num_qo_heads * self.head_size
271 + self.w = kwargs.get(
272 + "w", {
273 + "weight":
274 + torch.randint(256, (hidden_size, hidden_size // 2),
275 + dtype=FP4_DTYPE,
276 + device=self.device),
277 + "wscale_swizzled":
278 + torch.randn(hidden_size, hidden_size // 16).to(
279 + dtype=FP8_DTYPE, device=self.device),
280 + "wscale":
281 + torch.tensor([500], dtype=torch.float32, device=self.device),
282 + "scale":
283 + torch.tensor([0.002], dtype=torch.float32, device=self.device),
284 + })
285 +
286 + def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
287 + """Forward pass that creates the pattern to be fused."""
288 + attn_output = self.attn(q, k, v)
289 + quant_output, output_block_scale = scaled_fp4_quant(
290 + attn_output, 1 / self.w["scale"])
291 + return cutlass_scaled_fp4_mm(a=quant_output,
292 + b=self.w["weight"],
293 + block_scale_a=output_block_scale,
294 + block_scale_b=self.w["wscale_swizzled"],
295 + alpha=self.w["scale"] * self.w["wscale"],
296 + out_dtype=attn_output.dtype)
297
298
299 @pytest.mark.parametrize("num_qo_heads, num_kv_heads", [(64, 8), (40, 8)])
300 @pytest.mark.parametrize("head_size", [128])
301 @pytest.mark.parametrize("batch_size", [7, 256, 533])
302 @pytest.mark.parametrize("dtype", [torch.bfloat16])
247 -@pytest.mark.parametrize(
248 - "model_name, quant_key",
249 - [("nvidia/Llama-4-Scout-17B-16E-Instruct-FP8", kFp8StaticTensorSym)])
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")
310 @pytest.mark.skipif(not current_platform.supports_fp8(), reason="Need FP8")
@@ -255,8 +313,8 @@ class TestAttentionStaticQuantPatternModel(torch.nn.Module):
313 def test_attention_quant_pattern(num_qo_heads: int, num_kv_heads: int,
314 head_size: int, batch_size: int,
315 dtype: torch.dtype, model_name: str,
258 - quant_key: QuantKey, backend: _Backend,
259 - monkeypatch, dist_init):
316 + model_class: type[AttentionQuantPatternModel],
317 + backend: _Backend, monkeypatch, dist_init):
318 """Test AttentionStaticQuantPattern fusion pass"""
319
320 monkeypatch.setenv("VLLM_USE_V1", "1")
@@ -277,8 +335,10 @@ def test_attention_quant_pattern(num_qo_heads: int, num_kv_heads: int,
335 cache_config=CacheConfig(cache_dtype="fp8"))
336
337 # Create test inputs
280 - hidden_size = num_qo_heads * head_size
281 - q = torch.randn(batch_size, hidden_size, dtype=dtype, device=device)
338 + q = torch.randn(batch_size,
339 + num_qo_heads * head_size,
340 + dtype=dtype,
341 + device=device)
342 k = torch.randn(batch_size,
343 num_kv_heads * head_size,
344 dtype=dtype,
@@ -287,7 +347,6 @@ def test_attention_quant_pattern(num_qo_heads: int, num_kv_heads: int,
347 num_kv_heads * head_size,
348 dtype=dtype,
349 device=device)
290 - linear_w = torch.randn(hidden_size, hidden_size).to(FP8_DTYPE).t()
350
351 # Mark first dimension as dynamic for realistic testing
352 torch._dynamo.mark_dynamic(q, 0)
@@ -299,9 +358,12 @@ def test_attention_quant_pattern(num_qo_heads: int, num_kv_heads: int,
358 with set_current_vllm_config(vllm_config_unfused), set_forward_context(
359 attn_metadata=None, vllm_config=vllm_config_unfused
360 ), global_force_attn_backend_context_manager(backend):
302 - model_unfused = TestAttentionStaticQuantPatternModel(
303 - num_qo_heads, num_kv_heads, head_size, FP8_DTYPE, device,
304 - vllm_config_unfused)
361 + model_unfused = model_class(num_qo_heads=num_qo_heads,
362 + num_kv_heads=num_kv_heads,
363 + head_size=head_size,
364 + kv_cache_dtype=FP8_DTYPE,
365 + device=device,
366 + vllm_config=vllm_config_unfused)
367 model_unfused = model_unfused.to(device)
368
369 forward_ctx = get_forward_context()
@@ -309,7 +371,7 @@ def test_attention_quant_pattern(num_qo_heads: int, num_kv_heads: int,
371 batch_size)
372
373 # Run model directly without compilation and fusion
312 - result_unfused = model_unfused(q, k, v, linear_w)
374 + result_unfused = model_unfused(q, k, v)
375
376 # Run model with attn fusion enabled
377 vllm_config.compilation_config.pass_config = PassConfig(
@@ -317,9 +379,13 @@ def test_attention_quant_pattern(num_qo_heads: int, num_kv_heads: int,
379 with set_current_vllm_config(vllm_config), set_forward_context(
380 attn_metadata=None, vllm_config=vllm_config
381 ), global_force_attn_backend_context_manager(backend):
320 - model_fused = TestAttentionStaticQuantPatternModel(
321 - num_qo_heads, num_kv_heads, head_size, FP8_DTYPE, device,
322 - vllm_config)
382 + model_fused = model_class(num_qo_heads=num_qo_heads,
383 + num_kv_heads=num_kv_heads,
384 + head_size=head_size,
385 + kv_cache_dtype=FP8_DTYPE,
386 + device=device,
387 + vllm_config=vllm_config,
388 + w=model_unfused.w)
389 model_fused = model_fused.to(device)
390
391 forward_ctx = get_forward_context()
@@ -336,21 +402,20 @@ def test_attention_quant_pattern(num_qo_heads: int, num_kv_heads: int,
402 backend=test_backend,
403 fullgraph=True)
404 assert model_compiled.attn._o_scale_float is None
339 - result_fused_1 = model_compiled(q, k, v, linear_w)
405 + result_fused_1 = model_compiled(q, k, v)
406
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
345 - result_fused_2 = model_compiled(q, k, v, linear_w)
411 + result_fused_2 = model_compiled(q, k, v)
412 assert model_compiled.attn._o_scale_float is not None
413
414 # Check attn fusion support
415 + quant_key = model_class.quant_key
416 attn_fusion_supported = [
350 - layer.impl.fused_output_quant_supported(quant_key.dtype,
351 - quant_key.static,
352 - quant_key.group_shape) for key,
353 - layer in vllm_config.compilation_config.static_forward_context.items()
417 + layer.impl.fused_output_quant_supported(quant_key) for key, layer in
418 + vllm_config.compilation_config.static_forward_context.items()
419 ]
420 if any(attn_fusion_supported):
421 # Check quantization ops in the graph before and after fusion
@@ -370,6 +435,15 @@ def test_attention_quant_pattern(num_qo_heads: int, num_kv_heads: int,
435 assert attn_nodes_post[0].kwargs.get("output_scale") is not None, \
436 "Attention should have output_scale after fusion"
437
438 + assert attn_nodes_pre[0].kwargs.get("output_block_scale") is None, \
439 + "Attention should not have output_block_scale before fusion"
440 + if quant_key.dtype == FP8_DTYPE:
441 + assert attn_nodes_post[0].kwargs.get("output_block_scale") is None, \
442 + "Attention should not have output_block_scale after FP8 fusion"
443 + elif quant_key.dtype == FP4_DTYPE:
444 + assert attn_nodes_post[0].kwargs.get("output_block_scale") is not None, \
445 + "Attention should have output_block_scale after FP4 fusion" # noqa: E501
446 +
447 # Check that results are closed
448 torch.testing.assert_close(result_unfused,
449 result_fused_1,
tests/kernels/attention/test_flashinfer_trtllm_attention.py
+69 -7
@@ -6,7 +6,11 @@ import flashinfer
6 import pytest
7 import torch
8
9 +from tests.kernels.quantization.nvfp4_utils import (FLOAT4_E2M1_MAX,
10 + FLOAT8_E4M3_MAX,
11 + dequantize_nvfp4_to_dtype)
12 from vllm.platforms import current_platform
13 +from vllm.utils import round_up
14
15 if not current_platform.is_device_capability(100):
16 pytest.skip("This TRTLLM kernel requires NVIDIA Blackwell.",
@@ -14,6 +18,7 @@ if not current_platform.is_device_capability(100):
18
19 FLOAT32_BYTES = torch.finfo(torch.float).bits // 8
20 FP8_DTYPE = current_platform.fp8_dtype()
21 +FP4_DTYPE = torch.uint8
22
23
24 def to_float8(x, dtype=torch.float8_e4m3fn):
@@ -29,7 +34,9 @@ DTYPE = [torch.bfloat16]
34 QUANT_DTYPES = [
35 # (q_quant_dtype, kv_quant_dtype, o_quant_dtype)
36 (None, None, None),
37 + (None, FP8_DTYPE, None),
38 (FP8_DTYPE, FP8_DTYPE, FP8_DTYPE),
39 + (FP8_DTYPE, FP8_DTYPE, FP4_DTYPE),
40 ]
41 BATCH_SIZE = [4, 12]
42 MAX_SEQ_LENS = [(1024, 4096)]
@@ -153,11 +160,25 @@ def test_flashinfer_trtllm_decode_with_baseline(
160 output = torch.empty(ref_query.shape, dtype=dtype)
161 wrapper.run(ref_query, ref_kv_cache, out=output)
162 o_scale = 1.0
163 + o_sf_scale = None
164 if o_quant_dtype == FP8_DTYPE:
165 _, o_scale = to_float8(output)
166 + elif o_quant_dtype == FP4_DTYPE:
167 + o_sf_scale = ((FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) /
168 + torch.amax(output.flatten(), dim=-1)).to(torch.float32)
169
170 # TRTLLM Decode
160 - output_trtllm = torch.empty(query.shape, dtype=o_quant_dtype)
171 + if o_quant_dtype == FP4_DTYPE:
172 + output_trtllm = flashinfer.utils.FP4Tensor(
173 + torch.empty(query.shape[:-1] + (query.shape[-1] // 2, ),
174 + dtype=torch.uint8),
175 + torch.empty((round_up(query.shape[0], 128),
176 + round_up(query.shape[1] * query.shape[2] // 16, 4)),
177 + dtype=torch.float8_e4m3fn),
178 + )
179 + else:
180 + output_trtllm = torch.empty(query.shape, dtype=o_quant_dtype)
181 +
182 flashinfer.decode.trtllm_batch_decode_with_kv_cache(
183 query=query,
184 kv_cache=kv_cache,
@@ -167,15 +188,27 @@ def test_flashinfer_trtllm_decode_with_baseline(
188 max_seq_len=max_seq_len,
189 bmm1_scale=q_scale * k_scale * sm_scale,
190 bmm2_scale=v_scale / o_scale,
191 + o_sf_scale=o_sf_scale,
192 out=output_trtllm,
193 )
194 if o_quant_dtype == FP8_DTYPE:
195 output_trtllm = output_trtllm.to(dtype) * o_scale
174 -
175 - if q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP8_DTYPE:
196 + elif o_quant_dtype == FP4_DTYPE:
197 + output_trtllm.data = output_trtllm.data.reshape(
198 + -1, query.shape[1] * query.shape[2] // 2)
199 + output_trtllm = dequantize_nvfp4_to_dtype(output_trtllm.data,
200 + output_trtllm.scale,
201 + o_sf_scale, dtype,
202 + query.device)
203 + output_trtllm = output_trtllm.reshape(-1, query.shape[1],
204 + query.shape[2])
205 +
206 + if q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP4_DTYPE:
207 + rtol, atol = 3e-1, 1e0
208 + elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP8_DTYPE:
209 rtol, atol = 5e-2, 7e-2
210 else:
178 - rtol, atol = 1e-2, 1e-2
211 + rtol, atol = 1e-2, 2e-2
212
213 torch.testing.assert_close(output, output_trtllm, atol=atol, rtol=rtol), \
214 f"{torch.max(torch.abs(output - output_trtllm))}"
@@ -211,6 +244,9 @@ def test_flashinfer_trtllm_prefill_with_baseline(
244 kv_quant_dtype = kv_quant_dtype or dtype
245 o_quant_dtype = o_quant_dtype or dtype
246
247 + if q_quant_dtype != kv_quant_dtype:
248 + pytest.skip("Skipped mixed QKV dtypes for prefill")
249 +
250 max_q_len, max_kv_len = max_seq_lens
251
252 num_qo_heads, num_kv_heads = num_heads
@@ -303,11 +339,25 @@ def test_flashinfer_trtllm_prefill_with_baseline(
339 output = torch.empty(ref_query.shape, dtype=dtype)
340 wrapper.run(ref_query, ref_kv_cache, out=output)
341 o_scale = 1.0
342 + o_sf_scale = None
343 if o_quant_dtype == FP8_DTYPE:
344 _, o_scale = to_float8(output)
345 + elif o_quant_dtype == FP4_DTYPE:
346 + o_sf_scale = ((FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) /
347 + torch.amax(output.flatten(), dim=-1)).to(torch.float32)
348
349 # TRTLLM Prefill
310 - output_trtllm = torch.empty(query.shape, dtype=o_quant_dtype)
350 + if o_quant_dtype == FP4_DTYPE:
351 + output_trtllm = flashinfer.utils.FP4Tensor(
352 + torch.empty(query.shape[:-1] + (query.shape[-1] // 2, ),
353 + dtype=torch.uint8),
354 + torch.empty((round_up(query.shape[0], 128),
355 + round_up(query.shape[1] * query.shape[2] // 16, 4)),
356 + dtype=torch.float8_e4m3fn),
357 + )
358 + else:
359 + output_trtllm = torch.empty(query.shape, dtype=o_quant_dtype)
360 +
361 flashinfer.prefill.trtllm_batch_context_with_kv_cache(
362 query=query,
363 kv_cache=kv_cache,
@@ -321,12 +371,24 @@ def test_flashinfer_trtllm_prefill_with_baseline(
371 batch_size=batch_size,
372 cum_seq_lens_q=q_indptr,
373 cum_seq_lens_kv=kv_indptr,
374 + o_sf_scale=o_sf_scale,
375 out=output_trtllm,
376 )
377 if o_quant_dtype == FP8_DTYPE:
378 output_trtllm = output_trtllm.to(dtype) * o_scale
328 -
329 - if q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP8_DTYPE:
379 + elif o_quant_dtype == FP4_DTYPE:
380 + output_trtllm.data = output_trtllm.data.reshape(
381 + -1, query.shape[1] * query.shape[2] // 2)
382 + output_trtllm = dequantize_nvfp4_to_dtype(output_trtllm.data,
383 + output_trtllm.scale,
384 + o_sf_scale, dtype,
385 + query.device)
386 + output_trtllm = output_trtllm.reshape(-1, query.shape[1],
387 + query.shape[2])
388 +
389 + if q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP4_DTYPE:
390 + rtol, atol = 4e-1, 1e0
391 + elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP8_DTYPE:
392 rtol, atol = 5e-2, 7e-2
393 else:
394 rtol, atol = 1e-2, 1e-2
vllm/attention/backends/abstract.py
+5 -8
@@ -9,8 +9,7 @@ from typing import (TYPE_CHECKING, Any, Dict, Generic, List, Optional,
9
10 import torch
11
12 -from vllm.model_executor.layers.quantization.utils.quant_utils import (
13 - GroupShape)
12 +from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
13 from vllm.multimodal import MultiModalPlaceholderMap
14
15 if TYPE_CHECKING:
@@ -285,20 +284,17 @@ class AttentionImpl(ABC, Generic[T]):
284 attn_metadata: T,
285 output: Optional[torch.Tensor] = None,
286 output_scale: Optional[torch.Tensor] = None,
287 + output_block_scale: Optional[torch.Tensor] = None,
288 ) -> torch.Tensor:
289 raise NotImplementedError
290
291 - def fused_output_quant_supported(self, dtype: torch.dtype, static: bool,
292 - group_shape: GroupShape):
291 + def fused_output_quant_supported(self, quant_key: QuantKey):
292 """
293 Does this attention implementation support fused output quantization.
294 This is used by the AttnFusionPass to only fuse output quantization
295 onto implementations that support it.
296
298 - TODO(luka) merge parameters into QuantDescriptor
299 - :param dtype: quantized dtype
300 - :param static: static or dynamic quantization
301 - :param group_shape: quant group shape.
297 + :param quant_key: QuantKey object that describes the quantization op
298 :return: is fusion supported for this type of quantization
299 """
300 return False
@@ -317,6 +313,7 @@ class MLAAttentionImpl(AttentionImpl[T], Generic[T]):
313 attn_metadata: T,
314 output: Optional[torch.Tensor] = None,
315 output_scale: Optional[torch.Tensor] = None,
316 + output_block_scale: Optional[torch.Tensor] = None,
317 ) -> torch.Tensor:
318 raise NotImplementedError
319
vllm/attention/backends/differential_flash_attn.py
+6
@@ -800,6 +800,7 @@ class DifferentialFlashAttentionImpl(AttentionImpl):
800 attn_metadata: DifferentialFlashAttentionMetadata,
801 output: Optional[torch.Tensor] = None,
802 output_scale: Optional[torch.Tensor] = None,
803 + output_block_scale: Optional[torch.Tensor] = None,
804 ) -> torch.Tensor:
805 """Forward pass with FlashAttention.
806
@@ -817,6 +818,11 @@ class DifferentialFlashAttentionImpl(AttentionImpl):
818 {q,k,v}_descale to be (num_sequences, num_kv_heads).
819 We use torch's .expand() to avoid duplicating values
820 """
821 + if output_scale is not None or output_block_scale is not None:
822 + raise NotImplementedError(
823 + "fused output quantization is not yet supported"
824 + " for DifferentialFlashAttentionImpl")
825 +
826 if self.lambda_full is None:
827 self.lambda_init = self.differential_flash_attention_config[
828 "lambda_init"]
vllm/attention/backends/dual_chunk_flash_attn.py
+2 -1
@@ -371,6 +371,7 @@ class DualChunkFlashAttentionImpl(FlashAttentionImpl):
371 attn_metadata: DualChunkFlashAttentionMetadata,
372 output: Optional[torch.Tensor] = None,
373 output_scale: Optional[torch.Tensor] = None,
374 + output_block_scale: Optional[torch.Tensor] = None,
375 ) -> torch.Tensor:
376 """Forward pass with DualChunkFlashAttention.
377 Args:
@@ -386,7 +387,7 @@ class DualChunkFlashAttentionImpl(FlashAttentionImpl):
387 """
388 assert output is None, "Output tensor not supported for DualChunk"
389
389 - if output_scale is not None:
390 + if output_scale is not None or output_block_scale is not None:
391 raise NotImplementedError(
392 "fused output quantization is not yet supported"
393 " for FlashAttentionImpl")
vllm/attention/backends/flash_attn.py
+2 -1
@@ -596,6 +596,7 @@ class FlashAttentionImpl(AttentionImpl):
596 attn_metadata: FlashAttentionMetadata,
597 output: Optional[torch.Tensor] = None,
598 output_scale: Optional[torch.Tensor] = None,
599 + output_block_scale: Optional[torch.Tensor] = None,
600 ) -> torch.Tensor:
601 """Forward pass with FlashAttention.
602
@@ -615,7 +616,7 @@ class FlashAttentionImpl(AttentionImpl):
616 """
617 assert output is not None, "Output tensor must be provided."
618
618 - if output_scale is not None:
619 + if output_scale is not None or output_block_scale is not None:
620 raise NotImplementedError(
621 "fused output quantization is not yet supported"
622 " for FlashAttentionImpl")
vllm/attention/backends/mla/common.py
+2 -1
@@ -1238,12 +1238,13 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
1238 attn_metadata: T,
1239 output: Optional[torch.Tensor] = None,
1240 output_scale: Optional[torch.Tensor] = None,
1241 + output_block_scale: Optional[torch.Tensor] = None,
1242 ) -> torch.Tensor:
1243 if output is not None:
1244 raise NotImplementedError(
1245 "output is not yet supported for MLAImplBase")
1246
1246 - if output_scale is not None:
1247 + if output_scale is not None or output_block_scale is not None:
1248 raise NotImplementedError(
1249 "fused output quantization is not yet supported"
1250 " for MLAImplBase")
vllm/attention/backends/rocm_flash_attn.py
+9 -5
@@ -20,7 +20,7 @@ from vllm.attention.ops.paged_attn import (PagedAttention,
20 from vllm.config import get_current_vllm_config
21 from vllm.logger import init_logger
22 from vllm.model_executor.layers.quantization.utils.quant_utils import (
23 - GroupShape)
23 + QuantKey, kFp8StaticTensorSym)
24 from vllm.platforms import current_platform
25
26 logger = init_logger(__name__)
@@ -529,11 +529,9 @@ class ROCmFlashAttentionImpl(AttentionImpl):
529 head_dim).reshape(tokens, n_kv_heads * n_rep,
530 head_dim))
531
532 - def fused_output_quant_supported(self, dtype: torch.dtype, static: bool,
533 - group_shape: GroupShape):
532 + def fused_output_quant_supported(self, quant_key: QuantKey):
533 if self.use_triton_flash_attn:
535 - return dtype == current_platform.fp8_dtype(
536 - ) and static and group_shape == GroupShape.PER_TENSOR
534 + return quant_key == kFp8StaticTensorSym
535
536 # Only supported in the Triton backend
537 return False
@@ -548,6 +546,7 @@ class ROCmFlashAttentionImpl(AttentionImpl):
546 attn_metadata: ROCmFlashAttentionMetadata,
547 output: Optional[torch.Tensor] = None,
548 output_scale: Optional[torch.Tensor] = None,
549 + output_block_scale: Optional[torch.Tensor] = None,
550 ) -> torch.Tensor:
551 """Forward pass with FlashAttention and PagedAttention.
552
@@ -606,6 +605,11 @@ class ROCmFlashAttentionImpl(AttentionImpl):
605 "fused output quantization only supported for Triton"
606 " implementation in ROCMFlashAttentionImpl for now")
607
608 + if output_block_scale is not None:
609 + raise NotImplementedError(
610 + "fused nvfp4 output quantization is not supported"
611 + " for ROCMFlashAttentionImpl")
612 +
613 query = query.view(-1, self.num_heads, self.head_size)
614 if key is not None:
615 assert value is not None
vllm/attention/backends/xformers.py
+2 -1
@@ -432,6 +432,7 @@ class XFormersImpl(AttentionImpl[XFormersMetadata]):
432 attn_metadata: "XFormersMetadata",
433 output: Optional[torch.Tensor] = None,
434 output_scale: Optional[torch.Tensor] = None,
435 + output_block_scale: Optional[torch.Tensor] = None,
436 ) -> torch.Tensor:
437 """Forward pass with xFormers and PagedAttention.
438
@@ -484,7 +485,7 @@ class XFormersImpl(AttentionImpl[XFormersMetadata]):
485 Returns:
486 shape = [num_tokens, num_heads * head_size]
487 """
487 - if output_scale is not None:
488 + if output_scale is not None or output_block_scale is not None:
489 raise NotImplementedError(
490 "fused output quantization is not yet supported"
491 " for XFormersImpl")
vllm/attention/layer.py
+5 -2
@@ -495,6 +495,7 @@ def unified_attention_with_output(
495 output: torch.Tensor,
496 layer_name: str,
497 output_scale: Optional[torch.Tensor] = None,
498 + output_block_scale: Optional[torch.Tensor] = None,
499 ) -> None:
500 wait_for_kv_layer_from_connector(layer_name)
501 forward_context: ForwardContext = get_forward_context()
@@ -510,7 +511,8 @@ def unified_attention_with_output(
511 kv_cache,
512 attn_metadata,
513 output=output,
513 - output_scale=output_scale)
514 + output_scale=output_scale,
515 + output_block_scale=output_block_scale)
516
517 maybe_save_kv_layer_to_connector(layer_name, kv_cache)
518
@@ -522,6 +524,7 @@ def unified_attention_with_output_fake(
524 output: torch.Tensor,
525 layer_name: str,
526 output_scale: Optional[torch.Tensor] = None,
527 + output_block_scale: Optional[torch.Tensor] = None,
528 ) -> None:
529 return
530
@@ -529,7 +532,7 @@ def unified_attention_with_output_fake(
532 direct_register_custom_op(
533 op_name="unified_attention_with_output",
534 op_func=unified_attention_with_output,
532 - mutates_args=["output"],
535 + mutates_args=["output", "output_block_scale"],
536 fake_impl=unified_attention_with_output_fake,
537 dispatch_key=current_platform.dispatch_key,
538 )
vllm/compilation/fusion.py
+18 -48
@@ -12,7 +12,8 @@ from torch._ops import OpOverload
12 from vllm.config import VllmConfig
13 from vllm.logger import init_logger
14 from vllm.model_executor.layers.quantization.utils.quant_utils import (
15 - GroupShape)
15 + GroupShape, QuantKey, ScaleDesc, kFp8DynamicTensorSym, kFp8DynamicTokenSym,
16 + kFp8StaticTensorSym, kNvfp4Quant, kStaticTensorScale)
17 from vllm.platforms import current_platform
18
19 from .fx_utils import find_getitem_maybe
@@ -21,6 +22,7 @@ from .vllm_inductor_pass import VllmInductorPass
22
23 logger = init_logger(__name__)
24 FP8_DTYPE = current_platform.fp8_dtype()
25 +FP4_DTYPE = torch.uint8
26
27
28 def empty_bf16(*args, **kwargs):
@@ -31,41 +33,12 @@ def empty_fp32(*args, **kwargs):
33 return torch.empty(*args, **kwargs, dtype=torch.float32, device="cuda")
34
35
34 -RMS_OP = torch.ops._C.rms_norm.default
35 -RMS_ADD_OP = torch.ops._C.fused_add_rms_norm.default
36 -
37 -
38 -class QuantKey(NamedTuple):
39 - """
40 - Named tuple for identifying the type of quantization.
41 - dtype: quantized data type
42 - static: static quantization if True, dynamic if False
43 - group_shape: quantization group shape
44 - symmetric: symmetric if True, asymmetric if False
45 -
46 - TODO(luka) use QuantDescriptor once standardized:
47 - https://github.com/vllm-project/vllm/issues/8913
36 +def empty_i32(*args, **kwargs):
37 + return torch.empty(*args, **kwargs, dtype=torch.int32, device="cuda")
38
49 - """
50 - dtype: torch.dtype
51 - static: bool
52 - group_shape: GroupShape
53 - symmetric: bool = True
39
55 - def __str__(self):
56 - group_shape = ('per_tensor'
57 - if self.group_shape == GroupShape.PER_TENSOR else
58 - ('per_token' if self.group_shape == GroupShape.PER_TOKEN
59 - else str(self.group_shape)))
60 -
61 - return (f"QuantKey({'static' if self.static else 'dynamic'},"
62 - f"{fx.graph.dtype_abbrs[self.dtype]},{group_shape},"
63 - f"{'a' if not self.symmetric else ''}symmetric)")
64 -
65 -
66 -kFp8StaticTensorSym = QuantKey(FP8_DTYPE, True, GroupShape.PER_TENSOR, True)
67 -kFp8DynamicTensorSym = QuantKey(FP8_DTYPE, False, GroupShape.PER_TENSOR, True)
68 -kFp8DynamicTokenSym = QuantKey(FP8_DTYPE, False, GroupShape.PER_TOKEN, True)
40 +RMS_OP = torch.ops._C.rms_norm.default
41 +RMS_ADD_OP = torch.ops._C.fused_add_rms_norm.default
42
43 QUANT_OPS: dict[QuantKey, OpOverload] = {
44 kFp8StaticTensorSym:
@@ -74,6 +47,7 @@ QUANT_OPS: dict[QuantKey, OpOverload] = {
47 torch.ops._C.dynamic_scaled_fp8_quant.default, # noqa: E501
48 kFp8DynamicTokenSym:
49 torch.ops._C.dynamic_per_token_scaled_fp8_quant.default, # noqa: E501
50 + kNvfp4Quant: torch.ops._C.scaled_fp4_quant.default, # noqa: E501
51 }
52
53
@@ -187,11 +161,9 @@ class RMSNormStaticQuantPattern(RMSNormQuantPattern):
161 quant_dtype: torch.dtype,
162 symmetric=True):
163 fused_key = FusedRMSQuantKey(fused_add=False,
190 - quant=QuantKey(
191 - dtype=quant_dtype,
192 - static=True,
193 - group_shape=GroupShape.PER_TENSOR,
194 - symmetric=symmetric))
164 + quant=QuantKey(dtype=quant_dtype,
165 + scale=kStaticTensorScale,
166 + symmetric=symmetric))
167 super().__init__(epsilon, fused_key)
168
169 def register(self, pm_pass: PatternMatcherPass):
@@ -244,11 +216,9 @@ class FusedAddRMSNormStaticQuantPattern(RMSNormQuantPattern):
216 quant_dtype: torch.dtype,
217 symmetric=True):
218 key = FusedRMSQuantKey(fused_add=True,
247 - quant=QuantKey(
248 - dtype=quant_dtype,
249 - static=True,
250 - group_shape=GroupShape.PER_TENSOR,
251 - symmetric=symmetric))
219 + quant=QuantKey(dtype=quant_dtype,
220 + scale=kStaticTensorScale,
221 + symmetric=symmetric))
222 super().__init__(epsilon, key)
223
224 def register(self, pm_pass: PatternMatcherPass,
@@ -337,10 +307,10 @@ class RMSNormDynamicQuantPattern(RMSNormQuantPattern):
307 quant_dtype: torch.dtype,
308 group_shape: GroupShape = GroupShape.PER_TOKEN,
309 symmetric=True):
310 + scale = ScaleDesc(torch.float32, False, group_shape)
311 key = FusedRMSQuantKey(fused_add=False,
312 quant=QuantKey(dtype=quant_dtype,
342 - static=False,
343 - group_shape=group_shape,
313 + scale=scale,
314 symmetric=symmetric))
315 super().__init__(epsilon, key)
316
@@ -435,10 +405,10 @@ class FusedAddRMSNormDynamicQuantPattern(RMSNormQuantPattern):
405 quant_dtype: torch.dtype,
406 group_shape: GroupShape = GroupShape.PER_TOKEN,
407 symmetric=True):
408 + scale = ScaleDesc(torch.float32, False, group_shape)
409 key = FusedRMSQuantKey(fused_add=True,
410 quant=QuantKey(dtype=quant_dtype,
440 - static=False,
441 - group_shape=group_shape,
411 + scale=scale,
412 symmetric=symmetric))
413 super().__init__(epsilon, key)
414
vllm/compilation/fusion_attn.py
+151 -34
@@ -1,6 +1,8 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +from abc import ABC, abstractmethod
5 +
6 import torch
7 import torch._inductor.pattern_matcher as pm
8 from torch._higher_order_ops.auto_functionalize import auto_functionalized
@@ -11,44 +13,41 @@ from torch._subclasses.fake_tensor import (FakeTensorMode,
13 from vllm.attention import Attention
14 from vllm.config import VllmConfig, get_layers_from_vllm_config
15 from vllm.logger import init_logger
16 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
17 + QuantKey, kNvfp4Quant, kStaticTensorScale)
18 from vllm.platforms import current_platform
19 +from vllm.utils import round_up
20
16 -from .fusion import QUANT_OPS, GroupShape, QuantKey, empty_bf16, empty_fp32
21 +from .fusion import QUANT_OPS, empty_bf16, empty_fp32, empty_i32
22 from .vllm_inductor_pass import VllmInductorPass
23
24 logger = init_logger(__name__)
25
26 FP8_DTYPE = current_platform.fp8_dtype()
27 +FP4_DTYPE = torch.uint8
28
29 ATTN_OP = torch.ops.vllm.unified_attention_with_output.default
30 RESHAPE_OP = torch.ops.aten.reshape.default
31
32
27 -class AttentionStaticQuantPattern:
33 +class AttentionQuantPattern(ABC):
34 """
29 - Fusion for Attention+StaticQuant.
30 -
31 - Only triggers when the attention implementation returns True in
32 - `fused_output_quant_supported()`. If the pattern is found, the StaticQuant
33 - op will be removed from the graph, and its scale will be passed into
34 - Attention op as the `output_scale` argument.
35 + The base class for Attn+Quant fusions.
36 + Should not be used directly.
37 """
38
39 def __init__(
40 self,
41 layer: Attention,
40 - quant_dtype: torch.dtype,
41 - symmetric=True,
42 + quant_key: QuantKey,
43 ):
44 self.layer = layer
45 self.layer_name = layer.layer_name
46 self.num_heads = layer.num_heads
47 self.head_size = layer.head_size
47 - self.quant_dtype = quant_dtype
48 - self.quant_key = QuantKey(dtype=quant_dtype,
49 - static=True,
50 - group_shape=GroupShape.PER_TENSOR,
51 - symmetric=symmetric)
48 + self.quant_key = quant_key
49 + self.quant_dtype = quant_key.dtype
50 +
51 assert self.quant_key in QUANT_OPS, \
52 f"unsupported quantization scheme {self.quant_key}"
53 self.QUANT_OP = QUANT_OPS[self.quant_key]
@@ -57,12 +56,49 @@ class AttentionStaticQuantPattern:
56 kwargs = {'dtype': self.quant_dtype, 'device': "cuda", **kwargs}
57 return torch.empty(*args, **kwargs)
58
59 + @staticmethod
60 + def wrap_trace_fn(process_fx, trace_fn):
61 +
62 + def wrapped(*args, **kwargs):
63 + return process_fx(trace_fn(*args, **kwargs))
64 +
65 + return wrapped
66 +
67 + @staticmethod
68 + def fx_view_to_reshape(gm: torch.fx.GraphModule):
69 + from torch._inductor.fx_passes.post_grad import view_to_reshape
70 + view_to_reshape(gm)
71 + return gm
72 +
73 def register_if_supported(self, pm_pass: PatternMatcherPass):
61 - if self.layer.impl.fused_output_quant_supported(
62 - self.quant_dtype, self.quant_key.static,
63 - self.quant_key.group_shape):
74 + if self.layer.impl.fused_output_quant_supported(self.quant_key):
75 self._register(pm_pass)
76
77 + @abstractmethod
78 + def _register(self, pm_pass: PatternMatcherPass):
79 + raise NotImplementedError
80 +
81 +
82 +class AttentionFp8StaticQuantPattern(AttentionQuantPattern):
83 + """
84 + Fusion for Attention+Fp8StaticQuant.
85 +
86 + Only triggers when the attention implementation returns True in
87 + `fused_output_quant_supported()`. If the pattern is found, the
88 + Fp8StaticQuant op will be removed from the graph, and its scale
89 + will be passed into Attention op as the `output_scale` argument.
90 + """
91 +
92 + def __init__(
93 + self,
94 + layer: Attention,
95 + symmetric: bool = True,
96 + ):
97 + quant_key = QuantKey(dtype=FP8_DTYPE,
98 + scale=kStaticTensorScale,
99 + symmetric=symmetric)
100 + super().__init__(layer, quant_key)
101 +
102 def _register(self, pm_pass: PatternMatcherPass):
103
104 def pattern(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
@@ -74,9 +110,10 @@ class AttentionStaticQuantPattern:
110 value=v,
111 output=output_attn,
112 layer_name=self.layer_name,
77 - output_scale=None)
78 - attn_out_view = RESHAPE_OP(at1[1],
79 - [-1, self.num_heads * self.head_size])
113 + output_scale=None,
114 + output_block_scale=None)
115 + attn_out_view = RESHAPE_OP(
116 + at1[1], [q.shape[0], self.num_heads * self.head_size])
117 at2 = auto_functionalized(self.QUANT_OP,
118 result=output_quant,
119 input=attn_out_view,
@@ -98,7 +135,8 @@ class AttentionStaticQuantPattern:
135 value=v,
136 output=output_attn,
137 layer_name=self.layer_name,
101 - output_scale=scale)
138 + output_scale=scale,
139 + output_block_scale=None)
140 return RESHAPE_OP(at1[1], [-1, self.num_heads * self.head_size])
141
142 # Need custom fake mode, otherwise tracing happens with real tensors.
@@ -114,21 +152,94 @@ class AttentionStaticQuantPattern:
152 empty_fp32(1, 1) # scale
153 ]
154
117 - def wrap_trace_fn(process_fx, trace_fn):
155 + pm.register_replacement(
156 + pattern, replacement, inputs,
157 + AttentionQuantPattern.wrap_trace_fn(
158 + AttentionQuantPattern.fx_view_to_reshape, pm.fwd_only),
159 + pm_pass)
160 +
161
119 - def wrapped(*args, **kwargs):
120 - return process_fx(trace_fn(*args, **kwargs))
162 +class AttentionNvfp4QuantPattern(AttentionQuantPattern):
163 + """
164 + Fusion for Attention+Nvfp4Quant.
165 +
166 + Only triggers when the attention implementation returns True in
167 + `fused_output_quant_supported()`. If the pattern is found, the
168 + Nvfp4Quant op will be removed from the graph, and its scale
169 + will be passed into Attention op as the `output_scale` argument.
170 + """
171
122 - return wrapped
172 + def __init__(self, layer: Attention):
173 + super().__init__(layer, kNvfp4Quant)
174
124 - def fx_view_to_reshape(gm: torch.fx.GraphModule):
125 - from torch._inductor.fx_passes.post_grad import view_to_reshape
126 - view_to_reshape(gm)
127 - return gm
175 + def _register(self, pm_pass: PatternMatcherPass):
176 +
177 + def pattern(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
178 + output_attn: torch.Tensor, output_quant: torch.Tensor,
179 + output_scale: torch.Tensor, input_scale: torch.Tensor):
180 + at1 = auto_functionalized(ATTN_OP,
181 + query=q,
182 + key=k,
183 + value=v,
184 + output=output_attn,
185 + layer_name=self.layer_name,
186 + output_scale=None,
187 + output_block_scale=None)
188 + attn_out_view = RESHAPE_OP(
189 + at1[1], [q.shape[0], self.num_heads * self.head_size])
190 + at2 = auto_functionalized(self.QUANT_OP,
191 + output=output_quant,
192 + input=attn_out_view,
193 + output_scale=output_scale,
194 + input_scale=input_scale)
195 + output_scale_view = torch.ops.aten.view.dtype(at2[2], FP8_DTYPE)
196 + return at2[1], output_scale_view
197 +
198 + def replacement(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
199 + output_attn: torch.Tensor, output_quant: torch.Tensor,
200 + output_scale: torch.Tensor, input_scale: torch.Tensor):
201 + # attention output in quant_dtype
202 + output_attn = torch.ops.aten.full.default(
203 + [q.shape[0], self.num_heads, self.head_size // 2],
204 + 0.0,
205 + dtype=self.quant_dtype,
206 + device=q.device)
207 + # attention output block scale
208 + output_scale_view = torch.ops.aten.view.dtype(
209 + output_scale, FP8_DTYPE)
210 + at2 = auto_functionalized(ATTN_OP,
211 + query=q,
212 + key=k,
213 + value=v,
214 + output=output_attn,
215 + layer_name=self.layer_name,
216 + output_scale=input_scale,
217 + output_block_scale=output_scale_view)
218 + output = RESHAPE_OP(at2[1],
219 + [-1, self.num_heads * self.head_size // 2])
220 + return output, at2[2]
221 +
222 + # Need custom fake mode, otherwise tracing happens with real tensors.
223 + # That would not work for the unified_attention custom op.
224 + with unset_fake_temporarily(), FakeTensorMode():
225 + inputs = [
226 + empty_bf16(5, self.num_heads, self.head_size), # q
227 + empty_bf16(5, self.num_heads, self.head_size), # k
228 + empty_bf16(5, self.num_heads, self.head_size), # v
229 + empty_bf16(5, self.num_heads, self.head_size), # output_attn
230 + self.empty_quant(5, self.num_heads * self.head_size //
231 + 2), # output_quant
232 + empty_i32(128,
233 + round_up(self.num_heads * self.head_size // 16,
234 + 4)), # output_scale
235 + empty_fp32(1, 1), # input_scale
236 + ]
237
238 pm.register_replacement(
239 pattern, replacement, inputs,
131 - wrap_trace_fn(fx_view_to_reshape, pm.fwd_only), pm_pass)
240 + AttentionQuantPattern.wrap_trace_fn(
241 + AttentionQuantPattern.fx_view_to_reshape, pm.fwd_only),
242 + pm_pass)
243
244
245 class AttnFusionPass(VllmInductorPass):
@@ -151,8 +262,12 @@ class AttnFusionPass(VllmInductorPass):
262
263 attn_layers = get_layers_from_vllm_config(config, Attention)
264 for layer_name, layer in attn_layers.items():
154 - pattern = AttentionStaticQuantPattern(layer, FP8_DTYPE)
155 - pattern.register_if_supported(self.patterns)
265 + pattern_fp8 = AttentionFp8StaticQuantPattern(layer)
266 + pattern_fp8.register_if_supported(self.patterns)
267 +
268 + pattern_nvfp4 = AttentionNvfp4QuantPattern(layer)
269 + pattern_nvfp4.register_if_supported(self.patterns)
270 +
271 if len(attn_layers) == 0:
272 logger.warning(
273 "Attention + quant fusion is enabled, but no attention layers "
@@ -175,4 +290,6 @@ class AttnFusionPass(VllmInductorPass):
290 self.end_and_log()
291
292 def uuid(self):
178 - return VllmInductorPass.hash_source(self, AttentionStaticQuantPattern)
293 + return VllmInductorPass.hash_source(self, AttentionQuantPattern,
294 + AttentionFp8StaticQuantPattern,
295 + AttentionNvfp4QuantPattern)
vllm/model_executor/layers/quantization/utils/quant_utils.py
+63
@@ -2,16 +2,21 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """This file is used for /tests and /benchmarks"""
4 from collections.abc import Mapping
5 +from dataclasses import dataclass
6 from types import MappingProxyType
7 from typing import ClassVar, NamedTuple, Optional
8
9 import numpy
10 import torch
11 +from torch import fx
12
13 from vllm._custom_ops import cutlass_scaled_mm_supports_fp4
14 from vllm.platforms import current_platform
15 from vllm.scalar_type import ScalarType, scalar_types
16
17 +FP8_DTYPE = current_platform.fp8_dtype()
18 +FP4_DTYPE = torch.uint8
19 +
20
21 # Use proxy as NamedTuple direct subclasses cannot have static members
22 class _GroupShape(NamedTuple):
@@ -34,6 +39,64 @@ GroupShape.PER_TENSOR = GroupShape(-1, -1)
39 GroupShape.PER_TOKEN = GroupShape(1, -1)
40
41
42 +@dataclass(frozen=True)
43 +class ScaleDesc:
44 + """
45 + Class for describing a single quantization scaling factor.
46 + dtype: data type of the scale
47 + static: static scale if True, dynamic if False
48 + group_shape: group shape of the scale
49 + """
50 + dtype: torch.dtype
51 + static: bool
52 + group_shape: GroupShape
53 +
54 + def __str__(self):
55 + group_shape = ('per_tensor'
56 + if self.group_shape == GroupShape.PER_TENSOR else
57 + ('per_token' if self.group_shape == GroupShape.PER_TOKEN
58 + else str(self.group_shape)))
59 +
60 + return (f"{fx.graph.dtype_abbrs[self.dtype]},"
61 + f"{'static' if self.static else 'dynamic'},{group_shape}")
62 +
63 +
64 +@dataclass(frozen=True)
65 +class QuantKey:
66 + """
67 + Class for identifying the type of quantization.
68 + dtype: quantized data type
69 + scale: scale descriptor
70 + scale2: second-level scale descriptor
71 + symmetric: symmetric if True, asymmetric if False
72 + """
73 + dtype: torch.dtype
74 + scale: ScaleDesc
75 + scale2: Optional[ScaleDesc] = None
76 + symmetric: bool = True
77 +
78 + def __str__(self):
79 + scale2_str = f"scale2({self.scale2})," if self.scale2 else ""
80 + return (f"QuantKey({fx.graph.dtype_abbrs[self.dtype]},"
81 + f"scale({self.scale}),{scale2_str}"
82 + f"{'a' if not self.symmetric else ''}symmetric)")
83 +
84 +
85 +kStaticTensorScale = ScaleDesc(torch.float32, True, GroupShape.PER_TENSOR)
86 +kFp8StaticTensorSym = QuantKey(FP8_DTYPE, kStaticTensorScale, symmetric=True)
87 +
88 +kDynamicTensorScale = ScaleDesc(torch.float32, False, GroupShape.PER_TENSOR)
89 +kFp8DynamicTensorSym = QuantKey(FP8_DTYPE, kDynamicTensorScale, symmetric=True)
90 +
91 +kDynamicTokenScale = ScaleDesc(torch.float32, False, GroupShape.PER_TOKEN)
92 +kFp8DynamicTokenSym = QuantKey(FP8_DTYPE, kDynamicTokenScale, symmetric=True)
93 +
94 +kNvfp4GroupScale = ScaleDesc(FP8_DTYPE, False, GroupShape(1, 16))
95 +kNvfp4Quant = QuantKey(FP4_DTYPE,
96 + scale=kNvfp4GroupScale,
97 + scale2=kStaticTensorScale)
98 +
99 +
100 # Normalize the group_shape to the full extent for any dims that are -1
101 def _normalize_quant_group_shape(x: torch.Tensor, group_shape: GroupShape):
102 # -1 means full extent
vllm/v1/attention/backends/cpu_attn.py
+2 -1
@@ -483,6 +483,7 @@ class TorchSDPABackendImpl(AttentionImpl[TorchSDPAMetadata]):
483 attn_metadata: TorchSDPAMetadata, # type: ignore
484 output: Optional[torch.Tensor] = None,
485 output_scale: Optional[torch.Tensor] = None,
486 + output_block_scale: Optional[torch.Tensor] = None,
487 ) -> torch.Tensor:
488 """Forward pass with torch SDPA and PagedAttention.
489
@@ -497,7 +498,7 @@ class TorchSDPABackendImpl(AttentionImpl[TorchSDPAMetadata]):
498 Returns:
499 shape = [num_tokens, num_heads * head_size]
500 """
500 - if output_scale is not None:
501 + if output_scale is not None or output_block_scale is not None:
502 raise NotImplementedError(
503 "fused output quantization is not yet supported"
504 " for TorchSDPABackendImpl")
vllm/v1/attention/backends/flash_attn.py
+2 -1
@@ -430,6 +430,7 @@ class FlashAttentionImpl(AttentionImpl):
430 attn_metadata: FlashAttentionMetadata,
431 output: Optional[torch.Tensor] = None,
432 output_scale: Optional[torch.Tensor] = None,
433 + output_block_scale: Optional[torch.Tensor] = None,
434 ) -> torch.Tensor:
435 """Forward pass with FlashAttention.
436
@@ -447,7 +448,7 @@ class FlashAttentionImpl(AttentionImpl):
448 """
449 assert output is not None, "Output tensor must be provided."
450
450 - if output_scale is not None:
451 + if output_scale is not None or output_block_scale is not None:
452 raise NotImplementedError(
453 "fused output quantization is not yet supported"
454 " for FlashAttentionImpl")
vllm/v1/attention/backends/flashinfer.py
+49 -13
@@ -12,6 +12,7 @@ from flashinfer import (BatchDecodeWithPagedKVCacheWrapper,
12 MultiLevelCascadeAttentionWrapper)
13 from flashinfer.decode import _get_range_buf, trtllm_batch_decode_with_kv_cache
14 from flashinfer.prefill import trtllm_batch_context_with_kv_cache
15 +from flashinfer.utils import FP4Tensor
16
17 from vllm import _custom_ops as ops
18 from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
@@ -19,7 +20,7 @@ from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
20 from vllm.config import CUDAGraphMode, VllmConfig
21 from vllm.logger import init_logger
22 from vllm.model_executor.layers.quantization.utils.quant_utils import (
22 - GroupShape)
23 + QuantKey, kFp8StaticTensorSym, kNvfp4Quant)
24 from vllm.platforms import current_platform
25 from vllm.utils import cdiv, is_pin_memory_available
26 from vllm.utils.flashinfer import (supports_trtllm_attention,
@@ -40,6 +41,7 @@ from vllm.v1.kv_cache_interface import AttentionSpec
41 FLASHINFER_WORKSPACE_BUFFER_SIZE = 256 * 1024 * 1024
42
43 FP8_DTYPE = current_platform.fp8_dtype()
44 +FP4_DTYPE = torch.uint8
45
46 logger = init_logger(__name__)
47
@@ -653,14 +655,12 @@ class FlashInferImpl(AttentionImpl):
655 and num_heads % num_kv_heads == 0)
656 self.bmm1_scale: Optional[float] = None
657 self.bmm2_scale: Optional[float] = None
658 + self.o_sf_scale: Optional[float] = None
659
657 - def fused_output_quant_supported(self, dtype: torch.dtype, static: bool,
658 - group_shape: GroupShape):
659 - supported_quant_type = (dtype == FP8_DTYPE and static
660 - and group_shape == GroupShape.PER_TENSOR)
660 + def fused_output_quant_supported(self, quant_key: QuantKey):
661 return (self.support_trtllm_attn
662 and self.kv_cache_dtype.startswith("fp8")
663 - and supported_quant_type)
663 + and quant_key in (kFp8StaticTensorSym, kNvfp4Quant))
664
665 def forward(
666 self,
@@ -672,6 +672,7 @@ class FlashInferImpl(AttentionImpl):
672 attn_metadata: FlashInferMetadata,
673 output: Optional[torch.Tensor] = None,
674 output_scale: Optional[torch.Tensor] = None,
675 + output_block_scale: Optional[torch.Tensor] = None,
676 ) -> torch.Tensor:
677 """Forward pass with FlashInfer.
678
@@ -705,19 +706,32 @@ class FlashInferImpl(AttentionImpl):
706 if output_scale is None:
707 assert attn_metadata.q_data_type != FP8_DTYPE, \
708 "Query can only be FP8 if output fusion happened."
709 + assert output_block_scale is None, "output_block_scale "\
710 + "is not supported when fusion has not happened"
711 else:
712 assert attn_metadata.q_data_type == FP8_DTYPE, \
713 "Query must be FP8 when attn+quant fusion happened."
714 assert (attn_metadata.prefill_use_trtllm and
715 attn_metadata.decode_use_trtllm), "Must use TRT-LLM attn"
713 - assert output.dtype == FP8_DTYPE, \
714 - "Output must be FP8 when attn+quant fusion happened."
716
716 - # TRTLLM attn kernel requires o scale as a host scalar, store the
717 - # o scale to host scalar in warmup run with cuda graph not enabled
717 + if output.dtype == FP8_DTYPE:
718 + assert output_block_scale is None, \
719 + "output_block_scale should not be provided for fp8 output"
720 + elif output.dtype == FP4_DTYPE:
721 + assert output_block_scale is not None, \
722 + "output_block_scale is required for nvfp4 output"
723 + else:
724 + raise ValueError(f"Unsupported output dtype: {output.dtype}")
725 +
726 + # TRTLLM attn kernel requires o scale to pass as a host scalar,
727 + # store the o scale as a host scalar in warmup run with cuda graph
728 + # not enabled
729 if layer._o_scale_float is None:
730 layer._o_scale_float = output_scale.cpu().item()
720 - self.bmm2_scale = self.bmm2_scale / layer._o_scale_float
731 + if output.dtype == FP8_DTYPE:
732 + self.bmm2_scale = self.bmm2_scale / layer._o_scale_float
733 + elif output.dtype == FP4_DTYPE:
734 + self.o_sf_scale = layer._o_scale_float
735
736 # Insert FP8 quant for query
737 num_tokens, num_heads, head_size = query.shape
@@ -818,6 +832,16 @@ class FlashInferImpl(AttentionImpl):
832 assert block_tables_prefill.is_contiguous()
833 assert seq_lens_prefill.is_contiguous()
834
835 + if output.dtype == FP4_DTYPE:
836 + assert self.o_sf_scale is not None
837 + out = FP4Tensor(data=output[num_decode_tokens:],
838 + scale=output_block_scale,
839 + scale_start_index=num_decode_tokens,
840 + original_shape=prefill_query.shape)
841 + else:
842 + assert self.o_sf_scale is None
843 + out = output[num_decode_tokens:]
844 +
845 trtllm_batch_context_with_kv_cache(
846 query=prefill_query,
847 kv_cache=kv_cache_permute,
@@ -833,7 +857,8 @@ class FlashInferImpl(AttentionImpl):
857 cum_seq_lens_kv=attn_metadata.paged_kv_indptr_gpu,
858 window_left=self.window_left,
859 sinks=self.sinks,
836 - out=output[num_decode_tokens:],
860 + o_sf_scale=self.o_sf_scale,
861 + out=out,
862 )
863
864 if num_decode_tokens > 0:
@@ -870,6 +895,16 @@ class FlashInferImpl(AttentionImpl):
895 assert block_tables_decode.is_contiguous()
896 assert seq_lens_decode.is_contiguous()
897
898 + if output.dtype == FP4_DTYPE:
899 + assert self.o_sf_scale is not None
900 + out = FP4Tensor(data=output[:num_decode_tokens],
901 + scale=output_block_scale,
902 + scale_start_index=0,
903 + original_shape=decode_query.shape)
904 + else:
905 + assert self.o_sf_scale is None
906 + out = output[:num_decode_tokens]
907 +
908 trtllm_batch_decode_with_kv_cache(
909 query=decode_query,
910 kv_cache=kv_cache_permute,
@@ -881,7 +916,8 @@ class FlashInferImpl(AttentionImpl):
916 bmm2_scale=self.bmm2_scale,
917 window_left=self.window_left,
918 sinks=self.sinks,
884 - out=output[:num_decode_tokens],
919 + o_sf_scale=self.o_sf_scale,
920 + out=out,
921 )
922 return output_padded
923
vllm/v1/attention/backends/flex_attention.py
+2 -1
@@ -428,6 +428,7 @@ class FlexAttentionImpl(AttentionImpl):
428 attn_metadata: FlexAttentionMetadata,
429 output: Optional[torch.Tensor] = None,
430 output_scale: Optional[torch.Tensor] = None,
431 + output_block_scale: Optional[torch.Tensor] = None,
432 ) -> torch.Tensor:
433 """Forward pass with FLexAttention.
434
@@ -441,7 +442,7 @@ class FlexAttentionImpl(AttentionImpl):
442 shape = [num_tokens, num_heads * head_size]
443 """
444 assert output is not None, "Output tensor must be provided."
444 - if output_scale is not None:
445 + if output_scale is not None or output_block_scale is not None:
446 raise NotImplementedError(
447 "fused output quantization is not yet supported"
448 " for FlexAttentionImpl")
vllm/v1/attention/backends/mla/common.py
+2 -1
@@ -1138,10 +1138,11 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
1138 attn_metadata: M,
1139 output: Optional[torch.Tensor] = None,
1140 output_scale: Optional[torch.Tensor] = None,
1141 + output_block_scale: Optional[torch.Tensor] = None,
1142 ) -> torch.Tensor:
1143 assert output is not None, "Output tensor must be provided."
1144
1144 - if output_scale is not None:
1145 + if output_scale is not None or output_block_scale is not None:
1146 raise NotImplementedError(
1147 "fused output quantization is not yet supported"
1148 " for MLACommonImpl")
vllm/v1/attention/backends/pallas.py
+2 -1
@@ -227,6 +227,7 @@ class PallasAttentionBackendImpl(AttentionImpl):
227 attn_metadata: PallasMetadata,
228 output: Optional[torch.Tensor] = None,
229 output_scale: Optional[torch.Tensor] = None,
230 + output_block_scale: Optional[torch.Tensor] = None,
231 ) -> torch.Tensor:
232 """Forward pass with Pallas attention.
233
@@ -239,7 +240,7 @@ class PallasAttentionBackendImpl(AttentionImpl):
240 Returns:
241 shape = [num_tokens, num_heads * head_size]
242 """
242 - if output_scale is not None:
243 + if output_scale is not None or output_block_scale is not None:
244 raise NotImplementedError(
245 "fused output quantization is not yet supported"
246 " for PallasAttentionBackendImpl")
vllm/v1/attention/backends/rocm_aiter_fa.py
+2 -1
@@ -421,6 +421,7 @@ class AiterFlashAttentionImpl(AttentionImpl):
421 attn_metadata: AiterFlashAttentionMetadata,
422 output: Optional[torch.Tensor] = None,
423 output_scale: Optional[torch.Tensor] = None,
424 + output_block_scale: Optional[torch.Tensor] = None,
425 ) -> torch.Tensor:
426 """Forward pass with AiterFlashAttention.
427
@@ -438,7 +439,7 @@ class AiterFlashAttentionImpl(AttentionImpl):
439 """
440 assert output is not None, "Output tensor must be provided."
441
441 - if output_scale is not None:
442 + if output_scale is not None or output_block_scale is not None:
443 raise NotImplementedError(
444 "fused output quantization is not yet supported"
445 " for FlashAttentionImpl")
vllm/v1/attention/backends/tree_attn.py
+2 -1
@@ -354,6 +354,7 @@ class TreeAttentionImpl(AttentionImpl):
354 attn_metadata: TreeAttentionMetadata,
355 output: Optional[torch.Tensor] = None,
356 output_scale: Optional[torch.Tensor] = None,
357 + output_block_scale: Optional[torch.Tensor] = None,
358 ) -> torch.Tensor:
359 """Forward pass with TreeAttention.
360
@@ -368,7 +369,7 @@ class TreeAttentionImpl(AttentionImpl):
369 """
370 assert output is not None, "Output tensor must be provided."
371
371 - if output_scale is not None:
372 + if output_scale is not None or output_block_scale is not None:
373 raise NotImplementedError(
374 "fused output quantization is not yet supported"
375 " for TreeAttentionImpl")
vllm/v1/attention/backends/triton_attn.py
+2 -1
@@ -277,6 +277,7 @@ class TritonAttentionImpl(AttentionImpl):
277 attn_metadata: FlashAttentionMetadata,
278 output: Optional[torch.Tensor] = None,
279 output_scale: Optional[torch.Tensor] = None,
280 + output_block_scale: Optional[torch.Tensor] = None,
281 ) -> torch.Tensor:
282 """Forward pass with FlashAttention.
283
@@ -291,7 +292,7 @@ class TritonAttentionImpl(AttentionImpl):
292 """
293 assert output is not None, "Output tensor must be provided."
294
294 - if output_scale is not None:
295 + if output_scale is not None or output_block_scale is not None:
296 raise NotImplementedError(
297 "fused output quantization is not yet supported"
298 " for TritonAttentionImpl")
vllm/v1/attention/backends/xformers.py
+2 -1
@@ -322,6 +322,7 @@ class XFormersAttentionImpl(AttentionImpl):
322 attn_metadata: XFormersAttentionMetadata,
323 output: Optional[torch.Tensor] = None,
324 output_scale: Optional[torch.Tensor] = None,
325 + output_block_scale: Optional[torch.Tensor] = None,
326 ) -> torch.Tensor:
327 """Forward pass with XFormers.
328
@@ -336,7 +337,7 @@ class XFormersAttentionImpl(AttentionImpl):
337 """
338 assert output is not None, "Output tensor must be provided."
339
339 - if output_scale is not None:
340 + if output_scale is not None or output_block_scale is not None:
341 raise NotImplementedError(
342 "fused output quantization is not yet supported"
343 " for XFormersAttentionImpl")