A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 12:45 UTC
[Model Runner V2] Add ModelStateInterface [4/N] (#35621)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Feb 28, 2026 at 13:19 UTC
e3eb146f7ad4bc920e11e98cf88cee3839cf5f89
5 files changed
+90
-4
vllm/v1/worker/gpu/cudagraph_utils.py
+1
-1
@@ -22,7 +22,7 @@ from vllm.v1.worker.gpu.attn_utils import (
22
from vllm.v1.worker.gpu.block_table import BlockTables
23
from vllm.v1.worker.gpu.dp_utils import make_num_tokens_across_dp
24
from vllm.v1.worker.gpu.input_batch import InputBuffers
25
-from vllm.v1.worker.gpu.model_states import ModelState
25
+from vllm.v1.worker.gpu.model_states.interface import ModelState
26
from vllm.v1.worker.utils import AttentionGroup
27
28
vllm/v1/worker/gpu/model_runner.py
+2
-2
@@ -78,7 +78,7 @@ from vllm.v1.worker.gpu.kv_connector import (
78
)
79
from vllm.v1.worker.gpu.lora_utils import LoraState
80
from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
81
-from vllm.v1.worker.gpu.model_states import ModelState
81
+from vllm.v1.worker.gpu.model_states import init_model_state
82
from vllm.v1.worker.gpu.pool.pooling_runner import PoolingRunner
83
from vllm.v1.worker.gpu.pp_utils import pp_broadcast, pp_receive
84
from vllm.v1.worker.gpu.sample.output import SamplerOutput
@@ -267,7 +267,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
267
prepare_communication_buffer_for_model(self.speculator)
268
269
# Initialize the components that require the model.
270
- self.model_state = ModelState(
270
+ self.model_state = init_model_state(
271
self.vllm_config, self.model, self.encoder_cache, self.device
272
)
273
if self.is_pooling_model:
vllm/v1/worker/gpu/model_states/__init__.py
new
+18
@@ -0,0 +1,18 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
+import torch.nn as nn
5
+
6
+from vllm.config import VllmConfig
7
+from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
8
+
9
+
10
+def init_model_state(
11
+ vllm_config: VllmConfig,
12
+ model: nn.Module,
13
+ encoder_cache: EncoderCache | None,
14
+ device: torch.device,
15
+):
16
+ from vllm.v1.worker.gpu.model_states.default import DefaultModelState
17
+
18
+ return DefaultModelState(vllm_config, model, encoder_cache, device)
vllm/v1/worker/gpu/model_states/default.py
renamed
+2
-1
@@ -13,11 +13,12 @@ from vllm.v1.worker.gpu.input_batch import InputBatch
13
from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
14
from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
15
from vllm.v1.worker.gpu.mm.mrope_utils import MRopeState
16
+from vllm.v1.worker.gpu.model_states.interface import ModelState
17
from vllm.v1.worker.gpu.states import RequestState
18
from vllm.v1.worker.utils import AttentionGroup
19
20
20
-class ModelState:
21
+class DefaultModelState(ModelState):
22
def __init__(
23
self,
24
vllm_config: VllmConfig,
vllm/v1/worker/gpu/model_states/interface.py
new
+67
@@ -0,0 +1,67 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from abc import ABC, abstractmethod
4
+from typing import Any
5
+
6
+import torch
7
+import torch.nn as nn
8
+
9
+from vllm.config import VllmConfig
10
+from vllm.v1.core.sched.output import NewRequestData
11
+from vllm.v1.kv_cache_interface import KVCacheConfig
12
+from vllm.v1.worker.gpu.input_batch import InputBatch
13
+from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
14
+from vllm.v1.worker.gpu.states import RequestState
15
+from vllm.v1.worker.utils import AttentionGroup
16
+
17
+
18
+class ModelState(ABC):
19
+ @abstractmethod
20
+ def __init__(
21
+ self,
22
+ vllm_config: VllmConfig,
23
+ model: nn.Module,
24
+ encoder_cache: EncoderCache | None,
25
+ device: torch.device,
26
+ ) -> None:
27
+ raise NotImplementedError
28
+
29
+ @abstractmethod
30
+ def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
31
+ raise NotImplementedError
32
+
33
+ @abstractmethod
34
+ def apply_staged_writes(self) -> None:
35
+ raise NotImplementedError
36
+
37
+ @abstractmethod
38
+ def get_mm_embeddings(
39
+ self,
40
+ scheduled_encoder_inputs: dict[str, list[int]],
41
+ input_batch: InputBatch,
42
+ req_states: RequestState,
43
+ ) -> torch.Tensor:
44
+ raise NotImplementedError
45
+
46
+ @abstractmethod
47
+ def prepare_inputs(
48
+ self, input_batch: InputBatch, req_states: RequestState
49
+ ) -> dict[str, torch.Tensor | None]:
50
+ raise NotImplementedError
51
+
52
+ @abstractmethod
53
+ def prepare_dummy_inputs(
54
+ self, num_reqs: int, num_tokens: int
55
+ ) -> dict[str, torch.Tensor | None]:
56
+ raise NotImplementedError
57
+
58
+ @abstractmethod
59
+ def prepare_attn(
60
+ self,
61
+ input_batch: InputBatch,
62
+ block_tables: tuple[torch.Tensor, ...],
63
+ slot_mappings: torch.Tensor,
64
+ attn_groups: list[list[AttentionGroup]],
65
+ kv_cache_config: KVCacheConfig,
66
+ ) -> dict[str, Any]:
67
+ raise NotImplementedError