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

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 9, 20:45 UTC

[EC Connector] Optimize remote cache check in scheduler (#32585)

Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>

knlnguyen1802 committed Jan 22, 2026 at 11:30 UTC 378385b90cddbe8cbc6e51d4ed59ce83e499530a

5 files changed +76 -57

tests/v1/core/test_scheduler.py

+24 -16

@@ -2560,15 +2560,14 @@ def test_ec_connector_cache_hit_external_load(use_kv_connector):

2560 mm_positions=mm_positions,

2561 )[0]

2562

2563 - # Mock cache hit - encoder cache exists externally

2564 - scheduler.ec_connector.has_caches = Mock(return_value=[True])

2563 + # Mock cache hit - encoder cache has_exists externally

2564 + scheduler.ec_connector.has_cache_item = Mock(return_value=True)

2565 scheduler.ec_connector.update_state_after_alloc = Mock(

2566 wraps=scheduler.ec_connector.update_state_after_alloc

2567 )

2568

2569 scheduler.add_request(request)

2570 output = scheduler.schedule()

2571 -

2571 # Should schedule prompt tokens

2572 scheduled_tokens = output.num_scheduled_tokens[request.request_id]

2573 assert scheduled_tokens == NUM_TOKENS

@@ -2611,7 +2610,7 @@ def test_ec_connector_cache_miss_computes_locally(use_kv_connector):

2610 )[0]

2611

2612 # Mock cache miss - encoder cache doesn't exist externally

2614 - scheduler.ec_connector.has_caches = Mock(return_value=[False])

2613 + scheduler.ec_connector.has_cache_item = Mock(return_value=False)

2614

2615 scheduler.add_request(request_mm_missed)

2616 output = scheduler.schedule()

@@ -2665,7 +2664,7 @@ def test_ec_connector_with_partial_cache_hit_multi_round(use_kv_connector):

2664 PlaceholderRange(offset=250, length=NUM_ENCODER_TOKENS_1),

2665 ]

2666 ]

2668 -

2667 + has_cache_item_result_map_1 = {"hash1_A": False, "hash1_B": True, "hash1_F": True}

2668 # Create request with 4 MM items, with 2 identical items

2669 request1 = create_requests(

2670 num_requests=1,

@@ -2676,7 +2675,9 @@ def test_ec_connector_with_partial_cache_hit_multi_round(use_kv_connector):

2675 )[0]

2676

2677 # Mock partial cache hit: 1st and 3rd missing, 2nd and 4th exist

2679 - scheduler.ec_connector.has_caches = Mock(return_value=[False, True, False, True])

2678 + scheduler.ec_connector.has_cache_item = Mock(

2679 + side_effect=lambda hash_val: has_cache_item_result_map_1[hash_val]

2680 + )

2681 scheduler.ec_connector.update_state_after_alloc = Mock(

2682 wraps=scheduler.ec_connector.update_state_after_alloc

2683 )

@@ -2736,7 +2737,12 @@ def test_ec_connector_with_partial_cache_hit_multi_round(use_kv_connector):

2737 PlaceholderRange(offset=250, length=NUM_ENCODER_TOKENS_2),

2738 ]

2739 ]

2739 -

2740 + has_cache_item_result_map_2 = {

2741 + "hash1_C": True,

2742 + "hash1_D": False,

2743 + "hash1_E": False,

2744 + "hash1_A": True,

2745 + }

2746 request2 = create_requests(

2747 num_requests=1,

2748 num_tokens=NUM_TOKENS_2,

@@ -2746,7 +2752,9 @@ def test_ec_connector_with_partial_cache_hit_multi_round(use_kv_connector):

2752 )[0]

2753

2754 # Mock partial cache hit: only hash1_A and hash1_C exist in connector

2749 - scheduler.ec_connector.has_caches = Mock(return_value=[True, False, False, True])

2755 + scheduler.ec_connector.has_cache_item = Mock(

2756 + side_effect=lambda hash_val: has_cache_item_result_map_2[hash_val]

2757 + )

