From b226ddacfda921b8eef7b75794c33d49ae51186a Mon Sep 17 00:00:00 2001 From: bnellnm <49004751+bnellnm@users.noreply.github.com> Date: Tue, 26 May 2026 11:14:14 -0400 Subject: [PATCH] [MoE Refactor] Migrate ModelOptMxFp8FusedMoE to oracle (#42768) Signed-off-by: Bill Nell Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> --- .../model_executor/layers/fused_moe/config.py | 23 ++- .../layers/fused_moe/oracle/fp8.py | 23 +-- .../layers/fused_moe/oracle/int8.py | 6 + .../layers/quantization/modelopt.py | 156 ++++++++---------- .../layers/quantization/online/fp8.py | 7 +- .../layers/quantization/online/int8.py | 6 +- .../layers/quantization/online/moe_base.py | 16 -- .../layers/quantization/online/mxfp8.py | 7 +- 8 files changed, 114 insertions(+), 130 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index eec7d4b87f..17cd5592d1 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -888,20 +888,33 @@ def int4_w4a16_moe_quant_config( def fp8_w8a16_moe_quant_config( w1_scale: torch.Tensor, w2_scale: torch.Tensor, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, block_shape: list[int] | None = None, ) -> FusedMoEQuantConfig: """ Construct a quant config for 16-bit float activations and fp8 weights. """ group_shape = GroupShape(*block_shape) if block_shape is not None else None + fp8_dtype = current_platform.fp8_dtype() return FusedMoEQuantConfig( _a1=FusedMoEQuantDesc(), _a2=FusedMoEQuantDesc(), _w1=FusedMoEQuantDesc( - current_platform.fp8_dtype(), group_shape, w1_scale, None, None + fp8_dtype, + group_shape, + w1_scale, + None, + None, + w1_bias, ), _w2=FusedMoEQuantDesc( - current_platform.fp8_dtype(), group_shape, w2_scale, None, None + fp8_dtype, + group_shape, + w2_scale, + None, + None, + w2_bias, ), ) @@ -911,6 +924,8 @@ def int8_w8a16_moe_quant_config( w2_scale: torch.Tensor, w1_zp: torch.Tensor | None, w2_zp: torch.Tensor | None, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, block_shape: list[int] | None = None, ) -> FusedMoEQuantConfig: """ @@ -920,8 +935,8 @@ def int8_w8a16_moe_quant_config( return FusedMoEQuantConfig( _a1=FusedMoEQuantDesc(shape=group_shape), _a2=FusedMoEQuantDesc(shape=group_shape), - _w1=FusedMoEQuantDesc(torch.int8, group_shape, w1_scale, None, w1_zp), - _w2=FusedMoEQuantDesc(torch.int8, group_shape, w2_scale, None, w2_zp), + _w1=FusedMoEQuantDesc(torch.int8, group_shape, w1_scale, None, w1_zp, w1_bias), + _w2=FusedMoEQuantDesc(torch.int8, group_shape, w2_scale, None, w2_zp, w2_bias), ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index a64c99bdb9..562128d66d 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -421,6 +421,7 @@ def select_fp8_moe_backend( def convert_to_fp8_moe_kernel_format( fp8_backend: Fp8MoeBackend, + # TODO(bnell): replace layer with weight_block_size layer: torch.nn.Module, w13: torch.Tensor, w2: torch.Tensor, @@ -508,6 +509,8 @@ def make_fp8_moe_quant_config( w2_scale: torch.Tensor, a1_scale: torch.Tensor | None, a2_scale: torch.Tensor | None, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, block_shape: list[int] | None = None, per_act_token_quant: bool = False, per_out_ch_quant: bool = False, @@ -526,19 +529,13 @@ def make_fp8_moe_quant_config( a method of the modular kernel itself. """ - # MARLIN is mixed precision W8A16 config. - if fp8_backend == Fp8MoeBackend.MARLIN: - return fp8_w8a16_moe_quant_config( - w1_scale=w1_scale, - w2_scale=w2_scale, - block_shape=block_shape, - ) - - # CPU is mixed precision W8A16 config. - if fp8_backend == Fp8MoeBackend.CPU: + # MARLIN and CPU are mixed precision W8A16 config. + if fp8_backend == Fp8MoeBackend.MARLIN or fp8_backend == Fp8MoeBackend.CPU: return fp8_w8a16_moe_quant_config( w1_scale=w1_scale, w2_scale=w2_scale, + w1_bias=w1_bias, + w2_bias=w2_bias, block_shape=block_shape, ) @@ -549,6 +546,8 @@ def make_fp8_moe_quant_config( return fp8_w8a8_moe_quant_config( w1_scale=w1_scale, w2_scale=w2_scale, + w1_bias=w1_bias, + w2_bias=w2_bias, a1_scale=a1_scale, a2_scale=a2_scale, a1_gscale=(1.0 / a1_scale), @@ -566,6 +565,8 @@ def make_fp8_moe_quant_config( "mxfp8", w1_scale=w1_scale, w2_scale=w2_scale, + w1_bias=w1_bias, + w2_bias=w2_bias, a1_scale=a1_scale, a2_scale=a2_scale, block_shape=block_shape, @@ -577,6 +578,8 @@ def make_fp8_moe_quant_config( return fp8_w8a8_moe_quant_config( w1_scale=w1_scale, w2_scale=w2_scale, + w1_bias=w1_bias, + w2_bias=w2_bias, a1_scale=a1_scale, a2_scale=a2_scale, block_shape=block_shape, diff --git a/vllm/model_executor/layers/fused_moe/oracle/int8.py b/vllm/model_executor/layers/fused_moe/oracle/int8.py index f5fceef727..36c4ccdd00 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int8.py @@ -147,6 +147,8 @@ def make_int8_moe_quant_config( w2_scale: torch.Tensor, a1_scale: torch.Tensor | None = None, a2_scale: torch.Tensor | None = None, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, per_act_token_quant: bool = False, ) -> FusedMoEQuantConfig: assert (a1_scale is None and a2_scale is None) or ( @@ -159,6 +161,8 @@ def make_int8_moe_quant_config( w2_scale=w2_scale, w1_zp=None, w2_zp=None, + w1_bias=w1_bias, + w2_bias=w2_bias, ) return int8_w8a8_moe_quant_config( @@ -166,6 +170,8 @@ def make_int8_moe_quant_config( w2_scale=w2_scale, a1_scale=a1_scale, a2_scale=a2_scale, + w1_bias=w1_bias, + w2_bias=w2_bias, per_act_token_quant=per_act_token_quant, ) diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index c4036e2f18..a8540f2ab8 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -23,13 +23,10 @@ from vllm.model_executor.layers.fused_moe import ( FusedMoEMethodBase, FusedMoEQuantConfig, FusedMoeWeightScaleSupported, - MoEActivation, RoutedExperts, - RoutingMethodType, SharedExperts, ) from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( - Fp8MoeBackend, convert_to_fp8_moe_kernel_format, make_fp8_moe_kernel, make_fp8_moe_quant_config, @@ -70,7 +67,6 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( MXFP8_BLOCK_SIZE, MXFP8_SCALE_DTYPE, MXFP8_VALUE_DTYPE, - mxfp8_e4m3_quantize, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( GroupShape, @@ -93,7 +89,6 @@ from vllm.model_executor.parameter import ( PerTensorScaleParameter, ) from vllm.model_executor.utils import replace_parameter, set_weight_attrs -from vllm.utils.flashinfer import flashinfer_trtllm_fp8_block_scale_moe if TYPE_CHECKING: from vllm.model_executor.models.utils import WeightsMapper @@ -1860,10 +1855,11 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase): moe_config: FusedMoEConfig, ) -> None: super().__init__(moe_config) + self.weight_block_size = [1, MXFP8_BLOCK_SIZE] self.quant_config = quant_config assert self.quant_config.is_checkpoint_mxfp8_serialized - self.mxfp8_backend, _ = select_mxfp8_moe_backend(self.moe) + self.mxfp8_backend, self.experts_cls = select_mxfp8_moe_backend(config=self.moe) def create_weights( self, @@ -2059,12 +2055,41 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase): ) def process_weights_after_loading(self, layer: RoutedExperts) -> None: + # TODO(bnell): why is this required only for mxfp8? if getattr(layer, "_already_called_process_weights_after_loading", False): return + layer._already_called_process_weights_after_loading = True self._check_weight_dtypes(layer) - self._shuffle_weights_for_trtllm(layer) - layer._already_called_process_weights_after_loading = True + + layer.weight_block_size = self.weight_block_size + + w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format( + fp8_backend=self.mxfp8_backend, + layer=layer, + w13=layer.w13_weight, + w2=layer.w2_weight, + w13_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + w13_input_scale=None, + w2_input_scale=None, + ) + + replace_parameter(layer, "w13_weight", w13) + replace_parameter(layer, "w2_weight", w2) + replace_parameter(layer, "w13_weight_scale", w13_scale) + replace_parameter(layer, "w2_weight_scale", w2_scale) + + self.moe_quant_config = self.get_fused_moe_quant_config(layer) + assert self.moe_quant_config is not None + assert self.experts_cls is not None + self.moe_kernel = make_fp8_moe_kernel( + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + fp8_backend=self.mxfp8_backend, + experts_cls=self.experts_cls, + routing_tables=layer._expert_routing_tables(), + ) def maybe_make_prepare_finalize( self, @@ -2088,12 +2113,14 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase): def get_fused_moe_quant_config( self, layer: RoutedExperts ) -> FusedMoEQuantConfig | None: - # TRTLLM MXFP8 path is monolithic and does not use modular kernel config. - return None - - @property - def is_monolithic(self) -> bool: - return self.mxfp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM + return make_fp8_moe_quant_config( + fp8_backend=self.mxfp8_backend, + w1_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + a1_scale=None, + a2_scale=None, + block_shape=self.weight_block_size, + ) def apply_monolithic( self, @@ -2102,83 +2129,23 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase): router_logits: torch.Tensor, input_ids: torch.Tensor | None = None, ) -> torch.Tensor: - from flashinfer.fused_moe.core import ( - ActivationType, - Fp8QuantizationType, - ) - - assert self.mxfp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM - - if layer.eplb_state is not None: - raise NotImplementedError( - "EPLB is not supported for FlashInfer TRTLLM MXFP8 MoE backend." - ) - - supported_activations = [MoEActivation.SILU] - if layer.activation not in supported_activations: - raise NotImplementedError( - "FlashInfer TRTLLM MXFP8 MoE supports only " - f"{supported_activations}, got {layer.activation}." - ) - - # Map vLLM MoEActivation to FlashInfer ActivationType. - activation_map = { - MoEActivation.SILU: ActivationType.Swiglu, - MoEActivation.RELU2_NO_MUL: ActivationType.Relu2, - } - fi_activation_type: ActivationType = activation_map[layer.activation] - - # DeepSeekV3 routing requires float32 logits; others expect bfloat16. - if layer.routing_method_type == RoutingMethodType.DeepSeekV3: - assert router_logits.dtype == torch.float32, ( - "DeepSeekV3 routing requires float32 router_logits, " - f"got {router_logits.dtype}." - ) - else: - router_logits = router_logits.to(torch.bfloat16) - - # Treat 0 as "unset" for compatibility with ungrouped routing configs. - n_group = layer.num_expert_group or None - topk_group = layer.topk_group or None - - hidden_states_mxfp8, hidden_states_scale = mxfp8_e4m3_quantize( + assert self.is_monolithic + assert self.moe_kernel is not None + return self.moe_kernel.apply_monolithic( x, - is_sf_swizzled_layout=False, - ) - - kwargs: dict = dict( - routing_logits=router_logits, - routing_bias=layer.e_score_correction_bias, - hidden_states=hidden_states_mxfp8, - hidden_states_scale=hidden_states_scale, - gemm1_weights=layer.w13_weight, - gemm1_weights_scale=layer.w13_weight_scale, - gemm2_weights=layer.w2_weight, - gemm2_weights_scale=layer.w2_weight_scale, - num_experts=layer.global_num_experts, - top_k=layer.top_k, - # Keep Optional semantics: FlashInfer expects None for non-grouped - # routing (e.g. Qwen3 Renormalize), not 0. - n_group=n_group, - topk_group=topk_group, - intermediate_size=layer.intermediate_size_per_partition, - local_expert_offset=layer.ep_rank * layer.local_num_experts, - local_num_experts=layer.local_num_experts, + layer.w13_weight, + layer.w2_weight, + router_logits, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + num_expert_group=layer.num_expert_group, + topk_group=layer.topk_group, + e_score_correction_bias=layer.e_score_correction_bias, routed_scaling_factor=layer.routed_scaling_factor, - routing_method_type=layer.routing_method_type, - use_shuffled_weight=True, - weight_layout=0, - fp8_quantization_type=Fp8QuantizationType.MxFp8, ) - if fi_activation_type != ActivationType.Swiglu: - raise NotImplementedError( - "FlashInfer TRTLLM MXFP8 MoE supports only Swiglu activation, " - f"got {fi_activation_type}." - ) - - return flashinfer_trtllm_fp8_block_scale_moe(**kwargs) - def apply( self, layer: RoutedExperts, @@ -2189,8 +2156,19 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase): shared_experts_input: torch.Tensor | None, ) -> torch.Tensor: assert not self.is_monolithic - raise NotImplementedError( - "Non-monolithic MXFP8 MoE path is not yet implemented." + assert self.moe_kernel is not None + return self.moe_kernel.apply( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights, + topk_ids, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + shared_experts=shared_experts, + shared_experts_input=shared_experts_input, ) diff --git a/vllm/model_executor/layers/quantization/online/fp8.py b/vllm/model_executor/layers/quantization/online/fp8.py index 5e14eade26..18270cbfe5 100644 --- a/vllm/model_executor/layers/quantization/online/fp8.py +++ b/vllm/model_executor/layers/quantization/online/fp8.py @@ -371,19 +371,18 @@ class _Fp8OnlineMoEBase(OnlineMoEMethodBase): a1_scale = layer.w13_input_scale a2_scale = layer.w2_input_scale - quant_config = make_fp8_moe_quant_config( + return make_fp8_moe_quant_config( fp8_backend=self.fp8_backend, w1_scale=w1_scale, w2_scale=w2_scale, a1_scale=a1_scale, a2_scale=a2_scale, + w1_bias=getattr(layer, "w13_bias", None), + w2_bias=getattr(layer, "w2_bias", None), block_shape=self.weight_block_size, swiglu_limit=getattr(layer, "swiglu_limit", None), ) - self._maybe_inject_biases(quant_config, layer) - return quant_config - class Fp8PerTensorOnlineMoEMethod(_Fp8OnlineMoEBase): """Online tensorwise FP8 MoE quantization. diff --git a/vllm/model_executor/layers/quantization/online/int8.py b/vllm/model_executor/layers/quantization/online/int8.py index fc9c18e66c..4274b16d55 100644 --- a/vllm/model_executor/layers/quantization/online/int8.py +++ b/vllm/model_executor/layers/quantization/online/int8.py @@ -105,9 +105,9 @@ class Int8OnlineMoEMethod(OnlineMoEMethodBase): def get_fused_moe_quant_config( self, layer: torch.nn.Module ) -> "FusedMoEQuantConfig | None": - quant_config = make_int8_moe_quant_config( + return make_int8_moe_quant_config( w1_scale=layer.w13_scale, w2_scale=layer.w2_scale, + w1_bias=getattr(layer, "w13_bias", None), + w2_bias=getattr(layer, "w2_bias", None), ) - self._maybe_inject_biases(quant_config, layer) - return quant_config diff --git a/vllm/model_executor/layers/quantization/online/moe_base.py b/vllm/model_executor/layers/quantization/online/moe_base.py index 32eb81601b..fe6302fd11 100644 --- a/vllm/model_executor/layers/quantization/online/moe_base.py +++ b/vllm/model_executor/layers/quantization/online/moe_base.py @@ -8,7 +8,6 @@ import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm.model_executor.layers.fused_moe import ( FusedMoEMethodBase, - FusedMoEQuantConfig, RoutedExperts, SharedExperts, ) @@ -101,21 +100,6 @@ class OnlineMoEMethodBase(FusedMoEMethodBase): def process_weights_after_loading(self, layer: torch.nn.Module) -> None: pass - def _maybe_inject_biases( - self, - quant_config: FusedMoEQuantConfig, - layer: torch.nn.Module, - ) -> None: - """Inject biases into the quant config if the model has them - (e.g. GPT-OSS biased MoE).""" - if self.moe.has_bias: - w13_bias = getattr(layer, "w13_bias", None) - w2_bias = getattr(layer, "w2_bias", None) - if w13_bias is not None: - quant_config._w1.bias = w13_bias - if w2_bias is not None: - quant_config._w2.bias = w2_bias - def maybe_make_prepare_finalize( self, routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, diff --git a/vllm/model_executor/layers/quantization/online/mxfp8.py b/vllm/model_executor/layers/quantization/online/mxfp8.py index 28d0706b11..c197398a09 100644 --- a/vllm/model_executor/layers/quantization/online/mxfp8.py +++ b/vllm/model_executor/layers/quantization/online/mxfp8.py @@ -214,19 +214,18 @@ class Mxfp8OnlineMoEMethod(OnlineMoEMethodBase): a1_scale = layer.w13_input_scale a2_scale = layer.w2_input_scale - quant_config = make_fp8_moe_quant_config( + return make_fp8_moe_quant_config( fp8_backend=self.fp8_backend, w1_scale=w1_scale, w2_scale=w2_scale, a1_scale=a1_scale, a2_scale=a2_scale, + w1_bias=getattr(layer, "w13_bias", None), + w2_bias=getattr(layer, "w2_bias", None), block_shape=self.weight_block_size, swiglu_limit=getattr(layer, "swiglu_limit", None), ) - self._maybe_inject_biases(quant_config, layer) - return quant_config - def process_weights_after_loading(self, layer: Module) -> None: if getattr(layer, "_already_called_process_weights_after_loading", False): return