@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 95b1615ec9

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 2, 02:15 UTC

[Perf] Improve multimodal item handling from O(n) to O(log n) per step (#44212)

Signed-off-by: Andy Lo <andy@mistral.ai>

Andy Lo committed Jun 3, 2026 at 12:00 UTC 95b1615ec97dac31e670ed71e5cd2f3134d67ba2

6 files changed +92 -44

tests/v1/core/test_encoder_cache_manager.py

+30

@@ -43,6 +43,7 @@ def test_basic_allocate_and_reuse():

44 assert cache.check_and_update_cache(req, 0)

45 assert "r1" in cache.cached["imgA"]

46 + assert cache.request_cached_ids["r1"] == {0}

47 assert cache.num_free_slots == 6

49 # Free twice to bring refcount to 0.

@@ -50,6 +51,7 @@ def test_basic_allocate_and_reuse():

51 cache.free_encoder_input(req, 0)

53 assert not cache.cached["imgA"]

54 + assert "r1" not in cache.request_cached_ids

55 assert "imgA" in cache.freeable

56 assert cache.num_freeable_slots == 10

57 assert cache.num_free_slots == 6

@@ -63,10 +65,12 @@ def test_freeing_decreases_refcount_and_moves_to_freeable():

65 manager.allocate(req, 0)

67 assert len(manager.cached["img3"]) == 1

68 + assert manager.request_cached_ids["req2"] == {0}

70 manager.free_encoder_input(req, 0)

72 assert not manager.cached["img3"]

73 + assert "req2" not in manager.request_cached_ids

74 assert "img3" in manager.freeable

75 assert manager.num_freeable_slots == 10

@@ -83,11 +87,13 @@ def test_free_request_frees_all_inputs():

88 assert len(manager.cached["a"]) == 1

89 assert len(manager.cached["b"]) == 1

90 + assert manager.request_cached_ids["req3"] == {0, 1}

92 manager.free(req)

94 assert not manager.cached["a"]

95 assert not manager.cached["b"]

96 + assert "req3" not in manager.request_cached_ids

97 assert "a" in manager.freeable

98 assert "b" in manager.freeable

99 assert manager.num_freeable_slots == 10

@@ -108,6 +114,7 @@ def test_eviction_when_cache_is_full():

114

115 # 'x' should have been evicted.

116 assert "x" not in manager.cached

117 + assert "req1" not in manager.request_cached_ids

118 assert "x" in manager.get_freed_mm_hashes()

119

120

@@ -137,6 +144,7 @@ def test_has_cache_restores_from_freeable():

144 # Should restore from freeable.

145 assert manager.check_and_update_cache(req, 0)

146 assert len(manager.cached["imgZ"]) == 1

147 + assert manager.request_cached_ids["reqY"] == {0}

148 assert "imgZ" not in manager.freeable

149 assert manager.num_freeable_slots == 6

150

@@ -205,6 +213,7 @@ def test_encoder_cache_with_is_embed_mask():

213

214 assert manager.num_free_slots == 92

215 assert "img1" in manager.cached

216 + assert manager.request_cached_ids["r1"] == {0}

217

218 old_size = 100

219 new_size = request.mm_features[0].mm_position.get_num_embeds()

@@ -276,6 +285,7 @@ def test_reset_clears_all_state():

285 manager.reset()

286

287 assert len(manager.cached) == 0

288 + assert len(manager.request_cached_ids) == 0

289 assert len(manager.freeable) == 0

290 assert len(manager.freed) == 0

291 assert manager.num_free_slots == 20

@@ -298,6 +308,26 @@ def test_reset_allows_fresh_allocations():

308 assert manager.num_free_slots == 2

309 assert "img2" in manager.cached

310 assert "img1" not in manager.cached

311 + assert manager.request_cached_ids["req2"] == {0}

312 + assert "req1" not in manager.request_cached_ids

313 +

314 +

315 +def test_free_request_with_duplicate_mm_hashes():

316 + """Freeing a request whose two inputs share the same mm_hash must fully

317 + clean up request_cached_ids. After the first free_encoder_input call,

318 + cached[mm_hash] becomes empty; the second call must still remove the

319 + remaining input_id from request_cached_ids."""

320 + manager = EncoderCacheManager(cache_size=20)

321 +

322 + req = MockRequest("r1", ["imgA", "imgA"], [4, 4])

323 +

324 + manager.allocate(req, 0)

325 + # input 1 has the same hash, so it's already cached.

326 + assert manager.check_and_update_cache(req, 1)

327 + assert manager.request_cached_ids["r1"] == {0, 1}

328 +

329 + manager.free(req)

330 + assert "r1" not in manager.request_cached_ids

331

332

333 def test_encoder_decoder_cache_manager_reset():

vllm/multimodal/utils.py

+25

@@ -1,6 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 +import bisect

5 import mimetypes

6 from collections import defaultdict

7 from collections.abc import Generator, Sequence

@@ -18,6 +19,7 @@ from vllm.utils.import_utils import LazyLoader

19 from .hasher import MultiModalHasher

20 from .inputs import (

21 BatchedTensorInputs,

22 + MultiModalFeatureSpec,

23 MultiModalFieldElem,

24 MultiModalKwargsItem,

25 MultiModalSharedField,

@@ -109,6 +111,29 @@ def encode_video_url(

111 return f"data:{mimetype};base64,{video_b64}"

112

113

114 +def get_mm_features_in_window(

115 + mm_features: list[MultiModalFeatureSpec],

116 + start: int,

117 + end: int,

118 +) -> tuple[int, int]:

119 + """Return (lo, hi) indices for features overlapping [start, end).

120 +

121 + Assumes mm_features are sorted by offset and non-overlapping, so

122 + offset + length is also sorted.

123 + """

124 + lo = bisect.bisect_left(

125 + mm_features,

126 + start + 1,

127 + key=lambda f: f.mm_position.offset + f.mm_position.length,

128 + )

129 + hi = bisect.bisect_left(

130 + mm_features,

131 + end,

132 + key=lambda f: f.mm_position.offset,

133 + )

134 + return lo, hi

135 +

136 +

137 def argsort_mm_positions(

138 mm_positions: MultiModalPlaceholders,

139 ) -> list[tuple[str, int]]:

vllm/transformers_utils/processors/voxtral.py

+1 -1

@@ -44,7 +44,7 @@ class MistralCommonFeatureExtractor:

44 if not self.audio_encoder.audio_config.is_streaming:

45 audio = self.audio_encoder.pad(audio, self.sampling_rate)

47 - audios_processed.append(torch.tensor(audio))

47 + audios_processed.append(torch.from_numpy(audio))

49 return BatchFeature(

50 {"audio_arrays": audios_processed}, tensor_type=return_tensors

vllm/v1/core/encoder_cache_manager.py

+14 -14

@@ -71,6 +71,8 @@ class EncoderCacheManager:

72 # mm_hash of mm_data => ids of requests that reference the mm_data

73 self.cached: dict[str, set[str]] = {}

74 + # request_id => set of input_ids cached for that request

75 + self.request_cached_ids: dict[str, set[int]] = {}

77 # mm_hash of mm_data => num_encoder_embeds of the mm_data

78 self.freeable: OrderedDict[str, int] = OrderedDict()

@@ -83,6 +85,7 @@ class EncoderCacheManager:

85 Called when model weights are updated to invalidate stale embeddings.

86 """

87 self.cached.clear()

88 + self.request_cached_ids.clear()

89 self.freeable.clear()

90 self.freed.clear()

91 self.num_free_slots = self.cache_size

@@ -114,6 +117,7 @@ class EncoderCacheManager:

117 self.num_freeable_slots -= num_encoder_embeds

118

119 self.cached[mm_hash].add(request.request_id)

120 + self.request_cached_ids.setdefault(request.request_id, set()).add(input_id)

121 return True

122

123 def can_allocate(

@@ -201,22 +205,13 @@ class EncoderCacheManager:

205 assert self.num_freeable_slots >= num_encoder_embeds

206

207 self.cached[mm_hash].add(request_id)

208 + self.request_cached_ids.setdefault(request_id, set()).add(input_id)

209 self.num_free_slots -= num_encoder_embeds

210 self.num_freeable_slots -= num_encoder_embeds

211

212 def get_cached_input_ids(self, request: Request) -> set[int]:

208 - """Get all cached multimodal input IDs for a request.

209 -

210 - Returns the set of input IDs whose `mm_hash` exists in the cache map.

211 - This includes entries that are currently unreferenced (and thus present

212 - in `freeable`); for such entries, freeing for this request will be a

213 - no-op.

214 - """

215 - return {

216 - input_id

217 - for input_id in range(len(request.mm_features))

218 - if request.mm_features[input_id].identifier in self.cached

219 - }

213 + """Get all cached multimodal input IDs for a request."""

214 + return self.request_cached_ids.get(request.request_id, set())

215

216 def free_encoder_input(self, request: Request, input_id: int) -> None:

217 """Free the request's reference to the encoder input (`mm_data`)

@@ -230,6 +225,12 @@ class EncoderCacheManager:

225 """

226 req_id = request.request_id

227 mm_hash = request.mm_features[input_id].identifier

228 + # Always clean up request_cached_ids, even if the mm_hash was

229 + # already evicted from cache (e.g. by can_allocate).

230 + if req_id in self.request_cached_ids:

231 + self.request_cached_ids[req_id].discard(input_id)

232 + if not self.request_cached_ids[req_id]:

233 + del self.request_cached_ids[req_id]

234 # The mm_hash not in cache or the req_id set is empty

235 if not self.cached.get(mm_hash, None):

236 return

@@ -248,8 +249,7 @@ class EncoderCacheManager:

249

250 Typically called when a request is finished, cancelled, or aborted.

251 """

251 - input_ids = self.get_cached_input_ids(request)

252 - for input_id in input_ids:

252 + for input_id in list(self.get_cached_input_ids(request)):

253 self.free_encoder_input(request, input_id)

254

255 def get_freed_mm_hashes(self) -> list[str]:

vllm/v1/core/sched/scheduler.py

+13 -15

@@ -29,6 +29,7 @@ from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (

29 )

30 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry

31 from vllm.multimodal.encoder_budget import MultiModalBudget

32 +from vllm.multimodal.utils import get_mm_features_in_window

33 from vllm.v1.core.encoder_cache_manager import (

34 EncoderCacheManager,

35 EncoderDecoderCacheManager,

@@ -1163,22 +1164,23 @@ class Scheduler(SchedulerInterface):

1164 # trackers for accounting at the encoder input level.

1165 mm_hashes_to_schedule = set()

1166 num_embeds_to_schedule = 0

1166 - for i, mm_feature in enumerate(mm_features):

1167 +

1168 + lo, hi = get_mm_features_in_window(

1169 + mm_features,

1170 + start=num_computed_tokens,

1171 + end=num_computed_tokens + num_new_tokens + shift_computed_tokens,

1172 + )

1173 + # For encoder-decoder, all inputs sit at start_pos=0, so lo=0 always.

1174 + if self.is_encoder_decoder:

1175 + lo = 0

1176 +

1177 + for i in range(lo, hi):

1178 + mm_feature = mm_features[i]

1179 start_pos = mm_feature.mm_position.offset

1180 num_encoder_tokens = mm_feature.mm_position.length

1181 num_encoder_embeds = mm_feature.mm_position.get_num_embeds()

1182 item_identifier = mm_feature.identifier

1183

1172 - # The encoder output is needed if the two ranges overlap:

1173 - # [num_computed_tokens, num_computed_tokens + num_new_tokens) and

1174 - # [start_pos, start_pos + num_encoder_tokens)

1175 - if (

1176 - start_pos

1177 - >= num_computed_tokens + num_new_tokens + shift_computed_tokens

1178 - ):

1179 - # The encoder input is not needed in this step.

1180 - break

1181 -

1184 if self.is_encoder_decoder and num_computed_tokens > 0:

1185 assert start_pos == 0, (

1186 "Encoder input should be processed at the beginning of "

@@ -1194,10 +1196,6 @@ class Scheduler(SchedulerInterface):

1196 # decoder tokens (num_computed_tokens > 0), then we know we

1197 # already calculated encoder inputs and can skip here.

1198 continue

1197 - elif start_pos + num_encoder_tokens <= num_computed_tokens:

1198 - # The encoder input is already computed and stored

1199 - # in the decoder's KV cache.

1200 - continue

1199

1200 if not self.is_encoder_decoder:

1201 # We are not using the encoder cache for encoder-decoder models,

vllm/v1/worker/gpu_model_runner.py

+9 -14

@@ -104,7 +104,7 @@ from vllm.multimodal.inputs import (

104 MultiModalKwargsItem,

105 PlaceholderRange,

106 )

107 -from vllm.multimodal.utils import group_and_batch_mm_kwargs

107 +from vllm.multimodal.utils import get_mm_features_in_window, group_and_batch_mm_kwargs

108 from vllm.platforms import current_platform

109 from vllm.pooling_params import PoolingParams

110 from vllm.sampling_params import SamplingType

@@ -3100,23 +3100,18 @@ class GPUModelRunner(

3100 req_state = self.requests[req_id]

3101 num_computed_tokens = req_state.num_computed_tokens + shift_computed_tokens

3102

3103 - for mm_feature in req_state.mm_features:

3103 + mm_features = req_state.mm_features

3104 + lo, hi = get_mm_features_in_window(

3105 + mm_features,

3106 + start=num_computed_tokens,

3107 + end=num_computed_tokens + num_scheduled_tokens,

3108 + )

3109 + for i in range(lo, hi):

3110 + mm_feature = mm_features[i]

3111 pos_info = mm_feature.mm_position

3112 start_pos = pos_info.offset

3113 num_encoder_tokens = pos_info.length

3114

3108 - # The encoder output is needed if the two ranges overlap:

3109 - # [num_computed_tokens,

3110 - # num_computed_tokens + num_scheduled_tokens) and

3111 - # [start_pos, start_pos + num_encoder_tokens)

3112 - if start_pos >= num_computed_tokens + num_scheduled_tokens:

3113 - # The encoder output is not needed in this step.

3114 - break

3115 - if start_pos + num_encoder_tokens <= num_computed_tokens:

3116 - # The encoder output is already processed and stored

3117 - # in the decoder's KV cache.

3118 - continue

3119 -

3115 start_idx = max(num_computed_tokens - start_pos, 0)

3116 end_idx = min(

3117 num_computed_tokens - start_pos + num_scheduled_tokens,