[#14588][feat] AutoDeploy: DeepSeek-R1 optimization for low concurrency#14810
Merged
Conversation
taylor-yb-lee
force-pushed
the
taylor/dsr1_nvfp4_opt
branch
2 times, most recently
from
6172683 to
c967775
Compare
taylor-yb-lee
changed the title
taylor-yb-lee
force-pushed
the
taylor/dsr1_nvfp4_opt
branch
4 times, most recently
from
33d0a57 to
a31e9d1
Compare
Contributor
📝 WalkthroughWalkthroughAdds ChangesDeepSeek-R1 NVFP4-v2 support: custom ops, transforms, multi-stream scheduling, and model registry
Sequence Diagram(s)sequenceDiagram
participant ModelForward as DeepSeekV3Attention.forward
participant KVproj as kv_a_proj (computed first)
participant Qproj as q_a_proj (computed after KV)
participant AuxStream as Aux CUDA Stream
participant MLAop as trtllm_mla_with_cache
ModelForward->>KVproj: project hidden → compressed_kv + kpe
Note over KVproj,AuxStream: begin_aux_stream_passthrough (Pattern 2/3)
KVproj->>AuxStream: layernorm + reshape on aux stream
ModelForward->>Qproj: project hidden → q_nope + q_pe
AuxStream->>AuxStream: end_aux_stream_passthrough
Qproj->>MLAop: wait_aux_stream_passthrough before MLA
MLAop->>MLAop: fused MLA attention
sequenceDiagram
participant MoEForward as DeepSeekV3MoE.forward
participant RoutedExperts as EP-sharded routed experts
participant AllReduce as all_reduce(layer_type=moe)
participant SharedExpert as replicated shared expert
participant FusedNorm as trtllm_fused_allreduce_residual_rmsnorm
MoEForward->>RoutedExperts: compute routed expert output
MoEForward->>AllReduce: all_reduce(routed) via SYMM_MEM
MoEForward->>SharedExpert: compute shared expert (replicated)
AllReduce->>FusedNorm: routed_3d after AR
SharedExpert->>FusedNorm: shared + prev_residual
FusedNorm->>MoEForward: (normed_output, updated_residual)
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Suggested reviewers
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 11
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_attn.py (1)
716-724:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winUpdate the transform docstring to match the new fallback order.
The class docstring still says Pattern 0 falls straight through to Pattern 2, but
_apply()now tries Pattern 3 before Pattern 2. In a transform this pattern-heavy, stale ordering sends the next debugger to the wrong rewrite path.🤖 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/multi_stream_attn.py` around lines 716 - 724, Update the docstring of the multi-stream attention class to reflect the actual pattern trial order implemented in the _apply() method. The docstring currently states that Pattern 0 falls through to Pattern 2 if nothing is found, but the _apply() method now attempts Pattern 3 before Pattern 2. Revise the docstring to document the correct order: Pattern 0 is tried first, Pattern 3 is tried next, then Pattern 2, with Pattern 1 as the legacy fallback. This ensures the documented behavior matches the actual implementation and guides future debuggers through the correct rewrite path.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py (1)
199-207: ⚡ Quick winExtend bypass regression coverage to the other new
_fallback()branches.
tests/unittest/auto_deploy/multigpu/compile/test_bypass_captured_graphs.py:74-207covers thenot can_run_cuda_graph_allpath, but this change also routescg_batch_size is Noneandnot can_pad_allthrough_fallback(). Please add one scenario for each case so a future change cannot reintroduce mixed replay/eager execution on those branches.🤖 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/shim/ad_executor.py` around lines 199 - 207, Extend the bypass regression test coverage in the test file to include scenarios for the two new _fallback() branches introduced in ad_executor.py. Currently, only the not can_run_cuda_graph_all path is tested (lines 74-207). Add one test scenario that exercises the cg_batch_size is None branch and another test scenario that exercises the not can_pad_all branch, ensuring both paths properly fall back to eager execution instead of using captured graphs.
🤖 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.
Inline comments:
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/triton_routing.py`:
- Around line 257-289: The function ep_local_route_fn extracts K from
selected_experts.shape but does not validate that routing_weights has the same
shape. Add a shape validation check after the assertion and before the kernel
launch to ensure both selected_experts and routing_weights have identical
dimensions (T, K). This will prevent silent misindexing of routing_weights if
upstream code passes tensors with mismatched top_k values.
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/linear/linear.py`:
- Around line 80-86: The guard condition at lines 80-86 determines whether to
invoke the DSv3 fused GEMM operation but only checks the weight dtype and shape
along with whether bias is None. The targeted operation
torch.ops.auto_deploy.dsv3_fused_a_gemm requires bf16-precision inputs, but the
current guard does not validate input.dtype. Add a check for input.dtype ==
torch.bfloat16 to the conditional expression to prevent non-bf16 activations
from being routed to this path, which would cause a runtime failure.
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/linear/swiglu.py`:
- Around line 43-46: The exception handling in the try-except block that loads
the trtllm.silu_and_mul operation is too broad, catching all Exception types
instead of just the AttributeError that occurs when the operation is not
available. Replace the broad Exception catch with AttributeError specifically to
ensure only attribute lookup failures trigger the fallback to flashinfer,
preventing masked issues from other unexpected exceptions.
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/normalization/triton_rms_norm.py`:
- Around line 121-145: The row_stride is extracted only from hidden_states but
reused for all inputs and outputs in the add_rms_norm_kernel call. This causes
incorrect memory access if residual has a different row stride (stride(-2)) than
hidden_states. After checking and making hidden_states and residual contiguous
based on stride(-1), you must also verify that residual has the same row_stride
as hidden_states (check residual.stride(-2) against hidden_states.stride(-2)).
If the row strides don't match, make residual fully contiguous to align its
stride with hidden_states before extracting row_stride and calling the kernel.
In `@tensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek.py`:
- Around line 799-804: The code currently skips the _rope_deinterleave_load_hook
assuming that fuse_rope_into_trtllm_mla will always rewrite the graph, but
DeepSeekV3Attention.forward still calls torch_rope_with_explicit_cos_sin before
torch_mla. If the fusion is disabled or doesn't match the expected pattern, the
eager path will incorrectly apply RoPE rotation to native GPTJ-layout weights.
Either restore a de-interleaved fallback path in the model's load hook to handle
unfused execution safely, or add a runtime check that explicitly validates the
fused MLA rewrite was applied before forward pass execution to fail fast if
fusion is missing or disabled.
In
`@tensorrt_llm/_torch/auto_deploy/transform/library/fuse_moe_allreduce_residual_rmsnorm.py`:
- Around line 74-75: The `_trace_ar_path()` function extracts the all-reduce
strategy from the original node into the `_ar_strategy` variable, but when
emitting the fused operation, the code uses the global
`shared_config.dist_config.allreduce_strategy` instead. Replace the
`shared_config.dist_config.allreduce_strategy` reference on line 74 (and any
other location where the fused operation is created) with the extracted
`_ar_strategy` variable to preserve the original node's strategy in the fused
rewrite.
- Around line 42-56: The _trace_ar_path function extracts the structure of the
wait_aux_stream_passthrough node but does not preserve its kwargs, causing them
to be dropped during node rewriting and breaking stream ordering for graphs
using per-id events. Modify _trace_ar_path to extract the kwargs dictionary from
the node parameter (the wait_aux_stream_passthrough call_function node) and
include it in the return tuple alongside routed, view shape, and strategy. Then,
at the call site where the new wait node is created (referenced at line 111),
apply these extracted kwargs when constructing the new
wait_aux_stream_passthrough call to preserve the original event configuration.
In `@tensorrt_llm/_torch/auto_deploy/transform/library/fuse_rope_mla.py`:
- Around line 446-450: The exception handling in the GPTJ-layout detection block
starting at line 446 uses a broad except Exception clause that should be
narrowed to specific exception types. Examine the identical
factory._get_model_config() call at line 318 elsewhere in the file to see what
specific exceptions are already being caught there. Replace the broad except
Exception clause in the try-except block around _model_config assignment with
the same specific exception types used at line 318, maintaining the same logic
where already_gptj is set to False on exception.
In `@tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_attn.py`:
- Around line 415-439: The calls to begin_aux_stream_passthrough,
end_aux_stream_passthrough, and wait_aux_stream_passthrough are using the
default event_id=-1, which causes all MLA forks to share main/aux events instead
of getting unique per-fork event IDs from the new events_per_id pool. This
defeats the purpose of the new pooling mechanism and can cause duplicate-event
collapse across layers. Thread a stable per-fork event_id through all three
function calls in both affected locations (the primary rewrite at lines 415-439
involving kv_linear, last_kv_op, and mla_op, and the fused-path rewrite at lines
673-699). Use the same pattern for generating a unique, stable event ID per
match/fork that is already established elsewhere in the codebase to ensure the
events_per_id pool activates and maintains KV/Q overlap across layers.
In `@tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py`:
- Around line 274-285: The exclude_shard_node_filter configuration is documented
to match against weight/module names (like q_a_proj and kv_a_proj_with_mqa), but
the filtering logic around line 1429 checks against transform.target_node, which
contains FX-traced node names (torch_linear_*, torch_moe_*), causing the filter
to be silently ignored. Update the filtering logic to check against the actual
weight/module names instead of the FX node names. This may require accessing the
underlying module or weight name from the transform object rather than relying
on target_node for the comparison operation.
- Around line 2811-2832: The early return statement checking if
`num_simple_shards < len(a_proj_to_shard)` at lines 2828-2831 is preventing the
function from continuing to process q_b/kv_b/o_proj nodes when some a_proj nodes
were already sharded in prior passes. Remove this entire early abort condition
(the if block that returns 0) since the subsequent transform_container.add()
calls will safely handle duplicate nodes through no-op behavior, allowing the
full MLA processing pipeline to continue.
---
Outside diff comments:
In `@tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_attn.py`:
- Around line 716-724: Update the docstring of the multi-stream attention class
to reflect the actual pattern trial order implemented in the _apply() method.
The docstring currently states that Pattern 0 falls through to Pattern 2 if
nothing is found, but the _apply() method now attempts Pattern 3 before Pattern
2. Revise the docstring to document the correct order: Pattern 0 is tried first,
Pattern 3 is tried next, then Pattern 2, with Pattern 1 as the legacy fallback.
This ensures the documented behavior matches the actual implementation and
guides future debuggers through the correct rewrite path.
---
Nitpick comments:
In `@tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py`:
- Around line 199-207: Extend the bypass regression test coverage in the test
file to include scenarios for the two new _fallback() branches introduced in
ad_executor.py. Currently, only the not can_run_cuda_graph_all path is tested
(lines 74-207). Add one test scenario that exercises the cg_batch_size is None
branch and another test scenario that exercises the not can_pad_all branch,
ensuring both paths properly fall back to eager execution instead of using
captured graphs.
🪄 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: Enterprise
Run ID: 22572f7e-b2bb-46ae-be1e-6c66295a15fd
📒 Files selected for processing (28)
examples/auto_deploy/model_registry/configs/deepseek-r1.yamlexamples/auto_deploy/model_registry/models.yamltensorrt_llm/_torch/auto_deploy/config/default.yamltensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/triton_routing.pytensorrt_llm/_torch/auto_deploy/custom_ops/linear/__init__.pytensorrt_llm/_torch/auto_deploy/custom_ops/linear/dsv3_fused_a_gemm.pytensorrt_llm/_torch/auto_deploy/custom_ops/linear/dsv3_router_gemm.pytensorrt_llm/_torch/auto_deploy/custom_ops/linear/linear.pytensorrt_llm/_torch/auto_deploy/custom_ops/linear/swiglu.pytensorrt_llm/_torch/auto_deploy/custom_ops/mla/trtllm_mla.pytensorrt_llm/_torch/auto_deploy/custom_ops/normalization/triton_rms_norm.pytensorrt_llm/_torch/auto_deploy/models/custom/mla_rope_utils.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek.pytensorrt_llm/_torch/auto_deploy/shim/ad_executor.pytensorrt_llm/_torch/auto_deploy/transform/library/fuse_moe_allreduce_residual_rmsnorm.pytensorrt_llm/_torch/auto_deploy/transform/library/fuse_rope_mla.pytensorrt_llm/_torch/auto_deploy/transform/library/fused_add_rms_norm.pytensorrt_llm/_torch/auto_deploy/transform/library/fused_moe.pytensorrt_llm/_torch/auto_deploy/transform/library/fusion.pytensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_attn.pytensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.pytensorrt_llm/_torch/auto_deploy/transform/library/sharding.pytensorrt_llm/_torch/auto_deploy/utils/multi_stream_utils.pytests/integration/defs/accuracy/references/gsm8k.yamltests/integration/defs/accuracy/test_llm_api_autodeploy.pytests/integration/test_lists/qa/llm_function_core.txttests/integration/test_lists/test-db/l0_dgx_b200.ymltests/test_common/llm_data.py
1 task
taylor-yb-lee
changed the title
Collaborator
Author
|
/bot run --disable-fail-fast --stage-list "DGX_B200-8_GPUs-AutoDeploy-Post-Merge-1" |
Collaborator
|
PR_Github #54103 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #54103 [ run ] completed with state |
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #54150 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #54150 [ run ] completed with state
|
jieli-matrix
approved these changes
Jun 16, 2026
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #55000 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #55000 [ run ] completed with state
|
Collaborator
|
Could you add the perf number in the comment or PR description? |
Collaborator
Author
Added in the description. Thanks! |
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #57039 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #57039 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
… OOM) gemma3n_e2b_it.yaml was the only registry config relying on the AutoDeploy default free_gpu_memory_fraction=0.9. Gemma-3n uses VSWA + AltUp (4 parallel hidden streams): under the VSWA scheme the KV manager sizes the pool purely from free_gpu_memory_fraction and ignores max_num_tokens, so 0.9 grabs ~55 GiB of KV on an 80 GiB GPU. The AltUp prefill activation peak then has too little headroom and OOMs mid-run during long-output generation (GSM8K); the executor recovery path turns that OOM into a hang, so the accuracy test is reported as Timeout. Cap free_gpu_memory_fraction at 0.8 (matching the gemma4_e2b sibling). KV cache 55.37 -> 49.21 GiB; test_gemma3n_e2b_it now passes end-to-end locally on an H100 80GB (MMLU 59.36 >= 57.71, GSM8K 72.25 >= 68.97). Signed-off-by: Taylor Yeonbok Lee <249374542+taylor-yb-lee@users.noreply.github.com>
Collaborator
Author
|
/bot run |
1 similar comment
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #57219 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #57219 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #57234 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #57234 [ run ] completed with state
|
Signed-off-by: Taylor Yeonbok Lee <249374542+taylor-yb-lee@users.noreply.github.com>
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #57243 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #57243 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #57329 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #57329 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #57362 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #57362 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #57513 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #57513 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #57630 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #57630 [ run ] completed with state |
BrianLi23
pushed a commit
to BrianLi23/TensorRT-LLM
that referenced
this pull request
Jul 9, 2026
…currency (NVIDIA#14810) Signed-off-by: Taylor Yeonbok Lee <249374542+taylor-yb-lee@users.noreply.github.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary by CodeRabbit
Release Notes
New Features
Performance Improvements
Description
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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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.