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 14, 05:30 UTC
[Core][Multimodal] Track encode cache entries by mm_hash and enable embedding sharing between requests (#22711)
Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com> Signed-off-by: Roger Wang <hey@rogerw.io> Co-authored-by: knlnguyen1802 <knlnguyen1802@gmail.com> Co-authored-by: Roger Wang <hey@rogerw.io>
Chenguang Zheng committed
Aug 25, 2025 at 15:41 UTC
d765cf01fe574f94b85c2a49cc6e685a273245af
12 files changed
+366
-155
tests/v1/core/test_encoder_cache_manager.py
new
+144
@@ -0,0 +1,144 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from vllm.v1.core.encoder_cache_manager import EncoderCacheManager
5
+
6
+
7
+# ------------------ Mock Classes ------------------ #
8
+class MockRequest:
9
+
10
+ def __init__(self, request_id, mm_hashes, token_counts):
11
+ self.request_id = request_id
12
+ self.mm_hashes = mm_hashes
13
+ self._token_counts = token_counts
14
+
15
+ def get_num_encoder_tokens(self, input_id: int) -> int:
16
+ return self._token_counts[input_id]
17
+
18
+
19
+# ------------------ Unit Tests ------------------ #
20
+def test_basic_allocate_and_reuse():
21
+ cache = EncoderCacheManager(cache_size=10)
22
+ req = MockRequest("r1", ["imgA"], [4])
23
+
24
+ assert not cache.check_and_update_cache(req, 0)
25
+ assert cache.try_allocate(req, 0, int(1e9))
26
+
27
+ cache.allocate(req, 0)
28
+
29
+ assert cache.check_and_update_cache(req, 0)
30
+ assert "r1" in cache.cached["imgA"]
31
+ assert cache.num_free_slots == 6
32
+
33
+ # Free twice to bring refcount to 0.
34
+ cache.free_encoder_input(req, 0)
35
+ cache.free_encoder_input(req, 0)
36
+
37
+ assert not cache.cached["imgA"]
38
+ assert "imgA" in cache.freeable
39
+ assert cache.num_freeable_slots == 10
40
+ assert cache.num_free_slots == 6
41
+
42
+
43
+def test_freeing_decreases_refcount_and_moves_to_freeable():
44
+ manager = EncoderCacheManager(cache_size=10)
45
+ req = MockRequest("req2", ["img3"], [5])
46
+
47
+ assert manager.try_allocate(req, 0, int(1e9))
48
+ manager.allocate(req, 0)
49
+
50
+ assert len(manager.cached["img3"]) == 1
51
+
52
+ manager.free_encoder_input(req, 0)
53
+
54
+ assert not manager.cached["img3"]
55
+ assert "img3" in manager.freeable
56
+ assert manager.num_freeable_slots == 10
57
+
58
+
59
+def test_free_request_frees_all_inputs():
60
+ manager = EncoderCacheManager(cache_size=10)
61
+ req = MockRequest("req3", ["a", "b"], [2, 3])
62
+
63
+ assert manager.try_allocate(req, 0, int(1e9))
64
+ manager.allocate(req, 0)
65
+
66
+ assert manager.try_allocate(req, 1, int(1e9))
67
+ manager.allocate(req, 1)
68
+
69
+ assert len(manager.cached["a"]) == 1
70
+ assert len(manager.cached["b"]) == 1
71
+
72
+ manager.free(req)
73
+
74
+ assert not manager.cached["a"]
75
+ assert not manager.cached["b"]
76
+ assert "a" in manager.freeable
77
+ assert "b" in manager.freeable
78
+ assert manager.num_freeable_slots == 10
79
+
80
+
81
+def test_eviction_when_cache_is_full():
82
+ manager = EncoderCacheManager(cache_size=10)
83
+
84
+ req1 = MockRequest("req1", ["x"], [6])
85
+ req2 = MockRequest("req2", ["y"], [5])
86
+
87
+ assert manager.try_allocate(req1, 0, int(1e9))
88
+ manager.allocate(req1, 0)
89
+ manager.free_encoder_input(req1, 0)
90
+
91
+ assert manager.try_allocate(req2, 0, int(1e9))
92
+ manager.allocate(req2, 0)
93
+
94
+ # 'x' should have been evicted.
95
+ assert "x" not in manager.cached
96
+ assert "x" in manager.get_freed_mm_hashes()
97
+
98
+
99
+def test_get_cached_input_ids():
100
+ manager = EncoderCacheManager(cache_size=10)
101
+ req = MockRequest("reqX", ["m", "n", "o"], [2, 4, 3])
102
+
103
+ assert manager.try_allocate(req, 0, int(1e9))
104
+ manager.allocate(req, 0)
105
+
106
+ assert manager.try_allocate(req, 2, int(1e9))
107
+ manager.allocate(req, 2)
108
+
109
+ cached_ids = manager.get_cached_input_ids(req)
110
+ assert cached_ids == {0, 2}
111
+
112
+
113
+def test_has_cache_restores_from_freeable():
114
+ manager = EncoderCacheManager(cache_size=10)
115
+ req = MockRequest("reqY", ["imgZ"], [4])
116
+
117
+ assert manager.try_allocate(req, 0, int(1e9))
118
+ manager.allocate(req, 0)
119
+
120
+ manager.free_encoder_input(req, 0)
121
+
122
+ # Should restore from freeable.
123
+ assert manager.check_and_update_cache(req, 0)
124
+ assert len(manager.cached["imgZ"]) == 1
125
+ assert "imgZ" not in manager.freeable
126
+ assert manager.num_freeable_slots == 6
127
+
128
+
129
+def test_get_freed_mm_hashes_clears_freed_list():
130
+ manager = EncoderCacheManager(cache_size=10)
131
+ req1 = MockRequest("reqA", ["a"], [5])
132
+ req2 = MockRequest("reqB", ["b"], [6])
133
+
134
+ assert manager.try_allocate(req1, 0, int(1e9))
135
+ manager.allocate(req1, 0)
136
+ manager.free_encoder_input(req1, 0)
137
+
138
+ # Should trigger eviction of 'a'.
139
+ assert manager.try_allocate(req2, 0, int(1e9))
140
+ manager.allocate(req2, 0)
141
+
142
+ freed = manager.get_freed_mm_hashes()
143
+ assert "a" in freed
144
+ assert manager.get_freed_mm_hashes() == []
tests/v1/core/test_scheduler.py
+4
-4
@@ -338,7 +338,7 @@ def test_stop_via_update_from_output():
338
},
339
num_common_prefix_blocks=0,
340
finished_req_ids=set(),
341
- free_encoder_input_ids=[],
341
+ free_encoder_mm_hashes=[],
342
structured_output_request_ids={},
343
grammar_bitmask=None)
344
@@ -391,7 +391,7 @@ def test_stop_via_update_from_output():
391
},
392
num_common_prefix_blocks=0,
393
finished_req_ids=set(),
394
- free_encoder_input_ids=[],
394
+ free_encoder_mm_hashes=[],
395
structured_output_request_ids={},
396
grammar_bitmask=None,
397
)
@@ -443,7 +443,7 @@ def test_stop_via_update_from_output():
443
},
444
num_common_prefix_blocks=0,
445
finished_req_ids=set(),
446
- free_encoder_input_ids=[],
446
+ free_encoder_mm_hashes=[],
447
structured_output_request_ids={},
448
grammar_bitmask=None,
449
)
@@ -490,7 +490,7 @@ def test_stop_via_update_from_output():
490
},
491
num_common_prefix_blocks=0,
492
finished_req_ids=set(),
493
- free_encoder_input_ids=[],
493
+ free_encoder_mm_hashes=[],
494
structured_output_request_ids={},
495
grammar_bitmask=None)
496
tests/v1/core/utils.py
+5
-1
@@ -143,7 +143,11 @@ def create_requests(
143
mm_position = mm_positions[i]
144
mm_item = MultiModalKwargsItem.dummy("dummy_m")
145
mm_kwargs = [mm_item] * len(mm_position)
146
- mm_hashes = ["hash"] * len(mm_position)
146
+ # Dummy hash for each mm item should be unique
147
+ # since encoder cache tracks entries by hash
148
+ mm_hashes = [
149
+ "hash" + str(i) + "_" + str(j) for j in range(len(mm_position))
150
+ ]
151
else:
152
mm_position = None
153
mm_kwargs = None
tests/v1/tpu/worker/test_tpu_model_runner.py
+6
-6
@@ -85,7 +85,7 @@ def _schedule_new_request(*req_ids: str) -> SchedulerOutput:
85
scheduled_encoder_inputs={},
86
num_common_prefix_blocks=0,
87
finished_req_ids=set(),
88
- free_encoder_input_ids=[],
88
+ free_encoder_mm_hashes=[],
89
structured_output_request_ids={},
90
grammar_bitmask=None,
91
)
@@ -164,7 +164,7 @@ def test_update_states_request_finished(model_runner):
164
scheduled_encoder_inputs={},
165
num_common_prefix_blocks=0,
166
finished_req_ids={req_id},
167
- free_encoder_input_ids=[],
167
+ free_encoder_mm_hashes=[],
168
structured_output_request_ids={},
169
grammar_bitmask=None,
170
)
@@ -194,7 +194,7 @@ def test_update_states_request_resumed(model_runner):
194
scheduled_encoder_inputs={},
195
num_common_prefix_blocks=0,
196
finished_req_ids=set(),
197
- free_encoder_input_ids=[],
197
+ free_encoder_mm_hashes=[],
198
structured_output_request_ids={},
199
grammar_bitmask=None,
200
)
@@ -221,7 +221,7 @@ def test_update_states_request_resumed(model_runner):
221
scheduled_encoder_inputs={},
222
num_common_prefix_blocks=0,
223
finished_req_ids=set(),
224
- free_encoder_input_ids=[],
224
+ free_encoder_mm_hashes=[],
225
structured_output_request_ids={},
226
grammar_bitmask=None,
227
)
@@ -252,7 +252,7 @@ def test_update_states_no_changes(model_runner):
252
scheduled_encoder_inputs={},
253
num_common_prefix_blocks=0,
254
finished_req_ids=set(),
255
- free_encoder_input_ids=[],
255
+ free_encoder_mm_hashes=[],
256
structured_output_request_ids={},
257
grammar_bitmask=None,
258
)
@@ -287,7 +287,7 @@ def test_update_states_request_unscheduled(model_runner):
287
scheduled_encoder_inputs={},
288
num_common_prefix_blocks=0,
289
finished_req_ids=set(),
290
- free_encoder_input_ids=[],
290
+ free_encoder_mm_hashes=[],
291
structured_output_request_ids={},
292
grammar_bitmask=None,
293
)
tests/v1/worker/test_gpu_input_batch.py
+1
@@ -205,6 +205,7 @@ def _construct_cached_request_state(req_id_suffix: int):
205
pooling_params=None,
206
mm_kwargs=[],
207
mm_positions=[],
208
+ mm_hashes=[],
209
block_ids=([], ),
210
generator=None,
211
num_computed_tokens=len(output_token_ids),
tests/v1/worker/test_gpu_model_runner.py
+6
-6
@@ -141,7 +141,7 @@ def _schedule_new_request(*req_ids: str) -> SchedulerOutput:
141
scheduled_encoder_inputs={},
142
num_common_prefix_blocks=0,
143
finished_req_ids=set(),
144
- free_encoder_input_ids=[],
144
+ free_encoder_mm_hashes=[],
145
structured_output_request_ids={},
146
grammar_bitmask=None,
147
)
@@ -207,7 +207,7 @@ def test_update_states_request_finished(model_runner, dist_init):
207
scheduled_encoder_inputs={},
208
num_common_prefix_blocks=0,
209
finished_req_ids={req_id},
210
- free_encoder_input_ids=[],
210
+ free_encoder_mm_hashes=[],
211
structured_output_request_ids={},
212
grammar_bitmask=None,
213
)
@@ -239,7 +239,7 @@ def test_update_states_request_resumed(model_runner, dist_init):
239
scheduled_encoder_inputs={},
240
num_common_prefix_blocks=0,
241
finished_req_ids=set(),
242
- free_encoder_input_ids=[],
242
+ free_encoder_mm_hashes=[],
243
structured_output_request_ids={},
244
grammar_bitmask=None,
245
)
@@ -266,7 +266,7 @@ def test_update_states_request_resumed(model_runner, dist_init):
266
scheduled_encoder_inputs={},
267
num_common_prefix_blocks=0,
268
finished_req_ids=set(),
269
- free_encoder_input_ids=[],
269
+ free_encoder_mm_hashes=[],
270
structured_output_request_ids={},
271
grammar_bitmask=None,
272
)
@@ -347,7 +347,7 @@ def test_update_states_no_changes(model_runner, dist_init):
347
scheduled_encoder_inputs={},
348
num_common_prefix_blocks=0,
349
finished_req_ids=set(),
350
- free_encoder_input_ids=[],
350
+ free_encoder_mm_hashes=[],
351
structured_output_request_ids={},
352
grammar_bitmask=None,
353
)
@@ -384,7 +384,7 @@ def test_update_states_request_unscheduled(model_runner, dist_init):
384
scheduled_encoder_inputs={},
385
num_common_prefix_blocks=0,
386
finished_req_ids=set(),
387
- free_encoder_input_ids=[],
387
+ free_encoder_mm_hashes=[],
388
structured_output_request_ids={},
389
grammar_bitmask=None,
390
)
vllm/v1/core/encoder_cache_manager.py
+150
-87
@@ -1,5 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from collections import OrderedDict
5
from collections.abc import Mapping
6
from typing import TYPE_CHECKING
7
@@ -31,34 +33,52 @@ class EncoderCacheManager:
33
within requests, allowing for fine-grained memory management and enabling
34
chunked processing of multimodal inputs.
35
34
- Note that no caching is shared between requests at this time. If the same
35
- input is used across multiple requests, it will be reprocessed for each
36
- request.
36
+ Cache is enabled to share embeddings of same multimodal data
37
+ item (identified by their hash value) between different requests,
38
+ and eviction takes place at allocation time when there's no free
39
+ space for new embeddings.
40
+ Oldest cached embeddings with no request referenced will be first evicted.
41
42
Args:
43
cache_size: Limit the size of the cache, measured by the number of
44
tokens from the input sequence.
45
46
Attributes:
43
- cache_size: Total cache capacity in encoder tokens
44
- num_free_slots: Current available cache capacity in encoder tokens
45
- cached: Mapping from request_id to set of cached input_ids for that
46
- request
47
- freed: List of (request_id, input_id) pairs that were recently freed.
48
- This is cleared after every call to get_freed_ids().
47
+ cache_size: Total cache capacity in encoder tokens.
48
+ num_free_slots: Current available cache capacity in encoder tokens.
49
+ num_freeable_slots: Capacity that can be immediately reclaimed by
50
+ evicting entries with zero references (in encoder tokens).
51
+ cached: Mapping from mm_hash to a set of request IDs that currently
52
+ reference the cached entry. If the set is empty, the entry exists
53
+ but is not referenced by any request and is eligible for
54
+ reclamation.
55
+ freeable: List of tuples (mm_hash, num_tokens) representing entries
56
+ whose no current running request is needed and that can be freed to
57
+ make space when needed.
58
+ freed: List of mm_hash strings that were actually evicted since the
59
+ last call to get_freed_mm_hashes(). This list is cleared on return.
60
"""
61
62
def __init__(self, cache_size: int):
63
self.cache_size = cache_size
64
self.num_free_slots = cache_size
54
- # req_id -> cached input ids
55
- self.cached: dict[str, set[int]] = {}
56
- # list of [req_id, input_id]
57
- self.freed: list[tuple[str, int]] = []
65
+ self.num_freeable_slots = cache_size
66
+
67
+ # mm_hash of mm_data => ids of requests that reference the mm_data
68
+ self.cached: dict[str, set[str]] = {}
69
+
70
+ # mm_hash of mm_data => num_encoder_tokens of the mm_data
71
+ self.freeable: OrderedDict[str, int] = OrderedDict()
72
+ self.freed: list[str] = []
73
59
- def has_cache(self, request: Request, input_id: int) -> bool:
74
+ def check_and_update_cache(self, request: Request, input_id: int) -> bool:
75
"""Check if encoder output for a specific multimodal input is cached.
76
77
+ If the encoder output is cached, update `cached` to add the request id
78
+ to the set of request ids that reference the cached encoder output.
79
+ If the encoder output was previously not referenced by any request,
80
+ update `freeable` and `num_freeable_slots` accordingly.
81
+
82
Args:
83
request: The request containing the multimodal input
84
input_id: Index of the multimodal input within the request
@@ -66,103 +86,151 @@ class EncoderCacheManager:
86
Returns:
87
True if the encoder output for this input is already cached
88
"""
69
- req_id = request.request_id
70
- return req_id in self.cached and input_id in self.cached[req_id]
71
-
72
- def can_allocate(self, request: Request, input_id: int) -> bool:
73
- """Check if there's sufficient cache space for a multimodal input.
89
+ mm_hash = request.mm_hashes[input_id]
90
+ # Not cached at all
91
+ if mm_hash not in self.cached:
92
+ return False
93
+
94
+ # Cached but currently not referenced by any request
95
+ if not self.cached[mm_hash]:
96
+ num_tokens = self.freeable.pop(mm_hash)
97
+ self.num_freeable_slots -= num_tokens
98
+
99
+ self.cached[mm_hash].add(request.request_id)
100
+ return True
101
+
102
+ def try_allocate(self, request: Request, input_id: int,
103
+ encoder_budget: int) -> bool:
104
+ """Check if there's sufficient cache space for a multimodal input.
105
+ If there is, return True and update EncoderCacheManager state.
106
+
107
+ If there is not enough free space in `num_free_slots` but there is
108
+ enough reclaimable space in `num_freeable_slots`, entries will be
109
+ evicted from `freeable` (their mm_hash appended to `freed`) until
110
+ enough space is available, and then this method returns True.
111
+ Older entries are evicted first.
112
+
113
+ Returns False only if the requested number of tokens exceeds both
114
+ the free and reclaimable capacities combined.
115
116
Args:
76
- request: The request containing the multimodal input
77
- input_id: Index of the multimodal input within the request
117
+ request: The request containing the multimodal input.
118
+ input_id: Index of the multimodal input within the request.
119
120
Returns:
80
- True if there's enough free cache space to store the encoder output
81
- for this multimodal input
121
+ True if there's enough capacity to hold the encoder output for this
122
+ input (possibly after reclaiming `freeable` entries); otherwise
123
+ False.
124
+
125
+ Note: This method does not allocate physical memory for the encoder
126
+ output but only the state of EncoderCacheManager.
127
"""
128
num_tokens = request.get_num_encoder_tokens(input_id)
84
- return num_tokens <= self.num_free_slots
129
+
130
+ # Not enough compute budget
131
+ if num_tokens > encoder_budget:
132
+ return False
133
+
134
+ # Enough free slots
135
+ if num_tokens <= self.num_free_slots:
136
+ self.num_free_slots -= num_tokens
137
+ self.num_freeable_slots -= num_tokens
138
+ return True
139
+
140
+ # Not enough reclaimable slots
141
+ if num_tokens > self.num_freeable_slots:
142
+ return False
143
+
144
+ # Not enough free slots but enough reclaimable slots
145
+ # NOTE: Eviction takes place here, but physical memory is not freed
146
+ # until model runner is notified by the scheduler output.
147
+ while num_tokens > self.num_free_slots:
148
+ mm_hash, num_free_token = self.freeable.popitem(last=False)
149
+ del self.cached[mm_hash]
150
+ self.freed.append(mm_hash)
151
+ self.num_free_slots += num_free_token
152
+ self.num_free_slots -= num_tokens
153
+ self.num_freeable_slots -= num_tokens
154
+ return True
155
156
def allocate(self, request: Request, input_id: int) -> None:
157
"""Allocate cache space for a multimodal input's encoder output.
158
89
- This method reserves cache space for storing the encoder output of
90
- the specified multimodal input. The actual encoder output storage
91
- happens in the model runner, but this method ensures the cache
92
- manager tracks the allocation.
93
-
94
- Args:
95
- request: The request containing the multimodal input
96
- input_id: Index of the multimodal input within the request
159
+ This reserves cache space for storing the encoder output of the
160
+ specified multimodal input. The actual encoder output storage happens in
161
+ the model runner; this method updates the manager's bookkeeping.
162
163
Note:
99
- This method assumes can_allocate() returned True for the same
100
- request and input_id. It will reduce available cache space.
164
+ This method assumes try_allocate() returned True for the same input.
165
"""
102
- req_id = request.request_id
103
- if req_id not in self.cached:
104
- self.cached[req_id] = set()
105
- self.cached[req_id].add(input_id)
106
- self.num_free_slots -= request.get_num_encoder_tokens(input_id)
166
+ # Encoder cache space budget should be already updated for the
167
+ # multimodal input and non-negative after try_allocate() is called.
168
+ assert self.num_free_slots >= 0
169
+ assert self.num_freeable_slots >= 0
170
+
171
+ mm_hash = request.mm_hashes[input_id]
172
+ request_id = request.request_id
173
+ if mm_hash not in self.cached:
174
+ self.cached[mm_hash] = set()
175
+
176
+ self.cached[mm_hash].add(request_id)
177
178
def get_cached_input_ids(self, request: Request) -> set[int]:
179
"""Get all cached multimodal input IDs for a request.
180
111
- Args:
112
- request: The request to query
113
-
114
- Returns:
115
- Set of input_ids that have cached encoder outputs for this request.
116
- Returns empty set if no inputs are cached for this request.
181
+ Returns the set of input IDs whose `mm_hash` exists in the cache map.
182
+ This includes entries that are currently unreferenced (and thus present
183
+ in `freeable`); for such entries, freeing for this request will be a
184
+ no-op.
185
"""
118
- return self.cached.get(request.request_id, set())
186
+ return {
187
+ input_id
188
+ for input_id in range(len(request.mm_hashes))
189
+ if request.mm_hashes[input_id] in self.cached
190
+ }
191
192
def free_encoder_input(self, request: Request, input_id: int) -> None:
121
- """Free cache space for a single multimodal input's encoder output.
193
+ """Free the request's reference to the encoder input (`mm_data`)
194
123
- This method is called when:
124
- - The encoder output has been fully consumed by the decoder and is
125
- no longer needed (e.g., in vision-language models after image
126
- tokens are processed)
127
- - A request is being cancelled or aborted
195
+ When the reference set for the corresponding `mm_hash` becomes empty,
196
+ the entry is appended to `freeable` and `num_freeable_slots` is
197
+ increased by the number of encoder tokens for that input.
198
129
- Args:
130
- request: The request containing the multimodal input
131
- input_id: Index of the multimodal input to free from cache
199
+ The entry is NOT physically freed until capacity is needed (e.g., by
200
+ `can_allocate`).
201
"""
202
req_id = request.request_id
134
- if req_id not in self.cached:
203
+ mm_hash = request.mm_hashes[input_id]
204
+ # The mm_hash not in cache or the req_id set is empty
205
+ if not self.cached.get(mm_hash, None):
206
return
136
-
137
- self.cached[req_id].discard(input_id)
138
- if len(self.cached[req_id]) == 0:
139
- del self.cached[req_id]
140
- self.num_free_slots += request.get_num_encoder_tokens(input_id)
141
- self.freed.append((req_id, input_id))
207
+ self.cached[mm_hash].discard(req_id)
208
+ if not self.cached[mm_hash]:
209
+ num_tokens = request.get_num_encoder_tokens(input_id)
210
+ self.freeable[mm_hash] = num_tokens
211
+ self.num_freeable_slots += num_tokens
212
213
def free(self, request: Request) -> None:
144
- """Free all cached encoder outputs for a request.
214
+ """Free all encoder input cache reference held by *request*.
215
146
- This method is typically called when a request is finished, cancelled,
147
- or aborted, and all its encoder outputs should be freed from cache.
216
+ For each cached input ID, `free_encoder_input` is invoked.
217
+ The data stays in memory until eviction is triggered by a future
218
+ attempt allocation called by 'can_allocate'.
219
149
- Args:
150
- request: The request whose encoder outputs should be freed
220
+ Typically called when a request is finished, cancelled, or aborted.
221
"""
222
input_ids = self.get_cached_input_ids(request).copy()
223
for input_id in input_ids:
224
self.free_encoder_input(request, input_id)
225
156
- def get_freed_ids(self) -> list[tuple[str, int]]:
226
+ def get_freed_mm_hashes(self) -> list[str]:
227
"""Get and clear the list of recently freed encoder cache entries.
228
159
- This method returns all encoder cache entries that were freed since
160
- the last call to this method. It's used by the scheduler to notify
161
- workers about which encoder outputs can be removed from their caches.
162
-
229
Returns:
164
- List of (request_id, input_id) tuples that were freed since the
165
- last call. The internal freed list is cleared after this call.
230
+ List of mm_hash strings that were actually evicted since the last
231
+ call to be used by the scheduler to notify workers about which
232
+ encoder outputs can be removed from their caches. The internal
233
+ list is cleared after this call.
234
"""
235
freed = self.freed
236
self.freed = []
@@ -177,16 +245,11 @@ def compute_encoder_budget(
245
"""Compute the encoder cache budget based on the model and scheduler
246
configurations.
247
180
- Args:
181
- model_config: Model configuration.
182
- scheduler_config: Scheduler configuration.
183
- mm_registry: Provides information about the token cost.
184
-
248
Returns:
186
- - Compute budget for encoder execution, in unit of number of tokens
187
- in the input sequence.
188
- - Space budget for encoder cache size, in unit of number of tokens
189
- in the input sequence.
249
+ - Compute budget for encoder execution, measured in number of tokens
250
+ from the input sequence.
251
+ - Space budget for encoder cache size, measured in number of tokens
252
+ from the input sequence.
253
"""
254
if mm_registry.supports_multimodal_inputs(model_config):
255
max_tokens_by_modality = mm_registry \
@@ -231,10 +294,10 @@ def compute_mm_encoder_budget(
294
non-text modality.
295
296
Returns:
234
- - Compute budget for encoder execution, in unit of number of tokens
235
- in the input sequence.
236
- - Space budget for encoder cache size, in unit of number of tokens
237
- in the input sequence.
297
+ - Compute budget for encoder execution, measured in number of tokens
298
+ from the input sequence.
299
+ - Space budget for encoder cache size, measured in number of tokens
300
+ from the input sequence.
301
"""
302
303
if not max_tokens_by_modality:
vllm/v1/core/sched/output.py
+3
-3
@@ -143,9 +143,9 @@ class SchedulerOutput:
143
# steps. This is used to notify the workers about the finished requests
144
# so that they can free the cached states for those requests.
145
finished_req_ids: set[str]
146
- # list of (req_id, encoder_input_index) tuples.
147
- # Used to free the encoder cache.
148
- free_encoder_input_ids: list[tuple[str, int]]
146
+ # list of mm_hash strings associated with the encoder outputs to be
147
+ # freed from the encoder cache.
148
+ free_encoder_mm_hashes: list[str]
149
150
# Dict of request ids to their index within the batch
151
# for filling the next token bitmask
vllm/v1/core/sched/scheduler.py
+6
-4
@@ -252,6 +252,7 @@ class Scheduler(SchedulerInterface):
252
preempted_req = self.running.pop()
253
254
self.kv_cache_manager.free(preempted_req)
255
+ self.encoder_cache_manager.free(preempted_req)
256
preempted_req.status = RequestStatus.PREEMPTED
257
preempted_req.num_computed_tokens = 0
258
if self.log_stats:
@@ -550,7 +551,8 @@ class Scheduler(SchedulerInterface):
551
# It contains the request IDs that are finished in between
552
# the previous and the current steps.
553
finished_req_ids=self.finished_req_ids,
553
- free_encoder_input_ids=self.encoder_cache_manager.get_freed_ids(),
554
+ free_encoder_mm_hashes=self.encoder_cache_manager.
555
+ get_freed_mm_hashes(),
556
structured_output_request_ids=structured_output_request_ids,
557
grammar_bitmask=grammar_bitmask,
558
)
@@ -698,7 +700,7 @@ class Scheduler(SchedulerInterface):
700
# in the decoder's KV cache.
701
continue
702
701
- if self.encoder_cache_manager.has_cache(request, i):
703
+ if self.encoder_cache_manager.check_and_update_cache(request, i):
704
# The encoder input is already computed and cached.
705
continue
706
@@ -712,8 +714,8 @@ class Scheduler(SchedulerInterface):
714
num_new_tokens = start_pos - num_computed_tokens
715
break
716
715
- if (not self.encoder_cache_manager.can_allocate(request, i)
716
- or num_encoder_tokens > encoder_budget):
717
+ if not self.encoder_cache_manager.try_allocate(
718
+ request, i, encoder_budget):
719
# The encoder cache is full or the encoder budget is exhausted.
720
# NOTE(woosuk): We assume that the encoder input tokens should
721
# be processed altogether, as the encoder usually uses
vllm/v1/worker/gpu_input_batch.py
+1
@@ -33,6 +33,7 @@ class CachedRequestState:
33
prompt_token_ids: list[int]
34
mm_kwargs: list[MultiModalKwargsItem]
35
mm_positions: list[PlaceholderRange]
36
+ mm_hashes: list[str]
37
sampling_params: Optional[SamplingParams]
38
pooling_params: Optional[PoolingParams]
39
generator: Optional[torch.Generator]
vllm/v1/worker/gpu_model_runner.py
+21
-26
@@ -176,8 +176,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
176
self.attn_groups: list[list[AttentionGroup]] = []
177
# self.kv_cache_config: KVCacheConfig
178
179
- # req_id -> (input_id -> encoder_output)
180
- self.encoder_cache: dict[str, dict[int, torch.Tensor]] = {}
179
+ # mm_hash -> encoder_output
180
+ self.encoder_cache: dict[str, torch.Tensor] = {}
181
182
self.use_aux_hidden_state_outputs = False
183
# Set up speculative decoding.
@@ -436,7 +436,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
436
# Remove finished requests from the cached states.
437
for req_id in scheduler_output.finished_req_ids:
438
self.requests.pop(req_id, None)
439
- self.encoder_cache.pop(req_id, None)
439
# Remove the finished requests from the persistent batch.
440
# NOTE(woosuk): There could be an edge case where finished_req_ids and
441
# scheduled_req_ids overlap. This happens when a request is aborted and
@@ -447,12 +446,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
446
self.input_batch.remove_request(req_id)
447
448
# Free the cached encoder outputs.
450
- for req_id, input_id in scheduler_output.free_encoder_input_ids:
451
- encoder_outputs = self.encoder_cache.get(req_id)
452
- if encoder_outputs is not None:
453
- encoder_outputs.pop(input_id, None)
454
- if not encoder_outputs:
455
- self.encoder_cache.pop(req_id, None)
449
+ for mm_hash in scheduler_output.free_encoder_mm_hashes:
450
+ self.encoder_cache.pop(mm_hash, None)
451
452
# Remove the unscheduled requests from the persistent batch.
453
# NOTE(woosuk): The unscheduled requests are either preempted requests
@@ -496,6 +491,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
491
prompt_token_ids=new_req_data.prompt_token_ids,
492
mm_kwargs=new_req_data.mm_kwargs,
493
mm_positions=new_req_data.mm_positions,
494
+ mm_hashes=new_req_data.mm_hashes,
495
sampling_params=sampling_params,
496
pooling_params=pooling_params,
497
generator=generator,
@@ -1161,17 +1157,18 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1157
scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs
1158
if not scheduled_encoder_inputs:
1159
return
1164
-
1160
# Batch the multi-modal inputs.
1161
mm_kwargs = list[MultiModalKwargsItem]()
1167
- req_ids_pos = list[tuple[str, int, PlaceholderRange]]()
1162
+ # list of tuple (mm_hash, position_info)
1163
+ mm_hashes_pos = list[tuple[str, PlaceholderRange]]()
1164
for req_id, encoder_input_ids in scheduled_encoder_inputs.items():
1165
req_state = self.requests[req_id]
1166
1167
for mm_input_id in encoder_input_ids:
1168
+ mm_hash = req_state.mm_hashes[mm_input_id]
1169
mm_kwargs.append(req_state.mm_kwargs[mm_input_id])
1173
- req_ids_pos.append(
1174
- (req_id, mm_input_id, req_state.mm_positions[mm_input_id]))
1170
+ mm_hashes_pos.append(
1171
+ (mm_hash, req_state.mm_positions[mm_input_id]))
1172
1173
# Batch mm inputs as much as we can: if a request in the batch has
1174
# multiple modalities or a different modality than the previous one,
@@ -1204,15 +1201,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1201
for output in curr_group_outputs:
1202
encoder_outputs.append(output)
1203
1207
- # Cache the encoder outputs.
1208
- for (req_id, input_id, pos_info), output in zip(
1209
- req_ids_pos,
1210
- encoder_outputs,
1211
- ):
1212
- if req_id not in self.encoder_cache:
1213
- self.encoder_cache[req_id] = {}
1214
-
1215
- self.encoder_cache[req_id][input_id] = scatter_mm_placeholders(
1204
+ # Cache the encoder outputs by mm_hash
1205
+ for (mm_hash, pos_info), output in zip(mm_hashes_pos, encoder_outputs):
1206
+ self.encoder_cache[mm_hash] = scatter_mm_placeholders(
1207
output,
1208
is_embed=pos_info.is_embed,
1209
)
@@ -1230,6 +1221,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1221
num_computed_tokens = \
1222
req_state.num_computed_tokens + shift_computed_tokens
1223
mm_positions = req_state.mm_positions
1224
+ mm_hashes = req_state.mm_hashes
1225
for i, pos_info in enumerate(mm_positions):
1226
start_pos = pos_info.offset
1227
num_encoder_tokens = pos_info.length
@@ -1249,11 +1241,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1241
start_idx = max(num_computed_tokens - start_pos, 0)
1242
end_idx = min(
1243
num_computed_tokens - start_pos + num_scheduled_tokens,
1252
- num_encoder_tokens)
1244
+ num_encoder_tokens,
1245
+ )
1246
assert start_idx < end_idx
1254
- assert req_id in self.encoder_cache
1255
- assert i in self.encoder_cache[req_id]
1256
- encoder_output = self.encoder_cache[req_id][i]
1247
+
1248
+ mm_hash = mm_hashes[i]
1249
+ encoder_output = self.encoder_cache.get(mm_hash, None)
1250
+ assert encoder_output is not None,\
1251
+ f"Encoder cache miss for {mm_hash}."
1252
1253
if (is_embed := pos_info.is_embed) is not None:
1254
is_embed = is_embed[start_idx:end_idx]
vllm/v1/worker/tpu_model_runner.py
+19
-18
@@ -208,8 +208,8 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
208
# Lazy initialization
209
self.model: nn.Module # Set after load_model
210
self.kv_caches: list[torch.Tensor] = []
211
- # req_id -> (input_id -> encoder_output)
212
- self.encoder_cache: dict[str, dict[int, torch.Tensor]] = {}
211
+ # mm_hash -> encoder_output
212
+ self.encoder_cache: dict[str, torch.Tensor] = {}
213
214
# Request states.
215
self.requests: dict[str, CachedRequestState] = {}
@@ -342,7 +342,6 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
342
# Remove finished requests from the cached states.
343
for req_id in scheduler_output.finished_req_ids:
344
self.requests.pop(req_id, None)
345
- self.encoder_cache.pop(req_id, None)
345
346
# Remove the finished requests from the persistent batch.
347
# NOTE(woosuk): There could be an edge case where finished_req_ids and
@@ -357,12 +356,8 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
356
removed_req_indices.append(req_index)
357
358
# Free the cached encoder outputs.
360
- for req_id, input_id in scheduler_output.free_encoder_input_ids:
361
- encoder_outputs = self.encoder_cache.get(req_id)
362
- if encoder_outputs is not None:
363
- encoder_outputs.pop(input_id, None)
364
- if not encoder_outputs:
365
- self.encoder_cache.pop(req_id, None)
359
+ for mm_hash in scheduler_output.free_encoder_mm_hashes:
360
+ self.encoder_cache.pop(mm_hash, None)
361
362
# Remove the unscheduled requests from the persistent batch.
363
# NOTE(woosuk): The unscheduled requests are either preempted requests
@@ -394,6 +389,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
389
prompt_token_ids=new_req_data.prompt_token_ids,
390
mm_kwargs=new_req_data.mm_kwargs,
391
mm_positions=new_req_data.mm_positions,
392
+ mm_hashes=new_req_data.mm_hashes,
393
sampling_params=sampling_params,
394
pooling_params=None,
395
generator=None,
@@ -845,14 +841,16 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
841
842
# Batch the multi-modal inputs.
843
mm_kwargs = list[MultiModalKwargsItem]()
848
- req_ids_pos = list[tuple[str, int, PlaceholderRange]]()
844
+ # List of tuple (mm_hash, pos_info)
845
+ mm_hashes_pos = list[tuple[str, PlaceholderRange]]()
846
for req_id, encoder_input_ids in scheduled_encoder_inputs.items():
847
req_state = self.requests[req_id]
848
849
for mm_input_id in encoder_input_ids:
850
+ mm_hash = req_state.mm_hashes[mm_input_id]
851
mm_kwargs.append(req_state.mm_kwargs[mm_input_id])
854
- req_ids_pos.append(
855
- (req_id, mm_input_id, req_state.mm_positions[mm_input_id]))
852
+ mm_hashes_pos.append(
853
+ (mm_hash, req_state.mm_positions[mm_input_id]))
854
855
# Batch mm inputs as much as we can: if a request in the batch has
856
# multiple modalities or a different modality than the previous one,
@@ -895,15 +893,15 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
893
# NOTE (NickLucche) here we diverge from logic in other runners, as we
894
# assume to only have whole mm items to process. Hence we avoid the
895
# intrinsic dynamism that `scatter_mm_placeholders` introduces.
898
- for (req_id, input_id, pos_info), output in zip(
899
- req_ids_pos,
896
+ for (mm_hash, pos_info), output in zip(
897
+ mm_hashes_pos,
898
encoder_outputs,
899
):
900
if req_id not in self.encoder_cache:
901
self.encoder_cache[req_id] = {}
902
assert pos_info.is_embed is None, "Expected all positions to be"\
903
" contiguous and embeddings."
906
- self.encoder_cache[req_id][input_id] = output
904
+ self.encoder_cache[mm_hash] = output
905
906
def _gather_mm_embeddings(
907
self,
@@ -916,6 +914,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
914
req_state = self.requests[req_id]
915
num_computed_tokens = req_state.num_computed_tokens
916
mm_positions = req_state.mm_positions
917
+ mm_hashes = req_state.mm_hashes
918
# TODO unroll loop and assume/enforce --disable_chunked_mm_input
919
# NOTE (NickLucche) here we diverge from logic in other runners, as
920
# we assume to only have whole mm items to process. Hence we avoid
@@ -936,11 +935,13 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
935
# in the decoder's KV cache.
936
continue
937
939
- assert req_id in self.encoder_cache
940
- assert i in self.encoder_cache[req_id]
938
+ mm_hash = mm_hashes[i]
939
+ encoder_output = self.encoder_cache.get(mm_hash, None)
940
+ assert encoder_output is not None,\
941
+ f"Encoder cache miss for {mm_hash}."
942
assert pos_info.is_embed is None, "Expected all positions to"\
943
" be contiguous and embeddings."
943
- encoder_output = self.encoder_cache[req_id][i]
944
+ encoder_output = self.encoder_cache[mm_hash]
945
mm_embeds.append(encoder_output)
946
return mm_embeds
947