[yilin-void]

📌 Description

Expand the bounded tuning candidates for Hopper MXFP4 × FP8 MegaMoE, motivated by the latency observed on a Qwen3.8-Flash-Next-style routed MoE workload.

For an intermediate size of 640, K256 configurations are filtered out by alignment requirements, leaving three eligible configurations in the existing candidate pool. The pool misses useful K128 configurations across different token counts.

Changes

Add five measured M256/K128 configurations:

• An N16, cluster (1,1,1), group-330 configuration with two dispatch warps and FC1 store offload.
• Folded N64, cluster (1,1,1) configurations with group hints 132, 264, and 330.
• A folded N64, cluster (1,2,1), group-132 configuration with tail-pair scheduling.

All five use one scheduler stage. The folded configurations use one dispatch warp, early FC1 publication, and direct FC1 stores.

For the target shape, the eligible candidate pool grows from 3 to 8 configurations. Existing heuristic defaults are preserved; autotuning selects among the expanded candidates.

Benchmark configuration

• GPU: 2× NVIDIA H200, connected through NVLink.
• Input: [tokens_per_rank, 2560] per rank; EP=2.
• Experts: 512 total, 256 per rank; top-k=10; uniform expert load.
• Intermediate size: 640 after SwiGLU; gate/up projection width=1280.
• Precision: MXFP4 weights with K32 scaling, FP8 activations, BF16 input/output.

Performance

Both columns report tuned performance: the best of the original 3 candidates versus the best of the 8 candidates available with this PR.

Selection minimizes Compute latency; operator E2E is reported from the same selected configuration. Each configuration uses 10 warmups and 50 CUDA-event samples. The reported latency is the maximum of the two ranks’ median latencies. Before/after comparisons use measurements from the same tuning process at each token count.

All latencies are in µs. Latency reduction is calculated as (original − PR) / original × 100%.

Tokens/rank	Original Compute	PR Compute	Reduction	Original E2E	PR E2E	Reduction
128	339.216	316.272	6.76%	501.264	475.856	5.07%
256	348.656	337.296	3.26%	501.024	503.712	−0.54%
512	371.056	371.056	0.00%	560.656	560.656	0.00%
1024	593.264	488.128	17.72%	756.176	643.296	14.93%
2048	865.456	834.848	3.54%	1033.200	999.984	3.21%
4096	1405.728	1315.200	6.44%	1573.472	1475.216	6.24%
8192	2670.336	2343.024	12.26%	2891.616	2560.176	11.46%
16384	5223.024	4442.880	14.94%	5628.496	4873.776	13.41%

Compute includes prestaged MegaMoE execution, EP communication, and final top-k reduction. Operator E2E additionally includes input quantization/staging and output copy; weight preparation and JIT compilation are excluded. These are operator measurements, not model TTFT.

At 256 tokens/rank, the configuration selected by Compute is 2.688 µs slower in operator E2E. At 512 tokens/rank, the original configuration remains selected.

Validation

• All five added configurations passed numerical checks on both ranks at all eight token counts.
• Each token count’s selected winner also passed an independent replay.
• The existing host test suite passed all 23 tests.
• Pre-commit checks passed for the changed files.
• No new test files were added.

Measurements were collected with the exact registered configurations using repro_driver.py on pinned source 2394fa705; the final branch was not separately timed. The workload uses synthetic weights and balanced routing rather than a full model execution.

Summary by CodeRabbit

• Performance

  • Expanded MXFP4 tuning options to improve performance across additional model shapes and configurations.
  • Added support for more tile, warp, cluster, and tail-processing combinations while preserving existing heuristic defaults.

• Bug Fixes

  • Improved validation for folded and tail-processing MXFP4 execution paths, including partial tile routing scenarios.

• Tests

  • Added coverage for expert-parallel MXFP4 workloads with nonstandard token counts and both complete and partial tiles.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: defaults
• Review profile: CHILL
• Plan: Advanced
• Run ID: 0eb8ca1f-487a-4e7b-93f6-a5b756c87aeb

📥 Commits

Reviewing files that changed from the base of the PR and between 6131c5e and 3ed00bc.

📒 Files selected for processing: 1

• tests/moe_ep/test_moe_ep_sm90_pull_mxfp4_mega_multirank.py

Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 7 remain after this review.

📝 Walkthrough

