diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 78d1b6d921aa..2ee2838cbe3f 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, @@ -687,6 +694,7 @@ def __init__( model_engine) self._is_kv_cache_manager_v2 = issubclass(self._kv_cache_manager_cls, KVCacheManagerV2) + self._disable_overlap_scheduler = llm_args.disable_overlap_scheduler self._draft_config = draft_config self._skip_est = skip_est # Admission cap (tokens of summed context attended-KV) that the fp8 context-MLA workspace reservation @@ -1501,6 +1509,7 @@ def _create_kv_cache_manager( execution_stream=self._execution_stream, layer_mask=spec_dec_layer_mask, is_disagg=self._is_disagg, + disable_overlap_scheduler=self._disable_overlap_scheduler, kv_events_config=None if estimating_kv_cache or model_engine.is_draft_model else self._llm_args.kv_cache_config.kv_events_config, @@ -1721,6 +1730,7 @@ def _create_one_model_draft_kv_cache_manager( layer_mask=spec_dec_layer_mask, num_layers=num_draft_layers, is_disagg=self._is_disagg, + disable_overlap_scheduler=self._disable_overlap_scheduler, cold_page_codec_provider=cold_page_codec_provider, joint_kv_cache_reuse=self._joint_kv_cache_reuse, ) @@ -2097,6 +2107,7 @@ def _create_cross_kv_cache_manager( num_layers=num_layers, num_kv_heads=num_kv_heads, head_dim=head_dim, + disable_overlap_scheduler=self._disable_overlap_scheduler, kv_cache_type=tensorrt_llm.bindings.internal.batch_manager. CacheType.CROSS, ) @@ -2416,6 +2427,7 @@ def _create_kv_cache_manager( head_dim: Optional[int] = None, kv_cache_type=None, is_disagg: bool = False, + disable_overlap_scheduler: bool = False, cold_page_codec_provider: Optional[object] = None, kv_events_config: Optional[KVEventsConfig] = None, joint_kv_cache_reuse: bool = False) -> KVCacheManager: @@ -2569,6 +2581,8 @@ def _create_kv_cache_manager( "cold_page_codec_provider"] = cold_page_codec_provider manager_extra_kwargs["kv_events_config"] = kv_events_config manager_extra_kwargs["joint_kv_cache_reuse"] = joint_kv_cache_reuse + manager_extra_kwargs[ + "disable_overlap_scheduler"] = disable_overlap_scheduler # V2 builds the block-reuse cache key of a multimodal token run from # the vocabulary size. Resolve it here rather than per-branch: the # manager needs it whenever block reuse can meet multimodal input, @@ -3111,11 +3125,25 @@ def compute_max_num_sequences(mapping: Mapping, enable_overlap_headroom: bool = False) -> int: """Size the sequence-slot pool (and the sampler state it indexes). - ``enable_overlap_headroom`` is intentionally opt-in. Disaggregated - attention-DP needs a second non-PP slot set because the V2 scheduler can - backfill seats before the overlap scheduler releases the previous + ``enable_overlap_headroom`` is intentionally opt-in, and + ``should_enable_overlap_headroom`` is the only thing that should set it: a + non-PP overlap-enabled run needs a second 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``. + + Disaggregation deliberately does **not** appear here. Its ``2x`` is an + IndexMapper-local concern: a request awaiting its KV transfer holds an + *index* lease while holding no seat at all -- + ``SeqSlotManager.prepare_resources`` skips ``DISAGG_GENERATION_INIT`` + requests outright and only seats them once the transmission completes -- and + ``PyExecutor``'s admission arithmetic bounds ``active_requests`` at + ``max_batch_size * pp_size`` either way (``_fetch_new_requests``, which + subtracts ``len(active_requests)`` outside attention DP). So the index pool + legitimately runs ahead of the seat pool, and propagating that factor here + would double sampler state, the eagerly allocated + ``[seats, draft_len, vocab]`` draft-probability tensors and the pinned-host + block-offset tables for leases that never occupy a seat. """ if mapping.has_pp(): num_micro_batches = mapping.pp_size @@ -3125,6 +3153,37 @@ def compute_max_num_sequences(mapping: Mapping, return max_batch_size * num_micro_batches +def resolve_max_num_sequences(model_engine, + mapping: Mapping, + max_batch_size: int, + llm_args, + max_num_sequences: Optional[int] = None) -> int: + """Resolve the seat-pool size for a consumer, without re-deriving it. + + Order of preference, and the order matters: + + 1. an explicitly supplied value -- the caller already has the number the + engine published; + 2. ``model_engine.max_num_seq_slots`` -- the engine's own pool, which is + what every seat-keyed pool was sized against; + 3. only then a fresh ``compute_max_num_sequences``, reusing the engine's + headroom gate so the fallback cannot size the pool below the engine's. + """ + if max_num_sequences is not None: + return max_num_sequences + engine_seats = getattr(model_engine, "max_num_seq_slots", None) + if engine_seats is not None: + return engine_seats + # Engines that predate the attribute (unit-test stubs, mm-encoder-only + # engines): recompute, but with the same gate the engine would have used. + return compute_max_num_sequences(mapping, + max_batch_size, + llm_args.disable_overlap_scheduler, + enable_overlap_headroom=getattr( + model_engine, + "_enable_overlap_headroom", False)) + + def should_enable_adp_dummy_fixes(mapping: Mapping) -> bool: """Enable transactional ADP dummy handling while PP remains follow-up.""" return not mapping.has_pp() @@ -3149,15 +3208,107 @@ def should_enable_non_overlap_adp_forward_intent( and disable_overlap_scheduler) -def should_enable_disagg_adp_overlap_headroom( - mapping: Mapping, - 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) +def should_enable_overlap_headroom(mapping: Mapping, + disable_overlap_scheduler: bool, + kv_cache_manager_is_v2: bool, + is_hybrid: bool = False) -> bool: + """Gate the extra micro-batch of sequence slots (and its ADP counterpart). + + The predicate is deliberately the same shape as the one + ``KVCacheManagerV2`` uses for its index leases -- non-PP and overlap-on -- + so the seat pool and the index pool move together instead of the seat pool + tracking a narrower condition that has to be re-derived and kept in step. + The mechanism it pays for is ``ADPRouter``'s ``exclude_retiring_requests``: + it drops ``GENERATION_TO_COMPLETE`` requests from the per-rank active counts + that ``_fetch_new_requests`` subtracts from the admission capacity, so a + replacement cohort is admitted while the retiring cohort still holds its + seats. + + * ``not has_pp()`` -- pipeline parallelism already carries ``pp_size`` + micro-batches, and its interaction with the retirement filter is + unvalidated. Out of scope here. + * ``not disable_overlap_scheduler`` -- with the overlap scheduler off a + terminal request is torn down in-line, so it never coexists with its + replacement. + * ``kv_cache_manager_is_v2`` -- only ``KVCacheV2Scheduler`` stops counting a + request at ``GENERATION_TO_COMPLETE``. The V1 capacity schedulers (C++ + ``BindCapacityScheduler`` and ``PyCapacityScheduler``) hardcode + ``GENERATION_COMPLETE``, so under V1 the retiring cohort is still inside + the capacity budget and the headroom would only inflate seat-keyed pools. + * ``not is_hybrid`` -- SSM/Mamba state is sized by ``max_batch_size`` in + every manager, and ``_max_resident_sequences()`` is the floor a hybrid + manager can serve; extra seats there would let residency outrun the SSM + slots. (Hybrid state is not ``py_seq_slot``-indexed, so the failure is + exhaustion rather than an out-of-bounds write -- still a hang.) + + Attention DP is deliberately *not* a condition, even though it is the only + deployment where the extra seats are ever occupied: outside attention DP + ``_fetch_new_requests`` subtracts ``len(active_requests)`` with retirees + included, so residency stays at ``max_batch_size * pp_size`` and the surplus + seats sit idle. Including it would make the seat pool narrower than the index + pool on the ordinary TP overlap path, which is the asymmetry that has to be + tolerated by ``validate_seq_slot_pool_covers_admission`` and re-justified at + every reader. Disaggregation is not a condition either, for the opposite + reason: its ``2x`` is genuinely index-local, since a transfer-phase request + holds a lease and no seat at all. + """ + if is_hybrid or not kv_cache_manager_is_v2: + return False + return 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 KV index pool cannot cover the seat pool. + + One direction is a bug and the other is by design, so the check is + deliberately one-sided rather than an equality: + + * **index pool < seat pool -- always a bug.** ``_create_kv_cache`` warns + ``No free IndexMapper slots``, returns ``None`` and the scheduler defers + the request. Silent: costs throughput, raises nothing (nvbug 6627795). A + one-sided ``seats >= index pool`` guard is exactly what let this ship, and + this is the direction that guard could not see. + * **index pool > seat pool -- permitted.** The two predicates agree on the + overlap term (non-PP and overlap-on widens both), so the common case is + equality; where they differ the index pool is the more generous one, and + that is by design. Under disaggregation a request awaiting its KV transfer + holds an index lease and no seat at all + (``SeqSlotManager.prepare_resources`` skips ``DISAGG_GENERATION_INIT``), so + the manager's ``is_disagg`` term has no seat-pool counterpart. And a hybrid + model suppresses the seat headroom -- SSM state is sized from + ``max_batch_size`` -- while a V2 index pool behind it still widens. Spare + leases cost a few page-table rows; spare seats would cost sampler state, + the eagerly allocated ``[seats, draft_len, vocab]`` draft-probability + tensors and the pinned-host block-offset tables, which is why the surplus + is kept on this side. + + The cost of the asymmetry is that an *undersized seat pool* is no longer + distinguishable from a legitimately generous index pool, so it is not caught + here. ``compute_max_num_sequences`` being the single seat-pool definition is + what covers that side. + + Managers that do not publish ``max_admissible_sequences`` (V1, non-V2 + hybrid) are skipped rather than guessed at. ``isinstance`` rather than + ``is not None``: the V1/C++ manager does not publish the attribute at all, + and neither do the ``Mock`` stubs that stand in for a cache manager in other + modules' tests -- a ``Mock`` auto-creates it, so an ``is None`` opt-in would + fall through to the comparison and raise ``TypeError`` instead of skipping + the check. + """ + admissible = getattr(kv_cache_manager, "max_admissible_sequences", None) + if not isinstance(admissible, int): + return + if admissible >= max_num_sequences: + return + raise ValueError( + f"{type(kv_cache_manager).__name__} can lease KV cache indices for " + f"{admissible} concurrent sequences but the executor's sequence-slot " + f"pool holds {max_num_sequences}: the index pool is smaller than the " + "seat pool, so admitted requests would be silently deferred one at a " + "time (nvbug 6627795). The seat pool must come from " + "_util.compute_max_num_sequences and the index pool must cover it; a " + "shortfall means one of them was re-derived from max_batch_size.") def create_py_executor_instance( @@ -3195,15 +3346,18 @@ def create_py_executor_instance( spec_config = model_engine.spec_config - if max_num_sequences is None: - max_num_sequences = compute_max_num_sequences( - mapping, max_batch_size, llm_args.disable_overlap_scheduler) + is_disagg = is_disagg_enabled(cache_transceiver_config) + + max_num_sequences = resolve_max_num_sequences( + model_engine, + mapping, + max_batch_size, + llm_args, + max_num_sequences=max_num_sequences) 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( @@ -3361,6 +3515,11 @@ def create_py_executor_instance( if isinstance(model_engine, PyTorchModelEngine): model_engine._init_cuda_graph_lora_manager(lora_config) + # The seat pool and the KV index pool are sized independently and indexed by + # the same request identity, so an index pool narrower than the seat pool is + # always a bug. Check it here rather than at first use: a startup ValueError + # names both numbers, while the runtime symptom is a silent throughput loss. + validate_seq_slot_pool_covers_admission(max_num_sequences, kv_cache_manager) resources[ResourceManagerType.SEQ_SLOT_MANAGER] = SeqSlotManager( max_num_sequences) @@ -3392,8 +3551,10 @@ def create_py_executor_instance( # When scheduler_capacity == 1, attention dp dummy request will prevent the scheduling of DISAGG_GENERATION_INIT. # Enlarge scheduler capacity to avoid DISAGG_GENERATION_INIT stuck in the scheduler. - # V1 scheduler handles overlap via two_step_lookahead, so skip the - # slot-pool overlap factor here. + # V1 scheduler handles overlap via two_step_lookahead, so the capacity + # scheduler's budget stays at the pipeline-depth bound and deliberately does + # not follow the sequence-slot pool: the overlap headroom is spare seats for + # leases already held, not extra admission. scheduler_capacity = max_batch_size * mapping.pp_size if scheduler_capacity == 1 and mapping.enable_attention_dp and kv_cache_manager: scheduler_capacity += 1 @@ -3580,22 +3741,22 @@ def create_py_executor_instance( def create_torch_sampler_args( - mapping: Mapping, *, max_seq_len: int, - max_batch_size: int, speculative_config: SpeculativeConfig, max_beam_width: int, disable_overlap_scheduler: bool, enable_async_worker: bool, enable_speculative_beam_history_d2h: bool, - max_num_sequences: Optional[int] = None, + max_num_sequences: int, ): - # The sampler's per-slot state is indexed by sequence slots, so it must - # be sized identically to the executor's slot pool. - if max_num_sequences is None: - max_num_sequences = compute_max_num_sequences( - mapping, max_batch_size, disable_overlap_scheduler) + # The sampler's per-slot state is indexed by sequence slots, so it must be + # sized identically to the executor's slot pool. `max_num_sequences` is + # required, not optional: the old default recomputed the pool from + # `mapping`/`max_batch_size` *without* the overlap-headroom gate, so it could + # only ever produce a smaller number than the slots it indexes. Those two + # parameters are gone with it -- keeping them would leave the raw material + # for the same re-derivation lying next to the resolved value. max_draft_len = (0 if speculative_config is None else speculative_config.max_draft_len) max_total_draft_tokens = (0 if speculative_config is None else @@ -3627,10 +3788,15 @@ def instantiate_sampler( enable_async_worker = (confidential_compute_enabled() or llm_args.sampler_force_async_worker) - sampler_args = create_torch_sampler_args( + max_num_sequences = resolve_max_num_sequences( + engine, mapping, + max_batch_size, + llm_args, + max_num_sequences=max_num_sequences) + + sampler_args = create_torch_sampler_args( max_seq_len=engine.max_seq_len, - max_batch_size=max_batch_size, speculative_config=speculative_config, max_beam_width=max_beam_width, disable_overlap_scheduler=llm_args.disable_overlap_scheduler, 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 5ab91921d68a..37724b5dc72f 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 @@ -969,6 +969,7 @@ def __init__( is_disagg: bool = False, enable_stats: bool = False, num_reserved_index_slots: int = 1, + disable_overlap_scheduler: bool = False, kv_events_config: Optional[KVEventsConfig] = None, is_estimating_kv_cache: bool = False, cold_page_codec_provider: Optional[object] = None, @@ -1512,20 +1513,48 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: # simultaneously, so we need slots for all concurrent sequences. # Reserve stable slots for persistent request IDs such as CUDA-graph # padding requests. The default preserves the main-branch allocation. - # In disaggregated mode, use a coefficient of 2: at any moment up to - # `max_num_sequences` requests can be actively generating while another - # up to `max_num_sequences` requests are still in KV transfer - # (TRANS_IN_PROGRESS) and continue to hold their index slots. The 2x - # capacity lets the next batch of active requests acquire slots without - # waiting for the previous batch's transfers to finish. + # + # Two independent reasons to hold twice that many index leases, and they + # do not compose -- a request is either awaiting a transfer or retiring, + # never both -- so this is one doubling, not two: + # + # * Disaggregation: at any moment up to `max_num_sequences` requests can + # be actively generating while another up to `max_num_sequences` are + # still in KV transfer (TRANS_IN_PROGRESS) and continue to hold their + # index slots. The extra capacity lets the next batch acquire slots + # without waiting for the previous batch's transfers to finish. + # * The overlap scheduler: a terminal request's resources are released + # one iteration after it leaves the batch, so with overlap enabled a + # retiring cohort and its replacement hold index leases at the same + # time. Pipeline parallelism is excluded because `pp_size` micro-batches + # are already covered above. + # + # The overlap term is shared with the seat pool by construction: + # _util.should_enable_overlap_headroom uses the same `not has_pp() and + # overlap on` shape, so a retiring cohort that outlives its seat is + # covered on both sides and the two pools stay equal in the ordinary case. + # The disagg term has no seat-pool counterpart on purpose -- a + # transfer-phase request holds no sequence slot at all + # (SeqSlotManager.prepare_resources skips DISAGG_GENERATION_INIT) -- and + # the seat gate additionally suppresses itself for hybrid/SSM models and + # under the V1 capacity schedulers. So where the two disagree this pool is + # the more generous one: `validate_seq_slot_pool_covers_admission` requires + # it to cover the seat pool, not to equal it. Spare index leases cost a + # few page-table rows; spare seats would cost sampler and + # speculative-decoding state. max_num_sequences = max_batch_size * mapping.pp_size + needs_extra_leases = is_disagg or (not disable_overlap_scheduler and not mapping.has_pp()) 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 needs_extra_leases 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}, " + f"disable_overlap_scheduler={disable_overlap_scheduler}, " + f"pp_size={mapping.pp_size}, " f"num_reserved_index_slots={num_reserved_index_slots}, " f"max_beam_width={max_beam_width})" ) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 0fae89ad1b9b..1e904856d4dd 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -78,6 +78,7 @@ set_per_request_prefill_cuda_graph_flag, set_torch_compiling, with_model_extra_attrs) from .breakable_cuda_graph_runner import BreakableCUDAGraphRunner +from .config_utils import is_hybrid_linear from .cuda_graph_runner import (ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM, CUDAGraphRunner, CUDAGraphRunnerConfig, EncoderCUDAGraphRunner, @@ -448,24 +449,12 @@ 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, should_enable_adp_dummy_fixes, - should_enable_disagg_adp_overlap_headroom, should_enable_non_overlap_adp_forward_intent, + should_enable_overlap_headroom, should_enable_scheduler_aware_adp_dummy) - self._enable_disagg_adp_overlap_headroom = ( - should_enable_disagg_adp_overlap_headroom( - mapping, llm_args.cache_transceiver_config, - llm_args.disable_overlap_scheduler)) self._enable_adp_dummy_fixes = should_enable_adp_dummy_fixes(mapping) - self.max_num_seq_slots = compute_max_num_sequences( - mapping, - self.batch_size, - llm_args.disable_overlap_scheduler, - enable_overlap_headroom=self._enable_disagg_adp_overlap_headroom, - ) self.dist = dist if dist is not None: ExpertStatistic.create(self.dist.rank) @@ -567,6 +556,25 @@ def __init__( self._enable_non_overlap_adp_forward_intent = ( should_enable_non_overlap_adp_forward_intent( mapping, llm_args.disable_overlap_scheduler)) + # With the overlap scheduler on, a retiring request keeps its sequence + # slot for one more iteration while its replacement is admitted (under + # attention DP the router drops it from the per-rank counts admission + # subtracts from, so residency actually reaches the extra cohort). The + # gate needs the pretrained config -- hybrid SSM state is sized by + # max_batch_size and cannot absorb the extra cohort -- so it is resolved + # here rather than before the model load. + self._enable_overlap_headroom = should_enable_overlap_headroom( + mapping, + llm_args.disable_overlap_scheduler, + kv_cache_manager_is_v2=( + llm_args.kv_cache_config.use_kv_cache_manager_v2 is True), + is_hybrid=is_hybrid_linear(pretrained_config)) + self.max_num_seq_slots = compute_max_num_sequences( + mapping, + self.batch_size, + llm_args.disable_overlap_scheduler, + enable_overlap_headroom=self._enable_overlap_headroom, + ) self.sparse_attention_config = self.model.model_config.sparse_attention_config # In case that some tests use stub models and override `_load_model`. if not hasattr(self.model, 'extra_attrs'): @@ -1131,8 +1139,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 +3716,6 @@ 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) if self.spec_metadata is not None: return self.spec_metadata self.spec_metadata = get_spec_metadata( @@ -3724,7 +3726,7 @@ def _set_up_spec_metadata( spec_resource_manager=spec_resource_manager, is_draft_model=self.is_draft_model, max_seq_len=self.max_seq_len, - num_seq_slots=num_seq_slots) + num_seq_slots=self.max_num_seq_slots) return self.spec_metadata def cleanup(self) -> None: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 48ed6c771255..3d9f2393a567 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -110,7 +110,7 @@ from .scheduler import (RequestScheduler, ScheduledRequests, SerializableSchedulerOutput, WaitingQueue, create_waiting_queue) -from .scheduler.adp_router import ADPRouter +from .scheduler.adp_router import ADPRouter, count_retiring_requests if TYPE_CHECKING: from ray.actor import ActorHandle @@ -721,6 +721,8 @@ def __init__( # can receive the transfer-manager reference at construction time. self.adp_router: ADPRouter = ADPRouter.create( dist=self.dist, + has_seq_slot_headroom=getattr(model_engine, + "_enable_overlap_headroom", False), kv_cache_manager=self.kv_cache_manager, attention_dp_config=self.llm_args.attention_dp_config, async_transfer_manager=self.async_transfer_manager, @@ -5806,7 +5808,7 @@ def _validate_request(self, request: LlmRequest): self._validate_request_budget(request) def _fetch_and_enqueue_requests(self, waiting_queue: WaitingQueue, - total_num_active_requests: int) -> None: + total_num_live_requests: int) -> None: """Fetch requests from request_queue and enqueue to waiting_queue.""" # Block new requests while control requests are pending if len(self.control_requests) != 0: @@ -5817,7 +5819,7 @@ def _fetch_and_enqueue_requests(self, waiting_queue: WaitingQueue, # blocking would keep the loop from reaching the # `should_stop_processing` check that ends it, deadlocking shutdown() # on `shutdown_event`. - idle = (total_num_active_requests == 0 and len(waiting_queue) == 0 + idle = (total_num_live_requests == 0 and len(waiting_queue) == 0 and not self.is_shutdown) if idle: # In Ray path (TLLM_DISABLE_MPI=1), use a periodic heartbeat timeout so rank 0 @@ -6025,14 +6027,16 @@ def _fetch_new_requests( s.num_active_requests for s in all_rank_states ] total_num_active_requests = sum(all_ranks_num_active_requests) + total_num_live_requests = total_num_active_requests + sum( + s.num_retiring_requests for s in all_rank_states) else: total_num_active_requests = len(active_requests) + total_num_live_requests = total_num_active_requests all_ranks_num_active_requests = None all_rank_states = None # 2. Fetch and enqueue to waiting queue - self._fetch_and_enqueue_requests(waiting_queue, - total_num_active_requests) + self._fetch_and_enqueue_requests(waiting_queue, total_num_live_requests) # 3. Pop requests from waiting queue new_requests = self._pop_from_waiting_queue( @@ -7064,7 +7068,11 @@ def _pad_attention_dp_dummy_request(self): return expected_num_active_requests = self.expected_num_active_requests - if expected_num_active_requests < len(self.active_requests): + num_routable_active_requests = len(self.active_requests) + if self.adp_router.exclude_retiring_requests: + num_routable_active_requests -= count_retiring_requests( + self.active_requests) + if expected_num_active_requests < num_routable_active_requests: # Not fatal, and not a capacity violation. The router derives this # value as # min(max(ceil(multiplier * fair_share), max(per_rank_loads)), @@ -7085,11 +7093,12 @@ def _pad_attention_dp_dummy_request(self): # event loop on every affected rank at once, leaving the survivors # to HangDetector-abort. logger.warning( - f"active_requests ({len(self.active_requests)}) exceeds " + f"routable active_requests " + f"({num_routable_active_requests}) exceeds " f"expected_num_active_requests " f"({expected_num_active_requests}); tolerating (a busy rank " f"needs no attention-DP dummy).") - expected_num_active_requests = len(self.active_requests) + expected_num_active_requests = num_routable_active_requests num_active_request = self._count_schedulable_active_requests() diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 3f1a703891ec..e83ceb052395 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -38,7 +38,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, @@ -617,8 +618,13 @@ 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)) + 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, + enable_overlap_headroom=getattr(model_engine, + "_enable_overlap_headroom", False)) if is_mla(config): if model_engine.model.model_config.enable_flash_mla: tokens_per_block = 64 @@ -712,15 +718,12 @@ 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, + # The guided decoder's state is indexed by py_seq_slot + # (guided_decoder.py: grammar_matchers[req.seq_slot], the + # bitmask rows), so it must span the whole seat pool. + "max_num_sequences": max_num_seq_slots, "vocab_size_padded": model_engine.model.vocab_size_padded, "rank": mapping.rank, } @@ -831,8 +834,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/scheduler/adp_router.py b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py index bf9dd99d2e02..a42526ab1f27 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py @@ -27,6 +27,8 @@ from tensorrt_llm.logger import logger +from ..llm_request import LlmRequestState + if TYPE_CHECKING: from tensorrt_llm._torch.distributed.communicator import Distributed from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest @@ -47,6 +49,21 @@ def _num_input_tokens(request) -> int: return len(getattr(request, "input_token_ids", [])) +def is_retiring_request(request) -> bool: + """True if ``request`` has produced its final token and is being torn down.""" + return request.state == LlmRequestState.GENERATION_TO_COMPLETE + + +def build_active_requests_for_overlap(active_requests): + """Return ``active_requests`` without the requests that are already retiring.""" + return [req for req in active_requests if not is_retiring_request(req)] + + +def count_retiring_requests(active_requests) -> int: + """Count the retiring requests in ``active_requests``.""" + return sum(1 for req in active_requests if is_retiring_request(req)) + + @dataclass class RankIterStatsPayload: """Per-rank IterationStats payload piggybacked on the ADP allgather.""" @@ -99,6 +116,7 @@ class RankState: rank: int num_active_requests: int = 0 num_active_tokens: int = 0 + num_retiring_requests: int = 0 iter_stats: RankIterStatsPayload = field(default_factory=RankIterStatsPayload) def copy_iter_stats_from(self, iter_stats_payload: RankIterStatsPayload | None) -> None: @@ -112,6 +130,7 @@ def serialize(self) -> list[int]: self.rank, self.num_active_requests, self.num_active_tokens, + self.num_retiring_requests, *self.iter_stats.serialize(), ] @@ -119,7 +138,7 @@ def serialize(self) -> list[int]: def deserialize(cls, data: list[int]) -> RankState: """Deserialize from a flat list received via allgather.""" values = list(data) - rank_state_prefix_field_count = 3 + rank_state_prefix_field_count = 4 rank_state_fields = fields(cls)[:rank_state_prefix_field_count] max_field_count = rank_state_prefix_field_count + len(fields(RankIterStatsPayload)) if len(values) < 1: @@ -140,6 +159,7 @@ def deserialize(cls, data: list[int]) -> RankState: rank=rank_values[0], num_active_requests=rank_values[1], num_active_tokens=rank_values[2], + num_retiring_requests=rank_values[3], iter_stats=RankIterStatsPayload.deserialize(values[rank_state_prefix_field_count:]), ) @@ -163,13 +183,21 @@ class ADPRouter(ABC): needs_prefix_matches: bool = False - def __init__(self, dist: Distributed): + def __init__(self, dist: Distributed, has_seq_slot_headroom: bool = False): self.dist = dist + # Defaults off, and deliberately so: dropping retiring requests from the + # per-rank counts admission subtracts from its capacity lets residency + # exceed max_batch_size * pp_size, which is only safe when the seat pool + # was sized for it (_util.should_enable_overlap_headroom). A caller that + # forgets the flag should get the conservative accounting, not a pool + # overrun. + self.exclude_retiring_requests = has_seq_slot_headroom @classmethod def create( cls, dist: "Distributed", + has_seq_slot_headroom: bool, kv_cache_manager=None, attention_dp_config=None, async_transfer_manager=None, @@ -178,6 +206,9 @@ def create( Args: dist: Distributed communicator. + has_seq_slot_headroom: Whether the executor's sequence-slot pool was + sized with the extra overlap headroom. The retiring-request + correction is only applied when those extra seats exist. kv_cache_manager: KV cache manager instance (may be None). attention_dp_config: AttentionDpConfig instance (may be None). async_transfer_manager: PyExecutor's AsyncTransferManager, used by @@ -199,6 +230,7 @@ def create( # KV-cache-aware path and takes precedence when both are enabled. return ConversationAwareADPRouter( dist=dist, + has_seq_slot_headroom=has_seq_slot_headroom, max_sessions=attention_dp_config.kv_cache_routing_max_sessions, fair_share_multiplier=attention_dp_config.kv_cache_routing_fair_share_multiplier, new_conv_placement=attention_dp_config.kv_cache_routing_new_conv_placement, @@ -212,6 +244,7 @@ def create( ): return KVCacheAwareADPRouter( dist=dist, + has_seq_slot_headroom=has_seq_slot_headroom, kv_cache_manager=kv_cache_manager, load_balance_weight=attention_dp_config.kv_cache_routing_load_balance_weight, match_rate_threshold=attention_dp_config.kv_cache_routing_match_rate_threshold, @@ -221,7 +254,7 @@ def create( account_for_in_transfer=attention_dp_config.kv_cache_routing_account_for_in_transfer, ) - return DefaultADPRouter(dist=dist) + return DefaultADPRouter(dist=dist, has_seq_slot_headroom=has_seq_slot_headroom) @abstractmethod def create_rank_state( @@ -256,7 +289,19 @@ def gather_all_rank_states( iter_stats_payload: Completed previous-iteration stats payload to piggyback on this allgather, if one is pending. """ - local_state = self.create_rank_state(active_requests, new_requests or []) + # A request whose teardown the overlap scheduler has merely deferred is + # not load, and must not hold admission capacity that nothing can spend + # (nvbug 6627795). + if self.exclude_retiring_requests: + active_requests_for_overlap = build_active_requests_for_overlap(active_requests) + num_retiring_requests = len(active_requests) - len(active_requests_for_overlap) + else: + active_requests_for_overlap = active_requests + num_retiring_requests = 0 + local_state = self.create_rank_state(active_requests_for_overlap, new_requests or []) + # Reported separately so the executor loop can still tell that these + # ranks are not idle and keep its idle-fetch wait collective. + local_state.num_retiring_requests = num_retiring_requests local_state.copy_iter_stats_from(iter_stats_payload) responses = self.dist.tp_allgather(local_state.serialize()) return [RankState.deserialize(data=resp) for resp in responses] @@ -507,6 +552,7 @@ def __init__( self, dist: "Distributed", kv_cache_manager, + has_seq_slot_headroom: bool = False, load_balance_weight: float = 1.0, match_rate_threshold: float = 0.1, fair_share_multiplier: float = 2.0, @@ -514,7 +560,7 @@ def __init__( async_transfer_manager=None, account_for_in_transfer: bool = False, ): - super().__init__(dist) + super().__init__(dist, has_seq_slot_headroom=has_seq_slot_headroom) self.kv_cache_manager = kv_cache_manager self.load_balance_weight = load_balance_weight self.match_rate_threshold = match_rate_threshold @@ -812,11 +858,12 @@ class ConversationAwareADPRouter(ADPRouter): def __init__( self, dist: "Distributed", + has_seq_slot_headroom: bool = False, max_sessions: int = DEFAULT_MAX_SESSIONS, fair_share_multiplier: float = 2.0, new_conv_placement: str = "round_robin", ): - super().__init__(dist) + super().__init__(dist, has_seq_slot_headroom=has_seq_slot_headroom) self._conv_to_rank: "OrderedDict[str, int]" = OrderedDict() self._max_sessions = max(1, int(max_sessions)) self._fair_share_multiplier = max(1.0, float(fair_share_multiplier)) @@ -990,8 +1037,8 @@ def _next_rr(soft_cap: int) -> int: # Sticky returns use the hard cap, so a rank may now exceed the pre-loop # soft `expected`. Re-bump so the returned value covers the actual - # per-rank max -- _pad_attention_dp_dummy_request asserts - # expected >= len(active_requests) on every rank. + # per-rank max -- _pad_attention_dp_dummy_request compares `expected` + # against each rank's routable active count. expected_num_active_requests = max( expected_num_active_requests, max(all_ranks_num_active_requests) ) diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index 5d6e2011d844..69af428816ba 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -43,7 +43,8 @@ def __init__(self, max_num_requests: int, max_seq_len: int, max_num_tokens: int, - sa_manager=None): + sa_manager=None, + num_seq_slots: Optional[int] = None): self.dtype = dtype self.max_draft_len = config.max_draft_len self.hidden_size = hidden_size @@ -51,9 +52,17 @@ def __init__(self, self.max_seq_len = max_seq_len # Optional SA manager for EAGLE3+SA mode self.sa_manager = sa_manager + # ``slot_manager`` hands out slots keyed by request id and holds them for + # the request's whole lifetime, so the pool must span the executor's + # sequence-slot pool -- 2 * max_batch_size under the overlap headroom, + # where a retiring request keeps its slot for one more iteration while its + # replacement is admitted (nvbug-6627795). None means no headroom. + self.num_seq_slots = max(num_seq_slots or 0, max_num_requests) # There could be dummy request for padding batch when using CUDA graph. # Reserve one more slot for the dummy request. - slot_size = self.max_seq_len + 1 + # NOTE: max_seq_len is kept as a floor purely to preserve the historical + # (over-)sizing; it is a token count, not a slot count. + slot_size = max(self.num_seq_slots, self.max_seq_len) + 1 self.slot_manager = SlotManager(slot_size) # This class is reused by MTP_EAGLE from ...llmapi.llm_args import EagleDecodingConfig @@ -104,6 +113,7 @@ def __init__(self, max_total_draft_tokens=self.max_total_draft_tokens, eagle_choices=config.eagle_choices, dynamic_tree_max_topK=config.dynamic_tree_max_topK, + num_seq_slots=self.num_seq_slots, ) def prepare_resources(self, scheduled_batch: ScheduledRequests): @@ -165,8 +175,14 @@ class Eagle3OneModelDynamicTreeResourceManager(BaseResourceManager): hidden_states: Optional[torch.Tensor] = None batch_indices_cuda: Optional[torch.Tensor] = None - def __init__(self, config: "EagleDecodingConfig", max_num_requests: int): + def __init__(self, + config: "EagleDecodingConfig", + max_num_requests: int, + num_seq_slots: Optional[int] = None): self.max_num_requests = max_num_requests + # batch_indices_cuda is indexed by batch position, so it stays at + # max_batch_size; only the SpecTreeManager slot storage below is keyed by + # py_seq_slot and needs the executor's slot pool (nvbug-6627795). self.batch_indices_cuda = torch.empty( [max_num_requests], dtype=torch.int, @@ -179,6 +195,7 @@ def __init__(self, config: "EagleDecodingConfig", max_num_requests: int): max_total_draft_tokens=config.tokens_per_gen_step - 1, eagle_choices=config.eagle_choices, dynamic_tree_max_topK=config.dynamic_tree_max_topK, + num_seq_slots=num_seq_slots, ) def free_resources(self, request: LlmRequest): diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 4c271b5fdb9e..7ccbbdac6462 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -42,15 +42,26 @@ def __init__(self, dtype: torch.dtype, hidden_size: int, max_num_requests: int, - sa_manager=None): + sa_manager=None, + num_seq_slots: Optional[int] = None): self.dtype = dtype self.num_draft_slots = config.max_draft_len self.hidden_size = hidden_size self.max_num_requests = max_num_requests self.use_relaxed_acceptance_for_thinking = config.use_relaxed_acceptance_for_thinking + # These slots are keyed by live-request identity, not by batch position: + # add_slot runs on a request's first context chunk and the slot is only + # returned by free_resources. So the pool must cover every request that + # can be resident at once, which is the SeqSlotManager pool size + # (``num_seq_slots``) rather than max_batch_size -- under the overlap + # headroom the two differ by 2x, because a finished request holds its slot + # for one more iteration while its replacement is already admitted + # (nvbug-6627795). Sizing this at max_num_requests instead makes + # SlotManager.add_slot raise NoFreeSlotsError. Falls back to + # max_num_requests when the caller does not know the pool size. # Reserve one extra slot for the CUDA graph padding dummy request, # which is kept alive permanently and must not consume a real slot. - slot_pool_size = max_num_requests + 1 + slot_pool_size = (num_seq_slots or max_num_requests) + 1 self.slot_manager = SlotManager(slot_pool_size) # Optional SA manager for MTP+SA mode self.sa_manager = sa_manager diff --git a/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py b/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py index 86895f4c3558..021478a5fbec 100644 --- a/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py +++ b/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py @@ -1102,6 +1102,7 @@ def __init__( hidden_size: int, max_num_requests: int, sa_manager=None, + num_seq_slots: Optional[int] = None, ): from .spec_tree_manager import SpecTreeManager @@ -1113,10 +1114,18 @@ def __init__( max_total_draft_tokens=config.tokens_per_gen_step - 1, eagle_choices=None, dynamic_tree_max_topK=config.dynamic_tree_max_topK, + num_seq_slots=num_seq_slots, ) # MTP hidden-state slot pools (needed by MTPEagleWorker drafter inputs). + # num_seq_slots is forwarded because those pools are keyed by live-request + # identity; see MTPHiddenStatesManager.__init__. self._mtp_hidden_states_manager = MTPHiddenStatesManager( - config, dtype, hidden_size, max_num_requests, sa_manager=sa_manager + config, + dtype, + hidden_size, + max_num_requests, + sa_manager=sa_manager, + num_seq_slots=num_seq_slots, ) # Expose the MTPHiddenStatesManager surface MTPSpecMetadata expects. diff --git a/tensorrt_llm/_torch/speculative/spec_tree_manager.py b/tensorrt_llm/_torch/speculative/spec_tree_manager.py index 545b5bc7bb06..fc4ea5fe840c 100644 --- a/tensorrt_llm/_torch/speculative/spec_tree_manager.py +++ b/tensorrt_llm/_torch/speculative/spec_tree_manager.py @@ -1,7 +1,7 @@ import logging import math from itertools import accumulate -from typing import List +from typing import List, Optional import torch @@ -231,10 +231,14 @@ class SpecTreeManager: retrieve_next_sibling: torch.Tensor = None slot_storage: 'DynamicTreeSlotStorage | None' = None - def __init__(self, max_num_requests: int, use_dynamic_tree: bool, - max_total_draft_tokens: int, max_draft_len: int, + def __init__(self, + max_num_requests: int, + use_dynamic_tree: bool, + max_total_draft_tokens: int, + max_draft_len: int, eagle_choices: List[List[int]] | None, - dynamic_tree_max_topK: int): + dynamic_tree_max_topK: int, + num_seq_slots: Optional[int] = None): self.use_dynamic_tree = use_dynamic_tree self.max_total_draft_tokens = max_total_draft_tokens @@ -251,6 +255,15 @@ def __init__(self, max_num_requests: int, use_dynamic_tree: bool, self._internal_buf_dim = max_total_draft_tokens + 1 self.eagle_choices = eagle_choices self.num_trees = max_num_requests if use_dynamic_tree else 1 + # ``num_trees`` sizes the per-forward *work* buffers, which are indexed by + # batch position and so correctly stay at max_batch_size -- the micro-batch + # scheduler caps every forward there. ``num_slots`` sizes DynamicTreeSlotStorage, + # which is indexed by ``py_seq_slot`` and must therefore span the executor's + # sequence-slot pool: the overlap headroom makes that 2 * max_batch_size + # so a retiring request can keep its slot for one more iteration while + # its replacement is admitted (nvbug-6627795). + # None preserves the historical max_batch_size sizing. + self.num_slots = max(num_seq_slots or 0, max_num_requests) self.dynamic_tree_max_topK = dynamic_tree_max_topK self.cur_draft_layer_idx = 0 self.top_k_list = [] @@ -334,7 +347,7 @@ def init_tree_info_for_dynamic_tree(self): mask_width = math.ceil(num_draft_with_root / 32) self.slot_storage = DynamicTreeSlotStorage( - num_slots=self.num_trees, + num_slots=self.num_slots, n_dt=num_draft_with_root, mask_width=mask_width, top_k=self.dynamic_tree_max_topK, diff --git a/tensorrt_llm/_torch/speculative/suffix_automaton.py b/tensorrt_llm/_torch/speculative/suffix_automaton.py index b2d6ace23a9c..15170d9d4890 100644 --- a/tensorrt_llm/_torch/speculative/suffix_automaton.py +++ b/tensorrt_llm/_torch/speculative/suffix_automaton.py @@ -106,6 +106,7 @@ def __init__( config, max_num_requests: int, max_seq_len: int = 262144, + num_seq_slots: Optional[int] = None, ): if _sa_native is None: raise RuntimeError( @@ -144,14 +145,27 @@ def __init__( self.max_seq_len = sa_config.max_seq_len self.enable_global_pool = sa_config.enable_global_pool - # Pool sizing: effective_pool_size returns max_num_requests when - # global pool is off, or max(64, max_num_requests) / explicit - # value when on. All slot-indexed sizing uses pool_size. - self.pool_size = sa_config.effective_pool_size - if self.pool_size < max_num_requests: - raise ValueError( - f"global_pool_size ({self.pool_size}) must be >= " - f"max_batch_size ({max_num_requests})" + # A slot is held for the whole lifetime of a request id, so the pool has to + # cover every request that can be simultaneously live -- that is the + # executor's sequence-slot pool, which the overlap headroom raises to + # 2 * max_batch_size so a retiring request can keep its slot for one more + # iteration while its replacement is admitted (nvbug-6627795). + # None (no headroom) leaves this at max_batch_size, as before. + self._num_seq_slots = max(num_seq_slots or 0, max_num_requests) + + # Pool sizing: effective_pool_size returns max_slots when global pool is + # off, or max(64, max_slots) / the explicit value when on. All slot-indexed + # sizing uses pool_size, so the live-slot count is a floor on it. An + # explicit global_pool_size below that floor is grown rather than rejected: + # the same config is legal without the headroom, so failing here would turn + # enabling the overlap scheduler into a startup error for a value that + # TorchLlmArgs.validate_speculative_config already accepted. + self.pool_size = max(sa_config.effective_pool_size, self._num_seq_slots) + if sa_config.global_pool_size is not None and self.pool_size > sa_config.global_pool_size: + logger.warning( + f"Growing the SA pool from the configured global_pool_size " + f"({sa_config.global_pool_size}) to {self.pool_size} to cover the " + f"executor's sequence slots." ) # Calculate per-state size based on max_seq_len @@ -159,7 +173,7 @@ def __init__( logger.info( f"SA pool: {self.pool_size} slots " - f"({self.pool_size - max_num_requests} retained capacity, " + f"({self.pool_size - self._num_seq_slots} retained capacity, " f"{self.pool_size * self.state_size / 1024 / 1024:.1f} MB total)" ) diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index b4fc992e18da..4e82f6f9d32d 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -337,12 +337,26 @@ def get_spec_metadata(spec_config, max_num_tokens, spec_resource_manager=spec_resource_manager, is_draft_model=is_draft_model, - max_seq_len=max_seq_len, - num_seq_slots=num_seq_slots) + max_seq_len=max_seq_len) # Set here rather than in each branch below: every one-model mode needs it and # the per-mode constructors are easy to miss one of. if metadata is not None: metadata.enable_penalty = getattr(spec_config, "enable_penalty", False) + # Same reasoning for the sequence-slot pool size, which sizes every + # slot-indexed buffer (draft_probs, full_draft_probs, penalty_state) and + # the dummy scratch row appended after them. It used to be forwarded by + # the MTP-eagle branch alone, so every other one-engine mode -- vanilla + # MTP, Eagle3 one-model, PARD, DFlash/DSpark, draft-target one-model -- + # sized those buffers at max_num_requests while py_seq_slot ranged over + # the wider pool, indexing past the end of the allocation. + # + # Assigning after construction is in time: both consumers allocate + # lazily, from prepare() and from update_one_model_sampling_state, never + # from __post_init__. Leaving the field at its 0 default when the caller + # passes None keeps the established `num_seq_slots or max_num_requests` + # fallback in those two allocators. + if num_seq_slots is not None: + metadata.num_seq_slots = num_seq_slots return metadata @@ -352,14 +366,11 @@ def _build_spec_metadata(spec_config, max_num_tokens, spec_resource_manager=None, is_draft_model=False, - max_seq_len=262144, - num_seq_slots=None): + max_seq_len=262144): + """Construct the per-mode metadata. The slot-pool size is applied by the + caller (``get_spec_metadata``) so no branch can forget it.""" 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. - 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) # Draft-model vocab size, used to gate the d2t-expanded full_draft_probs # buffer allocation (see SpecMetadata.prepare_rejection_sampling_buffers). @@ -383,7 +394,6 @@ def _build_spec_metadata(spec_config, use_rejection_sampling=use_rejection_sampling, advanced_sampling_mode=spec_config.advanced_sampling_mode, vocab_size=vocab_size, - num_seq_slots=num_seq_slots, draft_vocab_size=draft_vocab_size, spec_resource_manager=spec_resource_manager, use_dynamic_tree=getattr(spec_config, 'use_dynamic_tree', False), @@ -524,6 +534,30 @@ def get_mtp_hidden_size(model_config) -> int: return hidden_size +def seat_pool_or_none(model_engine) -> Optional[int]: + """The engine's sequence-slot pool size, or None if it publishes none. + + Pools keyed by live-request identity must follow the executor's + SeqSlotManager pool rather than max_batch_size, because ``py_seq_slot`` -- + the index they are addressed by -- ranges over that whole pool. Buffers + indexed by *batch position* deliberately keep max_batch_size: the + micro-batch scheduler caps every forward at max_batch_size. + + Read unconditionally rather than behind the headroom gate. + ``max_num_seq_slots`` already exceeds max_batch_size for three independent + reasons -- pipeline depth, the overlap headroom (nvbug-6627795) and + disaggregation -- and the buffers here cannot tell them apart. Gating on + one of the three sized them at max_batch_size under the other two while + ``_set_up_spec_metadata`` sized the metadata at the full pool, so a + high-numbered slot indexed past the end of the allocation. + + Returns None only for engines that publish no pool at all (unit-test stubs, + mm-encoder-only engines), which keeps the established + ``num_seq_slots or max_num_requests`` fallback in the allocators. + """ + return getattr(model_engine, "max_num_seq_slots", None) + + def get_spec_resource_manager(model_engine, draft_model_engine=None): spec_config = model_engine.spec_config if spec_config is None: @@ -532,13 +566,16 @@ def get_spec_resource_manager(model_engine, draft_model_engine=None): max_num_requests = model_engine.batch_size max_seq_len = model_engine.max_seq_len max_num_tokens = model_engine.max_num_tokens + num_seq_slots = seat_pool_or_none(model_engine) spec_dec_mode = spec_config.spec_dec_mode if spec_dec_mode.is_mtp_eagle_one_model(): sa_manager = None sa_cfg = getattr(spec_config, 'sa_config', None) if sa_cfg is not None: - sa_manager = SuffixAutomatonManager(sa_cfg, max_num_requests, - max_seq_len) + sa_manager = SuffixAutomatonManager(sa_cfg, + max_num_requests, + max_seq_len, + num_seq_slots=num_seq_slots) # Dynamic tree combines SpecTreeManager with MTP hidden-state slots. if getattr(spec_config, 'use_dynamic_tree', False): return MTPEagleDynamicTreeResourceManager( @@ -547,6 +584,7 @@ def get_spec_resource_manager(model_engine, draft_model_engine=None): get_mtp_hidden_size(model_config), max_num_requests, sa_manager=sa_manager, + num_seq_slots=num_seq_slots, ) if spec_config.use_relaxed_acceptance_for_thinking or sa_manager is not None: # Unified resource manager: the unified worker reads @@ -560,6 +598,7 @@ def get_spec_resource_manager(model_engine, draft_model_engine=None): max_seq_len, max_num_tokens, sa_manager=sa_manager, + num_seq_slots=num_seq_slots, ) else: return None @@ -567,25 +606,30 @@ def get_spec_resource_manager(model_engine, draft_model_engine=None): sa_manager = None sa_cfg = getattr(spec_config, 'sa_config', None) if sa_cfg is not None: - sa_manager = SuffixAutomatonManager(sa_cfg, max_num_requests, - max_seq_len) + sa_manager = SuffixAutomatonManager(sa_cfg, + max_num_requests, + max_seq_len, + num_seq_slots=num_seq_slots) return MTPHiddenStatesManager( spec_config, model_config.torch_dtype, get_mtp_hidden_size(model_config), max_num_requests, sa_manager=sa_manager, + num_seq_slots=num_seq_slots, ) if spec_dec_mode.is_eagle3_one_model() and _is_effective_dynamic_tree( spec_config): - return Eagle3OneModelDynamicTreeResourceManager(spec_config, - max_num_requests) + return Eagle3OneModelDynamicTreeResourceManager( + spec_config, max_num_requests, num_seq_slots=num_seq_slots) if spec_dec_mode.is_eagle3_one_model(): sa_manager = None sa_cfg = getattr(spec_config, 'sa_config', None) if sa_cfg is not None: - sa_manager = SuffixAutomatonManager(sa_cfg, max_num_requests, - max_seq_len) + sa_manager = SuffixAutomatonManager(sa_cfg, + max_num_requests, + max_seq_len, + num_seq_slots=num_seq_slots) return Eagle3ResourceManager( spec_config, model_config.torch_dtype, @@ -594,6 +638,7 @@ def get_spec_resource_manager(model_engine, draft_model_engine=None): max_seq_len, max_num_tokens, sa_manager=sa_manager, + num_seq_slots=num_seq_slots, ) if spec_dec_mode.is_save_hidden_states(): return SaveHiddenStatesResourceManager( @@ -606,13 +651,18 @@ def get_spec_resource_manager(model_engine, draft_model_engine=None): if spec_dec_mode.is_parallel_draft(): sa_cfg = getattr(spec_config, 'sa_config', None) if sa_cfg is not None: - return SuffixAutomatonManager(sa_cfg, max_num_requests, max_seq_len) + return SuffixAutomatonManager(sa_cfg, + max_num_requests, + max_seq_len, + num_seq_slots=num_seq_slots) return None if spec_dec_mode.is_ngram(): return NGramPoolManager(spec_config, max_num_requests) if spec_dec_mode.is_sa(): - return SuffixAutomatonManager(spec_config, max_num_requests, - max_seq_len) + return SuffixAutomatonManager(spec_config, + max_num_requests, + max_seq_len, + num_seq_slots=num_seq_slots) if spec_dec_mode.is_user_provided(): return spec_config.resource_manager return None diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 5a883a81c248..b7f9bcb27dae 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/executor/test_profile_endpoints.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 diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index feb23ceb846a..a6f95a25a395 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -64,6 +64,7 @@ l0_h100: - unittest/_torch/speculative/test_eagle3.py - unittest/_torch/speculative/test_rejection_buffers_guard.py - unittest/_torch/speculative/test_sa_hybrid_state_promotion.py + - unittest/_torch/speculative/test_spec_slot_pool_sizing.py - unittest/_torch/speculative/hw_agnostic - unittest/_torch/speculative/test_capture_sampling_params.py - unittest/_torch/thop/parallel diff --git a/tests/unittest/_torch/disaggregation/test_benchmark_disagg.py b/tests/unittest/_torch/disaggregation/test_benchmark_disagg.py index c69a60c2d429..c702fbd126b4 100644 --- a/tests/unittest/_torch/disaggregation/test_benchmark_disagg.py +++ b/tests/unittest/_torch/disaggregation/test_benchmark_disagg.py @@ -757,6 +757,11 @@ def __init__( self.resource_manager = Mock() self.resource_manager.get_resource_manager.return_value = None + # The pad path reads the router's gate for "which requests count as + # routable load"; non-PP attention DP excludes the retiring ones + # (nvbug-6627795). + self.adp_router = Mock(exclude_retiring_requests=True) + from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor, _ADPForwardIntent _pad_attention_dp_dummy_request = PyExecutor._pad_attention_dp_dummy_request diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py index 6bad069afaa4..92aeeadb5805 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py @@ -1158,6 +1158,10 @@ def test_separate_one_model_draft_normalizes_target_pool_ratio() -> None: creator._is_disagg = False creator._mapping = Mock() creator._speculative_config = Mock() + # Every V2 manager is told whether the overlap scheduler is on, so it can size + # its index-lease pool; this hand-built creator has to carry the attribute the + # real one sets in __init__. The value does not move this test's assertions. + creator._disable_overlap_scheduler = False effective_draft_config = Mock() effective_draft_config.pretrained_config.torch_dtype = "bfloat16" diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py index 62b565ca5dac..2889aa53d23b 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py @@ -1709,3 +1709,147 @@ def test_disagg_role_mapper_kinds_default_to_indexed(): Role.ALL: MapperKind.INDEXED, Role.INDEX_KEY: MapperKind.REPLICATED, } + + +def _index_mapper_capacity_for( + *, + max_batch_size: int, + pp_size: int = 1, + is_disagg: bool = False, + num_reserved_index_slots: int = 1, + disable_overlap_scheduler: bool = True, +) -> tuple[int, int, int]: + """Construct a manager and return the three sizes it derives. + + ``(IndexMapper capacity, page-table capacity, max_admissible_sequences)``. + + The first two must agree: ``host_kv_cache_block_offsets`` is indexed by the + index the mapper hands out, so a page table sized below the mapper's capacity + would be an out-of-bounds write. The third is the published admission bound + that ``validate_seq_slot_pool_covers_admission`` checks the seat pool against + at startup, and it is the capacity minus the reserved dummy slots. + """ + module = "tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2" + fake_impl = Mock() + fake_impl.layer_grouping = [[0]] + fake_impl.pool_group_descs = [] + fake_impl.get_layer_group_id.side_effect = lambda _: 0 + + def build_base_config( + self: KVCacheManagerV2, + config: KvCacheConfig, + *, + tokens_per_block: int, + cache_tiers: list[object], + ) -> _FakeManagerConfig: + del self, config, tokens_per_block + return _FakeManagerConfig(cache_tiers=cache_tiers) + + with ( + patch(f"{module}.IndexMapper") as index_mapper_cls, + patch(f"{module}.KVCacheManagerPy", Mock(return_value=fake_impl)), + patch.object(KVCacheManagerV2, "_build_base_config", build_base_config), + patch.object(KVCacheManagerV2, "_build_cache_config", lambda self, config: config), + patch.object(KVCacheManagerV2, "get_num_available_tokens", return_value=MAX_SEQ_LEN), + patch.object(KVCacheManagerV2, "_prepare_page_table_tensor") as page_table, + patch.object(KVCacheManagerV2, "_log_kv_cache_pool_lifecycle_mapping"), + ): + manager = KVCacheManagerV2( + # A quota must be set or __init__ asserts before it sizes anything + # ("Quota not set. Check kv_cache_config.max_tokens or + # kv_cache_config.max_gpu_total_bytes"). The value is irrelevant to the + # index-mapper arithmetic, which reads only max_batch_size, pp_size, + # is_disagg, disable_overlap_scheduler and num_reserved_index_slots. + KvCacheConfig(max_gpu_total_bytes=16 << 20), + CacheType.SELFKONLY, + num_layers=1, + num_kv_heads=1, + head_dim=1, + tokens_per_block=TOKENS_PER_BLOCK, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=max_batch_size, + mapping=Mapping( + world_size=pp_size, + rank=0, + tp_size=1, + pp_size=pp_size, + ), + dtype=DataType.HALF, + vocab_size=16, + execution_stream=Mock(), + is_disagg=is_disagg, + num_reserved_index_slots=num_reserved_index_slots, + disable_overlap_scheduler=disable_overlap_scheduler, + ) + index_mapper_cls.assert_called_once() + page_table.assert_called_once() + return ( + index_mapper_cls.call_args.args[0], + page_table.call_args.args[0], + manager.max_admissible_sequences, + ) + + +# (max_batch_size, pp_size, disable_overlap_scheduler, is_disagg, reserved, expected) +# +# capacity == max_batch_size * pp_size +# * (2 if is_disagg or (overlap on and pp_size == 1) else 1) +# + reserved +# +# The overlap rows are the nvbug 6627795 case: the overlap scheduler defers a +# terminal request's teardown past the point where its replacement is admitted, +# so both cohorts hold index slots at once and a mapper sized at B+1 silently +# defers requests one at a time. The overlap half of this predicate is shared +# with _util.should_enable_overlap_headroom, which gates the *seat* pool, so the +# two pools are equal on those rows. The is_disagg half is index-local and has no +# seat-pool counterpart, which is why the startup check requires the index pool to +# cover the seat pool rather than to equal it +# (validate_seq_slot_pool_covers_admission). +_INDEX_MAPPER_CAPACITY_CASES = [ + # Overlap on, no PP: both cohorts are resident, so the mapper needs 2B. + pytest.param(2, 1, False, False, 1, 5, id="overlap_on"), + pytest.param(8, 1, False, False, 1, 17, id="overlap_on_b8"), + # Overlap off: the retiring request is torn down in-line, so B+1 is enough -- + # the pre-fix allocation. + pytest.param(2, 1, True, False, 1, 3, id="overlap_off"), + # Disagg already carried its own 2x; the two reasons cover the same extra + # cohort (a request is either transferring or retiring), so they do not + # compound. + pytest.param(2, 1, False, True, 1, 5, id="disagg_does_not_compound"), + pytest.param(2, 1, True, True, 1, 5, id="disagg_only"), + # Pipeline parallelism: max_batch_size * pp_size already covers the in-flight + # micro-batches, so the overlap term is excluded rather than compounded. + pytest.param(2, 4, True, False, 1, 9, id="pp4_plain"), + pytest.param(2, 4, False, False, 1, 9, id="pp4_overlap_excluded"), + # ... while the pre-existing disagg 2x under PP is left exactly as it was. + pytest.param(2, 4, True, True, 1, 17, id="pp4_disagg_unchanged"), + # Reserved slots are still added on top of the widened pool. + pytest.param(2, 1, False, False, 5, 9, id="reserved_slots_still_added"), +] + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "max_batch_size,pp_size,disable_overlap_scheduler,is_disagg,reserved,expected", + _INDEX_MAPPER_CAPACITY_CASES, +) +def test_index_mapper_capacity_covers_the_overlapping_cohorts( + max_batch_size: int, + pp_size: int, + disable_overlap_scheduler: bool, + is_disagg: bool, + reserved: int, + expected: int, +) -> None: + capacity, page_table_capacity, max_admissible_sequences = _index_mapper_capacity_for( + max_batch_size=max_batch_size, + pp_size=pp_size, + is_disagg=is_disagg, + num_reserved_index_slots=reserved, + disable_overlap_scheduler=disable_overlap_scheduler, + ) + assert capacity == expected + assert page_table_capacity == expected + # The published bound excludes the reserved slots, which only ever hold + # persistent dummies and are not available to admitted requests. + assert max_admissible_sequences == expected - reserved diff --git a/tests/unittest/_torch/executor/test_adp_router.py b/tests/unittest/_torch/executor/test_adp_router.py index 38e09e215f61..a0643bb74fe3 100644 --- a/tests/unittest/_torch/executor/test_adp_router.py +++ b/tests/unittest/_torch/executor/test_adp_router.py @@ -7,11 +7,13 @@ - Strict/relaxed attention-DP request routing while respecting rank capacity """ +import inspect from unittest.mock import MagicMock, Mock import pytest from tensorrt_llm._torch.pyexecutor.executor_request_queue import RequestQueueItem +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState from tensorrt_llm._torch.pyexecutor.request_utils import get_from_waiting_queue from tensorrt_llm._torch.pyexecutor.scheduler import FCFSWaitingQueue from tensorrt_llm._torch.pyexecutor.scheduler.adp_router import ( @@ -22,6 +24,9 @@ RankIterStatsPayload, RankState, _num_input_tokens, + build_active_requests_for_overlap, + count_retiring_requests, + is_retiring_request, ) from tensorrt_llm.conversation_params import ConversationParams from tensorrt_llm.scheduling_params import SchedulingParams @@ -40,12 +45,16 @@ def num_input_tokens(self): return len(self.input_token_ids) -def _mock_dist(tp_rank=0, tp_size=1, has_cp_helix=False): +def _mock_dist(tp_rank=0, tp_size=1, has_cp_helix=False, has_pp=False): """Create a mock Distributed object for testing.""" dist = MagicMock() dist.tp_rank = tp_rank dist.tp_size = tp_size dist.has_cp_helix = has_cp_helix + # ADPRouter reads this to decide whether to route on the overlap-corrected + # active list; a bare MagicMock would make has_pp() truthy and silently + # disable the correction in every test. + dist.mapping.has_pp.return_value = has_pp return dist @@ -130,6 +139,74 @@ def all_ranks_num_active_tokens(): return [10, 5, 15, 8] +def _retiring_request(prompt_len=100): + """Active request that has produced its final token (state 14).""" + return Mock( + py_orig_prompt_len=prompt_len, + cached_tokens=0, + state=LlmRequestState.GENERATION_TO_COMPLETE, + ) + + +class TestBuildActiveRequestsForOverlap: + # Retiring requests linger in active_requests for one extra iteration when + # the overlap scheduler is on. They must not be charged against ADP + # admission capacity, because no scheduler can ever forward them again + # (nvbug-6627795). Every other lingering state still owns a seat and KV, so + # it must keep counting. + def test_empty(self): + assert build_active_requests_for_overlap([]) == [] + assert count_retiring_requests([]) == 0 + + def test_drops_generation_to_complete(self): + reqs = [_retiring_request(), _retiring_request(), _retiring_request()] + assert build_active_requests_for_overlap(reqs) == [] + assert count_retiring_requests(reqs) == 3 + + @pytest.mark.parametrize( + "state", + [ + LlmRequestState.CONTEXT_INIT, + LlmRequestState.GENERATION_IN_PROGRESS, + LlmRequestState.DISAGG_CONTEXT_WAIT_SCHEDULER, + LlmRequestState.DISAGG_GENERATION_INIT, + LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS, + LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS, + LlmRequestState.DISAGG_TRANS_ERROR, + ], + ) + def test_other_states_still_count_as_load(self, state): + req = Mock(state=state) + assert is_retiring_request(req) is False + assert build_active_requests_for_overlap([req]) == [req] + assert count_retiring_requests([req]) == 0 + + def test_mixed_preserves_order_of_survivors(self): + keep_a = Mock(state=LlmRequestState.GENERATION_IN_PROGRESS) + keep_b = Mock(state=LlmRequestState.DISAGG_GENERATION_INIT) + reqs = [keep_a, _retiring_request(), keep_b, _retiring_request()] + assert build_active_requests_for_overlap(reqs) == [keep_a, keep_b] + assert count_retiring_requests(reqs) == 2 + + def test_returns_a_new_list(self): + # The filtered list is the ROUTER's view only; mutating it must never + # disturb PyExecutor.active_requests, whose teardown ordering is what + # created the bug in the first place. + reqs = [Mock(state=LlmRequestState.GENERATION_IN_PROGRESS)] + filtered = build_active_requests_for_overlap(reqs) + assert filtered is not reqs + filtered.clear() + assert len(reqs) == 1 + + def test_bare_mock_is_not_retiring(self): + # Router tests build requests with bare Mocks that never set `state`. + # Identity comparison against the enum keeps those routable; a truthy + # bound-property check would drop every one of them. + req = Mock(py_orig_prompt_len=10) + assert build_active_requests_for_overlap([req]) == [req] + assert count_retiring_requests([req]) == 0 + + class TestRankState: # RankState is the wire payload shared across attention-DP ranks. Keep its # serialization stable because iter-stats now ride on the same allgather. @@ -141,7 +218,7 @@ def test_creation(self): def test_serialize(self): state = RankState(rank=0, num_active_requests=5, num_active_tokens=100) - assert state.serialize() == [0, 5, 100, 0, -1, 0, 0, 0, 0, 0, 0, 0] + assert state.serialize() == [0, 5, 100, 0, 0, -1, 0, 0, 0, 0, 0, 0, 0] def test_deserialize(self): state = RankState.deserialize(data=[2, 3, 50]) @@ -154,10 +231,22 @@ def test_roundtrip(self): restored = RankState.deserialize(data=original.serialize()) assert original == restored + def test_roundtrip_with_retiring_requests(self): + original = RankState( + rank=1, + num_active_requests=10, + num_active_tokens=200, + num_retiring_requests=3, + ) + restored = RankState.deserialize(data=original.serialize()) + assert original == restored + assert restored.num_retiring_requests == 3 + def test_defaults(self): state = RankState(rank=0) assert state.num_active_requests == 0 assert state.num_active_tokens == 0 + assert state.num_retiring_requests == 0 assert state.iter_stats.has_iter_stats == 0 assert state.iter_stats.iter_stats_iter == -1 @@ -278,6 +367,158 @@ def test_create_rank_state_default(self): assert state.num_active_requests == 2 assert state.num_active_tokens == 300 + def test_create_rank_state_does_not_filter_retiring_itself(self): + # create_rank_state stays overlap-agnostic: it reports what it is given. + # The correction lives in gather_all_rank_states, so a router author + # cannot forget it. + dist = _mock_dist(tp_rank=0, has_cp_helix=False) + router = DefaultADPRouter(dist=dist) + active = [ + Mock(py_orig_prompt_len=100, state=LlmRequestState.GENERATION_IN_PROGRESS), + _retiring_request(prompt_len=200), + ] + state = router.create_rank_state(active_requests=active, new_requests=[]) + assert state.num_active_requests == 2 + assert state.num_active_tokens == 300 + + def test_gather_all_rank_states_excludes_retiring(self): + # Two of three requests are retiring, so only one is routable load -- + # and the tokens of the retiring pair go with them, because the filtered + # list is what create_rank_state sums. + dist = _mock_dist(tp_rank=0, has_cp_helix=False) + router = DefaultADPRouter(dist=dist, has_seq_slot_headroom=True) + active = [ + Mock(py_orig_prompt_len=100, state=LlmRequestState.GENERATION_IN_PROGRESS), + _retiring_request(prompt_len=200), + _retiring_request(prompt_len=300), + ] + dist.tp_allgather.side_effect = lambda payload: [payload] + + states = router.gather_all_rank_states(active_requests=active) + + assert len(states) == 1 + assert states[0].num_active_requests == 1 + assert states[0].num_retiring_requests == 2 + assert states[0].num_active_tokens == 100 + # The executor's own list is untouched -- only the router's view narrows. + assert len(active) == 3 + + def test_gather_all_rank_states_reports_zero_when_all_retiring(self): + # Nothing routable, but the rank is NOT idle: the retiring requests are + # still resident. num_retiring_requests carries that fact to every peer + # so the idle-fetch wait stays collective (nvbug-6627795). + dist = _mock_dist(tp_rank=0, has_cp_helix=False) + router = DefaultADPRouter(dist=dist, has_seq_slot_headroom=True) + active = [_retiring_request(), _retiring_request()] + dist.tp_allgather.side_effect = lambda payload: [payload] + + states = router.gather_all_rank_states(active_requests=active) + + assert states[0].num_active_requests == 0 + assert states[0].num_active_tokens == 0 + assert states[0].num_retiring_requests == 2 + + def test_gather_all_rank_states_keeps_retiring_without_headroom(self): + # Without seat headroom the correction is off, whatever the topology. + # The correction lets a rank hold more requests than it is charged for, + # so it is only sound when the sequence-slot pool was sized with the + # extra generation of seats -- e.g. hybrid/SSM architectures are excluded + # from the headroom because their state-slot pool is sized from + # max_batch_size alone, and the router must follow that exclusion rather + # than re-derive its own predicate. + dist = _mock_dist(tp_rank=0, has_cp_helix=False, has_pp=True) + router = DefaultADPRouter(dist=dist, has_seq_slot_headroom=False) + assert router.exclude_retiring_requests is False + active = [ + Mock(py_orig_prompt_len=100, state=LlmRequestState.GENERATION_IN_PROGRESS), + _retiring_request(prompt_len=200), + _retiring_request(prompt_len=300), + ] + dist.tp_allgather.side_effect = lambda payload: [payload] + + states = router.gather_all_rank_states(active_requests=active) + + # Same inputs as test_gather_all_rank_states_excludes_retiring, which + # asserts 1 / 2 / 100 -- here nothing is filtered. + assert states[0].num_active_requests == 3 + assert states[0].num_retiring_requests == 0 + assert states[0].num_active_tokens == 600 + + def test_exclude_retiring_requests_follows_the_seat_pool_headroom(self): + # The flag tracks the *engine's* headroom flag and nothing else, so the + # correction can only ever spend seats that were actually allocated. It + # used to re-derive the predicate as `not dist.mapping.has_pp()`, which + # is a second copy of a fact the sizing gate already owns. + for has_pp in (False, True): + dist = _mock_dist(has_pp=has_pp) + assert ( + DefaultADPRouter(dist=dist, has_seq_slot_headroom=True).exclude_retiring_requests + is True + ) + assert ( + DefaultADPRouter(dist=dist, has_seq_slot_headroom=False).exclude_retiring_requests + is False + ) + + def test_router_factory_requires_the_headroom_flag(self): + # Required rather than defaulted: a caller that silently got the + # aggressive behaviour would be re-introducing the skew this parameter + # exists to remove. There is exactly one production call site + # (PyExecutor.__init__), which passes the engine's flag. + params = inspect.signature(ADPRouter.create).parameters + assert params["has_seq_slot_headroom"].default is inspect.Parameter.empty + + def test_router_constructors_default_to_no_headroom(self): + # The constructors keep a default for direct instantiation, and it must + # be the conservative one: a router that drops retiring requests from + # the counts admission subtracts, without a seat pool sized for it, lets + # residency exceed max_batch_size * pp_size. + for cls in (DefaultADPRouter, ConversationAwareADPRouter, KVCacheAwareADPRouter): + default = inspect.signature(cls.__init__).parameters["has_seq_slot_headroom"].default + assert default is False, cls.__name__ + + @pytest.mark.parametrize("has_seq_slot_headroom", [True, False]) + def test_router_factory_propagates_the_headroom_flag(self, has_seq_slot_headroom): + # Every branch of the factory, not just the default one: the KV-cache- + # aware and conversation-affinity routers run the same admission + # correction and need the same gate. + dist = _mock_dist(tp_size=2) + mgr = Mock(enable_block_reuse=True) + configs = [ + None, + Mock( + kv_cache_routing_conversation_affinity=False, + enable_kv_cache_aware_routing=True, + kv_cache_routing_load_balance_weight=1.0, + kv_cache_routing_match_rate_threshold=0.1, + kv_cache_routing_fair_share_multiplier=2.0, + kv_cache_routing_cold_start_warmup=False, + kv_cache_routing_account_for_in_transfer=False, + ), + Mock( + kv_cache_routing_conversation_affinity=True, + kv_cache_routing_max_sessions=1 << 16, + kv_cache_routing_fair_share_multiplier=2.0, + kv_cache_routing_new_conv_placement="round_robin", + ), + ] + built = set() + for attention_dp_config in configs: + router = ADPRouter.create( + dist=dist, + has_seq_slot_headroom=has_seq_slot_headroom, + kv_cache_manager=mgr, + attention_dp_config=attention_dp_config, + ) + built.add(type(router).__name__) + assert router.exclude_retiring_requests is has_seq_slot_headroom + # Anti-vacuity: the three configs must actually reach three branches. + assert built == { + "DefaultADPRouter", + "KVCacheAwareADPRouter", + "ConversationAwareADPRouter", + } + def test_create_rank_state_cp_helix(self): dist = _mock_dist(tp_rank=1, has_cp_helix=True) router = DefaultADPRouter(dist=dist) @@ -1293,11 +1534,36 @@ def test_create_rank_state(self): assert state.num_active_requests == 2 assert state.num_active_tokens == 150 + def test_gather_all_rank_states_excludes_retiring(self): + # The filter sits in the shared ADPRouter.gather_all_rank_states, so it + # applies to this router without a line of router-specific code. + dist = _mock_dist(tp_rank=2) + router = ConversationAwareADPRouter(dist=dist) + active = [ + Mock(py_orig_prompt_len=100, state=LlmRequestState.GENERATION_IN_PROGRESS), + _retiring_request(prompt_len=50), + ] + dist.tp_allgather.side_effect = lambda payload: [payload] + + states = router.gather_all_rank_states(active_requests=active) + + assert states[0].num_active_requests == 1 + assert states[0].num_retiring_requests == 1 + assert states[0].num_active_tokens == 100 + def test_factory_selects_conversation_router(self): + # has_seq_slot_headroom is required rather than defaulted (see + # test_router_factory_requires_the_headroom_flag), so every call site states + # it -- including the ones that are only about which class gets built. cfg = MagicMock() cfg.kv_cache_routing_conversation_affinity = True cfg.kv_cache_routing_max_sessions = 8 - router = ADPRouter.create(dist=_mock_dist(), kv_cache_manager=None, attention_dp_config=cfg) + router = ADPRouter.create( + dist=_mock_dist(), + has_seq_slot_headroom=True, + kv_cache_manager=None, + attention_dp_config=cfg, + ) assert isinstance(router, ConversationAwareADPRouter) assert router._max_sessions == 8 # A mocked (non-string) placement value must fall back to round_robin. @@ -1352,7 +1618,12 @@ def test_new_conv_placement_config(self): cfg.kv_cache_routing_conversation_affinity = True cfg.kv_cache_routing_max_sessions = 8 cfg.kv_cache_routing_new_conv_placement = "least_queued" - router = ADPRouter.create(dist=_mock_dist(), kv_cache_manager=None, attention_dp_config=cfg) + router = ADPRouter.create( + dist=_mock_dist(), + has_seq_slot_headroom=True, + kv_cache_manager=None, + attention_dp_config=cfg, + ) assert router._new_conv_placement == "least_queued" bad = ConversationAwareADPRouter(dist=_mock_dist(tp_size=4), new_conv_placement="banana") assert bad._new_conv_placement == "round_robin" @@ -1361,7 +1632,12 @@ def test_factory_default_when_disabled(self): cfg = MagicMock() cfg.kv_cache_routing_conversation_affinity = False cfg.enable_kv_cache_aware_routing = False - router = ADPRouter.create(dist=_mock_dist(), kv_cache_manager=None, attention_dp_config=cfg) + router = ADPRouter.create( + dist=_mock_dist(), + has_seq_slot_headroom=True, + kv_cache_manager=None, + attention_dp_config=cfg, + ) assert isinstance(router, DefaultADPRouter) def test_returned_expected_covers_every_rank(self): diff --git a/tests/unittest/_torch/executor/test_kvcache_aware_router.py b/tests/unittest/_torch/executor/test_kvcache_aware_router.py index 1f173ed0da83..049cfe79a144 100644 --- a/tests/unittest/_torch/executor/test_kvcache_aware_router.py +++ b/tests/unittest/_torch/executor/test_kvcache_aware_router.py @@ -21,6 +21,7 @@ import pytest +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState from tensorrt_llm._torch.pyexecutor.scheduler.adp_router import ( ADPRouter, KVCacheAwareADPRouter, @@ -33,7 +34,7 @@ # ---- Helpers ---- -def _mock_dist(tp_rank=0, tp_size=1, has_cp_helix=False): +def _mock_dist(tp_rank=0, tp_size=1, has_cp_helix=False, has_pp=False): """Create a mock Distributed object for testing.""" dist = MagicMock() dist.tp_rank = tp_rank @@ -41,6 +42,10 @@ def _mock_dist(tp_rank=0, tp_size=1, has_cp_helix=False): # ADP scheduling assumes ``enable_attention_dp=True``, so ``dp_size`` # mirrors ``tp_size`` (see ``Mapping.dp_size``). dist.mapping.dp_size = tp_size + # ADPRouter reads this to decide whether to route on the overlap-corrected + # active list; a bare MagicMock would make has_pp() truthy and silently + # disable the correction in every test. + dist.mapping.has_pp.return_value = has_pp dist.has_cp_helix = has_cp_helix return dist @@ -116,10 +121,67 @@ def test_create_rank_state(self): assert state.num_active_requests == 2 assert state.num_active_tokens == 300 + def test_gather_all_rank_states_excludes_retiring(self): + dist = _mock_dist(tp_rank=0) + mgr = _mock_kv_cache_manager() + router = KVCacheAwareADPRouter(dist=dist, kv_cache_manager=mgr, has_seq_slot_headroom=True) + + req1 = Mock( + py_orig_prompt_len=100, + cached_tokens=0, + state=LlmRequestState.GENERATION_IN_PROGRESS, + ) + req2 = Mock( + py_orig_prompt_len=200, + cached_tokens=0, + state=LlmRequestState.GENERATION_TO_COMPLETE, + ) + dist.tp_allgather.side_effect = lambda payload: [payload] + + states = router.gather_all_rank_states([req1, req2]) + + assert states[0].num_active_requests == 1 + assert states[0].num_retiring_requests == 1 + assert states[0].num_active_tokens == 100 + + def test_gather_all_rank_states_retiring_and_in_transfer(self): + # In-transfer requests have already left active_requests, so this router + # adds them back as load; retiring requests are still in it and are + # filtered out. The two corrections are independent and must compose. + dist = _mock_dist(tp_rank=0) + mgr = _mock_kv_cache_manager() + transfer_mgr = MagicMock() + in_transfer_req = Mock(py_orig_prompt_len=70, cached_tokens=0) + transfer_mgr.requests_in_transfer.return_value = {1: in_transfer_req} + router = KVCacheAwareADPRouter( + dist=dist, + kv_cache_manager=mgr, + async_transfer_manager=transfer_mgr, + account_for_in_transfer=True, + ) + + req1 = Mock( + py_orig_prompt_len=100, + cached_tokens=0, + state=LlmRequestState.GENERATION_IN_PROGRESS, + ) + req2 = Mock( + py_orig_prompt_len=200, + cached_tokens=0, + state=LlmRequestState.GENERATION_TO_COMPLETE, + ) + dist.tp_allgather.side_effect = lambda payload: [payload] + + states = router.gather_all_rank_states([req1, req2]) + + assert states[0].num_active_requests == 2 # 1 routable + 1 in transfer + assert states[0].num_retiring_requests == 1 + assert states[0].num_active_tokens == 170 + def test_create_rank_state_cp_helix(self): dist = _mock_dist(tp_rank=1, has_cp_helix=True) mgr = _mock_kv_cache_manager() - router = KVCacheAwareADPRouter(dist=dist, kv_cache_manager=mgr) + router = KVCacheAwareADPRouter(dist=dist, kv_cache_manager=mgr, has_seq_slot_headroom=True) req1 = Mock(total_input_len_cp=150, cached_tokens=0) state = router.create_rank_state([req1], []) diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 64a6a01c2f41..0a4df59d5188 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -1972,6 +1972,7 @@ def __init__( enable_scheduler_aware_adp_dummy=None, enable_non_overlap_adp_forward_intent=None, peer_forward_intent=_ADPForwardIntent.GENERATION, + exclude_retiring_requests=True, ): self.enable_attention_dp = enable_attention_dp self.kv_cache_transceiver = kv_cache_transceiver @@ -2004,6 +2005,12 @@ def __init__( self.dist.tp_size = 1 self.dist.tp_allgather.side_effect = lambda value: [value] self.dist.tp_allreduce.side_effect = lambda value, op: max(value, int(peer_forward_intent)) + # The pad path reads the router's gate rather than re-deriving it, so + # the two views of "which requests are routable load" cannot drift. + # Default True models a non-PP attention-DP executor; under pipeline + # parallelism the router leaves retiring requests in the load vector + # and the pad path must not subtract them (nvbug-6627795). + self.adp_router = Mock(exclude_retiring_requests=exclude_retiring_requests) self.scheduler = Mock() self.scheduler.scheduling_state_range = ( @@ -2307,6 +2314,43 @@ def test_decoder_context_waiting_for_encoder_output_is_not_counted(): assert len(stub.active_requests) == 2 +def test_pad_does_not_warn_when_surplus_is_only_retiring_requests(): + # The router now excludes retiring requests from the per-rank loads that + # floor `expected` (nvbug-6627795), so `expected` can legitimately sit + # below len(active_requests). Measuring the surplus against the raw len() + # would log a warning on every iteration of a hot loop. + stub = _StubADPExecutor() + stub.active_requests = [ + _make_adp_request(_STATE_GENERATION_IN_PROGRESS, request_id=1), + _make_adp_request(_STATE_GENERATION_TO_COMPLETE, request_id=2), + _make_adp_request(_STATE_GENERATION_TO_COMPLETE, request_id=3), + ] + # What the router would have reported: 3 resident, 1 routable. + stub.expected_num_active_requests = 1 + + with patch("tensorrt_llm._torch.pyexecutor.py_executor.logger") as mock_logger: + _run_pad(stub) + + assert mock_logger.warning.call_count == 0 + # One routable request means the rank has real work; no dummy needed. + assert stub.add_dummy_calls == [] + + +def test_pad_still_warns_on_a_genuine_routable_surplus(): + stub = _StubADPExecutor() + stub.active_requests = [ + _make_adp_request(_STATE_GENERATION_IN_PROGRESS, request_id=1), + _make_adp_request(_STATE_GENERATION_IN_PROGRESS, request_id=2), + ] + stub.expected_num_active_requests = 1 + + with patch("tensorrt_llm._torch.pyexecutor.py_executor.logger") as mock_logger: + _run_pad(stub) + + assert mock_logger.warning.call_count == 1 + assert "exceeds expected_num_active_requests" in mock_logger.warning.call_args[0][0] + + def test_generic_disagg_adp_mixed_rank_states_stay_queueable(): # The generic non-PP path must give both ranks a non-empty scheduled batch: # one rank schedules its real request, while the terminal-only rank diff --git a/tests/unittest/_torch/executor/test_seq_slot_sizing.py b/tests/unittest/_torch/executor/test_seq_slot_sizing.py index 7db2c6ed74ac..1517cabeaa17 100644 --- a/tests/unittest/_torch/executor/test_seq_slot_sizing.py +++ b/tests/unittest/_torch/executor/test_seq_slot_sizing.py @@ -1,57 +1,123 @@ # 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 pools are sized here, and keeping them distinct is the whole point: + +The **seat pool** (``compute_max_num_sequences``) is how many sequences the +executor can seat at once: ``max_batch_size * pp_size`` under pipeline +parallelism, and otherwise ``max_batch_size`` doubled by the 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 capacity scheduler has already dropped them from its budget +(``no_schedule_after_state=GENERATION_TO_COMPLETE``) and the ADP router has +already excluded them from the counts admission subtracts from its capacity -- +so a replacement cohort is admitted while the retiring one is still resident. +``should_enable_overlap_headroom`` is the single gate for that. Its overlap term +matches ``KVCacheManagerV2``'s by construction -- non-PP and overlap-on -- so the +seat pool and the index pool move together; attention DP is deliberately *not* a +condition even though it is the only deployment where the extra seats are ever +occupied, because narrowing the seat gate below the manager's term is what forces +the two pools apart on the most ordinary path there is. It stays off for V1 (the +V1 capacity schedulers hardcode ``GENERATION_COMPLETE``) and for hybrid models +(SSM state is sized from ``max_batch_size``). + +The **index pool** (``KVCacheManagerV2.max_admissible_sequences``) therefore +equals the seat pool in the common case, and is the more generous of the two +wherever they differ. Its ``is_disagg`` term has no seat-pool counterpart: a +request awaiting its KV transfer holds an index lease and no seat at all, because +``SeqSlotManager.prepare_resources`` skips ``DISAGG_GENERATION_INIT``. And a +hybrid model suppresses the seat headroom while a V2 index pool behind it still +widens. Hence ``validate_seq_slot_pool_covers_admission`` requires the index pool +to *cover* the seat pool rather than to equal it. + +``compute_max_num_sequences`` is the single seat-pool implementation used both +for the executor's ``SeqSlotManager`` pool (``create_py_executor_instance``) and +for the sampler state (``create_torch_sampler_args``); +``resolve_max_num_sequences`` is how a consumer obtains it without re-deriving +it. Every other slot-indexed buffer follows the same number, since +``py_seq_slot`` indexes them all. """ +import inspect +from types import SimpleNamespace +from unittest.mock import Mock, patch + import pytest from tensorrt_llm._torch.pyexecutor._util import ( + KvCacheCreator, compute_max_num_sequences, create_torch_sampler_args, + is_disagg_enabled, + resolve_max_num_sequences, should_enable_adp_dummy_fixes, - should_enable_disagg_adp_overlap_headroom, should_enable_non_overlap_adp_forward_intent, + should_enable_overlap_headroom, 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 +# (pp_size, disable_overlap, enable_overlap_headroom, expected_factor) +# +# The two terms are mutually exclusive rather than combined: pipeline depth +# already puts pp_size micro-batches in flight, so the PP branch takes pp_size +# and ignores the headroom entirely. The gate never turns the headroom on under +# PP (see test_overlap_headroom_gate), so the pp>1 rows only pin that the branch +# cannot start compounding. +# +# Disaggregation is absent from this table on purpose -- see +# test_seat_pool_has_no_disagg_term. SIZING_CASES = [ - # (pp_size, disable_overlap, enable_overlap_headroom, expected_factor) - (1, False, True, 2), + # No PP: aggregated baseline, then the extra micro-batch. (1, False, False, 1), + (1, False, True, 2), + # The headroom is a no-op with the overlap scheduler off: the retiring cohort + # is torn down in-line, so it never coexists with its replacement. (1, True, True, 1), - # Existing PP sizing is preserved regardless of the DSv4 opt-in. + # PP: sized by pp_size, with or without the flag set. + (4, False, False, 4), + (4, True, True, 4), (2, False, True, 2), (4, False, True, 4), - (4, True, False, 4), ] @pytest.mark.parametrize( - "enable_attention_dp,is_disagg,pp_size,disable_overlap,expected", + "enable_attention_dp,pp_size,disable_overlap,is_v2,is_hybrid,expected", [ - (True, True, 1, False, True), - (False, True, 1, False, False), - (True, False, 1, False, False), - (True, True, 2, False, False), - (True, True, 1, True, False), + # The scenario this PR exists for: attention DP, no PP, overlap on, V2. + (True, 1, False, True, False, True), + # Attention DP is deliberately absent from the predicate, so a plain TP + # run of the same shape widens too. The extra seats are never occupied + # there -- admission subtracts len(active_requests) with retirees + # included -- but matching KVCacheManagerV2's own term keeps the two + # pools equal instead of making the index pool run ahead on the most + # ordinary path there is. + (False, 1, False, True, False, True), + # Overlap off: the retiring cohort is torn down before the next + # iteration's admission runs. + (True, 1, True, True, False, False), + # Pipeline parallelism: already sized by pp_size, and only the last stage + # marks a generation request GENERATION_TO_COMPLETE, so the router's + # correction is not rank-consistent. Out of scope. + (True, 2, False, True, False, False), + (True, 4, False, True, False, False), + # V1: BindCapacityScheduler and PyCapacityScheduler both hardcode + # no_schedule_after_state=GENERATION_COMPLETE, so the retiring cohort is + # still inside the capacity budget and no overcommit is possible. + (True, 1, False, False, False, False), + # Hybrid/SSM: recurrent state is sized from max_batch_size in every + # manager, so extra seats would let residency outrun the state slots. + (True, 1, False, True, True, False), + # Both exclusions at once, as in + # examples/configs/curated/qwen3.8-high-throughput-mtp3.yaml. + (True, 1, False, False, True, False), ], ) -def test_disagg_adp_overlap_headroom_gate( - enable_attention_dp, is_disagg, pp_size, disable_overlap, expected +def test_overlap_headroom_gate( + enable_attention_dp, pp_size, disable_overlap, is_v2, is_hybrid, expected ): mapping = Mapping( world_size=pp_size, @@ -59,14 +125,57 @@ def test_disagg_adp_overlap_headroom_gate( pp_size=pp_size, enable_attention_dp=enable_attention_dp, ) - cache_config = CacheTransceiverConfig(backend="NIXL") if is_disagg else None assert ( - should_enable_disagg_adp_overlap_headroom(mapping, cache_config, disable_overlap) + should_enable_overlap_headroom( + mapping, + disable_overlap, + kv_cache_manager_is_v2=is_v2, + is_hybrid=is_hybrid, + ) is expected ) +def test_overlap_headroom_gate_does_not_depend_on_disaggregation(): + """Disagg is not a reason to widen the *seat* pool, only the index pool. + + The gate used to take ``cache_transceiver_config`` and return True for a + disaggregated server regardless of attention DP. That bought seats that could + never be occupied: a request in KV transfer holds an index lease and no seat, + and outside attention DP admission subtracts ``len(active_requests)`` so + residency stays at ``max_batch_size * pp_size``. Asserting on the signature + states the intent that a value test cannot: the parameter's mere presence is + what invites the conflation back. + """ + params = inspect.signature(should_enable_overlap_headroom).parameters + assert "cache_transceiver_config" not in params + assert "is_disagg" not in params + + +@pytest.mark.parametrize("disable_overlap", [False, True]) +def test_overlap_headroom_gate_does_not_depend_on_attention_dp(disable_overlap): + """Attention DP is where the seats are *used*, not where they are sized. + + The gate deliberately ignores it so that its overlap term is the same shape + as ``KVCacheManagerV2``'s (``not has_pp() and overlap on``). Making the seat + pool the narrower of the two is what forces + ``validate_seq_slot_pool_covers_admission`` to tolerate a permanent gap on + plain TP with the overlap scheduler -- a default configuration -- and a + tolerated gap is one nobody can distinguish from an undersized seat pool + later. The cost of ignoring it is idle seats outside attention DP, where + admission subtracts ``len(active_requests)`` with retirees included; the + benefit is that the two pools are computed from the same condition. + """ + kwargs = dict(kv_cache_manager_is_v2=True, is_hybrid=False) + with_adp = Mapping(world_size=1, tp_size=1, pp_size=1, enable_attention_dp=True) + without_adp = Mapping(world_size=1, tp_size=1, pp_size=1, enable_attention_dp=False) + + assert should_enable_overlap_headroom( + with_adp, disable_overlap, **kwargs + ) is should_enable_overlap_headroom(without_adp, disable_overlap, **kwargs) + + @pytest.mark.parametrize("pp_size,expected", [(1, True), (2, False)]) def test_adp_dummy_fix_gate(pp_size, expected): mapping = Mapping(world_size=pp_size, tp_size=1, pp_size=pp_size) @@ -120,15 +229,253 @@ def test_compute_max_num_sequences_scopes_overlap_headroom( ) +def test_seat_pool_has_no_disagg_term(): + """The disaggregation 2x is confined to KVCacheManagerV2's index pool. + + A request awaiting its KV transfer holds an *index* lease and no seat at all: + ``SeqSlotManager.prepare_resources`` skips ``DISAGG_GENERATION_INIT`` + requests outright and only seats one once its transmission completes, while + admission stays at ``max_batch_size * pp_size``. So the index pool + legitimately runs ahead of the seat pool, and doubling the *seat* pool would + buy nothing while doubling everything keyed by seat -- sampler state, + ``[seats, draft_len, vocab]`` draft probabilities (~800 MB at 512 seats), the + penalty tensors and the pinned-host block-offset tables. + + Asserting on the signature rather than on a return value is deliberate: a + value test cannot distinguish "the parameter is gone" from "the parameter + defaults to False", and it is the parameter's *existence* that invites a + caller to propagate the factor. + """ + assert "is_disagg" not in inspect.signature(compute_max_num_sequences).parameters + assert "is_disagg" not in inspect.signature(resolve_max_num_sequences).parameters + + +@pytest.mark.parametrize("enable_attention_dp", [False, True]) +@pytest.mark.parametrize("pp_size", [1, 2]) +@pytest.mark.parametrize("is_disagg", [False, True]) +@pytest.mark.parametrize("disable_overlap_scheduler", [False, True]) +def test_sizing_matches_kv_manager_admission_bound( + enable_attention_dp, pp_size, is_disagg, disable_overlap_scheduler +): + """The two pools are computed in different modules; hold them in step. + + Asserting the two expressions against each other rather than against literals + means a drift in either one fails here rather than at startup in + ``validate_seq_slot_pool_covers_admission`` -- and the assertion is exactly + that validator's invariant. + + The overlap term is shared, so the two agree everywhere except where the + manager's ``is_disagg`` term fires without the shared term also firing -- + i.e. a disaggregated run that is either PP or overlap-off. There the index + pool runs ahead, which is permitted; what must never happen is the reverse. + """ + max_batch_size = 8 + mapping = Mapping( + world_size=pp_size, + tp_size=1, + pp_size=pp_size, + enable_attention_dp=enable_attention_dp, + ) + enable_overlap_headroom = should_enable_overlap_headroom( + mapping, + disable_overlap_scheduler, + kv_cache_manager_is_v2=True, + ) + + seats = compute_max_num_sequences( + mapping, + max_batch_size, + disable_overlap_scheduler, + enable_overlap_headroom=enable_overlap_headroom, + ) + + # Mirrors the arithmetic in KVCacheManagerV2.__init__. + overlap_term = not disable_overlap_scheduler and pp_size == 1 + extra_leases = is_disagg or overlap_term + admission_bound = max_batch_size * pp_size * (2 if extra_leases else 1) + + # The invariant the startup validator enforces. + assert admission_bound >= seats + # Equal unless the manager's disagg-only term is what widened it: the overlap + # term is shared with the seat gate, so it can never open a gap. + if is_disagg and not overlap_term: + assert admission_bound == 2 * seats + else: + assert admission_bound == seats + + +class _FakeManager: + """Stands in for a KV cache manager that publishes an admission bound.""" + + def __init__(self, max_admissible_sequences): + self.max_admissible_sequences = max_admissible_sequences + + +def test_validator_accepts_the_matching_pair(): + validate_seq_slot_pool_covers_admission(16, _FakeManager(16)) + + +def test_validator_rejects_an_index_pool_below_the_seat_pool(): + """The direction that shipped as nvbug 6627795, and it is never legitimate. + + A one-sided ``seats >= admissible`` guard is what let it through: the seat + pool grew to 2B while the index pool stayed at B+1, which satisfies the + one-sided form and silently defers admitted requests one at a time. + """ + with pytest.raises(ValueError, match="smaller than the seat"): + validate_seq_slot_pool_covers_admission(16, _FakeManager(8)) + + +def test_validator_permits_an_index_pool_above_the_seat_pool(): + """A surplus of index leases is the design, in two independent ways. + + Under disaggregation a request in KV transfer holds its index lease with no + seat (``SeqSlotManager.prepare_resources`` skips ``DISAGG_GENERATION_INIT``), + so the manager's ``is_disagg`` term deliberately has no seat-pool + counterpart. And a hybrid model suppresses the seat headroom -- SSM state is + sized from ``max_batch_size`` -- while a V2 index pool behind it still + widens. Admission is bounded independently at ``max_batch_size * pp_size``, + so the surplus is never extra concurrency -- and it is the cheap direction: + spare leases cost page-table rows, spare seats would cost sampler and + speculative-decoding state. + + The shared overlap term is what keeps this from being the *common* case + rather than the exception: were the seat gate narrowed back to attention DP, + every plain-TP overlap run would land here and the tolerated gap would stop + being evidence of anything. + + Asserting on the signature as well: ``is_disagg`` used to select which + surplus was tolerated, and its absence is what stops the validator from + growing a second copy of the manager's predicate. + """ + validate_seq_slot_pool_covers_admission(16, _FakeManager(32)) + assert "is_disagg" not in inspect.signature(validate_seq_slot_pool_covers_admission).parameters + + +def test_validate_seq_slot_pool_ignores_managers_without_a_bound(): + """The V1/C++ manager does not publish one; the check must not fire.""" + validate_seq_slot_pool_covers_admission(1, Mock(spec=[])) + validate_seq_slot_pool_covers_admission(1, None) + + +def test_validate_seq_slot_pool_ignores_a_non_integer_bound(): + """A bare ``Mock`` auto-creates the attribute, so ``is None`` is not enough. + + ``create_py_executor_instance`` is called with a ``Mock()`` cache manager in + other modules' tests (e.g. test_dual_pool_kv_cache). Keying the opt-in on + ``is None`` would let a ``Mock`` attribute reach the comparison and raise + ``TypeError`` from a startup validator. + """ + validate_seq_slot_pool_covers_admission(1, Mock()) + + +@pytest.mark.parametrize( + "cache_transceiver_config,expected", + [ + (None, False), + (SimpleNamespace(backend=None), False), + (SimpleNamespace(backend="UCX"), True), + ], +) +def test_is_disagg_enabled_is_the_single_definition(cache_transceiver_config, expected): + """One definition of "this is a disaggregated server". + + The ``backend is not None`` test used to be inlined at each use site, which is + how a derived fact acquires copies that then disagree. + """ + assert is_disagg_enabled(cache_transceiver_config) is expected + + +@pytest.mark.parametrize( + "explicit,engine_seats,expected", + [ + (24, 16, 24), # an explicit value wins + (None, 16, 16), # otherwise the engine's own pool + (None, None, 16), # only then recompute, *with* the engine's gate + ], +) +def test_resolve_max_num_sequences_prefers_the_published_pool(explicit, engine_seats, expected): + """The fallback must never be able to undercut the pool it indexes. + + The recomputing branch used to be the *first* branch and was called without + ``enable_overlap_headroom``, so a caller that omitted ``max_num_sequences`` + silently sized the sampler and the executor's SeqSlotManager below the index + pool they share indices with. The third row is that branch, and it must still + land on the headroom value. + """ + engine = SimpleNamespace( + max_num_seq_slots=engine_seats, + _enable_overlap_headroom=True, + ) + mapping = Mapping(world_size=1, tp_size=1, pp_size=1, enable_attention_dp=True) + llm_args = SimpleNamespace(disable_overlap_scheduler=False) + + assert ( + resolve_max_num_sequences( + engine, + mapping, + 8, + llm_args, + max_num_sequences=explicit, + ) + == expected + ) + + +def test_resolve_max_num_sequences_reads_llm_args_only_in_the_fallback(): + """The two short-circuit branches must not touch ``llm_args`` at all. + + Reading ``disable_overlap_scheduler`` at the *call site* made every caller + depend on a field only the third branch uses, which broke callers that hold a + lighter args object and pass ``max_num_sequences`` explicitly. An args object + that raises on attribute access is the only way to state that as a test: + asserting on the return value cannot distinguish "not used" from "used and + happened to agree". + """ + + class _Exploding: + def __getattr__(self, name): + raise AssertionError(f"llm_args.{name} read on a path that must not need it") + + engine_with_pool = SimpleNamespace(max_num_seq_slots=16) + mapping = Mapping(world_size=1, tp_size=1, pp_size=1, enable_attention_dp=True) + + # Branch 1: an explicit value wins, even with no pool published at all. + assert ( + resolve_max_num_sequences( + SimpleNamespace(), + mapping, + 8, + _Exploding(), + max_num_sequences=24, + ) + == 24 + ) + # Branch 2: the engine's published pool. + assert resolve_max_num_sequences(engine_with_pool, mapping, 8, _Exploding()) == 16 + + +def test_sampler_args_require_the_resolved_pool(): + """``max_num_sequences`` is required, and the raw material for re-deriving it + is gone from the signature. + + ``create_torch_sampler_args`` used to default it by recomputing from + ``mapping``/``max_batch_size`` without the headroom gate, i.e. it could only + ever produce a number smaller than the slots the sampler indexes. + """ + params = inspect.signature(create_torch_sampler_args).parameters + + assert params["max_num_sequences"].default is inspect.Parameter.empty + assert "mapping" not in params + assert "max_batch_size" not in params + + @pytest.mark.parametrize("slot_factor", [1, 2]) def test_sampler_uses_executor_slot_pool_capacity(slot_factor): max_batch_size = 8 - mapping = Mapping(world_size=1, tp_size=1, pp_size=1) max_num_sequences = max_batch_size * slot_factor args = create_torch_sampler_args( - mapping, max_seq_len=1024, - max_batch_size=max_batch_size, speculative_config=None, max_beam_width=1, disable_overlap_scheduler=False, @@ -137,3 +484,53 @@ def test_sampler_uses_executor_slot_pool_capacity(slot_factor): max_num_sequences=max_num_sequences, ) assert args.max_num_sequences == max_num_sequences + + +def _make_kv_cache_creator(disable_overlap_scheduler: bool, is_v2: bool = True) -> KvCacheCreator: + """Minimal creator whose only job is to reach _create_kv_cache_manager.""" + c = object.__new__(KvCacheCreator) + c._mapping = Mapping(world_size=1, tp_size=1, pp_size=1) + c._kv_cache_config = Mock() + c._tokens_per_block = 32 + c._max_seq_len = 1024 + c._max_batch_size = 8 + c._max_num_tokens = 8192 + c._max_beam_width = 1 + c._speculative_config = None + c._sparse_attention_config = None + c._kv_connector_manager = None + c._execution_stream = None + c._is_disagg = False + c._is_kv_cache_manager_v2 = is_v2 + c._disable_overlap_scheduler = disable_overlap_scheduler + # Short-circuit the post-construction max_seq_len fixup. + c._skip_est = True + c._get_model_kv_cache_manager_cls = Mock(return_value=Mock()) + c._should_create_separate_draft_kv_cache = Mock(return_value=False) + c._enable_kv_cache_stats = Mock(return_value=False) + return c + + +@pytest.mark.parametrize("disable_overlap_scheduler", [False, True]) +def test_kv_cache_manager_receives_the_overlap_scheduler_flag(disable_overlap_scheduler): + """The index pool needs the overlap flag, and only the creator can supply it. + + ``KVCacheManagerV2`` derives its own lease headroom from + ``is_disagg or (overlap on and not has_pp)`` -- deliberately looser than the + seat gate, since over-leasing is cheap and under-leasing hangs. That makes + ``disable_overlap_scheduler`` load-bearing on the constructor: drop it and + every aggregated V2 deployment silently reverts to a single cohort of leases, + which is the nvbug 6627795 shortfall. + """ + creator = _make_kv_cache_creator(disable_overlap_scheduler) + model_engine = SimpleNamespace( + model=SimpleNamespace(model_config=SimpleNamespace(is_generation=True)), + ) + + with patch( + "tensorrt_llm._torch.pyexecutor._util._create_kv_cache_manager", + return_value=None, + ) as create: + creator._create_kv_cache_manager(model_engine) + + assert create.call_args.kwargs["disable_overlap_scheduler"] is disable_overlap_scheduler diff --git a/tests/unittest/_torch/modeling/test_qwen4_exp_support.py b/tests/unittest/_torch/modeling/test_qwen4_exp_support.py index 3196cfde67fb..50e9285c7268 100644 --- a/tests/unittest/_torch/modeling/test_qwen4_exp_support.py +++ b/tests/unittest/_torch/modeling/test_qwen4_exp_support.py @@ -1110,6 +1110,7 @@ def make_manager( hidden_size: int, max_num_requests: int, sa_manager: object = None, + num_seq_slots: Optional[int] = None, ) -> object: captured.update( config=config, @@ -1117,6 +1118,7 @@ def make_manager( hidden_size=hidden_size, max_num_requests=max_num_requests, sa_manager=sa_manager, + num_seq_slots=num_seq_slots, ) return captured @@ -1146,6 +1148,9 @@ def make_manager( assert utils.get_spec_resource_manager(model_engine) is captured assert captured["hidden_size"] == 512 assert captured["max_num_requests"] == 16 + # This engine stub does not opt into the overlap seq-slot headroom, so the + # manager must fall back to the max_num_requests sizing. + assert captured["num_seq_slots"] is None def test_logits_processor_borrows_target_mixer_but_mtp_head_owns_one() -> None: diff --git a/tests/unittest/_torch/speculative/test_spec_slot_pool_sizing.py b/tests/unittest/_torch/speculative/test_spec_slot_pool_sizing.py new file mode 100644 index 000000000000..c7d6809eeb62 --- /dev/null +++ b/tests/unittest/_torch/speculative/test_spec_slot_pool_sizing.py @@ -0,0 +1,441 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Speculative-decoding state that is keyed by live-request identity must be +sized by the sequence-slot pool, not by max_batch_size. + +The two differ whenever ``compute_max_num_sequences`` widens the pool: under +pipeline depth, and under the overlap headroom where a finished request holds its +slot for one more iteration while its replacement is already admitted +(nvbug-6627795). Two distinct families follow from that, and only the first needs +the pool size: + +* keyed by ``py_seq_slot`` / a per-request ``SlotManager`` slot -- must span the + pool. ``SpecMetadata.num_seq_slots`` (draft_probs, full_draft_probs, + penalty_state), ``MTPHiddenStatesManager``'s hidden-state pools, + ``DynamicTreeSlotStorage``, ``Eagle3ResourceManager.slot_manager`` and the + ``SuffixAutomatonManager`` slot pool. +* keyed by *batch position* -- ``max_num_requests`` is correct and deliberately + unchanged, because the micro-batch scheduler caps every forward at + max_batch_size (its ``no_schedule_after_state=GENERATION_TO_COMPLETE`` default + keeps the retiring requests out of the batch entirely). ``SpecTreeManager``'s + per-forward work buffers and ``batch_indices_cuda`` are in this family. +""" + +import ast +import inspect +import textwrap +import types + +import pytest +import torch + +from tensorrt_llm._torch.speculative.eagle3 import ( + Eagle3OneModelDynamicTreeResourceManager, + Eagle3ResourceManager, +) +from tensorrt_llm._torch.speculative.mtp import MTPHiddenStatesManager +from tensorrt_llm._torch.speculative.mtp_dynamic_tree import MTPEagleDynamicTreeResourceManager +from tensorrt_llm._torch.speculative.spec_tree_manager import SpecTreeManager +from tensorrt_llm._torch.speculative.suffix_automaton import SAConfig, SuffixAutomatonManager +from tensorrt_llm._torch.speculative.utils import ( + _build_spec_metadata, + get_spec_metadata, + get_spec_resource_manager, +) + +R, POOL = 8, 16 # max_batch_size, 2 * max_batch_size (overlap headroom) + + +@pytest.mark.cpu_only +def test_slot_pool_size_is_applied_centrally(monkeypatch): + """``get_spec_metadata`` stamps the pool size onto whatever mode was built. + + This is the property that fixes the review finding: previously only the + MTP-eagle branch forwarded ``num_seq_slots``, so vanilla MTP, Eagle3 + one-model, PARD, DFlash/DSpark and draft-target one-model all sized their + slot-indexed buffers at ``max_num_requests``. Applying it once at the single + exit point makes it impossible for a mode -- including a future one -- to be + missed, so the assertion deliberately does not name any mode. + """ + built = types.SimpleNamespace() + monkeypatch.setattr( + "tensorrt_llm._torch.speculative.utils._build_spec_metadata", lambda *a, **k: built + ) + spec_config = types.SimpleNamespace(enable_penalty=False) + + out = get_spec_metadata( + spec_config, + model_config=object(), + max_num_requests=R, + max_num_tokens=128, + num_seq_slots=POOL, + ) + + assert out is built + assert out.num_seq_slots == POOL + + +@pytest.mark.cpu_only +def test_unknown_slot_pool_leaves_the_max_num_requests_fallback(monkeypatch): + """``num_seq_slots=None`` must not be written as a literal. + + Both allocators resolve the pool as ``self.num_seq_slots or + self.max_num_requests``, so leaving the dataclass default (0) in place is how + a caller that does not know the pool size keeps the old sizing. Writing + ``None`` would raise in the ``+ 1`` scratch-row arithmetic instead. + """ + built = types.SimpleNamespace() + monkeypatch.setattr( + "tensorrt_llm._torch.speculative.utils._build_spec_metadata", lambda *a, **k: built + ) + spec_config = types.SimpleNamespace(enable_penalty=False) + + get_spec_metadata( + spec_config, + model_config=object(), + max_num_requests=R, + max_num_tokens=128, + num_seq_slots=None, + ) + + assert not hasattr(built, "num_seq_slots") + + +@pytest.mark.cpu_only +def test_per_mode_builder_does_not_take_the_pool_size(): + """Guard the central-application invariant structurally. + + Re-plumbing ``num_seq_slots`` through the per-mode constructors is what let a + branch be forgotten in the first place; keep the builder free of it. + """ + assert "num_seq_slots" not in inspect.signature(_build_spec_metadata).parameters + + +def _mtp_config(): + return types.SimpleNamespace(max_draft_len=2, use_relaxed_acceptance_for_thinking=True) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="MTP hidden-state pools are CUDA tensors") +@pytest.mark.parametrize( + "num_seq_slots,expected_pool", + [ + (POOL, POOL + 1), + (None, R + 1), + ], +) +def test_mtp_hidden_states_pool_spans_the_slot_pool(num_seq_slots, expected_pool): + """The pool must cover every *resident* request, plus the CUDA-graph dummy. + + ``add_slot`` runs on a request's first context chunk and the slot is only + returned by ``free_resources``, which the overlap scheduler defers -- so at + ``max_num_requests + 1`` the replacement request raises ``NoFreeSlotsError``. + ``None`` keeps the pre-existing sizing for callers that do not know the pool. + """ + mgr = MTPHiddenStatesManager( + _mtp_config(), torch.float16, hidden_size=8, max_num_requests=R, num_seq_slots=num_seq_slots + ) + + assert mgr.slot_manager.max_num_requests == expected_pool + assert mgr.mtp_past_hidden_states_pool.shape[0] == expected_pool + assert mgr.mtp_past_tokens_pool.shape[0] == expected_pool + assert mgr.mtp_relaxed_delta_pool.shape[0] == expected_pool + # Batch-position state is unaffected: the forward batch is still capped at + # max_batch_size by the micro-batch scheduler. + assert mgr.get_max_resource_count() == R + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="MTP hidden-state pools are CUDA tensors") +def test_mtp_slot_pool_survives_a_full_overlap_turnover(): + """R retiring + R admitted must both hold slots at once. + + This is the exact interleaving the overlap scheduler produces and the one + that used to exhaust the pool. + """ + mgr = MTPHiddenStatesManager( + _mtp_config(), torch.float16, hidden_size=8, max_num_requests=R, num_seq_slots=POOL + ) + + retiring = [mgr.slot_manager.add_slot(rid) for rid in range(R)] + # Replacements are admitted before the deferred teardown frees the slots. + incoming = [mgr.slot_manager.add_slot(rid) for rid in range(R, 2 * R)] + + assert len(set(retiring) | set(incoming)) == 2 * R + assert all(0 <= slot < POOL + 1 for slot in retiring + incoming) + + +# --------------------------------------------------------------------------- +# Resource managers. Unlike SpecMetadata there is no single exit point to stamp +# the pool onto, so the plumbing is per-branch -- which is exactly how three +# managers were missed in a row. The AST guard below makes forgetting a branch a +# test failure instead of a runtime IndexError. +# --------------------------------------------------------------------------- + +#: Managers that legitimately do not take a slot pool. Adding a name here must be +#: a deliberate act with a reason, which is the point of the allow-list. +_MANAGERS_WITHOUT_A_SLOT_POOL = { + # The n-gram pool is keyed by pattern, not by request identity, and NGRAM is + # absent from SpeculativeDecodingMode.support_overlap_scheduler(), so + # py_executor_creator forces the overlap scheduler off and the headroom can + # never apply. + "NGramPoolManager", + # Hidden-state export path; no per-request slot pool. + "SaveHiddenStatesResourceManager", +} + +_MANAGERS_WITH_A_SLOT_POOL = ( + MTPHiddenStatesManager, + MTPEagleDynamicTreeResourceManager, + Eagle3ResourceManager, + Eagle3OneModelDynamicTreeResourceManager, + SuffixAutomatonManager, + SpecTreeManager, +) + + +@pytest.mark.cpu_only +def test_every_resource_manager_branch_forwards_the_slot_pool(): + """Mechanical guard: no branch of ``get_spec_resource_manager`` may omit it. + + ``num_seq_slots`` is computed once at the top of the function and then has to + reach every manager it builds. A new speculation mode -- or a new manager in + an existing mode's branch -- fails here rather than in production, where the + symptom is an out-of-range ``py_seq_slot`` write into a pool sized for + max_batch_size. + """ + tree = ast.parse(textwrap.dedent(inspect.getsource(get_spec_resource_manager))) + + missing = [] + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + name = getattr(node.func, "id", None) or getattr(node.func, "attr", None) + if name is None or not name.endswith("Manager") or name in _MANAGERS_WITHOUT_A_SLOT_POOL: + continue + if not any(kw.arg == "num_seq_slots" for kw in node.keywords): + missing.append(name) + + assert not missing, ( + f"get_spec_resource_manager builds {sorted(set(missing))} without forwarding " + "num_seq_slots; slot-keyed pools would be sized at max_batch_size. Either pass " + "it or justify the exemption in _MANAGERS_WITHOUT_A_SLOT_POOL." + ) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("manager", _MANAGERS_WITH_A_SLOT_POOL, ids=lambda m: m.__name__) +def test_slot_pool_managers_accept_an_optional_pool_size(manager): + """The receiving end of the same contract, with ``None`` as the default. + + ``None`` -- not ``max_num_requests`` -- has to be the default so that a caller + holding an engine that publishes no pool (unit-test stubs, mm-encoder-only + engines) keeps the established sizing without every call site restating it. + """ + param = inspect.signature(manager.__init__).parameters.get("num_seq_slots") + + assert param is not None, f"{manager.__name__} cannot be told its slot pool" + assert param.default is None, f"{manager.__name__} must default to None, got {param.default!r}" + + +def _tree_manager(num_seq_slots): + return SpecTreeManager( + max_num_requests=R, + use_dynamic_tree=True, + max_total_draft_tokens=3, + max_draft_len=3, + eagle_choices=None, + dynamic_tree_max_topK=2, + num_seq_slots=num_seq_slots, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="dynamic-tree slot storage is on CUDA") +@pytest.mark.parametrize("num_seq_slots,expected_slots", [(POOL, POOL), (None, R)]) +def test_dynamic_tree_slot_storage_spans_the_slot_pool(num_seq_slots, expected_slots): + """``DynamicTreeSlotStorage`` is documented as indexed by ``py_seq_slot``. + + It was nonetheless sized from ``num_trees`` (== max_batch_size), so the two + disagreed by 2x once the headroom was on. The dummy row sits one past the + pool, so every buffer is ``pool + 1`` deep. + """ + storage = _tree_manager(num_seq_slots).slot_storage + + assert storage.dummy_slot_id == expected_slots + for name in ( + "has_tree", + "packed_mask", + "position_offsets", + "retrieve_index", + "retrieve_next_token", + "retrieve_next_sibling", + ): + assert getattr(storage, name).shape[0] == expected_slots + 1, name + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="dynamic-tree slot storage is on CUDA") +def test_dynamic_tree_work_buffers_stay_at_max_batch_size(): + """The other family must not be widened along with it. + + ``num_trees`` indexes the build kernel's output by batch position, and the + micro-batch scheduler caps the forward at max_batch_size. Widening it would + waste memory quadratically in the tree dimensions for no benefit. + """ + mgr = _tree_manager(POOL) + + assert mgr.num_trees == R + assert mgr.retrieve_index.shape[0] == R + assert mgr.retrieve_next_token.shape[0] == R + assert mgr.retrieve_next_sibling.shape[0] == R + assert mgr.num_slots == POOL + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="dynamic-tree slot storage is on CUDA") +def test_marking_a_high_slot_invalid_needs_the_pool(): + """The concrete failure, plus a negative control that it was reachable. + + ``Eagle3OneModelDynamicTreeResourceManager.free_resources`` calls + ``mark_invalid(request.py_seq_slot)``, and with the headroom on ``py_seq_slot`` + ranges over the whole pool. Sized at max_batch_size the write is out of + range, so the second half of the assertion is what proves the first half is + not vacuous. + """ + _tree_manager(POOL).slot_storage.mark_invalid(POOL - 1) + + with pytest.raises(IndexError): + _tree_manager(None).slot_storage.mark_invalid(POOL - 1) + + +def _eagle_config(): + # Deliberately not an EagleDecodingConfig: that keeps max_total_draft_tokens + # on the max_draft_len branch and leaves spec_tree_manager unbuilt, so this + # exercises slot_manager sizing only. + return types.SimpleNamespace( + max_draft_len=2, + num_capture_layers=1, + use_relaxed_acceptance_for_thinking=True, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Eagle3 hidden states are CUDA tensors") +@pytest.mark.parametrize("num_seq_slots,expected_pool", [(POOL, POOL + 1), (None, R + 1)]) +def test_eagle3_slot_manager_spans_the_slot_pool(num_seq_slots, expected_pool): + """``Eagle3ResourceManager`` sized its ``SlotManager`` from ``max_seq_len``. + + That is a token count standing in for a slot count -- accidentally generous + for most configurations, but not for ``max_batch_size == max_seq_len``, where + the pool lands exactly one slot short of a full overlap turnover. + """ + mgr = Eagle3ResourceManager( + _eagle_config(), + torch.float16, + hidden_size=8, + max_num_requests=R, + max_seq_len=4, + max_num_tokens=64, + num_seq_slots=num_seq_slots, + ) + + assert mgr.slot_manager.max_num_requests == expected_pool + assert mgr.relaxed_delta_pool.shape[0] == expected_pool + assert len(mgr.seq_lens) == expected_pool + assert len(mgr.start_indices) == expected_pool + # Batch-position state is untouched. + assert mgr.batch_indices_cuda.shape[0] == R + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Eagle3 hidden states are CUDA tensors") +def test_eagle3_keeps_the_max_seq_len_floor(): + """Existing deployments must not shrink. + + ``max_seq_len`` stays a floor so that every configuration where it already + exceeded the slot pool allocates exactly what it did before this change. + """ + mgr = Eagle3ResourceManager( + _eagle_config(), + torch.float16, + hidden_size=8, + max_num_requests=R, + max_seq_len=1024, + max_num_tokens=64, + num_seq_slots=POOL, + ) + + assert mgr.slot_manager.max_num_requests == 1024 + 1 + + +def _sa_manager(num_seq_slots, **config_kwargs): + config = SAConfig(max_seq_len=1024, max_slots=R, **config_kwargs) + return SuffixAutomatonManager(config, R, 1024, num_seq_slots=num_seq_slots) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("num_seq_slots,expected_pool", [(POOL, POOL), (None, R)]) +def test_sa_pool_spans_the_slot_pool(num_seq_slots, expected_pool): + """SA slots are held for a request id's lifetime, so the pool follows it. + + The dummy slot index is derived from ``pool_size``, so it moves with the pool + rather than colliding with a real slot. + """ + mgr = _sa_manager(num_seq_slots) + + assert mgr.pool_size == expected_pool + assert len(mgr._free_slots) == expected_pool + assert mgr._dummy_slot_index == expected_pool + + +@pytest.mark.cpu_only +def test_sa_pool_survives_a_full_overlap_turnover(): + """2 * max_batch_size concurrent slots, with a negative control. + + Without the pool the (max_batch_size + 1)-th allocation has nothing free and + nothing retained to evict, which is a hard ``RuntimeError`` mid-run. + """ + mgr = _sa_manager(POOL) + slots = [mgr._allocate_slot() for _ in range(POOL)] + assert len(set(slots)) == POOL + + starved = _sa_manager(None) + for _ in range(R): + starved._allocate_slot() + with pytest.raises(RuntimeError, match="No free or retained slots"): + starved._allocate_slot() + + +@pytest.mark.cpu_only +def test_an_explicit_sa_pool_is_a_floor_not_a_rejection(): + """The seat count raises ``global_pool_size``; it never fails the run. + + ``TorchLlmArgs.validate_speculative_config`` accepts any + ``global_pool_size >= max_batch_size``, so rejecting a value between + max_batch_size and the seat count would make merely enabling the overlap + scheduler turn an already-validated config into a startup error. The last + assertion is the negative control: without the headroom the configured value + is used verbatim, so this is a floor and not an unconditional bump. + """ + assert _sa_manager(POOL, enable_global_pool=True, global_pool_size=64).pool_size == 64 + assert _sa_manager(POOL, enable_global_pool=True, global_pool_size=R).pool_size == POOL + assert _sa_manager(None, enable_global_pool=True, global_pool_size=R).pool_size == R + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="dynamic-tree slot storage is on CUDA") +def test_eagle3_one_model_dynamic_tree_forwards_the_slot_pool(): + """End-to-end for the manager whose ``free_resources`` triggers the write.""" + config = types.SimpleNamespace( + use_dynamic_tree=True, + max_draft_len=3, + tokens_per_gen_step=4, + eagle_choices=None, + dynamic_tree_max_topK=2, + ) + + mgr = Eagle3OneModelDynamicTreeResourceManager(config, R, num_seq_slots=POOL) + + assert mgr.spec_tree_manager.slot_storage.dummy_slot_id == POOL + assert mgr.spec_tree_manager.num_trees == R + assert mgr.batch_indices_cuda.shape[0] == R + mgr.free_resources(types.SimpleNamespace(py_seq_slot=POOL - 1)) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"]))