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, 22:00 UTC
[Feat][KVConnector] Support DSV4 in SimpleCPUOffloadBackend (#42296)
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Yifan Qiao committed
May 24, 2026 at 23:04 UTC
81252d4e2446552f93225b1d48873d45999f045e
3 files changed
+43
-19
tests/v1/simple_kv_offload/test_scheduler.py
+2
@@ -153,6 +153,8 @@ def make_scheduler(
153
vllm_config=vllm_config,
154
kv_cache_config=kv_cache_config,
155
cpu_capacity_bytes=cpu_capacity_bytes,
156
+ scheduler_block_size=BLOCK_SIZE,
157
+ hash_block_size=BLOCK_SIZE,
158
lazy_offload=lazy,
159
)
160
vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py
+8
@@ -96,10 +96,18 @@ class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA):
96
)
97
98
if role == KVConnectorRole.SCHEDULER:
99
+ from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes
100
+
101
+ assert kv_cache_config is not None
102
+ scheduler_block_size, hash_block_size = resolve_kv_cache_block_sizes(
103
+ kv_cache_config, vllm_config
104
+ )
105
self.scheduler_manager = SimpleCPUOffloadScheduler(
106
vllm_config,
107
kv_cache_config,
108
cpu_capacity_per_rank,
109
+ scheduler_block_size=scheduler_block_size,
110
+ hash_block_size=hash_block_size,
111
lazy_offload=lazy_offload,
112
)
113
elif role == KVConnectorRole.WORKER:
vllm/v1/simple_kv_offload/manager.py
+33
-19
@@ -72,6 +72,8 @@ class SimpleCPUOffloadScheduler:
72
vllm_config: VllmConfig,
73
kv_cache_config: "KVCacheConfig | None",
74
cpu_capacity_bytes: int,
75
+ scheduler_block_size: int,
76
+ hash_block_size: int,
77
lazy_offload: bool = False,
78
):
79
self.vllm_config = vllm_config
@@ -80,8 +82,9 @@ class SimpleCPUOffloadScheduler:
82
vllm_config.kv_events_config is not None
83
and vllm_config.kv_events_config.enable_kv_cache_events
84
)
83
- # NOTE: We use the same block size for both GPU and CPU.
84
- self.block_size = vllm_config.cache_config.block_size
85
+ self.block_size = scheduler_block_size
86
+ self.hash_block_size = hash_block_size
87
+ assert self.block_size % self.hash_block_size == 0
88
# Derive a CPU KVCacheConfig from the GPU config and build a coordinator
89
assert kv_cache_config is not None
90
self.cpu_kv_cache_config = self._derive_cpu_config(
@@ -95,6 +98,12 @@ class SimpleCPUOffloadScheduler:
98
self.fa_gidx = g_idx
99
break
100
assert 0 <= self.fa_gidx < len(self.cpu_kv_cache_config.kv_cache_groups)
101
+ # FA group's own block_size; divides scheduler_block_size (the LCM)
102
+ # but is NOT assumed to equal it.
103
+ self.fa_block_size: int = self.cpu_kv_cache_config.kv_cache_groups[
104
+ self.fa_gidx
105
+ ].kv_cache_spec.block_size
106
+ assert self.block_size % self.fa_block_size == 0
107
108
logger.info(
109
"SimpleCPUOffloadScheduler: Allocating %d CPU blocks (%.2f GB, mode=%s)",
@@ -118,7 +127,7 @@ class SimpleCPUOffloadScheduler:
127
enable_kv_cache_events=self.enable_kv_cache_events,
128
dcp_world_size=dcp_world_size,
129
pcp_world_size=pcp_world_size,
121
- hash_block_size=self.block_size,
130
+ hash_block_size=self.hash_block_size,
131
)
132
self.cpu_block_pool: BlockPool = self.cpu_coordinator.block_pool
133
@@ -218,19 +227,17 @@ class SimpleCPUOffloadScheduler:
227
def get_num_new_matched_tokens(
228
self, request: "Request", num_computed_tokens: int
229
) -> tuple[int | None, bool]:
221
- """Return (num_new_tokens, is_async) from consecutive CPU cache hits.
230
+ """Return (num_new_tokens, is_async) from consecutive CPU cache hits."""
231
223
- Pins found CPU blocks via touch() so they survive LRU eviction until
224
- update_state_after_alloc() consumes them. Any pin from an earlier
225
- call on the same request (e.g. retry after a failed allocate_slots)
226
- is dropped first.
227
- """
228
- stale = self._pending_cpu_hits.pop(request.request_id, None)
229
- if stale is not None:
232
+ # Pins found CPU blocks so they survive LRU eviction until
233
+ # update_state_after_alloc() consumes them. Any pin from an earlier
234
+ # call on the same request (e.g. retry after a failed allocate_slots)
235
+ # is dropped first.
236
+ if stale := self._pending_cpu_hits.pop(request.request_id, None):
237
self._free_pending_cpu_hit(stale)
238
232
- skipped = num_computed_tokens // self.block_size
233
- remaining_hashes = request.block_hashes[skipped:]
239
+ num_skipped_hashes = num_computed_tokens // self.hash_block_size
240
+ remaining_hashes = request.block_hashes[num_skipped_hashes:]
241
242
if not remaining_hashes:
243
return 0, False
@@ -306,11 +313,14 @@ class SimpleCPUOffloadScheduler:
313
314
cpu_hit_blocks_full, _ = pending
315
316
+ # ``num_external_tokens`` is LCM-aligned (checked per-group below),
317
+ # so this counts whole scheduler-aligned chunks of incoming tokens.
318
num_blocks_to_load = num_external_tokens // self.block_size
319
assert num_blocks_to_load > 0
311
-
312
- skipped = sum(blk.block_hash is not None for blk in blocks.blocks[self.fa_gidx])
313
- num_computed_tokens = skipped * self.block_size
320
+ num_cached_fa_blocks = sum(
321
+ blk.block_hash is not None for blk in blocks.blocks[self.fa_gidx]
322
+ )
323
+ num_computed_tokens = num_cached_fa_blocks * self.fa_block_size
324
325
# Build transfer pairs across all groups.
326
total_computed_tokens = num_computed_tokens + num_external_tokens
@@ -559,6 +569,8 @@ class SimpleCPUOffloadScheduler:
569
# Confirmed tokens: KV data written and visible to all streams.
570
req = state.request
571
confirmed_tokens = req.num_computed_tokens - req.num_output_placeholders
572
+ # Cap to blocks with confirmed KV data.
573
+ aligned_tokens = confirmed_tokens // self.block_size * self.block_size
574
575
for g in range(num_groups):
576
# FIXME (yifan): handle CPU cache eviction, where
@@ -567,9 +579,8 @@ class SimpleCPUOffloadScheduler:
579
already_stored_g = state.num_stored_blocks[g]
580
group_gpu_ids = block_ids_by_group[g]
581
570
- # Cap to blocks with confirmed KV data.
582
g_block_size = kv_cache_groups[g].kv_cache_spec.block_size
572
- ready_blocks_g = confirmed_tokens // g_block_size
583
+ ready_blocks_g = aligned_tokens // g_block_size
584
scannable = group_gpu_ids[already_stored_g:ready_blocks_g]
585
586
for gpu_block_id in scannable:
@@ -580,7 +591,10 @@ class SimpleCPUOffloadScheduler:
591
592
bhash_with_group = gpu_block.block_hash
593
if bhash_with_group is None:
583
- break
594
+ # Masked-out SWA position the coordinator chose not to
595
+ # hash; it can never serve a prefix-cache hit, so skip.
596
+ advanced_per_group[g] += 1
597
+ continue
598
599
# Skip if already scheduled for store or already cached in CPU.
600
if (