You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
[https://nvbugs/7012145][fix] Support non-power-of-2 expert counts in Triton FP8 MoE align - #21923
_moe_prefix_kernel now scans triton.next_power_of_2(num_experts) lanes. It masks loads and stores for lanes beyond NUM_EXPERTS; those lanes contribute zero to the padded-count sum. This removes the non-power-of-two range that Triton rejected. No public API changes.
The test shape list adds a 72-expert case for SM120 MoE correctness tests. The alignment reference test also covers non-power-of-two expert counts. Runtime and hardware validation results are not available.
QA Engineer Review
The test module adds test_moe_align_block_size_matches_reference, parameterized for 8, 60, 72, 256, and 288 experts. It checks the padded token total, expert block IDs, assigned token IDs, and padding slots against a Torch reference. The H100 CI list includes this test selector. No manual-QA list entry was found. Coverage verdict: sufficient for the alignment behavior changed here. Completed test results are not available.
Per-File QA Perspective
tensorrt_llm/_torch/moe/fused_moe/fused_moe_triton_fp8_block_scale.py: The prefix scan now handles non-power-of-two expert counts with a padded range and masked memory access. Verify offsets and total padded tokens for both power-of-two and non-power-of-two counts.
tests/unittest/_torch/moe/fused_moe/test_triton_fp8_block_scale.py: Adds the alignment reference test and an experts_72 MoE correctness shape. The alignment test is listed in the H100 CI test database; no manual-QA listing was found.
tests/integration/test_lists/test-db/l0_h100.yml: Adds the alignment test selector to the H100 CI list. Execution results are not available.
Description
moe_align_block_size in the Triton FP8 block-scale MoE fails to compile whenever the per-rank expert count is not a power of 2 (NVBug 7012145). _moe_prefix_kernel builds its expert range with tl.arange(0, NUM_EXPERTS), and Triton rejects that with arange's range must be a power of 2. DeepSeek-V3-Lite routes over 72 experts, so TestDeepSeekV3Lite::test_fp8_block_scales fails at model load on SM120/SM121, where TrtllmTritonFp8BlockScalesImpl serves FP8 block-scale MoE. The kernel came in with #13559 for Laguna-XS.2, which has 256 experts, so this path had never run with a non-power-of-2 count.
Without this PR, any model or EP split whose local expert count is not a power of 2 (60, 72, 160, 288, ...) cannot use this MoE path. With it, the prefix kernel pads its range to triton.next_power_of_2(num_experts) and masks both the load (other=0) and the store. Padded lanes add 0 to the prefix scan and to tl.sum, so offsets and num_post_pad are unchanged for power-of-2 counts. The other kernels in this file already index with fixed block sizes plus masks, so this kernel is the only one that needed the change.
Apart from one comment line, the kernel change is identical to the one in draft PR #21250, which also adds an SM90 FP8 path. It is split out here so the P0 bug can land on its own; rebasing #21250 over it only needs the test additions reconciled.
This fix targets main; no cherry-pick to release/1.3 is planned.
Test Coverage
New test_moe_align_block_size_matches_reference[8, 60, 72, 256, 288] compares num_post_pad, expert_ids, and each expert's sorted_token_ids segment and its padding against a torch reference. The alignment kernels do not depend on the architecture, so the test only requires CUDA and is added to the l0_h100.yml pre-merge MoE components list. The rest of test_triton_fp8_block_scale.py needs SM120 and is not in any CI list today. This PR does not add it to an SM120 pre-merge list: SM120 coverage of this path stays with QA, through TestDeepSeekV3Lite::test_fp8_block_scales in qa/llm_function_core.txt.
New experts_72 shape (DeepSeek-V3-Lite's expert count, top-6) in the SM120 end-to-end correctness tests.
Checked on H100 with Triton 3.8.0. On main, the new test body fails with CompilationError for 60, 72 and 288 experts and passes for 8 and 256; with this change all five pass. A standalone run_triton_fp8_block_scale_moe run against a BF16 reference with 60, 64, 72 and 160 experts gives cosine-diff of about 1.8e-5. Not run on SM120/SM121 hardware.
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.
Any new dependencies have been scanned for license and vulnerabilities
… Triton FP8 MoE align
_moe_prefix_kernel built its expert range with tl.arange(0, NUM_EXPERTS),
which Triton only accepts for powers of 2, so moe_align_block_size failed
to compile for models such as DeepSeek-V3-Lite (72 routed experts).
Pad the range to the next power of 2 and mask the load and store; padded
lanes read a count of 0 and do not change the prefix sum or the total.
Add an architecture-independent alignment test and run it on H100.
Signed-off-by: xxi <xxi@nvidia.com>
No actionable comments were generated in the recent review.
Walkthrough
The FP8 block-scale MoE prefix kernel now scans a power-of-two expert range and masks accesses to the actual expert count. Tests cover alignment for several expert counts, including 72, and add the reference test to the H100 pre-merge list.
The prefix kernel scans the next power-of-two expert range. It masks loads and stores to the actual expert count, and the caller passes the scan width.
Tests add a 72-expert case and validate alignment for expert counts 8, 60, 72, 256, and 288. The H100 pre-merge list adds the alignment reference test.
The change supports non-power-of-two expert counts, with reference coverage on H100. No actionable merge-blocking risk was identified.
Pre-merge checks: 5
Check name
Status
Explanation
Title check
The title names the main change: non-power-of-2 expert counts are now supported in the Triton FP8 MoE align path. It also includes the NVBugs ID and the [fix] type, which match the repository template…
Docstring Coverage
Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentat…
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.
Description check
The description follows the repository template closely. It covers the issue (non-power-of-2 expert counts fail to compile in _moe_prefix_kernel under tl.arange, NVBug 7012145), the fix (pad to `t…
Full details: Title check
Explanation
The title names the main change: non-power-of-2 expert counts are now supported in the Triton FP8 MoE align path. It also includes the NVBugs ID and the [fix] type, which match the repository template. The title is specific and concise, so it describes the primary change in the changeset.
Full details: Docstring Coverage
Explanation
Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 2 functions across 2 files.
Full details: Description check
Explanation
The description follows the repository template closely. It covers the issue (non-power-of-2 expert counts fail to compile in _moe_prefix_kernel under tl.arange, NVBug 7012145), the fix (pad to triton.next_power_of_2(num_experts) and mask load and store), and a detailed Test Coverage section with the new parametrized test, the experts_72 SM120 shape, the H100/Triton 3.8.0 results, and the hardware not tested (SM120/SM121). The PR Checklist is ticked. The title is not part of the description body, so its format cannot be judged here. Minor gaps: the Description section does not mention the draft PR #21250 relationship in the template's terms, and the checklist is ticked without stating which items were verified (for example, documentation or API label).
Finishing Touches
Generate unit tests (beta)
Create a new PR
Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.
The reason will be displayed to describe this comment to others. Learn more.
Thanks for splitting the P0 fix out. No blockers: the kernel change is correct. NUM_EXPERTS_POW2 covers every expert lane, padded lanes load a count of 0 so the scan and tl.sum are unchanged, the masked store stays inside out_offsets, power-of-2 counts behave exactly as before, and EP/EPLB splits (the leaf passes local slot ids, e.g. 72/4 = 18 experts per rank) are fixed by the same change. Counts: 0 Blocker, 0 Possibly blocking, 2 Required, 3 Nice to have (one Required and one Nice to have below, the rest inline).
[Required] Cherry-pick: release/1.3 still has the unfixed kernel. _moe_prefix_kernel there does e = tl.arange(0, NUM_EXPERTS) (fused_moe_triton_fp8_block_scale.py:45, launched with only NUM_EXPERTS), and fused_moe_cutlass.py:957 routes SM120 FP8 block-scale MoE to it (if self.has_deepseek_fp8_block_scales and get_sm_version() == 120:). So DeepSeek-V3-Lite FP8 (72 experts) on RTX Pro 6000 fails the same way on 1.3. Please say whether NVBug 7012145 also covers 1.3 and, if it does, open the cherry-pick; the kernel and launch hunks match the 1.3 code.
[Nice to have] Validation: the description says the fix was not run on SM120/SM121, and no /bot run has been triggered yet. Before closing the NVBug, could you post one SM120 (or SM121) run of the failing test, e.g. accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales[mtp=disable-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False], plus test_triton_fp8_block_scale.py on that GPU, and trigger /bot run so the new H100 entry has a passing pre-merge run?
Automatically added "ci: full pre-merge approved" because this PR has satisfied the required GitHub review approvals. Unresolved review conversations and other required checks remain independent merge requirements.
Explain the power-of-2 lane padding in _moe_prefix_kernel, keep the skip
markers together, and list moe_align_block_size in the test docstring.
Signed-off-by: xxi <xxi@nvidia.com>
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
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.
sunnyqgg
left a comment
Dev Engineer Review
_moe_prefix_kernelnow scanstriton.next_power_of_2(num_experts)lanes. It masks loads and stores for lanes beyondNUM_EXPERTS; those lanes contribute zero to the padded-count sum. This removes the non-power-of-two range that Triton rejected. No public API changes.The test shape list adds a 72-expert case for SM120 MoE correctness tests. The alignment reference test also covers non-power-of-two expert counts. Runtime and hardware validation results are not available.
QA Engineer Review
The test module adds
test_moe_align_block_size_matches_reference, parameterized for 8, 60, 72, 256, and 288 experts. It checks the padded token total, expert block IDs, assigned token IDs, and padding slots against a Torch reference. The H100 CI list includes this test selector. No manual-QA list entry was found. Coverage verdict: sufficient for the alignment behavior changed here. Completed test results are not available.Per-File QA Perspective
tensorrt_llm/_torch/moe/fused_moe/fused_moe_triton_fp8_block_scale.py: The prefix scan now handles non-power-of-two expert counts with a padded range and masked memory access. Verify offsets and total padded tokens for both power-of-two and non-power-of-two counts.tests/unittest/_torch/moe/fused_moe/test_triton_fp8_block_scale.py: Adds the alignment reference test and anexperts_72MoE correctness shape. The alignment test is listed in the H100 CI test database; no manual-QA listing was found.tests/integration/test_lists/test-db/l0_h100.yml: Adds the alignment test selector to the H100 CI list. Execution results are not available.Description
moe_align_block_sizein the Triton FP8 block-scale MoE fails to compile whenever the per-rank expert count is not a power of 2 (NVBug 7012145)._moe_prefix_kernelbuilds its expert range withtl.arange(0, NUM_EXPERTS), and Triton rejects that witharange's range must be a power of 2. DeepSeek-V3-Lite routes over 72 experts, soTestDeepSeekV3Lite::test_fp8_block_scalesfails at model load on SM120/SM121, whereTrtllmTritonFp8BlockScalesImplserves FP8 block-scale MoE. The kernel came in with #13559 for Laguna-XS.2, which has 256 experts, so this path had never run with a non-power-of-2 count.Without this PR, any model or EP split whose local expert count is not a power of 2 (60, 72, 160, 288, ...) cannot use this MoE path. With it, the prefix kernel pads its range to
triton.next_power_of_2(num_experts)and masks both the load (other=0) and the store. Padded lanes add 0 to the prefix scan and totl.sum, so offsets andnum_post_padare unchanged for power-of-2 counts. The other kernels in this file already index with fixed block sizes plus masks, so this kernel is the only one that needed the change.Apart from one comment line, the kernel change is identical to the one in draft PR #21250, which also adds an SM90 FP8 path. It is split out here so the P0 bug can land on its own; rebasing #21250 over it only needs the test additions reconciled.
This fix targets main; no cherry-pick to
release/1.3is planned.Test Coverage
test_moe_align_block_size_matches_reference[8, 60, 72, 256, 288]comparesnum_post_pad,expert_ids, and each expert'ssorted_token_idssegment and its padding against a torch reference. The alignment kernels do not depend on the architecture, so the test only requires CUDA and is added to thel0_h100.ymlpre-merge MoE components list. The rest oftest_triton_fp8_block_scale.pyneeds SM120 and is not in any CI list today. This PR does not add it to an SM120 pre-merge list: SM120 coverage of this path stays with QA, throughTestDeepSeekV3Lite::test_fp8_block_scalesinqa/llm_function_core.txt.experts_72shape (DeepSeek-V3-Lite's expert count, top-6) in the SM120 end-to-end correctness tests.CompilationErrorfor 60, 72 and 288 experts and passes for 8 and 256; with this change all five pass. A standalonerun_triton_fp8_block_scale_moerun against a BF16 reference with 60, 64, 72 and 160 experts gives cosine-diff of about 1.8e-5. Not run on SM120/SM121 hardware.PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.