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] Minor refactoring for EncoderRunner (#35628)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Mar 1, 2026 at 00:15 UTC
da543d1abe2468a1b79f230e91e8bbdc2bf6ee71
2 files changed
+13
-13
vllm/v1/worker/gpu/mm/encoder_runner.py
+4
-11
@@ -13,12 +13,14 @@ from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
13
class EncoderRunner:
14
def __init__(
15
self,
16
+ model: SupportsMultiModal,
17
max_num_tokens: int,
18
hidden_size: int,
19
encoder_cache: EncoderCache,
20
dtype: torch.dtype,
21
device: torch.device,
22
):
23
+ self.model = model
24
self.max_num_tokens = max_num_tokens
25
self.hidden_size = hidden_size
26
self.encoder_cache = encoder_cache
@@ -48,25 +50,17 @@ class EncoderRunner:
50
@torch.inference_mode()
51
def execute_mm_encoder(
52
self,
51
- model: SupportsMultiModal,
52
- mm_hashes: list[str],
53
mm_kwargs: list[tuple[str, MultiModalKwargsItem]],
54
) -> list[torch.Tensor]:
55
- if not mm_hashes:
56
- return []
57
-
55
encoder_outputs: list[torch.Tensor] = []
56
for modality, num_items, mm_kwargs_group in group_mm_kwargs_by_modality(
57
mm_kwargs, device=self.device, pin_memory=False
58
):
62
- curr_group_outputs = model.embed_multimodal(**mm_kwargs_group)
59
+ curr_group_outputs = self.model.embed_multimodal(**mm_kwargs_group)
60
sanity_check_mm_encoder_outputs(
61
curr_group_outputs, expected_num_items=num_items
62
)
63
encoder_outputs.extend(curr_group_outputs)
67
-
68
- # Cache the encoder outputs by mm_hash
69
- self.encoder_cache.encoder_outputs.update(zip(mm_hashes, encoder_outputs))
64
return encoder_outputs
65
66
def gather_mm_embeddings(
@@ -146,12 +140,11 @@ class EncoderRunner:
140
@torch.inference_mode()
141
def get_inputs_embeds(
142
self,
149
- model: SupportsMultiModal,
143
input_ids: torch.Tensor,
144
mm_embeds: list[torch.Tensor],
145
is_mm_embed: torch.Tensor,
146
) -> torch.Tensor:
154
- x = model.embed_input_ids(
147
+ x = self.model.embed_input_ids(
148
input_ids, multimodal_embeddings=mm_embeds, is_multimodal=is_mm_embed
149
)
150
# Copy to the pre-allocated buffer for CUDA graphs.
vllm/v1/worker/gpu/model_states/default.py
+9
-2
@@ -41,7 +41,9 @@ class DefaultModelState(ModelState):
41
42
if self.supports_mm_inputs:
43
assert encoder_cache is not None
44
+ self.encoder_cache = encoder_cache
45
self.encoder_runner = EncoderRunner(
46
+ model=self.model,
47
max_num_tokens=self.max_num_tokens,
48
hidden_size=self.inputs_embeds_size,
49
encoder_cache=encoder_cache,
@@ -82,7 +84,12 @@ class DefaultModelState(ModelState):
84
mm_hashes, mm_kwargs = self.encoder_runner.prepare_mm_inputs(
85
scheduled_encoder_inputs
86
)
85
- self.encoder_runner.execute_mm_encoder(self.model, mm_hashes, mm_kwargs)
87
+ if mm_kwargs:
88
+ # Execute the multimodal encoder.
89
+ encoder_outputs = self.encoder_runner.execute_mm_encoder(mm_kwargs)
90
+ # Cache the encoder outputs by mm_hash
91
+ self.encoder_cache.encoder_outputs.update(zip(mm_hashes, encoder_outputs))
92
+
93
mm_embeds, is_mm_embed = self.encoder_runner.gather_mm_embeddings(
94
input_batch.req_ids,
95
input_batch.num_tokens,
@@ -92,7 +99,7 @@ class DefaultModelState(ModelState):
99
req_states.num_computed_prefill_tokens[input_batch.idx_mapping_np],
100
)
101
inputs_embeds = self.encoder_runner.get_inputs_embeds(
95
- self.model, input_batch.input_ids, mm_embeds, is_mm_embed
102
+ input_batch.input_ids, mm_embeds, is_mm_embed
103
)
104
return inputs_embeds[: input_batch.num_tokens_after_padding]
105