[None][chore] AutoDeploy: Refactor finegrained FP8 scale sharding helpers#12999

galagam · 2026-04-13T13:03:28Z

@CodeRabbit description

Description

Technical debt PR:

Minor refactor, follow up on [AutoDeploy]: DeepSeek R1 onboarding + initial perf analysis #12593: Single public wrapper for compute_shard_bounds + expand_scale_per_row.
Added TestNemotronSuperV3::test_functional_small - a single-GPU smoke test that loads a layer-reduced (16-layer) Super V3 FP8 model using reduced_model_kwargs(), a reusable helper that truncates both num_hidden_layers and layers_block_type to fit large hybrid models on one device. Added to l0_b200.yml for pre-merge coverage.

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)
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.

galagam · 2026-04-13T13:04:06Z

/bot run --stage-list "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-Post-Merge-1,DGX_B200-8_GPUs-AutoDeploy-Post-Merge-1"

tensorrt-cicd · 2026-04-13T13:10:09Z

PR_Github #43042 [ run ] triggered by Bot. Commit: 0be96a1 Link to invocation

coderabbitai · 2026-04-13T13:12:44Z

📝 Walkthrough

Walkthrough

Refactored FineGrained FP8 weight sharding logic by extracting a centralized _get_sharded_scale() method in sharding.py to compute shard bounds and scale expansion. Updated existing methods and MoE TP scale sharding to use the new centralized approach. Corresponding test updated to verify the refactored logic.

Changes

Cohort / File(s)	Summary
FineGrained FP8 Sharding Refactoring `tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py`	Extracted `_get_sharded_scale()` method to centralize weight_scale_inv sharding computation. Updated `shard_scales()`, `shard_load_hook()`, and `_tp_shard_moe_scale()` to use the new centralized method. Added explicit return type annotation for `_compute_shard_bounds()`.
Test Update `tests/unittest/auto_deploy/multigpu/transformations/library/test_tp_sharding.py`	Renamed test from `test_finegrained_fp8_expand_scale_per_row` to `test_finegrained_fp8_get_sharded_scale`. Updated test to verify the new `_get_sharded_scale()` method instead of directly testing `_compute_shard_bounds()` and `_expand_scale_per_row()`. Adjusted assertions to match refactored behavior.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

🚥 Pre-merge checks | ✅ 1 | ❌ 2

❌ Failed checks (1 warning, 1 inconclusive)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 55.56% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Description check	❓ Inconclusive	PR description addresses key requirements but Test Coverage section is incomplete with only a comment placeholder and no actual test details listed.	Expand Test Coverage section to clearly list the specific tests that validate the refactoring changes, including the renamed test_finegrained_fp8_get_sharded_scale test and any other relevant test coverage.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: a refactor of finegrained FP8 scale sharding helpers in AutoDeploy, which aligns with the changeset.

_{✏️ 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 and usage tips.}

coderabbitai

🧹 Nitpick comments (1)

tests/unittest/auto_deploy/multigpu/transformations/library/test_tp_sharding.py (1)
128-177: Please add a dim=1 / MoE regression case.

This test only exercises dim=0, but the refactor also routes _tp_shard_moe_scale() through _get_sharded_scale() for row-sharded w_down experts. A bug in the dim=1 branch would still pass this suite.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In
`@tests/unittest/auto_deploy/multigpu/transformations/library/test_tp_sharding.py`
around lines 128 - 177, Add a new test variant invoking
FineGrainedFP8WeightShardingInfo._get_sharded_scale with dim=1 to cover
MoE/w_down expert row-sharded paths: reuse the existing test logic but transpose
the role of rows/columns (or create weight_original_m analog) and assert both
aligned (compact) and non-aligned (per-item expansion) behaviors, and include
the edge case where size is not divisible by block_n (clamping to last scale
row); also add a specific case that routes through _tp_shard_moe_scale by
constructing inputs that mimic row-sharded w_down experts so any bug in the
dim=1 branch fails the test.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Nitpick comments:
In
`@tests/unittest/auto_deploy/multigpu/transformations/library/test_tp_sharding.py`:
- Around line 128-177: Add a new test variant invoking
FineGrainedFP8WeightShardingInfo._get_sharded_scale with dim=1 to cover
MoE/w_down expert row-sharded paths: reuse the existing test logic but transpose
the role of rows/columns (or create weight_original_m analog) and assert both
aligned (compact) and non-aligned (per-item expansion) behaviors, and include
the edge case where size is not divisible by block_n (clamping to last scale
row); also add a specific case that routes through _tp_shard_moe_scale by
constructing inputs that mimic row-sharded w_down experts so any bug in the
dim=1 branch fails the test.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: ca7f5a1b-ba62-4449-a596-66267562a05d

📥 Commits

Reviewing files that changed from the base of the PR and between 68a1be1 and 0be96a1.

📒 Files selected for processing (2)

tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py
tests/unittest/auto_deploy/multigpu/transformations/library/test_tp_sharding.py

tensorrt-cicd · 2026-04-13T14:28:47Z

PR_Github #43042 [ run ] completed with state SUCCESS. Commit: 0be96a1
/LLM/main/L0_MergeRequest_PR pipeline #33688 (Partly Tested) 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

Link to invocation

galagam · 2026-04-19T07:09:29Z

/bot run --stage-list "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-Post-Merge-1,DGX_B200-8_GPUs-AutoDeploy-Post-Merge-1"

tensorrt-cicd · 2026-04-19T07:16:31Z

PR_Github #44167 [ run ] triggered by Bot. Commit: 1e552c2 Link to invocation

tensorrt-cicd · 2026-04-19T09:43:51Z

PR_Github #44167 [ run ] completed with state SUCCESS. Commit: 1e552c2
/LLM/main/L0_MergeRequest_PR pipeline #34594 (Partly Tested) 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

Link to invocation

galagam · 2026-04-19T18:10:45Z

/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-Post-Merge-1,DGX_B200-8_GPUs-AutoDeploy-Post-Merge-1"

tensorrt-cicd · 2026-04-19T18:17:25Z

PR_Github #44215 [ run ] triggered by Bot. Commit: 1e552c2 Link to invocation

tensorrt-cicd · 2026-04-19T19:39:21Z

PR_Github #44215 [ run ] completed with state SUCCESS. Commit: 1e552c2
/LLM/main/L0_MergeRequest_PR pipeline #34638 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

Link to invocation

galagam · 2026-04-20T05:27:29Z

/bot run --stage-list "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-Post-Merge-1,DGX_B200-8_GPUs-AutoDeploy-Post-Merge-1"

tensorrt-cicd · 2026-04-20T05:34:24Z

PR_Github #44317 [ run ] triggered by Bot. Commit: 1e552c2 Link to invocation

tensorrt-cicd · 2026-04-20T05:44:15Z

PR_Github #44317 [ run ] completed with state FAILURE. Commit: 1e552c2
/LLM/main/L0_MergeRequest_PR pipeline #34738 (Partly Tested) 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

Link to invocation

galagam · 2026-04-20T08:05:58Z

/bot run --stage-list "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-Post-Merge-1,DGX_B200-8_GPUs-AutoDeploy-Post-Merge-1"

tensorrt-cicd · 2026-04-20T08:11:55Z

PR_Github #44395 [ run ] triggered by Bot. Commit: 1e552c2 Link to invocation

tensorrt-cicd · 2026-04-20T08:35:47Z

PR_Github #44395 [ run ] completed with state FAILURE. Commit: 1e552c2
/LLM/main/L0_MergeRequest_PR pipeline #34810 (Partly Tested) 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

Link to invocation

galagam · 2026-04-20T12:01:50Z

/bot run --stage-list "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-Post-Merge-1,DGX_B200-8_GPUs-AutoDeploy-Post-Merge-1"

tensorrt-cicd · 2026-04-20T12:07:22Z

PR_Github #44449 [ run ] triggered by Bot. Commit: 1e552c2 Link to invocation

tensorrt-cicd · 2026-04-20T13:12:15Z

PR_Github #44449 [ run ] completed with state SUCCESS. Commit: 1e552c2
/LLM/main/L0_MergeRequest_PR pipeline #34853 (Partly Tested) 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

Link to invocation

galagam · 2026-04-20T17:03:42Z

/bot run --stage-list "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-Post-Merge-1,DGX_B200-8_GPUs-AutoDeploy-Post-Merge-1"

tensorrt-cicd · 2026-04-20T17:10:14Z

PR_Github #44480 [ run ] triggered by Bot. Commit: 1e552c2 Link to invocation

galagam · 2026-04-23T02:57:22Z

/bot run

tensorrt-cicd · 2026-04-23T03:04:07Z

PR_Github #45089 [ run ] triggered by Bot. Commit: 75b1a71 Link to invocation

tensorrt-cicd · 2026-04-23T09:56:52Z

PR_Github #45089 [ run ] completed with state SUCCESS. Commit: 75b1a71
/LLM/main/L0_MergeRequest_PR pipeline #35387 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

Link to invocation

galagam · 2026-04-23T10:12:55Z

/bot run

tensorrt-cicd · 2026-04-23T10:18:47Z

PR_Github #45173 [ run ] triggered by Bot. Commit: 75b1a71 Link to invocation

tensorrt-cicd · 2026-04-23T15:01:24Z

PR_Github #45173 [ run ] completed with state FAILURE. Commit: 75b1a71
/LLM/main/L0_MergeRequest_PR pipeline #35450 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

Link to invocation

galagam · 2026-04-23T17:02:17Z

/bot kill

tensorrt-cicd · 2026-04-23T17:09:13Z

PR_Github #45228 [ kill ] triggered by Bot. Commit: 75b1a71 Link to invocation

tensorrt-cicd · 2026-04-23T17:09:19Z

PR_Github #45228 [ kill ] completed with state SUCCESS. Commit: 75b1a71
Successfully killed previous jobs for commit 75b1a71

Link to invocation

Minor refactor, follow up on NVIDIA#12593 Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com>

Add test_functional_small to TestNemotronSuperV3 which overrides num_hidden_layers (and layers_block_type) via model_kwargs to fit the model on a single GPU. This provides pre-merge coverage for large models that otherwise only run in multi-GPU CI. Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com>

galagam · 2026-04-23T17:41:46Z

/bot run --stage-list "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-Post-Merge-1,DGX_B200-8_GPUs-AutoDeploy-Post-Merge-1"

tensorrt-cicd · 2026-04-23T17:48:09Z

PR_Github #45235 [ run ] triggered by Bot. Commit: 2cb3001 Link to invocation

tensorrt-cicd · 2026-04-24T01:29:34Z

PR_Github #45235 [ run ] completed with state SUCCESS. Commit: 2cb3001
/LLM/main/L0_MergeRequest_PR pipeline #35496 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

galagam · 2026-04-24T09:18:16Z

/bot run

tensorrt-cicd · 2026-04-24T09:23:55Z

PR_Github #45379 [ run ] triggered by Bot. Commit: 2cb3001 Link to invocation

tensorrt-cicd · 2026-04-24T14:42:57Z

PR_Github #45379 [ run ] completed with state SUCCESS. Commit: 2cb3001
/LLM/main/L0_MergeRequest_PR pipeline #35620 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

Link to invocation

galagam · 2026-04-24T20:12:21Z

/bot run

tensorrt-cicd · 2026-04-24T20:19:07Z

PR_Github #45441 [ run ] triggered by Bot. Commit: 2cb3001 Link to invocation

tensorrt-cicd · 2026-04-24T22:53:31Z

PR_Github #45441 [ run ] completed with state SUCCESS. Commit: 2cb3001
/LLM/main/L0_MergeRequest_PR pipeline #35673 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

Link to invocation

galagam · 2026-04-26T05:44:20Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-04-26T05:51:21Z

PR_Github #45554 [ run ] triggered by Bot. Commit: 2cb3001 Link to invocation

tensorrt-cicd · 2026-04-26T10:08:42Z

PR_Github #45554 [ run ] completed with state SUCCESS. Commit: 2cb3001
/LLM/main/L0_MergeRequest_PR pipeline #35772 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

Link to invocation

galagam · 2026-04-26T16:09:28Z

/bot run

tensorrt-cicd · 2026-04-26T16:17:53Z

PR_Github #45592 [ run ] triggered by Bot. Commit: 2cb3001 Link to invocation

tensorrt-cicd · 2026-04-26T19:22:15Z

PR_Github #45592 [ run ] completed with state SUCCESS. Commit: 2cb3001
/LLM/main/L0_MergeRequest_PR pipeline #35808 completed with status: 'SUCCESS'

CI Report

Link to invocation

…pers (NVIDIA#12999) Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com>

galagam requested a review from greg-kwasniewski1 April 13, 2026 13:03

galagam requested a review from a team as a code owner April 13, 2026 13:03

github-actions Bot assigned galagam Apr 13, 2026

coderabbitai Bot reviewed Apr 13, 2026

View reviewed changes

galagam force-pushed the gagam-minor-refactor branch from 0be96a1 to 1e552c2 Compare April 19, 2026 07:05

galagam requested review from a team as code owners April 19, 2026 07:05

galagam mentioned this pull request Apr 19, 2026

[https://nvbugs/6070421][fix] Prefer pre-compiled cubin over NVRTC in trtllm-gen FMHA run() #12969

Closed

1 task

galagam added the gagam-slack-watch Label for getting slack notifications via GitHub Slack app label Apr 20, 2026

galagam added 2 commits April 23, 2026 10:20

[None][chore] AutoDeploy: FineGrained FP8 scale sharding helper

cfe0d8e

Minor refactor, follow up on NVIDIA#12593 Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com>

galagam force-pushed the gagam-minor-refactor branch from 75b1a71 to 2cb3001 Compare April 23, 2026 17:41

greg-kwasniewski1 approved these changes Apr 23, 2026

View reviewed changes

galagam merged commit db01f01 into NVIDIA:main Apr 26, 2026
5 checks passed

yufeiwu-nv pushed a commit to yufeiwu-nv/TensorRT-LLM that referenced this pull request May 19, 2026

[None][chore] AutoDeploy: Refactor finegrained FP8 scale sharding hel…

42dd33c

…pers (NVIDIA#12999) Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com>

Uh oh!

Conversation

galagam commented Apr 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

galagam commented Apr 13, 2026

Uh oh!

tensorrt-cicd commented Apr 13, 2026

Uh oh!

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

Uh oh!

Walkthrough

Changes

Estimated code review effort

❌ Failed checks (1 warning, 1 inconclusive)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

tensorrt-cicd commented Apr 13, 2026

Uh oh!

galagam commented Apr 19, 2026

Uh oh!

tensorrt-cicd commented Apr 19, 2026

Uh oh!

tensorrt-cicd commented Apr 19, 2026

Uh oh!

galagam commented Apr 19, 2026

Uh oh!

tensorrt-cicd commented Apr 19, 2026

Uh oh!

tensorrt-cicd commented Apr 19, 2026

Uh oh!

galagam commented Apr 20, 2026

Uh oh!

tensorrt-cicd commented Apr 20, 2026

Uh oh!

tensorrt-cicd commented Apr 20, 2026

Uh oh!

galagam commented Apr 20, 2026

Uh oh!

tensorrt-cicd commented Apr 20, 2026

Uh oh!

tensorrt-cicd commented Apr 20, 2026

Uh oh!

galagam commented Apr 20, 2026

Uh oh!

tensorrt-cicd commented Apr 20, 2026

Uh oh!

tensorrt-cicd commented Apr 20, 2026

Uh oh!

galagam commented Apr 20, 2026

Uh oh!

tensorrt-cicd commented Apr 20, 2026

Uh oh!

galagam commented Apr 23, 2026

Uh oh!

tensorrt-cicd commented Apr 23, 2026

Uh oh!

tensorrt-cicd commented Apr 23, 2026

Uh oh!

galagam commented Apr 23, 2026

Uh oh!

tensorrt-cicd commented Apr 23, 2026

Uh oh!

tensorrt-cicd commented Apr 23, 2026

Uh oh!

galagam commented Apr 23, 2026

Uh oh!

tensorrt-cicd commented Apr 23, 2026

Uh oh!

tensorrt-cicd commented Apr 23, 2026

Uh oh!

galagam commented Apr 23, 2026

Uh oh!

tensorrt-cicd commented Apr 23, 2026

galagam commented Apr 13, 2026 •

edited

Loading

coderabbitai Bot commented Apr 13, 2026 •

edited

Loading