From 22dc81241684b5b645eb34ba70e5aaed4158659b Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Wed, 12 Aug 2026 07:33:24 +0000 Subject: [PATCH] [Quant] Let DeepGEMM consume prequantized activations Advertise the dynamic block-FP8 input contract and bypass duplicate activation quantization in block-scaled DeepGEMM linear kernels. Co-authored-by: OpenAI Codex Signed-off-by: Woosuk Kwon --- .../fusion/test_quant_activation_contract.py | 45 +++++++++++++++++++ .../scaled_mm/BlockScaledMMLinearKernel.py | 21 ++++++--- .../kernels/linear/scaled_mm/deep_gemm.py | 12 +++++ .../model_executor/layers/quantization/fp8.py | 7 ++- 4 files changed, 78 insertions(+), 7 deletions(-) diff --git a/tests/fusion/test_quant_activation_contract.py b/tests/fusion/test_quant_activation_contract.py index 48d492b8d2ee..87dc092f417e 100644 --- a/tests/fusion/test_quant_activation_contract.py +++ b/tests/fusion/test_quant_activation_contract.py @@ -5,6 +5,7 @@ import pytest import torch +import vllm.model_executor.kernels.linear.scaled_mm.deep_gemm as deep_gemm from vllm.model_executor.kernels.linear import ( _POSSIBLE_FP8_BLOCK_KERNELS, _POSSIBLE_FP8_KERNELS, @@ -19,9 +20,15 @@ FlashInferCutlassNvFp4LinearKernel, FlashInferTrtllmNvFp4LinearKernel, ) +from vllm.model_executor.kernels.linear.scaled_mm.BlockScaledMMLinearKernel import ( + Fp8BlockScaledMMLinearKernel, +) from vllm.model_executor.kernels.linear.scaled_mm.cutlass import ( CutlassFP8ScaledMMLinearKernel, ) +from vllm.model_executor.kernels.linear.scaled_mm.deep_gemm import ( + DeepGemmFp8BlockScaledMMKernel, +) from vllm.model_executor.kernels.linear.scaled_mm.flashinfer import ( FlashInferFP8ScaledMMLinearKernel, ) @@ -36,6 +43,8 @@ expose_input_quant_key, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8Dynamic128Sym, + kFp8Static128BlockSym, kFp8StaticTensorSym, kNvfp4Dynamic, ) @@ -44,6 +53,7 @@ # The only backends that consume a pre-quantized activation. SUPPORTING = { CutlassFP8ScaledMMLinearKernel, + DeepGemmFp8BlockScaledMMKernel, FlashInferFP8ScaledMMLinearKernel, FlashInferCutlassNvFp4LinearKernel, } @@ -73,6 +83,14 @@ def _probe(cls: type): obj.config = Int8ScaledMMLinearLayerConfig( is_static_input_scheme=True, is_channelwise=False, input_symmetric=True ) + elif issubclass(cls, Fp8BlockScaledMMLinearKernel): + obj.config = FP8ScaledMMLinearLayerConfig( + weight_quant_key=kFp8Static128BlockSym, + activation_quant_key=kFp8Dynamic128Sym, + weight_shape=(128, 128), + input_dtype=torch.bfloat16, + out_dtype=torch.bfloat16, + ) else: obj.config = FP8ScaledMMLinearLayerConfig( weight_quant_key=kFp8StaticTensorSym, @@ -96,6 +114,33 @@ def test_only_known_backends_support_prequantized_input(): assert declarers == SUPPORTING +def test_deepgemm_custom_op_hides_compiler_scale_storage_padding(monkeypatch): + q_input = torch.empty(3, 128) + input_scale = torch.empty_strided((4, 1), (1, 4), dtype=torch.int32) + weight = torch.empty(128, 128) + weight_scale = torch.empty(1, 1) + output = torch.empty(3, 128) + received_scale = None + + def fake_fp8_gemm_nt(a, b, out, *, is_deep_gemm_e8m0_used): + nonlocal received_scale + received_scale = a[1] + + monkeypatch.setattr(deep_gemm, "fp8_gemm_nt", fake_fp8_gemm_nt) + deep_gemm._fp8_gemm_nt_op( + q_input, + input_scale, + weight, + weight_scale, + output, + True, + ) + + assert received_scale is not None + assert received_scale.shape == (3, 1) + assert received_scale.stride() == input_scale.stride() + + def test_supporting_backend_declares_consume_via_helper(): for cls in SUPPORTING: fn = _resolved_apply_weights(cls) diff --git a/vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py b/vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py index c83885ca51c4..f39aa66d91d0 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py @@ -8,6 +8,10 @@ import torch from typing_extensions import Self +from vllm.model_executor.layers.fusion.quant_activation import ( + QuantizedActivation, + as_quantized_activation, +) from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 from vllm.model_executor.layers.quantization.utils.fp8_utils import ( process_fp8_weight_block_strategy, @@ -97,7 +101,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module): def apply_weights( self, layer: torch.nn.Module, - x: torch.Tensor, + x: torch.Tensor | QuantizedActivation, bias: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: @@ -112,15 +116,20 @@ def apply_weights( input_scale = params.input_scale scale_up = params.input_scale_ub - # View input as 2D matrix for fp8 methods - input_2d = x.view(-1, x.shape[-1]) - output_shape = [*x.shape[:-1], weight.shape[0]] + qa = as_quantized_activation(x, self.input_quant_key()) + if qa is not None: + q_input, input_scale = qa.data, qa.scale + output_shape = [*qa.orig_shape[:-1], weight.shape[0]] + else: + assert isinstance(x, torch.Tensor) + input_2d = x.view(-1, x.shape[-1]) + output_shape = [*x.shape[:-1], weight.shape[0]] - if self.apply_input_quant: + if qa is None and self.apply_input_quant: q_input, input_scale = self.quant_fp8( input_2d, input_scale, scale_up, use_triton=self.use_triton ) - else: + elif qa is None: q_input = input_2d # Provide a concrete placeholder so apply_block_scaled_mm args are # always Tensors. Subclasses with apply_input_quant=False must not diff --git a/vllm/model_executor/kernels/linear/scaled_mm/deep_gemm.py b/vllm/model_executor/kernels/linear/scaled_mm/deep_gemm.py index 70122f7b4ac6..587b31954ff1 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/deep_gemm.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/deep_gemm.py @@ -11,6 +11,8 @@ ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( GroupShape, + QuantKey, + kFp8Dynamic128Sym, ) from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform @@ -43,6 +45,11 @@ def __init__(self, config: FP8ScaledMMLinearLayerConfig): column_major_scales=True, ) + def input_quant_key(self) -> QuantKey | None: + if self.config.activation_quant_key == kFp8Dynamic128Sym: + return kFp8Dynamic128Sym + return None + @classmethod def is_supported(cls, compute_capability=None): if not current_platform.is_cuda(): @@ -131,6 +138,11 @@ def _fp8_gemm_nt_op( output: torch.Tensor, use_deep_gemm_e8m0: bool, ) -> None: + if use_deep_gemm_e8m0 and input_scale.dtype == torch.int32: + logical_rows = q_input.shape[0] + aligned_rows = (logical_rows + 3) // 4 * 4 + if input_scale.shape[0] == aligned_rows and aligned_rows != logical_rows: + input_scale = input_scale[:logical_rows] fp8_gemm_nt( (q_input, input_scale), (weight, weight_scale), diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 9c69ff93d305..8f972938d6be 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -33,6 +33,10 @@ make_fp8_moe_quant_config, select_fp8_moe_backend, ) +from vllm.model_executor.layers.fusion.quant_activation import ( + QuantizedActivation, + expose_input_quant_key, +) from vllm.model_executor.layers.linear import ( LinearBase, LinearMethodBase, @@ -364,6 +368,7 @@ def create_weights( out_dtype=self.out_dtype, module_name=self.__class__.__name__, ) + expose_input_quant_key(layer, self.fp8_linear) self.use_marlin = isinstance(self.fp8_linear, MarlinFP8ScaledMMLinearKernel) @@ -418,7 +423,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: def apply( self, layer: torch.nn.Module, - x: torch.Tensor, + x: torch.Tensor | QuantizedActivation, bias: torch.Tensor | None = None, ) -> torch.Tensor: # if batch invariant mode is enabled, prefer direct FP8 path