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

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 14, 05:30 UTC

[Bugfix] Fix scheduling when repeated images in one request (#23544)

Signed-off-by: Roger Wang <hey@rogerw.me> Signed-off-by: Roger Wang <hey@rogerw.io> Co-authored-by: Roger Wang <hey@rogerw.me> Co-authored-by: knlnguyen1802 <knlnguyen1802@gmail.com>

Roger Wang committed Aug 26, 2025 at 02:46 UTC b5d34af3286ee0334d9f7bd729774ac55c5805e9

3 files changed +96 -39

tests/v1/core/test_encoder_cache_manager.py

+38 -11

@@ -22,7 +22,7 @@ def test_basic_allocate_and_reuse():

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

24 assert not cache.check_and_update_cache(req, 0)

25 - assert cache.try_allocate(req, 0, int(1e9))

25 + assert cache.can_allocate(req, 0, int(1e9), 0)

27 cache.allocate(req, 0)

@@ -44,7 +44,7 @@ def test_freeing_decreases_refcount_and_moves_to_freeable():

44 manager = EncoderCacheManager(cache_size=10)

45 req = MockRequest("req2", ["img3"], [5])

47 - assert manager.try_allocate(req, 0, int(1e9))

47 + assert manager.can_allocate(req, 0, int(1e9), 0)

48 manager.allocate(req, 0)

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

@@ -60,10 +60,10 @@ def test_free_request_frees_all_inputs():

60 manager = EncoderCacheManager(cache_size=10)

61 req = MockRequest("req3", ["a", "b"], [2, 3])

63 - assert manager.try_allocate(req, 0, int(1e9))

63 + assert manager.can_allocate(req, 0, int(1e9), 0)

64 manager.allocate(req, 0)

66 - assert manager.try_allocate(req, 1, int(1e9))

66 + assert manager.can_allocate(req, 1, int(1e9), 0)

67 manager.allocate(req, 1)

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

@@ -84,11 +84,11 @@ def test_eviction_when_cache_is_full():

84 req1 = MockRequest("req1", ["x"], [6])

85 req2 = MockRequest("req2", ["y"], [5])

87 - assert manager.try_allocate(req1, 0, int(1e9))

87 + assert manager.can_allocate(req1, 0, int(1e9), 0)

88 manager.allocate(req1, 0)

89 manager.free_encoder_input(req1, 0)

91 - assert manager.try_allocate(req2, 0, int(1e9))

91 + assert manager.can_allocate(req2, 0, int(1e9), 0)

92 manager.allocate(req2, 0)

94 # 'x' should have been evicted.

@@ -100,10 +100,10 @@ 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))

103 + assert manager.can_allocate(req, 0, int(1e9), 0)

104 manager.allocate(req, 0)

105

106 - assert manager.try_allocate(req, 2, int(1e9))

106 + assert manager.can_allocate(req, 2, int(1e9), 0)

107 manager.allocate(req, 2)

108

109 cached_ids = manager.get_cached_input_ids(req)

@@ -114,7 +114,7 @@ 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))

117 + assert manager.can_allocate(req, 0, int(1e9), 0)

118 manager.allocate(req, 0)

119

120 manager.free_encoder_input(req, 0)

@@ -131,14 +131,41 @@ def test_get_freed_mm_hashes_clears_freed_list():

131 req1 = MockRequest("reqA", ["a"], [5])

132 req2 = MockRequest("reqB", ["b"], [6])

133

134 - assert manager.try_allocate(req1, 0, int(1e9))

134 + assert manager.can_allocate(req1, 0, int(1e9), 0)

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))

139 + assert manager.can_allocate(req2, 0, int(1e9), 0)

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() == []

145 +

146 +

147 +def test_schedule_request_multi_images_respect_space_limit():

148 + manager = EncoderCacheManager(cache_size=10)

149 + req = MockRequest("reqA", ["a", "b"], [5, 6])

150 + compute_budget = 100

151 +

152 + num_tokens_to_schedule = 0

153 + assert manager.can_allocate(req, 0, compute_budget, num_tokens_to_schedule)

154 + num_tokens_to_schedule += req.get_num_encoder_tokens(0)

155 + compute_budget -= req.get_num_encoder_tokens(0)

156 +

157 + assert not manager.can_allocate(req, 1, compute_budget,

158 + num_tokens_to_schedule)

159 +

160 +

161 +def test_schedule_request_multi_images_respect_compute_limit():

162 + manager = EncoderCacheManager(cache_size=100)

163 + req = MockRequest("reqA", ["a", "b"], [5, 6])

164 + compute_budget = 10

165 + num_tokens_to_schedule = 0

166 + assert manager.can_allocate(req, 0, compute_budget, num_tokens_to_schedule)

167 + num_tokens_to_schedule += req.get_num_encoder_tokens(0)

168 + compute_budget -= req.get_num_encoder_tokens(0)

169 +

170 + assert not manager.can_allocate(req, 1, compute_budget,

171 + num_tokens_to_schedule)

vllm/v1/core/encoder_cache_manager.py

+20 -12

@@ -99,8 +99,9 @@ class EncoderCacheManager:

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:

102 + def can_allocate(self, request: Request, input_id: int,

103 + encoder_compute_budget: int,

104 + num_tokens_to_schedule: int) -> bool:

105 """Check if there's sufficient cache space for a multimodal input.

106 If there is, return True and update EncoderCacheManager state.

107

@@ -116,6 +117,10 @@ class EncoderCacheManager:

117 Args:

118 request: The request containing the multimodal input.

119 input_id: Index of the multimodal input within the request.

120 + encoder_compute_budget: Number of encoder tokens allowed to be

121 + computed when this method is invoked.

122 + num_tokens_to_schedule: Number of tokens already scheduled to be

123 + allocated with cache space when this method is invoked.

124

125 Returns:

126 True if there's enough capacity to hold the encoder output for this

@@ -128,13 +133,13 @@ class EncoderCacheManager:

133 num_tokens = request.get_num_encoder_tokens(input_id)

134

135 # Not enough compute budget

131 - if num_tokens > encoder_budget:

136 + if num_tokens > encoder_compute_budget:

137 return False

138

139 + num_tokens += num_tokens_to_schedule

140 +

141 # Enough free slots

142 if num_tokens <= self.num_free_slots:

136 - self.num_free_slots -= num_tokens

137 - self.num_freeable_slots -= num_tokens

143 return True

144

145 # Not enough reclaimable slots

@@ -149,8 +154,6 @@ class EncoderCacheManager:

154 del self.cached[mm_hash]

155 self.freed.append(mm_hash)

156 self.num_free_slots += num_free_token

152 - self.num_free_slots -= num_tokens

153 - self.num_freeable_slots -= num_tokens

157 return True

158

159 def allocate(self, request: Request, input_id: int) -> None:

@@ -161,19 +164,24 @@ class EncoderCacheManager:

164 the model runner; this method updates the manager's bookkeeping.

165

166 Note:

164 - This method assumes try_allocate() returned True for the same input.

167 + This method assumes can_allocate() returned True for the same input.

168 """

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

169

170 mm_hash = request.mm_hashes[input_id]

171 request_id = request.request_id

172 if mm_hash not in self.cached:

173 self.cached[mm_hash] = set()

174

175 + num_encoder_tokens = request.get_num_encoder_tokens(input_id)

176 +

177 + # NOTE: Encoder cache should always have enough space for encoder inputs

178 + # that are scheduled since eviction takes place at can_allocate().

179 + assert self.num_free_slots >= num_encoder_tokens

180 + assert self.num_freeable_slots >= num_encoder_tokens

181 +

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

183 + self.num_free_slots -= num_encoder_tokens

184 + self.num_freeable_slots -= num_encoder_tokens

185

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

187 """Get all cached multimodal input IDs for a request.

vllm/v1/core/sched/scheduler.py

+38 -16

@@ -182,7 +182,7 @@ class Scheduler(SchedulerInterface):

182 token_budget = self.max_num_scheduled_tokens

183 # Encoder-related.

184 scheduled_encoder_inputs: dict[str, list[int]] = {}

185 - encoder_budget = self.max_num_encoder_input_tokens

185 + encoder_compute_budget = self.max_num_encoder_input_tokens

186 # Spec decode-related.

187 scheduled_spec_decode_tokens: dict[str, list[int]] = {}

188

@@ -211,12 +211,13 @@ class Scheduler(SchedulerInterface):

211

212 # Schedule encoder inputs.

213 encoder_inputs_to_schedule = None

214 - new_encoder_budget = encoder_budget

214 + new_encoder_compute_budget = encoder_compute_budget

215 if request.has_encoder_inputs:

216 (encoder_inputs_to_schedule, num_new_tokens,

217 - new_encoder_budget) = self._try_schedule_encoder_inputs(

217 + new_encoder_compute_budget

218 + ) = self._try_schedule_encoder_inputs(

219 request, request.num_computed_tokens, num_new_tokens,

219 - encoder_budget)

220 + encoder_compute_budget)

221

222 if num_new_tokens == 0:

223 # The request cannot be scheduled because one of the following

@@ -298,7 +299,7 @@ class Scheduler(SchedulerInterface):

299 # Allocate the encoder cache.

300 for i in encoder_inputs_to_schedule:

301 self.encoder_cache_manager.allocate(request, i)

301 - encoder_budget = new_encoder_budget

302 + encoder_compute_budget = new_encoder_compute_budget

303

304 # Record the LoRAs in scheduled_running_reqs

305 scheduled_loras: set[int] = set()

@@ -382,7 +383,7 @@ class Scheduler(SchedulerInterface):

383 num_computed_tokens = request.num_computed_tokens

384

385 encoder_inputs_to_schedule = None

385 - new_encoder_budget = encoder_budget

386 + new_encoder_compute_budget = encoder_compute_budget

387

388 # KVTransfer: loading remote KV, do not allocate for new work.

389 if load_kv_async:

@@ -413,10 +414,10 @@ class Scheduler(SchedulerInterface):

414 # Schedule encoder inputs.

415 if request.has_encoder_inputs:

416 (encoder_inputs_to_schedule, num_new_tokens,

416 - new_encoder_budget

417 + new_encoder_compute_budget

418 ) = self._try_schedule_encoder_inputs(

419 request, num_computed_tokens, num_new_tokens,

419 - encoder_budget)

420 + encoder_compute_budget)

421 if num_new_tokens == 0:

422 # The request cannot be scheduled.

423 break

@@ -495,7 +496,7 @@ class Scheduler(SchedulerInterface):

496 # Allocate the encoder cache.

497 for i in encoder_inputs_to_schedule:

498 self.encoder_cache_manager.allocate(request, i)

498 - encoder_budget = new_encoder_budget

499 + encoder_compute_budget = new_encoder_compute_budget

500

501 # Put back any skipped requests at the head of the waiting queue

502 if skipped_waiting_requests:

@@ -658,7 +659,7 @@ class Scheduler(SchedulerInterface):

659 request: Request,

660 num_computed_tokens: int,

661 num_new_tokens: int,

661 - encoder_budget: int,

662 + encoder_compute_budget: int,

663 ) -> tuple[list[int], int, int]:

664 """

