@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 739e5945dc

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 5, 06:00 UTC

[Quantization] [Refactor] Create special "GptOssMxfp4MoeMethod" (#39604)

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>

Yongye Zhu committed Apr 13, 2026 at 12:53 UTC 739e5945dc4b3ba30d84bdf6e637657abd4136b8

18 files changed +129 -34

docs/design/moe_kernel_features.md

+1 -1

@@ -59,7 +59,7 @@ Modular kernels are supported by the following `FusedMoEMethodBase` classes.

59 - [`Fp8MoEMethod`][vllm.model_executor.layers.quantization.fp8.Fp8MoEMethod]

60 - [`CompressedTensorsW4A4Nvfp4MoEMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w4a4_nvfp4.CompressedTensorsW4A4Nvfp4MoEMethod]

61 - [`CompressedTensorsW8A8Fp8MoEMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w8a8_fp8.CompressedTensorsW8A8Fp8MoEMethod]

62 -- [`Mxfp4MoEMethod`][vllm.model_executor.layers.quantization.mxfp4.Mxfp4MoEMethod]

62 +- [`GptOssMxfp4MoEMethod`][vllm.model_executor.layers.quantization.mxfp4.GptOssMxfp4MoEMethod]

63 - [`UnquantizedFusedMoEMethod`][vllm.model_executor.layers.fused_moe.layer.UnquantizedFusedMoEMethod]

65 ## Fused Experts Kernels

vllm/config/model.py

+2 -1

@@ -951,6 +951,7 @@ class ModelConfig:

951 # Ensure heavy backends are probed last to avoid unnecessary

952 # imports during override detection (e.g., MXFP4 imports Triton)

953 "mxfp4",

954 + "gpt_oss_mxfp4",

955 "cpu_awq",

956 "gguf",

957 ]

@@ -966,7 +967,7 @@ class ModelConfig:

967 for name in quantization_methods:

968 method = me_quant.get_quantization_config(name)

969 quantization_override = method.override_quantization_method(

969 - quant_cfg, self.quantization

970 + quant_cfg, self.quantization, hf_config=self.hf_config

971 )

972 if quantization_override is not None:

973 # Raise error if the override is not custom (custom would

vllm/model_executor/layers/fused_moe/layer.py

+1 -1

@@ -1063,7 +1063,7 @@ class FusedMoE(CustomOp):

1063 expert_id: int,

1064 return_success: bool = False,

1065 ) -> bool | None:

1066 - if self.quant_config and self.quant_config.get_name() == "mxfp4":

1066 + if self.quant_config and self.quant_config.get_name() == "gpt_oss_mxfp4":

1067 # (FIXME) for gpt-oss all experts are combined

1068 if "bias" in weight_name:

1069 dim1 = loaded_weight.shape[1]

vllm/model_executor/layers/fused_moe/oracle/mxfp4.py

+6 -3

@@ -194,7 +194,7 @@ def _backend_activation_key(backend: Mxfp4MoeBackend) -> QuantKey | None:

194 return None

195

196

