Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 3 additions & 5 deletions tensorrt_llm/_torch/pyexecutor/_util.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The mechanism you describe in the PR body — terminal requests holding slots while the V2 scheduler backfills their seats — is model-agnostic. Every non-PP model with overlap scheduling has the same latent 2× demand; DSv4 just happens to be the one where ADP dummy requests saturate the pool often enough to hit it. Keeping the gate model-scoped is a defensible risk call for a fix PR, but it means the same hang is still reachable elsewhere. Is there a follow-up ticket to validate and generalize? Worth naming it in the docstring so this doesn't sit as a permanent special case.

and spec_config is not None
and spec_config.spec_dec_mode.is_mtp_eagle_one_model()
and not disable_overlap_scheduler)


Expand Down
9 changes: 4 additions & 5 deletions tensorrt_llm/_torch/pyexecutor/model_engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = (

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This now doubles the slot pool for every non-PP DSv4 run, including plain no-spec-decode serving, which doubles the sampler's per-slot state (TorchSampler.Args.max_num_sequences) and the guided-decoder allocation. Previously only MTP configs paid that. Did you measure the host/device memory delta on a large max_batch_size config? If it's non-trivial, worth a line in the PR description so it isn't a surprise in perf triage.

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,
Expand Down
32 changes: 10 additions & 22 deletions tests/unittest/_torch/executor/test_seq_slot_sizing.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand All @@ -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(
Expand Down
Loading