iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[MoE] Free unused MXFP4 scales in OAI Triton Backend (#46549)

Browse Source 
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
This commit is contained in:
Woosuk Kwon
2026-06-24 00:06:41 -07:00
committed by GitHub
parent 9d6fdc2901
commit 96de8bb389
2 changed files with 20 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+10
View File
@@ -1146,8 +1146,13 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
weight_scale=w2_scale, flex_ctx=FlexCtx(rhs_data=w2_flex)
)
# The original mxfp4 block scales have been swizzled into the
# precision configs above and are no longer read by the kernel, so
# drop the now-dead weight/scale Parameters to free their memory.
del layer.w13_weight
del layer.w2_weight
del layer.w13_weight_scale
del layer.w2_weight_scale
return (
w13_weight,
@@ -1508,8 +1513,13 @@ def convert_weight_to_mxfp4_moe_kernel_format(
weight_scale=w2_scale, flex_ctx=FlexCtx(rhs_data=w2_flex)
)
# The original mxfp4 block scales have been swizzled into the
# precision configs above and are no longer read by the kernel, so
# drop the now-dead weight/scale Parameters to free their memory.
del layer.w13_weight
del layer.w2_weight
del layer.w13_weight_scale
del layer.w2_weight_scale
return (
w13_weight,
vllm/model_executor/layers/quantization/mxfp4.py
+10 -4
View File
@@ -393,16 +393,19 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):
def get_fused_moe_quant_config(
self, layer: RoutedExperts
) -> FusedMoEQuantConfig | None:
w1_scale = layer.w13_weight_scale
w2_scale = layer.w2_weight_scale
w1_bias = getattr(layer, "w13_bias", None)
w2_bias = getattr(layer, "w2_bias", None)
if self.mxfp4_backend in TRITON_BACKENDS:
# TRITON backends free w13/w2_weight_scale after swizzling; the
# swizzled scales live inside the precision configs instead.
assert self.w13_precision_config is not None
assert self.w2_precision_config is not None
w1_scale = self.w13_precision_config
w2_scale = self.w2_precision_config
else:
w1_scale = layer.w13_weight_scale
w2_scale = layer.w2_weight_scale
return make_mxfp4_moe_quant_config(
mxfp4_backend=self.mxfp4_backend,
@@ -738,17 +741,20 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
self,
layer: RoutedExperts,
) -> FusedMoEQuantConfig | None:
w1_scale = layer.w13_weight_scale
w2_scale = layer.w2_weight_scale
w1_bias = getattr(layer, "w13_bias", None)
w2_bias = getattr(layer, "w2_bias", None)
swiglu_limit = getattr(layer, "swiglu_limit", None)
if self.mxfp4_backend in TRITON_BACKENDS:
# TRITON backends free w13/w2_weight_scale after swizzling; the
# swizzled scales live inside the precision configs instead.
assert self.w13_precision_config is not None
assert self.w2_precision_config is not None
w1_scale = self.w13_precision_config
w2_scale = self.w2_precision_config
else:
w1_scale = layer.w13_weight_scale
w2_scale = layer.w2_weight_scale
return make_mxfp4_moe_quant_config(
mxfp4_backend=self.mxfp4_backend,