From ebf3e65857f8a5d74ab0a8350097fa44e2171004 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Mon, 27 Jul 2026 17:10:23 -0700 Subject: [PATCH 1/4] [NVBUG 6487039][fix] Generalize ADP dummy lifecycle Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/_util.py | 10 +++---- .../_torch/pyexecutor/model_engine.py | 11 ++++--- tensorrt_llm/_torch/pyexecutor/py_executor.py | 24 +++++++-------- .../_torch/executor/test_benchmark_disagg.py | 2 +- .../_torch/executor/test_py_executor.py | 29 +++++-------------- .../_torch/executor/test_seq_slot_sizing.py | 15 +++------- 6 files changed, 34 insertions(+), 57 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 9da5db99c10d..49e9d2b2345b 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2637,17 +2637,17 @@ def compute_max_num_sequences(mapping: Mapping, 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_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_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) + 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) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 17b2dc5a6c2a..47bc7bca4bec 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -359,7 +359,7 @@ 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_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) @@ -446,11 +446,10 @@ 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_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 2fa5a5d0607b..d4a4627bf5bd 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -569,8 +569,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_adp_dummy_fixes = getattr(model_engine, + "_enable_adp_dummy_fixes", False) self.enable_attention_dp = model_engine.enable_attention_dp self.dist = dist self.sampler = sampler @@ -3344,7 +3344,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 @@ -5765,16 +5765,15 @@ 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 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. 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_adp_dummy_fixes or self.kv_cache_transceiver is None): if self.kv_cache_transceiver is None: return len(self.active_requests) @@ -5910,7 +5909,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, @@ -5945,9 +5944,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/tests/unittest/_torch/executor/test_benchmark_disagg.py b/tests/unittest/_torch/executor/test_benchmark_disagg.py index b8cb399b1203..15da7c99e6a2 100644 --- a/tests/unittest/_torch/executor/test_benchmark_disagg.py +++ b/tests/unittest/_torch/executor/test_benchmark_disagg.py @@ -588,7 +588,7 @@ 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.max_num_tokens = None self.dist = Mock() diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index c9c05b39137f..e2671eba8d58 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -1306,7 +1306,7 @@ 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, ): self.enable_attention_dp = enable_attention_dp self.kv_cache_transceiver = kv_cache_transceiver @@ -1320,7 +1320,7 @@ 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.add_dummy_calls = [] self.model_engine = Mock(max_num_tokens=max_num_tokens, max_seq_len=max_seq_len) @@ -1446,9 +1446,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 @@ -1506,20 +1506,7 @@ 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 - ) - stub.kv_cache_manager.add_dummy_requests.assert_not_called() - - -def test_dsv4_pad_dummy_checks_full_context_capacity(): +def test_adp_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 @@ -1533,7 +1520,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 @@ -1546,7 +1533,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 diff --git a/tests/unittest/_torch/executor/test_seq_slot_sizing.py b/tests/unittest/_torch/executor/test_seq_slot_sizing.py index d42e6f4483c9..12601c7c1bb8 100644 --- a/tests/unittest/_torch/executor/test_seq_slot_sizing.py +++ b/tests/unittest/_torch/executor/test_seq_slot_sizing.py @@ -22,7 +22,7 @@ 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, ) from tensorrt_llm.mapping import Mapping @@ -65,17 +65,10 @@ def test_dsv4_overlap_headroom_gate( ) -@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): +@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_dsv4_adp_dummy_fixes(model_type, mapping) is expected + assert should_enable_adp_dummy_fixes(mapping) is expected @pytest.mark.parametrize( From 0dbfbf1c8a43820b1b78e2537c36b43779bcbd4a Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 28 Jul 2026 10:09:51 -0700 Subject: [PATCH 2/4] [NVBUG 6487039][test] Cover mixed-rank ADP padding Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../_torch/executor/test_py_executor.py | 36 +++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index e2671eba8d58..fc5295153896 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -1489,6 +1489,42 @@ def test_pad_dummy_added_when_only_wait_scheduler_requests_disagg(): assert len(stub.active_requests) == 2 +def test_non_dsv4_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 From b91ddfea07ff52ffc07b1806c70cd30ccd0ed0ef Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Wed, 29 Jul 2026 14:35:22 -0700 Subject: [PATCH 3/4] [NVBUG-6487039] Generalize ADP overlap lifecycle Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../_torch/models/modeling_qwen2vl.py | 13 +++- .../_torch/models/modeling_qwen3vl.py | 3 +- tensorrt_llm/_torch/pyexecutor/_util.py | 37 +++------- .../_torch/pyexecutor/model_engine.py | 61 ++++++++++------ .../_torch/pyexecutor/model_loader.py | 13 +++- tensorrt_llm/_torch/pyexecutor/py_executor.py | 16 ++-- .../_torch/pyexecutor/py_executor_creator.py | 8 +- .../_torch/pyexecutor/scheduler/scheduler.py | 38 +++++++++- .../pyexecutor/scheduler/scheduler_v2.py | 8 ++ tensorrt_llm/_torch/speculative/interface.py | 12 +-- tensorrt_llm/_torch/speculative/utils.py | 2 +- .../_torch/executor/test_benchmark_disagg.py | 15 +++- .../executor/test_dual_pool_kv_cache.py | 27 +++++++ .../_torch/executor/test_model_loader_gms.py | 10 +++ .../_torch/executor/test_py_executor.py | 43 +++++++++++ .../executor/test_pytorch_model_engine.py | 26 +++++-- .../_torch/executor/test_seq_slot_sizing.py | 73 +++++-------------- .../modeling/test_modeling_qwen2_5vl.py | 10 ++- .../test_rejection_buffers_guard.py | 2 +- 19 files changed, 278 insertions(+), 139 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py index d9cb0dac76fe..c365b1899457 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py @@ -200,6 +200,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): @@ -1757,8 +1766,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 e1d73297336f..a14daa36510a 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3vl.py @@ -58,6 +58,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, ) @@ -1227,7 +1228,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/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 49e9d2b2345b..375497bf70f7 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2616,24 +2616,18 @@ def create_kv_cache_compression_manager( return None -def compute_max_num_sequences(mapping: Mapping, - max_batch_size: int, - disable_overlap_scheduler: bool, - enable_overlap_headroom: bool = False) -> int: +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). - ``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``. + The overlap scheduler needs a second non-PP slot set because it can + backfill seats before releasing the previous iteration's terminal slots. + Pipeline parallelism already sizes the pool by ``pp_size``. """ if mapping.has_pp(): num_micro_batches = mapping.pp_size else: - num_micro_batches = (2 if enable_overlap_headroom - and not disable_overlap_scheduler else 1) + num_micro_batches = 1 if disable_overlap_scheduler else 2 return max_batch_size * num_micro_batches @@ -2642,17 +2636,6 @@ def should_enable_adp_dummy_fixes(mapping: Mapping) -> bool: return 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 (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) - - def create_py_executor_instance( *, dist, @@ -2976,8 +2959,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 47bc7bca4bec..b40782bbe966 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -356,11 +356,10 @@ def __init__( self.mapping = mapping if mapping.has_pp(): init_pp_comm(mapping) - # Start with the established pool size. Once the model is loaded we - # selectively enable headroom for the non-PP DeepSeek-V4 overlap path. + # The overlap scheduler can hold two iterations' requests at once. + # Every model-side buffer indexed by py_seq_slot must span this pool. from ._util import (compute_max_num_sequences, - should_enable_adp_dummy_fixes, - should_enable_dsv4_overlap_headroom) + should_enable_adp_dummy_fixes) self.max_num_seq_slots = compute_max_num_sequences( mapping, self.batch_size, llm_args.disable_overlap_scheduler) self.dist = dist @@ -434,6 +433,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, @@ -444,22 +444,11 @@ def __init__( setattr(self, "moe_load_balancer", moe_load_balancer) else: self.model = model - pretrained_config = self.model.model_config.pretrained_config - model_type = getattr(pretrained_config, "model_type", None) + self._validate_mrope_position_delta_cache_capacity() # 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. + # model. Sequence-slot capacity follows the independent overlap + # lifecycle invariant above. self._enable_adp_dummy_fixes = should_enable_adp_dummy_fixes(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 @@ -925,6 +914,33 @@ def set_guided_decoder(self, return success return False + 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 indexing past an undersized 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.") + @property def use_mrope(self): use_mrope = False @@ -2705,11 +2721,8 @@ 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 - # 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) + # Slot-indexed metadata must span the same pool as SeqSlotManager. + num_seq_slots = self.max_num_seq_slots if no_cache: return get_spec_metadata( spec_config, @@ -4092,7 +4105,7 @@ def _prepare_tp_inputs( # that carry no MRoPE metadata at all. The cache is zero-initialized and # the write path only ever targets real ``py_seq_slot``s, so this slot # permanently reads back a zero delta. - mrope_dummy_seq_slot = self.max_num_tokens * self.mapping.pp_size + mrope_dummy_seq_slot = self.max_num_seq_slots 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 2a6ea23076aa..5f7c01ca032d 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -363,7 +363,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. @@ -379,6 +380,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 @@ -386,6 +390,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 @@ -393,6 +398,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, @@ -1402,6 +1412,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 d4a4627bf5bd..1fd380ac5301 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -5765,10 +5765,10 @@ 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 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. 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. @@ -5783,12 +5783,8 @@ def _count_schedulable_active_requests(self) -> int: 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 schedule_from_value <= req.state_value < to_complete_value) + return sum(1 for req in self.active_requests + if self.scheduler.is_request_in_schedulable_state(req)) def _has_adp_dummy_kv_capacity(self, token_nums: Optional[List[int]]) -> bool: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 35f3ef5e6a42..c28f0b0acb48 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -784,14 +784,10 @@ 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, - # The scoped DeepSeek-V4 path follows the expanded slot - # pool. Other configurations retain max_batch_size. - "max_num_sequences": guided_decoder_slots, + # Guided-decoder state is indexed by sequence slot. + "max_num_sequences": max_num_seq_slots, "vocab_size_padded": model_engine.model.vocab_size_padded, "rank": mapping.rank, } diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py index caa8e3cb3de1..7e9b68d94cd4 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py @@ -216,6 +216,19 @@ 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 + 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 +405,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 +420,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 +449,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 +1896,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/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index bdb122ac4457..8897c3b0a09b 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -564,8 +564,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); DeepSeek-V4 overlap can use 2 * max_batch_size, - # larger than max_num_requests (== max_batch_size). + # [0, num_seq_slots); 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 @@ -614,10 +614,10 @@ def prepare_rejection_sampling_buffers(self): return # Slot-indexed buffers span the full SeqSlotManager pool: py_seq_slot - # can range over [0, num_seq_slots), which under DeepSeek-V4 overlap - # exceeds max_num_requests. Fall back to max_num_requests when the pool - # size is unknown (0). One extra scratch row at index ``slot_capacity`` - # absorbs CUDA-graph dummy/padding requests (``py_seq_slot is None``). + # can range over [0, num_seq_slots), which under overlap exceeds + # max_num_requests. Fall back to max_num_requests when the pool size is + # unknown (0). One extra scratch row at index ``slot_capacity`` absorbs + # CUDA-graph dummy/padding requests (``py_seq_slot is None``). slot_capacity = self.num_seq_slots or self.max_num_requests num_slot_rows = slot_capacity + 1 diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index a8d92fed4f46..50c081483933 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -92,7 +92,7 @@ def get_spec_metadata(spec_config, use_rejection_sampling = getattr(spec_config, "use_rejection_sampling", False) # Slot-indexed buffers (draft_probs) must span the SeqSlotManager pool; - # DeepSeek-V4 overlap can exceed max_num_requests. + # Overlap can make the sequence-slot pool 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) diff --git a/tests/unittest/_torch/executor/test_benchmark_disagg.py b/tests/unittest/_torch/executor/test_benchmark_disagg.py index 15da7c99e6a2..4a4000e92360 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 @@ -55,6 +55,8 @@ def _make_active_request( LlmRequestState.DISAGG_TRANS_ERROR if in_error else LlmRequestState.GENERATION_IN_PROGRESS ) req.is_attention_dp_dummy = False + req.is_context_init_state = False + req.py_encoder_output_ready_event = None return req @@ -595,6 +597,17 @@ def __init__( self.dist.tp_size = tp_size self.dist.tp_allgather.side_effect = lambda value: [value] + 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 self.kv_cache_manager.get_num_available_tokens.return_value = 1 << 30 diff --git a/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py b/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py index 06e0231a1a0b..ef516cb704a0 100644 --- a/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py +++ b/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py @@ -866,6 +866,29 @@ def test_cross_kv_cache_manager_and_until_state_are_forwarded(self): impl.assert_called_once_with([], kv_mgr, None, cross_mgr) +class TestBindMicroBatchSchedulerStateRange: + """C++-bound micro-batch scheduling exposes its configured state range.""" + + def test_encoder_state_range_is_forwarded(self): + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState + from tensorrt_llm._torch.pyexecutor.scheduler.scheduler import BindMicroBatchScheduler + + with patch( + "tensorrt_llm._torch.pyexecutor.scheduler.scheduler.tb_internal.algorithms.MicroBatchScheduler" + ) as micro_cls: + micro_cls.return_value = Mock() + scheduler = BindMicroBatchScheduler( + max_batch_size=8, + max_num_tokens=4096, + no_schedule_until_state=LlmRequestState.ENCODER_INIT, + ) + + kwargs = micro_cls.call_args.kwargs + assert kwargs["no_schedule_until_state"] == LlmRequestState.ENCODER_INIT + assert kwargs["no_schedule_after_state"] == LlmRequestState.GENERATION_TO_COMPLETE + assert scheduler.no_schedule_until_state == LlmRequestState.ENCODER_INIT + + class TestSimpleUnifiedSchedulerCrossParam: """V1 Python ``SimpleUnifiedScheduler`` exposes cross-KV wiring.""" @@ -895,6 +918,10 @@ def test_cross_kv_cache_manager_and_until_state_are_forwarded(self): assert ( scheduler.micro_batch_scheduler.no_schedule_until_state == LlmRequestState.ENCODER_INIT ) + assert scheduler.scheduling_state_range == ( + LlmRequestState.ENCODER_INIT, + LlmRequestState.GENERATION_TO_COMPLETE, + ) # --------------------------------------------------------------------------- diff --git a/tests/unittest/_torch/executor/test_model_loader_gms.py b/tests/unittest/_torch/executor/test_model_loader_gms.py index 3d073d95bf90..566bf38beb39 100644 --- a/tests/unittest/_torch/executor/test_model_loader_gms.py +++ b/tests/unittest/_torch/executor/test_model_loader_gms.py @@ -159,6 +159,16 @@ def _build_source_identity(_cls, *_args, **kwargs): return loader +def test_runtime_model_config_attrs_include_sequence_slot_capacity(): + loader = object.__new__(ModelLoader) + loader.max_num_seq_slots = 16 + config = SimpleNamespace(extra_attrs={}) + + loader._set_runtime_model_config_attrs(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 fc5295153896..9f04b0c28aa9 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -30,6 +30,7 @@ from tensorrt_llm._torch.pyexecutor.resource_manager import NoFreeSlotsError, ResourceManagerType from tensorrt_llm._torch.pyexecutor.scheduler import ( FCFSWaitingQueue, + RequestScheduler, ScheduledRequests, SerializableSchedulerOutput, ) @@ -1292,6 +1293,8 @@ 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 @@ -1328,6 +1331,17 @@ def __init__( self.dist.tp_size = 1 self.dist.tp_allgather.side_effect = lambda value: [value] + 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 kv_cache_manager.get_num_available_tokens.return_value = 1 << 30 @@ -1489,6 +1503,35 @@ def test_pad_dummy_added_when_only_wait_scheduler_requests_disagg(): assert len(stub.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 + + +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_non_dsv4_disagg_adp_mixed_rank_states_stay_queueable(): # The generic non-PP path must give both ranks a non-empty scheduled batch: # one rank schedules its real request, while the terminal-only rank diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 9a1242441523..a2a0925a9927 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -1699,6 +1699,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') @@ -1730,6 +1731,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 = [] @@ -1752,10 +1756,10 @@ def test_prepare_tp_inputs_with_partial_mrope_segments(self) -> None: [0]) # Read slots are dense w.r.t. the generation batch: the padded dummy # has no MRoPE metadata, so it resolves to the reserved zero slot - # (max_num_tokens * pp_size) rather than being dropped, which would + # (max_num_seq_slots) rather than being dropped, which would # shift every later request onto another request's delta. self.assertEqual(result["mrope_delta_read_seq_slots"].cpu().tolist(), - [0, 32]) + [0, model_engine.max_num_seq_slots]) self.assertNotIn("multimodal_embedding", multimodal_request.py_multimodal_data) kv_cache_manager.shutdown() @@ -1850,11 +1854,11 @@ def test_prepare_tp_inputs_mixed_text_only_keeps_mrope_deltas_dense( kv_cache_manager=kv_cache_manager, attn_metadata=attn_metadata) - # One entry per generation request, in batch order. Slot 32 is the - # reserved zero slot (max_num_tokens * pp_size) standing in for the - # text-only request's zero delta. + # One entry per generation request, in batch order. The reserved zero + # slot (max_num_seq_slots) stands in for the text-only request's zero + # delta. self.assertEqual(result["mrope_delta_read_seq_slots"].cpu().tolist(), - [0, 32, 2]) + [0, model_engine.max_num_seq_slots, 2]) # Only the two multimodal requests seed the seq-slot delta cache. self.assertEqual(result["mrope_delta_write_seq_slots"].cpu().tolist(), [0, 2]) @@ -1960,6 +1964,16 @@ def test_promoted_mrope_context_uses_decode_state_contract(self) -> None: self.assertEqual(model_engine.previous_request_ids, []) 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) -> None: """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 12601c7c1bb8..2ffced3e42c4 100644 --- a/tests/unittest/_torch/executor/test_seq_slot_sizing.py +++ b/tests/unittest/_torch/executor/test_seq_slot_sizing.py @@ -1,81 +1,48 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""DeepSeek-V4 seq-slot sizing includes overlap headroom. +"""Seq-slot pool and slot-indexed state 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. The headroom is intentionally limited to DeepSeek-V4; -other models preserve their established sizing pending separate validation. +2 * max_batch_size for every non-PP overlap configuration. 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_adp_dummy_fixes, - should_enable_dsv4_overlap_headroom, ) from tensorrt_llm.mapping import Mapping SIZING_CASES = [ - # (pp_size, disable_overlap, enable_overlap_headroom, expected_factor) - (1, False, True, 2), - (1, False, False, 1), - (1, True, True, 1), - # Existing PP sizing is preserved regardless of the DSv4 opt-in. - (2, False, True, 2), - (4, False, True, 4), - (4, True, False, 4), + # (pp_size, disable_overlap, expected_factor) + (1, False, 2), + (1, True, 1), + # PP already sizes the pool for its micro-batch count. + (2, False, 2), + (4, False, 4), + (4, True, 4), ] -@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 -): - 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 ( - should_enable_dsv4_overlap_headroom(model_type, spec_config, mapping, disable_overlap) - is expected - ) - - @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( - "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 +@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) @@ -84,26 +51,26 @@ def test_compute_max_num_sequences_scopes_overlap_headroom( 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): +@pytest.mark.parametrize("pp_size,disable_overlap,expected_factor", SIZING_CASES) +def test_sampler_uses_executor_slot_pool_capacity(pp_size, disable_overlap, expected_factor): max_batch_size = 8 - mapping = Mapping(world_size=1, tp_size=1, pp_size=1) - max_num_sequences = max_batch_size * slot_factor + 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=False, + disable_overlap_scheduler=disable_overlap, enable_async_worker=False, enable_speculative_beam_history_d2h=False, - max_num_sequences=max_num_sequences, ) - assert args.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 diff --git a/tests/unittest/_torch/modeling/test_modeling_qwen2_5vl.py b/tests/unittest/_torch/modeling/test_modeling_qwen2_5vl.py index 33dbccab0e25..f7fff190859c 100644 --- a/tests/unittest/_torch/modeling/test_modeling_qwen2_5vl.py +++ b/tests/unittest/_torch/modeling/test_modeling_qwen2_5vl.py @@ -25,7 +25,8 @@ Qwen2VLHfWeightMapper from tensorrt_llm._torch.models.modeling_qwen2vl import ( Qwen2_5_VisionModel, Qwen2_5_VLModel, Qwen2VisionModelBase, - Qwen2VLInputProcessorBase, Qwen2VLModel, _prepare_qwen_vl_mrope_config, + Qwen2VLInputProcessorBase, Qwen2VLModel, + _get_mrope_position_delta_cache_size, _prepare_qwen_vl_mrope_config, _prepare_qwen_vl_vision_attn_metadata) from tensorrt_llm._torch.models.modeling_qwen3vl import \ Qwen3VLInputProcessorBase @@ -428,6 +429,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 diff --git a/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py b/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py index a71b9f1e6640..f24fce926b65 100644 --- a/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py +++ b/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py @@ -67,7 +67,7 @@ def test_prepare_buffers_allocates_full_draft_probs_on_vocab_mismatch(): def test_prepare_buffers_span_seq_slot_pool(): - # Under DeepSeek-V4 overlap scheduling the SeqSlotManager pool + # Under overlap scheduling the SeqSlotManager pool # (num_seq_slots) can exceed max_num_requests; py_seq_slot then indexes past # max_num_requests. Slot-indexed buffers must span the full pool plus the # dummy scratch row, and dummy_slot_row must land on that last row so a real From 6408929bc1317a724cde2b28e7b4a3473ce4d7c6 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Thu, 6 Aug 2026 17:04:47 -0700 Subject: [PATCH 4/4] [NVBUG 6487039][fix] Narrow ADP dummy lifecycle scope Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../_torch/models/modeling_qwen2vl.py | 13 +- .../_torch/models/modeling_qwen3vl.py | 3 +- tensorrt_llm/_torch/pyexecutor/_util.py | 44 ++++++- .../_torch/pyexecutor/model_engine.py | 69 +++++----- .../_torch/pyexecutor/model_loader.py | 13 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 64 ++++++++- .../_torch/pyexecutor/py_executor_creator.py | 8 +- .../_torch/pyexecutor/scheduler/scheduler.py | 6 + tensorrt_llm/_torch/speculative/interface.py | 12 +- tensorrt_llm/_torch/speculative/utils.py | 2 +- .../_torch/executor/test_benchmark_disagg.py | 25 +++- .../executor/test_dual_pool_kv_cache.py | 27 ---- .../_torch/executor/test_model_loader_gms.py | 10 -- .../_torch/executor/test_py_executor.py | 121 +++++++++++++++--- .../executor/test_pytorch_model_engine.py | 26 +--- .../_torch/executor/test_seq_slot_sizing.py | 102 ++++++++++++--- .../modeling/test_modeling_qwen2_5vl.py | 10 +- .../test_rejection_buffers_guard.py | 2 +- 18 files changed, 362 insertions(+), 195 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py index c365b1899457..d9cb0dac76fe 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py @@ -200,15 +200,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): @@ -1766,8 +1757,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 a14daa36510a..e1d73297336f 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3vl.py @@ -58,7 +58,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, ) @@ -1228,7 +1227,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 375497bf70f7..8467e330273d 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2616,18 +2616,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: +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 overlap scheduler needs a second non-PP slot set because it can - backfill seats before releasing the previous iteration's terminal slots. - Pipeline parallelism already sizes the pool by ``pp_size``. + ``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 @@ -2636,6 +2642,32 @@ def should_enable_adp_dummy_fixes(mapping: Mapping) -> bool: return not mapping.has_pp() +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 == "deepseek_v4")) + + +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.""" + 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) + + def create_py_executor_instance( *, dist, diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index b40782bbe966..c5da4a7208c5 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -356,10 +356,13 @@ def __init__( self.mapping = mapping if mapping.has_pp(): init_pp_comm(mapping) - # The overlap scheduler can hold two iterations' requests at once. - # Every model-side buffer indexed by py_seq_slot must span this pool. + # 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_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) self.max_num_seq_slots = compute_max_num_sequences( mapping, self.batch_size, llm_args.disable_overlap_scheduler) self.dist = dist @@ -433,7 +436,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, @@ -444,11 +446,28 @@ 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) # Apply transactional dummy handling to every non-PP disaggregated ADP - # model. Sequence-slot capacity follows the independent overlap - # lifecycle invariant above. + # 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, + 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 @@ -914,33 +933,6 @@ def set_guided_decoder(self, return success return False - 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 indexing past an undersized 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.") - @property def use_mrope(self): use_mrope = False @@ -2721,8 +2713,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 must span the same pool as SeqSlotManager. - 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, @@ -4105,7 +4100,7 @@ def _prepare_tp_inputs( # that carry no MRoPE metadata at all. The cache is zero-initialized and # the write path only ever targets real ``py_seq_slot``s, so this slot # permanently reads back a zero delta. - mrope_dummy_seq_slot = self.max_num_seq_slots + 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 5f7c01ca032d..2a6ea23076aa 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -363,8 +363,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. @@ -380,9 +379,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 @@ -390,7 +386,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 @@ -398,11 +393,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, @@ -1412,7 +1402,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 1fd380ac5301..1304b0b35501 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 @@ -571,6 +578,10 @@ def __init__( self.model_engine = model_engine 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 @@ -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) @@ -5773,11 +5785,15 @@ def _count_schedulable_active_requests(self) -> int: Returns: The number of active requests eligible for scheduling. """ - if (not self._enable_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 @@ -5786,6 +5802,32 @@ def _count_schedulable_active_requests(self) -> int: return sum(1 for req in self.active_requests if self.scheduler.is_request_in_schedulable_state(req)) + 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: """Check the full dummy allocation before entering rank-local code. @@ -5871,8 +5913,18 @@ def _pad_attention_dp_dummy_request(self): if self._should_skip_dummy_for_benchmark_disagg(num_active_request): return - needs_dummy = (self.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 = (self.expected_num_active_requests > 0 + and num_active_request == 0) if not needs_dummy: return diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index c28f0b0acb48..35f3ef5e6a42 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -784,10 +784,14 @@ 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. - "max_num_sequences": max_num_seq_slots, + # 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, } diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py index 7e9b68d94cd4..05291aebce2b 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py @@ -226,6 +226,12 @@ 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 diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 8897c3b0a09b..bdb122ac4457 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -564,8 +564,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); overlap can use 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 @@ -614,10 +614,10 @@ def prepare_rejection_sampling_buffers(self): return # Slot-indexed buffers span the full SeqSlotManager pool: py_seq_slot - # can range over [0, num_seq_slots), which under overlap exceeds - # max_num_requests. Fall back to max_num_requests when the pool size is - # unknown (0). One extra scratch row at index ``slot_capacity`` absorbs - # CUDA-graph dummy/padding requests (``py_seq_slot is None``). + # can range over [0, num_seq_slots), which under DeepSeek-V4 overlap + # exceeds max_num_requests. Fall back to max_num_requests when the pool + # size is unknown (0). One extra scratch row at index ``slot_capacity`` + # absorbs CUDA-graph dummy/padding requests (``py_seq_slot is None``). slot_capacity = self.num_seq_slots or self.max_num_requests num_slot_rows = slot_capacity + 1 diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 50c081483933..a8d92fed4f46 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -92,7 +92,7 @@ def get_spec_metadata(spec_config, use_rejection_sampling = getattr(spec_config, "use_rejection_sampling", False) # Slot-indexed buffers (draft_probs) must span the SeqSlotManager pool; - # Overlap can make the sequence-slot pool exceed max_num_requests. + # 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) diff --git a/tests/unittest/_torch/executor/test_benchmark_disagg.py b/tests/unittest/_torch/executor/test_benchmark_disagg.py index 4a4000e92360..ac23d82ef21f 100644 --- a/tests/unittest/_torch/executor/test_benchmark_disagg.py +++ b/tests/unittest/_torch/executor/test_benchmark_disagg.py @@ -48,14 +48,21 @@ 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 @@ -591,11 +598,18 @@ def __init__( self._adp_dummy_is_gen = True self._pending_adp_dummy_request = None 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 = ( @@ -618,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_dual_pool_kv_cache.py b/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py index ef516cb704a0..06e0231a1a0b 100644 --- a/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py +++ b/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py @@ -866,29 +866,6 @@ def test_cross_kv_cache_manager_and_until_state_are_forwarded(self): impl.assert_called_once_with([], kv_mgr, None, cross_mgr) -class TestBindMicroBatchSchedulerStateRange: - """C++-bound micro-batch scheduling exposes its configured state range.""" - - def test_encoder_state_range_is_forwarded(self): - from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState - from tensorrt_llm._torch.pyexecutor.scheduler.scheduler import BindMicroBatchScheduler - - with patch( - "tensorrt_llm._torch.pyexecutor.scheduler.scheduler.tb_internal.algorithms.MicroBatchScheduler" - ) as micro_cls: - micro_cls.return_value = Mock() - scheduler = BindMicroBatchScheduler( - max_batch_size=8, - max_num_tokens=4096, - no_schedule_until_state=LlmRequestState.ENCODER_INIT, - ) - - kwargs = micro_cls.call_args.kwargs - assert kwargs["no_schedule_until_state"] == LlmRequestState.ENCODER_INIT - assert kwargs["no_schedule_after_state"] == LlmRequestState.GENERATION_TO_COMPLETE - assert scheduler.no_schedule_until_state == LlmRequestState.ENCODER_INIT - - class TestSimpleUnifiedSchedulerCrossParam: """V1 Python ``SimpleUnifiedScheduler`` exposes cross-KV wiring.""" @@ -918,10 +895,6 @@ def test_cross_kv_cache_manager_and_until_state_are_forwarded(self): assert ( scheduler.micro_batch_scheduler.no_schedule_until_state == LlmRequestState.ENCODER_INIT ) - assert scheduler.scheduling_state_range == ( - LlmRequestState.ENCODER_INIT, - LlmRequestState.GENERATION_TO_COMPLETE, - ) # --------------------------------------------------------------------------- diff --git a/tests/unittest/_torch/executor/test_model_loader_gms.py b/tests/unittest/_torch/executor/test_model_loader_gms.py index 566bf38beb39..3d073d95bf90 100644 --- a/tests/unittest/_torch/executor/test_model_loader_gms.py +++ b/tests/unittest/_torch/executor/test_model_loader_gms.py @@ -159,16 +159,6 @@ def _build_source_identity(_cls, *_args, **kwargs): return loader -def test_runtime_model_config_attrs_include_sequence_slot_capacity(): - loader = object.__new__(ModelLoader) - loader.max_num_seq_slots = 16 - config = SimpleNamespace(extra_attrs={}) - - loader._set_runtime_model_config_attrs(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 9f04b0c28aa9..a64b4f50dae7 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,7 +25,11 @@ 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, @@ -1293,7 +1296,9 @@ 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 @@ -1310,6 +1315,9 @@ def __init__( is_warmup=False, benchmark_req_queues_size=0, 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 @@ -1324,12 +1332,23 @@ def __init__( self._adp_dummy_is_gen = True self._pending_adp_dummy_request = None 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 = ( @@ -1370,6 +1389,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", ): @@ -1474,9 +1494,9 @@ def test_disabled_adp_dummy_fix_gate_preserves_pp_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 @@ -1491,8 +1511,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 @@ -1518,6 +1538,29 @@ def test_encoder_init_uses_encoder_decoder_scheduler_state_window(): 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) @@ -1532,7 +1575,7 @@ def test_decoder_context_waiting_for_encoder_output_is_not_counted(): assert len(stub.active_requests) == 2 -def test_non_dsv4_disagg_adp_mixed_rank_states_stay_queueable(): +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. @@ -1586,8 +1629,10 @@ def test_pad_dummy_allocation_failure_skips_padding(): def test_adp_pad_dummy_checks_full_context_capacity(): - stub = _StubADPExecutor(max_num_tokens=4096) - stub._adp_dummy_is_gen = False + stub = _StubADPExecutor( + max_num_tokens=4096, + peer_forward_intent=_ADPForwardIntent.CONTEXT, + ) stub.kv_cache_manager.get_num_available_tokens.return_value = 1024 _run_pad(stub) @@ -1728,8 +1773,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) @@ -1753,23 +1800,57 @@ def test_pad_dummy_gen_keeps_default_token_nums(): assert call["is_gen"] is True -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_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index a2a0925a9927..9a1242441523 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -1699,7 +1699,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') @@ -1731,9 +1730,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 = [] @@ -1756,10 +1752,10 @@ def test_prepare_tp_inputs_with_partial_mrope_segments(self) -> None: [0]) # Read slots are dense w.r.t. the generation batch: the padded dummy # has no MRoPE metadata, so it resolves to the reserved zero slot - # (max_num_seq_slots) rather than being dropped, which would + # (max_num_tokens * pp_size) rather than being dropped, which would # shift every later request onto another request's delta. self.assertEqual(result["mrope_delta_read_seq_slots"].cpu().tolist(), - [0, model_engine.max_num_seq_slots]) + [0, 32]) self.assertNotIn("multimodal_embedding", multimodal_request.py_multimodal_data) kv_cache_manager.shutdown() @@ -1854,11 +1850,11 @@ def test_prepare_tp_inputs_mixed_text_only_keeps_mrope_deltas_dense( kv_cache_manager=kv_cache_manager, attn_metadata=attn_metadata) - # One entry per generation request, in batch order. The reserved zero - # slot (max_num_seq_slots) stands in for the text-only request's zero - # delta. + # One entry per generation request, in batch order. Slot 32 is the + # reserved zero slot (max_num_tokens * pp_size) standing in for the + # text-only request's zero delta. self.assertEqual(result["mrope_delta_read_seq_slots"].cpu().tolist(), - [0, model_engine.max_num_seq_slots, 2]) + [0, 32, 2]) # Only the two multimodal requests seed the seq-slot delta cache. self.assertEqual(result["mrope_delta_write_seq_slots"].cpu().tolist(), [0, 2]) @@ -1964,16 +1960,6 @@ def test_promoted_mrope_context_uses_decode_state_contract(self) -> None: self.assertEqual(model_engine.previous_request_ids, []) 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) -> None: """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 2ffced3e42c4..3dc88156c210 100644 --- a/tests/unittest/_torch/executor/test_seq_slot_sizing.py +++ b/tests/unittest/_torch/executor/test_seq_slot_sizing.py @@ -1,48 +1,110 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Seq-slot pool and slot-indexed state 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 for every non-PP overlap configuration. +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_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 SIZING_CASES = [ - # (pp_size, disable_overlap, expected_factor) - (1, False, 2), - (1, True, 1), - # PP already sizes the pool for its micro-batch count. - (2, False, 2), - (4, False, 4), - (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( + "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 +): + 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 ( + should_enable_dsv4_overlap_headroom(model_type, spec_config, mapping, disable_overlap) + is expected + ) + + @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("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,pp_size,disable_overlap,expected", + [ + ("kimi_k2", 1, True, True), + ("kimi_k2", 1, False, False), + ("deepseek_v4", 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( + "pp_size,disable_overlap,expected", + [ + (1, True, True), + (1, False, False), + (2, True, False), + ], +) +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_non_overlap_adp_forward_intent(mapping, disable_overlap) 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) @@ -51,26 +113,26 @@ def test_compute_max_num_sequences_includes_overlap_headroom( mapping, max_batch_size, disable_overlap, + enable_overlap_headroom=enable_overlap_headroom, ) == max_batch_size * expected_factor ) -@pytest.mark.parametrize("pp_size,disable_overlap,expected_factor", SIZING_CASES) -def test_sampler_uses_executor_slot_pool_capacity(pp_size, disable_overlap, 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=pp_size, tp_size=1, pp_size=pp_size) + 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 f7fff190859c..33dbccab0e25 100644 --- a/tests/unittest/_torch/modeling/test_modeling_qwen2_5vl.py +++ b/tests/unittest/_torch/modeling/test_modeling_qwen2_5vl.py @@ -25,8 +25,7 @@ Qwen2VLHfWeightMapper from tensorrt_llm._torch.models.modeling_qwen2vl import ( Qwen2_5_VisionModel, Qwen2_5_VLModel, Qwen2VisionModelBase, - Qwen2VLInputProcessorBase, Qwen2VLModel, - _get_mrope_position_delta_cache_size, _prepare_qwen_vl_mrope_config, + Qwen2VLInputProcessorBase, Qwen2VLModel, _prepare_qwen_vl_mrope_config, _prepare_qwen_vl_vision_attn_metadata) from tensorrt_llm._torch.models.modeling_qwen3vl import \ Qwen3VLInputProcessorBase @@ -429,13 +428,6 @@ 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 diff --git a/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py b/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py index f24fce926b65..a71b9f1e6640 100644 --- a/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py +++ b/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py @@ -67,7 +67,7 @@ def test_prepare_buffers_allocates_full_draft_probs_on_vocab_mismatch(): def test_prepare_buffers_span_seq_slot_pool(): - # Under overlap scheduling the SeqSlotManager pool + # Under DeepSeek-V4 overlap scheduling the SeqSlotManager pool # (num_seq_slots) can exceed max_num_requests; py_seq_slot then indexes past # max_num_requests. Slot-indexed buffers must span the full pool plus the # dummy scratch row, and dummy_slot_row must land on that last row so a real