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

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

[Test][KV Connector] Add request_finished fence population tests for offloading scheduler (#45679)

Signed-off-by: Alex <alex.tech.lab@outlook.com> Signed-off-by: AlexHuang <jihuihuang@future.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>

AlexHuang committed Jun 18, 2026 at 13:13 UTC 554352a311eb2b106bd1a2fd02cbff27a6c36ed9

2 files changed +302 -9

tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py

+288 -6

@@ -906,9 +906,27 @@ def test_fence_at_update_state_after_alloc(request_runner):

906 runner.manager.prepare_store.side_effect = lambda keys, req_context: (

907 generate_store_output(keys)

908 )

909 - runner.run(decoded_tokens=[EOS_TOKEN_ID], complete_transfers=False)

909 +

910 + # Capture fence snapshots to verify block 0 is registered.

911 + fence_snapshots: list[dict] = []

912 +

913 + def capture_fence():

914 + fence_snapshots.append(

915 + dict(runner.connector_scheduler._block_id_to_pending_jobs)

916 + )

917 +

918 + runner.run(

919 + decoded_tokens=[EOS_TOKEN_ID],

920 + complete_transfers=False,

921 + post_step_fn=capture_fence,

922 + )

923 assert runner.connector_scheduler._block_id_to_pending_jobs

924

925 + # Verify fence was populated with the store job's block IDs.

926 + populated_fence = next((f for f in fence_snapshots if f), None)

927 + assert populated_fence is not None, "Fence was never populated"

928 + assert len(populated_fence) > 0, "Fence is empty"

929 +

930 runner.scheduler.reset_prefix_cache()

931 runner.new_request(token_ids=[0] * 4)

932 runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1

@@ -939,9 +957,27 @@ def test_fence_at_build_store_jobs(request_runner):

957 runner.manager.prepare_store.side_effect = lambda keys, req_context: (

958 generate_store_output(keys)

959 )

942 - runner.run(decoded_tokens=[EOS_TOKEN_ID], complete_transfers=False)

960 +

961 + # Capture fence snapshots to verify block 0 is registered.

962 + fence_snapshots: list[dict] = []

963 +

964 + def capture_fence():

965 + fence_snapshots.append(

966 + dict(runner.connector_scheduler._block_id_to_pending_jobs)

967 + )

968 +

969 + runner.run(

970 + decoded_tokens=[EOS_TOKEN_ID],

971 + complete_transfers=False,

972 + post_step_fn=capture_fence,

973 + )

974 assert runner.connector_scheduler._block_id_to_pending_jobs

975

976 + # Verify fence was populated with the store job's block IDs.

977 + populated_fence = next((f for f in fence_snapshots if f), None)

978 + assert populated_fence is not None, "Fence was never populated"

979 + assert len(populated_fence) > 0, "Fence is empty"

980 +

981 runner.scheduler.reset_prefix_cache()

982 runner.new_request(token_ids=[1] * 4)

983 runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 0

@@ -1021,8 +1057,8 @@ def test_max_offload_tokens_validation(request_runner, async_scheduling: bool):

1057 token_ids=[0] * offloaded_block_size * 3,

1058 kv_transfer_params={"max_offload_tokens": max_offload_tokens},

1059 )

1024 - r.manager.prepare_store.side_effect = (

1025 - lambda keys, req_context: generate_store_output(keys)

1060 + r.manager.prepare_store.side_effect = lambda keys, req_context: (

1061 + generate_store_output(keys)

1062 )

1063

1064 # Pending offloads drain via non-blocking stepping, not a flush, so no

@@ -1120,8 +1156,8 @@ def test_offload_prompt_only(request_runner, async_scheduling: bool):

1156 extra_config_overrides={"offload_prompt_only": True},

1157 )

1158

1123 - runner.manager.prepare_store.side_effect = (

1124 - lambda keys, req_context: generate_store_output(keys)

1159 + runner.manager.prepare_store.side_effect = lambda keys, req_context: (

1160 + generate_store_output(keys)

1161 )

1162

1163 runner.new_request(token_ids=[0] * offloaded_block_size * num_prompt_blocks)

@@ -2182,3 +2218,249 @@ class TestEagle:

2218 (1, 1),

2219 ),

2220 )

2221 +

2222 +

2223 +# ---------------------------------------------------------------------------

2224 +# Tests for request_finished fence population with in-flight pending stores.

2225 +# ---------------------------------------------------------------------------

