From f84317a5bf0b0e66f503c1e7906cb737cd04f79e Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Wed, 9 Sep 2026 16:32:28 -0700 Subject: [PATCH] [None][fix] Size seq-slot pool to cover disagg-gen KV admission On a disaggregated generation server KVCacheManagerV2 sizes its IndexMapper at 2 * max_batch_size * pp_size, so a batch still receiving KV can overlap a batch that is decoding, and the V2 scheduler admits DISAGG_GENERATION_INIT requests against that capacity rather than against the per-iteration request budget. The sequence-slot pool was still sized at max_batch_size, so the generation server could admit roughly twice what it was able to seat. Nothing detects the shortfall at admission time. The seat is claimed much later, once the KV transfer lands, in _prepare_disagg_gen_transmission_complete -> SeqSlotManager.prepare_resources -> add_slot. That runs on the executor's event-loop thread, so the raise kills the thread, the rank stops joining its TP collectives, and the peers block until the hang detector aborts the job 300 s later on an unrelated rank. Mirror the IndexMapper coefficient in compute_max_num_sequences. max() rather than another multiplication: the disagg factor and the attention-DP overlap headroom both cover one extra set of in-flight sequences, so they overlap rather than compose. Aggregated deployments keep their established sizing. Every seat holder on a generation server necessarily holds an IndexMapper slot first: KV is allocated during scheduling, the seat at transmission-complete, and both are released together in the same free_resources fan-out. So a pool matching the IndexMapper's non-reserved capacity can never come up short. Three slot-indexed buffers were still keyed to the attention-DP overlap opt-in rather than to the pool: spec-decode draft_probs via _set_up_spec_metadata and via the engine runner config, and GuidedDecoder's per-request state. Left alone, the wider pool would push py_seq_slot past the end of all three. They now follow max_num_seq_slots, which also closes the same overflow under pipeline parallelism, where slots already reached max_batch_size * pp_size. Finally, validate_seq_slot_pool_covers_admission compares the two bounds when the executor is built, so future drift between the coefficients fails at startup naming both numbers instead of deadlocking under load. It no-ops for the V1 manager, which publishes no admission bound. Exposure is any model running disagg generation with pipeline_parallel_size=1 and enough load to saturate one rank's pool. Attention DP raises the probability rather than changing the mechanism, since each rank owns a private pool and router skew concentrates the load on one of them. test_seq_slot_sizing.py pins this arithmetic but was in no test-db list, so it ran nowhere. Add it to l0_a10, which already carries the other CPU-only executor unit tests; these are pure sizing arithmetic over Mapping and Mock. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/_util.py | 68 ++++++-- .../kv_cache/kv_cache_manager_v2.py | 8 +- .../_torch/pyexecutor/model_engine.py | 20 +-- .../_torch/pyexecutor/py_executor_creator.py | 26 +-- .../_torch/pyexecutor/seq_slot_manager.py | 4 +- tensorrt_llm/_torch/speculative/interface.py | 11 +- tensorrt_llm/_torch/speculative/utils.py | 3 +- .../integration/test_lists/test-db/l0_a10.yml | 1 + .../_torch/executor/test_seq_slot_sizing.py | 153 ++++++++++++++++-- 9 files changed, 233 insertions(+), 61 deletions(-) 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