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

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 6, 14:30 UTC

[Quant][Feature] Support online MXFP8 quantization for MoE and dense models (#35448)

Signed-off-by: EdalatiAli <aliedalati@cohere.com>

EdalatiAli committed Mar 16, 2026 at 18:07 UTC e5b807607c8493155e6eccd665772d4c19b2114e
10 files changed +745 -54
tests/models/quantization/test_mxfp8.py new
+104
@@ -0,0 +1,104 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""E2E tests for online MXFP8 quantization.
5 +
6 +Loads a BF16 model with ``--quantization mxfp8`` (online quantization) and
7 +compares log-probabilities against the same model served in BF16 without
8 +quantization. This exercises the full pipeline: config parsing,
9 +``Mxfp8OnlineLinearMethod``, ``Mxfp8OnlineMoEMethod``, weight loading,
10 +online quantization / shuffling, and inference through ``apply_monolithic``.
11 +
12 +Layer skipping (``modules_to_not_convert``) is configured in the model's
13 +``config.json`` under ``quantization_config`` and is not tested here.
14 +
15 +``example_prompts`` is a pytest fixture (from conftest.py) that loads 8
16 +diverse prompts from ``tests/prompts/example.txt``.
17 +"""
18 +
19 +import pytest
20 +
21 +from tests.quantization.utils import is_quant_method_supported
22 +
23 +from ..utils import check_logprobs_close
24 +
25 +# A small MoE model that fits on a single GPU and has both linear + MoE layers.
26 +MOE_MODEL = "Qwen/Qwen3-30B-A3B"
27 +# A small dense model (no MoE) to validate the linear-only path.
28 +DENSE_MODEL = "Qwen/Qwen3-0.6B"
29 +
30 +MAX_MODEL_LEN = 1024
31 +MAX_TOKENS = 4
32 +NUM_LOG_PROBS = 8
33 +
34 +
35 +@pytest.mark.skipif(
36 + not is_quant_method_supported("mxfp8"),
37 + reason="mxfp8 is not supported on this GPU type (requires sm_100+).",
38 +)
39 +@pytest.mark.quant_model
40 +@pytest.mark.parametrize("model", [DENSE_MODEL, MOE_MODEL], ids=["dense", "moe"])
41 +def test_mxfp8_logprobs(
42 + vllm_runner,
43 + example_prompts,
44 + model: str,
45 + monkeypatch: pytest.MonkeyPatch,
46 +) -> None:
47 + """Compare BF16 baseline logprobs against online MXFP8-quantized model.
48 +
49 + Runs the same model twice -- once in BF16 (baseline) and once with
50 + online MXFP8 quantization -- then checks that the top log-probabilities
51 + are close. Only 4 tokens are generated to keep the test fast while
52 + still catching numerical divergence.
53 + """
54 + with monkeypatch.context() as m:
55 + m.setenv("TOKENIZERS_PARALLELISM", "true")
56 +
57 + with vllm_runner(
58 + model,
59 + max_model_len=MAX_MODEL_LEN,
60 + enforce_eager=True,
61 + ) as vllm_model:
62 + baseline_outputs = vllm_model.generate_greedy_logprobs(
63 + example_prompts, MAX_TOKENS, NUM_LOG_PROBS
64 + )
65 +
66 + with vllm_runner(
67 + model,
68 + max_model_len=MAX_MODEL_LEN,
69 + enforce_eager=True,
70 + quantization="mxfp8",
71 + ) as vllm_model:
72 + test_outputs = vllm_model.generate_greedy_logprobs(
73 + example_prompts, MAX_TOKENS, NUM_LOG_PROBS
74 + )
75 +
76 + check_logprobs_close(
77 + outputs_0_lst=baseline_outputs,
78 + outputs_1_lst=test_outputs,
79 + name_0="bf16",
80 + name_1="mxfp8",
81 + )
82 +
83 +
84 +@pytest.mark.skipif(
85 + not is_quant_method_supported("mxfp8"),
86 + reason="mxfp8 is not supported on this GPU type (requires sm_100+).",
87 +)
88 +@pytest.mark.quant_model
89 +@pytest.mark.parametrize("model", [DENSE_MODEL, MOE_MODEL], ids=["dense", "moe"])
90 +def test_mxfp8_generation(vllm_runner, model: str) -> None:
91 + """Smoke test: verify online MXFP8 model generates coherent text."""
92 + prompt = "1 2 3 4 5"
93 + with vllm_runner(
94 + model,
95 + enforce_eager=True,
96 + quantization="mxfp8",
97 + max_model_len=MAX_MODEL_LEN,
98 + ) as vllm_model:
99 + output = vllm_model.generate_greedy([prompt], max_tokens=5)
100 +
101 + generated = output[0][1]
102 + assert len(generated) > len(prompt), (
103 + f"MXFP8 model produced no new tokens. Output: {generated!r}"
104 + )
vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py
+90 -21
@@ -23,6 +23,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
23 kFp8Dynamic128Sym,
24 kFp8Static128BlockSym,
25 kFp8StaticTensorSym,
26 + kMxfp8Dynamic,
27 + kMxfp8Static,
28 )
29 from vllm.platforms import current_platform
30
@@ -67,11 +69,54 @@ class TrtLlmFp8ExpertsBase:
69 """Does not support non-gated MoE (i.e. Nanotron-3-Nano)."""
70 return True
71
72 + @staticmethod
73 + def _supports_quant_scheme(
74 + weight_key: QuantKey | None,
75 + activation_key: QuantKey | None,
76 + ) -> bool:
77 + """Supports Fp8 per-tensor, Fp8 block, and MXFP8."""
78 + SUPPORTED_W_A = [
79 + (kFp8Static128BlockSym, kFp8Dynamic128Sym),
80 + (kFp8StaticTensorSym, kFp8StaticTensorSym),
81 + (kMxfp8Static, kMxfp8Dynamic),
82 + ]
83 + return (weight_key, activation_key) in SUPPORTED_W_A
84 +
85 @staticmethod
86 def _supports_activation(activation: MoEActivation) -> bool:
87 """Supports only SiLU and RELU^2 non-gated activation."""
88 return activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]
89
90 + @staticmethod
91 + def _supports_routing_method(
92 + routing_method: RoutingMethodType,
93 + weight_key: QuantKey | None,
94 + activation_key: QuantKey | None,
95 + ) -> bool:
96 + """Monolithic kernels need to express router support."""
97 + # NOTE(dbari): TopK routing could also be enabled, but need to validate models
98 + # NOTE(dbari): Default is not implemented and should not be enabled until it is
99 + if (weight_key, activation_key) in [
100 + (kFp8Static128BlockSym, kFp8Dynamic128Sym),
101 + (kMxfp8Static, kMxfp8Dynamic),
102 + ]:
103 + # NOTE(rob): potentially allow others here. This is a conservative list.
104 + return routing_method in [
105 + RoutingMethodType.DeepSeekV3,
106 + RoutingMethodType.Renormalize,
107 + RoutingMethodType.RenormalizeNaive,
108 + ]
109 + elif (weight_key, activation_key) == (kFp8StaticTensorSym, kFp8StaticTensorSym):
110 + # NOTE(dbari): as above, potentially allow others here.
111 + return routing_method in [
112 + RoutingMethodType.DeepSeekV3,
113 + RoutingMethodType.Llama4,
114 + RoutingMethodType.Renormalize,
115 + RoutingMethodType.RenormalizeNaive,
116 + ]
117 + else:
118 + raise ValueError("Unsupported quantization scheme.")
119 +
120 @staticmethod
121 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
122 """Monolithic kernel so only use with naive DP/EP and TP."""
@@ -113,9 +158,10 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
158 weight_key: QuantKey | None,
159 activation_key: QuantKey | None,
160 ) -> bool:
116 - """Supports Fp8 block."""
161 + """Supports Fp8 block and MXFP8."""
162 SUPPORTED_W_A = [
163 (kFp8Static128BlockSym, kFp8Dynamic128Sym),
164 + (kMxfp8Static, kMxfp8Dynamic),
165 ]
166 return (weight_key, activation_key) in SUPPORTED_W_A
167
@@ -159,6 +205,7 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
205 apply_router_weight_on_input: bool,
206 ):
207 import flashinfer
208 + from flashinfer.fused_moe import Fp8QuantizationType
209
210 # Pack topk_ids and topk_weights into single tensor
211 # Format: (expert_id << 16) | (weight_bf16.view(int16))
@@ -175,6 +222,16 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
222
223 assert a1q_scale is not None
224
225 + is_mxfp8 = self.quant_config.block_shape == [1, 32]
226 + if is_mxfp8:
227 + fp8_quant_type = Fp8QuantizationType.MxFp8
228 + use_shuffled_weight = True
229 + hidden_states_scale = a1q_scale
230 + else:
231 + fp8_quant_type = Fp8QuantizationType.DeepSeekFp8
232 + use_shuffled_weight = False
233 + hidden_states_scale = a1q_scale.t().contiguous()
234 +
235 # `trtllm_fp8_block_scale_routed_moe` has a bug and does not write to the
236 # output tensor in-place so we need to manually copy the result to the
237 # output tensor
@@ -183,7 +240,7 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
240 topk_ids=packed_topk_ids,
241 routing_bias=None,
242 hidden_states=hidden_states,
186 - hidden_states_scale=a1q_scale.t().contiguous(), # type: ignore[union-attr]
243 + hidden_states_scale=hidden_states_scale,
244 gemm1_weights=w1,
245 gemm1_weights_scale=self.quant_config.w1_scale,
246 gemm2_weights=w2,
@@ -197,8 +254,9 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
254 local_num_experts=self.local_num_experts,
255 routed_scaling_factor=None,
256 routing_method_type=1,
200 - use_shuffled_weight=False,
257 + use_shuffled_weight=use_shuffled_weight,
258 weight_layout=0,
259 + fp8_quantization_type=fp8_quant_type,
260 # output=output,
261 )
262 output.copy_(result)
@@ -240,10 +298,11 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
298 weight_key: QuantKey | None,
299 activation_key: QuantKey | None,
300 ) -> bool:
243 - """Supports Fp8 per-tensor and Fp8 block."""
301 + """Supports Fp8 per-tensor, Fp8 block, and MXFP8."""
302 SUPPORTED_W_A = [
303 (kFp8Static128BlockSym, kFp8Dynamic128Sym),
304 (kFp8StaticTensorSym, kFp8StaticTensorSym),
305 + (kMxfp8Static, kMxfp8Dynamic),
306 ]
307 return (weight_key, activation_key) in SUPPORTED_W_A
308
@@ -256,7 +315,10 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
315 """Monolithic kernels need to express router support."""
316 # NOTE(dbari): TopK routing could also be enabled, but need to validate models
317 # NOTE(dbari): Default is not implemented and should not be enabled until it is
259 - if (weight_key, activation_key) == (kFp8Static128BlockSym, kFp8Dynamic128Sym):
318 + if (weight_key, activation_key) in [
319 + (kFp8Static128BlockSym, kFp8Dynamic128Sym),
320 + (kMxfp8Static, kMxfp8Dynamic),
321 + ]:
322 # NOTE(rob): potentially allow others here. This is a conservative list.
323 return routing_method in [
324 RoutingMethodType.DeepSeekV3,
@@ -274,7 +336,7 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
336 else:
337 raise ValueError("Unsupported quantization scheme.")
338
277 - def _apply_per_block(
339 + def _apply_block_scale(
340 self,
341 hidden_states: torch.Tensor,
342 w1: torch.Tensor,
@@ -291,32 +353,38 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
353 routed_scaling_factor: float | None = None,
354 topk_group: int | None = None,
355 ) -> torch.Tensor:
294 - # Delay import for non-CUDA.
356 import flashinfer
357 + from flashinfer.fused_moe import Fp8QuantizationType
358
359 assert not apply_router_weight_on_input
360 assert activation == MoEActivation.SILU
299 -
300 - if self.routing_method_type == RoutingMethodType.DeepSeekV3:
301 - router_logits = router_logits.to(torch.float32)
302 -
361 assert self.topk <= global_num_experts
362 assert self.topk <= 10
363 assert global_num_experts % 4 == 0
306 - assert self.quant_config.block_shape == [128, 128]
307 - # Routing kernel expects #experts <= #threads 512
364 + assert self.quant_config.block_shape in [[128, 128], [1, 32]]
365 + # Kernel expects #experts <= #threads 512
366 assert global_num_experts <= 512
309 -
310 - # Kernel requires transposed hidden state scales
367 # TODO: fuse into the quant kernel.
368 assert a1q_scale is not None
313 - a1q_scale_t = a1q_scale.t().contiguous()
369 +
370 + if self.routing_method_type == RoutingMethodType.DeepSeekV3:
371 + router_logits = router_logits.to(torch.float32)
372 +
373 + is_mxfp8 = self.quant_config.block_shape == [1, 32]
374 + if is_mxfp8:
375 + fp8_quant_type = Fp8QuantizationType.MxFp8
376 + use_shuffled_weight = True
377 + hidden_states_scale = a1q_scale
378 + else:
379 + fp8_quant_type = Fp8QuantizationType.DeepSeekFp8
380 + use_shuffled_weight = False
381 + hidden_states_scale = a1q_scale.t().contiguous()
382
383 return flashinfer.fused_moe.trtllm_fp8_block_scale_moe(
384 routing_logits=router_logits,
385 routing_bias=e_score_correction_bias,
386 hidden_states=hidden_states,
319 - hidden_states_scale=a1q_scale_t,
387 + hidden_states_scale=hidden_states_scale,
388 gemm1_weights=w1,
389 gemm1_weights_scale=self.quant_config.w1_scale,
390 gemm2_weights=w2,
@@ -330,7 +398,8 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
398 local_num_experts=self.local_num_experts,
399 routed_scaling_factor=routed_scaling_factor,
400 routing_method_type=self.routing_method_type,
333 - use_shuffled_weight=False,
401 + use_shuffled_weight=use_shuffled_weight,
402 + fp8_quantization_type=fp8_quant_type,
403 )
404
405 def _apply_per_tensor(
@@ -409,7 +478,7 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
478 topk_group: int | None = None,
479 ) -> torch.Tensor:
480 if self.quant_config.block_shape is not None:
412 - return self._apply_per_block(
481 + return self._apply_block_scale(
482 hidden_states,
483 w1,
484 w2,
@@ -441,6 +510,6 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
510 )
511 else:
512 raise NotImplementedError(
444 - "Only per-block and per-tensor quantization are supported in "
445 - f"{self.__class__.__name__}."
513 + "Only per-block, per-tensor, and MXFP8 quantization are "
514 + f"supported in {self.__class__.__name__}."
515 )
vllm/model_executor/layers/fused_moe/oracle/fp8.py
+16 -1
@@ -444,7 +444,7 @@ def convert_to_fp8_moe_kernel_format(
444 Fp8MoeBackend.FLASHINFER_CUTLASS,
445 Fp8MoeBackend.FLASHINFER_TRTLLM,
446 ]:
447 - w13, w2, w13_scale = prepare_fp8_moe_layer_for_fi(
447 + w13, w2, w13_scale, w2_scale = prepare_fp8_moe_layer_for_fi(
448 layer=layer,
449 w13=w13,
450 w2=w2,
@@ -512,6 +512,21 @@ def make_fp8_moe_quant_config(
512 g1_alphas=(w1_scale * a1_scale).squeeze(),
513 g2_alphas=(w2_scale * a2_scale).squeeze(),
514 )
515 + # MXFP8 uses "mxfp8" quant_dtype so the prepare step dispatches to
516 + # _mxfp8_e4m3_quantize rather than standard FP8 block quantization.
517 + # Non-swizzled layout is required since the TRTLLM kernel expects
518 + # scales in (num_tokens, hidden_dim // 32) format.
519 + if block_shape == [1, 32]:
520 + return FusedMoEQuantConfig.make(
521 + "mxfp8",
522 + w1_scale=w1_scale,
523 + w2_scale=w2_scale,
524 + a1_scale=a1_scale,
525 + a2_scale=a2_scale,
526 + block_shape=block_shape,
527 + is_nvfp4_scale_swizzled=False,
528 + )
529 +
530 # All other backends use normal config.
531 return fp8_w8a8_moe_quant_config(
532 w1_scale=w1_scale,
vllm/model_executor/layers/fused_moe/oracle/mxfp8.py
+66 -23
@@ -1,44 +1,87 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from enum import Enum
3
4 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
5 from vllm.logger import init_logger
6 from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
7 +from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
8 + Fp8MoeBackend,
9 + backend_to_kernel_cls,
10 +)
11 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
12 + kMxfp8Dynamic,
13 + kMxfp8Static,
14 +)
15
16 logger = init_logger(__name__)
17
18 +_SUPPORTED_BACKENDS: frozenset[Fp8MoeBackend] = frozenset(
19 + {
20 + Fp8MoeBackend.FLASHINFER_TRTLLM,
21 + }
22 +)
23
11 -class MxFp8MoeBackend(Enum):
12 - FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM"
24 +_BACKEND_NAME_MAP: dict[str, Fp8MoeBackend] = {
25 + "flashinfer_trtllm": Fp8MoeBackend.FLASHINFER_TRTLLM,
26 +}
27 +
28 +
29 +def _select_kernel_cls(
30 + backend: Fp8MoeBackend,
31 + config: FusedMoEConfig,
32 +) -> type[mk.FusedMoEExperts]:
33 + """Select the first supported expert class for the MXFP8 config."""
34 + activation_format = (
35 + mk.FusedMoEActivationFormat.BatchedExperts
36 + if config.moe_parallel_config.use_batched_activation_format
37 + else mk.FusedMoEActivationFormat.Standard
38 + )
39 + last_reason: str | None = None
40 + for cls in backend_to_kernel_cls(backend):
41 + supported, reason = cls.is_supported_config(
42 + cls,
43 + config,
44 + kMxfp8Static,
45 + kMxfp8Dynamic,
46 + activation_format,
47 + )
48 + if supported:
49 + return cls
50 + last_reason = reason
51 + raise ValueError(
52 + f"No supported MXFP8 expert class for {backend.value}: {last_reason}"
53 + )
54
55
56 def select_mxfp8_moe_backend(
57 config: FusedMoEConfig,
17 -) -> MxFp8MoeBackend:
58 +) -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts]]:
59 + """Select the MXFP8 MoE backend and the best expert class.
60 +
61 + Returns:
62 + A tuple of (fp8_backend, experts_cls).
63 + """
64 if config.is_lora_enabled:
65 raise NotImplementedError("LoRA is not supported for MXFP8 MoE.")
66
21 - AVAILABLE_BACKENDS = [
22 - MxFp8MoeBackend.FLASHINFER_TRTLLM,
23 - ]
24 -
67 runner_backend = config.moe_backend
68 if runner_backend != "auto":
27 - mapping = {
28 - "flashinfer_trtllm": MxFp8MoeBackend.FLASHINFER_TRTLLM,
29 - }
30 - if backend := mapping.get(runner_backend):
31 - logger.info_once(
32 - "Using '%s' MxFp8 MoE backend (user-requested).",
33 - backend.value,
69 + backend = _BACKEND_NAME_MAP.get(runner_backend)
70 + if backend is None:
71 + raise ValueError(
72 + f"moe_backend='{runner_backend}' is not supported for "
73 + f"MXFP8 MoE. Expected one of "
74 + f"{list(_BACKEND_NAME_MAP.keys())}."
75 )
35 - return backend
36 - raise ValueError(
37 - f"moe_backend='{runner_backend}' is not supported for MXFP8 MoE. "
38 - f"Expected one of {list(mapping.keys())}."
76 + logger.info_once(
77 + "Using '%s' MxFp8 MoE backend (user-requested).",
78 + backend.value,
79 )
80 + return backend, _select_kernel_cls(backend, config)
81 +
82 + # Auto-select: pick the first supported backend.
83 + for backend in _SUPPORTED_BACKENDS:
84 + logger.info_once("Using '%s' MxFp8 MoE backend.", backend.value)
85 + return backend, _select_kernel_cls(backend, config)
86
41 - # Auto-select: only one backend available for now.
42 - backend = AVAILABLE_BACKENDS[0]
43 - logger.info_once("Using '%s' MxFp8 MoE backend.", backend.value)
44 - return backend
87 + raise ValueError("No MXFP8 MoE backends available.")
vllm/model_executor/layers/fused_moe/utils.py
+1 -1
@@ -199,7 +199,7 @@ def _mxfp8_e4m3_quantize(
199 ) -> tuple[torch.Tensor, torch.Tensor]:
200 assert A_scale is None
201 assert not per_act_token_quant
202 - assert block_shape is None
202 + assert block_shape is None or block_shape == [1, 32]
203 return mxfp8_e4m3_quantize(A, is_sf_swizzled_layout)
204
205
vllm/model_executor/layers/quantization/__init__.py
+3
@@ -31,6 +31,7 @@ QuantizationMethods = Literal[
31 "torchao",
32 "inc",
33 "mxfp4",
34 + "mxfp8",
35 "petit_nvfp4",
36 "cpu_awq",
37 ]
@@ -129,6 +130,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
130 )
131 from .moe_wna16 import MoeWNA16Config
132 from .mxfp4 import Mxfp4Config
133 + from .mxfp8 import Mxfp8Config
134 from .petit import PetitNvFp4Config
135 from .ptpc_fp8 import PTPCFp8Config
136 from .torchao import TorchAOConfig
@@ -156,6 +158,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
158 "auto-round": INCConfig,
159 "inc": INCConfig,
160 "mxfp4": Mxfp4Config,
161 + "mxfp8": Mxfp8Config,
162 "petit_nvfp4": PetitNvFp4Config,
163 "cpu_awq": CPUAWQConfig,
164 }
vllm/model_executor/layers/quantization/modelopt.py
+4 -5
@@ -25,13 +25,13 @@ from vllm.model_executor.layers.fused_moe.layer import (
25 FusedMoeWeightScaleSupported,
26 )
27 from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
28 + Fp8MoeBackend,
29 convert_to_fp8_moe_kernel_format,
30 make_fp8_moe_kernel,
31 make_fp8_moe_quant_config,
32 select_fp8_moe_backend,
33 )
34 from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import (
34 - MxFp8MoeBackend,
35 select_mxfp8_moe_backend,
36 )
37 from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
@@ -1712,8 +1712,7 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
1712 self.quant_config = quant_config
1713 assert self.quant_config.is_checkpoint_mxfp8_serialized
1714
1715 - # Select MXFP8 MoE backend
1716 - self.mxfp8_backend = select_mxfp8_moe_backend(self.moe)
1715 + self.mxfp8_backend, _ = select_mxfp8_moe_backend(self.moe)
1716
1717 def create_weights(
1718 self,
@@ -1943,7 +1942,7 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
1942
1943 @property
1944 def is_monolithic(self) -> bool:
1946 - return self.mxfp8_backend == MxFp8MoeBackend.FLASHINFER_TRTLLM
1945 + return self.mxfp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM
1946
1947 def apply_monolithic(
1948 self,
@@ -1956,7 +1955,7 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
1955 Fp8QuantizationType,
1956 )
1957
1959 - assert self.mxfp8_backend == MxFp8MoeBackend.FLASHINFER_TRTLLM
1958 + assert self.mxfp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM
1959
1960 if layer.enable_eplb:
1961 raise NotImplementedError(
vllm/model_executor/layers/quantization/mxfp8.py new
+354
@@ -0,0 +1,354 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""Online MXFP8 (microscaling FP8, block-32) quantization config and methods."""
5 +
6 +from typing import Any
7 +
8 +import torch
9 +from torch.nn import Module
10 +
11 +from vllm.logger import init_logger
12 +from vllm.model_executor.layers.attention import Attention
13 +from vllm.model_executor.layers.fused_moe import (
14 + FusedMoE,
15 + FusedMoEMethodBase,
16 +)
17 +from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod
18 +from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import (
19 + select_mxfp8_moe_backend,
20 +)
21 +from vllm.model_executor.layers.linear import (
22 + LinearBase,
23 + UnquantizedLinearMethod,
24 +)
25 +from vllm.model_executor.layers.quantization import QuantizationMethods
26 +from vllm.model_executor.layers.quantization.base_config import (
27 + QuantizeMethodBase,
28 +)
29 +from vllm.model_executor.layers.quantization.fp8 import (
30 + Fp8Config,
31 + Fp8KVCacheMethod,
32 + Fp8OnlineLinearMethod,
33 + Fp8OnlineMoEMethod,
34 + _copy_missing_attrs,
35 +)
36 +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
37 + MXFP8_BLOCK_SIZE,
38 + Mxfp8LinearBackend,
39 + Mxfp8LinearOp,
40 + mxfp8_e4m3_quantize,
41 + swizzle_mxfp8_scale,
42 +)
43 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
44 + is_layer_skipped,
45 +)
46 +from vllm.model_executor.model_loader.weight_utils import (
47 + initialize_single_dummy_weight,
48 +)
49 +from vllm.model_executor.parameter import ModelWeightParameter
50 +from vllm.model_executor.utils import replace_parameter, set_weight_attrs
51 +from vllm.platforms import current_platform
52 +
53 +logger = init_logger(__name__)
54 +
55 +
56 +class Mxfp8Config(Fp8Config):
57 + """Config class for online MXFP8 MoE quantization."""
58 +
59 + def __init__(
60 + self,
61 + activation_scheme: str = "dynamic",
62 + ignored_layers: list[str] | None = None,
63 + ) -> None:
64 + if activation_scheme != "dynamic":
65 + raise ValueError("mxfp8 only supports dynamic activation scheme.")
66 + super().__init__(
67 + is_checkpoint_fp8_serialized=False,
68 + activation_scheme=activation_scheme,
69 + ignored_layers=ignored_layers,
70 + weight_block_size=None,
71 + )
72 +
73 + @classmethod
74 + def get_name(cls) -> QuantizationMethods:
75 + return "mxfp8"
76 +
77 + @classmethod
78 + def get_min_capability(cls) -> int:
79 + return 100
80 +
81 + @classmethod
82 + def from_config(cls, config: dict[str, Any]) -> "Mxfp8Config":
83 + activation_scheme = cls.get_from_keys_or(
84 + config, ["activation_scheme"], "dynamic"
85 + )
86 + ignored_layers = cls.get_from_keys_or(config, ["ignored_layers"], None)
87 + if not ignored_layers:
88 + ignored_layers = cls.get_from_keys_or(
89 + config, ["modules_to_not_convert"], None
90 + )
91 + return cls(
92 + activation_scheme=activation_scheme,
93 + ignored_layers=ignored_layers,
94 + )
95 +
96 + def get_quant_method(
97 + self, layer: torch.nn.Module, prefix: str
98 + ) -> "QuantizeMethodBase | None":
99 + if isinstance(layer, LinearBase):
100 + if is_layer_skipped(
101 + prefix=prefix,
102 + ignored_layers=self.ignored_layers,
103 + fused_mapping=self.packed_modules_mapping,
104 + skip_with_substr=True,
105 + ):
106 + return UnquantizedLinearMethod()
107 + return Mxfp8OnlineLinearMethod(self)
108 + elif isinstance(layer, FusedMoE):
109 + if is_layer_skipped(
110 + prefix=prefix,
111 + ignored_layers=self.ignored_layers,
112 + fused_mapping=self.packed_modules_mapping,
113 + skip_with_substr=True,
114 + ):
115 + return UnquantizedFusedMoEMethod(layer.moe_config)
116 + return Mxfp8OnlineMoEMethod(self, layer)
117 + elif isinstance(layer, Attention):
118 + return Fp8KVCacheMethod(self)
119 + return None
120 +
121 +
122 +class Mxfp8OnlineLinearMethod(Fp8OnlineLinearMethod):
123 + """Online MXFP8 linear method.
124 + Loads bf16/fp16 checkpoints and quantizes weights to MXFP8 (microscaling
125 + FP8 with block-32 scales) during weight loading.
126 +
127 + Args:
128 + quant_config: The MXFP8 quantization config.
129 + """
130 +
131 + uses_meta_device: bool = True
132 +
133 + def __init__(self, quant_config: "Mxfp8Config"):
134 + self.quant_config = quant_config
135 + self.out_dtype = torch.get_default_dtype()
136 + self.mxfp8_linear = Mxfp8LinearOp(self._select_backend())
137 + logger.info_once(
138 + "Using %s backend for MXFP8 GEMM", self.mxfp8_linear.backend.value
139 + )
140 +
141 + @staticmethod
142 + def _select_backend() -> Mxfp8LinearBackend:
143 + try:
144 + from vllm.utils import flashinfer as fi
145 +
146 + _ = fi.mm_mxfp8
147 + return Mxfp8LinearBackend.FLASHINFER_CUTLASS
148 + except Exception:
149 + logger.warning(
150 + "FlashInfer mm_mxfp8 not available, "
151 + "falling back to MXFP8 emulation backend."
152 + )
153 + return Mxfp8LinearBackend.EMULATION
154 +
155 + def create_weights(
156 + self,
157 + layer: torch.nn.Module,
158 + input_size_per_partition: int,
159 + output_partition_sizes: list[int],
160 + input_size: int,
161 + output_size: int,
162 + params_dtype: torch.dtype,
163 + **extra_weight_attrs,
164 + ):
165 + if input_size_per_partition % MXFP8_BLOCK_SIZE != 0:
166 + raise ValueError(
167 + f"MXFP8 requires input_size_per_partition "
168 + f"({input_size_per_partition}) to be divisible by "
169 + f"{MXFP8_BLOCK_SIZE}."
170 + )
171 +
172 + super().create_weights(
173 + layer,
174 + input_size_per_partition,
175 + output_partition_sizes,
176 + input_size,
177 + output_size,
178 + params_dtype,
179 + **extra_weight_attrs,
180 + )
181 +
182 + def process_weights_after_loading(self, layer: Module) -> None:
183 + if getattr(layer, "_already_called_process_weights_after_loading", False):
184 + return
185 +
186 + if layer.weight.device == torch.device("meta"):
187 + weight = ModelWeightParameter(
188 + data=torch.empty_like(layer.weight, device=layer._load_device),
189 + input_dim=1,
190 + output_dim=0,
191 + weight_loader=layer.weight.weight_loader,
192 + )
193 + _copy_missing_attrs(layer.weight, weight)
194 + layer.register_parameter("weight", weight)
195 + initialize_single_dummy_weight(layer.weight)
196 +
197 + weight_fp8, weight_scale = mxfp8_e4m3_quantize(layer.weight.contiguous())
198 +
199 + if self.mxfp8_linear.backend == Mxfp8LinearBackend.FLASHINFER_CUTLASS:
200 + N, K = layer.weight.shape[0], layer.weight.shape[1]
201 + weight_scale = swizzle_mxfp8_scale(weight_scale, N, K)
202 +
203 + layer.input_scale = None
204 + replace_parameter(layer, "weight", weight_fp8.data)
205 + replace_parameter(layer, "weight_scale", weight_scale.data)
206 +
207 + layer._already_called_process_weights_after_loading = True
208 +
209 + def apply(
210 + self,
211 + layer: torch.nn.Module,
212 + x: torch.Tensor,
213 + bias: torch.Tensor | None = None,
214 + ) -> torch.Tensor:
215 + return self.mxfp8_linear.apply(
216 + input=x,
217 + weight=layer.weight,
218 + weight_scale=layer.weight_scale,
219 + out_dtype=self.out_dtype,
220 + bias=bias,
221 + )
222 +
223 +
224 +class Mxfp8OnlineMoEMethod(Fp8OnlineMoEMethod):
225 + """MoE method for online MXFP8 (block) quantization."""
226 +
227 + uses_meta_device: bool = True
228 +
229 + def __init__(self, quant_config: Fp8Config, layer: torch.nn.Module):
230 + FusedMoEMethodBase.__init__(self, layer.moe_config)
231 + self.quant_config = quant_config
232 + assert not quant_config.is_checkpoint_fp8_serialized
233 + assert quant_config.activation_scheme == "dynamic"
234 +
235 + self.weight_block_size = [1, MXFP8_BLOCK_SIZE]
236 + self.block_quant = True
237 + self.weight_scale_name = "weight_scale"
238 +
239 + self.fp8_backend, self.experts_cls = select_mxfp8_moe_backend(config=self.moe)
240 +
241 + def create_weights(
242 + self,
243 + layer: Module,
244 + num_experts: int,
245 + hidden_size: int,
246 + intermediate_size_per_partition: int,
247 + params_dtype: torch.dtype,
248 + **extra_weight_attrs,
249 + ):
250 + if (
251 + hidden_size % MXFP8_BLOCK_SIZE != 0
252 + or intermediate_size_per_partition % MXFP8_BLOCK_SIZE != 0
253 + ):
254 + raise ValueError(
255 + "Online MXFP8 MoE requires hidden/intermediate sizes divisible "
256 + f"by {MXFP8_BLOCK_SIZE}."
257 + )
258 +
259 + super().create_weights(
260 + layer=layer,
261 + num_experts=num_experts,
262 + hidden_size=hidden_size,
263 + intermediate_size_per_partition=intermediate_size_per_partition,
264 + params_dtype=params_dtype,
265 + **extra_weight_attrs,
266 + )
267 +
268 + w13_weight_scale = torch.nn.Parameter(
269 + torch.zeros(
270 + num_experts,
271 + 2 * intermediate_size_per_partition,
272 + hidden_size // MXFP8_BLOCK_SIZE,
273 + dtype=torch.uint8,
274 + ),
275 + requires_grad=False,
276 + )
277 + w2_weight_scale = torch.nn.Parameter(
278 + torch.zeros(
279 + num_experts,
280 + hidden_size,
281 + intermediate_size_per_partition // MXFP8_BLOCK_SIZE,
282 + dtype=torch.uint8,
283 + ),
284 + requires_grad=False,
285 + )
286 + layer.register_parameter("w13_weight_scale", w13_weight_scale)
287 + layer.register_parameter("w2_weight_scale", w2_weight_scale)
288 + set_weight_attrs(w13_weight_scale, extra_weight_attrs)
289 + set_weight_attrs(w2_weight_scale, extra_weight_attrs)
290 + layer.weight_block_size = [1, MXFP8_BLOCK_SIZE]
291 +
292 + def _quantize_mxfp8_moe_weight(
293 + self, weight: torch.Tensor
294 + ) -> tuple[torch.Tensor, torch.Tensor]:
295 + """Batch quantization: bf16/fp16 weights -> MXFP8 (fp8 + uint8 scales)."""
296 + num_batches = weight.size(0)
297 + w_quant = []
298 + w_scales = []
299 + for i in range(num_batches):
300 + mx_fp8_quant, mx_fp8_scale = mxfp8_e4m3_quantize(
301 + weight[i], is_sf_swizzled_layout=False
302 + )
303 + w_quant.append(mx_fp8_quant)
304 + w_scales.append(mx_fp8_scale)
305 +
306 + return torch.stack(w_quant), torch.stack(w_scales)
307 +
308 + def process_weights_after_loading(self, layer: Module) -> None:
309 + if getattr(layer, "_already_called_process_weights_after_loading", False):
310 + return
311 +
312 + if layer.w13_weight.device == torch.device("meta"):
313 + w13_weight = torch.nn.Parameter(
314 + torch.empty_like(layer.w13_weight, device=layer._load_device),
315 + requires_grad=False,
316 + )
317 + set_weight_attrs(
318 + w13_weight, {"weight_loader": layer.w13_weight.weight_loader}
319 + )
320 + _copy_missing_attrs(layer.w13_weight, w13_weight)
321 + layer.register_parameter("w13_weight", w13_weight)
322 + initialize_single_dummy_weight(layer.w13_weight)
323 + if layer.w2_weight.device == torch.device("meta"):
324 + w2_weight = torch.nn.Parameter(
325 + torch.empty_like(layer.w2_weight, device=layer._load_device),
326 + requires_grad=False,
327 + )
328 + set_weight_attrs(
329 + w2_weight, {"weight_loader": layer.w2_weight.weight_loader}
330 + )
331 + _copy_missing_attrs(layer.w2_weight, w2_weight)
332 + layer.register_parameter("w2_weight", w2_weight)
333 + initialize_single_dummy_weight(layer.w2_weight)
334 +
335 + fp8_dtype = current_platform.fp8_dtype()
336 + w13 = torch.empty_like(layer.w13_weight, dtype=fp8_dtype)
337 + w2 = torch.empty_like(layer.w2_weight, dtype=fp8_dtype)
338 + w13_scale = layer.w13_weight_scale
339 + w2_scale = layer.w2_weight_scale
340 +
341 + w13, w13_scale = self._quantize_mxfp8_moe_weight(layer.w13_weight)
342 + w2, w2_scale = self._quantize_mxfp8_moe_weight(layer.w2_weight)
343 +
344 + self._setup_kernel(
345 + layer,
346 + w13,
347 + w2,
348 + w13_scale,
349 + w2_scale,
350 + layer.w13_input_scale,
351 + layer.w2_input_scale,
352 + )
353 +
354 + layer._already_called_process_weights_after_loading = True
vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+101 -3
@@ -305,6 +305,81 @@ def align_fp8_moe_weights_for_fi(
305 return padded_w13, padded_w2, padded_intermediate
306
307
308 +def _shuffle_mxfp8_moe_weights(
309 + w13: torch.Tensor,
310 + w2: torch.Tensor,
311 + w13_scale: torch.Tensor,
312 + w2_scale: torch.Tensor,
313 + is_gated: bool,
314 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
315 + """Preprocess MXFP8 weights and scales for the FlashInfer TRT-LLM kernel.
316 +
317 + Following flashinfer/tests/moe/test_trtllm_gen_fused_moe.py:
318 + 1. reorder_rows_for_gated_act_gemm (interleave gate/up rows)
319 + 2. shuffle_matrix_a (weight data layout shuffle)
320 + 3. shuffle_matrix_sf_a (scale factor layout shuffle)
321 + """
322 + from flashinfer import (
323 + reorder_rows_for_gated_act_gemm,
324 + shuffle_matrix_a,
325 + shuffle_matrix_sf_a,
326 + )
327 +
328 + epilogue_tile_m = 128
329 + num_experts = w13.shape[0]
330 + intermediate_size = w13.shape[1] // 2
331 + hidden_size = w13.shape[2]
332 +
333 + w13_interleaved: list[torch.Tensor] = []
334 + w13_scale_interleaved: list[torch.Tensor] = []
335 + for i in range(num_experts):
336 + if is_gated:
337 + w13_interleaved.append(
338 + reorder_rows_for_gated_act_gemm(
339 + w13[i].reshape(2 * intermediate_size, -1)
340 + )
341 + )
342 + w13_scale_interleaved.append(
343 + reorder_rows_for_gated_act_gemm(
344 + w13_scale[i].reshape(2 * intermediate_size, -1)
345 + )
346 + )
347 + else:
348 + w13_interleaved.append(w13[i])
349 + w13_scale_interleaved.append(w13_scale[i])
350 +
351 + w13_shuffled: list[torch.Tensor] = []
352 + w2_shuffled: list[torch.Tensor] = []
353 + w13_scale_shuffled: list[torch.Tensor] = []
354 + w2_scale_shuffled: list[torch.Tensor] = []
355 + for i in range(num_experts):
356 + w13_shuffled.append(
357 + shuffle_matrix_a(w13_interleaved[i].view(torch.uint8), epilogue_tile_m)
358 + )
359 + w2_shuffled.append(shuffle_matrix_a(w2[i].view(torch.uint8), epilogue_tile_m))
360 + w13_scale_shuffled.append(
361 + shuffle_matrix_sf_a(
362 + w13_scale_interleaved[i]
363 + .view(torch.uint8)
364 + .reshape(2 * intermediate_size, -1),
365 + epilogue_tile_m,
366 + )
367 + )
368 + w2_scale_shuffled.append(
369 + shuffle_matrix_sf_a(
370 + w2_scale[i].view(torch.uint8).reshape(hidden_size, -1),
371 + epilogue_tile_m,
372 + )
373 + )
374 +
375 + w13_out = torch.stack(w13_shuffled).view(torch.float8_e4m3fn)
376 + w2_out = torch.stack(w2_shuffled).view(torch.float8_e4m3fn)
377 + w13_scale_out = torch.stack(w13_scale_shuffled).reshape(w13_scale.shape)
378 + w2_scale_out = torch.stack(w2_scale_shuffled).reshape(w2_scale.shape)
379 +
380 + return w13_out, w2_out, w13_scale_out, w2_scale_out
381 +
382 +
383 def prepare_fp8_moe_layer_for_fi(
384 layer: torch.nn.Module,
385 w13: torch.Tensor,
@@ -314,7 +389,7 @@ def prepare_fp8_moe_layer_for_fi(
389 w2_scale: torch.Tensor,
390 w2_input_scale: torch.Tensor | None,
391 is_trtllm: bool = False,
317 -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
392 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
393 """
394 Convert Fp8 MoE weights to flashinfer kernel format
395
@@ -329,10 +404,33 @@ def prepare_fp8_moe_layer_for_fi(
404 block_quant = (
405 hasattr(layer, "weight_block_size") and layer.weight_block_size is not None
406 )
407 + is_mxfp8 = block_quant and w13_scale.dtype == torch.uint8
408 + is_gated = layer.activation.is_gated
409 +
410 + # MXFP8 TRT-LLM requires W31 swap + reorder + shuffle.
411 + if is_mxfp8 and is_trtllm:
412 + # FlashInfer TRT-LLM SwiGLU expects [up; gate] but vLLM stores
413 + # [gate; up]. Swap both weights and scales before interleaving.
414 + if layer.moe_config.is_act_and_mul:
415 + w13 = swap_w13_to_w31(w13)
416 + # Scales may be 2D [E, flat] from _quantize_mxfp8_moe_weight;
417 + # reshape to 3D so swap_w13_to_w31 can flip the two halves,
418 + # then flatten back.
419 + if w13_scale.ndim == 2:
420 + num_rows = w13.shape[1] # 2 * intermediate_size
421 + w13_scale = w13_scale.reshape(w13_scale.shape[0], num_rows, -1)
422 + w13_scale = swap_w13_to_w31(w13_scale)
423 + w13_scale = w13_scale.reshape(w13_scale.shape[0], -1)
424 + else:
425 + w13_scale = swap_w13_to_w31(w13_scale)
426 +
427 + w13, w2, w13_scale, w2_scale = _shuffle_mxfp8_moe_weights(
428 + w13, w2, w13_scale, w2_scale, is_gated
429 + )
430 + return w13, w2, w13_scale, w2_scale
431
432 # Some FI MoE kernels require internal alignment of 16
433 # for the gate-up proj. Pad the weights to respect this.
335 - is_gated = layer.activation.is_gated
434 if not block_quant:
435 min_alignment = 16 if is_gated else 128
436 w13, w2, new_intermediate = align_fp8_moe_weights_for_fi(
@@ -369,4 +467,4 @@ def prepare_fp8_moe_layer_for_fi(
467 w13_scale.clamp_(min=_FI_CUTLASS_MIN_BLOCK_SCALE)
468 w2_scale.clamp_(min=_FI_CUTLASS_MIN_BLOCK_SCALE)
469
372 - return w13, w2, w13_scale
470 + return w13, w2, w13_scale, w2_scale
vllm/model_executor/layers/quantization/utils/quant_utils.py
+6
@@ -149,6 +149,12 @@ kFp8Dynamic128Sym = QuantKey(FP8_DTYPE, kDynamic128Scale, symmetric=True)
149 kStatic128BlockScale = ScaleDesc(torch.float32, True, GroupShape(128, 128))
150 kFp8Static128BlockSym = QuantKey(FP8_DTYPE, kStatic128BlockScale, symmetric=True)
151
152 +kMxfp8StaticScale = ScaleDesc(torch.uint8, True, GroupShape(1, 32))
153 +kMxfp8Static = QuantKey(FP8_DTYPE, kMxfp8StaticScale, symmetric=True)
154 +
155 +kMxfp8DynamicScale = ScaleDesc(torch.uint8, False, GroupShape(1, 32))
156 +kMxfp8Dynamic = QuantKey(FP8_DTYPE, kMxfp8DynamicScale, symmetric=True)
157 +
158 kDynamic64Scale = ScaleDesc(torch.float32, False, GroupShape(1, 64))
159 kFp8Dynamic64Sym = QuantKey(FP8_DTYPE, kDynamic64Scale, symmetric=True)
160