2226 +

2227 +

2228 +def test_request_finished_with_pending_stores_populates_fence(request_runner):

2229 + """When a request finishes with in-flight store jobs, the fence index

2230 + (_block_id_to_pending_jobs) is correctly populated with the store jobs'

2231 + non_sliding_window_block_ids.

2232 +

2233 + This prevents data corruption when a subsequent request reuses the same

2234 + GPU blocks before the store completes.

2235 + """

2236 + block_size = 4

2237 + block_size_factor = 1

2238 + offloaded_block_size = block_size * block_size_factor

2239 +

2240 + # Use 2 GPU blocks so the second run reuses the same blocks,

2241 + # triggering a fence-based flush of the in-flight job from run 1.

2242 + runner = request_runner(

2243 + block_size=block_size,

2244 + num_gpu_blocks=2,

2245 + async_scheduling=False,

2246 + block_size_factor=block_size_factor,

2247 + )

2248 +

2249 + # 4 prompt tokens → 1 GPU block (block 0)

2250 + runner.new_request(token_ids=[0] * offloaded_block_size)

2251 + runner.manager.prepare_store.side_effect = lambda keys, req_context: (

2252 + generate_store_output(keys)

2253 + )

2254 +

2255 + # Capture fence state at each step to verify it was populated.

2256 + fence_snapshots: list[dict] = []

2257 + job_block_ids: set[int] = set()

2258 +

2259 + def capture_fence():

2260 + fence_snapshots.append(

2261 + dict(runner.connector_scheduler._block_id_to_pending_jobs)

2262 + )

2263 + for js in runner.connector_scheduler._jobs.values():

2264 + if js.is_store:

2265 + job_block_ids.update(js.non_sliding_window_block_ids or [])

2266 +

2267 + # Run 1: create store job, finish request, populate fence.

2268 + # With non-blocking drain (#45595), the job stays in-flight.

2269 + runner.run(

2270 + decoded_tokens=[EOS_TOKEN_ID],

2271 + complete_transfers=False,

2272 + post_step_fn=capture_fence,

2273 + )

2274 +

2275 + # Verify fence was populated at some point during the run.

2276 + assert len(job_block_ids) > 0, "No store job was created"

2277 + populated_fence = next((f for f in fence_snapshots if len(f) > 0), None)

2278 + assert populated_fence is not None, "Fence was never populated"

2279 +

2280 + # Verify fence contained the job's non-SW block IDs.

2281 + for bid in job_block_ids:

2282 + assert bid in populated_fence, f"Block {bid} not in fence: {populated_fence}"

2283 +

2284 + # Run 2: block reuse triggers fence-based flush → cleanup.

2285 + runner.scheduler.reset_prefix_cache()

2286 + runner.new_request(token_ids=[0] * offloaded_block_size)

2287 + runner.manager.prepare_store.side_effect = lambda keys, req_context: (

2288 + generate_store_output(keys)

2289 + )

2290 + runner.run(

2291 + decoded_tokens=[EOS_TOKEN_ID],

2292 + expected_stored=(0,),

2293 + expected_flushed=(0,),

2294 + )

2295 +

2296 + # Verify fence is empty after full lifecycle (cleanup happened).

2297 + assert runner.connector_scheduler._block_id_to_pending_jobs == {}

2298 + # req_status should be removed.

2299 + req_id = str(runner.req_id)

2300 + assert req_id not in runner.connector_scheduler._req_status

2301 +

2302 +

2303 +def test_multiple_in_flight_stores_all_flushed_by_fence(request_runner):

2304 + """When a request finishes with multiple in-flight store jobs,

2305 + ALL jobs are flushed when a new request reuses their blocks.

2306 +

2307 + Uses three runner.run() calls:

2308 + - Run 1: decode fills a block → job_0 created

2309 + - Run 2: decode fills another block + EOS → job_1 created, request finishes

2310 + - Run 3: block reuse → both jobs flushed via fence

2311 + """

2312 + block_size = 4

2313 + block_size_factor = 1

2314 + offloaded_block_size = block_size * block_size_factor

2315 +

2316 + # 4 GPU blocks: block 0 is null, blocks 1-3 are usable.

2317 + runner = request_runner(

2318 + block_size=block_size,

2319 + num_gpu_blocks=4,

2320 + async_scheduling=False,

2321 + block_size_factor=block_size_factor,

2322 + )

2323 +

