Repository navigation
Conversation
… tiles (gfx950) Speed up the compact-dispatch MegaMoE Stage1 kernel (EP8 dispatch producers + persistent FP8 x MXFP4 grouped GEMM1 + SwiGLU + MXFP8 quant). Outputs are bit-identical to main. Dispatch producer (every compact variant; the wire protocol is unchanged, so senders and receivers in different Stage1 variants interoperate): - One lane per task classifies a window of 64 tasks and loads its plan entries; empty tasks are skipped by a wave scan instead of being visited serially. - Windows whose tasks all have <= 16 rows run one task per wave, with no CTA barrier; route headers are written lane-parallel. - The fan-out pair is loaded once per producer instead of once per task, and every wave holds the destination's rows-per-tile so any wave can publish. - Remove the write-only PAYLOAD_CHUNK_DONE counter update (one system release, one atomic and two barriers per task). - One system release per published task instead of one per tile range, and no extra release fence before the ready-queue epoch store, which is already a system-scope release store. Consumer: the per-tile system acquire is done once by the thread that waited for the tile, before the CTA barrier, instead of by every thread after it. GEMM1 tile bodies, only for compact, route-major, FP8-in/FP8-out Stage1 (fixed-slot, A4W4, indexed payload and standalone GEMM1 keep main's code): - sort_block_m=128 prefill tile: A scales staged pre-packed for one ds_read_b128 per K step; XOR-swizzled C shuffle (no LDS bank conflicts); alias scopes so the A LDS-DMA does not force vmcnt(0) before A reads; next-step B packs issued up front with the K loop unrolled by two; ping-pong offsets kept opaque so addresses are not hoisted and spilled. - Other LDS-DMA variants: wave-uniform A tile base, so the LDS-DMA buffer descriptor stays in SGPRs. - Register-staged sbm32 variant (smallest decode buckets): LDS-DMA A with prefetch distance 2 through a 3-stage LDS ring. Every Stage1 kernel source changes, so the FlyDSL AOT cache must be regenerated (python -m aiter.aot.flydsl.mega_moe). Signed-off-by: Chushi Chen <chuschen@amd.com>
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
One backend per PR: PR title tags & labels: |
…ks sooner * quant is driven by tokens, not by rows: a token's bf16 row is read and quantised once and stored to its row in every destination's run, instead of once per destination (~4.6x fewer reads of the input) * a grid rendezvous separates the quant pass that fills the early chunks from the rest, so the second pass no longer races the first for bandwidth and delays the puts; the comm workgroup then puts those chunks with nothing to check, which drops the per-chunk completion counters entirely * `--early-chunks` sets how many chunks of each run prepare puts itself * clock stamps in both kernels behind `--trace`, plus `--stress` and `--stall-test` in the GEMM1 test, and a check that the fanout-group segments (which the route push does not cover) are empty At bs 8192 the first put moves from 177 to 106 us and the first tile's wait for data from 53 to 7 us; prepare pays ~25 us for the rendezvous. Flash, prepare + stage1 against aiter main without ROCm/aiter#7001: 1.36x at 512 tokens/rank, 1.17x at 1024, 1.14x at 2048, 1.06x at 4096, 1.00x at 8192. Two things measured and rejected rather than kept: writing through only the rows a put reads while the kernel runs (no change), and walking the early rows off `row_tok` instead of scanning the tokens (slower -- up to 75% of the emissions are duplicates at small batch). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
With --fuse-quant (MegaMoEStage1Sdma(fuse_quant=True)), aiter's prepare only plans and two kernels replace pack: - route: walks the production producer's rows as pack does, but only records each row's destination per (token, top-k slot) and writes the route headers to the destination over P2P; - qpack: one block per token quantises the bf16 row once, with aiter's per-1x32 MX fp8 arithmetic, stages it in LDS and stores it to each of the token's slab / landing rows with 16 contiguous bytes per lane. Every bf16 row is read once and every routed row written once, instead of prepare quantising into a buffer and pack copying each routed row from it. At 8192 tokens/rank route + qpack take 0.033 + 0.083 ms where pack took 0.188 ms. Stage1 incl. prepare improves by 0.02 / 0.04 / 0.08 ms at 2048 / 4096 / 8192 tokens/rank, still 0.09-0.16 ms slower than ROCm/aiter#7001. The layer output stays byte-identical to aiter's. Co-Authored-By: Claude <noreply@anthropic.com>
|
No activity for 15 days, so this is now labelled |
Summary
This PR speeds up the compact-dispatch MegaMoEV2 Stage1 kernel (EP8 dispatch producers + persistent FP8 x MXFP4 grouped GEMM1 + SwiGLU + MXFP8 quant) on gfx950. Output is bit-identical to main.
Where it applies: the producer change applies to every compact-dispatch Stage1 variant. The new GEMM1 tile bodies apply only to compact, route-major, FP8-activation/FP8-output Stage1. Fixed-slot dispatch, A4W4, the indexed payload (MTPR ≥ 32768) and standalone
compile_gemm1keep main's GEMM code.Motivation
On DeepSeek-V4-Pro with the InferenceX CI recipe, MegaMoE Stage1 runs in every decode and prefill step, and the profile put it at about 15 % of GPU time. On main it had three costs:
vmcnt(0)waits that serialized the A copy with the MFMAs, and to VGPR spills.Changes
All changes are in
aiter/ops/flydsl/kernels/mega_moe/{dispatch,gemm1,gemm_util,mega_moe_stage1}.py. The API, the workspace layout, the configs and the dispatch wire protocol are unchanged.Dispatch producer (
emit_dispatch_payload, every compact variant)PAYLOAD_CHUNK_DONEupdate. Nothing reads that counter, and each update cost one system release, one atomic and two barriers. Each task now issues one system release before it publishes all its tile ranges, instead of one per range. The extra release fence before the ready-queue epoch store is gone, since that store is already a system-scope release.Consumer. The per-tile system acquire is now done once by the thread that waited for the tile, before the CTA barrier. On main, every thread did it after the barrier, and an all-thread acquire is a CU-wide cache invalidate per wave. This relies on the acquire invalidating the CU and L2 caches, and on the barrier ordering the other waves' loads after it. Please look at this one closely.
GEMM1 tile bodies (
build_fused_gemm1(..., tuned=True))ds_read_b128instead of 16ds_read_u8plus shifts.ds_reads never alias the in-flight LDS-DMA, which always writes the idle ping-pong half.SIInsertWaitcntstherefore stops inserting avmcnt(0)before the first A read of each K step.readfirstlane), so the LDS-DMA buffer descriptor stays in SGPRs, with no waterfall loop around eachbuffer_load ... lds.Relation to #6127 (merged). #6127 added the 8-column GEMM1 quant epilogue, the compact fast copy, the fixed-slot paired-K loop and dedicated LDS broadcast slots. This PR keeps all of them. The prefill-tile epilogue is #6127's 8-column pass on a swizzled C tile, and the fast copy still runs for CTA-mode tasks. The rest is new: the prefilter, wave mode, removed counter and fence changes, consumer acquire placement, and all three GEMM bodies. The microbench below is measured on top of #6127.
Guard.
TUNED_GEMM1 = compact and not indexed_payload and a_dtype == out_dtype == "fp8"sits incompile_mega_moe_stage1.build_fused_gemm1picks a new tile body only when the variant matches it exactly, and every other variant (andtuned=False) runs main'sdo_tilecode path. There is no new environment variable. Rollback is a revert.An earlier, more aggressive revision was dropped. It added a static first work unit, decode split-K tails, work stealing and fast SwiGLU, and it was about 1.38x vs 1.23x on the same isolated test. Its decode variants skipped the ready-queue publish on the sender side, so the receive protocol depended on the sender's variant. With uneven DP batches (idle ranks at 1 token, skewed routing), a receiver that consumes in ready order then waits forever, and a live server deadlocked on its first real batches. The 8-rank unit test never mixes variants, so it did not catch this. In e2e, that revision gained no more than this one. This PR is the conservative version. It ran in live servers: on the earlier build in 4 × 60-min runs, and on main in a smoke run plus 6 scored 35-min e2e arms. The mixed-variant tests below cover the failure mode.
Performance
Setup. 8x MI355X (gfx950), ROCm 7.2.0. The Docker image is built from aiter main a29547a and sglang main d0a7565436 (plus a one-line import fix needed for sglang main to serve DeepSeek-V4 on gfx950, unrelated to MoE), FlyDSL 0.4.0. The stock and PR arms use the same image and differ only in these 4 files and their regenerated AOT cache. As of aiter 26f4f1f and sglang bf1a252953 (2026-10-24), nothing under
mega_moe/, the AOT tool, the FlyDSL pin or sglang'smega_moe_flydsl.pyhas changed, and this commit applies to that main unchanged.Microbench: prepare + Stage1 (µs), DeepSeek-V4-Pro a8w4, EP8
Method. I used my own 8-rank harness, which is not in the repo. It builds a real
MegaMoEV2fromtest_mega_moe_v2.py's inputs, and each rank picks its config from its own token count, as sglang does. It times_run_fused_stage1(main's prepare + Stage1) with CUDA events on the slowest rank, with a 1 GiB L2 flush before every call. Decode runs under CUDA graph and prefill eager, as served. "Stage1 kernel" is the profiler time of the Stage1 kernel alone (max over ranks of the per-rank median). Stock and PR runs are interleaved in fresh processes, 5 per arm at MTPR 8192 and 3 at MTPR 16384. Values are medians over processes, with the range in brackets; "worst" = slowest PR process vs fastest stock process.AITER_MEGA_FIXED_SLOT_MAX_MTPR=1023, the cookbook decode instance): 1.00x and bit-exact. Paths that keep main's GEMM but get the producer change, 1 run each: a4w4 1.07-1.11x, MTPR 32768 1.03-1.10x.kimi_k3_route1.19-1.28x, R1 1.08-1.32x, GLM-5.2 1.07-1.21x, all bit-exact.test_mega_moe_v2.py's own Stage1 timing (warm L2, graph, mean over ranks) shows the same picture: 1 token 168 → 138 µs, 8192 tokens 2461 → 2008 µs.End to end: DeepSeek-V4-Pro, TP8 / DP-attention 8 / EP8, MegaMoE a8w4, MTP, c128
Setup. Agentic trace replay with aiperf (InferenceX
dsv4/sglang/mi355x-fp4-mtp/agentic.yaml, variantoverride_tp8_c128_dp_ep8_megamoe_hicache, seed 42, simulated accept length 3.77). Each arm gets a fresh container and fresh caches. The PR arm ships only the 4 files plus the cache written bypython -m aiter.aot.flydsl.mega_moe --quant a8w4, and prepare is main's. Warm-up requests are dropped, then a 35-min profile runs. "[10,35]" scores the requests that start after minute 10 and end by minute 35; "full" is the whole profile. Interactivity = 1/ITL. Runs alternate PR, stock, PR, stock on one node per flag set. Flag sets:SGLANG_AMD_FLYDSL_MEGA_MOE_MTPR=8192). Every decode and prefill step runs compact Stage1.SGLANG_AMD_FLYDSL_MEGA_MOE_DECODE_MTPR=512,AITER_MEGA_FIXED_SLOT_MAX_MTPR=1023. Pure-decode steps run the fixed-slot instance, which this PR does not speed up.Accuracy
kimi_k3_route, Zipf, mixed and zero-token ranks).test_mega_moe_v2.pypasses in all 15 configurations below, with per-bs rel-L2 identical to stock (a8w4 0.0589-0.0593, a4w4 0.212). CUDA graph replay is exact.FlyDSL AOT cache
Every Stage1 kernel key changes (the source of
build_fused_gemm1andemit_dispatch_payloadis part of the key), so the shipped cache has to be regenerated:python -m aiter.aot.flydsl.mega_moe --quant a8w4(and--quant a4w4). On main this prints 192 ok, 0 failed, and adds 96 new Stage1 entries per quant.FLYDSL_RUNTIME_RUN_ONLY=1.How to reproduce
E2E. sglang d0a7565436, with server and router flags exactly as in the InferenceX recipe above. The CI flag set is the 2026-10-07 CI job of that recipe; cookbook is the recipe after InferenceX #4315. aiperf agentic replay at concurrency 128, 35-min profile after warm-up. Both arms use the same image; the PR arm bind-mounts the 4 files and the AOT cache from the command above.
Interaction with other PRs
git merge-treeis clean. [FlyDSL] [Perf] MegaMoE: shorten EP8 prepare chain, coalesce MXFP8 quant (gfx950) #6900 imports ninedispatch.pyhelpers (plan/group emitters, wave scans, fan-out pair loader), and this PR changes none of them; it changes only_publish_tile_rangeandemit_dispatch_payload. The prepare → Stage1 hand-off (PLAN_READY and the plan arrays) is unchanged on Stage1's side, and [FlyDSL] [Perf] MegaMoE: shorten EP8 prepare chain, coalesce MXFP8 quant (gfx950) #6900 keeps that contract. All numbers here use main's prepare kernel; I have not measured the two together. Whichever lands second, I will rerun the AOT tool, the repo tests and the mixed-variant tests on the merged tree.dispatch.pyand Stage1 helpers that this PR leaves unchanged. Calls it takes over do not run Stage1; all other calls get this PR.pair_kapplies to main'sdo_tile, which this PR keeps for non-matching variants.Known gaps
kimi_k3_routehave one microbench run each, plus the repo test.--unify-fieldsand a rank in the 8192 bucket, only that rank runs the extra aligned-pair forward in_run_size, which deadlocks. The mixed tests above cap that rank at 4917 unless they run in burst mode. I can send a separate fix.Testing
test_mega_moe_v2.pyandbench_mega_moe_v2.pyare inskip_testsof.github/scripts/aiter_test.sh, so aiter CI runs no test of this kernel. All runs above were manual, on 8x MI355X with main a29547a, which is identical to current main for every file involved.op_tests/multigpu_tests/test_mega_moe_v2.py(8 GPU, 15 configurations + 7 mixed-variant tests, see above)Documentation
Dependencies
Breaking Changes
None.
build_fused_gemm1gains an optionaltuned=Falseandgemm1.pygainsgemm1_a_lds_stages. ThePAYLOAD_CHUNK_DONEworkspace slot stays allocated but is no longer written.