[https://nvbugs/6478707][fix] detect v2 inside remove_functionalize_inner, resolve each mutates_args via…#16629
Open
trtllm-agent wants to merge 1 commit into
Open
Conversation
trtllm-agent
requested review from
asfiyab-nvidia,
brnguyen2,
pcastonguay,
schetlur-nv,
tburt-nv,
yingguo-trt and
yuanjingx87
July 20, 2026 13:46
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
WalkthroughThe MLA custom operation’s inplace metadata now records ChangesMLA inplace metadata
Estimated code review effort: 1 (Trivial) | ~3 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
mikeiovine
requested review from
mikeiovine
and removed request for
pcastonguay and
schetlur-nv
July 20, 2026 15:22
yingguo-trt
approved these changes
Jul 21, 2026
Contributor
JunyiXu-nv
approved these changes
Jul 22, 2026
Collaborator
asfiyab-nvidia
approved these changes
Jul 23, 2026
Torch 2.12 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 'output' are absent. Previously remove_copy_for_mutates_args never detected v2 nodes at the call site, so kwargs["output"] raised KeyError. The v2 branch itself keyed into _all_bases by mutates_args positional index (k-1) rather than by the _<name>_base_index kwarg, which is unrelated to op-arg position. Detect auto_functionalized_v2 inside the inner helper 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. Also register the two additional mutable args (dsv4_output, dsv4_output_sf) for mla_custom_op_inplace in inplace_info() so the reconstructed op call carries them. This unblocks TestDeepSeekV3Lite::test_nvfp4_4gpus with torch_compile under moe_backend=TRTLLM on B200. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
trtllm-agent
force-pushed
the
repair-bot-bug6478707
branch
from
July 24, 2026 08:48
d7dfd7f to
613e5eb
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
_all_bases_<arg>_base_indexinto_all_bases(None-safe), and register the two additional mla_custom_op_inplace mutable args (dsv4_output, dsv4_output_sf) in inplace_infoTest plan
Links
Dev Engineer Review
remove_copy_passhandling forauto_functionalized_v2nodes to resolve mutable arguments through_all_basesusing_<arg>_base_index, includingNonehandling.dsv4_outputanddsv4_output_sfas mutable outputs formla_custom_op_inplace.KeyErrorfailures caused by v2 nodes storing mutable arguments indirectly.QA Engineer Review
No test changes.