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

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, 06:00 UTC

[EC Connector] Non blocking EC Connector lookup (#41627)

Signed-off-by: omerpaz95 <omerpaz95@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

omerpaz95 committed Jun 2, 2026 at 11:48 UTC d247a9dc13a58fef5ce64dc3a8f4be051577bda5

3 files changed +154

tests/v1/core/test_scheduler.py

+125

@@ -4349,3 +4349,128 @@ def test_eagle3_mm_encoder_cache_with_shift():

4349 f"shifted_end={scheduled_end_with_shift}) overlapping MM at "

4350 f"{start_pos}. The fix must schedule encoder inputs."

4351 )

4352 +

4353 +

4354 +@pytest.mark.parametrize("use_kv_connector", [False, True])

4355 +def test_ec_connector_ensure_cache_available_defers_request(use_kv_connector):

4356 + """Test that ensure_cache_available() returning False defers the request.

4357 +

4358 + When the EC connector signals a prefetch is in progress (returns False),

4359 + the scheduler should:

4360 + 1. Not schedule the request (no KV cache or encoder cache allocated)

4361 + 2. Still schedule other requests behind the deferred one

4362 + 3. Schedule the deferred request on the next step when ensure_cache_available

4363 + returns True and has_cache_item returns True

4364 + """

4365 + scheduler = create_scheduler(

4366 + model="llava-hf/llava-1.5-7b-hf",

4367 + enable_prefix_caching=True,

4368 + use_kv_connector=use_kv_connector,

4369 + use_ec_connector=True,

4370 + ec_role="ec_consumer",

4371 + )

4372 +

4373 + NUM_TOKENS = 200

4374 + NUM_ENCODER_TOKENS = 100

4375 +

4376 + request_deferred = create_requests(

4377 + num_requests=1,

4378 + num_tokens=NUM_TOKENS,

4379 + mm_positions=[[PlaceholderRange(offset=0, length=NUM_ENCODER_TOKENS)]],

4380 + req_ids=["deferred"],

4381 + )[0]

4382 +

4383 + request_behind = create_requests(

4384 + num_requests=1,

4385 + num_tokens=20,

4386 + req_ids=["behind"],

4387 + )[0]

4388 +

4389 + # --- Step 1: ensure_cache_available returns False → request deferred ---

4390 + scheduler.ec_connector.ensure_cache_available = Mock(return_value=False)

4391 +

4392 + scheduler.add_request(request_deferred)

4393 + scheduler.add_request(request_behind)

4394 + output = scheduler.schedule()

4395 +

4396 + # ensure_cache_available must have been called with (request, num_computed_tokens=0)

4397 + # for a brand-new request that has no cached tokens yet.

4398 + scheduler.ec_connector.ensure_cache_available.assert_called_once_with(

4399 + request_deferred, 0

4400 + )

4401 + # Deferred request must NOT be scheduled

4402 + assert request_deferred.request_id not in output.num_scheduled_tokens

4403 + _assert_right_encoder_cache_allocated(scheduler, expected_total_allocated=0)

4404 + # No KV blocks allocated for the deferred request

4405 + for mgr in scheduler.kv_cache_manager.coordinator.single_type_managers:

4406 + assert request_deferred.request_id not in mgr.req_to_blocks

4407 +

4408 + # The text-only request behind the deferred one MUST still be scheduled

4409 + assert request_behind.request_id in output.num_scheduled_tokens

4410 + assert output.num_scheduled_tokens[request_behind.request_id] == 20

4411 +

4412 + # --- Step 2: prefetch done, cache exists → request scheduled ---

4413 + # has_cache_item is called inside _try_schedule_encoder_inputs (not during

4414 + # deferral), so it is only relevant here in step 2.

4415 + scheduler.ec_connector.ensure_cache_available = Mock(return_value=True)

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

4417 +

4418 + output = scheduler.schedule()

4419 +

4420 + # Now the deferred request should be scheduled

4421 + assert request_deferred.request_id in output.num_scheduled_tokens

4422 + assert output.num_scheduled_tokens[request_deferred.request_id] == NUM_TOKENS

4423 + _assert_right_encoder_cache_allocated(scheduler, requests=[request_deferred])

