diff --git a/atom/model_ops/fused_moe_triton.py b/atom/model_ops/fused_moe_triton.py index ac5d5b0f7..9f9ca475d 100644 --- a/atom/model_ops/fused_moe_triton.py +++ b/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 diff --git a/atom/model_ops/moe.py b/atom/model_ops/moe.py index bd943f549..74286dd46 100644 --- a/atom/model_ops/moe.py +++ b/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,