[#12784][feat] AutoDeploy: Optimize DeepSeek-R1 model performance

[taylor-yb-lee]

Summary by CodeRabbit

Description

• New Features
  • Enabled multi stream MLA/MoE for DeepSeekR1 model.
  • Enabled symm_mem AllGather which is faster than NCCL allgather
  • Enabled linear ops to use DeepGemm (b projections on MLA) on Blackwell (SM100f) so that we can optimize out redundant typecasts (Used for MLA and MLP)
  • Applied AllReduce + Add. + RMS fusion for attention
  • Enabled mlir eltwise fusion and applied required fixes
  • Update DeepSeek-R1 config

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.

[coderabbitai]

📝 Walkthrough

Walkthrough

The changes introduce specialized support for Blackwell (SM100f) GPU architecture with early detection paths in FP8 quantization, implement TMA-aligned tensor layout utilities for kernel fusion, refactor output-splitting logic to use torch.narrow with metadata propagation, extend quantization workflows with post-load hooks for format conversion, and enable new compilation modes in configuration. Comprehensive tests validate metadata preservation and numerical correctness.

Changes

Cohort / File(s)	Summary
Configuration Updates examples/auto_deploy/model_registry/configs/deepseek-r1.yaml	Switched runtime backends (trtllm, torch-cudagraph), expanded batch sizes and CUDA graph configurations, added KV cache tuning, enabled additional graph transformation passes, and introduced piecewise model compilation with token boundaries.
SM100F Quantization Support tensorrt_llm/_torch/auto_deploy/custom_ops/quantization/torch_quant.py	Added early Blackwell (SM100f) detection with direct dispatch to fp8_swap_ab_gemm when weight_scale.dtype == torch.int, bypassing existing casting and block-size fallback logic for this GPU variant.
Tensor Layout & Fusion Utilities tensorrt_llm/_torch/auto_deploy/transform/library/fuse_swiglu.py, tensorrt_llm/_torch/auto_deploy/transform/library/fusion.py	Introduced _ensure_tma_col_major helper for TMA-aligned UE8M0 tensor re-materialization; refactored output splitting from operator.getitem to torch.narrow with per-segment metadata propagation and contiguous-call logic updates.
FP8 Quantization Post-Load Workflow tensorrt_llm/_torch/auto_deploy/transform/library/quantization.py	Added conditional imports and post-load hooks for SM100F that detect 128×128 weight blocks, resmooth weights/scales to UE8M0, and apply DeepGEMM-required layout transformations in-place during module loading.
Integration & Correctness Tests tests/unittest/auto_deploy/singlegpu/transformations/library/test_gemm_fusion.py	Added parametrized test suite validating metadata preservation across fusion transformations, shape validity reporting, contiguous-call counts, layout restoration, and numerical equivalence for both standard and FP8 quantized scenarios.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

🚥 Pre-merge checks | ✅ 1 | ❌ 2

❌ Failed checks (1 warning, 1 inconclusive)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 51.52% 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 includes bullet points of new features and mentions test coverage and checklist, but lacks detailed explanation of the issue and solution.	Expand the Description section with more detail about the specific issues being solved and how the changes address them. Clarify which tests safeguard the changes in the 'Test Coverage' section.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Title check	✅ Passed	The PR title clearly identifies the main optimization focus (DeepSeek-R1 model performance) and includes proper ticket reference and feature type.

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

Actionable comments posted: 4

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (2)

tensorrt_llm/_torch/auto_deploy/transform/library/fusion.py (2)

581-584: ⚠️ Potential issue | 🟠 Major

Restore TMA layout after concatenating UE8M0 fine-grained scales.

On SM100f, FineGrainedFP8LinearQuantization.post_load_hook() converts weight_scale_inv to TMA-aligned torch.int, and trtllm_finegrained_fp8_linear() switches to fp8_swap_ab_gemm based only on that dtype. torch.cat(...) here drops the column-major/TMA layout, so the fused path can feed row-major UE8M0 scales straight into DeepGEMM. This needs the same layout re-materialization you added in fuse_swiglu.py, or the fused path should opt out of the DeepGEMM fast-path.

🤖 Prompt for AI Agents

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

In `@tensorrt_llm/_torch/auto_deploy/transform/library/fusion.py` around lines 581
- 584, The concatenation of weight_scale_inv via torch.cat drops the TMA-aligned
column-major layout produced by FineGrainedFP8LinearQuantization.post_load_hook,
so update the fused path after fused_weight_scale_inv =
torch.cat(weight_scale_inv, dim=0) to re-materialize the original
TMA/column-major layout and torch.int dtype (same approach used in
fuse_swiglu.py) so trtllm_finegrained_fp8_linear still sees the TMA-aligned
torch.int and chooses fp8_swap_ab_gemm; alternatively, if re-materializing is
undesirable, mark the fused path to opt out of the DeepGEMM fast-path by
ensuring the layout/dtype check used by trtllm_finegrained_fp8_linear fails.

---

442-446: ⚠️ Potential issue | 🟠 Major

Don't unconditionally claim shape metadata is valid.

_insert_fused_gemm() and _insert_fused_quant_gemm() only populate meta["val"] when the original linear node already had it. If a mixed-children fusion runs without that metadata, the new torch.narrow / contiguous nodes are still shape-less, but has_valid_shapes=True suppresses the fallback shape-prop pass and leaves downstream shape-based transforms with incomplete metadata.

🤖 Prompt for AI Agents

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

In `@tensorrt_llm/_torch/auto_deploy/transform/library/fusion.py` around lines 442
- 446, The TransformInfo currently sets has_valid_shapes=True unconditionally
which hides missing shape metadata created by _insert_fused_gemm() and
_insert_fused_quant_gemm(); change the code that builds the TransformInfo so
has_valid_shapes is computed (not hard-coded) by inspecting the fused outputs'
metadata (e.g., check meta.get("val") presence on the newly created nodes such
as the torch.narrow/contiguous results or verify that the original linear nodes
provided shape meta) and set has_valid_shapes=False if any fused node lacks
meta["val"] so the fallback shape-prop pass can run; use the same symbol names
(_insert_fused_gemm, _insert_fused_quant_gemm, TransformInfo, torch.narrow,
contiguous) to locate where to compute this flag.

🤖 Prompt for all review comments with AI agents

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

Inline comments:
In `@examples/auto_deploy/model_registry/configs/deepseek-r1.yaml`:
- Around line 35-37: The example config re-enables the piecewise CUDA-graph
splitting; change the compile_model setting so piecewise_enabled is false to
match the rollback and keep the regressing path disabled, leaving
piecewise_num_tokens intact for future use (update the compile_model:
piecewise_enabled value only).

In `@tensorrt_llm/_torch/auto_deploy/transform/library/quantization.py`:
- Around line 46-53: The imported FP8 helpers resmooth_to_fp8_e8m0 and
transform_sf_into_required_layout may be None; update the SM100f post-load hook
(the function that currently calls these helpers unconditionally) to check for
their presence before invoking them and either skip the FP8 conversion path or
raise a clear, descriptive error; apply the same None-guarding pattern to the
other occurrence around lines 934-946 so both call sites test "if
resmooth_to_fp8_e8m0 is not None and transform_sf_into_required_layout is not
None" (or handle each helper individually) and log or fallback appropriately
instead of calling a NoneType.
- Around line 948-957: The code currently replaces the quantization scale slot
with an nn.Parameter (using setattr on scale_attr), but weight_scale_inv must
remain a buffer so later fusion code (gm.get_buffer(...) in fusion.py) can find
it; modify the transform in the block that sets scale_attr so that when
scale_attr corresponds to the weight_scale_inv buffer you call
target_module.register_buffer(scale_attr, transformed_scale) (ensuring the
tensor is detached/converted to the correct device/dtype) instead of setting an
nn.Parameter, while leaving the weight replacement (attr_name) as an
nn.Parameter as before.

In
`@tests/unittest/auto_deploy/singlegpu/transformations/library/test_gemm_fusion.py`:
- Around line 643-645: In the forward method(s) in test_gemm_fusion.py where the
shape is unpacked (e.g., the forward function that sets batch_size, seq_len, _ =
x.shape), rename the intentionally unused bindings to start with an underscore
(e.g., _batch_size, _seq_len) so Ruff stops flagging them as unused; apply the
same change to the other occurrence(s) of the same unpacking later in the file
(the forward that currently uses batch_size and seq_len) by prefixing those
names with an underscore.

---

Outside diff comments:
In `@tensorrt_llm/_torch/auto_deploy/transform/library/fusion.py`:
- Around line 581-584: The concatenation of weight_scale_inv via torch.cat drops
the TMA-aligned column-major layout produced by
FineGrainedFP8LinearQuantization.post_load_hook, so update the fused path after
fused_weight_scale_inv = torch.cat(weight_scale_inv, dim=0) to re-materialize
the original TMA/column-major layout and torch.int dtype (same approach used in
fuse_swiglu.py) so trtllm_finegrained_fp8_linear still sees the TMA-aligned
torch.int and chooses fp8_swap_ab_gemm; alternatively, if re-materializing is
undesirable, mark the fused path to opt out of the DeepGEMM fast-path by
ensuring the layout/dtype check used by trtllm_finegrained_fp8_linear fails.
- Around line 442-446: The TransformInfo currently sets has_valid_shapes=True
unconditionally which hides missing shape metadata created by
_insert_fused_gemm() and _insert_fused_quant_gemm(); change the code that builds
the TransformInfo so has_valid_shapes is computed (not hard-coded) by inspecting
the fused outputs' metadata (e.g., check meta.get("val") presence on the newly
created nodes such as the torch.narrow/contiguous results or verify that the
original linear nodes provided shape meta) and set has_valid_shapes=False if any
fused node lacks meta["val"] so the fallback shape-prop pass can run; use the
same symbol names (_insert_fused_gemm, _insert_fused_quant_gemm, TransformInfo,
torch.narrow, contiguous) to locate where to compute this flag.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 5438b43f-9552-4332-930c-c6f38a64bef2

