From 224e73dad4e7fd5eae5807725a81c1cd8573d903 Mon Sep 17 00:00:00 2001 From: yadaish Date: Fri, 3 Jul 2026 14:47:33 +0800 Subject: [PATCH 1/2] update shufle weight for gfx1250 --- atom/model_ops/moe.py | 40 +++++++++++++++++++++++++++------------- 1 file changed, 27 insertions(+), 13 deletions(-) diff --git a/atom/model_ops/moe.py b/atom/model_ops/moe.py index 37b4bc8b9e..afcf2895f5 100644 --- a/atom/model_ops/moe.py +++ b/atom/model_ops/moe.py @@ -14,7 +14,11 @@ from aiter.jit.utils.chip_info import get_gfx from aiter.jit.utils.torch_guard import torch_compile_guard from aiter.ops.flydsl.moe_common import GateMode -from aiter.ops.shuffle import moe_shuffle_scale, shuffle_weight +from aiter.ops.shuffle import ( + interleave_gate_up_rows, + moe_shuffle_scale, + moe_shuffle_weight, +) from atom.config import ( Config, QuantizationConfig, @@ -800,13 +804,9 @@ def __init__(self, quant_config: LayerQuantConfig, moe: FusedMoEConfig): ) gfx = get_gfx() self.is_gfx1250 = gfx == "gfx1250" - # gfx1250 grouped a8w4 MoE kernel only supports the non-interleaved - # (gate|up separated) scale layout; reject is_guinterleave up front. - if self.is_gfx1250 and self.is_guinterleave: - raise NotImplementedError( - "gfx1250 MoE only supports is_guinterleave=False; " - "unset ATOM_MOE_GU_ITLV." - ) + # gfx1250 GUGU is handled arch-aware by moe_shuffle_weight / + # moe_shuffle_scale (row-level gate/up interleave + WMMA tile shuffle), + # so is_guinterleave is supported on gfx1250 as well. if envs.is_set("ATOM_USE_TRITON_MOE"): self.use_triton = envs.ATOM_USE_TRITON_MOE else: @@ -1022,20 +1022,30 @@ def process_weights_after_loading(self, layer): layer.w2_swizzle_layout = w2_swizzle_layout return - # shuffle weight - layer.w13_weight.data = shuffle_weight( + # shuffle weight (arch-aware: gfx1250 does the GUGU row interleave + + # WMMA tile shuffle internally, other archs use the lane-level path) + layer.w13_weight.data = moe_shuffle_weight( layer.w13_weight, + experts_cnt=self.num_experts, is_guinterleave=self.is_guinterleave, gate_up=True, ) - layer.w2_weight.data = shuffle_weight( + layer.w2_weight.data = moe_shuffle_weight( layer.w2_weight, + experts_cnt=self.num_experts, is_guinterleave=self.is_guinterleave, gate_up=False, ) layer.w13_weight.is_shuffled = True layer.w2_weight.is_shuffled = True + # GUGU (is_guinterleave) reorders the stage1 output rows to + # [g0, u0, g1, u1, ...]; moe_shuffle_weight interleaves the weight rows + # but not the bias, so interleave w13_bias to match. w2_bias (stage2, + # single N=hidden GEMM) has no gate/up concept and is left as-is. + if self.is_guinterleave and layer.w13_bias is not None: + layer.w13_bias.data = interleave_gate_up_rows(layer.w13_bias.data) + # shuffle scale w13_scale_2d = layer.w13_weight_scale.reshape( -1, layer.w13_weight_scale.shape[-1] @@ -1962,13 +1972,17 @@ def _process_block_quant(self, layer: nn.Module) -> None: # aiter's MXFP8 MoE kernels consume the same gate/up interleaved # layout used by their 1x32 shuffle helpers. Keep this branch # isolated so the existing 1x128 FP8 path still uses shuffle_weights. - layer.w13_weight.data = shuffle_weight( + # moe_shuffle_weight mirrors moe_shuffle_scale: arch-aware GUGU + # handling (row interleave + WMMA tile shuffle on gfx1250). + layer.w13_weight.data = moe_shuffle_weight( layer.w13_weight, + experts_cnt=self.num_experts, is_guinterleave=True, gate_up=True, ) - layer.w2_weight.data = shuffle_weight( + layer.w2_weight.data = moe_shuffle_weight( layer.w2_weight, + experts_cnt=self.num_experts, is_guinterleave=True, gate_up=False, ) From 8bbd6aeaec3ca013eceac75cee646ede25a4cb15 Mon Sep 17 00:00:00 2001 From: yadaish Date: Thu, 9 Jul 2026 19:54:08 +0800 Subject: [PATCH 2/2] address review: remove redundant gfx1250 GUGU comment Co-Authored-By: Claude Opus 4.8 (1M context) --- atom/model_ops/moe.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/atom/model_ops/moe.py b/atom/model_ops/moe.py index afcf2895f5..707e604d9f 100644 --- a/atom/model_ops/moe.py +++ b/atom/model_ops/moe.py @@ -804,9 +804,6 @@ def __init__(self, quant_config: LayerQuantConfig, moe: FusedMoEConfig): ) gfx = get_gfx() self.is_gfx1250 = gfx == "gfx1250" - # gfx1250 GUGU is handled arch-aware by moe_shuffle_weight / - # moe_shuffle_scale (row-level gate/up interleave + WMMA tile shuffle), - # so is_guinterleave is supported on gfx1250 as well. if envs.is_set("ATOM_USE_TRITON_MOE"): self.use_triton = envs.ATOM_USE_TRITON_MOE else: