[TRTLLM-14620][fix] Preserve decoder construction contracts

Keep input-preparer ownership in the native and encoder-decoder graph factories, and retain optional-config fixture construction without changing configured confidence behavior. Add focused factory and constructor regression coverage.

Signed-off-by: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com>

Author: zhaoyuanh-nvidia

tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py

@@ -357,15 +357,15 @@ def __init__(
         self._prepare_inputs_event: torch.cuda.Event | None = None
         self._gather_ids_buffer = (
             DeviceCopyBuffer(buffers.gather_ids_cuda, pin_memory=prefer_pinned())
-            if config.is_spec_decode
+            if config is not None and config.is_spec_decode
             else None
         )
         # Reuse device indices while the scheduled batch stays unchanged.
         self._staged_previous_batch_indices: list[int] | None = None
         self._staged_previous_pos_indices: list[int] | None = None
         self._step_input_gather: spec_step_copies.StepInputGather | None = (
             spec_step_copies.StepInputGather()
-            if config.is_spec_decode and spec_step_copies.is_supported()
+            if config is not None and config.is_spec_decode and spec_step_copies.is_supported()
             else None
         )
         # Steady-state generation-only prepare cache (non-speculative overlap
@@ -381,7 +381,7 @@ def __init__(
             os.getenv("TLLM_LOG_CACHED_KV_TOKENS_PER_REQ", "0") == "1"
         )
 
-        spec_config = self._config.spec_config
+        spec_config = self._config.spec_config if self._config is not None else None
         self._dspark_confidence_enabled = bool(
             spec_config is not None and getattr(spec_config, "enable_confidence_scheduling", False)
         )
@@ -1950,23 +1950,27 @@ def previous_seq_slots_device() -> torch.Tensor:
                     if host.previous_batch_token_start is not None
                     else previous_begin * runtime_tokens_per_gen_step
                 )