📥 Commits

Reviewing files that changed from the base of the PR and between a0b53e6 and fdb717b.

📒 Files selected for processing (6)

• examples/auto_deploy/model_registry/configs/deepseek-r1.yaml
• tensorrt_llm/_torch/auto_deploy/custom_ops/quantization/torch_quant.py
• tensorrt_llm/_torch/auto_deploy/transform/library/fuse_swiglu.py
• tensorrt_llm/_torch/auto_deploy/transform/library/fusion.py
• tensorrt_llm/_torch/auto_deploy/transform/library/quantization.py
• tests/unittest/auto_deploy/singlegpu/transformations/library/test_gemm_fusion.py

[taylor-yb-lee]

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

[tensorrt-cicd]

PR_Github #44763 [ run ] triggered by Bot. Commit: ed4e052 Link to invocation

[tensorrt-cicd]

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

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

[tensorrt-cicd]

PR_Github #45908 [ run ] triggered by Bot. Commit: 4550f61 Link to invocation

[taylor-yb-lee]

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

[tensorrt-cicd]

PR_Github #46542 [ run ] triggered by Bot. Commit: c1943cd Link to invocation

[taylor-yb-lee]

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

[tensorrt-cicd]

PR_Github #46546 [ run ] triggered by Bot. Commit: 7e0e791 Link to invocation

[tensorrt-cicd]

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