@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 07b7553465

[Feature][SimpleCPU] Load fine-grained hybrid prefix hits (#54736)

Signed-off-by: Yichao Zhu <Yichao.Zhu@amd.com>

YukioZzz committed Sep 11, 2026 at 09:36 UTC 07b75534651a0dd584c6f3be75d8c62fb8426397

6 files changed +982 -110

tests/v1/simple_kv_offload/test_kv_events.py

+9 -40

@@ -714,17 +714,8 @@ def test_disk_mode_block_removed_relabeled_storage() -> None:

714 assert ev.medium == MEDIUM_STORAGE

715

716

717 -def test_mamba_dcp_metadata_guard() -> None:

718 - """Mamba+DCP: token_ids guard emits [] when capture and emission sizes differ.

719 -

720 - The metadata capture at _prepare_eager_store_specs uses

721 - g_block_size = spec.block_size * cp_world_size = 32 * 2 = 64 to slice

722 - tokens. The event emission uses block_size = spec.block_size * 1 = 32

723 - for Mamba (SSM state is not CP-sharded). The guard in

724 - _process_store_completion detects len(token_ids) != block_size and emits

725 - token_ids=[] instead of misleading data. When #49962 fixes the capture

726 - path, the guard stops triggering and correct token_ids flow through.

727 - """

717 +def test_mamba_align_skips_positional_event_metadata() -> None:

718 + """Mamba align emits events only from explicit boundary handoffs."""

719 vbs = BLOCK_SIZE * 2

720 kv_cfg = _make_mixed_kv_cache_config(

721 num_blocks=16, with_mamba=True, mamba_block_size=32

@@ -739,8 +730,6 @@ def test_mamba_dcp_metadata_guard() -> None:

730 sched = fix.scheduler

731 gpu_pool = fix.gpu_block_pool

732

742 - # Mamba g_block_size = 32 * 2 = 64, so 2 Mamba blocks need 128 computed

743 - # tokens to be "ready"; FA g_block_size = 16 * 2 = 32 needs only 64.

733 req = _make_cp_request(num_blocks=4, virtual_block_size=vbs)

734 fa_blocks = _allocate_cp_gpu_blocks(gpu_pool, req, 2, vbs, group_id=0)

735 mamba_blocks = _allocate_cp_gpu_blocks(gpu_pool, req, 2, vbs, group_id=1)

@@ -759,32 +748,12 @@ def test_mamba_dcp_metadata_guard() -> None:

748

749 events = list(sched.take_events())

750 stored = [e for e in events if isinstance(e, BlockStored)]

762 - assert len(stored) == 4, (

763 - f"expected 4 BlockStored (2 FA + 2 Mamba), got {len(stored)}"

764 - )

765 -

766 - by_group: dict[int, list[BlockStored]] = {}

767 - for ev in stored:

768 - by_group.setdefault(ev.group_idx, []).append(ev)

769 -

770 - # FA group: capture size (64) == emission size (64) → token_ids present.

771 - fa_events = by_group[0]

772 - assert len(fa_events) == 2

773 - for ev in fa_events:

751 + assert len(stored) == 2

752 + assert {ev.group_idx for ev in stored} == {0}

753 + for idx, ev in enumerate(stored):

754 assert ev.block_size == vbs

775 - assert len(ev.token_ids) == vbs, (

776 - f"FA token_ids should be {vbs}, got {len(ev.token_ids)}"

777 - )

778 -

779 - # Mamba group: capture size (64) != emission size (32) → guard emits [].

780 - mamba_events = by_group[1]

781 - assert len(mamba_events) == 2

782 - for ev in mamba_events:

783 - assert ev.block_size == 32

784 - assert ev.token_ids == [], (

785 - f"Mamba token_ids should be [] (guard), got {ev.token_ids}"

786 - )

787 - assert ev.parent_block_hash is None, (

788 - "Mamba parent_block_hash should be None (guard)"

755 + assert ev.token_ids == req.prompt_token_ids[idx * vbs : (idx + 1) * vbs]

756 + expected_hash = maybe_convert_block_hash(

757 + get_block_hash(make_block_hash_with_group_id(req.block_hashes[idx], 0))

758 )

790 - assert ev.extra_keys is None, "Mamba extra_keys should be None (guard)"

759 + assert ev.block_hashes == [expected_hash]

tests/v1/simple_kv_offload/test_scheduler.py

+672

@@ -4,6 +4,7 @@

5 from __future__ import annotations

7 +import logging

8 from dataclasses import dataclass

10 import pytest

@@ -18,6 +19,7 @@ from vllm.config import (

19 SchedulerConfig,

20 VllmConfig,

21 )

22 +from vllm.config.kv_events import KVEventsConfig

23 from vllm.distributed.kv_transfer.kv_connector.v1.simple_cpu_offload_connector import (

24 SimpleCPUOffloadConnector,

25 )

@@ -31,6 +33,7 @@ from vllm.v1.core.kv_cache_utils import (

33 )

34 from vllm.v1.core.sched.output import (

35 CachedRequestData,

36 + KVConnectorBlockState,

37 NewRequestData,

38 SchedulerOutput,

39 )

@@ -42,6 +45,7 @@ from vllm.v1.kv_cache_interface import (

45 KVCacheConfig,

46 KVCacheGroupSpec,

47 KVCacheTensor,

48 + MambaSpec,

49 SlidingWindowSpec,

50 )

51 from vllm.v1.outputs import KVConnectorOutput

@@ -2128,3 +2132,671 @@ def test_cp_lazy_target_blocks_scaling(cp_world_size: int) -> None:

2132 f"cp_world_size={cp_world_size}: target_cp={target_cp} should be "

2133 f"less than target_base={target_base}"

2134 )

2135 +

2136 +

2137 +def _make_hybrid_attention_mamba_scheduler(

2138 + *,

2139 + num_cpu_blocks: int = 8,

2140 + num_gpu_blocks: int = 16,

2141 + attention_block_size: int = BLOCK_SIZE,

2142 + block_size: int = 4 * BLOCK_SIZE,

2143 + scheduler_block_size: int | None = None,

2144 + hash_block_size: int | None = None,

2145 + dcp_world_size: int = 4,

2146 + lazy: bool = False,

2147 + mamba_cache_mode: str = "align",

2148 + enable_kv_cache_events: bool = False,

2149 +) -> SchedulerFixture:

2150 + """Build a scheduler for one attention group plus one Mamba group."""

2151 + scheduler_block_size = scheduler_block_size or block_size

2152 + hash_block_size = hash_block_size or block_size

2153 + attention_spec = FullAttentionSpec(

2154 + block_size=attention_block_size,

2155 + num_kv_heads=NUM_KV_HEADS,

2156 + head_size=HEAD_SIZE,

2157 + dtype=DTYPE,

2158 + )

2159 + mamba_spec = MambaSpec(

2160 + block_size=block_size,

2161 + shapes=((1, 1),),

2162 + dtypes=(torch.float32,),

2163 + mamba_cache_mode=mamba_cache_mode,

2164 + )

2165 + groups = [

2166 + KVCacheGroupSpec(["attention"], attention_spec),

2167 + KVCacheGroupSpec(["mamba"], mamba_spec),

2168 + ]

2169 + tensors = [

2170 + KVCacheTensor(

2171 + size=spec.page_size_bytes * num_gpu_blocks,

2172 + layers=group.layer_names,

2173 + layer_stride=spec.page_size_bytes * num_gpu_blocks,

2174 + block_stride=spec.page_size_bytes,

2175 + )

2176 + for group, spec in zip(groups, (attention_spec, mamba_spec))

2177 + ]

2178 + kv_cache_config = KVCacheConfig(

2179 + num_blocks=num_gpu_blocks,

2180 + kv_cache_tensors=tensors,

2181 + kv_cache_groups=groups,

2182 + )

2183 + vllm_config = _make_cp_vllm_config(dcp_world_size=dcp_world_size)

2184 + vllm_config.cache_config.prefix_cache_retention_interval = 0

2185 + vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode

2186 + if enable_kv_cache_events:

2187 + vllm_config.kv_events_config = KVEventsConfig(enable_kv_cache_events=True)

2188 + # _derive_cpu_config() scales the requested capacity against

2189 + # kv_cache_tensors[0].size, so express it as a multiple of that tensor's

2190 + # per-block size to get exactly num_cpu_blocks.

2191 + cpu_capacity_bytes = tensors[0].size // num_gpu_blocks * num_cpu_blocks

2192 + sched = SimpleCPUOffloadScheduler(

2193 + vllm_config=vllm_config,

2194 + kv_cache_config=kv_cache_config,

2195 + cpu_capacity_bytes=cpu_capacity_bytes,

2196 + scheduler_block_size=scheduler_block_size,

2197 + hash_block_size=hash_block_size,

2198 + lazy_offload=lazy,

2199 + )

2200 + gpu_block_pool = BlockPool(

2201 + num_gpu_blocks=num_gpu_blocks,

2202 + enable_caching=True,

2203 + hash_block_size=hash_block_size,

2204 + )

2205 + sched.bind_gpu_block_pool(gpu_block_pool)

2206 + return SchedulerFixture(

2207 + scheduler=sched,

2208 + gpu_block_pool=gpu_block_pool,

2209 + vllm_config=vllm_config,

2210 + kv_cache_config=kv_cache_config,

2211 + )

2212 +

2213 +

2214 +def test_hybrid_store_uses_resolved_group_block_sizes() -> None:

2215 + """Replicated groups must not be scaled by the DCP world size.

2216 +

2217 + ``dcp_world_size_for_kv_cache_spec`` gives full attention the process DCP

2218 + size and every other spec 1. Reconstructing a group's geometry as

2219 + ``spec.block_size * cp_world_size`` over-scales the replicated groups, so

2220 + the eager store scan believes far fewer of their blocks are ready and

2221 + silently offloads only a fraction of them.

2222 + """

2223 + attention_block_size = BLOCK_SIZE

2224 + mamba_block_size = 4 * BLOCK_SIZE

2225 + dcp_world_size = 2

2226 + fix = _make_hybrid_attention_mamba_scheduler(

2227 + num_cpu_blocks=32,

2228 + num_gpu_blocks=32,

2229 + attention_block_size=attention_block_size,

2230 + block_size=mamba_block_size,

2231 + # Every prefix-cacheable group's resolved block size must be a multiple

2232 + # of the hash block, so it cannot exceed the smallest of them.

2233 + hash_block_size=BLOCK_SIZE,

2234 + dcp_world_size=dcp_world_size,

2235 + mamba_cache_mode="all",

2236 + )

2237 + sched = fix.scheduler

2238 + gpu_pool = fix.gpu_block_pool

2239 + attention_size, mamba_size = sched.group_block_sizes

2240 + assert attention_size == attention_block_size * dcp_world_size

2241 + assert mamba_size == mamba_block_size

2242 +

2243 + confirmed = 2 * sched.block_size

2244 + req = _make_cp_request(num_blocks=8, virtual_block_size=BLOCK_SIZE)

2245 + attention_blocks = _allocate_cp_gpu_blocks(

2246 + gpu_pool, req, confirmed // attention_size, attention_size, group_id=0

2247 + )

2248 + mamba_blocks = _allocate_cp_gpu_blocks(

2249 + gpu_pool, req, confirmed // mamba_size, mamba_size, group_id=1

2250 + )

2251 + kv_blocks = KVCacheBlocks(blocks=(attention_blocks, mamba_blocks))

2252 + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)

2253 + sched.build_connector_meta(

2254 + make_scheduler_output(

2255 + {req.request_id: confirmed},

2256 + new_reqs={req.request_id: kv_blocks.get_block_ids()},

2257 + )

2258 + )

2259 + req.num_computed_tokens = confirmed

2260 + meta = sched.build_connector_meta(

2261 + make_scheduler_output(

2262 + {req.request_id: 1},

2263 + cached_req_new_blocks={req.request_id: None},

2264 + )

2265 + )

2266 +

2267 + # Both groups are positionally stable here, so every confirmed block of

2268 + # both must be offloaded. Over-scaling the replicated group halves its

2269 + # count.

2270 + assert sched._reqs_to_store[req.request_id].num_stored_blocks == [

2271 + confirmed // attention_size,

2272 + confirmed // mamba_size,

2273 + ]

2274 + assert set(meta.store_gpu_blocks) == {

2275 + *(b.block_id for b in attention_blocks),

2276 + *(b.block_id for b in mamba_blocks),

2277 + }

2278 +

2279 +

2280 +def test_dcp_mixed_cache_loads_fine_grained_external_hit() -> None:

2281 + """Map a fine-grained external suffix after a nonzero local prefix."""

2282 + hash_block_size = BLOCK_SIZE

2283 + mamba_block_size = 4 * BLOCK_SIZE

2284 + local_tokens = 4 * BLOCK_SIZE

2285 + # Span two Mamba blocks and end partially in the second one. This exercises

2286 + # both cdiv(external_tokens, mamba_block_size) > 1 and the partial tail.

2287 + external_tokens = 6 * BLOCK_SIZE

2288 + total_cached_tokens = local_tokens + external_tokens

2289 + fix = _make_hybrid_attention_mamba_scheduler(

2290 + num_cpu_blocks=32,

2291 + num_gpu_blocks=32,

2292 + block_size=mamba_block_size,

2293 + hash_block_size=hash_block_size,

2294 + dcp_world_size=2,

2295 + )

2296 + sched = fix.scheduler

2297 + gpu_pool = fix.gpu_block_pool

2298 +

2299 + assert external_tokens % mamba_block_size != 0

2300 + assert sched.cpu_coordinator.enable_partial_hash_hits

2301 +

2302 + producer = _make_cp_request(num_blocks=10, virtual_block_size=hash_block_size)

2303 + attention_blocks = _allocate_cp_gpu_blocks(

2304 + gpu_pool,

2305 + producer,

2306 + num_blocks=5,

2307 + virtual_block_size=2 * BLOCK_SIZE,

2308 + group_id=0,

2309 + )

2310 + mamba_blocks = gpu_pool.get_new_blocks(3)

2311 + gpu_pool.cache_full_blocks(

2312 + request=producer,

2313 + blocks=mamba_blocks,

2314 + num_cached_blocks=0,

2315 + num_full_blocks=2,

2316 + block_size=mamba_block_size,

2317 + kv_cache_group_id=1,

2318 + )

2319 + gpu_pool.cache_partial_block(

2320 + request=producer,

2321 + block=mamba_blocks[2],

2322 + num_tokens=total_cached_tokens,

2323 + kv_cache_group_id=1,

2324 + block_size=mamba_block_size,

2325 + )

2326 + producer_blocks = KVCacheBlocks(blocks=(attention_blocks, mamba_blocks))

2327 + sched.update_state_after_alloc(producer, producer_blocks, num_external_tokens=0)

2328 +

2329 + # Positional scanning stores append-only FA blocks. Mamba align blocks use

2330 + # exact handoffs for every retained boundary, including the partial tail.

2331 + store_output = make_scheduler_output(

2332 + {producer.request_id: total_cached_tokens},

2333 + new_reqs={producer.request_id: producer_blocks.get_block_ids()},

2334 + )

2335 + store_output.kv_connector_block_state = KVConnectorBlockState(

2336 + req_ids=set(),

2337 + resolve_block_ids=lambda _: (),

2338 + boundary_state_offloads={

2339 + producer.request_id: [

2340 + (1, mamba_blocks[0].block_id, mamba_block_size),

2341 + (1, mamba_blocks[1].block_id, 2 * mamba_block_size),

2342 + (1, mamba_blocks[2].block_id, total_cached_tokens),

2343 + ]

2344 + },

2345 + )

2346 + store_meta = sched.build_connector_meta(store_output)

2347 + # The exact handoffs land in the step they arrive; the append-only FA scan

2348 + # picks its blocks up once the scheduler has committed the tokens.

2349 + assert set(store_meta.store_gpu_blocks) == {

2350 + block.block_id for block in mamba_blocks

2351 + }

2352 + simulate_store_completion(sched, store_meta.store_event)

2353 + producer.num_computed_tokens = total_cached_tokens

2354 + fa_meta = sched.build_connector_meta(

2355 + make_scheduler_output(

2356 + {producer.request_id: 1},

2357 + cached_req_new_blocks={producer.request_id: None},

2358 + )

2359 + )

2360 + assert set(fa_meta.store_gpu_blocks) == {

2361 + block.block_id for block in attention_blocks

2362 + }

2363 + simulate_store_completion(sched, fa_meta.store_event)

2364 +

2365 + consumer = Request(

2366 + request_id="req-dcp-mixed-fine-load",

2367 + prompt_token_ids=producer.prompt_token_ids,

2368 + sampling_params=producer.sampling_params,

2369 + pooling_params=None,

2370 + mm_features=None,

2371 + block_hasher=producer._block_hasher,

2372 + )

2373 + hit_tokens, is_async = sched.get_num_new_matched_tokens(

2374 + consumer, num_computed_tokens=local_tokens

2375 + )

2376 + assert hit_tokens == external_tokens

2377 + assert is_async is True

2378 + pending_blocks, pending_tokens, _ = sched._pending_cpu_hits[consumer.request_id]

2379 + assert pending_tokens == external_tokens

2380 + assert len(pending_blocks[1]) == 2

2381 + assert pending_blocks[1][0].is_null

2382 + assert not pending_blocks[1][1].is_null

2383 +

2384 + local_attention = attention_blocks[:2]

2385 + local_mamba = mamba_blocks[:1]

2386 + load_blocks = KVCacheBlocks(

2387 + blocks=(

2388 + local_attention + gpu_pool.get_new_blocks(3),

2389 + local_mamba + gpu_pool.get_new_blocks(2),

2390 + )

2391 + )

2392 + sched.update_state_after_alloc(

2393 + consumer,

2394 + load_blocks,

2395 + num_external_tokens=hit_tokens,

2396 + )

2397 + load_meta = sched.build_connector_meta(

2398 + make_scheduler_output(

2399 + {consumer.request_id: 1},

2400 + new_reqs={consumer.request_id: load_blocks.get_block_ids()},

2401 + )

2402 + )

2403 + assert load_meta.load_gpu_blocks == [

2404 + *(block.block_id for block in load_blocks.blocks[0][2:5]),

2405 + load_blocks.blocks[1][2].block_id,

2406 + ]

2407 + assert len(load_meta.load_cpu_blocks) == 4

2408 +

2409 +

2410 +def test_finished_eager_store_caches_fa_partial_tail_for_hybrid_hit() -> None:

2411 + """A finished non-LCM prompt needs the FA partial tail in CPU too."""

2412 + hash_block_size = BLOCK_SIZE

2413 + mamba_block_size = BLOCK_SIZE

2414 + attention_block_size = 4 * BLOCK_SIZE

2415 + dcp_world_size = 2

2416 + scheduler_block_size = attention_block_size * dcp_world_size

2417 + total_cached_tokens = 5 * hash_block_size

2418 + fix = _make_hybrid_attention_mamba_scheduler(

2419 + num_cpu_blocks=32,

2420 + num_gpu_blocks=32,

2421 + attention_block_size=attention_block_size,

2422 + block_size=mamba_block_size,

2423 + scheduler_block_size=scheduler_block_size,

2424 + hash_block_size=hash_block_size,

2425 + dcp_world_size=dcp_world_size,

2426 + )

2427 + sched = fix.scheduler

2428 + gpu_pool = fix.gpu_block_pool

2429 + assert sched.fa_block_size == scheduler_block_size

2430 + assert total_cached_tokens % sched.fa_block_size != 0

2431 +

2432 + producer = _make_cp_request(num_blocks=5, virtual_block_size=hash_block_size)

2433 + attention_blocks = gpu_pool.get_new_blocks(1)

2434 + gpu_pool.cache_partial_block(

2435 + request=producer,

2436 + block=attention_blocks[0],

2437 + num_tokens=total_cached_tokens,

2438 + kv_cache_group_id=0,

2439 + block_size=scheduler_block_size,

2440 + )

2441 + mamba_blocks = gpu_pool.get_new_blocks(5)

2442 + gpu_pool.cache_full_blocks(

2443 + request=producer,

2444 + blocks=mamba_blocks,

2445 + num_cached_blocks=0,

2446 + num_full_blocks=5,

2447 + kv_cache_group_id=1,

2448 + block_size=mamba_block_size,

2449 + )

2450 + producer_blocks = KVCacheBlocks(blocks=(attention_blocks, mamba_blocks))

2451 + sched.update_state_after_alloc(producer, producer_blocks, num_external_tokens=0)

2452 +

2453 + store_output = make_scheduler_output(

2454 + {producer.request_id: total_cached_tokens},

2455 + new_reqs={producer.request_id: producer_blocks.get_block_ids()},

2456 + )

2457 + store_output.kv_connector_block_state = KVConnectorBlockState(

2458 + req_ids=set(),

2459 + resolve_block_ids=lambda _: (),

2460 + boundary_state_offloads={

2461 + producer.request_id: [

2462 + (1, mamba_blocks[4].block_id, total_cached_tokens),

2463 + ]

2464 + },

2465 + )

2466 + store_meta = sched.build_connector_meta(store_output)

2467 + simulate_store_completion(sched, store_meta.store_event)

2468 +

2469 + producer.num_computed_tokens = total_cached_tokens

2470 + sched.request_finished_all_groups(producer, producer_blocks.get_block_ids())

2471 + finish_meta = sched.build_connector_meta(make_scheduler_output({}))

2472 + assert attention_blocks[0].block_id in finish_meta.store_gpu_blocks

2473 + simulate_store_completion(sched, finish_meta.store_event)

2474 +

2475 + consumer = Request(

2476 + request_id="req-finished-partial-tail-load",

2477 + prompt_token_ids=producer.prompt_token_ids,

2478 + sampling_params=producer.sampling_params,

2479 + pooling_params=None,

2480 + mm_features=None,

2481 + block_hasher=producer._block_hasher,

2482 + )

2483 + hit_tokens, is_async = sched.get_num_new_matched_tokens(

2484 + consumer, num_computed_tokens=0

2485 + )

2486 + assert hit_tokens == total_cached_tokens

2487 + assert is_async is True

2488 +

2489 +

2490 +def test_finished_eager_store_does_not_duplicate_completed_partial_tail() -> None:

2491 + """Decode can fill the boundary block; the scan then already covers it."""

2492 + hash_block_size = BLOCK_SIZE

2493 + mamba_block_size = BLOCK_SIZE

2494 + attention_block_size = 4 * BLOCK_SIZE

2495 + dcp_world_size = 2

2496 + scheduler_block_size = attention_block_size * dcp_world_size

2497 + prompt_tokens = 5 * hash_block_size

2498 + fix = _make_hybrid_attention_mamba_scheduler(

2499 + num_cpu_blocks=32,

2500 + num_gpu_blocks=32,

2501 + attention_block_size=attention_block_size,

2502 + block_size=mamba_block_size,

2503 + scheduler_block_size=scheduler_block_size,

2504 + hash_block_size=hash_block_size,

2505 + dcp_world_size=dcp_world_size,

2506 + )

2507 + sched = fix.scheduler

2508 + gpu_pool = fix.gpu_block_pool

2509 + assert prompt_tokens % sched.fa_block_size != 0

2510 +

2511 + producer = _make_cp_request(num_blocks=5, virtual_block_size=hash_block_size)

2512 + attention_blocks = gpu_pool.get_new_blocks(1)

2513 + gpu_pool.cache_partial_block(

2514 + request=producer,

2515 + block=attention_blocks[0],

2516 + num_tokens=prompt_tokens,

2517 + kv_cache_group_id=0,

2518 + block_size=scheduler_block_size,

2519 + )

2520 + mamba_blocks = gpu_pool.get_new_blocks(5)

2521 + gpu_pool.cache_full_blocks(

2522 + request=producer,

2523 + blocks=mamba_blocks,

2524 + num_cached_blocks=0,

2525 + num_full_blocks=5,

2526 + kv_cache_group_id=1,

2527 + block_size=mamba_block_size,

2528 + )

2529 + producer_blocks = KVCacheBlocks(blocks=(attention_blocks, mamba_blocks))

2530 + sched.update_state_after_alloc(producer, producer_blocks, num_external_tokens=0)

2531 +

2532 + # Decode carries the request past the physical attention block boundary, so

2533 + # the positional scan now considers the same block the partial tail names.

2534 + producer.num_computed_tokens = sched.fa_block_size

2535 + sched.request_finished_all_groups(producer, producer_blocks.get_block_ids())

2536 + finish_meta = sched.build_connector_meta(make_scheduler_output({}))

2537 +

2538 + stored = finish_meta.store_gpu_blocks

2539 + assert attention_blocks[0].block_id in stored

2540 + assert len(stored) == len(set(stored))

2541 + assert len(finish_meta.store_cpu_blocks) == len(stored)

2542 +

2543 +

2544 +def test_external_lookup_rejects_misaligned_local_prefix(caplog_vllm) -> None:

2545 + scheduler_block_size = 4 * BLOCK_SIZE

2546 + fix = _make_hybrid_attention_mamba_scheduler(block_size=scheduler_block_size)

2547 +

2548 + request = _make_cp_request(num_blocks=2, virtual_block_size=scheduler_block_size)

2549 + # The ``vllm`` logger sets propagate=False, so plain ``caplog`` can observe

2550 + # nothing; ``caplog_vllm`` re-enables propagation for the duration.

2551 + with caplog_vllm.at_level(logging.WARNING):

2552 + hit_tokens, is_async = fix.scheduler.get_num_new_matched_tokens(

2553 + request, num_computed_tokens=BLOCK_SIZE

2554 + )

2555 +

2556 + assert (hit_tokens, is_async) == (0, False)

2557 + assert "requires scheduler-block-aligned local tokens" in caplog_vllm.text

2558 +

2559 +

2560 +def test_fine_grained_external_hits_require_eager_offload() -> None:

2561 + hash_block_size = BLOCK_SIZE

2562 + scheduler_block_size = 4 * BLOCK_SIZE

2563 +

2564 + eager = _make_hybrid_attention_mamba_scheduler(

2565 + block_size=scheduler_block_size,

2566 + hash_block_size=hash_block_size,

2567 + )

2568 + lazy = _make_hybrid_attention_mamba_scheduler(

2569 + block_size=scheduler_block_size,

2570 + hash_block_size=hash_block_size,

2571 + lazy=True,

2572 + )

2573 +

2574 + assert eager.scheduler.cpu_coordinator.enable_partial_hash_hits

2575 + assert not lazy.scheduler.cpu_coordinator.enable_partial_hash_hits

2576 +

2577 +

2578 +def test_eager_store_does_not_scan_mamba_blocks_positionally() -> None:

2579 + """Mamba stores must use exact handoffs, not historical block positions."""

2580 + block_size = 4 * BLOCK_SIZE

2581 + fix = _make_hybrid_attention_mamba_scheduler(

2582 + num_cpu_blocks=16, num_gpu_blocks=24, block_size=block_size

2583 + )

2584 + sched = fix.scheduler

2585 + gpu_pool = fix.gpu_block_pool

2586 + attention_manager, mamba_manager = sched.cpu_coordinator.single_type_managers

2587 + assert attention_manager.has_positionally_stable_blocks

2588 + assert not mamba_manager.has_positionally_stable_blocks

2589 +

2590 + req = _make_cp_request(num_blocks=2, virtual_block_size=block_size)

2591 + attn_blocks = _allocate_cp_gpu_blocks(

2592 + gpu_pool, req, 2, virtual_block_size=block_size, group_id=0

2593 + )

2594 + gpu_blocks = gpu_pool.get_new_blocks(2)

2595 + gpu_blocks[1].set_block_hash(make_block_hash_with_group_id(req.block_hashes[1], 1))

2596 + kv_blocks = KVCacheBlocks(blocks=(attn_blocks, gpu_blocks))

2597 + req.num_computed_tokens = 0

2598 + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)

2599 +

2600 + output = make_scheduler_output(

2601 + {req.request_id: 2 * block_size},

2602 + new_reqs={req.request_id: kv_blocks.get_block_ids()},

2603 + )

2604 + output.kv_connector_block_state = KVConnectorBlockState(

2605 + req_ids=set(),

2606 + resolve_block_ids=lambda _: (),

2607 + boundary_state_offloads={

2608 + req.request_id: [(1, gpu_blocks[1].block_id, 2 * block_size)]

2609 + },

2610 + )

2611 + meta1 = sched.build_connector_meta(output)

2612 + assert meta1.store_event >= 0

2613 + # Only the exact handoff lands this step. build_connector_meta runs before

2614 + # the scheduler commits this step's tokens, so the positional scan has

2615 + # nothing confirmed yet.

2616 + assert set(meta1.store_gpu_blocks) == {gpu_blocks[1].block_id}

2617 + assert sched._reqs_to_store[req.request_id].num_stored_blocks == [0, 0]

2618 + simulate_store_completion(sched, meta1.store_event)

2619 +

2620 + # Next step the attention blocks are confirmed and scanned positionally. A

2621 + # later hash update still must not make the mamba align snapshot a valid

2622 + # source; the explicit boundary-handoff tests cover the supported path.

2623 + gpu_blocks[0].set_block_hash(make_block_hash_with_group_id(req.block_hashes[0], 1))

