Skip to content

[None][fix] piecewise cudagraph node name for gemma-3-1b-it#12131

Open
amukkara wants to merge 1 commit intoNVIDIA:mainfrom
amukkara:tc-gemma
Open

[None][fix] piecewise cudagraph node name for gemma-3-1b-it#12131
amukkara wants to merge 1 commit intoNVIDIA:mainfrom
amukkara:tc-gemma

Conversation

@amukkara
Copy link
Collaborator

@amukkara amukkara commented Mar 12, 2026

@coderabbitai summary

Description

  1. For gemma-3-1b-it, input_ids node name in fx graph is l_kwargs_input_ids_. For other models, the node name is l_input_ids_. This PR pattern matches against both names when estimating input_num_tokens.
  2. Fixes initialization of input_num_tokens to avoid "Undefined variable" error when pattern match in 1 fails.
  3. Common to all models: torch_inductor and piecewise_cuda_graph are incompatible according to our docs. This PR keeps the default value TorchCompileConfig.enable_inductor=False in example script to avoid conflict with piecewise_cuda_graph.

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.

Signed-off-by: Anurag Mukkara <134339030+amukkara@users.noreply.github.com>
@coderabbitai
Copy link
Contributor

coderabbitai bot commented Mar 12, 2026

📝 Walkthrough

Walkthrough

Changes remove a configuration parameter from an example file and refine backend compilation logic to initialize a variable, expand placeholder name matching to include an additional pattern, and add early termination after capturing input token counts.

Changes

Cohort / File(s) Summary
Example Configuration
examples/llm-api/quickstart_advanced.py
Removed enable_inductor parameter from TorchCompileConfig initialization call within setup_llm function.
Backend Compilation Logic
tensorrt_llm/_torch/compilation/backend.py
Initialize self.input_num_tokens to None at method start. Expanded placeholder name check to match both "l_input_ids_" and "l_kwargs_input_ids_" patterns. Added early break after capturing the first matching placeholder to prevent redundant iterations.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~8 minutes

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (2 passed)
Check name Status Explanation
Description check ✅ Passed The PR description is mostly complete with three clear points explaining what was changed and why, along with a filled PR checklist.
Title check ✅ Passed The title is partially related to the changeset, focusing on a specific fix for piecewise cudagraph node names (related to the backend.py changes), but doesn't address the removal of enable_inductor from quickstart_advanced.py, which is also a significant change.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Post copyable unit tests in a comment
📝 Coding Plan
  • Generate coding plan for human review comments

Comment @coderabbitai help to get the list of available commands and usage tips.

Tip

CodeRabbit can enforce grammar and style rules using `languagetool`.

Configure the reviews.tools.languagetool setting to enable/disable rules and categories. Refer to the LanguageTool Community to learn more.

@amukkara amukkara changed the title [None]fix] Input_ids node name for piecewise cudagraph backend [None]fix] piecewise cudagraphs input_ids node name for gemma-3-1b-it Mar 12, 2026
@amukkara amukkara changed the title [None]fix] piecewise cudagraphs input_ids node name for gemma-3-1b-it [None]fix] piecewise cudagraph node name for gemma-3-1b-it Mar 12, 2026
@amukkara amukkara changed the title [None]fix] piecewise cudagraph node name for gemma-3-1b-it [None][fix] piecewise cudagraph node name for gemma-3-1b-it Mar 12, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant