[None][feat] Optimize trtllmgen moe routing#15656
Merged
Merged
Conversation
Contributor
📝 WalkthroughWalkthroughRouting custom launch code now uses tier-specific launch configs to select cluster and histogram kernel variants. Top-k reduction now handles packed 32/64-bit comparisons differently, and non-power-of-two sorting uses a compile-time Batcher network. ChangesRouting launch dispatch
Top-k reduction and sorting
Sequence Diagram(s)sequenceDiagram
participant launchClusterKernel
participant launchClusterKernelForPolicy
participant launchClusterKernelForBlockDim
participant routingIndicesClusterKernel256
participant routingIndicesClusterKernel512
participant routingIndicesClusterKernel
participant routingIndicesClusterKernelBody
launchClusterKernel->>launchClusterKernelForPolicy: dispatch tier pair from mNumTokens
launchClusterKernelForPolicy->>launchClusterKernelForBlockDim: select block-size-specific launch
launchClusterKernelForBlockDim->>routingIndicesClusterKernel256: launch 256-thread wrapper
launchClusterKernelForBlockDim->>routingIndicesClusterKernel512: launch 512-thread wrapper
launchClusterKernelForBlockDim->>routingIndicesClusterKernel: launch 1024-thread wrapper
routingIndicesClusterKernel256->>routingIndicesClusterKernelBody: run shared cluster body
routingIndicesClusterKernel512->>routingIndicesClusterKernelBody: run shared cluster body
routingIndicesClusterKernel->>routingIndicesClusterKernelBody: run shared cluster body
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustom.cu (1)
1-1: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winApply clang-format changes from pre-commit.
The release check reports clang-format modified this file, so the committed version is not formatted as CI expects. As per coding guidelines, “Use the LLVM clang-format tool for formatting your changes prior to submitting the PR.”
🤖 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 `@cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustom.cu` at line 1, Apply the pending clang-format changes to RoutingCustom.cu so the committed version matches pre-commit/CI formatting expectations; reformat the file with the LLVM clang-format tool and ensure the resulting style is consistent throughout the blockScaleMoe routing kernel code, including any surrounding declarations in RoutingCustom.cu.Sources: Coding guidelines, Pipeline failures
cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustomPolicy.cuh (1)
1-1: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winApply clang-format changes from pre-commit.
The release check reports clang-format modified this file, so the committed version is not formatted as CI expects. As per coding guidelines, “Use the LLVM clang-format tool for formatting your changes prior to submitting the PR.”
🤖 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 `@cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustomPolicy.cuh` at line 1, The committed version of RoutingCustomPolicy.cuh is not matching the expected clang-format style, so apply the pre-commit clang-format changes and reformat the affected declarations/includes/comments in that header. Use the existing file content around RoutingCustomPolicy and any nearby blockScaleMoe routing symbols to locate the unformatted sections, then make sure the final version is fully clang-format clean before submitting.Sources: Coding guidelines, Pipeline failures
🧹 Nitpick comments (2)
cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernelTopK.cuh (2)
100-106: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winKeep the fast-redux temporaries
const.Line 100, Line 101, and Line 104 initialize values that are not modified before use.
Proposed cleanup
- uint32_t hi = static_cast<uint32_t>(compVal >> 32); - uint32_t lo = static_cast<uint32_t>(compVal & 0xffffffffu); + uint32_t const hi = static_cast<uint32_t>(compVal >> 32); + uint32_t const lo = static_cast<uint32_t>(compVal & 0xffffffffu); uint32_t maxHi; asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;\n" : "=r"(maxHi) : "r"(hi)); - uint32_t loContrib = (hi == maxHi) ? lo : 0u; + uint32_t const loContrib = (hi == maxHi) ? lo : 0u;As per coding guidelines, “Variables not modified after initialization should be declared as
const” and “Use east-const style.”🤖 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 `@cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernelTopK.cuh` around lines 100 - 106, The fast-redux temporaries in the RoutingKernelTopK.cuh block are never reassigned after initialization, so update the local declarations for hi, lo, maxHi, loContrib, and maxLo to use east-const style where applicable. Keep the existing redux.sync.max.u32 flow unchanged, but make the immutable values const to match the coding guideline and avoid non-const temporaries.Source: Coding guidelines
140-148: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse
k-prefixed camelCase names for new constexpr constants.The new constants
u,s1...s5,M,Half, andPaddedNshould follow the project constant naming convention.Example direction
- static constexpr unsigned u = static_cast<unsigned>(N - 1); - static constexpr unsigned s1 = u | (u >> 1); + static constexpr unsigned kInputMinusOne = static_cast<unsigned>(N - 1); + static constexpr unsigned kStage1 = kInputMinusOne | (kInputMinusOne >> 1);- constexpr int M = R * 2; - if constexpr (M < N) + constexpr int kMergeStride = R * 2; + if constexpr (kMergeStride < N)- constexpr int PaddedN = NextPow2<N>::value; - topkSortBatcher<0, PaddedN, N, RedType>(topK); + constexpr int kPaddedN = NextPow2<N>::value; + topkSortBatcher<0, kPaddedN, N, RedType>(topK);As per coding guidelines, “Constants should follow naming convention: camelCase with prefix 'k'.”
Also applies to: 186-191, 210-213, 269-270
🤖 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 `@cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernelTopK.cuh` around lines 140 - 148, The new constexpr constants in RoutingKernelTopK.cuh do not follow the project’s k-prefixed camelCase convention. Rename the internal helper constants in the affected templates and utilities, including u/s1/s2/s3/s4/s5, M, Half, and PaddedN, to k-prefixed camelCase names, and update all references within the relevant top-k/routing helpers so the code remains consistent and compiles.Source: Coding guidelines
🤖 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.
Inline comments:
In
`@cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernelTopK.cuh`:
- Line 1: The file is failing the Release Checks because its formatting does not
match clang-format output; rerun clang-format on the affected
RoutingKernelTopK.cuh content and commit the resulting formatting-only changes
so the pre-commit check passes.
---
Outside diff comments:
In
`@cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustom.cu`:
- Line 1: Apply the pending clang-format changes to RoutingCustom.cu so the
committed version matches pre-commit/CI formatting expectations; reformat the
file with the LLVM clang-format tool and ensure the resulting style is
consistent throughout the blockScaleMoe routing kernel code, including any
surrounding declarations in RoutingCustom.cu.
In
`@cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustomPolicy.cuh`:
- Line 1: The committed version of RoutingCustomPolicy.cuh is not matching the
expected clang-format style, so apply the pre-commit clang-format changes and
reformat the affected declarations/includes/comments in that header. Use the
existing file content around RoutingCustomPolicy and any nearby blockScaleMoe
routing symbols to locate the unformatted sections, then make sure the final
version is fully clang-format clean before submitting.
---
Nitpick comments:
In
`@cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernelTopK.cuh`:
- Around line 100-106: The fast-redux temporaries in the RoutingKernelTopK.cuh
block are never reassigned after initialization, so update the local
declarations for hi, lo, maxHi, loContrib, and maxLo to use east-const style
where applicable. Keep the existing redux.sync.max.u32 flow unchanged, but make
the immutable values const to match the coding guideline and avoid non-const
temporaries.
- Around line 140-148: The new constexpr constants in RoutingKernelTopK.cuh do
not follow the project’s k-prefixed camelCase convention. Rename the internal
helper constants in the affected templates and utilities, including
u/s1/s2/s3/s4/s5, M, Half, and PaddedN, to k-prefixed camelCase names, and
update all references within the relevant top-k/routing helpers so the code
remains consistent and compiles.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 3e98d90c-989a-4ff4-b484-0a496c148e51
📒 Files selected for processing (3)
cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustom.cucpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustomPolicy.cuhcpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernelTopK.cuh
jiahanc
changed the title
yweng0828
approved these changes
Jun 26, 2026
Collaborator
Author
|
/bot run --disable-fail-fast --test-backend "pytorch" |
Collaborator
|
PR_Github #56060 [ run ] triggered by Bot. Commit: |
tongyuantongyu
approved these changes
Jun 26, 2026
Collaborator
|
PR_Github #56060 [ run ] completed with state
|
Signed-off-by: jiahanc <173873397+jiahanc@users.noreply.github.com>
Signed-off-by: jiahanc <173873397+jiahanc@users.noreply.github.com>
jiahanc
force-pushed
the
opt-trtllngen-moe-routing
branch
from
f605ca5 to
e6c548b
Compare
Collaborator
Author
|
/bot run --disable-fail-fast --test-backend "pytorch" |
Collaborator
|
PR_Github #56175 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #56175 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #56250 [ run ] triggered by Bot. Commit: |
jiahanc
added a commit
to flashinfer-ai/flashinfer
that referenced
this pull request
Jun 29, 2026
Port routing launch policy and reduced-block histogram/cluster dispatch from NVIDIA/TensorRT-LLM#15656 into FlashInfer's TRTLLMGEN MoE path. AI-assisted.
Collaborator
|
PR_Github #56250 [ run ] completed with state |
jiahanc
added a commit
to flashinfer-ai/flashinfer
that referenced
this pull request
Jun 29, 2026
Port routing launch policy and reduced-block histogram/cluster dispatch from NVIDIA/TensorRT-LLM#15656 into FlashInfer's TRTLLMGEN MoE path. AI-assisted.
evezhier
pushed a commit
to evezhier/TensorRT-LLM
that referenced
this pull request
Jul 2, 2026
Signed-off-by: jiahanc <173873397+jiahanc@users.noreply.github.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
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
Optimize TRTLLM-Gen MoE routing to reduce register pressure and avoid heavy spilling in high-expert routing cases.
Test Coverage
./cpp/tests/unit_tests/kernels/routingKernelsTest
Perf Benchmark
noop_softmax
softmax_sum
sigmoid_bias_scaled
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.