[ModelRunner V2] Deduplicate ModelState init logic (#46776)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Jun 26, 2026 at 16:32 UTC
b94f212e37f4ddf4b5e1cc96cd87217f36e3ec0c
4 files changed
+26
-77
vllm/model_executor/models/diffusion_gemma.py
+1
-27
@@ -776,33 +776,7 @@ class DiffusionGemmaModelState(ModelState):
776
encoder_cache: Any,
777
device: torch.device,
778
) -> None:
779
- self.vllm_config = vllm_config
780
- self.model_config = vllm_config.model_config
781
- self.scheduler_config = vllm_config.scheduler_config
782
- self.model = model
783
- self.device = device
784
-
785
- self.supports_mm_inputs = encoder_cache is not None
786
- self.max_num_reqs = self.scheduler_config.max_num_seqs
787
- self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
788
- self.max_model_len = self.model_config.max_model_len
789
- self.inputs_embeds_size = self.model_config.get_inputs_embeds_size()
790
- self.dtype = self.model_config.dtype
791
-
792
- if self.supports_mm_inputs:
793
- from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
794
- from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
795
-
796
- assert isinstance(encoder_cache, EncoderCache)
797
- self.encoder_cache = encoder_cache
798
- self.encoder_runner = EncoderRunner(
799
- model=self.model,
800
- max_num_tokens=self.max_num_tokens,
801
- hidden_size=self.inputs_embeds_size,
802
- encoder_cache=encoder_cache,
803
- dtype=self.dtype,
804
- device=self.device,
805
- )
779
+ super().__init__(vllm_config, model, encoder_cache, device)
780
781
# Per-step MM data produced by get_mm_embeddings and consumed by
782
# prepare_inputs. Stored as raw (mm_embeds, is_mm_embed) so that
vllm/v1/worker/gpu/model_states/default.py
+1
-25
@@ -12,7 +12,6 @@ from vllm.v1.kv_cache_interface import KVCacheConfig
12
from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
13
from vllm.v1.worker.gpu.input_batch import InputBatch
14
from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
15
-from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
15
from vllm.v1.worker.gpu.mm.rope import get_rope_state
16
from vllm.v1.worker.gpu.model_states.interface import ModelState
17
from vllm.v1.worker.gpu.model_states.mm_pruning import maybe_create_mm_pruner
@@ -28,30 +27,7 @@ class DefaultModelState(ModelState):
27
encoder_cache: EncoderCache | None,
28
device: torch.device,
29
):
31
- self.vllm_config = vllm_config
32
- self.model_config = vllm_config.model_config
33
- self.scheduler_config = vllm_config.scheduler_config
34
- self.model = model
35
- self.device = device
36
-
37
- self.supports_mm_inputs = encoder_cache is not None
38
- self.max_model_len = self.model_config.max_model_len
39
- self.max_num_reqs = self.scheduler_config.max_num_seqs
40
- self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
41
- self.inputs_embeds_size = self.model_config.get_inputs_embeds_size()
42
- self.dtype = self.model_config.dtype
43
-
44
- if self.supports_mm_inputs:
45
- assert encoder_cache is not None
46
- self.encoder_cache = encoder_cache
47
- self.encoder_runner = EncoderRunner(
48
- model=self.model,
49
- max_num_tokens=self.max_num_tokens,
50
- hidden_size=self.inputs_embeds_size,
51
- encoder_cache=encoder_cache,
52
- dtype=self.dtype,
53
- device=self.device,
54
- )
30
+ super().__init__(vllm_config, model, encoder_cache, device)
31
32
self.rope_state = get_rope_state(
33
self.model_config,
vllm/v1/worker/gpu/model_states/encoder_decoder.py
+1
-19
@@ -13,7 +13,6 @@ from vllm.v1.kv_cache_interface import CrossAttentionSpec, KVCacheConfig
13
from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
14
from vllm.v1.worker.gpu.input_batch import InputBatch
15
from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
16
-from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
16
from vllm.v1.worker.gpu.model_states.interface import (
17
ModelSpecificAttnMetadata,
18
ModelState,
@@ -53,25 +52,8 @@ class EncoderDecoderModelState(ModelState):
52
encoder_cache: EncoderCache | None,
53
device: torch.device,
54
) -> None:
56
- self.vllm_config = vllm_config
57
- self.model_config = vllm_config.model_config
58
- self.scheduler_config = vllm_config.scheduler_config
59
- self.model = model
60
- self.max_num_reqs = vllm_config.scheduler_config.max_num_seqs
61
- self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
62
- self.max_model_len = self.model_config.max_model_len
63
- self.device = device
64
-
55
assert encoder_cache is not None
66
- self.encoder_cache = encoder_cache
67
- self.encoder_runner = EncoderRunner(
68
- model=self.model,
69
- max_num_tokens=self.max_num_tokens,
70
- hidden_size=self.model_config.get_inputs_embeds_size(),
71
- encoder_cache=self.encoder_cache,
72
- dtype=self.model_config.dtype,
73
- device=self.device,
74
- )
56
+ super().__init__(vllm_config, model, encoder_cache, device)
57
58
self.max_encoder_len = getattr(
59
self.model_config.hf_config,
vllm/v1/worker/gpu/model_states/interface.py
+23
-6
@@ -37,7 +37,6 @@ class ModelSpecificAttnMetadata:
37
38
39
class ModelState(ABC):
40
- @abstractmethod
40
def __init__(
41
self,
42
vllm_config: VllmConfig,
@@ -45,11 +44,29 @@ class ModelState(ABC):
44
encoder_cache: EncoderCache | None,
45
device: torch.device,
46
) -> None:
48
- raise NotImplementedError
49
-
50
- model: nn.Module
51
- # Set by mm-capable states; used by the default gather_mm_embeddings().
52
- encoder_runner: EncoderRunner
47
+ self.vllm_config = vllm_config
48
+ self.model_config = vllm_config.model_config
49
+ self.scheduler_config = vllm_config.scheduler_config
50
+ self.model = model
51
+ self.device = device
52
+
53
+ self.max_model_len = self.model_config.max_model_len
54
+ self.max_num_reqs = self.scheduler_config.max_num_seqs
55
+ self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
56
+ self.inputs_embeds_size = self.model_config.get_inputs_embeds_size()
57
+ self.dtype = self.model_config.dtype
58
+
59
+ self.supports_mm_inputs = encoder_cache is not None
60
+ if encoder_cache is not None:
61
+ self.encoder_cache = encoder_cache
62
+ self.encoder_runner = EncoderRunner(
63
+ model=self.model,
64
+ max_num_tokens=self.max_num_tokens,
65
+ hidden_size=self.inputs_embeds_size,
66
+ encoder_cache=encoder_cache,
67
+ dtype=self.dtype,
68
+ device=self.device,
69
+ )
70
71
def get_supported_generation_tasks(self) -> tuple[GenerationTask, ...]:
72
from vllm.model_executor.models.interfaces import (