197 -def select_mxfp4_moe_backend(

197 +def select_gpt_oss_mxfp4_moe_backend(

198 config: FusedMoEConfig,

199 ) -> tuple[Mxfp4MoeBackend, type[mk.FusedMoEExperts] | None]:

200 """

@@ -400,7 +400,7 @@ def mxfp4_round_up_hidden_size_and_intermediate_size(

400 return hidden_size, intermediate_size

401

402

403 -def convert_to_mxfp4_moe_kernel_format(

403 +def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(

404 mxfp4_backend: Mxfp4MoeBackend,

405 layer: torch.nn.Module,

406 w13_weight: torch.Tensor,

@@ -426,7 +426,10 @@ def convert_to_mxfp4_moe_kernel_format(

426

427 sf_block_size = 32 # mxfp4 block size

428

429 - if mxfp4_backend in (Mxfp4MoeBackend.MARLIN, Mxfp4MoeBackend.BATCHED_MARLIN):

429 + if mxfp4_backend in (

430 + Mxfp4MoeBackend.MARLIN,

431 + Mxfp4MoeBackend.BATCHED_MARLIN,

432 + ):

433 from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (

434 prepare_moe_mxfp4_layer_for_marlin,

435 )

vllm/model_executor/layers/quantization/__init__.py

+3 -1

@@ -30,6 +30,7 @@ QuantizationMethods = Literal[

30 "torchao",

31 "inc",

32 "mxfp4",

33 + "gpt_oss_mxfp4",

34 "mxfp8",

35 "cpu_awq",

36 "online",

@@ -133,7 +134,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:

134 ModelOptNvFp4Config,

135 )

136 from .moe_wna16 import MoeWNA16Config

136 - from .mxfp4 import Mxfp4Config

137 + from .mxfp4 import GptOssMxfp4Config, Mxfp4Config

138 from .mxfp8 import Mxfp8Config

139 from .online.base import OnlineQuantizationConfig

140 from .torchao import TorchAOConfig

@@ -160,6 +161,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:

161 "auto-round": INCConfig,

162 "inc": INCConfig,

163 "mxfp4": Mxfp4Config,

164 + "gpt_oss_mxfp4": GptOssMxfp4Config,

165 "mxfp8": Mxfp8Config,

166 "cpu_awq": CPUAWQConfig,

167 "online": OnlineQuantizationConfig,

vllm/model_executor/layers/quantization/awq_marlin.py

+1 -1

@@ -232,7 +232,7 @@ class AWQMarlinConfig(QuantizationConfig):

232

233 @classmethod

234 def override_quantization_method(

235 - cls, hf_quant_cfg, user_quant

235 + cls, hf_quant_cfg, user_quant, hf_config=None

236 ) -> "QuantizationMethods | None":

237 # Skip override to marlin kernels, as they are not

238 # batch invariant

vllm/model_executor/layers/quantization/base_config.py

+11 -2

@@ -110,13 +110,22 @@ class QuantizationConfig(ABC):

110

111 @classmethod

112 def override_quantization_method(

113 - cls, hf_quant_cfg, user_quant

113 + cls,

114 + hf_quant_cfg: dict[str, Any],

115 + user_quant: str | None,

116 + hf_config: Any = None,

117 ) -> QuantizationMethods | None:

118 """

119 Detects if this quantization method can support a given checkpoint

120 format by overriding the user specified quantization method --

121 this method should only be overwritten by subclasses in exceptional

119 - circumstances

122 + circumstances.

123 +

124 + Args:

125 + hf_quant_cfg: The checkpoint's quantization config dict.

126 + user_quant: The user-specified quantization method string.

127 + hf_config: The HuggingFace model config object (e.g. for

128 + model_type checks). May be None if not available.

129 """

130 return None

131

vllm/model_executor/layers/quantization/cpu_wna16.py

+1 -1

@@ -104,7 +104,7 @@ class CPUAWQConfig(QuantizationConfig):

104

105 @classmethod

106 def override_quantization_method(

107 - cls, hf_quant_cfg, user_quant

107 + cls, hf_quant_cfg, user_quant, hf_config=None

108 ) -> "QuantizationMethods | None":

109 quant_method = hf_quant_cfg.get("quant_method", "").lower()

110 if current_platform.is_cpu() and (quant_method == "awq"):

vllm/model_executor/layers/quantization/gguf.py

+1 -1

@@ -84,7 +84,7 @@ class GGUFConfig(QuantizationConfig):

85 @classmethod

86 def override_quantization_method(

87 - cls, hf_quant_cfg: dict[str, Any], user_quant: str | None

87 + cls, hf_quant_cfg: dict[str, Any], user_quant: str | None, hf_config=None

88 ) -> "QuantizationMethods | None":

89 # When user explicitly specifies --quantization gguf, override

90 # whatever quantization method is in the HF model config (e.g. fp8).

vllm/model_executor/layers/quantization/gptq_marlin.py

+1 -1

@@ -214,7 +214,7 @@ class GPTQMarlinConfig(QuantizationConfig):

214

215 @classmethod

216 def override_quantization_method(

217 - cls, hf_quant_cfg, user_quant

217 + cls, hf_quant_cfg, user_quant, hf_config=None

218 ) -> QuantizationMethods | None:

219 can_convert = cls.is_gptq_marlin_compatible(hf_quant_cfg)

220

vllm/model_executor/layers/quantization/inc.py

+1 -1

@@ -453,7 +453,7 @@ class INCConfig(QuantizationConfig):

453

454 @classmethod

455 def override_quantization_method(

456 - cls, hf_quant_cfg, user_quant

456 + cls, hf_quant_cfg, user_quant, hf_config=None

457 ) -> "QuantizationMethods | None":

458 """Override the `auto-round` method to `inc`."""

459 is_auto_round_format = hf_quant_cfg.get("quant_method", None) == "auto-round"

vllm/model_executor/layers/quantization/modelopt.py

+4 -4

@@ -406,7 +406,7 @@ class ModelOptFp8Config(ModelOptQuantConfigBase):

406

407 @classmethod

408 def override_quantization_method(

409 - cls, hf_quant_cfg, user_quant

409 + cls, hf_quant_cfg, user_quant, hf_config=None

410 ) -> QuantizationMethods | None:

411 algo = cls._extract_modelopt_quant_algo(hf_quant_cfg)

412 if algo is not None and algo == "FP8":

@@ -1028,7 +1028,7 @@ class ModelOptNvFp4Config(ModelOptQuantConfigBase):

1028

1029 @classmethod

1030 def override_quantization_method(

1031 - cls, hf_quant_cfg, user_quant

1031 + cls, hf_quant_cfg, user_quant, hf_config=None

1032 ) -> QuantizationMethods | None:

1033 algo = cls._extract_modelopt_quant_algo(hf_quant_cfg)

1034 if algo is not None and ("NVFP4" in algo or "FP4" in algo):

@@ -1525,7 +1525,7 @@ class ModelOptMxFp8Config(ModelOptQuantConfigBase):

1525

1526 @classmethod

1527 def override_quantization_method(

1528 - cls, hf_quant_cfg, user_quant

1528 + cls, hf_quant_cfg, user_quant, hf_config=None

1529 ) -> QuantizationMethods | None:

1530 algo = cls._extract_modelopt_quant_algo(hf_quant_cfg)

1531 if algo is not None and "MXFP8" in algo:

@@ -2052,7 +2052,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):

2052

2053 @classmethod

2054 def override_quantization_method(

2055 - cls, hf_quant_cfg, user_quant

2055 + cls, hf_quant_cfg, user_quant, hf_config=None

2056 ) -> QuantizationMethods | None:

2057 algo = cls._extract_modelopt_quant_algo(hf_quant_cfg)

2058 if algo is not None and algo == "MIXED_PRECISION":

vllm/model_executor/layers/quantization/moe_wna16.py

+1 -1

@@ -130,7 +130,7 @@ class MoeWNA16Config(QuantizationConfig):

130

131 @classmethod

132 def override_quantization_method(

133 - cls, hf_quant_cfg, user_quant

133 + cls, hf_quant_cfg, user_quant, hf_config=None

134 ) -> QuantizationMethods | None:

135 can_convert = cls.is_moe_wna16_compatible(hf_quant_cfg)

136 if can_convert and user_quant == "moe_wna16":

vllm/model_executor/layers/quantization/mxfp4.py

+48 -7

@@ -19,11 +19,11 @@ from vllm.model_executor.layers.fused_moe.config import (

19 from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (

20 TRITON_BACKENDS,

21 Mxfp4MoeBackend,

22 - convert_to_mxfp4_moe_kernel_format,

22 + convert_gpt_oss_weight_to_mxfp4_moe_kernel_format,

23 make_mxfp4_moe_kernel,

24 make_mxfp4_moe_quant_config,

25 mxfp4_round_up_hidden_size_and_intermediate_size,

26 - select_mxfp4_moe_backend,

26 + select_gpt_oss_mxfp4_moe_backend,

27 )

28 from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod

29 from vllm.model_executor.layers.quantization import QuantizationMethods

@@ -38,6 +38,12 @@ logger = init_logger(__name__)

40 class Mxfp4Config(QuantizationConfig):

41 + """Canonical base config for MXFP4 quantization.

42 +

43 + Subclasses override get_name() and override_quantization_method() to

44 + register themselves as the handler for a specific checkpoint format.

45 + """

46 +

47 def __init__(self, ignored_layers: list[str] | None = None):

48 super().__init__()

49 self.ignored_layers = ignored_layers

@@ -62,6 +68,8 @@ class Mxfp4Config(QuantizationConfig):

68 def get_config_filenames(cls) -> list[str]:

69 return []

71 + # TODO (zyongye) This is only temporaty fallback.

72 + # We should have `Mxfp4MoEMethod` after this migration is complete.

73 def get_quant_method(

74 self, layer: torch.nn.Module, prefix: str

75 ) -> "QuantizeMethodBase | None":

@@ -79,7 +87,7 @@ class Mxfp4Config(QuantizationConfig):

87 )

88 return UnquantizedLinearMethod()

89 elif isinstance(layer, FusedMoE):

82 - return Mxfp4MoEMethod(layer.moe_config)

90 + return GptOssMxfp4MoEMethod(layer.moe_config)

91 elif isinstance(layer, Attention):

92 logger.debug_once(

93 "MXFP4 attention layer is not implemented. "

@@ -93,13 +101,46 @@ class Mxfp4Config(QuantizationConfig):

101 return True

102

103

96 -class Mxfp4MoEMethod(FusedMoEMethodBase):

104 +class GptOssMxfp4Config(Mxfp4Config):

105 + """MXFP4 config for GPT-OSS checkpoints.

106 +

107 + Checkpoints carry ``"quant_method": "mxfp4"`` in their JSON config.

108 + override_quantization_method() maps that to the canonical internal name

109 + so that the rest of the loading path uses "gpt_oss_mxfp4" consistently.

110 + """

111 +

112 + @classmethod

113 + def get_name(cls) -> QuantizationMethods:

114 + return "gpt_oss_mxfp4"

115 +

116 + @classmethod

117 + def override_quantization_method(

118 + cls, hf_quant_cfg, user_quant, hf_config=None

119 + ) -> QuantizationMethods | None:

120 + # Match both "mxfp4" (original checkpoint value) and "gpt_oss_mxfp4"

121 + # (already normalized by verify_and_update_model_config) so that

122 + # explicit --quantization mxfp4 from the user doesn't cause a mismatch.

123 + if not (

124 + isinstance(hf_quant_cfg, dict)

125 + and hf_quant_cfg.get("quant_method") in ("mxfp4", "gpt_oss_mxfp4")

126 + ):

127 + return None

128 + # Require explicit confirmation that this is a GPT-OSS model.

129 + # Do NOT fall back to returning the override when hf_config is None,

130 + # as that would silently claim all mxfp4 checkpoints.

131 + model_type = getattr(hf_config, "model_type", None)

132 + if model_type != "gpt_oss":

133 + return None

134 + return "gpt_oss_mxfp4"

135 +

136 +

137 +class GptOssMxfp4MoEMethod(FusedMoEMethodBase):

138 """MXFP4 MoE quantization method."""

139

140 def __init__(self, moe: FusedMoEConfig):

141 super().__init__(moe)

101 - self.weight_dtype = "mxfp4"

102 - self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend(moe)

142 + self.weight_dtype = "gpt_oss_mxfp4"

143 + self.mxfp4_backend, self.experts_cls = select_gpt_oss_mxfp4_moe_backend(moe)

144

145 self.max_capture_size = (

146 get_current_vllm_config().compilation_config.max_cudagraph_capture_size

@@ -281,7 +322,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

322

323 # Convert weights to kernel format

324 w13, w2, w13_scale, w2_scale, w13_bias, w2_bias = (

284 - convert_to_mxfp4_moe_kernel_format(

325 + convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(

326 mxfp4_backend=self.mxfp4_backend,

327 layer=layer,

328 w13_weight=w13,

vllm/model_executor/layers/quantization/quark/quark_moe.py

+4 -4

@@ -30,11 +30,11 @@ from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_m

30 from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (

31 TRITON_BACKENDS,

32 Mxfp4MoeBackend,

33 - convert_to_mxfp4_moe_kernel_format,

33 + convert_gpt_oss_weight_to_mxfp4_moe_kernel_format,

34 make_mxfp4_moe_kernel,

35 make_mxfp4_moe_quant_config,

36 mxfp4_round_up_hidden_size_and_intermediate_size,

37 - select_mxfp4_moe_backend,

37 + select_gpt_oss_mxfp4_moe_backend,

38 )

39 from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (

40 prepare_fp8_moe_layer_for_marlin,

@@ -995,7 +995,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):

995 self.w2_precision_config = None

996

997 if self.ocp_mx_scheme == "w_mxfp4":

998 - self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend(moe)

998 + self.mxfp4_backend, self.experts_cls = select_gpt_oss_mxfp4_moe_backend(moe)

999 elif self.ocp_mx_scheme.startswith("w_mxfp4"):

1000 # TODO(bowenbao): refactor and introduce backends for other OCP MX schemes.

1001 self.mxfp4_backend = Mxfp4MoeBackend.NONE

@@ -1300,7 +1300,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):

1300

1301 # Convert weights to kernel format

1302 w13, w2, w13_scale, w2_scale, w13_bias, w2_bias = (

1303 - convert_to_mxfp4_moe_kernel_format(

1303 + convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(

1304 mxfp4_backend=self.mxfp4_backend,

1305 layer=layer,

1306 w13_weight=w13,

vllm/model_executor/models/config.py

+17

@@ -108,6 +108,23 @@ class Gemma4Config(VerifyAndUpdateConfig):

108

109

110 class GptOssForCausalLMConfig(VerifyAndUpdateConfig):

111 + @staticmethod

112 + def verify_and_update_model_config(model_config: "ModelConfig") -> None:

113 + quant_config = getattr(model_config.hf_config, "quantization_config", None)

114 + if quant_config is not None and quant_config.get("quant_method") == "mxfp4":

115 + model_config.hf_config.quantization_config["quant_method"] = "gpt_oss_mxfp4"

116 +

117 + hf_text_quant_config = getattr(

118 + model_config.hf_text_config, "quantization_config", None

119 + )

120 + if (

121 + hf_text_quant_config is not None

122 + and hf_text_quant_config.get("quant_method") == "mxfp4"

123 + ):

124 + model_config.hf_text_config.quantization_config["quant_method"] = (

125 + "gpt_oss_mxfp4"

126 + )

127 +

128 @staticmethod

129 def verify_and_update_config(vllm_config: "VllmConfig") -> None:

130 structured_outputs_config = vllm_config.structured_outputs_config

vllm/model_executor/models/gpt_oss.py

+25 -3

@@ -560,6 +560,14 @@ class GptOssModel(nn.Module, EagleModelMixin):

560 pcp_rank=get_pcp_group().rank_in_group,

561 )

562

563 + def _is_mxfp4(weight_dtype: str | None) -> bool:

564 + """Return True for any MXFP4 weight-dtype variant.

565 +

566 + Covers "gpt_oss_mxfp4" (GptOssMxfp4MoEMethod) and "mxfp4"

567 + (QuarkMoEMethod with fp4 weights) and any future variants.

568 + """

569 + return weight_dtype is not None and "mxfp4" in weight_dtype

570 +

571 def _get_moe_weight_dtype(layer_id: int = 0) -> str | None:

572 """Helper function to get MoE quantization weight dtype.

573

@@ -578,7 +586,7 @@ class GptOssModel(nn.Module, EagleModelMixin):

586

587 moe_weight_dtype = _get_moe_weight_dtype(layer_id=0)

588

581 - if moe_weight_dtype == "mxfp4":

589 + if _is_mxfp4(moe_weight_dtype):

590 # MXFP4 requires OCP_MX_BLOCK_SIZE alignment

591 intermediate_size_block = intermediate_size // OCP_MX_BLOCK_SIZE

592 per_rank_intermediate_size_block = cdiv(intermediate_size_block, tp_size)

@@ -682,7 +690,7 @@ class GptOssModel(nn.Module, EagleModelMixin):

690 continue

691

692 # Unified handler for mxfp4 weights and scales

685 - elif moe_quant_method == "mxfp4" and any(

693 + elif _is_mxfp4(moe_quant_method) and any(

694 name.endswith(suffix)

695 for suffix in [

696 ".w13_weight_scale",

@@ -1116,8 +1124,22 @@ class GptOssModel(nn.Module, EagleModelMixin):

1124 if hasattr(self.config, "quantization_config")

1125 else None

1126 )

1119 -

1127 + # Normalize the checkpoint's quant_method to the internal name.

1128 + # Note: there are three places where "mxfp4" -> "gpt_oss_mxfp4"

1129 + # normalization occurs, each serving a different data path:

1130 + # 1. GptOssMxfp4Config.override_quantization_method() — sets

1131 + # ModelConfig.quantization (used to select the QuantizationConfig

1132 + # class at model init time), reading from model_arch_config which

1133 + # is a snapshot taken before verify_and_update_model_config runs.

1134 + # 2. GptOssForCausalLMConfig.verify_and_update_model_config() —

1135 + # patches hf_config.quantization_config in-place (a separate copy

1136 + # of the dict from model_arch_config) for later hf_config lookups.

1137 + # 3. Here — reads directly from self.config (the raw HF config) which

1138 + # may still carry the original "mxfp4" string from the checkpoint.

1139 if quant_method == "mxfp4":

1140 + quant_method = "gpt_oss_mxfp4"

1141 +

1142 + if quant_method == "gpt_oss_mxfp4":

1143 return self._load_weights_mxfp4(

1144 ep_rank_end,

1145 ep_rank_start,

vllm/model_executor/models/transformers/base.py

+1 -1

@@ -150,7 +150,7 @@ class Base(

150 if self.quant_config:

151 quant_method_name = self.quant_config.get_name()

152 # Check for unsupported quantization methods.

153 - if quant_method_name == "mxfp4":

153 + if quant_method_name in ("mxfp4", "gpt_oss_mxfp4"):

154 raise NotImplementedError(

155 "Transformers modeling backend does "

156 "not support MXFP4 quantization yet."