[https://nvbugs/6473161][fix] Detect v2 at the call site (is_v2 = is_call_function(node…#16624
[https://nvbugs/6473161][fix] Detect v2 at the call site (is_v2 = is_call_function(node…#16624trtllm-agent wants to merge 2 commits into
Conversation
Torch 2.11 AOT-autograd emits auto_functionalized_v2 HOP nodes for
mutable-arg ops even when inductor's enable_auto_functionalized_v2 is
False (that flag only affects Inductor lowering). In v2 nodes, mutable
tensors live in _all_bases and each mutable arg has its own
_<name>_base_index kwarg; direct kwargs like 'query'/'key' are absent.
Previously remove_copy_for_mutates_args never set is_v2=True at the call
site, and the v2 branch itself keyed into _all_bases by the mutates_args
positional index (k-1) rather than by the _<name>_base_index kwarg,
which is unrelated to op-arg position. Both bugs together caused
KeyError('query') / IndexError on flashinfer_apply_rope_with_cos_sin_cache_inplace
and similar ops under torch.compile+FLASHINFER.
Detect auto_functionalized_v2 at the call site and, for each mutable
arg, resolve its base tensor via _<arg>_base_index into _all_bases,
mapping a missing/None index to None so optional mutable args are
preserved.
Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
WalkthroughThe copy-removal pass now derives ChangesFunctionalization copy removal
Estimated code review effort: 3 (Moderate) | ~15–30 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/compilation/remove_copy_pass.py (1)
20-20: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a precise return and mapping type annotation.
This non-returning helper lacks
-> None, anddicthides the contract that keys are output indices and values are mutable argument names. The coding guidelines require every function to be annotated.Suggested change
- def remove_functionalize_inner(node: Node, mutates_args: dict): + def remove_functionalize_inner( + node: Node, mutates_args: dict[int, str] + ) -> None:Based on learnings, this repository supports Python 3.10+ built-in generic annotations.
🤖 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 `@tensorrt_llm/_torch/compilation/remove_copy_pass.py` at line 20, Update remove_functionalize_inner with a complete type signature: annotate its return as None and replace the unparameterized mutates_args dict with a built-in generic mapping whose keys are output indices and values are mutable argument names, using the repository’s Python 3.10+ annotation style.Sources: Coding guidelines, Learnings
🤖 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.
Nitpick comments:
In `@tensorrt_llm/_torch/compilation/remove_copy_pass.py`:
- Line 20: Update remove_functionalize_inner with a complete type signature:
annotate its return as None and replace the unparameterized mutates_args dict
with a built-in generic mapping whose keys are output indices and values are
mutable argument names, using the repository’s Python 3.10+ annotation style.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 57317104-0dfe-4fea-9826-6c48cff080fb
📒 Files selected for processing (2)
tensorrt_llm/_torch/compilation/remove_copy_pass.pytests/integration/test_lists/waives.txt
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/waives.txt
BowenFu
left a comment
BowenFu
left a comment
There was a problem hiding this comment.
LGTM — the fix detects auto_functionalized_v2 in-function and rebuilds the mutable-arg kwargs; the sole caller never set is_v2, so the v2 branch was dead/crashing (KeyError under torch 2.11) before, and non-v2 (auto_functionalized) nodes are byte-identical. Strictly a fix (un-waives 4 FLASHINFER+torch_compile tests).
|
/bot run |
|
PR_Github #60851 [ run ] triggered by Bot. Commit: |
|
PR_Github #60851 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #60903 [ run ] triggered by Bot. Commit: |
|
PR_Github #60903 [ run ] completed with state
|
|
/bot run |
|
PR_Github #61191 [ run ] triggered by Bot. Commit: |
|
PR_Github #61191 [ run ] completed with state |
Signed-off-by: Lori Ren <lorir@nvidia.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tensorrt_llm/_torch/compilation/remove_copy_pass.py (1)
34-34: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd precise function annotations.
The nested helper is non-returning but lacks
-> None, and baredictdoes not describe its key/value types. Use precise Python 3.10+ annotations, such asdict[int, str].As per coding guidelines, every function must be annotated and non-returning functions must use
Noneas the return type.Proposed fix
- def remove_functionalize_inner(node: Node, mutates_args: dict): + def remove_functionalize_inner( + node: Node, mutates_args: dict[int, str] + ) -> None:🤖 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 `@tensorrt_llm/_torch/compilation/remove_copy_pass.py` at line 34, Update the nested helper remove_functionalize_inner with a Python 3.10+ return annotation of None and replace the bare mutates_args: dict annotation with the precise key/value type dict[int, str], preserving its existing behavior.Source: Coding guidelines
🤖 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 `@tensorrt_llm/_torch/compilation/remove_copy_pass.py`:
- Line 34: Remove the stale is_v2 keyword argument from every call to
remove_functionalize_inner in the functionalized-node handling path, especially
the calls around lines 76-80. Keep the existing node and mutates_args arguments
unchanged so recognized nodes proceed to the copy-removal logic without a
TypeError.
---
Nitpick comments:
In `@tensorrt_llm/_torch/compilation/remove_copy_pass.py`:
- Line 34: Update the nested helper remove_functionalize_inner with a Python
3.10+ return annotation of None and replace the bare mutates_args: dict
annotation with the precise key/value type dict[int, str], preserving its
existing behavior.
🪄 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: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 0a8dd450-b121-40dd-9ccc-9aeaf990606f
📒 Files selected for processing (1)
tensorrt_llm/_torch/compilation/remove_copy_pass.py
| nodes_to_remove: list[Node] = [] | ||
|
|
||
| def remove_functionalize_inner(node: Node, mutates_args: dict, is_v2=False): | ||
| def remove_functionalize_inner(node: Node, mutates_args: dict): |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
Remove the stale is_v2 keyword argument.
remove_functionalize_inner no longer accepts is_v2, but Lines 76-80 still pass it. Every recognized functionalized node will therefore fail with TypeError: unexpected keyword argument 'is_v2' before the copy-removal logic executes.
Proposed fix
remove_functionalize_inner(
node,
inplace_map[inplace_func],
- is_v2=node.target == auto_functionalized_v2,
)Also applies to: 76-80
🤖 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 `@tensorrt_llm/_torch/compilation/remove_copy_pass.py` at line 34, Remove the
stale is_v2 keyword argument from every call to remove_functionalize_inner in
the functionalized-node handling path, especially the calls around lines 76-80.
Keep the existing node and mutates_args arguments unchanged so recognized nodes
proceed to the copy-removal logic without a TypeError.
7 participants
Summary
Test plan
Links
Summary
Dev Engineer Review
tensorrt_llm/_torch/compilation/remove_copy_pass.pyto fixauto_functionalized_v2handling inremove_copy_for_mutates_args:remove_functionalize_innernow detectsauto_functionalized_v2directly and rebuilds the mutates-args replacement tensors usingnode.kwargs["_all_bases"]plus per-arg_{arg_name}_base_indexindices.Nonewhen their_{arg_name}_base_indexisNone, preventing mis-indexed/missing outputs (e.g.,query,key,output_sf) in Torch 2.11 AOT-autograd graphs.remove_functionalize_innersignature no longer accepts anis_v2parameter, but the call still passesis_v2=...(would raiseTypeErrorwhen the functionalization node path is executed). The callsite should be updated to stop passingis_v2.QA Engineer Review
tests/integration/test_lists/waives.txt: removed four waiver lines forfull:L40S/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus(and relatedtest_fp8coverage) withattn_backend=FLASHINFERandtorch_compile=True, forfp8kv=Falseandfp8kv=True.CI/Verification
#49439success).