[TRTLLM-17348][refactor] Reorganize engine and decoder tests around component ownership

.claude/skills/perf-host-optimization/SKILL.md

@@ -194,10 +194,10 @@ grep -rl "_prepare_tp_inputs\|prepare_inputs" tests/
 **Running tests:**
 ```bash
 # Run specific test file with stop-on-first-failure
-pytest tests/unittest/_torch/executor/test_pytorch_model_engine.py -v -x --timeout=120
+pytest tests/unittest/_torch/executor/test_model_engine.py -v -x --timeout=120
 
 # Run specific test method
-pytest tests/unittest/_torch/executor/test_pytorch_model_engine.py::PyTorchModelEngineTestCase::test_position_id_preparation -v -x
+pytest tests/unittest/_torch/executor/test_model_engine.py::ModelEngineTestCase::test_position_id_preparation -v -x
 ```
 
 For the full UT-to-file mapping, see [references/hot-path-files.md](references/hot-path-files.md).

.claude/skills/perf-host-optimization/references/hot-path-files.md

@@ -114,7 +114,8 @@ When a top-level profiled function shows >80% time in a single sub-call, drill d
 
 | Modified File | Primary Test File | Key Tests |
 |---|---|---|
-| `model_engine.py` | `tests/unittest/_torch/executor/test_pytorch_model_engine.py` | `test_position_id_preparation`, `test_pad_generation_requests`, `test_warmup` |
+| `model_engine.py` | `tests/unittest/_torch/executor/test_model_engine.py` | `test_position_id_preparation`, `test_warmup` |
+| `cuda_graph_runner.py` | `tests/unittest/_torch/executor/test_cuda_graph_runner_gpu.py` | `test_pad_generation_requests` |
 | `sampler.py` | `tests/unittest/_torch/executor/test_sampler.py` | Sampler-related tests |
 | `py_executor.py` | `tests/unittest/_torch/executor/test_py_executor.py` | Executor loop tests |
 | `resource_manager.py` | `tests/unittest/_torch/executor/test_resource_manager.py` | KV cache tests |

.pre-commit-config.yaml

@@ -530,7 +530,6 @@ common-files: &common_files |
         tests/unittest/_torch/executor/test_chunked_logits.py |
         tests/unittest/_torch/executor/test_executor_request_queue.py |
         tests/unittest/_torch/executor/test_overlap_scheduler.py |
-        tests/unittest/_torch/executor/test_pytorch_model_engine.py |
         tests/unittest/_torch/executor/test_resource_manager.py |
         tests/unittest/_torch/executor/test_router_dealer_ipc.py |
         tests/unittest/_torch/helpers.py |
@@ -1259,7 +1258,6 @@ legacy-files: &legacy_files |
         tests/unittest/_torch/executor/test_chunked_logits.py |
         tests/unittest/_torch/executor/test_executor_request_queue.py |
         tests/unittest/_torch/executor/test_overlap_scheduler.py |
-        tests/unittest/_torch/executor/test_pytorch_model_engine.py |
         tests/unittest/_torch/executor/test_resource_manager.py |
         tests/unittest/_torch/executor/test_router_dealer_ipc.py |
         tests/unittest/_torch/helpers.py |

legacy-files.txt

@@ -522,7 +522,6 @@ tests/unittest/_torch/debugger/test_debugger_addon.py
 tests/unittest/_torch/executor/test_chunked_logits.py
 tests/unittest/_torch/executor/test_executor_request_queue.py
 tests/unittest/_torch/executor/test_overlap_scheduler.py
-tests/unittest/_torch/executor/test_pytorch_model_engine.py
 tests/unittest/_torch/executor/test_resource_manager.py
 tests/unittest/_torch/executor/test_router_dealer_ipc.py
 tests/unittest/_torch/helpers.py

pyproject.toml

