[MoE] Port moe_gemm_2stage (stage1+stage2) to the new fx.* pipeline — fp8 (gfx942 + gfx950) - #947
Merged
Merged
Conversation
Add kernels/moe/moe_gemm_2stage/layout_helpers.py, the fxh helper layer the aiter layout-API reference kernel depends on, adapted to this repo's fx.* surface. Foundation for migrating gemm1/gemm2 off the legacy SmemAllocator/SmemPtr + buffer_ops API. Contains only reusable layout/copy-atom/fragment helpers (FlyObjCache, create_thr_mma, load/store_tiled_mma_frag*, get_tiled_copy_coalesced_mn, all_copy_atoms, all_elements, atom_tensor, split_works, etc.); aiter package couplings and host-side torch-view/debug shims dropped. No behavior change to existing kernels; new file only. Public API of the package is untouched. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add the native-fp8 (MFMA 16x16x32) gate-up helpers to layout_helpers.py, ported faithfully from the aiter reference prefill_1x4 path with compile-time closures (N,K,TOPK,BLOCK_M,weight_dtype) made explicit args: - make_1x4_tiled_mma: B-first (4,1,1) tiled_mma; fp8 k_perm (8,4,2):(1,16,8) - make_gateup_weight_view: shuffle_weight-order preshuffle + gate/up silu grouping - silu_pair_bf16: silu(gate)*up with optional fp8 weight/act dequant -> bf16 - read_sorted_index / make_tensor_with_index: MoE sorted-row gather + scatter - atomic_add_bf16: global_atomic_pk_add_bf16 for the down/atomic epilogue Validated: single-tile native-fp8 GEMM via make_1x4_tiled_mma reaches cos=1.0 cold (the 4x seen in an isolated dense readback is per-wave channel ownership double-counting, confirmed via per-wave channel-offset dump: waves 0/1/2/3 own channels 0/16/32/48). No wiring into gemm1.py yet (part 2). black+ruff clean. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…le configs New-pipeline (fx.*) MoE stage1 gate-up: - fp8: fixed silu dtype-reinterpret bug (silu_pair_bf16 out_dtype must match CShuffle store dtype); now correct across all tile configs and both out dtypes. - perf: 2-stage A LDS ping-pong (a_pong wired), B reg prefetch, LDS-read hoist one iter ahead -> 883->1193 TFLOPS (~92% of legacy on compute-bound; decode already parity). - tile configs: contiguous_n = max(tile_n//2, 64) supports tile_n=64; tile_m>=16. - bf16: builder made dtype-parametric (in_dtype param); gfx950 native MFMA(16,16,32); no dequant; num_records_bytes fixed to bytes (*elem_bytes). - fp16/f16 stays on legacy MFMA (unchanged). - Debug A/B toggles (default-off): MOE_FORCE_LEGACY_FP8 / MOE_FORCE_LEGACY_BF16. Verified cache-off: 60 passed (fp8/bf16/fp16 x S/M x atomic/reduce x f16/bf16). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…/f32) Add a native-fp8 stage2 down-projection builder `_build_moe_gemm2_fp8` (gemm2.py) mirroring stage1's B-first layout-API path, and route `compile_moe_gemm2` to it for in_dtype="fp8" (non-groupwise). Env toggle MOE_FORCE_LEGACY_G2_FP8=1 restores the legacy path for A/B. All other dtypes, fp16, groupwise, and the legacy body are untouched; compile_moe_gemm2_ex / _MoeGemm2ReduceWrapper wrap the new exe unchanged. New builder: fx.rocdl.make_buffer_tensor (num_records=M*K*1), SharedAllocator + @fx.struct/@fx.union LDS, 2-stage A LDS ping-pong (swz 3,4,3) + B register prefetch, per-token(row) x per-channel(model_dim) fp8 dequant, optional routed weight, and a CShuffle epilogue with three scatter modes: - atomic (accumulate=True), f16/bf16: packed buffer atomic-add into out[token, model_dim] - atomic, f32: scalar buffer atomic-add - reduce (accumulate=False), f16/bf16: plain store into out[token*topk+slot, model_dim] for the existing reduction kernel layout_helpers.py: - make_weight_view: plain (no gate/up silu) preshuffle weight view for stage2 - _TensorWithIndex.copy: add atomic scatter (buffer atomic-add) with an out-of-tile grid-slot guard (tall row-guard fake tensor) and value/offset width alignment, plus a token row_limit clamp for sentinel rows - _buffer_atomic_pk / _buffer_atomic_f32 helpers Key correctness fix: the A2-gather view is rank-3 [tokens, topk, inter] so the sorted id decodes to (token, slot); a rank-2 view silently read slot-0 of each token (masked on tame data, wrong on real routing). Verify (FLYDSL_RUNTIME_ENABLE_CACHE=0, gfx950): - test_moe_stage2_standalone -k fp8: 12 passed (atomic+reduce+torch/flydsl reduce+masked, incl tile_n=256, decode/prefill) - test_moe_gemm_2stage -k "fp8 and eager and nomask and (S or M)": 30 passed (S/M/L, atomic+reduce, f16/bf16/f32) - broader fp8 e2e (graph+mask): 56 passed Perf (dim 8192x8192, t=32768, e=16, k=4, tile 256/128, median-of-5, iters 60): new 1095 TFLOPS vs legacy 1094 TFLOPS (0.999x). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Generalize the native B-first stage2 down-projection builder
(_build_moe_gemm2_fp8) to also handle bf16, mirroring the fp8->bf16
extension already done on stage1:
- add in_dtype param; derive elem_t (BFloat16 vs Float8E4M3FNUZ),
elem_bytes (2 vs 1), _val_per_thr (8 vs 16), and A-LDS swizzle
((3,3,3) vs (3,4,3))
- size the A LDS buffers in bytes (BM*TILE_K*elem_bytes) so bf16 gets
2x the byte footprint; num_records stays M*K*elem_bytes
- skip the fp8 per-row/per-channel dequant for bf16 (unscaled inputs);
sentinel rows still contribute 0 via the OOB-zero A-gather
- keep the atomic (pk/f32) and reduce/f32 epilogues intact (bf16 pk
atomics go through the existing buffer_atomic_pk_add path on gfx950)
Route in_dtype in {fp8,bf16} (non-groupwise, not forced-legacy) to the
new builder in compile_moe_gemm2, and add a MOE_FORCE_LEGACY_G2_BF16
toggle mirroring MOE_FORCE_LEGACY_G2_FP8. compile_moe_gemm2_ex /
_MoeGemm2ReduceWrapper forward in_dtype unchanged, so reduce mode works
for bf16 too. fp8/f32/fp16/legacy/reduction paths are untouched.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Remove the dev-only MOE_FORCE_LEGACY_* A/B env toggles; fp8/bf16 (non-split-K, non-groupwise) now route unconditionally to the new-pipeline gemm1/gemm2 builders. Legacy body remains for the un-ported dtypes (int4/int8/int8smooth/ int4_bf16, fp16, fp8 split-K, groupwise). No behavior change with toggles unset. Verified cache-off: fp8+bf16 gemm1+gemm2 (atomic/reduce) green. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…peline CI (full matrix, RUN_TESTS_FULL) surfaced two bf16 new-path bugs: - gfx942 (CDNA3): compile abort — the builder emits gfx950-only bf16 MFMA(16,16,32); CDNA3 bf16 is MFMA(16,16,16). - gfx950: hard crash at in_dtype=bf16 + out_dtype=f32 at large tile_m (FP4-L). fp8 is green on both gfx942 and gfx950. Route only fp8 to the new gemm1/gemm2 builders; bf16 uses the proven legacy path until the two issues above are fixed. The dtype-parametric builder still accepts bf16 (kept for the follow-up), it is just not routed. No behavior change for fp8. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
coderfeli
changed the title
…es legacy The gfx942 (CDNA3) CI run aborted compiling the new fp8 gemm2 bf16-output atomic epilogue (op does not lower on CDNA3), and only ~10% of the fp8 matrix was ever exercised on gfx942 before the abort (reduce/f32/mask/graph untested there). Rather than ship an unvalidated CDNA3 path, gate the new gemm1/gemm2 builders to CDNA4 (gfx95*). gfx942 falls through to the proven legacy path entirely; gfx950 continues to use the new, validated pipeline. gfx950: fp8 still routes to the new path (verified, 24 passed). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
coderfeli
commented
Aug 1, 2026
| return x // y | ||
|
|
||
|
|
||
| def fly_ast_rewrite(member): |
coderfeli
commented
Aug 1, 2026
| r += 1 | ||
|
|
||
|
|
||
| def _encode_waitcnt(vmcnt=63, expcnt=7, lgkmcnt=63): |
Collaborator
Author
There was a problem hiding this comment.
there are existing apis in expr
coderfeli
commented
Aug 1, 2026
| return | ||
|
|
||
|
|
||
| def _as_ptr(p, dtype=None): |
coderfeli
commented
Aug 1, 2026
| # MLIR values are all SSA which is naturally different from each other and once | ||
| # defined stay unchanged for the rest of their lifetime, so they can be used | ||
| # safely as a cache key. | ||
| class FlyObjCache: |
coderfeli
commented
Aug 1, 2026
| # explicit args so the helpers are reusable across tile configs. | ||
|
|
||
|
|
||
| def atomic_add_bf16(ptr_base, reg_vec): |
Collaborator
Author
There was a problem hiding this comment.
not rocdl wrapper or fx.copy?
… LOC Behavior-preserving refactor of the NEW fp8 MoE 2-stage code (gfx950): - layout_helpers.py: remove confirmed-unused helper surface (FlyObjCache, all_elements/collect_nz_modes, split_works, load_fragment, all_copy_atoms, atom_tensor, div_up/div_e, get_d1_shape, inner_most_stride, sub_tensor, asm_mark, make_1d_coord_tensor) and the now-unused functools/inspect/types/ TargetAddressSpace imports. Add a tiny `reps(t, i)` helper. - gemm1/gemm2 fp8 builders: use fxh.reps for the repeated fx.size(fx.get_shape(...)[i]).to_py_value() idiom; trim restated comments. Legacy code paths, routing/arch-gate, and numerics untouched. fp8 2-stage and stage2-standalone tests pass; perf unchanged (stage1 ~1140 / stage2 reduce ~1524, atomic ~1101 TFLOPS). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…6-atomic carve-out Drop the CDNA4-only gate: fp8 now routes to the new gemm1/gemm2 builders on gfx94* (CDNA3) as well as gfx95* (CDNA4). gemm1 (fp8 MFMA + CShuffle store, no atomics) works on both. gemm2's only CDNA3-incompatible op is the packed bf16-output buffer atomic (`buffer_atomic_pk_add_bf16` does not exist on gfx942), so the single combo (atomic + out_dtype=bf16) stays on the legacy body on gfx942 (which uses global bf16 atomics); every other fp8 combo uses the new path. gfx950 is unchanged (all combos new path). gfx950: verified green (30 fp8 cases, cache-off). gfx942: compile-validated for the full fp8 matrix locally (this box can cross-compile gfx942 but not execute it, so correctness is validated by CI on real gfx942 hardware). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…code shapes The atomic scatter redirected out-of-tile grid lanes to element 0 (adding 0 to out[0]) instead of dropping them. On small-tile/decode grids (e.g. tile_m=16, E=128, topk=8) thousands of padding lanes atomic-add 0 to the same cache line, serializing the kernel: fp8 stage2 atomic on the decode bench shape was 79800us / 0.08 TFLOPS vs legacy 134us. Push invalid lanes to an OOB element index (== output element count) so the buffer resource's hardware bounds check drops the atomic (matches the plain-store path). Decode atomic 79800us -> 223us; compute-bound shape unchanged (~1097 TFLOPS). Correctness preserved (also drops the now-redundant per-lane value zeroing). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Profiled the ~1.6x gap between the new B-first fp8 stage2 pipeline and the legacy body on the decode/sparse-MoE shape (dim=6144x1024, E=128, k=8, tile_m=16, atomic) on gfx950. Faithful routing (~128 sparse blocks, ~4/16 rows valid) reproduces new=~220us vs legacy=~135us. ATT stall taxonomy shows the regime is memory-latency + synchronization bound, not compute/occupancy/epilogue bound: buffer_load ~45%, s_barrier ~23%, s_waitcnt ~21%; by source the gemm main loop (B weight load + sparse A gather) ~64% and the per-block sorted-id LDS-seed barrier ~22%; CShuffle epilogue only ~1.3%. Verified that raising occupancy (waves_per_eu=8; a single-buffered low-VGPR variant), shrinking the epilogue tile, and deepening the B prefetch to all-tiles-in-flight all leave perf unchanged (<=2%). Closing the gap requires porting legacy's X-major overlap schedule, a loop-structure rewrite kept off this shared path to protect the ~1100 TFLOPS compute-bound shape. Comment-only: records the findings so the ruled-out levers aren't re-tried. No behavior change; decode ~220us and compute-bound ~1091 TFLOPS unchanged; fp8 stage2 correctness suites pass (12 standalone + 24 2stage). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…ing/pong + CShuffle) The new gemm2 sized LDS as a_ping=max(a_lds,cshuf) + a_pong=a_lds. For large f32-output tiles (e.g. out_f32-L: tile_m=64,tile_n=256 -> cshuf=64KB) that is up to ~72-80KB, which fits CDNA4 (gfx950, 160KB LDS/CU) but overflows CDNA3 (gfx942, 64KB) -> launch failure (exit 1) on gfx942. Use one region sized max(2*a_lds_bytes, cshuf_bytes) with A ping/pong as offset sub-slices and the CShuffle epilogue reusing it after the main loop. This strictly reduces total LDS on both arches and keeps f32-L at exactly 64KB so it fits gfx942. gfx950: correctness (30 e2e + 12 standalone) and perf (compute-bound ~1093, decode ~220us) unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
coderfeli
changed the title
Profiled + ISA-diffed the fp8 stage2 ATOMIC epilogue slowdown vs REDUCE and vs the legacy body on gfx950 (cache-cleared, median-of-3): - new-atomic and new-reduce compile to instruction-identical bodies (same 4096 MFMA / 645 buffer_load / 68 barrier / 196 VGPR), differing only by 32 buffer_atomic_pk_add_f16 vs 8 buffer_store. - The per-element out-of-tile guard is constant-folded to ~2 cheap ops/group; atomics are already wide/coalesced (128b/lane, offset:4/8/12). So the leading guard-ALU / coalescing hypothesis is refuted. - The slowdown is the intrinsic atomic RMW cost + topk-collision L2 contention: compute-bound 1495->1087 TFLOPS, decode 46.5->29 TFLOPS (same main loop). - Legacy-atomic (1354 / 48) is faster despite the SAME 32 atomics only because its rolled main loop uses fewer VGPRs (170 vs 196) => 3 vs 2 resident blocks hide the atomic drain tail. That is a property of the shared main loop, not the epilogue, and the loop-structure rewrite is deliberately kept off this shared path to protect the compute-bound shape. Reduce mode remains the fast path. Comment-only; no behavior change. fp8 stage2 correctness (test_moe_stage2_standalone 12 + test_moe_gemm_2stage 24) stays green. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…ancy The new B-first fp8 stage2 main loop fully unrolled the k-tile loop (range_constexpr(num_tiles)), keeping both ping-pong prefetch buffers live across the whole (num_tiles-deep) unroll -> 196 VGPR -> only 2 resident blocks/CU. That could not hide the atomic RMW-drain tail, so the ATOMIC epilogue regressed vs the legacy X-major body (rolled, 170 VGPR, 3 blocks). Roll gemm2's k-tile loop into a single-buffer scf.for. All fragment/LDS state is fixed rmem accumulated in place, so nothing needs to be loop-carried; the targeted vmcnt wait still leaves the B register loads in flight across the ds_write + MFMA (intra-tile overlap kept), only the cross-tile prefetch is dropped. Two barriers per iteration enforce the single-buffer WAR/RAW ordering. Measured on gfx950 (cache-cleared median-of-3), new-rolled vs forced-legacy in-session: - gemm2 VGPR 196 -> 130 (0 spill), 2 -> 3 resident blocks/CU. - CB stage2-atomic (t32768 d8192x8192 e16k4): 1093 -> 1263 TFLOPS (legacy 1354, gap -19% -> -6.7%). - CB stage2-reduce: 1514 -> 1478 TFLOPS (still > legacy 1363, +8.4%). - decode stage2-atomic (t64 d6144x1024 e128k8): ~29 -> ~29 TFLOPS (unchanged; num_tiles=4, only 78 VGPR, NOT occupancy-bound -- its gap vs legacy 48 is the documented B-first vs X-major memory-pipeline difference, not fixable by the roll). gemm1 is left fully-unrolled (comment only): it is compute-bound and already low-VGPR (134 => 3 blocks/CU); rolling it to a single buffer removes the cross-tile overlap and regresses ~17% (1125 -> 938), so no change is warranted. Updates the now-inaccurate root-cause comment from b49e84d: the atomic gap was an occupancy property of the shared main loop and IS substantially closed here. Correctness (cache off) stays green: test_moe_gemm_2stage fp8 eager nomask out_{f16,bf16} S/M (24) + test_moe_stage2_standalone fp8 (12). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…ial why) Condense the 34-line profiling narrative into a 12-line perf note keeping the actionable facts (atomic was occupancy-bound, fixed by the k-tile roll; decode is memory/sync-bound and architectural; prefer reduce mode). No code change. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Condense the new-pipeline docstrings and multi-line inline comments to their essential why (builder docstrings, layout_helpers, routing/carve-out, rolled-loop, LDS-region, atomic-guard). Also fix a stale gemm1 routing comment that still said "CDNA4 only" after gfx942 was enabled. Comment-only; -133 net lines. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Resolve the modify/delete conflict from #948 (which deleted the whole moe_gemm_2stage/ package) by restoring the ported fp8 2-stage kernel: - Keep the ported gemm1.py / gemm2.py and new layout_helpers.py. - Restore __init__.py and reduction.py (gemm2 imports both), deleted on main but required by the port. - Restore test_moe_reduce.py (self-contained reduction-kernel test). - Add tests/kernels/test_moe_gemm_2stage.py: basic stage1/stage2 builder smoke tests (fp8 + bf16) guarding the port against API drift. Named to avoid the collision with main's repurposed test_moe_gemm.py. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…ion.py Ports the stage2 topk-reduce (aiter ROCm/aiter#4596) onto the tiled-copy layout API, replacing the legacy buffer_ops/scf.IfOp reduction.py. New kernels/moe/moe_gemm_2stage/moe_reduce.py: - make_layout_tv / make_tiled_copy / partition_S/D / zipped_divide / memref_load_vec, uniform soffset=k*row_stride, 64-bit base-ptr fold (voffsets stay i32-safe for X > 4 GiB). - Unifies dense (f16/bf16/f32) + fp8 MXFP8 route-out ([N fp8 | N/8 e8m0]) via fx.Constexpr params; fuses the EP gather valid = expert_mask[topk_ids[t,k]] != 0 (expert_mask + topk_ids, not a precomputed mask). New launcher takes fx.Pointer args. gemm2.py: _MoeGemm2ReduceWrapper.__call__ now takes expert_mask/topk_ids and dispatches the pointer launcher via tensor_shim._run_compiled + from_c_void_p; compile_moe_gemm2_ex switches valid_mask=None -> use_mask: bool and passes num_experts=experts. test_moe_reduce.py rewired to the new interface. Verified: 8/8 non-large test_moe_reduce cases pass on GPU (dense + masked EP-gather); fp8 path traces/compiles/runs. Bandwidth on par with the legacy kernel (~5 TB/s on prefill shapes; ~15% lower latency on small decode). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Collaborator
Author
The new layout-API fp8 stage2 atomic (accumulate) epilogue regressed vs the legacy c_shuffle_epilog: decode 34->11us baseline gap, prefill-S 3679us vs legacy 774us. rocprofv3 showed TCC_ATOMIC 4.25x higher for identical output bytes (all requests already 64B), i.e. pure lane->address coalescing loss in the CShuffle read/scatter TV layout. Root cause 1 (primary): the CShuffle read/scatter TV layout was row-straddling. A 64-lane wave spanned 16 rows x 32 channels (~4x atomic instructions). Replace with a CHANNEL-MAJOR layout mirroring legacy: 32 lanes walk contiguous channels of ONE row, row changes only every 32 lanes; the other 256/32=8 threads index the token/row dim. Parameterized by contiguous_n/BM, not hardcoded to 256. Root cause 2 (primary): atomic e_vec must be 2, not 8. _buffer_atomic_pk emits one pk-pair per instruction; with e_vec=8 the 4 pairs per lane are stride-8 channels apart, so each pk-pair SIMD instruction is lane-strided/uncoalesced. e_vec=2 makes the 32 lanes of one atomic instruction hit 64 contiguous channels (fully coalesced). This alone dropped TCC_ATOMIC from 234M to 58.7M, exactly matching legacy. Reduce mode keeps wide e_vec=8 for coalesced buffer stores. Root cause 3 (decode): padding-row atomics still counted as L2 TCC_ATOMIC traffic even when OOB-clamped. With the now row-uniform channel-major layout, skip the whole atomic emission per row via a (near) wave-uniform scf.IfOp(tok < row_limit), mirroring legacy's row predicate. Copy atoms (read + buffer) are sized to e_vec*out_width instead of a fixed 128b so narrow-e_vec paths stay consistent. bench_g2.py: fix reduce-mode correctness reference (kernel writes [tokens, topk, model_dim]; allocate that shape and sum over topk before the cos check) so reduce mode no longer reports cos=nan. Results (gfx950, cache disabled): decode atomic 36->11.3us (legacy 10.8), prefill-S atomic 3679->785us (legacy 774); reduce modes unchanged/improved; TCC_ATOMIC 234M->58.7M == legacy. cos=1.0 on all atomic configs. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
black line-length 120 + ruff (drop unused typing.Optional in bench_g2.py). No functional change. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Parameterize the fp8-only 2-stage layout-API pipeline on in_dtype ("fp8"|"int8").
int8 reuses the entire fp8 path -- same 1-byte element sizing, LDS swizzle, copy
atoms, A-gather/scatter, silu, CShuffle epilogue and atomic/reduce store -- and
only:
* swaps the MFMA atom to mfma_i32_16x16x32_i8 with an explicit i32 accumulator
(a default-f32 acc on the Int8 MFMA hard-aborts CDNA verification), and
* converts the i32 accumulator to f32 in the dequant step before applying the
per-token activation scale * per-channel weight scale (identical scale algebra
to fp8).
make_1x4_tiled_mma gains an acc_dtype arg. The dequant helpers now return the
(possibly fresh f32) fragment; for fp8 they return the same in-place-mutated
fragment, so the fp8 path is unchanged.
fp8 codegen is provably unaffected: 21_final_isa.s is md5-identical before/after
for gemm1 f16 and gemm2 f16-atomic. int8 numerics cos ~1.0 across stage1/stage2,
atomic/reduce, f16/bf16/f32, tile_m 16/64. int8 latency matches or beats the
legacy int8 kernels on prefill-S and decode.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ther) int8smooth is int8 with slot-major activations: X is pre-expanded to [topk*tokens, K] and scale_x to [topk*tokens], both indexed by row_ts = slot*tokens + token (fused sorted id: token = fused & 0xFFFFFF, slot = fused >> 24). This differs from int8 in stage1 ONLY. Stage1 (gemm1.py): route int8smooth through the int8 pipeline and extend the two localized touch points -- the A-gather token decode and the activation-scale load -- to the slot-major row. The A-gather record bound and scale_x view span topk*tokens rows. All new code is const_expr-gated on is_int8smooth so the fp8 and int8 paths are byte-identical (verified: 21_final_isa.s md5 unchanged vs the prior commit for fp8/int8 gemm1 f16 and gemm2 f16-atomic). Stage2 (gemm2.py): the smooth scale is applied host-side to A2 before quant, so stage2 just accepts the dtype name and routes to the existing int8 path -- no int8smooth branching. layout_helpers.py: _TensorWithIndex gains a token_slot_tokens knob for the slot-major rank-2 gather; the output-scatter (rank-3, plain token) is unchanged, preserving the sentinel/OOB-drop invariant. Tests: add int8smooth stage1 (slot-major gather guard) and stage2 cases. The stage1 tests fail (cosine ~0.96 < 0.99) if slot*tokens+token is swapped for plain token -- verified. Perf (gfx950, median us): int8smooth stage1 prefill-S 445 / decode 27.8 vs legacy 681 / 44.4 (1.53x / 1.60x faster); in the new-int8 neighbourhood (422 / 23). int8/fp8 unchanged. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Phase 1 probe (/tmp/w4a8_probe) proved the recorded A-operand blocker is GONE on the current toolchain (post-#945): feeding a computed nibble-unpacked value into the MFMA A-fragment lowers cleanly and is bit-exact. So W4A8 keeps the new pipeline's B-first layout (weight = MFMA A-operand) and unpacks in-register. This commit adds the int4 weight views + unpack helper (layout_helpers.py) and wires gemm1's weight load. Numerics not yet correct: the packed 64b tiled_copy splits the K/ki dims differently from int8's 128b load; fixing the packed load layout next. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Phase 1 (probe scripts in /tmp/w4a8_probe, not committed) conclusively shows the recorded W4A8 A-operand blocker is GONE on the current toolchain (post-#945): feeding a COMPUTED nibble-unpacked int4->int8 value into the MFMA A-fragment (tiled_mma.make_fragment_A + .store) lowers cleanly AND is bit-exact vs an int4xint8 reference (probe_correct.py, maxdiff 0). So no B-operand restructuring is needed for the compiler; W4A8 can keep the B-first weight=A layout and unpack in-register. Phase 2 is INCOMPLETE and int4 is gated OFF. Exhaustive probing shows the packed-int4 weight cannot be driven through the pipeline's generic flat_divide + make_tiled_copy_A in the int8 MFMA A-fragment's ki-separated kpack layout: a half-width packed load always addresses K contiguously, whereas int8's tiled_copy reads ki0/ki1 from kpack groups 1024 int8 apart. Groups 0,1 of the fragment come out bit-exact but groups 2,3 read the wrong K region, and no static reorder fixes it (the values differ, not just their order). The correct fix is an explicit ki-correct preshuffle-address loader mirroring the legacy load_b_pack_k32; that is the remaining work. This commit lands only the verified/reusable parts and keeps every existing dtype byte-identical: - gemm1 builder body is UNCHANGED vs 2c714ba (verified by diff); only the compile_moe_gemm1 arg validation gains an 'int4' branch that raises NotImplementedError with the diagnosis, plus a docstring note. - layout_helpers gains make_gateup_weight_view_int4 / make_weight_view_int4 (int8's view with byte strides halved, element_num=8) and unpack_int4_weight_frag (the 7-op even/odd sign-extend unpack), documented as scaffolding for the follow-up loader. - gemm2 and all tests untouched. Verified: 47/47 existing tests (test_moe_gemm_2stage + test_moe_reduce) pass; int4 fails fast; black+ruff clean. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Add fxh.make_preshuffle_b_layout_int4 / load_weight_int4_frag: an explicit ki-correct preshuffle-address loader that reproduces the int8 MFMA A-fragment's ki-separated kpack byte addressing, then 7-op nibble-unpacks each dword into the A-fragment (weight stays the A-operand). Wire it into gemm1 (gate+up) and gemm2 (plain); remove the int4 NotImplementedError gate. Replaces the incomplete half-width-view scaffolding. Numerics vs legacy int4 reference: cos=1.0 for gemm1 (tile_m 16/64, f16/bf16) and gemm2 (atomic/reduce, tile_m 16/64, f16/bf16). gemm1 even/odd nibble-swap perturbation drops cos to 0.002 (de-interleave order load-bearing). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…r perturbation) Parametrize test_moe_gemm1_numeric / test_moe_gemm2_numeric with in_dtype='int4' (W4A8) and add test_moe_gemm1_int4_perturb, which asserts the correct nibble packing matches (cos>0.99) while the even/odd de-interleave swap collapses (cos<0.5). Adds _pack_shuffled_int8_to_packed_int4 helper. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Read the packed weight as i32 dwords (one buffer_load = 4 packed bytes = 8 int4) instead of 4 scalar byte reads, and reuse one 8-byte kpack load for both MFMA k-halves. Drop the int4 weight ping-pong to a single buffer (the explicit load is serialized against its own dword loads, so the second buffer only inflated VGPR): gemm1 VGPR 246->158 (no spills), A-LDS keeps its ping-pong. Weight is loaded just-in-time before each MFMA. Perf (gfx950, median us, cos=1.0; legacy targets in parens): prefill gemm1 637 (755), gemm2 atomic 745 (758) decode gemm1 32.6 (58.7), gemm2 atomic 11.2 (14.3) fp8/int8/int8smooth ISA byte-identical to cf8c2c3 (int4 paths are const_expr-gated). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…unch helper Extract the identical compile+launch+arg-tuple block (repeated 4x across gemm1/gemm2 and their int8smooth variants) into a single _launch() helper. Same 52 test cases, same numerics, same assertions. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Cut historical narration, restatements of adjacent code, and over-long multi-line blocks in gemm1/gemm2/layout_helpers/moe_reduce/__init__. Preserved (and tagged BUG GUARD) the load-bearing constraint notes: e_vec narrow-on-atomic, channel-major CShuffle TV, sorted_lds sentinel seeding, global-atomic no-bounds-check, int8 i32-acc requirement, and the W4A8 ki-separated addressing + nibble order. ISA byte-identical for all 6 gate configs (fp8/int8/int4 x gemm1/gemm2-atomic). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Extract the identical 10-line routed-weight fragment loader into layout_helpers.load_sorted_weight_frag; both _apply_doweight closures now call it. ISA byte-identical for all 6 gate configs. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…, tiles) Add a FAST core (default CI) and an exhaustive large_shape sweep over ragged token counts (1,3,7,31,33,129), varied experts (4..128) and topk (1,2,6,8), and tile triples (tile_m 16/32/64/128, tile_n 64/128/256, tile_k 128/256), across fp8/int8/int4. Kernel constraints are derived from the builder asserts and invalid combos are skipped with explicit reasons. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Add output-canary/guard-region, input-poison, ragged-tail, sentinel-token0 regression, and num_valid_ids block-guard tests targeting the kernel's M-safety mechanisms head-on (num_records_bytes clamp, sentinel LDS seed, block guard, atomic OOB-redirect). All four dtypes where supported. Perturbation-verified (see PR report): disabling the sentinel seed collapses the sentinel/canary tests (token-0 cos ~0.62), and the bytes-vs-elements atomic output descriptor bug collapses the gemm2 canary/shape atomic tests (cos ~0.70). Docstrings note that the input-poison tests are defense-in-depth tripwires, not single-factor num_records isolations. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…erclaim Parametrize the per-dtype gaps in the moe_gemm_2stage suite: - fast shape sweeps: gemm1 crossed over all four dtypes (fp8/int8/int8smooth/ int4) on a trimmed ragged-M core; gemm2 over the three distinct paths (fp8/int8/int4). int8smooth stage2 is bit-identical to int8, pinned once by test_moe_gemm2_int8smooth_dispatch_equiv instead of being re-swept. - large_shape gemm1/gemm2 sweeps: add int8smooth. - output_canary / input_poison / sentinel_token0: add int8smooth to gemm1 (slot-major A/scale_x builder reused via _prep_gemm1 / _build_stage1_int8smooth); gemm2 gets single int8smooth dispatch-equivalence pins, not full re-sweeps. - num_valid_ids_guard: parametrize fp8/int8/int4 + one int8smooth pin. Trimmed the fast shape cross so the dtype dimension fits the CI budget (fast selection 118 -> 153 tests, ~46s -> ~51s), pushing the full cross into large_shape; ragged-M shapes stay in the fast set. Correct the TASK B header + output_canary docstrings: the historical num_records_bytes bytes-vs-elements bug was on gemm1's INPUT A descriptor and under-sized it 2x for BF16 input; BF16 input no longer exists and every surviving input dtype has elem_bytes==1, so that exact bug is unreproducible on the input descriptors. The demonstrated teeth are on the atomic OUTPUT descriptor (f16/bf16=2B, f32=4B), the same mistake class at a different site, isolated by test_moe_gemm2_output_canary. Tests only; no kernel behavior changed. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The gemm2 int4 fixtures built weights with torch.rand(...)*s (all-positive uniform), which masked nibble de-interleave errors: a within-group reorder of all-positive values barely changes the dot-product, so cosine stayed ~1.0 regardless of packing correctness. gemm2 shares load_weight_int4_frag with gemm1 but uses a different N/K mapping, so gemm1's perturb test did not cover it. Phase 1 verified (signed randn weights, dtypeMax=7) that gemm2's int4 path IS correct: cos=0.999997, max-rel-err=0.026 vs the torch reference. Swapping the nibble de-interleave collapses it to cos=-0.008, max-rel-err=13.1. - Add test_moe_gemm2_int4_perturb mirroring the gemm1 guard, with signed weights and a stricter max-abs-relative-error secondary check. - Switch the gemm2 fixtures (_run_gemm2, _run_gemm2_int8smooth, _prep_gemm2) from all-positive uniform to signed zero-mean randn weights: realistic (A2 silu output is genuinely signed) and strictly more sensitive to sign/ordering errors. All four dtypes stay green under signed data. No kernel code changed. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…eanup Root cause of the stage1 fp8 regression vs v0.3.0 (measured with rocprofv3 PMC counters, not ISA inference): the stage1 rewrite tiled N at contiguous_n = tile_n//2 = 64 channels/block, so the big shape launched 128 blocks/row instead of 64 -- 2x the workgroups (SQ_WAVES 10.6M vs v0.3.0's 5.3M) for byte-identical MFMA work (SQ_INSTS_MFMA equal). The 2x workgroups paid 2x per-block barrier/LDS-sync overhead (69 static barriers/block; SQ_WAIT_INST_LDS 4.9x, SQ_WAIT_ANY 1.3x) and drove both the ~16-30% median slowdown and the high variance / slow-dispatch tail. LDS bank conflicts were RULED OUT (NEW had fewer than v0.3.0); occupancy was RULED OUT (NEW uses fewer VGPRs, 134 vs 170, same 16KB LDS). Fix: fp8/int8 use a 128-channel block (num_acc_n=2), halving the big-shape grid back to 64 blocks/row. SQ_WAVES now matches v0.3.0 (5.3M) and SQ_WAIT_ANY drops below it. contiguous_n is capped at 128 so tile_n=256 keeps 2 blocks/row -- collapsing to a single 256-channel block regressed the inter_dim=256 shape ~40%. int4 keeps the narrower BN//2 tile: its single just-in-time weight buffer spills if fed a doubled per-block N (+63% measured). Measured (gfx950, cache off, cosine 1.0 all dtypes): big gemm1 fp8: 34.4ms -> 24.8ms median (min 22.2ms, now beats v0.3.0 24.8ms) big gemm1 int8: 29.4ms -> 23.1ms median big gemm1 int4: 22.4ms (unchanged), inter256 int8 433us / int4 644us (unchanged) small gemm1 (all dtypes) unchanged Also modernize the MMA loops to fx.gemm (project convention): gemm2's single accumulator and gemm1's dual gate/up accumulators. Perf-neutral and ISA-equivalent (identical MFMA/ds/barrier counts; gemm2 opcode histogram identical bar one fewer s_waitcnt). Removed the now-unused _m_reps/_n_reps. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…_dim=256) The gemm2 mma_atom_call -> fx.gemm swap is neutral on the official big shape (10612 vs 10638us) but costs ~4.5% on the inter_dim=256 prefill shape: 810-819us with fx.gemm vs 776-782us with the explicit loop, reproducible across 4 vs 3 runs. Isolated by reverting gemm2.py alone with the gemm1 tiling fix left in place. gemm1 keeps both its grid-doubling fix and its fx.gemm conversion -- those measure clean on both shape families. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
int4 stage1 at inter_dim=8192 is ~1.6% faster with a K-major serpentine fx.gemm traversal than the default schedule (22.05ms vs 22.40ms, reproducible with non-overlapping ranges), and neutral at inter_dim=256. fp8/int8 measure best at the default on both shape families, so the override is int4-only and compile-time gated. gemm2 keeps its explicit mma_atom_call loop: no single traversal_order wins across dtypes at the official big shape -- kmn_serpentine matches the loop for fp8 but costs int4 6.4%, while knm wins int4 8.2% and costs fp8 5.0%. Documented in-code. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Replace the last raw-dialect sites in the moe_gemm_2stage package with the current fx.* API. Zero codegen change (ISA byte-identical for fp8/int8/int4 gemm1, fp8/int8/int4 gemm2 f16-atomic, and gfx942 bf16-atomic compile). - gemm1/gemm2 launch wrappers: arith.index_cast(T.index, ...) -> fx.Int64(...) for the host-side grid math (matches moe_reduce.py). - gemm2 scf.for induction cast: arith.index_cast(T.i32, iv) -> fx.Int32(iv). - gemm2 atomic buffer resource: buffer_ops.create_buffer_resource(...) -> fx.rocdl.get_buffer_rsrc(fx.get_iter(fx.rocdl.make_buffer_tensor(...))). - layout_helpers atomic guards: scf.IfOp + _if_then -> plain Python `if` (the transformed _TensorWithIndex.copy traces both into the same scf.if; every atomic stays inside its predicate, BUG GUARD #4). - Drop now-dead imports (arith, scf, T, buffer_ops, _if_then). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Tighten the stage1 channel-block rationale (all measured numbers kept) and drop a restated parameter-narration line in gemm2. No code change; all 7 BUG GUARD comments and the corrected traversal_order rationale preserved. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Add an argparse benchmark CLI to tests/kernels/test_moe_gemm_2stage.py
(__main__ falls back to pytest when no args are passed) that times stage1
and stage2 in both atomic and reduce modes for fp8/int8/int8smooth/int4,
printing a correctness cosine and the log-line formats run_benchmark.sh
already parses (stage1 grep block + _emit_moe_s2_rows).
Wire a new `moe2stage` op into scripts/run_benchmark.sh: MOE_2STAGE_SHAPES
(reusing the official MOE_SHAPES entries), a RUN_MOE_2STAGE flag in the
--only parser and --list, and a benchmark section gated on IS_CDNA that
emits moe2stage_gemm1 / moe2stage_gemm2_{atomic,reduce} rows.
Harness-only; no kernel source changed.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
_bytes_stage1/2 multiplied the full weight matrix by the routed-row count, inflating TB/s ~4 orders of magnitude (749 TB/s reported, vs single-digit HBM peak). Weights are read once per expert and shared across a tile-block's rows, so count them unique -- matching test_moe_gemm.py's convention. Big shape gemm2 atomic now reads 0.381 TB/s against the harness's own documented example of 0.377 TB/s. TFLOPS and us were already correct. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The fp8 MOE_SHAPES rows have been silently dead since #948 rerouted tests/kernels/test_moe_gemm.py to mxfp_moe: the invocation still exits 0 but prints nothing, so run_benchmark.sh counted a success and emitted zero rows. Point the existing moe section back at the package's own CLI, restoring the original moe_gemm1 / moe_gemm2_atomic / moe_gemm2_reduce labels, shapes and flags. Drops the separate moe2stage op added a commit earlier -- it duplicated what the moe section is supposed to do. --skip_ref/--compare_aiter_ck now accept the value form ("--skip_ref false") the harness passes. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
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
1 participant
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.
Summary
Ports the MoE 2-stage MFMA kernels (
kernels/moe/moe_gemm_2stage/) from the legacy authoring API (SmemAllocator/SmemPtr+ rawbuffer_ops+ rawrocdl.mfma) to the currentfx.*surface (make_buffer_tensor+SharedAllocator+ tiled copy/MMA +fx.copy_atom_call), and makes the new pipeline the canonical path for fp8 (both stage1 and stage2).Scope: CDNA3 (gfx94, MI300) + CDNA4 (gfx95, MI350/MI355) — both CI-green.** fp8 routes to the new gemm1/gemm2 builders on both arches. Two CDNA3 accommodations: the packed bf16-output buffer atomic (
buffer_atomic_pk_add_bf16, gfx950-only) keeps the atomic+out=bf16 combo on the legacy global-atomic path on gfx942; and gemm2 LDS is packed into onemax(2*a_lds, cshuf)region so large f32-output tiles fit CDNA3's 64KB LDS (vs CDNA4's 160KB). Validated: mi325 (gfx942) + mi355 (gfx950) test jobs both pass.What's on the new pipeline (fp8)
_build_moe_gemm1_fp8_gateup(gemm1.py) and_build_moe_gemm2_fp8(gemm2.py) — B-first MFMA (weight=A-operand, activation=B-operand) with a 2-stage A-LDS ping-pong + B register prefetch + LDS-read hoist. The builders are dtype-parametric (fp8/bf16) but only fp8 is routed (see below).layout_helpers.py(tiled-MMA/copy factories, sorted gather/scatter_TensorWithIndex,silu_pair, pk-atomic epilogue).compile_moe_gemm1,compile_moe_gemm2,compile_moe_gemm2_ex,MoeGemm2Mode,_MoeGemm2ReduceWrapper,compile_moe_reduction.Performance (gfx950, median, ~14% clock noise)
Correctness (cache-off)
test_moe_gemm_2stagefp8 full matrix (gemm1+gemm2 x atomic/reduce x out f16/bf16/f32 x S/M/L, mask/nomask, eager/graph): 84 passed locally; green on gfx942 + gfx950 in CI.test_moe_stage2_standalone -k fp8(incl. tile_n=256, decode/prefill): 12 passed.Notable fixes / gotchas
out_elem); a bf16-vs-f16 mismatch silently reinterprets the bit pattern (0x4480 -> 1024.0 vs 4.5).num_records_bytes = M*K*elem_bytes(bytes, not elements).[tokens, topk, inter]so the sorted id decodes to(token, slot); a rank-2 view reads slot-0 only (invisible on tame data, wrong on real routing).buffer_atomic_addpath + a row-guard to drop out-of-tile grid slots (plain stores rely on buffer OOB clamp; atomics don't).Deferred to follow-ups
MFMA(16,16,16)variant (the builder currently emits the gfx950-onlyMFMA(16,16,32)bf16 op, which aborts compilation on gfx942); (2) a gfx950 crash atin_dtype=bf16 + out_dtype=f32at large tile_m. bf16 uses the proven legacy path meanwhile.ub.poisonafter rmem-SSA promotion / LLVMcastIsValid). Fix path: an A-first int4 variant (weight as MFMA B-operand, like legacy). int4 stays on legacy.reduction.pyremain on legacy.Notes
mainhas advanced since this branch point; may need a rebase before merge.Generated with Claude Code
Perf note (gfx950)
fp8 on par with legacy: stage1 compute-bound ~1266 (legacy ~1297) TFLOPS, stage2 atomic ~1100 (parity), reduce ~1520; decode stage1/reduce at parity. Fixed a ~600x atomic-scatter regression on small-tile/decode shapes (out-of-tile lanes were serializing on out[0]); decode-atomic residual ~1.6x vs legacy remains (small-tile issue-slot overhead, decode's preferred reduce mode is on par).