[MoE] Port moe_gemm_2stage (stage1+stage2) to the new fx.* pipeline — fp8 (gfx942 + gfx950) - #947
Open
coderfeli wants to merge 14 commits into
Open
[MoE] Port moe_gemm_2stage (stage1+stage2) to the new fx.* pipeline — fp8 (gfx942 + gfx950)#947coderfeli wants to merge 14 commits into
coderfeli wants to merge 14 commits into
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
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).