feat(moe): add unified block-scale FP8 support - #4026
Merged
aleozlx merged 8 commits intoJul 22, 2026
Merged
Conversation
Expose DeepSeek FP8 and MXFP8 through MoELayer with format-aware preparation, routing-mode support, and focused conformance/fuzz coverage. AI-assisted implementation and validation.
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds unified TRTLLM FP8 block MoE support for DeepSeek FP8 and MXFP8, including variant-aware preparation, runner dispatch, scale handling, architecture gating, routing support, and conformance and fuzz coverage. ChangesTRTLLM FP8 Block MoE
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Input
participant MoELayer
participant TrtllmFp8BlockRunner
participant MoERunner
participant FP8BlockKernel
Input->>MoELayer: provide BF16 activations and routing inputs
MoELayer->>TrtllmFp8BlockRunner: select FP8 block backend
TrtllmFp8BlockRunner->>MoERunner: pack routing, quantized tensors, and scales
MoERunner->>FP8BlockKernel: execute FP8 block-scale MoE
FP8BlockKernel-->>MoELayer: return computed output
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request adds support for TRTLLM block-FP8 MoE (covering DeepSeek FP8 and MXFP8 quantization variants) on Blackwell SM100+ architectures. It introduces the TrtllmFp8BlockRunner adapter, weight and activation preparation utilities, and comprehensive unit and fuzz tests. The feedback suggests ensuring that act.topk_weights is contiguous before calling .view(torch.int16) in runners.py to prevent potential runtime crashes on non-contiguous tensors.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Tighten calibrated conformance bounds, pin MXFP8 weight permutations, and track the shared DeepSeekV3 FromLogits wrong-answer path until its routing fix lands.
Select a deterministic curated seed that exercises production autotuning and document why block-FP8 preparation requires an explicit quant variant.
Use empty routing placeholders so the FP8 launcher consumes logits, reject unsupported SM110 early, and verify grouped routing selections exactly through replay.
feih-nv
marked this pull request as ready for review
July 20, 2026 11:40
feih-nv
requested review from
IwakuraRein,
aleozlx,
jiahanc,
nv-yunzheq,
samuellees and
yzh119
as code owners
July 20, 2026 11:40
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (2)
tests/moe/test_unified_moe_fuzz.py (1)
469-498: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win
prepare_weightsresult is unused for the MxFp8 path.
viewis computed unconditionally at Line 469 but only referenced in theDeepSeekFp8branch (Lines 480-485); theMxFp8branch (Lines 486-498) rebuilds weights via_mxfp8_quant_matrixand discardsview. Guarding the call avoids a redundant GPU weight preparation per MxFp8 reference call and removes a misleading dead computation.♻️ Move
prepare_weightsinto the DeepSeekFp8 branch- view = TrtllmFp8BlockConfig.prepare_weights( - w1, - w2, - variant=variant, - num_local_experts=w1.shape[0], - hidden_size=x.shape[1], - intermediate_size=intermediate_size, - device=x.device, - ) x32 = _block_fp8_dequant(x_q, x_sf, variant) if variant is QuantVariant.DeepSeekFp8: + view = TrtllmFp8BlockConfig.prepare_weights( + w1, + w2, + variant=variant, + num_local_experts=w1.shape[0], + hidden_size=x.shape[1], + intermediate_size=intermediate_size, + device=x.device, + ) w1_32 = _block_fp8_dequant( view["gemm1_weights"], view["gemm1_weights_scale"], variant ) w2_32 = _block_fp8_dequant( view["gemm2_weights"], view["gemm2_weights_scale"], variant )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/moe/test_unified_moe_fuzz.py` around lines 469 - 498, Move the TrtllmFp8BlockConfig.prepare_weights call into the DeepSeekFp8 branch, where its result is consumed by view["gemm1_weights"] and view["gemm2_weights"]. Leave the MxFp8 branch using _mxfp8_quant_matrix directly without computing or retaining view.tests/moe/test_unified_moe_fp8.py (1)
35-44: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse
flashinfer.utilsarch-support helpers instead of a manual arch tuple.
_is_trtllm_fp8_arch()hardcodesmajor*10+minor in (100, 103, 120, 121)and adds atorch.cuda.is_available()guard. Path instructions fortests/**/*.pycall for using canonical helpers likeis_sm100a_supported()/is_sm120a_supported()for arch gating, and this repo's tests assume CUDA is always available (no CPU guard needed in fixtures/skip helpers).♻️ Suggested refactor
-def _is_trtllm_fp8_arch() -> bool: - if not torch.cuda.is_available(): - return False - major, minor = get_compute_capability(torch.device("cuda")) - return major * 10 + minor in (100, 103, 120, 121) +def _is_trtllm_fp8_arch() -> bool: + device = torch.device("cuda") + return is_sm100a_supported(device) or is_sm120a_supported(device)As per path instructions, "Skip tests on unsupported CUDA architectures using flashinfer.utils functions like is_sm90a_supported(), is_sm100a_supported(), etc." Based on learnings, tests in this repo assume CUDA is available and avoid CPU-only guards outside of explicitly non-CUDA contexts.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/moe/test_unified_moe_fp8.py` around lines 35 - 44, Update _is_trtllm_fp8_arch() to use the canonical flashinfer.utils architecture-support helpers, such as is_sm100a_supported() and is_sm120a_supported(), instead of computing and comparing a hardcoded capability tuple. Remove the torch.cuda.is_available() guard and preserve support for the architectures required by the test skip condition.Sources: Path instructions, Learnings
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/moe/test_unified_moe_fp8.py`:
- Around line 35-44: Update _is_trtllm_fp8_arch() to use the canonical
flashinfer.utils architecture-support helpers, such as is_sm100a_supported() and
is_sm120a_supported(), instead of computing and comparing a hardcoded capability
tuple. Remove the torch.cuda.is_available() guard and preserve support for the
architectures required by the test skip condition.
In `@tests/moe/test_unified_moe_fuzz.py`:
- Around line 469-498: Move the TrtllmFp8BlockConfig.prepare_weights call into
the DeepSeekFp8 branch, where its result is consumed by view["gemm1_weights"]
and view["gemm2_weights"]. Leave the MxFp8 branch using _mxfp8_quant_matrix
directly without computing or retaining view.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 06bf879c-c6c7-48f7-8e15-9514665ee57e
📒 Files selected for processing (9)
flashinfer/fused_moe/__init__.pyflashinfer/fused_moe/api.pyflashinfer/fused_moe/core.pyflashinfer/fused_moe/layer.pyflashinfer/fused_moe/prepare.pyflashinfer/fused_moe/runners.pytests/moe/test_unified_moe.pytests/moe/test_unified_moe_fp8.pytests/moe/test_unified_moe_fuzz.py
Describe BF16, NVFP4, DeepSeek FP8, and MXFP8 payload/scale layouts and reflect block-FP8 FromLogits support.
Collaborator
Author
|
/bot run tests/moe |
Collaborator
Collaborator
|
[FAILED] Pipeline #58765989: 8/20 passed |
Avoid launching unsupported sm100f BMM cubins on SM120/121, align tests with the validated architecture family, and remove redundant MXFP8 reference preparation.
Resolve flashinfer-ai#3983 overlap by adopting runner-owned support validation and retaining b12x alongside the TRTLLM block-FP8 runner.
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
flashinfer/fused_moe/layer.py (1)
109-113: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMove runner instantiation inside the
tryblock.Moving
runner = runner_cls(config, device=self.device)inside thetryblock provides defense-in-depth. If a backend runner intentionally or inadvertently throws a validation error (ValueError,NotImplementedError, orRuntimeError) during its__init__phase, it will be safely caught and skipped instead of crashing the entire layer initialization.♻️ Proposed refactor
- runner = runner_cls(config, device=self.device) try: + runner = runner_cls(config, device=self.device) runner.check_support() except (NotImplementedError, ValueError, RuntimeError): continue🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/fused_moe/layer.py` around lines 109 - 113, Move the runner instantiation in the runner-selection loop inside the existing try block that calls runner.check_support(), so ValueError, NotImplementedError, and RuntimeError raised by runner_cls initialization are caught and skipped with unsupported backends.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@flashinfer/fused_moe/layer.py`:
- Around line 109-113: Move the runner instantiation in the runner-selection
loop inside the existing try block that calls runner.check_support(), so
ValueError, NotImplementedError, and RuntimeError raised by runner_cls
initialization are caught and skipped with unsupported backends.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: ff745c83-d8a9-4d4b-af18-38e8a21ffd78
📒 Files selected for processing (7)
flashinfer/fused_moe/__init__.pyflashinfer/fused_moe/api.pyflashinfer/fused_moe/core.pyflashinfer/fused_moe/layer.pyflashinfer/fused_moe/prepare.pyflashinfer/fused_moe/runners.pytests/moe/test_unified_moe.py
🚧 Files skipped from review as they are similar to previous changes (5)
- flashinfer/fused_moe/init.py
- flashinfer/fused_moe/core.py
- flashinfer/fused_moe/runners.py
- flashinfer/fused_moe/api.py
- flashinfer/fused_moe/prepare.py
Use the renamed unified routing fields so split BF16 and NVFP4 conformance tests execute again, and document the DeepSeek FP8 scale exception.
Collaborator
Author
|
/bot run tests/moe |
Collaborator
Collaborator
|
[SUCCESS] Pipeline #58896568: 14/20 passed |
Anerudhan
approved these changes
Jul 21, 2026
aleozlx
approved these changes
Jul 22, 2026
5 tasks
aleozlx
pushed a commit
that referenced
this pull request
Jul 24, 2026
## 📌 Description Add TRTLLM per-tensor FP8 execution to the unified `MoELayer` API. Also fixes the MXFP8 scale-layout issue reported in #4087. ### What changed * Added `TrtllmFp8PerTensorRunner` to the unified `MoELayer` API with: * `FromLogits` in-kernel routing * SM100 and SM103 support * Autotuning and CUDA graphs * Llama4 routing-scale-on-input behavior * Added per-tensor FP8 weight and activation preparation: * Per-expert E4M3 weight quantization * Calibrated per-tensor activation quantization * TRTLLM gated-row reorder and shuffled MajorK weights * Per-expert GEMM1 linear/gate and GEMM2 epilogue scales * Registered `QuantVariant.FP8PerTensor` with `MoELayer`. * Added independent-reference, routing-replay, CUDA-graph, architecture-gating, and unified-fuzzer coverage. * Fixed unified MXFP8 block-scale preparation by converting row-permuted GEMM1/GEMM2 scale tensors into TRTLLM’s required 128×4 interleaved layout. ## 🔍 Scope and follow-up Unified quantized MoE support is split into separate PRs because the formats use different kernels, scaling conventions, architectures, and execution contracts: * PR 1: block-scale DeepSeek FP8 and MXFP8 — merged in #4026 * PR 2: per-tensor FP8 — this PR * PR 3: CUTLASS W4A8 — follow-up This PR keeps per-tensor FP8 `FromLogits`-only, matching the existing TRTLLM kernel entry point. The legacy flat APIs remain unchanged. ## 🚀 Pull Request Checklist ### ✅ Pre-commit Checks * I have installed `pre-commit` by running `pip install pre-commit` (or used my 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. ## 🧪 Tests Validated on SM100: * `18 passed` — unified block-FP8 and per-tensor-FP8 conformance suites * `7 passed` — focused unified API/config validation * `1 passed` — curated per-tensor FP8 fuzzer profile with production autotuning * `1 passed` — MXFP8 seed `900013` regression, including valid-tactic coverage * `1 passed` — exact reported MXFP8 seed-100 configuration, including valid-tactic coverage ### MXFP8 #4087 validation The fuzzer and legacy MXFP8 references agree to within `1.16e-10`. Unified and legacy weight payloads were identical, but the unified scale tensors were missing the 128×4 physical interleave. With the corrected scale preparation: ```text before: 110 / 524288 elements over tolerance after: 0 / 524288 elements over tolerance ``` <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added TensorRT-LLM FP8 per-tensor quantization for Mixture-of-Experts, including integration into the cross-backend MoE layer. * Exposed a new FP8 per-tensor runner via the public API (FromLogits routing). * Added calibrated global scaling support for FP8 weights and activations. * **Bug Fixes** * Tightened FP8/MXFP8 preparation validation for MXFP8 and clarified the FP8 layout constraints. * Refined hardware support to SM100 family values (supported: 100/103; unsupported: 90/120). * **Tests** * Expanded FP8 per-tensor correctness and replay coverage, plus new negative/behavioral tests. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
3 tasks
feih-nv
added a commit
that referenced
this pull request
Jul 29, 2026
Expose the existing TRTLLM-gen `MxFP4xMxFP8` (W4A8) and `MxFP4xBf16` (W4A16) kernels through the unified MoE API. ## 📌 Description This is PR 3 in the unified MoE quantization series for FP8 support: 1. #4026 — unified block-scale FP8 (merged) 2. #4091 — unified per-tensor FP8 (merged) 3. This PR — unified TRTLLM MXFP4×MXFP8/W4A8 and MXFP4×BF16/W4A16 ### Changes - Generalize `TrtllmFp4RoutedRunner` beyond NVFP4: - `QuantVariant.MXFP4`: `MxE2m1` weights × `MxE4m3` activations - `QuantVariant.W4A16`: `MxE2m1` weights × BF16 activations - Add variant-aware TRTLLM FP4 preparation: - MXFP4 weights with 32-element UE8M0 scales - MXFP8 activation preparation for W4A8 - BF16 activation preparation for W4A16 - Add shape, dtype, and scale-layout validation. - Add unified conformance and fuzzer coverage for packed and `FromLogits` routing. ### Support matrix - NVFP4 and MXFP4/W4A8: SM100, SM103 - W4A16: SM100 only (remains disabled on SM103, matching upstream xfail #1754) - SM120/121: separate b12x backends where available - SM107 is unsupported after #4171 ### Scope - No CUDA/C++ kernel changes; both modes already exist in the TRTLLM flat API. - CUTLASS W4A8 is out of scope. ## 🚀 Pull Request Checklist ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. ## 🧪 Tests SM100, CUDA 13 CI container: - `tests/moe/test_unified_moe_mxfp4.py` - 21 passed - Unified fuzzer, packed routing: - seeds `900017,900018` - 2 passed - Unified fuzzer, `FromLogits`: - seeds `900019,900020` - 2 passed
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
4 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.
📌 Description
Adds first-class DeepSeek FP8 and MXFP8 block-scale execution through the unified
MoELayerAPI. These formats previously existed in the configuration schema and legacy flat APIs, but were not executable throughMoELayer.Fixed one issue in
tests/moe_ep/test_split_fused_moe_kernel_vs_reference.py.What changed
TrtllmFp8BlockRunnerto the unifiedMoELayerAPI with:🔍 Scope and follow-up
Unified FP8 support is split into two PRs because block-scale and per-tensor FP8 have distinct scaling and execution contracts. This PR covers block-scale FP8.
Follow-up work for per-tensor FP8
TrtllmFp8PerTensorRunnerQuantVariant.FP8PerTensorwithMoELayerFromLogits, matching the legacy APIAdditional future work
do_finalize=FalseExisting legacy tests remain under the Mirror → Bridge → Prune policy. The legacy flat APIs are unchanged. The unified runner delegates to the existing TRTLLM
MoERunnerand kernel entry points.🚀 Pull Request Checklist
✅ 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
Validated on SM100:
11 passed— unified block-FP8 conformance suite19 passed— unified API validation subset2 passed— curated DeepSeek FP8 and MXFP8 fuzzer cases with production autotuning1 passed— existing FP4FromLogitsfuzzer confirmation on SM100Summary by CodeRabbit