feat(moe_ep): add Hopper MXFP4 fused/split MegaMoE autotune - #4843
Draft
ruoqianguo wants to merge 17 commits into
Draft
ruoqianguo wants to merge 17 commits into
ruoqianguo wants to merge 17 commits into
Conversation
…token heuristics Kernel updates: * multi-CTA CGA execution (cluster 1x1/2x1/1x2/2x2, TMA multicast) * task-tile ping-pong scheduling across two WGMMA+epilogue warpgroups * two-stage FC2 R2S + TMA store pipeline * all six token-back x reduce mode combinations * guarded IKET ranges for MMA, TMA, epilogue, and dispatch * token-bucket launch heuristics with expanded token-sweep benchmarking Shim/backend: token_back_by_dispatch bool -> token_back_mode enum (bool kept as legacy alias); new pingpong and cluster_shape_mnk knobs; geometry knobs left unset resolve via the token heuristics; the SM90 mega benchmark gains --heuristic.
…chive results benchmarks/bench_moe_ep_sm90_mega.py: * replace the balanced routing with the drop runner's block-permutation algorithm (routing shape affects kernel time) * default to the token-bucket heuristic launch configs; --both-orders / --swap-ab / --no-swap-ab select fixed layouts * add --cooldown-s (default 5) to recover clocks before each timed series * generate fp8 payloads with the drop perf data recipe by default; --no-sparse-data switches to dense quantized-randn model data * extend the default token sweep to 8..32768 * archive results to benchmark_data/<date>/... by default (--output-csv), with the resolved heuristic launch-config columns appended TUNING.md: refresh the measured tables with the 2026-08-23 4xH200 heuristic sweep, document the compute-vs-e2e timed regions, and update the methodology and next levers. SKILL.md: record the current drop revision and the excluded Green Context commit.
…per-bucket token-back heuristics - Add shim tuner/knob_cache/autotune mirroring the SM100 flow: the config `knobs` field resolves the persistent knob cache and falls back to the token-bucket heuristic table; a dict applies explicit knobs; "auto" runs a collective online sweep at the first forward and persists the winner (32 candidates: 16 table geometries x 2 token-back modes). - Extend the offline tuning CLI (python -m flashinfer.moe_ep.tune) to the SM90 fp8 kernels (--fp8-scale-mode). - Add token_back_mode to the heuristic table (epi_warps small/mid buckets, reuse_dispatch_warps at the GEMM-bound tail: per_tensor >= 16384, blockwise >= 1024) and wire it through shim/backend/benchmark (--token-back defaults to the heuristic; CSV records the resolved mode). - Backend validates knobs values and knobs-vs-explicit-geometry conflicts. - Tests: tuner/knob-cache/backend-wiring unit tests; multirank reuse_dispatch_warps / standalone_warps correctness cases. - Refresh TUNING.md / SKILL.md.
- Update the vendored token-sweep tooling (run_token_sweep_benchmark, summarize_token_sweep, run_perf_test.sh) to the heuristic-aware flow; mega_runner now defaults token_back_mode to the heuristic table's per-bucket winner (an explicit --token_back_mode still wins). - test_heuristic_config gains per-bucket token-back assertions. - nvfp4 mega_runner aligns ranks after profiler startup so startup skew is not counted in the first timed kernel.
…ntized combine wire - Dispatch dedup (dedup_dispatch): sender-side carrier election per (token, destination rank) via warp shuffles (smallest expert carries), flags packed into the route word's spare bits; the receiver rendezvouses on a (src_rank, src_token) carrier table and bulk-copies duplicate pool rows locally instead of re-pulling over NVLink. Bit-exact; the wait graph is a DAG (duplicates only ever wait on strictly smaller experts), release/acquire paired with the existing fc1_ready contract. - Combine dedup (grouped_token_back): the LAST fc2 row of a (src_rank, src_token) group pre-reduces every member in fp32 and pushes ONE row per contributing rank into a [tokens][world_size] inbox; the standalone reducer becomes a rank-slot reduce gated by a per-token contributing-rank bitmask computed at dispatch. - Quantized combine wire (combine_format="32e4m3xe8m0"/"32e5m2xe8m0"): per-32 e8m0 + fp8 encoding inside the grouped reduction (single quantization; the receiver dequantizes to fp32 before accumulating). SM90 compatibility: bit-math e8m0 encode/decode, f16-hop fp8 decode, and a scalar-math mxfp8 path in the shared TopkReduce (slot_mask keeps stale inbox slots out of the sum). - All three knobs are opt-in and OFF by default; the disabled paths compile identically to before. dedup_dispatch is a perf knob (bit-exact); grouped_token_back/combine_format are correctness-class and must match on every EP rank. - Benchmark axes --dedup-dispatch / --grouped-token-back / --combine-format; ten new multirank cases (bit-exact dedup matrix incl. multi-waiter top-6 and in-kernel-reduce compositions; grouped bf16/fp8 gated at ~54 dB / ~31.5 dB SNR vs the exact reference).
Size the WORKING subset of the 4 dispatch warps (1/2/4); the physical layout stays at 4 (setmaxnreg is warpgroup-granular). Idle warps skip the whole dispatch body and rejoin at kernel_tail -- reserved for future in-kernel work. Barrier / grid-sync / reuse token-back walker counts follow the active count. A shallower NVLink read queue wins: 1 warp/SM already exceeds the H200 bandwidth-delay product. Measured vs 4 warps: EP4 +1.4% / EP8 +2.5% geomean, up to +11% at 16k tokens/rank per_tensor; 4-warp perf unchanged. Exposed via shim/backend config, tuner perf knob, and --active-dispatch-warps in the sweep benchmark.
Re-measured the heuristic sweep with the new active_dispatch_warps=1 default on a clock-locked (1830 MHz) 4x H200 node: per_tensor peak 841 -> 896 TFLOPS/rank, blockwise 568 -> 589; gains concentrate in the large-token reuse buckets, small-token points unchanged.
Contributor
|
Important Draft PR not reviewedDraft PRs are not automatically reviewed by default.
To automatically review draft PRs, update your CodeRabbit configuration: reviews:
auto_review:
drafts: trueThanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
…lish Hoist fc1_done publication to store-landed time so FC2 starts earlier. The empty warp runs an FC1 store server: the epilogue only R2S-stages FC1 output and hands (slot, dest, done-flag) over a per-WG smem mailbox FIFO; the server issues the TMA store, waits completion, and release-publishes fc1_done -- removing the consume_next / boundary barrier stall that otherwise defers publication by ~a tile at small token counts. Self-gates to non-ping-pong (ping-pong's retire section already publishes early). Covers both non-swap and swap-AB; swap-AB keeps the baseline store/consume overlap the lighter early-publish variant loses. The 2-WG paths (N256 / swap M256) use a dual-FIFO server with the fc2 spin threshold scaled x2. A dynamic register fit sizes the store server to the CTA's 64K budget (reclaims the 2-WG epilogue's 216->200 headroom) and falls back to in-epilogue early publication when it does not fit. fc1_early_done_publish kept as a tuner axis. Measured 4x H200 (1830 MHz) vs no-offload: non-swap bw512 +21~23%, bw1024-4096 +6~8%; swap-AB pt256 +15%, pt512 +13%, bw128 +10%. Peak per_tensor 896 -> 936 TFLOPS/rank. Multirank matches_reference / swap_ab / torch-oracle tests pass bit-exact.
…-calibrate heuristic With one active dispatch warp the other three dispatch slots idle, so fold the TMA-A / TMA-B / scheduler roles into them and drop the separate producer warpgroup: 128 fewer threads per CTA and the 2-WG register budget falls from the 65536 cap to 60160. dispatch_warp_id stays a 4-tuple so every TokenComm count is unchanged; one _apply_mega_warp_layout() drives both kernels. Without the epi_aux warp the FC1 store offload is replaced by early fc1_done publication; the freed budget is refit into the epilogue (216->232 on 2-WG kernels). The layout is what makes the two-warpgroup epilogue modes viable: under the old layout cooperative ran 14% and ping-pong 24% behind the basic one-WG tile on blockwise; under the fold both are ~43% faster. Every bucket of the heuristic table was re-measured against its two alternative epilogue modes (basic / ping-pong / cooperative, tile derived per bucket, cluster shape and token-back preserved) on one node: blockwise non-swap 512-32768 move to cooperative M64N256, per_tensor 8 to cooperative and per_tensor 64 to basic; the other 17 buckets keep their entries. Same-node full sweep vs the previous default: per_tensor +0.8%, blockwise +12.2% (bw32768 +44.5%); blockwise peak 641 -> 830 TFLOPS/rank. Standalone token-back loses everywhere and stays off. Also fixes the bench pingpong override to forward the bucket's full heuristic geometry so it flips only ping-pong, and adds --epi-mode to force basic / ping-pong / cooperative per bucket for twin sweeps. Multirank suite (29) + swap/oracle (5) + a per-row test pinning every re-calibrated heuristic entry (5) pass bit-exact.
ruoqianguo
force-pushed
the
sm90-mxfp4-megamoe-autotune-draft
branch
2 times, most recently
from
September 8, 2026 10:25
e3cf38b to
251f442
Compare
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.
Important
This is a stacked Draft PR on the latest #4688 snapshot at
b1b6b399b7a9885cbe7d543459d0d9a6797b61b4. It is published early for integration and performance validation while final heuristic promotion and cache hardening continue. Please do not merge yet.The branch will be rebased onto
mainafter #4688 lands.Summary
benchmarks/bench_moe_ep_sm90_mega.pyas the common FP8/MXFP4 workload, routing, cooldown, telemetry, and timing entry pointAutotune behavior
knobs="auto"times a compact 17-tactic runtime union: 8 H200 block-routing winners, 7 deduplicated published-exact winners, and 2 measured H20 anchorsknobs="auto"times 8 routing-profile-specific H200 tactics; split remains restricted to the certified standard H200, CC 9.0, 132-SM partitionknobs=Noneperforms a mode/device/workload/routing/provenance-scoped cache lookup, then falls back to the matching MXFP4 token-bucket heuristicValidation
unit_sm90_mxfp4test entry: 419 passedMAX(rank-local medians); aggregate audit: 2483 checks, 0 failuresCurrent WIP follow-ups
block_permutation_v1winner set into the shipped heuristic/provenance tablesDeliberately deferred
The general offline sweep launcher/domain, qualification helper, raw benchmark data, PIC-C artifacts, and temporary analysis scripts are intentionally excluded from this Draft and will be handled separately.