[MoE] Free unused MXFP4 scales in OAI Triton Backend (#46549)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Jun 24, 2026 at 00:06 UTC
96de8bb389ae2c32d5fb6ab42c430423efeb9486
2 files changed
+20
-4
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+10
@@ -1146,8 +1146,13 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
1146
weight_scale=w2_scale, flex_ctx=FlexCtx(rhs_data=w2_flex)
1147
)
1148
1149
+ # The original mxfp4 block scales have been swizzled into the
1150
+ # precision configs above and are no longer read by the kernel, so
1151
+ # drop the now-dead weight/scale Parameters to free their memory.
1152
del layer.w13_weight
1153
del layer.w2_weight
1154
+ del layer.w13_weight_scale
1155
+ del layer.w2_weight_scale
1156
1157
return (
1158
w13_weight,
@@ -1508,8 +1513,13 @@ def convert_weight_to_mxfp4_moe_kernel_format(
1513
weight_scale=w2_scale, flex_ctx=FlexCtx(rhs_data=w2_flex)
1514
)
1515
1516
+ # The original mxfp4 block scales have been swizzled into the
1517
+ # precision configs above and are no longer read by the kernel, so
1518
+ # drop the now-dead weight/scale Parameters to free their memory.
1519
del layer.w13_weight
1520
del layer.w2_weight
1521
+ del layer.w13_weight_scale
1522
+ del layer.w2_weight_scale
1523
1524
return (
1525
w13_weight,
vllm/model_executor/layers/quantization/mxfp4.py
+10
-4
@@ -393,16 +393,19 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):
393
def get_fused_moe_quant_config(
394
self, layer: RoutedExperts
395
) -> FusedMoEQuantConfig | None:
396
- w1_scale = layer.w13_weight_scale
397
- w2_scale = layer.w2_weight_scale
396
w1_bias = getattr(layer, "w13_bias", None)
397
w2_bias = getattr(layer, "w2_bias", None)
398
399
if self.mxfp4_backend in TRITON_BACKENDS:
400
+ # TRITON backends free w13/w2_weight_scale after swizzling; the
401
+ # swizzled scales live inside the precision configs instead.
402
assert self.w13_precision_config is not None
403
assert self.w2_precision_config is not None
404
w1_scale = self.w13_precision_config
405
w2_scale = self.w2_precision_config
406
+ else:
407
+ w1_scale = layer.w13_weight_scale
408
+ w2_scale = layer.w2_weight_scale
409
410
return make_mxfp4_moe_quant_config(
411
mxfp4_backend=self.mxfp4_backend,
@@ -738,17 +741,20 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
741
self,
742
layer: RoutedExperts,
743
) -> FusedMoEQuantConfig | None:
741
- w1_scale = layer.w13_weight_scale
742
- w2_scale = layer.w2_weight_scale
744
w1_bias = getattr(layer, "w13_bias", None)
745
w2_bias = getattr(layer, "w2_bias", None)
746
swiglu_limit = getattr(layer, "swiglu_limit", None)
747
748
if self.mxfp4_backend in TRITON_BACKENDS:
749
+ # TRITON backends free w13/w2_weight_scale after swizzling; the
750
+ # swizzled scales live inside the precision configs instead.
751
assert self.w13_precision_config is not None
752
assert self.w2_precision_config is not None
753
w1_scale = self.w13_precision_config
754
w2_scale = self.w2_precision_config
755
+ else:
756
+ w1_scale = layer.w13_weight_scale
757
+ w2_scale = layer.w2_weight_scale
758
759
return make_mxfp4_moe_quant_config(
760
mxfp4_backend=self.mxfp4_backend,