Skip to content

[RadeonFlow] flydsl mxfp4 a4w4 MoE: improv perf - #3861

Merged
coderfeli merged 9 commits into
ROCm:dev/randomflow_prfrom
RadeonFlow:radeonflow_pr_flydsl_perf
Jun 23, 2026
Merged

[RadeonFlow] flydsl mxfp4 a4w4 MoE: improv perf#3861
coderfeli merged 9 commits into
ROCm:dev/randomflow_prfrom
RadeonFlow:radeonflow_pr_flydsl_perf

Conversation

@ColorsWind

Copy link
Copy Markdown
Contributor

No description provided.

ColorsWind and others added 9 commits June 21, 2026 17:16
Pass bare data_ptr() device addresses (fx.Int64) for the global buffer args
instead of full memref descriptors. The kernels only need base pointers
(contiguity + compile-time sizes), so the addresses pack contiguously into
kernargs -> coalesced s_load prologue. ~7% faster at decode (M=4/8),
converging to parity by M=64. cos preserved.
- backend-managed K-loop waitcnt: drop the hand-tuned inline-asm vmcnt that
  the LLVM waitcnt pass double-inserted (the large-M scheduling bubble); let
  rocdl.barrier() own the load->ds_read wait.
- epilog output-address strength reduction: hoist the loop-invariant row*N_OUT
  i64 multiply out of the per-element store loop; per-element offsets become
  compile-time constants folded into the store address.
- unsigned index division (divui/remui) for the non-negative grid/tile/count
  index math, dropping signed-division sign-correction SALU.

M=4096 1.11 -> 1.00 (parity with HIP), M=16384 1.10 -> 1.03; faster/parity
through M=2048 unchanged. cos preserved.
Use rocdl.exp2 (v_exp_f32) in silu_mul instead of the software math.exp2
expansion (matches HIP __expf). Removes ~31% of non-MFMA VALU (the v_ldexp +
range-clamp v_cmp/v_cndmask the polynomial emitted); large-M ~2-6% faster
(M=4096 1.07 -> 1.05). cos preserved.
…IATE

Drop the CSV `_MXFP4OUT` kernel-name flag from the dispatch gate so the
mxfp4-out gemm2 path is controlled purely by AITER_MXFP4_INTERMEDIATE=1
(plus the Kimi/DSR shape guard). The CSV no longer needs a _MXFP4OUT row.
…3861 conflicts)

Base (dev/randomflow_pr) rewrote fused_moe.py (~1340 lines); this branch's
only fused_moe.py change is the "gate gemm2 mxfp4-out solely on
AITER_MXFP4_INTERMEDIATE" commit. Resolution:

- aiter/fused_moe.py: take base's rewritten structure, re-apply the single
  semantic change -- drop the `mxfp4out and` term from the mxfp4-out gate in
  _mxfp4_a4w4_stage2 so it gates on AITER_MXFP4_INTERMEDIATE=="1" alone.
- aiter/ops/flydsl/kernels/mxfp4_gemm2.py: keep this branch's abs-fold epilog
  (llvm.fabs.f32 + maxnumf amax); drop base's now-dead `i7fff` integer-mask
  declaration. The ROCm#3840 LDS pipelining is identical on both sides.
@coderfeli
coderfeli merged commit d2d28f0 into ROCm:dev/randomflow_pr Jun 23, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants