From 1ecb5973e374ee4e118f6e614c8660917e9c7d36 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/5] [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 | 14 ++------- .../_torch/pyexecutor/model_engine.py | 9 ++---- 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 | 18 +++--------- 6 files changed, 30 insertions(+), 66 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index c4fb115111fe..f79c46414480 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2637,17 +2637,9 @@ def compute_max_num_sequences(mapping: Mapping, return max_batch_size * num_micro_batches -# Model types whose disaggregated attention-DP path has been measured against -# the ADP dummy fixes. The gate stays an explicit list rather than a capability -# check (``enable_attention_dp and kv_cache_transceiver is not None``) so that -# each entry is added only after its disagg ADP behavior has been exercised. -_ADP_DUMMY_FIX_MODEL_TYPES = ("deepseek_v4", "qwen3_5_moe") - - -def should_enable_dsv4_adp_dummy_fixes(model_type: Optional[str], - mapping: Mapping) -> bool: - """Gate the ADP dummy fixes while PP remains follow-up scope.""" - return model_type in _ADP_DUMMY_FIX_MODEL_TYPES and not mapping.has_pp() +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_disagg_adp_overlap_headroom( diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index bd832625d042..2a416e8d05c3 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -391,12 +391,13 @@ def __init__( # Disaggregated attention-DP can backfill a batch before the overlap # scheduler releases the previous batch's terminal sequence slots. from ._util import (compute_max_num_sequences, - should_enable_disagg_adp_overlap_headroom, - should_enable_dsv4_adp_dummy_fixes) + should_enable_adp_dummy_fixes, + should_enable_disagg_adp_overlap_headroom) self._enable_disagg_adp_overlap_headroom = ( should_enable_disagg_adp_overlap_headroom( mapping, llm_args.cache_transceiver_config, llm_args.disable_overlap_scheduler)) + self._enable_adp_dummy_fixes = should_enable_adp_dummy_fixes(mapping) self.max_num_seq_slots = compute_max_num_sequences( mapping, self.batch_size, @@ -484,10 +485,6 @@ 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._enable_dsv4_adp_dummy_fixes = should_enable_dsv4_adp_dummy_fixes( - model_type, mapping) if drafting_loop_wrapper is not None: self.model = drafting_loop_wrapper(self.model) self.model_is_wrapped = True diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 62b7c6488ee9..1fb337f07bf9 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -584,8 +584,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 @@ -3376,7 +3376,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 @@ -6044,16 +6044,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) @@ -6216,7 +6215,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, @@ -6251,9 +6250,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 fa8441426dbf..02c88b115441 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -1582,7 +1582,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 @@ -1596,7 +1596,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) @@ -1722,9 +1722,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 @@ -1817,20 +1817,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 @@ -1844,7 +1831,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 @@ -1857,7 +1844,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 64cfb60dee71..86f81378c32b 100644 --- a/tests/unittest/_torch/executor/test_seq_slot_sizing.py +++ b/tests/unittest/_torch/executor/test_seq_slot_sizing.py @@ -20,8 +20,8 @@ from tensorrt_llm._torch.pyexecutor._util import ( compute_max_num_sequences, create_torch_sampler_args, + should_enable_adp_dummy_fixes, should_enable_disagg_adp_overlap_headroom, - should_enable_dsv4_adp_dummy_fixes, ) from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig from tensorrt_llm.mapping import Mapping @@ -65,20 +65,10 @@ def test_disagg_adp_overlap_headroom_gate( ) -@pytest.mark.parametrize( - "model_type,pp_size,expected", - [ - ("deepseek_v4", 1, True), - ("deepseek_v3", 1, False), - ("deepseek_v4", 2, False), - ("qwen3_5_moe", 1, True), - ("qwen3_5_moe", 2, False), - ("llama", 1, False), - ], -) -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 7144bf3b081999707eb720c2e53cca3080708a5f 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/5] [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 02c88b115441..23a43053cc03 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -1800,6 +1800,42 @@ def test_pad_dummy_still_added_when_surplus_requests_are_unschedulable() -> None assert stub.expected_num_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 4bc577ae0feb916c56a5f8700d3b4186e8a896be 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/5] [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 | 8 +++- .../_torch/pyexecutor/model_engine.py | 31 ++++++++++++- .../_torch/pyexecutor/model_loader.py | 13 +++++- tensorrt_llm/_torch/pyexecutor/py_executor.py | 16 +++---- .../_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 ++++++++--- .../modeling/test_modeling_qwen2_5vl.py | 10 ++++- .../test_rejection_buffers_guard.py | 2 +- 17 files changed, 243 insertions(+), 34 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 f79c46414480..9dd80d34163e 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2976,8 +2976,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 2a416e8d05c3..f05ee430f747 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -475,6 +475,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, @@ -485,6 +486,7 @@ def __init__( setattr(self, "moe_load_balancer", moe_load_balancer) else: self.model = model + self._validate_mrope_position_delta_cache_capacity() if drafting_loop_wrapper is not None: self.model = drafting_loop_wrapper(self.model) self.model_is_wrapped = True @@ -1008,6 +1010,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 @@ -4787,7 +4816,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 7a852bf9bfd1..2cf78d4ab44d 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -374,7 +374,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. @@ -390,6 +391,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 @@ -397,6 +401,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 @@ -404,6 +409,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, @@ -1413,6 +1423,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 1fb337f07bf9..dbc182d23c76 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -6044,10 +6044,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. @@ -6062,12 +6062,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/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 4bb3d6ea8938..5f7d02e3c48a 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 a64f9f7a204f..2f6a2d41719a 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -91,7 +91,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 23a43053cc03..40378b1bd8c2 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -35,6 +35,7 @@ from tensorrt_llm._torch.pyexecutor.resource_manager import NoFreeSlotsError, ResourceManagerType from tensorrt_llm._torch.pyexecutor.scheduler import ( FCFSWaitingQueue, + RequestScheduler, ScheduledRequests, SerializableSchedulerOutput, ) @@ -1568,6 +1569,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 @@ -1604,6 +1607,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 @@ -1800,6 +1814,35 @@ def test_pad_dummy_still_added_when_surplus_requests_are_unschedulable() -> None assert stub.expected_num_active_requests == 2 +def test_encoder_init_uses_encoder_decoder_scheduler_state_window(): + stub = _StubADPExecutor() + stub.scheduler.scheduling_state_range = ( + LlmRequestState.ENCODER_INIT, + LlmRequestState.GENERATION_TO_COMPLETE, + ) + stub.active_requests = [_make_adp_request(LlmRequestState.ENCODER_INIT)] + stub.expected_num_active_requests = 1 + + _run_pad(stub) + + assert stub.add_dummy_calls == [] + assert len(stub.active_requests) == 1 + + +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 03d76b437123..29172201f3b7 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -1844,6 +1844,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') @@ -1875,6 +1876,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 = [] @@ -1897,10 +1901,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() @@ -1995,11 +1999,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]) @@ -2105,6 +2109,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/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 b30e3548e907034f7b779334f9c3dd35b162da2e 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/5] [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 | 19 +++ .../_torch/pyexecutor/model_engine.py | 43 ++----- .../_torch/pyexecutor/model_loader.py | 13 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 64 +++++++++- .../_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 | 116 +++++++++++++++--- .../executor/test_pytorch_model_engine.py | 26 +--- .../_torch/executor/test_seq_slot_sizing.py | 30 +++++ .../modeling/test_modeling_qwen2_5vl.py | 10 +- .../test_rejection_buffers_guard.py | 2 +- 17 files changed, 261 insertions(+), 160 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 9dd80d34163e..f057ab159e2e 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2642,6 +2642,25 @@ def should_enable_adp_dummy_fixes(mapping: Mapping) -> bool: return not mapping.has_pp() +_VALIDATED_OVERLAP_ADP_DUMMY_MODEL_TYPES = ("deepseek_v4", "qwen3_5_moe") + + +def should_enable_scheduler_aware_adp_dummy( + model_type: Optional[str], mapping: Mapping, + disable_overlap_scheduler: bool) -> bool: + """Enable scheduler-aware padding for validated lifecycle configurations.""" + return (should_enable_adp_dummy_fixes(mapping) + and (disable_overlap_scheduler + or model_type in _VALIDATED_OVERLAP_ADP_DUMMY_MODEL_TYPES)) + + +def should_enable_non_overlap_adp_forward_intent( + mapping: Mapping, disable_overlap_scheduler: bool) -> bool: + """Enable fresh cross-rank dummy intent for the generic non-overlap path.""" + return (should_enable_adp_dummy_fixes(mapping) + and disable_overlap_scheduler) + + def should_enable_disagg_adp_overlap_headroom( mapping: Mapping, cache_transceiver_config: Optional[CacheTransceiverConfig], diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index f05ee430f747..b2998c194b93 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -392,7 +392,9 @@ def __init__( # scheduler releases the previous batch's terminal sequence slots. from ._util import (compute_max_num_sequences, should_enable_adp_dummy_fixes, - should_enable_disagg_adp_overlap_headroom) + should_enable_disagg_adp_overlap_headroom, + should_enable_non_overlap_adp_forward_intent, + should_enable_scheduler_aware_adp_dummy) self._enable_disagg_adp_overlap_headroom = ( should_enable_disagg_adp_overlap_headroom( mapping, llm_args.cache_transceiver_config, @@ -475,7 +477,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, @@ -486,7 +487,14 @@ 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) + 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)) if drafting_loop_wrapper is not None: self.model = drafting_loop_wrapper(self.model) self.model_is_wrapped = True @@ -1010,33 +1018,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 @@ -4816,7 +4797,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 2cf78d4ab44d..7a852bf9bfd1 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -374,8 +374,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. @@ -391,9 +390,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 @@ -401,7 +397,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 @@ -409,11 +404,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, @@ -1423,7 +1413,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 dbc182d23c76..b8de240e9759 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -98,6 +98,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 @@ -586,6 +593,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 @@ -727,8 +738,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 @@ -5159,7 +5170,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) @@ -6052,11 +6064,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 @@ -6065,6 +6081,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. @@ -6177,8 +6219,18 @@ def _pad_attention_dp_dummy_request(self): if self._should_skip_dummy_for_benchmark_disagg(num_active_request): return - needs_dummy = (expected_num_active_requests > 0 - and num_active_request == 0) + if (self._enable_non_overlap_adp_forward_intent + and self.kv_cache_transceiver is not None): + num_active_request, global_intent = ( + self._get_non_overlap_adp_forward_intent()) + if global_intent != _ADPForwardIntent.NONE: + self._adp_dummy_is_gen = ( + global_intent == _ADPForwardIntent.GENERATION) + needs_dummy = (global_intent != _ADPForwardIntent.NONE + and num_active_request == 0) + else: + needs_dummy = (expected_num_active_requests > 0 + and num_active_request == 0) if not needs_dummy: return 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 5f7d02e3c48a..4bb3d6ea8938 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 2f6a2d41719a..a64f9f7a204f 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -91,7 +91,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 40378b1bd8c2..2126bac1d664 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 @@ -31,6 +30,7 @@ DisaggTransferAdmissionController, EncoderStepResult, PyExecutor, + _ADPForwardIntent, ) from tensorrt_llm._torch.pyexecutor.resource_manager import NoFreeSlotsError, ResourceManagerType from tensorrt_llm._torch.pyexecutor.scheduler import ( @@ -1569,7 +1569,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 @@ -1586,6 +1588,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 @@ -1600,12 +1605,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 = ( @@ -1646,6 +1662,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", ): @@ -1750,9 +1767,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 @@ -1767,8 +1784,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 @@ -1829,6 +1846,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) @@ -1843,7 +1883,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. @@ -1897,8 +1937,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) @@ -2039,8 +2081,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) @@ -2064,23 +2108,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 29172201f3b7..03d76b437123 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -1844,7 +1844,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') @@ -1876,9 +1875,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 = [] @@ -1901,10 +1897,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() @@ -1999,11 +1995,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]) @@ -2109,16 +2105,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 86f81378c32b..7db2c6ed74ac 100644 --- a/tests/unittest/_torch/executor/test_seq_slot_sizing.py +++ b/tests/unittest/_torch/executor/test_seq_slot_sizing.py @@ -22,6 +22,8 @@ create_torch_sampler_args, should_enable_adp_dummy_fixes, should_enable_disagg_adp_overlap_headroom, + should_enable_non_overlap_adp_forward_intent, + should_enable_scheduler_aware_adp_dummy, ) from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig from tensorrt_llm.mapping import Mapping @@ -71,6 +73,34 @@ def test_adp_dummy_fix_gate(pp_size, expected): assert should_enable_adp_dummy_fixes(mapping) is expected +@pytest.mark.parametrize( + "model_type,pp_size,disable_overlap,expected", + [ + ("kimi_k2", 1, True, True), + ("kimi_k2", 1, False, False), + ("deepseek_v4", 1, False, True), + ("qwen3_5_moe", 1, False, True), + ("deepseek_v4", 2, True, False), + ], +) +def test_scheduler_aware_adp_dummy_scope(model_type, pp_size, disable_overlap, expected): + mapping = Mapping(world_size=pp_size, tp_size=1, pp_size=pp_size) + assert should_enable_scheduler_aware_adp_dummy(model_type, mapping, disable_overlap) is expected + + +@pytest.mark.parametrize( + "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 ) 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 From 7c9c56733a5afab08fd5a2253c0a789800492a8d Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Fri, 7 Aug 2026 16:42:27 -0700 Subject: [PATCH 5/5] [NVBUG 6487039][fix] Avoid duplicate V1 dummy registration Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../_torch/pyexecutor/resource_manager.py | 21 +++++- .../_torch/executor/test_py_executor.py | 7 +- .../_torch/executor/test_resource_manager.py | 74 +++++++++++++++++++ 3 files changed, 99 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 153fe93e82ae..5129b37faa3b 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -344,6 +344,11 @@ def __init__( self.indexer_k_cache_local_layer_mask = None self.kv_connector_manager = kv_connector_manager + # Dummy requests can reserve their V1 sequence before they enter the + # normal context prepare path. Track that ownership per manager so the + # same sequence is not registered twice, while a separate draft or + # cross-cache manager can still prepare its own copy. + self._preprepared_dummy_request_ids: set[int] = set() tp_size = mapping.tp_size if mapping.enable_attention_dp: @@ -785,6 +790,8 @@ def get_needed_resource_to_completion(self, request: LlmRequest) -> int: def _context_seq_len(self, req: LlmRequest, is_cross: bool, is_star_cp: bool) -> Optional[int]: """Return the sequence length to pass to add_sequence_batch, or None to skip this request.""" + if req.py_request_id in self._preprepared_dummy_request_ids: + return None if is_cross: if (getattr(req, "py_skip_cross_kv_projection", False) or not req.is_first_context_chunk @@ -1080,6 +1087,14 @@ def add_dummy_requests( raise cleanup_error raise + if batch_request_infos: + self._preprepared_dummy_request_ids.update( + req_id for req_id, _, _ in batch_request_infos) + if (draft_batch_request_infos + and isinstance(draft_kv_cache_manager, KVCacheManager)): + draft_kv_cache_manager._preprepared_dummy_request_ids.update( + req_id for req_id, _, _ in draft_batch_request_infos) + return requests def update_resources(self, @@ -1130,8 +1145,10 @@ def update_resources(self, self.impl.store_context_blocks(request) def free_resources(self, request: LlmRequest, pin_on_release: bool = False): - return self.impl.remove_sequence(request.py_request_id, request, - pin_on_release) + result = self.impl.remove_sequence(request.py_request_id, request, + pin_on_release) + self._preprepared_dummy_request_ids.discard(request.py_request_id) + return result def store_blocks_for_reuse(self, request: LlmRequest, diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 2126bac1d664..8b4c0a67077e 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -1645,8 +1645,11 @@ def __init__( def _add_dummy(**kwargs): self.add_dummy_calls.append(kwargs) + state = ( + _STATE_GENERATION_IN_PROGRESS if kwargs["is_gen"] else LlmRequestState.CONTEXT_INIT + ) req = _make_adp_request( - _STATE_GENERATION_IN_PROGRESS, + state, request_id=kwargs["request_ids"][0], is_dummy_request=True, ) @@ -2093,6 +2096,7 @@ def test_pad_dummy_ctx_pads_to_max_num_tokens(): call = stub.add_dummy_calls[0] assert call["token_nums"] == [4096] assert call["is_gen"] is False + assert stub.active_requests[-1].state == LlmRequestState.CONTEXT_INIT def test_pad_dummy_gen_keeps_default_token_nums(): @@ -2106,6 +2110,7 @@ def test_pad_dummy_gen_keeps_default_token_nums(): call = stub.add_dummy_calls[0] assert call["token_nums"] is None assert call["is_gen"] is True + assert stub.active_requests[-1].state == _STATE_GENERATION_IN_PROGRESS def test_overlap_adp_preserves_legacy_role_without_forward_intent_collective(): diff --git a/tests/unittest/_torch/executor/test_resource_manager.py b/tests/unittest/_torch/executor/test_resource_manager.py index f4ab82db1d54..0543fdf0ff5b 100644 --- a/tests/unittest/_torch/executor/test_resource_manager.py +++ b/tests/unittest/_torch/executor/test_resource_manager.py @@ -6,6 +6,7 @@ import subprocess import sys import unittest +from types import SimpleNamespace from typing import NamedTuple, Tuple from unittest.mock import MagicMock, patch @@ -19,6 +20,7 @@ from tensorrt_llm._torch.pyexecutor.resource_manager import ( KVCacheManager, PeftCacheManager, _merge_kv_cache_pool_pointers, _warn_if_unsupported_v1_kv_cache_event_hash_algo) +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.bindings import LayerType from tensorrt_llm.bindings import ModelConfig as ModelConfigCpp from tensorrt_llm.bindings import executor as tllm @@ -990,6 +992,8 @@ def test_add_dummy_requests_failure_frees_partial_allocation(self): token_nums=[64, 64], is_gen=True, max_num_draft_tokens=128) + self.assertEqual(kv_cache_manager._preprepared_dummy_request_ids, + set()) self.assertEqual(kv_cache_manager.get_num_free_blocks(), total_free) # The freed pool must serve follow-up allocations. requests = kv_cache_manager.add_dummy_requests([2], token_nums=[64]) @@ -1087,6 +1091,76 @@ def test_peft_cache_manager_with_execution_stream(self): self.assertTrue(peft_cache_manager.impl.enabled) +@pytest.mark.cpu_only +class TestKVCacheManagerPrepreparedDummies(unittest.TestCase): + + @staticmethod + def _make_manager(is_draft: bool = False) -> KVCacheManager: + manager = KVCacheManager.__new__(KVCacheManager) + manager.mapping = Mapping() + manager.impl = MagicMock() + manager.impl.get_kv_cache_stats.return_value = SimpleNamespace( + free_num_blocks=8) + manager.is_linear_attention = False + manager.is_vswa = False + manager.num_extra_kv_tokens = 0 + manager.kv_cache_type = tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF + manager.is_draft = is_draft + manager.kv_connector_manager = None + manager._kv_reserve_draft_tokens = 0 + manager._preprepared_dummy_request_ids = set() + return manager + + @staticmethod + def _context_batch(request: LlmRequest) -> ScheduledRequests: + batch = ScheduledRequests() + batch.context_requests_last_chunk = [request] + return batch + + def test_context_dummy_is_registered_once_and_id_can_be_reused(self): + manager = self._make_manager() + + requests = manager.add_dummy_requests([0], token_nums=[64]) + self.assertIsNotNone(requests) + request = requests[0] + self.assertEqual(manager._preprepared_dummy_request_ids, {0}) + self.assertEqual(manager.impl.add_sequence_batch.call_count, 1) + + manager.prepare_resources(self._context_batch(request)) + + self.assertEqual(manager.impl.add_sequence_batch.call_count, 1) + + manager.free_resources(request) + self.assertEqual(manager._preprepared_dummy_request_ids, set()) + + manager.add_dummy_requests([0], token_nums=[64]) + self.assertEqual(manager.impl.add_sequence_batch.call_count, 2) + + def test_preprepared_dummy_ownership_is_manager_local(self): + target_manager = self._make_manager() + draft_manager = self._make_manager(is_draft=True) + other_manager = self._make_manager() + + requests = target_manager.add_dummy_requests( + [0], + token_nums=[64], + draft_kv_cache_manager=draft_manager, + ) + self.assertIsNotNone(requests) + request = requests[0] + self.assertEqual(target_manager._preprepared_dummy_request_ids, {0}) + self.assertEqual(draft_manager._preprepared_dummy_request_ids, {0}) + self.assertEqual(other_manager._preprepared_dummy_request_ids, set()) + + target_manager.prepare_resources(self._context_batch(request)) + draft_manager.prepare_resources(self._context_batch(request)) + other_manager.prepare_resources(self._context_batch(request)) + + self.assertEqual(target_manager.impl.add_sequence_batch.call_count, 1) + self.assertEqual(draft_manager.impl.add_sequence_batch.call_count, 1) + self.assertEqual(other_manager.impl.add_sequence_batch.call_count, 1) + + @pytest.mark.cpu_only class TestKVCacheManagerConfigForwarding(unittest.TestCase):