The MXFP4 tuner adds five base tactics. The tuning guide updates candidate counts. A parameterized SM90 test validates folded and folded-tail tactics for a multirank configuration.

Changes: MXFP4 tuning

Layer / File(s)	Summary
Add MXFP4 base tactics flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/shim/mxfp4_tuner.py, flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/TUNING.md	The tuner adds five candidates with varied tile shapes, group hints, cluster configurations, dispatch settings, and tail-pair handling. The tuning guide updates the documented base and shape-specific candidate counts.
Validate folded MXFP4 tactics tests/moe_ep/test_moe_ep_sm90_pull_mxfp4_mega_multirank.py	A parameterized SM90 test covers folded and folded-tail tactics for the 2560×640 configuration with 129 tokens per rank. It compares full and partial N64 tile routing with an independent raw oracle.

Priority: ⬇️ Low

Estimated code review effort: 2 (Simple) | ~10 minutes

Change: Feature

Merge Risk: ⚪ Minimal · up to 3ed00

No identified issue blocks merging; the new GPU test was not independently run in this review.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 20.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 5 functions across 3 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Title check		The title clearly identifies the main change: adding a folded K128/N64 tuning candidate for Hopper MXFP4 × FP8 MegaMoE. It is concise and specific.
Description check		The description is detailed and on-topic. It explains the motivation, five candidate configurations, benchmark setup, performance results, validation, and measurement limitations. The Related Issues s…
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

Full details: Description check

Explanation

The description is detailed and on-topic. It explains the motivation, five candidate configurations, benchmark setup, performance results, validation, and measurement limitations. The Related Issues section is not completed, and checklist items remain unchecked, but these omissions are non-critical.

• Fix all pre-merge checks with AI

✨ Finishing Touches: 💡 1

🛠️ Fix failing CI checks: 💡

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[jiahanc]

/bot run tests/moe_ep/test_moe_ep_sm90_pull_mxfp4_mega_multirank.py

[flashinfer-bot]

GitLab MR !2291 has been created, and the CI pipeline #75210503 is currently running. I'll report back once the pipeline job completes.

[jiahanc]

@flashinfer-bot run tests/moe_ep/test_moe_ep_sm90_pull_mxfp4_mega_multirank.py

[flashinfer-bot]

[FAILED] Pipeline #75210503 — 23/31 executed test jobs passed

Compared with nightly #74968837.

Unit Tests

GPU	CUDA 12.9	CUDA 13.0	CUDA 13.4	Notes
B200	✅ Pass	⚠️ Infra	❔ Unknown	Timeout: job timed out before producing test results (2 jobs; CUDA 13.0) Unknown: script failed before producing a JUnit report (2 jobs; CUDA 13.4)
GB200	✅ Pass	✅ Pass	✅ Pass	—
GB300	✅ Pass	✅ Pass	✅ Pass	—
H100	—	⚠️ Infra	—	Timeout: job timed out before producing test results (1 job; CUDA 13.0)
RTX Pro 6000 Blackwell	⚠️ Infra	⚠️ Infra	⚠️ Infra	Infrastructure: CI infrastructure failure (3 jobs; CUDA 12.9, CUDA 13.0, CUDA 13.4)
Spark	—	✅ Pass	—	—
VR NVL72	—	—	✅ Pass	—

✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · ⚠️ Infrastructure · ❔ Unknown or unclassified · — Not run

Failure details

Timeouts, infrastructure, or incomplete jobs

• Infrastructure: CI infrastructure failure — RTX Pro 6000 Blackwell / CUDA 12.9, RTX Pro 6000 Blackwell / CUDA 13.0, RTX Pro 6000 Blackwell / CUDA 13.4
  • Jobs: unit_test_rtx-pro-6000-blackwell: [cu129], unit_test_rtx-pro-6000-blackwell: [cu134], unit_test_rtx-pro-6000-blackwell: [cu130]
• Timeout: job timed out before producing test results — B200 / CUDA 13.0, H100 / CUDA 13.0
  • Jobs: unit_test_h100: [0, cu130], unit_test_b200: [1, cu130], unit_test_b200: [0, cu130]
• Unknown: script failed before producing a JUnit report — B200 / CUDA 13.4
  • Jobs: unit_test_b200: [1, cu134], unit_test_b200: [0, cu134]