[https://nvbugs/6517842][fix] Handle mutable tensor lists in remove c… by liji-nv · Pull Request #16926 · NVIDIA/TensorRT-LLM

liji-nv · 2026-07-28T03:27:15Z

…opy pass

auto_functionalized_v2 represents mutable Tensor[] arguments with a list length and per-element base indices. The remove-copy pass assumed every mutable argument had a single base index, causing torch.compile to look up the nonexistent _tensors_base_index key for pipeline-parallel send and receive operations.

Reconstruct mutable tensor lists from their flattened bases, map each functionalized output back to the corresponding input tensor, and retain the existing scalar and optional-tensor behavior. Add regression coverage for both pp_send_tensors and pp_recv_tensors.

Unwaive the cases tracked by NVBUG 6517842, NVBUG 6445472, and NVBUG 6474888.

Tested on B200:

5 remove_copy_pass unit tests passed
DeepSeekV3-Lite TP2/PP2 torch.compile accuracy case passed
TensorRT-LLM build completed successfully

Dev Engineer Review

Updated remove_copy_for_mutates_args (v2 auto_functionalized_v2 path) to correctly reconstruct mutable tensor lists from flattened mutated outputs by building per-getitem replacement mappings from _all_bases and the inplace function schema write args.
Handles both list-form writes (when _{arg.name}_length is present, including length is None → kwargs[arg.name] = None) and the single-base-index fallback.
Updated getitem replacement logic so tensor-list consumers are mapped back to the correct input tensors (from the per-index replacement map when available), while preserving optional/None semantics via existing assertions.
Kept the overall pass behavior consistent: replaces getitem users, inserts the inplace call_function with reconstructed kwargs, and erases the auto_functionalized/auto_functionalized_v2 node and now-unneeded getitems.
Added regression coverage for tensor-list mutates in TRT-LLM ops:
- pp_send_tensors.default
- pp_recv_tensors.default
Updated tests/integration/test_lists/waives.txt to remove/adjust DeepSeekV3Lite skip waivers tied to NVBUG 6517842, 6445472, and 6474888, aligning with the intended “waive unwaived” behavior.

QA Engineer Review

Test list changes

Modified tests/integration/test_lists/waives.txt
- accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus:
  - Changed tp2pp2 waiver for mtp_nextn=2 overlap_scheduler=False → overlap_scheduler=True.
- test_nvfp4_4gpus:
  - Removed SKIP waiver for tp2pp2 with fp8kv=True and torch_compile=True.
- test_cute_dsl_nvfp4_4gpus:
  - Removed tp4 + fp8kv=False + torch_compile=True SKIPs for full:B300, full:GB200, full:GB300.
  - For full:B200, replaced torch_compile=True waiver with torch_compile=False.

Verdict: needs follow-up (CBTS coverage data not provided here).

Test code changes

Modified tests/unittest/_torch/compilation/test_remove_copy_pass.py
- Added test_remove_copy_for_mutates_tensor_list(inplace_func: Callable[..., object]) (parametrized over torch.ops.trtllm.pp_recv_tensors.default and torch.ops.trtllm.pp_send_tensors.default).
  - Validates: reconstruction of kwargs["tensors"], removal of auto_functionalized_v2, and preservation of inplace node behavior.
Coverage mapping:
- No explicit tests/integration/test_lists/ (test-db/qa) entry for this new unit test was identifiable from the provided context.

Verdict: needs follow-up (no CBTS/test-db/qa coverage linkage data available in the provided context).

Reported validation on B200:
- Five remove-copy-pass unit tests passed
- DeepSeekV3-Lite TP2/PP2 torch.compile accuracy case passed
- TensorRT-LLM build completed successfully

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)
If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.
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.

