From be92fdd1213dc91a599611ab7c5d512a49e03809 Mon Sep 17 00:00:00 2001 From: Leon Ling Date: Thu, 25 Jun 2026 02:24:23 +0000 Subject: [PATCH 1/5] Support SwiGLU-OAI in dense shared-expert GEMM (MiniMax-M3) Mxfp4MoEMethod._apply_shared_experts_dense hard-asserted the SiLU activation path, so MiniMax-M3 (ActivationType.Swiglu with fused shared experts) crashed with "dense shared-expert GEMM only supports the SiLU activation path". MiniMax-M3 uses SwiGLU-OAI (gate*sigmoid(alpha*gate)*(up+beta)) and does not interleave gate/up, so the dense GEMM output is split [gate|up] -- exactly what swiglu_oai_split consumes. The dense shared expert now mirrors MiniMaxM3MLP.forward and the routed experts' alpha / swiglu_add_residual=True path. - moe.py: drop the assert; branch the activation step. SiLU keeps fused_clamp_act_mul (DeepSeek, unchanged); SwiGLU uses swiglu_oai_split with alpha/beta/limit read from the layer. - minimax_m3.py: stash swiglu_alpha/swiglu_beta on self.experts (from config, defaults 1.702/1.0) next to swiglu_limit. - tests: numerical test that reuses the same kernel GEMM and varies only the activation, isolating the fix from mxfp4/bf16 GEMM precision. Verified on MI350X: fixed path matches the SwiGLU-OAI reference exactly, old SiLU behaviour diverged ~20%. Co-Authored-By: Claude Opus 4.8 (1M context) --- atom/model_ops/moe.py | 45 +++-- atom/models/minimax_m3.py | 6 + tests/test_mxfp4_shared_experts_swiglu.py | 222 ++++++++++++++++++++++ 3 files changed, 259 insertions(+), 14 deletions(-) create mode 100644 tests/test_mxfp4_shared_experts_swiglu.py diff --git a/atom/model_ops/moe.py b/atom/model_ops/moe.py index 37b4bc8b9e..3219d104c1 100644 --- a/atom/model_ops/moe.py +++ b/atom/model_ops/moe.py @@ -1303,15 +1303,21 @@ def _apply_shared_experts_dense(self, layer, x, activation): from aiter.ops.triton.fusions.fused_clamp_act_mul import fused_clamp_act_mul from aiter.ops.triton.gemm.basic.gemm_a16wfp4 import gemm_a16wfp4 - # The dense shared-expert GEMM only implements the SiLU activation - # path; SwiGLU models have no fused shared experts, so this assert - # documents the supported scope. - assert ( - activation != ActivationType.Swiglu - ), "dense shared-expert GEMM only supports the SiLU activation path" + from atom.model_ops.swiglu_oai import swiglu_oai_split + + # Two activation flavours are supported, matching the routed experts: + # * SiLU (DeepSeek): silu(gate) * clamp(up), split [gate | up] layout. + # * SwiGLU-OAI (MiniMax-M3 / gpt-oss): gate * sigmoid(alpha*gate) * + # (up + beta) with optional clamp. MiniMax-M3 does not interleave + # gate/up, so the dense GEMM output is split [gate | up] - exactly + # what swiglu_oai_split consumes (mirrors MiniMaxM3MLP.forward and + # the routed swiglu_add_residual=True / alpha path). + is_swiglu = activation == ActivationType.Swiglu M = x.shape[0] swiglu_limit = getattr(layer, "swiglu_limit", 0.0) + swiglu_alpha = getattr(layer, "swiglu_alpha", 1.702) + swiglu_beta = getattr(layer, "swiglu_beta", 1.0) use_a4w4 = self.act_quant == MoEActivationQuant.FP4 if use_a4w4: @@ -1337,14 +1343,25 @@ def _shared_expert_gemm(act, weight, weight_scale): if shared_w13_bias is not None: gate_up = gate_up + shared_w13_bias[e] half_n = gate_up.shape[-1] // 2 - intermediate = torch.empty((M, half_n), device=x.device, dtype=x.dtype) - fused_clamp_act_mul( - gate_up, - out=intermediate, - swiglu_limit=swiglu_limit, - activation="silu", - dtype_quant=None, - ) + if is_swiglu: + intermediate = swiglu_oai_split( + gate_up, + alpha=swiglu_alpha, + beta=swiglu_beta, + limit=swiglu_limit if swiglu_limit > 0 else None, + out_dtype=x.dtype, + ) + else: + intermediate = torch.empty( + (M, half_n), device=x.device, dtype=x.dtype + ) + fused_clamp_act_mul( + gate_up, + out=intermediate, + swiglu_limit=swiglu_limit, + activation="silu", + dtype_quant=None, + ) out_e = _shared_expert_gemm( intermediate, layer.shared_w2_weight[e], diff --git a/atom/models/minimax_m3.py b/atom/models/minimax_m3.py index 31efec5ef9..9405f076fa 100644 --- a/atom/models/minimax_m3.py +++ b/atom/models/minimax_m3.py @@ -291,6 +291,12 @@ def __init__( # padded intermediate avoids backend pad-skip precision issues. self.experts.quant_method.intermediate_pad = 0 self.experts.swiglu_limit = getattr(config, "swiglu_limit", 7.0) + # SwiGLU-OAI params for the standalone dense shared-expert GEMM + # (Mxfp4MoEMethod._apply_shared_experts_dense). The routed experts use + # alpha (default 1.702) and swiglu_add_residual=True (beta == 1.0); the + # dense shared experts must match. + self.experts.swiglu_alpha = getattr(config, "swiglu_alpha", 1.702) + self.experts.swiglu_beta = getattr(config, "swiglu_beta", 1.0) self.fuse_shared_experts = ( getattr(self.experts, "num_fused_shared_experts", 0) > 0 ) diff --git a/tests/test_mxfp4_shared_experts_swiglu.py b/tests/test_mxfp4_shared_experts_swiglu.py new file mode 100644 index 0000000000..dd8c8f5923 --- /dev/null +++ b/tests/test_mxfp4_shared_experts_swiglu.py @@ -0,0 +1,222 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Numerical test for the dense shared-expert path with SwiGLU-OAI activation. + +Regression: ``Mxfp4MoEMethod._apply_shared_experts_dense`` hard-asserted the +SiLU activation path, so MiniMax-M3 (``ActivationType.Swiglu`` *with* fused +shared experts) crashed with:: + + AssertionError: dense shared-expert GEMM only supports the SiLU activation path + +MiniMax-M3 does not interleave gate/up weights, so the dense GEMM output is +split ``[gate | up]`` — exactly what ``swiglu_oai_split`` consumes. The dense +shared expert must therefore replicate ``MiniMaxM3MLP.forward``: +gate_up GEMM -> swiglu_oai_split -> down GEMM, with the SwiGLU-OAI math +``gate * sigmoid(alpha*gate) * (up + beta)`` (alpha=1.702, beta=1.0), not SiLU. + +These tests run the *real* fixed code path (gemm_a16wfp4 + the activation +branch). The fix only changes the *activation*, so the reference reuses the +*same* ``gemm_a16wfp4`` for both matmuls and differs only in the activation +math (computed independently in plain torch). This isolates the activation and +avoids conflating it with the kernel's mxfp4/bf16 GEMM precision. The tests +prove: + * the SwiGLU branch matches the SwiGLU-OAI reference, and + * it is genuinely different from the SiLU reference (i.e. the fix changed + behaviour, it is not silently equivalent), and + * the SiLU branch (DeepSeek) is unchanged. +""" + +import types + +import pytest +import torch + +cuda_only = pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires an AMD GPU" +) + +SCALE_GROUP_SIZE = 32 + +# e2m1 (fp4) decode table: sign | 2-bit exp | 1-bit mantissa. +_MXFP4_TABLE = [ + 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, + -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0, +] + + +def _mxfp4_to_f32(packed: torch.Tensor) -> torch.Tensor: + """Decode a uint8 tensor of two packed e2m1 nibbles to f32 (last dim x2).""" + x = packed.repeat_interleave(2, dim=-1) + x[..., ::2] = x[..., ::2] & 0xF + x[..., 1::2] = x[..., 1::2] >> 4 + table = torch.tensor(_MXFP4_TABLE, dtype=torch.float32, device=packed.device) + return table[x.long()] + + +def _e8m0_to_f32(scale: torch.Tensor) -> torch.Tensor: + return 2.0 ** (scale.to(torch.float32) - 127) + + +def _make_weight(n: int, k: int, *, seed: int): + """Random fp4-packed weight (n, k//2) + e8m0 scales (n, k//32). + + The bf16 dequantization is not returned: the reference reuses the kernel's + own GEMM, so we never need a from-scratch dequant matmul. + """ + g = torch.Generator(device="cuda").manual_seed(seed) + low = torch.randint(0, 16, (n, k // 2), dtype=torch.uint8, device="cuda", generator=g) + high = torch.randint(0, 16, (n, k // 2), dtype=torch.uint8, device="cuda", generator=g) + packed = low | (high << 4) + # e8m0 scales near 1.0 (bias 127) keep the dequant range sane. + scales = torch.randint( + 125, 130, (n, k // SCALE_GROUP_SIZE), dtype=torch.uint8, device="cuda", generator=g + ) + return packed, scales + + +def _ref_swiglu_oai(gate_up, alpha, beta, limit): + n = gate_up.shape[-1] // 2 + gate = gate_up[:, :n].to(torch.float32) + up = gate_up[:, n:].to(torch.float32) + if limit is not None: + gate = torch.clamp(gate, max=limit) + up = torch.clamp(up, min=-limit, max=limit) + return (gate * torch.sigmoid(alpha * gate) * (up + beta)).to(gate_up.dtype) + + +def _ref_silu(gate_up, limit): + n = gate_up.shape[-1] // 2 + gate = gate_up[:, :n].to(torch.float32) + up = gate_up[:, n:].to(torch.float32) + if limit > 0: + gate = torch.clamp(gate, max=limit) + up = torch.clamp(up, min=-limit, max=limit) + return (gate * torch.sigmoid(gate) * up).to(gate_up.dtype) + + +def _build_method_and_layer(hidden, inter, *, alpha, beta, limit): + from atom.config import LayerQuantConfig + from atom.model_ops.moe import Mxfp4MoEMethod, MoEActivationQuant + from aiter import QuantType + from unittest.mock import MagicMock + + qc = LayerQuantConfig( + quant_type=QuantType.per_1x32, + quant_dtype=torch.float4_e2m1fn_x2, + quant_method="quark", + ) + method = Mxfp4MoEMethod(qc, MagicMock()) + # Exercise the a16w4 (bf16 activation) path so we feed plain bf16 inputs. + method.act_quant = MoEActivationQuant.BF16 + + w13, s13 = _make_weight(2 * inter, hidden, seed=1) # (2I, H) + w2, s2 = _make_weight(hidden, inter, seed=2) # (H, I) + + layer = types.SimpleNamespace( + num_fused_shared_experts=1, + shared_w13_weight=w13.unsqueeze(0), + shared_w13_weight_scale=s13.unsqueeze(0), + shared_w2_weight=w2.unsqueeze(0), + shared_w2_weight_scale=s2.unsqueeze(0), + shared_w13_bias=None, + shared_w2_bias=None, + swiglu_limit=limit, + swiglu_alpha=alpha, + swiglu_beta=beta, + ) + return method, layer, (w13, s13), (w2, s2) + + +def _kernel_gemm(act, packed_scale): + """Reuse the exact same GEMM the dense path uses, so the only difference + between the dense path and the reference is the activation.""" + from aiter.ops.triton.gemm.basic.gemm_a16wfp4 import gemm_a16wfp4 + + weight, scale = packed_scale + return gemm_a16wfp4(act, weight, scale, dtype=torch.bfloat16) + + +@cuda_only +def test_swiglu_shared_expert_matches_reference(): + from aiter import ActivationType + from atom.model_ops.moe import Mxfp4MoEMethod + + if not _fp4_available(): + pytest.skip("MXFP4 not supported on this architecture") + + hidden, inter, M = 256, 256, 64 + alpha, beta, limit = 1.702, 1.0, 7.0 + method, layer, w13, w2 = _build_method_and_layer( + hidden, inter, alpha=alpha, beta=beta, limit=limit + ) + + torch.manual_seed(0) + x = torch.randn(M, hidden, dtype=torch.bfloat16, device="cuda") * 0.5 + + out = Mxfp4MoEMethod._apply_shared_experts_dense( + method, layer, x, ActivationType.Swiglu + ) + + # Reference reuses the SAME kernel GEMM; only the activation is computed + # independently (plain torch), isolating the fix from GEMM precision. + gate_up = _kernel_gemm(x, w13) + inter_ref = _ref_swiglu_oai(gate_up, alpha, beta, limit) + out_ref = _kernel_gemm(inter_ref, w2) + + # SiLU on the same gate_up must be clearly different (proves the branch + # matters and the fix is not silently equivalent to the old code). + inter_silu = _ref_silu(gate_up, limit) + out_silu = _kernel_gemm(inter_silu, w2) + + err_swiglu = (out.float() - out_ref.float()).abs().mean().item() + err_vs_silu = (out_ref.float() - out_silu.float()).abs().mean().item() + + torch.testing.assert_close(out.float(), out_ref.float(), rtol=1e-2, atol=1e-2) + assert err_vs_silu > 10 * max(err_swiglu, 1e-6), ( + f"swiglu vs silu too close to distinguish " + f"(err_swiglu={err_swiglu}, err_vs_silu={err_vs_silu})" + ) + + +@cuda_only +def test_silu_shared_expert_unchanged(): + from aiter import ActivationType + from atom.model_ops.moe import Mxfp4MoEMethod + + if not _fp4_available(): + pytest.skip("MXFP4 not supported on this architecture") + + hidden, inter, M = 256, 256, 64 + limit = 7.0 + method, layer, w13, w2 = _build_method_and_layer( + hidden, inter, alpha=1.702, beta=1.0, limit=limit + ) + + torch.manual_seed(0) + x = torch.randn(M, hidden, dtype=torch.bfloat16, device="cuda") * 0.5 + + out = Mxfp4MoEMethod._apply_shared_experts_dense( + method, layer, x, ActivationType.Silu + ) + + gate_up = _kernel_gemm(x, w13) + inter_ref = _ref_silu(gate_up, limit) + out_ref = _kernel_gemm(inter_ref, w2) + + torch.testing.assert_close(out.float(), out_ref.float(), rtol=1e-2, atol=1e-2) + + +def _fp4_available(): + try: + import aiter.ops.triton.utils._triton.arch_info as arch_info + + return arch_info.is_fp4_avail() + except Exception: + return torch.cuda.is_available() + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) From 4ce80a69302bb2e21e4edefcc0868dabe87cb397 Mon Sep 17 00:00:00 2001 From: Leon Ling Date: Thu, 25 Jun 2026 06:40:00 +0000 Subject: [PATCH 2/5] Route MiniMax-M3 decode to unified_attention on gfx1250 (MI455) gfx1250 has no pa_decode_gluon kernel (gluon supports gfx942 / gfx950 only), so MiniMax-M3 decode crashed with "pa_decode_gluon only supports gfx942 (CDNA3) and gfx950 (CDNA4)". The compacted sparse block table / context lengths the runners already build are exactly the (block_table, seqused_k) contract unified_attention consumes over the same SHUFFLE KV cache, so route both the full-attn and sparse decode paths through the triton unified_attention on gfx1250. Full-attn (attention_mha.py): - paged_attention_triton: add a use_unified flag that includes get_gfx() == "gfx1250" so decode takes the unified_attention branch instead of run_pa_decode_gluon. Gluon retained for CDNA3/CDNA4. Sparse (minimax_m3/sparse_attn.py) -- note ATOM_USE_UNIFIED_ATTN does NOT gate this path; the sparse runners call run_pa_decode_gluon directly: - add _sparse_decode_unified_attention helper feeding the kv-head collapsed SHUFFLE cache + sparse_bt + sparse_ctx into unified_attention(shuffled_kv_cache=True), each token a length-1 causal sequence (mirrors gluon max_seqlen_q=1 per-token-as-decode). - gfx1250 branch in minimax_m3_sparse_attn_decode_asm and _run_prefill_fp8_gluon: bf16 -> helper; fp8 -> NotImplementedError (gluon per-page descale has no unified_attention equivalent yet). Caveat: validated to compile/import on gfx950; the sparse path's GQA/block_table semantics still need MI455 numerical validation against the gfx950 gluon reference, and fp8 KV cache on gfx1250 is unsupported. Co-Authored-By: Claude Opus 4.8 (1M context) --- atom/model_ops/attention_mha.py | 13 ++- atom/model_ops/minimax_m3/sparse_attn.py | 105 +++++++++++++++++++++++ 2 files changed, 116 insertions(+), 2 deletions(-) diff --git a/atom/model_ops/attention_mha.py b/atom/model_ops/attention_mha.py index 0ef0b38fe7..9af3f43993 100644 --- a/atom/model_ops/attention_mha.py +++ b/atom/model_ops/attention_mha.py @@ -487,14 +487,23 @@ def paged_attention_triton( attn_metadata = fwd_ctx.attn_metadata - if envs.ATOM_USE_UNIFIED_ATTN and self.kv_cache_dtype.startswith("fp8"): + # gfx1250 (MI455) has no pa_decode_gluon kernel (gluon supports gfx942 / + # gfx950 only), so route the decode through the triton unified_attention + # path, which is gfx1250-capable and reads the same SHUFFLE KV cache. + use_unified = ( + envs.ATOM_USE_UNIFIED_ATTN + or self.use_flash_layout + or get_gfx() == "gfx1250" + ) + + if use_unified and self.kv_cache_dtype.startswith("fp8"): o = torch.empty(*q.shape, dtype=torch.bfloat16, device=q.device) else: o = torch.empty_like(q) num_seqs = attn_metadata.context_lens.shape[0] - if envs.ATOM_USE_UNIFIED_ATTN or self.use_flash_layout: + if use_unified: # print(q.shape, k_cache.shape, v_cache.shape) sliding_window = ( (self.sliding_window - 1, 0) if self.sliding_window > 0 else (-1, -1) diff --git a/atom/model_ops/minimax_m3/sparse_attn.py b/atom/model_ops/minimax_m3/sparse_attn.py index 2fcf03a9ac..61f36cc738 100644 --- a/atom/model_ops/minimax_m3/sparse_attn.py +++ b/atom/model_ops/minimax_m3/sparse_attn.py @@ -14,6 +14,7 @@ import aiter # noqa: F401 (used by the gluon PA runners for aiter.dtypes.fp8) import torch +from aiter.jit.utils.chip_info import get_gfx try: from vllm.triton_utils import tl, triton @@ -125,6 +126,64 @@ def _is_fp8_kv_cache_tensor(kv_cache: torch.Tensor) -> bool: return kv_cache.dtype in {dtype for dtype in fp8_dtypes if dtype is not None} +def _sparse_decode_unified_attention( + q_view: torch.Tensor, # [num_seqs, gqa_group, head_dim] (kv-head collapsed) + out_view: torch.Tensor, # [num_seqs, gqa_group, head_dim] + k_cache_view: torch.Tensor, # SHUFFLE 5D, num_kv_heads collapsed to 1 + v_cache_view: torch.Tensor, + sparse_bt: torch.Tensor, # [num_seqs, max_pages] physical-16 block table + sparse_ctx: torch.Tensor, # [num_seqs] per-row effective context length + sm_scale: float, + num_seqs: int, +) -> None: + """gfx1250 fallback for the sparse per-token-as-decode gluon kernel. + + gfx1250 (MI455) has no ``pa_decode_gluon`` kernel (gluon supports gfx942 / + gfx950 only). The sparse runners have already compacted the indexer's + selected blocks into a dense physical-16 ``sparse_bt`` + exact ``sparse_ctx`` + over the (kv-head collapsed) SHUFFLE cache — which is exactly the + ``(block_table, seqused_k)`` contract ``unified_attention`` consumes with + ``shuffled_kv_cache=True``. Each token is a length-1 causal "sequence", + mirroring the gluon ``max_seqlen_q=1`` per-token-as-decode setup. + + bf16 KV cache only: fp8 sparse decode plumbs per-token (per-page) descales + into the gluon kernel, which does not map onto ``unified_attention``'s descale + contract here; the caller raises NotImplementedError for fp8 on gfx1250. + """ + from aiter.ops.triton.unified_attention import unified_attention + + # block_size (page granularity) from the SHUFFLE cache: + # key_cache: [num_blocks, num_kv_heads, head_size // x, block_size, x] + block_size = k_cache_view.shape[3] + # Each token is its own length-1 sequence (decode); cu_seqlens_q = 0..num_seqs. + cu_seqlens_q = torch.arange( + num_seqs + 1, dtype=torch.int32, device=q_view.device + ) + # Safe upper bound: full block table width * page size (>= every sparse_ctx). + max_seqlen_k = int(sparse_bt.shape[1]) * int(block_size) + + unified_attention( + q_view, + k_cache_view, + v_cache_view, + out_view, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=1, + seqused_k=sparse_ctx, + max_seqlen_k=max_seqlen_k, + softmax_scale=sm_scale, + causal=True, + window_size=(-1, -1), + block_table=sparse_bt, + softcap=0, + q_descale=None, + k_descale=None, + v_descale=None, + sinks=None, + shuffled_kv_cache=True, + ) + + # --------------------------------------------------------------------------- # GQA block-sparse attention. BLOCK_SIZE_K == 128, matching one selected block. # --------------------------------------------------------------------------- @@ -1177,6 +1236,29 @@ def minimax_m3_sparse_attn_decode_asm( v_cache_view = v_cache.view(nph16 * _hkv, 1, *v_cache.shape[2:]) num_seqs = T * num_kv_heads + + # gfx1250 (MI455): no gluon pa_decode kernel (gluon supports gfx942 / gfx950 + # only). Route through the triton unified_attention sparse fallback over the + # same SHUFFLE cache + compacted sparse block table. + if get_gfx() == "gfx1250": + if _is_fp8_kv_cache_tensor(k_cache): + raise NotImplementedError( + "MiniMax-M3 fp8 sparse decode is not yet supported on gfx1250 " + "(MI455): the gluon per-page descale path has no unified_attention " + "equivalent here. Use a bf16 KV cache on gfx1250." + ) + _sparse_decode_unified_attention( + q_view, + out_view, + k_cache_view, + v_cache_view, + sparse_bt, + sparse_ctx, + sm_scale, + num_seqs, + ) + return + num_kv_heads_view = 1 query_group_size = g max_context_partition_num = get_recommended_splits(num_seqs, num_kv_heads_view) @@ -1271,6 +1353,29 @@ def _run_prefill_fp8_gluon( v_cache_view = v_cache.view(nph16 * _hkv, 1, *v_cache.shape[2:]) num_seqs = T * num_kv_heads + + # gfx1250 (MI455): no gluon pa_decode kernel (gluon supports gfx942 / gfx950 + # only). Route through the triton unified_attention sparse fallback over the + # same SHUFFLE cache + compacted sparse block table. + if get_gfx() == "gfx1250": + if _is_fp8_kv_cache_tensor(k_cache): + raise NotImplementedError( + "MiniMax-M3 fp8 sparse decode is not yet supported on gfx1250 " + "(MI455): the gluon per-page descale path has no unified_attention " + "equivalent here. Use a bf16 KV cache on gfx1250." + ) + _sparse_decode_unified_attention( + q_view, + out_view, + k_cache_view, + v_cache_view, + sparse_bt, + sparse_ctx, + sm_scale, + num_seqs, + ) + return + num_kv_heads_view = 1 query_group_size = g max_context_partition_num = get_recommended_splits(num_seqs, num_kv_heads_view) From 19af56967c802e91cfa319c05a2e5824fddf29ac Mon Sep 17 00:00:00 2001 From: ganyi Date: Mon, 29 Jun 2026 11:59:12 +0000 Subject: [PATCH 3/5] Fix Triton Prefill Gemm weight layout(gguu/gugu) mismatch Signed-off-by: ganyi maybe acc right Signed-off-by: ganyi uint8 to view Signed-off-by: ganyi reduce memory consumption Signed-off-by: ganyi prefill correct Signed-off-by: ganyi Cleanup --- atom/model_ops/minimax_m3/sparse_attn.py | 4 +- atom/model_ops/moe.py | 105 +++++++++- tests/test_mxfp4_shared_experts_swiglu.py | 222 ---------------------- 3 files changed, 97 insertions(+), 234 deletions(-) delete mode 100644 tests/test_mxfp4_shared_experts_swiglu.py diff --git a/atom/model_ops/minimax_m3/sparse_attn.py b/atom/model_ops/minimax_m3/sparse_attn.py index 61f36cc738..f0938f81c5 100644 --- a/atom/model_ops/minimax_m3/sparse_attn.py +++ b/atom/model_ops/minimax_m3/sparse_attn.py @@ -156,9 +156,7 @@ def _sparse_decode_unified_attention( # key_cache: [num_blocks, num_kv_heads, head_size // x, block_size, x] block_size = k_cache_view.shape[3] # Each token is its own length-1 sequence (decode); cu_seqlens_q = 0..num_seqs. - cu_seqlens_q = torch.arange( - num_seqs + 1, dtype=torch.int32, device=q_view.device - ) + cu_seqlens_q = torch.arange(num_seqs + 1, dtype=torch.int32, device=q_view.device) # Safe upper bound: full block table width * page size (>= every sparse_ctx). max_seqlen_k = int(sparse_bt.shape[1]) * int(block_size) diff --git a/atom/model_ops/moe.py b/atom/model_ops/moe.py index 3219d104c1..8aef088a39 100644 --- a/atom/model_ops/moe.py +++ b/atom/model_ops/moe.py @@ -785,6 +785,70 @@ def rocm_aiter_fused_moe_fake( ) +def _interleave_gate_up_rows_(layer: torch.nn.Module) -> None: + """Reorder w13 gate/up rows from SEPARATED (gguu) to INTERLEAVED (gugu). + + Operates in place on ``layer.w13_weight``, ``layer.w13_weight_scale`` and + ``layer.w13_bias`` (if present) along their row axis (dim=1), which is + ``2 * intermediate_size`` laid out as ``[gate(0..I-1) | up(0..I-1)]``. The + new order is ``[gate0, up0, gate1, up1, ...]`` so that a downstream consumer + splitting even/odd rows (the triton a16w4 SwiGLU kernel) reads matching + gate/up pairs. ``w2`` (down_proj) has no gate/up split and is untouched. + + The reorder is on whole rows (the ``2I`` axis); the MXFP4-packed pairs live + on the LAST axis (``H//2`` bytes/row), so reordering never splits a packed + byte. For FP4 we reorder a ``uint8`` view (bit-exact, no dequant/requant, no + scale recompute — torch has no ``index_select`` for ``float4_e2m1fn_x2``). + + Memory: the permutation is applied **in place per expert**, reusing the + existing storage. A plain ``index_select(...).contiguous()`` over the whole + tensor would double peak memory (a full second copy of the ~GB-sized w13), + which OOMs across many layers at load time. Here the only transient is one + expert's rows (a few MB), so peak overhead is negligible. + + Idempotency guard: sets ``layer._w13_gate_up_interleaved`` so a double call + (e.g. process_weights_after_loading invoked twice) is a no-op. + """ + if getattr(layer, "_w13_gate_up_interleaved", False): + return + + def _interleave_inplace(t: torch.Tensor) -> None: + """In-place row reorder [g..|u..] -> [g0,u0,g1,u1,...], per expert. + + Only FP4 (float4_e2m1fn_x2) needs the uint8 view, because torch has no + index_select/reshape for it AND its bytes are row-contiguous so the row + reorder is exact. Other dtypes (uint8 e8m0 scale, bf16/float bias) are + reordered directly — viewing them as uint8 would reinterpret bytes and + corrupt the values. + + Per expert e: view its 2I rows as (2, I, *rest), transpose to + (I, 2, *rest) (the gugu order); reshape forces one small (single-expert) + temp, then copy_ it back into the same storage. No full-tensor duplicate. + """ + _fp4 = getattr(torch, "float4_e2m1fn_x2", None) + if _fp4 is not None and t.dtype == _fp4: + buf = t.view(torch.uint8) + else: + buf = t + + E, two_i = buf.shape[0], buf.shape[1] + assert two_i % 2 == 0, f"w13 row dim {two_i} not even" + i = two_i // 2 + rest = buf.shape[2:] + for e in range(E): + rows = buf[e] # view into storage, shape (2I, *rest) + gugu = ( + rows.view(2, i, *rest).transpose(0, 1).reshape(two_i, *rest) + ) # one (single-expert) temp + rows.copy_(gugu) # write back into the same storage + + _interleave_inplace(layer.w13_weight.data) + _interleave_inplace(layer.w13_weight_scale.data) + if getattr(layer, "w13_bias", None) is not None: + _interleave_inplace(layer.w13_bias.data) + layer._w13_gate_up_interleaved = True + + class Mxfp4MoEMethod(FusedMoEMethodBase): def __init__(self, quant_config: LayerQuantConfig, moe: FusedMoEConfig): super().__init__(moe) @@ -969,27 +1033,52 @@ def process_weights_after_loading(self, layer): # the MoE-kernel-only CDNA4 layout. n_shared = layer.num_fused_shared_experts if n_shared > 0: + # IMPORTANT: .clone() (not .contiguous()). A uint8 view of a + # contiguous slice is already contiguous, so .contiguous() returns + # a tensor that SHARES storage with w13_weight. The gguu->gugu + # interleave below mutates w13_weight in place, which would then + # corrupt these stashed shared weights (consumed by the gguu + # half-split swiglu_oai_split). Clone to fully detach. layer.shared_w13_weight = ( - layer.w13_weight.data[-n_shared:].view(torch.uint8).contiguous() + layer.w13_weight.data[-n_shared:].view(torch.uint8).clone() ) layer.shared_w13_weight_scale = layer.w13_weight_scale.data[ -n_shared: - ].contiguous() + ].clone() layer.shared_w2_weight = ( - layer.w2_weight.data[-n_shared:].view(torch.uint8).contiguous() + layer.w2_weight.data[-n_shared:].view(torch.uint8).clone() ) layer.shared_w2_weight_scale = layer.w2_weight_scale.data[ -n_shared: - ].contiguous() + ].clone() if layer.w13_bias is not None: - layer.shared_w13_bias = layer.w13_bias.data[-n_shared:].contiguous() + layer.shared_w13_bias = layer.w13_bias.data[-n_shared:].clone() else: layer.shared_w13_bias = None if layer.w2_bias is not None: - layer.shared_w2_bias = layer.w2_bias.data[-n_shared:].contiguous() + layer.shared_w2_bias = layer.w2_bias.data[-n_shared:].clone() else: layer.shared_w2_bias = None + # gguu -> gugu interleave for the routed triton SwiGLU kernel. + # + # The checkpoint stores w13 gate/up rows SEPARATED ("gguu": + # [all-gate | all-up]). The aiter default CK fused_moe path consumes + # that directly. The triton a16w4 SwiGLU kernel, however, fuses the + # activation into the GEMM epilogue and splits each tile as + # INTERLEAVED ("gugu": a[...,::2]=gate, a[...,1::2]=up). Feeding gguu + # weights to that kernel mixes gate with gate and misreads up, + # corrupting the output. We interleave the routed w13 rows once here + # so the existing (well-tested) interleaved kernel produces results + # identical to the default path. w2 (down_proj) has no gate/up split. + # + # NOTE: this runs AFTER the shared-expert stash above, which keeps the + # shared experts in gguu for the dense swiglu_oai_split path. Only the + # SwiGLU triton branch needs interleaved rows; the SiLU branch uses + # fused_clamp_act_mul, which half-splits gguu — so gate on activation. + if getattr(layer, "activation", None) == ActivationType.Swiglu: + _interleave_gate_up_rows_(layer) + ( w13_weight, w13_scale, @@ -1352,9 +1441,7 @@ def _shared_expert_gemm(act, weight, weight_scale): out_dtype=x.dtype, ) else: - intermediate = torch.empty( - (M, half_n), device=x.device, dtype=x.dtype - ) + intermediate = torch.empty((M, half_n), device=x.device, dtype=x.dtype) fused_clamp_act_mul( gate_up, out=intermediate, diff --git a/tests/test_mxfp4_shared_experts_swiglu.py b/tests/test_mxfp4_shared_experts_swiglu.py deleted file mode 100644 index dd8c8f5923..0000000000 --- a/tests/test_mxfp4_shared_experts_swiglu.py +++ /dev/null @@ -1,222 +0,0 @@ -# SPDX-License-Identifier: MIT -# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. - -"""Numerical test for the dense shared-expert path with SwiGLU-OAI activation. - -Regression: ``Mxfp4MoEMethod._apply_shared_experts_dense`` hard-asserted the -SiLU activation path, so MiniMax-M3 (``ActivationType.Swiglu`` *with* fused -shared experts) crashed with:: - - AssertionError: dense shared-expert GEMM only supports the SiLU activation path - -MiniMax-M3 does not interleave gate/up weights, so the dense GEMM output is -split ``[gate | up]`` — exactly what ``swiglu_oai_split`` consumes. The dense -shared expert must therefore replicate ``MiniMaxM3MLP.forward``: -gate_up GEMM -> swiglu_oai_split -> down GEMM, with the SwiGLU-OAI math -``gate * sigmoid(alpha*gate) * (up + beta)`` (alpha=1.702, beta=1.0), not SiLU. - -These tests run the *real* fixed code path (gemm_a16wfp4 + the activation -branch). The fix only changes the *activation*, so the reference reuses the -*same* ``gemm_a16wfp4`` for both matmuls and differs only in the activation -math (computed independently in plain torch). This isolates the activation and -avoids conflating it with the kernel's mxfp4/bf16 GEMM precision. The tests -prove: - * the SwiGLU branch matches the SwiGLU-OAI reference, and - * it is genuinely different from the SiLU reference (i.e. the fix changed - behaviour, it is not silently equivalent), and - * the SiLU branch (DeepSeek) is unchanged. -""" - -import types - -import pytest -import torch - -cuda_only = pytest.mark.skipif( - not torch.cuda.is_available(), reason="requires an AMD GPU" -) - -SCALE_GROUP_SIZE = 32 - -# e2m1 (fp4) decode table: sign | 2-bit exp | 1-bit mantissa. -_MXFP4_TABLE = [ - 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, - -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0, -] - - -def _mxfp4_to_f32(packed: torch.Tensor) -> torch.Tensor: - """Decode a uint8 tensor of two packed e2m1 nibbles to f32 (last dim x2).""" - x = packed.repeat_interleave(2, dim=-1) - x[..., ::2] = x[..., ::2] & 0xF - x[..., 1::2] = x[..., 1::2] >> 4 - table = torch.tensor(_MXFP4_TABLE, dtype=torch.float32, device=packed.device) - return table[x.long()] - - -def _e8m0_to_f32(scale: torch.Tensor) -> torch.Tensor: - return 2.0 ** (scale.to(torch.float32) - 127) - - -def _make_weight(n: int, k: int, *, seed: int): - """Random fp4-packed weight (n, k//2) + e8m0 scales (n, k//32). - - The bf16 dequantization is not returned: the reference reuses the kernel's - own GEMM, so we never need a from-scratch dequant matmul. - """ - g = torch.Generator(device="cuda").manual_seed(seed) - low = torch.randint(0, 16, (n, k // 2), dtype=torch.uint8, device="cuda", generator=g) - high = torch.randint(0, 16, (n, k // 2), dtype=torch.uint8, device="cuda", generator=g) - packed = low | (high << 4) - # e8m0 scales near 1.0 (bias 127) keep the dequant range sane. - scales = torch.randint( - 125, 130, (n, k // SCALE_GROUP_SIZE), dtype=torch.uint8, device="cuda", generator=g - ) - return packed, scales - - -def _ref_swiglu_oai(gate_up, alpha, beta, limit): - n = gate_up.shape[-1] // 2 - gate = gate_up[:, :n].to(torch.float32) - up = gate_up[:, n:].to(torch.float32) - if limit is not None: - gate = torch.clamp(gate, max=limit) - up = torch.clamp(up, min=-limit, max=limit) - return (gate * torch.sigmoid(alpha * gate) * (up + beta)).to(gate_up.dtype) - - -def _ref_silu(gate_up, limit): - n = gate_up.shape[-1] // 2 - gate = gate_up[:, :n].to(torch.float32) - up = gate_up[:, n:].to(torch.float32) - if limit > 0: - gate = torch.clamp(gate, max=limit) - up = torch.clamp(up, min=-limit, max=limit) - return (gate * torch.sigmoid(gate) * up).to(gate_up.dtype) - - -def _build_method_and_layer(hidden, inter, *, alpha, beta, limit): - from atom.config import LayerQuantConfig - from atom.model_ops.moe import Mxfp4MoEMethod, MoEActivationQuant - from aiter import QuantType - from unittest.mock import MagicMock - - qc = LayerQuantConfig( - quant_type=QuantType.per_1x32, - quant_dtype=torch.float4_e2m1fn_x2, - quant_method="quark", - ) - method = Mxfp4MoEMethod(qc, MagicMock()) - # Exercise the a16w4 (bf16 activation) path so we feed plain bf16 inputs. - method.act_quant = MoEActivationQuant.BF16 - - w13, s13 = _make_weight(2 * inter, hidden, seed=1) # (2I, H) - w2, s2 = _make_weight(hidden, inter, seed=2) # (H, I) - - layer = types.SimpleNamespace( - num_fused_shared_experts=1, - shared_w13_weight=w13.unsqueeze(0), - shared_w13_weight_scale=s13.unsqueeze(0), - shared_w2_weight=w2.unsqueeze(0), - shared_w2_weight_scale=s2.unsqueeze(0), - shared_w13_bias=None, - shared_w2_bias=None, - swiglu_limit=limit, - swiglu_alpha=alpha, - swiglu_beta=beta, - ) - return method, layer, (w13, s13), (w2, s2) - - -def _kernel_gemm(act, packed_scale): - """Reuse the exact same GEMM the dense path uses, so the only difference - between the dense path and the reference is the activation.""" - from aiter.ops.triton.gemm.basic.gemm_a16wfp4 import gemm_a16wfp4 - - weight, scale = packed_scale - return gemm_a16wfp4(act, weight, scale, dtype=torch.bfloat16) - - -@cuda_only -def test_swiglu_shared_expert_matches_reference(): - from aiter import ActivationType - from atom.model_ops.moe import Mxfp4MoEMethod - - if not _fp4_available(): - pytest.skip("MXFP4 not supported on this architecture") - - hidden, inter, M = 256, 256, 64 - alpha, beta, limit = 1.702, 1.0, 7.0 - method, layer, w13, w2 = _build_method_and_layer( - hidden, inter, alpha=alpha, beta=beta, limit=limit - ) - - torch.manual_seed(0) - x = torch.randn(M, hidden, dtype=torch.bfloat16, device="cuda") * 0.5 - - out = Mxfp4MoEMethod._apply_shared_experts_dense( - method, layer, x, ActivationType.Swiglu - ) - - # Reference reuses the SAME kernel GEMM; only the activation is computed - # independently (plain torch), isolating the fix from GEMM precision. - gate_up = _kernel_gemm(x, w13) - inter_ref = _ref_swiglu_oai(gate_up, alpha, beta, limit) - out_ref = _kernel_gemm(inter_ref, w2) - - # SiLU on the same gate_up must be clearly different (proves the branch - # matters and the fix is not silently equivalent to the old code). - inter_silu = _ref_silu(gate_up, limit) - out_silu = _kernel_gemm(inter_silu, w2) - - err_swiglu = (out.float() - out_ref.float()).abs().mean().item() - err_vs_silu = (out_ref.float() - out_silu.float()).abs().mean().item() - - torch.testing.assert_close(out.float(), out_ref.float(), rtol=1e-2, atol=1e-2) - assert err_vs_silu > 10 * max(err_swiglu, 1e-6), ( - f"swiglu vs silu too close to distinguish " - f"(err_swiglu={err_swiglu}, err_vs_silu={err_vs_silu})" - ) - - -@cuda_only -def test_silu_shared_expert_unchanged(): - from aiter import ActivationType - from atom.model_ops.moe import Mxfp4MoEMethod - - if not _fp4_available(): - pytest.skip("MXFP4 not supported on this architecture") - - hidden, inter, M = 256, 256, 64 - limit = 7.0 - method, layer, w13, w2 = _build_method_and_layer( - hidden, inter, alpha=1.702, beta=1.0, limit=limit - ) - - torch.manual_seed(0) - x = torch.randn(M, hidden, dtype=torch.bfloat16, device="cuda") * 0.5 - - out = Mxfp4MoEMethod._apply_shared_experts_dense( - method, layer, x, ActivationType.Silu - ) - - gate_up = _kernel_gemm(x, w13) - inter_ref = _ref_silu(gate_up, limit) - out_ref = _kernel_gemm(inter_ref, w2) - - torch.testing.assert_close(out.float(), out_ref.float(), rtol=1e-2, atol=1e-2) - - -def _fp4_available(): - try: - import aiter.ops.triton.utils._triton.arch_info as arch_info - - return arch_info.is_fp4_avail() - except Exception: - return torch.cuda.is_available() - - -if __name__ == "__main__": - import sys - - sys.exit(pytest.main([__file__, "-v"])) From 97ca9093d9d1c29ca2f3dc126fc382c3dd58f5c3 Mon Sep 17 00:00:00 2001 From: Leon Ling Date: Thu, 9 Jul 2026 03:59:02 +0000 Subject: [PATCH 4/5] Force paged attention to triton on 1250 --- atom/model_ops/attention_mha.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/atom/model_ops/attention_mha.py b/atom/model_ops/attention_mha.py index 9af3f43993..e492c245ef 100644 --- a/atom/model_ops/attention_mha.py +++ b/atom/model_ops/attention_mha.py @@ -886,7 +886,11 @@ def dispatch_backend( # _can_use_prefill_sink_asm is valid. if self._can_use_prefill_sink_asm(q, k, v, fwd_ctx): return self.prefill_attention - if envs.ATOM_USE_UNIFIED_ATTN or self.use_flash_layout: + if ( + envs.ATOM_USE_UNIFIED_ATTN + or self.use_flash_layout + or get_gfx() == "gfx1250" + ): return self.prefill_attention_triton return self.prefill_attention return self._dispatch_decode() From 111471ee1ce0424e8a4fcbd19b1328671d2e8b3a Mon Sep 17 00:00:00 2001 From: Leon Ling Date: Thu, 9 Jul 2026 10:09:41 +0000 Subject: [PATCH 5/5] Handle and skip torch compile disk save error --- atom/utils/compiler_inferface.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/atom/utils/compiler_inferface.py b/atom/utils/compiler_inferface.py index 62659beebf..36ba896381 100644 --- a/atom/utils/compiler_inferface.py +++ b/atom/utils/compiler_inferface.py @@ -626,11 +626,12 @@ def compile( compiled_graph.save(path=path, format="unpacked") compilation_counter.num_compiled_artifacts_saved += 1 handle = (key, path) - except AssertionError: + except (AssertionError, RuntimeError) as e: logger.warning( "Skipping standalone compiled graph save for %s because " - "PyTorch did not emit a complete unpacked artifact.", + "PyTorch did not emit a complete unpacked artifact (%s).", key, + e, ) # Post-process generated wrapper Python files: wrap regions between