wire GUGU - act+quant fusion into triton decode #1570
Open
Boss2002n wants to merge 1 commit into
Open
Conversation
Adds the GUGU (interleaved) fused decode path alongside the existing GGUU: - triton_kernel_fused_experts_a8w4_silu: GEMM1 fuses SiLU*up + MXFP8 quant (apply_swiglu + out_mx_quant), consuming the interleaved preshuffled weights moe_shuffle_weight already produces on gfx1250. - apply() dispatches is_guinterleave -> GUGU fused decode, else GGUU external. - __init__: allow GUGU with ATOM_USE_TRITON_MOE_DECODE (was GGUU-only). - decode a8w4 scale: interleave the w13 gate/up rows for GUGU to match the weight. Requires the aiter GUGU fusion (satya/gugu_act_quant: out_mx_quant MXFP8 emit in the gfx1250 gluon decode kernel). DECODE-ONLY: gfx1250 prefill (FlyDSL) is GGUU-only, so GUGU prefill on gfx1250 is out of scope here.
Contributor
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
Boss2002n
changed the title
Contributor
There was a problem hiding this comment.
Pull request overview
This PR extends the Triton decode path for MXFP4/A8W4 MoE to support the GUGU (gate/up interleaved) layout by using a fused act+quant GEMM1, alongside the existing GGUU (separated) decode path.
Changes:
- Enable Triton MoE decode for both GGUU and GUGU layouts (env-var gated).
- Interleave the
w13decode scales when GU-interleaving is enabled to match the decode weight layout. - Add a new decode-only Triton kernel (
triton_kernel_fused_experts_a8w4_silu) for the GUGU fused act+quant path and dispatch to it frommoe.py.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| atom/model_ops/moe.py | Enables Triton decode for GUGU and adjusts decode-scale preparation + kernel dispatch. |
| atom/model_ops/fused_moe_triton.py | Adds a GUGU decode-only A8W4 SiLU kernel using GEMM1 fused SiLU*up + MXFP8 out quant. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+1096
to
+1103
| 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() | ||
| ) |
| @@ -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: | |||
Comment on lines
+1187
to
+1191
| decode_experts = ( | ||
| triton_kernel_fused_experts_a8w4_silu | ||
| if self.is_guinterleave | ||
| else triton_kernel_fused_experts_a8w4_silu_gguu | ||
| ) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.