[ModelRunner V2] Deduplicate ModelState init logic

vllm/model_executor/models/diffusion_gemma.py

@@ -758,33 +758,7 @@ def __init__(
         encoder_cache: Any,
         device: torch.device,
     ) -> None:
-        self.vllm_config = vllm_config
-        self.model_config = vllm_config.model_config
-        self.scheduler_config = vllm_config.scheduler_config
-        self.model = model
-        self.device = device
-
-        self.supports_mm_inputs = encoder_cache is not None
-        self.max_num_reqs = self.scheduler_config.max_num_seqs
-        self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
-        self.max_model_len = self.model_config.max_model_len
-        self.inputs_embeds_size = self.model_config.get_inputs_embeds_size()
-        self.dtype = self.model_config.dtype
-
-        if self.supports_mm_inputs:
-            from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
-            from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
-
-            assert isinstance(encoder_cache, EncoderCache)
-            self.encoder_cache = encoder_cache
-            self.encoder_runner = EncoderRunner(
-                model=self.model,
-                max_num_tokens=self.max_num_tokens,
-                hidden_size=self.inputs_embeds_size,
-                encoder_cache=encoder_cache,
-                dtype=self.dtype,
-                device=self.device,
-            )
+        super().__init__(vllm_config, model, encoder_cache, device)
 
         # Per-step MM data produced by get_mm_embeddings and consumed by
         # prepare_inputs.  Stored as raw (mm_embeds, is_mm_embed) so that

vllm/v1/worker/gpu/model_states/default.py

@@ -12,7 +12,6 @@
 from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
 from vllm.v1.worker.gpu.input_batch import InputBatch
 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
-from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
 from vllm.v1.worker.gpu.mm.rope import get_rope_state
 from vllm.v1.worker.gpu.model_states.interface import ModelState
 from vllm.v1.worker.gpu.model_states.mm_pruning import maybe_create_mm_pruner
@@ -28,30 +27,7 @@ def __init__(
         encoder_cache: EncoderCache | None,
         device: torch.device,
     ):
-        self.vllm_config = vllm_config
-        self.model_config = vllm_config.model_config
-        self.scheduler_config = vllm_config.scheduler_config
-        self.model = model
-        self.device = device
-
-        self.supports_mm_inputs = encoder_cache is not None
-        self.max_model_len = self.model_config.max_model_len
-        self.max_num_reqs = self.scheduler_config.max_num_seqs
-        self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
-        self.inputs_embeds_size = self.model_config.get_inputs_embeds_size()
-        self.dtype = self.model_config.dtype
-
-        if self.supports_mm_inputs:
-            assert encoder_cache is not None
-            self.encoder_cache = encoder_cache
-            self.encoder_runner = EncoderRunner(
-                model=self.model,
-                max_num_tokens=self.max_num_tokens,
-                hidden_size=self.inputs_embeds_size,
-                encoder_cache=encoder_cache,
-                dtype=self.dtype,
-                device=self.device,
-            )
+        super().__init__(vllm_config, model, encoder_cache, device)
 
         self.rope_state = get_rope_state(
             self.model_config,

vllm/v1/worker/gpu/model_states/encoder_decoder.py

@@ -13,7 +13,6 @@
 from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
 from vllm.v1.worker.gpu.input_batch import InputBatch
 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
-from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
 from vllm.v1.worker.gpu.model_states.interface import (
     ModelSpecificAttnMetadata,
     ModelState,
@@ -53,25 +52,8 @@ def __init__(
         encoder_cache: EncoderCache | None,
         device: torch.device,
     ) -> None:
-        self.vllm_config = vllm_config
-        self.model_config = vllm_config.model_config
-        self.scheduler_config = vllm_config.scheduler_config
-        self.model = model
-        self.max_num_reqs = vllm_config.scheduler_config.max_num_seqs
-        self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
-        self.max_model_len = self.model_config.max_model_len
-        self.device = device
-
         assert encoder_cache is not None
-        self.encoder_cache = encoder_cache
-        self.encoder_runner = EncoderRunner(
-            model=self.model,
-            max_num_tokens=self.max_num_tokens,
-            hidden_size=self.model_config.get_inputs_embeds_size(),
-            encoder_cache=self.encoder_cache,
-            dtype=self.model_config.dtype,
-            device=self.device,
-        )
+        super().__init__(vllm_config, model, encoder_cache, device)
 
         self.max_encoder_len = getattr(
             self.model_config.hf_config,

vllm/v1/worker/gpu/model_states/interface.py

@@ -37,19 +37,36 @@ def get_extra_attn_kwargs(
 
 
 class ModelState(ABC):
-    @abstractmethod
     def __init__(
         self,
         vllm_config: VllmConfig,
         model: nn.Module,
         encoder_cache: EncoderCache | None,
         device: torch.device,
     ) -> None:
-        raise NotImplementedError
-
-    model: nn.Module
-    # Set by mm-capable states; used by the default gather_mm_embeddings().
-    encoder_runner: EncoderRunner
+        self.vllm_config = vllm_config
+        self.model_config = vllm_config.model_config
+        self.scheduler_config = vllm_config.scheduler_config
+        self.model = model
+        self.device = device
+
+        self.max_model_len = self.model_config.max_model_len
+        self.max_num_reqs = self.scheduler_config.max_num_seqs
+        self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
+        self.inputs_embeds_size = self.model_config.get_inputs_embeds_size()
+        self.dtype = self.model_config.dtype
+
+        self.supports_mm_inputs = encoder_cache is not None
+        if encoder_cache is not None:
+            self.encoder_cache = encoder_cache
+            self.encoder_runner = EncoderRunner(
+                model=self.model,
+                max_num_tokens=self.max_num_tokens,
+                hidden_size=self.inputs_embeds_size,
+                encoder_cache=encoder_cache,
+                dtype=self.dtype,
+                device=self.device,
+            )
 
     def get_supported_generation_tasks(self) -> tuple[GenerationTask, ...]:
         from vllm.model_executor.models.interfaces import (