Repository navigation
Fuse Rubin MegaMoE input staging for the MXFP8 family - #7601
Victor49152 wants to merge 2 commits into
Conversation
…on kinds The SM107 MX kinds (mxfp8_e4m3, mxfp8_e5m2, mxfp4_mxfp8) staged BF16 inputs with torch ops: an FP32 upcast, a dozen elementwise/reduce launches and per-tensor copies, 46-570 us per forward for DSv4.1-Flash (hidden 5120) at 16-8192 tokens/rank, against 14-39 us for the NVFP4 kind's fused stager. Route them through the same fused SM100 CuTe-DSL stager (fused_quant_stage) the NVFP4 kind uses; it already has mxfp8_e4m3 / mxfp8_e5m2 modes. Both paths round the E8M0 scale up from amax / fp8_max and the data to nearest even, so the staged data is identical; only an all-zero block's unused scale byte can differ. The torch path remains the fallback for unsupported layouts and for FLASHINFER_MEGA_FUSED_STAGE=0. AI-assisted. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Stage the same BF16 inputs through the fused stager and the torch path (FLASHINFER_MEGA_FUSED_STAGE=0) for mxfp8_e4m3, mxfp8_e5m2 and mxfp4_mxfp8 with int32/int64 routes, at live counts 64, 8, 0 and 32 in a 64-row buffer, including zero, tiny and large blocks. Data, routes, route weights and every non-zero block's scale must be bitwise equal; an all-zero block's scale byte is excluded since the two paths encode it differently. AI-assisted. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
No actionable comments were generated in the recent review.
 |
| Layer / File(s) | Summary |
|---|---|
Fused staging selection and parity flashinfer/moe_ep/backends/mega/kernel/sm107/mxfp8_mxfp8_bf16_cutedsl/staging.py, tests/moe_ep/test_sm107_preparation.py |
The staging function selects a quantization type and uses fused_quant_stage when the environment setting and input requirements allow it. Otherwise, it uses the existing Torch path. Tests compare outputs across quantization kinds, routing-ID dtypes, and live-row counts. Scale comparisons exclude all-zero input blocks. |
Priority:
Estimated code review effort: 3 (Moderate) | ~20 minutes
Change: Refactor
Possibly related issues
flashinfer-ai/flashinfer#7504— The issue tracks exact quantization and CuTe DSL MoE testing; this PR adds fused-versus-Torch staging comparisons.
Possibly related PRs
- flashinfer-ai/flashinfer#7062 — Adds the fused SM107 route and MXFP8 staging infrastructure used by this PR.
Suggested labels: arch: sm107
Suggested reviewers: akaashrp, jiahanc
Merge Risk: 
81fe6
-
No actionable merge-blocking issue is established; the change is ready for normal validation.
Pre-merge checks: 
-
Check name Status Explanation Resolution Docstring Coverage Docstring coverage is 25.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Documentation-only edits and deletions within surv… Write docstrings for the functions missing them to satisfy the coverage threshold. -
Check name Status Explanation Title check The title clearly identifies the main change: fusing Rubin MegaMoE input staging for the MXFP8 family. Description check The description is detailed and covers the implementation, fallback behavior, correctness considerations, performance results, and test updates. The related-issues section is empty, and the checklist … 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 covers the implementation, fallback behavior, correctness considerations, performance results, and test updates. The related-issues section is empty, and the checklist does not confirm that all tests or pre-commit checks passed, but the description is otherwise mostly complete.
-
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 25.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 4 functions across 2 files.
Resolution
Write docstrings for the functions missing them to satisfy the coverage threshold.
-
Finishing Touches: 💡 1
-
🧪 Generate unit tests (beta)
-
- Create a new PR
🛠️ Fix failing CI checks: 💡
-
- Commit to this branch
- 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.
Comment @coderabbitai help to get the list of available commands.
|
Closing as covered in #7062 |
📌 Description
On SM107 (Rubin), the MegaMoE kinds with MXFP8 activations (mxfp8_e4m3, mxfp8_e5m2 and mxfp4_mxfp8) still staged BF16 inputs with Torch ops. The path is an FP32 upcast, then per-block amax, scale and quantize ops, then per-tensor copies into the symmetric workspace. That costs 46–570 µs per forward at 16–8,192 tokens per rank. #6398 fused this preparation for nvfp4 only, which takes 14–39 µs.
This PR routes the MX kinds through the same fused CuTe-DSL stager (fused_quant_stage). That stager already supports mxfp8_e4m3 and mxfp8_e5m2 output. mxfp4_mxfp8 uses E4M3 activations.
mxfp8_mxfp8_bf16_cutedsl/staging.py: calls the fused stager when the layout allows, using the same conditions as the NVFP4 path:
BF16, contiguous, 16-byte-aligned activations and routes;
hidden size a multiple of 128;
an output buffer whose dtype matches the kind, and hidden / 32 scale columns;
FP16, BF16 or FP32 route weights.
It calls note_staged_tokens first, as the NVFP4 path does, so the shared stager's tail memo stays correct. The Torch path stays as the fallback for other layouts. FLASHINFER_MEGA_FUSED_STAGE=0 still forces it; this is the existing switch.
Same staged bytes. Both paths round the E8M0 scale up from amax / fp8_max and round the data to nearest even, so the staged data is identical. Only the scale byte of an all-zero block can differ: Torch clamps its exponent to 2^-126. That scale multiplies zero data, so the result is unchanged.
mxfp4_mxfp8 needs no change of its own. It shares this staging module, so it gets the fused path too.
MegaMoE forward, µs (lower is better)
mxfp4_mxfp8mxfp8_e4m3mxfp8_e5m2nvfp4(unchanged)🔍 Related Issues
🚀 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
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).🔬 Experimental Track
flashinfer/experimental/and/or an@flashinfer_experimental_api. Tracking issue: #tests/experimental/and were validated on the intended hardware; a runnable example is included.flashinfer/aot.py, and no experimental backend is reachable frombackend="auto"withoutFLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1. (Calling an@flashinfer_experimental_apior naming a backend explicitly is itself the opt-in and needs no environment variable.)Reviewer Notes
Summary by CodeRabbit