enable batch-M trtllm-gen MoE; refactor trtllm_fused_moe_kernel_launcher - #4281
Open
IwakuraRein wants to merge 10 commits into
Open
IwakuraRein wants to merge 10 commits into
IwakuraRein wants to merge 10 commits into
Conversation
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (4)
🚧 Files skipped from review as they are similar to previous changes (4)
📝 WalkthroughWalkthroughThe PR updates TRT-LLM batched GEMM options to use explicit activation and weight dtypes. Config matching, ranking, validation, workspace sizing, and execution now derive orientation and dimensions from candidate or selected configurations. The artifact pin and checksum are updated. ChangesTRT-LLM GEMM configuration
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Runner
participant CandidateConfig
participant ConfigValidator
participant Kernel
Runner->>CandidateConfig: derive orientation and tile dimensions
CandidateConfig->>ConfigValidator: build candidate GEMM problem
ConfigValidator-->>Runner: return selected config
Runner->>Kernel: size workspace and execute selected configuration
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 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 |
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
Signed-off-by: Siyuan Fu <siyuanf@nvidia.com>
Signed-off-by: Siyuan Fu <siyuanf@nvidia.com>
IwakuraRein
force-pushed
the
siyuanf/batch-m-3xfp4-moe
branch
from
August 3, 2026 20:58
0097931 to
64567d5
Compare
Collaborator
Author
|
/bot run tests/moe |
Collaborator
IwakuraRein
marked this pull request as ready for review
August 3, 2026 23:22
IwakuraRein
requested review from
aleozlx,
sricketts,
yongwww and
yzh119
as code owners
August 3, 2026 23:22
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
IwakuraRein
requested review from
Aneureka,
bkryu,
cyx-6,
dhiraj113,
feih-nv,
jiahanc,
jimmyzho,
kahyunnam,
nv-yunzheq,
qiching,
qsang-nv,
saltyminty,
samuellees and
yyihuang
as code owners
August 3, 2026 23:22
![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: 3
🤖 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 `@csrc/trtllm_batched_gemm_runner.cu`:
- Around line 367-368: Remove the unconditional selected-config printf near
transposeMmaOutput in csrc/trtllm_batched_gemm_runner.cu:367-368, and remove the
multiprocessor-count printf at csrc/trtllm_batched_gemm_runner.cu:401-407 while
preserving the computed multiprocessor value passed to bmm.run. Do not alter
GEMM execution behavior or add diagnostics unless routed through an explicit
debug facility.
- Around line 401-402: Check the return value of cudaDeviceGetAttribute in the
code initializing multiProcessorCount, using the repository’s established CUDA
error-checking macro. Ensure bmm.run is reached only after successful attribute
retrieval so multiProcessorCount is initialized.
- Around line 482-494: Update cmpFunc to build BatchedGemmData for both idx0 and
idx1 via buildGemmData, then compute each candidate’s rank using its own CTA
estimate before comparing optionsA and optionsB. Ensure the comparator remains
symmetric when mTransposeMmaOutput differs, and apply a deterministic
tie-breaker when the candidate ranks are equal.
🪄 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: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 10afd789-3ca9-4866-a2dc-7e6bff12add4
📒 Files selected for processing (4)
csrc/trtllm_batched_gemm_runner.cucsrc/trtllm_fused_moe_runner.cuflashinfer/artifacts.pyinclude/flashinfer/trtllm/batched_gemm/KernelRunner.h
Comment on lines
401
to
402
| int32_t multiProcessorCount; | ||
| cudaDeviceGetAttribute(&multiProcessorCount, cudaDevAttrMultiProcessorCount, device); |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
rg -n -C 2 '\bcudaDeviceGetAttribute\s*\(' csrc include
rg -n -C 2 '\bCUDACHECK\s*\(' csrc includeRepository: flashinfer-ai/flashinfer
Length of output: 49636
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
# Inspect the target function and nearby error-check macros before making a final claim.
wc -l csrc/trtllm_batched_gemm_runner.cu
sed -n '360,430p' csrc/trtllm_batched_gemm_runner.cu
sed -n '1,80p' include/flashinfer/trtllm/common.hRepository: flashinfer-ai/flashinfer
Length of output: 7168
Check the cudaDeviceGetAttribute result.
If cudaDeviceGetAttribute fails, multiProcessorCount is written only on success, so bmm.run can receive an undefined processor count. Use the repository CUDA error-checking macro.
🤖 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 `@csrc/trtllm_batched_gemm_runner.cu` around lines 401 - 402, Check the return
value of cudaDeviceGetAttribute in the code initializing multiProcessorCount,
using the repository’s established CUDA error-checking macro. Ensure bmm.run is
reached only after successful attribute retrieval so multiProcessorCount is
initialized.
Comment on lines
+482
to
+494
| auto buildGemmData = [&](int64_t configIndex) { | ||
| BatchedGemmData data{}; | ||
| bool const transposeMmaOutput = configs[configIndex].mOptions.mTransposeMmaOutput; | ||
| fillGemmProblemDimensions(data, transposeMmaOutput, m, n, k, batchedTokens, numTokens, | ||
| numBatches, maxNumCtasInBatchDim); | ||
| return data; | ||
| }; | ||
|
|
||
| auto cmpFunc = [&configs, &gemmData, &bmm, &multiProcessorCount](int64_t idx0, int64_t idx1) { | ||
| auto cmpFunc = [&configs, &bmm, &multiProcessorCount, &buildGemmData](int64_t idx0, | ||
| int64_t idx1) { | ||
| auto const& optionsA = configs[idx0].mOptions; | ||
| auto const& optionsB = configs[idx1].mOptions; | ||
| auto const gemmData = buildGemmData(idx0); |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Make cmpFunc symmetric across candidate configurations.
Line 494 builds gemmData only for idx0. Lines 529-535 then use that data to choose between optionsA and optionsB. If the candidates use different mTransposeMmaOutput values, reversing the comparator inputs can change the CTA estimate and make both comparison directions return true.
Build and evaluate problem data for both candidates. Rank each candidate with its own CTA estimate before applying a deterministic tie-breaker.
Proposed comparator fix
- auto const gemmData = buildGemmData(idx0);
- int32_t sizeK = gemmData.mProblemDimensions.mK;
+ auto const gemmDataA = buildGemmData(idx0);
+ auto const gemmDataB = buildGemmData(idx1);
+ int32_t sizeK = gemmDataA.mProblemDimensions.mK;
...
if (optionsA.mTileScheduler != optionsB.mTileScheduler) {
- auto options = bmm.getOptionsFromConfigAndData(configs[idx0], gemmData);
- auto numCtas = bmm.getNumCtas(options, gemmData.mProblemDimensions.mMaxNumCtasInTokenDim);
- if (numCtas > multiProcessorCount) {
- return optionsA.mTileScheduler == batchedGemm::gemm::TileScheduler::Persistent;
- } else {
- return optionsB.mTileScheduler == batchedGemm::gemm::TileScheduler::Persistent;
- }
+ auto const bmmOptionsA = bmm.getOptionsFromConfigAndData(configs[idx0], gemmDataA);
+ auto const bmmOptionsB = bmm.getOptionsFromConfigAndData(configs[idx1], gemmDataB);
+ auto const numCtasA =
+ bmm.getNumCtas(bmmOptionsA, gemmDataA.mProblemDimensions.mMaxNumCtasInTokenDim);
+ auto const numCtasB =
+ bmm.getNumCtas(bmmOptionsB, gemmDataB.mProblemDimensions.mMaxNumCtasInTokenDim);
+ auto const preferredA =
+ (numCtasA > multiProcessorCount) ==
+ (optionsA.mTileScheduler == batchedGemm::gemm::TileScheduler::Persistent);
+ auto const preferredB =
+ (numCtasB > multiProcessorCount) ==
+ (optionsB.mTileScheduler == batchedGemm::gemm::TileScheduler::Persistent);
+ if (preferredA != preferredB) return preferredA;
+ return optionsA.mTileScheduler == batchedGemm::gemm::TileScheduler::Persistent;
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| auto buildGemmData = [&](int64_t configIndex) { | |
| BatchedGemmData data{}; | |
| bool const transposeMmaOutput = configs[configIndex].mOptions.mTransposeMmaOutput; | |
| fillGemmProblemDimensions(data, transposeMmaOutput, m, n, k, batchedTokens, numTokens, | |
| numBatches, maxNumCtasInBatchDim); | |
| return data; | |
| }; | |
| auto cmpFunc = [&configs, &gemmData, &bmm, &multiProcessorCount](int64_t idx0, int64_t idx1) { | |
| auto cmpFunc = [&configs, &bmm, &multiProcessorCount, &buildGemmData](int64_t idx0, | |
| int64_t idx1) { | |
| auto const& optionsA = configs[idx0].mOptions; | |
| auto const& optionsB = configs[idx1].mOptions; | |
| auto const gemmData = buildGemmData(idx0); | |
| auto buildGemmData = [&](int64_t configIndex) { | |
| BatchedGemmData data{}; | |
| bool const transposeMmaOutput = configs[configIndex].mOptions.mTransposeMmaOutput; | |
| fillGemmProblemDimensions(data, transposeMmaOutput, m, n, k, batchedTokens, numTokens, | |
| numBatches, maxNumCtasInBatchDim); | |
| return data; | |
| }; | |
| auto cmpFunc = [&configs, &bmm, &multiProcessorCount, &buildGemmData](int64_t idx0, | |
| int64_t idx1) { | |
| auto const& optionsA = configs[idx0].mOptions; | |
| auto const& optionsB = configs[idx1].mOptions; | |
| auto const gemmDataA = buildGemmData(idx0); | |
| auto const gemmDataB = buildGemmData(idx1); |
🤖 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 `@csrc/trtllm_batched_gemm_runner.cu` around lines 482 - 494, Update cmpFunc to
build BatchedGemmData for both idx0 and idx1 via buildGemmData, then compute
each candidate’s rank using its own CTA estimate before comparing optionsA and
optionsB. Ensure the comparator remains symmetric when mTransposeMmaOutput
differs, and apply a deterministic tie-breaker when the candidate ranks are
equal.
5 tasks
Collaborator
|
[SUCCESS] Pipeline #60920011: 18/18 executed test jobs passed |
Signed-off-by: Siyuan Fu <siyuanf@nvidia.com>
Contributor
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
Collaborator
|
[SUCCESS] Pipeline #60920011: 18/18 executed test jobs passed |
1 similar comment
Collaborator
|
[SUCCESS] Pipeline #60920011: 18/18 executed test jobs 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
2 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
Refactor
csrc/trtllm_batched_gemm_runner.cu; construct batched gemm data according tomTransposeMmaOutput.🔍 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.).Reviewer Notes
Summary by CodeRabbit
Performance
Bug Fixes
API Updates