[#12230][fix] Add bounds checking in autotuner _find_nearest_profile for SM121 - #12310
Conversation
📝 WalkthroughWalkthroughDefensive bounds checks added to Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
📝 Coding Plan
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: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/autotuner.py (1)
1406-1435:⚠️ Potential issue | 🟡 MinorFormatting check is failing in this hunk (YAPF).
Release Checks report this region was reformatted by the hook. Please run and commit formatter output for this file before merge.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/autotuner.py` around lines 1406 - 1435, The formatting in tensorrt_llm/_torch/autotuner.py (around the loop handling DynamicTensorSpec and constraint_specs, i.e., symbols like DynamicTensorSpec, apply_map_to_tuning_buckets, tune_max_num_tokens, and constraint_specs) was changed by the pre-commit hook (YAPF); run the repository's YAPF formatter (or the configured formatter) on this file, review the resulting whitespace/line-wrapping for the shown hunk, stage and commit the formatted file so the reformatting is included in the PR.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tensorrt_llm/_torch/autotuner.py`:
- Around line 1398-1410: The _optimization_profiles function still indexes specs
without bounds checks and can raise IndexError when specs are malformed; update
_optimization_profiles to perform the same input/dimension guards used
elsewhere: before any access to base_profile[spec.input_idx] or
base_profile[spec.input_idx][spec.dim_idx] (and similarly for
spec.dim_idx+offset), verify spec.input_idx < len(base_profile) and spec.dim_idx
< len(base_profile[spec.input_idx]) (and handle negative indices if relevant),
log a debug message referencing DynamicTensorSpec/spec.input_idx/spec.dim_idx
and skip that spec (continue) when out of range, ensuring all spots where
base_profile is indexed (around the existing accesses you noted) are protected.
---
Outside diff comments:
In `@tensorrt_llm/_torch/autotuner.py`:
- Around line 1406-1435: The formatting in tensorrt_llm/_torch/autotuner.py
(around the loop handling DynamicTensorSpec and constraint_specs, i.e., symbols
like DynamicTensorSpec, apply_map_to_tuning_buckets, tune_max_num_tokens, and
constraint_specs) was changed by the pre-commit hook (YAPF); run the
repository's YAPF formatter (or the configured formatter) on this file, review
the resulting whitespace/line-wrapping for the shown hunk, stage and commit the
formatted file so the reformatting is included in the PR.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 17232aea-4e81-47a4-b8c0-3d24a43abe55
📒 Files selected for processing (1)
tensorrt_llm/_torch/autotuner.py
svc-trtllm-gh-bot
added
the
Community want to contribute
f5af0fe to
775662c
Compare
|
Hi @pamelap-nvidia , Could you take a look at this and see if it can fix #12230? Also @hyukn do you think we need such guard for autotuner? Thanks! |
775662c to
a894178
Compare
|
@pengbowang-nv Yes — this fixes #12230. The bounds checks cover both Tested on DGX Spark SM121 with Unit tests for the bounds checking are in commit 2 ( Will also test with |
…ofile for SM121 On DGX Spark (SM121), _find_nearest_profile() crashes with IndexError when spec.input_idx or spec.dim_idx exceeds the base_profile dimensions. This happens because some ops produce fewer or differently-shaped tensors than the specs expect. Add bounds checks before indexing into base_profile for both DynamicTensorSpec and ConstraintSpec loops. Out-of-bounds specs are skipped with a debug log message. All in-bounds behavior is unchanged. Tested with 8 cases covering normal operation, out-of-bounds indices, mixed valid/invalid specs, and edge cases. Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
…M121 Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
a894178 to
6d9c094
Compare
|
Follow-up on That said, the autotuner IndexError (Bug 1 in #12230) is hit during model initialization before the Triton kernel executes — so this PR would let users get past the autotuner crash. The stack trace in #12230 ( Validated with |
|
While running The failing kernels are TRT-LLM's CUTLASS TMA warp-specialized MoE grouped GEMM kernels compiled for
Root causeSM121 (DGX Spark / GB10) has less shared memory than SM120 (RTX 5090). These larger tile configurations (128x256x128B, 256x128x128B) exceed SM121's ~99KB SMEM limit for grouped GEMM TMA initialization. The smaller tile ImpactNon-blocking. The FlashInfer autotuner correctly skips the failing tactics and selects working ones. The model loads and serves successfully. These are warnings, not crashes. cc @eugr log |
|
/bot run |
|
PR_Github #40638 [ run ] triggered by Bot. Commit: |
|
PR_Github #40638 [ run ] completed with state |
|
@pengbowang-nv I think this is good to go. I see you tagged @pamelap-nvidia and @hyukn , any chance you all can take a look? happy to address any feedback. |
…om test docstring - Change all 8 bounds-check log messages from debug to warning level so users are notified when DynamicTensorSpec/ConstraintSpec entries are skipped due to out-of-bounds input_idx or dim_idx. - Remove SM121-specific reference from test docstring since the bounds checking is not architecture-specific. Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
|
@hyukn @pamelap-nvidia I fixed the suggestions and a couple of more small things. hopefully we're good this time 🤞 |
|
PR_Github #51158 [ run ] completed with state |
@hyukn @yizhang-nv this now has a clean CI pass. any objections to merging it? |
…ecks Per @hyukn review on PR NVIDIA#12310: deduplicate the input_idx / dim_idx bounds-check pattern between AutoTuner._optimization_profiles and AutoTuner._find_nearest_profile. Both sites now call a single helper. No behavioral change. Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
Per @hyukn review on PR NVIDIA#12310: collapse the parametrically-similar OOB cases into a single parametrized test, keeping the valid-spec case separate for readability. Same behavioral assertions, fewer test functions. Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
Completes the symmetric matrix of (spec type) x (oob input, oob dim, negative input, negative dim). One additional parametrize case in test_oob_spec_skipped; no other changes. Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
|
Could you please address the pre-commit failures, @mihai-chiorean ? |
Pre-commit yapf reformatted the new pytest.param entries and _find_nearest_profile calls into compact form. No behavioral change. Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
Oops. Done! |
|
/bot run --disable-fail-fast |
|
PR_Github #53262 [ run ] triggered by Bot. Commit: |
|
PR_Github #53262 [ run ] completed with state
|
hyukn
left a comment
hyukn
left a comment
There was a problem hiding this comment.
LGTM, only some minor changes required. Thanks for the effort!
Per @hyukn review: drop the standalone spec_kind: str parameter from _spec_in_bounds and derive the class name from the spec itself. Call sites are no longer required to pass the literal class name string. Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
…ntry points Per @hyukn review: - Drop the redundant 'from PR NVIDIA#12310' header comment. - Collapse TestFindNearestProfileBounds and TestOptimizationProfilesBounds into one TestSpecBoundsChecking class. - Single parametrized test_oob_spec_skipped covers both _find_nearest_profile and _optimization_profiles paths (entry x spec_config matrix). - Keep test_valid_specs_unaffected as a sanity check that valid specs are NOT skipped. Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
…ntry points Per @hyukn review: - Drop the redundant 'from PR NVIDIA#12310' header comment. - Collapse TestFindNearestProfileBounds and TestOptimizationProfilesBounds into one TestSpecBoundsChecking class. - Single parametrized test_oob_spec_skipped covers both _find_nearest_profile and _optimization_profiles paths (entry x spec_config matrix). - Keep test_valid_specs_unaffected as a sanity check that valid specs are NOT skipped. Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
a83d2f3 to
2711af0
Compare
Honor the 'compress to 1 test' request literally. TestSpecBoundsChecking now has a single parametrized test_oob_spec_skipped covering both _find_nearest_profile and _optimization_profiles. Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
done |
|
/bot run --disable-fail-fast |
|
PR_Github #53420 [ run ] triggered by Bot. Commit: |
|
PR_Github #53420 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #53534 [ run ] triggered by Bot. Commit: |
|
PR_Github #53534 [ run ] completed with state |
|
thanks @hyukn @pamelap-nvidia @karljang !! really appreciate your patience and reviews! looks like this is now ready. |
|
Yeah, finally! Thank you so much for your contribution! |
8 participants
Summary
Fixes #12230 (autotuner crash)
On DGX Spark (SM121),
_find_nearest_profile()crashes withIndexError: list assignment index out of rangeduring warmup, preventing the server from starting. This happens becauseDynamicTensorSpec.input_idxordim_idxexceeds thebase_profiledimensions when ops produce fewer or differently-shaped tensors than the specs expect.Adds bounds checks before indexing into
base_profilefor bothDynamicTensorSpecandConstraintSpecloops. Out-of-bounds specs are skipped with a debug log. All in-bounds behavior is unchanged.Test Coverage
Tested with 8 cases inside the Docker container:
input_idxout of range for DynamicTensorSpec — skipped cleanlydim_idxout of range for DynamicTensorSpec — skipped cleanlyinput_idx/dim_idxout of range for ConstraintSpec — skipped cleanlytune_max_num_tokenswith out-of-bounds spec — handledTest plan
Note: Full autotuner integration test requires model serving, deferred to NVIDIA CI.
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.
Summary by CodeRabbit