diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index dc7b753a89ef..b0843e91c02b 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2637,31 +2637,36 @@ def compute_max_num_sequences(mapping: Mapping, return max_batch_size * num_micro_batches -# Model types whose disaggregated attention-DP path has been measured against -# the ADP dummy fixes. The gate stays an explicit list rather than a capability -# check (``enable_attention_dp and kv_cache_transceiver is not None``) so that -# each entry is added only after its disagg ADP behavior has been exercised. -_ADP_DUMMY_FIX_MODEL_TYPES = ("deepseek_v4", "qwen3_5_moe") +def should_enable_adp_dummy_fixes(mapping: Mapping) -> bool: + """Enable transactional ADP dummy handling while PP remains follow-up.""" + return not mapping.has_pp() -def should_enable_dsv4_adp_dummy_fixes(model_type: Optional[str], - mapping: Mapping) -> bool: - """Gate the ADP dummy fixes while PP remains follow-up scope.""" - return model_type in _ADP_DUMMY_FIX_MODEL_TYPES and not mapping.has_pp() +_VALIDATED_OVERLAP_ADP_DUMMY_MODEL_TYPES = ("deepseek_v4", "qwen3_5_moe") + + +def should_enable_scheduler_aware_adp_dummy( + model_type: Optional[str], mapping: Mapping, + disable_overlap_scheduler: bool) -> bool: + """Enable scheduler-aware padding for validated lifecycle configurations.""" + return (should_enable_adp_dummy_fixes(mapping) + and (disable_overlap_scheduler + or model_type in _VALIDATED_OVERLAP_ADP_DUMMY_MODEL_TYPES)) + + +def should_enable_non_overlap_adp_forward_intent( + mapping: Mapping, disable_overlap_scheduler: bool) -> bool: + """Enable fresh cross-rank dummy intent for the generic non-overlap path.""" + return (should_enable_adp_dummy_fixes(mapping) + and disable_overlap_scheduler) 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. - - Deliberately NOT routed through ``should_enable_dsv4_adp_dummy_fixes``. - That gate now covers more model types, while this one doubles - ``max_num_sequences`` (see ``compute_max_num_sequences``) and therefore - changes the memory envelope; it must stay pinned to the one path it was - measured on. - """ - return (model_type == "deepseek_v4" and not mapping.has_pp() + """Gate extra sequence slots to the validated DSv4 MTP overlap path.""" + return (model_type == "deepseek_v4" + and should_enable_adp_dummy_fixes(mapping) and spec_config is not None and spec_config.spec_dec_mode.is_mtp_eagle_one_model() and not disable_overlap_scheduler) @@ -2990,8 +2995,12 @@ def create_py_executor_instance( enable_prefix_aware_scheduling=enable_prefix_aware_scheduling, ) - mb_scheduler = BindMicroBatchScheduler(max_batch_size, max_num_tokens, - ctx_chunk_config) + mb_scheduler = BindMicroBatchScheduler( + max_batch_size, + max_num_tokens, + ctx_chunk_config, + no_schedule_until_state=no_schedule_until_state, + ) reorder_policy_config = llm_args.reorder_policy_config if reorder_policy_config is not None: diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 390ea7e6f094..dc5f4b8071c6 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -360,8 +360,10 @@ def __init__( # 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) + should_enable_adp_dummy_fixes, + should_enable_dsv4_overlap_headroom, + should_enable_non_overlap_adp_forward_intent, + should_enable_scheduler_aware_adp_dummy) self.max_num_seq_slots = compute_max_num_sequences( mapping, self.batch_size, llm_args.disable_overlap_scheduler) self.dist = dist @@ -447,11 +449,16 @@ def __init__( self.model = model 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)) + # Apply transactional dummy handling to every non-PP disaggregated ADP + # model. The larger slot pool remains restricted to the validated + # DeepSeek-V4 MTP overlap configuration. + self._enable_adp_dummy_fixes = should_enable_adp_dummy_fixes(mapping) + self._enable_scheduler_aware_adp_dummy = ( + should_enable_scheduler_aware_adp_dummy( + model_type, mapping, llm_args.disable_overlap_scheduler)) + self._enable_non_overlap_adp_forward_intent = ( + should_enable_non_overlap_adp_forward_intent( + mapping, llm_args.disable_overlap_scheduler)) self._enable_dsv4_overlap_headroom = ( should_enable_dsv4_overlap_headroom( model_type, spec_config, mapping, diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 6cf98b759f02..def5d0e5dea6 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -97,6 +97,13 @@ _UNBOUNDED_STATS_MAX_LEN = -1 +class _ADPForwardIntent(IntEnum): + # MAX reduction gives context precedence when ADP ranks have mixed work. + NONE = 0 + GENERATION = 1 + CONTEXT = 2 + + def _stats_buffer_is_unbounded(max_stats_len: int) -> bool: return max_stats_len == _UNBOUNDED_STATS_MAX_LEN @@ -569,8 +576,12 @@ 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_adp_dummy_fixes = getattr(model_engine, + "_enable_adp_dummy_fixes", False) + self._enable_scheduler_aware_adp_dummy = getattr( + model_engine, "_enable_scheduler_aware_adp_dummy", False) + self._enable_non_overlap_adp_forward_intent = getattr( + model_engine, "_enable_non_overlap_adp_forward_intent", False) self.enable_attention_dp = model_engine.enable_attention_dp self.dist = dist self.sampler = sampler @@ -712,8 +723,8 @@ def __init__( # lifted to _handle_kv_transfer_timeouts_synced / _flush_iter_stats_synced. self._pending_timed_out_requests: List[LlmRequest] = [] self._pending_iter_stats_dict: Optional[Dict] = None - # ADP dummy role for _pad_attention_dp_dummy_request. Default is gen; - # updated from observed request types. + # Legacy ADP dummy role for overlap and PP fallback paths. The generic + # non-overlap path derives its role from fresh per-iteration intent. 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 @@ -3344,7 +3355,7 @@ 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: + if not self._enable_adp_dummy_fixes: return dummy_request = self._pending_adp_dummy_request @@ -5113,7 +5124,8 @@ def _fetch_new_requests( all_new_flat = [ req for reqs in all_ranks_new_requests.values() for req in reqs ] - self._update_adp_dummy_role(all_new_flat) + if not self._enable_non_overlap_adp_forward_intent: + self._update_adp_dummy_role(all_new_flat) # Update per-rank counter for DP self.num_fetch_requests_cur_rank += len(new_requests_cur_rank) @@ -5765,31 +5777,56 @@ def _check_disagg_ctx_schedulable_status(self, def _count_schedulable_active_requests(self) -> int: """Count active requests that are ready for scheduling. - 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. + The non-PP disaggregated ADP path uses the scheduler's state- + eligibility contract. This keeps decoder-only and encoder-decoder + boundaries and special exclusions aligned without duplicating them + here. PP eligibility remains follow-up scope. Returns: The number of active requests eligible for scheduling. """ - if (not self._enable_dsv4_adp_dummy_fixes + if (not self._enable_scheduler_aware_adp_dummy or self.kv_cache_transceiver is None): if self.kv_cache_transceiver is None: return len(self.active_requests) + # PP intentionally preserves its established ADP padding behavior + # until its dummy lifecycle is generalized. Keep this fallback on + # semantic request properties so enum reordering cannot silently + # change which transfer states it excludes. 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 self.scheduler.is_request_in_schedulable_state(req)) - return sum( - 1 for req in self.active_requests - if schedule_from_value <= req.state_value < to_complete_value) + def _get_non_overlap_adp_forward_intent( + self) -> tuple[int, _ADPForwardIntent]: + """Return local eligible-real count and fresh TP-wide forward role. + + This runs before capacity scheduling so the result is forward intent, + not a guarantee that every eligible request will be admitted. The + post-schedule queue vote commits or rolls back the tentative dummy. + """ + local_schedulable_count = 0 + local_intent = _ADPForwardIntent.NONE + for request in self.active_requests: + if (request.is_attention_dp_dummy or + not self.scheduler.is_request_in_schedulable_state(request) + ): + continue + + local_schedulable_count += 1 + if request.is_encoder_init_state or request.is_context_init_state: + local_intent = _ADPForwardIntent.CONTEXT + elif local_intent == _ADPForwardIntent.NONE: + local_intent = _ADPForwardIntent.GENERATION + + global_intent = self.dist.tp_allreduce(int(local_intent), + op=ReduceOp.MAX) + return local_schedulable_count, _ADPForwardIntent(global_intent) def _has_adp_dummy_kv_capacity(self, token_nums: Optional[List[int]]) -> bool: @@ -5903,8 +5940,18 @@ def _pad_attention_dp_dummy_request(self): if self._should_skip_dummy_for_benchmark_disagg(num_active_request): return - needs_dummy = (expected_num_active_requests > 0 - and num_active_request == 0) + if (self._enable_non_overlap_adp_forward_intent + and self.kv_cache_transceiver is not None): + num_active_request, global_intent = ( + self._get_non_overlap_adp_forward_intent()) + if global_intent != _ADPForwardIntent.NONE: + self._adp_dummy_is_gen = ( + global_intent == _ADPForwardIntent.GENERATION) + needs_dummy = (global_intent != _ADPForwardIntent.NONE + and num_active_request == 0) + else: + needs_dummy = (expected_num_active_requests > 0 + and num_active_request == 0) if not needs_dummy: return @@ -5937,7 +5984,7 @@ def _pad_attention_dp_dummy_request(self): key="attention_dp_dummy_insufficient_kv_capacity") return - if (not self._enable_dsv4_adp_dummy_fixes + if (not self._enable_adp_dummy_fixes or self.kv_cache_transceiver is None): llm_request = self.kv_cache_manager.add_dummy_requests( request_ids=dummy_request_ids, @@ -5972,9 +6019,8 @@ def _pad_attention_dp_dummy_request(self): except OutOfPagesError: dummy_requests = None if not dummy_requests: - logger.warning( - "Cannot allocate DeepSeek-V4 ADP pad dummy; rank schedules " - "an empty batch and the fleet will retry.") + logger.warning("Cannot allocate ADP pad dummy; rank schedules " + "an empty batch and the fleet will retry.") return dummy_request = dummy_requests[0] diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 153fe93e82ae..5129b37faa3b 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -344,6 +344,11 @@ def __init__( self.indexer_k_cache_local_layer_mask = None self.kv_connector_manager = kv_connector_manager + # Dummy requests can reserve their V1 sequence before they enter the + # normal context prepare path. Track that ownership per manager so the + # same sequence is not registered twice, while a separate draft or + # cross-cache manager can still prepare its own copy. + self._preprepared_dummy_request_ids: set[int] = set() tp_size = mapping.tp_size if mapping.enable_attention_dp: @@ -785,6 +790,8 @@ def get_needed_resource_to_completion(self, request: LlmRequest) -> int: def _context_seq_len(self, req: LlmRequest, is_cross: bool, is_star_cp: bool) -> Optional[int]: """Return the sequence length to pass to add_sequence_batch, or None to skip this request.""" + if req.py_request_id in self._preprepared_dummy_request_ids: + return None if is_cross: if (getattr(req, "py_skip_cross_kv_projection", False) or not req.is_first_context_chunk @@ -1080,6 +1087,14 @@ def add_dummy_requests( raise cleanup_error raise + if batch_request_infos: + self._preprepared_dummy_request_ids.update( + req_id for req_id, _, _ in batch_request_infos) + if (draft_batch_request_infos + and isinstance(draft_kv_cache_manager, KVCacheManager)): + draft_kv_cache_manager._preprepared_dummy_request_ids.update( + req_id for req_id, _, _ in draft_batch_request_infos) + return requests def update_resources(self, @@ -1130,8 +1145,10 @@ def update_resources(self, self.impl.store_context_blocks(request) def free_resources(self, request: LlmRequest, pin_on_release: bool = False): - return self.impl.remove_sequence(request.py_request_id, request, - pin_on_release) + result = self.impl.remove_sequence(request.py_request_id, request, + pin_on_release) + self._preprepared_dummy_request_ids.discard(request.py_request_id) + return result def store_blocks_for_reuse(self, request: LlmRequest, diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py index caa8e3cb3de1..05291aebce2b 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py @@ -216,6 +216,25 @@ def reset_context_requests(self, context_requests: RequestList | None = None) -> class RequestScheduler(ABC): + @property + @abstractmethod + def scheduling_state_range(self) -> tuple[LlmRequestState, LlmRequestState]: + """Return the half-open state range admitted to a forward batch.""" + raise NotImplementedError + + def is_request_in_schedulable_state(self, request: LlmRequest) -> bool: + """Return whether request state permits admission to a forward batch.""" + if is_decoder_context_request_waiting_for_encoder_output(request): + return False + if request.state in ( + LlmRequestState.DISAGG_CONTEXT_WAIT_SCHEDULER, + LlmRequestState.DISAGG_GENERATION_INIT, + LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS, + ): + return False + schedule_from, schedule_to = self.scheduling_state_range + return schedule_from.value <= request.state_value < schedule_to.value + @abstractmethod def schedule_request( self, active_requests: RequestList, inflight_request_ids: set[int] @@ -392,10 +411,14 @@ def __init__( max_batch_size: int, max_num_tokens: int = None, ctx_chunk_config: Optional[tuple[StrEnum, int]] = None, + no_schedule_until_state: LlmRequestState = LlmRequestState.CONTEXT_INIT, + no_schedule_after_state: LlmRequestState = LlmRequestState.GENERATION_TO_COMPLETE, ) -> None: super(BindMicroBatchScheduler, self).__init__() self.max_batch_size = max_batch_size self.max_num_tokens = max_num_tokens + self.no_schedule_until_state = no_schedule_until_state + self.no_schedule_after_state = no_schedule_after_state ctx_chunk_config_cpp = None if ctx_chunk_config is not None: @@ -403,7 +426,12 @@ def __init__( ctx_chunk_config[0]._to_pybind(), ctx_chunk_config[1] ) - self.impl = tb_internal.algorithms.MicroBatchScheduler(ctx_chunk_config_cpp, max_num_tokens) + self.impl = tb_internal.algorithms.MicroBatchScheduler( + ctx_chunk_config=ctx_chunk_config_cpp, + max_context_length=max_num_tokens, + no_schedule_until_state=no_schedule_until_state, + no_schedule_after_state=no_schedule_after_state, + ) def schedule( self, active_requests: RequestList, inflight_request_ids: set[int] @@ -427,6 +455,13 @@ def __init__( self.capacity_scheduler = capacity_scheduler self.micro_batch_scheduler = micro_batch_scheduler + @property + def scheduling_state_range(self) -> tuple[LlmRequestState, LlmRequestState]: + return ( + self.micro_batch_scheduler.no_schedule_until_state, + self.micro_batch_scheduler.no_schedule_after_state, + ) + def schedule_request( self, active_requests: RequestList, inflight_request_ids: set[int] ) -> SchedulerOutput: @@ -1867,6 +1902,13 @@ def __init__( no_schedule_until_state=no_schedule_until_state, ) + @property + def scheduling_state_range(self) -> tuple[LlmRequestState, LlmRequestState]: + return ( + self.micro_batch_scheduler.no_schedule_until_state, + self.micro_batch_scheduler.no_schedule_after_state, + ) + def schedule_request( self, active_requests: RequestList, inflight_request_ids: set[int] ) -> SchedulerOutput: diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index 958afc59ac04..56f353f936c0 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -205,6 +205,8 @@ def __init__( # MicroBatchScheduler. For encoder-decoder models, caller should pass # no_schedule_until_state=ENCODER_INIT to widen the range (same as # C++ trtEncoderModel which passes kENCODER_INIT). + self.no_schedule_until_state = no_schedule_until_state + self.no_schedule_after_state = no_schedule_after_state self._no_schedule_until_state_value = no_schedule_until_state.value self._no_schedule_after_state_value = no_schedule_after_state.value self._context_init_state_value = LlmRequestState.CONTEXT_INIT.value @@ -220,6 +222,12 @@ def __init__( os.environ.get("TLLM_DISAGG_GEN_PRIORITIZE_FIRST_TOKEN", "0") == "1" ) + @property + def scheduling_state_range( + self, + ) -> tuple[LlmRequestState, LlmRequestState]: + return self.no_schedule_until_state, self.no_schedule_after_state + def schedule_request( self, active_requests: RequestList, inflight_request_ids: set[int] ) -> SchedulerOutput: diff --git a/tests/unittest/_torch/executor/test_benchmark_disagg.py b/tests/unittest/_torch/executor/test_benchmark_disagg.py index b8cb399b1203..ac23d82ef21f 100644 --- a/tests/unittest/_torch/executor/test_benchmark_disagg.py +++ b/tests/unittest/_torch/executor/test_benchmark_disagg.py @@ -31,7 +31,7 @@ import pytest from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState -from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm._torch.pyexecutor.scheduler import RequestScheduler, ScheduledRequests pytestmark = pytest.mark.cpu_only @@ -48,13 +48,22 @@ 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 = ( - LlmRequestState.DISAGG_TRANS_ERROR if in_error else LlmRequestState.GENERATION_IN_PROGRESS - ) + if in_error: + req.state = LlmRequestState.DISAGG_TRANS_ERROR + elif in_transfer: + req.state = LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS + elif in_init: + req.state = LlmRequestState.DISAGG_GENERATION_INIT + else: + req.state = LlmRequestState.GENERATION_IN_PROGRESS + req.state_value = req.state.value req.is_attention_dp_dummy = False + req.is_encoder_init_state = False + req.is_context_init_state = False + req.is_generation_in_progress_state = req.state == LlmRequestState.GENERATION_IN_PROGRESS + req.py_encoder_output_ready_event = None return req @@ -588,12 +597,30 @@ 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._enable_adp_dummy_fixes = True + self._enable_scheduler_aware_adp_dummy = True + self._enable_non_overlap_adp_forward_intent = True self.max_num_tokens = None self.dist = Mock() self.dist.tp_size = tp_size self.dist.tp_allgather.side_effect = lambda value: [value] + # Simulate a peer rank with generation compute after the fill gate + # opens, so an empty local rank needs a generation dummy. + self.dist.tp_allreduce.side_effect = lambda value, op: max( + value, int(self._ADPForwardIntent.GENERATION) + ) + + self.scheduler = Mock() + self.scheduler.scheduling_state_range = ( + LlmRequestState.CONTEXT_INIT, + LlmRequestState.GENERATION_TO_COMPLETE, + ) + self.scheduler.is_request_in_schedulable_state.side_effect = ( + lambda request: RequestScheduler.is_request_in_schedulable_state( + self.scheduler, request + ) + ) self.kv_cache_manager = Mock() self.kv_cache_manager.mapping.has_cp_helix.return_value = False @@ -605,10 +632,11 @@ def __init__( self.resource_manager = Mock() self.resource_manager.get_resource_manager.return_value = None - from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor + from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor, _ADPForwardIntent _pad_attention_dp_dummy_request = PyExecutor._pad_attention_dp_dummy_request _count_schedulable_active_requests = PyExecutor._count_schedulable_active_requests + _get_non_overlap_adp_forward_intent = PyExecutor._get_non_overlap_adp_forward_intent _has_adp_dummy_kv_capacity = PyExecutor._has_adp_dummy_kv_capacity _should_skip_dummy_for_benchmark_disagg = PyExecutor._should_skip_dummy_for_benchmark_disagg diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index bd9f4ea2f634..07d02e296bd7 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -1,6 +1,5 @@ # 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 @@ -26,10 +25,15 @@ RequestQueueItem, ) 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.py_executor import ( + DisaggTransferAdmissionController, + PyExecutor, + _ADPForwardIntent, +) from tensorrt_llm._torch.pyexecutor.resource_manager import NoFreeSlotsError, ResourceManagerType from tensorrt_llm._torch.pyexecutor.scheduler import ( FCFSWaitingQueue, + RequestScheduler, ScheduledRequests, SerializableSchedulerOutput, ) @@ -1292,6 +1296,10 @@ def _make_adp_request( req.is_attention_dp_dummy = False req.llm_request_type = llm_request_type req.py_seq_slot = None + req.is_encoder_init_state = state == LlmRequestState.ENCODER_INIT + req.is_context_init_state = state == LlmRequestState.CONTEXT_INIT + req.is_generation_in_progress_state = state == _STATE_GENERATION_IN_PROGRESS + req.py_encoder_output_ready_event = None return req @@ -1306,7 +1314,10 @@ def __init__( kv_manager_max_seq_len=None, is_warmup=False, benchmark_req_queues_size=0, - enable_dsv4_adp_dummy_fixes=True, + enable_adp_dummy_fixes=True, + enable_scheduler_aware_adp_dummy=None, + enable_non_overlap_adp_forward_intent=None, + peer_forward_intent=_ADPForwardIntent.GENERATION, ): self.enable_attention_dp = enable_attention_dp self.kv_cache_transceiver = kv_cache_transceiver @@ -1320,13 +1331,35 @@ 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._enable_adp_dummy_fixes = enable_adp_dummy_fixes + self._enable_scheduler_aware_adp_dummy = ( + enable_adp_dummy_fixes + if enable_scheduler_aware_adp_dummy is None + else enable_scheduler_aware_adp_dummy + ) + self._enable_non_overlap_adp_forward_intent = ( + enable_adp_dummy_fixes + if enable_non_overlap_adp_forward_intent is None + else enable_non_overlap_adp_forward_intent + ) self.add_dummy_calls = [] self.model_engine = Mock(max_num_tokens=max_num_tokens, max_seq_len=max_seq_len) self.dist = Mock() self.dist.tp_size = 1 self.dist.tp_allgather.side_effect = lambda value: [value] + self.dist.tp_allreduce.side_effect = lambda value, op: max(value, int(peer_forward_intent)) + + self.scheduler = Mock() + self.scheduler.scheduling_state_range = ( + LlmRequestState.CONTEXT_INIT, + LlmRequestState.GENERATION_TO_COMPLETE, + ) + self.scheduler.is_request_in_schedulable_state.side_effect = ( + lambda request: RequestScheduler.is_request_in_schedulable_state( + self.scheduler, request + ) + ) kv_cache_manager = Mock() kv_cache_manager.mapping.has_cp_helix.return_value = False @@ -1339,8 +1372,11 @@ def __init__( def _add_dummy(**kwargs): self.add_dummy_calls.append(kwargs) + state = ( + _STATE_GENERATION_IN_PROGRESS if kwargs["is_gen"] else LlmRequestState.CONTEXT_INIT + ) req = _make_adp_request( - _STATE_GENERATION_IN_PROGRESS, + state, request_id=kwargs["request_ids"][0], is_dummy_request=True, ) @@ -1356,6 +1392,7 @@ def _add_dummy(**kwargs): def _run_pad(stub): for helper in ( "_count_schedulable_active_requests", + "_get_non_overlap_adp_forward_intent", "_has_adp_dummy_kv_capacity", "_should_skip_dummy_for_benchmark_disagg", ): @@ -1446,9 +1483,9 @@ def test_adp_dummy_role_unchanged_when_attention_dp_disabled(): 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) +def test_disabled_adp_dummy_fix_gate_preserves_pp_behavior(state): + # PP configurations remain on the established dummy path. + stub = _StubADPExecutor(enable_adp_dummy_fixes=False) stub.active_requests = [_make_adp_request(state)] stub.expected_num_active_requests = 1 @@ -1460,9 +1497,9 @@ def test_disabled_dsv4_gate_preserves_existing_disagg_behavior(state): 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. + # MicroBatchScheduler (no_schedule_after_state). When a peer has real + # generation work, a rank holding only terminal requests must receive a + # pad dummy or can_queue goes False fleet-wide. stub = _StubADPExecutor() stub.active_requests = [_make_adp_request(_STATE_GENERATION_TO_COMPLETE)] stub.expected_num_active_requests = 2 @@ -1477,8 +1514,8 @@ 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. + # schedules batch=0. A peer's generation intent therefore requires 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 @@ -1524,6 +1561,94 @@ def test_pad_dummy_still_added_when_surplus_requests_are_unschedulable() -> None assert stub.expected_num_active_requests == 2 +def test_encoder_init_uses_encoder_decoder_scheduler_state_window(): + stub = _StubADPExecutor() + stub.scheduler.scheduling_state_range = ( + LlmRequestState.ENCODER_INIT, + LlmRequestState.GENERATION_TO_COMPLETE, + ) + 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 + + +@pytest.mark.parametrize( + "state", + [ + LlmRequestState.DISAGG_CONTEXT_WAIT_SCHEDULER, + LlmRequestState.DISAGG_GENERATION_INIT, + LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS, + ], +) +def test_encoder_decoder_disagg_wait_and_transfer_states_are_not_schedulable(state): + stub = _StubADPExecutor() + stub.scheduler.scheduling_state_range = ( + LlmRequestState.ENCODER_INIT, + LlmRequestState.GENERATION_TO_COMPLETE, + ) + stub.active_requests = [_make_adp_request(state)] + stub.expected_num_active_requests = 2 + + _run_pad(stub) + + assert len(stub.add_dummy_calls) == 1 + assert len(stub.active_requests) == 2 + + +def test_decoder_context_waiting_for_encoder_output_is_not_counted(): + stub = _StubADPExecutor() + request = _make_adp_request(LlmRequestState.CONTEXT_INIT) + request.py_encoder_output_ready_event = Mock() + request.py_encoder_output_ready_event.query.return_value = False + stub.active_requests = [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_generic_disagg_adp_mixed_rank_states_stay_queueable(): + # The generic non-PP path must give both ranks a non-empty scheduled batch: + # one rank schedules its real request, while the terminal-only rank + # schedules the dummy inserted for the scheduler-excluded request. + busy_rank = _StubADPExecutor() + busy_rank.active_requests = [_make_adp_request(_STATE_GENERATION_IN_PROGRESS)] + busy_rank.expected_num_active_requests = 2 + terminal_rank = _StubADPExecutor() + terminal_rank.active_requests = [_make_adp_request(_STATE_GENERATION_TO_COMPLETE)] + terminal_rank.expected_num_active_requests = 2 + + _run_pad(busy_rank) + _run_pad(terminal_rank) + + assert busy_rank.add_dummy_calls == [] + assert len(terminal_rank.add_dummy_calls) == 1 + rank_batch_sizes = [ + busy_rank._count_schedulable_active_requests(), + terminal_rank._count_schedulable_active_requests(), + ] + assert rank_batch_sizes == [1, 1] + + for stub, batch_size in zip((busy_rank, terminal_rank), rank_batch_sizes, strict=True): + stub.dist.tp_allgather.side_effect = None + stub.dist.tp_allgather.return_value = rank_batch_sizes + can_queue, can_queue_this_rank = PyExecutor._can_queue( + stub, types.SimpleNamespace(batch_size=batch_size) + ) + + assert can_queue is True + assert can_queue_this_rank is True + PyExecutor._finalize_adp_dummy_allocation(stub, can_queue) + + assert terminal_rank._pending_adp_dummy_request is None + + 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 @@ -1541,22 +1666,11 @@ def test_pad_dummy_allocation_failure_skips_padding(): assert not any(r.is_attention_dp_dummy for r in stub.active_requests) -def test_disabled_dsv4_gate_checks_full_generation_capacity(): - stub = _StubADPExecutor(enable_dsv4_adp_dummy_fixes=False) - stub.max_total_draft_tokens = 4 - stub.kv_cache_manager.get_num_available_tokens.return_value = 4 - - _run_pad(stub) - - stub.kv_cache_manager.get_num_available_tokens.assert_called_once_with( - token_num_upper_bound=5, max_num_draft_tokens=4 +def test_adp_pad_dummy_checks_full_context_capacity(): + stub = _StubADPExecutor( + max_num_tokens=4096, + peer_forward_intent=_ADPForwardIntent.CONTEXT, ) - stub.kv_cache_manager.add_dummy_requests.assert_not_called() - - -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 _run_pad(stub) @@ -1568,7 +1682,7 @@ def test_dsv4_pad_dummy_checks_full_context_capacity(): assert stub._pending_adp_dummy_request is None -def test_dsv4_pad_dummy_checks_full_generation_capacity(): +def test_adp_pad_dummy_checks_full_generation_capacity(): stub = _StubADPExecutor() stub.kv_cache_manager.get_num_available_tokens.return_value = 0 @@ -1581,7 +1695,7 @@ def test_dsv4_pad_dummy_checks_full_generation_capacity(): assert stub._pending_adp_dummy_request is None -def test_dsv4_pad_dummy_capacity_includes_draft_reserve(): +def test_adp_pad_dummy_capacity_includes_draft_reserve(): stub = _StubADPExecutor() stub.max_total_draft_tokens = 3 stub.kv_cache_manager.get_num_available_tokens.return_value = 3 @@ -1697,8 +1811,10 @@ def test_pad_dummy_skips_when_active_request_present(): def test_pad_dummy_ctx_pads_to_max_num_tokens(): - stub = _StubADPExecutor(max_num_tokens=4096) - stub._adp_dummy_is_gen = False + stub = _StubADPExecutor( + max_num_tokens=4096, + peer_forward_intent=_ADPForwardIntent.CONTEXT, + ) stub.expected_num_active_requests = 1 _run_pad(stub) @@ -1707,6 +1823,7 @@ def test_pad_dummy_ctx_pads_to_max_num_tokens(): call = stub.add_dummy_calls[0] assert call["token_nums"] == [4096] assert call["is_gen"] is False + assert stub.active_requests[-1].state == LlmRequestState.CONTEXT_INIT def test_pad_dummy_gen_keeps_default_token_nums(): @@ -1720,25 +1837,60 @@ def test_pad_dummy_gen_keeps_default_token_nums(): call = stub.add_dummy_calls[0] assert call["token_nums"] is None assert call["is_gen"] is True + assert stub.active_requests[-1].state == _STATE_GENERATION_IN_PROGRESS -def test_pad_dummy_ctx_skips_padding_when_max_num_tokens_missing(): - stub = _StubADPExecutor(max_num_tokens=None) +def test_overlap_adp_preserves_legacy_role_without_forward_intent_collective(): + stub = _StubADPExecutor( + max_num_tokens=4096, + enable_scheduler_aware_adp_dummy=False, + enable_non_overlap_adp_forward_intent=False, + ) stub._adp_dummy_is_gen = False stub.expected_num_active_requests = 1 _run_pad(stub) + assert len(stub.add_dummy_calls) == 1 + assert stub.add_dummy_calls[0]["token_nums"] == [4096] + assert stub.add_dummy_calls[0]["is_gen"] is False + stub.dist.tp_allreduce.assert_not_called() + + +def test_pad_dummy_ctx_skips_padding_when_max_num_tokens_missing(): + stub = _StubADPExecutor( + max_num_tokens=None, + peer_forward_intent=_ADPForwardIntent.CONTEXT, + ) + stub.expected_num_active_requests = 1 + + _run_pad(stub) + assert len(stub.add_dummy_calls) == 1 assert stub.add_dummy_calls[0]["token_nums"] is None -def test_pad_dummy_ctx_added_for_disagg_rank_only_awaiting_kv_transfer(): - # Disagg ADP: a rank whose only request is awaiting KV transfer counts as - # idle (excluded by _count_schedulable), so a CTX dummy padded to - # max_num_tokens is added to keep it in the MoE all-to-all. - stub = _StubADPExecutor(max_num_tokens=4096) - stub._adp_dummy_is_gen = False +def test_pad_dummy_not_added_when_all_ranks_only_await_kv_transfer(): + stub = _StubADPExecutor( + max_num_tokens=4096, + peer_forward_intent=_ADPForwardIntent.NONE, + ) + stub.active_requests = [_make_adp_request(_STATE_DISAGG_GENERATION_INIT)] + stub.expected_num_active_requests = 1 + + _run_pad(stub) + + assert stub.add_dummy_calls == [] + assert stub._adp_dummy_is_gen is True + stub.dist.tp_allreduce.assert_called_once_with(int(_ADPForwardIntent.NONE), op=ReduceOp.MAX) + + +def test_pad_dummy_context_role_re_evaluated_while_local_rank_drains(): + stub = _StubADPExecutor( + max_num_tokens=4096, + peer_forward_intent=_ADPForwardIntent.CONTEXT, + ) + stub._adp_dummy_is_gen = True stub.active_requests = [_make_adp_request(_STATE_DISAGG_GENERATION_INIT)] stub.expected_num_active_requests = 1 diff --git a/tests/unittest/_torch/executor/test_resource_manager.py b/tests/unittest/_torch/executor/test_resource_manager.py index f4ab82db1d54..0543fdf0ff5b 100644 --- a/tests/unittest/_torch/executor/test_resource_manager.py +++ b/tests/unittest/_torch/executor/test_resource_manager.py @@ -6,6 +6,7 @@ import subprocess import sys import unittest +from types import SimpleNamespace from typing import NamedTuple, Tuple from unittest.mock import MagicMock, patch @@ -19,6 +20,7 @@ from tensorrt_llm._torch.pyexecutor.resource_manager import ( KVCacheManager, PeftCacheManager, _merge_kv_cache_pool_pointers, _warn_if_unsupported_v1_kv_cache_event_hash_algo) +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.bindings import LayerType from tensorrt_llm.bindings import ModelConfig as ModelConfigCpp from tensorrt_llm.bindings import executor as tllm @@ -990,6 +992,8 @@ def test_add_dummy_requests_failure_frees_partial_allocation(self): token_nums=[64, 64], is_gen=True, max_num_draft_tokens=128) + self.assertEqual(kv_cache_manager._preprepared_dummy_request_ids, + set()) self.assertEqual(kv_cache_manager.get_num_free_blocks(), total_free) # The freed pool must serve follow-up allocations. requests = kv_cache_manager.add_dummy_requests([2], token_nums=[64]) @@ -1087,6 +1091,76 @@ def test_peft_cache_manager_with_execution_stream(self): self.assertTrue(peft_cache_manager.impl.enabled) +@pytest.mark.cpu_only +class TestKVCacheManagerPrepreparedDummies(unittest.TestCase): + + @staticmethod + def _make_manager(is_draft: bool = False) -> KVCacheManager: + manager = KVCacheManager.__new__(KVCacheManager) + manager.mapping = Mapping() + manager.impl = MagicMock() + manager.impl.get_kv_cache_stats.return_value = SimpleNamespace( + free_num_blocks=8) + manager.is_linear_attention = False + manager.is_vswa = False + manager.num_extra_kv_tokens = 0 + manager.kv_cache_type = tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF + manager.is_draft = is_draft + manager.kv_connector_manager = None + manager._kv_reserve_draft_tokens = 0 + manager._preprepared_dummy_request_ids = set() + return manager + + @staticmethod + def _context_batch(request: LlmRequest) -> ScheduledRequests: + batch = ScheduledRequests() + batch.context_requests_last_chunk = [request] + return batch + + def test_context_dummy_is_registered_once_and_id_can_be_reused(self): + manager = self._make_manager() + + requests = manager.add_dummy_requests([0], token_nums=[64]) + self.assertIsNotNone(requests) + request = requests[0] + self.assertEqual(manager._preprepared_dummy_request_ids, {0}) + self.assertEqual(manager.impl.add_sequence_batch.call_count, 1) + + manager.prepare_resources(self._context_batch(request)) + + self.assertEqual(manager.impl.add_sequence_batch.call_count, 1) + + manager.free_resources(request) + self.assertEqual(manager._preprepared_dummy_request_ids, set()) + + manager.add_dummy_requests([0], token_nums=[64]) + self.assertEqual(manager.impl.add_sequence_batch.call_count, 2) + + def test_preprepared_dummy_ownership_is_manager_local(self): + target_manager = self._make_manager() + draft_manager = self._make_manager(is_draft=True) + other_manager = self._make_manager() + + requests = target_manager.add_dummy_requests( + [0], + token_nums=[64], + draft_kv_cache_manager=draft_manager, + ) + self.assertIsNotNone(requests) + request = requests[0] + self.assertEqual(target_manager._preprepared_dummy_request_ids, {0}) + self.assertEqual(draft_manager._preprepared_dummy_request_ids, {0}) + self.assertEqual(other_manager._preprepared_dummy_request_ids, set()) + + target_manager.prepare_resources(self._context_batch(request)) + draft_manager.prepare_resources(self._context_batch(request)) + other_manager.prepare_resources(self._context_batch(request)) + + self.assertEqual(target_manager.impl.add_sequence_batch.call_count, 1) + self.assertEqual(draft_manager.impl.add_sequence_batch.call_count, 1) + self.assertEqual(other_manager.impl.add_sequence_batch.call_count, 1) + + @pytest.mark.cpu_only class TestKVCacheManagerConfigForwarding(unittest.TestCase): diff --git a/tests/unittest/_torch/executor/test_seq_slot_sizing.py b/tests/unittest/_torch/executor/test_seq_slot_sizing.py index 6db765ac4a79..0da0f1c5fb1a 100644 --- a/tests/unittest/_torch/executor/test_seq_slot_sizing.py +++ b/tests/unittest/_torch/executor/test_seq_slot_sizing.py @@ -22,8 +22,10 @@ from tensorrt_llm._torch.pyexecutor._util import ( compute_max_num_sequences, create_torch_sampler_args, - should_enable_dsv4_adp_dummy_fixes, + should_enable_adp_dummy_fixes, should_enable_dsv4_overlap_headroom, + should_enable_non_overlap_adp_forward_intent, + should_enable_scheduler_aware_adp_dummy, ) from tensorrt_llm.mapping import Mapping @@ -48,9 +50,6 @@ ("deepseek_v4", True, False, 1, False, False), ("deepseek_v4", True, True, 2, False, False), ("deepseek_v4", True, True, 1, True, False), - # The widened ADP dummy gate must not leak into the headroom gate: - # doubling max_num_sequences changes the memory envelope and has only - # been validated on the DSv4 MTP overlap path. ("qwen3_5_moe", True, True, 1, False, False), ], ) @@ -69,20 +68,38 @@ def test_dsv4_overlap_headroom_gate( ) +@pytest.mark.parametrize("pp_size,expected", [(1, True), (2, False)]) +def test_adp_dummy_fix_gate(pp_size, expected): + mapping = Mapping(world_size=pp_size, tp_size=1, pp_size=pp_size) + assert should_enable_adp_dummy_fixes(mapping) is expected + + +@pytest.mark.parametrize( + "model_type,pp_size,disable_overlap,expected", + [ + ("kimi_k2", 1, True, True), + ("kimi_k2", 1, False, False), + ("deepseek_v4", 1, False, True), + ("qwen3_5_moe", 1, False, True), + ("deepseek_v4", 2, True, False), + ], +) +def test_scheduler_aware_adp_dummy_scope(model_type, pp_size, disable_overlap, expected): + mapping = Mapping(world_size=pp_size, tp_size=1, pp_size=pp_size) + assert should_enable_scheduler_aware_adp_dummy(model_type, mapping, disable_overlap) is expected + + @pytest.mark.parametrize( - "model_type,pp_size,expected", + "pp_size,disable_overlap,expected", [ - ("deepseek_v4", 1, True), - ("deepseek_v3", 1, False), - ("deepseek_v4", 2, False), - ("qwen3_5_moe", 1, True), - ("qwen3_5_moe", 2, False), - ("llama", 1, False), + (1, True, True), + (1, False, False), + (2, True, False), ], ) -def test_dsv4_adp_dummy_fix_gate(model_type, pp_size, expected): +def test_non_overlap_adp_forward_intent_scope(pp_size, disable_overlap, 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 + assert should_enable_non_overlap_adp_forward_intent(mapping, disable_overlap) is expected @pytest.mark.parametrize(