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 6, 14:30 UTC
Fix eplb nvfp4 experts hook (#37217)
Signed-off-by: Elvir Crncevic <elvircrn@gmail.com> Signed-off-by: Elvir Crncevic <elvir@anthropic.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Elvir Crnčević committed
Mar 16, 2026 at 23:03 UTC
fd4d96302a2999a8d773b1b331951d232e3f5e05
10 files changed
+57
-25
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+7
@@ -659,6 +659,13 @@ def run_cutlass_moe_fp4(
659
class CutlassExpertsFp4(mk.FusedMoEExpertsModular):
660
"""CUTLASS FP4 fused MoE expert implementation."""
661
662
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
663
+ # Fuse activation scales into w_scale_2 in-place so that
664
+ # g1/g2_alphas (which reference the same tensor) stay in sync
665
+ # when EPLB rearranges the parameter.
666
+ layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale)
667
+ layer.w2_weight_scale_2.data.mul_(layer.w2_input_scale)
668
+
669
@property
670
def expects_unquantized_inputs(self) -> bool:
671
return True
vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
+19
-4
@@ -56,10 +56,25 @@ class TrtLlmNvFp4ExpertsBase:
56
# g1_scale_c = a13_scale * w13_scale_2 / a2_scale
57
self.g1_scale_c = self.quant_config.g1_alphas * self.quant_config.a2_gscale
58
else:
59
- self.g1_scale_c = (
60
- torch.ones_like(self.quant_config.a1_gscale)
61
- * self.quant_config.a2_gscale
62
- )
59
+ self.g1_scale_c = self.quant_config.a2_gscale.clone()
60
+
61
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
62
+ layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale)
63
+ layer.w2_weight_scale_2.data.mul_(layer.w2_input_scale)
64
+ # Recompute g1_scale_c since g1_alphas was just fused in-place.
65
+ # Register as a layer parameter so EPLB rearranges it alongside
66
+ # other expert weights.
67
+ assert self.quant_config.g1_alphas is not None
68
+ assert self.quant_config.a2_gscale is not None
69
+ if self.moe_config.is_act_and_mul:
70
+ g1_scale_c = self.quant_config.g1_alphas * self.quant_config.a2_gscale
71
+ else:
72
+ g1_scale_c = self.quant_config.a2_gscale.clone()
73
+ layer.register_parameter(
74
+ "g1_scale_c",
75
+ torch.nn.Parameter(g1_scale_c, requires_grad=False),
76
+ )
77
+ self.g1_scale_c = layer.g1_scale_c
78
79
@staticmethod
80
def _supports_current_device() -> bool:
vllm/model_executor/layers/fused_moe/flashinfer_cutedsl_moe.py
+4
@@ -49,6 +49,10 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
49
)
50
self.out_dtype = moe_config.in_dtype
51
52
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
53
+ layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale)
54
+ layer.w2_weight_scale_2.data.mul_(layer.w2_input_scale)
55
+
56
@staticmethod
57
def activation_format() -> mk.FusedMoEActivationFormat:
58
return mk.FusedMoEActivationFormat.BatchedExperts
vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py
+5
@@ -61,6 +61,11 @@ def is_valid_flashinfer_cutlass_fused_moe(
61
62
63
class FlashInferExperts(mk.FusedMoEExpertsModular):
64
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
65
+ if self.quant_config.use_nvfp4_w4a4:
66
+ layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale)
67
+ layer.w2_weight_scale_2.data.mul_(layer.w2_input_scale)
68
+
69
def __init__(
70
self,
71
moe_config: mk.FusedMoEConfig,
vllm/model_executor/layers/fused_moe/layer.py
+11
-7
@@ -1421,19 +1421,23 @@ class FusedMoE(CustomOp):
1421
weights = list(self.named_parameters())
1422
weights = [(name, _maybe_make_contiguous(name, p)) for name, p in weights]
1423
1424
+ # `w13_input_scale` and `w2_input_scale` are global per-tensor
1425
+ # activation scales shared across all experts (e.g. NVFP4).
1426
+ # They are broadcast views (stride 0) from .expand() and are
1427
+ # not actual expert weights, so exclude them from EPLB.
1428
+ NON_EXPERT_WEIGHTS = {
1429
+ "e_score_correction_bias",
1430
+ "w13_input_scale",
1431
+ "w2_input_scale",
1432
+ }
1433
+
1434
assert all(
1435
weight.is_contiguous()
1436
for name, weight in weights
1437
if not (name.startswith("_shared_experts.") or name.startswith("_gate."))
1438
+ and name not in NON_EXPERT_WEIGHTS
1439
)
1440
1430
- # Filter out the non-expert weights.
1431
- # `e_score_correction_bias` is a bias for each logical expert,
1432
- # with shape (num_logical_experts,), not an expert weight.
1433
- NON_EXPERT_WEIGHTS = {
1434
- "e_score_correction_bias",
1435
- }
1436
-
1441
return [
1442
weight.view(self.local_num_experts, -1)
1443
for name, weight in weights
vllm/model_executor/layers/fused_moe/modular_kernel.py
+3
@@ -489,6 +489,9 @@ class FusedMoEExperts(ABC):
489
self.max_num_tokens = max_num_tokens
490
self.num_dispatchers = num_dispatchers
491
492
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # noqa: B027
493
+ pass
494
+
495
@staticmethod
496
def is_monolithic() -> bool:
497
raise NotImplementedError("Implemented by subclasses.")
vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
+6
-4
@@ -374,11 +374,13 @@ def make_nvfp4_moe_quant_config(
374
w2_scale=w2_scale,
375
)
376
377
- g1_alphas = a13_scale * w13_scale_2
378
- g2_alphas = a2_scale * w2_scale_2
377
+ # Pass w13_scale_2 / w2_scale_2 directly as g1/g2_alphas.
378
+ # The expert's process_weights_after_loading will fuse activation
379
+ # scales in-place. Since the quant config references the same tensor
380
+ # as the registered parameter, EPLB rearrangement stays in sync.
381
return nvfp4_moe_quant_config(
380
- g1_alphas=g1_alphas,
381
- g2_alphas=g2_alphas,
382
+ g1_alphas=w13_scale_2,
383
+ g2_alphas=w2_scale_2,
384
a1_gscale=(1.0 / a13_scale),
385
a2_gscale=(1.0 / a2_scale),
386
w1_scale=w13_scale,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+1
@@ -570,6 +570,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
570
shared_experts=layer.shared_experts,
571
routing_tables=layer._maybe_init_expert_routing_tables(),
572
)
573
+ self.moe_kernel.fused_experts.process_weights_after_loading(layer)
574
575
def maybe_make_prepare_finalize(
576
self,
vllm/model_executor/layers/quantization/modelopt.py
+1
@@ -1394,6 +1394,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1394
shared_experts=layer.shared_experts,
1395
routing_tables=layer._maybe_init_expert_routing_tables(),
1396
)
1397
+ self.moe_kernel.fused_experts.process_weights_after_loading(layer)
1398
1399
def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
1400
return make_nvfp4_moe_quant_config(
vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
-10
@@ -267,16 +267,6 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass(
267
num_experts=w13.size(0),
268
is_gated_activation=is_gated,
269
)
270
-
271
- # We do not need to make this a parameter, because
272
- # it is not used during the weight (re)-loading process.
273
- if is_gated:
274
- layer.g1_scale_c = a13_scale * w13_scale_2 / a2_scale
275
- else:
276
- layer.g1_scale_c = torch.ones_like(a13_scale) / a2_scale
277
- layer.a1_gscale = 1.0 / a13_scale
278
- layer.g1_alphas = a13_scale * w13_scale_2
279
- layer.g2_alphas = a2_scale * w2_scale_2
270
else:
271
# Swizzle the block scales for other FI NVFP4 MoE kernels.
272
w13_scale = swizzle_blockscale(w13_scale)