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 13, 13:15 UTC
[Rocm][fused_moe][fp4] view weight to torch.float4_e2m1fn_x2 when running aiter fused moe for fp4 model (#27474)
Signed-off-by: zejunchen-zejun <zejun.chen@amd.com>
zejunchen-zejun committed
Nov 10, 2025 at 23:38 UTC
b06b9470ca881f89feea72e4d89b20c213f360d4
1 file changed
+12
vllm/model_executor/layers/quantization/quark/quark_moe.py
+12
@@ -458,6 +458,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
458
459
self.weight_dtype = self.weight_quant["dtype"].replace("fp", "mxfp")
460
self.input_dtype = self.input_quant["dtype"].replace("fp", "mxfp")
461
+ self.fp4_dtype = getattr(torch, "float4_e2m1fn_x2", None)
462
463
self.ocp_mx_scheme = OCP_MX_Scheme.from_quant_dtype(
464
self.input_dtype, self.weight_dtype
@@ -581,6 +582,17 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
582
w2_weight_scale = layer.w2_weight_scale.view(s0 * s1, -1)
583
w2_weight_scale = e8m0_shuffle(w2_weight_scale)
584
layer.w2_weight_scale.data = w2_weight_scale.view(s0, s1, -1)
585
+
586
+ if self.fp4_dtype is not None:
587
+ layer.w13_weight = torch.nn.Parameter(
588
+ layer.w13_weight.view(self.fp4_dtype),
589
+ requires_grad=layer.w13_weight.requires_grad,
590
+ )
591
+ layer.w2_weight = torch.nn.Parameter(
592
+ layer.w2_weight.view(self.fp4_dtype),
593
+ requires_grad=layer.w2_weight.requires_grad,
594
+ )
595
+
596
torch.cuda.empty_cache()
597
598
def get_fused_moe_quant_config(