@@ -587,7 +587,6 @@ exclude = [
     "tests/unittest/_torch/executor/test_chunked_logits.py",
     "tests/unittest/_torch/executor/test_executor_request_queue.py",
     "tests/unittest/_torch/executor/test_overlap_scheduler.py",
-    "tests/unittest/_torch/executor/test_pytorch_model_engine.py",
     "tests/unittest/_torch/executor/test_resource_manager.py",
     "tests/unittest/_torch/executor/test_router_dealer_ipc.py",
     "tests/unittest/_torch/helpers.py",

ruff-legacy.toml

@@ -539,7 +539,6 @@ include = [
     "tests/unittest/_torch/executor/test_chunked_logits.py",
     "tests/unittest/_torch/executor/test_executor_request_queue.py",
     "tests/unittest/_torch/executor/test_overlap_scheduler.py",
-    "tests/unittest/_torch/executor/test_pytorch_model_engine.py",
     "tests/unittest/_torch/executor/test_resource_manager.py",
     "tests/unittest/_torch/executor/test_router_dealer_ipc.py",
     "tests/unittest/_torch/helpers.py",

tests/integration/defs/.test_durations

@@ -1110,8 +1110,8 @@
    "test_unittests.py::test_unittests_v2[unittest/_torch/executor/test_overlap_scheduler.py]": 449.18887196765496,
    "unittest/_torch/executor/test_profile_endpoints.py": 27.268500667556744,
    "test_unittests.py::test_unittests_v2[unittest/_torch/executor/test_profile_endpoints.py]": 27.268500667556744,
-   "unittest/_torch/executor/test_pytorch_model_engine.py -k \"test_spec_decode_graph_step_gather_kernel\"": 23.081429347826088,
-   "test_unittests.py::test_unittests_v2[unittest/_torch/executor/test_pytorch_model_engine.py -k \"test_spec_decode_graph_step_gather_kernel\"]": 23.081429347826088,
+   "unittest/_torch/executor/engine/runners/decoder/test_inputs.py -k \"test_spec_decode_graph_step_gather_kernel\"": 23.081429347826088,
+   "test_unittests.py::test_unittests_v2[unittest/_torch/executor/engine/runners/decoder/test_inputs.py -k \"test_spec_decode_graph_step_gather_kernel\"]": 23.081429347826088,
    "unittest/_torch/executor/test_resource_manager.py": 31.571532684283728,
    "test_unittests.py::test_unittests_v2[unittest/_torch/executor/test_resource_manager.py]": 31.571532684283728,
    "unittest/_torch/executor/test_seq_slot_sizing.py": 27.570640750670243,

tests/integration/test_lists/test-db/l0_b200.yml

@@ -342,7 +342,7 @@ l0_b200:
   # One-model speculative decoding step-copy kernels and the engine's decode
   # step that runs them (SM 100; they skip on other GPUs).
   - unittest/_torch/kernels/test_spec_step_copies.py
-  - unittest/_torch/executor/test_pytorch_model_engine.py -k "test_spec_decode_graph_step_gather_kernel"
+  - unittest/_torch/executor/engine/runners/decoder/test_inputs.py -k "test_spec_decode_graph_step_gather_kernel"
   # Kimi K3 DSpark drafter attention (CuTe DSL, SM 100) and its modeling_v2 catalog entries, plus the sm_100 cells of
   # fused_qk_norm_rope, which the drafter calls.
   - unittest/_torch/kernels/kimi_k3/test_k3_drafter_attn.py

tests/integration/test_lists/waives.txt

@@ -146,6 +146,9 @@ unittest/_torch/attention/multi_gpu/test_mha_helix.py::test_mha_helix_distribute
 unittest/_torch/attention/multi_gpu/test_mha_helix.py::test_mha_helix_distributed[nccl-scenario: Scenario(dtype=torch.bfloat16, kv_cache_dtype=torch.bfloat16, num_layers=1, num_heads=16, num_kv_heads=4, head_dim=128, hidden_size=2048, rope_theta=10000.0, kv_cache_tokens_per_block=32, bias=False, batch=1, ctx_len=512, atol=0.1, rtol=0.05)] SKIP (https://nvbugs/6991647)
 unittest/_torch/attention/multi_gpu/test_mha_helix.py::test_mha_helix_distributed[nccl-scenario: Scenario(dtype=torch.bfloat16, kv_cache_dtype=torch.bfloat16, num_layers=1, num_heads=16, num_kv_heads=4, head_dim=128, hidden_size=2048, rope_theta=10000.0, kv_cache_tokens_per_block=32, bias=False, batch=1, ctx_len=64, atol=0.1, rtol=0.05)] SKIP (https://nvbugs/6991647)
 unittest/_torch/attention/test_attention_mla.py::test_attention_mla SKIP (https://nvbugs/6862662)
+unittest/_torch/executor/engine/runners/decoder/test_inputs.py::DecoderInputPreparerTestCase::test_mrope_mixed_batched_and_fallback_positions_match_fallback SKIP (https://nvbugs/7003540)
+unittest/_torch/executor/engine/runners/decoder/test_inputs.py::DecoderInputPreparerTestCase::test_mrope_multi_beam_generation_positions SKIP (https://nvbugs/7003540)
+unittest/_torch/executor/engine/runners/decoder/test_inputs.py::DecoderInputPreparerTestCase::test_prepare_inputs_with_helix_parallelism SKIP (https://nvbugs/6818710)
 unittest/_torch/executor/kv_cache/test_kv_cache_v2_vocab_size.py::test_branch_only_enables_multimodal_keys_when_needed[dense-encoder_disable_noop] SKIP (https://nvbugs/6818710)
 unittest/_torch/executor/kv_cache/test_kv_cache_v2_vocab_size.py::test_branch_only_enables_multimodal_keys_when_needed[dense-multimodal] SKIP (https://nvbugs/6818710)
 unittest/_torch/executor/kv_cache/test_kv_cache_v2_vocab_size.py::test_branch_only_enables_multimodal_keys_when_needed[dense-text_only] SKIP (https://nvbugs/6818710)
@@ -157,9 +160,6 @@ unittest/_torch/executor/kv_cache/test_kv_cache_v2_vocab_size.py::test_branch_on
 unittest/_torch/executor/kv_cache/test_kv_cache_v2_vocab_size.py::test_branch_only_enables_multimodal_keys_when_needed[qwen3_5_hybrid-text_only] SKIP (https://nvbugs/6818710)
 unittest/_torch/executor/kv_cache/test_kv_cache_v2_vocab_size.py::test_unresolvable_vocab_size_only_warns_when_multimodal_keys_are_needed[False] SKIP (https://nvbugs/6818710)
 unittest/_torch/executor/kv_cache/test_kv_cache_v2_vocab_size.py::test_unresolvable_vocab_size_only_warns_when_multimodal_keys_are_needed[True] SKIP (https://nvbugs/6818710)
-unittest/_torch/executor/test_pytorch_model_engine.py::DecoderForwardTestCase::test_mrope_mixed_batched_and_fallback_positions_match_fallback SKIP (https://nvbugs/7003540)
-unittest/_torch/executor/test_pytorch_model_engine.py::DecoderForwardTestCase::test_mrope_multi_beam_generation_positions SKIP (https://nvbugs/7003540)
-unittest/_torch/executor/test_pytorch_model_engine.py::PyTorchModelEngineTestCase::test_prepare_inputs_with_helix_parallelism SKIP (https://nvbugs/6818710)
 unittest/_torch/kernels/parallel/test_fp4_linear.py::test_fp4_gemm_bias_per_backend[mnk2-cublaslt] SKIP (https://nvbugs/6581067)
 unittest/_torch/kernels/parallel/test_fp4_linear.py::test_fp4_gemm_bias_per_backend[mnk2-cutlass] SKIP (https://nvbugs/6581067)
 unittest/_torch/kernels/parallel/test_fp4_linear.py::test_fp4_gemm_bias_per_backend[mnk3-cublaslt] SKIP (https://nvbugs/6581067)

tests/unittest/_torch/executor/engine/runners/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0

tests/unittest/_torch/executor/engine/runners/decoder/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0

tests/unittest/_torch/executor/engine/runners/decoder/test_dummy_requests.py

@@ -0,0 +1,62 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+
+import pytest
+import torch
+from _torch.executor.engine.runners.decoder.utils import decoder_config
+from _torch.executor.engine_test_utils import DummyModel, kv_cache_manager
+
+from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.dummy_requests import DummyRequestFactory
+from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType
+from tensorrt_llm.llmapi.llm_args import (
+    DecodingBaseConfig,
+    DraftTargetDecodingConfig,
+    KvCacheConfig,
+    PARDDecodingConfig,
+)
+
+
+@pytest.mark.parametrize("config_cls", [DraftTargetDecodingConfig, PARDDecodingConfig])
+@pytest.mark.parametrize(
+    "draft_len", [None, 0, 3, 1], ids=["general", "cuda_graph_0", "cuda_graph_3", "cuda_graph_1"]
+)
+def test_warmup_builders_resynchronize_stale_draft_buffers(
+    config_cls: type[DecodingBaseConfig], draft_len: int | None
+) -> None:
+    spec = config_cls(max_draft_len=3, speculative_model="dummy", draft_len_schedule={1: 3, 4: 2})
+    config = decoder_config(spec_config=spec)
+    model = DummyModel(torch.half)
+    requests = DummyRequestFactory(
+        config=config,
+        model=model,
+        get_runtime_tokens_per_gen_step=spec.get_runtime_tokens_per_gen_step,
+        release_request_state=None,
+    )
+    with kv_cache_manager(
+        model=model,
+        config=KvCacheConfig(max_tokens=258, use_kv_cache_manager_v2=True),
+        max_num_tokens=258,
+    ) as manager:
+        resource_manager = ResourceManager({ResourceManagerType.KV_CACHE_MANAGER: manager})
+        batch_size = 2
+        if draft_len is None:
+            expected_draft_len = config.max_draft_len
+            warmup_request = requests.create_warmup_request(
+                resource_manager,
+                num_tokens=batch_size * spec.tokens_per_gen_step,
+                num_gen_requests=batch_size,
+            )
+        else:
+            expected_draft_len = draft_len
+            warmup_request = requests.create_cuda_graph_warmup_request(
+                resource_manager, batch_size, draft_len, max_seq_len=16
+            )
+        with requests.release_batch_context(warmup_request, resource_manager) as batch:
+            assert batch is not None
+            assert len(batch.generation_requests) == batch_size
+            expected_buffer_width = spec.get_runtime_tokens_per_gen_step(expected_draft_len) - 1
+            assert all(
+                len(request.py_draft_tokens) == expected_buffer_width
+                for request in batch.generation_requests
+            )