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

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

[Quantization] add online fp8 ptpc (#44132)

Signed-off-by: walterbm <walter.beller.morales@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Walter Beller-Morales committed Jun 8, 2026 at 10:42 UTC 753e9d55e6a37f7eaa698d11bcb15990845d4f08
6 files changed +385 -8
tests/models/quantization/test_fp8_per_channel.py new
+79
@@ -0,0 +1,79 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""E2E tests for online FP8 per-channel quantization.
5 +
6 +Loads a BF16 model with ``--quantization fp8_per_channel`` (online
7 +quantization) and compares log-probabilities against the same model served in
8 +BF16 without quantization. This exercises the full pipeline: config parsing,
9 +``Fp8PtpcOnlineLinearMethod``, ``Fp8PtpcOnlineMoEMethod``, weight
10 +loading, online quantization / shuffling, and inference.
11 +
12 +``example_prompts`` is a pytest fixture (from conftest.py) that loads 8
13 +diverse prompts from ``tests/prompts/example.txt``.
14 +"""
15 +
16 +import pytest
17 +
18 +from tests.quantization.utils import is_quant_method_supported
19 +
20 +from ..utils import check_logprobs_close
21 +
22 +# Small MoE model that fits on a single GPU and exercises both linear + MoE.
23 +MOE_MODEL = "allenai/OLMoE-1B-7B-0125-Instruct"
24 +# Small dense model (no MoE) to validate the linear-only path.
25 +DENSE_MODEL = "Qwen/Qwen3-0.6B"
26 +
27 +MAX_MODEL_LEN = 1024
28 +MAX_TOKENS = 4
29 +NUM_LOG_PROBS = 8
30 +
31 +
32 +@pytest.mark.skipif(
33 + not is_quant_method_supported("fp8"),
34 + reason="fp8 is not supported on this GPU type.",
35 +)
36 +@pytest.mark.quant_model
37 +@pytest.mark.parametrize("model", [DENSE_MODEL, MOE_MODEL], ids=["dense", "moe"])
38 +def test_fp8_per_channel_logprobs(
39 + vllm_runner,
40 + example_prompts,
41 + model: str,
42 + monkeypatch: pytest.MonkeyPatch,
43 +) -> None:
44 + """Compare BF16 baseline logprobs against online per-channel-quantized
45 + model.
46 +
47 + Runs the same model twice -- once in BF16 (baseline) and once with online
48 + FP8 per-channel quantization -- then checks that the top log-probabilities
49 + are close. Only 4 tokens are generated to keep the test fast while still
50 + catching numerical divergence beyond expected per-channel error.
51 + """
52 + with monkeypatch.context() as m:
53 + m.setenv("TOKENIZERS_PARALLELISM", "true")
54 +
55 + with vllm_runner(
56 + model,
57 + max_model_len=MAX_MODEL_LEN,
58 + enforce_eager=True,
59 + ) as vllm_model:
60 + baseline_outputs = vllm_model.generate_greedy_logprobs(
61 + example_prompts, MAX_TOKENS, NUM_LOG_PROBS
62 + )
63 +
64 + with vllm_runner(
65 + model,
66 + max_model_len=MAX_MODEL_LEN,
67 + enforce_eager=True,
68 + quantization="fp8_per_channel",
69 + ) as vllm_model:
70 + test_outputs = vllm_model.generate_greedy_logprobs(
71 + example_prompts, MAX_TOKENS, NUM_LOG_PROBS
72 + )
73 +
74 + check_logprobs_close(
75 + outputs_0_lst=baseline_outputs,
76 + outputs_1_lst=test_outputs,
77 + name_0="bf16",
78 + name_1="fp8_per_channel",
79 + )
tests/quantization/test_fp8_per_channel.py new
+103
@@ -0,0 +1,103 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for FP8 per-channel online quantization.
4 +
5 +Per-output-channel weight scale + dynamic per-token activation scale.
6 +bf16/fp16 checkpoints are quantized at load time with one fp32 scale per
7 +output channel for weights and one fp32 scale per token for activations
8 +(computed dynamically inside the kernel). Run via
9 +`pytest tests/quantization/test_fp8_per_channel.py --forked`.
10 +"""
11 +
12 +import pytest
13 +import torch
14 +
15 +from tests.quantization.utils import is_quant_method_supported
16 +from vllm import _custom_ops as ops
17 +from vllm.config.quantization import (
18 + _ONLINE_SHORTHANDS,
19 + QUANT_KEY_NAMES,
20 + QuantizationConfigArgs,
21 +)
22 +from vllm.model_executor.layers.quantization.online.base import (
23 + _ONLINE_LINEAR_METHODS,
24 + _ONLINE_MOE_METHODS,
25 +)
26 +from vllm.model_executor.layers.quantization.online.fp8 import (
27 + Fp8PtpcOnlineLinearMethod,
28 + Fp8PtpcOnlineMoEMethod,
29 +)
30 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
31 + kFp8StaticChannelSym,
32 +)
33 +from vllm.platforms import current_platform
34 +
35 +
36 +def test_fp8_per_channel_shorthand_registered() -> None:
37 + """The `fp8_per_channel` CLI shorthand must resolve to a config that
38 + dispatches the per-channel methods. Guards against regressions in
39 + `_ONLINE_SHORTHANDS` / `_ONLINE_LINEAR_METHODS` / `_ONLINE_MOE_METHODS`
40 + drifting out of sync.
41 + """
42 + args = _ONLINE_SHORTHANDS["fp8_per_channel"]
43 + assert isinstance(args, QuantizationConfigArgs)
44 + assert args.linear is not None
45 + assert args.moe is not None
46 + assert args.linear.weight is kFp8StaticChannelSym
47 + assert args.moe.weight is kFp8StaticChannelSym
48 +
49 + assert _ONLINE_LINEAR_METHODS[kFp8StaticChannelSym] is Fp8PtpcOnlineLinearMethod
50 + assert _ONLINE_MOE_METHODS[kFp8StaticChannelSym] is Fp8PtpcOnlineMoEMethod
51 +
52 + assert QUANT_KEY_NAMES["fp8_per_channel_static"] is kFp8StaticChannelSym
53 +
54 +
55 +@pytest.mark.skipif(
56 + not is_quant_method_supported("fp8"),
57 + reason="FP8 is not supported on this GPU type.",
58 +)
59 +def test_scaled_fp8_quant_per_channel_shape() -> None:
60 + """Verify the kernel call per-channel quant depends on: passing a 2D
61 + weight to `ops.scaled_fp8_quant` with `use_per_token_if_dynamic=True`
62 + yields one scale per output row -- a [N, 1] fp32 tensor.
63 + """
64 + x = (torch.randn(size=(96, 256), device="cuda") * 13).to(torch.bfloat16)
65 + y, s = ops.scaled_fp8_quant(x, scale=None, use_per_token_if_dynamic=True)
66 + assert y.shape == (96, 256)
67 + assert y.dtype == current_platform.fp8_dtype()
68 + assert s.shape == (96, 1)
69 + assert s.dtype == torch.float32
70 +
71 +
72 +@pytest.mark.skipif(
73 + not is_quant_method_supported("fp8"),
74 + reason="FP8 is not supported on this GPU type.",
75 +)
76 +def test_fp8_per_channel_online_quantization(
77 + vllm_runner,
78 + monkeypatch,
79 +) -> None:
80 + """End-to-end smoke: load `facebook/opt-125m` bf16 with
81 + `quantization='fp8_per_channel'`, check a dense Linear is wrapped by
82 + `Fp8PtpcOnlineLinearMethod`, its weights are fp8 with per-channel
83 + scales (shape `[N, 1]`), and a short greedy generation works.
84 + """
85 + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
86 +
87 + with vllm_runner(
88 + "facebook/opt-125m",
89 + quantization="fp8_per_channel",
90 + enforce_eager=True,
91 + ) as llm:
92 +
93 + def check_model(model):
94 + fc1 = model.model.decoder.layers[0].fc1
95 + assert isinstance(fc1.quant_method, Fp8PtpcOnlineLinearMethod)
96 + assert fc1.weight.dtype == current_platform.fp8_dtype()
97 + assert fc1.weight_scale.ndim == 2
98 + assert fc1.weight_scale.shape[-1] == 1
99 + assert fc1.input_scale is None
100 +
101 + llm.apply_model(check_model)
102 + outputs = llm.generate_greedy(["Hello my name is"], max_tokens=4)
103 + print(outputs[0][1])
vllm/config/quantization.py
+8
@@ -13,6 +13,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
13 kFp8DynamicTensorSym,
14 kFp8DynamicTokenSym,
15 kFp8Static128BlockSym,
16 + kFp8StaticChannelSym,
17 kFp8StaticTensorSym,
18 kInt8StaticChannelSym,
19 kMxfp4Dynamic,
@@ -24,6 +25,7 @@ QUANT_KEY_NAMES: dict[str, QuantKey] = {
25 "fp8_per_tensor_static": kFp8StaticTensorSym,
26 "fp8_per_tensor_dynamic": kFp8DynamicTensorSym,
27 "fp8_per_token": kFp8DynamicTokenSym,
28 + "fp8_per_channel_static": kFp8StaticChannelSym,
29 "fp8_per_block_static": kFp8Static128BlockSym,
30 "fp8_per_block_dynamic": kFp8Dynamic128Sym,
31 "mxfp8": kMxfp8Dynamic,
@@ -118,6 +120,12 @@ _ONLINE_SHORTHANDS: dict[str, QuantizationConfigArgs] = {
120 linear=QuantSpec(weight=kFp8Static128BlockSym),
121 moe=QuantSpec(weight=kFp8Static128BlockSym),
122 ),
123 + # Per-output-channel weight scale + dynamic per-token activation.
124 + # Same shape as llmcompressor's FP8_DYNAMIC recipe.
125 + "fp8_per_channel": QuantizationConfigArgs(
126 + linear=QuantSpec(weight=kFp8StaticChannelSym),
127 + moe=QuantSpec(weight=kFp8StaticChannelSym),
128 + ),
129 "mxfp8": QuantizationConfigArgs(
130 linear=QuantSpec(weight=kMxfp8Dynamic),
131 moe=QuantSpec(weight=kMxfp8Dynamic),
vllm/model_executor/layers/quantization/__init__.py
+1
@@ -40,6 +40,7 @@ QuantizationMethods = Literal[
40 # _ONLINE_SHORTHANDS by the assertion in get_quantization_config().
41 "fp8_per_tensor",
42 "fp8_per_block",
43 + "fp8_per_channel",
44 "int8_per_channel_weight_only",
45 "mxfp8",
46 ]
vllm/model_executor/layers/quantization/online/base.py
+5
@@ -30,6 +30,8 @@ from vllm.model_executor.layers.quantization.online.fp8 import (
30 Fp8PerBlockOnlineMoEMethod,
31 Fp8PerTensorOnlineLinearMethod,
32 Fp8PerTensorOnlineMoEMethod,
33 + Fp8PtpcOnlineLinearMethod,
34 + Fp8PtpcOnlineMoEMethod,
35 )
36 from vllm.model_executor.layers.quantization.online.int8 import (
37 Int8OnlineMoEMethod,
@@ -41,6 +43,7 @@ from vllm.model_executor.layers.quantization.online.mxfp8 import (
43 from vllm.model_executor.layers.quantization.utils.quant_utils import (
44 QuantKey,
45 kFp8Static128BlockSym,
46 + kFp8StaticChannelSym,
47 kFp8StaticTensorSym,
48 kInt8StaticChannelSym,
49 kMxfp8Dynamic,
@@ -55,12 +58,14 @@ logger = init_logger(__name__)
58 _ONLINE_LINEAR_METHODS: dict[QuantKey, type] = {
59 kFp8StaticTensorSym: Fp8PerTensorOnlineLinearMethod,
60 kFp8Static128BlockSym: Fp8PerBlockOnlineLinearMethod,
61 + kFp8StaticChannelSym: Fp8PtpcOnlineLinearMethod,
62 kMxfp8Dynamic: Mxfp8OnlineLinearMethod,
63 }
64
65 _ONLINE_MOE_METHODS: dict[QuantKey, type] = {
66 kFp8StaticTensorSym: Fp8PerTensorOnlineMoEMethod,
67 kFp8Static128BlockSym: Fp8PerBlockOnlineMoEMethod,
68 + kFp8StaticChannelSym: Fp8PtpcOnlineMoEMethod,
69 kMxfp8Dynamic: Mxfp8OnlineMoEMethod,
70 kInt8StaticChannelSym: Int8OnlineMoEMethod,
71 }
vllm/model_executor/layers/quantization/online/fp8.py
+189 -8
@@ -12,6 +12,7 @@ if TYPE_CHECKING:
12 FusedMoEQuantConfig,
13 )
14 from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend
15 + from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
16
17 import vllm.envs as envs
18 from vllm import _custom_ops as ops
@@ -19,6 +20,7 @@ from vllm.config import get_current_vllm_config
20 from vllm.model_executor.kernels.linear import init_fp8_linear_kernel
21 from vllm.model_executor.kernels.linear.scaled_mm import (
22 CutlassFP8ScaledMMLinearKernel,
23 + MarlinFP8ScaledMMLinearKernel,
24 )
25 from vllm.model_executor.layers.fused_moe import RoutedExperts
26 from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
@@ -37,6 +39,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
39 kFp8DynamicTensorSym,
40 kFp8DynamicTokenSym,
41 kFp8Static128BlockSym,
42 + kFp8StaticChannelSym,
43 kFp8StaticTensorSym,
44 )
45 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
@@ -270,6 +273,89 @@ class Fp8PerBlockOnlineLinearMethod(_Fp8OnlineLinearBase):
273 )
274
275
276 +class Fp8PtpcOnlineLinearMethod(_Fp8OnlineLinearBase):
277 + """Online PTPC FP8 linear quantization.
278 +
279 + Per-output-channel weight scale + dynamic per-token activation scale. The
280 + layout matches the llmcompressor's FP8_DYNAMIC recipe, so accuracy
281 + is comparable but no pre-quantized checkpoint is required.
282 + """
283 +
284 + weight_quant_key = kFp8StaticChannelSym
285 + activation_quant_key = kFp8DynamicTokenSym
286 +
287 + def create_weights(
288 + self,
289 + layer: torch.nn.Module,
290 + input_size_per_partition: int,
291 + output_partition_sizes: list[int],
292 + input_size: int,
293 + output_size: int,
294 + params_dtype: torch.dtype,
295 + **extra_weight_attrs,
296 + ):
297 + super().create_weights(
298 + layer,
299 + input_size_per_partition,
300 + output_partition_sizes,
301 + input_size,
302 + output_size,
303 + params_dtype,
304 + **extra_weight_attrs,
305 + )
306 +
307 + self.fp8_linear = init_fp8_linear_kernel(
308 + activation_quant_key=self.activation_quant_key,
309 + weight_quant_key=self.weight_quant_key,
310 + weight_shape=layer.weight.shape,
311 + input_dtype=self.input_dtype,
312 + out_dtype=self.out_dtype,
313 + module_name=self.__class__.__name__,
314 + )
315 + # PTPC requires per-token activation FP8; MarlinFP8 is W8A16 and
316 + # would silently produce a weight-only fp8 model.
317 + if isinstance(self.fp8_linear, MarlinFP8ScaledMMLinearKernel):
318 + raise ValueError(
319 + "FP8 PTPC online quant requires a kernel that honors "
320 + "per-token activation quantization; MarlinFP8 is W8A16 "
321 + "weight-only. Requires SM89+ for Cutlass FP8 or ROCm MI3xx "
322 + "for rowwise scaled_mm."
323 + )
324 +
325 + def process_weights_after_loading(self, layer: Module) -> None:
326 + if getattr(layer, "_already_called_process_weights_after_loading", False):
327 + return
328 +
329 + layer.input_scale = None
330 + qweight, weight_scale = ops.scaled_fp8_quant(
331 + layer.weight, scale=None, use_per_token_if_dynamic=True
332 + )
333 +
334 + replace_parameter(layer, "weight", qweight.t())
335 + replace_parameter(layer, "weight_scale", weight_scale)
336 +
337 + self.fp8_linear.process_weights_after_loading(layer)
338 +
339 + layer._already_called_process_weights_after_loading = True
340 +
341 + def apply(
342 + self,
343 + layer: torch.nn.Module,
344 + x: torch.Tensor,
345 + bias: torch.Tensor | None = None,
346 + ) -> torch.Tensor:
347 + # if batch invariant mode is enabled dequant
348 + if envs.VLLM_BATCH_INVARIANT and not isinstance(
349 + self.fp8_linear, CutlassFP8ScaledMMLinearKernel
350 + ):
351 + weight_dequant = (
352 + layer.weight.to(x.dtype) * layer.weight_scale.to(x.dtype).t()
353 + )
354 + return torch.nn.functional.linear(x, weight_dequant.t(), bias)
355 +
356 + return self.fp8_linear.apply_weights(layer, x, bias)
357 +
358 +
359 # ---------------------------------------------------------------------------
360 # Online FP8 MoE Methods
361 # ---------------------------------------------------------------------------
@@ -284,12 +370,17 @@ class _Fp8OnlineMoEBase(OnlineMoEMethodBase):
370 experts_cls: "type[mk.FusedMoEExperts] | None"
371 weight_scale_name: str
372 weight_block_size: list[int] | None
373 + per_act_token_quant: bool = False
374 + per_out_ch_quant: bool = False
375
376 def __init__(
377 self,
378 *,
379 weight_block_size: list[int] | None,
380 layer: torch.nn.Module,
381 + weight_key: "QuantKey | None" = None,
382 + activation_key: "QuantKey | None" = None,
383 + allow_vllm_cutlass: bool = False,
384 ):
385 super().__init__(layer.moe_config)
386 self.weight_block_size = weight_block_size
@@ -298,20 +389,22 @@ class _Fp8OnlineMoEBase(OnlineMoEMethodBase):
389 "weight_scale_inv" if self.block_quant else "weight_scale"
390 )
391
301 - # Set weight key and activation key for kernel compatibility
302 - if self.block_quant:
303 - weight_key = kFp8Static128BlockSym
304 - activation_key = kFp8Dynamic128Sym
305 - else:
306 - weight_key = kFp8StaticTensorSym
307 - activation_key = kFp8DynamicTensorSym
392 + # Subclasses may pass explicit kernel keys (PTPC needs channelwise +
393 + # per-token).
394 + if weight_key is None or activation_key is None:
395 + if self.block_quant:
396 + weight_key = kFp8Static128BlockSym
397 + activation_key = kFp8Dynamic128Sym
398 + else:
399 + weight_key = kFp8StaticTensorSym
400 + activation_key = kFp8DynamicTensorSym
401
402 # Select Fp8 MoE backend
403 self.fp8_backend, self.experts_cls = select_fp8_moe_backend(
404 config=self.moe,
405 weight_key=weight_key,
406 activation_key=activation_key,
314 - allow_vllm_cutlass=False,
407 + allow_vllm_cutlass=allow_vllm_cutlass,
408 )
409
410 def _setup_kernel(
@@ -380,6 +473,8 @@ class _Fp8OnlineMoEBase(OnlineMoEMethodBase):
473 w1_bias=getattr(layer, "w13_bias", None),
474 w2_bias=getattr(layer, "w2_bias", None),
475 block_shape=self.weight_block_size,
476 + per_act_token_quant=self.per_act_token_quant,
477 + per_out_ch_quant=self.per_out_ch_quant,
478 swiglu_limit=getattr(layer, "swiglu_limit", None),
479 )
480
@@ -511,3 +606,89 @@ class Fp8PerBlockOnlineMoEMethod(_Fp8OnlineMoEBase):
606
607 # Prevent duplicate processing (e.g., during weight reload)
608 layer._already_called_process_weights_after_loading = True
609 +
610 +
611 +class Fp8PtpcOnlineMoEMethod(_Fp8OnlineMoEBase):
612 + """Online PTPC FP8 MoE quantization.
613 +
614 + Quantizes each expert's weights per output channel during loading.
615 + Activations are quantized dynamically per token at runtime.
616 + """
617 +
618 + per_act_token_quant: bool = True
619 + per_out_ch_quant: bool = True
620 +
621 + def __init__(
622 + self,
623 + *,
624 + layer: torch.nn.Module,
625 + ):
626 + from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend
627 +
628 + super().__init__(
629 + weight_block_size=None,
630 + layer=layer,
631 + weight_key=kFp8StaticChannelSym,
632 + activation_key=kFp8DynamicTokenSym,
633 + allow_vllm_cutlass=True,
634 + )
635 + # Reject backends whose make_fp8_moe_quant_config branch silently
636 + # drops per_act_token_quant / per_out_ch_quant or collapses scales:
637 + # MARLIN / CPU route through fp8_w8a16_moe_quant_config; FLASHINFER_*
638 + # fold scales into a per-tensor alpha (oracle/fp8.py).
639 + if self.fp8_backend in (
640 + Fp8MoeBackend.MARLIN,
641 + Fp8MoeBackend.CPU,
642 + Fp8MoeBackend.FLASHINFER_CUTLASS,
643 + Fp8MoeBackend.FLASHINFER_TRTLLM,
644 + ):
645 + raise ValueError(
646 + f"FP8 PTPC online MoE quant is not supported with the "
647 + f"{self.fp8_backend.value} backend, which does not implement "
648 + "per-output-channel weight scales."
649 + )
650 +
651 + def process_weights_after_loading(self, layer: Module) -> None:
652 + if getattr(layer, "_already_called_process_weights_after_loading", False):
653 + return
654 +
655 + fp8_dtype = current_platform.fp8_dtype()
656 + w13 = torch.empty_like(layer.w13_weight, dtype=fp8_dtype)
657 + w2 = torch.empty_like(layer.w2_weight, dtype=fp8_dtype)
658 + # Scale's leading dim is taken from the fp8 weight tensor by
659 + # construction, so it cannot drift from the weight's expert count
660 + # under EP / padded MoE.
661 + n_w13 = layer.w13_weight.shape[1]
662 + n_w2 = layer.w2_weight.shape[1]
663 + w13_scale = torch.ones(
664 + w13.shape[0], n_w13, 1, device=w13.device, dtype=torch.float32
665 + )
666 + w2_scale = torch.ones(
667 + w2.shape[0], n_w2, 1, device=w2.device, dtype=torch.float32
668 + )
669 + layer.w13_input_scale = None
670 + layer.w2_input_scale = None
671 +
672 + for expert in range(layer.local_num_experts):
673 + w13[expert], w13_scale[expert] = ops.scaled_fp8_quant(
674 + layer.w13_weight[expert],
675 + scale=None,
676 + use_per_token_if_dynamic=True,
677 + )
678 + w2[expert], w2_scale[expert] = ops.scaled_fp8_quant(
679 + layer.w2_weight[expert],
680 + scale=None,
681 + use_per_token_if_dynamic=True,
682 + )
683 +
684 + self._setup_kernel(
685 + layer,
686 + w13,
687 + w2,
688 + w13_scale,
689 + w2_scale,
690 + w13_input_scale=None,
691 + w2_input_scale=None,
692 + )
693 +
694 + layer._already_called_process_weights_after_loading = True