[TRTLLM-15765][refactor] Remove the ModelEngine ABC and settle the executor interface

docs/source/llm-api/index.md

@@ -139,7 +139,7 @@ The `py_executor` property contains the timed scopes in `create_py_executor()`,
 | `worker_start_seconds` | Start the final PyExecutor worker. |
 | `total_py_executor_creation_seconds` | Total time for the `create_py_executor()` call, including all applicable scopes above. |
 
-The two model-engine properties contain times for various `PyTorchModelEngine` warmup stages.
+The two model-engine properties contain times for various `ModelEngine` warmup stages.
 `initial_model_engine` measures the timings for the initial model engine creation as part of the
 `py_executor.initial_py_executor_creation_seconds_for_kv_cache_estimation` timing.
 Likewise, `final_model_engine` measures the timings for the final model engine creation as part of

docs/source/torch/arch_overview.md

@@ -37,9 +37,14 @@ shared data contracts, process boundaries, and request ownership.
 
 ## Model Engine
 
-The core component of `PyExecutor` is the `ModelEngine`, responsible for executing the model's forward pass efficiently on the GPU.
-The key method of `ModelEngine` is `forward`, which handles the forward pass computation.
-For the PyTorch backend, the derived class is `PyTorchModelEngine`, declared in [model_engine.py](../../../tensorrt_llm/_torch/pyexecutor/model_engine.py).
+[`ModelEngine`](../../../tensorrt_llm/_torch/pyexecutor/model_engine.py) loads the model, resolves capacity and execution settings, constructs a model runner, and coordinates its lifecycle.
+`PyExecutor` calls `ModelEngine.forward` with scheduled requests, runtime resources, and the current speculative decoding state; the engine builds `ScheduledInputs` and delegates to a `ScheduledModelRunner`.
+`EncoderExecutor` instead passes `PackedInputs` through the engine to a `PackedModelRunner`.
+
+The concrete runner owns input preparation, model execution, CUDA graphs, metadata, and output processing.
+`EncoderDecoderRunner` combines decoder execution with a separate encoder stage, which `PyExecutor` schedules through explicit encoder entry points.
+Multimodal item encoding also has a separate scheduled entry point.
+The engine retains the shared services and coordinates warmup, late-installed hooks, and resource release.
 
 ## Decoder
 

docs/source/torch/kv_cache_manager.md

@@ -37,13 +37,13 @@ There are also two interfaces designed for `CapacityScheduler`:
   `CapacityScheduler` uses this to sum up the total resources needed and determine if new requests can be accommodated.
 
 In addition to the `BaseResourceManager` interfaces, `KVCacheManager` has interfaces related to the `ModelEngine` in use.
-For `PyTorchModelEngine`, common interfaces include:
+For `ModelEngine`, common interfaces include:
 
 - **get_batch_cache_indices**: Takes a list of `LlmRequest` and returns a `Dict[List[int]]`, indicating the block IDs for each request.
 - **get_buffers**: Returns the buffer of the KV Cache pool for a given layer, used by the attention backend. The shape might be [`num_blocks`, 2, `num_tokens_per_block`, `num_kv_heads`, `head_dim`].
 - **get_num_free_blocks**: Returns the number of free blocks available for allocation.
 
-There are also interfaces for warming up `PyTorchModelEngine`, especially when using CUDA graphs:
+There are also interfaces for warming up `ModelEngine`, especially when using CUDA graphs:
 
 - **add_padding_request**: Adds a sequence of context length 1 to KV Cache as a warmup request.
   This is optional if CUDA Graph is not used in your proof of concept.

tensorrt_llm/_torch/distributed/ops.py

