diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index faa8b064caf6..b81e62d15f11 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -121,6 +121,13 @@ def _resolve_disagg_transceiver_route( return backend, runtime +def is_disagg_enabled( + cache_transceiver_config: Optional[CacheTransceiverConfig]) -> bool: + """Whether this executor participates in disaggregated serving.""" + return (cache_transceiver_config is not None + and cache_transceiver_config.backend is not None) + + def get_kv_cache_manager_cls( model_config: ModelConfig, kv_cache_config: KvCacheConfig, @@ -3039,21 +3046,36 @@ def create_kv_cache_compression_manager( def compute_max_num_sequences(mapping: Mapping, max_batch_size: int, disable_overlap_scheduler: bool, - enable_overlap_headroom: bool = False) -> int: + enable_overlap_headroom: bool = False, + is_disagg: bool = False) -> int: """Size the sequence-slot pool (and the sampler state it indexes). - ``enable_overlap_headroom`` is intentionally opt-in. Disaggregated + The pool must seat every request the admission path can let through. On a + disaggregated generation server that bound is KVCacheManagerV2's + IndexMapper, sized at twice max_num_sequences so a batch in KV transfer + can overlap a batch that is generating. Seats below that bound let a + request be admitted that cannot be seated once its transfer lands, and + add_slot then raises on the executor's event-loop thread, killing the rank + mid-collective. + + enable_overlap_headroom covers a different case. Disaggregated attention-DP needs a second non-PP slot set because the V2 scheduler can backfill seats before the overlap scheduler releases the previous - iteration's terminal slots. Pipeline parallelism already sizes the pool - by ``pp_size``. + iteration's terminal slots. Pipeline parallelism already sizes the pool by + pp_size. """ if mapping.has_pp(): num_micro_batches = mapping.pp_size else: num_micro_batches = (2 if enable_overlap_headroom and not disable_overlap_scheduler else 1) - return max_batch_size * num_micro_batches + num_seats = max_batch_size * num_micro_batches + if is_disagg: + # max() rather than another multiplication: the disagg and + # overlap-headroom factors both cover one extra set of in-flight + # sequences, so they overlap rather than compose. + num_seats = max(num_seats, max_batch_size * mapping.pp_size * 2) + return num_seats def should_enable_adp_dummy_fixes(mapping: Mapping) -> bool: @@ -3085,10 +3107,28 @@ def should_enable_disagg_adp_overlap_headroom( cache_transceiver_config: Optional[CacheTransceiverConfig], disable_overlap_scheduler: bool) -> bool: """Gate extra sequence slots to non-PP disaggregated attention-DP.""" - is_disagg = (cache_transceiver_config is not None - and cache_transceiver_config.backend is not None) - return (mapping.enable_attention_dp and is_disagg and not mapping.has_pp() - and not disable_overlap_scheduler) + return (mapping.enable_attention_dp + and is_disagg_enabled(cache_transceiver_config) + and not mapping.has_pp() and not disable_overlap_scheduler) + + +def validate_seq_slot_pool_covers_admission(max_num_sequences: int, + kv_cache_manager) -> None: + """Fail at startup if the seat pool is smaller than what admission allows. + + Otherwise the shortfall stays invisible until a request that cannot be + seated arrives, and it then surfaces as a hang rather than an error. + No-op for managers that publish no admission bound. + """ + admission_bound = getattr(kv_cache_manager, "max_admissible_sequences", + None) + if admission_bound is None or max_num_sequences >= admission_bound: + return + raise ValueError( + f"Sequence-slot pool ({max_num_sequences} seats) is smaller than the " + f"number of sequences {type(kv_cache_manager).__name__} can admit " + f"({admission_bound}). Requests would be admitted that cannot be " + "seated; see compute_max_num_sequences.") def create_py_executor_instance( @@ -3126,15 +3166,18 @@ def create_py_executor_instance( spec_config = model_engine.spec_config + is_disagg = is_disagg_enabled(cache_transceiver_config) + if max_num_sequences is None: max_num_sequences = compute_max_num_sequences( - mapping, max_batch_size, llm_args.disable_overlap_scheduler) + mapping, + max_batch_size, + llm_args.disable_overlap_scheduler, + is_disagg=is_disagg) logger.info( f"max_seq_len={max_seq_len}, max_num_requests={max_num_sequences}, max_num_tokens={max_num_tokens}, max_batch_size={max_batch_size}" ) - is_disagg = (cache_transceiver_config is not None - and cache_transceiver_config.backend is not None) for key, value in llm_args.extra_resource_managers.items(): if key in resources: raise ValueError( @@ -3292,6 +3335,7 @@ def create_py_executor_instance( if isinstance(model_engine, PyTorchModelEngine): model_engine._init_cuda_graph_lora_manager(lora_config) + validate_seq_slot_pool_covers_admission(max_num_sequences, kv_cache_manager) resources[ResourceManagerType.SEQ_SLOT_MANAGER] = SeqSlotManager( max_num_sequences) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index cc598cbf67ca..45ac242342c6 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -1454,9 +1454,11 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: # waiting for the previous batch's transfers to finish. max_num_sequences = max_batch_size * mapping.pp_size assert num_reserved_index_slots >= 0, "num_reserved_index_slots must be non-negative" - index_mapper_capacity = ( - max_num_sequences * (2 if is_disagg else 1) + num_reserved_index_slots - ) + # Admission bound for real requests, excluding the reserved slots that + # only ever hold persistent dummies. Published so the sequence-slot + # pool can be checked against it at startup. + self.max_admissible_sequences = max_num_sequences * (2 if is_disagg else 1) + index_mapper_capacity = self.max_admissible_sequences + num_reserved_index_slots logger.info( f"KVCacheManagerV2: IndexMapper capacity={index_mapper_capacity} " f"(max_num_sequences={max_num_sequences}, is_disagg={is_disagg}, " diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index a5e76249dd74..886cbdd74eb6 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -448,13 +448,15 @@ def __init__( self.mapping = mapping if mapping.has_pp(): init_pp_comm(mapping) - # Disaggregated attention-DP can backfill a batch before the overlap - # scheduler releases the previous batch's terminal sequence slots. - from ._util import (compute_max_num_sequences, + # Both reasons the pool must exceed one forward batch are known from + # the runtime topology, before the model loads. + from ._util import (compute_max_num_sequences, is_disagg_enabled, should_enable_adp_dummy_fixes, should_enable_disagg_adp_overlap_headroom, should_enable_non_overlap_adp_forward_intent, should_enable_scheduler_aware_adp_dummy) + self._is_disagg = is_disagg_enabled( + getattr(llm_args, "cache_transceiver_config", None)) self._enable_disagg_adp_overlap_headroom = ( should_enable_disagg_adp_overlap_headroom( mapping, llm_args.cache_transceiver_config, @@ -465,6 +467,7 @@ def __init__( self.batch_size, llm_args.disable_overlap_scheduler, enable_overlap_headroom=self._enable_disagg_adp_overlap_headroom, + is_disagg=self._is_disagg, ) self.dist = dist if dist is not None: @@ -1132,8 +1135,7 @@ def _initialize_no_kv_cache_runner( mm_encoder_cache_enabled=self._mm_encoder_cache_enabled, spec_config=self.spec_config, is_draft_model=self.is_draft_model, - num_seq_slots=(self.max_num_seq_slots if - self._enable_disagg_adp_overlap_headroom else None), + num_seq_slots=self.max_num_seq_slots, original_max_draft_len=self.original_max_draft_len, original_max_total_draft_tokens=( self.original_max_total_draft_tokens), @@ -3709,11 +3711,9 @@ def forward_multimodal_encoder_items( def _set_up_spec_metadata( self, spec_resource_manager: Optional[BaseResourceManager]): spec_config = self.spec_config if self.enable_spec_decode else None - # The disaggregated attention-DP overlap path opts into larger metadata - # buffers. Passing None preserves the established max_num_requests - # fallback for other configurations, including PP. - num_seq_slots = (self.max_num_seq_slots - if self._enable_disagg_adp_overlap_headroom else None) + # Slot-indexed metadata buffers must span the whole pool, whatever + # sizes it. + num_seq_slots = self.max_num_seq_slots if self.spec_metadata is not None: return self.spec_metadata self.spec_metadata = get_spec_metadata( diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 8cdc9e5e46de..77da0671fe16 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -37,7 +37,8 @@ get_spec_resource_manager) from ..virtual_memory import scope as virtual_memory_scope from ._util import (KvCacheCreator, _adjust_torch_mem_fraction, - create_py_executor_instance, instantiate_sampler, is_mla, + compute_max_num_sequences, create_py_executor_instance, + instantiate_sampler, is_disagg_enabled, is_mla, validate_feature_combination) from .config_utils import (is_hybrid_linear, is_minimax_m3, resolve_cache_transceiver_config, @@ -662,8 +663,14 @@ def allocation_scope(current_stage: ExecutorMemoryType): resolve_cache_transceiver_config(cache_transceiver_config) config = model_engine.model.model_config.pretrained_config - max_num_seq_slots = getattr(model_engine, "max_num_seq_slots", - max_batch_size * getattr(mapping, "pp_size", 1)) + # Engines that are not PyTorchModelEngine do not size the pool themselves; + # fall back to the same sizing function so the two never disagree. + max_num_seq_slots = getattr( + model_engine, "max_num_seq_slots", None) or compute_max_num_sequences( + mapping, + max_batch_size, + llm_args.disable_overlap_scheduler, + is_disagg=is_disagg_enabled(cache_transceiver_config)) if is_mla(config): if model_engine.model.model_config.enable_flash_mla: tokens_per_block = 64 @@ -757,15 +764,11 @@ def allocation_scope(current_stage: ExecutorMemoryType): if guided_decoding_config is not None: with allocation_scope(ExecutorMemoryType.GUIDED_DECODER): if mapping.is_last_pp_rank(): - guided_decoder_slots = (max_num_seq_slots if getattr( - model_engine, "_enable_disagg_adp_overlap_headroom", False) - else max_batch_size) kwargs = { "guided_decoding_config": guided_decoding_config, - # The disaggregated attention-DP overlap path follows the - # expanded slot pool. Other configurations retain - # max_batch_size. - "max_num_sequences": guided_decoder_slots, + # Indexed by py_seq_slot, so it must span the whole pool + # rather than one forward batch. + "max_num_sequences": max_num_seq_slots, "vocab_size_padded": model_engine.model.vocab_size_padded, "rank": mapping.rank, } @@ -876,8 +879,7 @@ def allocation_scope(current_stage: ExecutorMemoryType): if model_engine.model.model_config.is_generation: #NOTE: non-generation models do not have kv cache - is_disagg = (cache_transceiver_config is not None - and cache_transceiver_config.backend is not None) + is_disagg = is_disagg_enabled(cache_transceiver_config) is_hybrid = is_hybrid_linear( model_engine.model.model_config.pretrained_config) diff --git a/tensorrt_llm/_torch/pyexecutor/seq_slot_manager.py b/tensorrt_llm/_torch/pyexecutor/seq_slot_manager.py index a3f11e564236..b56be5236b89 100644 --- a/tensorrt_llm/_torch/pyexecutor/seq_slot_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/seq_slot_manager.py @@ -1,3 +1,5 @@ +from tensorrt_llm.logger import logger + from .llm_request import LlmRequest from .resource_manager import BaseResourceManager, SlotManager from .scheduler import ScheduledRequests @@ -18,7 +20,7 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests) -> None: for llm_req in scheduled_batch.all_requests(): if llm_req.is_disagg_generation_init_state: logger.info( - f"Skip assigning sequence slot for DISAGG_GENERATION_INIT request." + "Skip assigning sequence slot for DISAGG_GENERATION_INIT request." ) continue if llm_req.seq_slot is None or llm_req.is_disagg_generation_transmission_complete: diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index add507f0c5cc..6ff0911244a9 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -669,8 +669,8 @@ def batch_uses_penalty(self, value: bool) -> None: # Vocab size used for draft_probs buffer allocation. vocab_size: int = 0 # Size of the SeqSlotManager pool. py_seq_slot values range over - # [0, num_seq_slots); DeepSeek-V4 overlap can use 2 * max_batch_size, - # larger than max_num_requests (== max_batch_size). + # [0, num_seq_slots), which exceeds max_num_requests (== max_batch_size) + # under PP, DeepSeek-V4 overlap headroom, and disaggregated serving. # Slot-indexed buffers (draft_probs) must span this full range. # 0 falls back to max_num_requests. num_seq_slots: int = 0 @@ -845,11 +845,8 @@ def prepare_rejection_sampling_buffers(self): if not self.use_rejection_sampling: return - # Slot-indexed buffers span the full SeqSlotManager pool: py_seq_slot - # can range over [0, num_seq_slots), which under DeepSeek-V4 overlap - # exceeds max_num_requests. Fall back to max_num_requests when the pool - # size is unknown (0). One extra scratch row at index ``slot_capacity`` - # absorbs CUDA-graph dummy/padding requests (``py_seq_slot is None``). + # One extra scratch row past the pool absorbs CUDA-graph dummy and + # padding requests, whose py_seq_slot is None. slot_capacity = self.num_seq_slots or self.max_num_requests num_slot_rows = slot_capacity + 1 diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 849a045108b2..100148418f9d 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -358,8 +358,7 @@ def _build_spec_metadata(spec_config, num_seq_slots=None): use_rejection_sampling = getattr(spec_config, "use_rejection_sampling", False) - # Slot-indexed buffers (draft_probs) must span the SeqSlotManager pool; - # DeepSeek-V4 overlap can exceed max_num_requests. + # See SpecMetadata.num_seq_slots for why draft_probs must span the pool. num_seq_slots = (num_seq_slots if num_seq_slots is not None else max_num_requests) vocab_size = getattr(model_config, "vocab_size", 0) diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index d4fdf3ddca7e..967064c97672 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -37,6 +37,7 @@ l0_a10: - unittest/_torch/executor/kv_cache/test_kv_cache_v2_capacity_only.py - unittest/_torch/executor/test_error_classification.py - unittest/_torch/executor/test_resource_manager.py + - unittest/_torch/executor/test_seq_slot_sizing.py - unittest/_torch/moe/test_communication_factory.py # NOTE: this is a CPU-only test, but we do not have a dedicated job for this (and therefore no # test list either). diff --git a/tests/unittest/_torch/executor/test_seq_slot_sizing.py b/tests/unittest/_torch/executor/test_seq_slot_sizing.py index 7db2c6ed74ac..5550bb988e3e 100644 --- a/tests/unittest/_torch/executor/test_seq_slot_sizing.py +++ b/tests/unittest/_torch/executor/test_seq_slot_sizing.py @@ -1,29 +1,41 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Disaggregated attention-DP seq-slot sizing includes overlap headroom. - -Under the overlap scheduler, requests finished in the previous iteration -still hold their sequence slots when the next iteration's -prepare_resources runs, while the V2 scheduler has already dropped them -from its budget (no_schedule_after_state=GENERATION_TO_COMPLETE) and -backfilled their seats. Transient slot demand is therefore -2 * max_batch_size, regardless of whether speculative decoding is enabled. -The headroom is selected from runtime topology rather than model architecture. - -compute_max_num_sequences is the single sizing implementation used both -for the executor's SeqSlotManager pool (create_py_executor_instance) and -for the sampler state (create_torch_sampler_args). +"""Seq-slot pool sizing. + +Two independent reasons the pool must exceed max_batch_size, both selected +from runtime topology rather than model architecture: + +Attention-DP overlap headroom. Requests finished in the previous iteration +still hold their sequence slots when the next iteration's prepare_resources +runs, while the V2 scheduler has already dropped them from its budget +(no_schedule_after_state=GENERATION_TO_COMPLETE) and backfilled their seats. +Transient slot demand is therefore twice max_batch_size, whether or not +speculative decoding is enabled. + +Disaggregated serving. On a generation server the admission bound is +KVCacheManagerV2's IndexMapper rather than the seat pool, and it is sized at +twice max_num_sequences. Unlike the headroom above, this holds regardless of +attention-DP, overlap, or PP. + +compute_max_num_sequences is the single sizing implementation, used for the +executor's SeqSlotManager pool and for the sampler state. Every other +slot-indexed buffer follows the same number, since py_seq_slot indexes them +all. """ +from unittest.mock import Mock + import pytest from tensorrt_llm._torch.pyexecutor._util import ( compute_max_num_sequences, create_torch_sampler_args, + is_disagg_enabled, should_enable_adp_dummy_fixes, should_enable_disagg_adp_overlap_headroom, should_enable_non_overlap_adp_forward_intent, should_enable_scheduler_aware_adp_dummy, + validate_seq_slot_pool_covers_admission, ) from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig from tensorrt_llm.mapping import Mapping @@ -33,7 +45,7 @@ (1, False, True, 2), (1, False, False, 1), (1, True, True, 1), - # Existing PP sizing is preserved regardless of the DSv4 opt-in. + # PP sizing is independent of the headroom opt-in. (2, False, True, 2), (4, False, True, 4), (4, True, False, 4), @@ -120,6 +132,119 @@ def test_compute_max_num_sequences_scopes_overlap_headroom( ) +@pytest.mark.parametrize( + "cache_transceiver_config,expected", + [ + (None, False), + (Mock(backend=None), False), + (Mock(backend="NIXL"), True), + ], +) +def test_is_disagg_enabled(cache_transceiver_config, expected): + assert is_disagg_enabled(cache_transceiver_config) is expected + + +# (pp_size, enable_overlap_headroom, expected_factor). The factor is relative +# to max_batch_size and, for disagg, must cover the IndexMapper's 2x. +DISAGG_SIZING_CASES = [ + # Disagg gets the coefficient on topology alone, with no headroom opt-in. + (1, False, 2), + # The two factors overlap rather than compose, so this stays 2x. + (1, True, 2), + # PP composes: the IndexMapper is likewise sized pp_size * 2. + (2, False, 4), + (4, False, 8), +] + + +@pytest.mark.parametrize("pp_size,enable_overlap_headroom,expected_factor", DISAGG_SIZING_CASES) +def test_disagg_seats_cover_index_mapper_capacity( + pp_size, enable_overlap_headroom, expected_factor +): + max_batch_size = 8 + mapping = Mapping(world_size=pp_size, tp_size=1, pp_size=pp_size) + + seats = compute_max_num_sequences( + mapping, + max_batch_size, + disable_overlap_scheduler=False, + enable_overlap_headroom=enable_overlap_headroom, + is_disagg=True, + ) + + assert seats == max_batch_size * expected_factor + + # Seats must cover every request admission can let through. Mirrors the + # expression in KVCacheManagerV2.__init__. + index_mapper_capacity = max_batch_size * pp_size * 2 + assert seats >= index_mapper_capacity + + +@pytest.mark.parametrize("disable_overlap", [False, True]) +def test_disagg_seats_do_not_depend_on_overlap_scheduler(disable_overlap): + """The IndexMapper is sized the same either way, so seats must be too. + + Turning the overlap scheduler off removes the terminal-slot race but not + the transfer/generate overlap that the 2x coefficient exists for. + """ + max_batch_size = 8 + mapping = Mapping(world_size=1, tp_size=1, pp_size=1) + + assert ( + compute_max_num_sequences(mapping, max_batch_size, disable_overlap, is_disagg=True) + == max_batch_size * 2 + ) + + +def test_aggregate_sizing_is_unchanged(): + """Aggregated deployments size the pool at one forward batch.""" + max_batch_size = 8 + mapping = Mapping(world_size=1, tp_size=1, pp_size=1) + + assert ( + compute_max_num_sequences( + mapping, max_batch_size, disable_overlap_scheduler=False, is_disagg=False + ) + == max_batch_size + ) + + +@pytest.mark.parametrize("pp_size", [1, 2]) +@pytest.mark.parametrize("is_disagg", [False, True]) +def test_sizing_matches_kv_manager_admission_bound(pp_size, is_disagg): + """The two coefficients are computed independently; hold them in step. + + KVCacheManagerV2 derives its own admission bound from max_batch_size, + pp_size and is_disagg. Assert the two expressions against each other + rather than against a literal, so a drift in either one fails here. + """ + max_batch_size = 8 + mapping = Mapping(world_size=pp_size, tp_size=1, pp_size=pp_size) + + seats = compute_max_num_sequences( + mapping, max_batch_size, disable_overlap_scheduler=False, is_disagg=is_disagg + ) + admission_bound = max_batch_size * pp_size * (2 if is_disagg else 1) + + assert seats >= admission_bound + + +def test_validate_seq_slot_pool_accepts_sufficient_pool(): + validate_seq_slot_pool_covers_admission(16, Mock(max_admissible_sequences=16)) + + +def test_validate_seq_slot_pool_rejects_undersized_pool(): + with pytest.raises(ValueError, match="smaller than the number of"): + validate_seq_slot_pool_covers_admission(8, Mock(max_admissible_sequences=16)) + + +def test_validate_seq_slot_pool_ignores_managers_without_a_bound(): + """The V1/C++ manager does not publish one; the check must not fire.""" + manager = Mock(spec=[]) + validate_seq_slot_pool_covers_admission(1, manager) + validate_seq_slot_pool_covers_admission(1, None) + + @pytest.mark.parametrize("slot_factor", [1, 2]) def test_sampler_uses_executor_slot_pool_capacity(slot_factor): max_batch_size = 8