Repository navigation
[FlyDSL] MegaMoEV2: decode/prefill optimizations and Stage1 LDS broadcast race fix - #6127
Conversation
- 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>
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
One backend per PR: PR title tags & labels: |
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>
What this PR fixesMegaMoEV2 ( 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 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 3. MTP decode graphs fell off the fixed-slot path (used by ROCm/ATOM#2463) With speculative decoding the DP-padded decode graphs have 4. Small-batch decode and prefill overheads (on by default)
The last three keep fallback env vars. Every Latest commits
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 |
Summary
MegaMoEV2 (FlyDSL fused dispatch + GEMM1 + GEMM2 + combine) optimizations and a fix, for DeepSeek-V4-Pro FP4 (MXFP8 x MXFP4) EP8:
forward(mask_invalid_slots=...), advertised byMegaMoEV2.supports_combine_mask), used for DP pad rowspool[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.AITER_MEGA_FIXED_SLOT_MAX_MTPR=511 / 1023 / 2047admits the 256 / 512 / 1024-token capacities (default unchanged: 255)AITER_MEGA_*) are read throughmega_moe/envs.py, which lists each with its default and purposetest_mega_moe_v2.py --pad-tail-rows N: whole -1 rows undermask_invalid_slots=Truecome back as exact zeros, real rows stay within rtol of the reference, CUDA graph replay is exactconst_expr/ deviceifguards keepnoqa: SIM102as elsewhere in the FlyDSL kernels)Companion ATOM PR: ROCm/ATOM#2463 (
--moe-backend mega: DP pad-row masking throughmask_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.pyon 8x MI355X, EP8 v4_pro: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-rowson 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--checkagainst epx + fused_moe with pad rows, fixed-slot MTPR 128/256 and compact MTPR 16384🤖 Generated with Claude Code