2758

2759 scheduler.add_request(request2)

2760 output = scheduler.schedule()

@@ -2821,9 +2829,9 @@ def test_ec_connector_schedule_multiple_requests(cache_exist, use_kv_connector):

2829

2830 if cache_exist == "connector_only":

2831 # Cache exist in ec_connector

2824 - scheduler.ec_connector.has_caches = Mock(return_value=[True])

2832 + scheduler.ec_connector.has_cache_item = Mock(return_value=True)

2833 elif cache_exist == "no_where":

2826 - scheduler.ec_connector.has_caches = Mock(return_value=[False])

2834 + scheduler.ec_connector.has_cache_item = Mock(return_value=False)

2835

2836 output = scheduler.schedule()

2837 assert len(output.scheduled_new_reqs) == len(requests)

@@ -2887,7 +2895,7 @@ def test_ec_connector_unable_to_allocate(use_kv_connector):

2895 )

2896

2897 # Mock ec_connector load external cache behavior

2890 - scheduler.ec_connector.has_caches = Mock(return_value=[True])

2898 + scheduler.ec_connector.has_cache_item = Mock(return_value=True)

2899 scheduler.ec_connector.update_state_after_alloc = Mock(

2900 wraps=scheduler.ec_connector.update_state_after_alloc

2901 )

@@ -2984,7 +2992,7 @@ def test_priority_scheduling_ec_connector_preemption_and_resumption(

2992 )

2993

2994 # Mock cache hit: Both cache exist in connector (at E->PD initially)

2987 - scheduler.ec_connector.has_caches = Mock(return_value=[True])

2995 + scheduler.ec_connector.has_cache_item = Mock(return_value=True)

2996 scheduler.ec_connector.update_state_after_alloc = Mock(

2997 wraps=scheduler.ec_connector.update_state_after_alloc

2998 )

