[#11423][feat] AutoDeploy: Basic Disagg Support by govind-ramnarayan · Pull Request #14057 · NVIDIA/TensorRT-LLM

govind-ramnarayan · 2026-05-12T23:00:25Z

Summary: Disaggregated serving support for the AutoDeploy backend. Since AutoDeploy is already integrated into the TRTLLM KvCacheManager, we just need to instantiate the KvCacheTransceiver when creating the PyExecutor for AutoDeploy. This PR implements support for disaggregation in AutoDeploy for MHA and MLA-based models.

All caches must be managed by the KvCacheManager in order for disagg to run. This is to proactively flag correctness issues where the prefill node does not register some cache with the KvCacheManager, so the gen node does not get it and crashes. Exceptions can be made by registering a cache as "Ephemeral". A good example is the hidden state cache we use for collecting target-model hidden states for Eagle / MTP. This is produced and consumed within a single forward pass, so does not need to be passed from prefill to generation workers.

Not Implemented:

Support for Mamba / Hybrid models. Disagg for these seems to require using the C++ CacheManager, which we do not yet support in AutoDeploy. Tracked in issue here: [Feature]: AutoDeploy: Integrate with CppMambaHybridCacheManager #14320

Testing:

Test Coverage:
We want to make sure disagg works with AutoDeploy for all the attention types that we believe we should support with disagg - currently MHA and MLA. So we want to test correctness for at least one MHA and one MLA model.

It also interacts with some other features - for instance with overlap scheduling enabled, the generation worker must not crash on the first iteration, and for Eagle, we should tolerate dropping the Hidden State Cache in the disaggregated handoff. For chunked prefill, we shouldn't do the handoff too early.

We lift some of the feature matrix for testing from https://github.com/NVIDIA/TensorRT-LLM/blob/main/tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py for testing.

Methodology / types of tests:

1. Smoke tests with reduced number of layers and random weights. They run context and generation workers sequentially on a single GPU. These are just meant to test that the disagg pipeline works without crashing. These should be cheap to run, but cannot test disagg vs agg output without using the same weights for both modes.
1. Sequential integration tests - run context and generation workers sequentially, usually on a single GPU, with real weights. We can check that the disagg pipeline runs and that the output matches aggregate output.
1. Async integration tests. These use an asynchronous messaging queue to more accurately simulate a real disagg workflow instead of running workers sequentially. Inspired by PyTorch backend tests here: https://github.com/NVIDIA/TensorRT-LLM/blob/main/tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py
1. TRTLLM Serve Smoke test - sets up trtllm-serve similar to how a customer would and makes sure disagg serving provides good output here.

We try to put most tests into 1) and 2) as they are cheaper and have simpler setup, and errors are more easily traceable to AutoDeploy's engine and executor. For 2) and 3) in particular, we can use reduced-layer counts with real weights and check that aggregate output matches disagg output (even if it is nonsense) - this requires less memory. We include some full-sized model tests to make sure the pipeline is working in realistic scenarios.

@coderabbitai summary

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

tensorrt-cicd · 2026-06-09T01:34:57Z

PR_Github #52889 [ run ] completed with state FAILURE. Commit: 029e77c
/LLM/main/L0_MergeRequest_PR pipeline #42139 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

suyoggupta · 2026-06-09T03:38:13Z

@juney-nvidia : can you please unblock. thanks

govind-ramnarayan · 2026-06-09T06:07:04Z

/bot run

tensorrt-cicd · 2026-06-09T06:12:51Z

PR_Github #52982 [ run ] triggered by Bot. Commit: 029e77c Link to invocation

tensorrt-cicd · 2026-06-09T06:37:52Z

PR_Github #52982 [ run ] completed with state FAILURE. Commit: 029e77c
/LLM/main/L0_MergeRequest_PR pipeline #42214 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

govind-ramnarayan · 2026-06-09T16:52:49Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-09T16:59:17Z