2624 + req.num_computed_tokens = 2 * block_size

2625 + meta2 = sched.build_connector_meta(

2626 + make_scheduler_output(

2627 + {req.request_id: 1},

2628 + cached_req_new_blocks={req.request_id: None},

2629 + )

2630 + )

2631 + assert set(meta2.store_gpu_blocks) == {block.block_id for block in attn_blocks}

2632 + assert sched._reqs_to_store[req.request_id].num_stored_blocks == [2, 0]

2633 + simulate_store_completion(sched, meta2.store_event)

2634 +

2635 + # Nothing further: mamba align is never picked up positionally.

2636 + meta3 = sched.build_connector_meta(

2637 + make_scheduler_output(

2638 + {req.request_id: 1},

2639 + cached_req_new_blocks={req.request_id: None},

2640 + )

2641 + )

2642 + assert meta3.store_event == -1

2643 + assert meta3.store_gpu_blocks == []

2644 + assert sched._reqs_to_store[req.request_id].num_stored_blocks == [2, 0]

2645 +

2646 +

2647 +def test_eager_store_scans_mamba_all_blocks_positionally() -> None:

2648 + """Mamba all-mode retains positional identity and needs no handoff."""

2649 + block_size = 4 * BLOCK_SIZE

2650 + fix = _make_hybrid_attention_mamba_scheduler(

2651 + num_cpu_blocks=16,

2652 + num_gpu_blocks=24,

2653 + block_size=block_size,

2654 + mamba_cache_mode="all",

2655 + )

2656 + sched = fix.scheduler

2657 + gpu_pool = fix.gpu_block_pool

2658 + attention_manager, mamba_manager = sched.cpu_coordinator.single_type_managers

2659 + assert attention_manager.has_positionally_stable_blocks

2660 + assert mamba_manager.has_positionally_stable_blocks

2661 +

2662 + req = _make_cp_request(num_blocks=2, virtual_block_size=block_size)

2663 + attention_blocks = _allocate_cp_gpu_blocks(

2664 + gpu_pool, req, 2, virtual_block_size=block_size, group_id=0

2665 + )

2666 + mamba_blocks = _allocate_cp_gpu_blocks(

2667 + gpu_pool, req, 2, virtual_block_size=block_size, group_id=1

2668 + )

2669 + kv_blocks = KVCacheBlocks(blocks=(attention_blocks, mamba_blocks))

2670 + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)

2671 +

2672 + # build_connector_meta runs before the scheduler commits this step's

2673 + # tokens, so the positional scan picks the blocks up on the next step.

2674 + first = sched.build_connector_meta(

2675 + make_scheduler_output(

2676 + {req.request_id: 2 * block_size},

2677 + new_reqs={req.request_id: kv_blocks.get_block_ids()},

2678 + )

2679 + )

2680 + assert first.store_gpu_blocks == []

2681 + req.num_computed_tokens = 2 * block_size

2682 + meta = sched.build_connector_meta(

2683 + make_scheduler_output(

2684 + {req.request_id: 1},

2685 + cached_req_new_blocks={req.request_id: None},

2686 + )

2687 + )

2688 +

2689 + assert set(meta.store_gpu_blocks) == {

2690 + *(block.block_id for block in attention_blocks),

2691 + *(block.block_id for block in mamba_blocks),

2692 + }

2693 + assert sched._reqs_to_store[req.request_id].num_stored_blocks == [2, 2]

2694 +

2695 +

2696 +@pytest.mark.parametrize("gone_via", ["preempted", "finished", "unregistered"])

2697 +def test_boundary_handoff_dropped_for_departing_request(gone_via: str) -> None:

2698 + """Handoffs for a request leaving this step must not be read.

2699 +

2700 + The scheduler drains boundary handoffs without filtering by request

2701 + liveness, so a preempted or finished request can still offer block IDs.

2702 + Those blocks are being released and may already back another request, so

2703 + copying them would publish unrelated KV under a valid hash.

2704 + """

2705 + block_size = 4 * BLOCK_SIZE

2706 + fix = _make_hybrid_attention_mamba_scheduler(

2707 + num_cpu_blocks=16, num_gpu_blocks=24, block_size=block_size

2708 + )

2709 + sched = fix.scheduler

2710 + gpu_pool = fix.gpu_block_pool

2711 +

2712 + req = _make_cp_request(num_blocks=2, virtual_block_size=block_size)

2713 + attn_blocks = _allocate_cp_gpu_blocks(

2714 + gpu_pool, req, 2, virtual_block_size=block_size, group_id=0

2715 + )

2716 + mamba_blocks = gpu_pool.get_new_blocks(2)

2717 + mamba_blocks[1].set_block_hash(

2718 + make_block_hash_with_group_id(req.block_hashes[1], 1)

2719 + )

2720 + kv_blocks = KVCacheBlocks(blocks=(attn_blocks, mamba_blocks))

2721 + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)

2722 +

2723 + output = make_scheduler_output(

2724 + {req.request_id: 2 * block_size},

2725 + new_reqs={req.request_id: kv_blocks.get_block_ids()},

2726 + )

2727 + output.kv_connector_block_state = KVConnectorBlockState(

2728 + req_ids=set(),

2729 + resolve_block_ids=lambda _: (),

2730 + boundary_state_offloads={

2731 + req.request_id: [(1, mamba_blocks[1].block_id, 2 * block_size)]

2732 + },

2733 + )

2734 + if gone_via == "preempted":

2735 + output.preempted_req_ids = {req.request_id}

2736 + elif gone_via == "finished":

2737 + output.finished_req_ids = {req.request_id}

2738 + else:

2739 + sched._reqs_to_store.pop(req.request_id)

2740 +

2741 + meta = sched.build_connector_meta(output)

2742 +

2743 + # The exact-handoff block is never a store source for a departing request.

2744 + assert mamba_blocks[1].block_id not in meta.store_gpu_blocks

2745 + stats = sched.get_boundary_store_stats()

2746 + assert stats.dropped_request_gone == 1

2747 + assert stats.stored == 0

2748 +

2749 +

2750 +def test_boundary_handoff_keeps_block_meta_index_parallel() -> None:

2751 + """Boundary stores must extend block_meta alongside the block ids.

2752 +

2753 + ``_process_store_completion`` indexes ``block_meta`` by position, so a

2754 + boundary handoff that appends only block ids shortens the list, trips its

2755 + length assertion and shifts every following entry's metadata.

2756 + """

2757 + block_size = 4 * BLOCK_SIZE

2758 + fix = _make_hybrid_attention_mamba_scheduler(

2759 + num_cpu_blocks=16,

2760 + num_gpu_blocks=24,

2761 + block_size=block_size,

2762 + enable_kv_cache_events=True,

2763 + )

2764 + sched = fix.scheduler

2765 + gpu_pool = fix.gpu_block_pool

2766 + assert sched.enable_kv_cache_events

2767 +

2768 + req = _make_cp_request(num_blocks=2, virtual_block_size=block_size)

2769 + attn_blocks = _allocate_cp_gpu_blocks(

2770 + gpu_pool, req, 2, virtual_block_size=block_size, group_id=0

2771 + )

2772 + mamba_blocks = gpu_pool.get_new_blocks(2)

2773 + mamba_blocks[1].set_block_hash(

2774 + make_block_hash_with_group_id(req.block_hashes[1], 1)

2775 + )

2776 + kv_blocks = KVCacheBlocks(blocks=(attn_blocks, mamba_blocks))

2777 + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)

2778 + # Make the FA blocks positionally eligible in the same event as the

2779 + # explicit Mamba boundary handoff.

2780 + req.num_computed_tokens = 2 * block_size

2781 +

2782 + output = make_scheduler_output(

2783 + {req.request_id: 2 * block_size},

2784 + new_reqs={req.request_id: kv_blocks.get_block_ids()},

2785 + )

2786 + output.kv_connector_block_state = KVConnectorBlockState(

2787 + req_ids=set(),

2788 + resolve_block_ids=lambda _: (),

2789 + boundary_state_offloads={

2790 + req.request_id: [(1, mamba_blocks[1].block_id, 2 * block_size)]

2791 + },

2792 + )

2793 +

2794 + gpu_ids, cpu_ids, req_ids, block_meta = sched.prepare_store_specs(output)

2795 + assert mamba_blocks[1].block_id in gpu_ids

2796 + assert set(gpu_ids) == {

2797 + *(block.block_id for block in attn_blocks),

2798 + mamba_blocks[1].block_id,

2799 + }

2800 + assert req_ids == [req.request_id]

2801 + assert block_meta is not None

2802 + assert len(block_meta) == len(gpu_ids) == len(cpu_ids)

vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py

@@ -19,6 +19,7 @@ from vllm.logger import init_logger

19 from vllm.v1.core.sched.output import SchedulerOutput

20 from vllm.v1.outputs import KVConnectorOutput

21 from vllm.v1.simple_kv_offload.manager import (

22 + BoundaryStoreStats,

23 SimpleCPUOffloadScheduler,

24 )

25 from vllm.v1.simple_kv_offload.metadata import (

@@ -300,6 +301,12 @@ class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA):

301 return self.scheduler_manager.has_pending_stores()

302 return False

303

304 + def get_boundary_store_stats(self) -> BoundaryStoreStats | None:

305 + """Return cumulative boundary-handoff store diagnostics."""

306 + if self.scheduler_manager is not None:

307 + return self.scheduler_manager.get_boundary_store_stats()

308 + return None

309 +

310 def take_events(self) -> Iterable[KVCacheEvent]:

311 if self.scheduler_manager is not None:

312 return self.scheduler_manager.take_events()

vllm/v1/core/kv_cache_coordinator.py

+12 -4

@@ -155,6 +155,9 @@ class KVCacheCoordinator(ABC):

155 for manager in self.single_type_managers:

156 if isinstance(manager, MambaManager):

157 manager.drop_eagle_checkpoint_block = True

158 + self.group_block_sizes = tuple(

159 + manager.block_size for manager in self.single_type_managers

160 + )

161

162 # A positive retention interval must be a multiple of the base hit granularity

163 # (``scheduler_block_size``) to land on real cache-hit boundaries.

@@ -605,6 +608,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

608 hash_block_size: int,

609 metrics_collector: KVCacheMetricsCollector | None = None,

610 num_prefill_lookahead: int = 0,

611 + allow_partial_hash_hits: bool = True,

612 ):

613 super().__init__(

614 kv_cache_config,

@@ -629,7 +633,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

633 # Only groups that participate in prefix caching must satisfy the

634 # divisibility constraint; groups that opt out (e.g. GLM-5.3-Flash kpool

635 # tail, block_size=kpool) are scratch buffers and excluded.

632 - group_block_sizes = [

636 + cacheable_block_sizes = [

637 manager.block_size

638 for manager, group in zip(

639 self.single_type_managers, kv_cache_config.kv_cache_groups

@@ -637,10 +641,10 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

641 if group.kv_cache_spec.prefix_cacheable

642 ]

643 assert all(

640 - block_size % hash_block_size == 0 for block_size in group_block_sizes

644 + block_size % hash_block_size == 0 for block_size in cacheable_block_sizes

645 ), (

646 "Each KV cache group's real block_size must be divisible by "

643 - f"hash_block_size. block_sizes={group_block_sizes}, "

647 + f"hash_block_size. block_sizes={cacheable_block_sizes}, "

648 f"hash_block_size={hash_block_size}"

649 )

650 assert pcp_world_size == 1, "PCP not support hybrid attn now."

@@ -669,7 +673,9 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

673 )

674 for g in kv_cache_config.kv_cache_groups

675 )

672 - self.enable_partial_hash_hits = has_partial_mamba_group

676 + self.enable_partial_hash_hits = (

677 + allow_partial_hash_hits and has_partial_mamba_group

678 + )

679 if self.enable_partial_hash_hits:

680 unsupported_partial_hit_managers = {

681 type(manager).__name__

@@ -987,6 +993,7 @@ def get_kv_cache_coordinator(

993 hash_block_size: int,

994 metrics_collector: KVCacheMetricsCollector | None = None,

995 num_prefill_lookahead: int = 0,

996 + allow_partial_hash_hits: bool = True,

997 ) -> KVCacheCoordinator:

998 if not enable_caching:

999 return KVCacheCoordinatorNoPrefixCache(

@@ -1030,4 +1037,5 @@ def get_kv_cache_coordinator(

1037 hash_block_size=hash_block_size,

1038 metrics_collector=metrics_collector,

1039 num_prefill_lookahead=num_prefill_lookahead,

1040 + allow_partial_hash_hits=allow_partial_hash_hits,

1041 )

vllm/v1/core/single_type_kv_cache_manager.py

+16

@@ -48,6 +48,16 @@ class SingleTypeKVCacheManager(ABC):

49 supports_fine_grained_hash_lookup: ClassVar[bool] = False

51 + @property

52 + def has_positionally_stable_blocks(self) -> bool:

53 + """Whether positional offload scans can follow this block table.

54 +

55 + True means already-scanned indices will not be reused for a different

56 + token range. Managers may still null blocks that left their retention

57 + window, as long as the cursor's positional history remains valid.

58 + """

59 + return True

60 +

61 def __init__(

62 self,

63 kv_cache_spec: KVCacheSpec,

@@ -1423,6 +1433,12 @@ class ChunkedLocalAttentionManager(SingleTypeKVCacheManager):

1433 class MambaManager(SingleTypeKVCacheManager):

1434 supports_fine_grained_hash_lookup: ClassVar[bool] = True

1435

1436 + @property

1437 + def has_positionally_stable_blocks(self) -> bool:

1438 + # Align-mode Mamba can null interior states and relocate speculative

1439 + # blocks in place. Other modes retain positional identity.

1440 + return self.mamba_cache_mode != "align"

1441 +

1442 def __init__(

1443 self, kv_cache_spec: MambaSpec, block_pool: BlockPool, **kwargs

1444 ) -> None:

vllm/v1/simple_kv_offload/manager.py

+266 -66

@@ -5,6 +5,7 @@

5 import contextlib

6 from collections.abc import Iterable

7 from dataclasses import dataclass, field, replace

8 +from enum import Enum, auto

9 from typing import TYPE_CHECKING, Any

11 from vllm.config import VllmConfig

@@ -26,6 +27,7 @@ from vllm.v1.core.kv_cache_coordinator import (

27 from vllm.v1.core.kv_cache_utils import (

28 BlockHashWithGroupId,

29 ExternalBlockHash,

30 + dcp_world_size_for_kv_cache_spec,

31 generate_block_hash_extra_keys,

32 get_block_hash,

33 get_group_id,

@@ -99,6 +101,36 @@ class StoreRequestState:

101 finished: bool = False

102

103

104 +class _StoreAdmission(Enum):

105 + READY = auto()

106 + NULL_BLOCK = auto()

107 + NOT_HASHED = auto()

108 + IN_FLIGHT = auto()

109 + ALREADY_CACHED = auto()

110 +

111 +

112 +@dataclass

113 +class BoundaryStoreStats:

114 + """Cumulative counters for boundary hand-off stores.

115 +

116 + Every branch that declines a hand-off is silent by design: the boundary is

117 + simply not offloaded and a later request misses. Without counters the only

118 + symptom is a lower cache hit rate with nothing in the logs, so each decline

119 + reason is counted and exposed through

120 + ``SimpleCPUOffloadConnector.get_boundary_store_stats()``. Reset by

121 + ``reset()``; not otherwise cleared.

122 + """

123 +

124 + published: int = 0

125 + stored: int = 0

126 + dropped_cpu_full: int = 0

127 + dropped_request_gone: int = 0

128 + skipped_already_cached: int = 0

129 + skipped_in_flight: int = 0

130 + dropped_null_block: int = 0

131 + dropped_not_hashed: int = 0

132 +

133 +

134 class SimpleCPUOffloadScheduler:

135 """Scheduler-side manager for CPU offloading."""

136

@@ -141,16 +173,6 @@ class SimpleCPUOffloadScheduler:

173 self.fa_gidx = g_idx

174 break

175 assert 0 <= self.fa_gidx < len(self.cpu_kv_cache_config.kv_cache_groups)

144 - # FA group's own block_size; divides scheduler_block_size (the LCM)

145 - # but is NOT assumed to equal it.

146 - self.fa_block_size: int = (

147 - self.cpu_kv_cache_config.kv_cache_groups[

148 - self.fa_gidx

149 - ].kv_cache_spec.block_size

150 - * self.cp_world_size

151 - )

152 - assert self.block_size % self.fa_block_size == 0

153 -

176 logger.info(

177 "SimpleCPUOffloadScheduler: Allocating %d offload blocks "

178 "(%.2f GB, mode=%s, backend=%s)",

@@ -175,7 +197,13 @@ class SimpleCPUOffloadScheduler:

197 pcp_world_size=1,

198 scheduler_block_size=self.block_size,

199 hash_block_size=self.hash_block_size,

200 + allow_partial_hash_hits=not lazy_offload,

201 )

202 + self.group_block_sizes = self.cpu_coordinator.group_block_sizes

203 + # FA group's own resolved block_size; divides scheduler_block_size (the

204 + # LCM) but is NOT assumed to equal it.

205 + self.fa_block_size: int = self.group_block_sizes[self.fa_gidx]

206 + assert self.block_size % self.fa_block_size == 0

207 self.cpu_block_pool: BlockPool = self.cpu_coordinator.block_pool

208 # GPU block pool reference - bound after scheduler builds kv_cache_manager

209 self._gpu_block_pool: BlockPool | None = None

@@ -186,10 +214,13 @@ class SimpleCPUOffloadScheduler:

214 # the worker reports completions by event index, not request id.

215 self._load_event_to_reqs: dict[int, list[str]] = {}

216

189 - # Pending (cpu_hit_blocks, hit_length) tuples from find_longest_cache_hit,

190 - # kept pinned via touch() while awaiting update_state_after_alloc().

217 + # Pending (cpu_hit_blocks, hit_length, num_computed_tokens) tuples from

218 + # find_longest_cache_hit, kept pinned via touch() while awaiting

219 + # update_state_after_alloc(). ``num_computed_tokens`` is the local

220 + # prefix the hit was resolved against; the load path needs it to place

221 + # external blocks and must not re-derive it from block metadata.

222 self._pending_cpu_hits: dict[

192 - str, tuple[tuple[list[KVCacheBlock], ...], int]

223 + str, tuple[tuple[list[KVCacheBlock], ...], int, int]

224 ] = {}

225

226 # Store metadata

@@ -216,6 +247,7 @@ class SimpleCPUOffloadScheduler:

247 # Event counters

248 self._load_event_counter: int = 0

249 self._store_event_counter: int = 0

250 + self.boundary_store_stats = BoundaryStoreStats()

251

252 # For TP/PP: track partial store completions across steps.

253 # Events must be reported by all world_size workers before considered complete.

@@ -267,7 +299,11 @@ class SimpleCPUOffloadScheduler:

299 target = 0

300 for g in kv_cache_config.kv_cache_groups:

301 spec = g.kv_cache_spec

270 - block_size = spec.block_size * cp_world_size

302 + # Only full attention is sharded across DCP ranks; replicated specs

303 + # (mamba, sliding window, chunked-local) keep their own block size.

304 + block_size = spec.block_size * dcp_world_size_for_kv_cache_spec(

305 + spec, cp_world_size

306 + )

307 if isinstance(spec, MambaSpec):

308 target += 2

309 elif isinstance(spec, SlidingWindowSpec):

@@ -296,6 +332,20 @@ class SimpleCPUOffloadScheduler:

332 if request.skip_reading_prefix_cache:

333 return 0, False

334

335 + if num_computed_tokens % self.block_size != 0:

336 + # Transfers are whole-block copies, so an external suffix cannot

337 + # start in the middle of a destination block. When the GPU-side

338 + # coordinator also serves fine-grained hits, the local prefix can

339 + # land on a hash boundary that is not a scheduler-block boundary;

340 + # such a request keeps its local hit and skips the external one.

341 + logger.warning_once(

342 + "SimpleCPU external lookup requires scheduler-block-aligned "

343 + "local tokens, got %d tokens with scheduler block size %d.",

344 + num_computed_tokens,

345 + self.block_size,

346 + )

347 + return 0, False

348 +

349 num_skipped_hashes = num_computed_tokens // self.hash_block_size

350 remaining_hashes = request.block_hashes[num_skipped_hashes:]

351

@@ -318,6 +368,7 @@ class SimpleCPUOffloadScheduler:

368 self._pending_cpu_hits[request.request_id] = (

369 cpu_hit_blocks,

370 hit_length,

371 + num_computed_tokens,

372 )

373 return hit_length, True

374 return 0, False

@@ -371,20 +422,19 @@ class SimpleCPUOffloadScheduler:

422 )

423 return

424

374 - cpu_hit_blocks_full, _ = pending

425 + cpu_hit_blocks_full, _, num_computed_tokens = pending

426

376 - # ``num_external_tokens`` is LCM-aligned (checked per-group below),

377 - # so this counts whole scheduler-aligned chunks of incoming tokens.

378 - num_blocks_to_load = num_external_tokens // self.block_size

379 - assert num_blocks_to_load > 0

380 - num_cached_fa_blocks = sum(

381 - blk.block_hash is not None for blk in blocks.blocks[self.fa_gidx]

427 + assert num_computed_tokens % self.block_size == 0, (

428 + "SimpleCPU external loads must start at a scheduler-block boundary"

429 + )

430 + # Fine-grained hits are hash-block aligned, not group-block aligned, so

431 + # the per-group block count below rounds up instead of dividing exactly.

432 + # Hash alignment is still required: it is what makes the accepted

433 + # external suffix start on a boundary every group can address.

434 + assert num_external_tokens % self.hash_block_size == 0, (

435 + f"num_external_tokens={num_external_tokens} is not aligned to "

436 + f"hash_block_size={self.hash_block_size}"

437 )

383 - num_computed_tokens = num_cached_fa_blocks * self.fa_block_size

384 -

385 - # Build transfer pairs across all groups.

386 - total_computed_tokens = num_computed_tokens + num_external_tokens

387 - kv_cache_groups = self.cpu_kv_cache_config.kv_cache_groups

438

439 # The scheduler may have accepted fewer blocks than

440 # get_num_new_matched_tokens() reported.

@@ -393,14 +443,8 @@ class SimpleCPUOffloadScheduler:

443 # the rest will be released along with the temp pin below.

444 cpu_hit_blocks: list[list[KVCacheBlock]] = []

445 for g in range(num_groups):

396 - g_block_size = (

397 - kv_cache_groups[g].kv_cache_spec.block_size * self.cp_world_size

398 - )

399 - assert num_external_tokens % g_block_size == 0, (

400 - f"num_external_tokens={num_external_tokens} not aligned to "

401 - f"group {g} block_size={g_block_size}"

402 - )

403 - n_take_g = num_external_tokens // g_block_size

446 + g_block_size = self.group_block_sizes[g]

447 + n_take_g = cdiv(num_external_tokens, g_block_size)

448 cpu_hit_blocks.append(cpu_hit_blocks_full[g][:n_take_g])

449

450 gpu_block_ids: list[int] = []

@@ -413,14 +457,8 @@ class SimpleCPUOffloadScheduler:

457 if n_ext_g == 0:

458 continue

459

416 - # Number of blocks in the computed range for this group.

417 - g_block_size = (

418 - kv_cache_groups[g].kv_cache_spec.block_size * self.cp_world_size

419 - )

420 - n_computed_g = cdiv(total_computed_tokens, g_block_size)

421 -

422 - # Back-trace: ext blocks sit at the tail of the computed range.

423 - gpu_ext_start = n_computed_g - n_ext_g

460 + g_block_size = self.group_block_sizes[g]

461 + gpu_ext_start = num_computed_tokens // g_block_size

462 group_gpu_ids = block_ids_by_group[g]

463

464 for i, cpu_blk in enumerate(cpu_blocks_g):

@@ -614,6 +652,72 @@ class SimpleCPUOffloadScheduler:

652 num_groups = len(self.cpu_kv_cache_config.kv_cache_groups)

653 # Dedup against blocks already scheduled.

654 in_flight = self._in_flight_store_gpu_blocks

655 + num_free = cpu_block_pool.get_num_free_blocks()

656 +

657 + block_state = scheduler_output.kv_connector_block_state

658 + boundary_offloads = (

659 + block_state.boundary_state_offloads if block_state is not None else {}

660 + )

661 + stats = self.boundary_store_stats

662 + preempted_req_ids = scheduler_output.preempted_req_ids or ()

663 + for req_id, entries in boundary_offloads.items():

664 + # The scheduler drains handoffs without filtering by request

665 + # liveness. A request that finished or was preempted in this step

666 + # is giving up its blocks, which may already have been reallocated

667 + # to another request; reading them would publish unrelated KV under

668 + # a valid hash. Drop conservatively, as the mooncake store does.

669 + store_state = self._reqs_to_store.get(req_id)

670 + if (

671 + store_state is None

672 + or store_state.finished

673 + or req_id in preempted_req_ids

674 + or req_id in scheduler_output.finished_req_ids

675 + ):

676 + stats.published += len(entries)

677 + stats.dropped_request_gone += len(entries)

678 + continue

679 + scheduled_for_req = False

680 + # ``boundary_tokens`` is not needed here: the cache key comes from

681 + # the handed-off block itself, not from the offered boundary.

682 + for _, gpu_block_id, _ in entries:

683 + stats.published += 1

684 + gpu_block = gpu_block_pool.blocks[gpu_block_id]

685 + admission = self._classify_store_candidate(gpu_block)

686 + if admission is _StoreAdmission.NULL_BLOCK:

687 + stats.dropped_null_block += 1

688 + continue

689 + if admission is _StoreAdmission.NOT_HASHED:

690 + stats.dropped_not_hashed += 1

691 + continue

692 + if admission is _StoreAdmission.IN_FLIGHT:

693 + stats.skipped_in_flight += 1

694 + continue

695 + if admission is _StoreAdmission.ALREADY_CACHED:

696 + stats.skipped_already_cached += 1

697 + continue

698 + if num_free <= 0:

699 + stats.dropped_cpu_full += 1

700 + logger.warning_once(

701 + "SimpleCPU dropped a boundary-state handoff because "

702 + "the CPU block pool is full; this boundary will not be retried."

703 + )

704 + continue

705 + cpu_block = cpu_block_pool.get_new_blocks(1)[0]

706 + merged_gpu_block_ids.append(gpu_block_id)

707 + merged_cpu_block_ids.append(cpu_block.block_id)

708 + if merged_block_meta is not None:

709 + # Keep the metadata list index-parallel with the block ids.

710 + # The hand-off block's own hashes are resolved at completion,

711 + # so there is nothing to capture here; emit the event without

712 + # a payload, as the size-mismatch fallback already does.

713 + merged_block_meta.append({})

714 + in_flight.add(gpu_block_id)

715 + gpu_block_pool.touch([gpu_block])

716 + num_free -= 1

717 + stats.stored += 1

718 + scheduled_for_req = True

719 + if scheduled_for_req:

720 + req_ids.append(req_id)

721

722 for req_id, new_block_id_groups, preempted in yield_req_data(scheduler_output):

723 state = self._reqs_to_store.get(req_id)

@@ -641,7 +745,7 @@ class SimpleCPUOffloadScheduler:

745 self._select_eager_blocks_to_store(state, block_ids_by_group)

746 )

747

644 - # --- Phase 2: Batch allocate CPU blocks ---

748 + # Batch allocate the CPU destinations.

749 n_to_alloc = len(gpu_block_ids)

750 if n_to_alloc > 0:

751 cpu_blocks_alloc = cpu_block_pool.get_new_blocks(n_to_alloc)

@@ -674,6 +778,9 @@ class SimpleCPUOffloadScheduler:

778 for g in range(num_groups):

779 state.num_stored_blocks[g] += advanced_per_group[g]

780

781 + # A request can contribute both boundary and positional blocks to the

782 + # same event. Completion tracking is set-based, so keep one request ID.

783 + req_ids = list(dict.fromkeys(req_ids))

784 return merged_gpu_block_ids, merged_cpu_block_ids, req_ids, merged_block_meta

785

786 def _select_eager_blocks_to_store(

@@ -694,19 +801,31 @@ class SimpleCPUOffloadScheduler:

801 )

802 request = state.request

803 confirmed_tokens = request.num_computed_tokens - request.num_output_placeholders

697 - aligned_tokens = confirmed_tokens // self.block_size * self.block_size

804 + # Truncate to the granularity a lookup can actually land on. With

805 + # fine-grained hits that is hash_block_size; otherwise hits only land on

806 + # the scheduler block (the group LCM). Using the LCM unconditionally

807 + # would drop whole blocks that sit between the last LCM boundary and a

808 + # reachable fine-grained boundary, so the other groups would hold that

809 + # boundary and this one would not, and the joint hybrid lookup would

810 + # reconcile to zero.

811 + store_alignment = (

812 + self.hash_block_size

813 + if self.cpu_coordinator.enable_partial_hash_hits

814 + else self.block_size

815 + )

816 + aligned_tokens = confirmed_tokens // store_alignment * store_alignment

817 num_free = self.cpu_block_pool.get_num_free_blocks()

818

819 for g, group_gpu_ids in enumerate(block_ids_by_group):

820 if len(gpu_block_ids) >= num_free:

821 break

822 + group_manager = self.cpu_coordinator.single_type_managers[g]

823 + if not group_manager.has_positionally_stable_blocks:

824 + continue

825 # FIXME (yifan): handle CPU cache eviction, where

826 # num_stored_blocks can be stale and omit evicted blocks in

827 # the middle of the request.

706 - group_size = (

707 - self.cpu_kv_cache_config.kv_cache_groups[g].kv_cache_spec.block_size

708 - * self.cp_world_size

709 - )

828 + group_size = self.group_block_sizes[g]

829 ready = min(len(group_gpu_ids), aligned_tokens // group_size)

830 resolved_hashes = (

831 resolve_block_hashes(

@@ -720,20 +839,16 @@ class SimpleCPUOffloadScheduler:

839 start = state.num_stored_blocks[g]

840 for i, gpu_block_id in enumerate(group_gpu_ids[start:ready], start=start):

841 gpu_block = self._gpu_block_pool.blocks[gpu_block_id]

723 - if gpu_block.is_null or gpu_block.block_hash is None:

724 - advanced_per_group[g] += 1

725 - continue

842 if (

727 - gpu_block_id in self._in_flight_store_gpu_blocks

728 - or self.cpu_block_pool.cached_block_hash_to_block.get_one_block(

729 - gpu_block.block_hash

730 - )

731 - is not None

843 + self._classify_store_candidate(gpu_block)

844 + is not _StoreAdmission.READY

845 ):

846 advanced_per_group[g] += 1

847 continue

848 if len(gpu_block_ids) >= num_free:

849 break

850 + primary_block_hash = gpu_block.block_hash

851 + assert primary_block_hash is not None

852 gpu_block_ids.append(gpu_block_id)

853 advanced_per_group[g] += 1

854 if block_meta is not None:

@@ -750,7 +865,7 @@ class SimpleCPUOffloadScheduler:

865 request, token_start, token_end, curr_mm_idx

866 )

867 meta_by_hash = {

753 - gpu_block.block_hash: BlockStoreMeta(

868 + primary_block_hash: BlockStoreMeta(

869 token_ids=list(

870 request.all_token_ids[token_start:token_end]

871 ),

@@ -766,7 +881,7 @@ class SimpleCPUOffloadScheduler:

881 block_hash = make_block_hash_with_group_id(

882 request.block_hashes[hash_idx], g

883 )

769 - if block_hash == gpu_block.block_hash:

884 + if block_hash is None or block_hash == primary_block_hash:

885 continue

886 hash_start = hash_idx * self.hash_block_size

887 hash_end = hash_start + self.hash_block_size

@@ -796,6 +911,25 @@ class SimpleCPUOffloadScheduler:

911

912 return gpu_block_ids, advanced_per_group, block_meta

913

914 + def _classify_store_candidate(self, gpu_block: "KVCacheBlock") -> _StoreAdmission:

915 + if gpu_block.is_null:

916 + return _StoreAdmission.NULL_BLOCK

917 + if gpu_block.block_hash is None:

918 + return _StoreAdmission.NOT_HASHED

919 + if gpu_block.block_id in self._in_flight_store_gpu_blocks:

920 + return _StoreAdmission.IN_FLIGHT

921 + if (

922 + self.cpu_block_pool.cached_block_hash_to_block.get_one_block(

923 + gpu_block.block_hash

924 + )

925 + is not None

926 + ):

927 + return _StoreAdmission.ALREADY_CACHED

928 + return _StoreAdmission.READY

929 +

930 + def get_boundary_store_stats(self) -> BoundaryStoreStats:

931 + return replace(self.boundary_store_stats)

932 +

933 def update_connector_output(self, connector_output: KVConnectorOutput) -> None:

934 """Handle async transfer completions from worker.

935

@@ -886,12 +1020,7 @@ class SimpleCPUOffloadScheduler:

1020 if self.enable_kv_cache_events:

1021 meta_by_hash = block_meta[i] if block_meta is not None else {}

1022 primary_group_idx = get_group_id(primary_hash)

889 - primary_spec = self.cpu_kv_cache_config.kv_cache_groups[

890 - primary_group_idx

891 - ].kv_cache_spec

892 - primary_block_size = primary_spec.block_size * (

893 - self.cp_world_size if not isinstance(primary_spec, MambaSpec) else 1

894 - )

1023 + primary_block_size = self.group_block_sizes[primary_group_idx]

1024 events_to_emit: list[

1025 tuple[BlockHashWithGroupId, BlockStoreMeta | None, int]

1026 ] = [

@@ -1028,6 +1157,21 @@ class SimpleCPUOffloadScheduler:

1157 if state is None or gpu_pool is None:

1158 return

1159 gpu_ids, _, block_meta = self._select_eager_blocks_to_store(state, block_ids)

1160 + partial_tail = self._find_fa_partial_tail_source(request, block_ids)

1161 + if (

1162 + partial_tail is not None

1163 + # Decode tokens can fill the boundary block, in which case the

1164 + # positional scan above already selected it. Appending again would

1165 + # allocate two CPU blocks for one GPU block.

1166 + and partial_tail not in gpu_ids

1167 + and len(gpu_ids) < self.cpu_block_pool.get_num_free_blocks()

1168 + ):

1169 + gpu_ids.append(partial_tail)

1170 + if block_meta is not None:

1171 + # The block's own hashes are registered at completion, so no

1172 + # capture is needed here; emit the event without a payload, as

1173 + # the size-mismatch fallback above already does.

1174 + block_meta.append({})

1175 if not gpu_ids:

1176 return

1177 cpu_blocks = self.cpu_block_pool.get_new_blocks(len(gpu_ids))

@@ -1037,9 +1181,64 @@ class SimpleCPUOffloadScheduler:

1181 self._in_flight_store_gpu_blocks.update(gpu_ids)

1182 gpu_pool.touch([gpu_pool.blocks[bid] for bid in gpu_ids])

1183

1040 - def _free_pending_cpu_hit(self, pending: tuple) -> None:

1184 + def _find_fa_partial_tail_source(

1185 + self, request: "Request", block_ids: tuple[list[int], ...]

1186 + ) -> int | None:

1187 + """Locate the full-attention prompt-tail block for a fine-grained hit.

1188 +

1189 + Mirrors ``FullAttentionManager._cache_partial_tail_block``: only the

1190 + final prompt hash boundary is eligible, and boundaries that land on a

1191 + physical block edge are already covered by the positional scan.

1192 +

1193 + Attention block tables are append-only, so the block is located

1194 + positionally, as the mooncake store and offloading connector also do;

1195 + no new core hand-off is required. The boundary key is registered at

1196 + completion along with the block's other hashes, so nothing has to be

1197 + captured here.

1198 + """

1199 + if not self.cpu_coordinator.enable_partial_hash_hits:

1200 + return None

1201 + assert self._gpu_block_pool is not None

1202 + boundary_tokens = (

1203 + request.num_prompt_tokens // self.hash_block_size * self.hash_block_size

1204 + )

1205 + if boundary_tokens == 0 or boundary_tokens > request.num_computed_tokens:

1206 + return None

1207 + if boundary_tokens % self.fa_block_size == 0:

1208 + return None

1209 + block_idx = boundary_tokens // self.fa_block_size

1210 + fa_block_ids = block_ids[self.fa_gidx]

1211 + if block_idx >= len(fa_block_ids):

1212 + return None

1213 + gpu_block_id = fa_block_ids[block_idx]

1214 + gpu_block = self._gpu_block_pool.blocks[gpu_block_id]

1215 + if gpu_block.is_null or gpu_block_id in self._in_flight_store_gpu_blocks:

1216 + return None

1217 + hash_idx = boundary_tokens // self.hash_block_size - 1

1218 + if hash_idx >= len(request.block_hashes):

1219 + return None

1220 + block_hash = make_block_hash_with_group_id(

1221 + request.block_hashes[hash_idx], self.fa_gidx

1222 + )

1223 + # Only worth copying when the boundary key is actually registered on

1224 + # the source block; completion replays the block's own hashes.

1225 + if (

1226 + gpu_block.block_hash != block_hash

1227 + and block_hash

1228 + not in self._gpu_block_pool.cached_block_hashes_by_block.get(

1229 + gpu_block_id, ()

1230 + )

1231 + ):

1232 + return None

1233 + if self.cpu_block_pool.cached_block_hash_to_block.get_one_block(block_hash):

1234 + return None

1235 + return gpu_block_id

1236 +

1237 + def _free_pending_cpu_hit(

1238 + self, pending: tuple[tuple[list["KVCacheBlock"], ...], int, int]

1239 + ) -> None:

1240 """Release the temporary CPU block pin taken in get_num_new_matched_tokens()."""

1042 - cpu_hit_blocks, _ = pending

1241 + cpu_hit_blocks, _hit_length, _num_computed_tokens = pending

1242 blocks_to_free = [

1243 blk for grp in cpu_hit_blocks for blk in grp if not blk.is_null

1244 ]

@@ -1140,6 +1339,7 @@ class SimpleCPUOffloadScheduler:

1339 if event_idx in self._abandoned_store_event_to_blocks

1340 }

1341 self._cursor = None

1342 + self.boundary_store_stats = BoundaryStoreStats()

1343 # NOTE: _load_event_counter / _store_event_counter are not

1344 # reset as they are monotonic and must stay ahead of the workers

1345 # high-water marks to avoid event index collisions