@@ -3139,9 +3147,9 @@ def test_priority_scheduling_ec_connector_preemption_and_resumption(

3147

3148 if cache_exist == "connector_only":

3149 # Cache exist in ec_connector

3142 - scheduler.ec_connector.has_caches = Mock(return_value=[True])

3150 + scheduler.ec_connector.has_cache_item = Mock(return_value=True)

3151 elif cache_exist == "no_where":

3144 - scheduler.ec_connector.has_caches = Mock(return_value=[False])

3152 + scheduler.ec_connector.has_cache_item = Mock(return_value=False)

3153

3154 # 4th Schedule - this should trigger req_low resumption from waiting

3155 output = scheduler.schedule()

@@ -3259,8 +3267,8 @@ def test_ec_connector_allocate_encoder_tokens_with_external_load(use_kv_connecto

3267 )[0]

3268

3269 # Mock cache hit: MM of request1 NOT cached remotely, request2 cached remotely

3262 - scheduler.ec_connector.has_caches = Mock(

3263 - side_effect=lambda req: [True, True, True] if req == request2 else [False]

3270 + scheduler.ec_connector.has_cache_item = Mock(

3271 + side_effect=lambda hash_value: hash_value in mm_hashes_list_2[0]

3272 )

3273 scheduler.ec_connector.update_state_after_alloc = Mock(

3274 wraps=scheduler.ec_connector.update_state_after_alloc

tests/v1/ec_connector/unit/test_ec_example_connector.py

+31 -17

@@ -123,15 +123,15 @@ class TestECExampleConnectorBasics:

123

124

125 class TestCacheExistence:

126 - """Test cache existence checking using has_caches() API."""

126 + """Test cache existence checking using has_cache_item() API."""

127

128 - def test_has_caches_all_exist_3_items(

128 + def test_has_cache_item_all_exist_3_items(

129 self,

130 mock_vllm_config_producer,

131 mock_vllm_config_consumer,

132 mock_request_with_3_mm,

133 ):

134 - """Test has_caches returns True when all 3 caches exist."""

134 + """Test has_cache_item returns True when all 3 caches exist."""

135 # Test for producer first

136 producer = ECExampleConnector(

137 vllm_config=mock_vllm_config_producer,

@@ -146,8 +146,11 @@ class TestCacheExistence:

146 encoder_cache[mm_hash] = torch.randn(10, 768)

147 producer.save_caches(encoder_cache, mm_hash)

148

149 - # Test using has_caches API

150 - producer_result = producer.has_caches(mock_request_with_3_mm)

149 + # Test using has_cache_item API

150 + producer_result = [

151 + producer.has_cache_item(mm_feature.identifier)

152 + for mm_feature in mock_request_with_3_mm.mm_features

153 + ]

154

155 # Assert

156 assert len(producer_result) == 3

@@ -159,14 +162,17 @@ class TestCacheExistence:

162 role=ECConnectorRole.SCHEDULER,

163 )

164

162 - # Test using has_caches API

163 - consumer_result = consumer.has_caches(mock_request_with_3_mm)

165 + # Test using has_cache_item API

166 + consumer_result = [

167 + consumer.has_cache_item(mm_feature.identifier)

168 + for mm_feature in mock_request_with_3_mm.mm_features

169 + ]

170

171 # Assert

172 assert len(consumer_result) == 3

173 assert all(consumer_result), f"Expected all True, got {consumer_result}"

174

169 - def test_has_caches_none_exist(

175 + def test_has_cache_item_none_exist(

176 self, mock_vllm_config_producer, mock_request_with_3_mm

177 ):

178 """Test has_caches returns False when no caches exist."""

@@ -176,13 +182,16 @@ class TestCacheExistence:

182 )

183

184 # Test without creating any files

179 - result = connector.has_caches(mock_request_with_3_mm)

185 + result = [

186 + connector.has_cache_item(mm_feature.identifier)

187 + for mm_feature in mock_request_with_3_mm.mm_features

188 + ]

189

190 # Assert

191 assert len(result) == 3

192 assert not any(result), f"Expected all False, got {result}"

193

185 - def test_has_caches_partial_exist(

194 + def test_has_cache_item_partial_exist(

195 self, mock_vllm_config_producer, mock_request_with_3_mm

196 ):

197 """Test has_caches with some caches existing (1 of 3)."""

@@ -197,7 +206,10 @@ class TestCacheExistence:

206 connector.save_caches(encoder_cache, mm_hash_second)

207

208 # Test

200 - result = connector.has_caches(mock_request_with_3_mm)

209 + result = [

210 + connector.has_cache_item(mm_feature.identifier)

211 + for mm_feature in mock_request_with_3_mm.mm_features

212 + ]

213

214 # Assert

215 assert len(result) == 3

@@ -323,8 +335,11 @@ class TestCacheSaving:

335 encoder_cache[mm_hash] = torch.randn(10, 768)

336 connector.save_caches(encoder_cache, mm_hash)

337

326 - # Verify all files exist using has_caches

327 - result = connector.has_caches(mock_request_with_3_mm)

338 + # Verify all files exist using has_cache_item

339 + result = [

340 + connector.has_cache_item(mm_feature.identifier)

341 + for mm_feature in mock_request_with_3_mm.mm_features

342 + ]

343 assert all(result), f"Not all caches were saved: {result}"

344

345 # Verify each file's content

@@ -347,10 +362,9 @@ class TestCacheSaving:

362 # Save should not raise but also not create file

363 connector.save_caches(encoder_cache, mm_hash)

364

350 - # Verify file doesn't exist using has_caches

351 - mock_request = MockRequest("req_consumer", [mm_hash], [10])

352 - result = connector.has_caches(mock_request)

353 - assert not result[0], "Consumer should not save caches"

365 + # Verify file doesn't exist using has_cache_item

366 + result = connector.has_cache_item(mm_hash)

367 + assert not result, "Consumer should not save caches"

368

369

370 class TestCacheLoading:

vllm/distributed/ec_transfer/ec_connector/base.py

+7 -7

@@ -182,19 +182,19 @@ class ECConnectorBase(ABC):

182 # ==============================

183

184 @abstractmethod

185 - def has_caches(

185 + def has_cache_item(

186 self,

187 - request: "Request",

188 - ) -> list[bool]:

187 + identifier: str,

188 + ) -> bool:

189 """

190 - Check if encoder cache exists for each mm data of requests

190 + Check if a single encoder cache exists

191

192 Args:

193 - request (Request): the request object.

193 + identifier (str): the identifier of the media.

194

195 Returns:

196 - A list bool where ith value is True if cache exist for

197 - ith mm_data of requests

196 + A bool where value is True if cache exist for

197 + the media

198 """

199 pass

200

vllm/distributed/ec_transfer/ec_connector/example_connector.py

+7 -10

@@ -117,23 +117,20 @@ class ECExampleConnector(ECConnectorBase):

117 safetensors.torch.save_file(tensors, filename)

118 logger.debug("Save cache successful for mm_hash %s", mm_hash)

119

120 - def has_caches(

120 + def has_cache_item(

121 self,

122 - request: "Request",

123 - ) -> list[bool]:

122 + identifier: str,

123 + ) -> bool:

124 """

125 - Check if cache exist externally for each mm_data of request

125 + Check if cache exist externally for the media

126

127 Args:

128 - request (Request): the request object.

128 + identifier (str): the identifier of the media.

129

130 Returns:

131 - List of bool indicate that ith mm_data exist in cache or not

131 + Bool indicate that media exists in cache or not

132 """

133 - result = []

134 - for feature in request.mm_features:

135 - result.append(self._found_match_for_mm_data(feature.identifier))

136 - return result

133 + return self._found_match_for_mm_data(identifier)

134

135 def update_state_after_alloc(

136 self,

vllm/v1/core/sched/scheduler.py

+7 -7

@@ -947,9 +947,6 @@ class Scheduler(SchedulerInterface):

947 assert len(mm_features) > 0

948 external_load_encoder_input = []

949

950 - # Check remote cache first

951 - if self.ec_connector is not None:

952 - remote_cache_has_item = self.ec_connector.has_caches(request)

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

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

952 # trackers for accounting at the encoder input level.

@@ -959,6 +956,7 @@ class Scheduler(SchedulerInterface):

956 start_pos = mm_feature.mm_position.offset

957 num_encoder_tokens = mm_feature.mm_position.length

958 num_encoder_embeds = mm_feature.mm_position.get_num_embeds

959 + item_identifier = mm_feature.identifier

960

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

962 # [num_computed_tokens, num_computed_tokens + num_new_tokens) and

@@ -993,7 +991,7 @@ class Scheduler(SchedulerInterface):

991 if not self.is_encoder_decoder:

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

993 # yet.

996 - if request.mm_features[i].identifier in mm_hashes_to_schedule:

994 + if item_identifier in mm_hashes_to_schedule:

995 # The same encoder input has already been scheduled in the

996 # current step.

997 continue

@@ -1051,15 +1049,17 @@ class Scheduler(SchedulerInterface):

1049 if curr_embeds_end - curr_embeds_start == 0:

1050 continue

1051

1054 - if self.ec_connector is not None and remote_cache_has_item[i]:

1055 - mm_hashes_to_schedule.add(request.mm_features[i].identifier)

1052 + if self.ec_connector is not None and self.ec_connector.has_cache_item(

1053 + item_identifier

1054 + ):

1055 + mm_hashes_to_schedule.add(item_identifier)

1056 external_load_encoder_input.append(i)

1057 num_embeds_to_schedule += num_encoder_embeds

1058 continue

1059

1060 num_embeds_to_schedule += num_encoder_embeds

1061 encoder_compute_budget -= num_encoder_embeds

1062 - mm_hashes_to_schedule.add(request.mm_features[i].identifier)

1062 + mm_hashes_to_schedule.add(item_identifier)

1063 encoder_inputs_to_schedule.append(i)

1064

1065 return (