[TRTLLM-15761][refactor] Parameterize decoder invocation state and separate decoder-only engine state

tensorrt_llm/_torch/pyexecutor/engine/runners/common.py

@@ -5,19 +5,25 @@
 
 import bisect
 import math
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import torch
 
 from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata
 from tensorrt_llm._torch.distributed import Distributed
 from tensorrt_llm._torch.models.modeling_multimodal_utils import filter_mm_token_from_input_ids
+from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests
 from tensorrt_llm._torch.speculative import SpecMetadata
 from tensorrt_llm._utils import maybe_pin_memory
 from tensorrt_llm.llmapi.llm_args import PrefillCudaGraphBackend
 from tensorrt_llm.logger import logger
 from tensorrt_llm.mapping import Mapping
 
+from .interface import ScheduledInputs
+
+if TYPE_CHECKING:
+    from tensorrt_llm._torch.pyexecutor.sampler.sampler import SampleStateTensors
+
 
 def get_all_rank_num_tokens(
     attn_metadata: AttentionMetadata,
@@ -207,3 +213,24 @@ def ship_multimodal_indices(
         text_token_indices_cpu = torch.cat([text_token_indices_cpu, extra_text])
     text_token_indices_cpu = maybe_pin_memory(text_token_indices_cpu)
     inputs["text_token_indices"] = text_token_indices_cpu.to("cuda", non_blocking=True)
+
+
+def make_scheduled_inputs(
+    batch: ScheduledRequests,
+    new_tensors_device: "SampleStateTensors | None",
+    cache_indirection_buffer: torch.Tensor | None,
+    *,
+    enable_spec_decode: bool,
+    runtime_draft_len: int,
+) -> ScheduledInputs:
+    """Build a scheduled record that gathers context logits when any request returns them."""
+    return ScheduledInputs(
+        batch=batch,
+        new_tensors_device=new_tensors_device,
+        cache_indirection_buffer=cache_indirection_buffer,
+        gather_context_logits=any(
+            request.py_return_context_logits for request in batch.context_requests
+        ),
+        enable_spec_decode=enable_spec_decode,
+        runtime_draft_len=runtime_draft_len,
+    )

tensorrt_llm/_torch/speculative/utils.py

@@ -978,44 +978,35 @@ def get_static_draft_len(model_engine: "ModelEngine") -> int:
     return model_engine.max_total_draft_tokens
 
 
-def update_draft_len(model_engine: "ModelEngine",
-                     scheduled_batch: "ScheduledRequests",
-                     *,
-                     draft_len: Optional[int] = None,
-                     speculation_permanently_disabled: bool = False) -> None:
-    """Resolve this batch's draft length and synchronize its draft buffers.
-
-    Normal iterations must call this before ``prepare_resources`` so KV cache
-    allocation uses the selected draft length. Dynamic and explicit lengths
-    pad or truncate generation-request buffers to a uniform width, as required
-    by CUDA graph replay and the attention kernel. Static normal decoding
-    preserves the drafter's proposals instead.
-
-    Warmup supplies the explicit length used to allocate its dummy batch,
-    including graph shapes that differ from the normal batch-size schedule.
+def resolve_draft_len(spec_config: Optional["DecodingBaseConfig"],
+                      scheduled_batch: "ScheduledRequests",
+                      *,
+                      max_draft_len: int,
+                      static_draft_len: int,
+                      draft_len: Optional[int] = None,
+                      speculation_permanently_disabled: bool = False) -> int:
+    """Synchronize this batch's draft buffers and return its draft length.
+
+    Dynamic and explicit lengths pad or truncate generation-request buffers to
+    a uniform width, as required by CUDA graph replay and the attention kernel.
+    Static decoding preserves the drafter's proposals instead.
     """
-    if not hasattr(model_engine, 'max_draft_len'):
-        return
-
     if speculation_permanently_disabled:
         for request in scheduled_batch.generation_requests:
             request.py_draft_tokens = []
-        model_engine.runtime_draft_len = 0
-        return
+        return 0
 
-    spec_config = model_engine.spec_config
     if draft_len is None:
         if (spec_config is not None
                 and spec_config.draft_len_schedule is not None
                 and spec_config.spec_dec_mode.support_dynamic_draft_len()):
             draft_len = get_draft_len_for_batch_size(
                 spec_config.draft_len_schedule, scheduled_batch.batch_size,
-                model_engine.max_draft_len)
+                max_draft_len)
         else:
             # Static decoding preserves the proposals produced by the drafter,
             # including requests intentionally entering with no draft tokens.
-            model_engine.runtime_draft_len = get_static_draft_len(model_engine)
-            return
+            return static_draft_len
 
     draft_buffer_pad = 0  # Buffer sentinel, not PARD mask_token_id.
     rejection_on = getattr(spec_config, "use_rejection_sampling", False)
@@ -1044,4 +1035,26 @@ def update_draft_len(model_engine: "ModelEngine",
         elif current_num_draft_tokens > draft_len:
             request.py_draft_tokens = request.py_draft_tokens[:draft_len]
 
-    model_engine.runtime_draft_len = draft_len
+    return draft_len
+
+
+def update_draft_len(model_engine: "ModelEngine",
+                     scheduled_batch: "ScheduledRequests",
+                     *,
+                     draft_len: Optional[int] = None,
+                     speculation_permanently_disabled: bool = False) -> None:
+    """Resolve this batch's draft length and store it on the engine.
+
+    Normal iterations must call this before ``prepare_resources`` so KV cache
+    allocation uses the selected draft length.
+    """
+    if not hasattr(model_engine, 'max_draft_len'):
+        return
+
+    model_engine.runtime_draft_len = resolve_draft_len(
+        model_engine.spec_config,
+        scheduled_batch,
+        max_draft_len=model_engine.max_draft_len,
+        static_draft_len=get_static_draft_len(model_engine),
+        draft_len=draft_len,
+        speculation_permanently_disabled=speculation_permanently_disabled)

tests/unittest/_torch/executor/engine/test_common.py

@@ -13,10 +13,12 @@
     get_all_rank_num_tokens,
     get_padding_params,
     get_top_level_model,
+    make_scheduled_inputs,
     prepare_multimodal_indices,
     set_spec_metadata_all_rank_num_tokens,
     ship_multimodal_indices,
 )
+from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests
 from tensorrt_llm.llmapi.llm_args import PrefillCudaGraphBackend
 
 pytestmark = pytest.mark.cpu_only
@@ -205,3 +207,26 @@ def test_ship_multimodal_indices_copies_and_extends_text_indices(
     else:
         arange.assert_not_called()
         cat.assert_not_called()
+
+
+@pytest.mark.parametrize("returns_context_logits", [False, True])
+def test_make_scheduled_inputs_gathers_context_logits_from_context_requests(
+    returns_context_logits: bool,
+) -> None:
+    batch = ScheduledRequests()
+    batch.context_requests_chunking = [SimpleNamespace(py_return_context_logits=False)]
+    batch.context_requests_last_chunk = [
+        SimpleNamespace(py_return_context_logits=returns_context_logits)
+    ]
+    batch.generation_requests = [SimpleNamespace(py_return_context_logits=True)]
+    new_tensors_device = object()
+
+    inputs = make_scheduled_inputs(
+        batch, new_tensors_device, None, enable_spec_decode=True, runtime_draft_len=3
+    )
+
+    assert inputs.batch is batch
+    assert inputs.new_tensors_device is new_tensors_device
+    assert inputs.cache_indirection_buffer is None
+    assert inputs.gather_context_logits is returns_context_logits
+    assert (inputs.enable_spec_decode, inputs.runtime_draft_len) == (True, 3)

tests/unittest/_torch/executor/engine/test_runners.py

@@ -648,24 +648,69 @@ def test_engine_consumes_optional_length_update_and_passes_call_state(effective_
     assert "runtime_draft_len" not in outputs
 
 
-def test_decoder_fallback_length_update_does_not_modify_graph_output_dictionary():
+@pytest.mark.parametrize("is_dummy", [False, True])
+def test_decoder_fallback_passes_call_state_without_engine_writes(is_dummy):
+    engine, resources = _model_engine_with_runner(None, kv_cache_manager=object())
+    engine._fallback_to_engine = True
+    # The dummy flag comes from the caller, not from the warmup scope.
+    engine._is_warmup = not is_dummy
+    engine.enable_spec_decode = False
+    engine.runtime_draft_len = 5
+    inputs = ScheduledInputs(
+        batch=ScheduledRequests(), enable_spec_decode=True, runtime_draft_len=2
+    )
+    expected = {"logits": object(), "runtime_draft_len": 4}
+    engine._forward_decoder = Mock(return_value=expected)
+
+    outputs = engine._forward_scheduled(inputs, resource_manager=resources, is_dummy=is_dummy)
+
+    assert outputs is expected
+    engine._forward_decoder.assert_called_once_with(inputs, resources, is_dummy=is_dummy)
+    assert (engine.enable_spec_decode, engine.runtime_draft_len) == (False, 5)
+    assert (inputs.enable_spec_decode, inputs.runtime_draft_len) == (True, 2)
+
+
+def test_engine_forward_keeps_call_state_when_decoder_fails():
     engine, resources = _model_engine_with_runner(None, kv_cache_manager=object())
     engine._fallback_to_engine = True
-    cached_outputs = {"logits": torch.tensor([1.0])}
-    inputs = ScheduledInputs(batch=ScheduledRequests(), runtime_draft_len=2)
+    engine._is_warmup = False
+    engine.enable_spec_decode = True
+    engine.runtime_draft_len = 2
+    engine._forward_decoder = Mock(side_effect=RuntimeError("decoder failure"))
+
+    with pytest.raises(RuntimeError, match="decoder failure"):
+        engine.forward(ScheduledRequests(), resources)
 
-    def decoder_forward(*args):
-        engine.runtime_draft_len = 4
-        return cached_outputs
+    inputs = engine._forward_decoder.call_args.args[0]
+    assert (inputs.enable_spec_decode, inputs.runtime_draft_len) == (True, 2)
+    assert (engine.enable_spec_decode, engine.runtime_draft_len) == (True, 2)
+
+
+def test_forward_warmup_runs_local_call_state_in_forward_context():
+    engine, resources = _model_engine_with_runner(None, kv_cache_manager=object())
+    engine._is_warmup = True
+    engine.enable_spec_decode = False
+    engine.runtime_draft_len = 5
+    batch = ScheduledRequests()
+
+    def decoder_forward(inputs, resource_manager, *, is_dummy):
+        assert torch.is_inference_mode_enabled()
+        assert get_model_extra_attrs() is engine.model.extra_attrs
+        return {"logits": object(), "runtime_draft_len": 3}
 
     engine._forward_decoder = Mock(side_effect=decoder_forward)
-    outputs = engine._forward_scheduled(inputs, resource_manager=resources)
 
-    assert outputs is not cached_outputs
-    assert outputs["logits"] is cached_outputs["logits"]
-    assert outputs.pop("runtime_draft_len") == 4
-    assert set(cached_outputs) == {"logits"}
-    assert inputs.runtime_draft_len == 2
+    outputs = engine._forward_warmup(batch, resources, enable_spec_decode=True, runtime_draft_len=1)
+
+    inputs = engine._forward_decoder.call_args.args[0]
+    assert inputs.batch is batch
+    assert inputs.new_tensors_device is None
+    assert inputs.cache_indirection_buffer is None
+    assert (inputs.enable_spec_decode, inputs.runtime_draft_len) == (True, 1)
+    assert engine._forward_decoder.call_args.kwargs == {"is_dummy": True}
+    assert outputs["runtime_draft_len"] == 3
+    assert (engine.enable_spec_decode, engine.runtime_draft_len) == (False, 5)
+    assert get_model_extra_attrs() is None
 
 
 def test_no_kv_forward_preserves_model_output_dictionary_without_length_update():
@@ -690,16 +735,12 @@ def test_engine_forward_preserves_raw_decoder_outputs_and_length(raw_output, is_
     engine._fallback_to_engine = True
     engine._is_warmup = is_dummy
     engine.runtime_draft_len = 2
-
-    def decoder_forward(inputs, resource_manager):
-        assert inputs.runtime_draft_len == 2
-        engine.runtime_draft_len = 4
-        return raw_output
-
-    engine._forward_decoder = Mock(side_effect=decoder_forward)
+    engine._forward_decoder = Mock(return_value=raw_output)
 
     assert engine.forward(ScheduledRequests(), resources) is raw_output
-    assert engine.runtime_draft_len == 4
+    assert engine.runtime_draft_len == 2
+    assert engine._forward_decoder.call_args.args[0].runtime_draft_len == 2
+    assert engine._forward_decoder.call_args.kwargs == {"is_dummy": is_dummy}
 
 
 def test_model_caller_uses_current_forward_and_restores_outer_attribute_context():

tests/unittest/_torch/executor/test_breakable_cuda_graph.py

@@ -258,10 +258,11 @@ def model_forward():
         batch = ScheduledRequests()
         resources = object()
 
-        def decoder_forward(forward_inputs, resource_manager):
+        def decoder_forward(forward_inputs, resource_manager, *, is_dummy):
             assert isinstance(forward_inputs, ScheduledInputs)
             assert forward_inputs.batch is batch
             assert resource_manager is resources
+            assert is_dummy
             return model_forward()
 
         engine = object.__new__(PyTorchModelEngine)

tests/unittest/_torch/executor/test_distributed_warmup_oom.py

@@ -70,6 +70,9 @@ def _allgather(value):
         tp_size=tp_size,
     )
     engine._reset_moe_alltoall_state = mock.Mock()
+    engine.is_spec_decode = False
+    engine.spec_config = None
+    engine.max_draft_len = 0
     return engine
 
 
@@ -400,7 +403,7 @@ def _general_warmup_engine(*, world_size: int, dwdp_size: int) -> PyTorchModelEn
 def test_general_warmup_oom_policy(world_size: int, dwdp_size: int, is_fatal: bool) -> None:
     engine = _general_warmup_engine(world_size=world_size, dwdp_size=dwdp_size)
     error = torch.OutOfMemoryError("asymmetric OOM")
-    engine.forward = mock.Mock(side_effect=error if is_fatal else [error, None])
+    engine._forward_warmup = mock.Mock(side_effect=error if is_fatal else [error, None])
 
     with _no_cuda_side_effects() as (empty_cache, _synchronize):
         if is_fatal:
@@ -413,9 +416,9 @@ def test_general_warmup_oom_policy(world_size: int, dwdp_size: int, is_fatal: bo
     if is_fatal:
         # The remaining shape is never attempted: this rank's peers are
         # already stuck in the failed forward's collectives.
-        engine.forward.assert_called_once()
+        engine._forward_warmup.assert_called_once()
     else:
-        assert engine.forward.call_count == 2
+        assert engine._forward_warmup.call_count == 2
         # A retry after an OOM between dispatch() and combine() has to start
         # from a clean MoE all-to-all state.
         engine._reset_moe_alltoall_state.assert_called_once_with()
@@ -474,7 +477,7 @@ def _run_mamba_warmup(engine: PyTorchModelEngine, resource_manager: object) -> N
 def test_mamba_preforward_error_policy_when_alone(error: Exception, recoverable: bool) -> None:
     engine, resource_manager = _mamba_engine()
     engine._create_warmup_request = mock.Mock(side_effect=error)
-    engine.forward = mock.Mock()
+    engine._forward_warmup = mock.Mock()
 
     with _no_cuda_side_effects():
         if recoverable:
@@ -483,7 +486,7 @@ def test_mamba_preforward_error_policy_when_alone(error: Exception, recoverable:
             with pytest.raises(RuntimeError, match="unexpected"):
                 _run_mamba_warmup(engine, resource_manager)
 
-    engine.forward.assert_not_called()
+    engine._forward_warmup.assert_not_called()
     # The failure predates dispatch(), so there is no half-finished MoE
     # all-to-all exchange to unwind.
     engine._reset_moe_alltoall_state.assert_not_called()
@@ -492,15 +495,15 @@ def test_mamba_preforward_error_policy_when_alone(error: Exception, recoverable:
 def test_mamba_midforward_runtime_error_recovers_when_alone() -> None:
     engine, resource_manager = _mamba_engine()
     engine._create_warmup_request = mock.Mock(return_value=object())
-    engine.forward = mock.Mock(side_effect=RuntimeError("mid-forward failure"))
+    engine._forward_warmup = mock.Mock(side_effect=RuntimeError("mid-forward failure"))
 
     with _no_cuda_side_effects():
         _run_mamba_warmup(engine, resource_manager)
 
     # Every shape is attempted, and each failed forward has to leave the MoE
     # all-to-all state clean for the shape that follows it.
-    assert engine.forward.call_count >= 1
-    assert engine._reset_moe_alltoall_state.call_count == engine.forward.call_count
+    assert engine._forward_warmup.call_count >= 1
+    assert engine._reset_moe_alltoall_state.call_count == engine._forward_warmup.call_count
 
 
 @pytest.mark.parametrize("phase", ["pre-forward", "mid-forward"])
@@ -509,10 +512,10 @@ def test_mamba_error_is_fatal_when_distributed(phase: str) -> None:
     error = RuntimeError(_KV_ALLOC_ERROR)
     if phase == "pre-forward":
         engine._create_warmup_request = mock.Mock(side_effect=error)
-        engine.forward = mock.Mock()
+        engine._forward_warmup = mock.Mock()
     else:
         engine._create_warmup_request = mock.Mock(return_value=object())
-        engine.forward = mock.Mock(side_effect=error)
+        engine._forward_warmup = mock.Mock(side_effect=error)
 
     with _no_cuda_side_effects():
         with pytest.raises(RuntimeError) as excinfo: