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
Move query quantization to attention layer for Flashinfer & Triton. (#26534)
Signed-off-by: adabeyta <aabeyta@redhat.com> Signed-off-by: Adrian Abeyta <aabeyta@redhat.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Adrian Abeyta committed
Oct 15, 2025 at 18:01 UTC
0a9ef0cfce131300fea219952b6e85c8dee45578
6 files changed
+43
-38
tests/compile/test_fusion_attn.py
+3
-1
@@ -421,7 +421,9 @@ def test_attention_quant_pattern(
421
]
422
if any(attn_fusion_supported):
423
# Check quantization ops in the graph before and after fusion
424
- test_backend.check_before_ops([QUANT_OPS[quant_key]], fully_replaced=True)
424
+ # Note: fully_replaced=False because query quant ops remain in graph.
425
+ # Only output quant ops are fused into attention.
426
+ test_backend.check_before_ops([QUANT_OPS[quant_key]], fully_replaced=False)
427
428
# access the underlying `AttnFusionPass` on the `LazyInitPass`
429
assert attn_pass.pass_.matched_count == sum(attn_fusion_supported)
vllm/attention/backends/abstract.py
+16
-8
@@ -41,14 +41,6 @@ class AttentionBackend(ABC):
41
# makes sure the output tensor is allocated inside the cudagraph.
42
accept_output_buffer: bool = False
43
44
- # Whether this backend supports receiving pre-quantized query input.
45
- # If True, the attention layer will handle query quantization instead
46
- # of the backend, allowing torch.compile to fuse quantization with
47
- # previous operations.
48
- # Needs to be worked through for all backends
49
- # https://github.com/vllm-project/vllm/issues/25584
50
- supports_quant_query_input: bool = False
51
-
44
@staticmethod
45
@abstractmethod
46
def get_name() -> str:
@@ -199,6 +191,22 @@ class AttentionImpl(ABC, Generic[T]):
191
"""
192
return False
193
194
+ def supports_quant_query_input(self) -> bool:
195
+ """
196
+ Check if this attention implementation supports pre-quantized query input.
197
+
198
+ When True, the attention layer will quantize queries before passing them
199
+ to this backend, allowing torch.compile to fuse the quantization with
200
+ previous operations. This is typically supported when using FP8 KV cache
201
+ with compatible attention kernels (e.g., TRT-LLM).
202
+ TODO add support to more backends:
203
+ https://github.com/vllm-project/vllm/issues/25584
204
+
205
+ Returns:
206
+ bool: True if the implementation can accept pre-quantized queries.
207
+ """
208
+ return False
209
+
210
211
class MLAAttentionImpl(AttentionImpl[T], Generic[T]):
212
@abstractmethod
vllm/attention/layer.py
+6
-3
@@ -36,6 +36,7 @@ from vllm.model_executor.models.vision import get_vit_attn_backend
36
from vllm.platforms import current_platform
37
from vllm.utils import GiB_bytes, direct_register_custom_op
38
39
+FP8_DTYPE = current_platform.fp8_dtype()
40
logger = init_logger(__name__)
41
USE_XFORMERS_OPS = None
42
@@ -304,7 +305,7 @@ class Attention(nn.Module, AttentionLayerBase):
305
self.query_quant = None
306
if (
307
self.kv_cache_dtype.startswith("fp8")
307
- and self.attn_backend.supports_quant_query_input
308
+ and self.impl.supports_quant_query_input()
309
):
310
self.query_quant = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR)
311
@@ -329,7 +330,6 @@ class Attention(nn.Module, AttentionLayerBase):
330
"""
331
if self.calculate_kv_scales:
332
torch.ops.vllm.maybe_calc_kv_scales(query, key, value, self.layer_name)
332
-
333
output_dtype = query.dtype
334
if self.query_quant is not None:
335
# quantizing with a simple torch operation enables
@@ -338,7 +338,10 @@ class Attention(nn.Module, AttentionLayerBase):
338
# Otherwise queries are quantized using custom ops
339
# which causes decoding overheads
340
assert self.kv_cache_dtype in {"fp8", "fp8_e4m3"}
341
- query, _ = self.query_quant(query, self._q_scale)
341
+
342
+ # check if query quantization is supported
343
+ if self.impl.supports_quant_query_input():
344
+ query, _ = self.query_quant(query, self._q_scale)
345
346
if self.use_output:
347
output_shape = output_shape if output_shape is not None else query.shape
vllm/v1/attention/backends/flash_attn.py
+3
-1
@@ -49,7 +49,6 @@ logger = init_logger(__name__)
49
50
class FlashAttentionBackend(AttentionBackend):
51
accept_output_buffer: bool = True
52
- supports_quant_query_input: bool = True
52
53
@classmethod
54
def get_supported_dtypes(cls) -> list[torch.dtype]:
@@ -494,6 +493,9 @@ class FlashAttentionImpl(AttentionImpl):
493
"heads in the layer"
494
)
495
496
+ def supports_quant_query_input(self) -> bool:
497
+ return True
498
+
499
def forward(
500
self,
501
layer: torch.nn.Module,
vllm/v1/attention/backends/flashinfer.py
+12
-10
@@ -16,7 +16,6 @@ from flashinfer.decode import _get_range_buf, trtllm_batch_decode_with_kv_cache
16
from flashinfer.prefill import trtllm_batch_context_with_kv_cache
17
from flashinfer.utils import FP4Tensor
18
19
-from vllm import _custom_ops as ops
19
from vllm.attention.backends.abstract import (
20
AttentionBackend,
21
AttentionImpl,
@@ -828,6 +827,12 @@ class FlashInferImpl(AttentionImpl):
827
and quant_key in (kFp8StaticTensorSym, kNvfp4Quant)
828
)
829
830
+ def supports_quant_query_input(self) -> bool:
831
+ if flashinfer_disable_q_quantization():
832
+ return False
833
+
834
+ return self.support_trtllm_attn
835
+
836
def forward(
837
self,
838
layer: torch.nn.Module,
@@ -859,6 +864,12 @@ class FlashInferImpl(AttentionImpl):
864
# Profiling run.
865
return output.fill_(0)
866
867
+ # Ensure query dtype matches the expected dtype from attention metadata
868
+ assert attn_metadata.q_data_type == query.dtype, (
869
+ f"Query dtype mismatch: expected {attn_metadata.q_data_type}, "
870
+ f"got {query.dtype}"
871
+ )
872
+
873
if self.bmm1_scale is None:
874
self.bmm1_scale = layer._q_scale_float * layer._k_scale_float * self.scale
875
@@ -899,15 +910,6 @@ class FlashInferImpl(AttentionImpl):
910
elif output.dtype == FP4_DTYPE:
911
self.o_sf_scale = layer._o_scale_float
912
902
- # Insert FP8 quant for query
903
- if attn_metadata.q_data_type == FP8_DTYPE:
904
- num_tokens, num_heads, head_size = query.shape
905
- query, _ = ops.scaled_fp8_quant(
906
- query.reshape((num_tokens, num_heads * head_size)).contiguous(),
907
- layer._q_scale,
908
- )
909
- query = query.reshape((num_tokens, num_heads, head_size))
910
-
913
# IMPORTANT!
914
# NOTE(woosuk): With piece-wise CUDA graphs, this method is executed in
915
# eager-mode PyTorch. Thus, we need to be careful about any CPU overhead
vllm/v1/attention/backends/triton_attn.py
+3
-15
@@ -32,11 +32,6 @@ from vllm.v1.attention.backends.utils import (
32
)
33
from vllm.v1.kv_cache_interface import AttentionSpec
34
35
-if current_platform.is_cuda_alike():
36
- from vllm import _custom_ops as ops
37
-elif current_platform.is_xpu():
38
- from vllm._ipex_ops import ipex_ops as ops
39
-
35
logger = init_logger(__name__)
36
37
@@ -210,6 +205,9 @@ class TritonAttentionImpl(AttentionImpl):
205
def fused_output_quant_supported(self, quant_key: QuantKey):
206
return quant_key == kFp8StaticTensorSym
207
208
+ def supports_quant_query_input(self) -> bool:
209
+ return current_platform.is_cuda()
210
+
211
def __init__(
212
self,
213
num_heads: int,
@@ -338,19 +336,9 @@ class TritonAttentionImpl(AttentionImpl):
336
if key_cache.dtype != self.fp8_dtype:
337
key_cache = key_cache.view(self.fp8_dtype)
338
value_cache = value_cache.view(self.fp8_dtype)
341
- num_tokens, num_heads, head_size = query.shape
339
assert layer._q_scale_float == 1.0, (
340
"A non 1.0 q_scale is not currently supported."
341
)
345
- if current_platform.is_cuda():
346
- # Skip Q quantization on ROCm and XPU, enable this on cuda
347
- # only, since dequantizing back to f32 in the attention kernel
348
- # is not supported.
349
- query, _ = ops.scaled_fp8_quant(
350
- query.reshape((num_tokens, num_heads * head_size)).contiguous(),
351
- layer._q_scale,
352
- )
353
- query = query.reshape((num_tokens, num_heads, head_size))
342
343
cu_seqlens_q = attn_metadata.query_start_loc
344
seqused_k = attn_metadata.seq_lens