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 4, 15:00 UTC
[kv_offload+HMA][11/N]: Support store with multiple KV groups (#39403)
Signed-off-by: Or Ozeri <oro@il.ibm.com>
Or Ozeri committed
Apr 25, 2026 at 20:00 UTC
60cd878a3beca91e63d9a34a9c60fd335e780182
1 file changed
+87
-42
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+87
-42
@@ -112,6 +112,13 @@ class RequestOffloadState:
112
for group_state, new_blocks in zip(self.group_states, new_block_id_groups):
113
group_state.block_ids.extend(new_blocks)
114
115
+ def advance_stored_idx(self, num_offloadable_tokens: int) -> None:
116
+ for group_config, group_state in zip(
117
+ self.config.kv_group_configs, self.group_states
118
+ ):
119
+ num_blocks = num_offloadable_tokens // group_config.offloaded_block_size
120
+ group_state.next_stored_block_idx = num_blocks
121
+
122
123
class OffloadingConnectorScheduler:
124
"""Implementation of Scheduler side methods"""
@@ -367,16 +374,16 @@ class OffloadingConnectorScheduler:
374
if self._blocks_being_loaded is not None:
375
self._blocks_being_loaded.update(req_blocks_being_loaded)
376
370
- def _get_reqs_to_store(self, scheduler_output: SchedulerOutput):
371
- # Below assertion will be removed once this function supports HMA
372
- assert len(self.config.kv_group_configs) == 1
373
- group_config = self.config.kv_group_configs[0]
374
-
377
+ def _get_reqs_to_store(
378
+ self, scheduler_output: SchedulerOutput
379
+ ) -> dict[ReqId, TransferSpec]:
380
+ block_size_factor = self.config.block_size_factor
381
reqs_to_store: dict[ReqId, TransferSpec] = {}
382
# iterate over both new and cached requests
383
for req_id, new_block_id_groups, preempted in yield_req_data(scheduler_output):
384
req_status = self._req_status[req_id]
385
req_status.update_offload_keys()
386
+ req = req_status.req
387
388
if preempted:
389
for group_state in req_status.group_states:
@@ -385,68 +392,106 @@ class OffloadingConnectorScheduler:
392
if new_block_id_groups:
393
req_status.update_block_id_groups(new_block_id_groups)
394
388
- # Below assertion will be removed once this function supports HMA
389
- assert len(req_status.group_states) == 1
390
- group_state = req_status.group_states[0]
395
+ num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id]
396
+ num_tokens_after_batch = req.num_computed_tokens + num_scheduled_tokens
397
+ # with async scheduling, some tokens may be missing
398
+ num_offloadable_tokens = min(num_tokens_after_batch, req.num_tokens)
399
392
- block_ids = group_state.block_ids
400
+ # Filter out blocks skipped due to sliding window attention / SSM
401
+ new_offload_keys: list[OffloadKey] = []
402
+ for group_config, group_state in zip(
403
+ self.config.kv_group_configs, req_status.group_states
404
+ ):
405
+ num_blocks = num_offloadable_tokens // group_config.offloaded_block_size
406
+ start_block_idx = group_state.next_stored_block_idx
407
+ if num_blocks <= start_block_idx:
408
+ continue
409
+ offload_keys = group_state.offload_keys[start_block_idx:num_blocks]
410
+ # For each block to offload, take the last corresponding GPU block.
411
+ # e.g. if block size factor is 3 and GPU block IDs are
412
+ # 1 5 6 7 2 4 9 3 8 then we'll take blocks 6 4 8.
413
+ # We will use these GPU blocks to determine if the block needs
414
+ # offloading, or (if the GPU block ID is 0) this block should
415
+ # be skipped due to sliding window attention / SSM.
416
+ # We know that if a block is skipped, then all the previous blocks
417
+ # are skipped as well. This is why we take the last of each block.
418
+ offload_block_ids = group_state.block_ids[
419
+ start_block_idx * block_size_factor
420
+ + block_size_factor
421
+ - 1 : num_blocks * block_size_factor : block_size_factor
422
+ ]
423
+ assert len(offload_keys) == len(offload_block_ids)
424
394
- req = req_status.req
395
- new_tokens = scheduler_output.num_scheduled_tokens[req_id]
396
- expected_tokens = req.num_computed_tokens + new_tokens
397
- # with async scheduling, some tokens may be missing
398
- total_tokens = min(expected_tokens, req.num_tokens)
399
- num_blocks = total_tokens // group_config.offloaded_block_size
400
- start_block_idx = group_state.next_stored_block_idx
401
- num_new_blocks = num_blocks - start_block_idx
425
+ for offload_key, block_id in zip(offload_keys, offload_block_ids):
426
+ if block_id != 0:
427
+ new_offload_keys.append(offload_key)
428
403
- if num_new_blocks <= 0:
429
+ if not new_offload_keys:
430
+ req_status.advance_stored_idx(num_offloadable_tokens)
431
continue
432
406
- num_gpu_blocks = num_blocks * self.config.block_size_factor
407
- assert len(req.block_hashes) >= num_gpu_blocks
408
-
409
- new_offload_keys = group_state.offload_keys[start_block_idx:num_blocks]
433
store_output = self.manager.prepare_store(
434
new_offload_keys, req_status.req_context
435
)
436
if store_output is None:
414
- logger.warning(
415
- "Request %s: cannot store %s blocks", req_id, num_new_blocks
416
- )
437
+ logger.warning("Request %s: cannot store blocks", req_id)
438
continue
439
419
- group_state.next_stored_block_idx = num_blocks
420
-
440
if not store_output.keys_to_store:
441
+ req_status.advance_stored_idx(num_offloadable_tokens)
442
continue
423
- keys_to_store = set(store_output.keys_to_store)
443
425
- self.manager.touch(group_state.offload_keys[:num_blocks])
444
+ for group_state in req_status.group_states:
445
+ self.manager.touch(group_state.offload_keys)
446
427
- dst_spec = store_output.store_spec
447
+ keys_to_store = set(store_output.keys_to_store)
448
+
449
+ group_sizes: list[int] = []
450
+ block_indices: list[int] = []
451
src_block_ids: list[int] = []
429
- for idx, key in enumerate(new_offload_keys):
430
- if key not in keys_to_store:
431
- continue
432
- offloaded_block_idx = start_block_idx + idx
433
- gpu_block_idx = offloaded_block_idx * self.config.block_size_factor
434
- for i in range(self.config.block_size_factor):
435
- src_block_ids.append(block_ids[gpu_block_idx + i])
452
+ for group_config, group_state in zip(
453
+ self.config.kv_group_configs, req_status.group_states
454
+ ):
455
+ num_blocks = num_offloadable_tokens // group_config.offloaded_block_size
456
+ start_block_idx = group_state.next_stored_block_idx
457
+ block_ids = group_state.block_ids
458
+ num_group_blocks = 0
459
+ start_gpu_block_idx: int | None = None
460
+ for idx, offload_key in enumerate(
461
+ group_state.offload_keys[start_block_idx:num_blocks]
462
+ ):
463
+ if offload_key not in keys_to_store:
464
+ continue
465
+
466
+ offloaded_block_idx = start_block_idx + idx
467
+ gpu_block_idx = offloaded_block_idx * block_size_factor
468
+ num_group_blocks += block_size_factor
469
+ for i in range(block_size_factor):
470
+ block_id = block_ids[gpu_block_idx + i]
471
+ if block_id == 0:
472
+ # skipped blocks cannot appear after non-skipped blocks
473
+ assert start_gpu_block_idx is None
474
+ continue
475
+ elif start_gpu_block_idx is None:
476
+ start_gpu_block_idx = gpu_block_idx + i
477
+ src_block_ids.append(block_id)
478
+ group_sizes.append(num_group_blocks)
479
+ block_indices.append(start_gpu_block_idx or 0)
480
+ group_state.next_stored_block_idx = num_blocks
481
+
482
src_spec = GPULoadStoreSpec(
437
- src_block_ids,
438
- group_sizes=(len(src_block_ids),),
439
- block_indices=(0,),
483
+ src_block_ids, group_sizes=group_sizes, block_indices=block_indices
484
)
485
+ dst_spec = store_output.store_spec
486
487
reqs_to_store[req_id] = (src_spec, dst_spec)
488
self._reqs_being_stored[req_id] |= keys_to_store
489
490
logger.debug(
446
- "Request %s offloading %s blocks starting from block #%d",
491
+ "Request %s offloading %s blocks upto %d tokens",
492
req_id,
493
len(keys_to_store),
449
- start_block_idx,
494
+ num_offloadable_tokens,
495
)
496
497
return reqs_to_store