@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f5f51e5931
Fork 0

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 11, 08:45 UTC

[Core][MM] Optimize encoder cache manager by operating with embeddings only (#30475)

Signed-off-by: Roger Wang <hey@rogerw.io> Co-authored-by: Sun Kim <sunytokki@gmail.com>

Roger Wang committed Dec 16, 2025 at 14:18 UTC f5f51e5931ffd99afe69696b60765b88d3eb13f2
14 files changed +306 -130
tests/models/multimodal/processing/test_mllama4.py
+2 -2
@@ -60,12 +60,12 @@ def test_profiling(model_id: str, max_model_len: int):
60 total_num_patches.item() + num_tiles.item() + 3
61 ) # image start, image, image end
62
63 - profiled_tokens = profiler.get_mm_max_contiguous_tokens(
63 + profiled_tokens = profiler.get_mm_max_tokens(
64 max_model_len,
65 mm_counts=mm_counts,
66 )
67
68 - assert total_tokens == profiled_tokens["image"]
68 + assert total_num_patches == profiled_tokens["image"]
69 assert total_tokens == sum(
70 placeholder.length
71 for placeholder in decoder_dummy_data.multi_modal_placeholders["image"]
tests/multimodal/test_utils.py
+92
@@ -9,6 +9,7 @@ from tempfile import NamedTemporaryFile, TemporaryDirectory
9
10 import numpy as np
11 import pytest
12 +import torch
13 from PIL import Image, ImageChops
14
15 from vllm.multimodal.image import convert_image_mode
@@ -410,6 +411,97 @@ def test_argsort_mm_positions(case):
411 assert modality_idxs == expected_modality_idxs
412
413
414 +@pytest.mark.parametrize(
415 + "is_embed,expected",
416 + [
417 + (None, 5),
418 + (torch.tensor([True, True, True, True, True]), 5),
419 + (torch.tensor([False, False, False, False, False]), 0),
420 + (torch.tensor([True, False, True, False, True]), 3),
421 + (torch.tensor([True]), 1),
422 + ],
423 +)
424 +def test_placeholder_range_get_num_embeds(is_embed, expected):
425 + length = len(is_embed) if is_embed is not None else 5
426 + pr = PlaceholderRange(offset=0, length=length, is_embed=is_embed)
427 + assert pr.get_num_embeds == expected
428 +
429 +
430 +@pytest.mark.parametrize(
431 + "is_embed,expected",
432 + [
433 + (None, None),
434 + (
435 + torch.tensor([False, True, False, True, True]),
436 + torch.tensor([0, 1, 1, 2, 3]),
437 + ),
438 + (torch.tensor([True, True, True]), torch.tensor([1, 2, 3])),
439 + ],
440 +)
441 +def test_placeholder_range_embeds_cumsum(is_embed, expected):
442 + length = len(is_embed) if is_embed is not None else 5
443 + pr = PlaceholderRange(offset=0, length=length, is_embed=is_embed)
444 +
445 + if expected is None:
446 + assert pr.embeds_cumsum is None
447 + return
448 +
449 + assert torch.equal(pr.embeds_cumsum, expected)
450 + # cached_property should return the same object on repeated access
451 + assert pr.embeds_cumsum is pr.embeds_cumsum
452 +
453 +
454 +@pytest.mark.parametrize(
455 + "is_embed,start_idx,end_idx,expected",
456 + [
457 + (None, 2, 4, (2, 4)),
458 + (
459 + torch.tensor([False, True, False, True, True]),
460 + 3,
461 + 5,
462 + (1, 3),
463 + ),
464 + (
465 + torch.tensor([False, True, False, True, True]),
466 + 0,
467 + 2,
468 + (0, 1),
469 + ),
470 + (
471 + torch.tensor([True, False, True, False]),
472 + 2,
473 + 2,
474 + (1, 1),
475 + ),
476 + ],
477 +)
478 +def test_placeholder_range_get_embeds_indices_in_range(
479 + is_embed, start_idx, end_idx, expected
480 +):
481 + length = len(is_embed) if is_embed is not None else 5
482 + pr = PlaceholderRange(offset=0, length=length, is_embed=is_embed)
483 + assert pr.get_embeds_indices_in_range(start_idx, end_idx) == expected
484 +
485 +
486 +@pytest.mark.parametrize(
487 + "offset,is_embed,expected",
488 + [
489 + (0, None, [(0, 4)]),
490 + (
491 + 2,
492 + torch.tensor([False, True, False, True, True]),
493 + [(3, 3), (5, 6)],
494 + ),
495 + (0, torch.tensor([True, True, True, True]), [(0, 3)]),
496 + (0, torch.tensor([False, False, False, False]), []),
497 + ],
498 +)
499 +def test_placeholder_range_extract_embeds_range(offset, is_embed, expected):
500 + length = len(is_embed) if is_embed is not None else 5
501 + pr = PlaceholderRange(offset=offset, length=length, is_embed=is_embed)
502 + assert pr.extract_embeds_range() == expected
503 +
504 +
505 @pytest.mark.asyncio
506 @pytest.mark.parametrize("video_url", TEST_VIDEO_URLS)
507 @pytest.mark.parametrize("num_frames", [-1, 32, 1800])
tests/v1/core/test_encoder_cache_manager.py
+74 -5
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import pytest
4 +import torch
5
6 from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
7 from vllm.v1.core.encoder_cache_manager import EncoderCacheManager
@@ -23,7 +24,7 @@ class MockRequest:
24 )
25 self.mm_features.append(feature)
26
26 - def get_num_encoder_tokens(self, input_id: int) -> int:
27 + def get_num_encoder_embeds(self, input_id: int) -> int:
28 return self._token_counts[input_id]
29
30
@@ -162,8 +163,8 @@ def test_schedule_request_multi_images_respect_space_limit():
163
164 num_tokens_to_schedule = 0
165 assert manager.can_allocate(req, 0, compute_budget, num_tokens_to_schedule)
165 - num_tokens_to_schedule += req.get_num_encoder_tokens(0)
166 - compute_budget -= req.get_num_encoder_tokens(0)
166 + num_tokens_to_schedule += req.get_num_encoder_embeds(0)
167 + compute_budget -= req.get_num_encoder_embeds(0)
168
169 assert not manager.can_allocate(req, 1, compute_budget, num_tokens_to_schedule)
170
@@ -174,7 +175,75 @@ def test_schedule_request_multi_images_respect_compute_limit():
175 compute_budget = 10
176 num_tokens_to_schedule = 0
177 assert manager.can_allocate(req, 0, compute_budget, num_tokens_to_schedule)
177 - num_tokens_to_schedule += req.get_num_encoder_tokens(0)
178 - compute_budget -= req.get_num_encoder_tokens(0)
178 + num_tokens_to_schedule += req.get_num_encoder_embeds(0)
179 + compute_budget -= req.get_num_encoder_embeds(0)
180
181 assert not manager.can_allocate(req, 1, compute_budget, num_tokens_to_schedule)
182 +
183 +
184 +def test_encoder_cache_with_is_embed_mask():
185 + class MockRequestWithMask(MockRequest):
186 + def get_num_encoder_embeds(self, input_id: int) -> int:
187 + return self.mm_features[input_id].mm_position.get_num_embeds
188 +
189 + is_embed = torch.zeros(100, dtype=torch.bool)
190 + is_embed[torch.tensor([5, 15, 25, 35, 45, 55, 65, 75])] = True
191 +
192 + request = MockRequestWithMask("r1", ["img1"], [100])
193 + request.mm_features[0] = MultiModalFeatureSpec(
194 + data=None,
195 + modality="image",
196 + identifier="img1",
197 + mm_position=PlaceholderRange(offset=0, length=100, is_embed=is_embed),
198 + )
199 +
200 + manager = EncoderCacheManager(cache_size=100)
201 + manager.allocate(request, 0)
202 +
203 + assert manager.num_free_slots == 92
204 + assert "img1" in manager.cached
205 +
206 + old_size = 100
207 + new_size = request.mm_features[0].mm_position.get_num_embeds
208 + assert new_size == 8
209 + savings_ratio = old_size / new_size
210 + assert savings_ratio == 12.5
211 +
212 +
213 +def test_encoder_cache_mask_based_retrieval():
214 + class MockRequestWithMask(MockRequest):
215 + def get_num_encoder_embeds(self, input_id: int) -> int:
216 + return self.mm_features[input_id].mm_position.get_num_embeds
217 +
218 + is_embed = torch.tensor(
219 + [False, False, True, True, False, True, True, True, False, False]
220 + )
221 +
222 + request = MockRequestWithMask("r1", ["img1"], [10])
223 + request.mm_features[0] = MultiModalFeatureSpec(
224 + data=None,
225 + modality="image",
226 + identifier="img1",
227 + mm_position=PlaceholderRange(offset=0, length=10, is_embed=is_embed),
228 + )
229 +
230 + manager = EncoderCacheManager(cache_size=50)
231 + manager.allocate(request, 0)
232 +
233 + assert request.mm_features[0].mm_position.get_num_embeds == 5
234 +
235 + start_idx = 2
236 + end_idx = 8
237 + num_embeds_before = is_embed[:start_idx].sum().item()
238 + num_embeds_in_range = is_embed[start_idx:end_idx].sum().item()
239 +
240 + assert num_embeds_before == 0
241 + assert num_embeds_in_range == 5
242 +
243 + start_idx = 0
244 + end_idx = 5
245 + num_embeds_before = is_embed[:start_idx].sum().item() if start_idx > 0 else 0
246 + num_embeds_in_range = is_embed[start_idx:end_idx].sum().item()
247 +
248 + assert num_embeds_before == 0
249 + assert num_embeds_in_range == 2
tests/v1/ec_connector/unit/test_ec_example_connector.py
+1 -1
@@ -38,7 +38,7 @@ class MockRequest:
38 )
39 self.mm_features.append(feature)
40
41 - def get_num_encoder_tokens(self, input_id: int) -> int:
41 + def get_num_encoder_embeds(self, input_id: int) -> int:
42 assert input_id < len(self._token_counts)
43 return self._token_counts[input_id]
44
vllm/distributed/ec_transfer/ec_connector/example_connector.py
+1 -1
@@ -144,7 +144,7 @@ class ECExampleConnector(ECConnectorBase):
144 Update ECConnector state after encoder cache allocation.
145 """
146 mm_hash = request.mm_features[index].identifier
147 - num_encoder_token = request.get_num_encoder_tokens(index)
147 + num_encoder_token = request.get_num_encoder_embeds(index)
148 # Insert mm_hash only if this block has not been recorded yet.
149 self._mm_datas_need_loads[mm_hash] = num_encoder_token
150
vllm/model_executor/models/qwen3_vl.py
+2 -6
@@ -713,17 +713,13 @@ class Qwen3VLProcessingInfo(Qwen2VLProcessingInfo):
713 mm_counts: Mapping[str, int],
714 ) -> int:
715 target_width, target_height = self.get_image_size_with_most_features()
716 - video_soft_tokens = self.get_num_video_tokens(
716 + num_video_soft_tokens = self.get_num_video_tokens(
717 image_width=target_width,
718 image_height=target_height,
719 num_frames=self.get_num_frames_with_most_features(seq_len, mm_counts),
720 image_processor=None,
721 )
722 -
723 - # NOTE: By default in Qwen3-VL, one video token is converted to
724 - # "<{timestamp} seconds>" (on average 9.5 tokens) + vision_start_token + video_token + vision_end_token # noqa: E501
725 - formatted_video_soft_tokens = video_soft_tokens * 12.5
726 - return int(formatted_video_soft_tokens)
722 + return num_video_soft_tokens
723
724 def _calculate_timestamps(
725 self, indices: list[int] | torch.Tensor, video_fps: float, merge_size: int
vllm/multimodal/inputs.py
+35 -4
@@ -5,7 +5,7 @@ from abc import ABC, abstractmethod
5 from collections import UserDict, defaultdict
6 from collections.abc import Mapping, Sequence
7 from dataclasses import dataclass
8 -from functools import partial
8 +from functools import cached_property, partial
9 from itertools import accumulate
10 from typing import (
11 TYPE_CHECKING,
@@ -169,11 +169,42 @@ class PlaceholderRange:
169 between `offset` and `offset + length` to assign embeddings to.
170 """
171
172 - def get_num_embeds(self) -> int:
172 + @cached_property
173 + def embeds_cumsum(self) -> torch.Tensor | None:
174 if self.is_embed is None:
175 + return None
176 +
177 + return self.is_embed.cumsum(dim=0)
178 +
179 + @cached_property
180 + def get_num_embeds(self) -> int:
181 + if self.embeds_cumsum is None:
182 return self.length
183
176 - return int(self.is_embed.sum().item())
184 + return int(self.embeds_cumsum[-1])
185 +
186 + def get_embeds_indices_in_range(
187 + self, start_idx: int, end_idx: int
188 + ) -> tuple[int, int]:
189 + """
190 + Returns the starting and ending indices of the embeddings of encoder outputs
191 + in the range of [start_idx, end_idx) in the placeholders.
192 +
193 + For example, given:
194 + PlaceholderRange(offset=2, length=5, is_embed=[False, True, False, True, True])
195 +
196 + If start_idx=3 and end_idx=5, the output is (1, 3) because we want to get
197 + the second and the third embeddings from the encoder output.
198 + """
199 + if self.embeds_cumsum is None:
200 + return start_idx, end_idx
201 +
202 + embeds_start_idx = (
203 + int(self.embeds_cumsum[start_idx - 1]) if start_idx > 0 else 0
204 + )
205 + embeds_end_idx = int(self.embeds_cumsum[end_idx - 1])
206 +
207 + return embeds_start_idx, embeds_end_idx
208
209 def extract_embeds_range(self) -> list[tuple[int, int]]:
210 """Extract the start and end indices of the embedded region in prompt.
@@ -188,7 +219,7 @@ class PlaceholderRange:
219 Returns full placeholder range if `is_embed` is `None`.
220 """
221 if self.is_embed is None:
191 - return [(self.offset, self.offset + self.length)]
222 + return [(self.offset, self.offset + self.length - 1)]
223
224 mask_i = self.is_embed.int()
225 starts = torch.nonzero(
vllm/multimodal/profiling.py
+7 -25
@@ -274,15 +274,11 @@ class MultiModalProfiler(Generic[_I]):
274 def _get_mm_num_tokens(
275 self,
276 mm_inputs: MultiModalInputs,
277 - mm_embeddings_only: bool = True,
277 ) -> Mapping[str, int]:
278 placeholders_by_modality = mm_inputs["mm_placeholders"]
279
280 return {
282 - modality: sum(
283 - item.get_num_embeds() if mm_embeddings_only else item.length
284 - for item in placeholders
285 - )
281 + modality: sum(item.get_num_embeds for item in placeholders)
282 for modality, placeholders in placeholders_by_modality.items()
283 }
284
@@ -328,12 +324,15 @@ class MultiModalProfiler(Generic[_I]):
324 multi_modal_placeholders=mm_inputs["mm_placeholders"],
325 )
326
331 - def _get_mm_max_tokens(
327 + def get_mm_max_tokens(
328 self,
329 seq_len: int,
330 mm_counts: Mapping[str, int] | None = None,
335 - mm_embeddings_only: bool = True,
331 ) -> Mapping[str, int]:
332 + """
333 + Returns the maximum number of embeddings per item of each modality, excluding
334 + any break/text tokens in-between multimodal embeddings/encoder outputs.
335 + """
336 if mm_counts is None:
337 mm_counts = self.get_mm_limits()
338
@@ -349,21 +348,4 @@ class MultiModalProfiler(Generic[_I]):
348 }
349
350 mm_inputs = self._get_dummy_mm_inputs(seq_len, mm_counts)
352 - return self._get_mm_num_tokens(mm_inputs, mm_embeddings_only=mm_embeddings_only)
353 -
354 - def get_mm_max_contiguous_tokens(
355 - self,
356 - seq_len: int,
357 - mm_counts: Mapping[str, int] | None = None,
358 - ) -> Mapping[str, int]:
359 - """
360 - Returns the maximum length of the multimodal (image placeholders+text)
361 - tokens, including any break/text tokens in-between image embeddings.
362 -
363 - `<im_start> [IMG] [IMG] [IMG] <row_break> [IMG] [IMG] [IMG] <im_end>`
364 - Returns 9, even when the number of image embeddings is 6.
365 -
366 - This is important to take into account when profiling and
367 - initializing the encoder cache size.
368 - """
369 - return self._get_mm_max_tokens(seq_len, mm_counts, mm_embeddings_only=False)
351 + return self._get_mm_num_tokens(mm_inputs)
vllm/multimodal/registry.py
+1 -1
@@ -164,7 +164,7 @@ class MultiModalRegistry:
164 profiler.get_mm_limits() if profiler_limits is None else profiler_limits
165 )
166
167 - return profiler.get_mm_max_contiguous_tokens(
167 + return profiler.get_mm_max_tokens(
168 seq_len,
169 {modality: 1 for modality, limit in profiler_limits.items() if limit > 0},
170 )
vllm/v1/core/encoder_cache_manager.py
+43 -37
@@ -39,20 +39,26 @@ class EncoderCacheManager:
39 space for new embeddings.
40 Oldest cached embeddings with no request referenced will be first evicted.
41
42 + NOTE: The EncoderCacheManager operates on the level of multimodal embeddings
43 + instead of encoder tokens (i.e. all tokens that represent the multimodal data
44 + in the input sequence). This means all break/text tokens in-between multimodal
45 + embeddings are not considered with respect to the cache size and the number
46 + of free slots.
47 +
48 Args:
49 cache_size: Limit the size of the cache, measured by the number of
44 - tokens from the input sequence.
50 + encoder embeddings from the input sequence.
51
52 Attributes:
47 - cache_size: Total cache capacity in encoder tokens.
48 - num_free_slots: Current available cache capacity in encoder tokens.
53 + cache_size: Total cache capacity in encoder embeddings.
54 + num_free_slots: Current available cache capacity in encoder embeddings.
55 num_freeable_slots: Capacity that can be immediately reclaimed by
50 - evicting entries with zero references (in encoder tokens).
56 + evicting entries with zero references (in encoder embeddings).
57 cached: Mapping from mm_hash to a set of request IDs that currently
58 reference the cached entry. If the set is empty, the entry exists
59 but is not referenced by any request and is eligible for
60 reclamation.
55 - freeable: List of tuples (mm_hash, num_tokens) representing entries
61 + freeable: List of tuples (mm_hash, num_encoder_embeds) representing entries
62 whose no current running request is needed and that can be freed to
63 make space when needed.
64 freed: List of mm_hash strings that were actually evicted since the
@@ -67,7 +73,7 @@ class EncoderCacheManager:
73 # mm_hash of mm_data => ids of requests that reference the mm_data
74 self.cached: dict[str, set[str]] = {}
75
70 - # mm_hash of mm_data => num_encoder_tokens of the mm_data
76 + # mm_hash of mm_data => num_encoder_embeds of the mm_data
77 self.freeable: OrderedDict[str, int] = OrderedDict()
78 self.freed: list[str] = []
79
@@ -93,8 +99,8 @@ class EncoderCacheManager:
99
100 # Cached but currently not referenced by any request
101 if not self.cached[mm_hash]:
96 - num_tokens = self.freeable.pop(mm_hash)
97 - self.num_freeable_slots -= num_tokens
102 + num_encoder_embeds = self.freeable.pop(mm_hash)
103 + self.num_freeable_slots -= num_encoder_embeds
104
105 self.cached[mm_hash].add(request.request_id)
106 return True
@@ -104,7 +110,7 @@ class EncoderCacheManager:
110 request: Request,
111 input_id: int,
112 encoder_compute_budget: int,
107 - num_tokens_to_schedule: int,
113 + num_embeds_to_schedule: int,
114 ) -> bool:
115 """Check if there's sufficient cache space for a multimodal input.
116 If there is, return True and update EncoderCacheManager state.
@@ -121,9 +127,9 @@ class EncoderCacheManager:
127 Args:
128 request: The request containing the multimodal input.
129 input_id: Index of the multimodal input within the request.
124 - encoder_compute_budget: Number of encoder tokens allowed to be
130 + encoder_compute_budget: Number of encoder embeddings allowed to be
131 computed when this method is invoked.
126 - num_tokens_to_schedule: Number of tokens already scheduled to be
132 + num_embeds_to_schedule: Number of encoder embeddings already scheduled to be
133 allocated with cache space when this method is invoked.
134
135 Returns:
@@ -134,30 +140,30 @@ class EncoderCacheManager:
140 Note: This method does not allocate physical memory for the encoder
141 output but only the state of EncoderCacheManager.
142 """
137 - num_tokens = request.get_num_encoder_tokens(input_id)
143 + num_embeds = request.get_num_encoder_embeds(input_id)
144
145 # Not enough compute budget
140 - if num_tokens > encoder_compute_budget:
146 + if num_embeds > encoder_compute_budget:
147 return False
148
143 - num_tokens += num_tokens_to_schedule
149 + num_embeds += num_embeds_to_schedule
150
151 # Enough free slots
146 - if num_tokens <= self.num_free_slots:
152 + if num_embeds <= self.num_free_slots:
153 return True
154
155 # Not enough reclaimable slots
150 - if num_tokens > self.num_freeable_slots:
156 + if num_embeds > self.num_freeable_slots:
157 return False
158
159 # Not enough free slots but enough reclaimable slots
160 # NOTE: Eviction takes place here, but physical memory is not freed
161 # until model runner is notified by the scheduler output.
156 - while num_tokens > self.num_free_slots:
157 - mm_hash, num_free_token = self.freeable.popitem(last=False)
162 + while num_embeds > self.num_free_slots:
163 + mm_hash, num_free_embeds = self.freeable.popitem(last=False)
164 del self.cached[mm_hash]
165 self.freed.append(mm_hash)
160 - self.num_free_slots += num_free_token
166 + self.num_free_slots += num_free_embeds
167 return True
168
169 def allocate(self, request: Request, input_id: int) -> None:
@@ -176,16 +182,16 @@ class EncoderCacheManager:
182 if mm_hash not in self.cached:
183 self.cached[mm_hash] = set()
184
179 - num_encoder_tokens = request.get_num_encoder_tokens(input_id)
185 + num_encoder_embeds = request.get_num_encoder_embeds(input_id)
186
187 # NOTE: Encoder cache should always have enough space for encoder inputs
188 # that are scheduled since eviction takes place at can_allocate().
183 - assert self.num_free_slots >= num_encoder_tokens
184 - assert self.num_freeable_slots >= num_encoder_tokens
189 + assert self.num_free_slots >= num_encoder_embeds
190 + assert self.num_freeable_slots >= num_encoder_embeds
191
192 self.cached[mm_hash].add(request_id)
187 - self.num_free_slots -= num_encoder_tokens
188 - self.num_freeable_slots -= num_encoder_tokens
193 + self.num_free_slots -= num_encoder_embeds
194 + self.num_freeable_slots -= num_encoder_embeds
195
196 def get_cached_input_ids(self, request: Request) -> set[int]:
197 """Get all cached multimodal input IDs for a request.
@@ -206,7 +212,7 @@ class EncoderCacheManager:
212
213 When the reference set for the corresponding `mm_hash` becomes empty,
214 the entry is appended to `freeable` and `num_freeable_slots` is
209 - increased by the number of encoder tokens for that input.
215 + increased by the number of encoder embeddings for that input.
216
217 The entry is NOT physically freed until capacity is needed (e.g., by
218 `can_allocate`).
@@ -218,9 +224,9 @@ class EncoderCacheManager:
224 return
225 self.cached[mm_hash].discard(req_id)
226 if not self.cached[mm_hash]:
221 - num_tokens = request.get_num_encoder_tokens(input_id)
222 - self.freeable[mm_hash] = num_tokens
223 - self.num_freeable_slots += num_tokens
227 + num_encoder_embeds = request.get_num_encoder_embeds(input_id)
228 + self.freeable[mm_hash] = num_encoder_embeds
229 + self.num_freeable_slots += num_encoder_embeds
230
231 def free(self, request: Request) -> None:
232 """Free all encoder input cache reference held by *request*.
@@ -361,20 +367,20 @@ class EncoderDecoderCacheManager(EncoderCacheManager):
367 request: Request,
368 input_id: int,
369 encoder_compute_budget: int,
364 - num_tokens_to_schedule: int,
370 + num_embeds_to_schedule: int,
371 ) -> bool:
366 - num_tokens = request.get_num_encoder_tokens(input_id)
372 + num_encoder_embeds = request.get_num_encoder_embeds(input_id)
373 # Not enough compute budget
368 - if num_tokens > encoder_compute_budget:
374 + if num_encoder_embeds > encoder_compute_budget:
375 return False
376
371 - num_tokens += num_tokens_to_schedule
377 + num_encoder_embeds += num_embeds_to_schedule
378 # Enough free slots
373 - return num_tokens <= self.num_free_slots
379 + return num_encoder_embeds <= self.num_free_slots
380
381 def allocate(self, request: Request, input_id: int) -> None:
376 - num_encoder_tokens = request.get_num_encoder_tokens(input_id)
377 - self.num_free_slots -= num_encoder_tokens
382 + num_encoder_embeds = request.get_num_encoder_embeds(input_id)
383 + self.num_free_slots -= num_encoder_embeds
384
385 mm_hash = request.mm_features[input_id].identifier
386 self.freed.append(mm_hash)
@@ -392,5 +398,5 @@ class EncoderDecoderCacheManager(EncoderCacheManager):
398 return freed
399
400 def free_encoder_input(self, request: Request, input_id: int) -> None:
395 - num_tokens = request.get_num_encoder_tokens(input_id)
396 - self.num_free_slots += num_tokens
401 + num_encoder_embeds = request.get_num_encoder_embeds(input_id)
402 + self.num_free_slots += num_encoder_embeds
vllm/v1/core/sched/scheduler.py
+26 -9
@@ -355,11 +355,11 @@ class Scheduler(SchedulerInterface):
355 if preempted_encoder_inputs:
356 # Restore encoder compute budget if the preempted
357 # request had encoder inputs scheduled in this step.
358 - num_tokens_to_restore = sum(
359 - preempted_req.get_num_encoder_tokens(i)
358 + num_embeds_to_restore = sum(
359 + preempted_req.get_num_encoder_embeds(i)
360 for i in preempted_encoder_inputs
361 )
362 - encoder_compute_budget += num_tokens_to_restore
362 + encoder_compute_budget += num_embeds_to_restore
363 req_index -= 1
364 else:
365 preempted_req = self.running.pop()
@@ -911,10 +911,11 @@ class Scheduler(SchedulerInterface):
911 # multiple encoder inputs per request), we need to create temporary
912 # trackers for accounting at the encoder input level.
913 mm_hashes_to_schedule = set()
914 - num_tokens_to_schedule = 0
914 + num_embeds_to_schedule = 0
915 for i, mm_feature in enumerate(mm_features):
916 start_pos = mm_feature.mm_position.offset
917 num_encoder_tokens = mm_feature.mm_position.length
918 + num_encoder_embeds = mm_feature.mm_position.get_num_embeds
919
920 # The encoder output is needed if the two ranges overlap:
921 # [num_computed_tokens, num_computed_tokens + num_new_tokens) and
@@ -970,9 +971,8 @@ class Scheduler(SchedulerInterface):
971 ):
972 num_new_tokens = start_pos - num_computed_tokens
973 break
973 -
974 if not self.encoder_cache_manager.can_allocate(
975 - request, i, encoder_compute_budget, num_tokens_to_schedule
975 + request, i, encoder_compute_budget, num_embeds_to_schedule
976 ):
977 # The encoder cache is full or the encoder budget is exhausted.
978 # NOTE(woosuk): We assume that the encoder input tokens should
@@ -992,14 +992,31 @@ class Scheduler(SchedulerInterface):
992 num_new_tokens = 0
993 break
994
995 + # Calculate the number of embeddings to schedule in the current range
996 + # of scheduled encoder placholder tokens.
997 + start_idx_rel = max(0, num_computed_tokens - start_pos)
998 + end_idx_rel = min(
999 + num_encoder_tokens, num_computed_tokens + num_new_tokens - start_pos
1000 + )
1001 + curr_embeds_start, curr_embeds_end = (
1002 + mm_feature.mm_position.get_embeds_indices_in_range(
1003 + start_idx_rel,
1004 + end_idx_rel,
1005 + )
1006 + )
1007 + # There's no embeddings in the current range of encoder placeholder tokens
1008 + # so we can skip the encoder input.
1009 + if curr_embeds_end - curr_embeds_start == 0:
1010 + continue
1011 +
1012 if self.ec_connector is not None and remote_cache_has_item[i]:
1013 mm_hashes_to_schedule.add(request.mm_features[i].identifier)
1014 external_load_encoder_input.append(i)
998 - num_tokens_to_schedule += num_encoder_tokens
1015 + num_embeds_to_schedule += num_encoder_embeds
1016 continue
1017
1001 - num_tokens_to_schedule += num_encoder_tokens
1002 - encoder_compute_budget -= num_encoder_tokens
1018 + num_embeds_to_schedule += num_encoder_embeds
1019 + encoder_compute_budget -= num_encoder_embeds
1020 mm_hashes_to_schedule.add(request.mm_features[i].identifier)
1021 encoder_inputs_to_schedule.append(i)
1022
vllm/v1/request.py
+3 -3
@@ -209,10 +209,10 @@ class Request:
209 def get_finished_reason(self) -> FinishReason | None:
210 return RequestStatus.get_finished_reason(self.status)
211
212 - def get_num_encoder_tokens(self, input_id: int) -> int:
212 + def get_num_encoder_embeds(self, input_id: int) -> int:
213 assert input_id < len(self.mm_features)
214 - num_tokens = self.mm_features[input_id].mm_position.length
215 - return num_tokens
214 + num_embeds = self.mm_features[input_id].mm_position.get_num_embeds
215 + return num_embeds
216
217 def record_event(
218 self,
vllm/v1/worker/gpu_model_runner.py
+13 -36
@@ -169,9 +169,7 @@ from .utils import (
169 MultiModalBudget,
170 add_kv_sharing_layers_to_kv_cache_groups,
171 bind_kv_cache,
172 - gather_mm_placeholders,
172 sanity_check_mm_encoder_outputs,
174 - scatter_mm_placeholders,
173 )
174
175 if TYPE_CHECKING:
@@ -2209,10 +2207,7 @@ class GPUModelRunner(
2207
2208 # Cache the encoder outputs by mm_hash
2209 for (mm_hash, pos_info), output in zip(mm_hashes_pos, encoder_outputs):
2212 - self.encoder_cache[mm_hash] = scatter_mm_placeholders(
2213 - output,
2214 - is_embed=pos_info.is_embed,
2215 - )
2210 + self.encoder_cache[mm_hash] = output
2211 logger.debug("Finish execute for mm hash %s", mm_hash)
2212 self.maybe_save_ec_to_connector(self.encoder_cache, mm_hash)
2213
@@ -2263,6 +2258,13 @@ class GPUModelRunner(
2258 num_encoder_tokens,
2259 )
2260 assert start_idx < end_idx
2261 + curr_embeds_start, curr_embeds_end = (
2262 + pos_info.get_embeds_indices_in_range(start_idx, end_idx)
2263 + )
2264 + # If there are no embeddings in the current range, we skip
2265 + # gathering the embeddings.
2266 + if curr_embeds_start == curr_embeds_end:
2267 + continue
2268
2269 mm_hash = mm_feature.identifier
2270 encoder_output = self.encoder_cache.get(mm_hash, None)
@@ -2270,16 +2272,14 @@ class GPUModelRunner(
2272
2273 if (is_embed := pos_info.is_embed) is not None:
2274 is_embed = is_embed[start_idx:end_idx]
2275 + mm_embeds_item = encoder_output[curr_embeds_start:curr_embeds_end]
2276 + else:
2277 + mm_embeds_item = encoder_output[start_idx:end_idx]
2278
2279 req_start_pos = req_start_idx + start_pos - num_computed_tokens
2280 is_mm_embed[req_start_pos + start_idx : req_start_pos + end_idx] = (
2281 True if is_embed is None else is_embed
2282 )
2278 -
2279 - mm_embeds_item = gather_mm_placeholders(
2280 - encoder_output[start_idx:end_idx],
2281 - is_embed=is_embed,
2282 - )
2283 mm_embeds_req.append(mm_embeds_item)
2284
2285 if self.is_multimodal_pruning_enabled and self.uses_mrope:
@@ -4508,31 +4508,8 @@ class GPUModelRunner(
4508 dummy_encoder_outputs,
4509 expected_num_items=max_mm_items_per_batch,
4510 )
4511 -
4512 - # NOTE: This happens when encoder cache needs to store
4513 - # the embeddings that encoder outputs are scattered onto.
4514 - # In this case we create dummy embeddings of size
4515 - # (max_tokens_for_modality, hidden_size) and scatter
4516 - # encoder output into it.
4517 - encoder_output_shape = dummy_encoder_outputs[0].shape
4518 - max_mm_tokens_per_item = mm_budget.max_tokens_by_modality[
4519 - dummy_modality
4520 - ]
4521 - if encoder_output_shape[0] < max_mm_tokens_per_item:
4522 - encoder_hidden_size = encoder_output_shape[-1]
4523 - expanded_outputs = []
4524 - for output in dummy_encoder_outputs:
4525 - expanded = output.new_zeros(
4526 - (max_mm_tokens_per_item, encoder_hidden_size)
4527 - )
4528 - num_tokens = output.shape[0]
4529 - expanded[:num_tokens].copy_(output)
4530 - expanded_outputs.append(expanded)
4531 -
4532 - dummy_encoder_outputs = expanded_outputs
4533 -
4534 - # Cache the dummy encoder outputs.
4535 - self.encoder_cache["tmp"] = dict(enumerate(dummy_encoder_outputs))
4511 + for i, output in enumerate(dummy_encoder_outputs):
4512 + self.encoder_cache[f"tmp_{i}"] = output
4513
4514 # Add `is_profile` here to pre-allocate communication buffers
4515 hidden_states, last_hidden_states = self._dummy_run(
vllm/v1/worker/utils.py
+6
@@ -4,10 +4,12 @@ from collections import defaultdict
4 from dataclasses import dataclass, field
5
6 import torch
7 +from typing_extensions import deprecated
8
9 from vllm.attention.backends.abstract import AttentionBackend
10 from vllm.attention.layer import Attention
11 from vllm.config import ModelConfig, SchedulerConfig, VllmConfig
12 +from vllm.logger import init_logger
13 from vllm.model_executor.models.interfaces import MultiModalEmbeddings
14 from vllm.model_executor.models.utils import extract_layer_index
15 from vllm.multimodal.cache import processor_only_cache_from_config
@@ -17,6 +19,8 @@ from vllm.v1.attention.backends.utils import AttentionMetadataBuilder
19 from vllm.v1.core.encoder_cache_manager import compute_mm_encoder_budget
20 from vllm.v1.kv_cache_interface import KVCacheGroupSpec, KVCacheSpec
21
22 +logger = init_logger(__name__)
23 +
24
25 class MultiModalBudget:
26 """Helper class to calculate budget information for multi-modal models."""
@@ -198,6 +202,7 @@ def sanity_check_mm_encoder_outputs(
202 )
203
204
205 +@deprecated("`scatter_mm_placeholders` is deprecated and will be removed in v0.15.0.")
206 def scatter_mm_placeholders(
207 embeds: torch.Tensor,
208 is_embed: torch.Tensor | None,
@@ -226,6 +231,7 @@ def scatter_mm_placeholders(
231 return placeholders
232
233
234 +@deprecated("`gather_mm_placeholders` is deprecated and will be removed in v0.15.0.")
235 def gather_mm_placeholders(
236 placeholders: torch.Tensor,
237 is_embed: torch.Tensor | None,