[TRTLLM-14701][feat] Update trtllm-gen batchedGemm kernel drop for Kimi K3 MoE - #17190
[TRTLLM-14701][feat] Update trtllm-gen batchedGemm kernel drop for Kimi K3 MoE#17190rosong11 wants to merge 1 commit into
Conversation
8f0076f to
34b6691
Compare
rosong11
changed the title
|
No actionable comments were generated in the recent review. 🎉 WalkthroughChangesThe PR adds SiTu gated-activation support across MoE runners, generated GEMM metadata, and tests. It adds cooperative block routing for eligible small elementwise workloads. Batched GEMM contracts now carry multicast, rank, scale-factor, and MoE finalize data while removing Slice-K options. MoE kernel updates
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant MoEInput
participant RoutingCustom
participant BatchedGemmInterface
participant GeneratedKernel
MoEInput->>RoutingCustom: submit routing data
RoutingCustom->>RoutingCustom: select cooperative block path
RoutingCustom->>BatchedGemmInterface: provide token mapping and expert weights
BatchedGemmInterface->>GeneratedKernel: set rank, barriers, and MoE finalize parameters
GeneratedKernel->>MoEInput: write routed GEMM output
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 10
🧹 Nitpick comments (15)
cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustom.cu (2)
784-789: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winAdd a host-side eligibility check in
launchCoopBlockKernel.
routingIndicesCoopBlockKernelcompiles to an empty body whenSupportedis false.run()guarantees eligibility, butRoutingKernel.hnow exportslaunchCoopBlockKernelfor tests. If a test calls it with a non-elementwise preprocess or with a tier aboveCoopBlockKernelMaxNumExperts, the launch succeeds and writes nothing.mPtrPermutedIdxSizeandmPtrNumNonExitingCtasthen keep stale values, and the failure appears later as wrong GEMM shapes instead of a clear error. Add the same conditions thatrun()uses as aTLLM_CHECK_WITH_INFOhere.🛡️ Proposed guard
void launchCoopBlockKernel(Data const& data, uint32_t numThreadsHist, void* stream) { + TLLM_CHECK_WITH_INFO(data.mPtrScores != nullptr, "launchCoopBlockKernel requires raw scores input."); + TLLM_CHECK_WITH_INFO(data.mNumTokens <= BlockKernelMaxNumTokens, + "launchCoopBlockKernel supports at most %d tokens, got %d", BlockKernelMaxNumTokens, data.mNumTokens); + TLLM_CHECK_WITH_INFO(queryDispatchedMaxExperts(data) <= CoopBlockKernelMaxNumExperts, + "launchCoopBlockKernel supports at most %d experts per dispatched tier", CoopBlockKernelMaxNumExperts); + TLLM_CHECK_WITH_INFO(data.mPreprocessType == RoutingPreprocessType::None + || data.mPreprocessType == RoutingPreprocessType::Sigmoid + || data.mPreprocessType == RoutingPreprocessType::SigmoidBias, + "launchCoopBlockKernel requires an elementwise preprocess policy"); LAUNCH_ROUTING_CUSTOM(data, false, routingIndicesCoopBlockKernel, 1, numThreadsHist, /*smemSize=*/0, // No dynamic smem stream); }🤖 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` around lines 784 - 789, Update launchCoopBlockKernel to perform the same eligibility checks as run() before invoking LAUNCH_ROUTING_CUSTOM: require elementwise preprocessing and a tier no greater than CoopBlockKernelMaxNumExperts. Use TLLM_CHECK_WITH_INFO with clear failure messages, then preserve the existing kernel launch for eligible inputs.
1636-1646: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDerive the elementwise condition from the policy trait and use
common::getBoolEnv.
dispatchRoutingPolicyalready exposes the selected preprocess type. Use itsIsElementwisetrait instead of duplicating the runtime enum list. IncludeenvUtils.hand calltensorrt_llm::common::getBoolEnv("TLLM_ROUTING_DISABLE_COOP_BLOCK"); unlike the local check, it requires the complete value"1".🤖 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` around lines 1636 - 1646, Update the routing logic around dispatchRoutingPolicy to derive the elementwise condition from its selected preprocess type’s IsElementwise trait instead of manually checking RoutingPreprocessType values. Include envUtils.h and replace the local std::getenv-based disableCoopBlock initialization with tensorrt_llm::common::getBoolEnv("TLLM_ROUTING_DISABLE_COOP_BLOCK"), preserving the existing useCoopBlock gating.cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustomPolicy.cuh (1)
305-312: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument the two new required aliases in the ExpertSelectPolicy contract.
routingIndicesCoopBlockKernelnamesExpertSelect::PreprocessPolicyandExpertSelect::PostprocessPolicyunconditionally, before itsif constexpr (Supported)branch. EveryExpertSelectPolicyused withLAUNCH_ROUTING_CUSTOMmust therefore expose both aliases, including custom policies that bypass the preprocess/topK/postprocess pattern. The contract comment at lines 285-301 does not list them, so a future custom policy fails to compile with a confusing error inside the kernel. Add the two aliases to that contract list.🤖 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` around lines 305 - 312, Update the ExpertSelectPolicy contract comment near TopKExpertSelect to explicitly require PreprocessPolicy and PostprocessPolicy aliases. Ensure the documented contract covers every policy used with LAUNCH_ROUTING_CUSTOM, including custom policies that bypass the standard preprocessing and postprocessing flow.cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernel.h (1)
368-373: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftAdd a direct A/B unit test for both launchers. Current routing tests call
routingCustom::run()only, so they do not validate the bit-identical output claim. Keep the declarations because production code calls them fromRoutingCustom.cuandRoutingFromTopKIds.cu; otherwise revise the comments to remove the unit-test justification.🤖 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/RoutingKernel.h` around lines 368 - 373, Add a unit test that directly invokes both launchBlockKernel and launchCoopBlockKernel with identical input data and launch parameters, then compares their outputs for bit-identical results. Keep the existing declarations and comments because the new test validates their stated purpose, while production callers in RoutingCustom.cu and RoutingFromTopKIds.cu remain unchanged.cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/KernelRunner.cpp (1)
65-79: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse an anonymous namespace for the file-local helper.
Move
toGemmGatedActTypeinto an anonymous namespace. Do not usestaticfor internal linkage in new C++ code.As per coding guidelines, “Prefer anonymous namespaces over
staticfor internal-linkage functions.”Proposed change
-static batchedGemm::gemmGatedAct::ActType toGemmGatedActType(ActType actType) +namespace { +batchedGemm::gemmGatedAct::ActType toGemmGatedActType(ActType actType) { ... } +} // namespace🤖 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/batchedGemm/KernelRunner.cpp` around lines 65 - 79, Move the file-local helper toGemmGatedActType into an anonymous namespace in KernelRunner.cpp and remove its static qualifier, preserving its existing mapping and error behavior.Source: Coding guidelines
cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu (2)
404-404: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse lower-camel-case names and
constfor the new local variables.
is_gated_activationuses snake_case instead of lower camel case. The four activation flags are not modified after initialization.intermediateSizeFactoris also not modified.Suggested refactor
- bool is_gated_activation = tensorrt_llm::kernels::isGatedActType(actType); + bool const isGatedActivation = tensorrt_llm::kernels::isGatedActType(actType); - bool is_gated_activation = tensorrt_llm::kernels::isGatedActType(mActType); - int32_t intermediateSizeFactor = (is_gated_activation ? 2 : 1); + bool const isGatedActivation = tensorrt_llm::kernels::isGatedActType(mActType); + int32_t const intermediateSizeFactor = isGatedActivation ? 2 : 1;Apply the identifier rename at all corresponding uses.
As per coding guidelines, C++ local variables must use lower camel case, and unmodified variables must be declared
const.Also applies to: 476-476, 501-501, 511-511
🤖 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/runner.cu` at line 404, Update the new activation flag locals around isGatedActType and the related activation checks to use lower-camel-case names, including renaming is_gated_activation at every use. Declare all four activation flags and intermediateSizeFactor as const, preserving their existing values and behavior.Source: Coding guidelines
727-727: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winAdd a regression test with unequal valid dimensions.
Keep
validHiddenSize, validIntermediateSizeorder in future declarations and callers. Both parameters areint32_t, so a positional swap compiles and selects an incorrect configuration.🤖 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/runner.cu` at line 727, Preserve the parameter order validHiddenSize, validIntermediateSize in the relevant declaration and all callers, including the function around the numLocalExperts and numTokens parameters. Add a regression test using unequal valid dimensions so a positional swap is detected.cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelTraits.h (3)
188-189: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer
[[maybe_unused]]on the parameters.
mmaKandisSparseAare used in no branch ofgetNumSmemBitsPerElt. The neighbouringdtypeNeedsPaddingintrtllm/gen/DtypeDecl.hmarks the same two parameters with[[maybe_unused]]in its signature. That form documents the intent at the declaration and removes the need for the void casts inside the branch.♻️ Proposed refactor
-inline int getNumSmemBitsPerElt(tg::Dtype dtype, tg::MmaKind mmaKind, int mmaK, bool isSparseA) +inline int getNumSmemBitsPerElt( + tg::Dtype dtype, tg::MmaKind mmaKind, [[maybe_unused]] int mmaK, [[maybe_unused]] bool isSparseA) { if (mmaKind == tg::MmaKind::Auto) { throw std::runtime_error("mmaKind != tg::MmaKind::Auto"); } if (mmaKind == tg::MmaKind::MxFp8Fp6Fp4) { - (void) mmaK; - (void) isSparseA; return 8; }🤖 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/batchedGemm/trtllmGen_bmm_export/KernelTraits.h` around lines 188 - 189, Update getNumSmemBitsPerElt to mark the mmaK and isSparseA parameters with [[maybe_unused]] in its signature, matching dtypeNeedsPadding, and remove their corresponding (void) casts from the function body.
96-107: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the repeated name lookup into one helper.
hasChunkByNameis the third linear scan overmSmemChunkNamesin this class.getChunkOffsetByNameandgetFirstChunkReuseFlagByNamecontain the same loop.getSmemOffsetBiasnow runs two of these scans back to back for a single lookup.A private
findChunkIndexthat returns an index or-1would let all three public methods share one implementation.♻️ Proposed refactor
+ // Returns whether a chunk with the given name was allocated. + bool hasChunkByName(std::string const& name) const + { + return findChunkIndex(name) >= 0; + }Add the private helper next to the other private methods, then reuse it:
private: // Returns the index of the chunk with the given name, or -1 when it does not exist. int32_t findChunkIndex(std::string const& name) const { for (size_t ii = 0; ii < mSmemChunkNames.size(); ++ii) { if (mSmemChunkNames[ii] == name) { return static_cast<int32_t>(ii); } } return -1; }🤖 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/batchedGemm/trtllmGen_bmm_export/KernelTraits.h` around lines 96 - 107, Extract the shared linear scan into a private findChunkIndex helper returning the matching index or -1, placed with the other private methods. Update hasChunkByName, getChunkOffsetByName, and getFirstChunkReuseFlagByName to reuse this helper, and ensure getSmemOffsetBias no longer performs redundant name lookups.
435-448: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMark this hand-edit so the next kernel re-export does not drop it.
This file is a generated drop under
trtllmGen_bmm_export. The comment at lines 437-439 justifies the single-chunk layout with "so we don't perturb SMEM offsets of unrelated tests", andgetSmemOffsetBiasat lines 706-710 adds a fallback for the same reason.BatchedGemmInterface.hlines 958-963 already carries an explicitNOTE(TRT-LLM local fix over producer output ...)marker with an instruction to report the deviation upstream.Apply the same marker convention here. Without it, the next generator run silently reverts the per-resource bias chunks.
Also note that the guard
useDeepSeekFp8 && isBiasTypeMn(biasType)is correct only whilecheckAndUpdateGemmOptionsrestricts DeepSeek FP8 toBiasType::MnandBiasType::None. State that dependency in the comment.🤖 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/batchedGemm/trtllmGen_bmm_export/KernelTraits.h` around lines 435 - 448, The generated-output hand-edit in the per-resource bias chunk logic must be marked for preservation. Add the established NOTE(TRT-LLM local fix over producer output ...) marker near the existing rationale in the numBiasResources block, instructing future re-export updates to report the deviation upstream. Extend the comment to state that the useDeepSeekFp8 && isBiasTypeMn(biasType) guard depends on checkAndUpdateGemmOptions restricting DeepSeek FP8 to BiasType::Mn or BiasType::None.cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/GemmOptions.h (1)
1442-1451: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReport an error instead of returning
falsesilently.When
mDtypeSfCisVoidandupdateOptionsis false, the function returnsfalsewith no diagnostic.isValidConfig()inBatchedGemmInterface.hcalls this path, so a config that shipsmDtypeSfC = Voidwith a block-formatdtypeCis filtered out with no explanation. Every other failure in this function emits a message. TheGEMM_UPDATE_OR_ERRORmacro used at line 1467 already implements this update-or-error pattern.♻️ Proposed refactor
- if (options.mDtypeSfC == tg::Dtype::Void) - { - if (updateOptions) - { - options.mDtypeSfC = defaultDtypeSfC; - } - else - { - return false; - } - } + if (options.mDtypeSfC == tg::Dtype::Void) + { + GEMM_UPDATE_OR_ERROR(options.mDtypeSfC, defaultDtypeSfC); + }Confirm that the generated configs already carry a non-
VoidmDtypeSfC:#!/bin/bash # Description: Check whether generated kernel metadata sets mDtypeSfC for block-format dtypeC. set -euo pipefail fd -t f 'KernelMetaInfo.h' -p 'trtllmGen_bmm_export' | while IFS= read -r f; do echo "=== $f ===" rg -c 'mDtypeSfC' "$f" || echo "mDtypeSfC absent" rg -n -m3 -o 'mDtypeSfC=[^,]*' "$f" done # Definition of the update-or-error macro. rg -n -C5 'define GEMM_UPDATE_OR_ERROR'🤖 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/batchedGemm/trtllmGen_bmm_export/GemmOptions.h` around lines 1442 - 1451, Update the mDtypeSfC handling in the relevant validation function so the updateOptions == false branch reports an error through the existing GEMM_UPDATE_OR_ERROR pattern instead of returning false silently. Preserve assigning defaultDtypeSfC when updateOptions is true and ensure the failure includes diagnostic context consistent with the other validation errors.tests/unittest/_torch/thop/serial/test_moe.py (2)
2351-2352: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument why SiTu needs reduced input scales and relaxed tolerances.
Two SiTu-specific adjustments weaken the assertion strength, and neither states a reason.
- Lines 2351-2352 reduce
input_scalefrom 2.0 to 0.5 andweight_scalefrom 1.0 to 0.1 for SiTu only. That narrows the dynamic range reaching the activation.- Lines 2667-2669 raise
atolfrom 0.0 to 0.1 andrtolfrom 0.10 to 0.15 for SiTu.The SiTu formula saturates through
tanh, so smaller inputs keep the operating point in the linear region where the fused kernel and the reference agree most closely. State that reasoning in a comment, and record which value drove each threshold.Also applies to: 2664-2676
🤖 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/unittest/_torch/thop/serial/test_moe.py` around lines 2351 - 2352, Add comments near the SiTu-specific input_scale and weight_scale assignments explaining that tanh saturation motivates the reduced scales, which keep activation inputs in the linear region for fused/reference agreement; also comment the SiTu-specific atol and rtol values near their assertions, documenting the observed value or comparison that drove each threshold.
2685-2705: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftExtract a shared helper instead of calling one test from another.
test_mxe4m3_mxe2m1_situ_matches_referencecallstest_moe_mxe2m1_weightsdirectly. Three consequences follow.
test_moe_mxe2m1_weightscallspytest.skip(...)in several branches. A skip raised inside the callee marks the caller as skipped. A SiTu regression can therefore be reported as a skip.- The caller cannot use fixtures, and it is bound to the callee's positional signature.
- Pytest reports two tests where one behavior is exercised.
Move the body of
test_moe_mxe2m1_weightsinto a plain helper, for example_run_moe_mxe2m1_weights(...), and call the helper from both tests.Also confirm that
routing_infowithn_groups=None,top_k_groups=None, androuted_scaling=Noneis accepted end to end.routed_scalingis forwarded totorch.ops.trtllm.mxe4m3_mxe2m1_block_scale_moe_runnerat line 2626.#!/bin/bash # Description: Check how routing_info None fields are handled and whether other fixtures use None. set -euo pipefail fd -t f 'test_moe.py' -p 'tests/unittest/_torch/thop' | while IFS= read -r f; do echo "=== $f ===" rg -n -C3 'def are_groups_valid' "$f" rg -n -C6 '"routed_scaling"' "$f" | head -60 done # Schema of the runner op for routed_scaling. rg -nP -C6 'mxe4m3_mxe2m1_block_scale_moe_runner' -g '*.cpp' -g '*.cu' -g '*.py'🤖 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/unittest/_torch/thop/serial/test_moe.py` around lines 2685 - 2705, Extract the implementation of test_moe_mxe2m1_weights into a plain helper such as _run_moe_mxe2m1_weights, preserving its behavior without pytest skips, and have both tests invoke that helper so pytest reports them independently. Update test_moe_mxe2m1_weights and test_mxe4m3_mxe2m1_situ_matches_reference to use the helper with explicit parameters rather than calling one test from another. Verify the helper and runner path accept routing_info values where n_groups, top_k_groups, and routed_scaling are None, forwarding routed_scaling correctly to mxe4m3_mxe2m1_block_scale_moe_runner.cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/config.json (1)
1503-1514: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the redundant standalone
fusedActkey.This config sets
"fusedAct": trueat line 1501 and then setsfusedActagain inside the composite key"fusedAct,act,biasType"at line 1503. The generator applies one of the two, and the merge order is not stated in the file. Both values agree today, so behavior does not change, but the duplication invites a future conflict.The FP4 configs, for example the
FP4_FC1_LowLatencyblock at line 293, use only the composite key. The same duplication appears at lines 1625/1627, 1712/1722, and 1800/1805.🤖 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/batchedGemm/trtllmGen_bmm_export/config.json` around lines 1503 - 1514, Remove the redundant standalone fusedAct entries from the affected configuration blocks, retaining fusedAct only within the composite fusedAct,act,biasType keys. Apply this consistently to the blocks around FP4_FC1_LowLatency and the additional duplicated entries near the identified composite-key definitions, without changing the composite values.cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelParamsDecl.h (1)
211-229: 🗄️ Data Integrity & Integration | 🔵 Trivial | 💤 Low valueRename the new pointer members to use the
ptrprefix. UseptrPermutedIdxToExpandedIdxandptrExpertWeights, including host assignments.🤖 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/batchedGemm/trtllmGen_bmm_export/KernelParamsDecl.h` around lines 211 - 229, Rename the new KernelParamsDecl pointer members permutedIdxToExpandedIdx and expertWeightsPtr to ptrPermutedIdxToExpandedIdx and ptrExpertWeights, and update all corresponding host-side assignments and references to use the new names consistently.
🤖 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/batchedGemm/trtllmGen_bmm_export/BatchedGemmInterface.h`:
- Around line 588-590: Align topology inputs across validation and launch:
update the KernelParams setup to use the constructor’s mRankId, and change
checkAndUpdateBatchedGemmOptions calls to pass
data.mProblemDimensions.mWorldSize instead of a hardcoded world size of 1. Apply
these changes at the relevant constructor, validation, and launch sites while
preserving BatchedGemmData’s world-size value as the validation source.
In
`@cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/BatchedGemmOptions.h`:
- Around line 228-229: Audit every call to checkAndUpdateBatchedGemmOptions,
especially the caller in BatchedGemmInterface.h, and ensure no positional third
argument intended for updateOptions is rebound to worldSize. Make the signature
non-silent by moving worldSize after updateOptions or requiring both parameters
explicitly, then update all affected callers while preserving validation-only
calls with updateOptions set to false.
In
`@cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/config.json`:
- Line 939: Remove the trailing comma from the composite key near sfLayoutB and
apply the same correction to the sfLayoutA composite key near line 989. Confirm
both keys contain exactly six field names matching the six-value tuples, without
changing the associated configuration values.
In
`@cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/GemmOptions.h`:
- Around line 2050-2051: Update the DeepSeek FP8 validation predicate in the
mUseDeepSeekFp8 check to accept both BiasType::None and BiasType::Mn. Preserve
rejection of all other bias types so plain DeepSeek FP8 and bias-enabled
DeepSeek FP8 configurations remain supported while invalid combinations are
still rejected.
- Around line 284-287: Update the comment above mDtypeSfC to list both Fp32 and
Bfloat16 as supported override types, matching the validation accepted by the
GemmOptions logic.
In
`@cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelParamsDecl.h`:
- Around line 584-591: Add default member initializers of {nullptr} to
ptrMulticastCompletionBarUc and ptrMulticastCompletionBarMc in KernelParams,
matching neighboring pointer members such as ptrRowMaxCompletionBars and
ptrDynamicTileCounter. Preserve the existing declarations and comments.
In
`@cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustomPolicy.cuh`:
- Around line 394-395: Add a Tier<1024, 32> entry to the Renormalize policy tier
list, placing it before the existing Tier<2048, 32> entry so configurations such
as 896 experts with topK 16 select the cooperative kernel instead of the
spilling classic kernel; keep the tier within the CoopBlockKernelMaxNumExperts
limit.
In `@cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu`:
- Around line 404-408: Update the DeepSeek FP8 validation near
is_gated_activation and the corresponding Runner::run dispatch so every
activation reaching moe::dev::activation::run is supported by its SwiGlu-only
implementation; reject Relu2 and Silu (and any other non-SwiGlu activation
accepted by getOptions), or route them through matching implementations while
preserving supported SwiGlu behavior.
In `@tests/unittest/_torch/thop/serial/test_moe.py`:
- Around line 49-53: Update test_act_type_enum_values_stable by adding a comment
that identifies
cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/KernelRunner.h as the C++
source of truth for the ActType numeric values consumed by act_type.value,
clarifying that the test protects Python/C++ enum alignment.
- Around line 2573-2576: Guard the two bias scaling statements in the
dtype_activation == "fp8" branch so they only multiply gemm1_bias_shuffled and
gemm2_bias_shuffled by their global scales when the corresponding value is not
None; preserve None for bias-free cases.
---
Nitpick comments:
In `@cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/KernelRunner.cpp`:
- Around line 65-79: Move the file-local helper toGemmGatedActType into an
anonymous namespace in KernelRunner.cpp and remove its static qualifier,
preserving its existing mapping and error behavior.
In
`@cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/config.json`:
- Around line 1503-1514: Remove the redundant standalone fusedAct entries from
the affected configuration blocks, retaining fusedAct only within the composite
fusedAct,act,biasType keys. Apply this consistently to the blocks around
FP4_FC1_LowLatency and the additional duplicated entries near the identified
composite-key definitions, without changing the composite values.
In
`@cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/GemmOptions.h`:
- Around line 1442-1451: Update the mDtypeSfC handling in the relevant
validation function so the updateOptions == false branch reports an error
through the existing GEMM_UPDATE_OR_ERROR pattern instead of returning false
silently. Preserve assigning defaultDtypeSfC when updateOptions is true and
ensure the failure includes diagnostic context consistent with the other
validation errors.
In
`@cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelParamsDecl.h`:
- Around line 211-229: Rename the new KernelParamsDecl pointer members
permutedIdxToExpandedIdx and expertWeightsPtr to ptrPermutedIdxToExpandedIdx and
ptrExpertWeights, and update all corresponding host-side assignments and
references to use the new names consistently.
In
`@cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelTraits.h`:
- Around line 188-189: Update getNumSmemBitsPerElt to mark the mmaK and
isSparseA parameters with [[maybe_unused]] in its signature, matching
dtypeNeedsPadding, and remove their corresponding (void) casts from the function
body.
- Around line 96-107: Extract the shared linear scan into a private
findChunkIndex helper returning the matching index or -1, placed with the other
private methods. Update hasChunkByName, getChunkOffsetByName, and
getFirstChunkReuseFlagByName to reuse this helper, and ensure getSmemOffsetBias
no longer performs redundant name lookups.
- Around line 435-448: The generated-output hand-edit in the per-resource bias
chunk logic must be marked for preservation. Add the established NOTE(TRT-LLM
local fix over producer output ...) marker near the existing rationale in the
numBiasResources block, instructing future re-export updates to report the
deviation upstream. Extend the comment to state that the useDeepSeekFp8 &&
isBiasTypeMn(biasType) guard depends on checkAndUpdateGemmOptions restricting
DeepSeek FP8 to BiasType::Mn or BiasType::None.
In
`@cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustom.cu`:
- Around line 784-789: Update launchCoopBlockKernel to perform the same
eligibility checks as run() before invoking LAUNCH_ROUTING_CUSTOM: require
elementwise preprocessing and a tier no greater than
CoopBlockKernelMaxNumExperts. Use TLLM_CHECK_WITH_INFO with clear failure
messages, then preserve the existing kernel launch for eligible inputs.
- Around line 1636-1646: Update the routing logic around dispatchRoutingPolicy
to derive the elementwise condition from its selected preprocess type’s
IsElementwise trait instead of manually checking RoutingPreprocessType values.
Include envUtils.h and replace the local std::getenv-based disableCoopBlock
initialization with
tensorrt_llm::common::getBoolEnv("TLLM_ROUTING_DISABLE_COOP_BLOCK"), preserving
the existing useCoopBlock gating.
In
`@cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustomPolicy.cuh`:
- Around line 305-312: Update the ExpertSelectPolicy contract comment near
TopKExpertSelect to explicitly require PreprocessPolicy and PostprocessPolicy
aliases. Ensure the documented contract covers every policy used with
LAUNCH_ROUTING_CUSTOM, including custom policies that bypass the standard
preprocessing and postprocessing flow.
In
`@cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernel.h`:
- Around line 368-373: Add a unit test that directly invokes both
launchBlockKernel and launchCoopBlockKernel with identical input data and launch
parameters, then compares their outputs for bit-identical results. Keep the
existing declarations and comments because the new test validates their stated
purpose, while production callers in RoutingCustom.cu and RoutingFromTopKIds.cu
remain unchanged.
In `@cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu`:
- Line 404: Update the new activation flag locals around isGatedActType and the
related activation checks to use lower-camel-case names, including renaming
is_gated_activation at every use. Declare all four activation flags and
intermediateSizeFactor as const, preserving their existing values and behavior.
- Line 727: Preserve the parameter order validHiddenSize, validIntermediateSize
in the relevant declaration and all callers, including the function around the
numLocalExperts and numTokens parameters. Add a regression test using unequal
valid dimensions so a positional swap is detected.
In `@tests/unittest/_torch/thop/serial/test_moe.py`:
- Around line 2351-2352: Add comments near the SiTu-specific input_scale and
weight_scale assignments explaining that tanh saturation motivates the reduced
scales, which keep activation inputs in the linear region for fused/reference
agreement; also comment the SiTu-specific atol and rtol values near their
assertions, documenting the observed value or comparison that drove each
threshold.
- Around line 2685-2705: Extract the implementation of test_moe_mxe2m1_weights
into a plain helper such as _run_moe_mxe2m1_weights, preserving its behavior
without pytest skips, and have both tests invoke that helper so pytest reports
them independently. Update test_moe_mxe2m1_weights and
test_mxe4m3_mxe2m1_situ_matches_reference to use the helper with explicit
parameters rather than calling one test from another. Verify the helper and
runner path accept routing_info values where n_groups, top_k_groups, and
routed_scaling are None, forwarding routed_scaling correctly to
mxe4m3_mxe2m1_block_scale_moe_runner.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
| BatchedGemmInterface(int32_t rankId = 0, bool const exportsCubin = false, int32_t const numRotations = 1) | ||
| : mRankId(rankId) | ||
| , mExportsCubin(exportsCubin) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Use one multi-GPU topology source for validation and launch.
mRankId is stored by the constructor but never reaches KernelParams. The launch instead uses BatchedGemmData::mRank. Validation always uses worldSize == 1, while launch uses BatchedGemmData::mWorldSize.
For a multi-GPU caller, validation can select a configuration for a different topology than the kernel receives. Use mRankId when setting kernel parameters, and pass data.mProblemDimensions.mWorldSize to checkAndUpdateBatchedGemmOptions. Remove mRankId instead if BatchedGemmData is the intended single source of rank.
Proposed change
- batchedGemmData.mProblemDimensions.mRank, batchedGemmData.mProblemDimensions.mWorldSize,
+ mRankId, batchedGemmData.mProblemDimensions.mWorldSize,
- /* worldSize */ 1, /* updateOptions */ false);
+ /* worldSize */ data.mProblemDimensions.mWorldSize, /* updateOptions */ false);Also applies to: 686-691, 961-968, 1055-1056
🤖 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/batchedGemm/trtllmGen_bmm_export/BatchedGemmInterface.h`
around lines 588 - 590, Align topology inputs across validation and launch:
update the KernelParams setup to use the constructor’s mRankId, and change
checkAndUpdateBatchedGemmOptions calls to pass
data.mProblemDimensions.mWorldSize instead of a hardcoded world size of 1. Apply
these changes at the relevant constructor, validation, and launch sites while
preserving BatchedGemmData’s world-size value as the validation source.
| inline bool checkAndUpdateBatchedGemmOptions( | ||
| BatchedGemmOptions& options, tg::CudaArch cudaArch, bool updateOptions = true) | ||
| BatchedGemmOptions& options, tg::CudaArch cudaArch, int worldSize = 1, bool updateOptions = true) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Inserting worldSize before updateOptions silently rebinds positional callers.
worldSize was added as the third parameter, ahead of updateOptions. Both parameters have defaults and both accept an implicit conversion from bool. A caller written as checkAndUpdateBatchedGemmOptions(options, sm, false) still compiles. It now binds false to worldSize (0) and takes the default updateOptions = true, which mutates the caller's options instead of validating them.
BatchedGemmInterface.h lines 955-969 already carries a local fix and a NOTE describing this exact failure. Confirm that no other caller is still positional.
Consider making the hazard non-silent by placing worldSize last, or by requiring both arguments explicitly.
#!/bin/bash
# Description: Find every call site of checkAndUpdateBatchedGemmOptions and check argument counts.
set -euo pipefail
rg -nPU --type=cpp --type=cuda -C6 'checkAndUpdateBatchedGemmOptions\s*\(' || true
rg -nPU -g '*.h' -g '*.hpp' -g '*.cu' -g '*.cuh' -g '*.cpp' -C6 'checkAndUpdateBatchedGemmOptions\s*\('
# Also check the gated-act and gemm variants for the same signature drift.
rg -nP -C3 'inline bool checkAndUpdate(Gemm|GemmGatedAct|BatchedGemm)Options'
rg -nP -C4 'checkAndUpdateGemmGatedActOptions\s*\(|checkAndUpdateGemmOptions\s*\('🤖 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/batchedGemm/trtllmGen_bmm_export/BatchedGemmOptions.h`
around lines 228 - 229, Audit every call to checkAndUpdateBatchedGemmOptions,
especially the caller in BatchedGemmInterface.h, and ensure no positional third
argument intended for updateOptions is rebound to worldSize. Make the signature
non-silent by moving worldSize after updateOptions or requiring both parameters
explicitly, then update all affected callers while preserving validation-only
calls with updateOptions set to false.
| // A barrier that is used to synchronize kernels running on multiple GPUs. | ||
| // UC is a unicast pointer in local GPU memory. | ||
| // MC is a multicast pointer to all other GPUs' memory. | ||
| // | ||
| // The value must be initialized to zero. The barrier is self-resetting. It | ||
| // does not have to be reset by another kernel after the grid has ended. | ||
| uint32_t* ptrMulticastCompletionBarUc; | ||
| uint32_t* ptrMulticastCompletionBarMc; |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Add default member initializers to the two barrier pointers.
ptrMulticastCompletionBarUc and ptrMulticastCompletionBarMc have no default initializer. Every other pointer member in KernelParams uses {nullptr}, including the neighbouring ptrRowMaxCompletionBars{nullptr} and ptrDynamicTileCounter{nullptr}. KernelParams is filled field by field by host setup code. If a setup path does not assign these two members, the kernel receives indeterminate pointers. The comment states the barrier "must be initialized to zero", which makes an unset pointer harder to detect.
As per coding guidelines, "prefer default member initializers and designated initializers for structs".
🛡️ Proposed fix
- uint32_t* ptrMulticastCompletionBarUc;
- uint32_t* ptrMulticastCompletionBarMc;
+ uint32_t* ptrMulticastCompletionBarUc{nullptr};
+ uint32_t* ptrMulticastCompletionBarMc{nullptr};Run the following script to check whether any setup path leaves these members unset:
#!/bin/bash
# Description: Locate KernelParams construction and check assignment of the new barrier pointers.
set -euo pipefail
rg -nP -C4 'ptrMulticastCompletionBar(Uc|Mc)'
rg -nP -C6 'KernelParams\s+\w+\s*(\{|;|=)' -g '*.h' -g '*.cuh' -g '*.cu' | head -100
rg -nP -C4 'setKernelParams|KernelParams::setKernelParams'🤖 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/batchedGemm/trtllmGen_bmm_export/KernelParamsDecl.h`
around lines 584 - 591, Add default member initializers of {nullptr} to
ptrMulticastCompletionBarUc and ptrMulticastCompletionBarMc in KernelParams,
matching neighboring pointer members such as ptrRowMaxCompletionBars and
ptrDynamicTileCounter. Preserve the existing declarations and comments.
Source: Coding guidelines
| // Cooperative block kernel: one thread per expert, so at most 1024 experts (1 CUDA block). | ||
| static constexpr int CoopBlockKernelMaxNumExperts = 1024; |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Determine which preprocess/postprocess policy large-expert models select at runtime.
set -euo pipefail
rg -nP --type=cpp -C 6 'RoutingPreprocessType::|RoutingPostprocessType::' \
cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu | head -80
rg -nP --type=cpp -C 3 'RoutingMethodType' \
cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernel.h | head -40Repository: NVIDIA/TensorRT-LLM
Length of output: 4301
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- cooperative-path selection and tier definitions ---'
rg -n -C 8 'useCoopBlock|queryDispatchedMaxExperts|Tier<576|Tier<1024|Tier<2048|CoopBlockKernelMaxNumExperts|DynBlockKernelMaxNumExperts' \
cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustomPolicy.cuh
printf '%s\n' '--- routing policy dispatch ---'
rg -n -C 10 'RoutingPreprocessType::None|RoutingPostprocessType::Softmax|NoOp|SoftmaxPostprocess|SigmoidBias|ScaledSumNormalize' \
cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routingRepository: NVIDIA/TensorRT-LLM
Length of output: 50375
Add a Tier<1024, 32> entry to the Renormalize policy
runner.cu selects NoOpPreprocess with SoftmaxPostprocess for Renormalize. For 896 experts and topK 16, the current tier list selects Tier<2048, 32>, so useCoopBlock skips the cooperative kernel and uses the spilling classic kernel.
🤖 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`
around lines 394 - 395, Add a Tier<1024, 32> entry to the Renormalize policy
tier list, placing it before the existing Tier<2048, 32> entry so configurations
such as 896 experts with topK 16 select the cooperative kernel instead of the
spilling classic kernel; keep the tier within the CoopBlockKernelMaxNumExperts
limit.
| bool is_gated_activation = tensorrt_llm::kernels::isGatedActType(actType); | ||
| // DeepSeek FP8 runs the gated activation as a standalone kernel (see moe::dev::activation), | ||
| // which only implements SwiGlu math; SiTu is available exclusively through fused FC1 cubins. | ||
| TLLM_CHECK_WITH_INFO(!(useDeepSeekFp8 && actType == ActType::SiTu), | ||
| "SiTu activation is not supported with DeepSeek FP8 (no standalone SiTu activation kernel)."); |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Reject every unsupported activation type in the DeepSeek FP8 standalone path.
The new check rejects only ActType::SiTu. getOptions also accepts ActType::Relu2 and ActType::Silu. Runner::run later invokes moe::dev::activation::run for every E4m3 DeepSeek FP8 run, while the comment states that this standalone kernel implements only SwiGlu math.
Reject every activation that the standalone kernel does not support, or route non-SwiGlu activations through a matching implementation.
Suggested validation
- TLLM_CHECK_WITH_INFO(!(useDeepSeekFp8 && actType == ActType::SiTu),
- "SiTu activation is not supported with DeepSeek FP8 (no standalone SiTu activation kernel).");
+ TLLM_CHECK_WITH_INFO(!useDeepSeekFp8 || actType == ActType::SwiGlu,
+ "Only SwiGlu activation is supported with DeepSeek FP8.");Based on the supplied non-gated options branch and DeepSeek FP8 activation dispatch.
Also applies to: 673-673
🤖 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/runner.cu` around
lines 404 - 408, Update the DeepSeek FP8 validation near is_gated_activation and
the corresponding Runner::run dispatch so every activation reaching
moe::dev::activation::run is supported by its SwiGlu-only implementation; reject
Relu2 and Silu (and any other non-SwiGlu activation accepted by getOptions), or
route them through matching implementations while preserving supported SwiGlu
behavior.
| def test_act_type_enum_values_stable(): | ||
| assert ActType.SwiGlu.value == 0 | ||
| assert ActType.Relu2.value == 1 | ||
| assert ActType.Silu.value == 2 | ||
| assert ActType.SiTu.value == 3 |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
This test does not detect C++ and Python divergence.
test_act_type_enum_values_stable asserts the values of the local Python ActType against literals in the same file. It fails only if someone edits both the enum and the assertions inconsistently. It does not detect a change to the C++ ActType in cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/KernelRunner.h, which is the value actually consumed by act_type.value at line 2627.
Add a comment naming the C++ header as the source of truth, so a future reader knows where the real contract lives.
🤖 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/unittest/_torch/thop/serial/test_moe.py` around lines 49 - 53, Update
test_act_type_enum_values_stable by adding a comment that identifies
cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/KernelRunner.h as the C++
source of truth for the ActType numeric values consumed by act_type.value,
clarifying that the test protects Python/C++ enum alignment.
There was a problem hiding this comment.
♻️ Duplicate comments (1)
cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelParamsDecl.h (1)
584-591: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winAdd default member initializers to the two barrier pointers (unresolved from prior review).
ptrMulticastCompletionBarUcandptrMulticastCompletionBarMcstill have no default initializer, while every neighboring pointer member (ptrRowMaxCompletionBars,ptrDynamicTileCounter) uses{nullptr}. If a host setup path does not explicitly assign these two members, the kernel receives an indeterminate pointer, and the comment's requirement that "the value must be initialized to zero" becomes harder to enforce.🛡️ Proposed fix
- uint32_t* ptrMulticastCompletionBarUc; - uint32_t* ptrMulticastCompletionBarMc; + uint32_t* ptrMulticastCompletionBarUc{nullptr}; + uint32_t* ptrMulticastCompletionBarMc{nullptr};🤖 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/batchedGemm/trtllmGen_bmm_export/KernelParamsDecl.h` around lines 584 - 591, Initialize both ptrMulticastCompletionBarUc and ptrMulticastCompletionBarMc with the same nullptr default member initializer used by neighboring pointer members, while preserving their existing types and declarations.
🧹 Nitpick comments (2)
cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelMetaInfo.h (1)
28-29: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winVerify the kernel drop is not built from a dirty tree.
TLLM_GEN_COMMITis set to"a2ad0544-dirty". The-dirtysuffix means the generator ran against a working tree with uncommitted changes. Checked-in generated cubins should trace to a clean, reproducible commit. Regenerate this header from a clean commit before merge, or confirm the dirty state is expected and intentional.🤖 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/batchedGemm/trtllmGen_bmm_export/KernelMetaInfo.h` around lines 28 - 29, Update TLLM_GEN_COMMIT in KernelMetaInfo.h to reference the clean generator commit without the “-dirty” suffix. Regenerate the checked-in kernel metadata from that clean commit, or explicitly validate and retain the suffix only if the dirty build is intentional.cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelParamsDecl.h (1)
217-230: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winNew fields break the established
ptrnaming convention.Every other pointer member in
KernelParamsuses aptrprefix:ptrA,ptrB,ptrScaleC,ptrRouteMap,ptrRowMaxCompletionBars,ptrDynamicTileCounter, and so on. The two new fields deviate from this pattern:permutedIdxToExpandedIdxhas no prefix at all, andexpertWeightsPtruses a suffix instead of a prefix. Rename both fields to match the local convention.♻️ Proposed rename for naming consistency
- int32_t const* permutedIdxToExpandedIdx{nullptr}; + int32_t const* ptrPermutedIdxToExpandedIdx{nullptr}; ... - void const* expertWeightsPtr{nullptr}; + void const* ptrExpertWeights{nullptr};🤖 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/batchedGemm/trtllmGen_bmm_export/KernelParamsDecl.h` around lines 217 - 230, Rename the new KernelParams pointer members permutedIdxToExpandedIdx and expertWeightsPtr to use the established ptr prefix convention, such as ptrPermutedIdxToExpandedIdx and ptrExpertWeights. Update all references to both fields throughout the affected code while preserving their types and behavior.
🤖 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.
Duplicate comments:
In
`@cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelParamsDecl.h`:
- Around line 584-591: Initialize both ptrMulticastCompletionBarUc and
ptrMulticastCompletionBarMc with the same nullptr default member initializer
used by neighboring pointer members, while preserving their existing types and
declarations.
---
Nitpick comments:
In
`@cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelMetaInfo.h`:
- Around line 28-29: Update TLLM_GEN_COMMIT in KernelMetaInfo.h to reference the
clean generator commit without the “-dirty” suffix. Regenerate the checked-in
kernel metadata from that clean commit, or explicitly validate and retain the
suffix only if the dirty build is intentional.
In
`@cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelParamsDecl.h`:
- Around line 217-230: Rename the new KernelParams pointer members
permutedIdxToExpandedIdx and expertWeightsPtr to use the established ptr prefix
convention, such as ptrPermutedIdxToExpandedIdx and ptrExpertWeights. Update all
references to both fields throughout the affected code while preserving their
types and behavior.
|
/bot run --disable-fail-fast |
|
PR_Github #63394 [ run ] triggered by Bot. Commit: |
|
PR_Github #63394 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63488 [ run ] triggered by Bot. Commit: |
|
PR_Github #63488 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63629 [ run ] triggered by Bot. Commit: |
|
PR_Github #63629 [ run ] completed with state
|
Signed-off-by: rosong11 <rosong@nvidia.com>
34b6691 to
f7b7244
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #63662 [ run ] triggered by Bot. Commit: |
3 participants
Dev Engineer Review
SiTu, and rejected unsupported activations.QA Engineer Review
test_mxe4m3_mxe2m1_situ_matches_reference(num_tokens).ActType.SiTu, including scaling, optional bias, validation, and accuracy handling.Description
Test Coverage
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.