[MiniMax-M3.1] Run the DSpark draft's MXFP8 linears on AITER's preshuffled GEMM

atom/model_ops/linear.py

@@ -684,7 +684,10 @@ def weight_is_stored_preshuffled(
     preshuffled group32 GEMM (ATOM_GROUP32_WEIGHT_PRESHUFFLE, on by default),
     always on gfx1250, whose only group32 GEMM is AITER's preshuffled FlyDSL
     one (it reads the 32x32 scales as stored), and always under
-    ``needs_preshuffled_weight`` (V4.1's grouped ``wo_a``).
+    ``needs_preshuffled_weight`` (V4.1's grouped ``wo_a``). A 1x32 (MXFP8) one
+    only when its module asks through ``needs_preshuffled_weight`` (the
+    MiniMax-M3.1 DSpark draft, whose shapes AITER's preshuffled group32 GEMM
+    runs), on gfx950 under the same switch; any other stays row-major.
 
     Says nothing about rank. Only 2D weights are shuffled -- Qwen3-Next's GDN
     conv1d expands its weight to 3D and must stay row-major -- so the caller
@@ -704,13 +707,21 @@ def weight_is_stored_preshuffled(
     side backwards.
     """
     if native_group_rows is not None:
-        if native_group_rows != 32 or params_dtype != dtypes.fp8:
+        if native_group_rows not in (1, 32) or params_dtype != dtypes.fp8:
             return False
-        if needs_preshuffled_weight:
-            return True
         from aiter.jit.utils.chip_info import get_gfx
 
         gfx = get_gfx()
+        if native_group_rows == 1:
+            # MXFP8 (1x32): only a module that asks (the MiniMax-M3.1 DSpark
+            # draft), on gfx950, where AITER's preshuffled group32 GEMM runs it.
+            return bool(
+                needs_preshuffled_weight
+                and envs.ATOM_GROUP32_WEIGHT_PRESHUFFLE
+                and gfx == "gfx950"
+            )
+        if needs_preshuffled_weight:
+            return True
         return gfx == "gfx1250" or (
             envs.ATOM_GROUP32_WEIGHT_PRESHUFFLE and gfx == "gfx950"
         )

atom/models/minimax_m3/dspark.py

@@ -478,6 +478,19 @@ def __init__(self, atom_config: "Config", layer_offset: int = 0) -> None:
                 for i in range(config.num_hidden_layers)
             ]
         )
+        # MXFP8 linears: AITER's preshuffled group32 GEMM (gfx950) is ~2.5x the
+        # Triton group32 one at these shapes and bit-identical; the loader
+        # preshuffles them where it runs (weight_is_stored_preshuffled).
+        if quant_config.quant_method:
+            for layer in self.layers:
+                attn, mlp = layer.self_attn, layer.mlp
+                for proj in (
+                    attn.qkv_proj,
+                    attn.o_proj,
+                    mlp.gate_up_proj,
+                    mlp.down_proj,
+                ):
+                    proj.needs_preshuffled_weight = True
         self.final_norm = GemmaRMSNorm(self.hidden_size, eps=config.rms_norm_eps)
         self.markov_head = DSparkMarkovHead(self.vocab_size, self.markov_rank)
         # 0 or absent: every block row attends the whole context.

docs/environment_variables.md

@@ -139,7 +139,7 @@ cannot affect a server that does not use the feature.
 |----------|------|---------|-------------|
 | **ATOM_USE_TRITON_GEMM** | bool | 0 (false) | If set to `1`, use AITER Triton FP4 weight preshuffled GEMM. Otherwise use AITER ASM FP4 weight preshuffled GEMM. |
 | **ATOM_FP8_BLOCKSCALE_USE_E8M0_SCALE** | bool | unset (per checkpoint) | E8M0 rather than FP32 128x128 FP8 block scales, for the weight and the activation quantized for it; on gfx950 E8M0 operands take the AITER FlyDSL GEMM. Unset: E8M0 on every arch but gfx942 when the checkpoint declares `scale_fmt: ue8m0` (DeepSeek-V4), whose power-of-two scales E8M0 restates exactly; FP32 otherwise. `1`/`0` force it on/off; forcing it on for a checkpoint with non-power-of-two FP32 scales rounds them. |
-| **ATOM_GROUP32_WEIGHT_PRESHUFFLE** | bool | 1 (true) | On gfx950, (16, 16)-shuffle native FP8 32x32 group32 weights (DeepSeek-V4.1) at load and after a weight sync, so AITER's preshuffled group32 GEMM (FlyDSL) reads them; that GEMM emits BF16 only. Set to `0` to keep checkpoint bytes and the row-major group32 GEMM. V4.1's grouped `wo_a` is shuffled either way. |
+| **ATOM_GROUP32_WEIGHT_PRESHUFFLE** | bool | 1 (true) | On gfx950, (16, 16)-shuffle native FP8 32x32 group32 weights (DeepSeek-V4.1), and the 1x32 (MXFP8) ones of the MiniMax-M3.1 DSpark draft, at load and after a weight sync, so AITER's preshuffled group32 GEMM (FlyDSL) reads them; that GEMM emits BF16 only. Set to `0` to keep checkpoint bytes and the row-major group32 GEMM. V4.1's grouped `wo_a` is shuffled either way. |
 | **ATOM_USE_FP4_NON_SHUFFLE_TRITON_GEMM** | bool | 0 (false) | If set to `1`, use AITER Triton FP4 GEMM with non-shuffled weights. Takes precedence over the FP4 preshuffled GEMM path selected by `ATOM_USE_TRITON_GEMM`. |
 | **ATOM_FP4_MXSCALE_USE_32X4_SCALE_SHUFFLE** | bool | 0 (false) | On gfx1250, if set to `1`, keep the 1x32 E8M0 weight scale of every 2D MXFP4 (`per_1x32`) linear with K % 128 == 0 in AITER's 32x4 layout, quantize the activation into the matching layout (or move the scale of one a fused norm/activation already quantized into it) and call AITER `gemm_a4w4`, which picks FlyDSL or ASM per shape. The other layers keep the Gluon FP4 preshuffled GEMM, as do those whose scale other kernels read directly (MLA `kv_b_proj`; DeepSeek `fused_qkv_a_proj` under `ATOM_USE_TRITON_GEMM=1`). No effect when disabled, on other archs, or when `ATOM_USE_FP4_NON_SHUFFLE_TRITON_GEMM=1`. |
 | **ATOM_MHC_USE_BF16** | bool | 1 (true) | Use AITER BF16 hi/lo mHC computation for attention, FFN and head. After loading, replace FP32 fn storage with `mhc_shuffle_fn` output; no FP32 copy is retained. Set to `0` for FP32 mHC. Takes effect at model load; restart to change modes. On gfx1250, AITER enables shuffled residuals only while its runtime `mhc_fused_post_pre` policy remains fused (`M < 1024`); larger M uses ordinary residual layout and the standalone post/pre fallback. |

tests/model_ops/test_blockscale_dispatch.py

@@ -63,3 +63,40 @@ def fail(*args, **kwargs):
             x_scale=torch.empty((3, 2), dtype=torch.float8_e8m0fnu),
         )
     assert raised.value is error
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="runs the FlyDSL and Triton GEMMs"
+)
+@pytest.mark.parametrize(
+    "n,k",
+    [(2304, 6144), (6144, 2048), (6144, 6144), (6144, 3072)],  # DSpark draft, TP4
+)
+@pytest.mark.parametrize("m", [1, 9, 200])
+def test_preshuffled_1x32_matches_the_row_major_gemm_bit_for_bit(n, k, m):
+    """The MiniMax-M3.1 DSpark draft's MXFP8 linears on AITER's preshuffled
+    group32 GEMM: same bytes, scales and products as the row-major (Triton)
+    path, so the output is identical."""
+    from aiter.jit.utils.chip_info import get_gfx
+    from aiter.ops.shuffle import shuffle_weight
+
+    if get_gfx() != "gfx950":
+        pytest.skip("the preshuffled 1x32 GEMM is gfx950's")
+    torch.manual_seed(0)
+    w = torch.randn(n, k, device="cuda") * 0.05
+    exp = torch.ceil(
+        torch.log2(w.view(n, k // 32, 32).abs().amax(-1).clamp_min(1e-12) / 448.0)
+    )
+    wq = (w.view(n, k // 32, 32) / torch.pow(2.0, exp)[..., None]).view(n, k)
+    wq = wq.clamp(-448, 448).to(torch.float8_e4m3fn)
+    ws = (exp + 127).to(torch.uint8).view(torch.float8_e8m0fnu)
+    x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
+    row_major = blockscale.native_quant_linear(x, wq, ws, weight_group_rows=1)
+    preshuffled = blockscale.native_quant_linear(
+        x,
+        shuffle_weight(wq, (16, 16)),
+        ws,
+        weight_group_rows=1,
+        weight_preshuffled=True,
+    )
+    torch.testing.assert_close(preshuffled, row_major, rtol=0, atol=0)

tests/test_weight_sync_shuffle_layout.py

@@ -111,6 +111,35 @@ def test_unquantized_and_per_tensor_are_not_shuffled():
     assert weight_is_stored_preshuffled(QuantType.per_Tensor, dtypes.fp8) is False
 
 
+@pytest.mark.parametrize(
+    "asked,env_value,gfx,expected",
+    [
+        (True, "1", "gfx950", True),
+        (False, "1", "gfx950", False),
+        (True, "0", "gfx950", False),
+        (True, "1", "gfx942", False),
+        (True, "1", "gfx1250", False),
+    ],
+)
+def test_native_1x32_fp8_is_shuffled_only_when_its_module_asks_on_gfx950(
+    monkeypatch, asked, env_value, gfx, expected
+):
+    """A native MXFP8 (1x32) weight is preshuffled only for a module that asks
+    (the MiniMax-M3.1 DSpark draft), and only where AITER's preshuffled group32
+    GEMM reads it: gfx950, under the group32 switch."""
+    monkeypatch.setenv("ATOM_GROUP32_WEIGHT_PRESHUFFLE", env_value)
+    monkeypatch.setattr("aiter.jit.utils.chip_info.get_gfx", lambda: gfx)
+    assert (
+        weight_is_stored_preshuffled(
+            QuantType.per_1x32,
+            dtypes.fp8,
+            needs_preshuffled_weight=asked,
+            native_group_rows=1,
+        )
+        is expected
+    )
+
+
 def test_mxfp4_preshuffle_selects_a_compatible_backend(monkeypatch):
     monkeypatch.setattr(linear_mod, "use_fp4_non_shuffle_triton_gemm", lambda: False)
     monkeypatch.setattr(linear_mod, "use_triton_gemm", lambda: False)
@@ -145,7 +174,9 @@ def test_forced_gfx1250_triton_uses_legacy_m_dependent_activation_layout(monkeyp
 # ── the property: load and sync agree ─────────────────────────────────────
 
 
-def _linear_double(quant_type, params_dtype, needs_preshuffled_weight, *, dim=2):
+def _linear_double(
+    quant_type, params_dtype, needs_preshuffled_weight, *, dim=2, native_group_rows=None
+):
     """The attributes `LinearBase.process_weights_after_loading` reads.
 
     A double rather than a real Linear so the case matrix does not need a
@@ -168,7 +199,7 @@ def _linear_double(quant_type, params_dtype, needs_preshuffled_weight, *, dim=2)
         needs_preshuffled_weight=needs_preshuffled_weight,
         # `None` = not the native group32 path, which returns before settling
         # any layout at all and so would answer every case below the same way.
-        native_a8_group_rows=None,
+        native_a8_group_rows=native_group_rows,
         _maybe_pad_a8w8_preshuffle_output=lambda: False,
         prefix="test",
     )
@@ -184,19 +215,21 @@ class _Updater(WeightUpdaterMixin):
     return _Updater()
 
 
-def _loader_shuffles(monkeypatch, case, *, dim=2):
+def _loader_shuffles(monkeypatch, case, *, dim=2, native_group_rows=None):
     calls = []
     monkeypatch.setattr(linear_mod, "shuffle_weights", lambda *a, **k: calls.append(a))
     # per_1x32 shuffles its scale through aiter at the end of the method; that
     # is not what is under test and it wants a real e8m0 tensor.
     monkeypatch.setattr(
         linear_mod.fp4_utils, "e8m0_shuffle", lambda scale: scale, raising=False
     )
-    LinearBase.process_weights_after_loading(_linear_double(*case, dim=dim))
+    LinearBase.process_weights_after_loading(
+        _linear_double(*case, dim=dim, native_group_rows=native_group_rows)
+    )
     return bool(calls)
 
 
-def _sync_shuffles(monkeypatch, case, *, dim=2):
+def _sync_shuffles(monkeypatch, case, *, dim=2, native_group_rows=None):
     import atom.model_ops.utils as utils_mod
 
     calls = []
@@ -212,6 +245,7 @@ def _sync_shuffles(monkeypatch, case, *, dim=2):
         needs_preshuffled_weight=needs_preshuffled_weight,
         weight_scale=None,
         need_normalize_e4m3fn_to_e4m3fnuz=False,
+        native_a8_group_rows=native_group_rows,
     )
     _updater_double()._post_process_fp8_weight(module, param)
     return bool(calls)
@@ -229,6 +263,25 @@ def test_load_and_sync_shuffle_in_the_same_cases(
     assert _loader_shuffles(monkeypatch, case) == _sync_shuffles(monkeypatch, case)
 
 
+@pytest.mark.parametrize("native_group_rows", [1, 32])
+@pytest.mark.parametrize("asked", [False, True])
+@pytest.mark.parametrize("preshuffle", ["0", "1"])
+@pytest.mark.parametrize("gfx", ["gfx950", "gfx1250"])
+def test_load_and_sync_shuffle_native_group32_the_same(
+    monkeypatch, native_group_rows, asked, preshuffle, gfx
+):
+    """The native group32 branch, 1x32 (MXFP8) included: the real loader and
+    the real sync shuffle in the same cases."""
+    monkeypatch.setenv("ATOM_GROUP32_WEIGHT_PRESHUFFLE", preshuffle)
+    monkeypatch.setattr("aiter.jit.utils.chip_info.get_gfx", lambda: gfx)
+    case = (QuantType.per_1x32, dtypes.fp8, asked)
+    loads = _loader_shuffles(monkeypatch, case, native_group_rows=native_group_rows)
+    syncs = _sync_shuffles(monkeypatch, case, native_group_rows=native_group_rows)
+    assert loads == syncs
+    if native_group_rows == 1:
+        assert loads is (asked and preshuffle == "1" and gfx == "gfx950")
+
+
 # ── one decision, one comparison ──────────────────────────────────────────