Refactor mxmoe#892
Open
coderfeli wants to merge 6 commits into
Open
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>
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