[akaashrp]

📌 Description

Rubin MegaMoE rejects NaN and infinite routing scores even when the expert ID is -1 and the slot is unused. Allow those scores for masked slots and stage them as zero without modifying the caller’s tensors. By default, active routes require finite scores, valid expert IDs and no duplicates.

Use a route-staging kernel for prequantized NVFP4/MXFP8 inputs. BF16-to-MXFP8 staging quantizes activations and writes routes directly into the workspace, with the same scale rounding as the Torch path. Unsupported layouts retain Torch staging.

Set FLASHINFER_SM107_MEGA_CHECK_ROUTES=0 to skip routing value checks when the caller guarantees valid routing. Metadata checks and masked-score zeroing still run. Set it before warmup and CUDA graph capture.

🔍 Related Issues

Related: #6398.

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

• Tests have been added or updated as needed.
• All tests are passing (unittest, etc.).

🔬 Experimental Track

• This PR is experimental: it adds or changes code under flashinfer/experimental/ and/or an @flashinfer_experimental_api. Tracking issue: #
  • The tracking issue names an owner, the reason for the experimental path, and a graduation plan with a target release.
  • Core changes are limited to a thin entry point (signature, shared validation, feature-gate check, backend selection, handoff).
  • Tests live in tests/experimental/ and were validated on the intended hardware; a runnable example is included.
  • Nothing is registered in flashinfer/aot.py, and no experimental backend is reachable from backend="auto" without FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1. (Calling an @flashinfer_experimental_api or naming a backend explicitly is itself the opt-in and needs no environment variable.)
  • Test scope declared below. The experimental CI lane runs exactly these targets, so keep them as narrow as the change allows.

# One target per line: a directory or a file. (A pytest ::selector is not
# supported -- the sharding runner cannot consume one.) Must be under
# tests/experimental/ and must exist. Delete these comment lines and add yours, e.g.
#
#   tests/experimental/test_my_backend.py
#   tests/experimental/my_backend/
#
# Declaring the whole tree (tests/experimental/) is allowed but means every
# experimental PR pays for every other feature's tests, in every matrix cell.

Reviewer Notes

None.

Summary by CodeRabbit

• New Features
  • Eligible BF16 inputs can use fused quantization and staging for NVFP4 or MXFP8. Unsupported layouts continue to use Torch staging, and fused staging can be disabled with a setting.
• Bug Fixes
  • Masked routes (expert ID -1) now accept NaN or infinite routing scores, which are treated as zero during staging without modifying caller tensors.
  • Non-finite scores on active routes remain invalid when routing checks are enabled. A setting can skip route-value checks while retaining metadata checks and masked-score zeroing.

[coderabbitai]

No actionable comments were generated in the recent review.

Recent review info

Run configuration

• Source mode: Auto
• Repository: Unavailable
• Branch: Unavailable
• File: No configuration file from this repository applied
• Commit: Unavailable
• Layers: defaults
• Review profile: CHILL
• Plan: Advanced
• Run ID: 001c98f0-fad7-49de-869a-7b58c4fbd9bd

Commits

Reviewing files that changed from the base of the PR and between bace296 and 25f286b.

Files selected for processing: 4

• CLAUDE.md
• docs/design_docs/moe_ep_sm107_qualification.md
• flashinfer/moe_ep/backends/mega/kernel/sm107/validation.py
• tests/moe_ep/test_sm107_block_scaled_contracts.py

Files skipped from review as they are similar to previous changes: 1

• docs/design_docs/moe_ep_sm107_qualification.md

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

SM107 MegaMoE routing validation now permits non-finite scores for masked routes. Staging writes zero scores for those routes using route-only or fused quantization paths when eligible. Tests cover staging behavior, quantization parity, and multirank graph replays.

Changes: SM107 MegaMoE routing and staging

