diff --git a/examples/auto_deploy/model_registry/configs/disagg_ctx.yaml b/examples/auto_deploy/model_registry/configs/disagg_ctx.yaml new file mode 100644 index 000000000000..8039084f6c3b --- /dev/null +++ b/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 diff --git a/examples/auto_deploy/model_registry/configs/disagg_gen.yaml b/examples/auto_deploy/model_registry/configs/disagg_gen.yaml new file mode 100644 index 000000000000..6a90950635a7 --- /dev/null +++ b/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 diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/flashinfer_attention.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/flashinfer_attention.py index 39042578f7e2..bacdc3d5ffa7 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/flashinfer_attention.py +++ b/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, ) } diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_attention.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_attention.py index 83e4911dea0e..87ab18b8d2c9 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_attention.py +++ b/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, ) } diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py index d56a76755f64..1d8cda373023 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py +++ b/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, ) } diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py index c7aa851b368d..e789a23236ea 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py +++ b/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 ) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py index 5e51d206174a..cc69d0ec722c 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py +++ b/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.""" diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_trtllm_mla.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_trtllm_mla.py index 4ac85d17f57a..d99dfef385b5 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_trtllm_mla.py +++ b/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( diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/trtllm_mla.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/trtllm_mla.py index 60d3f1df3a59..45b72e7873e2 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/trtllm_mla.py +++ b/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( diff --git a/tensorrt_llm/_torch/auto_deploy/llm_args.py b/tensorrt_llm/_torch/auto_deploy/llm_args.py index 195be8c79e0f..75dc34bb7527 100644 --- a/tensorrt_llm/_torch/auto_deploy/llm_args.py +++ b/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. diff --git a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek.py b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek.py index 39614903f11e..5325cdcdea58 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek.py +++ b/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 ) diff --git a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py index 285a8b6fe0dd..6cee4311beec 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py @@ -25,6 +25,10 @@ from tensorrt_llm._torch.pyexecutor._util import get_decoding_mode from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDA_GRAPH_DUMMY_REQUEST_ID from tensorrt_llm._torch.pyexecutor.guided_decoder import GuidedDecoder +from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import ( + AttentionTypeCpp, + create_kv_cache_transceiver, +) from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, get_draft_token_length from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import BaseMambaCacheManager from tensorrt_llm._torch.pyexecutor.model_engine import ModelEngine, PyTorchModelEngine @@ -51,6 +55,7 @@ from tensorrt_llm.llmapi.tokenizer import TokenizerBase from tensorrt_llm.mapping import Mapping +from ..custom_ops.attention_interface import AttentionType from ..distributed.common import initialize_or_skip from ..llm_args import LlmArgs from ..transform.optimizer import InferenceOptimizer @@ -59,6 +64,11 @@ from ..utils.logger import ad_logger from .interface import CachedSequenceInterface, GetInferenceModel +_ATTENTION_TYPE_TO_CPP = { + AttentionType.mha: AttentionTypeCpp.DEFAULT, + AttentionType.mla: AttentionTypeCpp.MLA, +} + # Non-model multimodal metadata consumed before the exported graph or ignored by AD. # These keys must NOT leak into the generic extra_args dict — entries there # are expected to be tensors, and these may be scalars, lists, or nested dicts. @@ -390,6 +400,7 @@ def build_from_config( vocab_size_padded=factory.vocab_size_padded, spec_config=ad_config.speculative_config, requires_uniform_kv_caches=ad_config.requires_uniform_kv_caches, + reject_unmanaged_persistent_caches=ad_config.reject_unmanaged_persistent_caches, ) reporting_info = ReportingInfo( @@ -457,6 +468,11 @@ def __init__( self.max_beam_width = ad_config.max_beam_width self.spec_config = ad_config.speculative_config self._disable_overlap_scheduler = ad_config.disable_overlap_scheduler + cache_transceiver_config = ad_config.cache_transceiver_config + self._cache_transceiver_enabled = ( + cache_transceiver_config is not None + and cache_transceiver_config.backend is not None + ) self.llm_args.max_stats_len = ad_config.max_stats_len self._enable_chunked_prefill = getattr(ad_config, "enable_chunked_prefill", False) else: @@ -468,6 +484,7 @@ def __init__( self.max_beam_width = 1 self.spec_config = None self._disable_overlap_scheduler = False + self._cache_transceiver_enabled = False self.llm_args.max_stats_len = 1000 self._enable_chunked_prefill = False @@ -700,12 +717,23 @@ def _prepare_inputs( gather_context_logits: bool = False, ) -> None: """Prepare inputs for AD Model from scheduled requests.""" + context_requests = scheduled_requests.context_requests + if ( + context_requests + and self._cache_transceiver_enabled + and not self._disable_overlap_scheduler + ): + raise RuntimeError( + "AutoDeploy disaggregated context workers do not support overlap scheduling. " + "Set disable_overlap_scheduler=True, or use " + "examples/auto_deploy/model_registry/configs/disagg_ctx.yaml when starting " + "a context worker with cache_transceiver_config." + ) + # cache manager kv_cache_manager = resource_manager.get_resource_manager( ResourceManagerType.KV_CACHE_MANAGER ) - # requests in order of context, generate - context_requests = scheduled_requests.context_requests extend_requests = [ r for r in scheduled_requests.generation_requests if get_draft_token_length(r) > 0 ] @@ -720,6 +748,7 @@ def _prepare_inputs( assert len(extend_requests) == 0 or len(generation_requests) == 0 gen_requests = extend_requests + generation_requests + # Requests in order of context, extend, generation. ordered_requests = context_requests + gen_requests # sequence information @@ -776,8 +805,15 @@ def _prepare_inputs( num_prefill_tokens = len(input_ids) for request in gen_requests: - # check if need overlap and draft length - is_overlap = not self._disable_overlap_scheduler and not request.is_dummy + # Use overlap only for non-dummy requests with a previous batch slot. + # Dummy requests do not need sampled tokens from the previous iteration. + # First-step disagg decode requests have not appeared in a previous batch yet, + # so their py_batch_idx is None. + is_overlap = ( + not self._disable_overlap_scheduler + and not request.is_dummy + and request.py_batch_idx is not None + ) # check draft length draft_len = get_draft_token_length(request) @@ -1207,6 +1243,42 @@ def create_autodeploy_executor( engine=engine, ) + cache_transceiver_config = ad_config.cache_transceiver_config + kv_cache_transceiver = None + if cache_transceiver_config is not None and cache_transceiver_config.backend is not None: + if isinstance(kv_cache_manager, BaseMambaCacheManager): + # See https://github.com/NVIDIA/TensorRT-LLM/issues/14320. + raise RuntimeError( + "AutoDeploy disaggregated serving does not currently support Mamba/hybrid cache " + "managers. A prerequisite for disaggregated serving of hybrid models is to use " + "the C++ MambaCacheManager, which is currently not supported in AutoDeploy." + ) + if cache_transceiver_config.max_tokens_in_buffer is None: + # The buffer must hold the prompt's KV state (full prefill length). + # We use max_seq_len as a safe upper bound on max ISL. + cache_transceiver_config.max_tokens_in_buffer = ( + engine.cache_seq_interface.info.max_seq_len + ) + + cache_attention_type = engine.cache_seq_interface.attention_type + if cache_attention_type is None: + raise RuntimeError( + "Cache transceiver is enabled, but AutoDeploy did not find a managed paged KV " + "resource to provide attention_type." + ) + if not isinstance(cache_attention_type, AttentionType): + raise TypeError(f"attention_type must be AttentionType, got {cache_attention_type!r}") + attention_type_cpp = _ATTENTION_TYPE_TO_CPP[cache_attention_type] + + kv_cache_transceiver = create_kv_cache_transceiver( + dist_mapping, + dist, + kv_cache_manager, + attention_type_cpp, + cache_transceiver_config, + mamba_cache_manager=None, + ) + # Guided (structured) decoding. guided_decoder = None if ( @@ -1251,6 +1323,7 @@ def create_autodeploy_executor( max_batch_size=ad_config.max_batch_size, max_beam_width=ad_config.max_beam_width, guided_decoder=guided_decoder, + kv_cache_transceiver=kv_cache_transceiver, resource_governor_queue=resource_governor_queue, garbage_collection_gen0_threshold=ad_config.garbage_collection_gen0_threshold, ) diff --git a/tensorrt_llm/_torch/auto_deploy/shim/interface.py b/tensorrt_llm/_torch/auto_deploy/shim/interface.py index e475959472ee..76451e46d70c 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/interface.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/interface.py @@ -49,7 +49,9 @@ torch_dtype_to_binding = None from ..custom_ops.attention_interface import ( + AttentionType, CausalConvResourceHandler, + EphemeralResourceHandler, IntermediateConvStateHandler, IntermediateSSMStateHandler, KVPagedResourceHandler, @@ -101,6 +103,7 @@ def __init__( vocab_size_padded: Optional[int] = None, spec_config=None, requires_uniform_kv_caches: bool = False, + reject_unmanaged_persistent_caches: bool = False, ) -> None: """Initialize the CachedSequenceInterface. @@ -117,6 +120,8 @@ def __init__( cache mapping. When True, KV layers incompatible with the managed KV cache reference raise during initialization, and managed KV layers must share a single page-stride multiplier. + reject_unmanaged_persistent_caches: Whether to reject non-ephemeral cache resources + that are not managed by cache managers. """ # TODO (lucaslie): this is somewhat circular/confusing. Here `device` denotes the desired # device and not the actual device unlike, e.g., in SequenceInfo. We rely on the attribute @@ -157,6 +162,7 @@ def __init__( self._unmanaged_resources: List[str] = [] self._spec_config = spec_config self._requires_uniform_kv_caches = requires_uniform_kv_caches + self._reject_unmanaged_persistent_caches = reject_unmanaged_persistent_caches # Propagate spec-dec config into BatchInfo so attention backends can read it # via the per-forward batch_info_host tensor without needing the Python config. @@ -350,6 +356,7 @@ def _identify_managed_kv_resources( pool_by_window: Dict[int, PoolConfiguration] = {} max_seq_len = self.info.max_seq_len + attention_type: Optional[AttentionType] = None for name, handler in self._resource_lookup.items(): if not isinstance(handler, KVPagedResourceHandler): @@ -358,6 +365,15 @@ def _identify_managed_kv_resources( # max_seq_len so the C++ side gets a single concrete window key. effective_window = handler.sliding_window if handler.sliding_window > 0 else max_seq_len + if attention_type is None: + attention_type = handler.attention_type + elif handler.attention_type != attention_type: + raise RuntimeError( + f"KV layer {name} has attention_type={handler.attention_type!r} but " + f"managed KV resources already use attention_type={attention_type!r}. " + "Disaggregated KV transfer requires a single attention type." + ) + kv_managed[name] = handler handler_dtype = torch_dtype_to_binding(handler.dtype) @@ -386,6 +402,7 @@ def _identify_managed_kv_resources( dtype=handler_dtype, ) + self.info.attention_type = attention_type pool_configurations: List[PoolConfiguration] = list(pool_by_window.values()) # If the runtime requires uniform KV caches (e.g. legacy single-pool @@ -893,6 +910,64 @@ def _assign_kv_cache_views(self, kv_managed: Dict[str, KVPagedResourceHandler]) return block_offset_multiplier + def _validate_no_unmanaged_persistent_caches( + self, + kv_managed: ResourceHandlerDict, + ssm_managed: list, + ssm_spec: list, + conv_managed: list, + conv_spec: list, + replay_old_x: list, + replay_old_B: list, + replay_old_dt: list, + replay_old_dA_cumsum: list, + replay_cache_buf_idx: list, + replay_prev_num_accepted: list, + ) -> None: + """Validate persistent cache resources are cache-manager backed. + + Speculative resources (intermediate SSM/conv states and replay buffers) are bound by the + cache manager only when speculative decoding is enabled (see _create_and_assign_state_views), + so they count as managed only under that condition. When spec decoding is off they are not + registered at all (see kvcache._suppress_spec_handlers_maybe), so the loop never encounters + them. + """ + if not self._reject_unmanaged_persistent_caches: + return + + managed_names = set(kv_managed) + managed_names.update(name for name, _ in ssm_managed) + managed_names.update(name for name, _ in conv_managed) + if self._spec_config is not None: + managed_names.update(name for name, _ in ssm_spec) + managed_names.update(name for name, _ in conv_spec) + for replay_resources in ( + replay_old_x, + replay_old_B, + replay_old_dt, + replay_old_dA_cumsum, + replay_cache_buf_idx, + replay_prev_num_accepted, + ): + managed_names.update(name for name, _ in replay_resources) + + unmanaged_transfer_resources = [] + for name, handler in self._resource_lookup.items(): + if isinstance(handler, EphemeralResourceHandler): + continue + if name in managed_names: + continue + unmanaged_transfer_resources.append(f"{name} ({type(handler).__name__})") + + if unmanaged_transfer_resources: + raise RuntimeError( + "Found unmanaged persistent cache resources while " + "reject_unmanaged_persistent_caches is enabled: " + f"{unmanaged_transfer_resources}. Persistent cache resources must be managed by " + "a cache manager for configurations that need cache transfer, such as " + "disaggregated serving." + ) + def _allocate_unmanaged_resources(self) -> None: """Allocate resources not managed by cache managers. @@ -998,6 +1073,7 @@ def _create_kv_cache_manager(self, max_tokens: Optional[int] = None) -> Dict: - SSMResourceHandler maps to MambaHybridCacheManager's ssm_states buffer - CausalConvResourceHandler maps to MambaHybridCacheManager's conv_states buffer - Generic StateResourceHandler and incompatible typed handlers are allocated locally + unless transfer policy requires persistent cache resources to be managed - When both SSM and Conv handlers exist, uses min(ssm_count, conv_count) layers Args: @@ -1110,20 +1186,35 @@ def _create_kv_cache_manager(self, max_tokens: Optional[int] = None) -> Dict: block_offset_multiplier=block_offset_multiplier, ) - # 7. Allocate remaining unmanaged resources + # 7. Validate persistent cache resources before allocating local fallbacks + self._validate_no_unmanaged_persistent_caches( + kv_managed, + ssm_managed, + ssm_spec, + conv_managed, + conv_spec, + replay_old_x, + replay_old_B, + replay_old_dt, + replay_old_dA_cumsum, + replay_cache_buf_idx, + replay_prev_num_accepted, + ) + + # 8. Allocate remaining unmanaged resources self._allocate_unmanaged_resources() - # 8. Patch shutdown + # 9. Patch shutdown self._kv_cache_manager.shutdown = with_pre_callback( self._kv_cache_manager.shutdown, self._clear_caches, ) - # 8. Compute final token count and cache statistics + # 10. Compute final token count and cache statistics max_resource_count = self._kv_cache_manager.get_max_resource_count() max_tokens_final = max_resource_count * self._kv_cache_manager.tokens_per_block - # 9. Collect statistics of different types of resources + # 11. Collect statistics of different types of resources num_state_total = sum( 1 for h in self._resource_lookup.values() if isinstance(h, StateResourceHandler) ) @@ -1345,6 +1436,10 @@ def kv_cache_config(self) -> KvCacheConfig: """Return the original KVCacheConfig as passed in.""" return self._kv_cache_config_original + @property + def attention_type(self) -> Optional[AttentionType]: + return self.info.attention_type + def _clear_caches(self) -> None: """Clear all caches and views before pool release.""" for k in self._caches: diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py b/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py index f86b7da8abbc..58cb9eddaac9 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py @@ -26,8 +26,8 @@ AttentionDescriptor, AttentionLayout, AttentionRegistry, + EphemeralResourceHandler, MHACallable, - ResourceHandler, ResourceHandlerDict, SequenceInfo, ) @@ -215,12 +215,17 @@ def _apply( ) -class HiddenStatesResourceHandler(ResourceHandler): +class HiddenStatesResourceHandler(EphemeralResourceHandler): """A resource handler for hidden states.""" def __init__(self, hidden_size: int, dtype: torch.dtype) -> None: """Initialize the HiddenStatesResourceHandler. + MTP/Eagle collects hidden states from the target model and reads them in the draft model + in the same forward pass. We store these resources in an EphemeralResourceHandler because + they do not need to persist between iterations, and can be dropped when transferring + resources between forward passes. + Args: hidden_size: The size of the hidden states resource. dtype: The dtype of the hidden states resource. diff --git a/tests/integration/defs/disaggregated/test_ad_disagg.py b/tests/integration/defs/disaggregated/test_ad_disagg.py new file mode 100644 index 000000000000..c74a38316066 --- /dev/null +++ b/tests/integration/defs/disaggregated/test_ad_disagg.py @@ -0,0 +1,1043 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +import os +import pickle +import sys +import traceback +import uuid +from contextlib import ExitStack, contextmanager +from dataclasses import replace + +import cloudpickle +import pytest +import torch +from defs.conftest import get_sm_version, skip_pre_hopper +from mpi4py import MPI +from mpi4py.futures import MPIPoolExecutor + +from tensorrt_llm import DisaggregatedParams, SamplingParams +from tensorrt_llm._torch.auto_deploy import LLM as AutoDeployLLM +from tensorrt_llm._utils import set_mpi_comm +from tensorrt_llm.llmapi import Eagle3DecodingConfig + +cloudpickle.register_pickle_by_value(sys.modules[__name__]) +MPI.pickle.__init__( + cloudpickle.dumps, + cloudpickle.loads, + pickle.HIGHEST_PROTOCOL, +) + +WORKER_READY = "ready" +REQUEST_MODE_AGGREGATE = "aggregate" +MPI_REQUEST = 9999 +MPI_RESULT = MPI_REQUEST + 1 +OMPI_COMM_WORLD_ENV_KEYS = ( + "OMPI_COMM_WORLD_SIZE", + "OMPI_COMM_WORLD_RANK", + "OMPI_COMM_WORLD_LOCAL_SIZE", + "OMPI_COMM_WORLD_LOCAL_RANK", + "OMPI_COMM_WORLD_NODE_RANK", + "OMPI_UNIVERSE_SIZE", +) +AUTODEPLOY_DISAGG_SEED = 1234 +REDUCED_TINYLLAMA_LAYERS = 2 +REDUCED_DEEPSEEK_LAYERS = 2 +LLAMA_EAGLE3_EXPECTED_TEXT = " Berlin\nWhat is the capital of France? Paris\nWhat is the capital of" +LLAMA_EAGLE3_EXPECTED_TOKEN_IDS = [ + 20437, + 198, + 3923, + 374, + 279, + 6864, + 315, + 9822, + 30, + 12366, + 198, + 3923, + 374, + 279, + 6864, + 315, +] + + +MODEL_PATHS = { + "EAGLE3-LLaMA3.1-Instruct-8B": "EAGLE3-LLaMA3.1-Instruct-8B", + "Llama-3.1-8B-Instruct": "llama-3.1-model/Llama-3.1-8B-Instruct/", + "TinyLlama-1.1B-Chat-v1.0": "llama-models-v2/TinyLlama-1.1B-Chat-v1.0", + "DeepSeek-V3-Lite": "DeepSeek-V3-Lite/bf16", +} + + +def model_path(model_name): + llm_models_root = os.environ["LLM_MODELS_ROOT"] + for name, path in MODEL_PATHS.items(): + if name in model_name: + return os.path.join(llm_models_root, path) + raise ValueError(f"Unknown model: {model_name}") + + +def response_summary(response): + """Summarize values returned by AutoDeploy test workers. + + Inputs: + response: Payload from an AutoDeploy worker. It can be a formatted exception + string or a list of normal LLM output objects. + + Outputs: + A string representing the input response. + + This is useful because subprocess workers send results through MPI, so + assertion failures otherwise lose the key fields + needed to debug the disaggregated handoff: generated text/tokens, request + type, context request id, draft-token count, and logits shape when present. + """ + if isinstance(response, str): + return f"error={response}" + if isinstance(response, list) and response and hasattr(response[0], "token_ids"): + summaries = [] + for idx, output in enumerate(response): + disaggregated_params = output.disaggregated_params + if disaggregated_params is None: + request_type = REQUEST_MODE_AGGREGATE + ctx_request_id = None + else: + request_type = disaggregated_params.request_type + ctx_request_id = disaggregated_params.ctx_request_id + draft_tokens = ( + len(disaggregated_params.draft_tokens) + if disaggregated_params is not None + and disaggregated_params.draft_tokens is not None + else 0 + ) + logits = output.generation_logits + logits_shape = tuple(logits.shape) if logits is not None else None + summaries.append( + f"{idx}: text={output.text!r}, token_ids={output.token_ids}, " + f"disagg_type={request_type}, ctx_request_id={ctx_request_id}, " + f"draft_tokens={draft_tokens}, logits_shape={logits_shape}" + ) + return "[" + "; ".join(summaries) + "]" + return repr(response) + + +def seed_disagg(): + torch.manual_seed(AUTODEPLOY_DISAGG_SEED) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(AUTODEPLOY_DISAGG_SEED) + + +def base_config(extra_config=None): + common_config = dict( + runtime="trtllm", + attn_backend="trtllm", + max_batch_size=4, + max_seq_len=2048, + max_num_tokens=512, + trust_remote_code=True, + kv_cache_config={"max_tokens": 2048}, + compile_backend="torch-cudagraph", + cuda_graph_config={"batch_sizes": [1, 2, 4]}, + ) + if extra_config: + common_config.update(extra_config) + + return common_config + + +def disagg_config(extra_config=None): + return dict( + base_config(extra_config), + cache_transceiver_config={"backend": "DEFAULT"}, + ) + + +def context_config(extra_config=None): + # Context-only transfer happens after the request completes its context phase, + # so keep the context worker on the non-overlap scheduling path. + return dict( + disagg_config(extra_config), + disable_overlap_scheduler=True, + ) + + +def generation_config(generation_overlap, extra_config=None): + config = disagg_config(extra_config) + if not generation_overlap: + config["disable_overlap_scheduler"] = True + + return config + + +def single_output(response): + if isinstance(response, str): + raise RuntimeError(response) + if not isinstance(response, list) or not response: + raise RuntimeError(f"Expected a non-empty output list, got {response_summary(response)}") + return response[0] + + +def first_output(responses): + if len(responses) != 1: + raise RuntimeError(f"Expected one response, got {response_summary(responses)}") + return single_output(responses[0]) + + +def generation_params_from_context(context_output): + context_params = context_output.disaggregated_params + if context_params is None: + raise RuntimeError( + f"Context output has no disaggregated params: {response_summary([context_output])}" + ) + return replace(context_params, request_type="generation_only") + + +def has_draft_tokens(output): + params = output.disaggregated_params + return params is not None and params.draft_tokens is not None and len(params.draft_tokens) > 0 + + +def has_handoff_transport_metadata(params): + # C++ transceiver carries handoff state in opaque_state; Python/native + # transceiver carries the context endpoint in ctx_info_endpoint. + return params.opaque_state is not None or params.ctx_info_endpoint is not None + + +def run_aggregate_generation( + model, + world_size, + prompt, + sampling_params_kwargs=None, + extra_config=None, +): + """Run one non-disaggregated AutoDeploy generation request.""" + if sampling_params_kwargs is None: + sampling_params_kwargs = {"max_tokens": 25, "ignore_eos": True} + + seed_disagg() + with AutoDeployLLM( + model=model_path(model), + world_size=world_size, + **base_config(extra_config), + ) as llm: + seed_disagg() + result = llm.generate( + prompt, + sampling_params=SamplingParams(**sampling_params_kwargs), + use_tqdm=False, + ) + + output = result.outputs[0] + print(f"[AD DISAGG TEST] aggregate output: {response_summary([output])}") + return output + + +# --------------------------------------------------------------------------- +# Sequential live-pair tests. +# +# These tests run context-only and generation-only AutoDeploy instances in one +# process, one after the other, while keeping the context instance alive for +# the generation handoff. They usually run in the 1-GPU stage. Unlike the unit +# smoke tests, these load real weights. Keep exact output comparisons to +# single-request cases; batched handoff uses semantic slot checks because IFB can +# make multi-request generation differ from aggregate output even when handoff is +# correct. +# --------------------------------------------------------------------------- + + +def reduced_tinyllama_config(extra_config=None): + config = { + "model_kwargs": {"num_hidden_layers": REDUCED_TINYLLAMA_LAYERS}, + "max_batch_size": 4, + "max_seq_len": 512, + "max_num_tokens": 256, + "kv_cache_config": {"max_tokens": 1024}, + } + if extra_config: + config.update(extra_config) + return config + + +def reduced_deepseek_v3_mla_config(): + return { + "model_kwargs": {"num_hidden_layers": REDUCED_DEEPSEEK_LAYERS}, + "max_batch_size": 4, + "max_seq_len": 512, + "max_num_tokens": 256, + "kv_cache_config": {"max_tokens": 1024, "free_gpu_memory_fraction": 0.05}, + "transforms": { + "insert_cached_mla_attention": {"backend": "trtllm_mla"}, + "fuse_rope_into_trtllm_mla": {"enabled": True}, + "multi_stream_mla_attn": {"stage": "compile", "enabled": False}, + }, + } + + +def long_context_prompt(): + return ( + "TensorRT-LLM disaggregated serving separates context prefill from token generation. " + "The context worker computes the prompt KV cache, sends the cache state to the " + "generation worker, and returns the first generated token metadata. " + ) + + +def capital_completion_prompts(): + return [ + "The capital of Germany is", + "The capital of France is", + "The capital of Italy is", + "The capital of Spain is", + ] + + +def assert_context_handoff_metadata(context_output, expect_logits=False): + context_params = context_output.disaggregated_params + assert context_params is not None + assert context_params.request_type == "context_only" + assert len(context_output.token_ids) == 1 + assert context_params.ctx_request_id is not None + assert context_params.first_gen_tokens is not None + if expect_logits: + assert context_params.first_gen_logits is not None + assert has_handoff_transport_metadata(context_params) + + +def run_sequential_handoff( + model, + generation_overlap, + prompt, + sampling_params_kwargs=None, + extra_config=None, +): + if sampling_params_kwargs is None: + sampling_params_kwargs = {"max_tokens": 25, "ignore_eos": True} + + model_name = model_path(model) + with AutoDeployLLM( + model=model_name, + world_size=1, + **context_config(extra_config), + ) as context_llm: + seed_disagg() + context_output = context_llm.generate( + prompt, + sampling_params=SamplingParams(**sampling_params_kwargs), + disaggregated_params=DisaggregatedParams(request_type="context_only"), + use_tqdm=False, + ).outputs[0] + print(f"[AD DISAGG TEST] context output: {response_summary([context_output])}") + generation_params = generation_params_from_context(context_output) + + # Keep the context-side sender alive while generation consumes the + # handoff params. + with AutoDeployLLM( + model=model_name, + world_size=1, + **generation_config(generation_overlap, extra_config), + ) as generation_llm: + seed_disagg() + generation_output = generation_llm.generate( + prompt, + sampling_params=SamplingParams(**sampling_params_kwargs), + disaggregated_params=generation_params, + use_tqdm=False, + ).outputs[0] + print(f"[AD DISAGG TEST] generation output: {response_summary([generation_output])}") + + return { + "context": context_output, + "generation": generation_output, + } + + +async def run_async_requests(llm, prompts, sampling_params_kwargs, disaggregated_params): + futures = [] + for prompt, params in zip(prompts, disaggregated_params, strict=True): + seed_disagg() + futures.append( + llm.generate_async( + prompt, + sampling_params=SamplingParams(**sampling_params_kwargs), + disaggregated_params=params, + ) + ) + + outputs = [] + for future in futures: + result = await future + outputs.append(result.outputs[0]) + return outputs + + +def run_sequential_batch_handoff( + model, + generation_overlap, + prompts, + sampling_params_kwargs=None, + extra_config=None, +): + if sampling_params_kwargs is None: + sampling_params_kwargs = {"max_tokens": 25, "ignore_eos": True} + + model_name = model_path(model) + context_params = [DisaggregatedParams(request_type="context_only") for _ in range(len(prompts))] + with AutoDeployLLM( + model=model_name, + world_size=1, + **context_config(extra_config), + ) as context_llm: + context_outputs = asyncio.run( + run_async_requests(context_llm, prompts, sampling_params_kwargs, context_params) + ) + print(f"[AD DISAGG TEST] context batch output: {response_summary(context_outputs)}") + generation_params = [ + generation_params_from_context(context_output) for context_output in context_outputs + ] + + # Submit all generation-only requests before awaiting them so this + # remains a batch slot-transfer test while avoiding the async queue + # infrastructure used by the multi-GPU tests. + with AutoDeployLLM( + model=model_name, + world_size=1, + **generation_config(generation_overlap, extra_config), + ) as generation_llm: + generation_outputs = asyncio.run( + run_async_requests( + generation_llm, + prompts, + sampling_params_kwargs, + generation_params, + ) + ) + print( + f"[AD DISAGG TEST] generation batch output: {response_summary(generation_outputs)}" + ) + + return { + "context": context_outputs, + "generation": generation_outputs, + } + + +def reduced_model_config(model, extra_config=None): + if "DeepSeek-V3-Lite" in model: + config = reduced_deepseek_v3_mla_config() + else: + config = reduced_tinyllama_config() + if extra_config: + config.update(extra_config) + return config + + +def reduced_model_cases(): + return [ + pytest.param( + "TinyLlama-1.1B-Chat-v1.0", + id="tinyllama", + ), + pytest.param( + "DeepSeek-V3-Lite", + id="deepseek_v3_mla", + marks=skip_pre_hopper, + ), + ] + + +@pytest.mark.parametrize( + "model", + reduced_model_cases(), +) +@pytest.mark.skip_less_device_memory(30000) +@pytest.mark.timeout(600) +def test_reduced_layer_handoff_matches_aggregate(model): + """Check single-request disaggregated handoff matches aggregate generation.""" + prompt = "What is the capital of Germany?" + sampling_params_kwargs = { + "max_tokens": 8, + "ignore_eos": True, + "top_k": 1, + "seed": AUTODEPLOY_DISAGG_SEED, + } + extra_config = reduced_model_config(model) + # Keep real weights loaded, but reduce the decoder stack so this still + # exercises the model-specific attention/cache path without full-model cost. + aggregate_output = run_aggregate_generation( + model, + world_size=1, + prompt=prompt, + sampling_params_kwargs=sampling_params_kwargs, + extra_config=extra_config, + ) + outputs = run_sequential_handoff( + model, + generation_overlap=True, + prompt=prompt, + sampling_params_kwargs=sampling_params_kwargs, + extra_config=extra_config, + ) + + context_output = outputs["context"] + generation_output = outputs["generation"] + assert_context_handoff_metadata(context_output) + assert context_output.token_ids == aggregate_output.token_ids[:1] + assert generation_output.text == aggregate_output.text + assert generation_output.token_ids == aggregate_output.token_ids + + +@pytest.mark.parametrize( + "model", + reduced_model_cases(), +) +@pytest.mark.skip_less_device_memory(30000) +@pytest.mark.timeout(600) +def test_disaggregated_logits(model): + # Keep weighted but reduced layers so the test focuses on logits + # transfer/equality rather than full-model compile and memory cost. + extra_config = reduced_model_config(model, {"gather_generation_logits": True}) + sampling_params_kwargs = { + "max_tokens": 10, + "ignore_eos": True, + "return_generation_logits": True, + } + prompt = "What is the capital of Germany?" + aggregate_output = run_aggregate_generation( + model, + world_size=1, + prompt=prompt, + sampling_params_kwargs=sampling_params_kwargs, + extra_config=extra_config, + ) + outputs = run_sequential_handoff( + model, + generation_overlap=True, + prompt=prompt, + sampling_params_kwargs=sampling_params_kwargs, + extra_config=extra_config, + ) + + context_output = outputs["context"] + generation_output = outputs["generation"] + assert_context_handoff_metadata(context_output, expect_logits=True) + assert context_output.token_ids == aggregate_output.token_ids[:1] + assert generation_output.text == aggregate_output.text + assert generation_output.token_ids == aggregate_output.token_ids + assert aggregate_output.generation_logits is not None + assert generation_output.generation_logits is not None + assert aggregate_output.generation_logits.shape == generation_output.generation_logits.shape + # The MLA generation worker reconstructs logits from the compressed KV latent + # through a different kernel/batching path than the single aggregate pass, so + # bf16 rounding yields ~1-ULP logit differences. Use a looser tolerance for the + # MLA (DeepSeek) case; MHA (tinyllama) stays tight. The functional checks above + # (text/token_ids equality) remain strict for both. + if "DeepSeek-V3-Lite" in model: + rtol, atol = 1e-1, 1e-1 + else: + rtol, atol = 1e-2, 1e-2 + torch.testing.assert_close( + generation_output.generation_logits, + aggregate_output.generation_logits, + rtol=rtol, + atol=atol, + ) + + +@pytest.mark.skip_less_device_memory(30000) +@pytest.mark.timeout(600) +def test_tinyllama_batch_handoff_semantic_slots(): + prompts = capital_completion_prompts() + expected_capitals = ["Berlin", "Paris", "Rome", "Madrid"] + sampling_params_kwargs = { + "max_tokens": 12, + "ignore_eos": True, + "top_k": 1, + "seed": AUTODEPLOY_DISAGG_SEED, + } + outputs = run_sequential_batch_handoff( + "TinyLlama-1.1B-Chat-v1.0", + generation_overlap=True, + prompts=prompts, + sampling_params_kwargs=sampling_params_kwargs, + ) + + for expected_capital, context_output, generation_output in zip( + expected_capitals, outputs["context"], outputs["generation"], strict=True + ): + assert_context_handoff_metadata(context_output) + assert expected_capital.lower() in generation_output.text.lower(), response_summary( + outputs["generation"] + ) + + +@pytest.mark.parametrize( + "model", + reduced_model_cases(), +) +@pytest.mark.skip_less_device_memory(30000) +@pytest.mark.timeout(600) +def test_chunked_prefill_handoff(model): + # Chunked prefill needs real weights for aggregate-vs-disaggregated + # comparison, but not a full decoder stack. Use reduced layers so the test + # focuses on chunk-boundary handoff behavior with cuda graph enabled. + extra_config = reduced_model_config( + model, + { + "enable_chunked_prefill": True, + "max_num_tokens": 96, + }, + ) + prompt = long_context_prompt() * 4 + sampling_params_kwargs = {"max_tokens": 8, "ignore_eos": True} + aggregate_output = run_aggregate_generation( + model, + world_size=1, + prompt=prompt, + sampling_params_kwargs=sampling_params_kwargs, + extra_config=extra_config, + ) + outputs = run_sequential_handoff( + model, + generation_overlap=True, + prompt=prompt, + sampling_params_kwargs=sampling_params_kwargs, + extra_config=extra_config, + ) + + context_output = outputs["context"] + generation_output = outputs["generation"] + assert_context_handoff_metadata(context_output) + assert generation_output.token_ids + assert context_output.token_ids == aggregate_output.token_ids[:1] + assert generation_output.text == aggregate_output.text + assert generation_output.token_ids == aggregate_output.token_ids + + +# --------------------------------------------------------------------------- +# Async MPI worker tests. +# +# These tests launch separate context and generation worker processes and pass +# requests through an MPI intercommunicator. They are closer to the real +# disaggregated deployment shape because context and generation models can live on +# different GPUs, and some cases shard each worker across multiple GPUs. +# --------------------------------------------------------------------------- + + +def llama_eagle3_config(): + return { + "speculative_config": Eagle3DecodingConfig( + max_draft_len=3, + speculative_model=model_path("EAGLE3-LLaMA3.1-Instruct-8B"), + eagle3_one_model=True, + eagle3_layers_to_capture={1, 15, 28}, + ), + # Force the Eagle3 draft to match the BF16 Llama 3.1 target. Shared KV + # cache management requires matching target and draft KV dtypes. + "speculative_model_kwargs": {"torch_dtype": "bfloat16"}, + } + + +def get_ucx_tls(): + if get_sm_version() < 90: + return "^cuda_ipc,ib,gdr_copy" + return "^ib,gdr_copy" + + +def worker_cuda_devices(worker_world_sizes, visible_devices): + required_devices = sum(worker_world_sizes) + if visible_devices: + devices = [device.strip() for device in visible_devices.split(",") if device.strip()] + if len(devices) < required_devices: + pytest.skip( + f"AutoDeploy disaggregated world sizes {worker_world_sizes} require " + f"{required_devices} visible GPUs, got {len(devices)}" + ) + else: + devices = [str(device) for device in range(required_devices)] + + cuda_visible_devices = [] + start = 0 + for world_size in worker_world_sizes: + end = start + world_size + cuda_visible_devices.append(",".join(devices[start:end])) + start = end + return cuda_visible_devices + + +def worker_error(error): + return f"{type(error).__name__}: {error}\n{traceback.format_exc()}" + + +def isolate_ad_worker_from_outer_mpi(): + """Hide pytest's MPI transport from AutoDeploy's distributed init.""" + rank = MPI.COMM_WORLD.Get_rank() + ad_comm = MPI.COMM_WORLD.Split(color=rank, key=0) + set_mpi_comm(ad_comm) + for key in OMPI_COMM_WORLD_ENV_KEYS: + os.environ.pop(key, None) + return ad_comm + + +async def run_worker( + config, + model_name, + world_size, + cuda_visible_devices, + service_name, +): + os.environ["CUDA_VISIBLE_DEVICES"] = cuda_visible_devices + os.environ.setdefault("UCX_TLS", get_ucx_tls()) + os.environ.setdefault("UCX_MM_ERROR_HANDLING", "y") + + intercomm = MPI.COMM_WORLD.Connect(MPI.Lookup_name(service_name)) + ad_comm = isolate_ad_worker_from_outer_mpi() + try: + seed_disagg() + with AutoDeployLLM( + model=model_name, + world_size=world_size, + **config, + ) as llm: + intercomm.send(WORKER_READY, dest=0, tag=MPI_RESULT) + while True: + requests = intercomm.recv(source=0, tag=MPI_REQUEST) + if requests is None: + break + + futures = [] + for request in requests: + seed_disagg() + try: + result = llm.generate_async( + request[0], + sampling_params=request[1], + disaggregated_params=request[2], + ) + futures.append(result) + except Exception as e: + intercomm.send(worker_error(e), dest=0, tag=MPI_RESULT) + + for result in futures: + try: + output = await result + intercomm.send(output.outputs, dest=0, tag=MPI_RESULT) + except Exception as e: + intercomm.send(worker_error(e), dest=0, tag=MPI_RESULT) + except Exception as e: + intercomm.send(worker_error(e), dest=0, tag=MPI_RESULT) + raise + finally: + intercomm.Disconnect() + ad_comm.Free() + + +def worker_entry_point( + config, + model_name, + world_size, + cuda_visible_devices, + service_name, +): + return asyncio.run( + run_worker( + config, + model_name, + world_size, + cuda_visible_devices, + service_name, + ) + ) + + +def mpi_publish_name(): + service_name = f"ad_disagg_{uuid.uuid4()}" + port_name = MPI.Open_port() + MPI.Publish_name(service_name, port_name) + return service_name, port_name + + +def send_requests_to_worker(requests, worker_rank, intercomms): + intercomm = intercomms[worker_rank] + intercomm.send(requests, dest=0, tag=MPI_REQUEST) + responses = [] + for _ in range(len(requests)): + responses.append(intercomm.recv(source=0, tag=MPI_RESULT)) + return responses + + +@contextmanager +def worker_pool(worker_configs, model_names, world_sizes): + """Start async MPI workers and always tear them down after the test body. + + MPI cloudpickle serialization keeps worker callables by value, so CI workers + do not re-import this pytest module from the source checkout before the + installed TensorRT-LLM wheel is on the import path. + """ + if len(worker_configs) != len(model_names) or len(model_names) != len(world_sizes): + raise ValueError("worker_configs, model_names, and world_sizes must have the same length") + visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES") + cuda_visible_devices = worker_cuda_devices(world_sizes, visible_devices) + services = [mpi_publish_name() for _ in world_sizes] + futures = [] + intercomms = [] + with ExitStack() as stack: + try: + for ( + config, + model_name, + world_size, + worker_cuda_visible_devices, + (service_name, _), + ) in zip( + worker_configs, + model_names, + world_sizes, + cuda_visible_devices, + services, + strict=True, + ): + executor = stack.enter_context( + MPIPoolExecutor( + max_workers=1, + path=sys.path, + env={ + "UCX_TLS": get_ucx_tls(), + "UCX_MM_ERROR_HANDLING": "y", + }, + ) + ) + futures.append( + executor.submit( + worker_entry_point, + config, + model_name, + world_size, + worker_cuda_visible_devices, + service_name, + ) + ) + + for _, port_name in services: + intercomms.append(MPI.COMM_SELF.Accept(port_name)) + for intercomm in intercomms: + ready_response = intercomm.recv(source=0, tag=MPI_RESULT) + if ready_response != WORKER_READY: + raise RuntimeError( + f"Unexpected AutoDeploy worker startup response: {ready_response}" + ) + yield intercomms + finally: + for intercomm in intercomms: + intercomm.send(None, dest=0, tag=MPI_REQUEST) + intercomm.Disconnect() + for service_name, port_name in services: + MPI.Unpublish_name(service_name, port_name) + MPI.Close_port(port_name) + for future in futures: + future.result() + + +def run_context_then_generation_handoff( + model, + worker_world_sizes, + generation_overlap, + prompt, + sampling_params_kwargs=None, + extra_config=None, +): + """Run one AutoDeploy disaggregated context-to-generation handoff. + + This launches a context worker and a generation worker. It sends one + context-only request to the context worker, turns the returned + ``DisaggregatedParams`` into a generation-only request, and sends that to + the generation worker. + + Returns: + dict with ``context`` and ``generation`` outputs. The caller owns all + behavioral assertions, including output text, handoff metadata, logits, + or draft-token checks. + """ + worker_configs = [ + context_config(extra_config), + generation_config(generation_overlap, extra_config), + ] + print( + "[AD DISAGG TEST] " + f"scenario start: model={model}, worker_world_sizes={worker_world_sizes}, " + f"generation_overlap={generation_overlap}, compile_backend=torch-cudagraph, " + ) + if sampling_params_kwargs is None: + sampling_params_kwargs = {"max_tokens": 25, "ignore_eos": True} + + model_names = [model_path(model) for _ in range(2)] + world_sizes = list(worker_world_sizes) + + with worker_pool(worker_configs, model_names, world_sizes) as intercomms: + context_requests = [ + ( + prompt, + SamplingParams(**sampling_params_kwargs), + DisaggregatedParams(request_type="context_only"), + ) + ] + context_responses = send_requests_to_worker(context_requests, 0, intercomms) + context_output = first_output(context_responses) + print( + f"[AD DISAGG TEST] context output: {response_summary([context_output])}", + ) + + generation_request_disagg_params = generation_params_from_context(context_output) + generation_requests = [ + (prompt, SamplingParams(**sampling_params_kwargs), generation_request_disagg_params) + ] + + generation_responses = send_requests_to_worker(generation_requests, 1, intercomms) + generation_output = first_output(generation_responses) + print( + f"[AD DISAGG TEST] generation output: {response_summary([generation_output])}", + ) + + return { + "context": context_output, + "generation": generation_output, + } + + +@pytest.mark.threadleak(enabled=False) +@pytest.mark.skip_less_device_memory(30000) +@pytest.mark.skip_less_device(2) +@pytest.mark.timeout(600) +def test_async_generation_matches_aggregate(): + aggregate_output = run_aggregate_generation( + "TinyLlama-1.1B-Chat-v1.0", + world_size=1, + prompt="What is the capital of Germany?", + sampling_params_kwargs={"max_tokens": 10, "ignore_eos": True}, + ) + outputs = run_context_then_generation_handoff( + "TinyLlama-1.1B-Chat-v1.0", + worker_world_sizes=(1, 1), + generation_overlap=True, + prompt="What is the capital of Germany?", + sampling_params_kwargs={"max_tokens": 10, "ignore_eos": True}, + ) + context_params = outputs["context"].disaggregated_params + assert context_params is not None + assert context_params.request_type == "context_only" + assert len(outputs["context"].token_ids) == 1 + assert context_params.ctx_request_id is not None + assert context_params.first_gen_tokens is not None + assert has_handoff_transport_metadata(context_params) + assert outputs["generation"].token_ids + assert outputs["context"].token_ids == aggregate_output.token_ids[:1] + assert outputs["generation"].text == aggregate_output.text + assert outputs["generation"].token_ids == aggregate_output.token_ids + + +@pytest.mark.threadleak(enabled=False) +@pytest.mark.skip_less_device_memory(30000) +@pytest.mark.skip_less_device(2) +@pytest.mark.timeout(600) +def test_async_generation_no_overlap_matches_aggregate(): + """Match aggregate generation with the generation worker on overlap=off. + + Same shape as test_async_generation_matches_aggregate but with the + generation worker on the non-overlap scheduling path. Covers MHA disagg + with overlap=off against the aggregate baseline. + """ + sampling_params_kwargs = {"max_tokens": 10, "ignore_eos": True} + aggregate_output = run_aggregate_generation( + "TinyLlama-1.1B-Chat-v1.0", + world_size=1, + prompt="What is the capital of Germany?", + sampling_params_kwargs=sampling_params_kwargs, + ) + outputs = run_context_then_generation_handoff( + "TinyLlama-1.1B-Chat-v1.0", + worker_world_sizes=(1, 1), + generation_overlap=False, + prompt="What is the capital of Germany?", + sampling_params_kwargs=sampling_params_kwargs, + ) + assert_context_handoff_metadata(outputs["context"]) + assert outputs["generation"].token_ids + assert outputs["context"].token_ids == aggregate_output.token_ids[:1] + assert outputs["generation"].text == aggregate_output.text + assert outputs["generation"].token_ids == aggregate_output.token_ids + + +@pytest.mark.threadleak(enabled=False) +@pytest.mark.skip_less_device_memory(30000) +@pytest.mark.skip_less_device(4) +@pytest.mark.timeout(900) +def test_async_sharded_generation_handoff(): + aggregate_output = run_aggregate_generation( + "TinyLlama-1.1B-Chat-v1.0", + world_size=2, + prompt="What is the capital of Germany?", + sampling_params_kwargs={"max_tokens": 10, "ignore_eos": True}, + ) + outputs = run_context_then_generation_handoff( + "TinyLlama-1.1B-Chat-v1.0", + worker_world_sizes=(2, 2), + generation_overlap=True, + prompt="What is the capital of Germany?", + sampling_params_kwargs={"max_tokens": 10, "ignore_eos": True}, + ) + assert_context_handoff_metadata(outputs["context"]) + assert outputs["generation"].token_ids + assert outputs["context"].token_ids == aggregate_output.token_ids[:1] + assert outputs["generation"].text == aggregate_output.text + assert outputs["generation"].token_ids == aggregate_output.token_ids + + +@skip_pre_hopper +@pytest.mark.threadleak(enabled=False) +@pytest.mark.skip_less_device_memory(80000) +@pytest.mark.skip_less_device(2) +@pytest.mark.timeout(900) +def test_async_eagle3_full_model_handoff(): + sampling_params_kwargs = { + "max_tokens": 16, + "ignore_eos": True, + "top_k": 1, + "seed": AUTODEPLOY_DISAGG_SEED, + } + extra_config = llama_eagle3_config() + outputs = run_context_then_generation_handoff( + "Llama-3.1-8B-Instruct", + worker_world_sizes=(1, 1), + generation_overlap=True, + prompt="What is the capital of Germany?", + sampling_params_kwargs=sampling_params_kwargs, + extra_config=extra_config, + ) + context_params = outputs["context"].disaggregated_params + assert context_params is not None + assert context_params.request_type == "context_only" + assert len(outputs["context"].token_ids) == 1 + assert context_params.ctx_request_id is not None + assert context_params.first_gen_tokens is not None + assert has_handoff_transport_metadata(context_params) + assert outputs["generation"].token_ids + assert has_draft_tokens(outputs["context"]) + assert has_draft_tokens(outputs["generation"]) + assert outputs["context"].text == " Berlin" + assert outputs["context"].token_ids == LLAMA_EAGLE3_EXPECTED_TOKEN_IDS[:1] + assert outputs["generation"].text == LLAMA_EAGLE3_EXPECTED_TEXT + assert outputs["generation"].token_ids == LLAMA_EAGLE3_EXPECTED_TOKEN_IDS diff --git a/tests/integration/defs/disaggregated/test_ad_disagg_trtllm_serve.py b/tests/integration/defs/disaggregated/test_ad_disagg_trtllm_serve.py new file mode 100644 index 000000000000..7c3dbc35f099 --- /dev/null +++ b/tests/integration/defs/disaggregated/test_ad_disagg_trtllm_serve.py @@ -0,0 +1,256 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +import os +from pathlib import Path + +import pytest +import requests +from defs.common import get_free_port_in_ci as get_free_port +from defs.conftest import llm_models_root +from disagg_test_utils import ( + CHECK_STATUS_INTERVAL, + HEARTBEAT_INTERVAL, + INACTIVE_TIMEOUT, + run_ctx_worker, + run_disagg_server, + run_gen_worker, + terminate, +) +from openai import OpenAI + +pytest_plugins = ["disagg_test_utils"] + +SERVER_START_TIMEOUT_S = 300 +SERVER_READY_REQUEST_TIMEOUT_S = 5 +OPENAI_REQUEST_TIMEOUT_S = 60 +PROXY_PORT_MAX_RETRIES = 5 +TINYLLAMA_MODEL_DIR = "llama-models-v2/TinyLlama-1.1B-Chat-v1.0" +AUTODEPLOY_BACKEND = "_autodeploy" +EXPECTED_COMPLETION_SUBSTRING = "Berlin" + + +def tinyllama_model_path(): + return str(Path(llm_models_root()) / TINYLLAMA_MODEL_DIR) + + +def worker_cuda_devices(num_workers): + visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES") + if visible_devices: + devices = [device.strip() for device in visible_devices.split(",") if device.strip()] + if len(devices) < num_workers: + pytest.skip( + f"AutoDeploy trtllm-serve disagg smoke requires {num_workers} " + f"visible GPUs, got {len(devices)}" + ) + return devices[:num_workers] + + return [str(device) for device in range(num_workers)] + + +def autodeploy_worker_config(disagg_cluster, disable_overlap_scheduler=False): + config = { + "backend": AUTODEPLOY_BACKEND, + "max_batch_size": 1, + "cuda_graph_config": {"batch_sizes": [1]}, + "cache_transceiver_config": {"backend": "DEFAULT"}, + "disagg_cluster": disagg_cluster, + } + if disable_overlap_scheduler: + config["disable_overlap_scheduler"] = True + + return config + + +def disagg_cluster_config(port): + """Create the service-discovery config shared by workers and proxy.""" + return { + "cluster_uri": f"http://localhost:{port}", + "cluster_name": "autodeploy_disagg_smoke", + "heartbeat_interval_sec": HEARTBEAT_INTERVAL, + "inactive_timeout_sec": INACTIVE_TIMEOUT, + "minimal_instances": { + "context_servers": 1, + "generation_servers": 1, + }, + } + + +def proxy_config(port, disagg_cluster): + """Create a disaggregated proxy config that discovers workers dynamically.""" + return { + "hostname": "localhost", + "port": port, + "backend": AUTODEPLOY_BACKEND, + "disagg_cluster": disagg_cluster, + "context_servers": {"router": {"type": "round_robin"}}, + "generation_servers": {"router": {"type": "round_robin"}}, + } + + +def _process_log(process_wrapper): + """Read captured subprocess output when the utility saved it to a file.""" + if process_wrapper is None or process_wrapper.log_path is None: + return "No process log was captured." + try: + with open(process_wrapper.log_path) as log_file: + return log_file.read() + except OSError as exc: + return f"Failed to read process log {process_wrapper.log_path}: {exc}" + + +async def wait_for_disagg_server_ready_or_exit(port, processes, timeout, request_timeout): + """Wait for proxy readiness, but fail fast if any subprocess exits.""" + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + last_readiness_error = "no readiness check was attempted" + while loop.time() < deadline: + for name, process_wrapper in processes.items(): + if ( + process_wrapper + and process_wrapper.process + and process_wrapper.process.poll() is not None + ): + # Process exited before the server became ready. + log = _process_log(process_wrapper) + startup_error = RuntimeError( + f"{name} process exited before disaggregated server became ready " + f"(returncode={process_wrapper.process.returncode}).\n{log}" + ) + raise startup_error + + try: + response = requests.get( + f"http://localhost:{port}/cluster_info", timeout=request_timeout + ) + if response.status_code == 200 and response.json().get("is_ready", False): + # Server is ready. + return + last_readiness_error = ( + f"last /cluster_info response: status={response.status_code}, body={response.text}" + ) + except requests.RequestException as exc: + last_readiness_error = f"last /cluster_info request failed: {exc}" + + await asyncio.sleep(CHECK_STATUS_INTERVAL) + + raise TimeoutError( + f"Timed out after {timeout}s waiting for disaggregated server on port {port}; " + f"{last_readiness_error}" + ) + + +@pytest.mark.skip_less_device_memory(30000) +@pytest.mark.skip_less_device(2) +@pytest.mark.timeout(900) +@pytest.mark.asyncio(loop_scope="module") +async def test_openai_completion(work_dir): + """Smoke test AutoDeploy disagg through trtllm-serve and the OpenAI API. + + The lower-level tests in ``test_ad_disagg.py`` drive AutoDeploy workers + directly and inspect context/generation handoff metadata. This test instead + verifies the trtllm-serve deployment shape: context worker, generation + worker, disaggregated proxy, and an OpenAI-compatible completion request. + """ + model = tinyllama_model_path() + ctx_device, gen_device = worker_cuda_devices(2) + + last_port_conflict = None + response = None + for attempt in range(PROXY_PORT_MAX_RETRIES): + disagg_port = get_free_port() + disagg_cluster = disagg_cluster_config(disagg_port) + ctx_worker = None + gen_worker = None + disagg_server = None + + try: + # Use the same service-discovery path as the broader PyTorch disagg + # tests for worker ports. Passing port=0 lets each trtllm-serve worker + # bind an OS-selected port in the child process and register that port + # with the disaggregated proxy. + ctx_worker = run_ctx_worker( + model, + autodeploy_worker_config(disagg_cluster, disable_overlap_scheduler=True), + work_dir, + port=0, + device=ctx_device, + ) + gen_worker = run_gen_worker( + model, + autodeploy_worker_config(disagg_cluster), + work_dir, + port=0, + device=gen_device, + ) + disagg_server = run_disagg_server( + proxy_config(disagg_port, disagg_cluster), + work_dir, + disagg_port, + save_log=True, + ) + try: + await wait_for_disagg_server_ready_or_exit( + disagg_port, + { + "context worker": ctx_worker, + "generation worker": gen_worker, + "disaggregated proxy": disagg_server, + }, + SERVER_START_TIMEOUT_S, + SERVER_READY_REQUEST_TIMEOUT_S, + ) + except RuntimeError as exc: + last_port_conflict = exc + if "disaggregated proxy" not in str(exc) or ( + "EADDRINUSE" not in str(exc) + and "address already in use" not in str(exc).lower() + ): + raise + print( + f"AutoDeploy disagg serve attempt {attempt + 1} of {PROXY_PORT_MAX_RETRIES} " + f"failed with proxy port conflict, retrying: {exc}" + ) + continue + + client = OpenAI( + api_key="tensorrt_llm", + base_url=f"http://localhost:{disagg_port}/v1", + timeout=OPENAI_REQUEST_TIMEOUT_S, + max_retries=0, + ) + response = client.completions.create( + model=model, + prompt="What is the capital of Germany?", + max_tokens=32, + temperature=0, + extra_body={"ignore_eos": True}, + ) + break + finally: + terminate(ctx_worker, gen_worker, disagg_server) + + if response is None: + raise RuntimeError( + f"Failed to start AutoDeploy disagg serve smoke after {PROXY_PORT_MAX_RETRIES} " + "proxy port attempts" + ) from last_port_conflict + + assert response.choices + response_text = response.choices[0].text + assert EXPECTED_COMPLETION_SUBSTRING in response_text, ( + f"expected {EXPECTED_COMPLETION_SUBSTRING!r} in response, got {response_text!r}" + ) diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 5bed3dfa9db5..8ac3c14f8607 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -814,6 +814,11 @@ accuracy/test_llm_api_pytorch_multimodal.py::TestStep3_7::test_nvfp4 TIMEOUT (12 accuracy/test_llm_api_pytorch_multimodal.py::TestVILA1_5_3B::test_auto_dtype accuracy/test_llm_api_pytorch_ray.py::TestLlama3_1_8BInstruct::test_pp2_ray unittest/disaggregated/test_openai_disagg_server.py +disaggregated/test_ad_disagg.py::test_async_eagle3_full_model_handoff +disaggregated/test_ad_disagg.py::test_async_generation_matches_aggregate +disaggregated/test_ad_disagg.py::test_async_generation_no_overlap_matches_aggregate +disaggregated/test_ad_disagg.py::test_async_sharded_generation_handoff +disaggregated/test_ad_disagg_trtllm_serve.py::test_openai_completion disaggregated/test_auto_scaling.py::test_disagg_server_restart[etcd-round_robin] disaggregated/test_auto_scaling.py::test_disagg_server_restart[http-round_robin] disaggregated/test_auto_scaling.py::test_minimal_instances[etcd-round_robin] diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index 7a516247f1d2..ced2f9403a38 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -378,6 +378,10 @@ l0_dgx_h100: - accuracy/test_llm_api_autodeploy.py::TestQwen3_5_397B_MoE::test_bf16_small[4] - accuracy/test_llm_api_autodeploy.py::TestGemma4MoE::test_bf16 - accuracy/test_llm_api_autodeploy.py::TestMiniMaxM2::test_finegrained_fp8 + - disaggregated/test_ad_disagg.py::test_async_generation_matches_aggregate + - disaggregated/test_ad_disagg.py::test_async_generation_no_overlap_matches_aggregate + - disaggregated/test_ad_disagg.py::test_async_sharded_generation_handoff + - disaggregated/test_ad_disagg.py::test_async_eagle3_full_model_handoff # ------------- AutoDeploy Backend Stages L1 / Nightly only --------------- - condition: ranges: @@ -394,6 +398,7 @@ l0_dgx_h100: auto_trigger: others orchestrator: mpi tests: + - disaggregated/test_ad_disagg_trtllm_serve.py::test_openai_completion - accuracy/test_llm_api_autodeploy.py::TestModelRegistryAccuracy::test_autodeploy_from_registry[google_gemma-3-1b-it-False] - accuracy/test_llm_api_autodeploy.py::TestModelRegistryAccuracy::test_autodeploy_from_registry[meta-llama_Llama-3.1-8B-Instruct-False] - accuracy/test_llm_api_autodeploy.py::TestModelRegistryAccuracy::test_autodeploy_from_registry[mistralai_Ministral-8B-Instruct-2410-False] diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index e32ea255c810..6fd8c15bf5bf 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -507,6 +507,13 @@ l0_h100: - examples/test_ad_speculative_decoding.py::test_eagle_wrapper_forward[2] - examples/test_ad_speculative_decoding.py::test_nemotron_mtp_model_with_weights - examples/test_ad_guided_decoding.py::test_autodeploy_guided_decoding_main_json + - disaggregated/test_ad_disagg.py::test_disaggregated_logits[tinyllama] + - disaggregated/test_ad_disagg.py::test_disaggregated_logits[deepseek_v3_mla] + - disaggregated/test_ad_disagg.py::test_reduced_layer_handoff_matches_aggregate[tinyllama] + - disaggregated/test_ad_disagg.py::test_reduced_layer_handoff_matches_aggregate[deepseek_v3_mla] + - disaggregated/test_ad_disagg.py::test_tinyllama_batch_handoff_semantic_slots + - disaggregated/test_ad_disagg.py::test_chunked_prefill_handoff[tinyllama] + - disaggregated/test_ad_disagg.py::test_chunked_prefill_handoff[deepseek_v3_mla] - condition: ranges: system_gpu_count: diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/test_resource_handlers.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/test_resource_handlers.py index b162dba723b5..e78be45a94bd 100644 --- a/tests/unittest/auto_deploy/singlegpu/custom_ops/test_resource_handlers.py +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/test_resource_handlers.py @@ -27,6 +27,7 @@ from tensorrt_llm._torch.auto_deploy.custom_ops.attention_interface import ( AttentionDescriptor, + AttentionType, CausalConvResourceHandler, IntermediateConvStateHandler, IntermediateSSMStateHandler, @@ -52,7 +53,9 @@ def test_paged_handler_with_nhd_layout(): """Test KVPagedResourceHandler with NHD layout.""" - handler = KVPagedResourceHandler(8, 64, dtype=torch.bfloat16, kv_layout="NHD") + handler = KVPagedResourceHandler( + 8, 64, dtype=torch.bfloat16, kv_layout="NHD", attention_type=AttentionType.mha + ) assert handler.num_kv_heads == 8 assert handler.head_dim == 64 assert handler.dtype == torch.bfloat16 @@ -61,7 +64,9 @@ def test_paged_handler_with_nhd_layout(): def test_paged_handler_with_hnd_layout(): """Test KVPagedResourceHandler with explicit HND layout.""" - handler = KVPagedResourceHandler(4, 128, dtype=torch.float32, kv_layout="HND") + handler = KVPagedResourceHandler( + 4, 128, dtype=torch.float32, kv_layout="HND", attention_type=AttentionType.mha + ) assert handler.num_kv_heads == 4 assert handler.head_dim == 128 assert handler.dtype == torch.float32 @@ -71,7 +76,9 @@ def test_paged_handler_with_hnd_layout(): @pytest.mark.parametrize("kv_layout", ["HND", "NHD"]) def test_paged_handler_allocate_with_blocks(kv_layout): """Verify KVPagedResourceHandler.allocate() returns correct shape.""" - handler = KVPagedResourceHandler(8, 64, dtype=torch.float16, kv_layout=kv_layout) + handler = KVPagedResourceHandler( + 8, 64, dtype=torch.float16, kv_layout=kv_layout, attention_type=AttentionType.mha + ) tokens_per_block = 32 seq_info = SequenceInfo( max_seq_len=128, @@ -109,7 +116,7 @@ def test_paged_handler_allocate_with_blocks(kv_layout): def test_paged_handler_is_resource_handler(): """Verify KVPagedResourceHandler is a ResourceHandler subclass.""" - handler = KVPagedResourceHandler(8, 64, dtype=torch.float16) + handler = KVPagedResourceHandler(8, 64, dtype=torch.float16, attention_type=AttentionType.mha) assert isinstance(handler, ResourceHandler) @@ -292,9 +299,13 @@ def test_resolve_cache_dtype_explicit_fp8(): def test_kv_paged_handler_eq_same_head_dim_dtype(): """Verify KVPagedResourceHandler __eq__ checks head_dim and dtype.""" - h1 = KVPagedResourceHandler(8, 64, dtype=torch.float16) - h2 = KVPagedResourceHandler(4, 64, dtype=torch.float16) # Different num_kv_heads - h3 = KVPagedResourceHandler(8, 64, dtype=torch.float16, kv_layout="NHD") # Different layout + h1 = KVPagedResourceHandler(8, 64, dtype=torch.float16, attention_type=AttentionType.mha) + h2 = KVPagedResourceHandler( + 4, 64, dtype=torch.float16, attention_type=AttentionType.mha + ) # Different num_kv_heads + h3 = KVPagedResourceHandler( + 8, 64, dtype=torch.float16, kv_layout="NHD", attention_type=AttentionType.mha + ) # Different layout # head_dim, kv_factor, dtype, kv_layout -> equal (num_kv_heads doesn't matter for compatibility) assert h1 == h2 @@ -303,14 +314,32 @@ def test_kv_paged_handler_eq_same_head_dim_dtype(): def test_kv_paged_handler_eq_different_head_dim_or_dtype(): """Verify KVPagedResourceHandler __eq__ returns False for different head_dim or dtype.""" - h1 = KVPagedResourceHandler(8, 64, dtype=torch.float16) - h2 = KVPagedResourceHandler(8, 128, dtype=torch.float16) # Different head_dim - h3 = KVPagedResourceHandler(8, 64, dtype=torch.bfloat16) # Different dtype + h1 = KVPagedResourceHandler(8, 64, dtype=torch.float16, attention_type=AttentionType.mha) + h2 = KVPagedResourceHandler( + 8, 128, dtype=torch.float16, attention_type=AttentionType.mha + ) # Different head_dim + h3 = KVPagedResourceHandler( + 8, 64, dtype=torch.bfloat16, attention_type=AttentionType.mha + ) # Different dtype assert h1 != h2 assert h1 != h3 +def test_kv_paged_handler_eq_different_attention_type(): + """Verify KVPagedResourceHandler __eq__ rejects different attention semantics.""" + default_handler = KVPagedResourceHandler( + 8, 64, dtype=torch.float16, kv_factor=1, attention_type=AttentionType.mha + ) + mla_handler = KVPagedResourceHandler( + 8, 64, dtype=torch.float16, kv_factor=1, attention_type=AttentionType.mla + ) + + assert default_handler.attention_type == AttentionType.mha + assert mla_handler.attention_type == AttentionType.mla + assert default_handler != mla_handler + + def test_ssm_handler_eq_same_params(): """Verify SSMResourceHandler __eq__ for same parameters.""" h1 = SSMResourceHandler(num_heads=8, head_dim=64, d_state=16, dtype=torch.bfloat16) @@ -424,7 +453,9 @@ def test_spec_conv_handler_from_base_none(): _NON_SPECULATIVE_HANDLERS = [ SSMResourceHandler(num_heads=4, head_dim=64, d_state=16, dtype=torch.bfloat16), CausalConvResourceHandler(conv_dim=128, d_conv=4, dtype=torch.float32), - KVPagedResourceHandler(num_kv_heads=4, head_dim=64, dtype=torch.bfloat16), + KVPagedResourceHandler( + num_kv_heads=4, head_dim=64, dtype=torch.bfloat16, attention_type=AttentionType.mha + ), ] diff --git a/tests/unittest/auto_deploy/singlegpu/shim/test_ad_executor_swa_eviction.py b/tests/unittest/auto_deploy/singlegpu/shim/test_ad_executor_swa_eviction.py index 210b12d4fa36..2dfb438d4240 100644 --- a/tests/unittest/auto_deploy/singlegpu/shim/test_ad_executor_swa_eviction.py +++ b/tests/unittest/auto_deploy/singlegpu/shim/test_ad_executor_swa_eviction.py @@ -25,7 +25,10 @@ from _model_test_utils import default_max_num_tokens from tensorrt_llm._torch.auto_deploy._compat import KvCacheConfig -from tensorrt_llm._torch.auto_deploy.custom_ops.attention_interface import KVPagedResourceHandler +from tensorrt_llm._torch.auto_deploy.custom_ops.attention_interface import ( + AttentionType, + KVPagedResourceHandler, +) from tensorrt_llm._torch.auto_deploy.shim.ad_executor import ( _compute_cyclic_full_view, _compute_window_local_view, @@ -68,9 +71,15 @@ def two_window_interface(): ), ) interface.add_resource( - "kv_swa", KVPagedResourceHandler(4, 32, dtype=torch.float16, sliding_window=SWA_WINDOW) + "kv_swa", + KVPagedResourceHandler( + 4, 32, dtype=torch.float16, attention_type=AttentionType.mha, sliding_window=SWA_WINDOW + ), + ) + interface.add_resource( + "kv_full", + KVPagedResourceHandler(4, 32, dtype=torch.float16, attention_type=AttentionType.mha), ) - interface.add_resource("kv_full", KVPagedResourceHandler(4, 32, dtype=torch.float16)) interface.initialize_resources() return interface @@ -334,9 +343,15 @@ def _build_two_pool_interface(requires_uniform_kv_caches: bool): requires_uniform_kv_caches=requires_uniform_kv_caches, ) interface.add_resource( - "kv_swa", KVPagedResourceHandler(4, 32, dtype=torch.float16, sliding_window=SWA_WINDOW) + "kv_swa", + KVPagedResourceHandler( + 4, 32, dtype=torch.float16, attention_type=AttentionType.mha, sliding_window=SWA_WINDOW + ), + ) + interface.add_resource( + "kv_full", + KVPagedResourceHandler(4, 32, dtype=torch.float16, attention_type=AttentionType.mha), ) - interface.add_resource("kv_full", KVPagedResourceHandler(4, 32, dtype=torch.float16)) return interface diff --git a/tests/unittest/auto_deploy/singlegpu/shim/test_cached_sequence_interface.py b/tests/unittest/auto_deploy/singlegpu/shim/test_cached_sequence_interface.py index 34b3a3fb2164..2dab50197afc 100644 --- a/tests/unittest/auto_deploy/singlegpu/shim/test_cached_sequence_interface.py +++ b/tests/unittest/auto_deploy/singlegpu/shim/test_cached_sequence_interface.py @@ -26,10 +26,18 @@ from tensorrt_llm._torch.auto_deploy._compat import KvCacheConfig from tensorrt_llm._torch.auto_deploy.custom_ops.attention_interface import ( + AttentionType, CausalConvResourceHandler, + EphemeralResourceHandler, IntermediateConvStateHandler, IntermediateSSMStateHandler, KVPagedResourceHandler, + ReplayCacheBufIdxHandler, + ReplayOldBHandler, + ReplayOldDAcumsumHandler, + ReplayOldDtHandler, + ReplayOldXHandler, + ReplayPrevNumAcceptedHandler, SequenceInfo, SSMResourceHandler, StateResourceHandler, @@ -56,6 +64,11 @@ def __init__(self, max_draft_len: int): self.spec_dec_mode = _SpecDecModeForStateBindingTest() +class _EphemeralResourceHandlerForTest(EphemeralResourceHandler): + def allocate(self, sequence_info: SequenceInfo) -> torch.Tensor: + return torch.empty(1, device=sequence_info.device) + + @pytest.fixture def default_kv_cache_config(): """KvCacheConfig with default settings.""" @@ -182,7 +195,9 @@ def test_add_resource_paged_handler(paged_kv_cache_config): kv_cache_config=paged_kv_cache_config, ) - handler = KVPagedResourceHandler(8, 64, dtype=torch.float16, kv_layout="HND") + handler = KVPagedResourceHandler( + 8, 64, dtype=torch.float16, kv_layout="HND", attention_type=AttentionType.mha + ) full_name = interface.add_resource("kv_cache_0", handler) assert full_name in interface._resource_lookup @@ -232,8 +247,12 @@ def test_add_multiple_resources(paged_kv_cache_config): kv_cache_config=paged_kv_cache_config, ) - kv_handler_0 = KVPagedResourceHandler(8, 64, dtype=torch.float16) - kv_handler_1 = KVPagedResourceHandler(8, 64, dtype=torch.float16) + kv_handler_0 = KVPagedResourceHandler( + 8, 64, dtype=torch.float16, attention_type=AttentionType.mha + ) + kv_handler_1 = KVPagedResourceHandler( + 8, 64, dtype=torch.float16, attention_type=AttentionType.mha + ) ssm_handler = SSMResourceHandler(num_heads=4, head_dim=64, d_state=16, dtype=torch.bfloat16) interface.add_resource("kv_cache_0", kv_handler_0) @@ -259,8 +278,14 @@ def test_initialize_resources_paged_only_creates_kv_cache_manager(paged_kv_cache ) # Add only paged resources (combined KV cache) - interface.add_resource("kv_cache_0", KVPagedResourceHandler(8, 64, dtype=torch.float16)) - interface.add_resource("kv_cache_1", KVPagedResourceHandler(8, 64, dtype=torch.float16)) + interface.add_resource( + "kv_cache_0", + KVPagedResourceHandler(8, 64, dtype=torch.float16, attention_type=AttentionType.mha), + ) + interface.add_resource( + "kv_cache_1", + KVPagedResourceHandler(8, 64, dtype=torch.float16, attention_type=AttentionType.mha), + ) num_caches = interface.initialize_resources() @@ -282,9 +307,70 @@ def test_initialize_resources_mixed_shape_pools_raise_when_uniform_managed_cache requires_uniform_kv_caches=True, ) - interface.add_resource("kv_cache_full", KVPagedResourceHandler(8, 64, dtype=torch.float16)) interface.add_resource( - "kv_cache_swa", KVPagedResourceHandler(8, 80, dtype=torch.float16, sliding_window=32) + "kv_cache_full", + KVPagedResourceHandler(8, 64, dtype=torch.float16, attention_type=AttentionType.mha), + ) + interface.add_resource( + "kv_cache_swa", + KVPagedResourceHandler( + 8, 80, dtype=torch.float16, attention_type=AttentionType.mha, sliding_window=32 + ), + ) + + with pytest.raises(RuntimeError): + interface.initialize_resources() + + +def test_initialize_resources_sets_attention_type_from_kv_reference( + paged_kv_cache_config, +): + """The managed KV reference owns the cache-level attention type.""" + interface = CachedSequenceInterface( + max_seq_len=128, + max_batch_size=4, + max_num_tokens=default_max_num_tokens(128, 4), + device="cuda", + kv_cache_config=paged_kv_cache_config, + ) + + interface.add_resource( + "kv_cache_0", + KVPagedResourceHandler( + 8, 64, dtype=torch.float16, kv_factor=1, attention_type=AttentionType.mla + ), + ) + + interface.initialize_resources() + + assert interface.info.attention_type == AttentionType.mla + assert interface.attention_type == AttentionType.mla + + +def test_initialize_resources_rejects_mixed_attention_types( + paged_kv_cache_config, +): + """Strict interface configurations reject mixed KV cache attention semantics.""" + interface = CachedSequenceInterface( + max_seq_len=128, + max_batch_size=4, + max_num_tokens=default_max_num_tokens(128, 4), + device="cuda", + kv_cache_config=paged_kv_cache_config, + requires_uniform_kv_caches=True, + ) + + interface.add_resource( + "kv_cache_0", + KVPagedResourceHandler( + 8, 64, dtype=torch.float16, kv_factor=1, attention_type=AttentionType.mha + ), + ) + interface.add_resource( + "kv_cache_1", + KVPagedResourceHandler( + 8, 64, dtype=torch.float16, kv_factor=1, attention_type=AttentionType.mla + ), ) with pytest.raises(RuntimeError): @@ -302,8 +388,14 @@ def test_initialize_resources_mixed_creates_mamba_hybrid_cache_manager(paged_kv_ ) # Add paged and state resources - interface.add_resource("kv_cache_0", KVPagedResourceHandler(8, 64, dtype=torch.float16)) - interface.add_resource("kv_cache_1", KVPagedResourceHandler(8, 64, dtype=torch.float16)) + interface.add_resource( + "kv_cache_0", + KVPagedResourceHandler(8, 64, dtype=torch.float16, attention_type=AttentionType.mha), + ) + interface.add_resource( + "kv_cache_1", + KVPagedResourceHandler(8, 64, dtype=torch.float16, attention_type=AttentionType.mha), + ) interface.add_resource( "ssm_state_0", SSMResourceHandler(num_heads=4, head_dim=64, d_state=16, dtype=torch.bfloat16), @@ -330,7 +422,13 @@ def test_initialize_resources_creates_cache_views_with_correct_shape(paged_kv_ca # Using HND layout (default) full_name = interface.add_resource( "kv_cache_0", - KVPagedResourceHandler(num_kv_heads, head_dim, dtype=torch.float16, kv_layout="HND"), + KVPagedResourceHandler( + num_kv_heads, + head_dim, + dtype=torch.float16, + kv_layout="HND", + attention_type=AttentionType.mha, + ), ) interface.initialize_resources() @@ -361,7 +459,10 @@ def test_initialize_resources_creates_state_views_with_correct_shape(paged_kv_ca num_heads = 4 head_dim = 64 ssm_state_size = 16 - interface.add_resource("kv_cache_0", KVPagedResourceHandler(8, 64, dtype=torch.float16)) + interface.add_resource( + "kv_cache_0", + KVPagedResourceHandler(8, 64, dtype=torch.float16, attention_type=AttentionType.mha), + ) ssm_name = interface.add_resource( "ssm_state_0", SSMResourceHandler( @@ -485,8 +586,14 @@ def test_is_paged_returns_true_for_paged_only(paged_kv_cache_config): kv_cache_config=paged_kv_cache_config, ) - interface.add_resource("kv_cache_0", KVPagedResourceHandler(8, 64, dtype=torch.float16)) - interface.add_resource("kv_cache_1", KVPagedResourceHandler(8, 64, dtype=torch.float16)) + interface.add_resource( + "kv_cache_0", + KVPagedResourceHandler(8, 64, dtype=torch.float16, attention_type=AttentionType.mha), + ) + interface.add_resource( + "kv_cache_1", + KVPagedResourceHandler(8, 64, dtype=torch.float16, attention_type=AttentionType.mha), + ) interface.initialize_resources() assert interface.kv_cache_config_tuned.enable_block_reuse is True @@ -502,7 +609,10 @@ def test_is_paged_returns_false_for_hybrid(paged_kv_cache_config): kv_cache_config=paged_kv_cache_config, ) - interface.add_resource("kv_cache_0", KVPagedResourceHandler(8, 64, dtype=torch.float16)) + interface.add_resource( + "kv_cache_0", + KVPagedResourceHandler(8, 64, dtype=torch.float16, attention_type=AttentionType.mha), + ) interface.add_resource( "ssm_state_0", SSMResourceHandler(num_heads=4, head_dim=64, d_state=16, dtype=torch.bfloat16), @@ -527,7 +637,10 @@ def test_needs_resize_returns_false_when_fraction_is_zero(paged_kv_cache_config) kv_cache_config=paged_kv_cache_config, ) - interface.add_resource("kv_cache_0", KVPagedResourceHandler(8, 64, dtype=torch.float16)) + interface.add_resource( + "kv_cache_0", + KVPagedResourceHandler(8, 64, dtype=torch.float16, attention_type=AttentionType.mha), + ) interface.initialize_resources() assert interface.needs_resize() is False @@ -543,7 +656,10 @@ def test_needs_resize_returns_true_when_fraction_is_positive(resizable_kv_cache_ kv_cache_config=resizable_kv_cache_config, ) - interface.add_resource("kv_cache_0", KVPagedResourceHandler(8, 64, dtype=torch.float16)) + interface.add_resource( + "kv_cache_0", + KVPagedResourceHandler(8, 64, dtype=torch.float16, attention_type=AttentionType.mha), + ) interface.initialize_resources() assert interface.needs_resize() is True @@ -563,9 +679,14 @@ def test_prepare_kv_cache_config_preserves_fraction_for_vswa_estimate(): ) interface.add_resource( "kv_swa", - KVPagedResourceHandler(8, 64, dtype=torch.float16, sliding_window=64), + KVPagedResourceHandler( + 8, 64, dtype=torch.float16, sliding_window=64, attention_type=AttentionType.mha + ), + ) + interface.add_resource( + "kv_full", + KVPagedResourceHandler(4, 128, dtype=torch.float16, attention_type=AttentionType.mha), ) - interface.add_resource("kv_full", KVPagedResourceHandler(4, 128, dtype=torch.float16)) kv_managed, _ = interface._identify_managed_kv_resources() kv_cache_config = interface._prepare_kv_cache_config( @@ -588,7 +709,10 @@ def test_resize_kv_cache_manager_skipped_when_not_needed(paged_kv_cache_config): kv_cache_config=paged_kv_cache_config, ) - interface.add_resource("kv_cache_0", KVPagedResourceHandler(8, 64, dtype=torch.float16)) + interface.add_resource( + "kv_cache_0", + KVPagedResourceHandler(8, 64, dtype=torch.float16, attention_type=AttentionType.mha), + ) interface.initialize_resources() # Get initial state @@ -616,8 +740,14 @@ def test_shutdown_clears_caches(paged_kv_cache_config): kv_cache_config=paged_kv_cache_config, ) - interface.add_resource("kv_cache_0", KVPagedResourceHandler(8, 64, dtype=torch.float16)) - interface.add_resource("kv_cache_1", KVPagedResourceHandler(8, 64, dtype=torch.float16)) + interface.add_resource( + "kv_cache_0", + KVPagedResourceHandler(8, 64, dtype=torch.float16, attention_type=AttentionType.mha), + ) + interface.add_resource( + "kv_cache_1", + KVPagedResourceHandler(8, 64, dtype=torch.float16, attention_type=AttentionType.mha), + ) interface.initialize_resources() assert len(interface._caches) == 2 @@ -637,7 +767,10 @@ def test_clear_caches_clears_all(paged_kv_cache_config): kv_cache_config=paged_kv_cache_config, ) - interface.add_resource("kv_cache_0", KVPagedResourceHandler(8, 64, dtype=torch.float16)) + interface.add_resource( + "kv_cache_0", + KVPagedResourceHandler(8, 64, dtype=torch.float16, attention_type=AttentionType.mha), + ) interface.add_resource( "ssm_state_0", SSMResourceHandler(num_heads=4, head_dim=64, d_state=16, dtype=torch.bfloat16), @@ -721,7 +854,8 @@ def test_named_args_includes_sequence_info_and_caches(paged_kv_cache_config): ) full_name = interface.add_resource( - "kv_cache_0", KVPagedResourceHandler(8, 64, dtype=torch.float16) + "kv_cache_0", + KVPagedResourceHandler(8, 64, dtype=torch.float16, attention_type=AttentionType.mha), ) interface.initialize_resources() @@ -745,7 +879,10 @@ def test_args_returns_tuple_of_tensors(paged_kv_cache_config): kv_cache_config=paged_kv_cache_config, ) - interface.add_resource("kv_cache_0", KVPagedResourceHandler(8, 64, dtype=torch.float16)) + interface.add_resource( + "kv_cache_0", + KVPagedResourceHandler(8, 64, dtype=torch.float16, attention_type=AttentionType.mha), + ) interface.initialize_resources() args = interface.args @@ -1128,6 +1265,212 @@ def test_generic_state_handler_allocated_locally(paged_kv_cache_config): assert isinstance(interface.kv_cache_manager, KVCacheManager) +def test_initialize_resources_rejects_unmanaged_state_handler( + paged_kv_cache_config, +): + """Disagg rejects persistent state resources that cannot be transferred.""" + interface = CachedSequenceInterface( + max_seq_len=128, + max_batch_size=4, + max_num_tokens=default_max_num_tokens(128, 4), + device="cuda", + kv_cache_config=paged_kv_cache_config, + reject_unmanaged_persistent_caches=True, + ) + + interface.add_resource("generic_state", StateResourceHandler(10, 20, dtype=torch.float32)) + + with pytest.raises(RuntimeError): + interface.initialize_resources() + + +def test_initialize_resources_rejects_unmanaged_incompatible_kv( + paged_kv_cache_config, +): + """Disagg rejects unmanaged paged KV resources that cannot be transferred.""" + interface = CachedSequenceInterface( + max_seq_len=128, + max_batch_size=4, + max_num_tokens=default_max_num_tokens(128, 4), + device="cuda", + kv_cache_config=paged_kv_cache_config, + reject_unmanaged_persistent_caches=True, + ) + + interface.add_resource( + "kv_cache_0", + KVPagedResourceHandler(8, 64, dtype=torch.float16, attention_type=AttentionType.mha), + ) + interface.add_resource( + "kv_cache_1", + KVPagedResourceHandler(8, 80, dtype=torch.float16, attention_type=AttentionType.mha), + ) + + with pytest.raises(RuntimeError): + interface.initialize_resources() + + +def test_initialize_resources_allows_ephemeral_handler( + paged_kv_cache_config, +): + """Disagg allows explicitly ephemeral resources to remain locally allocated.""" + interface = CachedSequenceInterface( + max_seq_len=128, + max_batch_size=4, + max_num_tokens=default_max_num_tokens(128, 4), + device="cuda", + kv_cache_config=paged_kv_cache_config, + reject_unmanaged_persistent_caches=True, + ) + + full_name = interface.add_resource("ephemeral", _EphemeralResourceHandlerForTest()) + + interface.initialize_resources() + + assert full_name in interface._unmanaged_resources + + +def _add_managed_spec_replay_resources(interface, num_layers=2): + """Register a valid spec-on replay-mode state resource set on ``interface``. + + Mirrors what a Mamba backend registers in replay mode: per layer, a base SSM + base/intermediate + conv state plus the six replay-buffer categories. The intermediate SSM state is intentionally + omitted because replay mode replaces it (it is dropped from get_cache_initializers). All of these + are bound by the cache manager under spec decoding, so the disagg validator must treat them as + managed. Returns the list of replay-buffer resource names. + """ + num_heads = 4 + head_dim = 64 + d_state = 16 + n_groups = 2 + conv_dim = head_dim * num_heads + 2 * n_groups * d_state + replay_names = [] + + for i in range(num_layers): + interface.add_resource( + f"ssm_state_{i}", + SSMResourceHandler( + num_heads=num_heads, head_dim=head_dim, d_state=d_state, dtype=torch.bfloat16 + ), + ) + interface.add_resource( + f"conv_state_{i}", + CausalConvResourceHandler(conv_dim=conv_dim, d_conv=4, dtype=torch.float32), + ) + interface.add_resource( + f"intermediate_conv_state_{i}", + IntermediateConvStateHandler(conv_dim=conv_dim, d_conv=4, dtype=torch.float32), + ) + # Per-layer replay buffers — manager-backed via get_replay_old_*(layer_idx). + replay_names.append( + interface.add_resource( + f"replay_old_x_{i}", + ReplayOldXHandler(num_heads=num_heads, head_dim=head_dim, dtype=torch.bfloat16), + ) + ) + replay_names.append( + interface.add_resource( + f"replay_old_B_{i}", + ReplayOldBHandler(n_groups=n_groups, d_state=d_state, dtype=torch.bfloat16), + ) + ) + replay_names.append( + interface.add_resource(f"replay_old_dt_{i}", ReplayOldDtHandler(num_heads=num_heads)) + ) + replay_names.append( + interface.add_resource( + f"replay_old_dA_cumsum_{i}", ReplayOldDAcumsumHandler(num_heads=num_heads) + ) + ) + # Global replay buffers — one resource entry per layer, all binding to the same tensor; + # the categorization requires len == len(ssm_managed) for every replay list. + replay_names.append( + interface.add_resource(f"replay_cache_buf_idx_{i}", ReplayCacheBufIdxHandler()) + ) + replay_names.append( + interface.add_resource(f"replay_prev_num_accepted_{i}", ReplayPrevNumAcceptedHandler()) + ) + + return replay_names + + +def test_validate_no_unmanaged_persistent_caches_accepts_managed_speculative_resources( + paged_kv_cache_config, +): + """_validate_no_unmanaged_persistent_caches must not flag manager-backed speculative resources. + + Speculative resources (intermediate state + replay buffers) are handled at two different levels + depending on whether speculation is active, and the validator relies on both: + + - Spec OFF: these resources are dropped at the kvcache-insert level + (_suppress_spec_handlers_maybe), so they are never registered and never reach this validator. + That is why the validator needs no speculative special-casing of its own — the dropping + happens at the correct, lower level. + - Spec ON (this test): the resources are genuinely allocated and bound by the cache manager, so + the validator must recognize them as managed rather than reporting them as "unmanaged" + persistent caches. + + We exercise the spec-on path end-to-end through initialize_resources (rather than calling the + validator directly) precisely so the regression guard covers the call-site wiring: the replay + buffer lists must be forwarded into _validate_no_unmanaged_persistent_caches. If they are not, + the manager-backed replay buffers are falsely flagged as unmanaged and disagg+spec+replay + breaks. (Spec OFF cannot be exercised here because there is no suppression in a hand-built + resource lookup; that path is covered by the kvcache-level gate tests.) + """ + interface = CachedSequenceInterface( + max_seq_len=128, + max_batch_size=4, + max_num_tokens=default_max_num_tokens(128, 4), + device="cuda", + kv_cache_config=paged_kv_cache_config, + spec_config=_SpecConfigForStateBindingTest(max_draft_len=2), + reject_unmanaged_persistent_caches=True, + ) + + replay_names = _add_managed_spec_replay_resources(interface) + + # Must not raise: every speculative resource is cache-manager backed under spec decoding. + interface.initialize_resources() + + for name in replay_names: + assert interface._caches[name] is not None + assert name not in interface._unmanaged_resources + + +def test_validate_no_unmanaged_persistent_caches_still_rejects_unmanaged_with_spec_on( + paged_kv_cache_config, +): + """The spec-on managed-name widening must not make the validator toothless. + + Widening managed_names to include the replay buffers (so they are not falsely flagged) must not + cause genuinely-unmanaged persistent resources to slip through. Here the full managed replay set + is present AND a generic unmanaged StateResourceHandler is registered; the validator must still + reject, and must flag the unmanaged resource rather than the manager-backed replay buffers. + """ + interface = CachedSequenceInterface( + max_seq_len=128, + max_batch_size=4, + max_num_tokens=default_max_num_tokens(128, 4), + device="cuda", + kv_cache_config=paged_kv_cache_config, + spec_config=_SpecConfigForStateBindingTest(max_draft_len=2), + reject_unmanaged_persistent_caches=True, + ) + + replay_names = _add_managed_spec_replay_resources(interface) + unmanaged_name = interface.add_resource( + "leaky_state", StateResourceHandler(10, 20, dtype=torch.float32) + ) + + with pytest.raises(RuntimeError) as exc_info: + interface.initialize_resources() + + # The unmanaged resource is the one reported, not the manager-backed replay buffers. + assert unmanaged_name in str(exc_info.value) + for name in replay_names: + assert name not in str(exc_info.value) + + # ============================================================================= # _requires_copy Tests # ============================================================================= @@ -1252,8 +1595,12 @@ def test_identify_managed_kv_resources_single_window(): tokens_per_block=32, max_tokens=1024, free_gpu_memory_fraction=0.0 ), ) - interface.add_resource("kv_0", KVPagedResourceHandler(8, 64, dtype=torch.float16)) - interface.add_resource("kv_1", KVPagedResourceHandler(8, 64, dtype=torch.float16)) + interface.add_resource( + "kv_0", KVPagedResourceHandler(8, 64, dtype=torch.float16, attention_type=AttentionType.mha) + ) + interface.add_resource( + "kv_1", KVPagedResourceHandler(8, 64, dtype=torch.float16, attention_type=AttentionType.mha) + ) kv_managed, pool_configurations = interface._identify_managed_kv_resources() @@ -1275,13 +1622,22 @@ def test_identify_managed_kv_resources_dual_window_gemma4_pattern(): ) # SWA window: head_dim=64 interface.add_resource( - "kv_0", KVPagedResourceHandler(8, 64, dtype=torch.float16, sliding_window=64) + "kv_0", + KVPagedResourceHandler( + 8, 64, dtype=torch.float16, attention_type=AttentionType.mha, sliding_window=64 + ), ) interface.add_resource( - "kv_1", KVPagedResourceHandler(8, 64, dtype=torch.float16, sliding_window=64) + "kv_1", + KVPagedResourceHandler( + 8, 64, dtype=torch.float16, attention_type=AttentionType.mha, sliding_window=64 + ), ) # Full-attention window: head_dim=128 - interface.add_resource("kv_2", KVPagedResourceHandler(4, 128, dtype=torch.float16)) + interface.add_resource( + "kv_2", + KVPagedResourceHandler(4, 128, dtype=torch.float16, attention_type=AttentionType.mha), + ) kv_managed, pool_configurations = interface._identify_managed_kv_resources() @@ -1303,10 +1659,15 @@ def test_identify_managed_kv_resources_rejects_mixed_head_dim_in_same_window(): ) # Both default to sliding_window=0 → effective window = max_seq_len. Different head_dims # in the same window are not representable by one C++ pool. - interface.add_resource("kv_0", KVPagedResourceHandler(8, 64, dtype=torch.float16)) - interface.add_resource("kv_1", KVPagedResourceHandler(4, 128, dtype=torch.float16)) + interface.add_resource( + "kv_0", KVPagedResourceHandler(8, 64, dtype=torch.float16, attention_type=AttentionType.mha) + ) + interface.add_resource( + "kv_1", + KVPagedResourceHandler(4, 128, dtype=torch.float16, attention_type=AttentionType.mha), + ) - with pytest.raises(RuntimeError, match="head_dim"): + with pytest.raises(RuntimeError): interface._identify_managed_kv_resources() @@ -1321,8 +1682,12 @@ def test_single_window_creates_plain_kv_cache_manager(): tokens_per_block=32, max_tokens=1024, free_gpu_memory_fraction=0.0 ), ) - interface.add_resource("kv_0", KVPagedResourceHandler(8, 64, dtype=torch.float16)) - interface.add_resource("kv_1", KVPagedResourceHandler(8, 64, dtype=torch.float16)) + interface.add_resource( + "kv_0", KVPagedResourceHandler(8, 64, dtype=torch.float16, attention_type=AttentionType.mha) + ) + interface.add_resource( + "kv_1", KVPagedResourceHandler(8, 64, dtype=torch.float16, attention_type=AttentionType.mha) + ) interface.initialize_resources() diff --git a/tests/unittest/auto_deploy/singlegpu/shim/test_create_ad_executor.py b/tests/unittest/auto_deploy/singlegpu/shim/test_create_ad_executor.py index 15958101684d..4387d3d2730c 100644 --- a/tests/unittest/auto_deploy/singlegpu/shim/test_create_ad_executor.py +++ b/tests/unittest/auto_deploy/singlegpu/shim/test_create_ad_executor.py @@ -14,13 +14,17 @@ # limitations under the License. from dataclasses import dataclass +from types import SimpleNamespace from typing import Any, Optional from unittest.mock import Mock, patch import pytest +from tensorrt_llm._torch.auto_deploy.custom_ops.attention_interface import AttentionType from tensorrt_llm._torch.auto_deploy.llm_args import LlmArgs from tensorrt_llm._torch.auto_deploy.shim.ad_executor import create_autodeploy_executor +from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import AttentionTypeCpp +from tensorrt_llm.llmapi import CacheTransceiverConfig class MockTokenizer: @@ -60,6 +64,7 @@ class MockPyExecutor: max_batch_size: int max_beam_width: int guided_decoder: Any + kv_cache_transceiver: Any = None resource_governor_queue: Any = None garbage_collection_gen0_threshold: Optional[int] = None @@ -74,6 +79,30 @@ class MockFactory: """Unit tests for create_autodeploy_executor function.""" +def make_mock_engine( + *, + max_batch_size: int = 4, + max_seq_len: int = 128, + max_num_tokens: int = 512, + vocab_size_padded: int = 1000, + attention_type: Optional[AttentionType] = AttentionType.mha, +): + kv_cache_manager = Mock() + kv_cache_manager.impl = Mock() + + mock_engine = Mock() + mock_engine.llm_args = SimpleNamespace(max_seq_len=max_seq_len) + mock_engine.cache_seq_interface.info.num_pages = 100 + mock_engine.cache_seq_interface.info.max_seq_len = max_seq_len + mock_engine.cache_seq_interface.info.max_num_tokens = max_num_tokens + mock_engine.cache_seq_interface.info.vocab_size_padded = vocab_size_padded + mock_engine.cache_seq_interface.max_num_state_slots = max_batch_size + mock_engine.cache_seq_interface.attention_type = attention_type + mock_engine.cache_seq_interface.kv_cache_manager = kv_cache_manager + mock_engine.cache_seq_interface.kv_cache_config_tuned = Mock() + return mock_engine, kv_cache_manager + + @pytest.mark.parametrize("guided_decoding_backend", ["xgrammar", "llguidance"]) @pytest.mark.parametrize("max_batch_size", [4, 8]) @pytest.mark.parametrize("vocab_size_padded", [42, 1000]) @@ -98,16 +127,9 @@ def test_create_autodeploy_executor_with_guided_decoding( guided_decoding_backend=guided_decoding_backend, tokenizer=mock_tokenizer ) - # Mock the engine attributes that are actually used by create_autodeploy_executor - mock_engine = Mock() - mock_engine.cache_seq_interface.info.num_pages = ( - 100 # placeholder to satisfy ADEngine.build_from_config - ) - mock_engine.cache_seq_interface.info.max_num_tokens = ( - 512 # placeholder to satisfy ADEngine.build_from_config + mock_engine, _ = make_mock_engine( + max_batch_size=max_batch_size, vocab_size_padded=vocab_size_padded ) - mock_engine.cache_seq_interface.info.vocab_size_padded = vocab_size_padded - mock_engine.cache_seq_interface.max_num_state_slots = max_batch_size # Mock the specific dependencies requested, plus minimal additional mocks to prevent errors with ( @@ -152,3 +174,214 @@ def test_create_autodeploy_executor_with_guided_decoding( assert guided_decoder.max_num_sequences == ad_config.max_batch_size assert guided_decoder.vocab_size_padded == vocab_size_padded assert result.resource_governor_queue is None + + +@pytest.mark.parametrize( + "cache_attention_type, expected_attention_type", + [ + (AttentionType.mha, AttentionTypeCpp.DEFAULT), + (AttentionType.mla, AttentionTypeCpp.MLA), + ], +) +def test_create_executor_uses_cache_transceiver(cache_attention_type, expected_attention_type): + """Test create_autodeploy_executor passes the configured KV cache transceiver to PyExecutor.""" + mock_tokenizer = MockTokenizer() + mock_transceiver = Mock() + + ad_config = LlmArgs( + model="test-model", + max_batch_size=4, + max_seq_len=128, + max_input_len=64, + backend="_autodeploy", + cuda_graph_config={"max_batch_size": 4}, + cache_transceiver_config=CacheTransceiverConfig(backend="DEFAULT"), + ) + + mock_engine, kv_cache_manager = make_mock_engine(attention_type=cache_attention_type) + + with ( + patch("tensorrt_llm._torch.auto_deploy.shim.ad_executor.PyExecutor") as py_executor_cls, + patch( + "tensorrt_llm._torch.auto_deploy.shim.ad_executor.ADEngine.build_from_config" + ) as mock_ad_engine, + patch( + "tensorrt_llm._torch.auto_deploy.shim.ad_executor.create_kv_cache_transceiver", + return_value=mock_transceiver, + ) as create_transceiver, + patch( + "tensorrt_llm._torch.auto_deploy.llm_args.LlmArgs.create_factory", + return_value=MockFactory(vocab_size_padded=1000), + ), + ): + mock_ad_engine.return_value = mock_engine + py_executor_cls.side_effect = MockPyExecutor + + result = create_autodeploy_executor(ad_config, mock_tokenizer) + + create_transceiver.assert_called_once() + _, _, passed_kv_cache_manager, attention_type, passed_config = create_transceiver.call_args.args + assert passed_kv_cache_manager is kv_cache_manager + assert attention_type == expected_attention_type + assert passed_config is ad_config.cache_transceiver_config + assert passed_config.max_tokens_in_buffer == mock_engine.cache_seq_interface.info.max_seq_len + assert create_transceiver.call_args.kwargs["mamba_cache_manager"] is None + assert result.kv_cache_transceiver is mock_transceiver + + +@pytest.mark.parametrize( + "cache_attention_type, expected_attention_type", + [ + (AttentionType.mha, AttentionTypeCpp.DEFAULT), + (AttentionType.mla, AttentionTypeCpp.MLA), + ], +) +def test_create_executor_preserves_explicit_transceiver_buffer_size( + cache_attention_type, expected_attention_type +): + """Test create_autodeploy_executor preserves explicit KV cache transceiver buffer sizing.""" + mock_tokenizer = MockTokenizer() + mock_transceiver = Mock() + + ad_config = LlmArgs( + model="test-model", + max_batch_size=4, + max_seq_len=128, + max_input_len=64, + backend="_autodeploy", + cuda_graph_config={"max_batch_size": 4}, + cache_transceiver_config=CacheTransceiverConfig( + backend="DEFAULT", max_tokens_in_buffer=1024 + ), + ) + + mock_engine, _ = make_mock_engine(attention_type=cache_attention_type) + + with ( + patch("tensorrt_llm._torch.auto_deploy.shim.ad_executor.PyExecutor") as py_executor_cls, + patch( + "tensorrt_llm._torch.auto_deploy.shim.ad_executor.ADEngine.build_from_config" + ) as mock_ad_engine, + patch( + "tensorrt_llm._torch.auto_deploy.shim.ad_executor.create_kv_cache_transceiver", + return_value=mock_transceiver, + ) as create_transceiver, + patch( + "tensorrt_llm._torch.auto_deploy.llm_args.LlmArgs.create_factory", + return_value=MockFactory(vocab_size_padded=1000), + ), + ): + mock_ad_engine.return_value = mock_engine + py_executor_cls.side_effect = MockPyExecutor + + create_autodeploy_executor(ad_config, mock_tokenizer) + + _, _, _, attention_type, passed_config = create_transceiver.call_args.args + assert attention_type == expected_attention_type + assert passed_config.max_tokens_in_buffer == 1024 + assert create_transceiver.call_args.kwargs["mamba_cache_manager"] is None + + +@pytest.mark.parametrize("cache_attention_type", ["mha", "unsupported"]) +def test_create_executor_rejects_non_enum_attention_type(cache_attention_type): + """Test create_autodeploy_executor requires enum KV cache attention semantics.""" + mock_tokenizer = MockTokenizer() + + ad_config = LlmArgs( + model="test-model", + max_batch_size=4, + max_seq_len=128, + max_input_len=64, + backend="_autodeploy", + cuda_graph_config={"max_batch_size": 4}, + cache_transceiver_config=CacheTransceiverConfig(backend="DEFAULT"), + ) + + mock_engine, _ = make_mock_engine() + mock_engine.cache_seq_interface.attention_type = cache_attention_type + + with ( + patch( + "tensorrt_llm._torch.auto_deploy.shim.ad_executor.ADEngine.build_from_config" + ) as mock_ad_engine, + patch( + "tensorrt_llm._torch.auto_deploy.llm_args.LlmArgs.create_factory", + return_value=MockFactory(vocab_size_padded=1000), + ), + ): + mock_ad_engine.return_value = mock_engine + + with pytest.raises(TypeError): + create_autodeploy_executor(ad_config, mock_tokenizer) + + +def test_create_executor_requires_attention_type(): + """Test create_autodeploy_executor requires KV cache attention semantics for disagg.""" + mock_tokenizer = MockTokenizer() + + ad_config = LlmArgs( + model="test-model", + max_batch_size=4, + max_seq_len=128, + max_input_len=64, + backend="_autodeploy", + cuda_graph_config={"max_batch_size": 4}, + cache_transceiver_config=CacheTransceiverConfig(backend="DEFAULT"), + ) + + mock_engine, _ = make_mock_engine(attention_type=None) + + with ( + patch( + "tensorrt_llm._torch.auto_deploy.shim.ad_executor.ADEngine.build_from_config" + ) as mock_ad_engine, + patch( + "tensorrt_llm._torch.auto_deploy.llm_args.LlmArgs.create_factory", + return_value=MockFactory(vocab_size_padded=1000), + ), + ): + mock_ad_engine.return_value = mock_engine + + with pytest.raises(RuntimeError): + create_autodeploy_executor(ad_config, mock_tokenizer) + + +def test_create_executor_rejects_mamba_cache_manager_for_transceiver(): + """Test create_autodeploy_executor rejects Mamba/hybrid cache transfer.""" + from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import BaseMambaCacheManager + + mock_tokenizer = MockTokenizer() + + ad_config = LlmArgs( + model="test-model", + max_batch_size=4, + max_seq_len=128, + max_input_len=64, + backend="_autodeploy", + cuda_graph_config={"max_batch_size": 4}, + cache_transceiver_config=CacheTransceiverConfig(backend="DEFAULT"), + ) + + mock_engine, _ = make_mock_engine(attention_type=AttentionType.mha) + mamba_cache_manager = Mock(spec=BaseMambaCacheManager) + mamba_cache_manager.impl = Mock() + mock_engine.cache_seq_interface.kv_cache_manager = mamba_cache_manager + + with ( + patch( + "tensorrt_llm._torch.auto_deploy.shim.ad_executor.ADEngine.build_from_config" + ) as mock_ad_engine, + patch( + "tensorrt_llm._torch.auto_deploy.shim.ad_executor.create_kv_cache_transceiver" + ) as create_transceiver, + patch( + "tensorrt_llm._torch.auto_deploy.llm_args.LlmArgs.create_factory", + return_value=MockFactory(vocab_size_padded=1000), + ), + ): + mock_ad_engine.return_value = mock_engine + + with pytest.raises(RuntimeError): + create_autodeploy_executor(ad_config, mock_tokenizer) + + create_transceiver.assert_not_called() diff --git a/tests/unittest/auto_deploy/singlegpu/shim/test_llm_config.py b/tests/unittest/auto_deploy/singlegpu/shim/test_llm_config.py index 598117396124..d8099ebf609a 100644 --- a/tests/unittest/auto_deploy/singlegpu/shim/test_llm_config.py +++ b/tests/unittest/auto_deploy/singlegpu/shim/test_llm_config.py @@ -66,6 +66,18 @@ def test_requires_uniform_kv_caches_follows_attention_backend(): ) +def test_cache_transceiver_rejects_unmanaged_persistent_caches(): + """Cache transceiver rejects unmanaged persistent cache resources.""" + args = LlmArgs( + model="test-model", + attn_backend="flashinfer", + cache_transceiver_config={"backend": "DEFAULT"}, + ) + + assert args.requires_uniform_kv_caches is False + assert args.reject_unmanaged_persistent_caches is True + + # ================================ # Config Flow Tests # ================================ diff --git a/tests/unittest/auto_deploy/singlegpu/smoke/test_disagg.py b/tests/unittest/auto_deploy/singlegpu/smoke/test_disagg.py new file mode 100644 index 000000000000..10b802fe7173 --- /dev/null +++ b/tests/unittest/auto_deploy/singlegpu/smoke/test_disagg.py @@ -0,0 +1,304 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +from dataclasses import replace + +import pytest +from _model_test_utils import get_small_model_config +from utils.util import skip_pre_hopper + +from tensorrt_llm import DisaggregatedParams, SamplingParams +from tensorrt_llm._torch.auto_deploy import LLM as AutoDeployLLM +from tensorrt_llm.llmapi import Eagle3DecodingConfig + +LLAMA_MODEL_ID = "meta-llama/Meta-Llama-3.1-8B-Instruct" +EAGLE3_MODEL_ID = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B" +DEEPSEEK_MODEL_ID = "deepseek-ai/DeepSeek-V3" +DEEPSEEK_DISAGG_TRANSFORMS = { + "insert_cached_attention": {"backend": "triton"}, + "insert_cached_mla_attention": {"backend": "trtllm_mla"}, + "fuse_rope_into_trtllm_mla": {"enabled": True}, + "compile_model": {"backend": "torch-simple"}, +} + + +def small_model_config_disagg(model_id, attn_backend, compile_backend, **overrides): + args = get_small_model_config(model_id)["args"] + kv_cache_config = args["kv_cache_config"] + kv_cache_config.update( + { + "tokens_per_block": 4, + "max_tokens": 64, + "free_gpu_memory_fraction": 0.001, + } + ) + args.update( + { + "world_size": 1, + "runtime": "trtllm", + "skip_tokenizer_init": True, + "attn_backend": attn_backend, + "compile_backend": compile_backend, + "cuda_graph_config": {"max_batch_size": 2} + if compile_backend == "torch-cudagraph" + else None, + "max_batch_size": 2, + "max_seq_len": 64, + "max_num_tokens": 16, + "cache_transceiver_config": {"backend": "DEFAULT"}, + } + ) + args.update(overrides) + return args + + +def _sampling_params(): + return SamplingParams( + max_tokens=4, + ignore_eos=True, + add_special_tokens=False, + end_id=2, + pad_id=0, + ) + + +def has_handoff_transport_metadata(params): + # C++ transceiver carries handoff state in opaque_state; Python/native + # transceiver carries the context endpoint in ctx_info_endpoint. + return params.opaque_state is not None or params.ctx_info_endpoint is not None + + +def context_output_valid(output): + params = output.disaggregated_params + return ( + params is not None + and params.request_type == "context_only" + and len(output.token_ids) == 1 + and params.ctx_request_id is not None + and params.first_gen_tokens is not None + and has_handoff_transport_metadata(params) + ) + + +def create_generation_params(context_output): + assert context_output_valid(context_output) + params = context_output.disaggregated_params + assert params is not None + return replace(params, request_type="generation_only") + + +def has_draft_tokens(output): + params = output.disaggregated_params + return params is not None and params.draft_tokens is not None and len(params.draft_tokens) > 0 + + +def run_live_disagg_smoke( + model_id, + attn_backend, + compile_backend, + common_config_overrides=None, + context_config_overrides=None, + generation_config_overrides=None, +): + common_config_overrides = common_config_overrides or {} + context_config_overrides = context_config_overrides or {} + generation_config_overrides = generation_config_overrides or {} + + with AutoDeployLLM( + **small_model_config_disagg( + model_id, + attn_backend, + compile_backend, + **common_config_overrides, + disable_overlap_scheduler=True, + **context_config_overrides, + ) + ) as context_llm: + context_output = context_llm.generate( + [1, 2, 3, 4], + sampling_params=_sampling_params(), + disaggregated_params=DisaggregatedParams(request_type="context_only"), + ).outputs[0] + assert context_output_valid(context_output) + disaggregated_params = create_generation_params(context_output) + + # Keep the context LLM alive while the generation LLM consumes the + # handoff params. The real cache transceiver uses the context-side + # sender endpoint, so this is the meaningful generation-only smoke. + with AutoDeployLLM( + **small_model_config_disagg( + model_id, + attn_backend, + compile_backend, + **common_config_overrides, + **generation_config_overrides, + ) + ) as generation_llm: + generation_output = generation_llm.generate( + [1, 2, 3, 4], + sampling_params=_sampling_params(), + disaggregated_params=disaggregated_params, + ).outputs[0] + assert generation_output.token_ids + return context_output, generation_output + + +async def run_async_requests(llm, prompts, sampling_params, disaggregated_params): + futures = [] + for prompt, params in zip(prompts, disaggregated_params, strict=True): + futures.append( + llm.generate_async( + prompt, + sampling_params=sampling_params, + disaggregated_params=params, + ) + ) + return [(await future).outputs[0] for future in futures] + + +def run_live_batch_disagg_smoke(model_id, attn_backend, compile_backend, config_overrides): + prompts = [[1, 2, 3, 4], [5, 6, 7, 8]] + context_params = [DisaggregatedParams(request_type="context_only") for _ in prompts] + + with AutoDeployLLM( + **small_model_config_disagg( + model_id, + attn_backend, + compile_backend, + **config_overrides, + disable_overlap_scheduler=True, + ) + ) as context_llm: + context_outputs = asyncio.run( + run_async_requests(context_llm, prompts, _sampling_params(), context_params) + ) + generation_params = [ + create_generation_params(context_output) for context_output in context_outputs + ] + + with AutoDeployLLM( + **small_model_config_disagg( + model_id, + attn_backend, + compile_backend, + **config_overrides, + ) + ) as generation_llm: + generation_outputs = asyncio.run( + run_async_requests( + generation_llm, + prompts, + _sampling_params(), + generation_params, + ) + ) + + for context_output, generation_output in zip(context_outputs, generation_outputs, strict=True): + assert context_output_valid(context_output) + assert generation_output.token_ids + + +GENERIC_DISAGG_SMOKE_CASES = [ + pytest.param(LLAMA_MODEL_ID, "trtllm", "torch-simple", {}, id="llama-trtllm-simple"), + pytest.param(LLAMA_MODEL_ID, "trtllm", "torch-cudagraph", {}, id="llama-trtllm-cudagraph"), + pytest.param(LLAMA_MODEL_ID, "flashinfer", "torch-simple", {}, id="llama-flashinfer-simple"), + pytest.param( + LLAMA_MODEL_ID, + "flashinfer", + "torch-cudagraph", + {}, + id="llama-flashinfer-cudagraph", + ), + pytest.param( + DEEPSEEK_MODEL_ID, + "trtllm", + "torch-simple", + {"transforms": DEEPSEEK_DISAGG_TRANSFORMS}, + marks=skip_pre_hopper, + id="deepseek-trtllm-simple", + ), +] + + +@pytest.mark.parametrize( + ("model_id", "attn_backend", "compile_backend", "config_overrides"), + GENERIC_DISAGG_SMOKE_CASES, +) +def test_autodeploy_disaggregated_smoke(model_id, attn_backend, compile_backend, config_overrides): + if model_id == DEEPSEEK_MODEL_ID: + pytest.importorskip( + "transformers.models.deepseek_v3.configuration_deepseek_v3", + reason="DeepseekV3Config requires a newer transformers version", + ) + + run_live_disagg_smoke(model_id, attn_backend, compile_backend, config_overrides) + + +@pytest.mark.parametrize( + ("model_id", "attn_backend", "compile_backend", "config_overrides"), + GENERIC_DISAGG_SMOKE_CASES, +) +def test_autodeploy_disaggregated_batch_smoke( + model_id, attn_backend, compile_backend, config_overrides +): + if model_id == DEEPSEEK_MODEL_ID: + pytest.importorskip( + "transformers.models.deepseek_v3.configuration_deepseek_v3", + reason="DeepseekV3Config requires a newer transformers version", + ) + + run_live_batch_disagg_smoke(model_id, attn_backend, compile_backend, config_overrides) + + +def test_autodeploy_disaggregated_eagle3_smoke(): + target_model_config = get_small_model_config(LLAMA_MODEL_ID) + eagle3_model_config = get_small_model_config(EAGLE3_MODEL_ID) + target_model_kwargs = { + **target_model_config["args"]["model_kwargs"], + "num_hidden_layers": 3, + } + speculative_config = Eagle3DecodingConfig( + max_draft_len=3, + speculative_model=eagle3_model_config["args"]["model"], + eagle3_one_model=True, + eagle3_layers_to_capture={0, 1, 2}, + ) + speculative_model_kwargs = { + **target_model_kwargs, + **eagle3_model_config["args"]["model_kwargs"], + "torch_dtype": "bfloat16", + } + + # This is intentionally a smoke test: small_model_config_disagg uses + # skip_loading_weights=True, so the meaningful assertions are that one-model + # Eagle builds with a reduced target/draft pair and carries draft-token + # metadata through the live disaggregated handoff. Force the draft dtype to + # match the BF16 Llama target because shared KV cache management requires + # target and draft KV resources to have the same dtype. Use three reduced + # target layers to match Llama Eagle3's default three-layer capture. + # Weighted acceptance and quality coverage belong in integration tests. + context_output, generation_output = run_live_disagg_smoke( + LLAMA_MODEL_ID, + "flashinfer", + "torch-simple", + common_config_overrides={ + "model_kwargs": target_model_kwargs, + "speculative_config": speculative_config, + "speculative_model_kwargs": speculative_model_kwargs, + }, + ) + assert has_draft_tokens(context_output) + assert has_draft_tokens(generation_output) diff --git a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_kv_cache.py b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_kv_cache.py index 83868014a7cc..dc214bfd4604 100644 --- a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_kv_cache.py +++ b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_kv_cache.py @@ -26,7 +26,10 @@ from tensorrt_llm._torch.auto_deploy._compat import KvCacheConfig # Initialize resources first (KVPagedResourceHandler is used within tests below) -from tensorrt_llm._torch.auto_deploy.custom_ops.attention_interface import KVPagedResourceHandler +from tensorrt_llm._torch.auto_deploy.custom_ops.attention_interface import ( + AttentionType, + KVPagedResourceHandler, +) from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm from tensorrt_llm._torch.auto_deploy.models.factory import ( FullModelExportInfo, @@ -553,7 +556,8 @@ def test_initialize_cache_transform_calls_initialize_resources(dummy_cached_inte ) dummy_cached_interface.add_resource( - "kv_cache_0", KVPagedResourceHandler(8, 64, dtype=torch.float16) + "kv_cache_0", + KVPagedResourceHandler(8, 64, dtype=torch.float16, attention_type=AttentionType.mha), ) # Mock the factory and shared_config @@ -574,7 +578,8 @@ def test_initialize_cache_transform_calls_initialize_resources(dummy_cached_inte def test_resize_kv_cache_transform_skipped_when_not_needed(dummy_cached_interface): """Verify ResizeKVCache transform is skipped when resize not needed.""" dummy_cached_interface.add_resource( - "kv_cache_0", KVPagedResourceHandler(8, 64, dtype=torch.float16) + "kv_cache_0", + KVPagedResourceHandler(8, 64, dtype=torch.float16, attention_type=AttentionType.mha), ) dummy_cached_interface.initialize_resources() @@ -615,7 +620,10 @@ def test_resize_kv_cache_transform_runs_when_needed(): kv_cache_config=kv_cache_config, ) - cm.add_resource("kv_cache_0", KVPagedResourceHandler(8, 64, dtype=torch.float16)) + cm.add_resource( + "kv_cache_0", + KVPagedResourceHandler(8, 64, dtype=torch.float16, attention_type=AttentionType.mha), + ) cm.initialize_resources() # Create the transform with a proper config @@ -1254,7 +1262,9 @@ def _non_speculative_handlers(): return [ SSMResourceHandler(num_heads=4, head_dim=64, d_state=16, dtype=torch.bfloat16), CausalConvResourceHandler(conv_dim=128, d_conv=4, dtype=torch.float32), - KVPagedResourceHandler(num_kv_heads=4, head_dim=64, dtype=torch.bfloat16), + KVPagedResourceHandler( + num_kv_heads=4, head_dim=64, dtype=torch.bfloat16, attention_type=AttentionType.mha + ), ]