[#11423][feat] AutoDeploy: Basic Disagg Support

examples/auto_deploy/model_registry/configs/disagg_ctx.yaml

@@ -0,0 +1,9 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+# Selects the KV-cache transport used to move cache blocks between disaggregated workers.
+# DEFAULT lets TensorRT-LLM choose; explicit backend values include UCX and NIXL.
+# See examples/disaggregated/README.md for backend details.
+cache_transceiver_config:
+  backend: DEFAULT
+# Overlap scheduling is currently unsupported for disaggregated context workers.
+disable_overlap_scheduler: true

examples/auto_deploy/model_registry/configs/disagg_gen.yaml

@@ -0,0 +1,7 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+# Selects the KV-cache transport used to move cache blocks between disaggregated workers.
+# DEFAULT lets TensorRT-LLM choose; explicit backend values include UCX and NIXL.
+# See examples/disaggregated/README.md for backend details.
+cache_transceiver_config:
+  backend: DEFAULT

tensorrt_llm/_torch/auto_deploy/custom_ops/attention/flashinfer_attention.py

@@ -40,6 +40,7 @@ def get_env_enable_pdl() -> bool:
     AttentionDescriptor,
     AttentionLayout,
     AttentionRegistry,
+    AttentionType,
     BatchInfo,
     Constant,
     KVPagedResourceHandler,
@@ -605,6 +606,7 @@ def get_cache_initializers(
                 kv_factor=2,
                 kv_layout=_GlobalFlashInferPlanner.kv_layout,
                 sliding_window=sliding_window,
+                attention_type=AttentionType.mha,
             )
         }
 

tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_attention.py

@@ -38,6 +38,7 @@
     AttentionDescriptor,
     AttentionLayout,
     AttentionRegistry,
+    AttentionType,
     Constant,
     KVPagedResourceHandler,
     MHACallable,
@@ -1559,6 +1560,7 @@ def get_cache_initializers(
                 kv_factor=2,
                 kv_layout=KV_LAYOUT,
                 sliding_window=sliding_window,
+                attention_type=AttentionType.mha,
             )
         }
 

tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py

@@ -45,6 +45,7 @@
     AttentionDescriptor,
     AttentionLayout,
     AttentionRegistry,
+    AttentionType,
     BatchInfo,
     Constant,
     KVPagedResourceHandler,
@@ -948,6 +949,7 @@ def get_cache_initializers(
                 kv_factor=2,
                 kv_layout="HND",
                 sliding_window=sliding_window,
+                attention_type=AttentionType.mha,
             )
         }
 

tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py

@@ -26,6 +26,7 @@
 
 import math
 from abc import ABC, abstractmethod
+from enum import Enum
 from typing import Dict, List, Literal, Optional, Protocol, Sequence, Set, Tuple, Type, Union
 
 import numpy as np
@@ -40,6 +41,12 @@
 
 Constant = Union[int, float, str, None]
 
