obscura/vllm
2
0
Fork 0
mirror of https://github.com/vllm-project/vllm.git synced 2026-08-24 22:50:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

[ROCm][Bugfix] Re-tag AITER MoE weights as preshuffled after replace_parameter (#42061)

Browse Source 
Signed-off-by: Markus Hartikainen <[email protected]>
Co-authored-by: TJian <[email protected]>
This commit is contained in:
Markus Hartikainen
2026-05-08 21:42:07 -07:00
committed by GitHub
co-authored by TJian
parent df2636a9d8
commit e8f9038ebd
3 changed files with 21 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+5
View File
@@ -171,6 +171,11 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
replace_parameter(layer, "w13_weight", w13_new, prefer_copy=is_weight_update)
replace_parameter(layer, "w2_weight", w2_new, prefer_copy=is_weight_update)
# AITER backend requires weights to be marked as shuffled.
if self.unquantized_backend == UnquantizedMoeBackend.AITER:
layer.w13_weight.is_shuffled = True
layer.w2_weight.is_shuffled = True
if not is_weight_update:
# Setup moe kernel only on the first call. For the unquantized
# method, moe_quant_config is either the constant
vllm/model_executor/layers/quantization/fp8.py
+6
View File
@@ -28,6 +28,7 @@ from vllm.model_executor.layers.fused_moe.config import (
)
from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod
from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
Fp8MoeBackend,
convert_to_fp8_moe_kernel_format,
make_fp8_moe_kernel,
make_fp8_moe_quant_config,
@@ -766,6 +767,11 @@ class Fp8MoEMethod(FusedMoEMethodBase):
replace_parameter(layer, f"w13_{self.weight_scale_name}", w13_scale)
replace_parameter(layer, f"w2_{self.weight_scale_name}", w2_scale)
# AITER backend requires weights to be marked as shuffled.
if self.fp8_backend == Fp8MoeBackend.AITER:
layer.w13_weight.is_shuffled = True
layer.w2_weight.is_shuffled = True
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
if self.moe_quant_config:
assert self.experts_cls is not None
vllm/model_executor/layers/quantization/mxfp4.py
+10
View File
@@ -350,6 +350,11 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):
self.w13_precision_config = w13_scale
self.w2_precision_config = w2_scale
# AITER backend requires weights to be marked as shuffled.
if self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16:
layer.w13_weight.is_shuffled = True
layer.w2_weight.is_shuffled = True
if w13_bias is not None and w2_bias is not None:
replace_parameter(layer, "w13_bias", w13_bias)
replace_parameter(layer, "w2_bias", w2_bias)
@@ -678,6 +683,11 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
self.w13_precision_config = w13_scale
self.w2_precision_config = w2_scale
# AITER backend requires weights to be marked as shuffled.
if self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16:
layer.w13_weight.is_shuffled = True
layer.w2_weight.is_shuffled = True
if w13_bias is not None and w2_bias is not None:
replace_parameter(layer, "w13_bias", w13_bias)
replace_parameter(layer, "w2_bias", w2_bias)