2324 + # Prompt: 4 tokens → block 1

2325 + runner.new_request(token_ids=[0] * offloaded_block_size)

2326 + runner.manager.prepare_store.side_effect = lambda keys, req_context: (

2327 + generate_store_output(keys)

2328 + )

2329 +

2330 + # Run 1: 4 decoded tokens → block 2 full → job_0 created for block 1.

2331 + runner.run(

2332 + decoded_tokens=[0] * offloaded_block_size,

2333 + complete_transfers=False,

2334 + )

2335 + assert len(runner.connector_scheduler._jobs) >= 1

2336 +

2337 + # Run 2: 4 more tokens + EOS → block 3 full → more jobs created.

2338 + # Request finishes → all jobs registered in fence.

2339 + runner.run(

2340 + decoded_tokens=[0] * offloaded_block_size + [EOS_TOKEN_ID],

2341 + complete_transfers=False,

2342 + )

2343 + num_jobs = len(runner.connector_scheduler._jobs)

2344 + assert num_jobs >= 2, f"Expected multiple in-flight jobs, got {num_jobs}"

2345 +

2346 + # Run 3: block reuse → fence flushes both jobs.

2347 + runner.scheduler.reset_prefix_cache()

2348 + runner.new_request(token_ids=[0] * offloaded_block_size * 3)

2349 + runner.manager.prepare_store.side_effect = lambda keys, req_context: (

2350 + generate_store_output(keys)

2351 + )

2352 + runner.run(

2353 + decoded_tokens=[EOS_TOKEN_ID],

2354 + expected_stored=(0, 1, 2),

2355 + expected_flushed=(0, 1, 2),

2356 + )

2357 +

2358 + # Post-condition: fence cleaned up, all jobs gone.

2359 + assert runner.connector_scheduler._block_id_to_pending_jobs == {}

2360 + assert len(runner.connector_scheduler._jobs) == 0

2361 +

2362 +

2363 +def test_request_finished_mixed_full_attn_and_sliding_window(

2364 + request_runner,

2365 +):

2366 + """With both FullAttention and SlidingWindow groups, a single store job

2367 + has both non_sliding_window_block_ids and sliding_window_block_ids.

2368 +

2369 + request_finished only registers non-SW blocks in the fence.

2370 + SW blocks were already registered at store creation time.

2371 + """

2372 + block_size = 4

2373 + sliding_window = 8 # 2 blocks

2374 +

2375 + kv_cache_groups = [

2376 + KVCacheGroupSpec(

2377 + ["layer0"],

2378 + FullAttentionSpec(

2379 + block_size=block_size,

2380 + num_kv_heads=1,

2381 + head_size=1,

2382 + dtype=torch.float32,

2383 + ),

2384 + ),

2385 + KVCacheGroupSpec(

2386 + ["layer1"],

2387 + SlidingWindowSpec(

2388 + block_size=block_size,

2389 + num_kv_heads=1,

2390 + head_size=1,

2391 + dtype=torch.float32,

2392 + sliding_window=sliding_window,

2393 + ),

2394 + ),

2395 + ]

2396 +

2397 + # Use 4 GPU blocks (2 per group) so run 2 reuses the same blocks,

2398 + # triggering a fence-based flush.

2399 + runner = request_runner(

2400 + block_size=block_size,

2401 + num_gpu_blocks=4,

2402 + async_scheduling=False,

2403 + kv_cache_groups=kv_cache_groups,

2404 + )

2405 +

2406 + # 1 block of prompt (4 tokens) — 1 block per group.

2407 + runner.new_request(token_ids=[0] * block_size)

2408 + runner.manager.prepare_store.side_effect = lambda keys, req_context: (

2409 + generate_store_output(keys)

2410 + )

2411 +

2412 + # Capture fence state and job block IDs at each step.

2413 + fence_snapshots: list[dict] = []

2414 + sw_block_ids: set[int] = set()

2415 + non_sw_block_ids: set[int] = set()

2416 +

2417 + def capture_fence():

2418 + fence_snapshots.append(

2419 + dict(runner.connector_scheduler._block_id_to_pending_jobs)

2420 + )

2421 + for js in runner.connector_scheduler._jobs.values():

2422 + if js.is_store:

2423 + sw_block_ids.update(js.sliding_window_block_ids or [])

2424 + non_sw_block_ids.update(js.non_sliding_window_block_ids or [])

2425 +

2426 + # Run 1: create store job, finish request, populate fence.

2427 + runner.run(

2428 + decoded_tokens=[EOS_TOKEN_ID],

2429 + complete_transfers=False,

2430 + post_step_fn=capture_fence,

2431 + )

2432 +

2433 + # Verify job had both SW and non-SW blocks.

2434 + assert len(sw_block_ids) > 0, "No SW blocks in store job"

2435 + assert len(non_sw_block_ids) > 0, "No non-SW blocks in store job"

2436 +

2437 + # Find the fence snapshot where both SW and non-SW blocks were present.

2438 + # SW blocks should appear at creation time, non-SW at request_finished.

2439 + populated_fence = None

2440 + for fence in fence_snapshots:

2441 + has_sw = all(bid in fence for bid in sw_block_ids)

2442 + has_non_sw = all(bid in fence for bid in non_sw_block_ids)

2443 + if has_sw and has_non_sw:

2444 + populated_fence = fence

2445 + break

2446 +

2447 + assert populated_fence is not None, (

2448 + f"Fence never contained both SW {sw_block_ids} and "

2449 + f"non-SW {non_sw_block_ids} blocks. Snapshots: {fence_snapshots}"

2450 + )

2451 +

2452 + # Run 2: block reuse triggers fence-based flush of the old job.

2453 + runner.scheduler.reset_prefix_cache()

2454 + runner.new_request(token_ids=[0] * block_size)

2455 + runner.manager.prepare_store.side_effect = lambda keys, req_context: (

2456 + generate_store_output(keys)

2457 + )

2458 + runner.run(

2459 + decoded_tokens=[EOS_TOKEN_ID],

2460 + expected_stored=((0, 0), (1, 0)),

2461 + expected_flushed=((1, 0),),

2462 + )

2463 +

2464 + # Verify fence is empty after full lifecycle (cleanup happened).

2465 + assert runner.connector_scheduler._block_id_to_pending_jobs == {}

2466 + assert len(runner.connector_scheduler._jobs) == 0

tests/v1/kv_connector/unit/offloading_connector/utils.py

+14 -3

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from collections.abc import Iterable, Iterator

3 +from collections.abc import Callable, Iterable, Iterator

4 from dataclasses import dataclass

5 from typing import Any

6 from unittest.mock import MagicMock

@@ -430,7 +430,12 @@ class RequestRunner:

430 for block_idx, block in enumerate(blocks):

431 self.gpu_blocks[block.block_id] = GPUBlock(group_idx, block_idx)

432

433 - def _run(self, decoded_tokens: list[int], complete_transfers: bool):

433 + def _run(

434 + self,

435 + decoded_tokens: list[int],

436 + complete_transfers: bool,

437 + post_step_fn: Callable[[], None] | None = None,

438 + ):

439 """

440 Runs multiple engine (scheduler + worker) steps.

441 Assumes a single request is running.

@@ -438,6 +443,8 @@ class RequestRunner:

443 Args:

444 decoded_tokens: the tokens to yield at each step.

445 complete_transfers: complete transfers immediately

446 + post_step_fn: optional callback invoked after each step's

447 + update_from_output(), before the next schedule().

448 """

449

450 tokens_iter = iter(decoded_tokens)

@@ -500,6 +507,9 @@ class RequestRunner:

507 else:

508 self.scheduler.update_from_output(scheduler_output, model_runner_output)

509

510 + if post_step_fn is not None:

511 + post_step_fn()

512 +

513 if (

514 prev_token_id == EOS_TOKEN_ID

515 and prev_token_id != token_id

@@ -545,6 +555,7 @@ class RequestRunner:

555 expected_stored: tuple[int | tuple[int, int], ...] = (),

556 expected_loaded: tuple[int | tuple[int, int], ...] = (),

557 expected_flushed: tuple[int | tuple[int, int], ...] = (),

558 + post_step_fn: Callable[[], None] | None = None,

559 ):

560 """

561 Runs multiple engine (scheduler + worker) steps.

@@ -570,7 +581,7 @@ class RequestRunner:

581 expected_flushed_gpu_blocks = self._to_gpu_blocks(expected_flushed)

582

583 self.manager.reset_mock()

573 - self._run(decoded_tokens, complete_transfers)

584 + self._run(decoded_tokens, complete_transfers, post_step_fn=post_step_fn)

585

586 loaded_gpu_blocks: set[GPUBlock] = set()

587 for transfer in self.completed_loads: