From dade718027db539537f2f82eb49cb2ce9847174f Mon Sep 17 00:00:00 2001 From: Aliasger Zaidy Date: Thu, 2 Jul 2026 03:32:23 +0000 Subject: [PATCH 1/5] Add triton MoE to DSv4 decode path --- atom/model_ops/fused_moe_triton.py | 96 +++++++++++++++++++++++++++++- atom/model_ops/moe.py | 92 ++++++++++++++++++++++++++++ atom/utils/envs.py | 1 + 3 files changed, 188 insertions(+), 1 deletion(-) diff --git a/atom/model_ops/fused_moe_triton.py b/atom/model_ops/fused_moe_triton.py index 507ba70b97..33956fa899 100644 --- a/atom/model_ops/fused_moe_triton.py +++ b/atom/model_ops/fused_moe_triton.py @@ -25,7 +25,7 @@ from aiter.ops.triton.utils._triton.arch_info import get_arch from atom.utils import envs -if envs.ATOM_USE_TRITON_GEMM or envs.ATOM_USE_TRITON_MOE: +if envs.ATOM_USE_TRITON_GEMM or envs.ATOM_USE_TRITON_MOE or envs.ATOM_USE_TRITON_MOE_DECODE: from aiter.ops.triton.moe.moe_routing.routing import routing from aiter.ops.triton.moe.moe_op_gemm_a8w4 import ( moe_gemm_a8w4, @@ -40,6 +40,7 @@ swizzle_scales as swizzle_scales_cdna4, ) from aiter.ops.triton.moe.quant_moe import downcast_to_static_fp8 + from aiter.ops.triton.moe.quant_moe import downcast_to_mxfp from atom.model_ops.moe import MoEActivationQuant @@ -387,3 +388,96 @@ def triton_kernel_fused_experts( output_tensor = output_tensor.view(M, K) return output_tensor + + +def triton_kernel_fused_experts_a8w4_silu( + output_tensor: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + routing_data, + gather_indx, + scatter_indx, + topk: int, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + w13_swizzle_layout, + w2_swizzle_layout, + a13_scale: torch.Tensor | None = None, + a2_scale: torch.Tensor | None = None, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + swiglu_limit: float = 10.0, + apply_router_weight_on_input: bool = False, + global_num_experts: int = -1, + expert_map: torch.Tensor | None = None, +) -> torch.Tensor: + """Decode-only A8W4 MoE for SiLU models (DSv4). + + Pipeline: MXFP8 quant -> GEMM1(a8w4) -> SiLU -> MXFP8 quant -> GEMM2(a8w4). + Weights must be in preshuffled layout (shuffle_weight_gfx1250). + Scales must be GFX1250-swizzled. + """ + assert hidden_states.ndim == 2 + assert hidden_states.dtype == torch.bfloat16 + + M, K = hidden_states.shape + N = w1.shape[-1] * 16 + half_N = N // 2 + + gammas = routing_data.gate_scal if routing_data else None + + x_fp8, x_scale = downcast_to_mxfp( + hidden_states, torch.float8_e4m3fn, axis=-1 + ) + + raw_intermediate = moe_gemm_a8w4( + x_fp8, + w1, + x_scale, + w13_scale, + a13_scale, + None, + w1_bias, + routing_data, + gather_indx=gather_indx, + gammas=gammas if apply_router_weight_on_input else None, + swizzle_mx_scale=w13_swizzle_layout, + apply_swiglu=False, + preshuffled=True, + ) + + raw_2d = raw_intermediate.view(M * topk, N) + intermediate = torch.empty( + (M * topk, half_N), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + fused_clamp_act_mul( + raw_2d, + out=intermediate, + swiglu_limit=swiglu_limit, + activation="silu", + dtype_quant=None, + ) + + interm_fp8, interm_scale = downcast_to_mxfp( + intermediate, torch.float8_e4m3fn, axis=-1 + ) + + output_tensor = moe_gemm_a8w4( + interm_fp8, + w2, + interm_scale, + w2_scale, + a2_scale, + None, + w2_bias, + routing_data, + scatter_indx=scatter_indx, + gammas=None if apply_router_weight_on_input else gammas, + swizzle_mx_scale=w2_swizzle_layout, + preshuffled=True, + ) + + return output_tensor diff --git a/atom/model_ops/moe.py b/atom/model_ops/moe.py index a92ce84491..a32937ac23 100644 --- a/atom/model_ops/moe.py +++ b/atom/model_ops/moe.py @@ -804,6 +804,10 @@ def __init__(self, quant_config: LayerQuantConfig, moe: FusedMoEConfig): gfx.startswith("gfx95") and envs.ATOM_USE_TRITON_GEMM ) self.act_quant = MoEActivationQuant.from_model_config(moe.a_quant_dtype) + if envs.is_set("ATOM_USE_TRITON_MOE_DECODE"): + self.use_triton_decode = envs.ATOM_USE_TRITON_MOE_DECODE + else: + self.use_triton_decode = False def create_weights( self, @@ -1012,6 +1016,10 @@ def process_weights_after_loading(self, layer): layer.w2_swizzle_layout = w2_swizzle_layout return + if self.use_triton_decode: + orig_w13_weight_scale = layer.w13_weight_scale.data.clone() + orig_w2_weight_scale = layer.w2_weight_scale.data.clone() + # shuffle weight layer.w13_weight.data = shuffle_weight( layer.w13_weight, @@ -1047,6 +1055,37 @@ def process_weights_after_loading(self, layer): layer.w13_weight_scale = atom_parameter(shuffled_w13_scale) layer.w2_weight_scale = atom_parameter(shuffled_w2_scale) + if self.use_triton_decode: + from aiter.ops.triton.moe.moe_op_gemm_a8w4 import ( + swizzle_scales as swizzle_scales_a8w4, + ) + + w13_u8 = layer.w13_weight.data + if w13_u8.dtype != torch.uint8: + w13_u8 = w13_u8.view(torch.uint8) + E_13, N_13, K_13 = w13_u8.shape + layer.w13_weight_preshuffled = w13_u8.view( + E_13, N_13 // 16, K_13 * 16 + ).transpose(-1, -2) + + w2_u8 = layer.w2_weight.data + if w2_u8.dtype != torch.uint8: + w2_u8 = w2_u8.view(torch.uint8) + E_2, N_2, K_2 = w2_u8.shape + layer.w2_weight_preshuffled = w2_u8.view( + E_2, N_2 // 16, K_2 * 16 + ).transpose(-1, -2) + + w13_scale_for_a8w4 = orig_w13_weight_scale.transpose(-2, -1) + w2_scale_for_a8w4 = orig_w2_weight_scale.transpose(-2, -1) + + layer.w13_weight_scale_a8w4, layer.w13_swizzle_layout_a8w4 = ( + swizzle_scales_a8w4(w13_scale_for_a8w4) + ) + layer.w2_weight_scale_a8w4, layer.w2_swizzle_layout_a8w4 = ( + swizzle_scales_a8w4(w2_scale_for_a8w4) + ) + def get_fused_moe_quant_config( self, layer: torch.nn.Module ) -> FusedMoEQuantConfig | None: @@ -1090,6 +1129,59 @@ def apply( fused_shared_experts_scoring_func: Optional[str] = None, activation: ActivationType = ActivationType.Silu, ) -> torch.Tensor: + if self.use_triton_decode and not get_forward_context().context.is_prefill: + from atom.model_ops.fused_moe_triton import ( + triton_kernel_fused_experts_a8w4_silu, + ) + from aiter.ops.triton.moe.moe_routing.routing import routing + + n_expts_act = top_k + + routing_data, gather_idx, scatter_idx = routing( + router_logits, + n_expts_act, + score_mode=scoring_func, + bias=( + e_score_correction_bias.to(torch.float32) + if e_score_correction_bias is not None + else None + ), + renorm=renormalize, + routed_scaling_factor=layer.routed_scaling_factor, + use_grouped_topk=use_grouped_topk, + num_expert_group=num_expert_group, + topk_group=topk_group, + ) + n_expts_act = routing_data.n_expts_act + + num_tokens, n_expts_tot = router_logits.shape + if global_num_experts > 0: + n_expts_tot = global_num_experts + + output = torch.empty_like(x) + return triton_kernel_fused_experts_a8w4_silu( + output, + x, + layer.w13_weight_preshuffled, + layer.w2_weight_preshuffled, + routing_data, + gather_idx, + scatter_idx, + topk=n_expts_act, + w13_scale=layer.w13_weight_scale_a8w4, + w2_scale=layer.w2_weight_scale_a8w4, + w13_swizzle_layout=layer.w13_swizzle_layout_a8w4, + w2_swizzle_layout=layer.w2_swizzle_layout_a8w4, + a13_scale=layer.w13_input_scale, + a2_scale=layer.w2_input_scale, + w1_bias=layer.w13_bias, + w2_bias=layer.w2_bias, + swiglu_limit=getattr(layer, "swiglu_limit", 0.0), + apply_router_weight_on_input=apply_router_weight_on_input, + global_num_experts=n_expts_tot, + expert_map=expert_map, + ) + if self.use_triton: from atom.model_ops.fused_moe_triton import ( triton_kernel_fused_experts, diff --git a/atom/utils/envs.py b/atom/utils/envs.py index 8baead10ce..7836327e81 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -42,6 +42,7 @@ os.getenv("ATOM_USE_TRITON_MLA_SHUFFLE_KV", "0") == "1" ), "ATOM_USE_TRITON_MOE": lambda: os.getenv("ATOM_USE_TRITON_MOE", "0") == "1", + "ATOM_USE_TRITON_MOE_DECODE": lambda: os.getenv("ATOM_USE_TRITON_MOE_DECODE", "0") == "1", "ATOM_MLA_PAGE_SIZE": lambda: int(os.getenv("ATOM_MLA_PAGE_SIZE", "1")), # --- Kernel Fusion Toggles --- # fused_compress_attn: switch between Triton (default historical) and a From 1800e3d79e563be4bbe4fcfd87d39254e4129314 Mon Sep 17 00:00:00 2001 From: Aliasger Zaidy Date: Sun, 5 Jul 2026 23:40:11 +0000 Subject: [PATCH 2/5] Enable moe+act+quant fusion for Triton DSv4 --- atom/model_ops/fused_moe_triton.py | 32 ++++++++---------------------- atom/model_ops/moe.py | 29 ++++++++++++++++++++++++++- 2 files changed, 36 insertions(+), 25 deletions(-) diff --git a/atom/model_ops/fused_moe_triton.py b/atom/model_ops/fused_moe_triton.py index 33956fa899..e33d67fc59 100644 --- a/atom/model_ops/fused_moe_triton.py +++ b/atom/model_ops/fused_moe_triton.py @@ -414,16 +414,14 @@ def triton_kernel_fused_experts_a8w4_silu( ) -> torch.Tensor: """Decode-only A8W4 MoE for SiLU models (DSv4). - Pipeline: MXFP8 quant -> GEMM1(a8w4) -> SiLU -> MXFP8 quant -> GEMM2(a8w4). - Weights must be in preshuffled layout (shuffle_weight_gfx1250). + Pipeline: MXFP8 quant -> GEMM1(a8w4, fused swiglu + mxfp8 out) -> GEMM2(a8w4). + Weights must be in preshuffled layout with w13 gate/up interleaved. Scales must be GFX1250-swizzled. """ assert hidden_states.ndim == 2 assert hidden_states.dtype == torch.bfloat16 M, K = hidden_states.shape - N = w1.shape[-1] * 16 - half_N = N // 2 gammas = routing_data.gate_scal if routing_data else None @@ -431,7 +429,7 @@ def triton_kernel_fused_experts_a8w4_silu( hidden_states, torch.float8_e4m3fn, axis=-1 ) - raw_intermediate = moe_gemm_a8w4( + interm_fp8, interm_scale = moe_gemm_a8w4( x_fp8, w1, x_scale, @@ -443,26 +441,12 @@ def triton_kernel_fused_experts_a8w4_silu( gather_indx=gather_indx, gammas=gammas if apply_router_weight_on_input else None, swizzle_mx_scale=w13_swizzle_layout, - apply_swiglu=False, + apply_swiglu=True, + alpha=1.0, + limit=swiglu_limit, + swiglu_add_residual=False, preshuffled=True, - ) - - raw_2d = raw_intermediate.view(M * topk, N) - intermediate = torch.empty( - (M * topk, half_N), - device=hidden_states.device, - dtype=hidden_states.dtype, - ) - fused_clamp_act_mul( - raw_2d, - out=intermediate, - swiglu_limit=swiglu_limit, - activation="silu", - dtype_quant=None, - ) - - interm_fp8, interm_scale = downcast_to_mxfp( - intermediate, torch.float8_e4m3fn, axis=-1 + out_mx_quant=True, ) output_tensor = moe_gemm_a8w4( diff --git a/atom/model_ops/moe.py b/atom/model_ops/moe.py index a32937ac23..025ba25bfd 100644 --- a/atom/model_ops/moe.py +++ b/atom/model_ops/moe.py @@ -1017,6 +1017,33 @@ def process_weights_after_loading(self, layer): return if self.use_triton_decode: + # Interleave w13 gate/up columns in-place BEFORE shuffle so + # both prefill (FlydSL, GateMode.INTERLEAVE) and decode (gluon, + # _swiglu) share one copy. [gate|up] → [g0,u0,g1,u1,...] + w13_raw = layer.w13_weight.data + w13_dtype = w13_raw.dtype + if w13_dtype != torch.uint8: + w13_raw = w13_raw.view(torch.uint8) + E_13, N_13, K_13 = w13_raw.shape + N_half = N_13 // 2 + layer.w13_weight.data = ( + w13_raw.view(E_13, 2, N_half, K_13) + .permute(0, 2, 1, 3) + .reshape(E_13, N_13, K_13) + .contiguous() + .view(w13_dtype) + ) + # Interleave w13 scales along N to match weight interleave. + # Scale shape is (E, N, K_scale). + w13_sc = layer.w13_weight_scale.data + E_s, N_s, K_s = w13_sc.shape + N_s_half = N_s // 2 + layer.w13_weight_scale.data = ( + w13_sc.view(E_s, 2, N_s_half, K_s) + .permute(0, 2, 1, 3) + .reshape(E_s, N_s, K_s) + .contiguous() + ) orig_w13_weight_scale = layer.w13_weight_scale.data.clone() orig_w2_weight_scale = layer.w2_weight_scale.data.clone() @@ -1313,7 +1340,7 @@ def apply( moe_extra_args = { "gate_mode": ( GateMode.INTERLEAVE.value - if self.is_guinterleave + if (self.is_guinterleave or self.use_triton_decode) else GateMode.SEPARATED.value ), "swiglu_limit": getattr(layer, "swiglu_limit", 0.0), From 02ef699f4a1ec02151497564094d4c9d2df03318 Mon Sep 17 00:00:00 2001 From: Aliasger Zaidy Date: Mon, 6 Jul 2026 04:20:06 +0000 Subject: [PATCH 3/5] Conform to new shuffle API --- atom/model_ops/fused_moe_triton.py | 14 ++++++++------ atom/model_ops/moe.py | 14 +++++--------- 2 files changed, 13 insertions(+), 15 deletions(-) diff --git a/atom/model_ops/fused_moe_triton.py b/atom/model_ops/fused_moe_triton.py index e33d67fc59..008a0864af 100644 --- a/atom/model_ops/fused_moe_triton.py +++ b/atom/model_ops/fused_moe_triton.py @@ -29,7 +29,6 @@ from aiter.ops.triton.moe.moe_routing.routing import routing from aiter.ops.triton.moe.moe_op_gemm_a8w4 import ( moe_gemm_a8w4, - swizzle_scales as swizzle_scales_a8w4, ) from aiter.ops.triton.moe.moe_op_gemm_a16w4 import ( moe_gemm_a16w4, @@ -37,8 +36,8 @@ from aiter.ops.triton.moe.moe_op_gemm_a4w4 import ( moe_gemm_a4w4, mxfp4_quant, - swizzle_scales as swizzle_scales_cdna4, ) + from aiter.ops.triton.utils.shuffle import shuffle_scale_moe from aiter.ops.triton.moe.quant_moe import downcast_to_static_fp8 from aiter.ops.triton.moe.quant_moe import downcast_to_mxfp @@ -51,11 +50,14 @@ def _swizzle_scales_for_kernel(scale, act_quant: MoEActivationQuant): FP8 (a8w4): arch-agnostic swizzle (CDNA4 on gfx950, GFX1250 on gfx1250). BF16/FP4 (a16w4/a4w4): CDNA4 swizzle on gfx942/gfx950, no swizzle elsewhere. """ + arch = get_arch() if act_quant == MoEActivationQuant.FP8: - return swizzle_scales_a8w4(scale) - # TODO: move arch dispatch into aiter's a4w4/a16w4 swizzle_scales (like a8w4) - if get_arch() in ("gfx942", "gfx950"): - return swizzle_scales_cdna4(scale), "CDNA4_SCALE" + layout = "GFX1250_SCALE" if arch == "gfx1250" else ( + "CDNA4_SCALE" if arch == "gfx950" else None + ) + return shuffle_scale_moe(scale), layout + if arch in ("gfx942", "gfx950"): + return shuffle_scale_moe(scale), "CDNA4_SCALE" return scale, None diff --git a/atom/model_ops/moe.py b/atom/model_ops/moe.py index 025ba25bfd..b2ae6def13 100644 --- a/atom/model_ops/moe.py +++ b/atom/model_ops/moe.py @@ -1083,9 +1083,7 @@ def process_weights_after_loading(self, layer): layer.w2_weight_scale = atom_parameter(shuffled_w2_scale) if self.use_triton_decode: - from aiter.ops.triton.moe.moe_op_gemm_a8w4 import ( - swizzle_scales as swizzle_scales_a8w4, - ) + from aiter.ops.triton.utils.shuffle import shuffle_scale_moe w13_u8 = layer.w13_weight.data if w13_u8.dtype != torch.uint8: @@ -1106,12 +1104,10 @@ def process_weights_after_loading(self, layer): w13_scale_for_a8w4 = orig_w13_weight_scale.transpose(-2, -1) w2_scale_for_a8w4 = orig_w2_weight_scale.transpose(-2, -1) - layer.w13_weight_scale_a8w4, layer.w13_swizzle_layout_a8w4 = ( - swizzle_scales_a8w4(w13_scale_for_a8w4) - ) - layer.w2_weight_scale_a8w4, layer.w2_swizzle_layout_a8w4 = ( - swizzle_scales_a8w4(w2_scale_for_a8w4) - ) + layer.w13_weight_scale_a8w4 = shuffle_scale_moe(w13_scale_for_a8w4) + layer.w13_swizzle_layout_a8w4 = "GFX1250_SCALE" + layer.w2_weight_scale_a8w4 = shuffle_scale_moe(w2_scale_for_a8w4) + layer.w2_swizzle_layout_a8w4 = "GFX1250_SCALE" def get_fused_moe_quant_config( self, layer: torch.nn.Module From 074a753ab460a4745f2f8c95705629b3ecef07ac Mon Sep 17 00:00:00 2001 From: satya Date: Wed, 8 Jul 2026 16:28:32 +0000 Subject: [PATCH 4/5] Clean up triton MoE decode integration move arch stuff to aiter --- atom/model_ops/fused_moe_triton.py | 54 ++++++++++-------------------- atom/model_ops/moe.py | 29 +++++++--------- 2 files changed, 31 insertions(+), 52 deletions(-) diff --git a/atom/model_ops/fused_moe_triton.py b/atom/model_ops/fused_moe_triton.py index 008a0864af..f11be07eb1 100644 --- a/atom/model_ops/fused_moe_triton.py +++ b/atom/model_ops/fused_moe_triton.py @@ -25,7 +25,11 @@ from aiter.ops.triton.utils._triton.arch_info import get_arch from atom.utils import envs -if envs.ATOM_USE_TRITON_GEMM or envs.ATOM_USE_TRITON_MOE or envs.ATOM_USE_TRITON_MOE_DECODE: +if ( + envs.ATOM_USE_TRITON_GEMM + or envs.ATOM_USE_TRITON_MOE + or envs.ATOM_USE_TRITON_MOE_DECODE +): from aiter.ops.triton.moe.moe_routing.routing import routing from aiter.ops.triton.moe.moe_op_gemm_a8w4 import ( moe_gemm_a8w4, @@ -44,23 +48,6 @@ from atom.model_ops.moe import MoEActivationQuant -def _swizzle_scales_for_kernel(scale, act_quant: MoEActivationQuant): - """Swizzle MoE weight scales for the kernel that act_quant selects. - - FP8 (a8w4): arch-agnostic swizzle (CDNA4 on gfx950, GFX1250 on gfx1250). - BF16/FP4 (a16w4/a4w4): CDNA4 swizzle on gfx942/gfx950, no swizzle elsewhere. - """ - arch = get_arch() - if act_quant == MoEActivationQuant.FP8: - layout = "GFX1250_SCALE" if arch == "gfx1250" else ( - "CDNA4_SCALE" if arch == "gfx950" else None - ) - return shuffle_scale_moe(scale), layout - if arch in ("gfx942", "gfx950"): - return shuffle_scale_moe(scale), "CDNA4_SCALE" - return scale, None - - def _swizzle_mxfp4( w1, w1_scale, @@ -72,9 +59,12 @@ def _swizzle_mxfp4( N_2, K_2, TP=1, - act_quant: MoEActivationQuant = MoEActivationQuant.BF16, ): - """Weight swizzle for mxfp4 moe, used for aiter triton mxfp4 moe kernels.""" + """Weight swizzle for mxfp4 moe, used for aiter triton mxfp4 moe kernels. + + The arch -> SWIZZLE_MX_SCALE label decision lives in aiter + (``shuffle_scale_moe(..., return_layout=True)``), so this stays arch-agnostic. + """ assert envs.ATOM_USE_TRITON_GEMM or envs.ATOM_USE_TRITON_MOE # Transposing for expected layout of aiter triton kernels @@ -84,15 +74,15 @@ def _swizzle_mxfp4( w2_scale_triton_layout = w2_scale.transpose(-2, -1) if N_1 % 32 == 0 and K_1 % (32 * 8) == 0: - w1_scale_triton_layout, w1_swizzle_layout = _swizzle_scales_for_kernel( - w1_scale_triton_layout, act_quant + w1_scale_triton_layout, w1_swizzle_layout = shuffle_scale_moe( + w1_scale_triton_layout, return_layout=True ) else: w1_swizzle_layout = None if N_2 % 32 == 0 and K_2 % (32 * 8) == 0: - w2_scale_triton_layout, w2_swizzle_layout = _swizzle_scales_for_kernel( - w2_scale_triton_layout, act_quant + w2_scale_triton_layout, w2_swizzle_layout = shuffle_scale_moe( + w2_scale_triton_layout, return_layout=True ) else: w2_swizzle_layout = None @@ -393,14 +383,12 @@ def triton_kernel_fused_experts( def triton_kernel_fused_experts_a8w4_silu( - output_tensor: torch.Tensor, hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, routing_data, gather_indx, scatter_indx, - topk: int, w13_scale: torch.Tensor, w2_scale: torch.Tensor, w13_swizzle_layout, @@ -411,25 +399,19 @@ def triton_kernel_fused_experts_a8w4_silu( w2_bias: torch.Tensor | None = None, swiglu_limit: float = 10.0, apply_router_weight_on_input: bool = False, - global_num_experts: int = -1, - expert_map: torch.Tensor | None = None, ) -> torch.Tensor: """Decode-only A8W4 MoE for SiLU models (DSv4). - Pipeline: MXFP8 quant -> GEMM1(a8w4, fused swiglu + mxfp8 out) -> GEMM2(a8w4). - Weights must be in preshuffled layout with w13 gate/up interleaved. - Scales must be GFX1250-swizzled. + Pipeline: MXFP8 quant -> GEMM1(a8w4, fused SiLU(gate)*up + mxfp8 out) -> + GEMM2(a8w4). Weights must be in the preshuffled a8w4 layout with w13 gate/up + interleaved; scales carry the swizzle label in w13/w2_swizzle_layout. """ assert hidden_states.ndim == 2 assert hidden_states.dtype == torch.bfloat16 - M, K = hidden_states.shape - gammas = routing_data.gate_scal if routing_data else None - x_fp8, x_scale = downcast_to_mxfp( - hidden_states, torch.float8_e4m3fn, axis=-1 - ) + x_fp8, x_scale = downcast_to_mxfp(hidden_states, torch.float8_e4m3fn, axis=-1) interm_fp8, interm_scale = moe_gemm_a8w4( x_fp8, diff --git a/atom/model_ops/moe.py b/atom/model_ops/moe.py index ff2f3ebab8..529f3f8836 100644 --- a/atom/model_ops/moe.py +++ b/atom/model_ops/moe.py @@ -1012,7 +1012,6 @@ def process_weights_after_loading(self, layer): self.hidden_size, # N_2, self.intermediate_size, # K_2, atom_config.tensor_parallel_size, - act_quant=self.act_quant, ) del layer.w13_weight del layer.w2_weight @@ -1095,6 +1094,10 @@ def process_weights_after_loading(self, layer): if self.use_triton_decode: from aiter.ops.triton.utils.shuffle import shuffle_scale_moe + # FlyDSL (prefill) and Triton/gluon (decode) are two different + # layouts. layer.w13_weight / w13_weight_scale stay in the FlyDSL + # layout shuffled above; the *_preshuffled / *_a8w4 tensors below + # hold the separate Triton gluon a8w4 layout for the decode kernel. w13_u8 = layer.w13_weight.data if w13_u8.dtype != torch.uint8: w13_u8 = w13_u8.view(torch.uint8) @@ -1114,10 +1117,15 @@ def process_weights_after_loading(self, layer): w13_scale_for_a8w4 = orig_w13_weight_scale.transpose(-2, -1) w2_scale_for_a8w4 = orig_w2_weight_scale.transpose(-2, -1) - layer.w13_weight_scale_a8w4 = shuffle_scale_moe(w13_scale_for_a8w4) - layer.w13_swizzle_layout_a8w4 = "GFX1250_SCALE" - layer.w2_weight_scale_a8w4 = shuffle_scale_moe(w2_scale_for_a8w4) - layer.w2_swizzle_layout_a8w4 = "GFX1250_SCALE" + # Arch -> SWIZZLE_MX_SCALE label decision lives in aiter, not here. + ( + layer.w13_weight_scale_a8w4, + layer.w13_swizzle_layout_a8w4, + ) = shuffle_scale_moe(w13_scale_for_a8w4, return_layout=True) + ( + layer.w2_weight_scale_a8w4, + layer.w2_swizzle_layout_a8w4, + ) = shuffle_scale_moe(w2_scale_for_a8w4, return_layout=True) def get_fused_moe_quant_config( self, layer: torch.nn.Module @@ -1185,22 +1193,13 @@ def apply( num_expert_group=num_expert_group, topk_group=topk_group, ) - n_expts_act = routing_data.n_expts_act - - num_tokens, n_expts_tot = router_logits.shape - if global_num_experts > 0: - n_expts_tot = global_num_experts - - output = torch.empty_like(x) return triton_kernel_fused_experts_a8w4_silu( - output, x, layer.w13_weight_preshuffled, layer.w2_weight_preshuffled, routing_data, gather_idx, scatter_idx, - topk=n_expts_act, w13_scale=layer.w13_weight_scale_a8w4, w2_scale=layer.w2_weight_scale_a8w4, w13_swizzle_layout=layer.w13_swizzle_layout_a8w4, @@ -1211,8 +1210,6 @@ def apply( w2_bias=layer.w2_bias, swiglu_limit=getattr(layer, "swiglu_limit", 0.0), apply_router_weight_on_input=apply_router_weight_on_input, - global_num_experts=n_expts_tot, - expert_map=expert_map, ) if self.use_triton: From d5d239a8f0b1579cba359313aea8deda3454ab06 Mon Sep 17 00:00:00 2001 From: satya Date: Wed, 8 Jul 2026 16:28:32 +0000 Subject: [PATCH 5/5] WIP: dsv4 MoE decode integration tweaks --- atom/model_ops/fused_moe_triton.py | 85 +++++++++++++++++++++++++- atom/model_ops/moe.py | 98 +++++++++++++++++------------- 2 files changed, 141 insertions(+), 42 deletions(-) diff --git a/atom/model_ops/fused_moe_triton.py b/atom/model_ops/fused_moe_triton.py index f11be07eb1..5456fc4cec 100644 --- a/atom/model_ops/fused_moe_triton.py +++ b/atom/model_ops/fused_moe_triton.py @@ -400,7 +400,7 @@ def triton_kernel_fused_experts_a8w4_silu( swiglu_limit: float = 10.0, apply_router_weight_on_input: bool = False, ) -> torch.Tensor: - """Decode-only A8W4 MoE for SiLU models (DSv4). + """Decode-only A8W4 MoE for SiLU models (DSv4), GUGU (interleaved gate/up). Pipeline: MXFP8 quant -> GEMM1(a8w4, fused SiLU(gate)*up + mxfp8 out) -> GEMM2(a8w4). Weights must be in the preshuffled a8w4 layout with w13 gate/up @@ -449,3 +449,86 @@ def triton_kernel_fused_experts_a8w4_silu( ) return output_tensor + + +def triton_kernel_fused_experts_a8w4_silu_gguu( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + routing_data, + gather_indx, + scatter_indx, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + w13_swizzle_layout, + w2_swizzle_layout, + a13_scale: torch.Tensor | None = None, + a2_scale: torch.Tensor | None = None, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + swiglu_limit: float = 10.0, + apply_router_weight_on_input: bool = False, +) -> torch.Tensor: + """Decode-only A8W4 MoE for SiLU models, GGUU (separated ``[gate|up]``). + + GGUU keeps gate and up as contiguous halves, so the per-block SiLU cannot be + fused into GEMM1's write-back (a tile spans only gate *or* only up). The + activation and quant therefore run as a separate step: + + MXFP8 quant -> GEMM1(a8w4, no swiglu, bf16 [gate|up]) -> + fused_clamp_act_mul(SiLU(gate)*up on the halves) -> + MXFP8 quant -> GEMM2(a8w4). + + The intermediate is re-quantized with ``downcast_to_mxfp`` (same op as the x + path) so GEMM2 sees the identical activation-scale format. Weights are in the + preshuffled a8w4 layout with w13 gate/up separated. + """ + assert hidden_states.ndim == 2 + assert hidden_states.dtype == torch.bfloat16 + + gammas = routing_data.gate_scal if routing_data else None + + x_fp8, x_scale = downcast_to_mxfp(hidden_states, torch.float8_e4m3fn, axis=-1) + + # GEMM1: raw bf16 [gate|up] output; no fused activation for the separated layout. + interm = moe_gemm_a8w4( + x_fp8, + w1, + x_scale, + w13_scale, + a13_scale, + None, + w1_bias, + routing_data, + gather_indx=gather_indx, + gammas=gammas if apply_router_weight_on_input else None, + swizzle_mx_scale=w13_swizzle_layout, + apply_swiglu=False, + out_dtype=torch.bfloat16, + preshuffled=True, + ) + + # Standalone SiLU(gate)*up over the contiguous halves, then MXFP8 quant. + interm_act = fused_clamp_act_mul( + interm, swiglu_limit=swiglu_limit, activation="silu" + ) + interm_fp8, interm_scale = downcast_to_mxfp( + interm_act, torch.float8_e4m3fn, axis=-1 + ) + + output_tensor = moe_gemm_a8w4( + interm_fp8, + w2, + interm_scale, + w2_scale, + a2_scale, + None, + w2_bias, + routing_data, + scatter_indx=scatter_indx, + gammas=None if apply_router_weight_on_input else gammas, + swizzle_mx_scale=w2_swizzle_layout, + preshuffled=True, + ) + + return output_tensor diff --git a/atom/model_ops/moe.py b/atom/model_ops/moe.py index 529f3f8836..e23737223d 100644 --- a/atom/model_ops/moe.py +++ b/atom/model_ops/moe.py @@ -800,13 +800,6 @@ 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." - ) if envs.is_set("ATOM_USE_TRITON_MOE"): self.use_triton = envs.ATOM_USE_TRITON_MOE else: @@ -818,6 +811,15 @@ def __init__(self, quant_config: LayerQuantConfig, moe: FusedMoEConfig): self.use_triton_decode = envs.ATOM_USE_TRITON_MOE_DECODE else: self.use_triton_decode = False + # gfx1250's shuffle_weight cannot apply is_guinterleave. The triton decode + # path instead produces GUGU via an explicit pre-interleave (is_guinterleave + # True) or leaves gate/up separated for GGUU (False), so only reject + # is_guinterleave up front when that path is not enabled. + if self.is_gfx1250 and self.is_guinterleave and not self.use_triton_decode: + raise NotImplementedError( + "gfx1250 MoE only supports is_guinterleave=False without " + "ATOM_USE_TRITON_MOE_DECODE; unset ATOM_MOE_GU_ITLV." + ) def create_weights( self, @@ -1026,45 +1028,51 @@ def process_weights_after_loading(self, layer): return if self.use_triton_decode: - # Interleave w13 gate/up columns in-place BEFORE shuffle so - # both prefill (FlydSL, GateMode.INTERLEAVE) and decode (gluon, - # _swiglu) share one copy. [gate|up] → [g0,u0,g1,u1,...] - w13_raw = layer.w13_weight.data - w13_dtype = w13_raw.dtype - if w13_dtype != torch.uint8: - w13_raw = w13_raw.view(torch.uint8) - E_13, N_13, K_13 = w13_raw.shape - N_half = N_13 // 2 - layer.w13_weight.data = ( - w13_raw.view(E_13, 2, N_half, K_13) - .permute(0, 2, 1, 3) - .reshape(E_13, N_13, K_13) - .contiguous() - .view(w13_dtype) - ) - # Interleave w13 scales along N to match weight interleave. - # Scale shape is (E, N, K_scale). - w13_sc = layer.w13_weight_scale.data - E_s, N_s, K_s = w13_sc.shape - N_s_half = N_s // 2 - layer.w13_weight_scale.data = ( - w13_sc.view(E_s, 2, N_s_half, K_s) - .permute(0, 2, 1, 3) - .reshape(E_s, N_s, K_s) - .contiguous() - ) + if self.is_guinterleave: + # GUGU: interleave w13 gate/up columns + scales in-place BEFORE + # shuffle. gfx1250 shuffle_weight can't apply is_guinterleave, so + # do the [gate|up] → [g0,u0,g1,u1,...] reorder explicitly here. + w13_raw = layer.w13_weight.data + w13_dtype = w13_raw.dtype + if w13_dtype != torch.uint8: + w13_raw = w13_raw.view(torch.uint8) + E_13, N_13, K_13 = w13_raw.shape + N_half = N_13 // 2 + layer.w13_weight.data = ( + w13_raw.view(E_13, 2, N_half, K_13) + .permute(0, 2, 1, 3) + .reshape(E_13, N_13, K_13) + .contiguous() + .view(w13_dtype) + ) + # Interleave w13 scales along N to match the weight interleave. + # Scale shape is (E, N, K_scale). + w13_sc = layer.w13_weight_scale.data + E_s, N_s, K_s = w13_sc.shape + N_s_half = N_s // 2 + layer.w13_weight_scale.data = ( + w13_sc.view(E_s, 2, N_s_half, K_s) + .permute(0, 2, 1, 3) + .reshape(E_s, N_s, K_s) + .contiguous() + ) + # else GGUU: leave gate/up as contiguous halves [gate|up]; the decode + # kernel consumes the separated layout via the unfused path. orig_w13_weight_scale = layer.w13_weight_scale.data.clone() orig_w2_weight_scale = layer.w2_weight_scale.data.clone() - # shuffle weight + # shuffle weight. On the decode path the gate/up ordering is already set + # (GUGU pre-interleaved above, or GGUU left separated) and gfx1250 + # shuffle_weight rejects is_guinterleave, so pass False there. + shuffle_guinterleave = False if self.use_triton_decode else self.is_guinterleave layer.w13_weight.data = shuffle_weight( layer.w13_weight, - is_guinterleave=self.is_guinterleave, + is_guinterleave=shuffle_guinterleave, gate_up=True, ) layer.w2_weight.data = shuffle_weight( layer.w2_weight, - is_guinterleave=self.is_guinterleave, + is_guinterleave=shuffle_guinterleave, gate_up=False, ) layer.w13_weight.is_shuffled = True @@ -1079,13 +1087,13 @@ def process_weights_after_loading(self, layer): shuffled_w13_scale = moe_shuffle_scale( w13_scale_2d, self.num_experts, - is_guinterleave=self.is_guinterleave, + is_guinterleave=shuffle_guinterleave, gate_up=True, ) shuffled_w2_scale = moe_shuffle_scale( w2_scale_2d, self.num_experts, - is_guinterleave=self.is_guinterleave, + is_guinterleave=shuffle_guinterleave, gate_up=False, ) layer.w13_weight_scale = atom_parameter(shuffled_w13_scale) @@ -1173,6 +1181,7 @@ def apply( if self.use_triton_decode and not get_forward_context().context.is_prefill: from atom.model_ops.fused_moe_triton import ( triton_kernel_fused_experts_a8w4_silu, + triton_kernel_fused_experts_a8w4_silu_gguu, ) from aiter.ops.triton.moe.moe_routing.routing import routing @@ -1193,7 +1202,14 @@ def apply( num_expert_group=num_expert_group, topk_group=topk_group, ) - return triton_kernel_fused_experts_a8w4_silu( + # GUGU uses the fused SiLU+quant GEMM1; GGUU (separated gate|up) can't + # fuse the per-block activation, so it takes the unfused path. + decode_experts = ( + triton_kernel_fused_experts_a8w4_silu + if self.is_guinterleave + else triton_kernel_fused_experts_a8w4_silu_gguu + ) + return decode_experts( x, layer.w13_weight_preshuffled, layer.w2_weight_preshuffled, @@ -1343,7 +1359,7 @@ def apply( moe_extra_args = { "gate_mode": ( GateMode.INTERLEAVE.value - if (self.is_guinterleave or self.use_triton_decode) + if self.is_guinterleave else GateMode.SEPARATED.value ), "swiglu_limit": getattr(layer, "swiglu_limit", 0.0),