+
+class AttentionType(Enum):
+    mha = "mha"
+    mla = "mla"
+
+
 # Torch dtype → numpy dtype for fast list-to-tensor conversion.
 # numpy's list→array conversion is ~2-3x faster than torch.tensor(list) for large lists.
 _TORCH_TO_NUMPY_DTYPE: Dict[torch.dtype, np.dtype] = {
@@ -706,6 +713,8 @@ def __init__(
         # will store num_blocks later...
         self._num_blocks = None
 
+        self.attention_type: Optional[AttentionType] = None
+
         # TODO (lucaslie): can we remove this eventually from this i/f?
         self.vocab_size_padded = vocab_size_padded
 
@@ -1907,6 +1916,20 @@ def allocate(self, sequence_info: SequenceInfo) -> torch.Tensor:
         """Initialize the resource for the given sequence info."""
 
 
+class EphemeralResourceHandler(ResourceHandler):
+    """Resources that are produced and consumed within one forward pass.
+
+    Examples include MTP/Eagle hidden-state resources, which are regenerated every
+    step and not needed across steps.
+
+    Used for judging whether resources can be safely dropped when transferring from one node
+    to another, e.g. for disagg. Ephemeral resources can be safely dropped if the transfer
+    happens between forward passes.
+
+    TODO: May need to revisit this notion for intra-forward resource transfers.
+    """
+
+
 class KVPagedResourceHandler(ResourceHandler):
     """Handler for paged KV cache resources.
 
@@ -1926,6 +1949,7 @@ class KVPagedResourceHandler(ResourceHandler):
         kv_layout: Memory layout for the KV cache. Either "HND" (head-num-dim) or
             "NHD" (num-head-dim). Default is "HND" which is the standard layout
             for flashinfer.
+        attention_type: Attention layout semantics for this cache resource, e.g. ``AttentionType.mha``.
         sliding_window: Sliding window size for this layer.  ``0`` means full
             attention; a positive value puts this layer in its own VSWA group.
     """
@@ -1940,6 +1964,7 @@ def __init__(
         num_kv_heads: int,
         head_dim: int,
         dtype: torch.dtype,
+        attention_type: AttentionType,
         kv_factor: int = 2,
         kv_layout: Literal["HND", "NHD"] = "HND",
         sliding_window: int = 0,
@@ -1952,6 +1977,7 @@ def __init__(
             dtype: The dtype of the KV cache.
             kv_factor: The factor of the KV cache. Default is 2.
             kv_layout: Memory layout - "HND" or "NHD". Default is "HND".
+            attention_type: Attention layout semantics for this cache resource, e.g. ``AttentionType.mha``.
             sliding_window: Sliding window size for this layer. 0 means full attention.
         """
         self.num_kv_heads = num_kv_heads
@@ -1960,6 +1986,9 @@ def __init__(
         self.kv_factor = kv_factor
         assert kv_factor in [1, 2], f"Invalid kv_factor: {kv_factor}"
         self.kv_layout = kv_layout
+        if not isinstance(attention_type, AttentionType):
+            raise TypeError(f"attention_type must be AttentionType, got {attention_type!r}")
+        self.attention_type = attention_type
         self.sliding_window = (
             sliding_window if isinstance(sliding_window, int) and sliding_window > 0 else 0
         )
@@ -1979,6 +2008,7 @@ def __eq__(self, other: Optional[ResourceHandler]) -> bool:
             and self.dtype == other.dtype
             and self.kv_factor == other.kv_factor
             and self.kv_layout == other.kv_layout
+            and self.attention_type == other.attention_type
             and self.sliding_window == other.sliding_window
         )
 

tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py

@@ -50,6 +50,7 @@
     AttentionDescriptor,
     AttentionLayout,
     AttentionRegistry,
+    AttentionType,
     BatchInfo,
     Constant,
     MHACallable,
@@ -847,6 +848,7 @@ def __init__(self, *token_shape: int, dtype: torch.dtype) -> None:
         """
         self.token_shape = token_shape
         self.dtype = dtype
+        self.attention_type = AttentionType.mla
 
     def _get_bytes_per_token(self) -> int:
         """The size of the resource per token in bytes."""

tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_trtllm_mla.py

@@ -39,6 +39,7 @@
     AttentionDescriptor,
     AttentionLayout,
     AttentionRegistry,
+    AttentionType,
     BatchInfo,
     Constant,
     MHACallable,
@@ -67,6 +68,7 @@ def is_paged(self) -> bool:
     def __init__(self, *token_shape: int, dtype: torch.dtype) -> None:
         self.token_shape = token_shape
         self.dtype = dtype
+        self.attention_type = AttentionType.mla
 
     def allocate(self, sequence_info: SequenceInfo) -> torch.Tensor:
         return torch.empty(

tensorrt_llm/_torch/auto_deploy/custom_ops/mla/trtllm_mla.py

@@ -76,6 +76,7 @@
     AttentionDescriptor,
     AttentionLayout,
     AttentionRegistry,
+    AttentionType,
     BatchInfo,
     Constant,
     KVPagedResourceHandler,
@@ -2206,15 +2207,15 @@ def get_cache_initializers(
 
         cache_dtype = cls.resolve_cache_dtype(cache_config.dtype, compressed_kv_fake.dtype)
 
-        return {
-            "kv_cache": KVPagedResourceHandler(
-                num_kv_heads=1,
-                head_dim=kv_lora_rank + qk_rope_head_dim,
-                dtype=cache_dtype,
-                kv_factor=1,
-                kv_layout="HND",
-            )
-        }
+        kv_handler = KVPagedResourceHandler(
+            num_kv_heads=1,
+            head_dim=kv_lora_rank + qk_rope_head_dim,
+            dtype=cache_dtype,
+            kv_factor=1,
+            kv_layout="HND",
+            attention_type=AttentionType.mla,
+        )
+        return {"kv_cache": kv_handler}
 
     @classmethod
     def get_host_prepare_metadata_function(

tensorrt_llm/_torch/auto_deploy/llm_args.py

@@ -459,6 +459,14 @@ def requires_uniform_kv_caches(self) -> bool:
         """
         return False
 
+    @property
+    def reject_unmanaged_persistent_caches(self) -> bool:
+        """Whether unmanaged persistent cache resources should be rejected."""
+        return (
+            self.cache_transceiver_config is not None
+            and self.cache_transceiver_config.backend is not None
+        )
+
     def create_factory(self) -> ModelFactory:
         """Create a model factory from the arguments.
 

tensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek.py

@@ -461,8 +461,11 @@ def __init__(self, config, layer_idx: Optional[int] = None):
         self.softmax_scale = self.q_head_dim ** (-0.5)
         if config.rope_scaling is not None:
             mscale_all_dim = config.rope_scaling.get("mscale_all_dim", 0)
-            scaling_factor = config.rope_scaling["factor"]
-            if mscale_all_dim:
+            # transformers 5.x populates rope_scaling to {"rope_type": "default"} when the
+            # checkpoint has no scaling (e.g. DeepSeek-V3-Lite), so "factor" may be absent.
+            # Only apply the YaRN mscale correction when an explicit factor is present.
+            scaling_factor = config.rope_scaling.get("factor")
+            if scaling_factor is not None and mscale_all_dim:
                 mscale = DeepSeekV3YarnRotaryEmbedding._yarn_get_mscale(
                     scaling_factor, mscale_all_dim
                 )