moe gemm optimization - #4730
Open
yadaish wants to merge 22 commits into
Open
Conversation
…fetch,
configurable wave grid, and main-loop scheduling
Squashes the dev/yadai_cluster_exp_v2 line of work into a single commit on top
of main. Five files, no overlap with anything main has changed since the
branch point.
Workgroup-cluster multicast (cluster_n)
cluster_n > 1 launches (cluster_n, 1, 1) workgroup clusters whose peers share
a tile's A rows and differ only in n_tile, so one A / A-scale tensor_load can
serve the whole cluster. A non-zero workgroup_mask turns the load into
CLUSTER_LOAD_ASYNC, which rendezvouses with the peers itself -- no cluster
barrier is emitted, and on a 1-D cluster with a contiguous mask none is
needed. The DeepGEMM contiguous-M swizzle runs on cluster granularity, with
n_tile reassembled from the cluster-local id. Selected host-side by
_pick_cluster_n (env AITER_FLYDSL_MXFP4_CLUSTER_N, default off, "auto" picks
the largest supported degree dividing the N-tile count); a partially filled
cluster would stall, so only divisors are usable.
Cross-tile prefetch (next_stage)
Computes a k-tile's last k128 while lds-loading the next tile's first one.
Env AITER_TDM_NEXT_STAGE_PREFETCH, default on; it only engages when
num_buffers leaves a TDM in flight after the rotated wait, so there is no
force-on. The k-tile loop is a rolled scf.for, so the tile boundary cannot
carry a Python value across iterations: one k128 of A/B/scales lives in two
fixed rmem slots alternating on ksl parity, shared by the within-tile
prefetch and the cross-tile carry rather than kept as a separate structure.
The carry targets slot 0, so it precedes the subtile's WMMA except when that
WMMA reads slot 0 itself (a single-k128 tile), where it follows instead.
Configurable wave grid
grouped_moe_gfx1250 gains m_warp / n_warp (and gemm2 overrides), validated by
get_wmma_m_rep(). The WMMA M-repeat fixes the preshuffled A-scale layout the
quant kernels must produce, so it has to be derived from the wave tile rather
than from tile_m whenever m_warp > 1.
Main-loop depctr and epilogue scheduling
Five changes, each measured on t256x256x256 with --ep-mode fake --const-init 0:
- mma_group keyed on stage1_act (4 with activation, 8 without): gemm2 -5.0%.
- LDS addresses split into a per-lane region base plus a compile-time offset
that folds into the ds 16-bit offset: field, with lds_addr_keepalive()
pinning the bases so the allocator cannot reuse a live base as a later
ds_load destination: depctr_vm_vsrc 5 -> 0.
- Bases hoisted and pinned at the tile top: depctr_va_vdst 6 -> 1. Any
dependence window crossing a WMMA needs va_vdst(0), a full matrix-pipe
drain, so the fix is to keep the window from crossing one.
- One scheduling region per k128 plus FENCE_COVER_MMA / FENCE_READ_FRONT,
putting WMMA between the last ds_load and the tail fence: -3.8%.
Net main loop 358 -> 344 instructions, s_wait_alu 4 -> 3, no spills.
- buf_ptr_opaque() on the drain tail's carry target, hiding its LDS stage
base from constant folding. Unrolled, the base is a literal and the backend
re-associates (base + lane) + off into (base + off) + lane, spending one
v_add and one depctr wait per ds_load instead of sharing a base register
with 16-bit offset: immediates -- 63 address registers against the main
loop's 8. Applied to the carry target only: the buffer a tile computes from
must stay foldable or it reads wrong LDS. Epilogue 4596 -> 4499, s_wait_alu
325 -> 271, -2.1%.
Retuned tuned_grouped_fmoe.csv for the new configuration space.
…g-pong from zan/a8w4_pipeline - C staging pads the row pitch to STORE_N+16 on the bf16 passthrough path. Unpadded the pitch is a multiple of 32 dwords, so the 16 rows one ds_store_b128 writes all land on one bank (16-way); +16 cols spreads them to 4-way, the b128 floor. The TDM store reads the padded width and lets the per-dim OOB extent clamp the inner axis, so pad cols never reach global. - The post-compute drain's last tile emits no fence; the previous tile's rotated fence already covers the only buffer it reads. - vgpr_keepalive() pins each row's b128 store data across the next row's cvt batch, forcing fresh registers and removing the WAR vm_vsrc wait. 4130 -> 3955 us (t256x256x256 cn=4, random activations).
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
Co-authored-by: zanzhang <zanzhang@amd.com>
Co-authored-by: zanzhang <zanzhang@amd.com>
yadaish
force-pushed
the
dev/yadai_cluster_exp_v2_rebase
branch
from
August 13, 2026 10:04
bd98c6b to
8c5245e
Compare
Co-authored-by: zanzhang <zanzhang@amd.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
2 participants
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.
Motivation
Technical Details
Test Plan
Test Result
Submission Checklist