665 Determine which encoder inputs need to be scheduled in the current step,

@@ -680,11 +681,17 @@ class Scheduler(SchedulerInterface):

681 blocks and externally cached blocks (via KVConnector).

682 """

683 if num_new_tokens == 0 or not request.has_encoder_inputs:

683 - return [], num_new_tokens, encoder_budget

684 + return [], num_new_tokens, encoder_compute_budget

685 encoder_inputs_to_schedule: list[int] = []

686 mm_positions = request.mm_positions

687 assert mm_positions is not None

688 assert len(mm_positions) > 0

689 +

690 + # NOTE: since scheduler operates on the request level (possibly with

691 + # multiple encoder inputs per request), we need to create temporary

692 + # trackers for accounting at the encoder input level.

693 + mm_hashes_to_schedule = set()

694 + num_tokens_to_schedule = 0

695 for i, pos_info in enumerate(mm_positions):

696 start_pos = pos_info.offset

697 num_encoder_tokens = pos_info.length

@@ -695,13 +702,20 @@ class Scheduler(SchedulerInterface):

702 if start_pos >= num_computed_tokens + num_new_tokens:

703 # The encoder input is not needed in this step.

704 break

705 +

706 if start_pos + num_encoder_tokens <= num_computed_tokens:

707 # The encoder input is already computed and stored

708 # in the decoder's KV cache.

709 continue

710

711 + # The same encoder input has already been scheduled in the current

712 + # step.

713 + if request.mm_hashes[i] in mm_hashes_to_schedule:

714 + continue

715 +

716 if self.encoder_cache_manager.check_and_update_cache(request, i):

704 - # The encoder input is already computed and cached.

717 + # The encoder input is already computed and cached from a

718 + # previous step.

719 continue

720

721 # If no encoder input chunking is allowed, we do not want to

@@ -714,8 +728,9 @@ class Scheduler(SchedulerInterface):

728 num_new_tokens = start_pos - num_computed_tokens

729 break

730

717 - if not self.encoder_cache_manager.try_allocate(

718 - request, i, encoder_budget):

731 + if not self.encoder_cache_manager.can_allocate(

732 + request, i, encoder_compute_budget,

733 + num_tokens_to_schedule):

734 # The encoder cache is full or the encoder budget is exhausted.

735 # NOTE(woosuk): We assume that the encoder input tokens should

736 # be processed altogether, as the encoder usually uses

@@ -732,9 +747,16 @@ class Scheduler(SchedulerInterface):

747 num_new_tokens = 0

748 break

749

735 - encoder_budget -= num_encoder_tokens

750 + num_tokens_to_schedule += num_encoder_tokens

751 + encoder_compute_budget -= num_encoder_tokens

752 + mm_hashes_to_schedule.add(request.mm_hashes[i])

753 encoder_inputs_to_schedule.append(i)

737 - return encoder_inputs_to_schedule, num_new_tokens, encoder_budget

754 +

755 + return (

756 + encoder_inputs_to_schedule,

757 + num_new_tokens,

758 + encoder_compute_budget,

759 + )

760

761 def get_grammar_bitmask(

762 self,