Refactor mxmoe#892
Merged
Merged
Conversation
The fused mxmoe stage2 epilog is coupled to tile_m (atomic for
{16,32,64}, reduce/nonatomic only at 128). Run the one applicable mode
in the run_benchmark __main__ loop instead of the fp8-style atomic/reduce
sweep (which picked an invalid combo and crashed via pytest.skip), and
add skip_ref plumbing to the a8w4/fp4 e2e path.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
gemm2: round-robin the non-persistent atomic launch index across the 8 XCDs (bijective) to balance HBM-channel use, and software-pipeline the B tiles (one-ahead prefetch) so f32 accumulators stay in ArchVGPRs (2 waves/SIMD instead of 1). gemm1: drop the non-temporal B-load hint for BM==32 once total_m_blocks >= NE, so W1 stays L2-resident when each expert owns >=1 M-block (closes the large-M regression: 8192-row stage1 380 -> 304 us; small-M win kept). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Delete the FP8 block-scale 2-stage MoE kernel (kernels/moe/moe_blockscale_2stage/) and its test; drop it from arch_compat CDNA_ONLY_TESTS and the docs. Also fixes the now-stale mixed_moe_gemm_2stage doc reference (replaced by mxmoe). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
coderfeli
force-pushed
the
worktree-mxmoe-a8w4
branch
from
July 23, 2026 15:22
331aab6 to
31d6cc3
Compare
- Rename the fused MoE package kernels/moe/mxmoe -> kernels/moe/mxfp_moe (imports, test helpers _run_mxfp_moe_e2e / test_mxfp_moe_variants, docstrings, error strings, api docs). - Add deprecation notes to kernels/moe/moe_gemm_2stage/* (legacy SmemAllocator + buffer_ops API; refactor to the fx.* surface / mxfp_moe). - Fix stale mixed_moe_gemm_2stage references in docs (-> mxfp_moe); reflow a few pre-existing >120c error strings touched by the above. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…re-graft
Port aiter PR #4349 ("port mxfp moe to layout APIs") gemm1 back into FlyDSL's
fused MoE stage1, replacing raw buffer_ops / raw_ptr_buffer_load_lds /
llvm.load|StoreOp / manual LDS byte-pointer math with the layout API:
- global B/B-scale/hidden loads -> fx.rocdl.make_buffer_tensor + fx.copy with
BufferCopy128b/32b atoms (use_nt preserved via BufferCopy128b(b_aux))
- A / A-scale global->LDS async DMA -> fx.copy with BufferCopyLDS128b/32b
- LDS A reads and the f32 accumulator (union region) -> UniversalCopy tiles
- global scalar reads (eids/cumsum/m_indices/a-scale) -> fx pointer / copy
- aqout / ascaleout stores -> layout-indexed UniversalCopy tiles (crd2idx)
aiter #4349 is a4w4-only; the a8w4 (fp8 e4m3 A x mxfp4 W) path is re-grafted in
the layout idiom: fp8 A-tile doubles (KH_TILE=BK, A_ROW_BYTES=K), issue_a_load_lds
second BufferCopyLDS pass at +128, issue_a_ds_read split-16@64 v8i32 assembly,
_iq_pack_store fp8 (cvt_scalef32_pk_fp8_bf16 -> vector<2xi16>), and MMA a_cbsz=0.
The MMA stays raw rocdl.mfma_scale_f32_16x16x128_f8f6f4 (as in aiter #4349;
no scaled-MMA atom wrapper is used there). The gemm1 xcd_swizzle grid stays.
Perf parity gate (gfx950, median-of-3, compute-bound 3072x3072/NE=128 shapes):
stage1 layout within +/-3% of the raw baseline for a4w4 and a8w4 across
bm32/bm64/bm128 -- within gfx950 clock noise. Cold-correct (cache disabled):
test_mxfp_moe_variants 6/6; test_moe_gemm_2stage FP4-M/L a4w4+a8w4 atomic+reduce
32/32.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Replace the raw rocdl.mfma_scale_f32_16x16x128_f8f6f4 intrinsic (9 issue sites in mfma_cluster) with fx.gemm over CDNA4 scaled-MMA atoms, for BOTH a4w4 and a8w4, mirroring kernels/gemm/mxfp4_preshuffle.py's proven idiom: - _scale_mma_atoms(a_dtype) builds the 16 (opsel_a, opsel_b) atoms via fx.make_mma_atom(fx.rocdl.cdna4.MFMA_Scale(16,16,128, elem_a, Float4E2M1FN, opsel_a=, opsel_b=)). elem_a = Float4E2M1FN (a4w4) or Float8E4M3FN (a8w4) encodes the f8f6f4 cbsz (fp4->4, fp8->0), so the raw a_cbsz arg is gone. - mfma_cluster now issues fx.gemm(atom, cf, a_frag, b_frag, cf, scale_a=sa, scale_b=sb) with the SAME opsel encoding as the raw call: kMChunks==1 uses opsel_a 0 then 2; kMChunks>1 uses 0,1 then 2,3; opsel_b = kh*2 + in_b. The per-subblock scale_a / per-mni scale_b words are unchanged. - A/B fragments are the layout rmem tensors from the fx.copy migration (issue_a_ds_read / issue_b_load_j now return tensors, not loaded Vectors; the fp8 split-16@64 lo++hi shuffle is stored into an i32[8] tensor). Accumulators are pre-seeded to 0 rmem tensors and accumulate in place (A*B+C every iter), so the raw "init" (C=0 on first K-iter) path is subsumed; init kept in the signature but unused. The fp8 cvt ABI bitcast (vector<2xi16>) is unchanged. grep: 0 raw mfma_scale, fx.gemm issue path only. Cold-correct (cache disabled): test_mxfp_moe_variants 6/6; test_moe_gemm_2stage FP4-M/L a4w4+a8w4 atomic+reduce 32/32. Perf (gfx950, median-of-3, compute-bound 3072x3072/NE=128) vs 8fa1ac1: stage1 within noise or faster for a4w4 and a8w4 across bm32/bm64/bm128 (e.g. fp4 bm64 -8%, fp8 bm128 +0.1%). Perf-neutral. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Consolidate helpers duplicated across gemm1/gemm2 into mxfp4_gemm_common
and drop dead code left by the layout-API port. Pure relocation/trim; no
codegen change (verified cold: variants 6/6 + a4w4/a8w4 FP4-M atomic +
FP4-L reduce pass).
- Hoist _scale_mma_atoms (+_A_ELEM) to common; gemm2 now calls
_scale_mma_atoms("fp4") instead of its inline fp4xfp4 atom dict.
- Hoist _global_i32_buffer_view/_global_i32_buffer_tiles to common;
remove gemm1's nested copies and gemm2's module-level copies.
- Remove gemm1's local _udiv/_umod; import from common (gemm2 already did).
- Drop gemm1 mfma_cluster's dead `init` param (del init) and the two
init=... call args.
- Trim redundant fx.Int32(bx_i32) re-wrap; drop now-unused Float4/Float8
imports from gemm1/gemm2.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Consolidate the two duplicate global-scalar-access styles used by the mxfp_moe stage kernels into one canonical set in mxfp4_gemm_common.py, shared by both gemm1 and gemm2. - Relocate stateless scalar-access helpers (_global_i32_ptr / _global_i32_at / _global_i32_load / _global_scalar_tiles / _scalar_store / _layout_idx) and stateless numeric helpers (_pkmax_u16 / _inline_e8m0 / _silu_mul_batch, plus LOG2E) from gemm1 into common; gemm1 imports them back. Pure relocation: gemm1 final ISA is byte-identical. - Migrate gemm2's raw scalar reads (cumsum0 x2, eids) from llvm.load(T.i32, _global_ptr1(...)) to common's _global_i32_at(...). The stids/sweights invariant scalar loads are left raw (they carry the !invariant.load hint that _global_i32_at does not emit). gemm2 final ISA is byte-identical. - Drop the now-dead _global_ptr1 helper (no remaining callers; _gep1 / _global_base_ptr1 stay for non-scalar pointer needs). Verified codegen-equivalent by diffing dumped final ISA (21_final_isa.s) for both kernels: identical vs 6ebcda9. Cold correctness (cache off, RUN_QUANT=1): test_mxfp_moe_variants 6/6; a4w4+a8w4 FP4-M/FP4-L e2e 46 passed / 14 skipped. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
#880 moved buffer_ops out of flydsl.expr into kernels.common; the fused mxfp_moe kernels (added after #880) still imported it from flydsl.expr, breaking a fresh-build import (CI). Point mxfp4_gemm_common at kernels.common.buffer_ops (same API: get_element_ptr / create_buffer_resource_from_addr). Local venv is pre-#880 so this was masked locally. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…builder sentinel The mixed_moe_gemm_2stage builders were replaced by the fused mxfp_moe kernels, but test_moe_gemm.py kept the old fp4/a8w4 host-prep branches alive by binding the removed builders to a raising _mixed_moe_removed sentinel and skip-guarding the runners. That left dead branches that still referenced the deleted builders. Drop the sentinel and route fp4/a8w4 to the new kernels: run_moe_stage1, run_moe_stage2, and test_moe_stage2_standalone now delegate these dtypes to _run_mxfp_moe_e2e (flydsl_mxfp4_gemm1/2), and the dead re-quant block in test_moe_gemm_2stage is removed. a8w4 is treated as fully covered (a_dtype="fp8"). Non-fp4 dtype paths are unchanged. Verified on gfx950: fp4/a8w4 2stage + standalone + variants pass (231 passed / 836 skipped), plus w4a16 and fp8 standalone. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Two fixes surfaced while validating the fp4/a8w4 routing: 1. Benchmark output regression. run_benchmark.sh drives fp4/a8w4 through the CLI -> test_moe_gemm_2stage -> _run_mxfp_moe_e2e, which was correctness-only and never emitted the "FlyDSL MoE stage1/stage2" lines the harness greps for, so the a4w4/a8w4 rows went blank. _run_mxfp_moe_e2e now optionally benchmarks gemm1/gemm2 with run_perftest (num_iters>0) and prints via a single _print_mxfp_moe_perf helper that mirrors the runners' table format. num_iters/ num_warmup/in_dtype_label are threaded through the three delegation sites. Atomic epilog is timed into its buffer, then zeroed + relaunched once for the correctness check; nonatomic overwrites so it needs no reset. 2. _run_compiled duplication. kernels/moe/mxfp_moe/host.py carried its own copy of _run_compiled that duplicated kernels/common/tensor_shim.py. host.py now imports the shared one; the MLIR-Context unwind-on-cold-compile-failure is folded into the shared tensor_shim._run_compiled so no behavior is lost. Verified on gfx950: fp4/a8w4 CLI now emits stage1+stage2 timing (atomic & reduce) and the harness parser extracts them; fp4/a8w4 2stage+standalone+variants pass (237 passed / 836 skipped); rmsnorm (shared _run_compiled) unaffected. Co-Authored-By: Claude Opus 4.8 (1M context) <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.
Refactor mxmoe