From 7d1db52c9879c1ae1fb8ca46004bed63db64e8a0 Mon Sep 17 00:00:00 2001 From: Xianjie <5410381+qiaoxj07@users.noreply.github.com> Date: Mon, 13 Jul 2026 16:19:08 +0800 Subject: [PATCH 1/4] [None][fix] ADP pad: count scheduler window states, tolerate dummy allocation failure Adopt the fix from PR #16161 (closed in favor of this PR): 1. _count_schedulable_active_requests (disaggregated mode): count only requests inside the scheduler's schedulability window [CONTEXT_INIT, GENERATION_TO_COMPLETE), mirroring no_schedule_until_state / no_schedule_after_state. States outside the window (awaiting KV transfer, DISAGG_CONTEXT_WAIT_SCHEDULER, GENERATION_TO_COMPLETE) cannot produce forward work; counting them suppresses the ADP pad dummy while the rank schedules batch=0, which turns can_queue False fleet-wide and stalls every rank (surfaces as the expected_num_active_requests assert or leaked pad dummies). Non-disaggregated behavior is unchanged. 2. _pad_attention_dp_dummy_request: add_dummy_requests returns None when no cache resources are free for even a 1-token dummy; skip padding for the iteration instead of crashing on an unchecked [0]. Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 45 ++++++++++---- .../_torch/executor/test_benchmark_disagg.py | 1 + .../_torch/executor/test_py_executor.py | 60 +++++++++++++++++++ 3 files changed, 96 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index f69041fa7ab8..f47b90ae3465 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -5291,10 +5291,15 @@ def _count_schedulable_active_requests(self) -> int: """Count active requests that are ready for scheduling. In non-disaggregated mode, all active requests are schedulable. - In disaggregated mode, requests still waiting for KV cache - transfer (in INIT or transmission-in-progress state) are - excluded because they cannot participate in the forward pass - until transfer completes. + In disaggregated mode, only requests inside the scheduler's + schedulability window [CONTEXT_INIT, GENERATION_TO_COMPLETE) are + counted, mirroring the MicroBatchScheduler bounds + (no_schedule_until_state / no_schedule_after_state). States + outside the window — e.g. waiting for KV transfer (below) or + awaiting completion bookkeeping (at or above) — cannot produce + forward work, so counting them here would suppress the pad dummy + while the rank still schedules batch=0, turning can_queue False + fleet-wide and stalling every rank. Returns: The number of active requests eligible for scheduling. @@ -5302,12 +5307,19 @@ def _count_schedulable_active_requests(self) -> int: if self.kv_cache_transceiver is None: return len(self.active_requests) - def _is_awaiting_kv_transfer(req) -> bool: - return (req.is_disagg_generation_init_state - or req.is_disagg_generation_transmission_in_progress) + # Both bounds must match the scheduler exactly. The lower bound + # covers e.g. DISAGG_GENERATION_INIT / TRANS_IN_PROGRESS and + # DISAGG_CONTEXT_WAIT_SCHEDULER (gen-first mode on the context + # server); the upper bound covers GENERATION_TO_COMPLETE and the + # terminal / post-completion disagg states. Disaggregated serving + # is decoder-only, so the scheduler's until-state is CONTEXT_INIT + # here (the ENCODER_INIT variant only applies to encoder-decoder + # models, which have no disagg path). + context_init_value = LlmRequestState.CONTEXT_INIT.value + to_complete_value = LlmRequestState.GENERATION_TO_COMPLETE.value return sum(1 for req in self.active_requests - if not _is_awaiting_kv_transfer(req)) + if context_init_value <= req.state_value < to_complete_value) def _should_skip_dummy_for_benchmark_disagg( self, num_schedulable_requests: int) -> bool: @@ -5382,13 +5394,26 @@ def _pad_attention_dp_dummy_request(self): and self.max_num_tokens is not None): token_nums = [self.max_num_tokens] dummy_request_ids = [ATTENTION_DP_DUMMY_REQUEST_ID] - llm_request = self.kv_cache_manager.add_dummy_requests( + dummy_requests = self.kv_cache_manager.add_dummy_requests( request_ids=dummy_request_ids, token_nums=token_nums, is_gen=self._adp_dummy_is_gen, prepare_resource=True, max_num_draft_tokens=self.max_total_draft_tokens, - )[0] + ) + if not dummy_requests: + # No free KV/mamba resources for even a 1-token dummy — + # possible while non-schedulable requests (e.g. awaiting + # completion bookkeeping or KV transfer) still hold theirs. + # Skip padding: this rank schedules an empty batch and the + # fleet skips one iteration; the resource-holding requests + # finish within an iteration or two and padding succeeds on + # the retry. Crashing here (the old unchecked [0]) is worse. + logger.warning( + "Cannot allocate ADP pad dummy (no free cache resources);" + " rank schedules an empty batch this iteration.") + return + llm_request = dummy_requests[0] llm_request.is_attention_dp_dummy = True spec_resource_manager = self.resource_manager.get_resource_manager( ResourceManagerType.SPEC_RESOURCE_MANAGER) diff --git a/tests/unittest/_torch/executor/test_benchmark_disagg.py b/tests/unittest/_torch/executor/test_benchmark_disagg.py index dabb2b910d55..d00fecb9bcbc 100644 --- a/tests/unittest/_torch/executor/test_benchmark_disagg.py +++ b/tests/unittest/_torch/executor/test_benchmark_disagg.py @@ -44,6 +44,7 @@ def _make_active_request( ) -> Mock: """Create an active request stub with disagg state flags.""" req = Mock() + req.state_value = LlmRequestState.GENERATION_IN_PROGRESS.value req.is_disagg_generation_init_state = in_init req.is_disagg_generation_transmission_in_progress = in_transfer req.state = ( diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 8d9b4c1e1ed8..490b86e61603 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -1162,6 +1162,13 @@ def test_pp_wait_exception_does_not_skip_notify(self): _STATE_DISAGG_GENERATION_INIT = "_disagg_init_sentinel" _STATE_DISAGG_GENERATION_TRANS_IN_PROGRESS = "_disagg_trans_sentinel" +# The sentinel mocks must expose the real state ints: the pad predicate +# compares state_value against the scheduler window bounds. +_SENTINEL_STATE_VALUES = { + _STATE_DISAGG_GENERATION_INIT: LlmRequestState.DISAGG_GENERATION_INIT.value, + _STATE_DISAGG_GENERATION_TRANS_IN_PROGRESS: LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS.value, +} + def _make_adp_request( state, @@ -1174,6 +1181,11 @@ def _make_adp_request( ): req = Mock() req.state = state + req.state_value = ( + state.value + if isinstance(state, LlmRequestState) + else _SENTINEL_STATE_VALUES.get(state, LlmRequestState.GENERATION_IN_PROGRESS.value) + ) req.py_request_id = request_id req.is_child = is_child req.parent_request_id = parent_request_id @@ -1320,6 +1332,54 @@ def test_pad_dummy_counts_generation_to_complete(): assert len(stub.active_requests) == 1 +def test_pad_dummy_added_when_only_to_complete_requests_disagg(): + # In disaggregated mode a GENERATION_TO_COMPLETE request is refused by + # MicroBatchScheduler (no_schedule_after_state), so a rank holding only + # such requests schedules batch=0. It must receive a pad dummy, or + # can_queue goes False fleet-wide and pad dummies leak on other ranks. + stub = _StubADPExecutor() + stub.active_requests = [_make_adp_request(_STATE_GENERATION_TO_COMPLETE)] + stub.expected_num_active_requests = 2 + + _run_pad(stub) + + assert len(stub.add_dummy_calls) == 1 + assert len(stub.active_requests) == 2 + + +def test_pad_dummy_added_when_only_wait_scheduler_requests_disagg(): + # Gen-first mode on the context server: DISAGG_CONTEXT_WAIT_SCHEDULER + # sits BELOW the scheduler's window [CONTEXT_INIT, GENERATION_TO_COMPLETE) + # (no_schedule_until_state), so a rank holding only such requests + # schedules batch=0 and must receive a pad dummy — the left-boundary + # mirror of the TO_COMPLETE case above. + stub = _StubADPExecutor() + stub.active_requests = [_make_adp_request(LlmRequestState.DISAGG_CONTEXT_WAIT_SCHEDULER)] + stub.expected_num_active_requests = 2 + + _run_pad(stub) + + assert len(stub.add_dummy_calls) == 1 + assert len(stub.active_requests) == 2 + + +def test_pad_dummy_allocation_failure_skips_padding(): + # add_dummy_requests returns None when the rank has no free cache + # resources for even a 1-token dummy (possible while non-schedulable + # requests still hold theirs). Padding must degrade to a skipped + # iteration, not crash on an unchecked [0]. + stub = _StubADPExecutor() + stub.active_requests = [_make_adp_request(_STATE_GENERATION_TO_COMPLETE)] + stub.expected_num_active_requests = 2 + stub.kv_cache_manager.add_dummy_requests.side_effect = None + stub.kv_cache_manager.add_dummy_requests.return_value = None + + _run_pad(stub) + + assert len(stub.active_requests) == 1 + assert not any(r.is_attention_dp_dummy for r in stub.active_requests) + + def test_pad_dummy_skips_when_active_request_present(): stub = _StubADPExecutor() stub.active_requests = [_make_adp_request(_STATE_GENERATION_IN_PROGRESS)] From 865e9c020abd4e253c494dd84765d4dab7bf518c Mon Sep 17 00:00:00 2001 From: Xianjie <5410381+qiaoxj07@users.noreply.github.com> Date: Mon, 13 Jul 2026 16:19:08 +0800 Subject: [PATCH 2/4] [None][fix] Seq-slot pool overlap headroom with consistent slot-indexed buffer sizing Under the overlap scheduler a finishing request holds its sequence slot until the next iteration's _process_previous_batch, while the V2 scheduler has already dropped it from the request budget (no_schedule_after_state=GENERATION_TO_COMPLETE) and backfilled its seat. Transient slot demand therefore reaches 2 * max_batch_size and SeqSlotManager raises ValueError("No free slots") (DSV4-Pro disagg GEN dep32/batch2/mtp3 dies at the first request burst). Size the pool with overlap headroom via a single-sourced compute_max_num_sequences (max_batch_size * num_micro_batches: pp_size with PP, else 2 under overlap), and size every buffer indexed by py_seq_slot from the same number so slot ids stay in range: - sampler state (create_torch_sampler_args) and the legacy TRTLLMSampler decoder state - guided-decoder per-slot buffers (py_executor_creator; AutoDeploy's guided decoder aligned to its own pool, which keeps pp_size headroom) - SpecMetadata.draft_probs for rejection sampling (new num_seq_slots field, threaded from model_engine) - eagle3 dynamic-tree DynamicTreeSlotStorage (new max_num_seq_slots, threaded via get_spec_resource_manager); batch-position work buffers keep num_trees sizing V1 scheduler capacity stays max_batch_size * pp_size (V1 retains GENERATION_TO_COMPLETE in capacity and handles overlap via two_step_lookahead); the V2 request budget stays max_batch_size. Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com> --- .../_torch/auto_deploy/shim/ad_executor.py | 4 +- tensorrt_llm/_torch/pyexecutor/_util.py | 30 +++++++- .../_torch/pyexecutor/model_engine.py | 13 +++- .../_torch/pyexecutor/py_executor_creator.py | 22 ++++-- .../_torch/pyexecutor/sampler/sampler.py | 6 +- tensorrt_llm/_torch/speculative/eagle3.py | 10 ++- tensorrt_llm/_torch/speculative/interface.py | 11 ++- .../_torch/speculative/spec_tree_manager.py | 18 +++-- tensorrt_llm/_torch/speculative/utils.py | 22 +++++- .../_torch/executor/test_seq_slot_sizing.py | 69 +++++++++++++++++++ 10 files changed, 181 insertions(+), 24 deletions(-) create mode 100644 tests/unittest/_torch/executor/test_seq_slot_sizing.py diff --git a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py index 86488bc36cda..c6bc6c2006b5 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py @@ -1310,7 +1310,9 @@ def create_autodeploy_executor( ) guided_decoder = GuidedDecoder( guided_decoding_config=guided_decoding_config, - max_num_sequences=ad_config.max_batch_size, + # Guided-decoder state is indexed by sequence slot; size it to the + # SeqSlotManager pool above (max_batch_size * pp_size). + max_num_sequences=max_num_sequences, vocab_size_padded=vocab_size_padded, ) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 97b2ca3a82f1..5e4caed62e43 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2051,6 +2051,24 @@ def create_kv_cache_compression_manager( return None +def compute_max_num_sequences(mapping: Mapping, max_batch_size: int, + disable_overlap_scheduler: bool) -> int: + """Size the sequence-slot pool (and the sampler state it indexes). + + The slot pool needs headroom for two iterations under overlap (prev-iter + finished requests still hold slots when next-iter prepare_resources + runs). Without it, the V2 scheduler can backfill the budget seat vacated + by a GENERATION_TO_COMPLETE request whose slot is not freed yet and + exhaust the pool ("No free slots"). With pipeline parallelism the + micro-batch count is pp_size, which covers overlap already. + """ + if mapping.has_pp(): + num_micro_batches = mapping.pp_size + else: + num_micro_batches = 1 if disable_overlap_scheduler else 2 + return max_batch_size * num_micro_batches + + def create_py_executor_instance( *, dist, @@ -2085,7 +2103,8 @@ def create_py_executor_instance( spec_config = model_engine.spec_config - max_num_sequences = max_batch_size * mapping.pp_size + max_num_sequences = compute_max_num_sequences( + mapping, max_batch_size, llm_args.disable_overlap_scheduler) 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}" @@ -2282,7 +2301,9 @@ 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. - scheduler_capacity = max_num_sequences + # V1 scheduler handles overlap via two_step_lookahead, so skip the + # slot-pool overlap factor here. + scheduler_capacity = max_batch_size * mapping.pp_size if scheduler_capacity == 1 and mapping.enable_attention_dp and kv_cache_manager: scheduler_capacity += 1 @@ -2439,7 +2460,10 @@ def create_torch_sampler_args( enable_async_worker: bool, enable_speculative_beam_history_d2h: bool, ): - max_num_sequences = max_batch_size * mapping.pp_size + # 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 = compute_max_num_sequences(mapping, max_batch_size, + disable_overlap_scheduler) 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 diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index a1048249d08d..2dfebdb65e7b 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -2154,7 +2154,14 @@ def _set_up_spec_metadata( self, spec_resource_manager: Optional[BaseResourceManager], no_cache=False): + # Lazy import: _util imports model_engine at module level. + from ._util import compute_max_num_sequences + spec_config = self.spec_config if self.enable_spec_decode else None + # Slot-indexed metadata buffers must span the SeqSlotManager pool, + # which exceeds batch_size under the overlap scheduler. + num_seq_slots = compute_max_num_sequences( + self.mapping, self.batch_size, self._disable_overlap_scheduler) if no_cache: return get_spec_metadata( spec_config, @@ -2163,7 +2170,8 @@ def _set_up_spec_metadata( max_num_tokens=self.max_num_tokens, spec_resource_manager=spec_resource_manager, is_draft_model=self.is_draft_model, - max_seq_len=self.max_seq_len) + max_seq_len=self.max_seq_len, + num_seq_slots=num_seq_slots) if self.spec_metadata is not None: return self.spec_metadata @@ -2174,7 +2182,8 @@ def _set_up_spec_metadata( max_num_tokens=self.max_num_tokens, spec_resource_manager=spec_resource_manager, is_draft_model=self.is_draft_model, - max_seq_len=self.max_seq_len) + max_seq_len=self.max_seq_len, + num_seq_slots=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 2ef732c166b4..6a5be5d0e4f3 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -35,8 +35,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) + compute_max_num_sequences, create_py_executor_instance, + instantiate_sampler, is_mla, validate_feature_combination) from .config_utils import is_hybrid_linear from .connectors.kv_cache_connector import KvCacheConnectorManager from .dwdp import DwdpManager @@ -757,10 +757,20 @@ def drafting_loop_wrapper(model): with allocation_scope(ExecutorMemoryType.GUIDED_DECODER): if mapping.is_last_pp_rank(): kwargs = { - "guided_decoding_config": guided_decoding_config, - "max_num_sequences": max_batch_size, - "vocab_size_padded": model_engine.model.vocab_size_padded, - "rank": mapping.rank, + "guided_decoding_config": + guided_decoding_config, + # Guided-decoder state is indexed by sequence slot; size + # it to the SeqSlotManager pool (2x max_batch_size under + # the overlap scheduler), matching the sampler's + # new_tokens width. + "max_num_sequences": + compute_max_num_sequences( + mapping, max_batch_size, + llm_args.disable_overlap_scheduler), + "vocab_size_padded": + model_engine.model.vocab_size_padded, + "rank": + mapping.rank, } if spec_config is not None: kwargs[ diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index e7fe112ace6f..1911f944c86f 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -4866,12 +4866,16 @@ def __init__( ) self.max_batch_size = max_batch_size self.max_beam_width = max_beam_width - self.max_num_sequences = mapping.pp_size * max_batch_size self.max_seq_idle_microseconds = 180 * 1000 * 1000 self.is_trt_overlap = not disable_overlap_scheduler self.num_micro_batches = ( mapping.pp_size if mapping.pp_size > 1 else (2 if self.is_trt_overlap else 1) ) + # Decoder state is indexed by sequence slot; size it to the + # SeqSlotManager pool (max_batch_size * num_micro_batches, matching + # compute_max_num_sequences), not pp_size alone, so py_seq_slot + # values under the overlap scheduler stay in range. + self.max_num_sequences = max_batch_size * self.num_micro_batches self.micro_batch_idx = 0 if mpi_disabled(): diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index 3978040724b1..ccc33b2f8087 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -40,7 +40,8 @@ def __init__(self, max_num_requests: int, max_seq_len: int, max_num_tokens: int, - sa_manager=None): + sa_manager=None, + max_num_seq_slots: int | None = None): self.dtype = dtype self.max_draft_len = config.max_draft_len self.hidden_size = hidden_size @@ -101,6 +102,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, + max_num_seq_slots=max_num_seq_slots, ) def prepare_resources(self, scheduled_batch: ScheduledRequests): @@ -156,7 +158,10 @@ 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, + max_num_seq_slots: int | None = None): self.max_num_requests = max_num_requests self.batch_indices_cuda = torch.empty( [max_num_requests], @@ -170,6 +175,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, + max_num_seq_slots=max_num_seq_slots, ) def free_resources(self, request: LlmRequest): diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index d0b0db6e3d32..6c208e4db3fe 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -533,10 +533,16 @@ class SpecMetadata: use_sampling_params_for_draft_tokens: bool = False # Vocab size used for draft_probs buffer allocation. vocab_size: int = 0 + # Size of the SeqSlotManager pool. py_seq_slot values range over + # [0, num_seq_slots); under the overlap scheduler this is + # 2 * max_batch_size, larger than max_num_requests (== max_batch_size). + # Slot-indexed buffers (draft_probs) must span this full range. + # 0 falls back to max_num_requests. + num_seq_slots: int = 0 # Draft probabilities buffer for rejection sampling, indexed by py_seq_slot # so per-request data is stable across iterations regardless of batch # composition shifts (chunking ctx, gen completion, new ctx joining). - # Shape: [max_num_requests, max_draft_len, vocab_size]. + # Shape: [num_seq_slots, max_draft_len, vocab_size]. draft_probs: Optional[torch.Tensor] = None draft_probs_vocab_size: int = 0 # Whether draft_probs contains valid data. @@ -573,8 +579,9 @@ def prepare(self): and self.vocab_size > 0): # 3D [slot, draft_step, vocab] so we can scatter/gather by slot id # and avoid the brittle "batch position == buffer position" mapping. + slot_capacity = self.num_seq_slots or self.max_num_requests self.draft_probs = torch.empty( - (self.max_num_requests, self.max_draft_len, self.vocab_size), + (slot_capacity, self.max_draft_len, self.vocab_size), dtype=torch.float32, device='cuda') self.draft_probs_vocab_size = self.vocab_size diff --git a/tensorrt_llm/_torch/speculative/spec_tree_manager.py b/tensorrt_llm/_torch/speculative/spec_tree_manager.py index 74dd622826f1..d8ef26498126 100644 --- a/tensorrt_llm/_torch/speculative/spec_tree_manager.py +++ b/tensorrt_llm/_torch/speculative/spec_tree_manager.py @@ -178,10 +178,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, + max_num_seq_slots: int | None = None): self.use_dynamic_tree = use_dynamic_tree self.max_total_draft_tokens = max_total_draft_tokens @@ -198,6 +202,12 @@ 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 + # Slot storage is indexed by py_seq_slot, whose range is the + # SeqSlotManager pool size (2x max_batch_size under the overlap + # scheduler) -- larger than num_trees, which sizes the + # batch-position work buffers. + self.max_num_seq_slots = (max_num_seq_slots if max_num_seq_slots + is not None else self.num_trees) self.dynamic_tree_max_topK = dynamic_tree_max_topK self.cur_draft_layer_idx = 0 self.top_k_list = [] @@ -281,7 +291,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.max_num_seq_slots, n_dt=num_draft_with_root, mask_width=mask_width, ) diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 164d6758ed4a..f98b67187504 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -45,9 +45,14 @@ def get_spec_metadata(spec_config, max_num_tokens, spec_resource_manager=None, is_draft_model=False, - max_seq_len=262144): + max_seq_len=262144, + num_seq_slots=None): use_rejection_sampling = getattr(spec_config, "use_rejection_sampling", False) + # Slot-indexed buffers (draft_probs) must span the SeqSlotManager pool, + # which exceeds max_num_requests under the overlap scheduler. + 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) if spec_config.spec_dec_mode.is_mtp_eagle_one_model(): # MTP Eagle one-model reuses Eagle3 one-model metadata for the @@ -67,6 +72,7 @@ def get_spec_metadata(spec_config, max_num_tokens=max_num_tokens, use_rejection_sampling=use_rejection_sampling, vocab_size=vocab_size, + num_seq_slots=num_seq_slots, spec_resource_manager=spec_resource_manager, ) if spec_config.spec_dec_mode.is_mtp_vanilla(): @@ -125,6 +131,7 @@ def get_spec_metadata(spec_config, layers_to_capture=spec_config.eagle3_layers_to_capture, use_rejection_sampling=use_rejection_sampling, vocab_size=vocab_size, + num_seq_slots=num_seq_slots, spec_resource_manager=spec_resource_manager, use_dynamic_tree=_is_effective_dynamic_tree(spec_config), eagle_choices=spec_config.eagle_choices, @@ -202,6 +209,9 @@ def get_mtp_hidden_size(model_config) -> int: def get_spec_resource_manager(model_engine, draft_model_engine=None): + # Lazy import: _util imports model_engine, which imports this module. + from ..pyexecutor._util import compute_max_num_sequences + spec_config = model_engine.spec_config if spec_config is None: return None @@ -209,6 +219,11 @@ 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 + # Slot-indexed storage must span the SeqSlotManager pool, which exceeds + # max_num_requests under the overlap scheduler. + max_num_seq_slots = compute_max_num_sequences( + model_engine.mapping, model_engine.batch_size, + model_engine._disable_overlap_scheduler) spec_dec_mode = spec_config.spec_dec_mode if spec_dec_mode.is_mtp_eagle_one_model(): sa_manager = None @@ -246,8 +261,8 @@ 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, max_num_seq_slots=max_num_seq_slots) if spec_dec_mode.is_eagle3_one_model(): sa_manager = None sa_cfg = getattr(spec_config, 'sa_config', None) @@ -272,6 +287,7 @@ def get_spec_resource_manager(model_engine, draft_model_engine=None): max_num_requests, max_seq_len, max_num_tokens, + max_num_seq_slots=max_num_seq_slots, ) if spec_dec_mode.is_save_hidden_states(): return SaveHiddenStatesResourceManager( diff --git a/tests/unittest/_torch/executor/test_seq_slot_sizing.py b/tests/unittest/_torch/executor/test_seq_slot_sizing.py new file mode 100644 index 000000000000..97649f01bde9 --- /dev/null +++ b/tests/unittest/_torch/executor/test_seq_slot_sizing.py @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Seq-slot pool / sampler sizing must include overlap headroom. + +Under the overlap scheduler, requests finished in the previous iteration +still hold their sequence slots when the next iteration's +prepare_resources runs, while the V2 scheduler has already dropped them +from its budget (no_schedule_after_state=GENERATION_TO_COMPLETE) and +backfilled their seats. Transient slot demand is therefore +2 * max_batch_size; sizing the pool (and the sampler state it indexes) +at max_batch_size alone exhausts it ("No free slots"). + +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). +""" + +import pytest + +from tensorrt_llm._torch.pyexecutor._util import ( + compute_max_num_sequences, + create_torch_sampler_args, +) +from tensorrt_llm.mapping import Mapping + +SIZING_CASES = [ + # (pp_size, disable_overlap_scheduler, expected_factor) + # Overlap without PP: headroom for two in-flight iterations. + (1, False, 2), + # No overlap: exact batch size suffices. + (1, True, 1), + # With PP the micro-batch count is pp_size, regardless of overlap. + (2, False, 2), + (4, True, 4), +] + + +@pytest.mark.parametrize("pp_size,disable_overlap,expected_factor", SIZING_CASES) +def test_compute_max_num_sequences_includes_overlap_headroom( + pp_size, disable_overlap, expected_factor +): + max_batch_size = 8 + mapping = Mapping(world_size=pp_size, tp_size=1, pp_size=pp_size) + assert ( + compute_max_num_sequences(mapping, max_batch_size, disable_overlap) + == max_batch_size * expected_factor + ) + + +@pytest.mark.parametrize("pp_size,disable_overlap,expected_factor", SIZING_CASES) +def test_sampler_max_num_sequences_matches_executor_slot_pool( + pp_size, disable_overlap, expected_factor +): + max_batch_size = 8 + mapping = Mapping(world_size=pp_size, tp_size=1, pp_size=pp_size) + 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=disable_overlap, + enable_async_worker=False, + enable_speculative_beam_history_d2h=False, + ) + assert args.max_num_sequences == compute_max_num_sequences( + mapping, max_batch_size, disable_overlap + ) + assert args.max_num_sequences == max_batch_size * expected_factor From 3895d4fe40dc87faca6f2543d4cf981995c93952 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Mon, 13 Jul 2026 14:20:19 -0700 Subject: [PATCH 3/4] [None][fix] Harden ADP padding and slot-indexed caches Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../_torch/models/modeling_qwen2vl.py | 13 +- .../_torch/models/modeling_qwen3vl.py | 3 +- .../_torch/pyexecutor/mamba_cache_manager.py | 49 ++- .../_torch/pyexecutor/model_engine.py | 49 ++- .../_torch/pyexecutor/model_loader.py | 13 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 236 +++++++++--- .../_torch/pyexecutor/resource_manager.py | 6 +- .../_torch/executor/test_benchmark_disagg.py | 16 + .../executor/test_mamba_cache_manager.py | 115 +++++- .../_torch/executor/test_model_loader_gms.py | 48 ++- .../_torch/executor/test_py_executor.py | 355 +++++++++++++++++- .../executor/test_pytorch_model_engine.py | 17 +- .../modeling/test_modeling_qwen2_5vl.py | 23 +- 13 files changed, 844 insertions(+), 99 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py index 0a4107bd3d9b..55f9d6ca478c 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py @@ -180,6 +180,15 @@ def _prepare_qwen_vl_mrope_config( _MAX_PIXELS_TOKEN_PROBE = 1 << 31 +def _get_mrope_position_delta_cache_size( + model_config: ModelConfig[PretrainedConfig]) -> int: + """Return real sequence-slot capacity plus one reserved dummy slot.""" + max_num_seq_slots = model_config.extra_attrs.get( + 'max_num_seq_slots', + model_config.max_num_tokens * model_config.mapping.pp_size) + return max_num_seq_slots + 1 + + class Qwen2VLInputProcessorBase(BaseMultimodalInputProcessor, BaseMultimodalDummyInputsBuilder): @@ -1705,8 +1714,8 @@ def __init__( if not disable_fuse_rope: self.init_mrope_embedding(model_config) # Extra slot is reserved for CUDA graph / warmup dummy requests. - max_mrope_delta_slots = ( - model_config.max_num_tokens * model_config.mapping.pp_size + 1) + max_mrope_delta_slots = _get_mrope_position_delta_cache_size( + model_config) self.register_buffer('mrope_position_deltas_cache', torch.zeros(max_mrope_delta_slots, dtype=torch.int32, diff --git a/tensorrt_llm/_torch/models/modeling_qwen3vl.py b/tensorrt_llm/_torch/models/modeling_qwen3vl.py index 95da74e5d3bd..e2015e8e8df6 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3vl.py @@ -56,6 +56,7 @@ from .modeling_qwen2vl import ( Qwen2_5_VLVisionAttention, Qwen2VLInputProcessorBase, + _get_mrope_position_delta_cache_size, _prepare_qwen_vl_mrope_config, _prepare_qwen_vl_vision_attn_metadata, ) @@ -1206,7 +1207,7 @@ def __init__( if not disable_fuse_rope: self.init_mrope_embedding(model_config) # Extra slot is reserved for CUDA graph / warmup dummy requests. - max_mrope_delta_slots = model_config.max_num_tokens * model_config.mapping.pp_size + 1 + max_mrope_delta_slots = _get_mrope_position_delta_cache_size(model_config) self.register_buffer( "mrope_position_deltas_cache", torch.zeros( diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 4d8dc3b717c1..350f135e6a45 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -1120,8 +1120,22 @@ def free_resources(self, request: LlmRequest, pin_on_release: bool = False): KVCacheManager.free_resources(self, request, pin_on_release) def add_dummy_requests(self, request_ids: List[int], **kwargs): - MambaCacheManager.add_dummy_requests(self, request_ids) - return KVCacheManager.add_dummy_requests(self, request_ids, **kwargs) + requests = KVCacheManager.add_dummy_requests(self, request_ids, + **kwargs) + if not requests: + return requests + try: + MambaCacheManager.add_dummy_requests(self, request_ids) + except RuntimeError: + # KV allocation completed first; restore it if the independent + # Mamba pool cannot reserve the matching dummy slot. + draft_kv_cache_manager = kwargs.get('draft_kv_cache_manager') + for request in requests: + if draft_kv_cache_manager is not None: + draft_kv_cache_manager.free_resources(request) + KVCacheManager.free_resources(self, request) + raise + return requests def shutdown(self): MambaCacheManager.shutdown(self) @@ -1831,23 +1845,29 @@ def get_num_available_tokens(self, result = super().get_num_available_tokens(token_num_upper_bound, max_num_draft_tokens, **kwargs) - # Also bound by the recurrent-state pool capacity: each request needs - # roughly ceil(N / states_snapshot_interval) recurrent-state blocks - # (+1 for the corner case where N is a multiple of tokens_per_block). - # When block reuse is disabled, only one snapshot is needed per - # request, so no additional capping is required here. - interval = (self.linear_attention_metadata.states_snapshot_interval - if self.linear_attention_metadata is not None else 0) + # Also bound by the recurrent-state pool capacity. With block reuse, + # each request needs roughly ceil(N / states_snapshot_interval) blocks + # (+1 for the final live state). Without reuse, it still needs one live + # recurrent-state block, so an exhausted pool must fail preflight. + if self.linear_attention_metadata is None: + # A PP rank with no local Mamba layers has no recurrent-state pool. + return max(result, 0) + + interval = self.linear_attention_metadata.states_snapshot_interval + stats = self.impl.get_kv_cache_stats() + rs_free = stats.num_free_blocks_per_window_size.get( + LinearCacheType.RECURRENT_STATES.value, 0) if interval and interval > 0: - stats = self.impl.get_kv_cache_stats() - rs_free = stats.num_free_blocks_per_window_size.get( - LinearCacheType.RECURRENT_STATES.value, 0) # Reserve 1 block for the always-allocated last block (corner case # / final live state) so we don't promise more tokens than the # pool can actually back at allocation time. usable_rs_blocks = max(0, rs_free - 1) - rs_token_cap = usable_rs_blocks * interval + rs_token_cap = max( + 0, usable_rs_blocks * interval - self.num_extra_kv_tokens - + max_num_draft_tokens) result = min(result, rs_token_cap) + elif rs_free == 0: + result = 0 return max(result, 0) def get_ssm_states(self, layer_idx: int) -> torch.Tensor: @@ -1968,8 +1988,7 @@ def _setup_state_indices(self, requests=None) -> None: self.cuda_state_indices.copy_(self._host_state_indices, non_blocking=True) - self._refresh_dummy_request_mask( - [req.is_dummy for req in self.requests]) + self._refresh_dummy_request_mask([req.is_dummy for req in requests]) # Build request_id → pool block offset mapping so that # get_state_indices can return indices in arbitrary request order. diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 2dfebdb65e7b..ecbcf30315e0 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import bisect import contextlib import functools @@ -296,6 +299,12 @@ def __init__( self.mapping = mapping if mapping.has_pp(): init_pp_comm(mapping) + # The overlap scheduler can hold two iterations' requests at once. + # Model-side buffers indexed by py_seq_slot must span that full pool, + # not the token budget or scheduler admission batch size. + from ._util import compute_max_num_sequences + self.max_num_seq_slots = compute_max_num_sequences( + mapping, self.batch_size, llm_args.disable_overlap_scheduler) self.dist = dist if dist is not None: ExpertStatistic.create(self.dist.rank) @@ -363,6 +372,7 @@ def __init__( sparse_attention_config=self.sparse_attention_config, max_num_tokens=self.max_num_tokens, max_seq_len=self.max_seq_len, + max_num_seq_slots=self.max_num_seq_slots, lora_config=lora_config, model_weights_memory_tag=model_weights_memory_tag, model_weights_restore_mode=model_weights_restore_mode, @@ -373,6 +383,7 @@ def __init__( setattr(self, "moe_load_balancer", moe_load_balancer) else: self.model = model + self._validate_mrope_position_delta_cache_capacity() if drafting_loop_wrapper is not None: self.model = drafting_loop_wrapper(self.model) self.model_is_wrapped = True @@ -877,6 +888,34 @@ def no_cuda_graph(self): finally: self.cuda_graph_runner.enabled = _run_cuda_graphs + def _validate_mrope_position_delta_cache_capacity(self) -> None: + """Validate slot-indexed MRoPE state on preconstructed models. + + Models created by ModelLoader receive ``max_num_seq_slots`` before + construction. A caller-supplied model bypasses that path, so fail + early instead of letting the overlap scheduler index past an old-sized + cache at runtime. + """ + mrope_position_deltas_cache = getattr(self.model, + "mrope_position_deltas_cache", + None) + if mrope_position_deltas_cache is None: + mrope_position_deltas_cache = getattr( + getattr(self.model, "draft_model", None), + "mrope_position_deltas_cache", None) + if mrope_position_deltas_cache is None: + return + + required_size = self.max_num_seq_slots + 1 + actual_size = mrope_position_deltas_cache.shape[0] + if actual_size < required_size: + raise ValueError( + "The supplied model's MRoPE position-delta cache has " + f"{actual_size} slots, but this executor requires at least " + f"{required_size} ({self.max_num_seq_slots} runtime sequence " + "slots plus one reserved dummy slot). Rebuild the model with " + "the executor's sequence-slot capacity.") + def _pad_batch_seed_mrope_delta_cache( self, padded_requests: ScheduledRequests) -> None: if not self.use_mrope or padded_requests.num_generation_requests == 0: @@ -2154,14 +2193,10 @@ def _set_up_spec_metadata( self, spec_resource_manager: Optional[BaseResourceManager], no_cache=False): - # Lazy import: _util imports model_engine at module level. - from ._util import compute_max_num_sequences - spec_config = self.spec_config if self.enable_spec_decode else None # Slot-indexed metadata buffers must span the SeqSlotManager pool, # which exceeds batch_size under the overlap scheduler. - num_seq_slots = compute_max_num_sequences( - self.mapping, self.batch_size, self._disable_overlap_scheduler) + num_seq_slots = self.max_num_seq_slots if no_cache: return get_spec_metadata( spec_config, @@ -3260,7 +3295,9 @@ def _prepare_tp_inputs( mrope_delta_read_seq_slots = [] # Extra model-side cache slot reserved for CUDA graph / warmup dummy # requests, whose outputs are discarded. - mrope_dummy_seq_slot = self.max_num_tokens * self.mapping.pp_size + mrope_dummy_seq_slot = getattr( + self, "max_num_seq_slots", + self.max_num_tokens * self.mapping.pp_size) num_accepted_draft_tokens = [] # per request is_enc_dec = self._is_encoder_decoder_model() cross_encoder_hidden_states: List[torch.Tensor] = [] diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 7596c2a68291..dfde27441422 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -310,7 +310,8 @@ def __init__(self, max_seq_len: Optional[int], lora_config: Optional[LoraConfig] = None, model_weights_memory_tag: Optional[ExecutorMemoryType] = None, - model_weights_restore_mode: Optional[RestoreMode] = None): + model_weights_restore_mode: Optional[RestoreMode] = None, + max_num_seq_slots: Optional[int] = None): """ Initializes the ModelLoader. @@ -326,6 +327,9 @@ def __init__(self, they can be released/materialized independently of buffers. model_weights_restore_mode: RestoreMode for the model weights virtual-memory scope. + max_num_seq_slots: Capacity of model buffers indexed by sequence + slot. This can exceed the scheduler admission batch size when + overlap scheduling is enabled. """ self.llm_args = llm_args self.mapping = mapping @@ -333,6 +337,7 @@ def __init__(self, self.sparse_attention_config = sparse_attention_config self.max_num_tokens = max_num_tokens self.max_seq_len = max_seq_len + self.max_num_seq_slots = max_num_seq_slots self.lora_config = lora_config self.model_weights_memory_tag = model_weights_memory_tag self.model_weights_restore_mode = model_weights_restore_mode @@ -340,6 +345,11 @@ def __init__(self, self._weight_pool_proxy = None self._gms_backend = None + def _set_runtime_model_config_attrs(self, config: ModelConfig) -> None: + """Attach executor-only allocation sizes before model construction.""" + if self.max_num_seq_slots is not None: + config.extra_attrs['max_num_seq_slots'] = self.max_num_seq_slots + @staticmethod def load_config_and_apply_defaults( checkpoint_dir: str, llm_args: TorchLlmArgs, @@ -1208,6 +1218,7 @@ def _load_and_validate_config( load_config_kwargs['model_kwargs'] = self.llm_args.model_kwargs config = checkpoint_loader.load_config(**load_config_kwargs) + self._set_runtime_model_config_attrs(config) # Store nvfp4 config in extra_attrs for Linear layer access config.extra_attrs[ diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index f47b90ae3465..9ed208071fcb 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import dataclasses import datetime import functools @@ -78,14 +81,16 @@ from .request_utils import (RequestBroadcaster, attach_py_objects_to_requests, derive_attention_dp_per_rank_request_cap, get_from_waiting_queue, merge_requests) -from .resource_manager import (ResourceManager, ResourceManagerType, - request_context) +from .resource_manager import (NoFreeSlotsError, ResourceManager, + ResourceManagerType, request_context) from .sampler import (AsyncWorkerMixin, Sampler, SamplerEvent, SampleState, SampleStateTensors, TRTLLMSampler) from .scheduler import (RequestScheduler, ScheduledRequests, SerializableSchedulerOutput, WaitingQueue, create_waiting_queue) from .scheduler.adp_router import ADPRouter +from .scheduler.scheduler import \ + is_decoder_context_request_waiting_for_encoder_output if TYPE_CHECKING: from ray.actor import ActorHandle @@ -709,6 +714,10 @@ def __init__( # ADP dummy role for _pad_attention_dp_dummy_request. Default is gen; # updated from observed request types. self._adp_dummy_is_gen: bool = True + # Dummy allocated by the current scheduling iteration. It is committed + # to the normal forward/termination lifecycle only after every ADP rank + # can queue; otherwise only this exact allocation is rolled back. + self._pending_adp_dummy_request: Optional[LlmRequest] = None self.async_transfer_manager = AsyncTransferManager( self.resource_manager, should_store_blocks=self.enable_disagg_partial_reuse_store) @@ -2522,6 +2531,8 @@ def _executor_loop_pp(self): can_queue, _ = self._can_queue(scheduled_batch) if not can_queue: self._revert_gen_alloc(scheduled_batch) + self._finalize_adp_dummy_allocation(can_queue) + if not can_queue: logger.debug( f"microbatch {microbatch_id} cannot be queued, skipping" ) @@ -3202,6 +3213,37 @@ def _revert_gen_alloc(self, scheduled_batch): for req in scheduled_batch.generation_requests: self.kv_cache_manager.revert_allocate_generation(req) + def _finalize_adp_dummy_allocation(self, can_queue: bool) -> None: + """Commit or roll back this iteration's tentative ADP dummy. + + Dummy allocation is rank-local, while ``can_queue`` is a TP-wide + decision. If one rank cannot allocate its dummy, peers that succeeded + must release theirs before retrying or the fixed dummy request ID leaks + cache resources on every skipped iteration. + """ + dummy_request = self._pending_adp_dummy_request + self._pending_adp_dummy_request = None + if dummy_request is None or can_queue: + return + + dummy_request.state = LlmRequestState.GENERATION_COMPLETE + if getattr(dummy_request, "py_seq_slot", None) is not None: + # Full resource preparation already ran before a connector changed + # the final queue decision. Every manager can use normal teardown. + self._terminate_request(dummy_request) + self.active_requests.remove(dummy_request) + return + + # Before prepare_resources(), only the KV and speculative managers have + # allocated anything for the dummy. Do not run the full teardown path: + # other resource managers may require a prior prepare call. + spec_resource_manager = self.resource_manager.get_resource_manager( + ResourceManagerType.SPEC_RESOURCE_MANAGER) + if spec_resource_manager is not None: + spec_resource_manager.free_resources(dummy_request) + self.kv_cache_manager.free_resources(dummy_request) + self.active_requests.remove(dummy_request) + def _revert_ctx_alloc(self, dropped_context_requests): """Revert V2 context KV growth for requests deferred after scheduling.""" for req in dropped_context_requests: @@ -3816,6 +3858,7 @@ def _executor_loop(self): for req in scheduled_batch.generation_requests: self.kv_cache_manager.revert_allocate_generation( req) + self._finalize_adp_dummy_allocation(False) continue if not self._is_kv_manager_v2: @@ -3865,6 +3908,7 @@ def _executor_loop(self): if not can_queue: self._revert_gen_alloc(scheduled_batch) + self._finalize_adp_dummy_allocation(can_queue) if can_queue: # init_disagg_gen_requests must be before drafter loop, otherwise draft requests do not have initialized matchers. @@ -4283,6 +4327,7 @@ def _executor_loop_overlap(self): for req in scheduled_batch.generation_requests: self.kv_cache_manager.revert_allocate_generation( req) + self._finalize_adp_dummy_allocation(False) continue if not self._is_kv_manager_v2: @@ -4331,6 +4376,7 @@ def _executor_loop_overlap(self): if not can_queue: self._revert_gen_alloc(scheduled_batch) + self._finalize_adp_dummy_allocation(can_queue) # If the batch is not empty on this rank, but empty on other ranks, # we need to delay the update of the previous batch's sample state, @@ -5290,36 +5336,57 @@ def _check_disagg_ctx_schedulable_status(self, def _count_schedulable_active_requests(self) -> int: """Count active requests that are ready for scheduling. - In non-disaggregated mode, all active requests are schedulable. - In disaggregated mode, only requests inside the scheduler's - schedulability window [CONTEXT_INIT, GENERATION_TO_COMPLETE) are - counted, mirroring the MicroBatchScheduler bounds - (no_schedule_until_state / no_schedule_after_state). States - outside the window — e.g. waiting for KV transfer (below) or - awaiting completion bookkeeping (at or above) — cannot produce - forward work, so counting them here would suppress the pad dummy - while the rank still schedules batch=0, turning can_queue False + Both scheduler implementations exclude requests already in the PP + inflight set and requests waiting for encoder output. They also apply + the schedulability window [CONTEXT_INIT, GENERATION_TO_COMPLETE), or + [ENCODER_INIT, GENERATION_TO_COMPLETE) for encoder-decoder models. + States outside that window — e.g. waiting for KV transfer (below) or + awaiting completion bookkeeping (at or above) — cannot produce forward + work. Counting any of these categories here would suppress the pad + dummy while the rank still schedules batch=0, turning can_queue False fleet-wide and stalling every rank. Returns: The number of active requests eligible for scheduling. """ - if self.kv_cache_transceiver is None: - return len(self.active_requests) - - # Both bounds must match the scheduler exactly. The lower bound + # Both bounds must match the scheduler exactly. The decoder-only lower # covers e.g. DISAGG_GENERATION_INIT / TRANS_IN_PROGRESS and # DISAGG_CONTEXT_WAIT_SCHEDULER (gen-first mode on the context - # server); the upper bound covers GENERATION_TO_COMPLETE and the - # terminal / post-completion disagg states. Disaggregated serving - # is decoder-only, so the scheduler's until-state is CONTEXT_INIT - # here (the ENCODER_INIT variant only applies to encoder-decoder - # models, which have no disagg path). - context_init_value = LlmRequestState.CONTEXT_INIT.value + # server); encoder-decoder widens it to ENCODER_INIT. The upper bound + # covers GENERATION_TO_COMPLETE and all terminal/post-completion states. + cross_kv_cache_manager = self.resource_manager.get_resource_manager( + ResourceManagerType.CROSS_KV_CACHE_MANAGER) + schedule_from_value = (LlmRequestState.ENCODER_INIT.value + if cross_kv_cache_manager is not None else + LlmRequestState.CONTEXT_INIT.value) to_complete_value = LlmRequestState.GENERATION_TO_COMPLETE.value - return sum(1 for req in self.active_requests - if context_init_value <= req.state_value < to_complete_value) + return sum( + 1 for req in self.active_requests + if req.request_id not in self.inflight_req_ids + and not is_decoder_context_request_waiting_for_encoder_output(req) + and schedule_from_value <= req.state_value < to_complete_value) + + def _has_adp_dummy_kv_capacity(self, + token_nums: Optional[List[int]]) -> bool: + """Check the full dummy allocation before entering rank-local code. + + V1's dummy allocator historically checked only that one block was + free. A context-side ADP dummy can be ``max_num_tokens`` long, while a + generation dummy also reserves draft tokens, so that check can still + let the C++ batch allocation throw before PP/CP peers reach their vote. + """ + token_num = (token_nums[0] if token_nums is not None else 1 + + self.max_total_draft_tokens) + mapping = getattr(self.kv_cache_manager, "mapping", None) + has_cp_helix = getattr(mapping, "has_cp_helix", None) + if callable(has_cp_helix) and has_cp_helix(): + token_num = max(token_num, 2) + draft_reserve = (self.max_total_draft_tokens + if self._adp_dummy_is_gen else 0) + available_tokens = self.kv_cache_manager.get_num_available_tokens( + token_num_upper_bound=token_num, max_num_draft_tokens=draft_reserve) + return available_tokens >= token_num def _should_skip_dummy_for_benchmark_disagg( self, num_schedulable_requests: int) -> bool: @@ -5371,6 +5438,33 @@ def _update_adp_dummy_role(self, candidates: List[LlmRequest]) -> None: elif has_gen: self._adp_dummy_is_gen = True + def _can_commit_adp_dummy_candidate( + self, needs_dummy: bool, + dummy_candidate: Optional[LlmRequest]) -> bool: + """Agree on dummy creation across one PP/CP scheduling lane. + + The first PP stage propagates request IDs to later stages, so a dummy + may only enter ``active_requests`` when every stage (and every CP shard) + needs it and allocated its local cache resources. TP/ADP lanes remain + independent here; their existing ``_can_queue`` vote decides whether + the fleet can run and finalizes or rolls back the committed candidate. + """ + local_status = (needs_dummy, not needs_dummy + or dummy_candidate is not None) + lane_statuses = [local_status] + + if self._dist_size(self.dist, "cp_size") > 1: + lane_statuses = self.dist.cp_allgather(local_status) + if self._dist_size(self.dist, "pp_size") > 1: + lane_statuses = [ + status + for stage_statuses in self.dist.pp_allgather(lane_statuses) + for status in stage_statuses + ] + + return (all(status[0] for status in lane_statuses) + and all(status[1] for status in lane_statuses)) + @nvtx_range("_pad_attention_dp_dummy_request") def _pad_attention_dp_dummy_request(self): """ @@ -5379,47 +5473,79 @@ def _pad_attention_dp_dummy_request(self): if not self.enable_attention_dp: return + assert self._pending_adp_dummy_request is None assert self.expected_num_active_requests >= len(self.active_requests) num_active_request = self._count_schedulable_active_requests() - if self._should_skip_dummy_for_benchmark_disagg(num_active_request): - return - - # Other ranks have work but this rank is idle — insert a dummy so - # it can participate in collective operations during the forward pass. - if num_active_request == 0 and self.expected_num_active_requests > 0: + skip_dummy = self._should_skip_dummy_for_benchmark_disagg( + num_active_request) + + # Allocate tentatively. PP/CP peers must agree before the request ID is + # exposed to scheduling; otherwise schedule propagation can reference a + # dummy that a later stage failed to create. + # Every stage/shard must enter the vote, including ranks that expect no + # active requests or suppress padding during benchmark fill. Returning + # early on those ranks would leave a peer blocked in pp/cp_allgather. + needs_dummy = (not skip_dummy and self.expected_num_active_requests > 0 + and num_active_request == 0) + dummy_request = None + dummy_request_ids = [ATTENTION_DP_DUMMY_REQUEST_ID] + has_live_adp_dummy = any( + request.py_request_id == ATTENTION_DP_DUMMY_REQUEST_ID + for request in self.active_requests) + spec_resource_manager = None + if needs_dummy and not has_live_adp_dummy: token_nums = None if (not self._adp_dummy_is_gen and self.kv_cache_transceiver is not None and self.max_num_tokens is not None): token_nums = [self.max_num_tokens] - dummy_request_ids = [ATTENTION_DP_DUMMY_REQUEST_ID] - dummy_requests = self.kv_cache_manager.add_dummy_requests( - request_ids=dummy_request_ids, - token_nums=token_nums, - is_gen=self._adp_dummy_is_gen, - prepare_resource=True, - max_num_draft_tokens=self.max_total_draft_tokens, - ) - if not dummy_requests: - # No free KV/mamba resources for even a 1-token dummy — - # possible while non-schedulable requests (e.g. awaiting - # completion bookkeeping or KV transfer) still hold theirs. - # Skip padding: this rank schedules an empty batch and the - # fleet skips one iteration; the resource-holding requests - # finish within an iteration or two and padding succeeds on - # the retry. Crashing here (the old unchecked [0]) is worse. + dummy_requests = None + if self._has_adp_dummy_kv_capacity(token_nums): + try: + dummy_requests = self.kv_cache_manager.add_dummy_requests( + request_ids=dummy_request_ids, + token_nums=token_nums, + is_gen=self._adp_dummy_is_gen, + prepare_resource=True, + max_num_draft_tokens=self.max_total_draft_tokens, + ) + except OutOfPagesError: + # V2 can lose a capacity race between its estimate and + # resize. Treat that known resource condition exactly like + # its normal ``None`` result and join the lane vote. + dummy_requests = None + if dummy_requests: + dummy_request = dummy_requests[0] + spec_resource_manager = self.resource_manager.get_resource_manager( + ResourceManagerType.SPEC_RESOURCE_MANAGER) + if spec_resource_manager is not None: + try: + spec_resource_manager.add_dummy_requests( + dummy_request_ids) + except NoFreeSlotsError: + # SlotManager checks capacity before inserting the one + # ADP ID, so this specific built-in failure has no + # partial speculative state to unwind. Unexpected + # ValueErrors and all RuntimeErrors remain fatal. + self.kv_cache_manager.free_resources(dummy_request) + dummy_request = None + + if not self._can_commit_adp_dummy_candidate(needs_dummy, dummy_request): + if dummy_request is not None: + if spec_resource_manager is not None: + spec_resource_manager.free_resources(dummy_request) + self.kv_cache_manager.free_resources(dummy_request) + if needs_dummy: logger.warning( - "Cannot allocate ADP pad dummy (no free cache resources);" - " rank schedules an empty batch this iteration.") - return - llm_request = dummy_requests[0] - llm_request.is_attention_dp_dummy = True - spec_resource_manager = self.resource_manager.get_resource_manager( - ResourceManagerType.SPEC_RESOURCE_MANAGER) - if spec_resource_manager is not None: - spec_resource_manager.add_dummy_requests(dummy_request_ids) - self.active_requests.append(llm_request) + "Cannot commit ADP pad dummy on every PP/CP lane rank; " + "rolled back local allocation and will retry.") + return + + assert dummy_request is not None + dummy_request.is_attention_dp_dummy = True + self.active_requests.append(dummy_request) + self._pending_adp_dummy_request = dummy_request @nvtx_range("_prepare_disagg_gen_init") def _prepare_disagg_gen_init(self, fitting_disagg_gen_init_requests): diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index aadf90a584c7..73e1df624f26 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -2223,6 +2223,10 @@ def reset_reuse_state(self): self.impl.reset_reuse_state() +class NoFreeSlotsError(ValueError): + """A slot-backed resource manager has no capacity for another request.""" + + class SlotManager: def __init__(self, max_num_requests: int): @@ -2251,7 +2255,7 @@ def add_slot(self, request_id: int): return self.slot_mapping[request_id] if len(self.free_slots) == 0: - raise ValueError("No free slots") + raise NoFreeSlotsError("No free slots") slot = self.free_slots.pop() self.slot_mapping[request_id] = slot return slot diff --git a/tests/unittest/_torch/executor/test_benchmark_disagg.py b/tests/unittest/_torch/executor/test_benchmark_disagg.py index d00fecb9bcbc..1a5952a43cd0 100644 --- a/tests/unittest/_torch/executor/test_benchmark_disagg.py +++ b/tests/unittest/_torch/executor/test_benchmark_disagg.py @@ -44,7 +44,11 @@ def _make_active_request( ) -> Mock: """Create an active request stub with disagg state flags.""" req = Mock() + req.request_id = 1 + req.py_request_id = 1 req.state_value = LlmRequestState.GENERATION_IN_PROGRESS.value + req.is_context_init_state = False + req.py_encoder_output_ready_event = None req.is_disagg_generation_init_state = in_init req.is_disagg_generation_transmission_in_progress = in_transfer req.state = ( @@ -509,12 +513,21 @@ def __init__( self.benchmark_req_queues_size = benchmark_req_queues_size self.max_total_draft_tokens = 0 self._adp_dummy_is_gen = True + self._pending_adp_dummy_request = None self.max_num_tokens = None + self.inflight_req_ids = set() self.dist = Mock() self.dist.tp_size = tp_size + self.dist.pp_size = 1 + self.dist.cp_size = 1 + self.dist.tp_allgather.side_effect = lambda value: [value] + self.dist.pp_allgather.side_effect = lambda value: [value] + self.dist.cp_allgather.side_effect = lambda value: [value] self.kv_cache_manager = Mock() + self.kv_cache_manager.mapping.has_cp_helix.return_value = False + self.kv_cache_manager.get_num_available_tokens.return_value = 1 << 30 dummy_req = Mock() dummy_req.is_attention_dp_dummy = True self.kv_cache_manager.add_dummy_requests.return_value = [dummy_req] @@ -526,7 +539,10 @@ def __init__( _pad_attention_dp_dummy_request = PyExecutor._pad_attention_dp_dummy_request _count_schedulable_active_requests = PyExecutor._count_schedulable_active_requests + _has_adp_dummy_kv_capacity = PyExecutor._has_adp_dummy_kv_capacity _should_skip_dummy_for_benchmark_disagg = PyExecutor._should_skip_dummy_for_benchmark_disagg + _can_commit_adp_dummy_candidate = PyExecutor._can_commit_adp_dummy_candidate + _dist_size = staticmethod(PyExecutor._dist_size) class TestPadAttentionDpDummyBenchmarkDisagg: diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index 30eaa2442424..74f668a60038 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -5,6 +5,7 @@ import os from types import SimpleNamespace +from unittest.mock import MagicMock import pytest import torch @@ -14,10 +15,12 @@ from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( MIN_REPLAY_HISTORY_SIZE, CppMambaHybridCacheManager, + MambaCacheManager, + MixedMambaHybridCacheManager, PythonMambaCacheManager, _get_mamba_hybrid_pool_size, ) -from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp +from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp, KVCacheManager from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.bindings.internal.batch_manager import LinearCacheType from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MTPDecodingConfig @@ -360,6 +363,7 @@ def _build_hybrid_with_mamba_layer( enable_block_reuse=False, mamba_state_cache_interval=256, is_estimating_kv_cache=False, + use_replay_state_update=False, ): """Construct a real CppMambaHybridCacheManager with one mamba layer + one full-attention layer so the parent KVCacheManager goes through the @@ -396,6 +400,7 @@ def _build_hybrid_with_mamba_layer( spec_config=spec_config, layer_mask=attn_mask, is_estimating_kv_cache=is_estimating_kv_cache, + use_replay_state_update=use_replay_state_update, ) @@ -433,6 +438,114 @@ def test_cpp_hybrid_recurrent_pool_reserves_draft_len_sentinel_slots(): ) +@skip_no_cuda +def test_cpp_hybrid_adp_dummy_mask_excludes_stale_requests(): + """The replay mask must describe the batch passed to state setup. + + ADP dummy allocation happens before the previous overlap batch has been + released, so ``self.requests`` can contain a full stale batch in addition + to the newly allocated dummy. The dummy-only setup must not size its mask + from that combined lifecycle list. + """ + max_batch_size = 2 + mgr = _build_hybrid_with_mamba_layer( + spec_config=MTPDecodingConfig(max_draft_len=1), + max_batch_size=max_batch_size, + use_replay_state_update=True, + ) + stale_requests = [SimpleNamespace(is_dummy=False) for _ in range(max_batch_size)] + mgr.requests = stale_requests.copy() + + dummy_requests = mgr.add_dummy_requests([ATTENTION_DP_DUMMY_REQUEST_ID]) + + assert dummy_requests + assert mgr.requests == stale_requests + dummy_requests + assert mgr._dummy_request_mask_host.tolist() == [True, False] + + +def test_mixed_hybrid_does_not_allocate_mamba_when_dummy_kv_fails(monkeypatch): + mgr = object.__new__(MixedMambaHybridCacheManager) + kv_add_dummy = MagicMock(return_value=None) + mamba_add_dummy = MagicMock() + monkeypatch.setattr(KVCacheManager, "add_dummy_requests", kv_add_dummy) + monkeypatch.setattr(MambaCacheManager, "add_dummy_requests", mamba_add_dummy) + + requests = mgr.add_dummy_requests([ATTENTION_DP_DUMMY_REQUEST_ID]) + + assert requests is None + mamba_add_dummy.assert_not_called() + + +def test_mixed_hybrid_rolls_back_kv_when_mamba_dummy_fails(monkeypatch): + mgr = object.__new__(MixedMambaHybridCacheManager) + dummy_request = SimpleNamespace(py_request_id=ATTENTION_DP_DUMMY_REQUEST_ID) + draft_kv_cache_manager = MagicMock() + kv_add_dummy = MagicMock(return_value=[dummy_request]) + kv_free = MagicMock() + mamba_add_dummy = MagicMock(side_effect=RuntimeError("no mamba slot")) + monkeypatch.setattr(KVCacheManager, "add_dummy_requests", kv_add_dummy) + monkeypatch.setattr(KVCacheManager, "free_resources", kv_free) + monkeypatch.setattr(MambaCacheManager, "add_dummy_requests", mamba_add_dummy) + + with pytest.raises(RuntimeError, match="no mamba slot"): + mgr.add_dummy_requests( + [ATTENTION_DP_DUMMY_REQUEST_ID], draft_kv_cache_manager=draft_kv_cache_manager + ) + + draft_kv_cache_manager.free_resources.assert_called_once_with(dummy_request) + kv_free.assert_called_once_with(mgr, dummy_request) + + +def test_cpp_hybrid_available_tokens_reserves_extra_and_draft_tokens(monkeypatch): + mgr = object.__new__(CppMambaHybridCacheManager) + mgr.num_extra_kv_tokens = 1 + mgr.linear_attention_metadata = SimpleNamespace(states_snapshot_interval=16) + mgr.impl = MagicMock() + mgr.impl.get_kv_cache_stats.return_value = SimpleNamespace( + num_free_blocks_per_window_size={LinearCacheType.RECURRENT_STATES.value: 3} + ) + monkeypatch.setattr(KVCacheManager, "get_num_available_tokens", MagicMock(return_value=100)) + + available_tokens = mgr.get_num_available_tokens( + token_num_upper_bound=100, max_num_draft_tokens=2 + ) + + # One recurrent block is reserved for the final live state. The remaining + # two cover 32 tokens, of which one extra KV token and two draft tokens are + # consumed immediately after the initial dummy allocation. + assert available_tokens == 29 + + +@pytest.mark.parametrize("rs_free, expected", [(0, 0), (1, 100)]) +def test_cpp_hybrid_available_tokens_requires_live_state_without_reuse( + monkeypatch, rs_free, expected +): + mgr = object.__new__(CppMambaHybridCacheManager) + mgr.num_extra_kv_tokens = 1 + mgr.linear_attention_metadata = SimpleNamespace(states_snapshot_interval=0) + mgr.impl = MagicMock() + mgr.impl.get_kv_cache_stats.return_value = SimpleNamespace( + num_free_blocks_per_window_size={LinearCacheType.RECURRENT_STATES.value: rs_free} + ) + monkeypatch.setattr(KVCacheManager, "get_num_available_tokens", MagicMock(return_value=100)) + + available_tokens = mgr.get_num_available_tokens(token_num_upper_bound=100) + + assert available_tokens == expected + + +def test_cpp_hybrid_available_tokens_without_local_mamba_uses_attention_cap(monkeypatch): + mgr = object.__new__(CppMambaHybridCacheManager) + mgr.linear_attention_metadata = None + mgr.impl = MagicMock() + monkeypatch.setattr(KVCacheManager, "get_num_available_tokens", MagicMock(return_value=100)) + + available_tokens = mgr.get_num_available_tokens(token_num_upper_bound=100) + + assert available_tokens == 100 + mgr.impl.get_kv_cache_stats.assert_not_called() + + def _build_hybrid_with_mamba_layer_pp( spec_config=None, max_batch_size=4, enable_block_reuse=False, pp_size=2 ): diff --git a/tests/unittest/_torch/executor/test_model_loader_gms.py b/tests/unittest/_torch/executor/test_model_loader_gms.py index a1a6e1b9b700..9f19a3bfaeef 100644 --- a/tests/unittest/_torch/executor/test_model_loader_gms.py +++ b/tests/unittest/_torch/executor/test_model_loader_gms.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Unit tests for GMS-specific branches in ``ModelLoader``.""" +"""Unit tests for ``ModelLoader`` configuration and GMS-specific branches.""" from contextlib import contextmanager, nullcontext from types import SimpleNamespace @@ -112,6 +112,52 @@ def _make_loader(monkeypatch, *, events, spec_config=None): return loader +def test_load_config_propagates_runtime_seq_slot_capacity(monkeypatch): + llm_args = MagicMock() + llm_args.model_kwargs = None + llm_args.cuda_graph_config = None + llm_args.moe_config = MagicMock() + llm_args.nvfp4_gemm_config = MagicMock() + llm_args.multimodal_config = MagicMock() + llm_args.kv_cache_config = MagicMock() + loader = ModelLoader( + llm_args=llm_args, + mapping=MagicMock(), + spec_config=None, + sparse_attention_config=None, + max_num_tokens=128, + max_seq_len=128, + max_num_seq_slots=16, + ) + text_config = SimpleNamespace(hidden_size=64) + pretrained_config = SimpleNamespace( + get_text_config=lambda: text_config, + torch_dtype=torch.float16, + model_type="test", + ) + config = SimpleNamespace( + extra_attrs={}, + nvfp4_gemm_allowed_backends=None, + max_num_tokens=128, + pretrained_config=pretrained_config, + ) + checkpoint_loader = MagicMock() + checkpoint_loader.load_config.return_value = config + monkeypatch.setattr(model_loader_mod, "validate_encoder_decoder_tp_scope", lambda *_args: None) + monkeypatch.setattr( + model_loader_mod, "validate_encoder_decoder_kv_cache_config", lambda *_args: None + ) + monkeypatch.setattr(model_loader_mod, "validate_and_set_kv_cache_quant", lambda *_args: None) + monkeypatch.setattr( + model_loader_mod, "validate_and_set_mamba_ssm_cache_dtype", lambda *_args: None + ) + + loaded_config = loader._load_and_validate_config("/ckpt", checkpoint_loader) + + assert loaded_config is config + assert config.extra_attrs["max_num_seq_slots"] == 16 + + def _build_gms_backend(*, is_rw, events): backend = MagicMock() backend.connect.return_value = True diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 490b86e61603..6dd4167a1338 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + """Tests for PyExecutor request handling functionality. This module tests the request handling logic that was moved from ExecutorRequestQueue @@ -22,9 +25,14 @@ SHUTDOWN_REQUEST_ID, RequestQueueItem, ) -from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState, SamplingConfig +from tensorrt_llm._torch.pyexecutor.llm_request import ( + ATTENTION_DP_DUMMY_REQUEST_ID, + LlmRequest, + LlmRequestState, + SamplingConfig, +) from tensorrt_llm._torch.pyexecutor.py_executor import DisaggTransferAdmissionController, PyExecutor -from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType +from tensorrt_llm._torch.pyexecutor.resource_manager import NoFreeSlotsError, ResourceManagerType from tensorrt_llm._torch.pyexecutor.scheduler import ( FCFSWaitingQueue, ScheduledRequests, @@ -1186,6 +1194,7 @@ def _make_adp_request( if isinstance(state, LlmRequestState) else _SENTINEL_STATE_VALUES.get(state, LlmRequestState.GENERATION_IN_PROGRESS.value) ) + req.request_id = request_id req.py_request_id = request_id req.is_child = is_child req.parent_request_id = parent_request_id @@ -1196,6 +1205,9 @@ def _make_adp_request( req.is_dummy_request = is_dummy_request req.is_attention_dp_dummy = False req.llm_request_type = llm_request_type + req.py_seq_slot = None + req.is_context_init_state = state == LlmRequestState.CONTEXT_INIT + req.py_encoder_output_ready_event = None return req @@ -1220,13 +1232,29 @@ def __init__( self.max_total_draft_tokens = 0 self.max_num_tokens = max_num_tokens self._adp_dummy_is_gen = True + self._pending_adp_dummy_request = None self.add_dummy_calls = [] + self.inflight_req_ids = set() + + self.dist = Mock() + self.dist.tp_size = 1 + self.dist.pp_size = 1 + self.dist.cp_size = 1 + self.dist.tp_allgather.side_effect = lambda value: [value] + self.dist.pp_allgather.side_effect = lambda value: [value] + self.dist.cp_allgather.side_effect = lambda value: [value] kv_cache_manager = Mock() + kv_cache_manager.mapping.has_cp_helix.return_value = False + kv_cache_manager.get_num_available_tokens.return_value = 1 << 30 def _add_dummy(**kwargs): self.add_dummy_calls.append(kwargs) - req = _make_adp_request(_STATE_GENERATION_IN_PROGRESS, is_dummy_request=True) + req = _make_adp_request( + _STATE_GENERATION_IN_PROGRESS, + request_id=kwargs["request_ids"][0], + is_dummy_request=True, + ) return [req] kv_cache_manager.add_dummy_requests.side_effect = _add_dummy @@ -1237,8 +1265,14 @@ def _add_dummy(**kwargs): def _run_pad(stub): - for helper in ("_count_schedulable_active_requests", "_should_skip_dummy_for_benchmark_disagg"): + for helper in ( + "_count_schedulable_active_requests", + "_has_adp_dummy_kv_capacity", + "_should_skip_dummy_for_benchmark_disagg", + "_can_commit_adp_dummy_candidate", + ): setattr(stub, helper, types.MethodType(getattr(PyExecutor, helper), stub)) + stub._dist_size = PyExecutor._dist_size PyExecutor._pad_attention_dp_dummy_request(stub) @@ -1318,20 +1352,57 @@ def test_adp_dummy_role_unchanged_when_attention_dp_disabled(): assert stub._adp_dummy_is_gen is True -def test_pad_dummy_counts_generation_to_complete(): - # Regression for test_ptp_quickstart_advanced_deepseek_v3_lite_4gpus_adp_balance: - # a GENERATION_TO_COMPLETE request still counts as active, so no dummy is added - # on top of it (which would overflow max_batch_size=1). +def test_pad_dummy_added_when_only_to_complete_requests_non_disagg(): + # GENERATION_TO_COMPLETE is outside the MicroBatchScheduler upper bound in + # every mode. Counting it would suppress padding even though this rank's + # next scheduled batch is empty. stub = _StubADPExecutor(kv_cache_transceiver=None) stub.active_requests = [_make_adp_request(_STATE_GENERATION_TO_COMPLETE)] stub.expected_num_active_requests = 1 _run_pad(stub) + assert len(stub.add_dummy_calls) == 1 + assert len(stub.active_requests) == 2 + + +def test_pad_dummy_counts_encoder_init_for_encoder_decoder(): + stub = _StubADPExecutor(kv_cache_transceiver=None) + cross_kv_cache_manager = object() + stub.resource_manager.get_resource_manager.side_effect = ( + lambda resource_type: cross_kv_cache_manager + if resource_type == ResourceManagerType.CROSS_KV_CACHE_MANAGER + else None + ) + stub.active_requests = [_make_adp_request(LlmRequestState.ENCODER_INIT)] + stub.expected_num_active_requests = 1 + + _run_pad(stub) + assert stub.add_dummy_calls == [] assert len(stub.active_requests) == 1 +def test_pad_dummy_ignores_context_waiting_for_encoder_output(): + stub = _StubADPExecutor(kv_cache_transceiver=None) + cross_kv_cache_manager = object() + stub.resource_manager.get_resource_manager.side_effect = ( + lambda resource_type: cross_kv_cache_manager + if resource_type == ResourceManagerType.CROSS_KV_CACHE_MANAGER + else None + ) + waiting_request = _make_adp_request(LlmRequestState.CONTEXT_INIT) + waiting_request.py_encoder_output_ready_event = Mock() + waiting_request.py_encoder_output_ready_event.query.return_value = False + stub.active_requests = [waiting_request] + stub.expected_num_active_requests = 2 + + _run_pad(stub) + + assert len(stub.add_dummy_calls) == 1 + assert len(stub.active_requests) == 2 + + def test_pad_dummy_added_when_only_to_complete_requests_disagg(): # In disaggregated mode a GENERATION_TO_COMPLETE request is refused by # MicroBatchScheduler (no_schedule_after_state), so a rank holding only @@ -1373,13 +1444,281 @@ def test_pad_dummy_allocation_failure_skips_padding(): stub.expected_num_active_requests = 2 stub.kv_cache_manager.add_dummy_requests.side_effect = None stub.kv_cache_manager.add_dummy_requests.return_value = None + stub.dist.pp_size = 2 + stub.dist.pp_allgather.side_effect = None + stub.dist.pp_allgather.return_value = [[(True, False)], [(True, True)]] _run_pad(stub) + stub.dist.pp_allgather.assert_called_once_with([(True, False)]) assert len(stub.active_requests) == 1 assert not any(r.is_attention_dp_dummy for r in stub.active_requests) +def test_pad_dummy_context_capacity_failure_joins_lane_consensus(): + stub = _StubADPExecutor(max_num_tokens=4096) + stub._adp_dummy_is_gen = False + stub.kv_cache_manager.get_num_available_tokens.return_value = 1024 + stub.dist.pp_size = 2 + stub.dist.pp_allgather.side_effect = None + stub.dist.pp_allgather.return_value = [[(True, False)], [(True, True)]] + + _run_pad(stub) + + stub.kv_cache_manager.get_num_available_tokens.assert_called_once_with( + token_num_upper_bound=4096, max_num_draft_tokens=0 + ) + stub.kv_cache_manager.add_dummy_requests.assert_not_called() + stub.dist.pp_allgather.assert_called_once_with([(True, False)]) + assert stub._pending_adp_dummy_request is None + + +def test_pad_dummy_generation_capacity_failure_joins_lane_consensus(): + stub = _StubADPExecutor() + stub.kv_cache_manager.get_num_available_tokens.return_value = 0 + stub.dist.pp_size = 2 + stub.dist.pp_allgather.side_effect = None + stub.dist.pp_allgather.return_value = [[(True, False)], [(True, True)]] + + _run_pad(stub) + + stub.kv_cache_manager.get_num_available_tokens.assert_called_once_with( + token_num_upper_bound=1, max_num_draft_tokens=0 + ) + stub.kv_cache_manager.add_dummy_requests.assert_not_called() + stub.dist.pp_allgather.assert_called_once_with([(True, False)]) + assert stub._pending_adp_dummy_request is None + + +def test_pad_dummy_pp_peer_allocation_failure_rolls_back_candidate(): + stub = _StubADPExecutor() + terminal_request = _make_adp_request(_STATE_GENERATION_TO_COMPLETE) + stub.active_requests = [terminal_request] + stub.expected_num_active_requests = 2 + stub.dist.pp_size = 2 + stub.dist.pp_allgather.side_effect = None + stub.dist.pp_allgather.return_value = [[(True, True)], [(True, False)]] + dummy_request = _make_adp_request( + _STATE_GENERATION_IN_PROGRESS, + request_id=ATTENTION_DP_DUMMY_REQUEST_ID, + is_dummy_request=True, + ) + spec_resource_manager = Mock() + stub.resource_manager.get_resource_manager.return_value = spec_resource_manager + stub.kv_cache_manager.add_dummy_requests.side_effect = None + stub.kv_cache_manager.add_dummy_requests.return_value = [dummy_request] + + _run_pad(stub) + + assert stub.active_requests == [terminal_request] + spec_resource_manager.add_dummy_requests.assert_called_once_with( + [ATTENTION_DP_DUMMY_REQUEST_ID] + ) + spec_resource_manager.free_resources.assert_called_once_with(dummy_request) + stub.kv_cache_manager.free_resources.assert_called_once_with(dummy_request) + assert stub._pending_adp_dummy_request is None + + +@pytest.mark.parametrize( + "cp_statuses, pp_statuses", + [ + ([(True, True), (True, False)], None), + ( + [(True, True), (True, True)], + [ + [(True, True), (True, True)], + [(True, True), (True, False)], + ], + ), + ], + ids=["cp", "cp-and-pp"], +) +def test_adp_dummy_lane_consensus_includes_every_cp_shard(cp_statuses, pp_statuses): + stub = _StubADPExecutor() + stub._dist_size = PyExecutor._dist_size + stub.dist.cp_size = 2 + stub.dist.cp_allgather.side_effect = None + stub.dist.cp_allgather.return_value = cp_statuses + if pp_statuses is not None: + stub.dist.pp_size = 2 + stub.dist.pp_allgather.side_effect = None + stub.dist.pp_allgather.return_value = pp_statuses + + can_commit = PyExecutor._can_commit_adp_dummy_candidate( + stub, needs_dummy=True, dummy_candidate=object() + ) + + assert can_commit is False + stub.dist.cp_allgather.assert_called_once_with((True, True)) + if pp_statuses is not None: + stub.dist.pp_allgather.assert_called_once_with(cp_statuses) + + +def test_pad_dummy_busy_pp_rank_participates_in_lane_consensus(): + stub = _StubADPExecutor() + active_request = _make_adp_request(_STATE_GENERATION_IN_PROGRESS) + stub.active_requests = [active_request] + stub.expected_num_active_requests = 1 + stub.dist.pp_size = 2 + stub.dist.pp_allgather.side_effect = None + stub.dist.pp_allgather.return_value = [[(False, True)], [(True, True)]] + + _run_pad(stub) + + stub.dist.pp_allgather.assert_called_once_with([(False, True)]) + assert stub.active_requests == [active_request] + stub.kv_cache_manager.add_dummy_requests.assert_not_called() + + +@pytest.mark.parametrize("skip_dummy", [False, True]) +def test_pad_dummy_nonparticipating_pp_rank_still_enters_lane_consensus( + skip_dummy, +): + stub = _StubADPExecutor() + stub.expected_num_active_requests = 0 + if skip_dummy: + stub._benchmark_fill_phase_active = True + stub.dist.pp_size = 2 + stub.dist.pp_allgather.side_effect = None + stub.dist.pp_allgather.return_value = [[(False, True)], [(True, True)]] + + _run_pad(stub) + + stub.dist.pp_allgather.assert_called_once_with([(False, True)]) + stub.kv_cache_manager.add_dummy_requests.assert_not_called() + assert stub._pending_adp_dummy_request is None + + +def test_pad_dummy_added_when_only_active_request_is_inflight(): + stub = _StubADPExecutor(kv_cache_transceiver=None) + active_request = _make_adp_request(_STATE_GENERATION_IN_PROGRESS) + stub.active_requests = [active_request] + stub.expected_num_active_requests = 1 + stub.inflight_req_ids.add(active_request.request_id) + + _run_pad(stub) + + assert len(stub.add_dummy_calls) == 1 + assert len(stub.active_requests) == 2 + + +def test_pad_dummy_does_not_reallocate_live_inflight_dummy_id(): + stub = _StubADPExecutor(kv_cache_transceiver=None) + prior_dummy = _make_adp_request( + _STATE_GENERATION_IN_PROGRESS, + request_id=ATTENTION_DP_DUMMY_REQUEST_ID, + is_dummy_request=True, + ) + prior_dummy.is_attention_dp_dummy = True + stub.active_requests = [prior_dummy] + stub.expected_num_active_requests = 1 + stub.inflight_req_ids.add(ATTENTION_DP_DUMMY_REQUEST_ID) + stub.dist.pp_size = 2 + stub.dist.pp_allgather.side_effect = None + stub.dist.pp_allgather.return_value = [[(True, False)], [(True, True)]] + + _run_pad(stub) + + stub.kv_cache_manager.add_dummy_requests.assert_not_called() + stub.dist.pp_allgather.assert_called_once_with([(True, False)]) + assert stub.active_requests == [prior_dummy] + assert stub._pending_adp_dummy_request is None + + +def test_pad_dummy_spec_allocation_failure_rolls_back_kv_candidate(): + stub = _StubADPExecutor() + terminal_request = _make_adp_request(_STATE_GENERATION_TO_COMPLETE) + stub.active_requests = [terminal_request] + stub.expected_num_active_requests = 2 + spec_resource_manager = Mock() + spec_resource_manager.add_dummy_requests.side_effect = NoFreeSlotsError("No free slots") + stub.resource_manager.get_resource_manager.return_value = spec_resource_manager + stub.dist.pp_size = 2 + stub.dist.pp_allgather.side_effect = None + stub.dist.pp_allgather.return_value = [[(True, False)], [(True, True)]] + + _run_pad(stub) + + assert stub.active_requests == [terminal_request] + assert stub._pending_adp_dummy_request is None + stub.dist.pp_allgather.assert_called_once_with([(True, False)]) + stub.kv_cache_manager.free_resources.assert_called_once() + + +def test_adp_dummy_peer_empty_rolls_back_and_retry_succeeds(): + stub = _StubADPExecutor() + spec_resource_manager = Mock() + stub.resource_manager.get_resource_manager.return_value = spec_resource_manager + terminal_request = _make_adp_request(_STATE_GENERATION_TO_COMPLETE) + stub.active_requests = [terminal_request] + stub.expected_num_active_requests = 2 + + _run_pad(stub) + first_dummy = stub._pending_adp_dummy_request + assert first_dummy is not None + + stub.dist.tp_allgather.side_effect = None + stub.dist.tp_allgather.return_value = [1, 0] + can_queue, _ = PyExecutor._can_queue(stub, types.SimpleNamespace(batch_size=1)) + assert can_queue is False + PyExecutor._finalize_adp_dummy_allocation(stub, can_queue) + + assert stub.active_requests == [terminal_request] + spec_resource_manager.free_resources.assert_called_once_with(first_dummy) + stub.kv_cache_manager.free_resources.assert_called_once_with(first_dummy) + + stub.dist.tp_allgather.return_value = [1, 1] + _run_pad(stub) + second_dummy = stub._pending_adp_dummy_request + assert second_dummy is not None + assert second_dummy is not first_dummy + + can_queue, _ = PyExecutor._can_queue(stub, types.SimpleNamespace(batch_size=1)) + assert can_queue is True + PyExecutor._finalize_adp_dummy_allocation(stub, can_queue) + + assert stub._pending_adp_dummy_request is None + assert stub.active_requests == [terminal_request, second_dummy] + assert spec_resource_manager.add_dummy_requests.call_count == 2 + spec_resource_manager.free_resources.assert_called_once_with(first_dummy) + stub.kv_cache_manager.free_resources.assert_called_once_with(first_dummy) + + +def test_adp_dummy_rollback_only_frees_pending_candidate(): + stub = _StubADPExecutor() + prior_dummy = _make_adp_request( + _STATE_GENERATION_IN_PROGRESS, request_id=10, is_dummy_request=True + ) + prior_dummy.is_attention_dp_dummy = True + current_dummy = _make_adp_request( + _STATE_GENERATION_IN_PROGRESS, request_id=11, is_dummy_request=True + ) + current_dummy.is_attention_dp_dummy = True + stub.active_requests = [prior_dummy, current_dummy] + stub._pending_adp_dummy_request = current_dummy + + PyExecutor._finalize_adp_dummy_allocation(stub, can_queue=False) + + assert stub.active_requests == [prior_dummy] + stub.kv_cache_manager.free_resources.assert_called_once_with(current_dummy) + + +def test_adp_dummy_post_prepare_rollback_uses_normal_termination(): + stub = _StubADPExecutor() + dummy_request = _make_adp_request(_STATE_GENERATION_IN_PROGRESS, is_dummy_request=True) + dummy_request.is_attention_dp_dummy = True + dummy_request.py_seq_slot = 3 + stub.active_requests = [dummy_request] + stub._pending_adp_dummy_request = dummy_request + stub._terminate_request = Mock() + + PyExecutor._finalize_adp_dummy_allocation(stub, can_queue=False) + + stub._terminate_request.assert_called_once_with(dummy_request) + stub.kv_cache_manager.free_resources.assert_not_called() + assert stub.active_requests == [] + + def test_pad_dummy_skips_when_active_request_present(): stub = _StubADPExecutor() stub.active_requests = [_make_adp_request(_STATE_GENERATION_IN_PROGRESS)] diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index e5497d999e23..f4ce3a7e8899 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -3,6 +3,7 @@ import unittest from dataclasses import dataclass +from types import SimpleNamespace from unittest.mock import Mock import torch @@ -594,6 +595,7 @@ def test_prepare_tp_inputs_with_partial_mrope_segments(self) -> None: attn_metadata.is_cuda_graph = False model_engine.max_num_tokens = 32 + model_engine.max_num_seq_slots = 8 model_engine.input_ids_cuda = torch.zeros(32, dtype=torch.int32, device='cuda') @@ -625,6 +627,9 @@ def test_prepare_tp_inputs_with_partial_mrope_segments(self) -> None: dummy_request.sampling_config.beam_width = 1 dummy_request.py_multimodal_data = {} dummy_request.is_cuda_graph_dummy = True + dummy_request.py_mrope_position_delta = torch.tensor([[0]], + dtype=torch.int32, + device='cuda') scheduled_requests = ScheduledRequests() scheduled_requests.context_requests_last_chunk = [] @@ -646,11 +651,21 @@ def test_prepare_tp_inputs_with_partial_mrope_segments(self) -> None: self.assertEqual(result["mrope_delta_write_seq_slots"].cpu().tolist(), [0]) self.assertEqual(result["mrope_delta_read_seq_slots"].cpu().tolist(), - [0]) + [0, model_engine.max_num_seq_slots]) self.assertNotIn("multimodal_embedding", multimodal_request.py_multimodal_data) kv_cache_manager.shutdown() + def test_preconstructed_mrope_model_requires_runtime_seq_slot_capacity( + self) -> None: + model_engine = object.__new__(PyTorchModelEngine) + model_engine.max_num_seq_slots = 8 + model_engine.model = SimpleNamespace( + mrope_position_deltas_cache=torch.zeros(8, dtype=torch.int32)) + + with self.assertRaisesRegex(ValueError, "requires at least 9"): + model_engine._validate_mrope_position_delta_cache_capacity() + def test_kv_cache_manager_with_execution_stream(self): """Test that KVCacheManager uses the provided execution_stream. """ diff --git a/tests/unittest/_torch/modeling/test_modeling_qwen2_5vl.py b/tests/unittest/_torch/modeling/test_modeling_qwen2_5vl.py index 3f6bc50706ac..a5e177a2c03f 100644 --- a/tests/unittest/_torch/modeling/test_modeling_qwen2_5vl.py +++ b/tests/unittest/_torch/modeling/test_modeling_qwen2_5vl.py @@ -24,7 +24,8 @@ Qwen2VLHfWeightMapper from tensorrt_llm._torch.models.modeling_qwen2vl import ( Qwen2_5_VLModel, Qwen2VisionModelBase, Qwen2VLInputProcessorBase, - Qwen2VLModel, _prepare_qwen_vl_mrope_config) + Qwen2VLModel, _get_mrope_position_delta_cache_size, + _prepare_qwen_vl_mrope_config) from tensorrt_llm._torch.models.modeling_qwen3vl import \ Qwen3VLInputProcessorBase from tensorrt_llm._utils import get_sm_version @@ -426,13 +427,20 @@ def _mrope_param(delta: int) -> MultimodalParams: }) +def test_mrope_delta_cache_size_uses_runtime_seq_slot_capacity(): + model_config = ModelConfig(max_num_tokens=32) + model_config.extra_attrs['max_num_seq_slots'] = 8 + + assert _get_mrope_position_delta_cache_size(model_config) == 9 + + def test_prepare_qwen_vl_mrope_config_mixed_context_generation(): rotary_dim = 2 num_tokens = 5 position_ids = torch.arange(3 * num_tokens, device="cuda", dtype=torch.int32).reshape(3, 1, num_tokens) - mrope_position_deltas_cache = torch.zeros(8, + mrope_position_deltas_cache = torch.zeros(9, device="cuda", dtype=torch.int32) mrope_rotary_cos_sin_workspace = torch.empty((1, 8 * rotary_dim * 2), @@ -441,24 +449,25 @@ def test_prepare_qwen_vl_mrope_config_mixed_context_generation(): config = _prepare_qwen_vl_mrope_config( multimodal_params=[_mrope_param(11), _mrope_param(22)], - num_generation_requests=2, + num_generation_requests=3, position_ids=position_ids, rotary_emb=_FakeRotaryEmbedding(rotary_dim), mrope_position_deltas_cache=mrope_position_deltas_cache, mrope_rotary_cos_sin_workspace=mrope_rotary_cos_sin_workspace, - mrope_delta_write_seq_slots=torch.tensor([2, 5], + mrope_delta_write_seq_slots=torch.tensor([4, 7], device="cuda", dtype=torch.long), - mrope_delta_read_seq_slots=torch.tensor([2, 5], + mrope_delta_read_seq_slots=torch.tensor([4, 7, 8], device="cuda", dtype=torch.long)) torch.testing.assert_close( - mrope_position_deltas_cache[[2, 5]], + mrope_position_deltas_cache[[4, 7]], torch.tensor([11, 22], device="cuda", dtype=torch.int32)) + assert mrope_position_deltas_cache[8].item() == 0 torch.testing.assert_close( config["mrope_position_deltas"], - torch.tensor([[11], [22]], device="cuda", dtype=torch.int32)) + torch.tensor([[11], [22], [0]], device="cuda", dtype=torch.int32)) packed = config["mrope_rotary_cos_sin"].view(1, num_tokens, rotary_dim, 2) cos, sin = _FakeRotaryEmbedding(rotary_dim).get_cos_sin(position_ids) From 25a2422c12ca1d2cef8a613a9ff19b5f6e96cafc Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Mon, 13 Jul 2026 17:06:07 -0700 Subject: [PATCH 4/4] [None][fix] Scope overlap fixes to DeepSeek-V4 Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../_torch/auto_deploy/shim/ad_executor.py | 4 +- .../_torch/models/modeling_qwen2vl.py | 13 +- .../_torch/models/modeling_qwen3vl.py | 3 +- tensorrt_llm/_torch/pyexecutor/_util.py | 54 +++-- .../_torch/pyexecutor/mamba_cache_manager.py | 49 ++-- .../_torch/pyexecutor/model_engine.py | 68 +++--- .../_torch/pyexecutor/model_loader.py | 13 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 189 ++++++--------- .../_torch/pyexecutor/py_executor_creator.py | 35 +-- .../_torch/pyexecutor/sampler/sampler.py | 12 +- tensorrt_llm/_torch/speculative/eagle3.py | 10 +- tensorrt_llm/_torch/speculative/interface.py | 4 +- .../_torch/speculative/spec_tree_manager.py | 18 +- tensorrt_llm/_torch/speculative/utils.py | 21 +- .../_torch/executor/test_benchmark_disagg.py | 12 +- .../executor/test_mamba_cache_manager.py | 115 +-------- .../_torch/executor/test_model_loader_gms.py | 48 +--- .../_torch/executor/test_py_executor.py | 224 ++---------------- .../executor/test_pytorch_model_engine.py | 17 +- .../_torch/executor/test_seq_slot_sizing.py | 99 ++++++-- .../modeling/test_modeling_qwen2_5vl.py | 23 +- 21 files changed, 313 insertions(+), 718 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py index c6bc6c2006b5..86488bc36cda 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py @@ -1310,9 +1310,7 @@ def create_autodeploy_executor( ) guided_decoder = GuidedDecoder( guided_decoding_config=guided_decoding_config, - # Guided-decoder state is indexed by sequence slot; size it to the - # SeqSlotManager pool above (max_batch_size * pp_size). - max_num_sequences=max_num_sequences, + max_num_sequences=ad_config.max_batch_size, vocab_size_padded=vocab_size_padded, ) diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py index 55f9d6ca478c..0a4107bd3d9b 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py @@ -180,15 +180,6 @@ def _prepare_qwen_vl_mrope_config( _MAX_PIXELS_TOKEN_PROBE = 1 << 31 -def _get_mrope_position_delta_cache_size( - model_config: ModelConfig[PretrainedConfig]) -> int: - """Return real sequence-slot capacity plus one reserved dummy slot.""" - max_num_seq_slots = model_config.extra_attrs.get( - 'max_num_seq_slots', - model_config.max_num_tokens * model_config.mapping.pp_size) - return max_num_seq_slots + 1 - - class Qwen2VLInputProcessorBase(BaseMultimodalInputProcessor, BaseMultimodalDummyInputsBuilder): @@ -1714,8 +1705,8 @@ def __init__( if not disable_fuse_rope: self.init_mrope_embedding(model_config) # Extra slot is reserved for CUDA graph / warmup dummy requests. - max_mrope_delta_slots = _get_mrope_position_delta_cache_size( - model_config) + max_mrope_delta_slots = ( + model_config.max_num_tokens * model_config.mapping.pp_size + 1) self.register_buffer('mrope_position_deltas_cache', torch.zeros(max_mrope_delta_slots, dtype=torch.int32, diff --git a/tensorrt_llm/_torch/models/modeling_qwen3vl.py b/tensorrt_llm/_torch/models/modeling_qwen3vl.py index e2015e8e8df6..95da74e5d3bd 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3vl.py @@ -56,7 +56,6 @@ from .modeling_qwen2vl import ( Qwen2_5_VLVisionAttention, Qwen2VLInputProcessorBase, - _get_mrope_position_delta_cache_size, _prepare_qwen_vl_mrope_config, _prepare_qwen_vl_vision_attn_metadata, ) @@ -1207,7 +1206,7 @@ def __init__( if not disable_fuse_rope: self.init_mrope_embedding(model_config) # Extra slot is reserved for CUDA graph / warmup dummy requests. - max_mrope_delta_slots = _get_mrope_position_delta_cache_size(model_config) + max_mrope_delta_slots = model_config.max_num_tokens * model_config.mapping.pp_size + 1 self.register_buffer( "mrope_position_deltas_cache", torch.zeros( diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index dc7b86e480c9..c31ef6430c0c 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2046,24 +2046,43 @@ def create_kv_cache_compression_manager( return None -def compute_max_num_sequences(mapping: Mapping, max_batch_size: int, - disable_overlap_scheduler: bool) -> int: +def compute_max_num_sequences(mapping: Mapping, + max_batch_size: int, + disable_overlap_scheduler: bool, + enable_overlap_headroom: bool = False) -> int: """Size the sequence-slot pool (and the sampler state it indexes). - The slot pool needs headroom for two iterations under overlap (prev-iter - finished requests still hold slots when next-iter prepare_resources - runs). Without it, the V2 scheduler can backfill the budget seat vacated - by a GENERATION_TO_COMPLETE request whose slot is not freed yet and - exhaust the pool ("No free slots"). With pipeline parallelism the - micro-batch count is pp_size, which covers overlap already. + ``enable_overlap_headroom`` is intentionally opt-in. DeepSeek-V4 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. + Other models retain their established sizing until that behavior is + validated independently. Pipeline parallelism already sizes the pool by + ``pp_size``. """ if mapping.has_pp(): num_micro_batches = mapping.pp_size else: - num_micro_batches = 1 if disable_overlap_scheduler else 2 + num_micro_batches = (2 if enable_overlap_headroom + and not disable_overlap_scheduler else 1) return max_batch_size * num_micro_batches +def should_enable_dsv4_adp_dummy_fixes(model_type: Optional[str], + mapping: Mapping) -> bool: + """Gate DSv4 ADP dummy behavior while PP remains follow-up scope.""" + return model_type == "deepseek_v4" and not mapping.has_pp() + + +def should_enable_dsv4_overlap_headroom( + model_type: Optional[str], spec_config: Optional[SpeculativeConfig], + mapping: Mapping, disable_overlap_scheduler: bool) -> bool: + """Gate extra sequence slots to the validated DSv4 MTP overlap path.""" + return (should_enable_dsv4_adp_dummy_fixes(model_type, mapping) + and spec_config is not None + and spec_config.spec_dec_mode.is_mtp_eagle_one_model() + and not disable_overlap_scheduler) + + def create_py_executor_instance( *, dist, @@ -2090,6 +2109,7 @@ def create_py_executor_instance( virtual_memory_pools: Optional[dict] = None, execution_stream: Optional[torch.cuda.Stream] = None, dwdp_manager: Optional[DwdpManager] = None, + max_num_sequences: Optional[int] = None, ) -> PyExecutor: set_low_latency_dispatch( getattr(llm_args, 'enable_low_latency_host_dispatch', False)) @@ -2098,8 +2118,9 @@ def create_py_executor_instance( spec_config = model_engine.spec_config - max_num_sequences = compute_max_num_sequences( - mapping, max_batch_size, llm_args.disable_overlap_scheduler) + if max_num_sequences is None: + max_num_sequences = compute_max_num_sequences( + mapping, max_batch_size, llm_args.disable_overlap_scheduler) 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}" @@ -2454,11 +2475,13 @@ def create_torch_sampler_args( disable_overlap_scheduler: bool, enable_async_worker: bool, enable_speculative_beam_history_d2h: bool, + max_num_sequences: Optional[int] = None, ): # 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 = compute_max_num_sequences(mapping, max_batch_size, - disable_overlap_scheduler) + if max_num_sequences is None: + max_num_sequences = compute_max_num_sequences( + mapping, max_batch_size, disable_overlap_scheduler) 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 @@ -2488,6 +2511,7 @@ def instantiate_sampler( speculative_config: SpeculativeConfig, decoding_config: trtllm.DecodingConfig, kv_cache_config: KvCacheConfig, + max_num_sequences: Optional[int] = None, ): enable_async_worker = (confidential_compute_enabled() or llm_args.sampler_force_async_worker) @@ -2502,6 +2526,7 @@ def instantiate_sampler( enable_async_worker=enable_async_worker, enable_speculative_beam_history_d2h=llm_args. enable_speculative_beam_history_d2h, + max_num_sequences=max_num_sequences, ) decoding_mode = get_decoding_mode(decoding_config=decoding_config, max_beam_width=max_beam_width) @@ -2530,7 +2555,8 @@ def instantiate_sampler( max_beam_width=max_beam_width, decoding_config=decoding_config, kv_cache_config=kv_cache_config, - enable_async_worker=enable_async_worker) + enable_async_worker=enable_async_worker, + max_num_sequences=max_num_sequences) if not engine.model.model_config.is_generation: # NOTE: choose sampler based on model type return EarlyStopSampler() diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 350f135e6a45..4d8dc3b717c1 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -1120,22 +1120,8 @@ def free_resources(self, request: LlmRequest, pin_on_release: bool = False): KVCacheManager.free_resources(self, request, pin_on_release) def add_dummy_requests(self, request_ids: List[int], **kwargs): - requests = KVCacheManager.add_dummy_requests(self, request_ids, - **kwargs) - if not requests: - return requests - try: - MambaCacheManager.add_dummy_requests(self, request_ids) - except RuntimeError: - # KV allocation completed first; restore it if the independent - # Mamba pool cannot reserve the matching dummy slot. - draft_kv_cache_manager = kwargs.get('draft_kv_cache_manager') - for request in requests: - if draft_kv_cache_manager is not None: - draft_kv_cache_manager.free_resources(request) - KVCacheManager.free_resources(self, request) - raise - return requests + MambaCacheManager.add_dummy_requests(self, request_ids) + return KVCacheManager.add_dummy_requests(self, request_ids, **kwargs) def shutdown(self): MambaCacheManager.shutdown(self) @@ -1845,29 +1831,23 @@ def get_num_available_tokens(self, result = super().get_num_available_tokens(token_num_upper_bound, max_num_draft_tokens, **kwargs) - # Also bound by the recurrent-state pool capacity. With block reuse, - # each request needs roughly ceil(N / states_snapshot_interval) blocks - # (+1 for the final live state). Without reuse, it still needs one live - # recurrent-state block, so an exhausted pool must fail preflight. - if self.linear_attention_metadata is None: - # A PP rank with no local Mamba layers has no recurrent-state pool. - return max(result, 0) - - interval = self.linear_attention_metadata.states_snapshot_interval - stats = self.impl.get_kv_cache_stats() - rs_free = stats.num_free_blocks_per_window_size.get( - LinearCacheType.RECURRENT_STATES.value, 0) + # Also bound by the recurrent-state pool capacity: each request needs + # roughly ceil(N / states_snapshot_interval) recurrent-state blocks + # (+1 for the corner case where N is a multiple of tokens_per_block). + # When block reuse is disabled, only one snapshot is needed per + # request, so no additional capping is required here. + interval = (self.linear_attention_metadata.states_snapshot_interval + if self.linear_attention_metadata is not None else 0) if interval and interval > 0: + stats = self.impl.get_kv_cache_stats() + rs_free = stats.num_free_blocks_per_window_size.get( + LinearCacheType.RECURRENT_STATES.value, 0) # Reserve 1 block for the always-allocated last block (corner case # / final live state) so we don't promise more tokens than the # pool can actually back at allocation time. usable_rs_blocks = max(0, rs_free - 1) - rs_token_cap = max( - 0, usable_rs_blocks * interval - self.num_extra_kv_tokens - - max_num_draft_tokens) + rs_token_cap = usable_rs_blocks * interval result = min(result, rs_token_cap) - elif rs_free == 0: - result = 0 return max(result, 0) def get_ssm_states(self, layer_idx: int) -> torch.Tensor: @@ -1988,7 +1968,8 @@ def _setup_state_indices(self, requests=None) -> None: self.cuda_state_indices.copy_(self._host_state_indices, non_blocking=True) - self._refresh_dummy_request_mask([req.is_dummy for req in requests]) + self._refresh_dummy_request_mask( + [req.is_dummy for req in self.requests]) # Build request_id → pool block offset mapping so that # get_state_indices can return indices in arbitrary request order. diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index ecbcf30315e0..513625e87cc5 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -299,10 +299,11 @@ def __init__( self.mapping = mapping if mapping.has_pp(): init_pp_comm(mapping) - # The overlap scheduler can hold two iterations' requests at once. - # Model-side buffers indexed by py_seq_slot must span that full pool, - # not the token budget or scheduler admission batch size. - from ._util import compute_max_num_sequences + # Start with the established pool size. Once the model is loaded we + # selectively enable headroom for the non-PP DeepSeek-V4 overlap path. + from ._util import (compute_max_num_sequences, + should_enable_dsv4_adp_dummy_fixes, + should_enable_dsv4_overlap_headroom) self.max_num_seq_slots = compute_max_num_sequences( mapping, self.batch_size, llm_args.disable_overlap_scheduler) self.dist = dist @@ -372,7 +373,6 @@ def __init__( sparse_attention_config=self.sparse_attention_config, max_num_tokens=self.max_num_tokens, max_seq_len=self.max_seq_len, - max_num_seq_slots=self.max_num_seq_slots, lora_config=lora_config, model_weights_memory_tag=model_weights_memory_tag, model_weights_restore_mode=model_weights_restore_mode, @@ -383,7 +383,23 @@ def __init__( setattr(self, "moe_load_balancer", moe_load_balancer) else: self.model = model - self._validate_mrope_position_delta_cache_capacity() + pretrained_config = self.model.model_config.pretrained_config + model_type = getattr(pretrained_config, "model_type", None) + # Keep the scheduler/dummy fix model-scoped, while the larger slot pool + # is restricted to the validated MTP overlap configuration. PP remains + # on its established path for follow-up changes. + self._enable_dsv4_adp_dummy_fixes = (should_enable_dsv4_adp_dummy_fixes( + model_type, mapping)) + self._enable_dsv4_overlap_headroom = ( + should_enable_dsv4_overlap_headroom( + model_type, spec_config, mapping, + llm_args.disable_overlap_scheduler)) + self.max_num_seq_slots = compute_max_num_sequences( + mapping, + self.batch_size, + llm_args.disable_overlap_scheduler, + enable_overlap_headroom=self._enable_dsv4_overlap_headroom, + ) if drafting_loop_wrapper is not None: self.model = drafting_loop_wrapper(self.model) self.model_is_wrapped = True @@ -888,34 +904,6 @@ def no_cuda_graph(self): finally: self.cuda_graph_runner.enabled = _run_cuda_graphs - def _validate_mrope_position_delta_cache_capacity(self) -> None: - """Validate slot-indexed MRoPE state on preconstructed models. - - Models created by ModelLoader receive ``max_num_seq_slots`` before - construction. A caller-supplied model bypasses that path, so fail - early instead of letting the overlap scheduler index past an old-sized - cache at runtime. - """ - mrope_position_deltas_cache = getattr(self.model, - "mrope_position_deltas_cache", - None) - if mrope_position_deltas_cache is None: - mrope_position_deltas_cache = getattr( - getattr(self.model, "draft_model", None), - "mrope_position_deltas_cache", None) - if mrope_position_deltas_cache is None: - return - - required_size = self.max_num_seq_slots + 1 - actual_size = mrope_position_deltas_cache.shape[0] - if actual_size < required_size: - raise ValueError( - "The supplied model's MRoPE position-delta cache has " - f"{actual_size} slots, but this executor requires at least " - f"{required_size} ({self.max_num_seq_slots} runtime sequence " - "slots plus one reserved dummy slot). Rebuild the model with " - "the executor's sequence-slot capacity.") - def _pad_batch_seed_mrope_delta_cache( self, padded_requests: ScheduledRequests) -> None: if not self.use_mrope or padded_requests.num_generation_requests == 0: @@ -2194,9 +2182,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 - # Slot-indexed metadata buffers must span the SeqSlotManager pool, - # which exceeds batch_size under the overlap scheduler. - num_seq_slots = self.max_num_seq_slots + # Only the scoped DeepSeek-V4 overlap path opts into larger metadata + # buffers. Passing None preserves the established max_num_requests + # fallback for every other model, including MTP-Eagle with PP. + num_seq_slots = (self.max_num_seq_slots + if self._enable_dsv4_overlap_headroom else None) if no_cache: return get_spec_metadata( spec_config, @@ -3295,9 +3285,7 @@ def _prepare_tp_inputs( mrope_delta_read_seq_slots = [] # Extra model-side cache slot reserved for CUDA graph / warmup dummy # requests, whose outputs are discarded. - mrope_dummy_seq_slot = getattr( - self, "max_num_seq_slots", - self.max_num_tokens * self.mapping.pp_size) + mrope_dummy_seq_slot = self.max_num_tokens * self.mapping.pp_size num_accepted_draft_tokens = [] # per request is_enc_dec = self._is_encoder_decoder_model() cross_encoder_hidden_states: List[torch.Tensor] = [] diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index dfde27441422..7596c2a68291 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -310,8 +310,7 @@ def __init__(self, max_seq_len: Optional[int], lora_config: Optional[LoraConfig] = None, model_weights_memory_tag: Optional[ExecutorMemoryType] = None, - model_weights_restore_mode: Optional[RestoreMode] = None, - max_num_seq_slots: Optional[int] = None): + model_weights_restore_mode: Optional[RestoreMode] = None): """ Initializes the ModelLoader. @@ -327,9 +326,6 @@ def __init__(self, they can be released/materialized independently of buffers. model_weights_restore_mode: RestoreMode for the model weights virtual-memory scope. - max_num_seq_slots: Capacity of model buffers indexed by sequence - slot. This can exceed the scheduler admission batch size when - overlap scheduling is enabled. """ self.llm_args = llm_args self.mapping = mapping @@ -337,7 +333,6 @@ def __init__(self, self.sparse_attention_config = sparse_attention_config self.max_num_tokens = max_num_tokens self.max_seq_len = max_seq_len - self.max_num_seq_slots = max_num_seq_slots self.lora_config = lora_config self.model_weights_memory_tag = model_weights_memory_tag self.model_weights_restore_mode = model_weights_restore_mode @@ -345,11 +340,6 @@ def __init__(self, self._weight_pool_proxy = None self._gms_backend = None - def _set_runtime_model_config_attrs(self, config: ModelConfig) -> None: - """Attach executor-only allocation sizes before model construction.""" - if self.max_num_seq_slots is not None: - config.extra_attrs['max_num_seq_slots'] = self.max_num_seq_slots - @staticmethod def load_config_and_apply_defaults( checkpoint_dir: str, llm_args: TorchLlmArgs, @@ -1218,7 +1208,6 @@ def _load_and_validate_config( load_config_kwargs['model_kwargs'] = self.llm_args.model_kwargs config = checkpoint_loader.load_config(**load_config_kwargs) - self._set_runtime_model_config_attrs(config) # Store nvfp4 config in extra_attrs for Linear layer access config.extra_attrs[ diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 9ed208071fcb..8cdf1183f0d1 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -89,8 +89,6 @@ SerializableSchedulerOutput, WaitingQueue, create_waiting_queue) from .scheduler.adp_router import ADPRouter -from .scheduler.scheduler import \ - is_decoder_context_request_waiting_for_encoder_output if TYPE_CHECKING: from ray.actor import ActorHandle @@ -567,6 +565,8 @@ def __init__( self.resource_manager = resource_manager self.scheduler = scheduler self.model_engine = model_engine + self._enable_dsv4_adp_dummy_fixes = getattr( + model_engine, "_enable_dsv4_adp_dummy_fixes", False) self.enable_attention_dp = model_engine.enable_attention_dp self.dist = dist self.sampler = sampler @@ -2531,7 +2531,6 @@ def _executor_loop_pp(self): can_queue, _ = self._can_queue(scheduled_batch) if not can_queue: self._revert_gen_alloc(scheduled_batch) - self._finalize_adp_dummy_allocation(can_queue) if not can_queue: logger.debug( f"microbatch {microbatch_id} cannot be queued, skipping" @@ -3221,6 +3220,9 @@ def _finalize_adp_dummy_allocation(self, can_queue: bool) -> None: must release theirs before retrying or the fixed dummy request ID leaks cache resources on every skipped iteration. """ + if not self._enable_dsv4_adp_dummy_fixes: + return + dummy_request = self._pending_adp_dummy_request self._pending_adp_dummy_request = None if dummy_request is None or can_queue: @@ -5336,36 +5338,31 @@ def _check_disagg_ctx_schedulable_status(self, def _count_schedulable_active_requests(self) -> int: """Count active requests that are ready for scheduling. - Both scheduler implementations exclude requests already in the PP - inflight set and requests waiting for encoder output. They also apply - the schedulability window [CONTEXT_INIT, GENERATION_TO_COMPLETE), or - [ENCODER_INIT, GENERATION_TO_COMPLETE) for encoder-decoder models. - States outside that window — e.g. waiting for KV transfer (below) or - awaiting completion bookkeeping (at or above) — cannot produce forward - work. Counting any of these categories here would suppress the pad - dummy while the rank still schedules batch=0, turning can_queue False - fleet-wide and stalling every rank. + The non-PP DeepSeek-V4 disaggregated ADP path mirrors the decoder + scheduler's state window [CONTEXT_INIT, GENERATION_TO_COMPLETE). This + covers generation-first context requests below the lower bound and + terminal requests at the upper bound. Other configurations retain the + established ADP behavior; PP eligibility remains follow-up scope. Returns: The number of active requests eligible for scheduling. """ - # Both bounds must match the scheduler exactly. The decoder-only lower - # covers e.g. DISAGG_GENERATION_INIT / TRANS_IN_PROGRESS and - # DISAGG_CONTEXT_WAIT_SCHEDULER (gen-first mode on the context - # server); encoder-decoder widens it to ENCODER_INIT. The upper bound - # covers GENERATION_TO_COMPLETE and all terminal/post-completion states. - cross_kv_cache_manager = self.resource_manager.get_resource_manager( - ResourceManagerType.CROSS_KV_CACHE_MANAGER) - schedule_from_value = (LlmRequestState.ENCODER_INIT.value - if cross_kv_cache_manager is not None else - LlmRequestState.CONTEXT_INIT.value) + if (not self._enable_dsv4_adp_dummy_fixes + or self.kv_cache_transceiver is None): + if self.kv_cache_transceiver is None: + return len(self.active_requests) + + return sum( + 1 for req in self.active_requests + if not (req.is_disagg_generation_init_state + or req.is_disagg_generation_transmission_in_progress)) + + schedule_from_value = LlmRequestState.CONTEXT_INIT.value to_complete_value = LlmRequestState.GENERATION_TO_COMPLETE.value return sum( 1 for req in self.active_requests - if req.request_id not in self.inflight_req_ids - and not is_decoder_context_request_waiting_for_encoder_output(req) - and schedule_from_value <= req.state_value < to_complete_value) + if schedule_from_value <= req.state_value < to_complete_value) def _has_adp_dummy_kv_capacity(self, token_nums: Optional[List[int]]) -> bool: @@ -5374,7 +5371,7 @@ def _has_adp_dummy_kv_capacity(self, V1's dummy allocator historically checked only that one block was free. A context-side ADP dummy can be ``max_num_tokens`` long, while a generation dummy also reserves draft tokens, so that check can still - let the C++ batch allocation throw before PP/CP peers reach their vote. + let the C++ batch allocation fail after a peer rank has succeeded. """ token_num = (token_nums[0] if token_nums is not None else 1 + self.max_total_draft_tokens) @@ -5438,33 +5435,6 @@ def _update_adp_dummy_role(self, candidates: List[LlmRequest]) -> None: elif has_gen: self._adp_dummy_is_gen = True - def _can_commit_adp_dummy_candidate( - self, needs_dummy: bool, - dummy_candidate: Optional[LlmRequest]) -> bool: - """Agree on dummy creation across one PP/CP scheduling lane. - - The first PP stage propagates request IDs to later stages, so a dummy - may only enter ``active_requests`` when every stage (and every CP shard) - needs it and allocated its local cache resources. TP/ADP lanes remain - independent here; their existing ``_can_queue`` vote decides whether - the fleet can run and finalizes or rolls back the committed candidate. - """ - local_status = (needs_dummy, not needs_dummy - or dummy_candidate is not None) - lane_statuses = [local_status] - - if self._dist_size(self.dist, "cp_size") > 1: - lane_statuses = self.dist.cp_allgather(local_status) - if self._dist_size(self.dist, "pp_size") > 1: - lane_statuses = [ - status - for stage_statuses in self.dist.pp_allgather(lane_statuses) - for status in stage_statuses - ] - - return (all(status[0] for status in lane_statuses) - and all(status[1] for status in lane_statuses)) - @nvtx_range("_pad_attention_dp_dummy_request") def _pad_attention_dp_dummy_request(self): """ @@ -5473,75 +5443,74 @@ def _pad_attention_dp_dummy_request(self): if not self.enable_attention_dp: return - assert self._pending_adp_dummy_request is None assert self.expected_num_active_requests >= len(self.active_requests) num_active_request = self._count_schedulable_active_requests() - skip_dummy = self._should_skip_dummy_for_benchmark_disagg( - num_active_request) + if self._should_skip_dummy_for_benchmark_disagg(num_active_request): + return - # Allocate tentatively. PP/CP peers must agree before the request ID is - # exposed to scheduling; otherwise schedule propagation can reference a - # dummy that a later stage failed to create. - # Every stage/shard must enter the vote, including ranks that expect no - # active requests or suppress padding during benchmark fill. Returning - # early on those ranks would leave a peer blocked in pp/cp_allgather. - needs_dummy = (not skip_dummy and self.expected_num_active_requests > 0 + needs_dummy = (self.expected_num_active_requests > 0 and num_active_request == 0) - dummy_request = None + if not needs_dummy: + return + dummy_request_ids = [ATTENTION_DP_DUMMY_REQUEST_ID] + token_nums = None + if (not self._adp_dummy_is_gen and self.kv_cache_transceiver is not None + and self.max_num_tokens is not None): + token_nums = [self.max_num_tokens] + + if (not self._enable_dsv4_adp_dummy_fixes + or self.kv_cache_transceiver is None): + llm_request = self.kv_cache_manager.add_dummy_requests( + request_ids=dummy_request_ids, + token_nums=token_nums, + is_gen=self._adp_dummy_is_gen, + prepare_resource=True, + max_num_draft_tokens=self.max_total_draft_tokens, + )[0] + llm_request.is_attention_dp_dummy = True + spec_resource_manager = self.resource_manager.get_resource_manager( + ResourceManagerType.SPEC_RESOURCE_MANAGER) + if spec_resource_manager is not None: + spec_resource_manager.add_dummy_requests(dummy_request_ids) + self.active_requests.append(llm_request) + return + + assert self._pending_adp_dummy_request is None has_live_adp_dummy = any( request.py_request_id == ATTENTION_DP_DUMMY_REQUEST_ID for request in self.active_requests) - spec_resource_manager = None - if needs_dummy and not has_live_adp_dummy: - token_nums = None - if (not self._adp_dummy_is_gen - and self.kv_cache_transceiver is not None - and self.max_num_tokens is not None): - token_nums = [self.max_num_tokens] + if has_live_adp_dummy or not self._has_adp_dummy_kv_capacity( + token_nums): + return + + try: + dummy_requests = self.kv_cache_manager.add_dummy_requests( + request_ids=dummy_request_ids, + token_nums=token_nums, + is_gen=self._adp_dummy_is_gen, + prepare_resource=True, + max_num_draft_tokens=self.max_total_draft_tokens, + ) + except OutOfPagesError: dummy_requests = None - if self._has_adp_dummy_kv_capacity(token_nums): - try: - dummy_requests = self.kv_cache_manager.add_dummy_requests( - request_ids=dummy_request_ids, - token_nums=token_nums, - is_gen=self._adp_dummy_is_gen, - prepare_resource=True, - max_num_draft_tokens=self.max_total_draft_tokens, - ) - except OutOfPagesError: - # V2 can lose a capacity race between its estimate and - # resize. Treat that known resource condition exactly like - # its normal ``None`` result and join the lane vote. - dummy_requests = None - if dummy_requests: - dummy_request = dummy_requests[0] - spec_resource_manager = self.resource_manager.get_resource_manager( - ResourceManagerType.SPEC_RESOURCE_MANAGER) - if spec_resource_manager is not None: - try: - spec_resource_manager.add_dummy_requests( - dummy_request_ids) - except NoFreeSlotsError: - # SlotManager checks capacity before inserting the one - # ADP ID, so this specific built-in failure has no - # partial speculative state to unwind. Unexpected - # ValueErrors and all RuntimeErrors remain fatal. - self.kv_cache_manager.free_resources(dummy_request) - dummy_request = None - - if not self._can_commit_adp_dummy_candidate(needs_dummy, dummy_request): - if dummy_request is not None: - if spec_resource_manager is not None: - spec_resource_manager.free_resources(dummy_request) - self.kv_cache_manager.free_resources(dummy_request) - if needs_dummy: - logger.warning( - "Cannot commit ADP pad dummy on every PP/CP lane rank; " - "rolled back local allocation and will retry.") + if not dummy_requests: + logger.warning( + "Cannot allocate DeepSeek-V4 ADP pad dummy; rank schedules " + "an empty batch and the fleet will retry.") return + dummy_request = dummy_requests[0] + spec_resource_manager = self.resource_manager.get_resource_manager( + ResourceManagerType.SPEC_RESOURCE_MANAGER) + if spec_resource_manager is not None: + try: + spec_resource_manager.add_dummy_requests(dummy_request_ids) + except NoFreeSlotsError: + self.kv_cache_manager.free_resources(dummy_request) + return + assert dummy_request is not None dummy_request.is_attention_dp_dummy = True self.active_requests.append(dummy_request) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index dc7eef5548b5..a4c2e48ddc53 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import copy import gc import importlib @@ -35,8 +38,8 @@ get_spec_resource_manager) from ..virtual_memory import scope as virtual_memory_scope from ._util import (KvCacheCreator, _adjust_torch_mem_fraction, - compute_max_num_sequences, create_py_executor_instance, - instantiate_sampler, is_mla, validate_feature_combination) + create_py_executor_instance, instantiate_sampler, is_mla, + validate_feature_combination) from .config_utils import is_hybrid_linear from .connectors.kv_cache_connector import KvCacheConnectorManager from .dwdp import DwdpManager @@ -670,6 +673,8 @@ def drafting_loop_wrapper(model): cache_transceiver_config.max_tokens_in_buffer = net_max_seq_len 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)) if is_hybrid_linear(config) and kv_cache_config.enable_block_reuse and ( cache_transceiver_config is not None and cache_transceiver_config.backend is not None @@ -756,21 +761,16 @@ def drafting_loop_wrapper(model): 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_dsv4_overlap_headroom", False) else + max_batch_size) kwargs = { - "guided_decoding_config": - guided_decoding_config, - # Guided-decoder state is indexed by sequence slot; size - # it to the SeqSlotManager pool (2x max_batch_size under - # the overlap scheduler), matching the sampler's - # new_tokens width. - "max_num_sequences": - compute_max_num_sequences( - mapping, max_batch_size, - llm_args.disable_overlap_scheduler), - "vocab_size_padded": - model_engine.model.vocab_size_padded, - "rank": - mapping.rank, + "guided_decoding_config": guided_decoding_config, + # The scoped DeepSeek-V4 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, } if spec_config is not None: kwargs[ @@ -806,6 +806,7 @@ def drafting_loop_wrapper(model): speculative_config=spec_config, decoding_config=decoding_config, kv_cache_config=kv_cache_config, + max_num_sequences=max_num_seq_slots, ) logger.info(f"Using Sampler: {type(sampler).__name__}") @@ -994,6 +995,7 @@ def drafting_loop_wrapper(model): cache_transceiver_config=cache_transceiver_config, virtual_memory_pools=vm_pools if not estimating_kv_cache else None, execution_stream=execution_stream, + max_num_sequences=max_num_seq_slots, ) # Originally, peft_cache_config might be mutated inside @@ -1069,6 +1071,7 @@ def drafting_loop_wrapper(model): virtual_memory_pools=vm_pools, execution_stream=execution_stream, dwdp_manager=dwdp_manager, + max_num_sequences=max_num_seq_slots, ) _adjust_torch_mem_fraction() diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 1911f944c86f..6c0f1140f971 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -4848,6 +4848,7 @@ def __init__( decoding_config: Optional[DecodingConfig] = None, kv_cache_config: Optional[KvCacheConfig] = None, enable_async_worker: bool = False, + max_num_sequences: Optional[int] = None, ): assert model.config is not None vocab_size = model.config.vocab_size @@ -4871,11 +4872,12 @@ def __init__( self.num_micro_batches = ( mapping.pp_size if mapping.pp_size > 1 else (2 if self.is_trt_overlap else 1) ) - # Decoder state is indexed by sequence slot; size it to the - # SeqSlotManager pool (max_batch_size * num_micro_batches, matching - # compute_max_num_sequences), not pp_size alone, so py_seq_slot - # values under the overlap scheduler stay in range. - self.max_num_sequences = max_batch_size * self.num_micro_batches + # Decoder state is indexed by sequence slot and must match the + # executor's SeqSlotManager. The fallback preserves the established + # sizing for direct callers outside the PyExecutor creator. + self.max_num_sequences = ( + max_num_sequences if max_num_sequences is not None else mapping.pp_size * max_batch_size + ) self.micro_batch_idx = 0 if mpi_disabled(): diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index ccc33b2f8087..3978040724b1 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -40,8 +40,7 @@ def __init__(self, max_num_requests: int, max_seq_len: int, max_num_tokens: int, - sa_manager=None, - max_num_seq_slots: int | None = None): + sa_manager=None): self.dtype = dtype self.max_draft_len = config.max_draft_len self.hidden_size = hidden_size @@ -102,7 +101,6 @@ 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, - max_num_seq_slots=max_num_seq_slots, ) def prepare_resources(self, scheduled_batch: ScheduledRequests): @@ -158,10 +156,7 @@ 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, - max_num_seq_slots: int | None = None): + def __init__(self, config: "EagleDecodingConfig", max_num_requests: int): self.max_num_requests = max_num_requests self.batch_indices_cuda = torch.empty( [max_num_requests], @@ -175,7 +170,6 @@ def __init__(self, max_total_draft_tokens=config.tokens_per_gen_step - 1, eagle_choices=config.eagle_choices, dynamic_tree_max_topK=config.dynamic_tree_max_topK, - max_num_seq_slots=max_num_seq_slots, ) def free_resources(self, request: LlmRequest): diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 6c208e4db3fe..6244873efc43 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -534,8 +534,8 @@ class SpecMetadata: # Vocab size used for draft_probs buffer allocation. vocab_size: int = 0 # Size of the SeqSlotManager pool. py_seq_slot values range over - # [0, num_seq_slots); under the overlap scheduler this is - # 2 * max_batch_size, larger than max_num_requests (== max_batch_size). + # [0, num_seq_slots); DeepSeek-V4 overlap can use 2 * max_batch_size, + # larger than max_num_requests (== max_batch_size). # Slot-indexed buffers (draft_probs) must span this full range. # 0 falls back to max_num_requests. num_seq_slots: int = 0 diff --git a/tensorrt_llm/_torch/speculative/spec_tree_manager.py b/tensorrt_llm/_torch/speculative/spec_tree_manager.py index d8ef26498126..74dd622826f1 100644 --- a/tensorrt_llm/_torch/speculative/spec_tree_manager.py +++ b/tensorrt_llm/_torch/speculative/spec_tree_manager.py @@ -178,14 +178,10 @@ 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, - max_num_seq_slots: int | None = None): + dynamic_tree_max_topK: int): self.use_dynamic_tree = use_dynamic_tree self.max_total_draft_tokens = max_total_draft_tokens @@ -202,12 +198,6 @@ def __init__(self, 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 - # Slot storage is indexed by py_seq_slot, whose range is the - # SeqSlotManager pool size (2x max_batch_size under the overlap - # scheduler) -- larger than num_trees, which sizes the - # batch-position work buffers. - self.max_num_seq_slots = (max_num_seq_slots if max_num_seq_slots - is not None else self.num_trees) self.dynamic_tree_max_topK = dynamic_tree_max_topK self.cur_draft_layer_idx = 0 self.top_k_list = [] @@ -291,7 +281,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.max_num_seq_slots, + num_slots=self.num_trees, n_dt=num_draft_with_root, mask_width=mask_width, ) diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index f98b67187504..d7a0e8185da5 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + from bisect import bisect_left from dataclasses import dataclass from typing import TYPE_CHECKING, Dict, Optional @@ -49,8 +52,8 @@ def get_spec_metadata(spec_config, num_seq_slots=None): use_rejection_sampling = getattr(spec_config, "use_rejection_sampling", False) - # Slot-indexed buffers (draft_probs) must span the SeqSlotManager pool, - # which exceeds max_num_requests under the overlap scheduler. + # 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) @@ -131,7 +134,6 @@ def get_spec_metadata(spec_config, layers_to_capture=spec_config.eagle3_layers_to_capture, use_rejection_sampling=use_rejection_sampling, vocab_size=vocab_size, - num_seq_slots=num_seq_slots, spec_resource_manager=spec_resource_manager, use_dynamic_tree=_is_effective_dynamic_tree(spec_config), eagle_choices=spec_config.eagle_choices, @@ -209,9 +211,6 @@ def get_mtp_hidden_size(model_config) -> int: def get_spec_resource_manager(model_engine, draft_model_engine=None): - # Lazy import: _util imports model_engine, which imports this module. - from ..pyexecutor._util import compute_max_num_sequences - spec_config = model_engine.spec_config if spec_config is None: return None @@ -219,11 +218,6 @@ 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 - # Slot-indexed storage must span the SeqSlotManager pool, which exceeds - # max_num_requests under the overlap scheduler. - max_num_seq_slots = compute_max_num_sequences( - model_engine.mapping, model_engine.batch_size, - model_engine._disable_overlap_scheduler) spec_dec_mode = spec_config.spec_dec_mode if spec_dec_mode.is_mtp_eagle_one_model(): sa_manager = None @@ -261,8 +255,8 @@ 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, max_num_seq_slots=max_num_seq_slots) + return Eagle3OneModelDynamicTreeResourceManager(spec_config, + max_num_requests) if spec_dec_mode.is_eagle3_one_model(): sa_manager = None sa_cfg = getattr(spec_config, 'sa_config', None) @@ -287,7 +281,6 @@ def get_spec_resource_manager(model_engine, draft_model_engine=None): max_num_requests, max_seq_len, max_num_tokens, - max_num_seq_slots=max_num_seq_slots, ) if spec_dec_mode.is_save_hidden_states(): return SaveHiddenStatesResourceManager( diff --git a/tests/unittest/_torch/executor/test_benchmark_disagg.py b/tests/unittest/_torch/executor/test_benchmark_disagg.py index 1a5952a43cd0..b9e30f296476 100644 --- a/tests/unittest/_torch/executor/test_benchmark_disagg.py +++ b/tests/unittest/_torch/executor/test_benchmark_disagg.py @@ -44,11 +44,7 @@ def _make_active_request( ) -> Mock: """Create an active request stub with disagg state flags.""" req = Mock() - req.request_id = 1 - req.py_request_id = 1 req.state_value = LlmRequestState.GENERATION_IN_PROGRESS.value - req.is_context_init_state = False - req.py_encoder_output_ready_event = None req.is_disagg_generation_init_state = in_init req.is_disagg_generation_transmission_in_progress = in_transfer req.state = ( @@ -514,16 +510,12 @@ def __init__( self.max_total_draft_tokens = 0 self._adp_dummy_is_gen = True self._pending_adp_dummy_request = None + self._enable_dsv4_adp_dummy_fixes = True self.max_num_tokens = None - self.inflight_req_ids = set() self.dist = Mock() self.dist.tp_size = tp_size - self.dist.pp_size = 1 - self.dist.cp_size = 1 self.dist.tp_allgather.side_effect = lambda value: [value] - self.dist.pp_allgather.side_effect = lambda value: [value] - self.dist.cp_allgather.side_effect = lambda value: [value] self.kv_cache_manager = Mock() self.kv_cache_manager.mapping.has_cp_helix.return_value = False @@ -541,8 +533,6 @@ def __init__( _count_schedulable_active_requests = PyExecutor._count_schedulable_active_requests _has_adp_dummy_kv_capacity = PyExecutor._has_adp_dummy_kv_capacity _should_skip_dummy_for_benchmark_disagg = PyExecutor._should_skip_dummy_for_benchmark_disagg - _can_commit_adp_dummy_candidate = PyExecutor._can_commit_adp_dummy_candidate - _dist_size = staticmethod(PyExecutor._dist_size) class TestPadAttentionDpDummyBenchmarkDisagg: diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index 74f668a60038..30eaa2442424 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -5,7 +5,6 @@ import os from types import SimpleNamespace -from unittest.mock import MagicMock import pytest import torch @@ -15,12 +14,10 @@ from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( MIN_REPLAY_HISTORY_SIZE, CppMambaHybridCacheManager, - MambaCacheManager, - MixedMambaHybridCacheManager, PythonMambaCacheManager, _get_mamba_hybrid_pool_size, ) -from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp, KVCacheManager +from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.bindings.internal.batch_manager import LinearCacheType from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MTPDecodingConfig @@ -363,7 +360,6 @@ def _build_hybrid_with_mamba_layer( enable_block_reuse=False, mamba_state_cache_interval=256, is_estimating_kv_cache=False, - use_replay_state_update=False, ): """Construct a real CppMambaHybridCacheManager with one mamba layer + one full-attention layer so the parent KVCacheManager goes through the @@ -400,7 +396,6 @@ def _build_hybrid_with_mamba_layer( spec_config=spec_config, layer_mask=attn_mask, is_estimating_kv_cache=is_estimating_kv_cache, - use_replay_state_update=use_replay_state_update, ) @@ -438,114 +433,6 @@ def test_cpp_hybrid_recurrent_pool_reserves_draft_len_sentinel_slots(): ) -@skip_no_cuda -def test_cpp_hybrid_adp_dummy_mask_excludes_stale_requests(): - """The replay mask must describe the batch passed to state setup. - - ADP dummy allocation happens before the previous overlap batch has been - released, so ``self.requests`` can contain a full stale batch in addition - to the newly allocated dummy. The dummy-only setup must not size its mask - from that combined lifecycle list. - """ - max_batch_size = 2 - mgr = _build_hybrid_with_mamba_layer( - spec_config=MTPDecodingConfig(max_draft_len=1), - max_batch_size=max_batch_size, - use_replay_state_update=True, - ) - stale_requests = [SimpleNamespace(is_dummy=False) for _ in range(max_batch_size)] - mgr.requests = stale_requests.copy() - - dummy_requests = mgr.add_dummy_requests([ATTENTION_DP_DUMMY_REQUEST_ID]) - - assert dummy_requests - assert mgr.requests == stale_requests + dummy_requests - assert mgr._dummy_request_mask_host.tolist() == [True, False] - - -def test_mixed_hybrid_does_not_allocate_mamba_when_dummy_kv_fails(monkeypatch): - mgr = object.__new__(MixedMambaHybridCacheManager) - kv_add_dummy = MagicMock(return_value=None) - mamba_add_dummy = MagicMock() - monkeypatch.setattr(KVCacheManager, "add_dummy_requests", kv_add_dummy) - monkeypatch.setattr(MambaCacheManager, "add_dummy_requests", mamba_add_dummy) - - requests = mgr.add_dummy_requests([ATTENTION_DP_DUMMY_REQUEST_ID]) - - assert requests is None - mamba_add_dummy.assert_not_called() - - -def test_mixed_hybrid_rolls_back_kv_when_mamba_dummy_fails(monkeypatch): - mgr = object.__new__(MixedMambaHybridCacheManager) - dummy_request = SimpleNamespace(py_request_id=ATTENTION_DP_DUMMY_REQUEST_ID) - draft_kv_cache_manager = MagicMock() - kv_add_dummy = MagicMock(return_value=[dummy_request]) - kv_free = MagicMock() - mamba_add_dummy = MagicMock(side_effect=RuntimeError("no mamba slot")) - monkeypatch.setattr(KVCacheManager, "add_dummy_requests", kv_add_dummy) - monkeypatch.setattr(KVCacheManager, "free_resources", kv_free) - monkeypatch.setattr(MambaCacheManager, "add_dummy_requests", mamba_add_dummy) - - with pytest.raises(RuntimeError, match="no mamba slot"): - mgr.add_dummy_requests( - [ATTENTION_DP_DUMMY_REQUEST_ID], draft_kv_cache_manager=draft_kv_cache_manager - ) - - draft_kv_cache_manager.free_resources.assert_called_once_with(dummy_request) - kv_free.assert_called_once_with(mgr, dummy_request) - - -def test_cpp_hybrid_available_tokens_reserves_extra_and_draft_tokens(monkeypatch): - mgr = object.__new__(CppMambaHybridCacheManager) - mgr.num_extra_kv_tokens = 1 - mgr.linear_attention_metadata = SimpleNamespace(states_snapshot_interval=16) - mgr.impl = MagicMock() - mgr.impl.get_kv_cache_stats.return_value = SimpleNamespace( - num_free_blocks_per_window_size={LinearCacheType.RECURRENT_STATES.value: 3} - ) - monkeypatch.setattr(KVCacheManager, "get_num_available_tokens", MagicMock(return_value=100)) - - available_tokens = mgr.get_num_available_tokens( - token_num_upper_bound=100, max_num_draft_tokens=2 - ) - - # One recurrent block is reserved for the final live state. The remaining - # two cover 32 tokens, of which one extra KV token and two draft tokens are - # consumed immediately after the initial dummy allocation. - assert available_tokens == 29 - - -@pytest.mark.parametrize("rs_free, expected", [(0, 0), (1, 100)]) -def test_cpp_hybrid_available_tokens_requires_live_state_without_reuse( - monkeypatch, rs_free, expected -): - mgr = object.__new__(CppMambaHybridCacheManager) - mgr.num_extra_kv_tokens = 1 - mgr.linear_attention_metadata = SimpleNamespace(states_snapshot_interval=0) - mgr.impl = MagicMock() - mgr.impl.get_kv_cache_stats.return_value = SimpleNamespace( - num_free_blocks_per_window_size={LinearCacheType.RECURRENT_STATES.value: rs_free} - ) - monkeypatch.setattr(KVCacheManager, "get_num_available_tokens", MagicMock(return_value=100)) - - available_tokens = mgr.get_num_available_tokens(token_num_upper_bound=100) - - assert available_tokens == expected - - -def test_cpp_hybrid_available_tokens_without_local_mamba_uses_attention_cap(monkeypatch): - mgr = object.__new__(CppMambaHybridCacheManager) - mgr.linear_attention_metadata = None - mgr.impl = MagicMock() - monkeypatch.setattr(KVCacheManager, "get_num_available_tokens", MagicMock(return_value=100)) - - available_tokens = mgr.get_num_available_tokens(token_num_upper_bound=100) - - assert available_tokens == 100 - mgr.impl.get_kv_cache_stats.assert_not_called() - - def _build_hybrid_with_mamba_layer_pp( spec_config=None, max_batch_size=4, enable_block_reuse=False, pp_size=2 ): diff --git a/tests/unittest/_torch/executor/test_model_loader_gms.py b/tests/unittest/_torch/executor/test_model_loader_gms.py index 9f19a3bfaeef..a1a6e1b9b700 100644 --- a/tests/unittest/_torch/executor/test_model_loader_gms.py +++ b/tests/unittest/_torch/executor/test_model_loader_gms.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Unit tests for ``ModelLoader`` configuration and GMS-specific branches.""" +"""Unit tests for GMS-specific branches in ``ModelLoader``.""" from contextlib import contextmanager, nullcontext from types import SimpleNamespace @@ -112,52 +112,6 @@ def _make_loader(monkeypatch, *, events, spec_config=None): return loader -def test_load_config_propagates_runtime_seq_slot_capacity(monkeypatch): - llm_args = MagicMock() - llm_args.model_kwargs = None - llm_args.cuda_graph_config = None - llm_args.moe_config = MagicMock() - llm_args.nvfp4_gemm_config = MagicMock() - llm_args.multimodal_config = MagicMock() - llm_args.kv_cache_config = MagicMock() - loader = ModelLoader( - llm_args=llm_args, - mapping=MagicMock(), - spec_config=None, - sparse_attention_config=None, - max_num_tokens=128, - max_seq_len=128, - max_num_seq_slots=16, - ) - text_config = SimpleNamespace(hidden_size=64) - pretrained_config = SimpleNamespace( - get_text_config=lambda: text_config, - torch_dtype=torch.float16, - model_type="test", - ) - config = SimpleNamespace( - extra_attrs={}, - nvfp4_gemm_allowed_backends=None, - max_num_tokens=128, - pretrained_config=pretrained_config, - ) - checkpoint_loader = MagicMock() - checkpoint_loader.load_config.return_value = config - monkeypatch.setattr(model_loader_mod, "validate_encoder_decoder_tp_scope", lambda *_args: None) - monkeypatch.setattr( - model_loader_mod, "validate_encoder_decoder_kv_cache_config", lambda *_args: None - ) - monkeypatch.setattr(model_loader_mod, "validate_and_set_kv_cache_quant", lambda *_args: None) - monkeypatch.setattr( - model_loader_mod, "validate_and_set_mamba_ssm_cache_dtype", lambda *_args: None - ) - - loaded_config = loader._load_and_validate_config("/ckpt", checkpoint_loader) - - assert loaded_config is config - assert config.extra_attrs["max_num_seq_slots"] == 16 - - def _build_gms_backend(*, is_rw, events): backend = MagicMock() backend.connect.return_value = True diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 6dd4167a1338..a0cd574b196d 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -25,12 +25,7 @@ SHUTDOWN_REQUEST_ID, RequestQueueItem, ) -from tensorrt_llm._torch.pyexecutor.llm_request import ( - ATTENTION_DP_DUMMY_REQUEST_ID, - LlmRequest, - LlmRequestState, - SamplingConfig, -) +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState, SamplingConfig from tensorrt_llm._torch.pyexecutor.py_executor import DisaggTransferAdmissionController, PyExecutor from tensorrt_llm._torch.pyexecutor.resource_manager import NoFreeSlotsError, ResourceManagerType from tensorrt_llm._torch.pyexecutor.scheduler import ( @@ -1194,7 +1189,6 @@ def _make_adp_request( if isinstance(state, LlmRequestState) else _SENTINEL_STATE_VALUES.get(state, LlmRequestState.GENERATION_IN_PROGRESS.value) ) - req.request_id = request_id req.py_request_id = request_id req.is_child = is_child req.parent_request_id = parent_request_id @@ -1206,8 +1200,6 @@ def _make_adp_request( req.is_attention_dp_dummy = False req.llm_request_type = llm_request_type req.py_seq_slot = None - req.is_context_init_state = state == LlmRequestState.CONTEXT_INIT - req.py_encoder_output_ready_event = None return req @@ -1220,6 +1212,7 @@ def __init__( max_num_tokens=8192, is_warmup=False, benchmark_req_queues_size=0, + enable_dsv4_adp_dummy_fixes=True, ): self.enable_attention_dp = enable_attention_dp self.kv_cache_transceiver = kv_cache_transceiver @@ -1233,16 +1226,12 @@ def __init__( self.max_num_tokens = max_num_tokens self._adp_dummy_is_gen = True self._pending_adp_dummy_request = None + self._enable_dsv4_adp_dummy_fixes = enable_dsv4_adp_dummy_fixes self.add_dummy_calls = [] - self.inflight_req_ids = set() self.dist = Mock() self.dist.tp_size = 1 - self.dist.pp_size = 1 - self.dist.cp_size = 1 self.dist.tp_allgather.side_effect = lambda value: [value] - self.dist.pp_allgather.side_effect = lambda value: [value] - self.dist.cp_allgather.side_effect = lambda value: [value] kv_cache_manager = Mock() kv_cache_manager.mapping.has_cp_helix.return_value = False @@ -1269,10 +1258,8 @@ def _run_pad(stub): "_count_schedulable_active_requests", "_has_adp_dummy_kv_capacity", "_should_skip_dummy_for_benchmark_disagg", - "_can_commit_adp_dummy_candidate", ): setattr(stub, helper, types.MethodType(getattr(PyExecutor, helper), stub)) - stub._dist_size = PyExecutor._dist_size PyExecutor._pad_attention_dp_dummy_request(stub) @@ -1352,29 +1339,17 @@ def test_adp_dummy_role_unchanged_when_attention_dp_disabled(): assert stub._adp_dummy_is_gen is True -def test_pad_dummy_added_when_only_to_complete_requests_non_disagg(): - # GENERATION_TO_COMPLETE is outside the MicroBatchScheduler upper bound in - # every mode. Counting it would suppress padding even though this rank's - # next scheduled batch is empty. - stub = _StubADPExecutor(kv_cache_transceiver=None) - stub.active_requests = [_make_adp_request(_STATE_GENERATION_TO_COMPLETE)] - stub.expected_num_active_requests = 1 - - _run_pad(stub) - - assert len(stub.add_dummy_calls) == 1 - assert len(stub.active_requests) == 2 - - -def test_pad_dummy_counts_encoder_init_for_encoder_decoder(): - stub = _StubADPExecutor(kv_cache_transceiver=None) - cross_kv_cache_manager = object() - stub.resource_manager.get_resource_manager.side_effect = ( - lambda resource_type: cross_kv_cache_manager - if resource_type == ResourceManagerType.CROSS_KV_CACHE_MANAGER - else None - ) - stub.active_requests = [_make_adp_request(LlmRequestState.ENCODER_INIT)] +@pytest.mark.parametrize( + "state", + [ + _STATE_GENERATION_TO_COMPLETE, + LlmRequestState.DISAGG_CONTEXT_WAIT_SCHEDULER, + ], +) +def test_disabled_dsv4_gate_preserves_existing_disagg_behavior(state): + # The disabled gate covers non-DSv4 and PP configurations. + stub = _StubADPExecutor(enable_dsv4_adp_dummy_fixes=False) + stub.active_requests = [_make_adp_request(state)] stub.expected_num_active_requests = 1 _run_pad(stub) @@ -1383,26 +1358,6 @@ def test_pad_dummy_counts_encoder_init_for_encoder_decoder(): assert len(stub.active_requests) == 1 -def test_pad_dummy_ignores_context_waiting_for_encoder_output(): - stub = _StubADPExecutor(kv_cache_transceiver=None) - cross_kv_cache_manager = object() - stub.resource_manager.get_resource_manager.side_effect = ( - lambda resource_type: cross_kv_cache_manager - if resource_type == ResourceManagerType.CROSS_KV_CACHE_MANAGER - else None - ) - waiting_request = _make_adp_request(LlmRequestState.CONTEXT_INIT) - waiting_request.py_encoder_output_ready_event = Mock() - waiting_request.py_encoder_output_ready_event.query.return_value = False - stub.active_requests = [waiting_request] - stub.expected_num_active_requests = 2 - - _run_pad(stub) - - assert len(stub.add_dummy_calls) == 1 - assert len(stub.active_requests) == 2 - - def test_pad_dummy_added_when_only_to_complete_requests_disagg(): # In disaggregated mode a GENERATION_TO_COMPLETE request is refused by # MicroBatchScheduler (no_schedule_after_state), so a rank holding only @@ -1444,24 +1399,17 @@ def test_pad_dummy_allocation_failure_skips_padding(): stub.expected_num_active_requests = 2 stub.kv_cache_manager.add_dummy_requests.side_effect = None stub.kv_cache_manager.add_dummy_requests.return_value = None - stub.dist.pp_size = 2 - stub.dist.pp_allgather.side_effect = None - stub.dist.pp_allgather.return_value = [[(True, False)], [(True, True)]] _run_pad(stub) - stub.dist.pp_allgather.assert_called_once_with([(True, False)]) assert len(stub.active_requests) == 1 assert not any(r.is_attention_dp_dummy for r in stub.active_requests) -def test_pad_dummy_context_capacity_failure_joins_lane_consensus(): +def test_dsv4_pad_dummy_checks_full_context_capacity(): stub = _StubADPExecutor(max_num_tokens=4096) stub._adp_dummy_is_gen = False stub.kv_cache_manager.get_num_available_tokens.return_value = 1024 - stub.dist.pp_size = 2 - stub.dist.pp_allgather.side_effect = None - stub.dist.pp_allgather.return_value = [[(True, False)], [(True, True)]] _run_pad(stub) @@ -1469,16 +1417,12 @@ def test_pad_dummy_context_capacity_failure_joins_lane_consensus(): token_num_upper_bound=4096, max_num_draft_tokens=0 ) stub.kv_cache_manager.add_dummy_requests.assert_not_called() - stub.dist.pp_allgather.assert_called_once_with([(True, False)]) assert stub._pending_adp_dummy_request is None -def test_pad_dummy_generation_capacity_failure_joins_lane_consensus(): +def test_dsv4_pad_dummy_checks_full_generation_capacity(): stub = _StubADPExecutor() stub.kv_cache_manager.get_num_available_tokens.return_value = 0 - stub.dist.pp_size = 2 - stub.dist.pp_allgather.side_effect = None - stub.dist.pp_allgather.return_value = [[(True, False)], [(True, True)]] _run_pad(stub) @@ -1486,142 +1430,20 @@ def test_pad_dummy_generation_capacity_failure_joins_lane_consensus(): token_num_upper_bound=1, max_num_draft_tokens=0 ) stub.kv_cache_manager.add_dummy_requests.assert_not_called() - stub.dist.pp_allgather.assert_called_once_with([(True, False)]) - assert stub._pending_adp_dummy_request is None - - -def test_pad_dummy_pp_peer_allocation_failure_rolls_back_candidate(): - stub = _StubADPExecutor() - terminal_request = _make_adp_request(_STATE_GENERATION_TO_COMPLETE) - stub.active_requests = [terminal_request] - stub.expected_num_active_requests = 2 - stub.dist.pp_size = 2 - stub.dist.pp_allgather.side_effect = None - stub.dist.pp_allgather.return_value = [[(True, True)], [(True, False)]] - dummy_request = _make_adp_request( - _STATE_GENERATION_IN_PROGRESS, - request_id=ATTENTION_DP_DUMMY_REQUEST_ID, - is_dummy_request=True, - ) - spec_resource_manager = Mock() - stub.resource_manager.get_resource_manager.return_value = spec_resource_manager - stub.kv_cache_manager.add_dummy_requests.side_effect = None - stub.kv_cache_manager.add_dummy_requests.return_value = [dummy_request] - - _run_pad(stub) - - assert stub.active_requests == [terminal_request] - spec_resource_manager.add_dummy_requests.assert_called_once_with( - [ATTENTION_DP_DUMMY_REQUEST_ID] - ) - spec_resource_manager.free_resources.assert_called_once_with(dummy_request) - stub.kv_cache_manager.free_resources.assert_called_once_with(dummy_request) assert stub._pending_adp_dummy_request is None -@pytest.mark.parametrize( - "cp_statuses, pp_statuses", - [ - ([(True, True), (True, False)], None), - ( - [(True, True), (True, True)], - [ - [(True, True), (True, True)], - [(True, True), (True, False)], - ], - ), - ], - ids=["cp", "cp-and-pp"], -) -def test_adp_dummy_lane_consensus_includes_every_cp_shard(cp_statuses, pp_statuses): - stub = _StubADPExecutor() - stub._dist_size = PyExecutor._dist_size - stub.dist.cp_size = 2 - stub.dist.cp_allgather.side_effect = None - stub.dist.cp_allgather.return_value = cp_statuses - if pp_statuses is not None: - stub.dist.pp_size = 2 - stub.dist.pp_allgather.side_effect = None - stub.dist.pp_allgather.return_value = pp_statuses - - can_commit = PyExecutor._can_commit_adp_dummy_candidate( - stub, needs_dummy=True, dummy_candidate=object() - ) - - assert can_commit is False - stub.dist.cp_allgather.assert_called_once_with((True, True)) - if pp_statuses is not None: - stub.dist.pp_allgather.assert_called_once_with(cp_statuses) - - -def test_pad_dummy_busy_pp_rank_participates_in_lane_consensus(): +def test_dsv4_pad_dummy_capacity_includes_draft_reserve(): stub = _StubADPExecutor() - active_request = _make_adp_request(_STATE_GENERATION_IN_PROGRESS) - stub.active_requests = [active_request] - stub.expected_num_active_requests = 1 - stub.dist.pp_size = 2 - stub.dist.pp_allgather.side_effect = None - stub.dist.pp_allgather.return_value = [[(False, True)], [(True, True)]] + stub.max_total_draft_tokens = 3 + stub.kv_cache_manager.get_num_available_tokens.return_value = 3 _run_pad(stub) - stub.dist.pp_allgather.assert_called_once_with([(False, True)]) - assert stub.active_requests == [active_request] - stub.kv_cache_manager.add_dummy_requests.assert_not_called() - - -@pytest.mark.parametrize("skip_dummy", [False, True]) -def test_pad_dummy_nonparticipating_pp_rank_still_enters_lane_consensus( - skip_dummy, -): - stub = _StubADPExecutor() - stub.expected_num_active_requests = 0 - if skip_dummy: - stub._benchmark_fill_phase_active = True - stub.dist.pp_size = 2 - stub.dist.pp_allgather.side_effect = None - stub.dist.pp_allgather.return_value = [[(False, True)], [(True, True)]] - - _run_pad(stub) - - stub.dist.pp_allgather.assert_called_once_with([(False, True)]) - stub.kv_cache_manager.add_dummy_requests.assert_not_called() - assert stub._pending_adp_dummy_request is None - - -def test_pad_dummy_added_when_only_active_request_is_inflight(): - stub = _StubADPExecutor(kv_cache_transceiver=None) - active_request = _make_adp_request(_STATE_GENERATION_IN_PROGRESS) - stub.active_requests = [active_request] - stub.expected_num_active_requests = 1 - stub.inflight_req_ids.add(active_request.request_id) - - _run_pad(stub) - - assert len(stub.add_dummy_calls) == 1 - assert len(stub.active_requests) == 2 - - -def test_pad_dummy_does_not_reallocate_live_inflight_dummy_id(): - stub = _StubADPExecutor(kv_cache_transceiver=None) - prior_dummy = _make_adp_request( - _STATE_GENERATION_IN_PROGRESS, - request_id=ATTENTION_DP_DUMMY_REQUEST_ID, - is_dummy_request=True, + stub.kv_cache_manager.get_num_available_tokens.assert_called_once_with( + token_num_upper_bound=4, max_num_draft_tokens=3 ) - prior_dummy.is_attention_dp_dummy = True - stub.active_requests = [prior_dummy] - stub.expected_num_active_requests = 1 - stub.inflight_req_ids.add(ATTENTION_DP_DUMMY_REQUEST_ID) - stub.dist.pp_size = 2 - stub.dist.pp_allgather.side_effect = None - stub.dist.pp_allgather.return_value = [[(True, False)], [(True, True)]] - - _run_pad(stub) - stub.kv_cache_manager.add_dummy_requests.assert_not_called() - stub.dist.pp_allgather.assert_called_once_with([(True, False)]) - assert stub.active_requests == [prior_dummy] assert stub._pending_adp_dummy_request is None @@ -1633,15 +1455,11 @@ def test_pad_dummy_spec_allocation_failure_rolls_back_kv_candidate(): spec_resource_manager = Mock() spec_resource_manager.add_dummy_requests.side_effect = NoFreeSlotsError("No free slots") stub.resource_manager.get_resource_manager.return_value = spec_resource_manager - stub.dist.pp_size = 2 - stub.dist.pp_allgather.side_effect = None - stub.dist.pp_allgather.return_value = [[(True, False)], [(True, True)]] _run_pad(stub) assert stub.active_requests == [terminal_request] assert stub._pending_adp_dummy_request is None - stub.dist.pp_allgather.assert_called_once_with([(True, False)]) stub.kv_cache_manager.free_resources.assert_called_once() diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index f4ce3a7e8899..e5497d999e23 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -3,7 +3,6 @@ import unittest from dataclasses import dataclass -from types import SimpleNamespace from unittest.mock import Mock import torch @@ -595,7 +594,6 @@ def test_prepare_tp_inputs_with_partial_mrope_segments(self) -> None: attn_metadata.is_cuda_graph = False model_engine.max_num_tokens = 32 - model_engine.max_num_seq_slots = 8 model_engine.input_ids_cuda = torch.zeros(32, dtype=torch.int32, device='cuda') @@ -627,9 +625,6 @@ def test_prepare_tp_inputs_with_partial_mrope_segments(self) -> None: dummy_request.sampling_config.beam_width = 1 dummy_request.py_multimodal_data = {} dummy_request.is_cuda_graph_dummy = True - dummy_request.py_mrope_position_delta = torch.tensor([[0]], - dtype=torch.int32, - device='cuda') scheduled_requests = ScheduledRequests() scheduled_requests.context_requests_last_chunk = [] @@ -651,21 +646,11 @@ def test_prepare_tp_inputs_with_partial_mrope_segments(self) -> None: self.assertEqual(result["mrope_delta_write_seq_slots"].cpu().tolist(), [0]) self.assertEqual(result["mrope_delta_read_seq_slots"].cpu().tolist(), - [0, model_engine.max_num_seq_slots]) + [0]) self.assertNotIn("multimodal_embedding", multimodal_request.py_multimodal_data) kv_cache_manager.shutdown() - def test_preconstructed_mrope_model_requires_runtime_seq_slot_capacity( - self) -> None: - model_engine = object.__new__(PyTorchModelEngine) - model_engine.max_num_seq_slots = 8 - model_engine.model = SimpleNamespace( - mrope_position_deltas_cache=torch.zeros(8, dtype=torch.int32)) - - with self.assertRaisesRegex(ValueError, "requires at least 9"): - model_engine._validate_mrope_position_delta_cache_capacity() - def test_kv_cache_manager_with_execution_stream(self): """Test that KVCacheManager uses the provided execution_stream. """ diff --git a/tests/unittest/_torch/executor/test_seq_slot_sizing.py b/tests/unittest/_torch/executor/test_seq_slot_sizing.py index 97649f01bde9..d42e6f4483c9 100644 --- a/tests/unittest/_torch/executor/test_seq_slot_sizing.py +++ b/tests/unittest/_torch/executor/test_seq_slot_sizing.py @@ -1,69 +1,116 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Seq-slot pool / sampler sizing must include overlap headroom. +"""DeepSeek-V4 seq-slot sizing includes overlap headroom. Under the overlap scheduler, requests finished in the previous iteration still hold their sequence slots when the next iteration's prepare_resources runs, while the V2 scheduler has already dropped them from its budget (no_schedule_after_state=GENERATION_TO_COMPLETE) and -backfilled their seats. Transient slot demand is therefore -2 * max_batch_size; sizing the pool (and the sampler state it indexes) -at max_batch_size alone exhausts it ("No free slots"). +backfilled their seats. Transient slot demand is therefore +2 * max_batch_size. The headroom is intentionally limited to DeepSeek-V4; +other models preserve their established sizing pending separate validation. 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). """ +from unittest.mock import Mock + import pytest from tensorrt_llm._torch.pyexecutor._util import ( compute_max_num_sequences, create_torch_sampler_args, + should_enable_dsv4_adp_dummy_fixes, + should_enable_dsv4_overlap_headroom, ) from tensorrt_llm.mapping import Mapping SIZING_CASES = [ - # (pp_size, disable_overlap_scheduler, expected_factor) - # Overlap without PP: headroom for two in-flight iterations. - (1, False, 2), - # No overlap: exact batch size suffices. - (1, True, 1), - # With PP the micro-batch count is pp_size, regardless of overlap. - (2, False, 2), - (4, True, 4), + # (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 DSv4 opt-in. + (2, False, True, 2), + (4, False, True, 4), + (4, True, False, 4), ] -@pytest.mark.parametrize("pp_size,disable_overlap,expected_factor", SIZING_CASES) -def test_compute_max_num_sequences_includes_overlap_headroom( - pp_size, disable_overlap, expected_factor +@pytest.mark.parametrize( + "model_type,has_spec,is_mtp_one_model,pp_size,disable_overlap,expected", + [ + ("deepseek_v4", True, True, 1, False, True), + ("deepseek_v3", True, True, 1, False, False), + ("deepseek_v4", False, False, 1, False, False), + ("deepseek_v4", True, False, 1, False, False), + ("deepseek_v4", True, True, 2, False, False), + ("deepseek_v4", True, True, 1, True, False), + ], +) +def test_dsv4_overlap_headroom_gate( + model_type, has_spec, is_mtp_one_model, pp_size, disable_overlap, expected ): - max_batch_size = 8 + spec_config = None + if has_spec: + spec_config = Mock() + spec_config.spec_dec_mode.is_mtp_eagle_one_model.return_value = is_mtp_one_model mapping = Mapping(world_size=pp_size, tp_size=1, pp_size=pp_size) + assert ( - compute_max_num_sequences(mapping, max_batch_size, disable_overlap) - == max_batch_size * expected_factor + should_enable_dsv4_overlap_headroom(model_type, spec_config, mapping, disable_overlap) + is expected ) -@pytest.mark.parametrize("pp_size,disable_overlap,expected_factor", SIZING_CASES) -def test_sampler_max_num_sequences_matches_executor_slot_pool( - pp_size, disable_overlap, expected_factor +@pytest.mark.parametrize( + "model_type,pp_size,expected", + [ + ("deepseek_v4", 1, True), + ("deepseek_v3", 1, False), + ("deepseek_v4", 2, False), + ], +) +def test_dsv4_adp_dummy_fix_gate(model_type, pp_size, expected): + mapping = Mapping(world_size=pp_size, tp_size=1, pp_size=pp_size) + assert should_enable_dsv4_adp_dummy_fixes(model_type, mapping) is expected + + +@pytest.mark.parametrize( + "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, expected_factor ): max_batch_size = 8 mapping = Mapping(world_size=pp_size, tp_size=1, pp_size=pp_size) + assert ( + compute_max_num_sequences( + mapping, + max_batch_size, + disable_overlap, + enable_overlap_headroom=enable_overlap_headroom, + ) + == max_batch_size * expected_factor + ) + + +@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=disable_overlap, + disable_overlap_scheduler=False, enable_async_worker=False, enable_speculative_beam_history_d2h=False, + max_num_sequences=max_num_sequences, ) - assert args.max_num_sequences == compute_max_num_sequences( - mapping, max_batch_size, disable_overlap - ) - assert args.max_num_sequences == max_batch_size * expected_factor + assert args.max_num_sequences == max_num_sequences diff --git a/tests/unittest/_torch/modeling/test_modeling_qwen2_5vl.py b/tests/unittest/_torch/modeling/test_modeling_qwen2_5vl.py index a5e177a2c03f..3f6bc50706ac 100644 --- a/tests/unittest/_torch/modeling/test_modeling_qwen2_5vl.py +++ b/tests/unittest/_torch/modeling/test_modeling_qwen2_5vl.py @@ -24,8 +24,7 @@ Qwen2VLHfWeightMapper from tensorrt_llm._torch.models.modeling_qwen2vl import ( Qwen2_5_VLModel, Qwen2VisionModelBase, Qwen2VLInputProcessorBase, - Qwen2VLModel, _get_mrope_position_delta_cache_size, - _prepare_qwen_vl_mrope_config) + Qwen2VLModel, _prepare_qwen_vl_mrope_config) from tensorrt_llm._torch.models.modeling_qwen3vl import \ Qwen3VLInputProcessorBase from tensorrt_llm._utils import get_sm_version @@ -427,20 +426,13 @@ def _mrope_param(delta: int) -> MultimodalParams: }) -def test_mrope_delta_cache_size_uses_runtime_seq_slot_capacity(): - model_config = ModelConfig(max_num_tokens=32) - model_config.extra_attrs['max_num_seq_slots'] = 8 - - assert _get_mrope_position_delta_cache_size(model_config) == 9 - - def test_prepare_qwen_vl_mrope_config_mixed_context_generation(): rotary_dim = 2 num_tokens = 5 position_ids = torch.arange(3 * num_tokens, device="cuda", dtype=torch.int32).reshape(3, 1, num_tokens) - mrope_position_deltas_cache = torch.zeros(9, + mrope_position_deltas_cache = torch.zeros(8, device="cuda", dtype=torch.int32) mrope_rotary_cos_sin_workspace = torch.empty((1, 8 * rotary_dim * 2), @@ -449,25 +441,24 @@ def test_prepare_qwen_vl_mrope_config_mixed_context_generation(): config = _prepare_qwen_vl_mrope_config( multimodal_params=[_mrope_param(11), _mrope_param(22)], - num_generation_requests=3, + num_generation_requests=2, position_ids=position_ids, rotary_emb=_FakeRotaryEmbedding(rotary_dim), mrope_position_deltas_cache=mrope_position_deltas_cache, mrope_rotary_cos_sin_workspace=mrope_rotary_cos_sin_workspace, - mrope_delta_write_seq_slots=torch.tensor([4, 7], + mrope_delta_write_seq_slots=torch.tensor([2, 5], device="cuda", dtype=torch.long), - mrope_delta_read_seq_slots=torch.tensor([4, 7, 8], + mrope_delta_read_seq_slots=torch.tensor([2, 5], device="cuda", dtype=torch.long)) torch.testing.assert_close( - mrope_position_deltas_cache[[4, 7]], + mrope_position_deltas_cache[[2, 5]], torch.tensor([11, 22], device="cuda", dtype=torch.int32)) - assert mrope_position_deltas_cache[8].item() == 0 torch.testing.assert_close( config["mrope_position_deltas"], - torch.tensor([[11], [22], [0]], device="cuda", dtype=torch.int32)) + torch.tensor([[11], [22]], device="cuda", dtype=torch.int32)) packed = config["mrope_rotary_cos_sin"].view(1, num_tokens, rotary_dim, 2) cos, sin = _FakeRotaryEmbedding(rotary_dim).get_cos_sin(position_ids)