wire GUGU - act+quant fusion into triton decode

atom/model_ops/fused_moe_triton.py

@@ -463,3 +463,73 @@ def triton_kernel_fused_experts_a8w4_silu_gguu(
     )
 
     return output_tensor
+
+
+def triton_kernel_fused_experts_a8w4_silu(
+    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, GUGU (interleaved gate/up).
+
+    GEMM1 fuses SiLU(gate)*up + MXFP8 quant into the write-back (out_mx_quant),
+    valid only for the interleaved layout where gate/up pairs are in one tile:
+        MXFP8 quant -> GEMM1(a8w4, fused SiLU*up + mxfp8 out) -> GEMM2(a8w4).
+    Weights must be in the preshuffled a8w4 layout with w13 gate/up interleaved.
+    """
+    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)
+
+    interm_fp8, interm_scale = 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=True,
+        alpha=1.0,
+        limit=swiglu_limit,
+        swiglu_add_residual=False,
+        preshuffled=True,
+        out_mx_quant=True,
+    )
+
+    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

atom/model_ops/moe.py

@@ -811,7 +811,9 @@ 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") and not self.is_guinterleave:
+        # Triton decode supports both GGUU (separated, external act+quant) and
+        # GUGU (interleaved, fused act+quant in GEMM1 via out_mx_quant).
+        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
@@ -1088,11 +1090,22 @@ def process_weights_after_loading(self, layer):
             layer.w13_weight_preshuffled = moe_weight_decode_view(layer.w13_weight.data)
             layer.w2_weight_preshuffled = moe_weight_decode_view(layer.w2_weight.data)
 
+            # GUGU: interleave the w13 gate/up scale rows to match the interleaved
+            # decode weight (moe_shuffle_weight interleaves the weight but the a8w4
+            # decode scale is built from the raw snapshot). GGUU stays separated.
+            if self.is_guinterleave:
+                E_s, N_s, K_s = orig_w13_weight_scale.shape
+                orig_w13_weight_scale = (
+                    orig_w13_weight_scale.view(E_s, 2, N_s // 2, K_s)
+                    .permute(0, 2, 1, 3)
+                    .reshape(E_s, N_s, K_s)
+                    .contiguous()
+                )
+
             w13_scale_for_a8w4 = orig_w13_weight_scale.transpose(-2, -1)
             w2_scale_for_a8w4 = orig_w2_weight_scale.transpose(-2, -1)
 
             # Arch -> SWIZZLE_MX_SCALE label decision lives in aiter, not here.
-            # GGUU keeps gate/up separated, so no interleave on the decode scales.
             (
                 layer.w13_weight_scale_a8w4,
                 layer.w13_swizzle_layout_a8w4,
@@ -1146,8 +1159,10 @@ def apply(
         activation: ActivationType = ActivationType.Silu,
     ) -> torch.Tensor:
         if self.use_triton_decode and not get_forward_context().context.is_prefill:
-            # Triton decode is GGUU-only; GUGU uses the FlyDSL path.
+            # Triton decode: GUGU (interleaved) uses the fused act+quant GEMM1;
+            # GGUU (separated) uses the external act+quant path.
             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
@@ -1169,7 +1184,12 @@ def apply(
                 num_expert_group=num_expert_group,
                 topk_group=topk_group,
             )
-            return triton_kernel_fused_experts_a8w4_silu_gguu(
+            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,