4424 + # EC connector metadata should carry the deferred request's MM data

4425 + _assert_right_ec_connector_metadata(

4426 + output, mm_features_list=request_deferred.mm_features

4427 + )

4428 + # No local encoder compute — all loaded externally

4429 + _assert_right_encoder_inputs(output, expected_total_reqs=0)

4430 +

4431 +

4432 +def test_ec_connector_pending_prefetch_only_checks_future_mm_features():

4433 + """Test that future mm feature filtering only yields features beyond

4434 + the computed token frontier.

4435 +

4436 + Features already within num_computed_tokens (past/boundary) must be

4437 + filtered out; only features that extend beyond the frontier (future) should

4438 + be yielded so that connector implementations know which items to prefetch.

4439 +

4440 + Filter cases:

4441 + "past": end = 0 + 16 = 16 < 32 → filtered OUT

4442 + "boundary": end = 16 + 16 = 32 == 32 → filtered OUT (condition is >, not >=)

4443 + "future": end = 48 + 32 = 80 > 32 → yielded

4444 + """

4445 + BLOCK_SIZE = 16

4446 + NUM_COMPUTED_TOKENS = BLOCK_SIZE * 2 # 32

4447 + NUM_TOKENS = BLOCK_SIZE * 8 # 128

4448 +

4449 + HASH_PAST = "hash_past"

4450 + HASH_BOUNDARY = "hash_boundary"

4451 + HASH_FUTURE = "hash_future"

4452 +

4453 + request = create_requests(

4454 + num_requests=1,

4455 + num_tokens=NUM_TOKENS,

4456 + mm_hashes_list=[[HASH_PAST, HASH_BOUNDARY, HASH_FUTURE]],

4457 + mm_positions=[

4458 + [

4459 + PlaceholderRange(offset=0, length=BLOCK_SIZE), # end=16 (past)

4460 + PlaceholderRange(offset=16, length=BLOCK_SIZE), # end=32 (boundary)

4461 + PlaceholderRange(offset=48, length=BLOCK_SIZE * 2), # end=80 (future)

4462 + ]

4463 + ],

4464 + block_size=BLOCK_SIZE,

4465 + )[0]

4466 +

4467 + future_hashes = [

4468 + f.identifier

4469 + for f in request.mm_features

4470 + if f.mm_position.offset + f.mm_position.length > NUM_COMPUTED_TOKENS

4471 + ]

4472 +

4473 + assert future_hashes == [HASH_FUTURE], (

4474 + f"Expected only {HASH_FUTURE!r} from future mm feature filtering, "

4475 + f"got {future_hashes!r}. Past/boundary features must be filtered out."

4476 + )

vllm/distributed/ec_transfer/ec_connector/base.py

+17

@@ -211,6 +211,23 @@ class ECConnectorBase(ABC):

211 """

212 pass

213

214 + def ensure_cache_available(

215 + self, request: "Request", num_computed_tokens: int

216 + ) -> bool:

217 + """

218 + Ensure encoder cache items are available for the given request.

219 + May initiate asynchronous transfers for items not yet local.

220 +

221 + Args:

222 + request: the request whose multimodal features to check.

223 + num_computed_tokens: tokens already covered by cached KV blocks.

224 +

225 + Returns:

226 + True if all items are ready or no transfer is needed.

227 + False if any items are still in transit (request should be deferred).

228 + """

229 + return True

230 +

231 @abstractmethod

232 def update_state_after_alloc(self, request: "Request", index: int):

233 """

vllm/v1/core/sched/scheduler.py

+12

@@ -629,6 +629,18 @@ class Scheduler(SchedulerInterface):

629 )

630 assert num_computed_tokens <= request.num_tokens

631

632 + # Skip request with pending mm encoding prefetches

633 + if (

634 + self.ec_connector is not None

635 + and request.mm_features

636 + and not self.ec_connector.ensure_cache_available(

637 + request, num_computed_tokens

638 + )

639 + ):

640 + request_queue.pop_request()

641 + step_skipped_waiting.prepend_request(request)

642 + continue

643 +

644 # Track first scheduled prefill, not post-preemption repeat prefills

645 if request.prefill_stats is not None:

646 assert num_computed_tokens <= request.num_prompt_tokens