PR_Github #53104 [ run ] triggered by Bot. Commit: 1aa2c14 Link to invocation

Squashed branch work for AutoDeploy disagg serving: KV cache transport config, ad_executor/interface changes for disagg, MLA/attention custom op updates, DeepSeek model tweaks, plus integration tests (test_ad_disagg, trtllm_serve), unit tests, smoke tests, and test-list registrations. Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>

KVPagedResourceHandler now requires attention_type; update the two remaining base-handler fixtures (_non_speculative_handlers and _NON_SPECULATIVE_HANDLERS) that still constructed it without one. Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>

govind-ramnarayan · 2026-06-09T19:20:47Z

/bot kill

tensorrt-cicd · 2026-06-09T19:27:53Z

PR_Github #53132 [ kill ] triggered by Bot. Commit: fcad4dd Link to invocation

tensorrt-cicd · 2026-06-09T19:32:42Z

PR_Github #53104 [ run ] completed with state ABORTED. Commit: 1aa2c14

Link to invocation

tensorrt-cicd · 2026-06-09T19:32:47Z

PR_Github #53132 [ kill ] completed with state SUCCESS. Commit: fcad4dd
Successfully killed previous jobs for commit fcad4dd

Link to invocation

govind-ramnarayan · 2026-06-09T20:08:27Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-09T20:15:36Z

PR_Github #53136 [ run ] triggered by Bot. Commit: fcad4dd Link to invocation

tensorrt-cicd · 2026-06-10T01:49:44Z

PR_Github #53136 [ run ] completed with state SUCCESS. Commit: fcad4dd
/LLM/main/L0_MergeRequest_PR pipeline #42340 completed with status: 'SUCCESS'

CI Report

Link to invocation

govind-ramnarayan · 2026-06-10T01:50:08Z

/bot run --stage-list DGX_H100-4_GPUs-AutoDeploy-Post-Merge-1

tensorrt-cicd · 2026-06-10T01:55:54Z

PR_Github #53193 [ run ] triggered by Bot. Commit: fcad4dd Link to invocation

tensorrt-cicd · 2026-06-10T03:19:34Z

PR_Github #53193 [ run ] completed with state SUCCESS. Commit: fcad4dd
/LLM/main/L0_MergeRequest_PR pipeline #42390 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

govind-ramnarayan · 2026-06-10T04:46:02Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-10T04:52:32Z

PR_Github #53230 [ run ] triggered by Bot. Commit: fcad4dd Link to invocation

tensorrt-cicd · 2026-06-10T05:17:55Z

PR_Github #53230 [ run ] completed with state SUCCESS. Commit: fcad4dd
/LLM/main/L0_MergeRequest_PR pipeline #42425 completed with status: 'SUCCESS'

CI Report

Link to invocation

ADEngine subclasses the abstract ModelEngine and does not run PyTorchModelEngine.__init__, so it never set `enable_spec_decode`. After NVIDIA#14546 added an unguarded `self.model_engine.enable_spec_decode` read in `_prepare_disagg_gen_transmission_complete` (the disagg generation handoff path that ADEngine traverses via NVIDIA#14057 AutoDeploy Basic Disagg Support), AutoDeploy disaggregated runs crash with: AttributeError: 'ADEngine' object has no attribute 'enable_spec_decode' NVIDIA#14546 and NVIDIA#14057 each passed CI independently but conflict semantically once both are on main. Set `is_spec_decode`/`enable_spec_decode` in ADEngine.__init__, mirroring PyTorchModelEngine (enable_spec_decode == spec_config is not None), so ADEngine satisfies the ModelEngine attribute contract that shared PyExecutor code relies on. Signed-off-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com>

Shixiaowei02 · 2026-06-11T11:21:09Z

I submitted a pr #15260 for the conflict fix of #14057 and #14546 .

github-actions Bot assigned govind-ramnarayan May 12, 2026

govind-ramnarayan force-pushed the gramnarayan/disagg branch from 1daf15e to 842ef6f Compare May 13, 2026 00:27