Layer / File(s)	Summary
Masked-route validation contract flashinfer/moe_ep/backends/mega/kernel/sm107/validation.py, flashinfer/moe_ep/backends/mega/kernel/sm107/routing_validation.py, tests/moe_ep/test_sm107_block_scaled_contracts.py, tests/moe_ep/test_sm107_preparation.py, docs/design_docs/moe_ep_sm107_qualification.md, CLAUDE.md	Validation requires finite scores only for active routes. FLASHINFER_SM107_MEGA_CHECK_ROUTES=0 skips route-value checks while metadata checks remain active. Tests cover masked and active routes, fused checking, and the setting.
Route staging and backend wiring flashinfer/moe_ep/backends/mega/kernel/sm107/routing_staging.py, flashinfer/moe_ep/backends/mega/kernel/sm107/*_cutedsl/backend.py, flashinfer/moe_ep/backends/mega/kernel/sm107/*_cutedsl/staging.py, tests/moe_ep/test_sm107_preparation.py	stage_routes stages IDs as int32 and scores as float32, setting masked-route scores to zero. Backend paths use route staging. Tests cover layouts, fallback behavior, graph replay, and unchanged caller scores.
Fused quantization and staging flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/shim/quant_stage.py, flashinfer/moe_ep/backends/mega/kernel/sm107/mxfp8_mxfp8_bf16_cutedsl/staging.py, tests/moe_ep/test_sm107_preparation.py, docs/design_docs/moe_ep_sm107_qualification.md	Eligible MXFP8 inputs use fused quantization and staging; unsupported inputs use Torch staging. fused_quant_stage adds an MXFP8 normal-scale option. Tests compare fused outputs and scales with reference paths.
Multirank and graph verification tests/moe_ep/test_moe_ep_sm107_block_scaled_mega_multirank.py	Multirank tests compare outputs for masked scores containing NaN or infinities with outputs for zeroed scores, including direct calls and graph replays.

Priority: Normal

Estimated code review effort: 4 (Complex) | ~45 minutes

Change: Bug fix

Sequence Diagram(s)

sequenceDiagram
  participant MegaMoEBackend
  participant stage_mega_moe_inputs
  participant fused_quant_stage
  participant Workspace
  MegaMoEBackend->>stage_mega_moe_inputs: Stage eligible inputs and routes
  stage_mega_moe_inputs->>fused_quant_stage: Pass inputs and route buffers
  fused_quant_stage->>Workspace: Write quantized inputs and staged routes

Loading

Possibly related issues

• [CAKE] MoE & Routing Kernel Progress Tracker #6597: This tracker covers MegaMoE work, including the routing and staging changes in this PR.

Suggested labels: op: moe

Merge Risk: Minimal · up to 25f28

No actionable merge-blocking issue is established by the supplied evidence; the change is mergeable after normal checks.

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 35 functions across 11 files. (2 skipped:…	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 and concisely describes the primary change: allowing non-finite scores on masked Rubin routes.
Description check		The description includes the required sections, explains the routing and staging changes, links a related issue, and records pre-commit and test checklist status. The experimental section remains unto…
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 includes the required sections, explains the routing and staging changes, links a related issue, and records pre-commit and test checklist status. The experimental section remains untouched, as appropriate for this non-experimental change.

Full details: Docstring Coverage

Explanation

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 35 functions across 11 files. (2 skipped: 2 unsupported.)

Resolution

Write docstrings for the functions missing them to satisfy the coverage threshold.

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.}

[Anerudhan]

Codex bot review · model: gpt-6-astra

Reviewed bace29645155 against base 3f398977d75b.

The masked-route contract and staging paths are consistent: inactive scores are zeroed without changing caller inputs, while active routes keep the finite-score and uniqueness checks. Reviewed validation, both backend call paths, fused quantization, route staging, cache specialization, graph bindings, capacity tails, tests, and documentation. No approval-blocking defect found.

[P3] Record before/after staging measurements

This also enables a new fused MXFP8 staging path by default and replaces prequantized route copies with a kernel, but the description has no before/after measurements. Add a reproducible SM107 staging or full-forward benchmark with absolute times, shapes, dtype, and graph mode, as required by the repository review guidance for performance changes.

Evidence: The MXFP8 stage selects fused_quant_stage by default when eligible; the description and qualification document describe the path without measured staging or forward costs. Attribution: New finding from this review.

Validation: 119 focused host tests passed; 231 SM107 preparation cases skipped on the CPU-only test environment. 384 additional cases passed through the actual backend staging functions on CPU: four quantization kinds, BF16/prequantized inputs, four score dtypes, contiguous/strided/offset routes, and live counts 0/1/8/16. Checked masked NaN/+Inf/-Inf scores, input preservation, activation/scale bytes, live counts, and capacity tails. Offline CuTe-DSL 4.8.1/CUDA 12.9 compilation generated SM100a GPU binaries for the new MXFP8 kernels at hidden widths 16384 (E4M3) and 7168 (E5M2), plus six route-staging variants spanning FP16/BF16/FP32 scores and contiguous/strided inputs. No kernel was launched. Checked the current discussions and the existing failed CI summary: its test jobs were skipped pending authorization; it did not report a test assertion failure.

Limitations: No local GPU execution: the shared GPU was occupied. Offline compilation targeted SM100a because the installed CUDA 12.9 compiler does not support SM107a; this does not validate native SM107 behavior, graph replay, multi-rank execution, or performance. The available GitHub focused CI matrix uses H100 for these paths, so its SM107-only tests will skip. Native SM107 staging parity and pooled multi-rank graph tests remain unverified by this review.

Approved: no P0 or high-risk P1 found. Remaining findings stay with the owner; merge timing stays with the owner.

[Anerudhan]

Requesting targeted CI for bace29645155 with @flashinfer-bot run tests/moe_ep/test_sm107_block_scaled_contracts.py tests/moe_ep/test_sm107_preparation.py.

Codex bot review — model gpt-6-astra.

[Anerudhan]

@flashinfer-bot run tests/moe_ep/test_sm107_block_scaled_contracts.py tests/moe_ep/test_sm107_preparation.py

[Anerudhan]

Focused CI completed for bace29645155b6fadbecc37eda0d4888e621ec15: run 39991622539. H100 passed 119 contract tests with no failures. All 231 SM107 preparation tests skipped, as expected on H100. All six AOT build/import jobs passed.

The setup and H100 logs both identify merge 491633db8baa493da644827f2590a4e41753974f (this head into base 8ec9464ce81d81c9e3670aa09acd1a8780f83695). This run does not validate SM107 kernel execution, CUDA graph parity, multi-rank behavior, or performance.

Codex bot review — model gpt-6-astra.

[Anerudhan]

Codex bot review · model: gpt-6-astra

Reviewed 25f286b7852d against base 872554187e41.

The explicit route-check opt-out skips only value checks. Default validation, metadata checks, masked-score zeroing and caller preservation remain intact. Graph behavior is documented correctly: the switch must be set before capture, and changing captured checks requires recapture. No new approval-blocking defect found.

• [P3] Still open: Record before/after staging measurements.

Validation: 125 focused host tests passed, including all six default/enabled/disabled and fused/fallback opt-out cases; 231 SM107 preparation tests skipped in the CPU environment. With FLASHINFER_SM107_MEGA_CHECK_ROUTES=0, 384 cases passed through the actual backend staging functions on CPU: four quantization kinds, BF16/prequantized inputs, four score dtypes, contiguous/strided/offset routes, and live counts 0/1/8/16. Masked NaN/Inf scores were zeroed; caller tensors, activation/scale bytes, live counts and capacity tails were checked. Reviewed the complete four-file follow-up, both backend stage_inputs callers and the metadata-validation boundary. All GPU kernels, quantization, route staging, cache keys and execution bindings are byte-identical to the prior bace296 review. No redundant compilation sweep was run. Current-head run 40405319989 skipped setup/tests pending authorization; its failed summary is not a test assertion failure. Trusted main CI files at 8725541 match the previously read 215f83c versions. A focused contract-file CI request is appropriate after publication.

Limitations: No local GPU, sanitizer, multi-rank or timing run; shared GPUs remain occupied. CPU staging exercises the Torch path, not fused GPU kernels. Prior offline compilation targeted SM100a, not native SM107, and prior H100 CI skipped SM107 preparation. The opt-out assumes caller-validated IDs and active scores; it intentionally does not sanitize invalid active routes.

Comment only: this PR already has an active approval; preserving existing approvals.

[Anerudhan]

Requesting targeted CI for 25f286b7852d with @flashinfer-bot run tests/moe_ep/test_sm107_block_scaled_contracts.py.

Codex bot review — model gpt-6-astra.

[Anerudhan]

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

[Anerudhan]

CI follow-up for 25f286b7852d: the focused GitHub run passed. All six AOT build jobs passed.

The H100 job passed all 125 tests in tests/moe_ep/test_sm107_block_scaled_contracts.py, with no skips or failures. These validate host contracts; this H100 run does not provide SM107 kernel execution coverage.

Codex bot review — model gpt-6-astra.