Skip to content

[FlyDSL] MegaMoEV2: decode/prefill optimizations and Stage1 LDS broadcast race fix - #6127

Merged
yhl-amd merged 18 commits into
ROCm:mainfrom
yhl-amd:perf/megamoe-v2-decode-opts
Oct 5, 2026
Merged

yhl-amd merged 18 commits into
ROCm:mainfrom
yhl-amd:perf/megamoe-v2-decode-opts

Conversation

@yhl-amd

@yhl-amd yhl-amd commented Oct 3, 2026 •

Copy link
Copy Markdown
Contributor

Summary

MegaMoEV2 (FlyDSL fused dispatch + GEMM1 + GEMM2 + combine) optimizations and a fix, for DeepSeek-V4-Pro FP4 (MXFP8 x MXFP4) EP8:

  • opt-in combine masking of top-k slots with id -1 (forward(mask_invalid_slots=...), advertised by MegaMoEV2.supports_combine_mask), used for DP pad rows
  • Stage1 single-variant launcher keyed on its kernel (the JIT cache key ignored the spec object, so different Stage1 configs could share a compiled kernel)
  • 8-column GEMM1 quant epilogue (one 8-byte store per lane)
  • compact dispatch fast copy (two rows in flight) on by default
  • persistent Stage2 (128 CTAs) for fixed-slot buckets < 128
  • paired K loop + LDS-DMA A for small fixed-slot GEMM1 buckets (the rolled loop copied the prefetched B tile into loop-carried VGPRs and drained vmcnt every K step)
  • fix: Stage1 block broadcasts (work index, ready-order tile, fixed-slot ticket) get dedicated LDS slots. They lived in pool[0], which the A-tile LDS DMA and the C tile reuse; with skewed routing (ready-order scheduling, e.g. hash-routed layers on identical warmup tokens) a lagging wave read A bytes as its tile index and faulted loading W13 from a garbage expert.
  • non-temporal W2 loads for the fixed-slot 256 bucket (DP-padded decode graphs)
  • fixed-slot dispatch is validated through MTPR 1024: AITER_MEGA_FIXED_SLOT_MAX_MTPR=511 / 1023 / 2047 admits the 256 / 512 / 1024-token capacities (default unchanged: 255)
  • MegaMoEV2's environment variables (AITER_MEGA_*) are read through mega_moe/envs.py, which lists each with its default and purpose
  • test_mega_moe_v2.py --pad-tail-rows N: whole -1 rows under mask_invalid_slots=True come back as exact zeros, real rows stay within rtol of the reference, CUDA graph replay is exact
  • black and ruff on the touched files (formatting and lint only; the nested const_expr / device if guards keep noqa: SIM102 as elsewhere in the FlyDSL kernels)

Companion ATOM PR: ROCm/ATOM#2463 (--moe-backend mega: DP pad-row masking through mask_invalid_slots, and a configurable decode capacity that relies on the larger fixed-slot capacities).

Test plan

  • op_tests/multigpu_tests/test_mega_moe_v2.py on 8x MI355X, EP8 v4_pro:
    • fixed-slot with AITER_MEGA_FIXED_SLOT_MAX_MTPR=1023 (MTPR 512, bs 1..512) and =2047 (MTPR 1024, bs 1..1024), uniform and skewed per-rank token counts: output relL2 against the reference equals the compact path's at every size (the a8w4 quantization floor, ~0.053), CUDA graph replay exact
    • --pad-tail-rows on fixed-slot MTPR 128 / 512 / 1024 and compact MTPR 1024, including steps where every row is padding: pad rows exactly zero, real rows within rtol, graph replay exact
  • --check against epx + fused_moe with pad rows, fixed-slot MTPR 128/256 and compact MTPR 16384
  • skewed (same-expert) routing: 150 back-to-back calls x 4 sizes clean after the LDS fix (faulted within a few calls before)
  • End-to-end DeepSeek-V4-Pro DSpark, TP8 + DP attention + EP8 through ATOM (companion PR), 3600 s agentic replay at several concurrencies with zero request errors; GSM8K 5-shot (1319) with real speculative acceptance: 0.9606 strict / 0.9613 flexible (stock image 0.9530)
  • black and ruff (0.16.0) clean on the touched files

🤖 Generated with Claude Code

