[https://nvbugs/6098442][fix] Add fix for IMA with TRTLLM-Gen GmemReductionWithSeparateKernel#13541
Conversation
Signed-off-by: Pengbo Wang <221450789+pengbowang-nv@users.noreply.github.com>
|
/bot run --disable-fail-fast |
📝 WalkthroughWalkthroughThe hotfix that prevents crashes from null Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h (1)
312-320:⚠️ Potential issue | 🔴 CriticalNull-pointer undefined behavior still occurs in
setFmhaData()despite mode guard.The guard at lines 307–315 disables
MultiCtasKvModewhen pointers are null, butsetFmhaData()(called at line 338) unconditionally performs pointer arithmetic on potentially nullparams.multiCtasKvScratchPtrat lines 712–715:
- Line 712:
reinterpret_cast<float2*>(params.multiCtasKvScratchPtr)on null pointer- Line 715: pointer arithmetic on the result
Both operations constitute undefined behavior under the C++ standard. The mode being disabled prevents kernel execution but does not prevent the UB in
setFmhaData().Conditionally initialize pointers to
nullptrand only set them when GmemReduction modes are enabled with valid pointers, usingTLLM_CHECK_WITH_INFOto enforce the invariant.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h` around lines 312 - 320, The code currently casts and does pointer arithmetic on params.multiCtasKvScratchPtr inside setFmhaData(), causing undefined behavior when that pointer is null even if options.mMultiCtasKvMode was later forced to Disabled; to fix, change setFmhaData() so it initializes any local multi-cta pointer variables to nullptr and only perform reinterpret_cast<float2*>(...) and subsequent pointer arithmetic when options.mMultiCtasKvMode is one of GmemReduction or GmemReductionWithSeparateKernel AND params.multiCtasKvScratchPtr (and multiCtasKvCounterPtr if used) are non-null, and add TLLM_CHECK_WITH_INFO guards to assert those invariants (referencing options.mMultiCtasKvMode, params.multiCtasKvScratchPtr, params.multiCtasKvCounterPtr, and setFmhaData) so the function never performs pointer arithmetic on a null pointer.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Outside diff comments:
In `@cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h`:
- Around line 312-320: The code currently casts and does pointer arithmetic on
params.multiCtasKvScratchPtr inside setFmhaData(), causing undefined behavior
when that pointer is null even if options.mMultiCtasKvMode was later forced to
Disabled; to fix, change setFmhaData() so it initializes any local multi-cta
pointer variables to nullptr and only perform reinterpret_cast<float2*>(...) and
subsequent pointer arithmetic when options.mMultiCtasKvMode is one of
GmemReduction or GmemReductionWithSeparateKernel AND
params.multiCtasKvScratchPtr (and multiCtasKvCounterPtr if used) are non-null,
and add TLLM_CHECK_WITH_INFO guards to assert those invariants (referencing
options.mMultiCtasKvMode, params.multiCtasKvScratchPtr,
params.multiCtasKvCounterPtr, and setFmhaData) so the function never performs
pointer arithmetic on a null pointer.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: fd7863a8-3222-4f98-a302-8b7c88f4c6df
📒 Files selected for processing (1)
cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h
|
PR_Github #45838 [ run ] triggered by Bot. Commit: |
|
PR_Github #45838 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #46272 [ run ] triggered by Bot. Commit: |
|
PR_Github #46272 [ run ] completed with state |
…uctionWithSeparateKernel (NVIDIA#13541) Signed-off-by: Pengbo Wang <221450789+pengbowang-nv@users.noreply.github.com>
4 participants
Summary by CodeRabbit
Bug Fixes
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)
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.