[Quant] Let DeepGEMM consume prequantized activations

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)

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

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),

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