[zhuyuhua-v]

Summary

The MiniMax-M3.1 DSpark draft keeps its linears as shipped: MXFP8 (FP8 E4M3, one E8M0 scale per 32). Until now they were held row-major and ran on the Triton group32 GEMM. This PR (16, 16)-preshuffles them at load, so native_quant_linear's existing AITER call runs them on FlyDSL's MXFP8 bmm. AITER already ships MiniMax-M3 1x32 tuned rows for these shapes (ROCm/aiter#6060).

Only the layout changes; the output is bit-identical. Only the draft is affected: every other 1x32 weight, and every other arch, stays row-major.

Change

• atom/models/minimax_m3/dspark.py: an MXFP8 draft sets needs_preshuffled_weight on each layer's qkv_proj, o_proj, gate_up_proj and down_proj.
• atom/model_ops/linear.py: weight_is_stored_preshuffled preshuffles a native 1x32 FP8 weight only when its module asks, on gfx950, under the existing ATOM_GROUP32_WEIGHT_PRESHUFFLE switch. The load and the weight sync both go through it.

Performance

Draft linears, TP4 per-rank shapes, one MI355X. Timings are per layer (2 x qkv + o + gate_up + down), with activation quantization included:

rows	before (Triton)	this PR	speedup
8	106.5 us	38.7 us	2.75x
64	133.6 us	50.3 us	2.66x
256	195.7 us	87.5 us	2.24x
512	319.4 us	121.2 us	2.64x

End to end: MiniMax-M3.1-Flash, TP4, AgentX c32, 3600 s, one run each, on the ganyi/m31-mono-fp4-pr stack (#2812), with the same draft weights preshuffled.

	before	this PR
P90 interactivity	140.6	144.9
P50 interactivity	213.7	220.8
total tok/s/GPU	57,943	57,943

gsm8k (5-shot, 1319 samples, real DSpark acceptance): strict-match 0.9666, against 0.9621 / 0.9629 before.

Tests

• tests/test_weight_sync_shuffle_layout.py:
  • the 1x32 truth table;
  • load-vs-sync parity on the native group32 branch.
• tests/model_ops/test_blockscale_dispatch.py: on gfx950, a bit-exact check of the preshuffled 1x32 GEMM against the row-major path at the draft's shapes.
• 93 passed on MI355X with aiter 1053d8d18. Both files skip on CI's unit runner (no aiter / GPU).

Notes

• Needs an AITER with [HIP] [OPUS] [FlyDSL] gfx950 MXFP8 bmm: 1x32 weight scales, grouped XCD order, MiniMax-M3 configs aiter#6060 (in main since 2026-10-27).
• ATOM_GROUP32_WEIGHT_PRESHUFFLE=0 restores the row-major path.

🤖 Generated with Claude Code

[github-actions]

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

• ✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

• ✅ Run after the PR is approved and Pre Checkin passes
• ✅ Run immediately when an approval review is submitted
• ✅ Can be requested before approval with labels

Label	Tests
ci:full	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom	Run native ATOM model accuracy tests
ci:vllm	Run ATOM vLLM OOT model accuracy tests
ci:sglang	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 2876 --add-label <label>

[Copilot]

🟡 Changes recommended

The default gfx950 path selects preshuffling for supported native weight shapes that the shuffle cannot tile, causing weight loading to fail unless the shared load/sync decision checks shape compatibility.

1 open issue

• Keep unsupported native 1x32 weight shapes row-major · New

What changed in this PR

This PR extends ATOM’s native FP8 preshuffling to 1x32 MXFP8 weights on gfx950, allowing MiniMax-M3.1 DSpark draft linears and other eligible checkpoints to use AITER’s preshuffled group32 GEMM.

Changes:

• Enables native 1x32 FP8 preshuffling under the existing ATOM_GROUP32_WEIGHT_PRESHUFFLE switch.
• Adds eligibility, dispatch, and GPU bit-exact comparison tests.
• Documents the switch’s expanded scope.

File	Description
tests/​test_weight_sync_shuffle_layout.py	Tests architecture and switch eligibility, plus unchanged FP4 behavior.
tests/​model_ops/​test_blockscale_dispatch.py	Tests AITER dispatch and bit-exact results against row-major GEMM.
docs/​environment_variables.md	Documents native 1x32 MXFP8 preshuffling.
atom/​model_ops/​linear.py	Extends the shared load/sync shuffle policy to native 1x32 FP8 weights.

_{Review effort: Balanced}

[valar-cop]

some changes needed..

ci green.

[Copilot]

🟡 Changes recommended

The default-enabled gfx950 guard admits unsupported weight shapes that fail on the first forward pass because shuffle compatibility does not ensure AITER kernel support.

1 open issue

• Shuffle guard enables unsupported gfx950 BMM shapes · New

1 resolved since last review

• Keep unsupported native 1x32 weight shapes row-major

_{Review effort: Balanced}

[valar-cop]

some changes needed..

ci green.

[Copilot]

🟡 Changes recommended

The default-on gfx950 path can fail weight loading when optional dependencies are missing and breaks previously supported FP32 output for preshuffled native 1x32 weights.

2 open issues

• Missing AITER dependencies abort native weight loading · New
• Default preshuffling breaks FP32 output for native 1x32 weights · New

1 resolved since last review

• Shuffle guard enables unsupported gfx950 BMM shapes

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

The snapshot lacks the claimed guards and boundary tests, while the AITER #6060-dependent gfx950 behavior remains unverified because the affected tests skip on CPU and GPU validation was unavailable.

1 open issue

• Default preshuffling breaks FP32 output for native 1x32 weights

1 resolved since last review

• Missing AITER dependencies abort native weight loading

_{Review effort: Balanced}