@@ -377,7 +377,7 @@ def _get_or_scale_allreduce_mnnvl_workspace(
                 f"[MNNVL] AllReduce workspace must grow to {buffer_size_bytes} "
                 "bytes inside a torch.compile'd region, but allocation cannot "
                 "run under tracing. Pre-scale the workspace before compiling "
-                "via MNNVLAllReduce.prescale_workspace (PyTorchModelEngine "
+                "via MNNVLAllReduce.prescale_workspace (ModelEngine "
                 "does this for max_num_tokens x hidden_size).")
         # Initial buffer to be large enough to support 1024 tokens * 8192 hidden_dim
         init_buffer_size_bytes = max(1024 * 8192 * elem_size, buffer_size_bytes

tensorrt_llm/_torch/memory/gpu_memory_backend.py

@@ -531,7 +531,7 @@ def cleanup(self) -> None:
         in upstream `client.close()` or `evict_gms_client_memory_manager`
         is logged (warning for the outer error, debug for `close()`)
         and swallowed so callers — including `__del__` paths in
-        `ModelLoader` and `PyTorchModelEngine` — never raise from
+        `ModelLoader` and `ModelEngine` — never raise from
         teardown.
 
         After return:

tensorrt_llm/_torch/models/modeling_multimodal_encoder.py

@@ -132,7 +132,7 @@ class MultimodalEncoderMixin:
     """Encoder-side counterpart to ``MultimodalModelMixin``.
 
     Marker + default ``setup_attn_metadata`` for multimodal encoders whose
-    ``AttentionMetadata`` is built by ``PyTorchModelEngine`` after model load
+    ``AttentionMetadata`` is built by ``ModelEngine`` after model load
     using the runtime encoder token budget (``max_num_tokens``).
 
     Subclasses set ``metadata_cls`` in their own ``__init__`` (typically from

tensorrt_llm/_torch/models/modeling_vila.py

@@ -1313,7 +1313,7 @@ def get_llm(self):
 
     @property
     def llm_checkpoint_dir(self):
-        """Return the directory of the LLM checkpoint. Workaround for multimodal models that config is in root directory but llm checkpoint is in subdirectory. Without this, PyTorchModelEngine will try to load LLM checkpoint from root."""
+        """Return the directory of the LLM checkpoint. Workaround for multimodal models that config is in root directory but llm checkpoint is in subdirectory. Without this, ModelEngine will try to load LLM checkpoint from root."""
         return self.llm_path
 
     def load_weights(self, weights):

tensorrt_llm/_torch/modules/linear.py

@@ -3422,7 +3422,7 @@ def __init__(self) -> None:
         super().__init__()
         self.use_cutlass = _mxfp8_cutlass_op_available()
         self.backend = os.environ.get("TRTLLM_MXFP8_GEMM_BACKEND", "trtllm")
-        # Only PyTorchModelEngine owns the startup tuning lifecycle. Keep
+        # Only ModelEngine owns the startup tuning lifecycle. Keep
         # standalone modules and engine paths that skip warmup (for example,
         # Helix CP) on the direct native op instead of leaving Python
         # AutoTuner dispatch armed indefinitely.

tensorrt_llm/_torch/moe/fused_moe/cutlass/lora.py

@@ -284,7 +284,7 @@ def _extract_moe_lora_tensors(self, lora_params: dict | None) -> dict[str, objec
 
         # Gather (ranks, weight_ptrs) per kernel slot. weight_pointers is built
         # flat ([num_seqs * 3], row-major (A, B, DoRA) per seq) in
-        # PyTorchModelEngine._build_lora_params; the op expects [num_seqs, 3].
+        # ModelEngine._build_lora_params; the op expects [num_seqs, 3].
         active_max_rank = 0
 
         def _source(module_type: LoraModuleType) -> _RanksAndPtrs | None:

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -73,7 +73,7 @@
                                            MixedMambaHybridCacheManager,
                                            use_py_mamba_cache_manager)
 from .llm_request import ExecutorResponse, LlmRequestState
-from .model_engine import PyTorchModelEngine
+from .model_engine import ModelEngine
 from .py_executor import PyExecutor
 from .resource_manager import (KVCacheCompressionManager, KVCacheManager,
                                PeftCacheManager, ResourceManager,
@@ -920,7 +920,7 @@ class KvCacheCreator:
     def __init__(
         self,
         *,
-        model_engine: PyTorchModelEngine,
+        model_engine: ModelEngine,
         mapping: Mapping,
         net_max_seq_len: int,
         kv_connector_manager: Optional[KvCacheConnectorManager],
@@ -977,7 +977,7 @@ def __init__(
 
     def _get_model_kv_cache_manager_cls(
         self,
-        model_engine: PyTorchModelEngine,
+        model_engine: ModelEngine,
         kv_cache_config_override: Optional[KvCacheConfig] = None,
     ):
         kv_cache_config = (kv_cache_config_override if kv_cache_config_override
@@ -2040,7 +2040,7 @@ def configure_kv_cache_capacity(self,
 
     def _create_kv_cache_manager(
         self,
-        model_engine: PyTorchModelEngine,
+        model_engine: ModelEngine,
         estimating_kv_cache: bool = False,
         kv_cache_config_override: Optional[KvCacheConfig] = None,
         cold_page_codec_provider: Optional[object] = None,
@@ -2581,10 +2581,8 @@ def _warn_if_admitted_concurrency_outruns_graph_capture(self) -> None:
         throughput than the extra concurrency buys; still valid at low load,
         hence a warning that names the knobs rather than a refusal.
         """
-        capture_sizes = getattr(self._model_engine, "_cuda_graph_batch_sizes",
-                                None)
-        padding = getattr(self._model_engine, "_cuda_graph_padding_enabled",
-                          None)
+        capture_sizes = self._model_engine.cuda_graph_batch_sizes
+        padding = self._model_engine.cuda_graph_padding_enabled
         if (not isinstance(capture_sizes, (list, tuple)) or not capture_sizes
                 or not all(isinstance(s, int) for s in capture_sizes)):
             # CUDA graphs disabled, or the engine does not expose its capture
@@ -3244,7 +3242,7 @@ def _mamba_conv_layout_kwargs(kv_cache_manager_cls: type,
 
 
 def _create_kv_cache_manager(
-        model_engine: Optional[PyTorchModelEngine],
+        model_engine: Optional[ModelEngine],
         kv_cache_manager_cls,
         mapping: Mapping,
         kv_cache_config: KvCacheConfig,
@@ -3808,7 +3806,7 @@ def validate_kv_cache_compression_compatibility(
 def create_kv_cache_compression_manager(
     config: Optional[KvCacheCompressionConfig],
     *,
-    model_engine: PyTorchModelEngine,
+    model_engine: ModelEngine,
     kv_cache_config: KvCacheConfig,
     estimating_kv_cache: bool = False,
 ) -> Optional[KVCacheCompressionManager]:
@@ -4467,7 +4465,7 @@ def create_torch_sampler_args(
 
 
 def instantiate_sampler(
-    engine: PyTorchModelEngine,
+    engine: ModelEngine,
     llm_args: TorchLlmArgs,
     mapping: Mapping,
     *,

tensorrt_llm/_torch/pyexecutor/encoder_executor.py

@@ -12,14 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict
+from typing import TYPE_CHECKING, Any, Dict
 
 import torch
 
 from tensorrt_llm.logger import logger
 
 from .engine.runners.interface import PackedInputs
 
+if TYPE_CHECKING:
+    from ..distributed import Distributed
+    from .model_engine import ModelEngine
+
 
 class EncoderExecutor:
     """Executor for models using the encode-only path.
@@ -31,7 +35,7 @@ class EncoderExecutor:
     and no request queue. It runs entirely on the calling thread.
     """
 
-    def __init__(self, model_engine, dist):
+    def __init__(self, model_engine: "ModelEngine", dist: "Distributed") -> None:
         self.model_engine = model_engine
         self.dist = dist
 

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

@@ -1,7 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-"""Decoder execution moved out of ``PyTorchModelEngine``."""
+"""Decoder execution moved out of ``ModelEngine``."""
 
 import contextlib
 import inspect

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

@@ -10,7 +10,7 @@
 from collections.abc import Callable, Iterator
 from dataclasses import dataclass, field
 from functools import partial
-from typing import Any, cast
+from typing import Any
 
 import torch
 from torch import nn
@@ -61,7 +61,7 @@
 
 
 @dataclass(frozen=True, kw_only=True)
-class EncoderConfigMixin:
+class EncoderConfig(RunnerConfig):
     """Encoder-specific fields shared by runner configuration types."""
 
     enable_autotuner: bool
@@ -220,7 +220,7 @@ def create(
 
 
 @dataclass(frozen=True)
-class EncoderRunnerConfig(EncoderConfigMixin, RunnerConfig):
+class EncoderRunnerConfig(EncoderConfig):
     """Configuration for the EncoderOnly model runner."""
 
     @classmethod
@@ -272,7 +272,7 @@ class EncoderMixin:
     def _initialize_encoder(
         self,
         model: nn.Module,
-        config: EncoderConfigMixin,
+        config: EncoderConfig,
         *,
         mapping: Mapping,
         dist: Distributed | None,
@@ -284,7 +284,7 @@ def _initialize_encoder(
         self._dist = dist
         self._moe_load_balancer = moe_load_balancer
         self._attention_runtime_features = attention_runtime_features
-        self._config = cast(RunnerConfig, config)
+        self._config = config
         self._encoder_config = config
         self._initialize_encoder_cuda_graph()
 

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

@@ -53,12 +53,12 @@
 from .decoder.dummy_requests import DummyRequestFactory
 from .decoder.graphs import DecoderCudaGraphs
 from .decoder.inputs import DecoderInputPreparer
-from .encoder import EncoderConfigMixin, EncoderMixin, EncoderPreparedInputs
-from .interface import PreparedInputs, RunnerConfig, ScheduledInputs
+from .encoder import EncoderConfig, EncoderMixin, EncoderPreparedInputs
+from .interface import PreparedInputs, ScheduledInputs
 
 
 @dataclass(frozen=True)
-class EncoderStageConfig(EncoderConfigMixin, RunnerConfig):
+class EncoderStageConfig(EncoderConfig):
     """Configuration for the encoder stage of an encoder-decoder model."""
 
 

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -4,7 +4,6 @@
 import functools
 import math
 import os
-from abc import ABC, abstractmethod
 from collections import defaultdict
 from contextlib import contextmanager
 from typing import (Any, Callable, Dict, Hashable, Iterator, List, Optional,
@@ -79,30 +78,6 @@
 from .sampler import SampleStateTensors
 from .scheduler import ScheduledRequests
 
-
-class ModelEngine(ABC):
-
-    @abstractmethod
-    def get_max_num_sequences(self) -> int:
-        raise NotImplementedError
-
-    @abstractmethod
-    def forward(self,
-                scheduled_requests: ScheduledRequests,
-                resource_manager: Optional[ResourceManager],
-                new_tensors_device: Optional[SampleStateTensors],
-                cache_indirection_buffer: Optional[torch.Tensor] = None):
-        raise NotImplementedError
-
-    def warmup(self, resource_manager: Optional[ResourceManager]) -> None:
-        """
-        This method is called after runtime resources are initialized. The
-        resource manager is absent for drivers that allocate none. Override to
-        perform warmup actions: instantiating CUDA graphs, torch.compile, etc.
-        """
-        return
-
-
 # Refit only needs the compile wrapper removed before weights are loaded: the
 # wrapper inserts "_orig_mod" into every parameter path, and load_weights()
 # matches checkpoint tensors by dotted path (modeling_utils.py:1141-1160). With
@@ -136,7 +111,8 @@ def _configure_deep_gemm_pdl() -> None:
     _DEEP_GEMM_PDL_CONFIGURED = True
 
 
-class PyTorchModelEngine(ModelEngine):
+class ModelEngine:
+    """Own model setup and coordinate runner execution and resource lifetimes."""
 
     def __init__(
         self,
@@ -579,7 +555,7 @@ def _init_execution_services(self) -> None:
             self.max_total_draft_tokens = 0
 
         self._metrics: dict[str, float] = defaultdict(float)
-        self.iter_states = {}
+        self.iter_states: dict[str, Any] = {}
         self.lora_model_config: Optional[LoraModelConfig] = None
         self._lora_owner: Optional[DecoderLoraOwner] = None
         self._lora_model: Optional[torch.nn.Module] = None
@@ -793,28 +769,41 @@ def cuda_graph_runner(self) -> Optional[CUDAGraphRunner]:
             return None
         return self._runner.cuda_graph_runner
 
+    @property
+    def cuda_graph_batch_sizes(self) -> list[int]:
+        """Return the resolved decode graph batch sizes."""
+        return self._cuda_graph_batch_sizes
+
+    @property
+    def cuda_graph_padding_enabled(self) -> bool:
+        """Return whether decode graph batches are padded to capture sizes."""
+        return self._cuda_graph_padding_enabled
+
     @property
     def metrics(self) -> dict[str, float]:
         """Return model-engine warmup time metrics."""
         return self._metrics
 
-    def register_forward_pass_callable(self, callable: Callable):
+    def register_forward_pass_callable(self, callable: Callable) -> None:
+        """Install the KV connector hook on the decoder runner."""
         if isinstance(self._runner, DecoderRunner):
             self._runner.forward_pass_callable = callable
 
-    def register_sample_in_graph_callable(self, callable: Optional[Callable]):
+    def register_sample_in_graph_callable(self,
+                                          callable: Optional[Callable]) -> None:
         """Register the hook that samples at the tail of the forward graph."""
         if isinstance(self._runner, DecoderRunner):
             self._runner.sample_in_graph_callable = callable
 
-    def register_beam_histories_callable(self, callable: Optional[Callable]):
+    def register_beam_histories_callable(self,
+                                         callable: Optional[Callable]) -> None:
         """Register where logits post-processors get live beam histories."""
         if isinstance(self._runner, DecoderRunner):
             self._runner.start_beam_histories_callable = callable
 
     def register_sample_type_resolver(self,
                                       resolver: Optional[Callable],
-                                      stage: Optional[Callable] = None):
+                                      stage: Optional[Callable] = None) -> None:
         """Register how a batch maps to its sampling tier, for the graph key."""
         assert isinstance(self._runner, DecoderRunner)
         self._runner.register_sample_type_resolver(resolver, stage)
@@ -1104,8 +1093,8 @@ def __del__(self) -> None:
         try:
             self.cleanup()
         except (RuntimeError, AttributeError) as e:
-            logger.warning(
-                "PyTorchModelEngine cleanup failed during destruction: %s", e)
+            logger.warning("ModelEngine cleanup failed during destruction: %s",
+                           e)
 
     def _init_max_seq_len(self):
         # Allow user to override the inferred max_seq_len with a warning.
@@ -1445,7 +1434,7 @@ def _init_userbuffers(self, hidden_size):
 
         return True
 
-    def wait_for_input_copy(self):
+    def wait_for_input_copy(self) -> None:
         """
         Wait for input preparation and H2D copy of previous iteration before modifying host input,
         otherwise the input of previous iteration will be overwritten.