obscura/vllm
2
0
Fork 0
mirror of https://github.com/vllm-project/vllm.git synced 2026-08-21 13:10:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Bill Nell <[email protected]>
Co-authored-by: Robert Shaw <[email protected]>
This commit is contained in:
bnellnm
2026-05-26 11:14:14 -04:00
committed by GitHub
co-authored by Robert Shaw
parent 6ab6ffb428
commit b226ddacfd
8 changed files with 114 additions and 130 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/model_executor/layers/fused_moe/config.py
+19 -4
View File
@@ -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),
)
vllm/model_executor/layers/fused_moe/oracle/fp8.py
+13 -10
View File
@@ -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,
vllm/model_executor/layers/fused_moe/oracle/int8.py
+6
View File
@@ -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,
)
vllm/model_executor/layers/quantization/modelopt.py
+67 -89
View File
@@ -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,
)
vllm/model_executor/layers/quantization/online/fp8.py
+3 -4
View File
@@ -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.
vllm/model_executor/layers/quantization/online/int8.py
+3 -3
View File
@@ -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
vllm/model_executor/layers/quantization/online/moe_base.py
-16
View File
@@ -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,
vllm/model_executor/layers/quantization/online/mxfp8.py
+3 -4
View File
@@ -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