From e2f7e712078dd8effab72003de76a40340635408 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou Date: Tue, 4 Aug 2026 21:39:21 -0700 Subject: [PATCH] [None][fix] enable DSv4 overlap slot headroom without MTP Signed-off-by: Lizhi Zhou --- tensorrt_llm/_torch/pyexecutor/_util.py | 8 ++--- .../_torch/pyexecutor/model_engine.py | 9 +++--- .../_torch/executor/test_seq_slot_sizing.py | 32 ++++++------------- 3 files changed, 17 insertions(+), 32 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index ec26f5a22855..35627f34f835 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2526,12 +2526,10 @@ def should_enable_dsv4_adp_dummy_fixes(model_type: Optional[str], 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.""" + model_type: Optional[str], mapping: Mapping, + disable_overlap_scheduler: bool) -> bool: + """Gate extra sequence slots to the non-PP DSv4 overlap path.""" return (should_enable_dsv4_adp_dummy_fixes(model_type, mapping) - and spec_config is not None - and spec_config.spec_dec_mode.is_mtp_eagle_one_model() and not disable_overlap_scheduler) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 02dc5da52e6a..b6da0726ba49 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -446,15 +446,14 @@ 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. + # Keep the scheduler/dummy fix and larger slot pool model-scoped. The + # overlap lifetime requires headroom with or without speculative + # decoding. PP remains on its established path. self._enable_dsv4_adp_dummy_fixes = (should_enable_dsv4_adp_dummy_fixes( model_type, mapping)) self._enable_dsv4_overlap_headroom = ( should_enable_dsv4_overlap_headroom( - model_type, spec_config, mapping, - llm_args.disable_overlap_scheduler)) + model_type, mapping, llm_args.disable_overlap_scheduler)) self.max_num_seq_slots = compute_max_num_sequences( mapping, self.batch_size, diff --git a/tests/unittest/_torch/executor/test_seq_slot_sizing.py b/tests/unittest/_torch/executor/test_seq_slot_sizing.py index d42e6f4483c9..dcef1acd0fb9 100644 --- a/tests/unittest/_torch/executor/test_seq_slot_sizing.py +++ b/tests/unittest/_torch/executor/test_seq_slot_sizing.py @@ -7,16 +7,15 @@ 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, regardless of whether speculative decoding is enabled. +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 ( @@ -40,29 +39,18 @@ @pytest.mark.parametrize( - "model_type,has_spec,is_mtp_one_model,pp_size,disable_overlap,expected", + "model_type,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), + ("deepseek_v4", 1, False, True), + ("deepseek_v3", 1, False, False), + ("deepseek_v4", 2, False, False), + ("deepseek_v4", 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 +def test_dsv4_overlap_headroom_gate(model_type, pp_size, disable_overlap, expected): 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 - ) + assert should_enable_dsv4_overlap_headroom(model_type, mapping, disable_overlap) is expected @pytest.mark.parametrize(