…opy pass auto_functionalized_v2 represents mutable Tensor[] arguments with a list length and per-element base indices. The remove-copy pass assumed every mutable argument had a single base index, causing torch.compile to look up the nonexistent _tensors_base_index key for pipeline-parallel send and receive operations. Reconstruct mutable tensor lists from their flattened bases, map each functionalized output back to the corresponding input tensor, and retain the existing scalar and optional-tensor behavior. Add regression coverage for both pp_send_tensors and pp_recv_tensors. Unwaive the cases tracked by NVBUG 6517842, NVBUG 6445472, and NVBUG 6474888. Tested on B200: - 5 remove_copy_pass unit tests passed - DeepSeekV3-Lite TP2/PP2 torch.compile accuracy case passed - TensorRT-LLM build completed successfully Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>

liji-nv · 2026-07-28T03:27:59Z

/bot run --disable-fail-fast

coderabbitai · 2026-07-28T03:30:24Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: a56cf0c9-b0dc-4c62-9857-dc57d96ddb2a

📥 Commits

Reviewing files that changed from the base of the PR and between 5be7bee and 2b88c51.

📒 Files selected for processing (1)

tests/integration/test_lists/waives.txt

💤 Files with no reviewable changes (1)

tests/integration/test_lists/waives.txt

Walkthrough

The remove-copy pass now reconstructs replacements for flattened mutated tensor-list outputs in the v2 auto-functionalization path. Unit coverage validates TRT-LLM tensor-list inplace operations, and selected DeepSeekV3Lite integration waivers are updated.

Changes

Tensor-list mutation handling

Layer / File(s)	Summary
Reconstruct flattened tensor-list replacements `tensorrt_llm/_torch/compilation/remove_copy_pass.py`	The v2 path maps flattened mutated outputs to tensor-list write arguments and uses the mappings when replacing getitem users.
Validate tensor-list inplace transformation `tests/unittest/_torch/compilation/test_remove_copy_pass.py`	Parametrized coverage verifies rewritten tensor arguments, removal of `auto_functionalized_v2`, and FX graph validity.

Integration waiver updates

Layer / File(s)	Summary
Update DeepSeekV3Lite waivers `tests/integration/test_lists/waives.txt`	Selected bfloat16, nvfp4, and cute-dsl nvfp4 waiver entries are removed or replaced across execution groups.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant auto_functionalized_v2
  participant remove_copy_for_mutates_args
  participant FXGraph
  auto_functionalized_v2->>remove_copy_for_mutates_args: provide flattened mutated outputs
  remove_copy_for_mutates_args->>remove_copy_for_mutates_args: build tensor_list_replacements
  remove_copy_for_mutates_args->>FXGraph: replace getitem users and insert inplace call

Possibly related PRs

NVIDIA/TensorRT-LLM#16934: Updates overlapping DeepSeekV3Lite SKIP waivers in tests/integration/test_lists/waives.txt.
NVIDIA/TensorRT-LLM#16949: Changes related DeepSeekV3Lite and B200 waiver entries in the same file.
NVIDIA/TensorRT-LLM#16950: Updates overlapping cute-dsl nvfp4 DeepSeekV3Lite waivers.

Suggested reviewers: qijune, stanleysun639, xinhe-nv, zhanruisunch, cascade812

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title matches the main fix and follows the required [ticket][type] pattern.
Description check	✅ Passed	The PR description covers the issue, solution, tests, and checklist, though it uses summary text instead of the exact template headings.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

tensorrt-cicd · 2026-07-28T03:34:13Z

PR_Github #62083 [ run ] triggered by Bot. Commit: 5be7bee Link to invocation

BowenFu

Regression-safe as far as I can tell. The non-list path is byte-identical to base 89c6635, and the new branch only fires when the _<arg>_length kwarg is present, which only happens for a mutable Tensor[] — so every node shape that worked before still takes the old path. min(mutates_args) can't raise because every inplace_info() entry is non-empty. And _all_bases indices are absolute across scalar and list-element bases, so first_mutated_output + base_index is the right output slot.