Honglie Yi and others added 11 commits October 3, 2026 08:38
- combine_no_stage1 can take the top-k ids and skip slots whose id is -1.
  Those slots are never dispatched, so their partial is whatever an earlier
  call left in the shared combine buffer. MegaMoEV2.forward(
  mask_invalid_slots=True) enables it for callers that pad with -1 ids (DP pad
  rows come back as exact zeros, partially masked rows are correct);
  supports_combine_mask advertises the argument. The id check costs ~1 us per
  decode layer but 9-12% of a prefill layer (8 x MI355X, compact MTPR 16384:
  1550 -> 1683 us at 2048 tokens/rank, 4678 -> 5248 at 8192), so it is off
  by default (AITER_MEGA_COMBINE_MASK=1 turns it on for every call).
- AITER_MEGA_CFG_OVERRIDE: per-path/per-bucket Stage1/Stage2 config overrides
  (JSON) for tuning, e.g. persistent Stage2 for the fixed-slot buckets 1..64
  (the default launches ~86k mostly empty CTAs; -10..-12 us/layer at 7..28
  tokens/rank).
- AITER_MEGA_FIXED_SLOT_MAX_MTPR=511 admits a fixed-slot MTPR-256 instance
  (224 tokens/rank: 436 us/layer vs 589 on the 16384 compact path, with a
  BF16 instead of FP8 combine wire).
- AITER_MEGA_S1_LAUNCH_HANDSHAKE=0 drops the fixed-slot LAUNCH_READY round
  trip, which the combine barrier already orders when every Stage1 is followed
  by Stage2 + combine_no_stage1 (-1..-2 us/layer).
- AITER_MEGA_COMBINE_GEOM: combine launch geometry by token count (tuning
  knob; without masking the tuning table is within 0.5% of the best measured).
- AITER_MEGA_TRACE=1: per-CTA s_memrealtime timeline of prepare, Stage1,
  Stage2 and combine (debug aid).

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: yihonglie <hyi@amd.com>
…ter fixed-slot sync knobs

- compile_mega_moe_stage1()'s launch() closed over a _Stage1KernelSpec object,
  which FlyDSL's JIT cache key does not inspect, so every Stage1 config shared
  one key and a cached launcher of another config could be reused. The
  launcher now captures the kernel's own cache key and launch geometry.
- AITER_MEGA_S1_OWNER_LIGHT=1: the fixed-slot owner publishes its counter
  resets with one agent release and a relaxed gate store instead of two agent
  releases plus a system-scope release store (each writes back L2).
- AITER_MEGA_S1_PRODUCER_FENCE=0: fixed-slot producers rely on store
  acknowledgement and only the last producer of a group issues the system
  release before COUNT_DONE.
  Both are correct (8 x MI355X, pad rows exactly 0) but worth only ~1 us per
  layer (T=56: 331.9 -> 331.3 us), so they stay off by default.
- AITER_MEGA_TRACE also records the owner's reset phase and each producer's
  atomic/copy completion.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: yihonglie <hyi@amd.com>
SiluQuantEpilogue's store pass handled 2 columns per lane: a 16-lane
shuffle-max per pair of values and one 16-bit store. With 8 columns per lane a
1x32 group spans 4 lanes (2 shuffles) and each lane stores its 8 FP8 values
at once; scales and rounding are unchanged (outputs bit-identical). 8 x MI355X
compact MTPR 16384: 1550 -> 1535 us/layer at 2048 tokens/rank, 2473 -> 2436 at
4096, 4685 -> 4623 at 8192; decode unchanged. AITER_MEGA_S1_EPI_EVEC=2 restores
the old pass.

AITER_MEGA_S2_SCATTER_VEC (8 default, or 16) sets the values per lane in the
Stage2 peer scatter; 16 measured neutral at prefill.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: yihonglie <hyi@amd.com>
AITER_MEGA_DISPATCH_FAST_COPY (compact producers keep two rows in flight),
AITER_MEGA_S2_FAST_FP8 (DPP group max, one packed scale store per row),
AITER_MEGA_S2_STATIC_K (fully unrolled GEMM2 K loop) and AITER_MEGA_S2_WPE
(Stage2 waves-per-EU hint). Only the fast copy paid off (next commit turns it
on); the others measured neutral or slower and stay off.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: yihonglie <hyi@amd.com>
Producers read two route ids back to back and keep both rows' loads in flight.
Prefill -0.5..-1.2% per layer (4096: 2436->2423 us, 8192: 4623->4577 us); decode (fixed-slot) unaffected.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: yihonglie <hyi@amd.com>
… 128

