@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / c90650088d
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Add the QuantizedActivation linear-kernel contract (#44260)

Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>

Michael Goin committed Jun 12, 2026 at 21:48 UTC c90650088dafc8ad5fc372b412b67170c5ad3f4a
13 files changed +327 -27
.buildkite/test_areas/quantization.yaml
+12
@@ -21,6 +21,18 @@ steps:
21 - uv pip install --system conch-triton-kernels
22 - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py
23
24 +- label: Quantized Fusions
25 + key: quantized-fusions
26 + timeout_in_minutes: 30
27 + source_file_dependencies:
28 + - tests/fusion
29 + - vllm/model_executor/layers/fusion
30 + - vllm/model_executor/kernels/linear
31 + - vllm/model_executor/layers/quantization/compressed_tensors
32 + - vllm/model_executor/layers/quantization/modelopt.py
33 + commands:
34 + - pytest -v -s fusion/
35 +
36 - label: Quantized MoE Test (B200)
37 key: quantized-moe-test-b200
38 timeout_in_minutes: 60
tests/fusion/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
tests/fusion/test_quant_activation_contract.py new
+131
@@ -0,0 +1,131 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Contract tests for the QuantizedActivation linear-kernel integration."""
4 +
5 +import pytest
6 +import torch
7 +
8 +from vllm.model_executor.kernels.linear import (
9 + _POSSIBLE_FP8_BLOCK_KERNELS,
10 + _POSSIBLE_FP8_KERNELS,
11 + _POSSIBLE_INT8_KERNELS,
12 + _POSSIBLE_NVFP4_KERNELS,
13 +)
14 +from vllm.model_executor.kernels.linear.nvfp4.base import (
15 + NvFp4LinearKernel,
16 + NvFp4LinearLayerConfig,
17 +)
18 +from vllm.model_executor.kernels.linear.nvfp4.flashinfer import (
19 + FlashInferCutlassNvFp4LinearKernel,
20 + FlashInferTrtllmNvFp4LinearKernel,
21 +)
22 +from vllm.model_executor.kernels.linear.scaled_mm.cutlass import (
23 + CutlassFP8ScaledMMLinearKernel,
24 +)
25 +from vllm.model_executor.kernels.linear.scaled_mm.flashinfer import (
26 + FlashInferFP8ScaledMMLinearKernel,
27 +)
28 +from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import (
29 + FP8ScaledMMLinearLayerConfig,
30 + Int8ScaledMMLinearKernel,
31 + Int8ScaledMMLinearLayerConfig,
32 +)
33 +from vllm.model_executor.layers.fusion.quant_activation import (
34 + QuantizedActivation,
35 + as_quantized_activation,
36 + expose_input_quant_key,
37 +)
38 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
39 + kFp8StaticTensorSym,
40 + kNvfp4Dynamic,
41 +)
42 +from vllm.platforms import current_platform
43 +
44 +# The only backends that consume a pre-quantized activation.
45 +SUPPORTING = {
46 + CutlassFP8ScaledMMLinearKernel,
47 + FlashInferFP8ScaledMMLinearKernel,
48 + FlashInferCutlassNvFp4LinearKernel,
49 +}
50 +
51 +
52 +def _all_kernel_classes() -> list[type]:
53 + seen: dict[type, None] = {}
54 + for registry in (
55 + _POSSIBLE_FP8_KERNELS,
56 + _POSSIBLE_FP8_BLOCK_KERNELS,
57 + _POSSIBLE_INT8_KERNELS,
58 + _POSSIBLE_NVFP4_KERNELS,
59 + ):
60 + for kernels in registry.values():
61 + for cls in kernels:
62 + seen.setdefault(cls, None)
63 + return list(seen)
64 +
65 +
66 +def _probe(cls: type):
67 + """A bare kernel instance with a plausible config, so input_quant_key()
68 + can be queried without the hardware-gated constructor."""
69 + obj = cls.__new__(cls) # type: ignore[call-overload]
70 + if issubclass(cls, NvFp4LinearKernel):
71 + obj.config = NvFp4LinearLayerConfig()
72 + elif issubclass(cls, Int8ScaledMMLinearKernel):
73 + obj.config = Int8ScaledMMLinearLayerConfig(
74 + is_static_input_scheme=True, is_channelwise=False, input_symmetric=True
75 + )
76 + else:
77 + obj.config = FP8ScaledMMLinearLayerConfig(
78 + weight_quant_key=kFp8StaticTensorSym,
79 + activation_quant_key=kFp8StaticTensorSym,
80 + weight_shape=(16, 16),
81 + input_dtype=torch.bfloat16,
82 + out_dtype=torch.bfloat16,
83 + )
84 + return obj
85 +
86 +
87 +def _resolved_apply_weights(cls: type):
88 + for base in cls.__mro__:
89 + if "apply_weights" in base.__dict__:
90 + return base.__dict__["apply_weights"]
91 + raise AssertionError(f"{cls.__name__} has no apply_weights in its MRO")
92 +
93 +
94 +def test_only_known_backends_support_prequantized_input():
95 + declarers = {c for c in _all_kernel_classes() if _probe(c).input_quant_key()}
96 + assert declarers == SUPPORTING
97 +
98 +
99 +def test_supporting_backend_declares_consume_via_helper():
100 + for cls in SUPPORTING:
101 + fn = _resolved_apply_weights(cls)
102 + assert "as_quantized_activation" in fn.__code__.co_names, cls.__name__
103 +
104 +
105 +def test_bridge_marks_supporting_and_skips_others():
106 + supported = _probe(FlashInferCutlassNvFp4LinearKernel)
107 + layer = torch.nn.Module()
108 + expose_input_quant_key(layer, supported)
109 + assert layer.input_quant_key == kNvfp4Dynamic
110 +
111 + unsupported = _probe(FlashInferTrtllmNvFp4LinearKernel)
112 + assert unsupported.input_quant_key() is None
113 + layer = torch.nn.Module()
114 + expose_input_quant_key(layer, unsupported)
115 + assert not hasattr(layer, "input_quant_key")
116 +
117 +
118 +def test_as_quantized_activation_validates_key():
119 + qa = QuantizedActivation(
120 + data=torch.zeros(2, 4, dtype=current_platform.fp8_dtype()),
121 + scale=torch.tensor(1.0),
122 + orig_dtype=torch.bfloat16,
123 + orig_shape=torch.Size([2, 4]),
124 + quant_key=kFp8StaticTensorSym,
125 + )
126 + with pytest.raises(AssertionError):
127 + as_quantized_activation(qa, kNvfp4Dynamic)
128 + with pytest.raises(AssertionError):
129 + as_quantized_activation(qa, None)
130 + assert as_quantized_activation(torch.zeros(2, 4), kFp8StaticTensorSym) is None
131 + assert as_quantized_activation(qa, kFp8StaticTensorSym) is qa
vllm/model_executor/kernels/linear/base.py
+8
@@ -8,6 +8,8 @@ from typing import Any, ClassVar, Generic, TypeVar
8 import torch
9 from typing_extensions import Self
10
11 +from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
12 +
13
14 @dataclass
15 class MMLinearLayerConfig: ...
@@ -237,6 +239,12 @@ class MMLinearKernel(ABC, Generic[_ConfigT, _ParamsT]):
239 """
240 self.config = config
241
242 + def input_quant_key(self) -> QuantKey | None:
243 + """Return the input quantization key supported by this kernel. If the kernel
244 + does not support input quantization outside of the kernel, return None.
245 + """
246 + return None
247 +
248 @abstractmethod
249 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
250 """Process and transform weights after loading from checkpoint.
vllm/model_executor/kernels/linear/nvfp4/base.py
+8
@@ -6,6 +6,8 @@ from dataclasses import dataclass
6
7 import torch
8
9 +from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
10 +
11
12 @dataclass
13 class NvFp4LinearLayerConfig:
@@ -33,6 +35,12 @@ class NvFp4LinearKernel(ABC):
35 assert self.is_supported()[0]
36 self.config = config
37
38 + def input_quant_key(self) -> QuantKey | None:
39 + """Return the input quantization key supported by this kernel. If the kernel
40 + does not support input quantization outside of the kernel, return None.
41 + """
42 + return None
43 +
44 @classmethod
45 @abstractmethod
46 def is_supported(
vllm/model_executor/kernels/linear/nvfp4/flashinfer.py
+31 -10
@@ -4,12 +4,20 @@
4 import torch
5
6 from vllm._custom_ops import scaled_fp4_quant
7 +from vllm.model_executor.layers.fusion.quant_activation import (
8 + QuantizedActivation,
9 + as_quantized_activation,
10 +)
11 from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
12 pad_nvfp4_activation_for_cutlass,
13 pad_nvfp4_weight_for_cutlass,
14 slice_nvfp4_output,
15 swizzle_blockscale,
16 )
17 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
18 + QuantKey,
19 + kNvfp4Dynamic,
20 +)
21 from vllm.platforms import current_platform
22 from vllm.utils.flashinfer import (
23 flashinfer_scaled_fp4_mm,
@@ -23,6 +31,11 @@ from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig
31 class FlashInferCutlassNvFp4LinearKernel(NvFp4LinearKernel):
32 """NVFP4 GEMM via FlashInfer's CUTLASS wrapper."""
33
34 + def input_quant_key(self) -> QuantKey | None:
35 + """This kernel supports dynamic quantization of the input. By
36 + convention, pre-quantized blockscales must use the swizzled layout."""
37 + return kNvfp4Dynamic
38 +
39 @classmethod
40 def is_supported(
41 cls, compute_capability: int | None = None
@@ -56,21 +69,29 @@ class FlashInferCutlassNvFp4LinearKernel(NvFp4LinearKernel):
69 def apply_weights(
70 self,
71 layer: torch.nn.Module,
59 - x: torch.Tensor,
72 + x: torch.Tensor | QuantizedActivation,
73 bias: torch.Tensor | None = None,
74 ) -> torch.Tensor:
75 output_size = layer.output_size_per_partition
63 - output_dtype = x.dtype
64 - output_shape = [*x.shape[:-1], output_size]
76 weights_padding_bytes = getattr(layer, "weights_padding_cols", 0)
77
67 - x_fp4, x_blockscale = scaled_fp4_quant(
68 - x,
69 - layer.input_global_scale_inv,
70 - is_sf_swizzled_layout=True,
71 - backend="flashinfer-cutlass",
72 - padded_n=x.shape[-1] + weights_padding_bytes * 2,
73 - )
78 + qa = as_quantized_activation(x, self.input_quant_key())
79 + if qa is not None:
80 + x_fp4, x_blockscale = qa.data, qa.scale
81 + x_fp4 = pad_nvfp4_activation_for_cutlass(x_fp4, weights_padding_bytes)
82 + output_dtype = qa.orig_dtype
83 + output_shape = [*qa.orig_shape[:-1], output_size]
84 + else:
85 + assert isinstance(x, torch.Tensor)
86 + output_dtype = x.dtype
87 + output_shape = [*x.shape[:-1], output_size]
88 + x_fp4, x_blockscale = scaled_fp4_quant(
89 + x,
90 + layer.input_global_scale_inv,
91 + is_sf_swizzled_layout=True,
92 + backend="flashinfer-cutlass",
93 + padded_n=x.shape[-1] + weights_padding_bytes * 2,
94 + )
95
96 out = flashinfer_scaled_fp4_mm(
97 x_fp4,
vllm/model_executor/kernels/linear/scaled_mm/ScaledMMLinearKernel.py
+31 -16
@@ -8,6 +8,10 @@ from typing import Generic, TypeVar
8
9 import torch
10
11 +from vllm.model_executor.layers.fusion.quant_activation import (
12 + QuantizedActivation,
13 + as_quantized_activation,
14 +)
15 from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
16 from vllm.model_executor.layers.quantization.utils.quant_utils import (
17 QuantKey,
@@ -71,6 +75,17 @@ class ScaledMMLinearKernel(Generic[_ConfigT, _ParamsT], ABC):
75 self.config = c
76 self.layer_param_names = layer_param_names
77
78 + def input_quant_key(self) -> QuantKey | None:
79 + """The activation quant key this kernel can consume pre-quantized.
80 +
81 + Manual fusion uses this to decide whether to hoist activation
82 + quantization out of apply_weights into an upstream fused kernel.
83 + Return None when the kernel needs in-kernel quantization (custom
84 + padding or swizzling, dynamic scales, etc.). Kernels that return a
85 + key must consume the activation via as_quantized_activation.
86 + """
87 + return None
88 +
89 @abstractmethod
90 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
91 raise NotImplementedError
@@ -120,30 +135,30 @@ class FP8ScaledMMLinearKernel(
135 def apply_weights(
136 self,
137 layer: torch.nn.Module,
123 - x: torch.Tensor,
138 + x: torch.Tensor | QuantizedActivation,
139 bias: torch.Tensor | None = None,
140 ) -> torch.Tensor:
141 fp8_dtype = self.fp8_dtype
142 maybe_out_dtype = self.config.out_dtype
143 w, w_s, x_s, x_s_ub = self._get_layer_params(layer)
144
130 - # ops.scaled_fp8_quant supports both dynamic and static quant.
131 - # If dynamic, layer.input_scale is None and x_s computed from x.
132 - # If static, layer.input_scale is scalar and x_s is input_scale.
133 - # View input as 2D matrix for fp8 methods
134 - x_2d = x.view(-1, x.shape[-1])
135 - output_shape = [*x.shape[:-1], w.shape[1]]
136 - out_dtype = x.dtype if maybe_out_dtype is None else maybe_out_dtype
145 + qa = as_quantized_activation(x, self.input_quant_key())
146 + if qa is not None:
147 + x_data, x_s = qa.data, qa.scale
148 + orig_shape, orig_dtype = qa.orig_shape, qa.orig_dtype
149 + assert x_data.dtype == fp8_dtype
150 + else:
151 + assert isinstance(x, torch.Tensor)
152 + x_data = x
153 + orig_shape, orig_dtype = x.shape, x.dtype
154 +
155 + x_2d = x_data.view(-1, x_data.shape[-1])
156 + output_shape = [*orig_shape[:-1], w.shape[1]]
157 + out_dtype = orig_dtype if maybe_out_dtype is None else maybe_out_dtype
158
138 - # If input not quantized
139 - # TODO(luka) remove this path if not used anymore
159 x_2d_q = x_2d
141 - if x.dtype != fp8_dtype:
142 - x_2d_q, x_s = self.quant_fp8(
143 - x_2d,
144 - x_s,
145 - x_s_ub,
146 - )
160 + if qa is None:
161 + x_2d_q, x_s = self.quant_fp8(x_2d, x_s, x_s_ub)
162 return self.apply_scaled_mm(
163 A=x_2d_q,
164 B=w,
vllm/model_executor/kernels/linear/scaled_mm/cutlass.py
+9
@@ -11,6 +11,8 @@ from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
11 from vllm.model_executor.layers.quantization.utils import replace_parameter
12 from vllm.model_executor.layers.quantization.utils.quant_utils import (
13 GroupShape,
14 + QuantKey,
15 + kFp8StaticTensorSym,
16 )
17 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
18 CUTLASS_BLOCK_FP8_SUPPORTED,
@@ -171,6 +173,13 @@ class CutlassFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
173 def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
174 return True, None
175
176 + def input_quant_key(self) -> QuantKey | None:
177 + """Only static per-tensor activation quantization is supported for external
178 + quantization."""
179 + if self.config.activation_quant_key == kFp8StaticTensorSym:
180 + return kFp8StaticTensorSym
181 + return None
182 +
183 @staticmethod
184 def _pad_to_alignment(
185 x: torch.Tensor, dim: int, alignment: int, value: float = 0.0
vllm/model_executor/kernels/linear/scaled_mm/flashinfer.py
+7
@@ -12,6 +12,8 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
12 )
13 from vllm.model_executor.layers.quantization.utils.quant_utils import (
14 GroupShape,
15 + QuantKey,
16 + kFp8StaticTensorSym,
17 )
18 from vllm.platforms import current_platform
19 from vllm.utils.flashinfer import (
@@ -62,6 +64,11 @@ class FlashInferFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
64
65 return True, None
66
67 + def input_quant_key(self) -> QuantKey | None:
68 + if self.config.activation_quant_key == kFp8StaticTensorSym:
69 + return kFp8StaticTensorSym
70 + return None
71 +
72 def apply_scaled_mm(
73 self,
74 *,
vllm/model_executor/layers/fusion/quant_activation.py new
+71
@@ -0,0 +1,71 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +A QuantizedActivation is a pre-quantized activation produced by a fused kernel
5 +and consumed directly by a linear layer, letting the layer skip its own input
6 +quantization. A linear advertises the key its kernel can consume via
7 +expose_input_quant_key; the kernel validates and reads the activation via
8 +as_quantized_activation.
9 +"""
10 +
11 +from dataclasses import dataclass
12 +
13 +import torch
14 +
15 +from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
16 +
17 +
18 +@dataclass
19 +class QuantizedActivation:
20 + """A quantized activation paired with its scale and original metadata.
21 +
22 + The quant_key describes how data and scale are to be interpreted (dtype,
23 + scale granularity, value packing). Details the key does not capture, such
24 + as blockscale layout or activation padding, must follow the consumer
25 + kernel's convention.
26 +
27 + TODO(mgoin): Encode layout and padding requirements in the contract so
28 + producers can match consumer kernels without relying on convention.
29 + """
30 +
31 + data: torch.Tensor
32 + scale: torch.Tensor
33 + orig_dtype: torch.dtype
34 + orig_shape: torch.Size
35 + quant_key: QuantKey
36 +
37 +
38 +def expose_input_quant_key(layer: torch.nn.Module, kernel) -> None:
39 + """Advertise the kernel's pre-quantized input key on the layer, if any.
40 +
41 + This is the bridge from a kernel's input_quant_key() to the
42 + layer.input_quant_key attribute that fusion call sites read. The attribute
43 + is left unset when the kernel quantizes its own input, so non-supporting
44 + backends never receive a QuantizedActivation.
45 +
46 + TODO(mgoin): Producers also need the consumer's quantization scales (e.g.
47 + static input scale, global scale). Expose those here as well so producers
48 + do not reach into kernel-specific layer attributes.
49 + """
50 + key = kernel.input_quant_key()
51 + if key is not None:
52 + layer.input_quant_key = key
53 +
54 +
55 +def as_quantized_activation(
56 + x: "torch.Tensor | QuantizedActivation", expected_key: QuantKey | None
57 +) -> "QuantizedActivation | None":
58 + """Validate and narrow a pre-quantized activation for a consumer kernel.
59 +
60 + Returns the QuantizedActivation when x is one whose key matches the
61 + kernel's declared expected_key, and None when x is a plain tensor (the
62 + caller quantizes in-kernel). Raises on a key mismatch so a wrongly routed
63 + activation fails loudly instead of being silently re-quantized.
64 + """
65 + if not isinstance(x, QuantizedActivation):
66 + return None
67 + assert x.quant_key == expected_key, (
68 + f"QuantizedActivation key {x.quant_key} != consumer kernel "
69 + f"input_quant_key {expected_key}"
70 + )
71 + return x
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py
+5
@@ -7,6 +7,9 @@ from torch.nn.parameter import Parameter
7
8 from vllm.logger import init_logger
9 from vllm.model_executor.kernels.linear import init_nvfp4_linear_kernel
10 +from vllm.model_executor.layers.fusion.quant_activation import (
11 + expose_input_quant_key,
12 +)
13 from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
14 CompressedTensorsScheme,
15 )
@@ -87,6 +90,8 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):
90 )
91 layer.register_parameter("input_global_scale", input_global_scale)
92
93 + expose_input_quant_key(layer, self.kernel)
94 +
95 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
96 # Rename CT checkpoint names to standardized names
97 layer.weight = layer.weight_packed
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py
+7 -1
@@ -13,6 +13,10 @@ from vllm.logger import init_logger
13 from vllm.model_executor.kernels.linear import (
14 init_fp8_linear_kernel,
15 )
16 +from vllm.model_executor.layers.fusion.quant_activation import (
17 + QuantizedActivation,
18 + expose_input_quant_key,
19 +)
20 from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
21 CompressedTensorsScheme,
22 )
@@ -143,6 +147,8 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):
147 module_name=self.__class__.__name__,
148 )
149
150 + expose_input_quant_key(layer, self.fp8_linear)
151 +
152 def process_weights_after_loading(self, layer) -> None:
153 if self.strategy == QuantizationStrategy.TENSOR:
154 weight, weight_scale, input_scale = process_fp8_weight_tensor_strategy(
@@ -191,7 +197,7 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):
197 def apply_weights(
198 self,
199 layer: torch.nn.Module,
194 - x: torch.Tensor,
200 + x: torch.Tensor | QuantizedActivation,
201 bias: torch.Tensor | None = None,
202 ) -> torch.Tensor:
203 return self.fp8_linear.apply_weights(layer, x, bias)
vllm/model_executor/layers/quantization/modelopt.py
+5
@@ -42,6 +42,9 @@ from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
42 make_nvfp4_moe_quant_config,
43 select_nvfp4_moe_backend,
44 )
45 +from vllm.model_executor.layers.fusion.quant_activation import (
46 + expose_input_quant_key,
47 +)
48 from vllm.model_executor.layers.linear import (
49 LinearBase,
50 LinearMethodBase,
@@ -1191,6 +1194,8 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
1194
1195 layer.register_parameter("weight_scale", weight_scale)
1196
1197 + expose_input_quant_key(layer, self.kernel)
1198 +
1199 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
1200 if (
1201 torch.unique(layer.input_scale).numel() != 1