[None][fix] AutoDeploy: return fused_weight_dims so fused QKV split sizes are rescaled under TP

[CodersAcademy006]

Summary

Fixes #11220.

_determine_fused_weight_dims() computes fused projection dimensions for fused QKV layouts but never returns the computed value. As a result, _process_column_sharding() always receives None and skips downstream split-size rescaling during tensor parallel sharding.

Root Cause

The helper populates fused_weight_dims but falls through without a return statement despite callers expecting the value.

Fix

• Change return type from None to Optional[list]
• Return computed fused_weight_dims
• Add regression tests covering fused and non-fused paths

Validation

Added CPU-only regression tests for:

• fused QKV split detection
• unfused linear projections
• multiple linear inputs

The fused QKV test reproduces the regression (the helper returns None before the fix); the other two pin existing behavior.

Summary by CodeRabbit

• Bug Fixes
  • Improved handling of fused weight dimension detection so the calculated split sizes are now returned correctly when available.
• Tests
  • Added regression coverage for fused QKV split detection.
  • Added checks for cases where fused dimensions should not be returned, including unfused linear paths and multiple-linear scenarios.

[coderabbitai]

📝 Walkthrough

Walkthrough

A missing return statement in _determine_fused_weight_dims is fixed so it now returns computed fused weight dimensions instead of implicitly returning None, addressing a regression affecting Phi-4 model sharding. A new test module validates the fix across fused, unfused, and multi-linear graph scenarios.

Changes

Fused Weight Dims Fix and Tests

Layer / File(s)	Summary
Fix missing return statement tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py	_determine_fused_weight_dims return type annotation updated to Optional[list], and an explicit return fused_weight_dims is added so the computed value propagates to callers instead of being silently dropped.
Regression tests for fused weight dims tests/unittest/auto_deploy/singlegpu/transformations/library/test_determine_fused_weight_dims.py	New CPU-only test module with a _linear_node helper and three tests validating: fused Phi-4 QKV split sizes are returned, unfused linear returns None, and multiple linear nodes return None.

Estimated code review effort: 2 (Simple) | ~10 minutes

Related issues: Fixes #11220 — [Bug]: microsoft/phi-4 models regression, caused by _determine_fused_weight_dims not returning its computed value, leading to a matmul reduction-dim mismatch during sharding.

Suggested labels: bug, auto-deploy, tests

Suggested reviewers: (repository maintainers familiar with auto_deploy sharding transforms)

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	The fix matches #11220 by returning fused split sizes and adding regression tests for fused, unfused, and multi-linear paths.
Out of Scope Changes check	✅ Passed	The changes stay tightly scoped to the AutoDeploy sharding bug fix and its regression tests.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Title check	✅ Passed	The title matches the main fix and follows the required [None][fix] pattern.
Description check	✅ Passed	The description clearly states the issue, root cause, fix, and validation, covering the required content well.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

🧹 Nitpick comments (2)

tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py (1)

2709-2711: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Use list[int] | None instead of Optional[list].

Per repo typing conventions, prefer built-in generics with the | syntax over typing.Optional, and specify the element type rather than a bare list.

♻️ Proposed fix

 def _determine_fused_weight_dims(
     linear_nodes: List[Node],
-) -> Optional[list]:
+) -> list[int] | None:

As per coding guidelines, "Prefer built-in types list, dict, tuple over typing.List, typing.Dict, typing.Tuple; use | syntax instead of typing.Union" and "Always annotate functions with return types... annotate class members and variables when necessary."

🤖 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/auto_deploy/transform/library/sharding.py` around lines
2709 - 2711, Update the return type annotation on _determine_fused_weight_dims
to use the repo’s preferred built-in generic syntax: replace Optional[list] with
list[int] | None. Keep the change localized to the _determine_fused_weight_dims
function signature in sharding.py, and make sure any related type hints in
nearby code follow the same built-in generics style.

Source: Coding guidelines

tests/unittest/auto_deploy/singlegpu/transformations/library/test_determine_fused_weight_dims.py (1)

45-102: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Good regression coverage for the split path; slice/chunk branches remain untested.

The three tests solidly cover the split_with_sizes regression path and two None-return guards. However, _determine_fused_weight_dims also has elif branches for torch.ops.aten.slice users (fused QKV via slicing, sharding.py lines 2731-2746) and torch.ops.aten.chunk users (sharding.py lines 2748-2754), both of which are also affected by the same missing-return bug and are not exercised here.

Coverage is sufficient for the immediate Phi-4 regression (split-based fusion) but insufficient for the helper as a whole. Consider adding test_fused_qkv_slice_returns_split_sizes and test_fused_qkv_chunk_returns_split_sizes to this same file to close the gap, or track as a fast follow-up if out of scope for this PR.

🤖 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
`@tests/unittest/auto_deploy/singlegpu/transformations/library/test_determine_fused_weight_dims.py`
around lines 45 - 102, Add regression tests for the untested fused paths in
`_determine_fused_weight_dims`: create one test that feeds a linear node into
`torch.ops.aten.slice` and another into `torch.ops.aten.chunk`, then assert the
helper returns the expected fused sizes instead of None. Reuse the existing test
setup pattern in `test_determine_fused_weight_dims.py` and the helper
`_linear_node` so the new coverage matches the current split-based test and
exercises the `slice` and `chunk` branches in `_determine_fused_weight_dims`.

Source: Path instructions

🤖 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/auto_deploy/transform/library/sharding.py`:
- Around line 2709-2711: Update the return type annotation on
_determine_fused_weight_dims to use the repo’s preferred built-in generic
syntax: replace Optional[list] with list[int] | None. Keep the change localized
to the _determine_fused_weight_dims function signature in sharding.py, and make
sure any related type hints in nearby code follow the same built-in generics
style.

In
`@tests/unittest/auto_deploy/singlegpu/transformations/library/test_determine_fused_weight_dims.py`:
- Around line 45-102: Add regression tests for the untested fused paths in
`_determine_fused_weight_dims`: create one test that feeds a linear node into
`torch.ops.aten.slice` and another into `torch.ops.aten.chunk`, then assert the
helper returns the expected fused sizes instead of None. Reuse the existing test
setup pattern in `test_determine_fused_weight_dims.py` and the helper
`_linear_node` so the new coverage matches the current split-based test and
exercises the `slice` and `chunk` branches in `_determine_fused_weight_dims`.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 14b2a658-40df-45e2-8b75-0b889269e4db

📥 Commits

Reviewing files that changed from the base of the PR and between 45fae3e and 6eac3ea.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py
• tests/unittest/auto_deploy/singlegpu/transformations/library/test_determine_fused_weight_dims.py

[greg-kwasniewski1]

Good catch!

[greg-kwasniewski1]

/bot run

[tensorrt-cicd]

PR_Github #58495 [ run ] triggered by Bot. Commit: bbea605 Link to invocation

[tensorrt-cicd]

PR_Github #58495 [ run ] completed with state SUCCESS. Commit: bbea605
/LLM/main/L0_MergeRequest_PR pipeline #47103 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation