[https://nvbugs/6240561][fix] Fix AutoDeploy DeepSeek-R1 accuracy drop by taylor-yb-lee · Pull Request #14793 · NVIDIA/TensorRT-LLM

taylor-yb-lee · 2026-05-31T02:15:10Z

Fixed accuracy drop in DeepSeek-R1-0528 AutoDeploy GSM8K 90.0 → 95.489 after PR #13478

Bug: After PR13478, kv_a in MLA became unsharded (i.e., replicated) and a hidden issue on the finegrained fp8 bf16 fallback dequant happened.
Before PR13478, kv_a weight was shareded as 72x7168 per rank, and kv-a's scale was per row, so the dequant_block_size was 1 and it was okay.
However after PR13478, kv_a is to be unsharded, (i.e., replicated), and for a weight 576x7168, kv_a_scale was to be a coarse block size (5 rows), and then the calculation of the dequant_block became wrong, i.e., 576/5 = 116, while it should be 128.
Fix: use the canonical 128 when the scale grid is coarse
Impact: DeepSeek-R1-0528 AutoDeploy GSM8K 90.0 → 95.489 (kv_a left replicated, multi_stream + deepgemm on).

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.

coderabbitai · 2026-05-31T05:34:43Z

📝 Walkthrough

Walkthrough

This PR updates MLA latent down-projections in DeepSeek to use simple_shard mode (colwise sharding + all_gather) instead of replication, and extends the sharding IR to implement simple-shard for linear layers with FP8 scale support using centralized FP8 logic and backend-aware distribution selection.

Changes

DeepSeek MLA Simple-Shard and IR Support

Layer / File(s)	Summary
DeepSeek MLA sharding-hint updates `tensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek.py`	Module and class documentation, forward-pass comments, and operator hints for `q_a_proj` and `kv_a_proj_with_mqa` are updated from `tp_mode="none"` (replicated) to `tp_mode="simple_shard"` (colwise + all_gather), with corresponding comment adjustments in the Q and KV projection paths.
FP8 scale sharding infrastructure refactoring `tensorrt_llm/_torch/auto_deploy/transform/library/sharding_ir.py`	`FineGrainedFP8WeightShardingInfo` is imported and `_split_fp8_block_scale` is refactored to delegate scale splitting to the centralized sharding info logic, accepting a `weight_n` parameter to handle misaligned FP8 scale/weight relationships.
LinearShardableNode simple_shard implementation and all_gather `tensorrt_llm/_torch/auto_deploy/transform/library/sharding_ir.py`	`LinearShardableNode.apply` gains a `tp_mode == "simple_shard"` branch that column-shards linear weights and biases, shards FP8 scales, and inserts a backend-aware `all_gather` node to restore full activations; `_apply_simple_shard` is updated to select `all_gather` via `_get_dist_ops("auto")` instead of hardcoding torch backend.
FP8 Linear and SwiGLU scale sharding for simple_shard `tensorrt_llm/_torch/auto_deploy/transform/library/sharding_ir.py`	`FineGrainedFP8LinearShardableNode._shard_scales` computes `weight_n` from the weight tensor's sharded dimension and passes it to `_split_fp8_block_scale`; `FineGrainedFP8SwiGLUShardableNode._shard_scales` maps each scale to its corresponding weight tensor and computes per-scale `weight_n` for accurate expansion.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~75 minutes

Possibly related PRs

NVIDIA/TensorRT-LLM#13478: Both PRs modify DeepSeekV3Attention's MLA sharding hints and tp_mode setup for latent down-projections.

Suggested reviewers

galagam
bmarimuthu-nv
MrGeva
tcherckez-nvidia

🚥 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 title clearly and specifically identifies the bug being fixed (DeepSeek-R1 accuracy drop) with proper NVBugs ticket reference and [fix] type tag.
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.
Description check	✅ Passed	PR description provides detailed context on the bug, root cause, and fix, though Test Coverage section is missing.

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

…tions (q_a/kv_a) via YAML hint Restores the DSR1 GSM8K accuracy that dropped around 92c5030 ("Switch ... to sharding-IR canonical models"): 90.0 -> 95.0 (>= legacy 94.92). - Before 92c5030: the legacy heuristic sharder simple-sharded MLA q_a/kv_a, so their misaligned per-rank FP8 scale hit the kernel BF16 dequant fallback -- the latent projections effectively ran in BF16. - After 92c5030 (sharding-IR): the canonical model hints q_a/kv_a tp_mode="none" (replicate), so they ran in true FP8. This side effect of the new sharding is the ~5-point regression. - Fix (no model edit): add a tp_mode="simple_shard" mode (colwise + all_gather) and a YAML-driven tp_mode_overrides (fnmatch weight-name -> tp_mode) to apply_sharding_hints. deepseek-r1.yaml maps q_a/kv_a to simple_shard; the model keeps tp_mode="none" so other models/linears are unaffected. Also adds per-row FP8 scale expansion (fixes kv_a N=576 -> 5-row scale < world_size IndexError) and backend-aware all_gather. - Gotcha: only manifests with deepgemm ON; forcing those linears to BF16 via the finegrained kernel gate needed deepgemm-off, which dropped the baseline and confounded the bisection. MLA latent down-projection (q_a_proj N=1536, kv_a_proj N=576), FP8 weights: tp_mode="none" (replicate): full weight on every rank --> 128-aligned --> true FP8 tp_mode="simple_shard": colwise split --> per-rank 192/72 rows (not 128-mult) --> BF16 fallback --> all_gather Signed-off-by: Taylor Yeonbok Lee <249374542+taylor-yb-lee@users.noreply.github.com>

galagam

LGTM, but I would rather @greg-kwasniewski1 also takes a look

…ment for non-128-multiple N - BF16-fallback dequant expanded block scales by ceil(N/scale_n); for N not a multiple of 128 (MLA kv_a N=576 -> 116) the 128-block scales land on the wrong rows -> ~40% weight error. - Use the canonical 128 when the scale grid is coarse (scale_n == ceil(N/128)); keep ceil only for per-row scales (scale_n == N) produced by TP sharding. - DSR1-0528 GSM8K 90.0 -> 95.489 with kv_a left replicated; no sharding, no model change, perf-neutral. - Reverts the prior simple_shard workaround (sharding_ir.py + deepseek-r1.yaml) now that the root kernel-dequant bug is fixed generally. Scale row -> weight row mapping (kv_a N=576, 5 scale rows): before block_n=ceil(576/5)=116: s0->[0:116] s1->[116:232] s2->[232:348] ... (mis-aligned) after block_n=128 : s0->[0:128] s1->[128:256] ... s4->[512:576] (correct) Signed-off-by: Taylor Yeonbok Lee <249374542+taylor-yb-lee@users.noreply.github.com>

taylor-yb-lee · 2026-05-31T22:33:26Z

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

tensorrt-cicd · 2026-05-31T22:38:45Z

PR_Github #51265 [ run ] triggered by Bot. Commit: 6b86e2d Link to invocation

tensorrt-cicd · 2026-06-01T03:59:04Z

PR_Github #51265 [ run ] completed with state ABORTED. Commit: 6b86e2d

Link to invocation

taylor-yb-lee · 2026-06-01T19:15:27Z

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

tensorrt-cicd · 2026-06-01T19:21:22Z

PR_Github #51408 [ run ] triggered by Bot. Commit: 6b86e2d Link to invocation

tensorrt-cicd · 2026-06-01T21:29:30Z

PR_Github #51408 [ run ] completed with state FAILURE. Commit: 6b86e2d
/LLM/main/L0_MergeRequest_PR pipeline #40817 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

taylor-yb-lee · 2026-06-01T23:07:09Z

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

tensorrt-cicd · 2026-06-01T23:12:42Z

PR_Github #51439 [ run ] triggered by Bot. Commit: 6b86e2d Link to invocation

tensorrt-cicd · 2026-06-02T01:09:00Z

PR_Github #51439 [ run ] completed with state SUCCESS. Commit: 6b86e2d
/LLM/main/L0_MergeRequest_PR pipeline #40848 completed with status: 'SUCCESS'

CI Report

Link to invocation

github-actions Bot assigned taylor-yb-lee May 31, 2026

taylor-yb-lee force-pushed the taylor/dsr_acc branch from e749a8c to 4cfca57 Compare May 31, 2026 02:53

taylor-yb-lee marked this pull request as ready for review May 31, 2026 05:30

taylor-yb-lee requested a review from a team as a code owner May 31, 2026 05:30

taylor-yb-lee requested a review from galagam May 31, 2026 05:30

galagam requested a review from greg-kwasniewski1 May 31, 2026 05:49

galagam reviewed May 31, 2026

View reviewed changes

taylor-yb-lee force-pushed the taylor/dsr_acc branch from 4cfca57 to 7713e3b Compare May 31, 2026 08:01

taylor-yb-lee requested a review from galagam May 31, 2026 08:43

taylor-yb-lee force-pushed the taylor/dsr_acc branch from 7713e3b to 1a6260a Compare May 31, 2026 08:44

taylor-yb-lee commented May 31, 2026

View reviewed changes

Comment thread tests/unittest/auto_deploy/singlegpu/utils/test_fp8_dequant.py Outdated

taylor-yb-lee force-pushed the taylor/dsr_acc branch from 1a6260a to f8651ed Compare May 31, 2026 08:49

taylor-yb-lee force-pushed the taylor/dsr_acc branch from f8651ed to 6b86e2d Compare May 31, 2026 08:51

galagam approved these changes May 31, 2026

View reviewed changes

taylor-yb-lee merged commit c482b81 into NVIDIA:main Jun 2, 2026
7 checks passed

taylor-yb-lee mentioned this pull request Jun 2, 2026

[https://nvbugs/6240561][fix] Unwaive DeepSeek R1 accuracy test #14870

Merged

1 task

coderabbitai Bot mentioned this pull request Jun 10, 2026

[None][feat] AutoDeploy: Qwen3.5: Apply whielist based sharding and apply lm_head sharding #15185

Merged

1 task

Uh oh!

Conversation

taylor-yb-lee commented May 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

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

Uh oh!

Walkthrough

Changes

Estimated code review effort

Possibly related PRs

Suggested reviewers

❌ Failed checks (1 warning)

Uh oh!

galagam left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

taylor-yb-lee commented May 31, 2026

Uh oh!

tensorrt-cicd commented May 31, 2026

Uh oh!

tensorrt-cicd commented Jun 1, 2026

Uh oh!

taylor-yb-lee commented Jun 1, 2026

Uh oh!

tensorrt-cicd commented Jun 1, 2026

Uh oh!

tensorrt-cicd commented Jun 1, 2026

Uh oh!

taylor-yb-lee commented Jun 1, 2026

Uh oh!

tensorrt-cicd commented Jun 1, 2026

Uh oh!

tensorrt-cicd commented Jun 2, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

taylor-yb-lee commented May 31, 2026 •

edited

Loading

coderabbitai Bot commented May 31, 2026 •

edited

Loading