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

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 2, 19:30 UTC

[MoE Refactor] Migrate ModelOptMxFp8FusedMoE to oracle (#42768)

Signed-off-by: Bill Nell <bnell@redhat.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>

bnellnm committed May 26, 2026 at 11:14 UTC b226ddacfda921b8eef7b75794c33d49ae51186a

8 files changed +114 -130

vllm/model_executor/layers/fused_moe/config.py

+19 -4

@@ -888,20 +888,33 @@ def int4_w4a16_moe_quant_config(

888 def fp8_w8a16_moe_quant_config(

889 w1_scale: torch.Tensor,

890 w2_scale: torch.Tensor,

891 + w1_bias: torch.Tensor | None = None,

892 + w2_bias: torch.Tensor | None = None,

893 block_shape: list[int] | None = None,

894 ) -> FusedMoEQuantConfig:

895 """

896 Construct a quant config for 16-bit float activations and fp8 weights.

897 """

898 group_shape = GroupShape(*block_shape) if block_shape is not None else None

899 + fp8_dtype = current_platform.fp8_dtype()

900 return FusedMoEQuantConfig(

901 _a1=FusedMoEQuantDesc(),

902 _a2=FusedMoEQuantDesc(),

903 _w1=FusedMoEQuantDesc(

901 - current_platform.fp8_dtype(), group_shape, w1_scale, None, None

904 + fp8_dtype,

905 + group_shape,

906 + w1_scale,

907 + None,

908 + None,

909 + w1_bias,

910 ),

911 _w2=FusedMoEQuantDesc(

904 - current_platform.fp8_dtype(), group_shape, w2_scale, None, None

912 + fp8_dtype,

913 + group_shape,

914 + w2_scale,

915 + None,

916 + None,

917 + w2_bias,

918 ),

919 )

920

@@ -911,6 +924,8 @@ def int8_w8a16_moe_quant_config(

924 w2_scale: torch.Tensor,

925 w1_zp: torch.Tensor | None,

926 w2_zp: torch.Tensor | None,

927 + w1_bias: torch.Tensor | None = None,

928 + w2_bias: torch.Tensor | None = None,

929 block_shape: list[int] | None = None,

930 ) -> FusedMoEQuantConfig:

931 """

@@ -920,8 +935,8 @@ def int8_w8a16_moe_quant_config(

935 return FusedMoEQuantConfig(

936 _a1=FusedMoEQuantDesc(shape=group_shape),

937 _a2=FusedMoEQuantDesc(shape=group_shape),

923 - _w1=FusedMoEQuantDesc(torch.int8, group_shape, w1_scale, None, w1_zp),

924 - _w2=FusedMoEQuantDesc(torch.int8, group_shape, w2_scale, None, w2_zp),

938 + _w1=FusedMoEQuantDesc(torch.int8, group_shape, w1_scale, None, w1_zp, w1_bias),

939 + _w2=FusedMoEQuantDesc(torch.int8, group_shape, w2_scale, None, w2_zp, w2_bias),

940 )

941

942

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

+13 -10

@@ -421,6 +421,7 @@ def select_fp8_moe_backend(

421

422 def convert_to_fp8_moe_kernel_format(

423 fp8_backend: Fp8MoeBackend,

424 + # TODO(bnell): replace layer with weight_block_size

425 layer: torch.nn.Module,

426 w13: torch.Tensor,

427 w2: torch.Tensor,

@@ -508,6 +509,8 @@ def make_fp8_moe_quant_config(

509 w2_scale: torch.Tensor,

510 a1_scale: torch.Tensor | None,

511 a2_scale: torch.Tensor | None,

512 + w1_bias: torch.Tensor | None = None,

513 + w2_bias: torch.Tensor | None = None,

514 block_shape: list[int] | None = None,

515 per_act_token_quant: bool = False,

516 per_out_ch_quant: bool = False,

@@ -526,19 +529,13 @@ def make_fp8_moe_quant_config(

529 a method of the modular kernel itself.

530 """

531

529 - # MARLIN is mixed precision W8A16 config.

530 - if fp8_backend == Fp8MoeBackend.MARLIN:

531 - return fp8_w8a16_moe_quant_config(

532 - w1_scale=w1_scale,

533 - w2_scale=w2_scale,

534 - block_shape=block_shape,

535 - )

536 -

537 - # CPU is mixed precision W8A16 config.

538 - if fp8_backend == Fp8MoeBackend.CPU:

532 + # MARLIN and CPU are mixed precision W8A16 config.

533 + if fp8_backend == Fp8MoeBackend.MARLIN or fp8_backend == Fp8MoeBackend.CPU:

534 return fp8_w8a16_moe_quant_config(

535 w1_scale=w1_scale,

536 w2_scale=w2_scale,

537 + w1_bias=w1_bias,

538 + w2_bias=w2_bias,

539 block_shape=block_shape,

540 )

541

@@ -549,6 +546,8 @@ def make_fp8_moe_quant_config(

546 return fp8_w8a8_moe_quant_config(

547 w1_scale=w1_scale,

548 w2_scale=w2_scale,

549 + w1_bias=w1_bias,

550 + w2_bias=w2_bias,

551 a1_scale=a1_scale,

552 a2_scale=a2_scale,

553 a1_gscale=(1.0 / a1_scale),

@@ -566,6 +565,8 @@ def make_fp8_moe_quant_config(

565 "mxfp8",

566 w1_scale=w1_scale,

567 w2_scale=w2_scale,

568 + w1_bias=w1_bias,

569 + w2_bias=w2_bias,

570 a1_scale=a1_scale,

571 a2_scale=a2_scale,

572 block_shape=block_shape,

@@ -577,6 +578,8 @@ def make_fp8_moe_quant_config(

578 return fp8_w8a8_moe_quant_config(

579 w1_scale=w1_scale,

580 w2_scale=w2_scale,

581 + w1_bias=w1_bias,

582 + w2_bias=w2_bias,

583 a1_scale=a1_scale,

584 a2_scale=a2_scale,

585 block_shape=block_shape,

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

@@ -147,6 +147,8 @@ def make_int8_moe_quant_config(

147 w2_scale: torch.Tensor,

148 a1_scale: torch.Tensor | None = None,

149 a2_scale: torch.Tensor | None = None,

150 + w1_bias: torch.Tensor | None = None,

151 + w2_bias: torch.Tensor | None = None,

152 per_act_token_quant: bool = False,

153 ) -> FusedMoEQuantConfig:

154 assert (a1_scale is None and a2_scale is None) or (

@@ -159,6 +161,8 @@ def make_int8_moe_quant_config(

161 w2_scale=w2_scale,

162 w1_zp=None,

163 w2_zp=None,

164 + w1_bias=w1_bias,

165 + w2_bias=w2_bias,

166 )

167

168 return int8_w8a8_moe_quant_config(

@@ -166,6 +170,8 @@ def make_int8_moe_quant_config(

170 w2_scale=w2_scale,

171 a1_scale=a1_scale,

172 a2_scale=a2_scale,

173 + w1_bias=w1_bias,

174 + w2_bias=w2_bias,

175 per_act_token_quant=per_act_token_quant,

176 )

177

vllm/model_executor/layers/quantization/modelopt.py

+67 -89

@@ -23,13 +23,10 @@ from vllm.model_executor.layers.fused_moe import (

23 FusedMoEMethodBase,

24 FusedMoEQuantConfig,

25 FusedMoeWeightScaleSupported,

26 - MoEActivation,

26 RoutedExperts,

28 - RoutingMethodType,

27 SharedExperts,

28 )

29 from vllm.model_executor.layers.fused_moe.oracle.fp8 import (

32 - Fp8MoeBackend,

30 convert_to_fp8_moe_kernel_format,

31 make_fp8_moe_kernel,

32 make_fp8_moe_quant_config,

@@ -70,7 +67,6 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (

67 MXFP8_BLOCK_SIZE,

68 MXFP8_SCALE_DTYPE,

69 MXFP8_VALUE_DTYPE,

73 - mxfp8_e4m3_quantize,

70 )

71 from vllm.model_executor.layers.quantization.utils.quant_utils import (

72 GroupShape,

@@ -93,7 +89,6 @@ from vllm.model_executor.parameter import (

89 PerTensorScaleParameter,

90 )

91 from vllm.model_executor.utils import replace_parameter, set_weight_attrs

96 -from vllm.utils.flashinfer import flashinfer_trtllm_fp8_block_scale_moe

93 if TYPE_CHECKING:

94 from vllm.model_executor.models.utils import WeightsMapper

@@ -1860,10 +1855,11 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):

1855 moe_config: FusedMoEConfig,

1856 ) -> None:

1857 super().__init__(moe_config)

1858 + self.weight_block_size = [1, MXFP8_BLOCK_SIZE]

1859 self.quant_config = quant_config

1860 assert self.quant_config.is_checkpoint_mxfp8_serialized

1861

1866 - self.mxfp8_backend, _ = select_mxfp8_moe_backend(self.moe)

1862 + self.mxfp8_backend, self.experts_cls = select_mxfp8_moe_backend(config=self.moe)

1863

1864 def create_weights(

1865 self,

@@ -2059,12 +2055,41 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):

2055 )

2056

2057 def process_weights_after_loading(self, layer: RoutedExperts) -> None:

2058 + # TODO(bnell): why is this required only for mxfp8?

2059 if getattr(layer, "_already_called_process_weights_after_loading", False):

2060 return

2061 + layer._already_called_process_weights_after_loading = True

2062

2063 self._check_weight_dtypes(layer)

2066 - self._shuffle_weights_for_trtllm(layer)

2067 - layer._already_called_process_weights_after_loading = True

2064 +

2065 + layer.weight_block_size = self.weight_block_size

2066 +

2067 + w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format(

2068 + fp8_backend=self.mxfp8_backend,

2069 + layer=layer,

2070 + w13=layer.w13_weight,

2071 + w2=layer.w2_weight,

2072 + w13_scale=layer.w13_weight_scale,

2073 + w2_scale=layer.w2_weight_scale,

2074 + w13_input_scale=None,

2075 + w2_input_scale=None,

2076 + )

2077 +

2078 + replace_parameter(layer, "w13_weight", w13)

2079 + replace_parameter(layer, "w2_weight", w2)

2080 + replace_parameter(layer, "w13_weight_scale", w13_scale)

2081 + replace_parameter(layer, "w2_weight_scale", w2_scale)

2082 +

2083 + self.moe_quant_config = self.get_fused_moe_quant_config(layer)

2084 + assert self.moe_quant_config is not None

2085 + assert self.experts_cls is not None

2086 + self.moe_kernel = make_fp8_moe_kernel(

2087 + moe_quant_config=self.moe_quant_config,

2088 + moe_config=self.moe,

2089 + fp8_backend=self.mxfp8_backend,

2090 + experts_cls=self.experts_cls,

2091 + routing_tables=layer._expert_routing_tables(),

2092 + )

2093

2094 def maybe_make_prepare_finalize(

2095 self,

@@ -2088,12 +2113,14 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):

2113 def get_fused_moe_quant_config(

2114 self, layer: RoutedExperts

2115 ) -> FusedMoEQuantConfig | None:

2091 - # TRTLLM MXFP8 path is monolithic and does not use modular kernel config.

2092 - return None

2093 -

2094 - @property

2095 - def is_monolithic(self) -> bool:

2096 - return self.mxfp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM

2116 + return make_fp8_moe_quant_config(

2117 + fp8_backend=self.mxfp8_backend,

2118 + w1_scale=layer.w13_weight_scale,

2119 + w2_scale=layer.w2_weight_scale,

2120 + a1_scale=None,

2121 + a2_scale=None,

2122 + block_shape=self.weight_block_size,

2123 + )

2124

2125 def apply_monolithic(

2126 self,

@@ -2102,83 +2129,23 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):

2129 router_logits: torch.Tensor,

2130 input_ids: torch.Tensor | None = None,

2131 ) -> torch.Tensor:

2105 - from flashinfer.fused_moe.core import (

2106 - ActivationType,

2107 - Fp8QuantizationType,

2108 - )

2109 -

2110 - assert self.mxfp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM

2111 -

2112 - if layer.eplb_state is not None:

2113 - raise NotImplementedError(

2114 - "EPLB is not supported for FlashInfer TRTLLM MXFP8 MoE backend."

2115 - )

2116 -

2117 - supported_activations = [MoEActivation.SILU]

2118 - if layer.activation not in supported_activations:

2119 - raise NotImplementedError(

2120 - "FlashInfer TRTLLM MXFP8 MoE supports only "

2121 - f"{supported_activations}, got {layer.activation}."

2122 - )

2123 -

2124 - # Map vLLM MoEActivation to FlashInfer ActivationType.

2125 - activation_map = {

2126 - MoEActivation.SILU: ActivationType.Swiglu,

2127 - MoEActivation.RELU2_NO_MUL: ActivationType.Relu2,

2128 - }

2129 - fi_activation_type: ActivationType = activation_map[layer.activation]

2130 -

2131 - # DeepSeekV3 routing requires float32 logits; others expect bfloat16.

2132 - if layer.routing_method_type == RoutingMethodType.DeepSeekV3:

2133 - assert router_logits.dtype == torch.float32, (

2134 - "DeepSeekV3 routing requires float32 router_logits, "

2135 - f"got {router_logits.dtype}."

2136 - )

2137 - else:

2138 - router_logits = router_logits.to(torch.bfloat16)

2139 -

2140 - # Treat 0 as "unset" for compatibility with ungrouped routing configs.

2141 - n_group = layer.num_expert_group or None

2142 - topk_group = layer.topk_group or None

2143 -

2144 - hidden_states_mxfp8, hidden_states_scale = mxfp8_e4m3_quantize(

2132 + assert self.is_monolithic

2133 + assert self.moe_kernel is not None

2134 + return self.moe_kernel.apply_monolithic(

2135 x,

2146 - is_sf_swizzled_layout=False,

2147 - )

2148 -

2149 - kwargs: dict = dict(

2150 - routing_logits=router_logits,

2151 - routing_bias=layer.e_score_correction_bias,

2152 - hidden_states=hidden_states_mxfp8,

2153 - hidden_states_scale=hidden_states_scale,

2154 - gemm1_weights=layer.w13_weight,

2155 - gemm1_weights_scale=layer.w13_weight_scale,

2156 - gemm2_weights=layer.w2_weight,

2157 - gemm2_weights_scale=layer.w2_weight_scale,

2158 - num_experts=layer.global_num_experts,

2159 - top_k=layer.top_k,

2160 - # Keep Optional semantics: FlashInfer expects None for non-grouped

2161 - # routing (e.g. Qwen3 Renormalize), not 0.

2162 - n_group=n_group,

2163 - topk_group=topk_group,

2164 - intermediate_size=layer.intermediate_size_per_partition,

2165 - local_expert_offset=layer.ep_rank * layer.local_num_experts,

2166 - local_num_experts=layer.local_num_experts,

2136 + layer.w13_weight,

2137 + layer.w2_weight,

2138 + router_logits,

2139 + activation=layer.activation,

2140 + global_num_experts=layer.global_num_experts,

2141 + expert_map=layer.expert_map,

2142 + apply_router_weight_on_input=layer.apply_router_weight_on_input,

2143 + num_expert_group=layer.num_expert_group,

2144 + topk_group=layer.topk_group,

2145 + e_score_correction_bias=layer.e_score_correction_bias,

2146 routed_scaling_factor=layer.routed_scaling_factor,

2168 - routing_method_type=layer.routing_method_type,

2169 - use_shuffled_weight=True,

2170 - weight_layout=0,

2171 - fp8_quantization_type=Fp8QuantizationType.MxFp8,

2147 )

2148

2174 - if fi_activation_type != ActivationType.Swiglu:

2175 - raise NotImplementedError(

2176 - "FlashInfer TRTLLM MXFP8 MoE supports only Swiglu activation, "

2177 - f"got {fi_activation_type}."

2178 - )

2179 -

2180 - return flashinfer_trtllm_fp8_block_scale_moe(**kwargs)

2181 -

2149 def apply(

2150 self,

2151 layer: RoutedExperts,

@@ -2189,8 +2156,19 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):

2156 shared_experts_input: torch.Tensor | None,

2157 ) -> torch.Tensor:

2158 assert not self.is_monolithic

2192 - raise NotImplementedError(

2193 - "Non-monolithic MXFP8 MoE path is not yet implemented."

2159 + assert self.moe_kernel is not None

2160 + return self.moe_kernel.apply(

2161 + x,

2162 + layer.w13_weight,

2163 + layer.w2_weight,

2164 + topk_weights,

2165 + topk_ids,

2166 + activation=layer.activation,

2167 + global_num_experts=layer.global_num_experts,

2168 + expert_map=layer.expert_map,

2169 + apply_router_weight_on_input=layer.apply_router_weight_on_input,

2170 + shared_experts=shared_experts,

2171 + shared_experts_input=shared_experts_input,

2172 )

2173

2174

vllm/model_executor/layers/quantization/online/fp8.py

+3 -4

@@ -371,19 +371,18 @@ class _Fp8OnlineMoEBase(OnlineMoEMethodBase):

371 a1_scale = layer.w13_input_scale

372 a2_scale = layer.w2_input_scale

373

374 - quant_config = make_fp8_moe_quant_config(

374 + return make_fp8_moe_quant_config(

375 fp8_backend=self.fp8_backend,

376 w1_scale=w1_scale,

377 w2_scale=w2_scale,

378 a1_scale=a1_scale,

379 a2_scale=a2_scale,

380 + w1_bias=getattr(layer, "w13_bias", None),

381 + w2_bias=getattr(layer, "w2_bias", None),

382 block_shape=self.weight_block_size,

383 swiglu_limit=getattr(layer, "swiglu_limit", None),

384 )

385

384 - self._maybe_inject_biases(quant_config, layer)

385 - return quant_config

386 -

386

387 class Fp8PerTensorOnlineMoEMethod(_Fp8OnlineMoEBase):

388 """Online tensorwise FP8 MoE quantization.

vllm/model_executor/layers/quantization/online/int8.py

+3 -3

@@ -105,9 +105,9 @@ class Int8OnlineMoEMethod(OnlineMoEMethodBase):

105 def get_fused_moe_quant_config(

106 self, layer: torch.nn.Module

107 ) -> "FusedMoEQuantConfig | None":

108 - quant_config = make_int8_moe_quant_config(

108 + return make_int8_moe_quant_config(

109 w1_scale=layer.w13_scale,

110 w2_scale=layer.w2_scale,

111 + w1_bias=getattr(layer, "w13_bias", None),

112 + w2_bias=getattr(layer, "w2_bias", None),

113 )

112 - self._maybe_inject_biases(quant_config, layer)

113 - return quant_config

vllm/model_executor/layers/quantization/online/moe_base.py

-16

@@ -8,7 +8,6 @@ import torch

8 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

9 from vllm.model_executor.layers.fused_moe import (

10 FusedMoEMethodBase,

11 - FusedMoEQuantConfig,

11 RoutedExperts,

12 SharedExperts,

13 )

@@ -101,21 +100,6 @@ class OnlineMoEMethodBase(FusedMoEMethodBase):

100 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

101 pass

102

104 - def _maybe_inject_biases(

105 - self,

106 - quant_config: FusedMoEQuantConfig,

107 - layer: torch.nn.Module,

108 - ) -> None:

109 - """Inject biases into the quant config if the model has them

110 - (e.g. GPT-OSS biased MoE)."""

111 - if self.moe.has_bias:

112 - w13_bias = getattr(layer, "w13_bias", None)

113 - w2_bias = getattr(layer, "w2_bias", None)

114 - if w13_bias is not None:

115 - quant_config._w1.bias = w13_bias

116 - if w2_bias is not None:

117 - quant_config._w2.bias = w2_bias

118 -

103 def maybe_make_prepare_finalize(

104 self,

105 routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,

vllm/model_executor/layers/quantization/online/mxfp8.py

+3 -4

@@ -214,19 +214,18 @@ class Mxfp8OnlineMoEMethod(OnlineMoEMethodBase):

214 a1_scale = layer.w13_input_scale

215 a2_scale = layer.w2_input_scale

216

217 - quant_config = make_fp8_moe_quant_config(

217 + return make_fp8_moe_quant_config(

218 fp8_backend=self.fp8_backend,

219 w1_scale=w1_scale,

220 w2_scale=w2_scale,

221 a1_scale=a1_scale,

222 a2_scale=a2_scale,

223 + w1_bias=getattr(layer, "w13_bias", None),

224 + w2_bias=getattr(layer, "w2_bias", None),

225 block_shape=self.weight_block_size,

226 swiglu_limit=getattr(layer, "swiglu_limit", None),

227 )

228

227 - self._maybe_inject_biases(quant_config, layer)

228 - return quant_config

229 -

229 def process_weights_after_loading(self, layer: Module) -> None:

230 if getattr(layer, "_already_called_process_weights_after_loading", False):

231 return