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] Make Query Quantization Fusable (#24914)
Signed-off-by: Jonas Kuebler <kuebj@amazon.com>
Jonas M. Kübler committed
Sep 25, 2025 at 15:25 UTC
69a8c8e99ab9e16028453efd5e7860063f002bf8
3 files changed
+32
-8
vllm/attention/backends/abstract.py
+8
@@ -31,6 +31,14 @@ class AttentionBackend(ABC):
31
# makes sure the output tensor is allocated inside the cudagraph.
32
accept_output_buffer: bool = False
33
34
+ # Whether this backend supports receiving pre-quantized query input.
35
+ # If True, the attention layer will handle query quantization instead
36
+ # of the backend, allowing torch.compile to fuse quantization with
37
+ # previous operations.
38
+ # Needs to be worked through for all backends
39
+ # https://github.com/vllm-project/vllm/issues/25584
40
+ supports_quant_query_input: bool = False
41
+
42
@staticmethod
43
@abstractmethod
44
def get_name() -> str:
vllm/attention/layer.py
+22
-1
@@ -22,7 +22,10 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
22
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
23
from vllm.model_executor.layers.quantization.base_config import (
24
QuantizationConfig)
25
+from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
26
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
27
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
28
+ GroupShape)
29
from vllm.model_executor.models.vision import get_vit_attn_backend
30
from vllm.platforms import _Backend, current_platform
31
from vllm.utils import GiB_bytes, direct_register_custom_op
@@ -247,6 +250,13 @@ class Attention(nn.Module, AttentionLayerBase):
250
"This may be caused by insufficient memory to allocate "
251
"kv cache.") from e
252
253
+ # for attn backends supporting query quantization
254
+ self.query_quant = None
255
+ if self.kv_cache_dtype.startswith(
256
+ "fp8") and self.attn_backend.supports_quant_query_input:
257
+ self.query_quant = QuantFP8(static=True,
258
+ group_shape=GroupShape.PER_TENSOR)
259
+
260
def forward(
261
self,
262
query: torch.Tensor,
@@ -270,11 +280,22 @@ class Attention(nn.Module, AttentionLayerBase):
280
attn_metadata = get_forward_context().attn_metadata
281
if attn_metadata.enable_kv_scales_calculation:
282
self.calc_kv_scales(query, key, value)
283
+
284
+ output_dtype = query.dtype
285
+ if self.query_quant is not None:
286
+ # quantizing with a simple torch operation enables
287
+ # torch.compile to fuse this into previous ops
288
+ # which reduces overheads during decoding.
289
+ # Otherwise queries are quantized using custom ops
290
+ # which causes decoding overheads
291
+ assert self.kv_cache_dtype in {"fp8", "fp8_e4m3"}
292
+ query, _ = self.query_quant(query, self._q_scale)
293
+
294
if self.use_output:
295
output_shape = (output_shape
296
if output_shape is not None else query.shape)
297
output = torch.zeros(output_shape,
277
- dtype=query.dtype,
298
+ dtype=output_dtype,
299
device=query.device)
300
hidden_size = output_shape[-1]
301
# We skip reshaping query, key and value tensors for the MLA
vllm/v1/attention/backends/flash_attn.py
+2
-7
@@ -7,7 +7,6 @@ from typing import Optional
7
import numpy as np
8
import torch
9
10
-from vllm import _custom_ops as ops
10
from vllm import envs
11
from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
12
AttentionMetadata, AttentionType,
@@ -38,6 +37,7 @@ logger = init_logger(__name__)
37
class FlashAttentionBackend(AttentionBackend):
38
39
accept_output_buffer: bool = True
40
+ supports_quant_query_input: bool = True
41
42
@classmethod
43
def get_supported_dtypes(cls) -> list[torch.dtype]:
@@ -506,16 +506,11 @@ class FlashAttentionImpl(AttentionImpl):
506
)
507
508
if self.kv_cache_dtype.startswith("fp8"):
509
+ # queries are quantized in the attention layer
510
dtype = FlashAttentionBackend.get_fp8_dtype_for_flashattn(
511
self.kv_cache_dtype)
512
key_cache = key_cache.view(dtype)
513
value_cache = value_cache.view(dtype)
513
- num_tokens, num_heads, head_size = query.shape
514
- query, _ = ops.scaled_fp8_quant(
515
- query.reshape(
516
- (num_tokens, num_heads * head_size)).contiguous(),
517
- layer._q_scale)
518
- query = query.reshape((num_tokens, num_heads, head_size))
514
515
if not attn_metadata.use_cascade:
516
cu_seqlens_q = attn_metadata.query_start_loc