One residual, non-blocking: a mutable scalar arg positioned after a mutable list would have its mutates_args key shifted by the list expansion and isn't reconstructed here. No registered op in inplace_info() has that shape today, and it never worked before either, so it's not a regression — just worth a comment if someone adds such an op later.

tensorrt-cicd · 2026-07-28T08:10:05Z

PR_Github #62083 [ run ] completed with state SUCCESS. Commit: 5be7bee
/LLM/main/L0_MergeRequest_PR pipeline #50266 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

liji-nv · 2026-07-28T12:44:42Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-28T12:51:33Z

PR_Github #62198 [ run ] triggered by Bot. Commit: 5be7bee Link to invocation

tensorrt-cicd · 2026-07-28T13:32:11Z

PR_Github #62198 [ run ] completed with state SUCCESS. Commit: 5be7bee
/LLM/main/L0_MergeRequest_PR pipeline #50368 completed with status: 'SUCCESS'

CI Report

Link to invocation

Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>

liji-nv · 2026-07-29T04:57:56Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-29T05:03:32Z

PR_Github #62419 [ run ] triggered by Bot. Commit: 2b88c51 Link to invocation

tensorrt-cicd · 2026-07-29T08:47:10Z

PR_Github #62419 [ run ] completed with state SUCCESS. Commit: 2b88c51
/LLM/main/L0_MergeRequest_PR pipeline #50578 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

liji-nv · 2026-07-30T02:53:45Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-30T02:59:03Z

PR_Github #62666 [ run ] triggered by Bot. Commit: 2b88c51 Link to invocation

tensorrt-cicd · 2026-07-30T04:35:18Z

PR_Github #62666 [ run ] completed with state SUCCESS. Commit: 2b88c51
/LLM/main/L0_MergeRequest_PR pipeline #50806 completed with status: 'SUCCESS'

CI Report

Link to invocation

liji-nv requested review from a team as code owners July 28, 2026 03:27

liji-nv requested review from BowenFu, QiJune, ZhanruiSunCh, cascade812, ruodil, tburt-nv, tongyuantongyu and xinhe-nv July 28, 2026 03:27

github-actions Bot assigned liji-nv Jul 28, 2026

BowenFu approved these changes Jul 28, 2026

View reviewed changes

coderabbitai Bot mentioned this pull request Jul 28, 2026

[TRTLLM-13579][feat] BREAKING: Support BCG in Prefill #16609

Open

tburt-nv approved these changes Jul 28, 2026

View reviewed changes

Merge branch 'main' into unwaive-nvbugs-6445472-6474888

2b88c51

Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>

liji-nv requested a review from a team as a code owner July 29, 2026 04:55

liji-nv requested a review from StanleySun639 July 29, 2026 04:55

QiJune approved these changes Jul 29, 2026

View reviewed changes

liji-nv merged commit eed4d5a into NVIDIA:main Jul 30, 2026
10 checks passed

Uh oh!

Conversation

liji-nv commented Jul 28, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Dev Engineer Review

QA Engineer Review

Test list changes

Test code changes

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

liji-nv commented Jul 28, 2026

Uh oh!

coderabbitai Bot commented Jul 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Sequence Diagram(s)

Uh oh!

tensorrt-cicd commented Jul 28, 2026

Uh oh!

BowenFu left a comment

Choose a reason for hiding this comment

Uh oh!

tensorrt-cicd commented Jul 28, 2026

Uh oh!

liji-nv commented Jul 28, 2026

Uh oh!

tensorrt-cicd commented Jul 28, 2026

Uh oh!

tensorrt-cicd commented Jul 28, 2026

Uh oh!

liji-nv commented Jul 29, 2026

Uh oh!

tensorrt-cicd commented Jul 29, 2026

Uh oh!

tensorrt-cicd commented Jul 29, 2026

Uh oh!

liji-nv commented Jul 30, 2026

Uh oh!

tensorrt-cicd commented Jul 30, 2026

Uh oh!

tensorrt-cicd commented Jul 30, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

liji-nv commented Jul 28, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Jul 28, 2026 •

edited

Loading