[#14679][fix] Fix fused-QKV TP sharding for Phi-3/Phi-4

[guan404ming]

Summary by CodeRabbit

Release Notes

• New Features

  • Enabled Microsoft Phi-4 model variants (microsoft/phi-4, microsoft/Phi-4-reasoning, microsoft/Phi-4-reasoning-plus) for automatic deployment.

• Bug Fixes

  • Fixed tensor parallel sharding computation for fused weight operations to ensure correct dimension calculations.

Description

close #14679

_determine_fused_weight_dims computed the q/k/v split sizes but never returned them, so tensor-parallel column sharding of a fused qkv_proj got None and broke Phi-3/Phi-4 at TP≥2 (reduction-dim mismatch [s44*s70, 3840] X [2560, 5120]).

Test Coverage

• test_determine_fused_weight_dims_qkv -> New regression test. Exports a fused-QKV block and asserts _determine_fused_weight_dims returns the [q, k, v] split sizes (not None).
• test_tp_sharding.py -> Guards that the broader column/row TP-sharding path still produces correct sharded outputs
• manually tested on local L4 GPU and worked as expected

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.

[coderabbitai]

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: 27bdac07-4361-4aab-8911-16473de8c9db

📥 Commits

Reviewing files that changed from the base of the PR and between beb922f and 302dcbe.

📒 Files selected for processing (3)

• examples/auto_deploy/model_registry/models.yaml
• tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py
• tests/unittest/auto_deploy/multigpu/transformations/library/test_tp_sharding.py

---

📝 Walkthrough

Walkthrough

Fixes a bug in _determine_fused_weight_dims where the function fell through without returning its computed fused_weight_dims, yielding None to callers and causing shape mismatches. The fix corrects the return annotation and adds an explicit return. A regression test is added, and the three disabled Phi-4 model registry entries are re-enabled.

Changes

Phi-4 Shape Mismatch Fix

Layer / File(s)	Summary
_determine_fused_weight_dims bug fix and regression test tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py, tests/unittest/auto_deploy/multigpu/transformations/library/test_tp_sharding.py	Return type annotation changed from None to Optional[List[int]] and an explicit return fused_weight_dims added. A new _FusedQKVProj model and test_determine_fused_weight_dims_qkv regression test export an FX graph, locate QKV slice nodes, and assert the returned split sizes [q, k, v] are correct.
Re-enable Phi-4 model registry entries examples/auto_deploy/model_registry/models.yaml	microsoft/phi-4, microsoft/Phi-4-reasoning, and microsoft/Phi-4-reasoning-plus are uncommented and activated with config_id: default_ws_2 and yaml_extra pointing to dashboard_default.yaml and world_size_2.yaml.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 42.86% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The PR title clearly identifies the fix for fused-QKV TP sharding issue affecting Phi-3/Phi-4 models, directly corresponding to the main change in the changeset.
Description check	✅ Passed	The PR description includes a clear explanation of the issue, the solution, comprehensive test coverage details, and a completed PR checklist following the template requirements.
Linked Issues check	✅ Passed	The PR fully addresses issue #14679 by fixing the _determine_fused_weight_dims function to return split sizes, enabling proper TP sharding for Phi-4 models, and re-enabling the three disabled model entries.
Out of Scope Changes check	✅ Passed	All changes are directly related to fixing the fused-QKV TP sharding bug: the core fix in sharding.py, regression test in test_tp_sharding.py, and model registry updates are all within scope.

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

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[guan404ming]

Hi @govind-ramnarayan could you help take a look at this one, thanks!

[govind-ramnarayan]

Hi @guan404ming ! Thank you for the contribution! We're currently in the process of moving over from our heuristic-based sharder to a model based on agentically-generated modeling hints. The idea is that we will have modeling code provide hints as to how each operation would like to be sharded, which provides more transparency and doesn't have as many corner cases as the purely structural approach. See e.g.

https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_nemotron_h.py#L10

and https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/auto_deploy/transform/library/sharding_ir.py

https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_nemotron_h.py#L152

With this in mind, a preferable solution would be to see if we can port the Microsoft-Phi models over to the new sharding IR without issue by adding similar "sharding hints" to the model. If it is not possible, then this solution to fix the legacy sharder looks fine to me, but I would like @greg-kwasniewski1 to confirm. He has a PR up to change our default mode of sharding to the new model-hints approach here: #14786 , which I think is fine with your current PR because it preserves the legacy approach as a fall-back, but I'd like to make sure the combination does not cause issues.

[guan404ming]

Thanks @govind-ramnarayan! Just ported Phi-3/Phi-4 to the sharding IR instead of patching the legacy sharder. New modeling_phi3.py adds hints on the fused qkv_proj/gate_up_proj (output_sizes + split_with_sizes, mirroring modeling_nemotron_h), and the Phi-4 registry entries now run with enable_sharder_ir.yaml. Dropped the legacy sharding.py change.

Validated on L4: single-GPU numerics match HF (MLP exact, full-model rmse ratio 0.0016) and the model exports cleanly; microsoft/phi-4 builds at TP=2 through the IR sharder, so the fused-qkv split no longer hits the #14679 shape mismatch. PTAL.

[guan404ming]

Hi @govind-ramnarayan could you help take another look at this one, really appreciate for your help!

[guan404ming]

Hi @greg-kwasniewski1, could you help take a look at this one, thanks!

[govind-ramnarayan]

Thanks! Sorry about the delay; we are currently in the process of moving AutoDeploy out of TRTLLM. This code change will still be useful for us so I'll see if I can merge it before the move.

[govind-ramnarayan]

/bot run

[tensorrt-cicd]

PR_Github #59313 [ run ] triggered by Bot. Commit: 792dc7c Link to invocation

[tensorrt-cicd]

PR_Github #59313 [ run ] completed with state SUCCESS. Commit: 792dc7c
/LLM/main/L0_MergeRequest_PR pipeline #47796 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

[greg-kwasniewski1]

/bot run

[tensorrt-cicd]

PR_Github #59430 [ run ] triggered by Bot. Commit: 792dc7c Link to invocation

[greg-kwasniewski1]

@govind-ramnarayan @guan404ming

1. the fused weights shape was already fixed in [None][fix] AutoDeploy: return fused_weight_dims so fused QKV split sizes are rescaled under TP #15351
2. Govind is right that anyway legacy path (sharding.py) is marked as obsolete in favor of sharding-aware modeling code + sharding_ir.py

So let's merge this PR as-is, I agree

[tensorrt-cicd]

PR_Github #59430 [ run ] completed with state SUCCESS. Commit: 792dc7c
/LLM/main/L0_MergeRequest_PR pipeline #47898 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

[govind-ramnarayan]

/bot run

[tensorrt-cicd]

PR_Github #59510 [ run ] triggered by Bot. Commit: ce5f6b6 Link to invocation

[tensorrt-cicd]

PR_Github #59510 [ run ] completed with state SUCCESS. Commit: ce5f6b6
/LLM/main/L0_MergeRequest_PR pipeline #47967 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

[govind-ramnarayan]

/bot run

[tensorrt-cicd]

PR_Github #59530 [ run ] triggered by Bot. Commit: ce5f6b6 Link to invocation

[tensorrt-cicd]

PR_Github #59530 [ run ] completed with state SUCCESS. Commit: ce5f6b6
/LLM/main/L0_MergeRequest_PR pipeline #47982 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

[govind-ramnarayan]

/bot run

[tensorrt-cicd]

PR_Github #59565 [ run ] triggered by Bot. Commit: 4610f53 Link to invocation

[tensorrt-cicd]

PR_Github #59565 [ run ] completed with state SUCCESS. Commit: 4610f53
/LLM/main/L0_MergeRequest_PR pipeline #48010 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

[guan404ming]

Thanks for you all, feel free to let me know if there is anything need to refine!

[govind-ramnarayan]

/bot run

[tensorrt-cicd]

PR_Github #60080 [ run ] triggered by Bot. Commit: 4610f53 Link to invocation

[tensorrt-cicd]

PR_Github #60080 [ run ] completed with state SUCCESS. Commit: 4610f53
/LLM/main/L0_MergeRequest_PR pipeline #48467 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