The non-persistent grid is sized for the MTPR capacity; its idle CTAs delay the
real tiles. T=7/14/28: -10..-11 us/layer, T=56 -1; --check PASS.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: yihonglie <hyi@amd.com>
The rolled GEMM1 K loop copied the prefetched B tile into loop-carried
registers, which forced a vmcnt drain of the prefetch every K step; with
register A copies the A load was also waited for in the same step. Fixed-slot
tiles with async A now run the K loop two steps per iteration (ping-pong B
registers; VGPR 138/174), and buckets <= 8 use async A. T=3 -11, T=7 -2
us/layer; 14..112 and compact prefill unchanged; --check PASS.
AITER_MEGA_S1_UNROLL_K (0/1/2) stays as an experiment override.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: yihonglie <hyi@amd.com>
The consumer loop broadcast first_work, then scheduled_first (ready-order
path), through pool[0], which the A-tile LDS DMA and the GEMM1 C tile also
overwrite. A wave that had not yet read the broadcast value then read the next
value or A bytes as its tile index and loaded W13 from a garbage expert
(memory access fault). Needed skewed routing (ready-order scheduling: DSV4 hash
layers on identical warmup tokens) plus a timing shift; the 8-column epilogue
made it reliable in ATOM warmup. first_work, scheduled_first and the
fixed-slot ticket now use their own LDS slots (+48 B, occupancy unchanged).
Also bound the 8-column epilogue's full-output view like the 2-column path.
Same-expert routing: 150 back-to-back calls x 4 sizes clean; prefill 2048/4096/
8192 1534/2415/4512 us, decode MTPR 256 --check PASS.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: yihonglie <hyi@amd.com>
C128 decode graphs are DP-padded to 224 rows/rank with ~84 real: Stage2 with
use_nt -14 us/layer (364 -> 350); fully real 224 rows -2 us; 160/56 unchanged;
--check PASS. persist_cu (64..256) made no difference.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: yihonglie <hyi@amd.com>
Formatting only; the pre-checks black job flags them.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: yihonglie <hyi@amd.com>
…to MTPR 1024

test_mega_moe_v2.py --pad-tail-rows N routes each rank's last N rows to -1
after the unmasked passes (so the combine input still holds their values) and
checks that forward(mask_invalid_slots=True) returns those rows as exact
zeros, keeps the real rows within rtol of the reference, and replays the same
output from a CUDA graph. This is the shape ATOM's DP pad-row mask produces.

Note in mega_moe_config that AITER_MEGA_FIXED_SLOT_MAX_MTPR=1023 / 2047 extends
fixed-slot to the 512 / 1024-token capacities: on EP8 v4_pro, bs 1..1024,
uniform and skewed per-rank counts, and with -1 tail rows, the fixed-slot
output matches the reference with the same relL2 as the compact path, and CUDA
graph replay is exact.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: yihonglie <hyi@amd.com>
@github-actions

github-actions Bot commented Oct 3, 2026

Copy link
Copy Markdown
Contributor

🏷️ CI Guide

Runs automatically on every PR:

  • ✅ Pre-checks (submodule verification, code formatting)
  • ✅ Aiter op tests (gfx942 + gfx950)
  • ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label Tests
ci:gfx1250-ffm-triton Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x Run an additional Triton test job on MI300X in PRs (added automatically when gfx942 configs change); main branch always runs both MI35X and MI300X
ci:triton-355 Run the full Triton test suite on MI35X, not only the tests the change affects
multigpu Aiter multi-GPU tests on the 8-GPU runner
ci:sglang SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 6127 --add-label <label>

One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.

yhl-amd and others added 7 commits October 3, 2026 08:52
Sort the trace imports, merge the two persist_cu=128 arms, and mark the
nested `if const_expr(TRACE_BASE)` / runtime-`if` guards `noqa: SIM102` as
elsewhere in the FlyDSL kernels: the outer test is resolved at trace time and
the inner one is a device branch, so they must stay separate statements.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: yihonglie <hyi@amd.com>
None of them is set by any caller, and each was measured neutral, slower or
within noise, or its finding is already the default:

- AITER_MEGA_S2_FAST_FP8, AITER_MEGA_S2_STATIC_K, AITER_MEGA_S2_WPE: neutral
  or slower; removes the DPP scale path, the unrolled GEMM2 K loop and the
  waves-per-EU hint (gemm2.py is back to main).
- AITER_MEGA_S1_UNROLL_K: the paired K loop the fixed-slot buckets use stays;
  only the env override and the fully unrolled mode it alone reached go.
- AITER_MEGA_S1_LAUNCH_HANDSHAKE, AITER_MEGA_S1_PRODUCER_FENCE,
  AITER_MEGA_S1_OWNER_LIGHT: within noise; the handshake, fences and owner
  resets are main's again (the handshake-free variant also relied on every
  Stage1 being followed by combine).
- AITER_MEGA_CFG_OVERRIDE: its finding (persistent fixed-slot Stage2) is the
  default now.
- AITER_MEGA_COMBINE_GEOM and the per-call combine geometry: the tuning table
  is already within noise of the best geometry, and the op's cfg can pin one.
- AITER_MEGA_S2_SCATTER_VEC: 16 was neutral; the scatter width is 8 again.

With the knobs unset the kernels and their cache keys are unchanged.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: yihonglie <hyi@amd.com>
The per-CTA s_memrealtime timeline was a debugging aid for tuning the
series. Remove trace.py and its hooks in prepare, Stage1, the fixed-slot
dispatch, Stage2 and the fused combine; mega_moe_prepare.py and
mega_moe_stage2.py are back to main. With tracing off (the default) the
kernels and their cache keys are unchanged.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: yihonglie <hyi@amd.com>
mega_moe/envs.py lists every AITER_MEGA_* variable MegaMoEV2 reads, with its
default and what it does; config, Stage1, GEMM1 and MegaMoEV2 read them
through it instead of os.environ. Values are still read when accessed, so
each keeps its point of effect (import, Stage1 compile, instance build).

GEMM1's epilogue width is now a build_fused_gemm1() argument that Stage1
passes from the same read it keys the kernel on, instead of a second
os.environ read inside GEMM1.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: yihonglie <hyi@amd.com>
- Comments state the mechanism without benchmark numbers.
- envs.py names each variable once: name -> (parser, default).
- test_mega_moe_v2.py: one _rel_l2() helper for the relL2 check in
  _run_size and _check_pad_tail (which also guards an all-zero reference),
  and _check_pad_tail's rank-wide flags go through _reduce_float.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: yihonglie <hyi@amd.com>
…the instance

- forward() hands the top-k ids to combine as an argument through
  _run_stage2 / _run_fused_stage2 / _run_aligned_pair_stage2. Nothing is
  kept between calls, so a direct Stage2 call no longer inherits the mask of
  the previous forward and the instance holds no reference to the caller's
  ids.
- supports_combine_mask is a class attribute.
- One loop preloads both combine variants.
- _run_combine_kernel takes the tok_map value; its signature is main's again
  apart from the new argument.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: yihonglie <hyi@amd.com>
… store helpers

- The single-variant Stage1 launcher closes over its kernel and launch ints
  instead of the opaque spec, so FlyDSL keys it itself; the hand-built
  _key_kernel / _key_geometry tuple goes.
- Stage1 knobs: FAST_COPY already folds in compact dispatch and PAIR_K is a
  bool (was k_unroll 0/2); kernel names are unchanged.
- One GEMM1 output descriptor in the epilogue's store width instead of an
  Int16 one plus out_vec_rsrc; evec is a required argument, and standalone
  compile_gemm1 states it (8 with a tile resource, else 2), replacing the
  silent override in SiluQuantEpilogue.
- The 2- and 8-column store passes share _tile_out, _row_target,
  _group_absmax and _e8m0_scale.
- emit_dispatch_payload: fast_copy is required and the two-row store is a
  _store_token helper next to _load_token.

