Update unquantized_fused_moe_method.py

Change assignment of unquantized moe weights when using aiter on rocm, making it safer for reloading the weights. Solve the random output case after wake-up and reloading weights in reinforcement learning.
This commit is contained in:
Zhaodong Bing 2025-12-24 11:22:37 +08:00 committed by bingzhaodong
parent 3ce791ac77
commit cc7d87c6dd

View File

@ -211,16 +211,16 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
super().process_weights_after_loading(layer)
# Padding the weight for better performance on ROCm
layer.w13_weight.data = self._maybe_pad_weight(layer.w13_weight.data)
layer.w2_weight.data = self._maybe_pad_weight(layer.w2_weight.data)
layer.w13_weight.data.copy_(self._maybe_pad_weight(layer.w13_weight.data))
layer.w2_weight.data.copy_(self._maybe_pad_weight(layer.w2_weight.data))
if self.rocm_aiter_moe_enabled:
shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights(
layer.w13_weight.data, layer.w2_weight.data
)
layer.w13_weight.data = shuffled_w13
layer.w2_weight.data = shuffled_w2
layer.w13_weight.data.copy_(shuffled_w13)
layer.w2_weight.data.copy_(shuffled_w2)
if self.flashinfer_cutlass_moe_enabled:
# Swap halves to arrange as [w3; w1] (kernel expectation)