-                if is_ragged_gen or step_gather is None or not step_gather.gather(
-                    new_tokens_device,
-                    next_draft_tokens_device,
-                    new_tokens_lens_device,
-                    self._buffers.previous_batch_indices_cuda,
-                    self._buffers.previous_pos_indices_cuda,
-                    previous_batch_len,
-                    runtime_tokens_per_gen_step,
-                    runtime_draft_token_buffer_width,
-                    self._buffers.input_ids_cuda,
-                    num_tokens,
-                    self._buffers.draft_tokens_cuda,
-                    num_draft_tokens,
-                    self._buffers.previous_pos_id_offsets_cuda,
-                    pos_offsets_start,
-                    self._buffers.previous_kv_lens_offsets_cuda,
-                    previous_begin,
+                if (
+                    is_ragged_gen
+                    or step_gather is None
+                    or not step_gather.gather(
+                        new_tokens_device,
+                        next_draft_tokens_device,
+                        new_tokens_lens_device,
+                        self._buffers.previous_batch_indices_cuda,
+                        self._buffers.previous_pos_indices_cuda,
+                        previous_batch_len,
+                        runtime_tokens_per_gen_step,
+                        runtime_draft_token_buffer_width,
+                        self._buffers.input_ids_cuda,
+                        num_tokens,
+                        self._buffers.draft_tokens_cuda,
+                        num_draft_tokens,
+                        self._buffers.previous_pos_id_offsets_cuda,
+                        pos_offsets_start,
+                        self._buffers.previous_kv_lens_offsets_cuda,
+                        previous_begin,
+                    )
                 ):
                     # previous input ids
                     if is_ragged_gen:

tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py

@@ -232,7 +232,7 @@ def _make_input_preparer(
 
     def _make_cuda_graphs(self) -> DecoderCudaGraphs:
         """Build the CUDA graph component; call after the input preparer."""
-        return DecoderCudaGraphs(**self._cuda_graphs_kwargs())
+        return DecoderCudaGraphs(**self._cuda_graphs_kwargs(), inputs=self._inputs)
 
     def _dummy_requests_kwargs(self) -> dict[str, Any]:
         return dict(
@@ -281,7 +281,6 @@ def _cuda_graphs_kwargs(self) -> dict[str, Any]:
             torch_compile_bypass=self._bypass_torch_compile,
             context_promotion_supported=self._context_graph_promotion_supported,
             warmup_inputs=self._buffers.input_ids_cuda,
-            inputs=self._inputs,
         )
 
     @property

tests/unittest/_torch/executor/test_pytorch_model_engine.py

@@ -39,6 +39,8 @@
     PrefillCompiledModel
 from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.warmup import \
     DecoderWarmup
+from tensorrt_llm._torch.pyexecutor.engine.runners.encoder_decoder import \
+    EncoderDecoderRunner
 from tensorrt_llm._torch.pyexecutor.engine.runners.interface import \
     PreparedInputs
 from tensorrt_llm._torch.pyexecutor.engine.runners.pooling import PoolingRunner
@@ -78,17 +80,80 @@
 _TWO_MODEL_SPEC_MODE = SimpleNamespace(use_one_engine=lambda: False)
 
 
-def _input_preparer() -> DecoderInputPreparer:
-    return DecoderInputPreparer(buffers=None,
-                                config=SimpleNamespace(is_spec_decode=False),
-                                mapping=None,
-                                dist=None,
-                                model=None,
-                                get_runtime_tokens_per_gen_step=None,
-                                iter_states={},
-                                lora=None,
-                                use_beam_search=False,
-                                steady_gen_cache_supported=True)
+def _input_preparer(*, without_config: bool = False) -> DecoderInputPreparer:
+    return DecoderInputPreparer(
+        buffers=None,
+        config=(None if without_config else SimpleNamespace(
+            is_spec_decode=False, spec_config=None)),
+        mapping=None,
+        dist=None,
+        model=None,
+        get_runtime_tokens_per_gen_step=None,
+        iter_states={},
+        lora=None,
+        use_beam_search=False,
+        steady_gen_cache_supported=True)
+
+
+@pytest.mark.parametrize("runner_type,graphs_path", [
+    (DecoderRunner,
+     "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner.DecoderCudaGraphs"
+     ),
+    (EncoderDecoderRunner,
+     "tensorrt_llm._torch.pyexecutor.engine.runners.encoder_decoder.EncoderDecoderCudaGraphs"
+     )
+])
+def test_cuda_graph_factory_preserves_input_preparer(runner_type, graphs_path):
+    runner = Mock()
+    runner._inputs = object()
+    runner._cuda_graphs_kwargs.side_effect = lambda: DecoderRunner._cuda_graphs_kwargs(
+        runner)
+
+    with patch(graphs_path) as graphs_type:
+        graphs = runner_type._make_cuda_graphs(runner)
+
+    graphs_type.assert_called_once()
+    assert graphs_type.call_args.kwargs["inputs"] is runner._inputs
+    assert graphs_type.call_args.kwargs[
+        "warmup_inputs"] is runner._buffers.input_ids_cuda
+    assert graphs is graphs_type.return_value
+
+
+@pytest.mark.parametrize("enabled,fused", [(None, False), (False, True),
+                                           (True, False), (True, True)])
+def test_input_preparer_configured_confidence_flags(enabled, fused):
+    spec_config = (None if enabled is None else SimpleNamespace(
+        enable_confidence_scheduling=enabled,
+        enable_fused_confidence_scheduler=fused))
+    config = SimpleNamespace(is_spec_decode=False, spec_config=spec_config)
+    inputs = DecoderInputPreparer(buffers=None,
+                                  config=config,
+                                  mapping=None,
+                                  dist=None,
+                                  model=None,
+                                  get_runtime_tokens_per_gen_step=None,
+                                  iter_states={},
+                                  lora=None,
+                                  use_beam_search=False,
+                                  steady_gen_cache_supported=True)
+
+    assert inputs._config is config
+    assert inputs._dspark_confidence_enabled is bool(enabled)
+    assert inputs._dspark_trims_submitted_tokens is bool(enabled)
+    assert inputs._dspark_device_windows is bool(enabled and fused)
+    assert inputs.previous_verify_lens == ([] if enabled else None)
+    assert inputs.has_previous_device_draft is False
+
+
+def test_input_preparer_without_config_keeps_confidence_disabled():
+    inputs = _input_preparer(without_config=True)
+
+    assert inputs._config is None
+    assert inputs._dspark_confidence_enabled is False
+    assert inputs._dspark_trims_submitted_tokens is False
+    assert inputs._dspark_device_windows is False
+    assert inputs.previous_verify_lens is None
+    assert inputs.has_previous_device_draft is False
 
 
 class TestCachedKvTokenLogging(unittest.TestCase):

tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph_runtime.py

@@ -42,7 +42,7 @@ def num_generation_requests(self):
 
 def _runner_stub(runner_type=DSparkCUDAGraphRunner):
     runner = runner_type.__new__(runner_type)
-    runner.clear = lambda: None
+    runner.clear = lambda **_kwargs: None
     return runner
 
 
@@ -220,11 +220,11 @@ def test_decoder_selects_confidence_graph_subclass_only_when_enabled(
     decoder = SimpleNamespace(
         _dspark_confidence_enabled=confidence_enabled,
         _cuda_graph_runner_config=lambda: config,
-        model=SimpleNamespace(register_cuda_graph_pre_replay_hooks=registered.append),
+        _model=SimpleNamespace(register_cuda_graph_pre_replay_hooks=registered.append),
     )
 
     runner = DecoderRunner._initialize_cuda_graph_runner(decoder)
-    runner.clear = lambda: None
+    runner.clear = lambda **_kwargs: None
 
     expected_type = DSparkCUDAGraphRunner if confidence_enabled else CUDAGraphRunner
     assert type(runner) is expected_type

tests/unittest/_torch/speculative/hw_agnostic/test_dspark_model_engine_runtime.py

@@ -373,7 +373,7 @@ def test_model_engine_delegates_ragged_state_to_decoder_owner() -> None:
     inputs._dspark_exact_identity_words = tuple(range(8))
     inputs.fit_ragged_verify_lens = lambda *args, **kwargs: (args, kwargs)
     worker = object()
-    runner.model = SimpleNamespace(spec_worker=worker)
+    runner._model = SimpleNamespace(spec_worker=worker)
     engine = object.__new__(PyTorchModelEngine)
     engine._cleanup_done = True
     engine._runner = runner
@@ -440,7 +440,7 @@ def __init__(self, **kwargs):
     metadata_module.build_attention_metadata(model, **kwargs)
 
     assert "enable_ragged_verification" not in seen[0]
-    assert seen[0]["num_heads_per_kv"] == 4
+    assert seen[0]["max_num_heads_per_kv"] == 4
 
 
 def test_metadata_helper_enables_only_supported_ragged_metadata():
@@ -481,6 +481,7 @@ def __init__(self, **kwargs):
         **kwargs, spec_config=SimpleNamespace(enable_confidence_scheduling=enabled)
     )
     decoder.mapping = kwargs["mapping"]
+    decoder._attention_runtime_features = kwargs["attention_runtime_features"]
     decoder._buffers = SimpleNamespace(cache_indirection=None)
     monkeypatch.setattr(decoder_metadata_module, "resolve_mamba_metadata_cls", lambda _model: None)
     kv = object()
@@ -498,7 +499,7 @@ def __init__(self, **kwargs):
 @pytest.mark.parametrize("graph_type", [CUDAGraphRunner, DSparkCUDAGraphRunner])
 def test_current_main_warmup_keeps_serving_timeout_only_for_capture(monkeypatch, graph_type):
     graph = graph_type.__new__(graph_type)
-    graph.clear = lambda: None
+    graph.clear = lambda **_kwargs: None
     graph._capture_allowed = False
     graph.padding_dummy_requests = {5: object()}
     if graph_type is DSparkCUDAGraphRunner:
@@ -569,12 +570,16 @@ def test_confidence_without_an_sps_table_is_eager_only() -> None:
         graphs._get_graphs_to_capture([1, 4])
 
 
+@pytest.mark.parametrize("use_fused_gather", [False, True])
 @pytest.mark.parametrize(
     ("trims", "prefix_window", "previous_windows"),
     [(True, 1, [3, 3]), (True, 1, [1, 2]), (False, None, [None, None])],
 )
 def test_overlap_offsets_follow_the_packed_prefix(
-    trims: bool, prefix_window: int | None, previous_windows: list[int | None]
+    trims: bool,
+    prefix_window: int | None,
+    previous_windows: list[int | None],
+    use_fused_gather: bool,
 ) -> None:
     """Exercise the full input packer with CPU buffers and metadata-only leaves."""
     decoder = object.__new__(DecoderInputPreparer)
@@ -618,6 +623,27 @@ def test_overlap_offsets_follow_the_packed_prefix(
     decoder._mm_item_scheduler = None
     decoder._pinned_host_cache = {}
     decoder._pinned_host_active = {}
+    decoder._staged_previous_batch_indices = None
+    decoder._staged_previous_pos_indices = None
+    gather_calls = []
+
+    def gather(*args) -> bool:
+        # A CPU kernel leaf verifies the production caller's packed offsets.
+        new, draft, lens, slots, positions, rows, width, draft_width = args[:8]
+        ids, input_begin, drafts, draft_begin, pos, pos_begin, kv, kv_begin = args[8:]
+        selected_slots = slots[:rows].long()
+        gather_calls.append((input_begin, pos_begin, kv_begin))
+        ids[input_begin : input_begin + rows * width].copy_(
+            new.transpose(0, 1)[selected_slots, :width].flatten()
+        )
+        drafts[draft_begin : draft_begin + rows * draft_width].copy_(
+            draft[selected_slots, :draft_width].flatten()
+        )
+        pos[pos_begin : pos_begin + rows * width].copy_(lens[positions[: rows * width].long()])
+        kv[kv_begin : kv_begin + rows].copy_(lens[selected_slots] - width)
+        return True
+
+    decoder._step_input_gather = SimpleNamespace(gather=gather) if use_fused_gather else None
     decoder._ragged_gather_indices = lambda slots, lengths: tuple(
         torch.tensor(indices, dtype=torch.long)
         for indices in ragged_gather_index_lists(slots, lengths)
@@ -626,7 +652,6 @@ def test_overlap_offsets_follow_the_packed_prefix(
     for name in (
         "input_ids_cuda",
         "draft_tokens_cuda",
-        "num_accepted_draft_tokens_cuda",
         "previous_pos_id_offsets_cuda",
         "previous_kv_lens_offsets_cuda",
         "previous_pos_indices_cuda",
@@ -635,6 +660,7 @@ def test_overlap_offsets_follow_the_packed_prefix(
         "gather_ids_cuda",
     ):
         setattr(decoder._buffers, name, torch.full((32,), -99, dtype=torch.int32))
+    decoder._gather_ids_buffer = decoder_module.DeviceCopyBuffer(decoder._buffers.gather_ids_cuda)
     windows = [prefix_window, *previous_windows]
     requests = [
         SimpleNamespace(
@@ -699,3 +725,87 @@ def test_overlap_offsets_follow_the_packed_prefix(
     assert metadata.kv_cache_params.num_extra_kv_tokens == 2
     assert gather_ids.tolist() == list(range(sum(lengths)))
     assert draft_len == 3
+    assert gather_calls == (
+        [(prefix_len, prefix_len, 1)] if use_fused_gather and previous_lengths == [4, 4] else []
+    )
+
+
+@pytest.mark.parametrize("trims", [False, True])
+def test_graph_overlap_offset_zeroing_preserves_ragged_bucket_padding(trims: bool) -> None:
+    decoder = object.__new__(DecoderInputPreparer)
+    decoder._config = SimpleNamespace(
+        max_num_tokens=16, max_beam_width=1, disable_overlap_scheduler=False
+    )
+    decoder._dspark_trims_submitted_tokens = trims
+    decoder._staged_previous_batch_indices = None
+    decoder._staged_previous_pos_indices = None
+    decoder._step_input_gather = None
+    decoder.get_runtime_tokens_per_gen_step = lambda _length: 4
+    decoder._ragged_gather_indices = lambda slots, lengths: tuple(
+        torch.tensor(indices, dtype=torch.long)
+        for indices in ragged_gather_index_lists(slots, lengths)
+    )
+    decoder._buffers = SimpleNamespace(
+        **{
+            name: torch.full((16,), -99, dtype=torch.int32)
+            for name in (
+                "input_ids_cuda",
+                "draft_tokens_cuda",
+                "previous_batch_indices_cuda",
+                "previous_pos_indices_cuda",
+                "previous_pos_id_offsets_cuda",
+                "previous_kv_lens_offsets_cuda",
+            )
+        }
+    )
+    lengths = [1, 2] if trims else [4, 4]
+    host = SimpleNamespace(
+        lists=SimpleNamespace(
+            input_ids=[],
+            draft_tokens=[],
+            draft_lens=[],
+            position_ids=list(range(sum(lengths))),
+            previous_batch_indices=[0, 1],
+            previous_pos_indices=[0] * lengths[0] + [1] * lengths[1],
+            extend_requests=[object(), object()],
+            extend_dummy_requests=[],
+            generation_requests=[],
+            unwritten_input_id_rows=[],
+        ),
+        previous_batch_tokens_per_request=lengths if trims else None,
+        previous_batch_token_start=0,
+    )
+    decoder._copy_token_ids_to_device(
+        host,
+        SimpleNamespace(is_attention_dp_phase_idle=False, num_context_requests=0),
+        torch.arange(8, dtype=torch.int32).reshape(4, 2, 1),
+        torch.tensor([1, 2], dtype=torch.int32),
+        torch.arange(6, dtype=torch.int32).reshape(2, 3),
+        3,
+        attn_metadata=SimpleNamespace(is_cuda_graph=True),
+        enable_spec_decode=True,
+    )
+    expected_tail = 0 if trims else -99
+    assert decoder._buffers.previous_pos_id_offsets_cuda[sum(lengths) :].tolist() == [
+        expected_tail
+    ] * (16 - sum(lengths))
+    assert decoder._buffers.previous_kv_lens_offsets_cuda[2:].tolist() == [expected_tail] * 14
+
+
+@pytest.mark.parametrize("release_owners", [False, True])
+def test_graph_release_preserves_nccl_window_owner_choice(release_owners: bool) -> None:
+    calls = []
+    graphs = object.__new__(DecoderCudaGraphs)
+    graphs.cuda_graph_runner = SimpleNamespace(
+        clear=lambda **kwargs: calls.append(("native", kwargs))
+    )
+    graphs.breakable_cuda_graph_runner = SimpleNamespace(
+        clear=lambda **kwargs: calls.append(("breakable", kwargs))
+    )
+
+    graphs.release(release_nccl_window_owners=release_owners)
+
+    assert calls == [
+        ("native", {"release_nccl_window_owners": release_owners}),
+        ("breakable", {"release_nccl_window_owners": release_owners}),
+    ]