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 4, 01:15 UTC
[ROCm][Bugfix] Re-tag AITER MoE weights as preshuffled after replace_parameter (#42061)
Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
Markus Hartikainen committed
May 9, 2026 at 07:42 UTC
e8f9038ebd2b72e8fea24642ca8b4c3500264145
3 files changed
+21
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+5
@@ -171,6 +171,11 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
171
replace_parameter(layer, "w13_weight", w13_new, prefer_copy=is_weight_update)
172
replace_parameter(layer, "w2_weight", w2_new, prefer_copy=is_weight_update)
173
174
+ # AITER backend requires weights to be marked as shuffled.
175
+ if self.unquantized_backend == UnquantizedMoeBackend.AITER:
176
+ layer.w13_weight.is_shuffled = True
177
+ layer.w2_weight.is_shuffled = True
178
+
179
if not is_weight_update:
180
# Setup moe kernel only on the first call. For the unquantized
181
# method, moe_quant_config is either the constant
vllm/model_executor/layers/quantization/fp8.py
+6
@@ -28,6 +28,7 @@ from vllm.model_executor.layers.fused_moe.config import (
28
)
29
from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod
30
from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
31
+ Fp8MoeBackend,
32
convert_to_fp8_moe_kernel_format,
33
make_fp8_moe_kernel,
34
make_fp8_moe_quant_config,
@@ -766,6 +767,11 @@ class Fp8MoEMethod(FusedMoEMethodBase):
767
replace_parameter(layer, f"w13_{self.weight_scale_name}", w13_scale)
768
replace_parameter(layer, f"w2_{self.weight_scale_name}", w2_scale)
769
770
+ # AITER backend requires weights to be marked as shuffled.
771
+ if self.fp8_backend == Fp8MoeBackend.AITER:
772
+ layer.w13_weight.is_shuffled = True
773
+ layer.w2_weight.is_shuffled = True
774
+
775
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
776
if self.moe_quant_config:
777
assert self.experts_cls is not None
vllm/model_executor/layers/quantization/mxfp4.py
+10
@@ -350,6 +350,11 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):
350
self.w13_precision_config = w13_scale
351
self.w2_precision_config = w2_scale
352
353
+ # AITER backend requires weights to be marked as shuffled.
354
+ if self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16:
355
+ layer.w13_weight.is_shuffled = True
356
+ layer.w2_weight.is_shuffled = True
357
+
358
if w13_bias is not None and w2_bias is not None:
359
replace_parameter(layer, "w13_bias", w13_bias)
360
replace_parameter(layer, "w2_bias", w2_bias)
@@ -678,6 +683,11 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
683
self.w13_precision_config = w13_scale
684
self.w2_precision_config = w2_scale
685
686
+ # AITER backend requires weights to be marked as shuffled.
687
+ if self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16:
688
+ layer.w13_weight.is_shuffled = True
689
+ layer.w2_weight.is_shuffled = True
690
+
691
if w13_bias is not None and w2_bias is not None:
692
replace_parameter(layer, "w13_bias", w13_bias)
693
replace_parameter(layer, "w2_bias", w2_bias)