[#13718][feat] AutoDeploy MoE all-to-all: cache + runtime max-tokens

[greg-kwasniewski1]

Summary

Closes #13718.
Closes #13360.

Implementation

• Module-level _GlobalMoeA2ACache in tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/trtllm_moe.py caches MoeAlltoAll instances and parsed DistConfig mappings, eliminating per-call Python ctor + JSON deserialization (mirrors _GlobalTrtllmPlanner in trtllm_attention.py).
• Runtime max-tokens-per-rank via BatchInfo slot 13: pre-forward tp_allgather in ADEngine writes the cross-rank max(total_num_tokens) to the slot; the MoE custom op reads it at capture/eager time. Replaces the static max_num_tokens over-padding with the per-iteration cross-rank max — matches base TRT-LLM's runtime_max_tokens_per_rank from fused_moe_cutlass.py:764.

Files touched:

• tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py — BatchInfo._NUM_ELEMENTS = 14, slot 13 helpers
• tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py — pre-forward tp_allgather override
• tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py — thread batch_info_host graph input into MoE op call sites
• tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/trtllm_moe.py — _GlobalMoeA2ACache + slot 13 read in dispatch path
• tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/torch_moe.py — batch_info_host kwarg threading for sharding-time targets

Why this is safe under CUDA-graph capture

The int(batch_info_host[13].item()) read is on a pinned-host CPU tensor — pure Python read, no GPU sync. Under graph capture the value is baked in at capture time, so each captured graph corresponds to the cg_batch_size bucket it was warmed at. AD's existing maybe_pad_for_cuda_graph (ad_executor.py) syncs batch_size across ranks via tp_allgather and pads up to cg_batch_size — so for any decode replay path, all DP ranks have the same total_num_tokens, and the captured constant equals the runtime cross-rank max.

For eager paths (prefill, mixed batches, fallback when sync gate fails), the .item() re-runs Python and reads the just-written cross-rank max — also correct.

Test plan

• CI: /bot run --extra-stage "DGX_H100-4_GPUs-AutoDeploy-1, DGX_B200-4_GPUs-AutoDeploy-1" — multi-GPU AutoDeploy stages exercise this path
• Local correctness: pytest tests/integration/defs/accuracy/test_llm_api_autodeploy.py -k "attn_dp" (attention-DP variants)
• Local perf: bench-sweep comparing main vs this branch on Nemotron-Nano + attention-DP — expect dispatch padding to shrink from max_num_tokens=2048 to runtime_max ≈ batch_size * (1 + draft_len)

Validation already done locally

• tests/unittest/auto_deploy/multigpu/transformations/library/test_ep_sharding.py::test_ep_shard (4/4 PASS) on the cache-only stage
• Smoke run: examples/auto_deploy/build_and_run_ad.py --yaml-extra _repro_logs/nano_attn_dp_smoke.yaml on Nemotron-Nano-30B-FP8 + attn-DP + 4×H100, 10/10 prompts, coherent generation, no NaN. With diagnostic instrumentation:
  • Override branch entered on all 4 ranks (enable_attention_dp=True, tp_size=4)
  • tp_allgather returns proper 4-element list (e.g., [53, 56, 56, 77]); all ranks compute and write the same max=77
  • MoE op reads the post-override value (e.g., all ranks read slot13=77 while their local_num_tokens differ: 53/56/56/77)

Summary by CodeRabbit

Release Notes

• New Features

  • Added support for distributed parallel token tracking across ranks in batch metadata, enabling improved handling in distributed inference scenarios.
  • Extended MoE custom operations to support optional batch information from the host, enabling more efficient all-to-all communication patterns.

• Performance

  • Implemented caching mechanism for MoE all-to-all instances to reduce redundant initialization overhead.
  • Optimized distributed inference by computing cross-rank maximum tokens for better resource utilization.

[greg-kwasniewski1]

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

[tensorrt-cicd]

PR_Github #46644 [ run ] triggered by Bot. Commit: a17bc36 Link to invocation

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR implements a stateful MoE all-to-all cache and runtime max-tokens feature. BatchInfo gains a 14th element storing DP-aware max tokens. MoeAlltoAll instances are now cached per-constructor-key instead of re-instantiated per layer call. Cross-rank token counts are gathered in the AD executor and passed to MoE ops to replace static config-level padding.

Changes

DP-Aware Batch Info and MoE All-to-All Caching

Layer / File(s)	Summary
Data Shape tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py	BatchInfo._NUM_ELEMENTS increases from 13 to 14; slot 13 holds max_dp_num_tokens. New update_max_dp_num_tokens() and get_max_dp_num_tokens() methods added. SequenceInfo.nest_sequences() and switch_to_generate_() initialize slot 13 with local token counts.
Core Logic tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/trtllm_moe.py	Introduces _MoeA2ACache singleton to cache parsed DistConfig/Mapping and lazily constructed MoeAlltoAll instances keyed by constructor parameters. _run_moe_with_alltoall() and _run_trtllm_gen_nvfp4_moe_with_alltoall() now accept batch_info_host, derive runtime_max_tokens_per_rank from batch_info_host[13] (fallback to max_num_tokens), and retrieve cached MoeAlltoAll instances instead of reconstructing them per call.
Aggregation tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py	ADEngine.forward() now computes cross-rank max total_num_tokens via tp_allgather when Attention-DP is enabled and TP size > 1, then updates batch_info via update_max_dp_num_tokens().
Wiring tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/torch_moe.py, tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/trtllm_moe.py	All MoE custom op signatures and fake variants extended with batch_info_host: Optional[torch.Tensor] | None = None parameter. All-to-all enabled code paths forward batch_info_host into corresponding runners.
Integration tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py	ShardingTransformExecutor._apply() ensures batch_info_host placeholder node exists in graph before EP transforms. _insert_sharded_moe() retrieves the placeholder and passes it to MoE op wiring so all-to-all path can consume runtime max tokens when available.

Sequence Diagram

sequenceDiagram
    participant AD as ADEngine
    participant Gath as tp_allgather
    participant Batch as BatchInfo
    participant Cache as _GlobalMoeA2ACache
    participant Runner as _run_moe_with_alltoall
    participant A2A as MoeAlltoAll

    AD->>AD: Collect local_num_tokens per rank
    AD->>Gath: tp_allgather(local_num_tokens)
    Gath-->>AD: all_rank_num_tokens[]
    AD->>AD: max_dp = max(all_rank_num_tokens)
    AD->>Batch: update_max_dp_num_tokens(max_dp)
    
    Note over Batch: batch_info_host[13] = max_dp

    AD->>Runner: Call with batch_info_host tensor
    Runner->>Runner: Extract runtime_max = batch_info_host[13]
    Runner->>Cache: get_alltoall(mapping, ...)
    alt Cache Hit
        Cache-->>A2A: Return cached instance
    else Cache Miss
        Cache->>A2A: Create new MoeAlltoAll(...)
        Cache-->>A2A: Cache & return
    end
    
    Runner->>A2A: run_dispatch(runtime_max_tokens_per_rank=runtime_max)
    A2A-->>Runner: Padded output
    Runner-->>AD: Result

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~50 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% 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 concisely summarizes the main changes: AutoDeploy MoE all-to-all optimization through caching and runtime max-tokens, directly addressing the feature described in #13718.
Linked Issues check	✅ Passed	All coding requirements from issue #13718 are met: Phase 1 (_GlobalMoeA2ACache singleton to eliminate per-call ctor overhead) and Phase 2 (runtime max-tokens via BatchInfo slot 13 with tp_allgather) are fully implemented with proper CUDA-graph safety handling.
Out of Scope Changes check	✅ Passed	All changes are directly aligned with #13718 objectives: attention_interface.py (BatchInfo slot 13), ad_executor.py (tp_allgather override), sharding.py (batch_info_host threading), trtllm_moe.py (cache + slot read), torch_moe.py (kwarg threading). No unrelated modifications detected.
Description check	✅ Passed	PR description comprehensively covers implementation details, design rationale, safety considerations, and validation results across all required template sections.

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

🧹 Nitpick comments (2)

tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py (1)

966-978: ⚡ Quick win

Add targeted unit coverage for the slot-13 override gate.

Please add tests for both branches: (enable_attention_dp=True, tp_size>1) writes cross-rank max, and fallback path keeps local total. This is the key correctness contract for the new runtime padding behavior.

🤖 Prompt for AI Agents

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

In `@tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py` around lines 966 - 978,
Add unit tests in the ad_executor test suite that exercise the slot-13 override
gate around attention DP: create one test where
ad_executor.enable_attention_dp=True and ad_executor.dist_config.tp_size>1 and
mock ad_executor.dist.tp_allgather to return different per-rank totals so that
cache_seq_interface.info.batch_info.update_max_dp_num_tokens is called with the
cross-rank max (verify call and stored value using the batch_info getter);
create another test where either enable_attention_dp=False or tp_size==1 and
assert that update_max_dp_num_tokens is not invoked and the original local total
from cache_seq_interface.info.batch_info.get_total_num_tokens remains unchanged.
Ensure mocks/stubs target the methods dist.tp_allgather,
cache_seq_interface.info.batch_info.get_total_num_tokens, and
cache_seq_interface.info.batch_info.update_max_dp_num_tokens to observe
behavior.

tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/torch_moe.py (1)

276-291: ⚡ Quick win

Thread batch_info_host through _template_moe() too.

The new arg is accepted by the public op signature, but this path still drops it before _template_moe() / _template_moe_alltoall(), so the torch all-to-all backend keeps padding from static max_num_tokens. That leaves eager/reference runs on the old behavior even though sharding now wires the host-side runtime value into the graph.

Also applies to: 343-351

🤖 Prompt for AI Agents

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

In `@tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/torch_moe.py` around
lines 276 - 291, The public op torch_moe now accepts batch_info_host but drops
it before calling the internal runners; update torch_moe to pass batch_info_host
through to _template_moe (and where applicable to _template_moe_alltoall) so the
internal implementations receive the host-side runtime value instead of using
static max_num_tokens; locate calls to _template_moe/_template_moe_alltoall in
torch_moe and append the batch_info_host argument (and mirror the change for the
similar call paths around the other occurrence noted at lines 343-351) so both
all-to-all and non-all-to-all branches forward the parameter.

🤖 Prompt for all review comments with AI agents

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

Inline comments:
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/trtllm_moe.py`:
- Around line 220-245: After reading runtime_max_tokens_per_rank from
batch_info_host, validate and clamp it to the workspace upper bound
(max_num_tokens) to avoid indexing/size mismatches with the cached MoeAlltoAll
instance; i.e., if runtime_max_tokens_per_rank <= 0 set to max_num_tokens as
already done, and additionally if runtime_max_tokens_per_rank > max_num_tokens
reset it to max_num_tokens (optionally emit a warning). Apply the same
clamping/validation wherever runtime_max_tokens_per_rank is derived and later
passed into dispatch()/combine() or used for padding (the code paths around the
runtime_max_tokens_per_rank variable and the call to
_GlobalMoeA2ACache.get_alltoall and subsequent dispatch/combine).

---

Nitpick comments:
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/torch_moe.py`:
- Around line 276-291: The public op torch_moe now accepts batch_info_host but
drops it before calling the internal runners; update torch_moe to pass
batch_info_host through to _template_moe (and where applicable to
_template_moe_alltoall) so the internal implementations receive the host-side
runtime value instead of using static max_num_tokens; locate calls to
_template_moe/_template_moe_alltoall in torch_moe and append the batch_info_host
argument (and mirror the change for the similar call paths around the other
occurrence noted at lines 343-351) so both all-to-all and non-all-to-all
branches forward the parameter.

In `@tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py`:
- Around line 966-978: Add unit tests in the ad_executor test suite that
exercise the slot-13 override gate around attention DP: create one test where
ad_executor.enable_attention_dp=True and ad_executor.dist_config.tp_size>1 and
mock ad_executor.dist.tp_allgather to return different per-rank totals so that
cache_seq_interface.info.batch_info.update_max_dp_num_tokens is called with the
cross-rank max (verify call and stored value using the batch_info getter);
create another test where either enable_attention_dp=False or tp_size==1 and
assert that update_max_dp_num_tokens is not invoked and the original local total
from cache_seq_interface.info.batch_info.get_total_num_tokens remains unchanged.
Ensure mocks/stubs target the methods dist.tp_allgather,
cache_seq_interface.info.batch_info.get_total_num_tokens, and
cache_seq_interface.info.batch_info.update_max_dp_num_tokens to observe
behavior.

🪄 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: ce766138-1807-4fef-bd54-530e18fa1873

📥 Commits

Reviewing files that changed from the base of the PR and between abe5570 and a17bc36.

📒 Files selected for processing (5)

• tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py
• tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/torch_moe.py
• tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/trtllm_moe.py
• tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
• tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py

[tensorrt-cicd]

PR_Github #46644 [ run ] completed with state SUCCESS. Commit: a17bc36
/LLM/main/L0_MergeRequest_PR pipeline #36687 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 --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" --disable-fail-fast

[tensorrt-cicd]

PR_Github #46801 [ run ] triggered by Bot. Commit: a17bc36 Link to invocation

[tensorrt-cicd]

PR_Github #46801 [ run ] completed with state ABORTED. Commit: a17bc36

Link to invocation

[greg-kwasniewski1]

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

[tburt-nv]

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

[tensorrt-cicd]

PR_Github #47035 [ run ] triggered by Bot. Commit: 1fe8a39 Link to invocation

[tensorrt-cicd]

PR_Github #47035 [ run ] completed with state SUCCESS. Commit: 1fe8a39
/LLM/main/L0_MergeRequest_PR pipeline #37011 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 --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" --disable-fail-fast

[tensorrt-cicd]

PR_Github #47189 [ run ] triggered by Bot. Commit: 1fe8a39 Link to invocation

[suyoggupta]

could we also test with deepseek R1?

[tensorrt-cicd]

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

[greg-kwasniewski1]

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

[tensorrt-cicd]

PR_Github #52050 [ run ] triggered by Bot. Commit: 7e6e920 Link to invocation

[greg-kwasniewski1]

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

[tensorrt-cicd]

PR_Github #52074 [ run ] triggered by Bot. Commit: 72db9f0 Link to invocation

[tensorrt-cicd]

PR_Github #52050 [ run ] completed with state ABORTED. Commit: 7e6e920

Link to invocation

[tensorrt-cicd]

PR_Github #52074 [ run ] completed with state FAILURE. Commit: 72db9f0
/LLM/main/L0_MergeRequest_PR pipeline #41405 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]

as discussed, @MrGeva , ok with the current status?

[greg-kwasniewski1]

/bot skip --comment "AutoDeploy multi-GPU CI clean on pipeline #41405: 0 AD failures across DGX_H100-4_GPUs-AutoDeploy-{1,Post-Merge-1} and DGX_B200-4_GPUs-AutoDeploy-1. MoE+ADP validation: TestNemotronNanoV3::test_accuracy[fp8-4-attn_dp_on-trtllm] PASSED"

[tensorrt-cicd]

PR_Github #52395 [ skip ] triggered by Bot. Commit: 0900e7e Link to invocation

[tensorrt-cicd]

PR_Github #52395 [ skip ] completed with state SUCCESS. Commit: 0900e7e
Release Check Pipeline #4115 failed

Link to invocation

[greg-kwasniewski1]

/bot skip --comment "AutoDeploy multi-GPU CI clean on pipeline #41405: 0 AD failures across DGX_H100-4_GPUs-AutoDeploy-{1,Post-Merge-1} and DGX_B200-4_GPUs-AutoDeploy-1. MoE+ADP validation: TestNemotronNanoV3::test_accuracy[fp8-4-attn_dp_on-trtllm] PASSED"

[tensorrt-cicd]

PR_Github #52570 [ skip ] triggered by Bot. Commit: 2f5e6af Link to invocation

[tensorrt-cicd]

PR_Github #52570 [ skip ] completed with state SUCCESS. Commit: 2f5e6af
Skipping testing for commit 2f5e6af

Link to invocation