From 04fe30a9f89f8aabd72d632580f15c752972d954 Mon Sep 17 00:00:00 2001 From: Chenfei Zhang Date: Mon, 31 Aug 2026 08:31:07 -0700 Subject: [PATCH 01/22] [https://nvbugs/6627795][fix] stop charging retiring requests against ADP admission and capacity PR #17390 flipped `disable_overlap_scheduler` true->false (overlap ENABLED) on several perf-sanity worker configs and cost disagg-e2e-gb300_glm-5-fp4_8k1k_con1024_ctx1_dep2_gen1_dep8_eplb256_mtp1_ccb-NIXL 20.40% throughput. With overlap enabled a finished request's teardown is deferred by one iteration: `_process_previous_batch` -- the only thing that removes a finished request from `PyExecutor.active_requests` -- runs ~200 lines AFTER `_fetch_new_requests` in the same `_executor_loop_overlap` body. So requests in GENERATION_TO_COMPLETE are still in the active list when the next batch is admitted, and were charged against it three times over: 1. the ADP router balanced load on them, so `_expected_num_active_requests` floored `expected` at a phantom per-rank load and its heap filter then excluded the "loaded" rank entirely -- one rank idle every iteration; 2. `_pop_from_waiting_queue` spent global admission budget on them (`admission_capacity - total_num_active_requests`); 3. the C++ capacity scheduler counted them toward `mMaxNumRequests`: the `numAdmittedRequests >= mMaxNumRequests` break sits after the state gate and before classification, and `isGenerationInProgressState()` includes kGENERATION_TO_COMPLETE. Charge 3 is the binding one, and it needs sequence-slot headroom to be actionable, so all three are fixed together: * `adp_router.py`: filter the retiring requests out of the active list once, in `gather_all_rank_states`, and route on that. One choke point corrects `num_active_requests` and `num_active_tokens` for all three routers and keeps `create_rank_state` overlap-agnostic. The count is reported in a new `RankState.num_retiring_requests` field. * `py_executor.py`: fold that count back in for the idle-fetch liveness test only. Liveness is collective -- a rank reporting zero routable work would block on the untimed request-queue wait while its peers blocked in the broadcast, and end-of-run drain hits exactly that state. Also measure the dummy-request pad surplus against the routable count, so its warning does not fire every iteration. * `scheduler.py`: `BindCapacityScheduler` now passes `no_schedule_after_state=GENERATION_TO_COMPLETE`, matching every micro-batch scheduler. The KV cache of a retiring request is released by the teardown that is already queued, so keeping it inside the capacity window bought nothing. * `_util.py`: `should_enable_disagg_adp_overlap_headroom` -> `should_enable_adp_overlap_seq_slot_headroom`, no longer gated on disaggregation. The regression reproduced on an aggregated context-only run with no cache transceiver configured, and without the headroom the capacity change has no free slot to backfill into (it raises NoFreeSlotsError on a pool sized 1x max_batch_size). Measured on the ctx worker of the regressing glm-5 case, four arms at a6ea52f7ea, matched nodes, no nsys, ADP-router tracing on all of them: | arm | tput | vs bug | fwd batch | |---------------------------------------|----------|---------|-----------| | overlap disabled (pre-#17390) | 34672.51 | +25.8% | 1.999 | | overlap enabled (#17390, the bug) | 27556.82 | -- | 1.000 | | + charges 1+2 only | 27635.17 | +0.28% | 1.000 | | + charges 1+2+3 and slot headroom | 35502.08 | +28.8% | 2.000 | The fixed arm admits 2.00 requests/rank/iteration (the configured max_batch_size) on 2559/2563 iterations, versus 0.75 for the bug, and finishes the same 10240 requests in 2563 iterations instead of 6828 -- slightly ahead of the overlap-disabled arm, so the regression is recovered rather than merely reduced. Run-to-run spread on this rig is +/-0.3%. Follow-up, deliberately not in this change: `batch_size_input = len(self.active_requests)` feeding `drafter.get_draft_len_for_batch_size` is reachable only with spec-dec plus an explicit `draft_len_schedule` and has the same staleness. Signed-off-by: Chenfei Zhang --- tensorrt_llm/_torch/pyexecutor/_util.py | 26 ++- .../_torch/pyexecutor/model_engine.py | 25 ++- tensorrt_llm/_torch/pyexecutor/py_executor.py | 41 ++++- .../_torch/pyexecutor/py_executor_creator.py | 9 +- .../_torch/pyexecutor/scheduler/adp_router.py | 86 +++++++++- .../_torch/pyexecutor/scheduler/scheduler.py | 11 +- .../_torch/executor/test_adp_router.py | 154 +++++++++++++++++- .../executor/test_kvcache_aware_router.py | 58 +++++++ .../_torch/executor/test_py_executor.py | 37 +++++ .../_torch/executor/test_seq_slot_sizing.py | 35 ++-- 10 files changed, 422 insertions(+), 60 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index b4422c084463..aa289e643970 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2852,14 +2852,24 @@ 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() +def should_enable_adp_overlap_seq_slot_headroom( + mapping: Mapping, disable_overlap_scheduler: bool) -> bool: + """Gate extra sequence slots to non-PP attention-DP with overlap enabled. + + The overlap scheduler defers a finished request's teardown by one iteration, + so its sequence slot is still held when the ADP router admits the batch that + replaces it. Without spare slots the router cannot backfill and the forward + batch runs short (nvbug-6627795); a second set of slots lets admission reach + max_batch_size on every rank every iteration. + + Requiring attention DP -- rather than merely overlap -- is deliberate: with + a single scheduling domain the executor's own admission bound already tracks + the pool, whereas ADP admits against an allgathered load vector that the + deferred teardown desynchronizes. Not gated on disaggregation: the mechanism + is a property of overlap plus ADP admission, and was measured on an + aggregated context-only run with no cache transceiver configured. + """ + return (mapping.enable_attention_dp and not mapping.has_pp() and not disable_overlap_scheduler) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index cbd37669a7fa..65c734cbdf14 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -451,23 +451,22 @@ 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. + # 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_adp_overlap_seq_slot_headroom, should_enable_non_overlap_adp_forward_intent, 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_overlap_seq_slot_headroom = ( + should_enable_adp_overlap_seq_slot_headroom( + mapping, 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, + enable_overlap_headroom=self._enable_adp_overlap_seq_slot_headroom, ) self.dist = dist if dist is not None: @@ -3764,11 +3763,11 @@ def _set_up_spec_metadata( spec_resource_manager: Optional[BaseResourceManager], no_cache=False): 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) + # The 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_adp_overlap_seq_slot_headroom else None) if no_cache: return get_spec_metadata( spec_config, diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index c12227a2fcff..26406233d669 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -100,7 +100,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 @@ -5592,8 +5592,16 @@ 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: - """Fetch requests from request_queue and enqueue to waiting_queue.""" + total_num_live_requests: int) -> None: + """Fetch requests from request_queue and enqueue to waiting_queue. + + `total_num_live_requests` counts every request still resident on any + rank, including the retiring ones that `num_active_requests` excludes. + The idle decision below must be taken on that figure and it must be + identical on every rank: it selects a blocking versus a zero timeout, + and a rank that blocks on the untimed queue wait while its peers reach + `dist.broadcast(root=0)` deadlocks the iteration. + """ # Block new requests while control requests are pending if len(self.control_requests) != 0: return @@ -5603,7 +5611,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 @@ -5811,14 +5819,21 @@ def _fetch_new_requests( s.num_active_requests for s in all_rank_states ] total_num_active_requests = sum(all_ranks_num_active_requests) + # Retiring requests are excluded from num_active_requests (they + # cannot be scheduled, so they must not consume admission + # capacity -- nvbug-6627795) but they are still resident, so the + # loop is NOT idle while any of them exists. Fold them back in + # for the liveness test only. + 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( @@ -7000,7 +7015,14 @@ 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): + # Compare against the same routable count the router balanced on: + # create_rank_state excludes retiring requests from the per-rank loads + # that floor `expected` (nvbug-6627795), so measuring against the raw + # len() here would make the warning below fire every iteration. + num_routable_active_requests = ( + len(self.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)), @@ -7021,11 +7043,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 039f81b06ff1..29ac2079a94c 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -745,13 +745,12 @@ def allocation_scope(current_stage: ExecutorMemoryType): 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) + model_engine, "_enable_adp_overlap_seq_slot_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. + # The attention-DP overlap path follows the expanded slot + # pool. Other configurations retain max_batch_size. "max_num_sequences": guided_decoder_slots, "vocab_size_padded": model_engine.model.vocab_size_padded, "rank": mapping.rank, diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py index 220136a4b6d3..83e08bdc393d 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,58 @@ 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. + + A request in ``GENERATION_TO_COMPLETE`` has produced its final token; the + next ``_update_request_states`` tears it down, and no scheduler will ever + put it in a forward batch again (every micro-batch scheduler bounds its + window at ``GENERATION_TO_COMPLETE``). Without the overlap scheduler that + teardown happens in-line, so such requests are gone before the ADP router + next runs. With overlap enabled it is deferred by one iteration, leaving + them in ``active_requests`` when the router builds its load vector -- where + they inflate per-rank load *and* consume global admission capacity that + nothing can use (nvbug-6627795). + + ``GENERATION_TO_COMPLETE`` is the only state treated this way. The disagg + limbo states -- ``DISAGG_CONTEXT_WAIT_SCHEDULER``, + ``DISAGG_GENERATION_INIT``, ``DISAGG_GENERATION_TRANS_IN_PROGRESS``, + ``DISAGG_CONTEXT_TRANS_IN_PROGRESS`` and ``DISAGG_TRANS_ERROR`` -- also + linger in ``active_requests``, but they still own a sequence slot and KV + cache, so they must keep counting as load. + """ + return request.state == LlmRequestState.GENERATION_TO_COMPLETE + + +def build_active_requests_for_overlap(active_requests): + """Return ``active_requests`` minus the requests that are already retiring. + + This is the ADP router's view of the active list, and *only* the router's: + the requests are dropped from the load vector, not from the executor. Under + the overlap scheduler ``_process_previous_batch`` -- the only thing that + removes a finished request from ``PyExecutor.active_requests`` -- runs some + two hundred lines *after* ``_fetch_new_requests`` in the same + ``_executor_loop_overlap`` body, so the router would otherwise route + against a list that is one teardown stale (nvbug-6627795). + + Filtering the list at the single ``gather_all_rank_states`` choke point + rather than adjusting a count inside each ``create_rank_state`` corrects + ``num_active_requests`` *and* ``num_active_tokens`` for every router + implementation at once, and keeps the routers ignorant of overlap. + """ + return [req for req in active_requests if not is_retiring_request(req)] + + +def count_retiring_requests(active_requests) -> int: + """Count the requests that ``build_active_requests_for_overlap`` filters out. + + Used where only the size of the correction is needed, not the list itself + (the dummy-request pad path, which compares the router's ``expected`` + against a rank's routable active count). + """ + 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.""" @@ -97,8 +151,16 @@ class RankState: """ rank: int + # Routable load only. ``gather_all_rank_states`` hands ``create_rank_state`` + # the overlap-corrected list, so retiring requests are absent from both + # counts below (see ``build_active_requests_for_overlap``). This is what the + # router balances on and what bounds admission. num_active_requests: int = 0 num_active_tokens: int = 0 + # Requests filtered out of the two counts above because they are retiring. + # Reported so the inference loop can tell "nothing routable" from "nothing + # at all" and keep its idle-fetch wait collective (nvbug-6627795). + 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 +174,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 +182,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 +203,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:]), ) @@ -256,7 +320,20 @@ 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 []) + # Route on the overlap-corrected list: 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). Applied + # here rather than in each create_rank_state so every router -- and both + # num_active_requests and num_active_tokens -- is corrected at once. + active_requests_for_overlap = build_active_requests_for_overlap(active_requests) + num_retiring_requests = len(active_requests) - len(active_requests_for_overlap) + local_state = self.create_rank_state(active_requests_for_overlap, new_requests or []) + # The retiring requests are still resident, so the executor loop is NOT + # idle while any of them exists. Report the count so the idle-fetch wait + # stays collective: liveness is a global property, and a rank that + # reported zero would block on the untimed request-queue wait while its + # peers blocked in the allgather. + 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] @@ -985,8 +1062,9 @@ 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 (retiring requests excluded, + # matching create_rank_state) and warns if it comes up short. expected_num_active_requests = max( expected_num_active_requests, max(all_ranks_num_active_requests) ) diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py index 99a1710936f4..a0963245c6b3 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py @@ -452,7 +452,16 @@ def __init__( has_kv_cache_manager=kv_cache_manager is not None, two_step_lookahead=two_step_lookahead, no_schedule_until_state=no_schedule_until_state, - no_schedule_after_state=LlmRequestState.GENERATION_COMPLETE, + # Stop the window one state early, at GENERATION_TO_COMPLETE. That + # state means "final token produced, teardown deferred by the overlap + # scheduler": no micro-batch scheduler will ever forward such a + # request again, yet the C++ capacity loop's + # `numAdmittedRequests >= mMaxNumRequests` break sits after the state + # gate and before classification, so each one consumed a slot in + # mMaxNumRequests and shortened the real forward batch by one + # (nvbug-6627795). Their KV cache is released by the teardown that is + # already queued, so keeping them inside the window bought nothing. + no_schedule_after_state=LlmRequestState.GENERATION_TO_COMPLETE, enable_prefix_aware_scheduling=enable_prefix_aware_scheduling, ) diff --git a/tests/unittest/_torch/executor/test_adp_router.py b/tests/unittest/_torch/executor/test_adp_router.py index 250b214c86e1..837a59610c79 100644 --- a/tests/unittest/_torch/executor/test_adp_router.py +++ b/tests/unittest/_torch/executor/test_adp_router.py @@ -12,6 +12,7 @@ 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 +23,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 @@ -130,6 +134,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 +213,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 +226,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 +362,57 @@ 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) + 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) + 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_create_rank_state_cp_helix(self): dist = _mock_dist(tp_rank=1, has_cp_helix=True) router = DefaultADPRouter(dist=dist) @@ -1254,6 +1389,23 @@ 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): cfg = MagicMock() cfg.kv_cache_routing_conversation_affinity = True diff --git a/tests/unittest/_torch/executor/test_kvcache_aware_router.py b/tests/unittest/_torch/executor/test_kvcache_aware_router.py index 1f173ed0da83..0cb94c3932db 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, @@ -116,6 +117,63 @@ 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) + + 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() diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 48d6f7a01434..6a8355411a81 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -2136,6 +2136,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..236a8c4c6fe8 100644 --- a/tests/unittest/_torch/executor/test_seq_slot_sizing.py +++ b/tests/unittest/_torch/executor/test_seq_slot_sizing.py @@ -1,11 +1,11 @@ # 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. +"""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 +prepare_resources runs, while the capacity 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. @@ -21,11 +21,10 @@ compute_max_num_sequences, create_torch_sampler_args, should_enable_adp_dummy_fixes, - should_enable_disagg_adp_overlap_headroom, + should_enable_adp_overlap_seq_slot_headroom, should_enable_non_overlap_adp_forward_intent, should_enable_scheduler_aware_adp_dummy, ) -from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig from tensorrt_llm.mapping import Mapping SIZING_CASES = [ @@ -33,7 +32,7 @@ (1, False, True, 2), (1, False, False, 1), (1, True, True, 1), - # Existing PP sizing is preserved regardless of the DSv4 opt-in. + # Existing PP sizing is preserved regardless of the headroom opt-in. (2, False, True, 2), (4, False, True, 4), (4, True, False, 4), @@ -41,17 +40,19 @@ @pytest.mark.parametrize( - "enable_attention_dp,is_disagg,pp_size,disable_overlap,expected", + "enable_attention_dp,pp_size,disable_overlap,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), + # No cache-transceiver term: the gate no longer looks at disaggregation + # at all, because nvbug-6627795 reproduced on an aggregated context-only + # run with no transceiver configured. + (True, 1, False, True), + (False, 1, False, False), + (True, 2, False, False), + (True, 1, True, False), ], ) -def test_disagg_adp_overlap_headroom_gate( - enable_attention_dp, is_disagg, pp_size, disable_overlap, expected +def test_adp_overlap_seq_slot_headroom_gate( + enable_attention_dp, pp_size, disable_overlap, expected ): mapping = Mapping( world_size=pp_size, @@ -59,12 +60,8 @@ 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) - is expected - ) + assert should_enable_adp_overlap_seq_slot_headroom(mapping, disable_overlap) is expected @pytest.mark.parametrize("pp_size,expected", [(1, True), (2, False)]) From 0dcebd7c3a06ae5aaa4a3185f934c198fda5e286 Mon Sep 17 00:00:00 2001 From: Chenfei Zhang Date: Wed, 2 Sep 2026 02:10:58 -0700 Subject: [PATCH 02/22] [https://nvbugs/6627795][fix] gate the retiring-request filter on pipeline parallelism The overlap correction in gather_all_rank_states was not gated on pipeline parallelism, while the sequence-slot headroom that makes it safe is. Gate both on the same condition. Two independent reasons the correction must stay off under PP: - The slot pool is sized pp_size * max_batch_size with no headroom for requests a rank holds but is not charged for, so excluding retiring requests from admission can exhaust it and raise NoFreeSlotsError. - GENERATION_TO_COMPLETE is marked on the last pipeline stage only, while every rank pops from its own copy of the waiting queue. A corrected count would therefore differ per stage and the stages would admit different numbers of requests. _pad_attention_dp_dummy_request now reads the router's flag instead of re-deriving the predicate, so the two views of which requests count as routable load cannot drift. Under PP it no longer subtracts requests the router still counted, which would have under-reported the rank's load. No behaviour change for attention DP without PP, which is where the fix was measured. All 43 gb300 disaggregated perf-sanity configs with a context worker are either attention DP with pipeline_parallel_size 1 (40, unchanged) or non-attention-DP (3, router unused); none is attention DP with PP. Tests: the mock Distributed helpers now return a real bool from mapping.has_pp(), since a bare MagicMock is truthy and would have silently disabled the correction everywhere; the two stub executors that call the pad path unbound gained an adp_router. Added coverage for the PP path and for the gate itself. Signed-off-by: Chenfei Zhang --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 17 ++++++--- .../_torch/pyexecutor/scheduler/adp_router.py | 27 ++++++++++++- .../_torch/executor/test_adp_router.py | 38 ++++++++++++++++++- .../_torch/executor/test_benchmark_disagg.py | 5 +++ .../executor/test_kvcache_aware_router.py | 6 ++- .../_torch/executor/test_py_executor.py | 7 ++++ 6 files changed, 90 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 26406233d669..a852dfa05e9b 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -7016,12 +7016,17 @@ def _pad_attention_dp_dummy_request(self): expected_num_active_requests = self.expected_num_active_requests # Compare against the same routable count the router balanced on: - # create_rank_state excludes retiring requests from the per-rank loads - # that floor `expected` (nvbug-6627795), so measuring against the raw - # len() here would make the warning below fire every iteration. - num_routable_active_requests = ( - len(self.active_requests) - - count_retiring_requests(self.active_requests)) + # gather_all_rank_states excludes retiring requests from the per-rank + # loads that floor `expected` (nvbug-6627795), so measuring against the + # raw len() here would make the warning below fire every iteration. + # Read the router's own flag rather than re-deriving the gate, so the + # two can never disagree -- it is off under pipeline parallelism, where + # subtracting requests the router still counted would under-report this + # rank's load instead. + 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 diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py index 83e08bdc393d..df23c77b971d 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py @@ -229,6 +229,24 @@ class ADPRouter(ABC): def __init__(self, dist: Distributed): self.dist = dist + # Whether to route on the overlap-corrected active list (nvbug-6627795). + # + # Gated off under pipeline parallelism, deliberately matching + # ``should_enable_adp_overlap_seq_slot_headroom`` in ``_util.py``: the + # correction lets a rank hold more requests than it is charged for, so + # it is only safe where the sequence-slot pool has the matching + # headroom, and that headroom is sized for non-PP only. Beyond the slot + # accounting, ``GENERATION_TO_COMPLETE`` is marked on the last pipeline + # stage alone, while every rank pops from its own copy of the waiting + # queue -- so a PP-enabled correction would have the stages admit + # different numbers of requests and diverge. + # + # Not additionally gated on ``disable_overlap_scheduler``: without + # overlap the retire is not deferred, so no request is ever in + # ``GENERATION_TO_COMPLETE`` when the router runs and the filter is + # arithmetically a no-op (measured: zero such requests on 5126/5126 + # routing records with overlap disabled). + self.exclude_retiring_requests = not dist.mapping.has_pp() @classmethod def create( @@ -325,8 +343,13 @@ def gather_all_rank_states( # admission capacity that nothing can spend (nvbug-6627795). Applied # here rather than in each create_rank_state so every router -- and both # num_active_requests and num_active_tokens -- is corrected at once. - active_requests_for_overlap = build_active_requests_for_overlap(active_requests) - num_retiring_requests = len(active_requests) - len(active_requests_for_overlap) + # Disabled under pipeline parallelism; see exclude_retiring_requests. + 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 []) # The retiring requests are still resident, so the executor loop is NOT # idle while any of them exists. Report the count so the idle-fetch wait diff --git a/tests/unittest/_torch/executor/test_adp_router.py b/tests/unittest/_torch/executor/test_adp_router.py index 837a59610c79..478fad7d4bdb 100644 --- a/tests/unittest/_torch/executor/test_adp_router.py +++ b/tests/unittest/_torch/executor/test_adp_router.py @@ -44,12 +44,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 @@ -413,6 +417,38 @@ def test_gather_all_rank_states_reports_zero_when_all_retiring(self): assert states[0].num_active_tokens == 0 assert states[0].num_retiring_requests == 2 + def test_gather_all_rank_states_keeps_retiring_under_pp(self): + # Under pipeline parallelism the correction is off, matching the + # sequence-slot headroom gate in _util.py. Two reasons it must stay off: + # the slot pool is sized pp_size * max_batch_size with no headroom for + # requests a rank holds but is not charged for, and + # GENERATION_TO_COMPLETE is marked on the last stage only while every + # rank pops from its own copy of the waiting queue -- so a corrected + # count would have the stages admit different numbers of requests. + dist = _mock_dist(tp_rank=0, has_cp_helix=False, has_pp=True) + router = DefaultADPRouter(dist=dist) + 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_pipeline_parallelism(self): + # The flag is the single gate; _pad_attention_dp_dummy_request reads it + # rather than re-deriving the predicate, so the two cannot drift. + assert DefaultADPRouter(dist=_mock_dist(has_pp=False)).exclude_retiring_requests is True + assert DefaultADPRouter(dist=_mock_dist(has_pp=True)).exclude_retiring_requests is False + def test_create_rank_state_cp_helix(self): dist = _mock_dist(tp_rank=1, has_cp_helix=True) router = DefaultADPRouter(dist=dist) diff --git a/tests/unittest/_torch/executor/test_benchmark_disagg.py b/tests/unittest/_torch/executor/test_benchmark_disagg.py index e4b1aa14a7c4..c3bc74191809 100644 --- a/tests/unittest/_torch/executor/test_benchmark_disagg.py +++ b/tests/unittest/_torch/executor/test_benchmark_disagg.py @@ -643,6 +643,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/test_kvcache_aware_router.py b/tests/unittest/_torch/executor/test_kvcache_aware_router.py index 0cb94c3932db..43bdcf72ec85 100644 --- a/tests/unittest/_torch/executor/test_kvcache_aware_router.py +++ b/tests/unittest/_torch/executor/test_kvcache_aware_router.py @@ -34,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 @@ -42,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 diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 6a8355411a81..d26759e12268 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -1801,6 +1801,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 @@ -1833,6 +1834,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 = ( From 399890b7a18969551c735153231cad367112cc55 Mon Sep 17 00:00:00 2001 From: Chenfei Zhang Date: Wed, 2 Sep 2026 02:30:42 -0700 Subject: [PATCH 03/22] [https://nvbugs/6627795][fix] size one-model spec-decode slot pools by the seq-slot pool The attention-DP overlap headroom makes the executor's SeqSlotManager pool 2 * max_batch_size, so a finished request can hold its slot for one more iteration while its replacement is already admitted. Speculative-decoding state splits into two families under that change and only one of them was updated: - keyed by live-request identity (py_seq_slot, or a per-request SlotManager slot) -- must span the pool. - keyed by batch position -- max_num_requests is correct and is left alone, because the micro-batch scheduler caps every forward at max_batch_size and no_schedule_after_state=GENERATION_TO_COMPLETE keeps the retiring requests out of the batch entirely. Two gaps in the first family: SpecMetadata.num_seq_slots was forwarded by the MTP-eagle branch of _build_spec_metadata alone, so every other one-engine mode -- vanilla MTP, Eagle3 one-model, PARD, DFlash/DSpark, draft-target one-model -- sized draft_probs, full_draft_probs and penalty_state at max_num_requests while py_seq_slot ranged over the wider pool. Rather than adding the argument to six more constructors (the pattern that let a branch be missed in the first place), apply it once in get_spec_metadata next to enable_penalty, and drop the parameter from _build_spec_metadata so no branch can forget it. This is in time because both consumers allocate lazily, from prepare() and from the one-model sampling scan, never from __post_init__; leaving the field at its 0 default preserves the existing `num_seq_slots or max_num_requests` fallback. MTPHiddenStatesManager sized its SlotManager at max_num_requests + 1. Those slots are keyed by live-request identity -- add_slot runs on a request's first context chunk and the slot is only returned by free_resources, which overlap defers -- so vanilla MTP and MTP-eagle with use_dynamic_tree could exhaust the pool and raise NoFreeSlotsError. Take num_seq_slots and size the pool by it. Both sites are gated on the same _enable_adp_overlap_seq_slot_headroom flag _set_up_spec_metadata uses, so the resource manager and the metadata can never disagree about the pool size. Every other topology, PP included, keeps the established max_num_requests sizing. Signed-off-by: Chenfei Zhang --- tensorrt_llm/_torch/speculative/mtp.py | 15 +- .../_torch/speculative/mtp_dynamic_tree.py | 10 +- tensorrt_llm/_torch/speculative/utils.py | 43 +++-- .../speculative/test_spec_slot_pool_sizing.py | 152 ++++++++++++++++++ 4 files changed, 208 insertions(+), 12 deletions(-) create mode 100644 tests/unittest/_torch/speculative/test_spec_slot_pool_sizing.py diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 5496ef78ce1d..dad095c9e7fa 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -38,15 +38,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 attention-DP + # 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 7e81d22b2f02..f1b63e82a772 100644 --- a/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py +++ b/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py @@ -1095,6 +1095,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 @@ -1108,8 +1109,15 @@ def __init__( dynamic_tree_max_topK=config.dynamic_tree_max_topK, ) # 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/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 8efbac572925..3b3dd4b5098e 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -339,12 +339,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 @@ -354,14 +368,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). @@ -385,7 +396,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), @@ -565,6 +575,19 @@ 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 + # Pools keyed by live-request identity must follow the executor's + # SeqSlotManager pool rather than max_batch_size: the attention-DP overlap + # headroom makes it 2 * max_batch_size so a finished request can hold its slot + # for one more iteration while its replacement is admitted (nvbug-6627795). + # Buffers indexed by *batch position* deliberately keep max_num_requests -- + # the micro-batch scheduler caps every forward at max_batch_size. + # + # Opted into by the same flag ``_set_up_spec_metadata`` uses, so the manager + # and the metadata never disagree about the pool. None (the other topologies, + # PP included) preserves the established max_num_requests sizing. + num_seq_slots = None + if getattr(model_engine, "_enable_adp_overlap_seq_slot_headroom", False): + num_seq_slots = getattr(model_engine, "max_num_seq_slots", None) spec_dec_mode = spec_config.spec_dec_mode if spec_dec_mode.is_mtp_eagle_one_model(): sa_manager = None @@ -580,6 +603,7 @@ def get_spec_resource_manager(model_engine, draft_model_engine=None): model_config.hidden_size, 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 @@ -608,6 +632,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_dec_mode.is_eagle3_one_model() and _is_effective_dynamic_tree( spec_config): 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..fc5de82ab01d --- /dev/null +++ b/tests/unittest/_torch/speculative/test_spec_slot_pool_sizing.py @@ -0,0 +1,152 @@ +# 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. + +Under the attention-DP overlap headroom the two differ by 2x +(``compute_max_num_sequences``): 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) and ``MTPHiddenStatesManager``'s hidden-state pools. +* keyed by *batch position* -- ``max_num_requests`` is correct and deliberately + unchanged, because the micro-batch scheduler caps every forward at + max_batch_size (``no_schedule_after_state=GENERATION_TO_COMPLETE`` also keeps + the retiring requests out of the batch entirely). +""" + +import inspect +import types + +import pytest +import torch + +from tensorrt_llm._torch.speculative.mtp import MTPHiddenStatesManager +from tensorrt_llm._torch.speculative.utils import _build_spec_metadata, get_spec_metadata + +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) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) From e440aa56320dd60a89ffa7a930e608f27e5b772f Mon Sep 17 00:00:00 2001 From: Chenfei Zhang Date: Wed, 2 Sep 2026 03:36:52 -0700 Subject: [PATCH 04/22] [https://nvbugs/6627795][fix] charge LoRA pages for retiring requests in the C++ capacity scheduler Ending the capacity window at kGENERATION_TO_COMPLETE stops a retiring request from consuming a slot in mMaxNumRequests, which is the point of the fix. But the same `continue` also skips claimPeftPagesForRequest, and that part was not intended: the request's LoRA adapter is still resident on device, because LoraCache::markTaskDone runs as part of the teardown the overlap scheduler has deferred, and LoraCache::claimPagesWithEvict only ever evicts tasks in mDoneTasks. Leaving it uncharged overstates availablePeftPages, so a pending request carrying a *different* adapter passes the `neededPeftPages <= availablePeftPages` gate and then dies in PeftCacheManager::ensureBatch with LoraCacheFullException ("Cache is full. There are no done tasks to evict") -- turning a perf fix into a crash under LoRA. This is LoRA-only (NoOpPeftCacheManager::determineNumPages returns 0), which is why the perf A/B never surfaced it. Charge those adapters up front, before the admission loop, in both GuaranteedNoEvictScheduler::impl (which kSTATIC_BATCH also reaches) and MaxUtilizationScheduler. A retiring request must count against adapter *residency* while staying out of the forward batch. This mirrors KVCacheV2Scheduler.pre_claim_peft, which the Python V2 path already needed for the same reason. The pre-claim is idempotent with respect to the loop: claimPeftPagesForRequest dedupes on uniqTaskIds, so a window that still admits kGENERATION_TO_COMPLETE charges exactly the same total as before, and a pending request reusing a retiring request's adapter is still correctly charged zero new pages. The new test asserts both directions across all three policies -- under-charging is the bug being fixed, but over-charging would be a new one, so the same-adapter case must still be admitted for free. KV is sized generously so PEFT pages are the only binding constraint. Signed-off-by: Chenfei Zhang --- .../batch_manager/capacityScheduler.cpp | 33 +++++++ .../batch_manager/capacitySchedulerTest.cpp | 86 +++++++++++++++++++ 2 files changed, 119 insertions(+) diff --git a/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp b/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp index 335e30403d0f..7b507688957e 100644 --- a/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp +++ b/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp @@ -141,6 +141,37 @@ void claimPeftPagesForRequest(std::shared_ptr const& req, } } +//! @brief Charge PEFT pages for requests the state window excludes but whose adapters are still resident. +//! +//! kGENERATION_TO_COMPLETE means "final token produced, teardown deferred by the overlap scheduler". +//! Such a request is never forwarded again, so the admission loop skips it on the +//! getNoScheduleAfterState() gate -- but its LoRA adapter is still on device, because +//! LoraCache::markTaskDone runs as part of the teardown that has not happened yet and +//! LoraCache::claimPagesWithEvict only ever evicts tasks in mDoneTasks. Leaving it uncharged +//! overstates availablePeftPages, so a pending request carrying a *different* adapter can be admitted +//! and then fail in PeftCacheManager::ensureBatch with LoraCacheFullException ("Cache is full. There +//! are no done tasks to evict") -- turning the nvbug-6627795 admission fix into a crash under LoRA. +//! +//! Charging without consuming a request slot or a token budget is the point: a retiring request must +//! count against adapter *residency* while staying out of the forward batch. This mirrors +//! KVCacheV2Scheduler's pre_claim_peft, which the Python V2 path already needed for the same reason. +//! +//! Idempotent with respect to the admission loop: claimPeftPagesForRequest dedupes on uniqTaskIds, so +//! a window that still admits kGENERATION_TO_COMPLETE charges the same total, and a pending request +//! reusing a retiring request's adapter is still correctly charged zero new pages. +void preClaimPeftPagesForRetiringRequests(RequestList const& activeRequests, + OptionalRef peftCacheManager, SizeType32& claimedPeftPages, + std::unordered_set& uniqTaskIds) +{ + for (auto const& req : activeRequests) + { + if (req->isGenerationToCompleteState()) + { + claimPeftPagesForRequest(req, peftCacheManager, claimedPeftPages, uniqTaskIds); + } + } +} + } // namespace MaxRequestsScheduler::MaxRequestsScheduler( @@ -253,6 +284,7 @@ std::tuple GuaranteedNoEvictScheduler::impl( : std::nullopt; SizeType32 claimedPeftPages{0}; std::unordered_set uniqTaskIds{}; + preClaimPeftPagesForRetiringRequests(activeRequests, peftCacheManager, claimedPeftPages, uniqTaskIds); std::size_t numAdmittedRequests{0}; RequestVector pendingRequests; RequestVector pendingDisGenInitRequests; @@ -471,6 +503,7 @@ std::tuple MaxUtilizationScheduler::operator()( } SizeType32 numScheduledPeftPages{0}; std::unordered_set seenTaskIds; + preClaimPeftPagesForRetiringRequests(activeRequests, peftCacheManager, numScheduledPeftPages, seenTaskIds); // Keep track of blocks contributed by requests in context phase std::unordered_set newlyContributedContextBlocks; diff --git a/cpp/tests/unit_tests/batch_manager/capacitySchedulerTest.cpp b/cpp/tests/unit_tests/batch_manager/capacitySchedulerTest.cpp index 2c9757b847b2..7e7ea9e45e6f 100644 --- a/cpp/tests/unit_tests/batch_manager/capacitySchedulerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/capacitySchedulerTest.cpp @@ -659,6 +659,92 @@ TEST_F(CapacitySchedulerTest, SimpleLoraDoesntFitDuplicateTask) } } +// nvbug-6627795: the capacity window now ends at kGENERATION_TO_COMPLETE so a retiring request stops +// consuming a slot in mMaxNumRequests. Its LoRA adapter, however, is still resident on device -- +// LoraCache::markTaskDone runs as part of the teardown the overlap scheduler has deferred, and +// claimPagesWithEvict only evicts tasks in mDoneTasks. So the adapter must still be charged, or a +// request carrying a different adapter is admitted against a budget that only looks free and then dies +// in ensureBatch with LoraCacheFullException. +// +// Both directions are asserted. Under-charging is the bug; over-charging would be a new one, so the +// same-adapter case must still be admitted for free. +TEST_F(CapacitySchedulerTest, RetiringRequestStillChargesItsLoraAdapter) +{ + SizeType32 constexpr kvCacheTokensPerBlock = 10; + // Deliberately generous: PEFT pages must be the only binding constraint, otherwise a KV shortfall + // could make this test pass for the wrong reason. + SizeType32 constexpr kvCacheMaxNumTokens = 2000; + SizeType32 constexpr kvCacheMaxNumTokensPerSeq = 90; + SizeType32 constexpr maxNumRequests = 8; + // 20 pages per distinct adapter against 30 on device: one adapter fits, two never can. + SizeType32 constexpr pagesPerTask = 20; + SizeType32 constexpr maxDevicePages = 30; + uint64_t constexpr retiringTaskId = 1234; + uint64_t constexpr otherTaskId = 5678; + + auto const capacitySchedulerPolicies + = std::vector{CapacitySchedulerPolicy::kGUARANTEED_NO_EVICT, + CapacitySchedulerPolicy::kMAX_UTILIZATION, CapacitySchedulerPolicy::kSTATIC_BATCH}; + + for (auto capacitySchedulerPolicy : capacitySchedulerPolicies) + { + // true -> the incoming request reuses the retiring adapter, so it needs no new pages. + // false -> it brings a second adapter, which cannot fit alongside the resident one. + for (bool reuseRetiringAdapter : {false, true}) + { + auto kvCacheManager = getKvCacheManager( + maxNumRequests, kvCacheTokensPerBlock, kvCacheMaxNumTokens, kvCacheMaxNumTokensPerSeq); + std::shared_ptr peftCacheManager + = std::make_shared(pagesPerTask, maxDevicePages, maxDevicePages); + // The window this fix installs: kGENERATION_TO_COMPLETE is outside it. + auto capacityScheduler = CapacityScheduler(maxNumRequests, capacitySchedulerPolicy, + static_cast(kvCacheManager), /*twoStepsLookAhead=*/false, LlmRequestState::kCONTEXT_INIT, + LlmRequestState::kGENERATION_TO_COMPLETE); + + int32_t constexpr maxNewTokens = 50; + int32_t constexpr promptLen = 10; + + RequestList activeRequests; + // Final token produced, teardown deferred: not schedulable, but its adapter is still loaded. + activeRequests.push_back(createRequest(promptLen, maxNewTokens, 0, retiringTaskId, + tensorrt_llm::executor::Request::kDefaultPriority, LlmRequestState::kGENERATION_TO_COMPLETE)); + activeRequests.push_back( + createRequest(promptLen, maxNewTokens, 1, reuseRetiringAdapter ? retiringTaskId : otherTaskId)); + + auto [scheduledRequests, scheduledDisaggGenInitRequests, pausedRequests] + = capacityScheduler(activeRequests, kvCacheManager, peftCacheManager); + + // The retiring request is outside the window, so it must never be scheduled either way. + for (auto const& req : scheduledRequests) + { + EXPECT_NE(req->mRequestId, 0u) + << "retiring request was scheduled, policy " << static_cast(capacitySchedulerPolicy); + } + + bool incomingScheduled = false; + for (auto const& req : scheduledRequests) + { + incomingScheduled |= (req->mRequestId == 1u); + } + + if (reuseRetiringAdapter) + { + EXPECT_TRUE(incomingScheduled) + << "a request reusing the resident adapter needs no new pages and must still be " + "admitted; policy " + << static_cast(capacitySchedulerPolicy); + } + else + { + EXPECT_FALSE(incomingScheduled) + << "a second adapter cannot fit beside the retiring request's resident one; " + "admitting it would fail in ensureBatch. policy " + << static_cast(capacitySchedulerPolicy); + } + } + } +} + TEST_F(CapacitySchedulerTest, SimpleShouldFitInChunk) { SizeType32 kvCacheMaxNumTokens = 200; From 154f606ffdd05df2f035ec50af23fb051a3286fe Mon Sep 17 00:00:00 2001 From: Chenfei Zhang Date: Wed, 2 Sep 2026 09:15:31 -0700 Subject: [PATCH 05/22] [https://nvbugs/6627795][fix] Apply the retiring-request fix to the pure-Python scheduler There are three scheduler paths in _util.py:3145-3193: KVCacheV2Scheduler, SimpleUnifiedScheduler (when scheduler_config.use_python_scheduler), and SimpleScheduler. The earlier commits in this PR fixed the first and third; this commit fixes the second, which had the identical defect. PyCapacityScheduler defaulted no_schedule_after_state to GENERATION_COMPLETE, so its window spanned GENERATION_TO_COMPLETE. A request in that state has already produced its final token and is only awaiting the teardown the overlap scheduler defers by one iteration -- PyMicroBatchScheduler will never forward it. But the `len(scheduled_requests) >= max_num_requests` break in each policy sits after the state gate, so every such request consumed a capacity slot and shortened the real forward batch by one. Because it is a break and not a continue, later admissible requests were never examined at all. Dropping them from the window reserves nothing that was needed: a retiring request's get_remaining_blocks_to_completion is ~0, since it generates no further tokens. Its LoRA adapter, however, is still resident -- LoraCache::markTaskDone runs in the deferred teardown and claimPagesWithEvict only evicts from mDoneTasks. Narrowing the window alone would therefore drop the adapter's page charge, letting a pending request with a *different* adapter pass `neededPeftPages <= availablePeftPages` and then die in ensureBatch with LoraCacheFullException. So the second half of the change pre-claims PEFT pages for retiring requests before the main loop, mirroring preClaimPeftPagesForRetiringRequests on the C++ side and pre_claim_peft in the V2 scheduler. The claim is idempotent by dedupe: seen_task_ids means an adapter charged up front is a no-op if the main loop reaches it again, so the total is unchanged from before this PR. test_generation_to_complete_scheduled asserted the buggy behavior (that such a request is scheduled), so it is inverted and renamed to test_generation_to_complete_filtered. Two new tests cover the capacity slot and the PEFT charge, the latter parametrized over all three policies and over whether the incoming request reuses the retiring adapter. Verified by execution, with a negative control: against the unpatched scheduler.py exactly the 8 new/inverted tests fail (8 failed, 144 deselected, including `assert 0 not in [0]`); against the patched file they pass (8 passed); the whole file is green (152 passed). Caveat: no bindings exist for this PR's base, so the run used a self-consistent tree at 3253b640 with these two files rebased onto it (git apply clean, line offsets only). The C++ test added earlier in this PR is type-checked only; CI will be the first to build it. Signed-off-by: Chenfei Zhang --- .../_torch/pyexecutor/scheduler/scheduler.py | 62 ++++++++++++- .../_torch/executor/test_py_scheduler.py | 87 +++++++++++++++++-- 2 files changed, 139 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py index a0963245c6b3..dbcedfb673be 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py @@ -1471,9 +1471,14 @@ def schedule( reserved_cross_blocks = NoEvictScheduledBlocksManager(scheduler.cross_kv_cache_manager) # PEFT state - only used when has_peft - claimed_peft_pages = 0 available_peft_pages = scheduler._get_max_peft_pages() if has_peft else 0 uniq_task_ids: set[int] = set() if has_peft else None + # Retiring requests are outside the state window but still hold their + # adapters on device; charge them before the loop so the budget below is + # honest. See _pre_claim_peft_pages_for_retiring_requests. + claimed_peft_pages = scheduler._pre_claim_peft_pages_for_retiring_requests( + active_requests, uniq_task_ids + ) pending_requests: RequestList = [] pending_dis_gen_init_requests: RequestList = [] @@ -1639,8 +1644,12 @@ def schedule( scheduler.cross_kv_cache_manager, scheduler.two_step_lookahead ) - num_scheduled_peft_pages = 0 seen_task_ids: set[int] = set() + # Same reasoning as GuaranteedNoEvictPolicy: retiring requests are skipped + # by the state gate but their adapters are still resident. + num_scheduled_peft_pages = scheduler._pre_claim_peft_pages_for_retiring_requests( + active_requests, seen_task_ids + ) newly_contributed_context_blocks, _ = scheduler._prefill_contributed_blocks(active_requests) # Summary cache populated lazily by _beneficial_to_skip; consumed by @@ -1939,7 +1948,18 @@ def __init__( cross_kv_cache_manager: object | None = None, two_step_lookahead: bool = False, no_schedule_until_state: LlmRequestState = LlmRequestState.CONTEXT_INIT, - no_schedule_after_state: LlmRequestState = LlmRequestState.GENERATION_COMPLETE, + # Ends one state early, at GENERATION_TO_COMPLETE, matching + # BindCapacityScheduler and PyMicroBatchScheduler. That state means "final + # token produced, teardown deferred by the overlap scheduler": the + # micro-batch scheduler will never forward such a request, yet the + # ``len(scheduled_requests) >= scheduler.max_num_requests`` break in each + # policy sits after the state gate, so every one of them used to consume a + # capacity slot and shorten the real forward batch by one (nvbug-6627795). + # Keeping them in the window reserved nothing either: their + # get_remaining_blocks_to_completion is ~0, since they generate no further + # tokens. Their LoRA adapters *are* still resident, which is why the + # policies pre-claim PEFT pages for them separately. + no_schedule_after_state: LlmRequestState = LlmRequestState.GENERATION_TO_COMPLETE, enable_prefix_aware_scheduling: bool = True, ) -> None: """ @@ -1968,6 +1988,7 @@ def __init__( # Cache state values to avoid repeated .value access (optimization) self._no_schedule_until_state_value = no_schedule_until_state.value self._no_schedule_after_state_value = no_schedule_after_state.value + self._gen_to_complete_state_value = LlmRequestState.GENERATION_TO_COMPLETE.value # Initialize the appropriate policy self._policy = self._create_policy() @@ -2166,6 +2187,41 @@ def _get_peft_task_info( required_pages = self._get_peft_pages_for_request(req) if is_new_task else 0 return lora_task_id, is_new_task, required_pages + def _pre_claim_peft_pages_for_retiring_requests( + self, active_requests: RequestList, seen_task_ids: Optional[set[int]] + ) -> int: + """Charge PEFT pages for requests the state window excludes but whose adapters are resident. + + ``GENERATION_TO_COMPLETE`` requests sit outside ``no_schedule_after_state``, + so the policy loops skip them before reaching their PEFT claim. Their LoRA + adapters are still on device, though: the adapter is released by the + teardown the overlap scheduler has deferred, and the cache can only evict + tasks already marked done. Leaving them uncharged overstates the free + budget, so a pending request carrying a *different* adapter is admitted + against space that only looks free and then fails in ``ensure_batch``. + + Charging without consuming a capacity slot or a token budget is the point: + a retiring request must count against adapter *residency* while staying out + of the forward batch. Mirrors ``preClaimPeftPagesForRetiringRequests`` in + capacityScheduler.cpp and ``KVCacheV2Scheduler``'s ``pre_claim_peft``. + + Idempotent with respect to the policy loops: ``_get_peft_task_info`` dedupes + on ``seen_task_ids``, so a window that still admits + ``GENERATION_TO_COMPLETE`` charges the same total, and a pending request + reusing a retiring request's adapter is still charged zero new pages. + """ + if self.peft_cache_manager is None or seen_task_ids is None: + return 0 + claimed_pages = 0 + for req in active_requests: + if req.state_value != self._gen_to_complete_state_value: + continue + lora_task_id, is_new_task, peft_pages = self._get_peft_task_info(req, seen_task_ids) + if is_new_task: + claimed_pages += peft_pages + seen_task_ids.add(lora_task_id) + return claimed_pages + def _can_be_scheduled_with_disagg_exception(self, req: LlmRequest) -> bool: """ Check if request can be scheduled, with exception for disagg generation init state. diff --git a/tests/unittest/_torch/executor/test_py_scheduler.py b/tests/unittest/_torch/executor/test_py_scheduler.py index d323ee9d5898..04eed0b20fd7 100644 --- a/tests/unittest/_torch/executor/test_py_scheduler.py +++ b/tests/unittest/_torch/executor/test_py_scheduler.py @@ -2236,12 +2236,17 @@ def test_generation_complete_filtered(self): fitting, disagg, paused = scheduler.schedule_request([make_completed_request(0)]) assert len(fitting) == 0 - def test_generation_to_complete_scheduled(self): - """GENERATION_TO_COMPLETE is schedulable in PyCapacityScheduler. - PyCapacityScheduler uses no_schedule_after=GENERATION_COMPLETE (20), - so GENERATION_TO_COMPLETE (14) passes state gating. The real C++ binding's - is_generation_in_progress_state includes GENERATION_TO_COMPLETE, so the - MaxRequestsPolicy schedules it.""" + def test_generation_to_complete_filtered(self): + """GENERATION_TO_COMPLETE is at no_schedule_after, filtered out (nvbug-6627795). + + The state means "final token produced, teardown deferred by the overlap + scheduler". PyMicroBatchScheduler has always excluded it, so admitting it + here only burned a capacity slot -- the + ``len(scheduled_requests) >= max_num_requests`` break sits after the state + gate -- and shortened the real forward batch by one. Note the request is + still ``is_generation_in_progress_state`` (that predicate spans 13 and 14), + so only the state window keeps it out. + """ scheduler = PyCapacityScheduler( max_num_requests=4, kv_cache_manager=None, @@ -2250,8 +2255,27 @@ def test_generation_to_complete_scheduled(self): request_id=0, state=LlmRequestState.GENERATION_TO_COMPLETE, ) + assert req.is_generation_in_progress_state fitting, disagg, paused = scheduler.schedule_request([req]) - assert len(fitting) == 1 + assert len(fitting) == 0 + + def test_retiring_request_does_not_consume_capacity(self): + """A retiring request must not displace a schedulable one. + + This is the nvbug-6627795 mechanism in miniature: with capacity 1, the + state-14 request used to be admitted first and the `break` then shut the + loop before the real generation request was ever considered. + """ + kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=5) + scheduler = PyCapacityScheduler( + max_num_requests=1, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, + ) + retiring = _make_request(0, state=LlmRequestState.GENERATION_TO_COMPLETE) + runnable = _make_request(1, state=LlmRequestState.GENERATION_IN_PROGRESS) + fitting, _disagg, _paused = scheduler.schedule_request([retiring, runnable]) + assert [r.request_id for r in fitting] == [1] # ############################################################################ @@ -2326,6 +2350,55 @@ def test_max_utilization_peft_page_accumulation(self): # 2 tasks x 10 pages = 20 <= 25; 3rd task would push to 30 > 25 assert len(fitting) == 2 + @pytest.mark.parametrize( + "scheduler_policy", + [ + CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, + CapacitySchedulerPolicy.MAX_UTILIZATION, + CapacitySchedulerPolicy.STATIC_BATCH, + ], + ) + @pytest.mark.parametrize("reuse_retiring_adapter", [False, True]) + def test_retiring_request_still_charges_its_lora_adapter( + self, scheduler_policy, reuse_retiring_adapter + ): + """nvbug-6627795: a retiring request leaves the window but keeps its adapter. + + GENERATION_TO_COMPLETE is now outside no_schedule_after_state, so the policy + loops skip it before reaching their PEFT claim. The adapter is still resident + -- it is released by the teardown the overlap scheduler deferred, and the + cache can only evict tasks already marked done -- so it must still be + charged, or a request carrying a different adapter is admitted against a + budget that only looks free and then dies in ensure_batch. + + Both directions are asserted: under-charging is the bug, but over-charging + would be a new one, so the same-adapter case must still be admitted free. + C++ ref: CapacitySchedulerTest.RetiringRequestStillChargesItsLoraAdapter. + """ + # 10 pages per distinct adapter against 15 on device: one fits, two never can. + kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=5) + peft = MockPeftCacheManager(max_pages=15, pages_per_request=10) + scheduler = PyCapacityScheduler( + max_num_requests=4, + kv_cache_manager=kv, + peft_cache_manager=peft, + scheduler_policy=scheduler_policy, + ) + retiring = _make_request(0, state=LlmRequestState.GENERATION_TO_COMPLETE, lora_task_id=1) + incoming = _make_request(1, lora_task_id=1 if reuse_retiring_adapter else 2) + + fitting, _disagg, _paused = scheduler.schedule_request([retiring, incoming]) + + scheduled_ids = [r.request_id for r in fitting] + # Outside the window either way. + assert 0 not in scheduled_ids + if reuse_retiring_adapter: + # Needs no new pages, so the pre-claim must not lock it out. + assert scheduled_ids == [1] + else: + # 10 already charged + 10 needed > 15 available. + assert scheduled_ids == [] + # ############################################################################ # From 9ef2f5eab55ad20a20fc765a64350c960ccd3a16 Mon Sep 17 00:00:00 2001 From: Chenfei Zhang Date: Wed, 2 Sep 2026 21:47:56 -0700 Subject: [PATCH 06/22] [https://nvbugs/6627795][fix] revert the V1 capacity-scheduler changes, keep the ADP fix This PR carried two independent layers. The one that the measurements are about is ADP admission: adp_router.py stops reporting requests in GENERATION_TO_COMPLETE as active, so the router neither balances load on them nor charges them against the global pop budget. That layer is scheduler-agnostic -- it contains no reference to use_python_scheduler, KVCacheV2Scheduler, BindCapacityScheduler, PyCapacityScheduler or capacityScheduler -- and it is what the ctx_only table in the PR body measured. The second layer narrowed the *capacity* scheduler's window by one state on the two V1 paths, plus the PEFT pre-claim needed to keep that narrowing LoRA-safe. This commit removes it, restoring the four files to their pre-PR contents byte for byte: cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp cpp/tests/unit_tests/batch_manager/capacitySchedulerTest.cpp tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py tests/unittest/_torch/executor/test_py_scheduler.py Reasons, in order of weight: 1. It is a separate concern from ADP admission and it has no perf evidence in this PR. All 15 ctx_only cases in the table resolved use_kv_cache_manager_v2="auto" to V2 (confirmed in the engine logs), and scheduler_v2.py already had both halves before this PR, so the table validates the ADP layer only. The V1 capacity change had unit tests and nothing else. 2. The C++ half was never built here. Its new unit test is type-checked only, and the LoRA hazard it guards against is unmeasured. 3. Dropping it costs correctness nothing. The V1 defect is a throughput defect: a retiring request occupies a max_num_requests slot that the micro-batch scheduler then discards, so a real request that could have run waits one iteration. It does not accumulate -- the deferred teardown frees the slot next iteration unconditionally -- and GUARANTEED_NO_EVICT still prevents KV over-commit because state 14 keeps its reservation exactly as before. With the ADP layer admitting more while the V1 window stays wide, the effect is a smaller forward batch, never a crash. 4. The window narrowing and the PEFT pre-claim are a matched pair: the pre-claim exists only to replace the page charge the narrowing would otherwise drop for an adapter that LoraCache::markTaskDone has not yet released. Reverting them together is internally consistent; keeping either one alone would be dead code or a LoraCacheFullException. The residual cost is that V1 paths get only the ADP part of the fix. V1 is still the default for models whose class does not declare a preference, and disagg demotes a V2-preferring model to V1 whenever the transceiver is not NIXL + Python. Severity scales as retiring-requests-per-iteration over max_batch_size, so it is negligible for a long-output decode batch and material only for a ctx worker with a very small max_batch_size. Left as a follow-up, to be filed with its own measurement rather than shipped here on unit tests alone. scheduler_v2.py is untouched by this PR and keeps its own window default and pre_claim_peft, so V2 behaviour is unchanged throughout. Signed-off-by: Chenfei Zhang --- .../batch_manager/capacityScheduler.cpp | 33 ------- .../batch_manager/capacitySchedulerTest.cpp | 86 ------------------ .../_torch/pyexecutor/scheduler/scheduler.py | 73 +--------------- .../_torch/executor/test_py_scheduler.py | 87 ++----------------- 4 files changed, 11 insertions(+), 268 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp b/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp index 7b507688957e..335e30403d0f 100644 --- a/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp +++ b/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp @@ -141,37 +141,6 @@ void claimPeftPagesForRequest(std::shared_ptr const& req, } } -//! @brief Charge PEFT pages for requests the state window excludes but whose adapters are still resident. -//! -//! kGENERATION_TO_COMPLETE means "final token produced, teardown deferred by the overlap scheduler". -//! Such a request is never forwarded again, so the admission loop skips it on the -//! getNoScheduleAfterState() gate -- but its LoRA adapter is still on device, because -//! LoraCache::markTaskDone runs as part of the teardown that has not happened yet and -//! LoraCache::claimPagesWithEvict only ever evicts tasks in mDoneTasks. Leaving it uncharged -//! overstates availablePeftPages, so a pending request carrying a *different* adapter can be admitted -//! and then fail in PeftCacheManager::ensureBatch with LoraCacheFullException ("Cache is full. There -//! are no done tasks to evict") -- turning the nvbug-6627795 admission fix into a crash under LoRA. -//! -//! Charging without consuming a request slot or a token budget is the point: a retiring request must -//! count against adapter *residency* while staying out of the forward batch. This mirrors -//! KVCacheV2Scheduler's pre_claim_peft, which the Python V2 path already needed for the same reason. -//! -//! Idempotent with respect to the admission loop: claimPeftPagesForRequest dedupes on uniqTaskIds, so -//! a window that still admits kGENERATION_TO_COMPLETE charges the same total, and a pending request -//! reusing a retiring request's adapter is still correctly charged zero new pages. -void preClaimPeftPagesForRetiringRequests(RequestList const& activeRequests, - OptionalRef peftCacheManager, SizeType32& claimedPeftPages, - std::unordered_set& uniqTaskIds) -{ - for (auto const& req : activeRequests) - { - if (req->isGenerationToCompleteState()) - { - claimPeftPagesForRequest(req, peftCacheManager, claimedPeftPages, uniqTaskIds); - } - } -} - } // namespace MaxRequestsScheduler::MaxRequestsScheduler( @@ -284,7 +253,6 @@ std::tuple GuaranteedNoEvictScheduler::impl( : std::nullopt; SizeType32 claimedPeftPages{0}; std::unordered_set uniqTaskIds{}; - preClaimPeftPagesForRetiringRequests(activeRequests, peftCacheManager, claimedPeftPages, uniqTaskIds); std::size_t numAdmittedRequests{0}; RequestVector pendingRequests; RequestVector pendingDisGenInitRequests; @@ -503,7 +471,6 @@ std::tuple MaxUtilizationScheduler::operator()( } SizeType32 numScheduledPeftPages{0}; std::unordered_set seenTaskIds; - preClaimPeftPagesForRetiringRequests(activeRequests, peftCacheManager, numScheduledPeftPages, seenTaskIds); // Keep track of blocks contributed by requests in context phase std::unordered_set newlyContributedContextBlocks; diff --git a/cpp/tests/unit_tests/batch_manager/capacitySchedulerTest.cpp b/cpp/tests/unit_tests/batch_manager/capacitySchedulerTest.cpp index 7e7ea9e45e6f..2c9757b847b2 100644 --- a/cpp/tests/unit_tests/batch_manager/capacitySchedulerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/capacitySchedulerTest.cpp @@ -659,92 +659,6 @@ TEST_F(CapacitySchedulerTest, SimpleLoraDoesntFitDuplicateTask) } } -// nvbug-6627795: the capacity window now ends at kGENERATION_TO_COMPLETE so a retiring request stops -// consuming a slot in mMaxNumRequests. Its LoRA adapter, however, is still resident on device -- -// LoraCache::markTaskDone runs as part of the teardown the overlap scheduler has deferred, and -// claimPagesWithEvict only evicts tasks in mDoneTasks. So the adapter must still be charged, or a -// request carrying a different adapter is admitted against a budget that only looks free and then dies -// in ensureBatch with LoraCacheFullException. -// -// Both directions are asserted. Under-charging is the bug; over-charging would be a new one, so the -// same-adapter case must still be admitted for free. -TEST_F(CapacitySchedulerTest, RetiringRequestStillChargesItsLoraAdapter) -{ - SizeType32 constexpr kvCacheTokensPerBlock = 10; - // Deliberately generous: PEFT pages must be the only binding constraint, otherwise a KV shortfall - // could make this test pass for the wrong reason. - SizeType32 constexpr kvCacheMaxNumTokens = 2000; - SizeType32 constexpr kvCacheMaxNumTokensPerSeq = 90; - SizeType32 constexpr maxNumRequests = 8; - // 20 pages per distinct adapter against 30 on device: one adapter fits, two never can. - SizeType32 constexpr pagesPerTask = 20; - SizeType32 constexpr maxDevicePages = 30; - uint64_t constexpr retiringTaskId = 1234; - uint64_t constexpr otherTaskId = 5678; - - auto const capacitySchedulerPolicies - = std::vector{CapacitySchedulerPolicy::kGUARANTEED_NO_EVICT, - CapacitySchedulerPolicy::kMAX_UTILIZATION, CapacitySchedulerPolicy::kSTATIC_BATCH}; - - for (auto capacitySchedulerPolicy : capacitySchedulerPolicies) - { - // true -> the incoming request reuses the retiring adapter, so it needs no new pages. - // false -> it brings a second adapter, which cannot fit alongside the resident one. - for (bool reuseRetiringAdapter : {false, true}) - { - auto kvCacheManager = getKvCacheManager( - maxNumRequests, kvCacheTokensPerBlock, kvCacheMaxNumTokens, kvCacheMaxNumTokensPerSeq); - std::shared_ptr peftCacheManager - = std::make_shared(pagesPerTask, maxDevicePages, maxDevicePages); - // The window this fix installs: kGENERATION_TO_COMPLETE is outside it. - auto capacityScheduler = CapacityScheduler(maxNumRequests, capacitySchedulerPolicy, - static_cast(kvCacheManager), /*twoStepsLookAhead=*/false, LlmRequestState::kCONTEXT_INIT, - LlmRequestState::kGENERATION_TO_COMPLETE); - - int32_t constexpr maxNewTokens = 50; - int32_t constexpr promptLen = 10; - - RequestList activeRequests; - // Final token produced, teardown deferred: not schedulable, but its adapter is still loaded. - activeRequests.push_back(createRequest(promptLen, maxNewTokens, 0, retiringTaskId, - tensorrt_llm::executor::Request::kDefaultPriority, LlmRequestState::kGENERATION_TO_COMPLETE)); - activeRequests.push_back( - createRequest(promptLen, maxNewTokens, 1, reuseRetiringAdapter ? retiringTaskId : otherTaskId)); - - auto [scheduledRequests, scheduledDisaggGenInitRequests, pausedRequests] - = capacityScheduler(activeRequests, kvCacheManager, peftCacheManager); - - // The retiring request is outside the window, so it must never be scheduled either way. - for (auto const& req : scheduledRequests) - { - EXPECT_NE(req->mRequestId, 0u) - << "retiring request was scheduled, policy " << static_cast(capacitySchedulerPolicy); - } - - bool incomingScheduled = false; - for (auto const& req : scheduledRequests) - { - incomingScheduled |= (req->mRequestId == 1u); - } - - if (reuseRetiringAdapter) - { - EXPECT_TRUE(incomingScheduled) - << "a request reusing the resident adapter needs no new pages and must still be " - "admitted; policy " - << static_cast(capacitySchedulerPolicy); - } - else - { - EXPECT_FALSE(incomingScheduled) - << "a second adapter cannot fit beside the retiring request's resident one; " - "admitting it would fail in ensureBatch. policy " - << static_cast(capacitySchedulerPolicy); - } - } - } -} - TEST_F(CapacitySchedulerTest, SimpleShouldFitInChunk) { SizeType32 kvCacheMaxNumTokens = 200; diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py index dbcedfb673be..99a1710936f4 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py @@ -452,16 +452,7 @@ def __init__( has_kv_cache_manager=kv_cache_manager is not None, two_step_lookahead=two_step_lookahead, no_schedule_until_state=no_schedule_until_state, - # Stop the window one state early, at GENERATION_TO_COMPLETE. That - # state means "final token produced, teardown deferred by the overlap - # scheduler": no micro-batch scheduler will ever forward such a - # request again, yet the C++ capacity loop's - # `numAdmittedRequests >= mMaxNumRequests` break sits after the state - # gate and before classification, so each one consumed a slot in - # mMaxNumRequests and shortened the real forward batch by one - # (nvbug-6627795). Their KV cache is released by the teardown that is - # already queued, so keeping them inside the window bought nothing. - no_schedule_after_state=LlmRequestState.GENERATION_TO_COMPLETE, + no_schedule_after_state=LlmRequestState.GENERATION_COMPLETE, enable_prefix_aware_scheduling=enable_prefix_aware_scheduling, ) @@ -1471,14 +1462,9 @@ def schedule( reserved_cross_blocks = NoEvictScheduledBlocksManager(scheduler.cross_kv_cache_manager) # PEFT state - only used when has_peft + claimed_peft_pages = 0 available_peft_pages = scheduler._get_max_peft_pages() if has_peft else 0 uniq_task_ids: set[int] = set() if has_peft else None - # Retiring requests are outside the state window but still hold their - # adapters on device; charge them before the loop so the budget below is - # honest. See _pre_claim_peft_pages_for_retiring_requests. - claimed_peft_pages = scheduler._pre_claim_peft_pages_for_retiring_requests( - active_requests, uniq_task_ids - ) pending_requests: RequestList = [] pending_dis_gen_init_requests: RequestList = [] @@ -1644,12 +1630,8 @@ def schedule( scheduler.cross_kv_cache_manager, scheduler.two_step_lookahead ) + num_scheduled_peft_pages = 0 seen_task_ids: set[int] = set() - # Same reasoning as GuaranteedNoEvictPolicy: retiring requests are skipped - # by the state gate but their adapters are still resident. - num_scheduled_peft_pages = scheduler._pre_claim_peft_pages_for_retiring_requests( - active_requests, seen_task_ids - ) newly_contributed_context_blocks, _ = scheduler._prefill_contributed_blocks(active_requests) # Summary cache populated lazily by _beneficial_to_skip; consumed by @@ -1948,18 +1930,7 @@ def __init__( cross_kv_cache_manager: object | None = None, two_step_lookahead: bool = False, no_schedule_until_state: LlmRequestState = LlmRequestState.CONTEXT_INIT, - # Ends one state early, at GENERATION_TO_COMPLETE, matching - # BindCapacityScheduler and PyMicroBatchScheduler. That state means "final - # token produced, teardown deferred by the overlap scheduler": the - # micro-batch scheduler will never forward such a request, yet the - # ``len(scheduled_requests) >= scheduler.max_num_requests`` break in each - # policy sits after the state gate, so every one of them used to consume a - # capacity slot and shorten the real forward batch by one (nvbug-6627795). - # Keeping them in the window reserved nothing either: their - # get_remaining_blocks_to_completion is ~0, since they generate no further - # tokens. Their LoRA adapters *are* still resident, which is why the - # policies pre-claim PEFT pages for them separately. - no_schedule_after_state: LlmRequestState = LlmRequestState.GENERATION_TO_COMPLETE, + no_schedule_after_state: LlmRequestState = LlmRequestState.GENERATION_COMPLETE, enable_prefix_aware_scheduling: bool = True, ) -> None: """ @@ -1988,7 +1959,6 @@ def __init__( # Cache state values to avoid repeated .value access (optimization) self._no_schedule_until_state_value = no_schedule_until_state.value self._no_schedule_after_state_value = no_schedule_after_state.value - self._gen_to_complete_state_value = LlmRequestState.GENERATION_TO_COMPLETE.value # Initialize the appropriate policy self._policy = self._create_policy() @@ -2187,41 +2157,6 @@ def _get_peft_task_info( required_pages = self._get_peft_pages_for_request(req) if is_new_task else 0 return lora_task_id, is_new_task, required_pages - def _pre_claim_peft_pages_for_retiring_requests( - self, active_requests: RequestList, seen_task_ids: Optional[set[int]] - ) -> int: - """Charge PEFT pages for requests the state window excludes but whose adapters are resident. - - ``GENERATION_TO_COMPLETE`` requests sit outside ``no_schedule_after_state``, - so the policy loops skip them before reaching their PEFT claim. Their LoRA - adapters are still on device, though: the adapter is released by the - teardown the overlap scheduler has deferred, and the cache can only evict - tasks already marked done. Leaving them uncharged overstates the free - budget, so a pending request carrying a *different* adapter is admitted - against space that only looks free and then fails in ``ensure_batch``. - - Charging without consuming a capacity slot or a token budget is the point: - a retiring request must count against adapter *residency* while staying out - of the forward batch. Mirrors ``preClaimPeftPagesForRetiringRequests`` in - capacityScheduler.cpp and ``KVCacheV2Scheduler``'s ``pre_claim_peft``. - - Idempotent with respect to the policy loops: ``_get_peft_task_info`` dedupes - on ``seen_task_ids``, so a window that still admits - ``GENERATION_TO_COMPLETE`` charges the same total, and a pending request - reusing a retiring request's adapter is still charged zero new pages. - """ - if self.peft_cache_manager is None or seen_task_ids is None: - return 0 - claimed_pages = 0 - for req in active_requests: - if req.state_value != self._gen_to_complete_state_value: - continue - lora_task_id, is_new_task, peft_pages = self._get_peft_task_info(req, seen_task_ids) - if is_new_task: - claimed_pages += peft_pages - seen_task_ids.add(lora_task_id) - return claimed_pages - def _can_be_scheduled_with_disagg_exception(self, req: LlmRequest) -> bool: """ Check if request can be scheduled, with exception for disagg generation init state. diff --git a/tests/unittest/_torch/executor/test_py_scheduler.py b/tests/unittest/_torch/executor/test_py_scheduler.py index 04eed0b20fd7..d323ee9d5898 100644 --- a/tests/unittest/_torch/executor/test_py_scheduler.py +++ b/tests/unittest/_torch/executor/test_py_scheduler.py @@ -2236,17 +2236,12 @@ def test_generation_complete_filtered(self): fitting, disagg, paused = scheduler.schedule_request([make_completed_request(0)]) assert len(fitting) == 0 - def test_generation_to_complete_filtered(self): - """GENERATION_TO_COMPLETE is at no_schedule_after, filtered out (nvbug-6627795). - - The state means "final token produced, teardown deferred by the overlap - scheduler". PyMicroBatchScheduler has always excluded it, so admitting it - here only burned a capacity slot -- the - ``len(scheduled_requests) >= max_num_requests`` break sits after the state - gate -- and shortened the real forward batch by one. Note the request is - still ``is_generation_in_progress_state`` (that predicate spans 13 and 14), - so only the state window keeps it out. - """ + def test_generation_to_complete_scheduled(self): + """GENERATION_TO_COMPLETE is schedulable in PyCapacityScheduler. + PyCapacityScheduler uses no_schedule_after=GENERATION_COMPLETE (20), + so GENERATION_TO_COMPLETE (14) passes state gating. The real C++ binding's + is_generation_in_progress_state includes GENERATION_TO_COMPLETE, so the + MaxRequestsPolicy schedules it.""" scheduler = PyCapacityScheduler( max_num_requests=4, kv_cache_manager=None, @@ -2255,27 +2250,8 @@ def test_generation_to_complete_filtered(self): request_id=0, state=LlmRequestState.GENERATION_TO_COMPLETE, ) - assert req.is_generation_in_progress_state fitting, disagg, paused = scheduler.schedule_request([req]) - assert len(fitting) == 0 - - def test_retiring_request_does_not_consume_capacity(self): - """A retiring request must not displace a schedulable one. - - This is the nvbug-6627795 mechanism in miniature: with capacity 1, the - state-14 request used to be admitted first and the `break` then shut the - loop before the real generation request was ever considered. - """ - kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=5) - scheduler = PyCapacityScheduler( - max_num_requests=1, - kv_cache_manager=kv, - scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, - ) - retiring = _make_request(0, state=LlmRequestState.GENERATION_TO_COMPLETE) - runnable = _make_request(1, state=LlmRequestState.GENERATION_IN_PROGRESS) - fitting, _disagg, _paused = scheduler.schedule_request([retiring, runnable]) - assert [r.request_id for r in fitting] == [1] + assert len(fitting) == 1 # ############################################################################ @@ -2350,55 +2326,6 @@ def test_max_utilization_peft_page_accumulation(self): # 2 tasks x 10 pages = 20 <= 25; 3rd task would push to 30 > 25 assert len(fitting) == 2 - @pytest.mark.parametrize( - "scheduler_policy", - [ - CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, - CapacitySchedulerPolicy.MAX_UTILIZATION, - CapacitySchedulerPolicy.STATIC_BATCH, - ], - ) - @pytest.mark.parametrize("reuse_retiring_adapter", [False, True]) - def test_retiring_request_still_charges_its_lora_adapter( - self, scheduler_policy, reuse_retiring_adapter - ): - """nvbug-6627795: a retiring request leaves the window but keeps its adapter. - - GENERATION_TO_COMPLETE is now outside no_schedule_after_state, so the policy - loops skip it before reaching their PEFT claim. The adapter is still resident - -- it is released by the teardown the overlap scheduler deferred, and the - cache can only evict tasks already marked done -- so it must still be - charged, or a request carrying a different adapter is admitted against a - budget that only looks free and then dies in ensure_batch. - - Both directions are asserted: under-charging is the bug, but over-charging - would be a new one, so the same-adapter case must still be admitted free. - C++ ref: CapacitySchedulerTest.RetiringRequestStillChargesItsLoraAdapter. - """ - # 10 pages per distinct adapter against 15 on device: one fits, two never can. - kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=5) - peft = MockPeftCacheManager(max_pages=15, pages_per_request=10) - scheduler = PyCapacityScheduler( - max_num_requests=4, - kv_cache_manager=kv, - peft_cache_manager=peft, - scheduler_policy=scheduler_policy, - ) - retiring = _make_request(0, state=LlmRequestState.GENERATION_TO_COMPLETE, lora_task_id=1) - incoming = _make_request(1, lora_task_id=1 if reuse_retiring_adapter else 2) - - fitting, _disagg, _paused = scheduler.schedule_request([retiring, incoming]) - - scheduled_ids = [r.request_id for r in fitting] - # Outside the window either way. - assert 0 not in scheduled_ids - if reuse_retiring_adapter: - # Needs no new pages, so the pre-claim must not lock it out. - assert scheduled_ids == [1] - else: - # 10 already charged + 10 needed > 15 available. - assert scheduled_ids == [] - # ############################################################################ # From eed14605aea1ff20f264ce5965435e4e5c1a73e0 Mon Sep 17 00:00:00 2001 From: Chenfei Zhang Date: Thu, 3 Sep 2026 20:19:22 -0700 Subject: [PATCH 07/22] [https://nvbugs/6627795][fix] size the remaining spec-dec identity pools by the sequence-slot pool Follow-up to 399890b7a1, which fixed SpecMetadata and MTPHiddenStatesManager but left three managers behind. Under the attention-DP overlap headroom compute_max_num_sequences() returns 2 * max_batch_size, so py_seq_slot ranges over the whole pool while these stores were still sized at max_batch_size: * DynamicTreeSlotStorage was sized from SpecTreeManager.num_trees. Its own docstring says it is indexed by py_seq_slot, and Eagle3OneModelDynamicTreeResourceManager.free_resources() writes mark_invalid(request.py_seq_slot) into it. * Eagle3ResourceManager.slot_manager was sized max_seq_len + 1 -- a token count standing in for a slot count. Accidentally generous for most configurations, but exactly one slot short when max_batch_size - 1 <= max_seq_len < 2 * max_batch_size - 1 (e.g. max_batch_size: 1024, max_seq_len: 1024 gives 1025 slots where 2048 are needed). max_seq_len is kept as a floor so no existing deployment shrinks. * The SuffixAutomatonManager slot pool holds a slot for a request id's whole lifetime, so max_batch_size + 1 concurrent ids hit RuntimeError("No free or retained slots available") mid-run. Also forwards num_seq_slots at mtp_dynamic_tree.py:1103, eleven lines above a pool 399890b7a1 already fixed, and at every other manager construction in get_spec_resource_manager. The two-family distinction is preserved: only identity-keyed stores are widened. Buffers keyed by batch position (SpecTreeManager.num_trees and its per-forward work buffers, batch_indices_cuda) deliberately stay at max_batch_size, because the micro-batch scheduler caps every forward there -- its no_schedule_after_state=GENERATION_TO_COMPLETE default (scheduler.py:492, :889) is untouched by this series and keeps retiring requests out of the batch entirely. Widening them would cost memory quadratically in the tree dimensions for nothing. num_seq_slots defaults to None everywhere so a caller that does not know the pool (pipeline parallelism, no attention DP, overlap disabled) keeps the established sizing without every call site restating it; the receiving sites use max(num_seq_slots or 0, max_num_requests) so a smaller value can never shrink a buffer that must also cover batch positions. SA pool sizing splits on intent rather than growing silently: a default pool grows to cover the slot pool, while an explicit global_pool_size is a user memory contract, so it is honoured and validated instead. That turns a mid-run slot exhaustion into a startup ValueError naming the real bound, and only newly rejects configurations that were already broken (explicit pool == max_batch_size with the ADP headroom on). Rather than a fourth patch of the same shape, the plumbing is now guarded mechanically. test_every_resource_manager_branch_forwards_the_slot_pool parses get_spec_resource_manager's AST and fails if any *Manager( construction omits num_seq_slots, with an allow-list that requires a written justification per entry; NGramPoolManager and SaveHiddenStatesResourceManager are exempt because their pools are not keyed by request identity and neither mode appears in SpeculativeDecodingMode.support_overlap_scheduler(), so py_executor_creator forces the overlap scheduler off and the headroom can never apply. test_slot_pool_managers_accept_an_optional_pool_size checks the receiving end of the same contract, including that the default is None. The behavioural tests carry negative controls: mark_invalid(pool - 1) and the (max_batch_size + 1)-th SA allocation must raise without the pool, otherwise the positive assertion proves nothing. Testing: both structural guards were run locally against three tree states to show they are not vacuous. On the worktree neither reports anything; at the previous HEAD the AST guard names exactly ['Eagle3OneModelDynamicTreeResourceManager', 'Eagle3ResourceManager', 'SuffixAutomatonManager'] and the signature guard reports four of six classes missing the parameter; at the merge base the AST guard additionally names MTPEagleDynamicTreeResourceManager and MTPHiddenStatesManager, the two 399890b7a1 fixed. The SA and CUDA-gated tests in this file were NOT executed here: this host has no x86_64 build, and the only built tree available is 2546 lines of drift in tensorrt_llm/_torch/speculative/ away from the merge base, so running against it would have measured that drift rather than this change. They run in CI. Signed-off-by: Chenfei Zhang --- tensorrt_llm/_torch/speculative/eagle3.py | 23 +- .../_torch/speculative/mtp_dynamic_tree.py | 1 + .../_torch/speculative/spec_tree_manager.py | 23 +- .../_torch/speculative/suffix_automaton.py | 29 +- tensorrt_llm/_torch/speculative/utils.py | 36 ++- .../speculative/test_spec_slot_pool_sizing.py | 294 +++++++++++++++++- 6 files changed, 375 insertions(+), 31 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index d497862e0c8e..f9b23ec4c9c2 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -42,7 +42,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 @@ -50,9 +51,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 attention-DP 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 @@ -103,6 +112,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): @@ -164,8 +174,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, @@ -178,6 +194,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_dynamic_tree.py b/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py index f1b63e82a772..8843fc62849f 100644 --- a/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py +++ b/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py @@ -1107,6 +1107,7 @@ 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 diff --git a/tensorrt_llm/_torch/speculative/spec_tree_manager.py b/tensorrt_llm/_torch/speculative/spec_tree_manager.py index 545b5bc7bb06..d82c34cc7479 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 attention-DP 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 4eba0781d46f..ef86c23b3d5f 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,26 @@ 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: + # 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 attention-DP 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. An explicit global_pool_size is a user contract + # about memory, so it is honoured and validated rather than grown silently. + if sa_config.global_pool_size is not None: + self.pool_size = sa_config.global_pool_size + else: + self.pool_size = max(sa_config.effective_pool_size, self._num_seq_slots) + if self.pool_size < self._num_seq_slots: raise ValueError( - f"global_pool_size ({self.pool_size}) must be >= " - f"max_batch_size ({max_num_requests})" + f"global_pool_size ({self.pool_size}) must be >= the number of " + f"sequence slots ({self._num_seq_slots})" ) # Calculate per-state size based on max_seq_len @@ -159,7 +172,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 3b3dd4b5098e..407e991f82b7 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -593,8 +593,10 @@ 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) # Dynamic tree combines SpecTreeManager with MTP hidden-state slots. if getattr(spec_config, 'use_dynamic_tree', False): return MTPEagleDynamicTreeResourceManager( @@ -617,6 +619,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 @@ -624,8 +627,10 @@ 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, @@ -636,14 +641,16 @@ def get_spec_resource_manager(model_engine, draft_model_engine=None): ) 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, @@ -652,6 +659,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_eagle3() or spec_dec_mode.is_mtp_eagle(): assert draft_model_engine is not None, "Draft model engine is required for Eagle3 and MTP Eagle two model flow." @@ -662,6 +670,7 @@ def get_spec_resource_manager(model_engine, draft_model_engine=None): max_num_requests, max_seq_len, max_num_tokens, + num_seq_slots=num_seq_slots, ) if spec_dec_mode.is_save_hidden_states(): return SaveHiddenStatesResourceManager( @@ -674,13 +683,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/unittest/_torch/speculative/test_spec_slot_pool_sizing.py b/tests/unittest/_torch/speculative/test_spec_slot_pool_sizing.py index fc5de82ab01d..c4422df334f3 100644 --- a/tests/unittest/_torch/speculative/test_spec_slot_pool_sizing.py +++ b/tests/unittest/_torch/speculative/test_spec_slot_pool_sizing.py @@ -10,21 +10,37 @@ * 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) and ``MTPHiddenStatesManager``'s hidden-state pools. + 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 (``no_schedule_after_state=GENERATION_TO_COMPLETE`` also keeps - the retiring requests out of the batch entirely). + 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.utils import _build_spec_metadata, get_spec_metadata +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) @@ -146,6 +162,276 @@ def test_mtp_slot_pool_survives_a_full_overlap_turnover(): 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 + which does not know the pool (PP, no attention DP, overlap disabled) keeps the + established sizing without every call site having to restate 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_honoured_but_validated(): + """``global_pool_size`` is a memory contract, so it is never grown silently. + + It is validated against the sequence-slot pool instead, turning what would be + a mid-run slot exhaustion into a startup error that names the real bound. + """ + grown = _sa_manager(POOL, enable_global_pool=True, global_pool_size=64) + assert grown.pool_size == 64 + + with pytest.raises(ValueError, match="sequence slots"): + _sa_manager(POOL, enable_global_pool=True, global_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 From aa0c2a1f42abdc6066174cfb10677a9d17960458 Mon Sep 17 00:00:00 2001 From: Chenfei Zhang Date: Sun, 6 Sep 2026 22:54:50 -0700 Subject: [PATCH 08/22] [https://nvbugs/6627795][fix] accept num_seq_slots in the qwen4 dynamic-tree test stub get_spec_resource_manager now forwards num_seq_slots to MTPEagleDynamicTreeResourceManager, but this test monkeypatches that class with a stub whose signature predates the kwarg, so the call raised TypeError: make_manager() got an unexpected keyword argument 'num_seq_slots'. Mirror the real constructor signature in the stub, capture the new kwarg alongside the others, and assert it stays None for an engine that does not opt into the attention-DP overlap seq-slot headroom. Signed-off-by: Chenfei Zhang --- tests/unittest/_torch/modeling/test_qwen4_exp_support.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/unittest/_torch/modeling/test_qwen4_exp_support.py b/tests/unittest/_torch/modeling/test_qwen4_exp_support.py index 5da980b03c4c..cb98535d9942 100644 --- a/tests/unittest/_torch/modeling/test_qwen4_exp_support.py +++ b/tests/unittest/_torch/modeling/test_qwen4_exp_support.py @@ -1008,6 +1008,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, @@ -1015,6 +1016,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 @@ -1044,6 +1046,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 attention-DP 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: From 9f3c3cf8b83ac47f24c04be5f2c1e47daccbc81a Mon Sep 17 00:00:00 2001 From: Chenfei Zhang Date: Mon, 7 Sep 2026 08:54:54 -0700 Subject: [PATCH 09/22] [https://nvbugs/6627795][fix] unify the sequence-slot coefficient and plumb it into every seat-indexed pool The number of simultaneously-live sequences was re-derived from max_batch_size in seven places with different formulas, and every skew between any two of them is a bug. Two have already shipped, in opposite directions: * index pool smaller than the seat pool -- _create_kv_cache warns "No free IndexMapper slots", returns None and the scheduler defers the request. This is nvbug 6627795: widening the seat pool to 2 * max_batch_size for aggregated attention-DP with the overlap scheduler left KVCacheManagerV2's index mapper at max_batch_size + 1, so the admission the fix recovered was handed straight back, silently. * index pool larger than the seat pool -- a request is admitted that cannot be seated and SlotManager.add_slot raises on the executor's event-loop thread, killing the rank mid-collective (PR #18742). So compute_max_num_sequences becomes the single definition, and the consumers receive it instead of recomputing it: KVCacheManagerV2 (target, draft and cross managers all take the *target* engine's pool, since there is one SeqSlotManager per executor), the guided decoder, the sampler, and the two-model drafter's own SeqSlotManager. The coefficient itself is additive rather than multiplicative. Pipeline depth costs pp_size micro-batches of seats; the overlap scheduler defers a finished request's teardown by exactly one iteration, which costs one more generation on top -- not one more per stage. So the pool is (pp_size + 1) * max_batch_size. At pp_size == 1 the additive and multiplicative readings coincide at 2 * max_batch_size, which is why the headroom used to be expressible as a factor of 2; that coincidence is what made the multiplicative form look general. It is not, and it is not free: the pools scaled by this number include eagerly allocated [seats, draft_len, vocab] fp32 tensors, so the multiplicative form costs +100% at pp_size == 4 where the additive one costs +25%. Every reachable cell keeps its current value. The (pp, headroom) cell changes from B * pp to (pp + 1) * B but stays unreachable here: the gate still excludes pipeline parallelism, now for the router's sake rather than the sizing's. Also in this commit: * is_disagg_enabled() replaces three inlined copies of "cache_transceiver_config.backend is not None", one of which fed the index pool's factor of 2 while another fed the seat pool. * validate_seq_slot_pool_covers_admission() fails at startup on any skew, two-sided. A one-sided "seats >= admissible" guard is exactly what let nvbug 6627795 through. * resolve_max_num_sequences() replaces two fallbacks that recomputed the pool *without* the headroom gate -- they could only ever produce a number smaller than the slots they index. create_torch_sampler_args now requires the resolved value and no longer accepts the raw material for re-deriving it. * the guided decoder is sized by the seat pool unconditionally. Its state is indexed by py_seq_slot over the whole pool, so max_batch_size was already an IndexError waiting under pipeline parallelism, where admission permits max_batch_size * pp_size live requests. Pre-existing, and no CI coverage: none of the 39 guided-decoding entries in the QA lists uses PP. * hybrid/SSM architectures are withheld from the headroom, and the ADP router now takes the engine's headroom flag instead of re-deriving a predicate for it. MambaHybridCacheManagerV2 sizes its state-index pool from max_batch_size alone, so an extra seat would have no state slot behind it; the router must not credit a rank with seats the engine never allocated. Fixing that pool is a separate change. Verified with a CPU negative-control ladder (nsc cpu partition, whole-file baseline vs patched sources over 5 test files): 21 declared tests fail on the unpatched sources and pass on the patched ones, while 142 invariant tests pass on both -- including 6 of the 10 index-mapper capacity rows, which is the evidence that the refactor is a no-op outside the cells it claims to move. Signed-off-by: Chenfei Zhang --- tensorrt_llm/_torch/pyexecutor/_util.py | 250 +++++++++++++-- .../kv_cache/kv_cache_manager_v2.py | 31 +- .../kv_cache/mamba_cache_manager.py | 21 ++ .../_torch/pyexecutor/model_engine.py | 44 +-- tensorrt_llm/_torch/pyexecutor/py_executor.py | 8 + .../_torch/pyexecutor/py_executor_creator.py | 25 +- .../_torch/pyexecutor/scheduler/adp_router.py | 69 +++-- tensorrt_llm/_torch/speculative/utils.py | 45 ++- .../kv_cache/test_kv_cache_manager_v2.py | 174 +++++++++++ .../_torch/executor/test_adp_router.py | 114 ++++++- .../_torch/executor/test_seq_slot_sizing.py | 285 ++++++++++++++++-- .../speculative/test_spec_slot_pool_sizing.py | 110 +++++++ 12 files changed, 1042 insertions(+), 134 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index a5a1538c0f8e..af428e7fcb29 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -1414,6 +1414,17 @@ def configure_kv_cache_capacity(self, self._profiling_stage_data["activation_bytes"] = activation_bytes # ---------------------------handle max_gpu_total_bytes--------------------------------- + def _target_max_num_seq_slots(self) -> Optional[int]: + """Size of the executor's sequence-slot pool, or None if unavailable. + + There is exactly one SeqSlotManager per executor, sized by + compute_max_num_sequences from the *target* engine (see + create_py_executor_instance). Every KV cache manager -- target, draft and + cross -- must be able to hand out an index for each seat, otherwise the + pools drift apart and requests are silently deferred (nvbug 6627795). + """ + return getattr(self._model_engine, "max_num_seq_slots", None) + def _create_kv_cache_manager( self, model_engine: PyTorchModelEngine, @@ -1457,6 +1468,10 @@ def _create_kv_cache_manager( execution_stream=self._execution_stream, layer_mask=spec_dec_layer_mask, is_disagg=self._is_disagg, + # Always the target engine's seat pool, even when building the draft + # manager: the executor has a single SeqSlotManager, sized from the + # target engine, and every manager must cover it. + max_num_seq_slots=self._target_max_num_seq_slots(), cold_page_codec_provider=cold_page_codec_provider, ) @@ -1651,6 +1666,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, + max_num_seq_slots=self._target_max_num_seq_slots(), cold_page_codec_provider=cold_page_codec_provider, ) @@ -2026,6 +2042,7 @@ def _create_cross_kv_cache_manager( num_layers=num_layers, num_kv_heads=num_kv_heads, head_dim=head_dim, + max_num_seq_slots=self._target_max_num_seq_slots(), kv_cache_type=tensorrt_llm.bindings.internal.batch_manager. CacheType.CROSS, ) @@ -2336,6 +2353,7 @@ def _create_kv_cache_manager( head_dim: Optional[int] = None, kv_cache_type=None, is_disagg: bool = False, + max_num_seq_slots: Optional[int] = None, cold_page_codec_provider: Optional[object] = None) -> KVCacheManager: """ Returns: @@ -2475,6 +2493,12 @@ def _create_kv_cache_manager( manager_extra_kwargs["enable_stats"] = enable_kv_cache_stats manager_extra_kwargs[ "cold_page_codec_provider"] = cold_page_codec_provider + # Hybrid managers size their SSM state-slot pool from max_batch_size + # (see MambaHybridCacheManager._max_resident_sequences), so growing only + # the index mapper would let a request hold an index with no state slot. + # Sizing both pools together is left as a follow-up. + if not issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2): + manager_extra_kwargs["max_num_seq_slots"] = max_num_seq_slots if issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2): manager_extra_kwargs["is_disagg"] = is_disagg @@ -2994,24 +3018,142 @@ def create_kv_cache_compression_manager( return None +def is_disagg_enabled(cache_transceiver_config) -> bool: + """True when a cache transceiver backend is configured. + + Single definition of "this is a disaggregated server", so the seat pool and + the KV cache managers cannot disagree about it. The test was previously + inlined at each use site, which is how a derived fact acquires copies. + """ + return (cache_transceiver_config is not None + and cache_transceiver_config.backend is not None) + + def compute_max_num_sequences(mapping: Mapping, max_batch_size: int, disable_overlap_scheduler: bool, - enable_overlap_headroom: bool = False) -> int: + enable_overlap_headroom: bool = False, + is_disagg: bool = False) -> int: """Size the sequence-slot pool (and the sampler state it indexes). - ``enable_overlap_headroom`` is intentionally opt-in. Disaggregated - attention-DP needs a second non-PP slot set because the V2 scheduler can - backfill seats before the overlap scheduler releases the previous - iteration's terminal slots. Pipeline parallelism already sizes the pool - by ``pp_size``. + This is *the* definition of how many sequences can be simultaneously live. + Every pool keyed by live-request identity -- the sampler's per-slot state, + the guided decoder, the spec-decoding managers, KVCacheManagerV2's index + mapper -- must consume this number rather than re-deriving it from + ``max_batch_size``, because any two formulas that disagree are a bug: too + few index slots silently defers admitted requests (nvbug 6627795), too few + seats raises ``NoFreeSlotsError`` on the executor's event-loop thread. + + The terms are **additive, not multiplicative**: + + * ``pp_size`` micro-batches are structurally in flight under pipeline + parallelism, so the pool scales with pipeline depth. + * The overlap scheduler defers a finished request's teardown by exactly + **one** iteration, so at most one extra generation of seats is held while + the ADP router admits the batch that replaces it. That is ``+1`` + micro-batch worth of seats, not a doubling of every pipeline stage. + + At ``pp_size == 1`` the additive and multiplicative forms coincide at + ``2 * max_batch_size``, which is why the overlap headroom used to be + expressible as a factor of 2. It is a coincidence of ``pp_size == 1``, and + the multiplicative reading costs ``+100%`` at ``pp_size == 4`` where the + additive one costs ``+25%`` -- on pools that include eagerly-allocated + ``[seats, draft_len, vocab]`` tensors. """ - if mapping.has_pp(): - num_micro_batches = mapping.pp_size - else: - num_micro_batches = (2 if enable_overlap_headroom - and not disable_overlap_scheduler else 1) - return max_batch_size * num_micro_batches + # Pipeline depth: pp_size micro-batches are structurally in flight. + num_seats = max_batch_size * mapping.pp_size + if enable_overlap_headroom and not disable_overlap_scheduler: + num_seats += max_batch_size + if is_disagg: + # KVCacheManagerV2 sizes its own index pool max_batch_size * pp_size * 2 + # under disagg (the lease outlives the request while the KV transfer + # drains), and the seat pool must cover the index pool. max() rather + # than a further multiplication because the disagg and overlap + # coefficients each cover one extra cohort of in-flight sequences -- + # they overlap rather than compose. + num_seats = max(num_seats, max_batch_size * mapping.pp_size * 2) + return num_seats + + +def validate_seq_slot_pool_covers_admission(max_num_sequences: int, + kv_cache_manager) -> None: + """Fail at startup if the seat pool and the KV index pool disagree. + + Both are leases held for a request's whole lifetime, so the two counts must + be *equal*: every sequence the executor can seat needs an index, and every + sequence the KV cache manager can index needs a seat. The check is + deliberately two-sided, because each direction has already shipped as a + separate bug and a one-sided ``>=`` guard is what let the first one through: + + * index pool < seat pool -- ``_create_kv_cache`` warns ``No free IndexMapper + slots``, returns ``None`` and the scheduler defers the request. Silent: + costs throughput, no error (nvbug 6627795). + * index pool > seat pool -- a request is admitted that cannot be seated and + ``SlotManager.add_slot`` raises on the executor's event-loop thread, + killing the rank mid-collective. + + Managers that do not publish ``max_admissible_sequences`` (V1, hybrid) are + skipped rather than guessed at; adding the attribute is how a manager opts + into the check. + """ + if kv_cache_manager is None: + return + admissible = getattr(kv_cache_manager, "max_admissible_sequences", None) + if admissible is None: + return + if admissible == max_num_sequences: + return + direction = ("smaller" if admissible < max_num_sequences else "larger") + consequence = ( + "admitted requests would be silently deferred one at a time " + "(nvbug 6627795)" if admissible < max_num_sequences else + "a request could be admitted with no sequence slot to seat it, and " + "SlotManager.add_slot would raise on the executor's event loop") + raise ValueError( + f"{type(kv_cache_manager).__name__} can seat {admissible} concurrent " + f"sequences but the executor's sequence-slot pool holds " + f"{max_num_sequences}: the index pool is {direction} than the seat " + f"pool, so {consequence}. Both must come from " + "_util.compute_max_num_sequences; a mismatch means one of them was " + "re-derived from max_batch_size.") + + +def resolve_max_num_sequences(model_engine, + mapping: Mapping, + max_batch_size: int, + disable_overlap_scheduler: bool, + is_disagg: bool, + 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 the KV cache managers were sized against; + 3. only then a fresh ``compute_max_num_sequences``, reusing the engine's + headroom gate. + + Step 3 used to be step 1, called *without* ``enable_overlap_headroom``. That + silently sized the sampler and the executor's ``SeqSlotManager`` below the + index pool they share indices with, i.e. it reintroduced the very skew this + number exists to eliminate, in the one code path that has no test coverage. + """ + 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, + disable_overlap_scheduler, + enable_overlap_headroom=getattr( + model_engine, "_enable_adp_overlap_seq_slot_headroom", False), + is_disagg=is_disagg) def should_enable_adp_dummy_fixes(mapping: Mapping) -> bool: @@ -3039,14 +3181,16 @@ def should_enable_non_overlap_adp_forward_intent( def should_enable_adp_overlap_seq_slot_headroom( - mapping: Mapping, disable_overlap_scheduler: bool) -> bool: - """Gate extra sequence slots to non-PP attention-DP with overlap enabled. + mapping: Mapping, + disable_overlap_scheduler: bool, + is_hybrid: bool = False) -> bool: + """Gate extra sequence slots to attention-DP with the overlap scheduler on. The overlap scheduler defers a finished request's teardown by one iteration, so its sequence slot is still held when the ADP router admits the batch that replaces it. Without spare slots the router cannot backfill and the forward - batch runs short (nvbug-6627795); a second set of slots lets admission reach - max_batch_size on every rank every iteration. + batch runs short (nvbug-6627795); one extra micro-batch worth of slots lets + admission reach max_batch_size on every rank every iteration. Requiring attention DP -- rather than merely overlap -- is deliberate: with a single scheduling domain the executor's own admission bound already tracks @@ -3054,7 +3198,25 @@ def should_enable_adp_overlap_seq_slot_headroom( deferred teardown desynchronizes. Not gated on disaggregation: the mechanism is a property of overlap plus ADP admission, and was measured on an aggregated context-only run with no cache transceiver configured. + + Pipeline parallelism is still excluded here, but no longer because of the + sizing: compute_max_num_sequences now expresses the headroom additively, so + ``(pp_size + 1) * max_batch_size`` is a well-defined pool under PP. What is + missing is the *consumer*: ADPRouter's retiring-request correction requires + every pipeline stage to agree on which requests are retiring, and today only + the last stage marks generation requests GENERATION_TO_COMPLETE. Opening the + gate before that is fixed would allocate headroom no rank ever spends. + + Hybrid (Mamba/SSM) architectures are excluded. MambaHybridCacheManagerV2 + sizes its state-index pool from max_batch_size alone + (mamba_cache_manager.py: state_index_capacity), with neither pp_size nor the + seat count, so extra seats would let a request hold a seat with no SSM state + slot behind it. That pool is already inconsistent with its own + _max_resident_sequences() under PP; fixing it is a separate change, and + until then the headroom must not reach these models. """ + if is_hybrid: + return False return (mapping.enable_attention_dp and not mapping.has_pp() and not disable_overlap_scheduler) @@ -3094,15 +3256,26 @@ 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.disable_overlap_scheduler, + is_disagg, + max_num_sequences=max_num_sequences) + + # The seat pool and the KV index pool are sized independently and indexed by + # the same request identity, so any skew between them is a bug. Check it here + # rather than at first use: a startup ValueError names both numbers, while + # the runtime symptoms are a silent throughput loss in one direction and a + # raise inside the event loop in the other. + validate_seq_slot_pool_covers_admission(max_num_sequences, kv_cache_manager) 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( @@ -3291,8 +3464,13 @@ 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 pool is larger than this under the + # attention-DP overlap headroom -- (pp_size + 1) * max_batch_size -- because + # it must also cover the retiring cohort whose teardown the overlap scheduler + # defers by one iteration. Those seats are headroom for leases already held, + # not extra admission, so growing the budget with them would over-admit. scheduler_capacity = max_batch_size * mapping.pp_size if scheduler_capacity == 1 and mapping.enable_attention_dp and kv_cache_manager: scheduler_capacity += 1 @@ -3479,22 +3657,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 @@ -3526,10 +3704,16 @@ 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.disable_overlap_scheduler, + is_disagg_enabled(llm_args.cache_transceiver_config), + 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 d4a93d8de9e5..5fb16a914808 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 @@ -896,6 +896,7 @@ def __init__( is_disagg: bool = False, enable_stats: bool = False, num_reserved_index_slots: int = 1, + max_num_seq_slots: Optional[int] = None, is_estimating_kv_cache: bool = False, cold_page_codec_provider: Optional[object] = None, **kwargs, @@ -1399,17 +1400,45 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: # (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. + # + # `max_num_seq_slots` is the executor's sequence-slot pool size + # (`PyTorchModelEngine.max_num_seq_slots`, sized by + # `_util.compute_max_num_sequences`). An index slot and a sequence slot are + # both leases held for the whole lifetime of a request, so whenever the + # executor can admit N concurrent sequences the index mapper must be able + # to hand out N indices. Under the overlap scheduler with attention DP the + # seat pool carries an extra factor of 2 because teardown of the retiring + # batch (`_process_previous_batch`) happens *after* the replacement batch + # has already been scheduled, so both cohorts hold their leases at once. + # Without this term the index mapper runs dry and `_create_kv_cache` + # silently defers requests one at a time (nvbug 6627795). + # + # Take the larger of the two bounds rather than multiplying them: the + # disagg and overlap coefficients each cover one extra set of in-flight + # sequences, so they overlap rather than compose. max_num_sequences = max_batch_size * mapping.pp_size assert num_reserved_index_slots >= 0, "num_reserved_index_slots must be non-negative" index_mapper_capacity = ( - max_num_sequences * (2 if is_disagg else 1) + num_reserved_index_slots + max(max_num_sequences * (2 if is_disagg else 1), max_num_seq_slots or 0) + + 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"max_num_seq_slots={max_num_seq_slots}, " f"num_reserved_index_slots={num_reserved_index_slots}, " f"max_beam_width={max_beam_width})" ) + # Concurrent sequences this manager can seat: the index pool net of the + # slots reserved for padding/dummy requests, which are not available to + # real sequences. Published so that + # `_util.validate_seq_slot_pool_covers_admission` can compare it against + # the executor's sequence-slot pool without re-deriving either number -- + # every skew between the two is a bug, in both directions (a smaller + # index pool defers requests one at a time, a larger one lets a request + # be admitted that cannot be seated and `SlotManager.add_slot` then + # raises on the executor's event loop). + self.max_admissible_sequences = index_mapper_capacity - num_reserved_index_slots self.index_mapper = IndexMapper(index_mapper_capacity, max_beam_width) self._early_freed_index_requests: set[int] = set() self._prepare_page_table_tensor(index_mapper_capacity) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index 71ed3d3f3fb5..ab2fa6af9587 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -3172,6 +3172,11 @@ def __init__( self._request_id_to_state_index = {} self._request_id_to_is_dummy = {} + # Sized by *scheduled batch* position, not by live sequence: + # _setup_state_indices fills [0, len(requests)) for one iteration's + # requests. So this is deliberately max_batch_size and must NOT grow with + # pp_size or with the executor's sequence-slot pool -- unlike the SSM + # state slots checked below, which are per-live-sequence leases. state_index_capacity = (self.max_batch_size + self._num_reserved_dummy_slots) self.cuda_state_indices = torch.zeros([state_index_capacity], @@ -3192,6 +3197,22 @@ def __init__( LayerId(first_mamba_local_layer), MambaRole.SSM_STATE) num_ssm_slots = ((num_ssm_pages + self._ssm_page_index_scale - 1) // self._ssm_page_index_scale) + # Per-live-sequence leases, so this floor tracks the number of + # sequences that can be resident at once -- max_batch_size * pp_size. + # It is deliberately *not* raised to `max_admissible_sequences` (the + # index-mapper pool net of reserved slots): under disaggregation that + # pool carries a 2x for requests still draining their KV transfer, + # and whether such a request also retains its SSM state slot is + # unresolved. Requiring the larger number would turn a possibly + # adequate Mamba pool into a startup failure, so the count is rounded + # down here on purpose. + # + # This manager is also excluded from the attention-DP overlap seat + # headroom (`should_enable_adp_overlap_seq_slot_headroom` returns + # False for hybrid architectures, and `_create_kv_cache_manager` + # withholds `max_num_seq_slots`), so the seat pool it is sized against + # is exactly `_max_resident_sequences()` and cannot outgrow this + # floor. required_live_slots = (self._max_resident_sequences() + self._num_reserved_dummy_slots) if num_ssm_slots < required_live_slots: diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 1cd8d8e7041f..2fc6ddca9632 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -80,7 +80,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_mla +from .config_utils import is_hybrid_linear, is_mla from .cuda_graph_runner import (ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM, CUDAGraphRunner, CUDAGraphRunnerConfig, EncoderCUDAGraphRunner, @@ -432,23 +432,6 @@ def __init__( self.mapping = mapping if mapping.has_pp(): init_pp_comm(mapping) - # 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_adp_overlap_seq_slot_headroom, - should_enable_non_overlap_adp_forward_intent, - should_enable_scheduler_aware_adp_dummy) - self._enable_adp_overlap_seq_slot_headroom = ( - should_enable_adp_overlap_seq_slot_headroom( - mapping, 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_adp_overlap_seq_slot_headroom, - ) self.dist = dist if dist is not None: ExpertStatistic.create(self.dist.rank) @@ -526,6 +509,31 @@ def __init__( self._validate_breakable_cuda_graph_compatibility() pretrained_config = self.model.model_config.pretrained_config model_type = getattr(pretrained_config, "model_type", None) + # Attention-DP can backfill a batch before the overlap scheduler + # releases the previous batch's terminal sequence slots, so the seat + # pool needs one extra generation of slots. Both the gate and the + # sizing depend on the architecture -- hybrid/SSM cache managers size + # their state-slot pool from max_batch_size alone, so they cannot use + # the headroom -- which is why this runs after the model is loaded + # rather than next to `self.mapping`. + from ._util import (compute_max_num_sequences, is_disagg_enabled, + should_enable_adp_dummy_fixes, + should_enable_adp_overlap_seq_slot_headroom, + should_enable_non_overlap_adp_forward_intent, + should_enable_scheduler_aware_adp_dummy) + self._enable_adp_overlap_seq_slot_headroom = ( + should_enable_adp_overlap_seq_slot_headroom( + mapping, + llm_args.disable_overlap_scheduler, + is_hybrid=is_hybrid_linear(pretrained_config))) + 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_adp_overlap_seq_slot_headroom, + is_disagg=is_disagg_enabled(llm_args.cache_transceiver_config), + ) self._enable_scheduler_aware_adp_dummy = ( should_enable_scheduler_aware_adp_dummy( model_type, mapping, llm_args.disable_overlap_scheduler)) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index c2c43dced4fd..b9a8e80665a5 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -684,8 +684,16 @@ def __init__( # Router is built after async_transfer_manager so KVCacheAwareADPRouter # can receive the transfer-manager reference at construction time. + # + # The router's overlap correction spends sequence-slot headroom, so it is + # handed the engine's headroom flag rather than re-deriving the + # predicate: the flag is what sized the seat pool, and it withholds the + # headroom for architectures (hybrid/SSM) whose state-slot pool is not + # sized from that number. Absent flag => no headroom => no correction. self.adp_router: ADPRouter = ADPRouter.create( dist=self.dist, + has_seq_slot_headroom=getattr( + model_engine, "_enable_adp_overlap_seq_slot_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, diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 0b9a9317d926..cf42d7bb99be 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -37,8 +37,8 @@ get_spec_resource_manager) from ..virtual_memory import scope as virtual_memory_scope from ._util import (KvCacheCreator, _adjust_torch_mem_fraction, - create_py_executor_instance, instantiate_sampler, is_mla, - validate_feature_combination) + 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, uses_vswa_kv_cache_layout) @@ -756,14 +756,17 @@ 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_adp_overlap_seq_slot_headroom", - False) else max_batch_size) kwargs = { "guided_decoding_config": guided_decoding_config, - # The attention-DP overlap path follows the expanded slot - # pool. Other configurations retain max_batch_size. - "max_num_sequences": guided_decoder_slots, + # Unconditionally the seat pool. The guided decoder's state + # is indexed by py_seq_slot (guided_decoder.py: grammar_matchers + # [req.seq_slot], the bitmask rows), and py_seq_slot ranges over + # the whole pool -- so sizing this at max_batch_size is an + # IndexError under pipeline parallelism, where admission already + # permits max_batch_size * pp_size live requests. The previous + # conditional made the correct size depend on an unrelated + # attention-DP flag. + "max_num_sequences": max_num_seq_slots, "vocab_size_padded": model_engine.model.vocab_size_padded, "rank": mapping.rank, } @@ -874,8 +877,10 @@ 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) + # Same helper the model engine sizes its seat pool with: this predicate + # feeds KVCacheManagerV2's index-pool coefficient, so an independent copy + # here could disagree with the seat pool by a factor of 2. + 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 df23c77b971d..43499ed294ed 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py @@ -227,31 +227,50 @@ class ADPRouter(ABC): needs_prefix_matches: bool = False - def __init__(self, dist: Distributed): + def __init__(self, dist: Distributed, has_seq_slot_headroom: bool = True): self.dist = dist # Whether to route on the overlap-corrected active list (nvbug-6627795). # - # Gated off under pipeline parallelism, deliberately matching - # ``should_enable_adp_overlap_seq_slot_headroom`` in ``_util.py``: the - # correction lets a rank hold more requests than it is charged for, so - # it is only safe where the sequence-slot pool has the matching - # headroom, and that headroom is sized for non-PP only. Beyond the slot - # accounting, ``GENERATION_TO_COMPLETE`` is marked on the last pipeline - # stage alone, while every rank pops from its own copy of the waiting - # queue -- so a PP-enabled correction would have the stages admit - # different numbers of requests and diverge. + # This is the *model engine's* headroom flag, passed in rather than + # re-derived here, because the two must never disagree: # - # Not additionally gated on ``disable_overlap_scheduler``: without - # overlap the retire is not deferred, so no request is ever in - # ``GENERATION_TO_COMPLETE`` when the router runs and the filter is - # arithmetically a no-op (measured: zero such requests on 5126/5126 - # routing records with overlap disabled). - self.exclude_retiring_requests = not dist.mapping.has_pp() + # 1. The correction lets a rank hold more requests than it is charged + # for, so it is only sound where the sequence-slot pool has matching + # headroom. That pool is sized by + # ``_util.compute_max_num_sequences`` from the same flag + # (``should_enable_adp_overlap_seq_slot_headroom``), which withholds + # the headroom from hybrid/SSM architectures whose state-slot pool is + # sized from ``max_batch_size`` alone. Re-deriving the predicate here + # -- as ``not dist.mapping.has_pp()`` used to -- is exactly how the + # router comes to credit a rank with seats the engine never + # allocated. + # 2. Pipeline parallelism is excluded by that same flag, and the reason + # is no longer the sizing -- the seat pool is well defined under PP at + # ``(pp_size + 1) * max_batch_size``. It is that the stages must agree + # on *which* requests are retiring: each rank pops from its own copy + # of the waiting queue, so a per-stage disagreement makes them admit + # different numbers of requests and diverge. Only the last stage marks + # ``GENERATION_TO_COMPLETE`` for generation requests today. Note the + # *context* path in ``_update_request_states_tp`` already evaluates the + # same predicate on every rank, so the asymmetry is only ever in the + # generation path. + # + # Not additionally gated on ``disable_overlap_scheduler`` here: the flag + # already is, and without overlap the retire is not deferred, so no + # request is ever in ``GENERATION_TO_COMPLETE`` when the router runs and + # the filter is arithmetically a no-op (measured: zero such requests on + # 5126/5126 routing records with overlap disabled). + # + # Erring off is the safe direction: excluding fewer retirees admits + # fewer requests (a missed optimization), while excluding more than the + # seat pool covers is a slot exhaustion or a stage divergence. + 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, @@ -260,6 +279,14 @@ def create( Args: dist: Distributed communicator. + has_seq_slot_headroom: Whether the executor's sequence-slot pool was + sized with the overlap headroom + (``should_enable_adp_overlap_seq_slot_headroom``). Required, not + defaulted, because the router's retiring-request correction is + only sound when those extra seats exist -- see + ``__init__``. Passed through from + ``model_engine._enable_adp_overlap_seq_slot_headroom`` so the + sizing and the routing decision cannot drift apart. 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 @@ -281,6 +308,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, @@ -294,6 +322,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, @@ -303,7 +332,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( @@ -607,6 +636,7 @@ def __init__( self, dist: "Distributed", kv_cache_manager, + has_seq_slot_headroom: bool = True, load_balance_weight: float = 1.0, match_rate_threshold: float = 0.1, fair_share_multiplier: float = 2.0, @@ -614,7 +644,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 @@ -912,11 +942,12 @@ class ConversationAwareADPRouter(ADPRouter): def __init__( self, dist: "Distributed", + has_seq_slot_headroom: bool = True, 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)) diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 79be1459a043..2c5cd4f0c467 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -572,6 +572,25 @@ 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 to keep max_batch_size. + + Pools keyed by live-request identity must follow the executor's + SeqSlotManager pool rather than max_batch_size: the overlap scheduler holds a + finished request's slot for one more iteration while its replacement is + admitted, so the transient demand exceeds max_batch_size (nvbug-6627795). + Buffers indexed by *batch position* deliberately keep max_batch_size -- the + micro-batch scheduler caps every forward at max_batch_size. + + Gated on the same flag ``_set_up_spec_metadata`` reads, so every spec-decoding + pool agrees with the metadata about which number it is indexed by. Returning + None (headroom off) preserves the established max_batch_size sizing. + """ + if not getattr(model_engine, "_enable_adp_overlap_seq_slot_headroom", False): + return None + 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: @@ -580,19 +599,7 @@ 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 - # Pools keyed by live-request identity must follow the executor's - # SeqSlotManager pool rather than max_batch_size: the attention-DP overlap - # headroom makes it 2 * max_batch_size so a finished request can hold its slot - # for one more iteration while its replacement is admitted (nvbug-6627795). - # Buffers indexed by *batch position* deliberately keep max_num_requests -- - # the micro-batch scheduler caps every forward at max_batch_size. - # - # Opted into by the same flag ``_set_up_spec_metadata`` uses, so the manager - # and the metadata never disagree about the pool. None (the other topologies, - # PP included) preserves the established max_num_requests sizing. - num_seq_slots = None - if getattr(model_engine, "_enable_adp_overlap_seq_slot_headroom", False): - num_seq_slots = getattr(model_engine, "max_num_seq_slots", None) + 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 @@ -760,6 +767,16 @@ def get_spec_drafter(model_engine, return spec_config.drafter max_num_requests = model_engine.batch_size + # The draft loop runs its own slot pool, but the indices it hands out address + # buffers sized by the *target* engine's seat pool: the shared sampler + # (instantiate_sampler), the draft KV cache manager's IndexMapper + # (KvCacheCreator._target_max_num_seq_slots) and spec_resource_manager above. + # It must therefore be sized from the same number. It is also load-bearing, + # not merely tidy: the previous draft batch's slots are released by + # cleanup_previous_draft_resources a full iteration later + # (py_executor.py:5347), so a pool of max_batch_size raises NoFreeSlotsError + # precisely when the overlap headroom is doing its job. + draft_slots = seat_pool_or_none(model_engine) or max_num_requests if spec_config.spec_dec_mode.is_draft_target( ) or spec_config.spec_dec_mode.is_eagle3( ) or spec_config.spec_dec_mode.is_mtp_eagle(): @@ -767,7 +784,7 @@ def get_spec_drafter(model_engine, draft_model_engine, spec_config.max_draft_len, spec_config.tokens_per_gen_step - 1, - SeqSlotManager(max_num_requests), + SeqSlotManager(draft_slots), sampler, spec_resource_manager=spec_resource_manager, guided_decoder=guided_decoder) 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 c359d2eed1c1..6cff111d3ae8 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 @@ -1140,3 +1140,177 @@ 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, + max_num_seq_slots: int | None = None, +) -> tuple[int, int, int | None]: + """Construct a manager and return (IndexMapper capacity, page-table capacity, + published ``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 same number net of the + reserved slots, published for the startup validator so that the seat pool and + the index pool can be compared without either being re-derived. + + The third is read with ``getattr(..., None)`` rather than as an attribute so + that a manager which does not publish it at all fails only the test that is + *about* publishing it. Asserting it here would make every capacity row fail + for one and the same trivial reason, which would destroy the negative + control's ability to say *which* topologies the coefficient actually moved. + """ + 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, num_reserved_index_slots and max_num_seq_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, + max_num_seq_slots=max_num_seq_slots, + ) + 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], + getattr(manager, "max_admissible_sequences", None), + ) + + +# (max_batch_size, pp_size, is_disagg, reserved, max_num_seq_slots, expected capacity) +# +# The seat pool (`PyTorchModelEngine.max_num_seq_slots`) and the index mapper are +# both per-request leases, so the mapper must cover whatever the executor can +# admit. Rows 1-2 are the nvbug 6627795 case: aggregated attention DP under the +# overlap scheduler doubles the seat pool, and before the fix the mapper stayed +# at B+1 and silently deferred requests. +_INDEX_MAPPER_CAPACITY_CASES = [ + # aggregated + ADP + overlap: seats are 2B, so the mapper must be 2B too. + pytest.param(2, 1, False, 1, 4, 5, id="agg_adp_overlap"), + pytest.param(8, 1, False, 1, 16, 17, id="agg_adp_overlap_b8"), + # No seat headroom (no ADP, or overlap off): unchanged at B+1. + pytest.param(2, 1, False, 1, 2, 3, id="agg_no_headroom"), + # Negative control: callers that pass nothing keep the pre-fix allocation. + pytest.param(2, 1, False, 1, None, 3, id="seats_unset_is_unchanged"), + # Negative control: disagg does not compound with the seat headroom. The two + # coefficients cover the same extra cohort, so this is max(4, 4)+1, not 4B+1. + pytest.param(2, 1, True, 1, 4, 5, id="disagg_does_not_compound"), + # Disagg without seat headroom keeps its own 2x. + pytest.param(2, 1, True, 1, None, 5, id="disagg_only"), + # PP without the headroom: seats are B*pp, which the mapper already matched. + pytest.param(2, 4, False, 1, 8, 9, id="pp4"), + # PP with the headroom: seats are (pp+1)*B == 10, additive rather than + # 2*B*pp == 16. The mapper follows the seat pool verbatim, so this row is what + # makes the extra generation of seats usable under pipeline parallelism. + pytest.param(2, 4, False, 1, 10, 11, id="pp4_adp_overlap"), + # Reserved slots are still added on top of the widened pool. + pytest.param(2, 1, False, 5, 4, 9, id="reserved_slots_still_added"), + # A seat pool smaller than the mapper's own floor must never shrink it. + pytest.param(4, 1, False, 1, 1, 5, id="small_seat_pool_does_not_shrink"), +] + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "max_batch_size,pp_size,is_disagg,reserved,max_num_seq_slots,expected", + _INDEX_MAPPER_CAPACITY_CASES, +) +def test_index_mapper_capacity_covers_seq_slot_pool( + max_batch_size: int, + pp_size: int, + is_disagg: bool, + reserved: int, + max_num_seq_slots: int | None, + expected: int, +) -> None: + capacity, page_table_capacity, _ = _index_mapper_capacity_for( + max_batch_size=max_batch_size, + pp_size=pp_size, + is_disagg=is_disagg, + num_reserved_index_slots=reserved, + max_num_seq_slots=max_num_seq_slots, + ) + assert capacity == expected + assert page_table_capacity == expected + + +# The rows where the *published* number is worth stating separately from the +# capacity arithmetic: the aggregated cell nvbug 6627795 was filed for, the new +# pipeline-parallel cell, and a caller that supplies no seat pool at all. +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "max_batch_size,pp_size,reserved,max_num_seq_slots,expected_admissible", + [ + pytest.param(2, 1, 1, 4, 4, id="agg_adp_overlap"), + pytest.param(2, 4, 1, 10, 10, id="pp4_adp_overlap"), + pytest.param(2, 1, 5, None, 2, id="seats_unset_reserved_excluded"), + ], +) +def test_index_mapper_publishes_max_admissible_sequences( + max_batch_size: int, + pp_size: int, + reserved: int, + max_num_seq_slots: int | None, + expected_admissible: int, +) -> None: + """The manager publishes what it can seat, so nobody has to re-derive it. + + ``_util.validate_seq_slot_pool_covers_admission`` compares this against the + executor's sequence-slot pool at startup. It is the pool *net of* the reserved + padding/dummy slots, which no real sequence can take, and on every row where a + seat pool was plumbed through it equals that seat pool exactly -- that + equality is the invariant the validator enforces, in both directions. + """ + _, _, admissible = _index_mapper_capacity_for( + max_batch_size=max_batch_size, + pp_size=pp_size, + num_reserved_index_slots=reserved, + max_num_seq_slots=max_num_seq_slots, + ) + assert admissible == expected_admissible + if max_num_seq_slots is not None: + assert admissible == max_num_seq_slots diff --git a/tests/unittest/_torch/executor/test_adp_router.py b/tests/unittest/_torch/executor/test_adp_router.py index 478fad7d4bdb..1adb9d591bf4 100644 --- a/tests/unittest/_torch/executor/test_adp_router.py +++ b/tests/unittest/_torch/executor/test_adp_router.py @@ -7,6 +7,7 @@ - Strict/relaxed attention-DP request routing while respecting rank capacity """ +import inspect from unittest.mock import MagicMock, Mock import pytest @@ -417,16 +418,16 @@ def test_gather_all_rank_states_reports_zero_when_all_retiring(self): assert states[0].num_active_tokens == 0 assert states[0].num_retiring_requests == 2 - def test_gather_all_rank_states_keeps_retiring_under_pp(self): - # Under pipeline parallelism the correction is off, matching the - # sequence-slot headroom gate in _util.py. Two reasons it must stay off: - # the slot pool is sized pp_size * max_batch_size with no headroom for - # requests a rank holds but is not charged for, and - # GENERATION_TO_COMPLETE is marked on the last stage only while every - # rank pops from its own copy of the waiting queue -- so a corrected - # count would have the stages admit different numbers of requests. + 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) + 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), @@ -443,11 +444,76 @@ def test_gather_all_rank_states_keeps_retiring_under_pp(self): assert states[0].num_retiring_requests == 0 assert states[0].num_active_tokens == 600 - def test_exclude_retiring_requests_follows_pipeline_parallelism(self): + def test_exclude_retiring_requests_follows_the_seat_pool_headroom(self): # The flag is the single gate; _pad_attention_dp_dummy_request reads it # rather than re-deriving the predicate, so the two cannot drift. - assert DefaultADPRouter(dist=_mock_dist(has_pp=False)).exclude_retiring_requests is True - assert DefaultADPRouter(dist=_mock_dist(has_pp=True)).exclude_retiring_requests is False + # + # It tracks the *engine's* headroom flag and nothing else. It used to be + # `not dist.mapping.has_pp()`, which was a second derivation of the same + # fact: correct only as long as the sizing gate happened to exclude + # exactly PP. Once the sizing gate also excluded hybrid architectures the + # two disagreed, and the router credited a rank with seats that were + # never allocated. Pipeline parallelism is now in scope on both sides. + 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 + + @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) @@ -1443,10 +1509,18 @@ def test_gather_all_rank_states_excludes_retiring(self): 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. @@ -1501,7 +1575,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" @@ -1510,7 +1589,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_seq_slot_sizing.py b/tests/unittest/_torch/executor/test_seq_slot_sizing.py index 236a8c4c6fe8..d3ba178b985f 100644 --- a/tests/unittest/_torch/executor/test_seq_slot_sizing.py +++ b/tests/unittest/_torch/executor/test_seq_slot_sizing.py @@ -6,53 +6,98 @@ still hold their sequence slots when the next iteration's prepare_resources runs, while the capacity 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. +backfilled their seats. Transient slot demand is therefore one extra +micro-batch worth of slots, regardless of whether speculative decoding is +enabled. The headroom is selected from runtime topology, not model +architecture -- except that hybrid/SSM architectures are excluded, because +their state-slot pool is not sized from this number. + +That extra generation is **additive** in pp_size, not multiplicative: +pipeline depth already accounts for the pp_size micro-batches structurally +in flight, and the overlap deferral is one iteration on top of them. So the +pool is (pp_size + 1) * max_batch_size, which at pp_size == 1 coincides with +the historical 2 * max_batch_size. 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). +for the sampler state (create_torch_sampler_args); resolve_max_num_sequences +is how a consumer obtains it without re-deriving it. """ +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_adp_overlap_seq_slot_headroom, should_enable_non_overlap_adp_forward_intent, should_enable_scheduler_aware_adp_dummy, + validate_seq_slot_pool_covers_admission, ) from tensorrt_llm.mapping import Mapping +# (pp_size, disable_overlap, enable_overlap_headroom, is_disagg, expected_factor) +# +# The terms are additive, not multiplicative: pipeline depth costs pp_size +# micro-batches of seats, and the overlap deferral costs exactly one more +# generation on top -- not one more per stage. At pp_size == 1 the two readings +# coincide at 2x, which is why the headroom used to be expressible as a factor of +# 2; the pp>1 rows are where they part company (5x, not 8x, at pp=4). SIZING_CASES = [ - # (pp_size, disable_overlap, enable_overlap_headroom, expected_factor) - (1, False, True, 2), - (1, False, False, 1), - (1, True, True, 1), - # Existing PP sizing is preserved regardless of the headroom opt-in. - (2, False, True, 2), - (4, False, True, 4), - (4, True, False, 4), + # No PP. Every cell here is unchanged by this PR. + (1, False, False, False, 1), + (1, False, True, False, 2), + (1, True, True, False, 1), + (1, False, False, True, 2), + (1, False, True, True, 2), + # PP without the headroom: unchanged. + (4, False, False, False, 4), + (4, True, True, False, 4), + (2, False, False, True, 4), + (4, False, False, True, 8), + # PP with the headroom: the one intended behaviour change (was pp_size). + (2, False, True, False, 3), + (4, False, True, False, 5), + # Disagg dominates the additive term rather than compounding with it: the two + # coefficients each cover one extra cohort of in-flight sequences. + (4, False, True, True, 8), ] @pytest.mark.parametrize( - "enable_attention_dp,pp_size,disable_overlap,expected", + "enable_attention_dp,pp_size,disable_overlap,is_hybrid,expected", [ # No cache-transceiver term: the gate no longer looks at disaggregation # at all, because nvbug-6627795 reproduced on an aggregated context-only # run with no transceiver configured. - (True, 1, False, True), - (False, 1, False, False), - (True, 2, False, False), - (True, 1, True, False), + (True, 1, False, False, True), + (False, 1, False, False, False), + (True, 1, True, False, False), + # Pipeline parallelism is still out of scope, but no longer because the + # sizing cannot express it -- compute_max_num_sequences is additive, so + # (pp_size + 1) * max_batch_size is well defined. The missing piece is the + # consumer: only the last pipeline stage marks generation requests + # GENERATION_TO_COMPLETE, so the ADP router's retiring-request correction + # would not be rank-consistent. + (True, 2, False, False, False), + (True, 4, False, False, False), + # Hybrid/SSM architectures are excluded: MambaHybridCacheManagerV2 sizes + # its state-index pool from max_batch_size alone, so an extra seat would + # have no state slot behind it. + (True, 1, False, True, False), + (True, 4, False, True, False), ], ) def test_adp_overlap_seq_slot_headroom_gate( - enable_attention_dp, pp_size, disable_overlap, expected + enable_attention_dp, pp_size, disable_overlap, is_hybrid, expected ): mapping = Mapping( world_size=pp_size, @@ -61,7 +106,12 @@ def test_adp_overlap_seq_slot_headroom_gate( enable_attention_dp=enable_attention_dp, ) - assert should_enable_adp_overlap_seq_slot_headroom(mapping, disable_overlap) is expected + assert ( + should_enable_adp_overlap_seq_slot_headroom( + mapping, disable_overlap, is_hybrid=is_hybrid + ) + is expected + ) @pytest.mark.parametrize("pp_size,expected", [(1, True), (2, False)]) @@ -99,10 +149,10 @@ def test_non_overlap_adp_forward_intent_scope(pp_size, disable_overlap, expected @pytest.mark.parametrize( - "pp_size,disable_overlap,enable_overlap_headroom,expected_factor", SIZING_CASES + "pp_size,disable_overlap,enable_overlap_headroom,is_disagg,expected_factor", SIZING_CASES ) def test_compute_max_num_sequences_scopes_overlap_headroom( - pp_size, disable_overlap, enable_overlap_headroom, expected_factor + pp_size, disable_overlap, enable_overlap_headroom, is_disagg, expected_factor ): max_batch_size = 8 mapping = Mapping(world_size=pp_size, tp_size=1, pp_size=pp_size) @@ -112,20 +162,86 @@ def test_compute_max_num_sequences_scopes_overlap_headroom( max_batch_size, disable_overlap, enable_overlap_headroom=enable_overlap_headroom, + is_disagg=is_disagg, ) == max_batch_size * expected_factor ) +@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 very skew this number exists to remove. + The third row is that branch, and it must still land on the headroom value. + """ + engine = SimpleNamespace( + max_num_seq_slots=engine_seats, + _enable_adp_overlap_seq_slot_headroom=True, + ) + mapping = Mapping(world_size=1, tp_size=1, pp_size=1, enable_attention_dp=True) + + assert ( + resolve_max_num_sequences( + engine, + mapping, + 8, + False, + False, + max_num_sequences=explicit, + ) + == expected + ) + + +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, @@ -134,3 +250,124 @@ def test_sampler_uses_executor_slot_pool_capacity(slot_factor): max_num_sequences=max_num_sequences, ) assert args.max_num_sequences == max_num_sequences + + +class _FakeManager: + """Stands in for a KV cache manager that publishes its seating capacity.""" + + 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)) + + +@pytest.mark.parametrize("admissible", [8, 32]) +def test_validator_is_two_sided(admissible): + """Both directions of the skew are bugs, and both have shipped. + + A one-sided ``seats >= admissible`` guard is what let nvbug 6627795 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. The + other direction -- index pool larger than the seat pool -- admits a request + that cannot be seated and raises inside the executor's event loop. + """ + with pytest.raises(ValueError, match="sequence-slot pool"): + validate_seq_slot_pool_covers_admission(16, _FakeManager(admissible)) + + +@pytest.mark.parametrize( + "manager", + [ + None, # non-generation models have no KV cache manager + SimpleNamespace(), # V1 sizes its index pool by an unrelated rule + ], +) +def test_validator_skips_managers_that_do_not_publish_capacity(manager): + """Opt-in, not guessed at. + + A manager whose index pool is sized by some other rule -- V1, which receives + none of this plumbing and re-derives from bare max_batch_size -- must not be + compared against a number it never consumed. + + Hybrid is deliberately *not* an example here: MambaHybridCacheManagerV2 + derives from KVCacheManagerV2 and therefore does publish + max_admissible_sequences, so it *is* validated. That comparison holds + because the headroom is withheld from hybrid on both sides -- seats are + B*pp (or 2*B*pp under disagg) and, with max_num_seq_slots withheld, its + index pool lands on exactly the same number. What hybrid does not receive + is the seat pool, not the check. + """ + validate_seq_slot_pool_covers_admission(16, manager) + + +def _make_kv_cache_creator(max_num_seq_slots) -> 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 + # 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) + c._model_engine = SimpleNamespace(max_num_seq_slots=max_num_seq_slots) + return c + + +@pytest.mark.parametrize("max_num_seq_slots", [8, 16, None]) +def test_kv_cache_manager_receives_executor_seq_slot_pool(max_num_seq_slots): + """The seat pool size must reach the KV cache manager verbatim. + + The manager sizes its IndexMapper from this number so that every sequence + the executor can admit is guaranteed an index. Recomputing the coefficient + inside the manager would let the two pools drift apart (nvbug 6627795). + """ + creator = _make_kv_cache_creator(max_num_seq_slots) + model_engine = SimpleNamespace( + model=SimpleNamespace(model_config=SimpleNamespace(is_generation=True)), + max_num_seq_slots=max_num_seq_slots, + ) + + 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["max_num_seq_slots"] == max_num_seq_slots + + +def test_draft_manager_uses_the_target_engines_seq_slot_pool(): + """A draft engine's own seat count must not size the draft index pool. + + The executor has a single SeqSlotManager, sized from the target engine. Two- + model speculative decoding builds the draft KV cache manager by passing the + *draft* engine to the same helper; if that engine's (smaller) number were + used, the draft manager's IndexMapper would become the new bottleneck and + reintroduce the silent deferral this sizing exists to prevent. + """ + creator = _make_kv_cache_creator(16) + draft_engine = SimpleNamespace( + model=SimpleNamespace(model_config=SimpleNamespace(is_generation=True)), + max_num_seq_slots=8, + ) + + with patch( + "tensorrt_llm._torch.pyexecutor._util._create_kv_cache_manager", + return_value=None, + ) as create: + creator._create_kv_cache_manager(draft_engine) + + assert create.call_args.kwargs["max_num_seq_slots"] == 16 diff --git a/tests/unittest/_torch/speculative/test_spec_slot_pool_sizing.py b/tests/unittest/_torch/speculative/test_spec_slot_pool_sizing.py index c4422df334f3..b3591513c00f 100644 --- a/tests/unittest/_torch/speculative/test_spec_slot_pool_sizing.py +++ b/tests/unittest/_torch/speculative/test_spec_slot_pool_sizing.py @@ -36,8 +36,10 @@ 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.pyexecutor.resource_manager import NoFreeSlotsError from tensorrt_llm._torch.speculative.utils import ( _build_spec_metadata, + get_spec_drafter, get_spec_metadata, get_spec_resource_manager, ) @@ -220,6 +222,114 @@ def test_every_resource_manager_branch_forwards_the_slot_pool(): ) +# --------------------------------------------------------------------------- +# The drafter's own slot pool. Two-model speculation runs a *second* +# SeqSlotManager (get_spec_drafter), and the guard above cannot see it: it walks +# get_spec_resource_manager only, and matches names ending in "Manager" that take +# a num_seq_slots *keyword*, while SeqSlotManager takes its size positionally. +# That blind spot is why this pool stayed at max_batch_size while every pool +# around it moved to the seat count. +# --------------------------------------------------------------------------- + + +def _drafter_engine(headroom: bool, seats: int = POOL): + """A stub target engine for get_spec_drafter's draft-target branch.""" + spec_dec_mode = types.SimpleNamespace( + is_user_provided=lambda: False, + is_draft_target=lambda: True, + is_eagle3=lambda: False, + is_mtp_eagle=lambda: False, + is_ngram=lambda: False, + ) + spec_config = types.SimpleNamespace( + spec_dec_mode=spec_dec_mode, + max_draft_len=2, + tokens_per_gen_step=3, + max_concurrency=None, + draft_len_schedule=None, + ) + return types.SimpleNamespace( + spec_config=spec_config, + batch_size=R, + max_num_seq_slots=seats, + _enable_adp_overlap_seq_slot_headroom=headroom, + ) + + +def _drafter(headroom: bool, seats: int = POOL): + return get_spec_drafter( + _drafter_engine(headroom, seats), + draft_model_engine=object(), + sampler=object(), + spec_resource_manager=None, + ) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("headroom,expected_pool", [(True, POOL), (False, R)]) +def test_draft_slot_pool_follows_the_target_seat_pool(headroom, expected_pool): + """The draft pool is sized by the *target* engine's seat count. + + The indices it hands out address buffers sized by that seat count -- the + sampler shared with PyExecutor, the draft KV cache manager's IndexMapper, and + spec_resource_manager -- and ``_create_draft_request`` even carries the + target's ``py_seq_slot`` across as ``target_seq_slot``. Sizing this pool + independently is what makes the two disagree. Headroom off keeps + max_batch_size, so nothing changes for the other topologies. + """ + assert _drafter(headroom).draft_seq_slot_manager.slot_manager.max_num_requests == expected_pool + + +@pytest.mark.cpu_only +def test_draft_slot_pool_survives_a_full_overlap_turnover(): + """R retiring + R admitted, with the negative control that proves it bites. + + ``cleanup_previous_draft_resources`` releases the previous draft batch's slots + a full iteration later (py_executor.py:5347), so with the headroom on both + cohorts hold draft slots at once. At max_batch_size the (R+1)-th lease raises + NoFreeSlotsError -- inside the draft loop, mid-iteration. + """ + pool = _drafter(True).draft_seq_slot_manager.slot_manager + slots = [pool.add_slot(rid) for rid in range(2 * R)] + assert len(set(slots)) == 2 * R + + starved = _drafter(False).draft_seq_slot_manager.slot_manager + for rid in range(R): + starved.add_slot(rid) + with pytest.raises(NoFreeSlotsError): + starved.add_slot(R) + + +@pytest.mark.cpu_only +def test_the_drafter_slot_pool_is_not_re_derived_from_max_batch_size(): + """Structural guard, because the behavioural test above can be satisfied by + an accident of the stub: assert the source never passes a batch-size symbol + to SeqSlotManager. A new speculation mode that adds a second drafter branch + fails here rather than in a draft loop. + """ + tree = ast.parse(textwrap.dedent(inspect.getsource(get_spec_drafter))) + + offenders = [] + seen = 0 + for node in ast.walk(tree): + if not isinstance(node, ast.Call) or getattr(node.func, "id", None) != "SeqSlotManager": + continue + seen += 1 + for arg in list(node.args) + [kw.value for kw in node.keywords]: + if isinstance(arg, ast.Name) and arg.id in ("max_num_requests", "max_batch_size"): + offenders.append(arg.id) + elif isinstance(arg, ast.Attribute) and arg.attr in ("batch_size", "max_batch_size"): + offenders.append(arg.attr) + + # Without this the guard passes vacuously the day the call is renamed away. + assert seen, "get_spec_drafter no longer builds a SeqSlotManager; retarget this guard" + assert not offenders, ( + f"get_spec_drafter sizes its SeqSlotManager from {sorted(set(offenders))}; it must " + "follow the target engine's seat pool (seat_pool_or_none) or the draft lease for a " + "request seated above max_batch_size raises NoFreeSlotsError." + ) + + @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): From ac72cda1ba9f21ad88fe16e4b6bc1e22f5777fca Mon Sep 17 00:00:00 2001 From: Chenfei Zhang Date: Mon, 7 Sep 2026 08:55:14 -0700 Subject: [PATCH 10/22] [https://nvbugs/6627795][fix] derive the retiring-request count identically on every pipeline stage The ADP router subtracts retiring requests from each rank's reported load. Under pipeline parallelism that correction is only safe if every stage agrees on *which* requests are retiring: each rank pops from its own copy of the waiting queue and applies the router's decision locally, so a per-stage disagreement means the stages admit different numbers of requests and diverge -- a hang, not a wrong number. That is why the correction is gated off whenever pp_size > 1. The disagreement is in the call site, not the predicate. LlmRequest::willCompleteNextIteration is pure arithmetic on token counts, with no EOS check, no stop words and no sampler state, and those counts are replicated to every stage in the same iteration: the last stage's sample state is ring-broadcast by _ring_broadcast_sample_state and applied on all ranks by _handle_executed_batch, with the per-iteration batch count itself ring-broadcast from rank 0. The *context* path in _update_request_states_tp already evaluates the same predicate on every rank. Only the generation-path marking is asymmetric, and only because it sits inside the last-stage branch of _executor_loop_pp. So _forward_step_inter_pp makes the same call the last stage makes, at the structurally identical point -- immediately after _update_request_states, under the same overlap guard. Deliberately not hoisted into a shared location such as _handle_executed_batch: that would change *when* GENERATION_TO_COMPLETE is set on the existing PP path, which also feeds set_exclude_last_generation_logits and the capacity scheduler's no_schedule_after_state. Adding the call to the stage that is missing one has the smaller blast radius. This is inert on its own -- the headroom gate still excludes PP, so exclude_retiring_requests stays False there and nothing reads the marking. The next commit opens the gate. The new test drives the stage-local arithmetic with one request object per simulated stage, as the real thing has, and asserts the marked *set* agrees rather than merely the count. Its negative control is the pre-change behaviour: if only the last stage marks, the stages must be detectably inconsistent -- without that case a test asserting "the counts agree" would keep passing after a regression that removed the marking from every stage. Two structural assertions pin the call itself, including that it stays ordered after the state update: marking first would read token counts from before this micro-batch and disagree with a stage that marks after. Signed-off-by: Chenfei Zhang --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 13 ++ .../test_pp_retiring_rank_consistency.py | 192 ++++++++++++++++++ 2 files changed, 205 insertions(+) create mode 100644 tests/unittest/_torch/executor/test_pp_retiring_rank_consistency.py diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index b9a8e80665a5..efc4851ec749 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -5601,6 +5601,19 @@ def _forward_step_inter_pp(self, sampler_event = torch.cuda.Event() sampler_event.record() self._update_request_states(scheduled_batch) + # Mark retiring generation requests here too, not only on the last + # pipeline stage (_executor_loop_pp). The predicate is pure arithmetic on + # token counts (LlmRequest::willCompleteNextIteration) and those counts are + # replicated to every stage in the same iteration: the last stage's sample + # state is ring-broadcast (_ring_broadcast_sample_state) and applied on all + # ranks by _handle_executed_batch, with the per-iteration batch count + # itself ring-broadcast from rank 0. Without this call the stages disagree + # about which requests are retiring, and ADPRouter's admission correction + # -- which subtracts them from each rank's load -- would let the stages + # admit different numbers of requests and diverge. + if not self.disable_overlap_scheduler: + self._update_generation_requests_that_will_complete_next_iteration( + scheduled_batch.generation_requests) sampling_requests = scheduled_batch.context_requests_last_chunk + scheduled_batch.generation_requests return self.sampler.SampleState( requests=sampling_requests, diff --git a/tests/unittest/_torch/executor/test_pp_retiring_rank_consistency.py b/tests/unittest/_torch/executor/test_pp_retiring_rank_consistency.py new file mode 100644 index 000000000000..47eb13e977e8 --- /dev/null +++ b/tests/unittest/_torch/executor/test_pp_retiring_rank_consistency.py @@ -0,0 +1,192 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Every pipeline stage must derive the same retiring-request count. + +The ADP router subtracts retiring requests from each rank's reported load +(nvbug-6627795). Under pipeline parallelism that correction is only safe if every +stage agrees on *which* requests are retiring: each rank pops from its own copy of +the waiting queue and applies the router's decision locally, so a per-stage +disagreement means the stages admit different numbers of requests and diverge -- +a hang, not a wrong number. + +Before this change only the last stage marked ``GENERATION_TO_COMPLETE`` for +generation requests (``_executor_loop_pp``), which is why the correction was gated +off whenever ``pp_size > 1``. ``_forward_step_inter_pp`` now makes the same call at +the structurally identical point, so the counts agree. + +The marking is safe to replicate because the predicate +(``LlmRequest::willCompleteNextIteration``) is pure arithmetic on token counts, and +those counts are replicated to every stage in the same iteration: the last stage's +sample state is ring-broadcast and applied on all ranks by +``_handle_executed_batch``, with the per-iteration batch count itself ring-broadcast +from rank 0. +""" + +import ast +import inspect +import textwrap +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor +from tensorrt_llm._torch.pyexecutor.scheduler.adp_router import count_retiring_requests + +pytestmark = pytest.mark.cpu_only + +# (num_generated_tokens, max_new_tokens): the replicated token counts every stage +# holds for one request. The first two retire on the next iteration, the rest do +# not -- so a stage that skips the marking under-counts by exactly two. +REPLICATED_TOKEN_COUNTS = [(15, 16), (7, 8), (3, 16), (1, 8), (15, 64)] + + +class _StageLocalRequest: + """One pipeline stage's own object for a request. + + Each stage has a distinct ``LlmRequest`` instance; only the token counts are + replicated. Mirroring that here is the point of the test: the marking must be + reproducible from the replicated fields alone, with no reference to sampler + state that lives on the last stage. + """ + + def __init__(self, num_generated_tokens: int, max_new_tokens: int, tokens_per_iteration: int = 1): + self.num_generated_tokens = num_generated_tokens + self.max_new_tokens = max_new_tokens + self.tokens_per_iteration = tokens_per_iteration + self.state = LlmRequestState.GENERATION_IN_PROGRESS + self.exclude_last_generation_logits = True + + def will_complete_next_iteration(self) -> bool: + # Mirrors LlmRequest::willCompleteNextIteration (llmRequest.h): pure + # arithmetic on counts that are identical on every stage. No EOS check, no + # stop words, no sampler state. + return self.num_generated_tokens + self.tokens_per_iteration >= self.max_new_tokens + + def set_exclude_last_generation_logits(self, value: bool) -> None: + self.exclude_last_generation_logits = value + + +def _stage_local_batch(): + return [_StageLocalRequest(generated, limit) for generated, limit in REPLICATED_TOKEN_COUNTS] + + +def _mark(requests) -> None: + """Run the real marking method against a stub ``self``. + + The method reads nothing off ``self``, which is exactly why it can be called + from both loops without any stage-local context. + """ + PyExecutor._update_generation_requests_that_will_complete_next_iteration( + SimpleNamespace(), requests) + + +EXPECTED_RETIRING = 2 + + +def test_the_batch_actually_contains_retiring_requests(): + """Anti-vacuity: without this the consistency assertions hold trivially at 0.""" + batch = _stage_local_batch() + _mark(batch) + assert count_retiring_requests(batch) == EXPECTED_RETIRING + assert EXPECTED_RETIRING < len(batch) + + +@pytest.mark.parametrize("pp_size", [2, 4]) +def test_every_stage_derives_the_same_retiring_count(pp_size): + """All stages mark: the counts agree, which is what makes the correction safe.""" + stages = [_stage_local_batch() for _ in range(pp_size)] + for batch in stages: + _mark(batch) + + counts = [count_retiring_requests(batch) for batch in stages] + assert counts == [EXPECTED_RETIRING] * pp_size + + # Stronger than the count: the same *requests* are marked on every stage, so + # each rank subtracts the same load, not merely the same amount of it. + marked = { + tuple(i for i, req in enumerate(batch) if req.state == LlmRequestState.GENERATION_TO_COMPLETE) + for batch in stages + } + assert len(marked) == 1 + + +@pytest.mark.parametrize("pp_size", [2, 4]) +def test_last_stage_only_marking_is_detectably_inconsistent(pp_size): + """Negative control: the pre-change behaviour must fail this test's premise. + + If only the last stage marks, the stages disagree by exactly the number of + retiring requests. Without this case a test that merely asserts "the counts + agree" would keep passing after a regression that removes the marking from + *every* stage. + """ + stages = [_stage_local_batch() for _ in range(pp_size)] + _mark(stages[-1]) + + counts = [count_retiring_requests(batch) for batch in stages] + assert counts[-1] == EXPECTED_RETIRING + assert counts[:-1] == [0] * (pp_size - 1) + assert len(set(counts)) > 1 + + +def test_marking_is_idempotent_across_repeated_stage_passes(): + """A stage that marks twice in an iteration must not drift from one that marks once. + + ``_forward_step_inter_pp`` runs once per micro-batch, and a request can appear + in consecutive micro-batches, so the operation has to be a no-op on an + already-marked request. + """ + once, twice = _stage_local_batch(), _stage_local_batch() + _mark(once) + _mark(twice) + _mark(twice) + assert count_retiring_requests(once) == count_retiring_requests(twice) + + +def test_already_complete_requests_are_left_alone(): + """GENERATION_COMPLETE is terminal; re-marking it would resurrect a torn-down + request into the retiring set on one stage only.""" + batch = _stage_local_batch() + batch[0].state = LlmRequestState.GENERATION_COMPLETE + _mark(batch) + assert batch[0].state == LlmRequestState.GENERATION_COMPLETE + assert count_retiring_requests(batch) == EXPECTED_RETIRING - 1 + + +def _calls_in(func) -> set: + """Names of every method called on ``self`` inside ``func``.""" + tree = ast.parse(textwrap.dedent(inspect.getsource(func))) + names = set() + for node in ast.walk(tree): + if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute): + names.add(node.func.attr) + return names + + +def test_inter_pp_forward_marks_retiring_requests(): + """The structural half of the fix, and the part a unit test can pin. + + The stage-local arithmetic above is only reached if the non-last-stage path + actually calls the marking method. This is the assertion that fails if that + call is dropped -- i.e. it is what makes the seat-headroom PP cell honest. + """ + calls = _calls_in(PyExecutor._forward_step_inter_pp) + assert "_update_generation_requests_that_will_complete_next_iteration" in calls + # Same point in the sequence as the last stage: right after the state update. + assert "_update_request_states" in calls + + +def test_marking_stays_paired_with_the_state_update_on_both_paths(): + """Both loops must mark, and neither may mark without first updating state. + + The count is only rank-consistent if the two operations stay adjacent: marking + before ``_update_request_states`` would read token counts from before this + micro-batch and disagree with a stage that marks after. + """ + for func in (PyExecutor._forward_step_inter_pp, PyExecutor._executor_loop_pp): + source = inspect.getsource(func) + assert "_update_generation_requests_that_will_complete_next_iteration" in source, ( + f"{func.__name__} does not mark retiring generation requests; the ADP " + "router's admission correction would then diverge between stages") + assert source.index("_update_request_states") < source.index( + "_update_generation_requests_that_will_complete_next_iteration") From 85b9ae1eca9a1aa570e4bf4facc0f93ff3277dbf Mon Sep 17 00:00:00 2001 From: Chenfei Zhang Date: Mon, 7 Sep 2026 08:55:49 -0700 Subject: [PATCH 11/22] [https://nvbugs/6627795][feat] extend the attention-DP overlap seat headroom to pipeline parallelism With the retiring-request count now derived identically on every pipeline stage, the two reasons the headroom gate excluded PP are both gone: * the sizing can express it -- compute_max_num_sequences is additive, so the pool is (pp_size + 1) * max_batch_size rather than pp_size * max_batch_size; * the consumer is rank-consistent -- ADPRouter's correction subtracts the same requests on every stage, so the stages admit the same number of requests. This is the one behaviour change in the series, and the only cell of the sizing table whose value moves. It costs +1/pp_size seats: +25% at pp_size == 4, against +100% had the overlap term been multiplicative. Growing the pool cannot over-admit. ModelEngine.get_max_num_sequences() returns mapping.pp_size * batch_size and is computed independently of the headroom flag; py_executor sets max_num_active_requests from it and thereafter only ever reduces it. So the extra seats are headroom for leases already held, not extra admission capacity -- which is also why the capacity scheduler's own budget deliberately stays at max_batch_size * pp_size. should_enable_adp_dummy_fixes stays non-PP. It is an independent concern with an independent failure mode: the ADP dummy is a singleton fixed request ID while pp_size micro-batches are in flight, and _finalize_adp_dummy_allocation is never called from _executor_loop_pp, so a skipped iteration leaks the dummy. Widening both gates in one change would conflate them; with this one left alone the PP dummy path behaves exactly as it does today. Hybrid/SSM architectures remain excluded, since MambaHybridCacheManagerV2 still sizes its state-index pool from max_batch_size alone. The two gate rows that pinned the PP exclusion as intended contract are updated here rather than in the commit that introduced the additive coefficient, so that the behaviour change and the tests that assert it move together. Signed-off-by: Chenfei Zhang --- tensorrt_llm/_torch/pyexecutor/_util.py | 27 ++++++++++------- .../_torch/pyexecutor/scheduler/adp_router.py | 30 ++++++++++++------- .../_torch/executor/test_seq_slot_sizing.py | 13 ++++---- 3 files changed, 42 insertions(+), 28 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index af428e7fcb29..6c3fd677f163 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -3157,7 +3157,17 @@ def resolve_max_num_sequences(model_engine, def should_enable_adp_dummy_fixes(mapping: Mapping) -> bool: - """Enable transactional ADP dummy handling while PP remains follow-up.""" + """Enable transactional ADP dummy handling; still non-PP only. + + Independent of ``should_enable_adp_overlap_seq_slot_headroom``, which *is* now + enabled under pipeline parallelism. The dummy path has its own PP obstacles, + unrelated to seat capacity: the ADP dummy is a singleton fixed request ID + while ``pp_size`` micro-batches are in flight, and + ``_finalize_adp_dummy_allocation`` is never called from ``_executor_loop_pp``, + so a skipped iteration leaks the dummy. Widening this gate is a separate + change with a separate failure mode, so the two are deliberately not moved + together. + """ return not mapping.has_pp() @@ -3199,13 +3209,11 @@ def should_enable_adp_overlap_seq_slot_headroom( is a property of overlap plus ADP admission, and was measured on an aggregated context-only run with no cache transceiver configured. - Pipeline parallelism is still excluded here, but no longer because of the - sizing: compute_max_num_sequences now expresses the headroom additively, so - ``(pp_size + 1) * max_batch_size`` is a well-defined pool under PP. What is - missing is the *consumer*: ADPRouter's retiring-request correction requires - every pipeline stage to agree on which requests are retiring, and today only - the last stage marks generation requests GENERATION_TO_COMPLETE. Opening the - gate before that is fixed would allocate headroom no rank ever spends. + Pipeline parallelism is included. The extra seats are additive in pp_size + (see compute_max_num_sequences), and they are only *usable* because the + retiring-request count is now derived identically on every pipeline stage -- + ADPRouter.exclude_retiring_requests documents that contract. Enabling this + gate without that fix would allocate headroom no rank ever spends. Hybrid (Mamba/SSM) architectures are excluded. MambaHybridCacheManagerV2 sizes its state-index pool from max_batch_size alone @@ -3217,8 +3225,7 @@ def should_enable_adp_overlap_seq_slot_headroom( """ if is_hybrid: return False - return (mapping.enable_attention_dp and not mapping.has_pp() - and not disable_overlap_scheduler) + return (mapping.enable_attention_dp and not disable_overlap_scheduler) def create_py_executor_instance( diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py index 43499ed294ed..6eeb8eac116d 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py @@ -244,16 +244,26 @@ def __init__(self, dist: Distributed, has_seq_slot_headroom: bool = True): # -- as ``not dist.mapping.has_pp()`` used to -- is exactly how the # router comes to credit a rank with seats the engine never # allocated. - # 2. Pipeline parallelism is excluded by that same flag, and the reason - # is no longer the sizing -- the seat pool is well defined under PP at - # ``(pp_size + 1) * max_batch_size``. It is that the stages must agree - # on *which* requests are retiring: each rank pops from its own copy - # of the waiting queue, so a per-stage disagreement makes them admit - # different numbers of requests and diverge. Only the last stage marks - # ``GENERATION_TO_COMPLETE`` for generation requests today. Note the - # *context* path in ``_update_request_states_tp`` already evaluates the - # same predicate on every rank, so the asymmetry is only ever in the - # generation path. + # 2. Pipeline parallelism is now in scope: the seat pool is + # ``(pp_size + 1) * max_batch_size`` (additive, because pipeline depth + # is already paid for by the ``pp_size`` term and the overlap deferral + # is one iteration on top of it), and every rank must agree on which + # requests are retiring or the stages admit different numbers of + # requests and diverge. Previously only the last stage marked + # ``GENERATION_TO_COMPLETE`` for generation requests; + # ``_forward_step_inter_pp`` now makes the same call the last stage + # makes, at the structurally identical point (right after + # ``_update_request_states``, under the same overlap guard). That is + # sound rather than merely symmetric because the predicate + # (``LlmRequest::willCompleteNextIteration``) is pure arithmetic on + # token counts that are replicated to every stage in the same + # iteration by ``_ring_broadcast_sample_state`` / + # ``_handle_executed_batch``, with the per-iteration batch count itself + # ring-broadcast from rank 0. + # + # Note the *context* path in ``_update_request_states_tp`` already + # evaluated the same predicate on every rank, so the asymmetry was only + # ever in the generation path. # # Not additionally gated on ``disable_overlap_scheduler`` here: the flag # already is, and without overlap the retire is not deferred, so no diff --git a/tests/unittest/_torch/executor/test_seq_slot_sizing.py b/tests/unittest/_torch/executor/test_seq_slot_sizing.py index d3ba178b985f..bd7a20139f03 100644 --- a/tests/unittest/_torch/executor/test_seq_slot_sizing.py +++ b/tests/unittest/_torch/executor/test_seq_slot_sizing.py @@ -81,14 +81,11 @@ (True, 1, False, False, True), (False, 1, False, False, False), (True, 1, True, False, False), - # Pipeline parallelism is still out of scope, but no longer because the - # sizing cannot express it -- compute_max_num_sequences is additive, so - # (pp_size + 1) * max_batch_size is well defined. The missing piece is the - # consumer: only the last pipeline stage marks generation requests - # GENERATION_TO_COMPLETE, so the ADP router's retiring-request correction - # would not be rank-consistent. - (True, 2, False, False, False), - (True, 4, False, False, False), + # Pipeline parallelism is in scope. The extra seats are additive in + # pp_size and are only spendable because the ADP router now derives the + # retiring-request count identically on every stage. + (True, 2, False, False, True), + (True, 4, False, False, True), # Hybrid/SSM architectures are excluded: MambaHybridCacheManagerV2 sizes # its state-index pool from max_batch_size alone, so an extra seat would # have no state slot behind it. From 166cf0ffeedbfae57c3aeb1aa77f8d4fce7e44bc Mon Sep 17 00:00:00 2001 From: Chenfei Zhang Date: Mon, 7 Sep 2026 09:50:51 -0700 Subject: [PATCH 12/22] [https://nvbugs/6627795][chore] satisfy yapf, ruff and ruff-format on the changed files Pre-commit runs over the files changed against the base, so a file's existing formatting only comes into scope once a commit touches it. Three of these four were already non-conforming before this branch and were never checked; the fourth is a genuinely new import in the wrong sort position. speculative/utils.py yapf 0.43.0, one wrapped getattr call test_pp_retiring_rank_consistency.py ruff-format 0.9.4 test_seq_slot_sizing.py ruff-format 0.9.4 test_spec_slot_pool_sizing.py ruff I001 -- NoFreeSlotsError sort order No behaviour change, and not asserted by eye: for the first three the AST is identical to the parent commit both with and without docstrings, and for the fourth the multiset of import nodes and the entire non-import body are identical, the diff being a single import line moving up. The CPU negative-control ladder result therefore still describes this tree. Tool versions match the pinned pre-commit revs exactly (yapf v0.43.0, ruff v0.9.4), so this is what CI will compute rather than an approximation of it. Signed-off-by: Chenfei Zhang --- tensorrt_llm/_torch/speculative/utils.py | 3 ++- .../test_pp_retiring_rank_consistency.py | 17 ++++++++++++----- .../_torch/executor/test_seq_slot_sizing.py | 4 +--- .../speculative/test_spec_slot_pool_sizing.py | 2 +- 4 files changed, 16 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 2c5cd4f0c467..02a0faac355a 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -586,7 +586,8 @@ def seat_pool_or_none(model_engine) -> Optional[int]: pool agrees with the metadata about which number it is indexed by. Returning None (headroom off) preserves the established max_batch_size sizing. """ - if not getattr(model_engine, "_enable_adp_overlap_seq_slot_headroom", False): + if not getattr(model_engine, "_enable_adp_overlap_seq_slot_headroom", + False): return None return getattr(model_engine, "max_num_seq_slots", None) diff --git a/tests/unittest/_torch/executor/test_pp_retiring_rank_consistency.py b/tests/unittest/_torch/executor/test_pp_retiring_rank_consistency.py index 47eb13e977e8..fe27165ad059 100644 --- a/tests/unittest/_torch/executor/test_pp_retiring_rank_consistency.py +++ b/tests/unittest/_torch/executor/test_pp_retiring_rank_consistency.py @@ -50,7 +50,9 @@ class _StageLocalRequest: state that lives on the last stage. """ - def __init__(self, num_generated_tokens: int, max_new_tokens: int, tokens_per_iteration: int = 1): + def __init__( + self, num_generated_tokens: int, max_new_tokens: int, tokens_per_iteration: int = 1 + ): self.num_generated_tokens = num_generated_tokens self.max_new_tokens = max_new_tokens self.tokens_per_iteration = tokens_per_iteration @@ -78,7 +80,8 @@ def _mark(requests) -> None: from both loops without any stage-local context. """ PyExecutor._update_generation_requests_that_will_complete_next_iteration( - SimpleNamespace(), requests) + SimpleNamespace(), requests + ) EXPECTED_RETIRING = 2 @@ -105,7 +108,9 @@ def test_every_stage_derives_the_same_retiring_count(pp_size): # Stronger than the count: the same *requests* are marked on every stage, so # each rank subtracts the same load, not merely the same amount of it. marked = { - tuple(i for i, req in enumerate(batch) if req.state == LlmRequestState.GENERATION_TO_COMPLETE) + tuple( + i for i, req in enumerate(batch) if req.state == LlmRequestState.GENERATION_TO_COMPLETE + ) for batch in stages } assert len(marked) == 1 @@ -187,6 +192,8 @@ def test_marking_stays_paired_with_the_state_update_on_both_paths(): source = inspect.getsource(func) assert "_update_generation_requests_that_will_complete_next_iteration" in source, ( f"{func.__name__} does not mark retiring generation requests; the ADP " - "router's admission correction would then diverge between stages") + "router's admission correction would then diverge between stages" + ) assert source.index("_update_request_states") < source.index( - "_update_generation_requests_that_will_complete_next_iteration") + "_update_generation_requests_that_will_complete_next_iteration" + ) diff --git a/tests/unittest/_torch/executor/test_seq_slot_sizing.py b/tests/unittest/_torch/executor/test_seq_slot_sizing.py index bd7a20139f03..2cf72ce4a066 100644 --- a/tests/unittest/_torch/executor/test_seq_slot_sizing.py +++ b/tests/unittest/_torch/executor/test_seq_slot_sizing.py @@ -104,9 +104,7 @@ def test_adp_overlap_seq_slot_headroom_gate( ) assert ( - should_enable_adp_overlap_seq_slot_headroom( - mapping, disable_overlap, is_hybrid=is_hybrid - ) + should_enable_adp_overlap_seq_slot_headroom(mapping, disable_overlap, is_hybrid=is_hybrid) is expected ) diff --git a/tests/unittest/_torch/speculative/test_spec_slot_pool_sizing.py b/tests/unittest/_torch/speculative/test_spec_slot_pool_sizing.py index b3591513c00f..ba77148da059 100644 --- a/tests/unittest/_torch/speculative/test_spec_slot_pool_sizing.py +++ b/tests/unittest/_torch/speculative/test_spec_slot_pool_sizing.py @@ -28,6 +28,7 @@ import pytest import torch +from tensorrt_llm._torch.pyexecutor.resource_manager import NoFreeSlotsError from tensorrt_llm._torch.speculative.eagle3 import ( Eagle3OneModelDynamicTreeResourceManager, Eagle3ResourceManager, @@ -36,7 +37,6 @@ 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.pyexecutor.resource_manager import NoFreeSlotsError from tensorrt_llm._torch.speculative.utils import ( _build_spec_metadata, get_spec_drafter, From 4030e8e1bea075ea4d018e971085b8041f3ff913 Mon Sep 17 00:00:00 2001 From: Chenfei Zhang Date: Mon, 7 Sep 2026 21:50:24 -0700 Subject: [PATCH 13/22] Revert "[https://nvbugs/6627795][feat] extend the attention-DP overlap seat headroom to pipeline parallelism" This reverts commit 85b9ae1eca9a1aa570e4bf4facc0f93ff3277dbf. Signed-off-by: Chenfei Zhang --- tensorrt_llm/_torch/pyexecutor/_util.py | 27 +++++++---------- .../_torch/pyexecutor/scheduler/adp_router.py | 30 +++++++------------ .../_torch/executor/test_seq_slot_sizing.py | 13 ++++---- 3 files changed, 28 insertions(+), 42 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 3118c2f41ae0..51ab19f2bbee 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -3188,17 +3188,7 @@ def resolve_max_num_sequences(model_engine, def should_enable_adp_dummy_fixes(mapping: Mapping) -> bool: - """Enable transactional ADP dummy handling; still non-PP only. - - Independent of ``should_enable_adp_overlap_seq_slot_headroom``, which *is* now - enabled under pipeline parallelism. The dummy path has its own PP obstacles, - unrelated to seat capacity: the ADP dummy is a singleton fixed request ID - while ``pp_size`` micro-batches are in flight, and - ``_finalize_adp_dummy_allocation`` is never called from ``_executor_loop_pp``, - so a skipped iteration leaks the dummy. Widening this gate is a separate - change with a separate failure mode, so the two are deliberately not moved - together. - """ + """Enable transactional ADP dummy handling while PP remains follow-up.""" return not mapping.has_pp() @@ -3240,11 +3230,13 @@ def should_enable_adp_overlap_seq_slot_headroom( is a property of overlap plus ADP admission, and was measured on an aggregated context-only run with no cache transceiver configured. - Pipeline parallelism is included. The extra seats are additive in pp_size - (see compute_max_num_sequences), and they are only *usable* because the - retiring-request count is now derived identically on every pipeline stage -- - ADPRouter.exclude_retiring_requests documents that contract. Enabling this - gate without that fix would allocate headroom no rank ever spends. + Pipeline parallelism is still excluded here, but no longer because of the + sizing: compute_max_num_sequences now expresses the headroom additively, so + ``(pp_size + 1) * max_batch_size`` is a well-defined pool under PP. What is + missing is the *consumer*: ADPRouter's retiring-request correction requires + every pipeline stage to agree on which requests are retiring, and today only + the last stage marks generation requests GENERATION_TO_COMPLETE. Opening the + gate before that is fixed would allocate headroom no rank ever spends. Hybrid (Mamba/SSM) architectures are excluded. MambaHybridCacheManagerV2 sizes its state-index pool from max_batch_size alone @@ -3256,7 +3248,8 @@ def should_enable_adp_overlap_seq_slot_headroom( """ if is_hybrid: return False - return (mapping.enable_attention_dp and not disable_overlap_scheduler) + return (mapping.enable_attention_dp and not mapping.has_pp() + and not disable_overlap_scheduler) def create_py_executor_instance( diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py index 6eeb8eac116d..43499ed294ed 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py @@ -244,26 +244,16 @@ def __init__(self, dist: Distributed, has_seq_slot_headroom: bool = True): # -- as ``not dist.mapping.has_pp()`` used to -- is exactly how the # router comes to credit a rank with seats the engine never # allocated. - # 2. Pipeline parallelism is now in scope: the seat pool is - # ``(pp_size + 1) * max_batch_size`` (additive, because pipeline depth - # is already paid for by the ``pp_size`` term and the overlap deferral - # is one iteration on top of it), and every rank must agree on which - # requests are retiring or the stages admit different numbers of - # requests and diverge. Previously only the last stage marked - # ``GENERATION_TO_COMPLETE`` for generation requests; - # ``_forward_step_inter_pp`` now makes the same call the last stage - # makes, at the structurally identical point (right after - # ``_update_request_states``, under the same overlap guard). That is - # sound rather than merely symmetric because the predicate - # (``LlmRequest::willCompleteNextIteration``) is pure arithmetic on - # token counts that are replicated to every stage in the same - # iteration by ``_ring_broadcast_sample_state`` / - # ``_handle_executed_batch``, with the per-iteration batch count itself - # ring-broadcast from rank 0. - # - # Note the *context* path in ``_update_request_states_tp`` already - # evaluated the same predicate on every rank, so the asymmetry was only - # ever in the generation path. + # 2. Pipeline parallelism is excluded by that same flag, and the reason + # is no longer the sizing -- the seat pool is well defined under PP at + # ``(pp_size + 1) * max_batch_size``. It is that the stages must agree + # on *which* requests are retiring: each rank pops from its own copy + # of the waiting queue, so a per-stage disagreement makes them admit + # different numbers of requests and diverge. Only the last stage marks + # ``GENERATION_TO_COMPLETE`` for generation requests today. Note the + # *context* path in ``_update_request_states_tp`` already evaluates the + # same predicate on every rank, so the asymmetry is only ever in the + # generation path. # # Not additionally gated on ``disable_overlap_scheduler`` here: the flag # already is, and without overlap the retire is not deferred, so no diff --git a/tests/unittest/_torch/executor/test_seq_slot_sizing.py b/tests/unittest/_torch/executor/test_seq_slot_sizing.py index 2cf72ce4a066..01b95fa6f126 100644 --- a/tests/unittest/_torch/executor/test_seq_slot_sizing.py +++ b/tests/unittest/_torch/executor/test_seq_slot_sizing.py @@ -81,11 +81,14 @@ (True, 1, False, False, True), (False, 1, False, False, False), (True, 1, True, False, False), - # Pipeline parallelism is in scope. The extra seats are additive in - # pp_size and are only spendable because the ADP router now derives the - # retiring-request count identically on every stage. - (True, 2, False, False, True), - (True, 4, False, False, True), + # Pipeline parallelism is still out of scope, but no longer because the + # sizing cannot express it -- compute_max_num_sequences is additive, so + # (pp_size + 1) * max_batch_size is well defined. The missing piece is the + # consumer: only the last pipeline stage marks generation requests + # GENERATION_TO_COMPLETE, so the ADP router's retiring-request correction + # would not be rank-consistent. + (True, 2, False, False, False), + (True, 4, False, False, False), # Hybrid/SSM architectures are excluded: MambaHybridCacheManagerV2 sizes # its state-index pool from max_batch_size alone, so an extra seat would # have no state slot behind it. From f498e0f49b17d5542aecf067ccf44da41d90e935 Mon Sep 17 00:00:00 2001 From: Chenfei Zhang Date: Mon, 7 Sep 2026 21:50:34 -0700 Subject: [PATCH 14/22] Revert "[https://nvbugs/6627795][fix] derive the retiring-request count identically on every pipeline stage" This reverts commit ac72cda1ba9f21ad88fe16e4b6bc1e22f5777fca. Signed-off-by: Chenfei Zhang --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 13 -- .../test_pp_retiring_rank_consistency.py | 199 ------------------ 2 files changed, 212 deletions(-) delete mode 100644 tests/unittest/_torch/executor/test_pp_retiring_rank_consistency.py diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index d59b38ea7ab0..545342c725dd 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -5621,19 +5621,6 @@ def _forward_step_inter_pp(self, sampler_event = torch.cuda.Event() sampler_event.record() self._update_request_states(scheduled_batch) - # Mark retiring generation requests here too, not only on the last - # pipeline stage (_executor_loop_pp). The predicate is pure arithmetic on - # token counts (LlmRequest::willCompleteNextIteration) and those counts are - # replicated to every stage in the same iteration: the last stage's sample - # state is ring-broadcast (_ring_broadcast_sample_state) and applied on all - # ranks by _handle_executed_batch, with the per-iteration batch count - # itself ring-broadcast from rank 0. Without this call the stages disagree - # about which requests are retiring, and ADPRouter's admission correction - # -- which subtracts them from each rank's load -- would let the stages - # admit different numbers of requests and diverge. - if not self.disable_overlap_scheduler: - self._update_generation_requests_that_will_complete_next_iteration( - scheduled_batch.generation_requests) sampling_requests = scheduled_batch.context_requests_last_chunk + scheduled_batch.generation_requests return self.sampler.SampleState( requests=sampling_requests, diff --git a/tests/unittest/_torch/executor/test_pp_retiring_rank_consistency.py b/tests/unittest/_torch/executor/test_pp_retiring_rank_consistency.py deleted file mode 100644 index fe27165ad059..000000000000 --- a/tests/unittest/_torch/executor/test_pp_retiring_rank_consistency.py +++ /dev/null @@ -1,199 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -"""Every pipeline stage must derive the same retiring-request count. - -The ADP router subtracts retiring requests from each rank's reported load -(nvbug-6627795). Under pipeline parallelism that correction is only safe if every -stage agrees on *which* requests are retiring: each rank pops from its own copy of -the waiting queue and applies the router's decision locally, so a per-stage -disagreement means the stages admit different numbers of requests and diverge -- -a hang, not a wrong number. - -Before this change only the last stage marked ``GENERATION_TO_COMPLETE`` for -generation requests (``_executor_loop_pp``), which is why the correction was gated -off whenever ``pp_size > 1``. ``_forward_step_inter_pp`` now makes the same call at -the structurally identical point, so the counts agree. - -The marking is safe to replicate because the predicate -(``LlmRequest::willCompleteNextIteration``) is pure arithmetic on token counts, and -those counts are replicated to every stage in the same iteration: the last stage's -sample state is ring-broadcast and applied on all ranks by -``_handle_executed_batch``, with the per-iteration batch count itself ring-broadcast -from rank 0. -""" - -import ast -import inspect -import textwrap -from types import SimpleNamespace - -import pytest - -from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState -from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor -from tensorrt_llm._torch.pyexecutor.scheduler.adp_router import count_retiring_requests - -pytestmark = pytest.mark.cpu_only - -# (num_generated_tokens, max_new_tokens): the replicated token counts every stage -# holds for one request. The first two retire on the next iteration, the rest do -# not -- so a stage that skips the marking under-counts by exactly two. -REPLICATED_TOKEN_COUNTS = [(15, 16), (7, 8), (3, 16), (1, 8), (15, 64)] - - -class _StageLocalRequest: - """One pipeline stage's own object for a request. - - Each stage has a distinct ``LlmRequest`` instance; only the token counts are - replicated. Mirroring that here is the point of the test: the marking must be - reproducible from the replicated fields alone, with no reference to sampler - state that lives on the last stage. - """ - - def __init__( - self, num_generated_tokens: int, max_new_tokens: int, tokens_per_iteration: int = 1 - ): - self.num_generated_tokens = num_generated_tokens - self.max_new_tokens = max_new_tokens - self.tokens_per_iteration = tokens_per_iteration - self.state = LlmRequestState.GENERATION_IN_PROGRESS - self.exclude_last_generation_logits = True - - def will_complete_next_iteration(self) -> bool: - # Mirrors LlmRequest::willCompleteNextIteration (llmRequest.h): pure - # arithmetic on counts that are identical on every stage. No EOS check, no - # stop words, no sampler state. - return self.num_generated_tokens + self.tokens_per_iteration >= self.max_new_tokens - - def set_exclude_last_generation_logits(self, value: bool) -> None: - self.exclude_last_generation_logits = value - - -def _stage_local_batch(): - return [_StageLocalRequest(generated, limit) for generated, limit in REPLICATED_TOKEN_COUNTS] - - -def _mark(requests) -> None: - """Run the real marking method against a stub ``self``. - - The method reads nothing off ``self``, which is exactly why it can be called - from both loops without any stage-local context. - """ - PyExecutor._update_generation_requests_that_will_complete_next_iteration( - SimpleNamespace(), requests - ) - - -EXPECTED_RETIRING = 2 - - -def test_the_batch_actually_contains_retiring_requests(): - """Anti-vacuity: without this the consistency assertions hold trivially at 0.""" - batch = _stage_local_batch() - _mark(batch) - assert count_retiring_requests(batch) == EXPECTED_RETIRING - assert EXPECTED_RETIRING < len(batch) - - -@pytest.mark.parametrize("pp_size", [2, 4]) -def test_every_stage_derives_the_same_retiring_count(pp_size): - """All stages mark: the counts agree, which is what makes the correction safe.""" - stages = [_stage_local_batch() for _ in range(pp_size)] - for batch in stages: - _mark(batch) - - counts = [count_retiring_requests(batch) for batch in stages] - assert counts == [EXPECTED_RETIRING] * pp_size - - # Stronger than the count: the same *requests* are marked on every stage, so - # each rank subtracts the same load, not merely the same amount of it. - marked = { - tuple( - i for i, req in enumerate(batch) if req.state == LlmRequestState.GENERATION_TO_COMPLETE - ) - for batch in stages - } - assert len(marked) == 1 - - -@pytest.mark.parametrize("pp_size", [2, 4]) -def test_last_stage_only_marking_is_detectably_inconsistent(pp_size): - """Negative control: the pre-change behaviour must fail this test's premise. - - If only the last stage marks, the stages disagree by exactly the number of - retiring requests. Without this case a test that merely asserts "the counts - agree" would keep passing after a regression that removes the marking from - *every* stage. - """ - stages = [_stage_local_batch() for _ in range(pp_size)] - _mark(stages[-1]) - - counts = [count_retiring_requests(batch) for batch in stages] - assert counts[-1] == EXPECTED_RETIRING - assert counts[:-1] == [0] * (pp_size - 1) - assert len(set(counts)) > 1 - - -def test_marking_is_idempotent_across_repeated_stage_passes(): - """A stage that marks twice in an iteration must not drift from one that marks once. - - ``_forward_step_inter_pp`` runs once per micro-batch, and a request can appear - in consecutive micro-batches, so the operation has to be a no-op on an - already-marked request. - """ - once, twice = _stage_local_batch(), _stage_local_batch() - _mark(once) - _mark(twice) - _mark(twice) - assert count_retiring_requests(once) == count_retiring_requests(twice) - - -def test_already_complete_requests_are_left_alone(): - """GENERATION_COMPLETE is terminal; re-marking it would resurrect a torn-down - request into the retiring set on one stage only.""" - batch = _stage_local_batch() - batch[0].state = LlmRequestState.GENERATION_COMPLETE - _mark(batch) - assert batch[0].state == LlmRequestState.GENERATION_COMPLETE - assert count_retiring_requests(batch) == EXPECTED_RETIRING - 1 - - -def _calls_in(func) -> set: - """Names of every method called on ``self`` inside ``func``.""" - tree = ast.parse(textwrap.dedent(inspect.getsource(func))) - names = set() - for node in ast.walk(tree): - if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute): - names.add(node.func.attr) - return names - - -def test_inter_pp_forward_marks_retiring_requests(): - """The structural half of the fix, and the part a unit test can pin. - - The stage-local arithmetic above is only reached if the non-last-stage path - actually calls the marking method. This is the assertion that fails if that - call is dropped -- i.e. it is what makes the seat-headroom PP cell honest. - """ - calls = _calls_in(PyExecutor._forward_step_inter_pp) - assert "_update_generation_requests_that_will_complete_next_iteration" in calls - # Same point in the sequence as the last stage: right after the state update. - assert "_update_request_states" in calls - - -def test_marking_stays_paired_with_the_state_update_on_both_paths(): - """Both loops must mark, and neither may mark without first updating state. - - The count is only rank-consistent if the two operations stay adjacent: marking - before ``_update_request_states`` would read token counts from before this - micro-batch and disagree with a stage that marks after. - """ - for func in (PyExecutor._forward_step_inter_pp, PyExecutor._executor_loop_pp): - source = inspect.getsource(func) - assert "_update_generation_requests_that_will_complete_next_iteration" in source, ( - f"{func.__name__} does not mark retiring generation requests; the ADP " - "router's admission correction would then diverge between stages" - ) - assert source.index("_update_request_states") < source.index( - "_update_generation_requests_that_will_complete_next_iteration" - ) From 8ffd77b416bf022c33bbb603c58946fbdc78775e Mon Sep 17 00:00:00 2001 From: Chenfei Zhang Date: Mon, 7 Sep 2026 21:55:20 -0700 Subject: [PATCH 15/22] [https://nvbugs/6627795][fix] make the new seat-pool plumbing tolerate test doubles Four unit tests in CPU-Generic-x86-1 failed on the previous head, all in pre-existing files this branch does not otherwise touch. Each is the same mistake in a different guise: new code that reads an attribute the real object always has, from a hand-built object that does not. * validate_seq_slot_pool_covers_admission ordered max_admissible_sequences against an int. A Mock auto-creates the attribute, so "absent" arrived as a Mock and "== int" was False, which fell through to "< int" and raised TypeError in test_factory_forwards_v2_scheduler_gates. Non-integral now means the same thing as absent -- this manager did not opt into the check. * resolve_max_num_sequences took disable_overlap_scheduler, so both call sites evaluated llm_args.disable_overlap_scheduler eagerly as an argument -- including on the two branches that never use it. A caller passing max_num_sequences explicitly with a lighter args object then died on attribute access rather than short-circuiting. It now takes llm_args whole and reads the field only inside the fallback, and a new test drives the two short-circuit branches with an args object that raises on any attribute read; asserting on the return value alone could not distinguish "not used" from "used and happened to agree". * the draft KV cache manager is now sized from the target engine's published seat pool, which means _create_kv_cache_manager reads self._model_engine. One hand-built KvCacheCreator in the estimation tests did not set it; setting it there is already the convention in three of the four files that build a creator via object.__new__. Signed-off-by: Chenfei Zhang --- tensorrt_llm/_torch/pyexecutor/_util.py | 37 ++++++++++++---- .../kv_cache/test_kv_cache_estimation.py | 5 +++ .../_torch/executor/test_seq_slot_sizing.py | 44 ++++++++++++++++++- 3 files changed, 77 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 51ab19f2bbee..64dcf09a10b4 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -3124,13 +3124,15 @@ def validate_seq_slot_pool_covers_admission(max_num_sequences: int, killing the rank mid-collective. Managers that do not publish ``max_admissible_sequences`` (V1, hybrid) are - skipped rather than guessed at; adding the attribute is how a manager opts - into the check. + skipped rather than guessed at; publishing the attribute *as an int* is how a + manager opts into the check. Anything non-integral means the same thing as + absent -- this manager did not opt in -- so the check is skipped rather than + attempted against a value it cannot order. """ if kv_cache_manager is None: return admissible = getattr(kv_cache_manager, "max_admissible_sequences", None) - if admissible is None: + if not isinstance(admissible, int): return if admissible == max_num_sequences: return @@ -3152,7 +3154,7 @@ def validate_seq_slot_pool_covers_admission(max_num_sequences: int, def resolve_max_num_sequences(model_engine, mapping: Mapping, max_batch_size: int, - disable_overlap_scheduler: bool, + llm_args, is_disagg: bool, max_num_sequences: Optional[int] = None) -> int: """Resolve the seat-pool size for a consumer, without re-deriving it. @@ -3170,6 +3172,12 @@ def resolve_max_num_sequences(model_engine, silently sized the sampler and the executor's ``SeqSlotManager`` below the index pool they share indices with, i.e. it reintroduced the very skew this number exists to eliminate, in the one code path that has no test coverage. + + ``llm_args`` is taken whole rather than as an already-read + ``disable_overlap_scheduler`` flag so that the fallback's inputs are read + only when the fallback runs. Passing the flag made the caller evaluate it + eagerly, which reintroduced a hard dependency on the field in the branch + that never uses it -- and broke callers holding a lighter args object. """ if max_num_sequences is not None: return max_num_sequences @@ -3181,7 +3189,7 @@ def resolve_max_num_sequences(model_engine, return compute_max_num_sequences( mapping, max_batch_size, - disable_overlap_scheduler, + llm_args.disable_overlap_scheduler, enable_overlap_headroom=getattr( model_engine, "_enable_adp_overlap_seq_slot_headroom", False), is_disagg=is_disagg) @@ -3238,6 +3246,19 @@ def should_enable_adp_overlap_seq_slot_headroom( the last stage marks generation requests GENERATION_TO_COMPLETE. Opening the gate before that is fixed would allocate headroom no rank ever spends. + Measured, so that the next attempt starts from the number rather than the + argument: a generation-bearing pp=2 + ADP + overlap A/B (GLM-5 NVFP4, 4-way + GB300, max_batch_size=4, 1920 requests per arm) produced *byte-identical* + scheduling on both arms -- mean scheduled batch 3.997 of a cap of 4, the same + per-iteration histogram, and zero index-pool exhaustions. The seats are + unspendable at pp>1 because admission is capped independently at + ``pp_size * max_batch_size`` by ``get_max_num_sequences()``, which is already + below the unpatched index pool of ``pp_size * max_batch_size + 1``: the + per-micro-batch batch size, not the seat pool, is the limiter. So the PP cell + needs a case where admission is the binding constraint before it can show a + benefit; the throughput difference between those two arms (-2.34%) is this + configuration's noise floor, not a result. + Hybrid (Mamba/SSM) architectures are excluded. MambaHybridCacheManagerV2 sizes its state-index pool from max_batch_size alone (mamba_cache_manager.py: state_index_capacity), with neither pp_size nor the @@ -3293,7 +3314,7 @@ def create_py_executor_instance( model_engine, mapping, max_batch_size, - llm_args.disable_overlap_scheduler, + llm_args, is_disagg, max_num_sequences=max_num_sequences) @@ -3739,8 +3760,8 @@ def instantiate_sampler( engine, mapping, max_batch_size, - llm_args.disable_overlap_scheduler, - is_disagg_enabled(llm_args.cache_transceiver_config), + llm_args, + is_disagg_enabled(getattr(llm_args, "cache_transceiver_config", None)), max_num_sequences=max_num_sequences) sampler_args = create_torch_sampler_args( 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..b93d29f0ad10 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,11 @@ def test_separate_one_model_draft_normalizes_target_pool_ratio() -> None: creator._is_disagg = False creator._mapping = Mock() creator._speculative_config = Mock() + # The draft manager is now sized from the *target* engine's published seat + # pool, so this hand-built creator has to carry the attribute the real one + # sets in __init__. None means "no pool published", which is the pre-existing + # behaviour this test asserts. + creator._model_engine = None effective_draft_config = Mock() effective_draft_config.pretrained_config.torch_dtype = "bfloat16" diff --git a/tests/unittest/_torch/executor/test_seq_slot_sizing.py b/tests/unittest/_torch/executor/test_seq_slot_sizing.py index 01b95fa6f126..84a7d8abd28e 100644 --- a/tests/unittest/_torch/executor/test_seq_slot_sizing.py +++ b/tests/unittest/_torch/executor/test_seq_slot_sizing.py @@ -205,13 +205,14 @@ def test_resolve_max_num_sequences_prefers_the_published_pool(explicit, engine_s _enable_adp_overlap_seq_slot_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, - False, + llm_args, False, max_num_sequences=explicit, ) @@ -219,6 +220,40 @@ def test_resolve_max_num_sequences_prefers_the_published_pool(explicit, engine_s ) +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(), + False, + max_num_sequences=24, + ) + == 24 + ) + # Branch 2: the engine's published pool. + assert resolve_max_num_sequences(engine_with_pool, mapping, 8, _Exploding(), False) == 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. @@ -280,6 +315,13 @@ def test_validator_is_two_sided(admissible): [ None, # non-generation models have no KV cache manager SimpleNamespace(), # V1 sizes its index pool by an unrelated rule + # A test double auto-creates every attribute, so "absent" reaches the + # validator as a non-integral value rather than as None. Ordering it + # against an int raises TypeError, which is a crash in an unrelated + # caller's test rather than a finding about this PR -- so non-integral + # has to mean the same thing as absent. + Mock(), + _FakeManager(None), ], ) def test_validator_skips_managers_that_do_not_publish_capacity(manager): From c1c619e9540ef3fdfc36e18c98167a5328656ce3 Mon Sep 17 00:00:00 2001 From: Chenfei Zhang Date: Tue, 8 Sep 2026 02:44:17 -0700 Subject: [PATCH 16/22] [https://nvbugs/6627795][fix] confine the disagg 2x to the KV index pool Review feedback: the disaggregation multiplier belongs to KVCacheManagerV2's IndexMapper capacity, not to the sequence-slot pool. 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 one once its transmission completes -- while admission stays at max_batch_size * pp_size either way. So the index pool legitimately runs ahead of the seat pool, and propagating the 2x into the seat pool bought nothing while doubling everything keyed by seat: sampler state, the eager [seats, draft_len, vocab] draft-probability tensors (~800 MB at 512 seats), the penalty tensors and the pinned-host block-offset tables. - drop the is_disagg term (and parameter) from compute_max_num_sequences and the pass-through parameter from resolve_max_num_sequences; the seat pool is now max_batch_size * pp_size, plus one micro-batch under ADP + overlap. - make validate_seq_slot_pool_covers_admission asymmetric instead of an equality: an index pool below the seat pool is always a bug (nvbug 6627795), above it is a bug only when aggregated, and expected under disaggregation. - KVCacheManagerV2's arithmetic is unchanged; its comments now say why the 2x is local to that pool. - tests: drop the is_disagg column from SIZING_CASES, add a signature guard so the parameter cannot come back, split the validator's two directions, and add the index-pool > seat-pool disagg pairing to the capacity cases. Signed-off-by: Chenfei Zhang --- tensorrt_llm/_torch/pyexecutor/_util.py | 114 ++++++++++-------- .../kv_cache/kv_cache_manager_v2.py | 38 +++--- .../_torch/pyexecutor/model_engine.py | 3 +- .../kv_cache/test_kv_cache_manager_v2.py | 34 ++++-- .../_torch/executor/test_seq_slot_sizing.py | 106 ++++++++++------ 5 files changed, 185 insertions(+), 110 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 64dcf09a10b4..41df8d8a5f9b 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -3063,17 +3063,16 @@ def is_disagg_enabled(cache_transceiver_config) -> bool: def compute_max_num_sequences(mapping: Mapping, max_batch_size: int, disable_overlap_scheduler: bool, - enable_overlap_headroom: bool = False, - is_disagg: bool = False) -> int: + enable_overlap_headroom: bool = False) -> int: """Size the sequence-slot pool (and the sampler state it indexes). - This is *the* definition of how many sequences can be simultaneously live. - Every pool keyed by live-request identity -- the sampler's per-slot state, - the guided decoder, the spec-decoding managers, KVCacheManagerV2's index - mapper -- must consume this number rather than re-deriving it from - ``max_batch_size``, because any two formulas that disagree are a bug: too - few index slots silently defers admitted requests (nvbug 6627795), too few - seats raises ``NoFreeSlotsError`` on the executor's event-loop thread. + This is *the* definition of how many sequences can be simultaneously + **seated**. Every pool keyed by seat identity -- the sampler's per-slot + state, the guided decoder, the spec-decoding managers -- must consume this + number rather than re-deriving it from ``max_batch_size``, because any two + formulas that disagree are a bug: too few seats raises ``NoFreeSlotsError`` + on the executor's event-loop thread, and a KV index pool narrower than the + seat pool silently defers admitted requests (nvbug 6627795). The terms are **additive, not multiplicative**: @@ -3090,38 +3089,50 @@ def compute_max_num_sequences(mapping: Mapping, the multiplicative reading costs ``+100%`` at ``pp_size == 4`` where the additive one costs ``+25%`` -- on pools that include eagerly-allocated ``[seats, draft_len, vocab]`` tensors. + + 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 -- so + the index pool legitimately runs ahead of the seat pool while admission + stays at ``max_batch_size * pp_size``. Propagating that factor into the seat + pool would double sampler state, ``[seats, draft_len, vocab]`` draft + probabilities and the pinned-host block-offset tables for leases that never + occupy a seat. """ # Pipeline depth: pp_size micro-batches are structurally in flight. num_seats = max_batch_size * mapping.pp_size if enable_overlap_headroom and not disable_overlap_scheduler: num_seats += max_batch_size - if is_disagg: - # KVCacheManagerV2 sizes its own index pool max_batch_size * pp_size * 2 - # under disagg (the lease outlives the request while the KV transfer - # drains), and the seat pool must cover the index pool. max() rather - # than a further multiplication because the disagg and overlap - # coefficients each cover one extra cohort of in-flight sequences -- - # they overlap rather than compose. - num_seats = max(num_seats, max_batch_size * mapping.pp_size * 2) return num_seats def validate_seq_slot_pool_covers_admission(max_num_sequences: int, - kv_cache_manager) -> None: + kv_cache_manager, + is_disagg: bool = False) -> None: """Fail at startup if the seat pool and the KV index pool disagree. - Both are leases held for a request's whole lifetime, so the two counts must - be *equal*: every sequence the executor can seat needs an index, and every - sequence the KV cache manager can index needs a seat. The check is - deliberately two-sided, because each direction has already shipped as a - separate bug and a one-sided ``>=`` guard is what let the first one through: - - * index pool < seat pool -- ``_create_kv_cache`` warns ``No free IndexMapper - slots``, returns ``None`` and the scheduler defers the request. Silent: - costs throughput, no error (nvbug 6627795). - * index pool > seat pool -- a request is admitted that cannot be seated and - ``SlotManager.add_slot`` raises on the executor's event-loop thread, - killing the rank mid-collective. + The invariant is ``index pool >= seat pool``, and the two bounds have + different characters -- which is why the check is asymmetric rather than a + plain 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, no error (nvbug 6627795). A + one-sided ``seats >= index pool`` guard is exactly what let this ship. + * index pool > seat pool -- a bug **when aggregated**, because there every + indexed sequence is also a seated one, so the surplus can only come from + one of the two numbers having been re-derived; a request would be admitted + that cannot be seated and ``SlotManager.add_slot`` would raise on the + executor's event-loop thread. Under disaggregation it is **by design**: a + request awaiting its KV transfer holds an index lease and no seat + (``SeqSlotManager.prepare_resources`` skips ``DISAGG_GENERATION_INIT``), + so the index pool carries a ``2x`` that the seat pool must not. + + Admission is bounded independently at ``max_batch_size * pp_size`` by + ``ModelEngine.get_max_num_sequences()`` in both cases, so the surplus is + index headroom, never extra concurrency. Managers that do not publish ``max_admissible_sequences`` (V1, hybrid) are skipped rather than guessed at; publishing the attribute *as an int* is how a @@ -3136,26 +3147,31 @@ def validate_seq_slot_pool_covers_admission(max_num_sequences: int, return if admissible == max_num_sequences: return - direction = ("smaller" if admissible < max_num_sequences else "larger") - consequence = ( - "admitted requests would be silently deferred one at a time " - "(nvbug 6627795)" if admissible < max_num_sequences else - "a request could be admitted with no sequence slot to seat it, and " - "SlotManager.add_slot would raise on the executor's event loop") + if admissible > max_num_sequences: + if is_disagg: + # Expected: transfer-phase requests hold index leases, not seats. + return + direction, consequence = ( + "larger", + "a request could be admitted with no sequence slot to seat it, and " + "SlotManager.add_slot would raise on the executor's event loop") + else: + direction, consequence = ( + "smaller", "admitted requests would be silently deferred one at a " + "time (nvbug 6627795)") raise ValueError( f"{type(kv_cache_manager).__name__} can seat {admissible} concurrent " f"sequences but the executor's sequence-slot pool holds " f"{max_num_sequences}: the index pool is {direction} than the seat " - f"pool, so {consequence}. Both must come from " - "_util.compute_max_num_sequences; a mismatch means one of them was " - "re-derived from max_batch_size.") + f"pool, so {consequence}. The seat pool must come from " + "_util.compute_max_num_sequences and the index pool must be sized from " + "it; a mismatch means one of them was re-derived from max_batch_size.") def resolve_max_num_sequences(model_engine, mapping: Mapping, max_batch_size: int, llm_args, - is_disagg: bool, max_num_sequences: Optional[int] = None) -> int: """Resolve the seat-pool size for a consumer, without re-deriving it. @@ -3191,8 +3207,7 @@ def resolve_max_num_sequences(model_engine, max_batch_size, llm_args.disable_overlap_scheduler, enable_overlap_headroom=getattr( - model_engine, "_enable_adp_overlap_seq_slot_headroom", False), - is_disagg=is_disagg) + model_engine, "_enable_adp_overlap_seq_slot_headroom", False)) def should_enable_adp_dummy_fixes(mapping: Mapping) -> bool: @@ -3315,15 +3330,17 @@ def create_py_executor_instance( mapping, max_batch_size, llm_args, - is_disagg, max_num_sequences=max_num_sequences) # The seat pool and the KV index pool are sized independently and indexed by - # the same request identity, so any skew between them is a bug. Check it here - # rather than at first use: a startup ValueError names both numbers, while - # the runtime symptoms are a silent throughput loss in one direction and a - # raise inside the event loop in the other. - validate_seq_slot_pool_covers_admission(max_num_sequences, kv_cache_manager) + # the same request identity, so an index pool narrower than the seat pool is + # always a bug and a wider one is a bug unless disaggregation explains it. + # Check it here rather than at first use: a startup ValueError names both + # numbers, while the runtime symptoms are a silent throughput loss in one + # direction and a raise inside the event loop in the other. + validate_seq_slot_pool_covers_admission(max_num_sequences, + kv_cache_manager, + is_disagg=is_disagg) logger.info( f"max_seq_len={max_seq_len}, max_num_requests={max_num_sequences}, max_num_tokens={max_num_tokens}, max_batch_size={max_batch_size}" @@ -3761,7 +3778,6 @@ def instantiate_sampler( mapping, max_batch_size, llm_args, - is_disagg_enabled(getattr(llm_args, "cache_transceiver_config", None)), max_num_sequences=max_num_sequences) sampler_args = create_torch_sampler_args( 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 17599d716904..458ba4409900 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 @@ -1456,15 +1456,24 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: # # `max_num_seq_slots` is the executor's sequence-slot pool size # (`PyTorchModelEngine.max_num_seq_slots`, sized by - # `_util.compute_max_num_sequences`). An index slot and a sequence slot are - # both leases held for the whole lifetime of a request, so whenever the - # executor can admit N concurrent sequences the index mapper must be able - # to hand out N indices. Under the overlap scheduler with attention DP the - # seat pool carries an extra factor of 2 because teardown of the retiring - # batch (`_process_previous_batch`) happens *after* the replacement batch - # has already been scheduled, so both cohorts hold their leases at once. - # Without this term the index mapper runs dry and `_create_kv_cache` - # silently defers requests one at a time (nvbug 6627795). + # `_util.compute_max_num_sequences`). Every sequence the executor can + # seat must be able to obtain an index, so the index pool has to cover + # the seat pool: `index pool >= seat pool`. Under the overlap scheduler + # with attention DP the seat pool carries an extra micro-batch because + # teardown of the retiring batch (`_process_previous_batch`) happens + # *after* the replacement batch has already been scheduled, so both + # cohorts hold their leases at once. Without consuming that number the + # index mapper runs dry and `_create_kv_cache` silently defers requests + # one at a time (nvbug 6627795). + # + # The disagg 2x is deliberately *local to this pool* and is not + # propagated back into the seat pool: a request in TRANS_IN_PROGRESS + # holds an index lease while holding no sequence slot at all + # (`SeqSlotManager.prepare_resources` skips DISAGG_GENERATION_INIT and + # only seats a request once its transmission completes). Doubling the + # seat pool instead would double sampler state, the + # `[seats, draft_len, vocab]` draft-probability tensors and the + # pinned-host block-offset tables for leases that never occupy a seat. # # Take the larger of the two bounds rather than multiplying them: the # disagg and overlap coefficients each cover one extra set of in-flight @@ -1486,11 +1495,12 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: # slots reserved for padding/dummy requests, which are not available to # real sequences. Published so that # `_util.validate_seq_slot_pool_covers_admission` can compare it against - # the executor's sequence-slot pool without re-deriving either number -- - # every skew between the two is a bug, in both directions (a smaller - # index pool defers requests one at a time, a larger one lets a request - # be admitted that cannot be seated and `SlotManager.add_slot` then - # raises on the executor's event loop). + # the executor's sequence-slot pool without re-deriving either number. + # An index pool *smaller* than the seat pool is always a bug (it defers + # requests one at a time); a *larger* one is a bug only when aggregated, + # where it would let a request be admitted that cannot be seated and + # `SlotManager.add_slot` would then raise on the executor's event loop. + # Under disagg the surplus is the 2x above, and is expected. self.max_admissible_sequences = index_mapper_capacity - num_reserved_index_slots self.index_mapper = IndexMapper(index_mapper_capacity, max_beam_width) self._early_freed_index_requests: set[int] = set() diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index de18f75ad189..f515806e5ce6 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -518,7 +518,7 @@ def __init__( # their state-slot pool from max_batch_size alone, so they cannot use # the headroom -- which is why this runs after the model is loaded # rather than next to `self.mapping`. - from ._util import (compute_max_num_sequences, is_disagg_enabled, + from ._util import (compute_max_num_sequences, should_enable_adp_dummy_fixes, should_enable_adp_overlap_seq_slot_headroom, should_enable_non_overlap_adp_forward_intent, @@ -534,7 +534,6 @@ def __init__( self.batch_size, llm_args.disable_overlap_scheduler, enable_overlap_headroom=self._enable_adp_overlap_seq_slot_headroom, - is_disagg=is_disagg_enabled(llm_args.cache_transceiver_config), ) self._enable_scheduler_aware_adp_dummy = ( should_enable_scheduler_aware_adp_dummy( 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 ba93fbeff568..fade285df8fc 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 @@ -1643,6 +1643,12 @@ def build_base_config( pytest.param(2, 1, True, 1, 4, 5, id="disagg_does_not_compound"), # Disagg without seat headroom keeps its own 2x. pytest.param(2, 1, True, 1, None, 5, id="disagg_only"), + # The same cell as it actually occurs once the seat pool is plumbed: disagg + # without ADP seats B*pp == 2, and the mapper still carries its own 2x. The + # surplus is deliberate and is the reason the 2x lives *here* and not in + # compute_max_num_sequences -- a request in KV transfer holds an index lease + # while SeqSlotManager.prepare_resources skips it, so it occupies no seat. + pytest.param(2, 1, True, 1, 2, 5, id="disagg_index_pool_exceeds_seat_pool"), # PP without the headroom: seats are B*pp, which the mapper already matched. pytest.param(2, 4, False, 1, 8, 9, id="pp4"), # PP with the headroom: seats are (pp+1)*B == 10, additive rather than @@ -1685,34 +1691,44 @@ def test_index_mapper_capacity_covers_seq_slot_pool( # pipeline-parallel cell, and a caller that supplies no seat pool at all. @pytest.mark.cpu_only @pytest.mark.parametrize( - "max_batch_size,pp_size,reserved,max_num_seq_slots,expected_admissible", + "max_batch_size,pp_size,is_disagg,reserved,max_num_seq_slots,expected_admissible", [ - pytest.param(2, 1, 1, 4, 4, id="agg_adp_overlap"), - pytest.param(2, 4, 1, 10, 10, id="pp4_adp_overlap"), - pytest.param(2, 1, 5, None, 2, id="seats_unset_reserved_excluded"), + pytest.param(2, 1, False, 1, 4, 4, id="agg_adp_overlap"), + pytest.param(2, 4, False, 1, 10, 10, id="pp4_adp_overlap"), + pytest.param(2, 1, False, 5, None, 2, id="seats_unset_reserved_excluded"), + # Disagg is the one case where the published number legitimately exceeds + # the seat pool it is compared against. + pytest.param(2, 1, True, 1, 2, 4, id="disagg_exceeds_seat_pool"), ], ) def test_index_mapper_publishes_max_admissible_sequences( max_batch_size: int, pp_size: int, + is_disagg: bool, reserved: int, max_num_seq_slots: int | None, expected_admissible: int, ) -> None: - """The manager publishes what it can seat, so nobody has to re-derive it. + """The manager publishes what it can index, so nobody has to re-derive it. ``_util.validate_seq_slot_pool_covers_admission`` compares this against the executor's sequence-slot pool at startup. It is the pool *net of* the reserved - padding/dummy slots, which no real sequence can take, and on every row where a - seat pool was plumbed through it equals that seat pool exactly -- that - equality is the invariant the validator enforces, in both directions. + padding/dummy slots, which no real sequence can take. Aggregated, it equals + the seat pool exactly on every row where one was plumbed through, and the + validator enforces that in both directions. Under disaggregation it may + exceed the seat pool, because an index lease outlives the seat while the KV + transfer drains -- so there the validator enforces only ``>=``. """ _, _, admissible = _index_mapper_capacity_for( max_batch_size=max_batch_size, pp_size=pp_size, + is_disagg=is_disagg, num_reserved_index_slots=reserved, max_num_seq_slots=max_num_seq_slots, ) assert admissible == expected_admissible if max_num_seq_slots is not None: - assert admissible == max_num_seq_slots + if is_disagg: + assert admissible >= max_num_seq_slots + else: + assert admissible == max_num_seq_slots diff --git a/tests/unittest/_torch/executor/test_seq_slot_sizing.py b/tests/unittest/_torch/executor/test_seq_slot_sizing.py index 84a7d8abd28e..fb4ff971d08e 100644 --- a/tests/unittest/_torch/executor/test_seq_slot_sizing.py +++ b/tests/unittest/_torch/executor/test_seq_slot_sizing.py @@ -44,31 +44,27 @@ ) from tensorrt_llm.mapping import Mapping -# (pp_size, disable_overlap, enable_overlap_headroom, is_disagg, expected_factor) +# (pp_size, disable_overlap, enable_overlap_headroom, expected_factor) # # The terms are additive, not multiplicative: pipeline depth costs pp_size # micro-batches of seats, and the overlap deferral costs exactly one more # generation on top -- not one more per stage. At pp_size == 1 the two readings # coincide at 2x, which is why the headroom used to be expressible as a factor of # 2; the pp>1 rows are where they part company (5x, not 8x, at pp=4). +# +# Disaggregation is absent from this table on purpose -- see +# test_seat_pool_has_no_disagg_term. SIZING_CASES = [ # No PP. Every cell here is unchanged by this PR. - (1, False, False, False, 1), - (1, False, True, False, 2), - (1, True, True, False, 1), - (1, False, False, True, 2), - (1, False, True, True, 2), + (1, False, False, 1), + (1, False, True, 2), + (1, True, True, 1), # PP without the headroom: unchanged. - (4, False, False, False, 4), - (4, True, True, False, 4), - (2, False, False, True, 4), - (4, False, False, True, 8), + (4, False, False, 4), + (4, True, True, 4), # PP with the headroom: the one intended behaviour change (was pp_size). - (2, False, True, False, 3), - (4, False, True, False, 5), - # Disagg dominates the additive term rather than compounding with it: the two - # coefficients each cover one extra cohort of in-flight sequences. - (4, False, True, True, 8), + (2, False, True, 3), + (4, False, True, 5), ] @@ -147,10 +143,10 @@ def test_non_overlap_adp_forward_intent_scope(pp_size, disable_overlap, expected @pytest.mark.parametrize( - "pp_size,disable_overlap,enable_overlap_headroom,is_disagg,expected_factor", SIZING_CASES + "pp_size,disable_overlap,enable_overlap_headroom,expected_factor", SIZING_CASES ) def test_compute_max_num_sequences_scopes_overlap_headroom( - pp_size, disable_overlap, enable_overlap_headroom, is_disagg, expected_factor + pp_size, disable_overlap, enable_overlap_headroom, expected_factor ): max_batch_size = 8 mapping = Mapping(world_size=pp_size, tp_size=1, pp_size=pp_size) @@ -160,12 +156,32 @@ def test_compute_max_num_sequences_scopes_overlap_headroom( max_batch_size, disable_overlap, enable_overlap_headroom=enable_overlap_headroom, - is_disagg=is_disagg, ) == max_batch_size * expected_factor ) +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( "cache_transceiver_config,expected", [ @@ -213,7 +229,6 @@ def test_resolve_max_num_sequences_prefers_the_published_pool(explicit, engine_s mapping, 8, llm_args, - False, max_num_sequences=explicit, ) == expected @@ -245,13 +260,12 @@ def __getattr__(self, name): mapping, 8, _Exploding(), - False, max_num_sequences=24, ) == 24 ) # Branch 2: the engine's published pool. - assert resolve_max_num_sequences(engine_with_pool, mapping, 8, _Exploding(), False) == 16 + assert resolve_max_num_sequences(engine_with_pool, mapping, 8, _Exploding()) == 16 def test_sampler_args_require_the_resolved_pool(): @@ -292,22 +306,42 @@ 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)) +@pytest.mark.parametrize("is_disagg", [False, True]) +def test_validator_accepts_the_matching_pair(is_disagg): + validate_seq_slot_pool_covers_admission(16, _FakeManager(16), is_disagg=is_disagg) -@pytest.mark.parametrize("admissible", [8, 32]) -def test_validator_is_two_sided(admissible): - """Both directions of the skew are bugs, and both have shipped. +@pytest.mark.parametrize("is_disagg", [False, True]) +def test_validator_always_rejects_an_index_pool_below_the_seat_pool(is_disagg): + """The direction that shipped as nvbug 6627795, and it is never legitimate. - A one-sided ``seats >= admissible`` guard is what let nvbug 6627795 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. The - other direction -- index pool larger than the seat pool -- admits a request - that cannot be seated and raises inside the executor's event loop. + 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. + Disaggregation is no excuse here -- its 2x makes the index pool *larger*, so + a shortfall under disagg means the two numbers were derived separately. """ - with pytest.raises(ValueError, match="sequence-slot pool"): - validate_seq_slot_pool_covers_admission(16, _FakeManager(admissible)) + with pytest.raises(ValueError, match="smaller than the seat"): + validate_seq_slot_pool_covers_admission(16, _FakeManager(8), is_disagg=is_disagg) + + +def test_validator_rejects_a_larger_index_pool_only_when_aggregated(): + """The other direction is a bug when aggregated and by design under disagg. + + Aggregated, every indexed sequence is also a seated one, so a surplus can + only mean one of the two numbers was re-derived -- and a request would be + admitted that cannot be seated, with ``SlotManager.add_slot`` raising on the + executor's event-loop thread. Under disaggregation the surplus *is* the + mechanism: a request in KV transfer holds its index lease with no seat + (``SeqSlotManager.prepare_resources`` skips ``DISAGG_GENERATION_INIT``), so + the index pool carries a 2x that must not reach the seat pool. Admission is + bounded independently at ``max_batch_size * pp_size`` either way, so the + surplus is never extra concurrency. + """ + with pytest.raises(ValueError, match="larger than the seat"): + validate_seq_slot_pool_covers_admission(16, _FakeManager(32), is_disagg=False) + + validate_seq_slot_pool_covers_admission(16, _FakeManager(32), is_disagg=True) @pytest.mark.parametrize( @@ -335,9 +369,9 @@ def test_validator_skips_managers_that_do_not_publish_capacity(manager): derives from KVCacheManagerV2 and therefore does publish max_admissible_sequences, so it *is* validated. That comparison holds because the headroom is withheld from hybrid on both sides -- seats are - B*pp (or 2*B*pp under disagg) and, with max_num_seq_slots withheld, its - index pool lands on exactly the same number. What hybrid does not receive - is the seat pool, not the check. + B*pp and, with max_num_seq_slots withheld, its index pool lands on exactly + the same number (2*B*pp under disagg, which the validator allows). What + hybrid does not receive is the seat pool, not the check. """ validate_seq_slot_pool_covers_admission(16, manager) From b01d2084f1753dc06288258028623090581fbe4d Mon Sep 17 00:00:00 2001 From: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> Date: Wed, 9 Sep 2026 00:44:44 -0700 Subject: [PATCH 17/22] [https://nvbugs/6627795][fix] scope the overlap headroom to non-PP attention DP Size the KV index pool from the manager's own inputs instead of importing the executor's seat-pool size: the coefficient becomes 2 if is_disagg or (attention_dp and overlap and not pp) else 1 which keeps the manager deriving its capacity from max_batch_size * pp_size and leaves every non-ADP topology byte-identical. Because the two pools are no longer expressed in the same terms -- the seat pool carries the PP multiplier and the index pool does not -- drop validate_seq_slot_pool_covers_admission and the max_admissible_sequences it compared, and stop plumbing max_num_seq_slots into the manager. Keep should_enable_disagg_adp_overlap_headroom under its original name and widen only its predicate to attention DP without PP, with either disaggregation or the overlap scheduler; compute_max_num_sequences goes back to pp_size micro-batches under PP and the headroom factor otherwise. Pipeline parallelism is out of scope on both sides. Revert the model_engine.py and mamba_cache_manager.py changes: neither needed one. In SuffixAutomatonManager, treat the sequence-slot count as a floor on an explicit global_pool_size rather than a new startup failure, so a config TorchLlmArgs already accepted does not start raising once attention DP is on. Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/_util.py | 250 +++--------------- .../kv_cache/kv_cache_manager_v2.py | 53 +--- .../kv_cache/mamba_cache_manager.py | 21 -- .../_torch/pyexecutor/model_engine.py | 54 ++-- tensorrt_llm/_torch/pyexecutor/py_executor.py | 30 +-- .../_torch/pyexecutor/py_executor_creator.py | 14 +- .../_torch/pyexecutor/scheduler/adp_router.py | 113 +------- .../_torch/speculative/suffix_automaton.py | 21 +- tensorrt_llm/_torch/speculative/utils.py | 20 +- .../kv_cache/test_kv_cache_estimation.py | 9 +- .../kv_cache/test_kv_cache_manager_v2.py | 156 ++++------- .../_torch/executor/test_adp_router.py | 13 +- .../_torch/executor/test_seq_slot_sizing.py | 235 +++++----------- .../speculative/test_spec_slot_pool_sizing.py | 24 +- 14 files changed, 253 insertions(+), 760 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 41df8d8a5f9b..8fa818d20feb 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -642,6 +642,7 @@ def __init__( self._dummy_encoder_inputs: List[MultimodalParams] = [] self._profiling_stage_data = profiling_stage_data self._is_disagg = is_disagg + self._enable_overlap_scheduler = not llm_args.disable_overlap_scheduler self._cache_transceiver_config = llm_args.cache_transceiver_config self._execution_stream = execution_stream self._kv_cache_manager_cls = self._get_model_kv_cache_manager_cls( @@ -1419,17 +1420,6 @@ def configure_kv_cache_capacity(self, self._profiling_stage_data["activation_bytes"] = activation_bytes # ---------------------------handle max_gpu_total_bytes--------------------------------- - def _target_max_num_seq_slots(self) -> Optional[int]: - """Size of the executor's sequence-slot pool, or None if unavailable. - - There is exactly one SeqSlotManager per executor, sized by - compute_max_num_sequences from the *target* engine (see - create_py_executor_instance). Every KV cache manager -- target, draft and - cross -- must be able to hand out an index for each seat, otherwise the - pools drift apart and requests are silently deferred (nvbug 6627795). - """ - return getattr(self._model_engine, "max_num_seq_slots", None) - def _create_kv_cache_manager( self, model_engine: PyTorchModelEngine, @@ -1473,10 +1463,7 @@ def _create_kv_cache_manager( execution_stream=self._execution_stream, layer_mask=spec_dec_layer_mask, is_disagg=self._is_disagg, - # Always the target engine's seat pool, even when building the draft - # manager: the executor has a single SeqSlotManager, sized from the - # target engine, and every manager must cover it. - max_num_seq_slots=self._target_max_num_seq_slots(), + enable_overlap_scheduler=self._enable_overlap_scheduler, cold_page_codec_provider=cold_page_codec_provider, joint_kv_cache_reuse=self._joint_kv_cache_reuse, ) @@ -1683,7 +1670,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, - max_num_seq_slots=self._target_max_num_seq_slots(), + enable_overlap_scheduler=self._enable_overlap_scheduler, cold_page_codec_provider=cold_page_codec_provider, joint_kv_cache_reuse=self._joint_kv_cache_reuse, ) @@ -2060,7 +2047,7 @@ def _create_cross_kv_cache_manager( num_layers=num_layers, num_kv_heads=num_kv_heads, head_dim=head_dim, - max_num_seq_slots=self._target_max_num_seq_slots(), + enable_overlap_scheduler=self._enable_overlap_scheduler, kv_cache_type=tensorrt_llm.bindings.internal.batch_manager. CacheType.CROSS, ) @@ -2380,7 +2367,7 @@ def _create_kv_cache_manager( head_dim: Optional[int] = None, kv_cache_type=None, is_disagg: bool = False, - max_num_seq_slots: Optional[int] = None, + enable_overlap_scheduler: bool = False, cold_page_codec_provider: Optional[object] = None, joint_kv_cache_reuse: bool = False) -> KVCacheManager: """ @@ -2522,12 +2509,8 @@ def _create_kv_cache_manager( manager_extra_kwargs[ "cold_page_codec_provider"] = cold_page_codec_provider manager_extra_kwargs["joint_kv_cache_reuse"] = joint_kv_cache_reuse - # Hybrid managers size their SSM state-slot pool from max_batch_size - # (see MambaHybridCacheManager._max_resident_sequences), so growing only - # the index mapper would let a request hold an index with no state slot. - # Sizing both pools together is left as a follow-up. - if not issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2): - manager_extra_kwargs["max_num_seq_slots"] = max_num_seq_slots + manager_extra_kwargs[ + "enable_overlap_scheduler"] = enable_overlap_scheduler if issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2): manager_extra_kwargs["is_disagg"] = is_disagg @@ -3050,12 +3033,7 @@ def create_kv_cache_compression_manager( def is_disagg_enabled(cache_transceiver_config) -> bool: - """True when a cache transceiver backend is configured. - - Single definition of "this is a disaggregated server", so the seat pool and - the KV cache managers cannot disagree about it. The test was previously - inlined at each use site, which is how a derived fact acquires copies. - """ + """True when a cache transceiver backend is configured.""" return (cache_transceiver_config is not None and cache_transceiver_config.backend is not None) @@ -3066,106 +3044,18 @@ def compute_max_num_sequences(mapping: Mapping, enable_overlap_headroom: bool = False) -> int: """Size the sequence-slot pool (and the sampler state it indexes). - This is *the* definition of how many sequences can be simultaneously - **seated**. Every pool keyed by seat identity -- the sampler's per-slot - state, the guided decoder, the spec-decoding managers -- must consume this - number rather than re-deriving it from ``max_batch_size``, because any two - formulas that disagree are a bug: too few seats raises ``NoFreeSlotsError`` - on the executor's event-loop thread, and a KV index pool narrower than the - seat pool silently defers admitted requests (nvbug 6627795). - - The terms are **additive, not multiplicative**: - - * ``pp_size`` micro-batches are structurally in flight under pipeline - parallelism, so the pool scales with pipeline depth. - * The overlap scheduler defers a finished request's teardown by exactly - **one** iteration, so at most one extra generation of seats is held while - the ADP router admits the batch that replaces it. That is ``+1`` - micro-batch worth of seats, not a doubling of every pipeline stage. - - At ``pp_size == 1`` the additive and multiplicative forms coincide at - ``2 * max_batch_size``, which is why the overlap headroom used to be - expressible as a factor of 2. It is a coincidence of ``pp_size == 1``, and - the multiplicative reading costs ``+100%`` at ``pp_size == 4`` where the - additive one costs ``+25%`` -- on pools that include eagerly-allocated - ``[seats, draft_len, vocab]`` tensors. - - 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 -- so - the index pool legitimately runs ahead of the seat pool while admission - stays at ``max_batch_size * pp_size``. Propagating that factor into the seat - pool would double sampler state, ``[seats, draft_len, vocab]`` draft - probabilities and the pinned-host block-offset tables for leases that never - occupy a seat. - """ - # Pipeline depth: pp_size micro-batches are structurally in flight. - num_seats = max_batch_size * mapping.pp_size - if enable_overlap_headroom and not disable_overlap_scheduler: - num_seats += max_batch_size - return num_seats - - -def validate_seq_slot_pool_covers_admission(max_num_sequences: int, - kv_cache_manager, - is_disagg: bool = False) -> None: - """Fail at startup if the seat pool and the KV index pool disagree. - - The invariant is ``index pool >= seat pool``, and the two bounds have - different characters -- which is why the check is asymmetric rather than a - plain 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, no error (nvbug 6627795). A - one-sided ``seats >= index pool`` guard is exactly what let this ship. - * index pool > seat pool -- a bug **when aggregated**, because there every - indexed sequence is also a seated one, so the surplus can only come from - one of the two numbers having been re-derived; a request would be admitted - that cannot be seated and ``SlotManager.add_slot`` would raise on the - executor's event-loop thread. Under disaggregation it is **by design**: a - request awaiting its KV transfer holds an index lease and no seat - (``SeqSlotManager.prepare_resources`` skips ``DISAGG_GENERATION_INIT``), - so the index pool carries a ``2x`` that the seat pool must not. - - Admission is bounded independently at ``max_batch_size * pp_size`` by - ``ModelEngine.get_max_num_sequences()`` in both cases, so the surplus is - index headroom, never extra concurrency. - - Managers that do not publish ``max_admissible_sequences`` (V1, hybrid) are - skipped rather than guessed at; publishing the attribute *as an int* is how a - manager opts into the check. Anything non-integral means the same thing as - absent -- this manager did not opt in -- so the check is skipped rather than - attempted against a value it cannot order. + ``enable_overlap_headroom`` is intentionally opt-in; see + ``should_enable_disagg_adp_overlap_headroom`` for when it is set. It buys one + extra micro-batch worth of slots, because a finished request's teardown is + deferred by one iteration and its slot is still held while the replacement + batch is admitted (nvbug 6627795). Pipeline parallelism already sizes the + pool by ``pp_size``. """ - if kv_cache_manager is None: - return - admissible = getattr(kv_cache_manager, "max_admissible_sequences", None) - if not isinstance(admissible, int): - return - if admissible == max_num_sequences: - return - if admissible > max_num_sequences: - if is_disagg: - # Expected: transfer-phase requests hold index leases, not seats. - return - direction, consequence = ( - "larger", - "a request could be admitted with no sequence slot to seat it, and " - "SlotManager.add_slot would raise on the executor's event loop") + if mapping.has_pp(): + num_micro_batches = mapping.pp_size else: - direction, consequence = ( - "smaller", "admitted requests would be silently deferred one at a " - "time (nvbug 6627795)") - raise ValueError( - f"{type(kv_cache_manager).__name__} can seat {admissible} concurrent " - f"sequences but the executor's sequence-slot pool holds " - f"{max_num_sequences}: the index pool is {direction} than the seat " - f"pool, so {consequence}. The seat pool must come from " - "_util.compute_max_num_sequences and the index pool must be sized from " - "it; a mismatch means one of them was re-derived from max_batch_size.") + num_micro_batches = (2 if enable_overlap_headroom else 1) + return max_batch_size * num_micro_batches def resolve_max_num_sequences(model_engine, @@ -3180,20 +3070,9 @@ def resolve_max_num_sequences(model_engine, 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 the KV cache managers were sized against; + what every seat-keyed pool was sized against; 3. only then a fresh ``compute_max_num_sequences``, reusing the engine's - headroom gate. - - Step 3 used to be step 1, called *without* ``enable_overlap_headroom``. That - silently sized the sampler and the executor's ``SeqSlotManager`` below the - index pool they share indices with, i.e. it reintroduced the very skew this - number exists to eliminate, in the one code path that has no test coverage. - - ``llm_args`` is taken whole rather than as an already-read - ``disable_overlap_scheduler`` flag so that the fallback's inputs are read - only when the fallback runs. Passing the flag made the caller evaluate it - eagerly, which reintroduced a hard dependency on the field in the branch - that never uses it -- and broke callers holding a lighter args object. + headroom gate so the fallback cannot size the pool below the engine's. """ if max_num_sequences is not None: return max_num_sequences @@ -3202,12 +3081,13 @@ def resolve_max_num_sequences(model_engine, 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_adp_overlap_seq_slot_headroom", False)) + return compute_max_num_sequences(mapping, + max_batch_size, + llm_args.disable_overlap_scheduler, + enable_overlap_headroom=getattr( + model_engine, + "_enable_disagg_adp_overlap_headroom", + False)) def should_enable_adp_dummy_fixes(mapping: Mapping) -> bool: @@ -3234,58 +3114,27 @@ def should_enable_non_overlap_adp_forward_intent( and disable_overlap_scheduler) -def should_enable_adp_overlap_seq_slot_headroom( +def should_enable_disagg_adp_overlap_headroom( mapping: Mapping, - disable_overlap_scheduler: bool, - is_hybrid: bool = False) -> bool: - """Gate extra sequence slots to attention-DP with the overlap scheduler on. + cache_transceiver_config: Optional[CacheTransceiverConfig], + disable_overlap_scheduler: bool) -> bool: + """Gate extra sequence slots to non-PP attention DP. The overlap scheduler defers a finished request's teardown by one iteration, so its sequence slot is still held when the ADP router admits the batch that replaces it. Without spare slots the router cannot backfill and the forward - batch runs short (nvbug-6627795); one extra micro-batch worth of slots lets - admission reach max_batch_size on every rank every iteration. - - Requiring attention DP -- rather than merely overlap -- is deliberate: with - a single scheduling domain the executor's own admission bound already tracks - the pool, whereas ADP admits against an allgathered load vector that the - deferred teardown desynchronizes. Not gated on disaggregation: the mechanism - is a property of overlap plus ADP admission, and was measured on an - aggregated context-only run with no cache transceiver configured. - - Pipeline parallelism is still excluded here, but no longer because of the - sizing: compute_max_num_sequences now expresses the headroom additively, so - ``(pp_size + 1) * max_batch_size`` is a well-defined pool under PP. What is - missing is the *consumer*: ADPRouter's retiring-request correction requires - every pipeline stage to agree on which requests are retiring, and today only - the last stage marks generation requests GENERATION_TO_COMPLETE. Opening the - gate before that is fixed would allocate headroom no rank ever spends. - - Measured, so that the next attempt starts from the number rather than the - argument: a generation-bearing pp=2 + ADP + overlap A/B (GLM-5 NVFP4, 4-way - GB300, max_batch_size=4, 1920 requests per arm) produced *byte-identical* - scheduling on both arms -- mean scheduled batch 3.997 of a cap of 4, the same - per-iteration histogram, and zero index-pool exhaustions. The seats are - unspendable at pp>1 because admission is capped independently at - ``pp_size * max_batch_size`` by ``get_max_num_sequences()``, which is already - below the unpatched index pool of ``pp_size * max_batch_size + 1``: the - per-micro-batch batch size, not the seat pool, is the limiter. So the PP cell - needs a case where admission is the binding constraint before it can show a - benefit; the throughput difference between those two arms (-2.34%) is this - configuration's noise floor, not a result. - - Hybrid (Mamba/SSM) architectures are excluded. MambaHybridCacheManagerV2 - sizes its state-index pool from max_batch_size alone - (mamba_cache_manager.py: state_index_capacity), with neither pp_size nor the - seat count, so extra seats would let a request hold a seat with no SSM state - slot behind it. That pool is already inconsistent with its own - _max_resident_sequences() under PP; fixing it is a separate change, and - until then the headroom must not reach these models. + batch runs short (nvbug 6627795). Disaggregation needs the same headroom even + with overlap off, because a request awaiting its KV transfer keeps its lease. + + Pipeline parallelism is excluded: admission is capped independently at + ``pp_size * max_batch_size``, so the extra seats are unspendable, and + ADPRouter's retiring-request correction would additionally require every + pipeline stage to agree on which requests are retiring. """ - if is_hybrid: - return False + is_disagg = is_disagg_enabled(cache_transceiver_config) + enable_overlap_scheduler = not disable_overlap_scheduler return (mapping.enable_attention_dp and not mapping.has_pp() - and not disable_overlap_scheduler) + and (is_disagg or enable_overlap_scheduler)) def create_py_executor_instance( @@ -3332,16 +3181,6 @@ def create_py_executor_instance( llm_args, max_num_sequences=max_num_sequences) - # 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 and a wider one is a bug unless disaggregation explains it. - # Check it here rather than at first use: a startup ValueError names both - # numbers, while the runtime symptoms are a silent throughput loss in one - # direction and a raise inside the event loop in the other. - validate_seq_slot_pool_covers_admission(max_num_sequences, - kv_cache_manager, - is_disagg=is_disagg) - logger.info( f"max_seq_len={max_seq_len}, max_num_requests={max_num_sequences}, max_num_tokens={max_num_tokens}, max_batch_size={max_batch_size}" ) @@ -3535,11 +3374,8 @@ def create_py_executor_instance( # Enlarge scheduler capacity to avoid DISAGG_GENERATION_INIT stuck in the scheduler. # 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 pool is larger than this under the - # attention-DP overlap headroom -- (pp_size + 1) * max_batch_size -- because - # it must also cover the retiring cohort whose teardown the overlap scheduler - # defers by one iteration. Those seats are headroom for leases already held, - # not extra admission, so growing the budget with them would over-admit. + # 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 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 458ba4409900..8b84f862eb88 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 @@ -931,7 +931,7 @@ def __init__( is_disagg: bool = False, enable_stats: bool = False, num_reserved_index_slots: int = 1, - max_num_seq_slots: Optional[int] = None, + enable_overlap_scheduler: bool = False, is_estimating_kv_cache: bool = False, cold_page_codec_provider: Optional[object] = None, joint_kv_cache_reuse: bool = False, @@ -1454,54 +1454,29 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: # capacity lets the next batch of active requests acquire slots without # waiting for the previous batch's transfers to finish. # - # `max_num_seq_slots` is the executor's sequence-slot pool size - # (`PyTorchModelEngine.max_num_seq_slots`, sized by - # `_util.compute_max_num_sequences`). Every sequence the executor can - # seat must be able to obtain an index, so the index pool has to cover - # the seat pool: `index pool >= seat pool`. Under the overlap scheduler - # with attention DP the seat pool carries an extra micro-batch because - # teardown of the retiring batch (`_process_previous_batch`) happens - # *after* the replacement batch has already been scheduled, so both - # cohorts hold their leases at once. Without consuming that number the - # index mapper runs dry and `_create_kv_cache` silently defers requests - # one at a time (nvbug 6627795). - # - # The disagg 2x is deliberately *local to this pool* and is not - # propagated back into the seat pool: a request in TRANS_IN_PROGRESS - # holds an index lease while holding no sequence slot at all - # (`SeqSlotManager.prepare_resources` skips DISAGG_GENERATION_INIT and - # only seats a request once its transmission completes). Doubling the - # seat pool instead would double sampler state, the - # `[seats, draft_len, vocab]` draft-probability tensors and the - # pinned-host block-offset tables for leases that never occupy a seat. - # - # Take the larger of the two bounds rather than multiplying them: the - # disagg and overlap coefficients each cover one extra set of in-flight - # sequences, so they overlap rather than compose. + # Attention DP with the overlap scheduler needs the same coefficient: + # teardown of the retiring batch (`_process_previous_batch`) runs *after* + # the replacement batch has already been scheduled, so both cohorts hold + # their index slots at once. Without the extra capacity the index mapper + # runs dry and `_create_kv_cache` silently defers requests one at a time + # (nvbug 6627795). Pipeline parallelism is excluded because + # `max_num_sequences` already scales with the number of in-flight + # microbatches. + needs_extra_index_slots = is_disagg or ( + mapping.enable_attention_dp and enable_overlap_scheduler and not mapping.has_pp() + ) max_num_sequences = max_batch_size * mapping.pp_size assert num_reserved_index_slots >= 0, "num_reserved_index_slots must be non-negative" index_mapper_capacity = ( - max(max_num_sequences * (2 if is_disagg else 1), max_num_seq_slots or 0) - + num_reserved_index_slots + max_num_sequences * (2 if needs_extra_index_slots else 1) + 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"max_num_seq_slots={max_num_seq_slots}, " + f"enable_overlap_scheduler={enable_overlap_scheduler}, " f"num_reserved_index_slots={num_reserved_index_slots}, " f"max_beam_width={max_beam_width})" ) - # Concurrent sequences this manager can seat: the index pool net of the - # slots reserved for padding/dummy requests, which are not available to - # real sequences. Published so that - # `_util.validate_seq_slot_pool_covers_admission` can compare it against - # the executor's sequence-slot pool without re-deriving either number. - # An index pool *smaller* than the seat pool is always a bug (it defers - # requests one at a time); a *larger* one is a bug only when aggregated, - # where it would let a request be admitted that cannot be seated and - # `SlotManager.add_slot` would then raise on the executor's event loop. - # Under disagg the surplus is the 2x above, and is expected. - self.max_admissible_sequences = index_mapper_capacity - num_reserved_index_slots self.index_mapper = IndexMapper(index_mapper_capacity, max_beam_width) self._early_freed_index_requests: set[int] = set() self._prepare_page_table_tensor(index_mapper_capacity) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index dcaddbdfeffc..e242959e5117 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -3213,11 +3213,6 @@ def __init__( self._request_id_to_state_index = {} self._request_id_to_is_dummy = {} - # Sized by *scheduled batch* position, not by live sequence: - # _setup_state_indices fills [0, len(requests)) for one iteration's - # requests. So this is deliberately max_batch_size and must NOT grow with - # pp_size or with the executor's sequence-slot pool -- unlike the SSM - # state slots checked below, which are per-live-sequence leases. state_index_capacity = (self.max_batch_size + self._num_reserved_dummy_slots) self.cuda_state_indices = torch.zeros([state_index_capacity], @@ -3238,22 +3233,6 @@ def __init__( LayerId(first_mamba_local_layer), MambaRole.SSM_STATE) num_ssm_slots = ((num_ssm_pages + self._ssm_page_index_scale - 1) // self._ssm_page_index_scale) - # Per-live-sequence leases, so this floor tracks the number of - # sequences that can be resident at once -- max_batch_size * pp_size. - # It is deliberately *not* raised to `max_admissible_sequences` (the - # index-mapper pool net of reserved slots): under disaggregation that - # pool carries a 2x for requests still draining their KV transfer, - # and whether such a request also retains its SSM state slot is - # unresolved. Requiring the larger number would turn a possibly - # adequate Mamba pool into a startup failure, so the count is rounded - # down here on purpose. - # - # This manager is also excluded from the attention-DP overlap seat - # headroom (`should_enable_adp_overlap_seq_slot_headroom` returns - # False for hybrid architectures, and `_create_kv_cache_manager` - # withholds `max_num_seq_slots`), so the seat pool it is sized against - # is exactly `_max_resident_sequences()` and cannot outgrow this - # floor. required_live_slots = (self._max_resident_sequences() + self._num_reserved_dummy_slots) if num_ssm_slots < required_live_slots: diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index f515806e5ce6..82543b67b43c 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -80,7 +80,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, is_mla +from .config_utils import is_mla from .cuda_graph_runner import (ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM, CUDAGraphRunner, CUDAGraphRunnerConfig, EncoderCUDAGraphRunner, @@ -434,6 +434,24 @@ 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_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) @@ -511,30 +529,6 @@ def __init__( self._validate_breakable_cuda_graph_compatibility() pretrained_config = self.model.model_config.pretrained_config model_type = getattr(pretrained_config, "model_type", None) - # Attention-DP can backfill a batch before the overlap scheduler - # releases the previous batch's terminal sequence slots, so the seat - # pool needs one extra generation of slots. Both the gate and the - # sizing depend on the architecture -- hybrid/SSM cache managers size - # their state-slot pool from max_batch_size alone, so they cannot use - # the headroom -- which is why this runs after the model is loaded - # rather than next to `self.mapping`. - from ._util import (compute_max_num_sequences, - should_enable_adp_dummy_fixes, - should_enable_adp_overlap_seq_slot_headroom, - should_enable_non_overlap_adp_forward_intent, - should_enable_scheduler_aware_adp_dummy) - self._enable_adp_overlap_seq_slot_headroom = ( - should_enable_adp_overlap_seq_slot_headroom( - mapping, - llm_args.disable_overlap_scheduler, - is_hybrid=is_hybrid_linear(pretrained_config))) - 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_adp_overlap_seq_slot_headroom, - ) self._enable_scheduler_aware_adp_dummy = ( should_enable_scheduler_aware_adp_dummy( model_type, mapping, llm_args.disable_overlap_scheduler)) @@ -3577,11 +3571,11 @@ def _set_up_spec_metadata( spec_resource_manager: Optional[BaseResourceManager], no_cache=False): spec_config = self.spec_config if self.enable_spec_decode else None - # The 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_adp_overlap_seq_slot_headroom 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 no_cache: return get_spec_metadata( spec_config, diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 545342c725dd..356e4ef7f267 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -692,14 +692,12 @@ def __init__( # can receive the transfer-manager reference at construction time. # # The router's overlap correction spends sequence-slot headroom, so it is - # handed the engine's headroom flag rather than re-deriving the - # predicate: the flag is what sized the seat pool, and it withholds the - # headroom for architectures (hybrid/SSM) whose state-slot pool is not - # sized from that number. Absent flag => no headroom => no correction. + # handed the engine's headroom flag rather than re-deriving the predicate: + # that flag is what sized the seat pool the correction spends. self.adp_router: ADPRouter = ADPRouter.create( dist=self.dist, has_seq_slot_headroom=getattr( - model_engine, "_enable_adp_overlap_seq_slot_headroom", False), + model_engine, "_enable_disagg_adp_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,11 +5804,9 @@ def _fetch_and_enqueue_requests(self, waiting_queue: WaitingQueue, """Fetch requests from request_queue and enqueue to waiting_queue. `total_num_live_requests` counts every request still resident on any - rank, including the retiring ones that `num_active_requests` excludes. - The idle decision below must be taken on that figure and it must be - identical on every rank: it selects a blocking versus a zero timeout, - and a rank that blocks on the untimed queue wait while its peers reach - `dist.broadcast(root=0)` deadlocks the iteration. + rank, including the retiring ones that `num_active_requests` excludes: + the idle decision below selects a blocking versus a zero timeout and must + be identical on every rank. """ # Block new requests while control requests are pending if len(self.control_requests) != 0: @@ -6029,11 +6025,9 @@ def _fetch_new_requests( s.num_active_requests for s in all_rank_states ] total_num_active_requests = sum(all_ranks_num_active_requests) - # Retiring requests are excluded from num_active_requests (they - # cannot be scheduled, so they must not consume admission - # capacity -- nvbug-6627795) but they are still resident, so the - # loop is NOT idle while any of them exists. Fold them back in - # for the liveness test only. + # Retiring requests are excluded from num_active_requests but are + # still resident, so fold them back in for the liveness test only + # (nvbug 6627795). total_num_live_requests = total_num_active_requests + sum( s.num_retiring_requests for s in all_rank_states) else: @@ -7250,12 +7244,8 @@ def _pad_attention_dp_dummy_request(self): expected_num_active_requests = self.expected_num_active_requests # Compare against the same routable count the router balanced on: # gather_all_rank_states excludes retiring requests from the per-rank - # loads that floor `expected` (nvbug-6627795), so measuring against the + # loads that floor `expected` (nvbug 6627795), so measuring against the # raw len() here would make the warning below fire every iteration. - # Read the router's own flag rather than re-deriving the gate, so the - # two can never disagree -- it is off under pipeline parallelism, where - # subtracting requests the router still counted would under-report this - # rank's load instead. num_routable_active_requests = len(self.active_requests) if self.adp_router.exclude_retiring_requests: num_routable_active_requests -= count_retiring_requests( diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index cf42d7bb99be..adfdcd62c7e1 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -758,14 +758,9 @@ def allocation_scope(current_stage: ExecutorMemoryType): if mapping.is_last_pp_rank(): kwargs = { "guided_decoding_config": guided_decoding_config, - # Unconditionally the seat pool. The guided decoder's state - # is indexed by py_seq_slot (guided_decoder.py: grammar_matchers - # [req.seq_slot], the bitmask rows), and py_seq_slot ranges over - # the whole pool -- so sizing this at max_batch_size is an - # IndexError under pipeline parallelism, where admission already - # permits max_batch_size * pp_size live requests. The previous - # conditional made the correct size depend on an unrelated - # attention-DP flag. + # 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, @@ -877,9 +872,6 @@ def allocation_scope(current_stage: ExecutorMemoryType): if model_engine.model.model_config.is_generation: #NOTE: non-generation models do not have kv cache - # Same helper the model engine sizes its seat pool with: this predicate - # feeds KVCacheManagerV2's index-pool coefficient, so an independent copy - # here could disagree with the seat pool by a factor of 2. 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 43499ed294ed..048543427f68 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py @@ -50,54 +50,17 @@ def _num_input_tokens(request) -> int: def is_retiring_request(request) -> bool: - """True if ``request`` has produced its final token and is being torn down. - - A request in ``GENERATION_TO_COMPLETE`` has produced its final token; the - next ``_update_request_states`` tears it down, and no scheduler will ever - put it in a forward batch again (every micro-batch scheduler bounds its - window at ``GENERATION_TO_COMPLETE``). Without the overlap scheduler that - teardown happens in-line, so such requests are gone before the ADP router - next runs. With overlap enabled it is deferred by one iteration, leaving - them in ``active_requests`` when the router builds its load vector -- where - they inflate per-rank load *and* consume global admission capacity that - nothing can use (nvbug-6627795). - - ``GENERATION_TO_COMPLETE`` is the only state treated this way. The disagg - limbo states -- ``DISAGG_CONTEXT_WAIT_SCHEDULER``, - ``DISAGG_GENERATION_INIT``, ``DISAGG_GENERATION_TRANS_IN_PROGRESS``, - ``DISAGG_CONTEXT_TRANS_IN_PROGRESS`` and ``DISAGG_TRANS_ERROR`` -- also - linger in ``active_requests``, but they still own a sequence slot and KV - cache, so they must keep counting as load. - """ + """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`` minus the requests that are already retiring. - - This is the ADP router's view of the active list, and *only* the router's: - the requests are dropped from the load vector, not from the executor. Under - the overlap scheduler ``_process_previous_batch`` -- the only thing that - removes a finished request from ``PyExecutor.active_requests`` -- runs some - two hundred lines *after* ``_fetch_new_requests`` in the same - ``_executor_loop_overlap`` body, so the router would otherwise route - against a list that is one teardown stale (nvbug-6627795). - - Filtering the list at the single ``gather_all_rank_states`` choke point - rather than adjusting a count inside each ``create_rank_state`` corrects - ``num_active_requests`` *and* ``num_active_tokens`` for every router - implementation at once, and keeps the routers ignorant of overlap. - """ + """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 requests that ``build_active_requests_for_overlap`` filters out. - - Used where only the size of the correction is needed, not the list itself - (the dummy-request pad path, which compares the router's ``expected`` - against a rank's routable active count). - """ + """Count the retiring requests in ``active_requests``.""" return sum(1 for req in active_requests if is_retiring_request(req)) @@ -151,15 +114,8 @@ class RankState: """ rank: int - # Routable load only. ``gather_all_rank_states`` hands ``create_rank_state`` - # the overlap-corrected list, so retiring requests are absent from both - # counts below (see ``build_active_requests_for_overlap``). This is what the - # router balances on and what bounds admission. num_active_requests: int = 0 num_active_tokens: int = 0 - # Requests filtered out of the two counts above because they are retiring. - # Reported so the inference loop can tell "nothing routable" from "nothing - # at all" and keep its idle-fetch wait collective (nvbug-6627795). num_retiring_requests: int = 0 iter_stats: RankIterStatsPayload = field(default_factory=RankIterStatsPayload) @@ -229,41 +185,6 @@ class ADPRouter(ABC): def __init__(self, dist: Distributed, has_seq_slot_headroom: bool = True): self.dist = dist - # Whether to route on the overlap-corrected active list (nvbug-6627795). - # - # This is the *model engine's* headroom flag, passed in rather than - # re-derived here, because the two must never disagree: - # - # 1. The correction lets a rank hold more requests than it is charged - # for, so it is only sound where the sequence-slot pool has matching - # headroom. That pool is sized by - # ``_util.compute_max_num_sequences`` from the same flag - # (``should_enable_adp_overlap_seq_slot_headroom``), which withholds - # the headroom from hybrid/SSM architectures whose state-slot pool is - # sized from ``max_batch_size`` alone. Re-deriving the predicate here - # -- as ``not dist.mapping.has_pp()`` used to -- is exactly how the - # router comes to credit a rank with seats the engine never - # allocated. - # 2. Pipeline parallelism is excluded by that same flag, and the reason - # is no longer the sizing -- the seat pool is well defined under PP at - # ``(pp_size + 1) * max_batch_size``. It is that the stages must agree - # on *which* requests are retiring: each rank pops from its own copy - # of the waiting queue, so a per-stage disagreement makes them admit - # different numbers of requests and diverge. Only the last stage marks - # ``GENERATION_TO_COMPLETE`` for generation requests today. Note the - # *context* path in ``_update_request_states_tp`` already evaluates the - # same predicate on every rank, so the asymmetry is only ever in the - # generation path. - # - # Not additionally gated on ``disable_overlap_scheduler`` here: the flag - # already is, and without overlap the retire is not deferred, so no - # request is ever in ``GENERATION_TO_COMPLETE`` when the router runs and - # the filter is arithmetically a no-op (measured: zero such requests on - # 5126/5126 routing records with overlap disabled). - # - # Erring off is the safe direction: excluding fewer retirees admits - # fewer requests (a missed optimization), while excluding more than the - # seat pool covers is a slot exhaustion or a stage divergence. self.exclude_retiring_requests = has_seq_slot_headroom @classmethod @@ -280,13 +201,8 @@ def create( Args: dist: Distributed communicator. has_seq_slot_headroom: Whether the executor's sequence-slot pool was - sized with the overlap headroom - (``should_enable_adp_overlap_seq_slot_headroom``). Required, not - defaulted, because the router's retiring-request correction is - only sound when those extra seats exist -- see - ``__init__``. Passed through from - ``model_engine._enable_adp_overlap_seq_slot_headroom`` so the - sizing and the routing decision cannot drift apart. + 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 @@ -367,12 +283,9 @@ def gather_all_rank_states( iter_stats_payload: Completed previous-iteration stats payload to piggyback on this allgather, if one is pending. """ - # Route on the overlap-corrected list: 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). Applied - # here rather than in each create_rank_state so every router -- and both - # num_active_requests and num_active_tokens -- is corrected at once. - # Disabled under pipeline parallelism; see exclude_retiring_requests. + # 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) @@ -380,11 +293,8 @@ def gather_all_rank_states( active_requests_for_overlap = active_requests num_retiring_requests = 0 local_state = self.create_rank_state(active_requests_for_overlap, new_requests or []) - # The retiring requests are still resident, so the executor loop is NOT - # idle while any of them exists. Report the count so the idle-fetch wait - # stays collective: liveness is a global property, and a rank that - # reported zero would block on the untimed request-queue wait while its - # peers blocked in the allgather. + # 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()) @@ -1117,8 +1027,7 @@ 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 compares `expected` - # against each rank's routable active count (retiring requests excluded, - # matching create_rank_state) and warns if it comes up short. + # 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/suffix_automaton.py b/tensorrt_llm/_torch/speculative/suffix_automaton.py index ef86c23b3d5f..af413526a8b3 100644 --- a/tensorrt_llm/_torch/speculative/suffix_automaton.py +++ b/tensorrt_llm/_torch/speculative/suffix_automaton.py @@ -155,16 +155,17 @@ def __init__( # 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. An explicit global_pool_size is a user contract - # about memory, so it is honoured and validated rather than grown silently. - if sa_config.global_pool_size is not None: - self.pool_size = sa_config.global_pool_size - else: - self.pool_size = max(sa_config.effective_pool_size, self._num_seq_slots) - if self.pool_size < self._num_seq_slots: - raise ValueError( - f"global_pool_size ({self.pool_size}) must be >= the number of " - f"sequence slots ({self._num_seq_slots})" + # 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 attention DP 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 diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 02a0faac355a..194cb4ea5fbd 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -583,11 +583,9 @@ def seat_pool_or_none(model_engine) -> Optional[int]: micro-batch scheduler caps every forward at max_batch_size. Gated on the same flag ``_set_up_spec_metadata`` reads, so every spec-decoding - pool agrees with the metadata about which number it is indexed by. Returning - None (headroom off) preserves the established max_batch_size sizing. + pool agrees with the metadata about which number it is indexed by. """ - if not getattr(model_engine, "_enable_adp_overlap_seq_slot_headroom", - False): + if not getattr(model_engine, "_enable_disagg_adp_overlap_headroom", False): return None return getattr(model_engine, "max_num_seq_slots", None) @@ -769,14 +767,12 @@ def get_spec_drafter(model_engine, max_num_requests = model_engine.batch_size # The draft loop runs its own slot pool, but the indices it hands out address - # buffers sized by the *target* engine's seat pool: the shared sampler - # (instantiate_sampler), the draft KV cache manager's IndexMapper - # (KvCacheCreator._target_max_num_seq_slots) and spec_resource_manager above. - # It must therefore be sized from the same number. It is also load-bearing, - # not merely tidy: the previous draft batch's slots are released by - # cleanup_previous_draft_resources a full iteration later - # (py_executor.py:5347), so a pool of max_batch_size raises NoFreeSlotsError - # precisely when the overlap headroom is doing its job. + # buffers sized by the *target* engine's seat pool (the shared sampler and + # spec_resource_manager above), so it must be sized from the same number. The + # previous draft batch's slots are also released by + # cleanup_previous_draft_resources a full iteration later, so a pool of + # max_batch_size raises NoFreeSlotsError precisely when the overlap headroom + # is doing its job. draft_slots = seat_pool_or_none(model_engine) or max_num_requests if spec_config.spec_dec_mode.is_draft_target( ) or spec_config.spec_dec_mode.is_eagle3( 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 b93d29f0ad10..7c995283b5dc 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,11 +1158,10 @@ def test_separate_one_model_draft_normalizes_target_pool_ratio() -> None: creator._is_disagg = False creator._mapping = Mock() creator._speculative_config = Mock() - # The draft manager is now sized from the *target* engine's published seat - # pool, so this hand-built creator has to carry the attribute the real one - # sets in __init__. None means "no pool published", which is the pre-existing - # behaviour this test asserts. - creator._model_engine = None + # Every manager is now told whether the overlap scheduler is on, so this + # hand-built creator has to carry the attribute the real one sets in + # __init__. False is the pre-existing sizing this test asserts. + creator._enable_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 fade285df8fc..e170b3fd54d1 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 @@ -1550,22 +1550,14 @@ def _index_mapper_capacity_for( pp_size: int = 1, is_disagg: bool = False, num_reserved_index_slots: int = 1, - max_num_seq_slots: int | None = None, -) -> tuple[int, int, int | None]: - """Construct a manager and return (IndexMapper capacity, page-table capacity, - published ``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 same number net of the - reserved slots, published for the startup validator so that the seat pool and - the index pool can be compared without either being re-derived. - - The third is read with ``getattr(..., None)`` rather than as an attribute so - that a manager which does not publish it at all fails only the test that is - *about* publishing it. Asserting it here would make every capacity row fail - for one and the same trivial reason, which would destroy the negative - control's ability to say *which* topologies the coefficient actually moved. + enable_attention_dp: bool = False, + enable_overlap_scheduler: bool = False, +) -> tuple[int, int]: + """Construct a manager and return (IndexMapper capacity, page-table capacity). + + The 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. """ module = "tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2" fake_impl = Mock() @@ -1592,12 +1584,13 @@ def build_base_config( patch.object(KVCacheManagerV2, "_prepare_page_table_tensor") as page_table, patch.object(KVCacheManagerV2, "_log_kv_cache_pool_lifecycle_mapping"), ): - manager = KVCacheManagerV2( + 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, num_reserved_index_slots and max_num_seq_slots. + # enable_attention_dp, is_disagg, enable_overlap_scheduler and + # num_reserved_index_slots. KvCacheConfig(max_gpu_total_bytes=16 << 20), CacheType.SELFKONLY, num_layers=1, @@ -1606,129 +1599,78 @@ def build_base_config( 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), + mapping=Mapping( + world_size=pp_size, + rank=0, + tp_size=1, + pp_size=pp_size, + enable_attention_dp=enable_attention_dp, + ), dtype=DataType.HALF, vocab_size=16, execution_stream=Mock(), is_disagg=is_disagg, num_reserved_index_slots=num_reserved_index_slots, - max_num_seq_slots=max_num_seq_slots, + enable_overlap_scheduler=enable_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], - getattr(manager, "max_admissible_sequences", None), ) -# (max_batch_size, pp_size, is_disagg, reserved, max_num_seq_slots, expected capacity) +# (max_batch_size, pp_size, adp, is_disagg, overlap, reserved, expected capacity) # -# The seat pool (`PyTorchModelEngine.max_num_seq_slots`) and the index mapper are -# both per-request leases, so the mapper must cover whatever the executor can -# admit. Rows 1-2 are the nvbug 6627795 case: aggregated attention DP under the -# overlap scheduler doubles the seat pool, and before the fix the mapper stayed -# at B+1 and silently deferred requests. +# capacity == max_batch_size * pp_size * (2 if is_disagg or (adp and overlap and +# not pp) else 1) + reserved. Rows 1-2 are the nvbug 6627795 case: under +# attention DP the overlap scheduler defers the retiring batch'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. _INDEX_MAPPER_CAPACITY_CASES = [ - # aggregated + ADP + overlap: seats are 2B, so the mapper must be 2B too. - pytest.param(2, 1, False, 1, 4, 5, id="agg_adp_overlap"), - pytest.param(8, 1, False, 1, 16, 17, id="agg_adp_overlap_b8"), - # No seat headroom (no ADP, or overlap off): unchanged at B+1. - pytest.param(2, 1, False, 1, 2, 3, id="agg_no_headroom"), - # Negative control: callers that pass nothing keep the pre-fix allocation. - pytest.param(2, 1, False, 1, None, 3, id="seats_unset_is_unchanged"), - # Negative control: disagg does not compound with the seat headroom. The two - # coefficients cover the same extra cohort, so this is max(4, 4)+1, not 4B+1. - pytest.param(2, 1, True, 1, 4, 5, id="disagg_does_not_compound"), - # Disagg without seat headroom keeps its own 2x. - pytest.param(2, 1, True, 1, None, 5, id="disagg_only"), - # The same cell as it actually occurs once the seat pool is plumbed: disagg - # without ADP seats B*pp == 2, and the mapper still carries its own 2x. The - # surplus is deliberate and is the reason the 2x lives *here* and not in - # compute_max_num_sequences -- a request in KV transfer holds an index lease - # while SeqSlotManager.prepare_resources skips it, so it occupies no seat. - pytest.param(2, 1, True, 1, 2, 5, id="disagg_index_pool_exceeds_seat_pool"), - # PP without the headroom: seats are B*pp, which the mapper already matched. - pytest.param(2, 4, False, 1, 8, 9, id="pp4"), - # PP with the headroom: seats are (pp+1)*B == 10, additive rather than - # 2*B*pp == 16. The mapper follows the seat pool verbatim, so this row is what - # makes the extra generation of seats usable under pipeline parallelism. - pytest.param(2, 4, False, 1, 10, 11, id="pp4_adp_overlap"), + # ADP + overlap, no PP: both cohorts are resident, so the mapper needs 2B. + pytest.param(2, 1, True, False, True, 1, 5, id="adp_overlap"), + pytest.param(8, 1, True, False, True, 1, 17, id="adp_overlap_b8"), + # Either half of the conjunction missing leaves the pre-fix allocation. + pytest.param(2, 1, True, False, False, 1, 3, id="adp_no_overlap"), + pytest.param(2, 1, False, False, True, 1, 3, id="overlap_no_adp"), + # Disagg already carried its own 2x; the two coefficients cover the same + # extra cohort, so they do not compound. + pytest.param(2, 1, True, True, True, 1, 5, id="disagg_does_not_compound"), + pytest.param(2, 1, False, True, False, 1, 5, id="disagg_only"), + # Pipeline parallelism is out of scope: max_batch_size * pp_size already + # covers the in-flight micro-batches, so the ADP coefficient stays off. + pytest.param(2, 4, True, False, True, 1, 9, id="pp4_adp_overlap"), + pytest.param(2, 4, False, False, False, 1, 9, id="pp4_plain"), + # ... while the pre-existing disagg 2x under PP is left exactly as it was. + pytest.param(2, 4, False, True, False, 1, 17, id="pp4_disagg_unchanged"), # Reserved slots are still added on top of the widened pool. - pytest.param(2, 1, False, 5, 4, 9, id="reserved_slots_still_added"), - # A seat pool smaller than the mapper's own floor must never shrink it. - pytest.param(4, 1, False, 1, 1, 5, id="small_seat_pool_does_not_shrink"), + pytest.param(2, 1, True, False, True, 5, 9, id="reserved_slots_still_added"), ] @pytest.mark.cpu_only @pytest.mark.parametrize( - "max_batch_size,pp_size,is_disagg,reserved,max_num_seq_slots,expected", + "max_batch_size,pp_size,adp,is_disagg,overlap,reserved,expected", _INDEX_MAPPER_CAPACITY_CASES, ) -def test_index_mapper_capacity_covers_seq_slot_pool( +def test_index_mapper_capacity_covers_the_overlapping_cohorts( max_batch_size: int, pp_size: int, + adp: bool, is_disagg: bool, + overlap: bool, reserved: int, - max_num_seq_slots: int | None, expected: int, ) -> None: - capacity, page_table_capacity, _ = _index_mapper_capacity_for( + capacity, page_table_capacity = _index_mapper_capacity_for( max_batch_size=max_batch_size, pp_size=pp_size, is_disagg=is_disagg, num_reserved_index_slots=reserved, - max_num_seq_slots=max_num_seq_slots, + enable_attention_dp=adp, + enable_overlap_scheduler=overlap, ) assert capacity == expected assert page_table_capacity == expected - - -# The rows where the *published* number is worth stating separately from the -# capacity arithmetic: the aggregated cell nvbug 6627795 was filed for, the new -# pipeline-parallel cell, and a caller that supplies no seat pool at all. -@pytest.mark.cpu_only -@pytest.mark.parametrize( - "max_batch_size,pp_size,is_disagg,reserved,max_num_seq_slots,expected_admissible", - [ - pytest.param(2, 1, False, 1, 4, 4, id="agg_adp_overlap"), - pytest.param(2, 4, False, 1, 10, 10, id="pp4_adp_overlap"), - pytest.param(2, 1, False, 5, None, 2, id="seats_unset_reserved_excluded"), - # Disagg is the one case where the published number legitimately exceeds - # the seat pool it is compared against. - pytest.param(2, 1, True, 1, 2, 4, id="disagg_exceeds_seat_pool"), - ], -) -def test_index_mapper_publishes_max_admissible_sequences( - max_batch_size: int, - pp_size: int, - is_disagg: bool, - reserved: int, - max_num_seq_slots: int | None, - expected_admissible: int, -) -> None: - """The manager publishes what it can index, so nobody has to re-derive it. - - ``_util.validate_seq_slot_pool_covers_admission`` compares this against the - executor's sequence-slot pool at startup. It is the pool *net of* the reserved - padding/dummy slots, which no real sequence can take. Aggregated, it equals - the seat pool exactly on every row where one was plumbed through, and the - validator enforces that in both directions. Under disaggregation it may - exceed the seat pool, because an index lease outlives the seat while the KV - transfer drains -- so there the validator enforces only ``>=``. - """ - _, _, admissible = _index_mapper_capacity_for( - max_batch_size=max_batch_size, - pp_size=pp_size, - is_disagg=is_disagg, - num_reserved_index_slots=reserved, - max_num_seq_slots=max_num_seq_slots, - ) - assert admissible == expected_admissible - if max_num_seq_slots is not None: - if is_disagg: - assert admissible >= max_num_seq_slots - else: - assert admissible == max_num_seq_slots diff --git a/tests/unittest/_torch/executor/test_adp_router.py b/tests/unittest/_torch/executor/test_adp_router.py index 1adb9d591bf4..38b9bae8dd15 100644 --- a/tests/unittest/_torch/executor/test_adp_router.py +++ b/tests/unittest/_torch/executor/test_adp_router.py @@ -445,15 +445,10 @@ def test_gather_all_rank_states_keeps_retiring_without_headroom(self): assert states[0].num_active_tokens == 600 def test_exclude_retiring_requests_follows_the_seat_pool_headroom(self): - # The flag is the single gate; _pad_attention_dp_dummy_request reads it - # rather than re-deriving the predicate, so the two cannot drift. - # - # It tracks the *engine's* headroom flag and nothing else. It used to be - # `not dist.mapping.has_pp()`, which was a second derivation of the same - # fact: correct only as long as the sizing gate happened to exclude - # exactly PP. Once the sizing gate also excluded hybrid architectures the - # two disagreed, and the router credited a rank with seats that were - # never allocated. Pipeline parallelism is now in scope on both sides. + # 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 ( diff --git a/tests/unittest/_torch/executor/test_seq_slot_sizing.py b/tests/unittest/_torch/executor/test_seq_slot_sizing.py index fb4ff971d08e..29e0943e2555 100644 --- a/tests/unittest/_torch/executor/test_seq_slot_sizing.py +++ b/tests/unittest/_torch/executor/test_seq_slot_sizing.py @@ -9,14 +9,12 @@ backfilled their seats. Transient slot demand is therefore one extra micro-batch worth of slots, regardless of whether speculative decoding is enabled. The headroom is selected from runtime topology, not model -architecture -- except that hybrid/SSM architectures are excluded, because -their state-slot pool is not sized from this number. +architecture. -That extra generation is **additive** in pp_size, not multiplicative: -pipeline depth already accounts for the pp_size micro-batches structurally -in flight, and the overlap deferral is one iteration on top of them. So the -pool is (pp_size + 1) * max_batch_size, which at pp_size == 1 coincides with -the historical 2 * max_batch_size. +Pipeline parallelism is out of scope: the pool is already sized by pp_size +there, and the ADP router's retiring-request correction is not rank-consistent +because only the last pipeline stage marks generation requests +GENERATION_TO_COMPLETE. compute_max_num_sequences is the single sizing implementation used both for the executor's SeqSlotManager pool (create_py_executor_instance) and @@ -37,63 +35,50 @@ is_disagg_enabled, resolve_max_num_sequences, should_enable_adp_dummy_fixes, - should_enable_adp_overlap_seq_slot_headroom, + should_enable_disagg_adp_overlap_headroom, should_enable_non_overlap_adp_forward_intent, should_enable_scheduler_aware_adp_dummy, - validate_seq_slot_pool_covers_admission, ) from tensorrt_llm.mapping import Mapping -# (pp_size, disable_overlap, enable_overlap_headroom, expected_factor) -# -# The terms are additive, not multiplicative: pipeline depth costs pp_size -# micro-batches of seats, and the overlap deferral costs exactly one more -# generation on top -- not one more per stage. At pp_size == 1 the two readings -# coincide at 2x, which is why the headroom used to be expressible as a factor of -# 2; the pp>1 rows are where they part company (5x, not 8x, at pp=4). +_UCX = SimpleNamespace(backend="UCX") + +# (pp_size, enable_overlap_headroom, expected_factor) # # Disaggregation is absent from this table on purpose -- see # test_seat_pool_has_no_disagg_term. SIZING_CASES = [ - # No PP. Every cell here is unchanged by this PR. - (1, False, False, 1), - (1, False, True, 2), - (1, True, True, 1), - # PP without the headroom: unchanged. - (4, False, False, 4), - (4, True, True, 4), - # PP with the headroom: the one intended behaviour change (was pp_size). - (2, False, True, 3), - (4, False, True, 5), + # No PP: the headroom buys one extra micro-batch worth of seats. + (1, False, 1), + (1, True, 2), + # PP sizes the pool by pipeline depth and ignores the headroom. + (4, False, 4), + (2, True, 2), + (4, True, 4), ] @pytest.mark.parametrize( - "enable_attention_dp,pp_size,disable_overlap,is_hybrid,expected", + "enable_attention_dp,pp_size,cache_transceiver_config,disable_overlap,expected", [ - # No cache-transceiver term: the gate no longer looks at disaggregation - # at all, because nvbug-6627795 reproduced on an aggregated context-only - # run with no transceiver configured. - (True, 1, False, False, True), - (False, 1, False, False, False), - (True, 1, True, False, False), - # Pipeline parallelism is still out of scope, but no longer because the - # sizing cannot express it -- compute_max_num_sequences is additive, so - # (pp_size + 1) * max_batch_size is well defined. The missing piece is the - # consumer: only the last pipeline stage marks generation requests - # GENERATION_TO_COMPLETE, so the ADP router's retiring-request correction - # would not be rank-consistent. - (True, 2, False, False, False), - (True, 4, False, False, False), - # Hybrid/SSM architectures are excluded: MambaHybridCacheManagerV2 sizes - # its state-index pool from max_batch_size alone, so an extra seat would - # have no state slot behind it. - (True, 1, False, True, False), - (True, 4, False, True, False), + # Aggregated ADP with overlap on: nvbug 6627795 reproduced here, on a + # context-only run with no cache transceiver configured. + (True, 1, None, False, True), + (False, 1, None, False, False), + # Aggregated ADP with overlap off: teardown is in-line, no headroom. + (True, 1, None, True, False), + # Disagg needs the headroom even with overlap off: a request awaiting its + # KV transfer keeps its lease. + (True, 1, _UCX, True, True), + (True, 1, _UCX, False, True), + (False, 1, _UCX, False, False), + # Pipeline parallelism is out of scope in both directions. + (True, 2, None, False, False), + (True, 4, _UCX, False, False), ], ) -def test_adp_overlap_seq_slot_headroom_gate( - enable_attention_dp, pp_size, disable_overlap, is_hybrid, expected +def test_disagg_adp_overlap_headroom_gate( + enable_attention_dp, pp_size, cache_transceiver_config, disable_overlap, expected ): mapping = Mapping( world_size=pp_size, @@ -103,7 +88,9 @@ def test_adp_overlap_seq_slot_headroom_gate( ) assert ( - should_enable_adp_overlap_seq_slot_headroom(mapping, disable_overlap, is_hybrid=is_hybrid) + should_enable_disagg_adp_overlap_headroom( + mapping, cache_transceiver_config, disable_overlap + ) is expected ) @@ -142,11 +129,9 @@ def test_non_overlap_adp_forward_intent_scope(pp_size, disable_overlap, expected assert should_enable_non_overlap_adp_forward_intent(mapping, disable_overlap) is expected -@pytest.mark.parametrize( - "pp_size,disable_overlap,enable_overlap_headroom,expected_factor", SIZING_CASES -) +@pytest.mark.parametrize("pp_size,enable_overlap_headroom,expected_factor", SIZING_CASES) def test_compute_max_num_sequences_scopes_overlap_headroom( - pp_size, disable_overlap, enable_overlap_headroom, expected_factor + pp_size, enable_overlap_headroom, expected_factor ): max_batch_size = 8 mapping = Mapping(world_size=pp_size, tp_size=1, pp_size=pp_size) @@ -154,7 +139,7 @@ def test_compute_max_num_sequences_scopes_overlap_headroom( compute_max_num_sequences( mapping, max_batch_size, - disable_overlap, + disable_overlap_scheduler=False, enable_overlap_headroom=enable_overlap_headroom, ) == max_batch_size * expected_factor @@ -162,16 +147,13 @@ 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. + """The disaggregation 2x reaches the seat pool only through the gate. 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. + requests outright and only seats one once its transmission completes. So the + sizing function itself must not carry a disaggregation term -- the single + ``enable_overlap_headroom`` flag is the only way in. Asserting on the signature rather than on a return value is deliberate: a value test cannot distinguish "the parameter is gone" from "the parameter @@ -213,12 +195,12 @@ def test_resolve_max_num_sequences_prefers_the_published_pool(explicit, engine_s 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 very skew this number exists to remove. - The third row is that branch, and it must still land on the headroom value. + 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_adp_overlap_seq_slot_headroom=True, + _enable_disagg_adp_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) @@ -299,84 +281,7 @@ def test_sampler_uses_executor_slot_pool_capacity(slot_factor): assert args.max_num_sequences == max_num_sequences -class _FakeManager: - """Stands in for a KV cache manager that publishes its seating capacity.""" - - def __init__(self, max_admissible_sequences): - self.max_admissible_sequences = max_admissible_sequences - - -@pytest.mark.parametrize("is_disagg", [False, True]) -def test_validator_accepts_the_matching_pair(is_disagg): - validate_seq_slot_pool_covers_admission(16, _FakeManager(16), is_disagg=is_disagg) - - -@pytest.mark.parametrize("is_disagg", [False, True]) -def test_validator_always_rejects_an_index_pool_below_the_seat_pool(is_disagg): - """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. - Disaggregation is no excuse here -- its 2x makes the index pool *larger*, so - a shortfall under disagg means the two numbers were derived separately. - """ - with pytest.raises(ValueError, match="smaller than the seat"): - validate_seq_slot_pool_covers_admission(16, _FakeManager(8), is_disagg=is_disagg) - - -def test_validator_rejects_a_larger_index_pool_only_when_aggregated(): - """The other direction is a bug when aggregated and by design under disagg. - - Aggregated, every indexed sequence is also a seated one, so a surplus can - only mean one of the two numbers was re-derived -- and a request would be - admitted that cannot be seated, with ``SlotManager.add_slot`` raising on the - executor's event-loop thread. Under disaggregation the surplus *is* the - mechanism: a request in KV transfer holds its index lease with no seat - (``SeqSlotManager.prepare_resources`` skips ``DISAGG_GENERATION_INIT``), so - the index pool carries a 2x that must not reach the seat pool. Admission is - bounded independently at ``max_batch_size * pp_size`` either way, so the - surplus is never extra concurrency. - """ - with pytest.raises(ValueError, match="larger than the seat"): - validate_seq_slot_pool_covers_admission(16, _FakeManager(32), is_disagg=False) - - validate_seq_slot_pool_covers_admission(16, _FakeManager(32), is_disagg=True) - - -@pytest.mark.parametrize( - "manager", - [ - None, # non-generation models have no KV cache manager - SimpleNamespace(), # V1 sizes its index pool by an unrelated rule - # A test double auto-creates every attribute, so "absent" reaches the - # validator as a non-integral value rather than as None. Ordering it - # against an int raises TypeError, which is a crash in an unrelated - # caller's test rather than a finding about this PR -- so non-integral - # has to mean the same thing as absent. - Mock(), - _FakeManager(None), - ], -) -def test_validator_skips_managers_that_do_not_publish_capacity(manager): - """Opt-in, not guessed at. - - A manager whose index pool is sized by some other rule -- V1, which receives - none of this plumbing and re-derives from bare max_batch_size -- must not be - compared against a number it never consumed. - - Hybrid is deliberately *not* an example here: MambaHybridCacheManagerV2 - derives from KVCacheManagerV2 and therefore does publish - max_admissible_sequences, so it *is* validated. That comparison holds - because the headroom is withheld from hybrid on both sides -- seats are - B*pp and, with max_num_seq_slots withheld, its index pool lands on exactly - the same number (2*B*pp under disagg, which the validator allows). What - hybrid does not receive is the seat pool, not the check. - """ - validate_seq_slot_pool_covers_admission(16, manager) - - -def _make_kv_cache_creator(max_num_seq_slots) -> KvCacheCreator: +def _make_kv_cache_creator(enable_overlap_scheduler: bool) -> 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) @@ -391,27 +296,29 @@ def _make_kv_cache_creator(max_num_seq_slots) -> KvCacheCreator: c._kv_connector_manager = None c._execution_stream = None c._is_disagg = False + c._enable_overlap_scheduler = enable_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) - c._model_engine = SimpleNamespace(max_num_seq_slots=max_num_seq_slots) return c -@pytest.mark.parametrize("max_num_seq_slots", [8, 16, None]) -def test_kv_cache_manager_receives_executor_seq_slot_pool(max_num_seq_slots): - """The seat pool size must reach the KV cache manager verbatim. +@pytest.mark.parametrize("enable_overlap_scheduler", [False, True]) +def test_kv_cache_manager_receives_the_overlap_flag(enable_overlap_scheduler): + """The manager sizes its own index pool, but needs the overlap flag to do it. - The manager sizes its IndexMapper from this number so that every sequence - the executor can admit is guaranteed an index. Recomputing the coefficient - inside the manager would let the two pools drift apart (nvbug 6627795). + The index pool must cover both the retiring cohort and its replacement when + attention DP runs with the overlap scheduler, otherwise ``_create_kv_cache`` + silently defers admitted requests one at a time (nvbug 6627795). The flag is + passed rather than the seat-pool size so the manager keeps deriving its + capacity from ``max_batch_size * pp_size``, which is not comparable with a + seat pool that also carries the PP multiplier. """ - creator = _make_kv_cache_creator(max_num_seq_slots) + creator = _make_kv_cache_creator(enable_overlap_scheduler) model_engine = SimpleNamespace( model=SimpleNamespace(model_config=SimpleNamespace(is_generation=True)), - max_num_seq_slots=max_num_seq_slots, ) with patch( @@ -420,28 +327,4 @@ def test_kv_cache_manager_receives_executor_seq_slot_pool(max_num_seq_slots): ) as create: creator._create_kv_cache_manager(model_engine) - assert create.call_args.kwargs["max_num_seq_slots"] == max_num_seq_slots - - -def test_draft_manager_uses_the_target_engines_seq_slot_pool(): - """A draft engine's own seat count must not size the draft index pool. - - The executor has a single SeqSlotManager, sized from the target engine. Two- - model speculative decoding builds the draft KV cache manager by passing the - *draft* engine to the same helper; if that engine's (smaller) number were - used, the draft manager's IndexMapper would become the new bottleneck and - reintroduce the silent deferral this sizing exists to prevent. - """ - creator = _make_kv_cache_creator(16) - draft_engine = SimpleNamespace( - model=SimpleNamespace(model_config=SimpleNamespace(is_generation=True)), - max_num_seq_slots=8, - ) - - with patch( - "tensorrt_llm._torch.pyexecutor._util._create_kv_cache_manager", - return_value=None, - ) as create: - creator._create_kv_cache_manager(draft_engine) - - assert create.call_args.kwargs["max_num_seq_slots"] == 16 + assert create.call_args.kwargs["enable_overlap_scheduler"] is enable_overlap_scheduler diff --git a/tests/unittest/_torch/speculative/test_spec_slot_pool_sizing.py b/tests/unittest/_torch/speculative/test_spec_slot_pool_sizing.py index ba77148da059..86c18220b885 100644 --- a/tests/unittest/_torch/speculative/test_spec_slot_pool_sizing.py +++ b/tests/unittest/_torch/speculative/test_spec_slot_pool_sizing.py @@ -252,7 +252,7 @@ def _drafter_engine(headroom: bool, seats: int = POOL): spec_config=spec_config, batch_size=R, max_num_seq_slots=seats, - _enable_adp_overlap_seq_slot_headroom=headroom, + _enable_disagg_adp_overlap_headroom=headroom, ) @@ -510,17 +510,19 @@ def test_sa_pool_survives_a_full_overlap_turnover(): @pytest.mark.cpu_only -def test_an_explicit_sa_pool_is_honoured_but_validated(): - """``global_pool_size`` is a memory contract, so it is never grown silently. - - It is validated against the sequence-slot pool instead, turning what would be - a mid-run slot exhaustion into a startup error that names the real bound. +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 attention DP 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. """ - grown = _sa_manager(POOL, enable_global_pool=True, global_pool_size=64) - assert grown.pool_size == 64 - - with pytest.raises(ValueError, match="sequence slots"): - _sa_manager(POOL, enable_global_pool=True, global_pool_size=R) + 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") From 4088a73c3c5c1c47c08c29cf9be983882f943e89 Mon Sep 17 00:00:00 2001 From: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> Date: Wed, 9 Sep 2026 01:02:56 -0700 Subject: [PATCH 18/22] [https://nvbugs/6627795][fix] name the new flag disable_overlap_scheduler The plumbing added for the attention-DP overlap headroom carried the flag as enable_overlap_scheduler, which inverts the config key users actually set. Thread disable_overlap_scheduler through instead so the plumbing, LlmArgs and the YAML all agree on one polarity. Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/_util.py | 15 +++---- .../kv_cache/kv_cache_manager_v2.py | 6 +-- .../kv_cache/test_kv_cache_estimation.py | 6 +-- .../kv_cache/test_kv_cache_manager_v2.py | 45 ++++++++++--------- .../_torch/executor/test_seq_slot_sizing.py | 12 ++--- 5 files changed, 42 insertions(+), 42 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 8fa818d20feb..501eb63be102 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -642,7 +642,7 @@ def __init__( self._dummy_encoder_inputs: List[MultimodalParams] = [] self._profiling_stage_data = profiling_stage_data self._is_disagg = is_disagg - self._enable_overlap_scheduler = not llm_args.disable_overlap_scheduler + self._disable_overlap_scheduler = llm_args.disable_overlap_scheduler self._cache_transceiver_config = llm_args.cache_transceiver_config self._execution_stream = execution_stream self._kv_cache_manager_cls = self._get_model_kv_cache_manager_cls( @@ -1463,7 +1463,7 @@ def _create_kv_cache_manager( execution_stream=self._execution_stream, layer_mask=spec_dec_layer_mask, is_disagg=self._is_disagg, - enable_overlap_scheduler=self._enable_overlap_scheduler, + disable_overlap_scheduler=self._disable_overlap_scheduler, cold_page_codec_provider=cold_page_codec_provider, joint_kv_cache_reuse=self._joint_kv_cache_reuse, ) @@ -1670,7 +1670,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, - enable_overlap_scheduler=self._enable_overlap_scheduler, + disable_overlap_scheduler=self._disable_overlap_scheduler, cold_page_codec_provider=cold_page_codec_provider, joint_kv_cache_reuse=self._joint_kv_cache_reuse, ) @@ -2047,7 +2047,7 @@ def _create_cross_kv_cache_manager( num_layers=num_layers, num_kv_heads=num_kv_heads, head_dim=head_dim, - enable_overlap_scheduler=self._enable_overlap_scheduler, + disable_overlap_scheduler=self._disable_overlap_scheduler, kv_cache_type=tensorrt_llm.bindings.internal.batch_manager. CacheType.CROSS, ) @@ -2367,7 +2367,7 @@ def _create_kv_cache_manager( head_dim: Optional[int] = None, kv_cache_type=None, is_disagg: bool = False, - enable_overlap_scheduler: bool = False, + disable_overlap_scheduler: bool = False, cold_page_codec_provider: Optional[object] = None, joint_kv_cache_reuse: bool = False) -> KVCacheManager: """ @@ -2510,7 +2510,7 @@ def _create_kv_cache_manager( "cold_page_codec_provider"] = cold_page_codec_provider manager_extra_kwargs["joint_kv_cache_reuse"] = joint_kv_cache_reuse manager_extra_kwargs[ - "enable_overlap_scheduler"] = enable_overlap_scheduler + "disable_overlap_scheduler"] = disable_overlap_scheduler if issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2): manager_extra_kwargs["is_disagg"] = is_disagg @@ -3132,9 +3132,8 @@ def should_enable_disagg_adp_overlap_headroom( pipeline stage to agree on which requests are retiring. """ is_disagg = is_disagg_enabled(cache_transceiver_config) - enable_overlap_scheduler = not disable_overlap_scheduler return (mapping.enable_attention_dp and not mapping.has_pp() - and (is_disagg or enable_overlap_scheduler)) + and (is_disagg or not disable_overlap_scheduler)) def create_py_executor_instance( 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 8b84f862eb88..42c3981fb627 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 @@ -931,7 +931,7 @@ def __init__( is_disagg: bool = False, enable_stats: bool = False, num_reserved_index_slots: int = 1, - enable_overlap_scheduler: bool = False, + disable_overlap_scheduler: bool = False, is_estimating_kv_cache: bool = False, cold_page_codec_provider: Optional[object] = None, joint_kv_cache_reuse: bool = False, @@ -1463,7 +1463,7 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: # `max_num_sequences` already scales with the number of in-flight # microbatches. needs_extra_index_slots = is_disagg or ( - mapping.enable_attention_dp and enable_overlap_scheduler and not mapping.has_pp() + mapping.enable_attention_dp and not disable_overlap_scheduler and not mapping.has_pp() ) max_num_sequences = max_batch_size * mapping.pp_size assert num_reserved_index_slots >= 0, "num_reserved_index_slots must be non-negative" @@ -1473,7 +1473,7 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: logger.info( f"KVCacheManagerV2: IndexMapper capacity={index_mapper_capacity} " f"(max_num_sequences={max_num_sequences}, is_disagg={is_disagg}, " - f"enable_overlap_scheduler={enable_overlap_scheduler}, " + f"disable_overlap_scheduler={disable_overlap_scheduler}, " f"num_reserved_index_slots={num_reserved_index_slots}, " f"max_beam_width={max_beam_width})" ) 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 7c995283b5dc..efe1732f9c13 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,10 +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 manager is now told whether the overlap scheduler is on, so this + # Every manager is now told whether the overlap scheduler is off, so this # hand-built creator has to carry the attribute the real one sets in - # __init__. False is the pre-existing sizing this test asserts. - creator._enable_overlap_scheduler = False + # __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 e170b3fd54d1..5c0ff295f729 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 @@ -1551,7 +1551,7 @@ def _index_mapper_capacity_for( is_disagg: bool = False, num_reserved_index_slots: int = 1, enable_attention_dp: bool = False, - enable_overlap_scheduler: bool = False, + disable_overlap_scheduler: bool = False, ) -> tuple[int, int]: """Construct a manager and return (IndexMapper capacity, page-table capacity). @@ -1589,7 +1589,7 @@ def build_base_config( # ("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, - # enable_attention_dp, is_disagg, enable_overlap_scheduler and + # enable_attention_dp, is_disagg, disable_overlap_scheduler and # num_reserved_index_slots. KvCacheConfig(max_gpu_total_bytes=16 << 20), CacheType.SELFKONLY, @@ -1611,7 +1611,7 @@ def build_base_config( execution_stream=Mock(), is_disagg=is_disagg, num_reserved_index_slots=num_reserved_index_slots, - enable_overlap_scheduler=enable_overlap_scheduler, + disable_overlap_scheduler=disable_overlap_scheduler, ) index_mapper_cls.assert_called_once() page_table.assert_called_once() @@ -1621,38 +1621,39 @@ def build_base_config( ) -# (max_batch_size, pp_size, adp, is_disagg, overlap, reserved, expected capacity) +# (max_batch_size, pp_size, adp, is_disagg, disable_overlap, reserved, expected) # -# capacity == max_batch_size * pp_size * (2 if is_disagg or (adp and overlap and -# not pp) else 1) + reserved. Rows 1-2 are the nvbug 6627795 case: under -# attention DP the overlap scheduler defers the retiring batch'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. +# capacity == max_batch_size * pp_size * (2 if is_disagg or (adp and not +# disable_overlap and not pp) else 1) + reserved. Rows 1-2 are the nvbug 6627795 +# case: under attention DP the overlap scheduler defers the retiring batch'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. _INDEX_MAPPER_CAPACITY_CASES = [ # ADP + overlap, no PP: both cohorts are resident, so the mapper needs 2B. - pytest.param(2, 1, True, False, True, 1, 5, id="adp_overlap"), - pytest.param(8, 1, True, False, True, 1, 17, id="adp_overlap_b8"), + pytest.param(2, 1, True, False, False, 1, 5, id="adp_overlap"), + pytest.param(8, 1, True, False, False, 1, 17, id="adp_overlap_b8"), # Either half of the conjunction missing leaves the pre-fix allocation. - pytest.param(2, 1, True, False, False, 1, 3, id="adp_no_overlap"), - pytest.param(2, 1, False, False, True, 1, 3, id="overlap_no_adp"), + pytest.param(2, 1, True, False, True, 1, 3, id="adp_no_overlap"), + pytest.param(2, 1, False, False, False, 1, 3, id="overlap_no_adp"), # Disagg already carried its own 2x; the two coefficients cover the same # extra cohort, so they do not compound. - pytest.param(2, 1, True, True, True, 1, 5, id="disagg_does_not_compound"), - pytest.param(2, 1, False, True, False, 1, 5, id="disagg_only"), + pytest.param(2, 1, True, True, False, 1, 5, id="disagg_does_not_compound"), + pytest.param(2, 1, False, True, True, 1, 5, id="disagg_only"), # Pipeline parallelism is out of scope: max_batch_size * pp_size already # covers the in-flight micro-batches, so the ADP coefficient stays off. - pytest.param(2, 4, True, False, True, 1, 9, id="pp4_adp_overlap"), - pytest.param(2, 4, False, False, False, 1, 9, id="pp4_plain"), + pytest.param(2, 4, True, False, False, 1, 9, id="pp4_adp_overlap"), + pytest.param(2, 4, False, False, True, 1, 9, id="pp4_plain"), # ... while the pre-existing disagg 2x under PP is left exactly as it was. - pytest.param(2, 4, False, True, False, 1, 17, id="pp4_disagg_unchanged"), + pytest.param(2, 4, False, True, True, 1, 17, id="pp4_disagg_unchanged"), # Reserved slots are still added on top of the widened pool. - pytest.param(2, 1, True, False, True, 5, 9, id="reserved_slots_still_added"), + pytest.param(2, 1, True, False, False, 5, 9, id="reserved_slots_still_added"), ] @pytest.mark.cpu_only @pytest.mark.parametrize( - "max_batch_size,pp_size,adp,is_disagg,overlap,reserved,expected", + "max_batch_size,pp_size,adp,is_disagg,disable_overlap,reserved,expected", _INDEX_MAPPER_CAPACITY_CASES, ) def test_index_mapper_capacity_covers_the_overlapping_cohorts( @@ -1660,7 +1661,7 @@ def test_index_mapper_capacity_covers_the_overlapping_cohorts( pp_size: int, adp: bool, is_disagg: bool, - overlap: bool, + disable_overlap: bool, reserved: int, expected: int, ) -> None: @@ -1670,7 +1671,7 @@ def test_index_mapper_capacity_covers_the_overlapping_cohorts( is_disagg=is_disagg, num_reserved_index_slots=reserved, enable_attention_dp=adp, - enable_overlap_scheduler=overlap, + disable_overlap_scheduler=disable_overlap, ) assert capacity == expected assert page_table_capacity == expected diff --git a/tests/unittest/_torch/executor/test_seq_slot_sizing.py b/tests/unittest/_torch/executor/test_seq_slot_sizing.py index 29e0943e2555..e24ab54963e4 100644 --- a/tests/unittest/_torch/executor/test_seq_slot_sizing.py +++ b/tests/unittest/_torch/executor/test_seq_slot_sizing.py @@ -281,7 +281,7 @@ def test_sampler_uses_executor_slot_pool_capacity(slot_factor): assert args.max_num_sequences == max_num_sequences -def _make_kv_cache_creator(enable_overlap_scheduler: bool) -> KvCacheCreator: +def _make_kv_cache_creator(disable_overlap_scheduler: bool) -> 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) @@ -296,7 +296,7 @@ def _make_kv_cache_creator(enable_overlap_scheduler: bool) -> KvCacheCreator: c._kv_connector_manager = None c._execution_stream = None c._is_disagg = False - c._enable_overlap_scheduler = enable_overlap_scheduler + 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()) @@ -305,8 +305,8 @@ def _make_kv_cache_creator(enable_overlap_scheduler: bool) -> KvCacheCreator: return c -@pytest.mark.parametrize("enable_overlap_scheduler", [False, True]) -def test_kv_cache_manager_receives_the_overlap_flag(enable_overlap_scheduler): +@pytest.mark.parametrize("disable_overlap_scheduler", [False, True]) +def test_kv_cache_manager_receives_the_overlap_flag(disable_overlap_scheduler): """The manager sizes its own index pool, but needs the overlap flag to do it. The index pool must cover both the retiring cohort and its replacement when @@ -316,7 +316,7 @@ def test_kv_cache_manager_receives_the_overlap_flag(enable_overlap_scheduler): capacity from ``max_batch_size * pp_size``, which is not comparable with a seat pool that also carries the PP multiplier. """ - creator = _make_kv_cache_creator(enable_overlap_scheduler) + creator = _make_kv_cache_creator(disable_overlap_scheduler) model_engine = SimpleNamespace( model=SimpleNamespace(model_config=SimpleNamespace(is_generation=True)), ) @@ -327,4 +327,4 @@ def test_kv_cache_manager_receives_the_overlap_flag(enable_overlap_scheduler): ) as create: creator._create_kv_cache_manager(model_engine) - assert create.call_args.kwargs["enable_overlap_scheduler"] is enable_overlap_scheduler + assert create.call_args.kwargs["disable_overlap_scheduler"] is disable_overlap_scheduler From 5d34836345c16536ef2596d0f9cb749a97776ab2 Mon Sep 17 00:00:00 2001 From: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> Date: Wed, 9 Sep 2026 01:07:55 -0700 Subject: [PATCH 19/22] [https://nvbugs/6627795][fix] drop the added comments in kv_cache_manager_v2 The index-mapper sizing block already carries the pre-existing note about the disaggregated 2x coefficient; the added paragraph restated it for the attention-DP case without saying anything the predicate does not. Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> --- .../_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py | 9 --------- 1 file changed, 9 deletions(-) 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 42c3981fb627..c9caafd4302a 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 @@ -1453,15 +1453,6 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: # (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. - # - # Attention DP with the overlap scheduler needs the same coefficient: - # teardown of the retiring batch (`_process_previous_batch`) runs *after* - # the replacement batch has already been scheduled, so both cohorts hold - # their index slots at once. Without the extra capacity the index mapper - # runs dry and `_create_kv_cache` silently defers requests one at a time - # (nvbug 6627795). Pipeline parallelism is excluded because - # `max_num_sequences` already scales with the number of in-flight - # microbatches. needs_extra_index_slots = is_disagg or ( mapping.enable_attention_dp and not disable_overlap_scheduler and not mapping.has_pp() ) From 206dd2d19c2e97101cf92ae004c69c1ce45c8552 Mon Sep 17 00:00:00 2001 From: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> Date: Wed, 9 Sep 2026 01:21:25 -0700 Subject: [PATCH 20/22] [https://nvbugs/6627795][fix] drop the added comments in _util and py_executor should_enable_disagg_adp_overlap_headroom keeps the one-line docstring it had before this PR, and the four notes added around the router wiring and the liveness count in py_executor go away. Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/_util.py | 14 +------------- tensorrt_llm/_torch/pyexecutor/py_executor.py | 19 +------------------ 2 files changed, 2 insertions(+), 31 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 501eb63be102..76014baa531a 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -3118,19 +3118,7 @@ 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 attention DP. - - The overlap scheduler defers a finished request's teardown by one iteration, - so its sequence slot is still held when the ADP router admits the batch that - replaces it. Without spare slots the router cannot backfill and the forward - batch runs short (nvbug 6627795). Disaggregation needs the same headroom even - with overlap off, because a request awaiting its KV transfer keeps its lease. - - Pipeline parallelism is excluded: admission is capped independently at - ``pp_size * max_batch_size``, so the extra seats are unspendable, and - ADPRouter's retiring-request correction would additionally require every - pipeline stage to agree on which requests are retiring. - """ + """Gate extra sequence slots to non-PP attention DP.""" is_disagg = is_disagg_enabled(cache_transceiver_config) return (mapping.enable_attention_dp and not mapping.has_pp() and (is_disagg or not disable_overlap_scheduler)) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 356e4ef7f267..ec85f029e2f8 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -690,10 +690,6 @@ def __init__( # Router is built after async_transfer_manager so KVCacheAwareADPRouter # can receive the transfer-manager reference at construction time. - # - # The router's overlap correction spends sequence-slot headroom, so it is - # handed the engine's headroom flag rather than re-deriving the predicate: - # that flag is what sized the seat pool the correction spends. self.adp_router: ADPRouter = ADPRouter.create( dist=self.dist, has_seq_slot_headroom=getattr( @@ -5801,13 +5797,7 @@ def _validate_request(self, request: LlmRequest): def _fetch_and_enqueue_requests(self, waiting_queue: WaitingQueue, total_num_live_requests: int) -> None: - """Fetch requests from request_queue and enqueue to waiting_queue. - - `total_num_live_requests` counts every request still resident on any - rank, including the retiring ones that `num_active_requests` excludes: - the idle decision below selects a blocking versus a zero timeout and must - be identical on every rank. - """ + """Fetch requests from request_queue and enqueue to waiting_queue.""" # Block new requests while control requests are pending if len(self.control_requests) != 0: return @@ -6025,9 +6015,6 @@ def _fetch_new_requests( s.num_active_requests for s in all_rank_states ] total_num_active_requests = sum(all_ranks_num_active_requests) - # Retiring requests are excluded from num_active_requests but are - # still resident, so fold them back in for the liveness test only - # (nvbug 6627795). total_num_live_requests = total_num_active_requests + sum( s.num_retiring_requests for s in all_rank_states) else: @@ -7242,10 +7229,6 @@ def _pad_attention_dp_dummy_request(self): return expected_num_active_requests = self.expected_num_active_requests - # Compare against the same routable count the router balanced on: - # gather_all_rank_states excludes retiring requests from the per-rank - # loads that floor `expected` (nvbug 6627795), so measuring against the - # raw len() here would make the warning below fire every iteration. num_routable_active_requests = len(self.active_requests) if self.adp_router.exclude_retiring_requests: num_routable_active_requests -= count_retiring_requests( From f8c5080f3f191c99bf6bfcda3b9931657f6c55e8 Mon Sep 17 00:00:00 2001 From: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> Date: Thu, 10 Sep 2026 09:08:15 -0700 Subject: [PATCH 21/22] [None][fix] Align seq-slot sizing with PR 18983 Bring PR 18457 in line with the seat-pool sizing in https://github.com/NVIDIA/TensorRT-LLM/pull/18983 so the two can land in either order: - compute_max_num_sequences takes 18983's is_disagg term. Disaggregation and the attention-DP overlap headroom both cover one extra set of in-flight sequences, so they are combined with max() rather than multiplied. - KVCacheManagerV2 publishes max_admissible_sequences, the IndexMapper capacity minus the reserved dummy slots, keeping 18457's widened coefficient (disagg, or attention DP with overlap on and no PP). - validate_seq_slot_pool_covers_admission fails at startup when the seat pool is smaller than what admission allows, called just before the SeqSlotManager is built. Keyed on isinstance(int) rather than "is not None" so a Mock cache manager in another module's tests skips the check instead of raising TypeError from a comparison. - Slot-indexed spec-decoding buffers follow model_engine.max_num_seq_slots unconditionally, in _set_up_spec_metadata, _initialize_no_kv_cache_runner and seat_pool_or_none. The pool already exceeds max_batch_size for three independent reasons -- pipeline depth, the overlap headroom and disaggregation -- and those buffers cannot tell them apart, so gating on one of the three sized them at max_batch_size while py_seq_slot ranged over the full pool. - py_executor_creator's max_num_seq_slots fallback recomputes with the disagg term instead of assuming max_batch_size * pp_size. Also merges main (05838ceb) and registers test_seq_slot_sizing.py in l0_a10. Co-Authored-By: Claude Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/_util.py | 93 ++++++--- .../kv_cache/kv_cache_manager_v2.py | 8 +- .../_torch/pyexecutor/model_engine.py | 15 +- .../_torch/pyexecutor/py_executor_creator.py | 13 +- tensorrt_llm/_torch/speculative/utils.py | 24 ++- .../integration/test_lists/test-db/l0_a10.yml | 1 + .../kv_cache/test_kv_cache_manager_v2.py | 22 +- .../_torch/executor/test_seq_slot_sizing.py | 189 +++++++++++++++--- .../speculative/test_spec_slot_pool_sizing.py | 54 ++--- 9 files changed, 307 insertions(+), 112 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 659d984b489d..7782e26018f5 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, @@ -3063,30 +3070,39 @@ def create_kv_cache_compression_manager( return None -def is_disagg_enabled(cache_transceiver_config) -> bool: - """True when a cache transceiver backend is configured.""" - return (cache_transceiver_config is not None - and cache_transceiver_config.backend is not None) - - def compute_max_num_sequences(mapping: Mapping, max_batch_size: int, disable_overlap_scheduler: bool, - enable_overlap_headroom: bool = False) -> int: + enable_overlap_headroom: bool = False, + is_disagg: bool = False) -> int: """Size the sequence-slot pool (and the sampler state it indexes). - ``enable_overlap_headroom`` is intentionally opt-in; see - ``should_enable_disagg_adp_overlap_headroom`` for when it is set. It buys one - extra micro-batch worth of slots, because a finished request's teardown is - deferred by one iteration and its slot is still held while the replacement - batch is admitted (nvbug 6627795). Pipeline parallelism already sizes the - pool by ``pp_size``. + The pool must seat every request the admission path can let through. On a + disaggregated generation server that bound is KVCacheManagerV2's + IndexMapper, sized at twice max_num_sequences so a batch in KV transfer + can overlap a batch that is generating. Seats below that bound let a + request be admitted that cannot be seated once its transfer lands, and + add_slot then raises on the executor's event-loop thread, killing the rank + mid-collective. + + enable_overlap_headroom covers a different case: attention DP can backfill + seats before the overlap scheduler releases the previous iteration's + terminal slots (nvbug 6627795). See + should_enable_disagg_adp_overlap_headroom for when it is set. Pipeline + parallelism already sizes the pool by pp_size. """ if mapping.has_pp(): num_micro_batches = mapping.pp_size else: - num_micro_batches = (2 if enable_overlap_headroom else 1) - return max_batch_size * num_micro_batches + num_micro_batches = (2 if enable_overlap_headroom + and not disable_overlap_scheduler else 1) + num_seats = max_batch_size * num_micro_batches + if is_disagg: + # max() rather than another multiplication: the disagg and + # overlap-headroom factors both cover one extra set of in-flight + # sequences, so they overlap rather than compose. + num_seats = max(num_seats, max_batch_size * mapping.pp_size * 2) + return num_seats def resolve_max_num_sequences(model_engine, @@ -3112,13 +3128,15 @@ def resolve_max_num_sequences(model_engine, 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_disagg_adp_overlap_headroom", - False)) + return compute_max_num_sequences( + mapping, + max_batch_size, + llm_args.disable_overlap_scheduler, + enable_overlap_headroom=getattr(model_engine, + "_enable_disagg_adp_overlap_headroom", + False), + is_disagg=is_disagg_enabled( + getattr(llm_args, "cache_transceiver_config", None))) def should_enable_adp_dummy_fixes(mapping: Mapping) -> bool: @@ -3150,9 +3168,35 @@ def should_enable_disagg_adp_overlap_headroom( cache_transceiver_config: Optional[CacheTransceiverConfig], disable_overlap_scheduler: bool) -> bool: """Gate extra sequence slots to non-PP attention DP.""" - is_disagg = is_disagg_enabled(cache_transceiver_config) return (mapping.enable_attention_dp and not mapping.has_pp() - and (is_disagg or not disable_overlap_scheduler)) + and (is_disagg_enabled(cache_transceiver_config) + or not disable_overlap_scheduler)) + + +def validate_seq_slot_pool_covers_admission(max_num_sequences: int, + kv_cache_manager) -> None: + """Fail at startup if the seat pool is smaller than what admission allows. + + Otherwise the shortfall stays invisible until a request that cannot be + seated arrives, and it then surfaces as a hang rather than an error. + No-op for managers that publish no admission bound. + + ``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. + """ + admission_bound = getattr(kv_cache_manager, "max_admissible_sequences", + None) + if not isinstance(admission_bound, + int) or max_num_sequences >= admission_bound: + return + raise ValueError( + f"Sequence-slot pool ({max_num_sequences} seats) is smaller than the " + f"number of sequences {type(kv_cache_manager).__name__} can admit " + f"({admission_bound}). Requests would be admitted that cannot be " + "seated; see compute_max_num_sequences.") def create_py_executor_instance( @@ -3359,6 +3403,7 @@ def create_py_executor_instance( if isinstance(model_engine, PyTorchModelEngine): model_engine._init_cuda_graph_lora_manager(lora_config) + validate_seq_slot_pool_covers_admission(max_num_sequences, kv_cache_manager) resources[ResourceManagerType.SEQ_SLOT_MANAGER] = SeqSlotManager( max_num_sequences) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 8b8673ea0db6..042973f4edf0 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 @@ -1531,9 +1531,11 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: ) max_num_sequences = max_batch_size * mapping.pp_size assert num_reserved_index_slots >= 0, "num_reserved_index_slots must be non-negative" - index_mapper_capacity = ( - max_num_sequences * (2 if needs_extra_index_slots 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_index_slots else 1) + index_mapper_capacity = self.max_admissible_sequences + num_reserved_index_slots logger.info( f"KVCacheManagerV2: IndexMapper capacity={index_mapper_capacity} " f"(max_num_sequences={max_num_sequences}, is_disagg={is_disagg}, " diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index b52149765636..b57721cfa9fd 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -450,11 +450,13 @@ def __init__( 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, + from ._util import (compute_max_num_sequences, is_disagg_enabled, should_enable_adp_dummy_fixes, should_enable_disagg_adp_overlap_headroom, should_enable_non_overlap_adp_forward_intent, should_enable_scheduler_aware_adp_dummy) + self._is_disagg = is_disagg_enabled( + getattr(llm_args, "cache_transceiver_config", None)) self._enable_disagg_adp_overlap_headroom = ( should_enable_disagg_adp_overlap_headroom( mapping, llm_args.cache_transceiver_config, @@ -465,6 +467,7 @@ def __init__( self.batch_size, llm_args.disable_overlap_scheduler, enable_overlap_headroom=self._enable_disagg_adp_overlap_headroom, + is_disagg=self._is_disagg, ) self.dist = dist if dist is not None: @@ -1132,8 +1135,7 @@ def _initialize_no_kv_cache_runner( mm_encoder_cache_enabled=self._mm_encoder_cache_enabled, spec_config=self.spec_config, is_draft_model=self.is_draft_model, - num_seq_slots=(self.max_num_seq_slots if - self._enable_disagg_adp_overlap_headroom else None), + num_seq_slots=self.max_num_seq_slots, original_max_draft_len=self.original_max_draft_len, original_max_total_draft_tokens=( self.original_max_total_draft_tokens), @@ -3709,11 +3711,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 +3721,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_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 209f4778b53b..56d5c5dbe389 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -37,8 +37,9 @@ 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_disagg_enabled, is_mla, validate_feature_combination) + 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, uses_vswa_kv_cache_layout) @@ -662,8 +663,12 @@ 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, + is_disagg=is_disagg_enabled(cache_transceiver_config)) if is_mla(config): if model_engine.model.model_config.enable_flash_mla: tokens_per_block = 64 diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 194cb4ea5fbd..4bd42e07459d 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -573,20 +573,26 @@ def get_mtp_hidden_size(model_config) -> int: def seat_pool_or_none(model_engine) -> Optional[int]: - """The engine's sequence-slot pool size, or None to keep max_batch_size. + """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: the overlap scheduler holds a - finished request's slot for one more iteration while its replacement is - admitted, so the transient demand exceeds max_batch_size (nvbug-6627795). - Buffers indexed by *batch position* deliberately keep max_batch_size -- the + 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. - Gated on the same flag ``_set_up_spec_metadata`` reads, so every spec-decoding - pool agrees with the metadata about which number it is indexed by. + Read unconditionally rather than behind the attention-DP headroom gate. + ``max_num_seq_slots`` already exceeds max_batch_size for three independent + reasons -- pipeline depth, the attention-DP 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. """ - if not getattr(model_engine, "_enable_disagg_adp_overlap_headroom", False): - return None return getattr(model_engine, "max_num_seq_slots", None) diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index d4fdf3ddca7e..967064c97672 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -37,6 +37,7 @@ l0_a10: - unittest/_torch/executor/kv_cache/test_kv_cache_v2_capacity_only.py - unittest/_torch/executor/test_error_classification.py - unittest/_torch/executor/test_resource_manager.py + - unittest/_torch/executor/test_seq_slot_sizing.py - unittest/_torch/moe/test_communication_factory.py # NOTE: this is a CPU-only test, but we do not have a dedicated job for this (and therefore no # test list either). diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py index 5c0ff295f729..e52b1e9f32fa 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 @@ -1552,12 +1552,16 @@ def _index_mapper_capacity_for( num_reserved_index_slots: int = 1, enable_attention_dp: bool = False, disable_overlap_scheduler: bool = False, -) -> tuple[int, int]: - """Construct a manager and return (IndexMapper capacity, page-table capacity). +) -> tuple[int, int, int]: + """Construct a manager and return the three sizes it derives. - The 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. + ``(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() @@ -1584,7 +1588,7 @@ def build_base_config( patch.object(KVCacheManagerV2, "_prepare_page_table_tensor") as page_table, patch.object(KVCacheManagerV2, "_log_kv_cache_pool_lifecycle_mapping"), ): - KVCacheManagerV2( + 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 @@ -1618,6 +1622,7 @@ def build_base_config( return ( index_mapper_cls.call_args.args[0], page_table.call_args.args[0], + manager.max_admissible_sequences, ) @@ -1665,7 +1670,7 @@ def test_index_mapper_capacity_covers_the_overlapping_cohorts( reserved: int, expected: int, ) -> None: - capacity, page_table_capacity = _index_mapper_capacity_for( + 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, @@ -1675,3 +1680,6 @@ def test_index_mapper_capacity_covers_the_overlapping_cohorts( ) 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_seq_slot_sizing.py b/tests/unittest/_torch/executor/test_seq_slot_sizing.py index e24ab54963e4..7b54783ecd82 100644 --- a/tests/unittest/_torch/executor/test_seq_slot_sizing.py +++ b/tests/unittest/_torch/executor/test_seq_slot_sizing.py @@ -1,25 +1,31 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""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 capacity scheduler has already dropped -them from its budget (no_schedule_after_state=GENERATION_TO_COMPLETE) and -backfilled their seats. Transient slot demand is therefore one extra -micro-batch worth of slots, regardless of whether speculative decoding is -enabled. The headroom is selected from runtime topology, not model -architecture. - -Pipeline parallelism is out of scope: the pool is already sized by pp_size +"""Seq-slot pool sizing. + +Two independent reasons the pool must exceed max_batch_size, both selected from +runtime topology rather than model architecture: + +Attention-DP overlap headroom. 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 capacity scheduler has already dropped them +from its budget (no_schedule_after_state=GENERATION_TO_COMPLETE) and backfilled +their seats. Transient slot demand is therefore one extra micro-batch worth of +slots, regardless of whether speculative decoding is enabled. Pipeline +parallelism is out of scope for this term: the pool is already sized by pp_size there, and the ADP router's retiring-request correction is not rank-consistent because only the last pipeline stage marks generation requests GENERATION_TO_COMPLETE. -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); resolve_max_num_sequences -is how a consumer obtains it without re-deriving it. +Disaggregated serving. On a generation server the admission bound is +KVCacheManagerV2's IndexMapper rather than the seat pool, and it is sized at +twice max_num_sequences. Unlike the headroom above, this holds regardless of +attention DP, overlap, or PP. + +compute_max_num_sequences is the single sizing implementation used 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 @@ -38,15 +44,14 @@ should_enable_disagg_adp_overlap_headroom, should_enable_non_overlap_adp_forward_intent, should_enable_scheduler_aware_adp_dummy, + validate_seq_slot_pool_covers_admission, ) from tensorrt_llm.mapping import Mapping _UCX = SimpleNamespace(backend="UCX") -# (pp_size, enable_overlap_headroom, expected_factor) -# -# Disaggregation is absent from this table on purpose -- see -# test_seat_pool_has_no_disagg_term. +# (pp_size, enable_overlap_headroom, expected_factor) for an aggregated server. +# The disaggregated counterpart is DISAGG_SIZING_CASES below. SIZING_CASES = [ # No PP: the headroom buys one extra micro-batch worth of seats. (1, False, 1), @@ -146,22 +151,140 @@ def test_compute_max_num_sequences_scopes_overlap_headroom( ) -def test_seat_pool_has_no_disagg_term(): - """The disaggregation 2x reaches the seat pool only through the gate. +# (pp_size, enable_overlap_headroom, expected_factor). The factor is relative to +# max_batch_size and, for disagg, must cover the IndexMapper's 2x. +DISAGG_SIZING_CASES = [ + # Disagg gets the coefficient on topology alone, with no headroom opt-in. + (1, False, 2), + # The two factors overlap rather than compose, so this stays 2x. + (1, True, 2), + # PP composes: the IndexMapper is likewise sized pp_size * 2. + (2, False, 4), + (4, False, 8), +] + + +@pytest.mark.parametrize("pp_size,enable_overlap_headroom,expected_factor", DISAGG_SIZING_CASES) +def test_disagg_seats_cover_index_mapper_capacity( + pp_size, enable_overlap_headroom, expected_factor +): + max_batch_size = 8 + mapping = Mapping(world_size=pp_size, tp_size=1, pp_size=pp_size) + + seats = compute_max_num_sequences( + mapping, + max_batch_size, + disable_overlap_scheduler=False, + enable_overlap_headroom=enable_overlap_headroom, + is_disagg=True, + ) + + assert seats == max_batch_size * expected_factor + + # Seats must cover every request admission can let through. Mirrors the + # expression in KVCacheManagerV2.__init__. + index_mapper_capacity = max_batch_size * pp_size * 2 + assert seats >= index_mapper_capacity + + +@pytest.mark.parametrize("disable_overlap", [False, True]) +def test_disagg_seats_do_not_depend_on_overlap_scheduler(disable_overlap): + """The IndexMapper is sized the same either way, so seats must be too. + + Turning the overlap scheduler off removes the terminal-slot race but not the + transfer/generate overlap that the 2x coefficient exists for. + """ + max_batch_size = 8 + mapping = Mapping(world_size=1, tp_size=1, pp_size=1) + + assert ( + compute_max_num_sequences(mapping, max_batch_size, disable_overlap, is_disagg=True) + == max_batch_size * 2 + ) + + +def test_aggregate_sizing_is_unchanged(): + """Aggregated deployments size the pool at one forward batch.""" + max_batch_size = 8 + mapping = Mapping(world_size=1, tp_size=1, pp_size=1) + + assert ( + compute_max_num_sequences( + mapping, max_batch_size, disable_overlap_scheduler=False, is_disagg=False + ) + == max_batch_size + ) + + +@pytest.mark.parametrize("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 coefficients are computed independently; hold them in step. + + KVCacheManagerV2 derives its own admission bound from max_batch_size, + pp_size, is_disagg and the attention-DP overlap condition. Assert the two + expressions against each other rather than against a literal, so a drift in + either one fails here rather than at startup in + validate_seq_slot_pool_covers_admission. + """ + max_batch_size = 8 + mapping = Mapping( + world_size=pp_size, + tp_size=1, + pp_size=pp_size, + enable_attention_dp=enable_attention_dp, + ) + + seats = compute_max_num_sequences( + mapping, + max_batch_size, + disable_overlap_scheduler, + enable_overlap_headroom=should_enable_disagg_adp_overlap_headroom( + mapping, + _UCX if is_disagg else None, + disable_overlap_scheduler, + ), + is_disagg=is_disagg, + ) + + # Mirrors the coefficient in KVCacheManagerV2.__init__. + needs_extra_index_slots = is_disagg or ( + enable_attention_dp and not disable_overlap_scheduler and pp_size == 1 + ) + admission_bound = max_batch_size * pp_size * (2 if needs_extra_index_slots else 1) + + assert seats >= admission_bound + + +def test_validate_seq_slot_pool_accepts_sufficient_pool(): + validate_seq_slot_pool_covers_admission(16, Mock(max_admissible_sequences=16)) + + +def test_validate_seq_slot_pool_rejects_undersized_pool(): + with pytest.raises(ValueError, match="smaller than the number of"): + validate_seq_slot_pool_covers_admission(8, Mock(max_admissible_sequences=16)) + + +def test_validate_seq_slot_pool_ignores_managers_without_a_bound(): + """The V1/C++ manager does not publish one; the check must not fire.""" + manager = Mock(spec=[]) + validate_seq_slot_pool_covers_admission(1, manager) + validate_seq_slot_pool_covers_admission(1, None) + - 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. So the - sizing function itself must not carry a disaggregation term -- the single - ``enable_overlap_headroom`` flag is the only way in. +def test_validate_seq_slot_pool_ignores_a_non_integer_bound(): + """A bare ``Mock`` auto-creates the attribute, so ``is None`` is not enough. - 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. + ``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. """ - assert "is_disagg" not in inspect.signature(compute_max_num_sequences).parameters - assert "is_disagg" not in inspect.signature(resolve_max_num_sequences).parameters + validate_seq_slot_pool_covers_admission(1, Mock()) @pytest.mark.parametrize( diff --git a/tests/unittest/_torch/speculative/test_spec_slot_pool_sizing.py b/tests/unittest/_torch/speculative/test_spec_slot_pool_sizing.py index 86c18220b885..6d135c0d5906 100644 --- a/tests/unittest/_torch/speculative/test_spec_slot_pool_sizing.py +++ b/tests/unittest/_torch/speculative/test_spec_slot_pool_sizing.py @@ -3,10 +3,11 @@ """Speculative-decoding state that is keyed by live-request identity must be sized by the sequence-slot pool, not by max_batch_size. -Under the attention-DP overlap headroom the two differ by 2x -(``compute_max_num_sequences``): 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: +The two differ whenever ``compute_max_num_sequences`` widens the pool: under +pipeline depth, under disaggregation, and under the attention-DP 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, @@ -232,8 +233,12 @@ def test_every_resource_manager_branch_forwards_the_slot_pool(): # --------------------------------------------------------------------------- -def _drafter_engine(headroom: bool, seats: int = POOL): - """A stub target engine for get_spec_drafter's draft-target branch.""" +def _drafter_engine(seats): + """A stub target engine for get_spec_drafter's draft-target branch. + + ``seats=None`` stands for an engine that publishes no pool at all, which is + the only case that still falls back to max_batch_size. + """ spec_dec_mode = types.SimpleNamespace( is_user_provided=lambda: False, is_draft_target=lambda: True, @@ -248,17 +253,15 @@ def _drafter_engine(headroom: bool, seats: int = POOL): max_concurrency=None, draft_len_schedule=None, ) - return types.SimpleNamespace( - spec_config=spec_config, - batch_size=R, - max_num_seq_slots=seats, - _enable_disagg_adp_overlap_headroom=headroom, - ) + engine = types.SimpleNamespace(spec_config=spec_config, batch_size=R) + if seats is not None: + engine.max_num_seq_slots = seats + return engine -def _drafter(headroom: bool, seats: int = POOL): +def _drafter(seats): return get_spec_drafter( - _drafter_engine(headroom, seats), + _drafter_engine(seats), draft_model_engine=object(), sampler=object(), spec_resource_manager=None, @@ -266,18 +269,23 @@ def _drafter(headroom: bool, seats: int = POOL): @pytest.mark.cpu_only -@pytest.mark.parametrize("headroom,expected_pool", [(True, POOL), (False, R)]) -def test_draft_slot_pool_follows_the_target_seat_pool(headroom, expected_pool): +@pytest.mark.parametrize("seats,expected_pool", [(POOL, POOL), (R, R), (None, R)]) +def test_draft_slot_pool_follows_the_target_seat_pool(seats, expected_pool): """The draft pool is sized by the *target* engine's seat count. The indices it hands out address buffers sized by that seat count -- the sampler shared with PyExecutor, the draft KV cache manager's IndexMapper, and spec_resource_manager -- and ``_create_draft_request`` even carries the target's ``py_seq_slot`` across as ``target_seq_slot``. Sizing this pool - independently is what makes the two disagree. Headroom off keeps - max_batch_size, so nothing changes for the other topologies. + independently is what makes the two disagree. + + Read from the published pool rather than from the reason it is wide: seats + exceed max_batch_size under pipeline depth and under disaggregation as well as + under the attention-DP headroom, and this pool cannot tell them apart. An + engine whose pool already equals max_batch_size (row 2) is therefore + unchanged, which is what keeps aggregated non-ADP deployments untouched. """ - assert _drafter(headroom).draft_seq_slot_manager.slot_manager.max_num_requests == expected_pool + assert _drafter(seats).draft_seq_slot_manager.slot_manager.max_num_requests == expected_pool @pytest.mark.cpu_only @@ -289,11 +297,11 @@ def test_draft_slot_pool_survives_a_full_overlap_turnover(): cohorts hold draft slots at once. At max_batch_size the (R+1)-th lease raises NoFreeSlotsError -- inside the draft loop, mid-iteration. """ - pool = _drafter(True).draft_seq_slot_manager.slot_manager + pool = _drafter(POOL).draft_seq_slot_manager.slot_manager slots = [pool.add_slot(rid) for rid in range(2 * R)] assert len(set(slots)) == 2 * R - starved = _drafter(False).draft_seq_slot_manager.slot_manager + starved = _drafter(R).draft_seq_slot_manager.slot_manager for rid in range(R): starved.add_slot(rid) with pytest.raises(NoFreeSlotsError): @@ -336,8 +344,8 @@ 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 - which does not know the pool (PP, no attention DP, overlap disabled) keeps the - established sizing without every call site having to restate it. + 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") From c17e82f7a0119f7bd1d7d1686b5699bd3f60d054 Mon Sep 17 00:00:00 2001 From: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> Date: Fri, 11 Sep 2026 06:26:18 -0700 Subject: [PATCH 22/22] [None][fix] Register test_spec_slot_pool_sizing.py in l0_h100 The file was added without an entry on any test-db list, so CI never collected it. Files directly under tests/unittest/_torch/speculative/ are registered per-file -- only the hw_agnostic/ subdirectory has a directory-level entry -- so the 558 lines of slot-pool sizing guards, including the NoFreeSlotsError and IndexError negative controls, would never have run. l0_h100 already carries this file's siblings per-file, and a GPU stage covers both halves: the eleven cpu_only cases and the nine that are skipif(not torch.cuda.is_available()) because the pools they size are CUDA tensors. Co-Authored-By: Claude Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> --- tests/integration/test_lists/test-db/l0_h100.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index a91ef915d8aa..f5a5ef008564 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