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 8, 02:30 UTC
[MoE Refactor] Convert mxfp4 marlin into modular kernel format (#34588)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
Yongye Zhu committed
Feb 18, 2026 at 17:37 UTC
40da9625a106222f53e27761497fcf960b249a67
2 files changed
+42
-23
vllm/model_executor/layers/fused_moe/fused_marlin_moe.py
+4
@@ -30,6 +30,7 @@ from vllm.model_executor.layers.fused_moe.utils import (
30
disable_inplace,
31
)
32
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
33
+ get_marlin_input_dtype,
34
marlin_make_workspace_new,
35
marlin_moe_intermediate_size,
36
marlin_quant_input,
@@ -550,6 +551,8 @@ class MarlinExpertsBase(mk.FusedMoEPermuteExpertsUnpermute):
551
self.w13_g_idx_sort_indices = w13_g_idx_sort_indices
552
self.w2_g_idx_sort_indices = w2_g_idx_sort_indices
553
self.is_k_full = is_k_full
554
+ self.input_dtype = get_marlin_input_dtype()
555
+
556
super().__init__(
557
moe_config=moe_config,
558
quant_config=quant_config,
@@ -736,6 +739,7 @@ class MarlinExperts(MarlinExpertsBase):
739
sort_indices1=self.w13_g_idx_sort_indices,
740
sort_indices2=self.w2_g_idx_sort_indices,
741
is_k_full=self.is_k_full,
742
+ input_dtype=self.input_dtype,
743
)
744
745
def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None:
vllm/model_executor/layers/quantization/mxfp4.py
+38
-23
@@ -16,6 +16,9 @@ from vllm.model_executor.layers.fused_moe import (
16
MoEActivation,
17
)
18
from vllm.model_executor.layers.fused_moe import modular_kernel as mk
19
+from vllm.model_executor.layers.fused_moe.all2all_utils import (
20
+ maybe_make_prepare_finalize,
21
+)
22
from vllm.model_executor.layers.fused_moe.config import (
23
FusedMoEQuantConfig,
24
mxfp4_mxfp8_moe_quant_config,
@@ -25,7 +28,6 @@ from vllm.model_executor.layers.fused_moe.config import (
28
from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
29
BatchedMarlinExperts,
30
MarlinExperts,
28
- fused_marlin_moe,
31
)
32
from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
33
OAITritonExperts,
@@ -52,7 +54,6 @@ from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
54
from vllm.model_executor.layers.quantization.utils.quant_utils import is_layer_skipped
55
from vllm.model_executor.utils import set_weight_attrs
56
from vllm.platforms import current_platform
55
-from vllm.scalar_type import scalar_types
57
from vllm.utils.flashinfer import has_flashinfer
58
from vllm.utils.import_utils import has_triton_kernels
59
from vllm.utils.math_utils import round_up
@@ -219,7 +220,6 @@ class Mxfp4Config(QuantizationConfig):
220
return XpuMxfp4MoEMethod(layer.moe_config)
221
else:
222
quant_method = Mxfp4MoEMethod(layer.moe_config)
222
- quant_method.marlin_input_dtype = get_marlin_input_dtype(prefix)
223
return quant_method
224
elif isinstance(layer, Attention):
225
# TODO: Add support for MXFP4 Attention.
@@ -243,7 +243,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
243
self.weight_dtype = "mxfp4"
244
self.mxfp4_backend = get_mxfp4_backend(moe.is_lora_enabled)
245
246
- self.marlin_input_dtype = None
246
self.max_capture_size = (
247
get_current_vllm_config().compilation_config.max_cudagraph_capture_size
248
)
@@ -254,6 +253,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
253
"Please check your environment and try again."
254
)
255
self._cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
256
+ self.moe_mk: mk.FusedMoEModularKernel | None = None
257
258
def create_weights(
259
self,
@@ -408,7 +408,30 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
408
409
def process_weights_after_loading(self, layer):
410
if self.mxfp4_backend == Mxfp4Backend.MARLIN:
411
- prepare_moe_fp4_layer_for_marlin(layer, input_dtype=self.marlin_input_dtype)
411
+ prepare_moe_fp4_layer_for_marlin(
412
+ layer, input_dtype=get_marlin_input_dtype()
413
+ )
414
+
415
+ self.moe_quant_config = self.get_fused_moe_quant_config(layer)
416
+ assert self.moe_quant_config is not None
417
+
418
+ prepare_finalize = maybe_make_prepare_finalize(
419
+ moe=self.moe,
420
+ quant_config=self.moe_quant_config,
421
+ routing_tables=layer._maybe_init_expert_routing_tables(),
422
+ allow_new_interface=True,
423
+ )
424
+ assert prepare_finalize is not None
425
+
426
+ self.moe_mk = mk.FusedMoEModularKernel(
427
+ prepare_finalize,
428
+ MarlinExperts(
429
+ self.moe,
430
+ self.moe_quant_config,
431
+ ),
432
+ inplace=not self.moe.disable_inplace,
433
+ shared_experts=None,
434
+ )
435
elif (
436
self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM
437
or self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_BF16
@@ -910,27 +933,19 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
933
raise NotImplementedError("EPLB is not supported for mxfp4")
934
935
if self.mxfp4_backend == Mxfp4Backend.MARLIN:
913
- return fused_marlin_moe(
914
- x,
915
- layer.w13_weight,
916
- layer.w2_weight,
917
- layer.w13_bias,
918
- layer.w2_bias,
919
- layer.w13_weight_scale,
920
- layer.w2_weight_scale,
921
- topk_weights,
922
- topk_ids,
923
- global_scale1=None,
924
- global_scale2=None,
925
- quant_type_id=scalar_types.float4_e2m1f.id,
926
- apply_router_weight_on_input=layer.apply_router_weight_on_input,
927
- global_num_experts=layer.global_num_experts,
936
+ assert self.moe_mk is not None
937
+
938
+ return self.moe_mk(
939
+ hidden_states=x,
940
+ w1=layer.w13_weight,
941
+ w2=layer.w2_weight,
942
+ topk_weights=topk_weights,
943
+ topk_ids=topk_ids,
944
activation=layer.activation,
945
+ global_num_experts=layer.global_num_experts,
946
expert_map=layer.expert_map,
930
- input_dtype=self.marlin_input_dtype,
931
- inplace=not self.moe.disable_inplace,
947
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
948
)
933
-
949
assert _can_support_mxfp4(
950
layer.use_grouped_topk,
951
layer.topk_group,