Compile-only build of every MegaMoE AOT bundle (MTPR 128/512/1024/16384), the
single-variant Stage1 launchers and standalone GEMM1, against the previous
head: the final ISA is identical, with and without the fallback env knobs.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: yihonglie <hyi@amd.com>
@yhl-amd
yhl-amd marked this pull request as ready for review October 4, 2026 15:22
@yhl-amd
yhl-amd requested a review from a team October 4, 2026 15:22
@github-actions github-actions Bot changed the title MegaMoEV2: decode/prefill optimizations and Stage1 LDS broadcast race fix [FlyDSL] MegaMoEV2: decode/prefill optimizations and Stage1 LDS broadcast race fix Oct 4, 2026
@github-actions github-actions Bot added the FlyDSL label Oct 4, 2026
@yhl-amd

yhl-amd commented Oct 4, 2026

Copy link
Copy Markdown
Contributor Author

What this PR fixes

MegaMoEV2 (--moe-backend mega) for DeepSeek-V4-Pro DEP (TP8 + DP attention + EP8). Four problems:

1. GPU memory access fault during DSV4 warmup (bug)

Stage1 broadcasts block-wide scalars (work index, ready-order tile, fixed-slot ticket) through LDS: tid 0 writes, every wave reads after a barrier. They lived in pool[0], which the A-tile LDS DMA and the GEMM1 C tile reuse. With skewed routing (e.g. hash-routed layers on identical warmup tokens) the waves drift apart, a lagging wave reads A bytes as its tile index and loads W13 from an expert that does not exist. Each broadcast now has its own LDS slot.

Also fixed: the single-variant Stage1 launcher closed over an opaque spec object that FlyDSL's JIT key ignores, so different Stage1 configs could reuse one compiled launcher. It now closes over the kernel and its launch geometry directly.

2. DP pad rows paid dispatch, GEMM and combine (used by ROCm/ATOM#2463)

A padded DEP decode step runs running_tokens rows per rank; only the leading scheduled_tokens carry a request, and the identical pad rows all land on the same experts. MegaMoEV2 already never dispatches or computes a top-k id of -1, but the fused combine summed every top-k slot, so a -1 slot read a stale (possibly non-finite) partial. forward(mask_invalid_slots=True), advertised by MegaMoEV2.supports_combine_mask, makes combine skip -1 slots, so whole -1 rows come back as exact zeros. It is off by default because the id check costs prefill time; ATOM asks for it only on padded passes.

3. MTP decode graphs fell off the fixed-slot path (used by ROCm/ATOM#2463)

With speculative decoding the DP-padded decode graphs have running_bs * (k + 1) rows per rank, above fixed-slot's 128-token limit, so they ran the prefill-sized compact instance. AITER_MEGA_FIXED_SLOT_MAX_MTPR=511 / 1023 / 2047 admits 256 / 512 / 1024-token fixed-slot instances (default unchanged).

4. Small-batch decode and prefill overheads (on by default)

Where Cause Change
fixed-slot GEMM1, buckets <= 8 the rolled K loop copied the prefetched B tile into loop-carried registers and drained vmcnt every step paired K loop + LDS-DMA A copies
fixed-slot Stage2, buckets < 128 the grid is sized for the capacity; idle CTAs delay the real tiles persistent Stage2 with 128 CTAs
fixed-slot Stage2, bucket 256 W2 loads pollute the cache on DP-padded decode graphs non-temporal W2 loads
GEMM1 quant epilogue 2 columns per lane: a 16-lane max and a 16-bit store per pair 8 columns per lane, one 8-byte store
compact dispatch producers one dependent route-id read plus several load round trips per row two rows in flight

The last three keep fallback env vars. Every AITER_MEGA_* variable MegaMoEV2 reads is listed, with its default, in aiter/ops/flydsl/kernels/mega_moe/envs.py.

Latest commits

7635e5ff..984210c3 drop the experiment / tuning env knobs and the trace instrumentation, read the env vars through mega_moe/envs.py, pass the combine mask ids as an argument instead of instance state, and share the 2-/8-column epilogue store helpers.

They do not change generated code: a compile-only build (FlyDSL 0.3.4.1, gfx950) of every MegaMoE AOT bundle (MTPR 128 / 512 / 1024 / 16384), the same with the fallback env vars flipped, the single-variant Stage1 launchers, standalone GEMM1 and the masked combine gives final ISA byte-identical to the pre-cleanup head a3ee77e1 (109 kernel modules, register and LDS metadata included). The end-to-end and GSM8K results in the description were measured on that earlier head.

@yhl-amd
yhl-amd merged commit f4be7f1 into ROCm:main Oct 5, 2026
96 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants