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, 17:30 UTC
[kv_offload+HMA][10/N]: Support load with multiple KV groups (#39402)
Signed-off-by: Or Ozeri <oro@il.ibm.com>
Or Ozeri committed
Apr 23, 2026 at 20:00 UTC
5ef33ab250b3904da375ecb18bdda00a4a73c3a8
1 file changed
+51
-29
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+51
-29
@@ -14,6 +14,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
14
ReqId,
15
)
16
from vllm.logger import init_logger
17
+from vllm.utils.math_utils import cdiv
18
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
19
from vllm.v1.core.sched.output import SchedulerOutput
20
from vllm.v1.kv_offload.abstract import (
@@ -271,45 +272,66 @@ class OffloadingConnectorScheduler:
272
return
273
274
req_status = self._req_status[request.request_id]
274
- block_groups = blocks.get_block_ids()
275
276
- # Below assertions will be removed once this function supports HMA
277
- assert len(self.config.kv_group_configs) == 1
278
- assert len(req_status.group_states) == 1
279
- assert len(block_groups) == 1
280
- block_ids = block_groups[0]
281
- group_config = self.config.kv_group_configs[0]
282
- group_state = req_status.group_states[0]
283
-
284
- num_computed_gpu_blocks = sum(
285
- block.block_hash is not None for block in blocks.blocks[0]
286
- )
287
- num_computed_tokens = num_computed_gpu_blocks * group_config.gpu_block_size
288
- full_block_tokens = num_computed_tokens + num_external_tokens
289
- assert full_block_tokens % group_config.offloaded_block_size == 0
276
+ num_locally_computed_tokens = req_status.num_locally_computed_tokens
277
+ num_cached_tokens = num_locally_computed_tokens + num_external_tokens
278
+
279
+ keys_to_load: list[OffloadKey] = []
280
+ dst_block_ids: list[int] = []
281
+ # per group
282
+ group_sizes: list[int] = []
283
+ block_indices: list[int] = []
284
+ for group_config, group_state, group_blocks in zip(
285
+ self.config.kv_group_configs,
286
+ req_status.group_states,
287
+ blocks.blocks,
288
+ ):
289
+ gpu_block_size = group_config.gpu_block_size
290
+ offloaded_block_size = group_config.offloaded_block_size
291
+ offload_keys = group_state.offload_keys
292
+ num_gpu_blocks = cdiv(num_cached_tokens, gpu_block_size)
293
+
294
+ assert len(group_blocks) >= num_gpu_blocks
295
+ num_locally_computed_gpu_blocks = num_gpu_blocks
296
+ # Skip null placeholder blocks (used for sliding window or mamba padding).
297
+ for i, block in enumerate(group_blocks[:num_gpu_blocks]):
298
+ if not block.is_null and block.block_hash is None:
299
+ num_locally_computed_gpu_blocks = i
300
+ break
301
+
302
+ assert (
303
+ num_locally_computed_tokens
304
+ <= num_locally_computed_gpu_blocks * gpu_block_size
305
+ )
306
+ num_pending_gpu_blocks = num_gpu_blocks - num_locally_computed_gpu_blocks
307
291
- num_pending_gpu_blocks = len(block_ids) - num_computed_gpu_blocks
292
- assert (
293
- num_external_tokens == num_pending_gpu_blocks * group_config.gpu_block_size
294
- )
308
+ num_blocks = cdiv(num_cached_tokens, offloaded_block_size)
309
+ assert len(offload_keys) >= num_blocks
310
+ if num_pending_gpu_blocks:
311
+ start_block_idx = (
312
+ num_locally_computed_gpu_blocks // self.config.block_size_factor
313
+ )
314
+ keys_to_load.extend(offload_keys[start_block_idx:num_blocks])
315
296
- start_block_idx = num_computed_tokens // group_config.offloaded_block_size
297
- num_blocks = full_block_tokens // group_config.offloaded_block_size
316
+ dst_block_ids.extend(
317
+ block.block_id
318
+ for block in group_blocks[
319
+ num_locally_computed_gpu_blocks:num_gpu_blocks
320
+ ]
321
+ )
322
+ group_sizes.append(num_pending_gpu_blocks)
323
+ block_indices.append(num_locally_computed_gpu_blocks)
324
299
- assert len(request.block_hashes) // self.config.block_size_factor >= num_blocks
300
- offload_keys = group_state.offload_keys[start_block_idx:num_blocks]
325
+ group_state.next_stored_block_idx = num_blocks
326
302
- src_spec = self.manager.prepare_load(offload_keys, req_status.req_context)
327
+ src_spec = self.manager.prepare_load(keys_to_load, req_status.req_context)
328
dst_spec = GPULoadStoreSpec(
304
- block_ids[num_computed_gpu_blocks:],
305
- group_sizes=(num_pending_gpu_blocks,),
306
- block_indices=(num_computed_gpu_blocks,),
329
+ dst_block_ids, group_sizes=group_sizes, block_indices=block_indices
330
)
331
332
self._reqs_to_load[request.request_id] = (src_spec, dst_spec)
333
req_blocks_being_loaded = self._reqs_being_loaded[request.request_id]
311
- req_blocks_being_loaded.update(offload_keys)
312
- group_state.next_stored_block_idx = num_blocks
334
+ req_blocks_being_loaded.update(keys_to_load)
335
336
if self._blocks_being_loaded is not None:
337
self._blocks_being_loaded.update(req_blocks_being_loaded)