From 5ef33ab250b3904da375ecb18bdda00a4a73c3a8 Mon Sep 17 00:00:00 2001 From: Or Ozeri Date: Thu, 23 Apr 2026 20:00:45 +0300 Subject: [PATCH] [kv_offload+HMA][10/N]: Support load with multiple KV groups (#39402) Signed-off-by: Or Ozeri --- .../kv_connector/v1/offloading/scheduler.py | 78 ++++++++++++------- 1 file changed, 50 insertions(+), 28 deletions(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index bff512815a6..5cee750811e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -14,6 +14,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import ( ReqId, ) from vllm.logger import init_logger +from vllm.utils.math_utils import cdiv from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_offload.abstract import ( @@ -271,45 +272,66 @@ class OffloadingConnectorScheduler: return req_status = self._req_status[request.request_id] - block_groups = blocks.get_block_ids() - # Below assertions will be removed once this function supports HMA - assert len(self.config.kv_group_configs) == 1 - assert len(req_status.group_states) == 1 - assert len(block_groups) == 1 - block_ids = block_groups[0] - group_config = self.config.kv_group_configs[0] - group_state = req_status.group_states[0] + num_locally_computed_tokens = req_status.num_locally_computed_tokens + num_cached_tokens = num_locally_computed_tokens + num_external_tokens - num_computed_gpu_blocks = sum( - block.block_hash is not None for block in blocks.blocks[0] - ) - num_computed_tokens = num_computed_gpu_blocks * group_config.gpu_block_size - full_block_tokens = num_computed_tokens + num_external_tokens - assert full_block_tokens % group_config.offloaded_block_size == 0 + keys_to_load: list[OffloadKey] = [] + dst_block_ids: list[int] = [] + # per group + group_sizes: list[int] = [] + block_indices: list[int] = [] + for group_config, group_state, group_blocks in zip( + self.config.kv_group_configs, + req_status.group_states, + blocks.blocks, + ): + gpu_block_size = group_config.gpu_block_size + offloaded_block_size = group_config.offloaded_block_size + offload_keys = group_state.offload_keys + num_gpu_blocks = cdiv(num_cached_tokens, gpu_block_size) - num_pending_gpu_blocks = len(block_ids) - num_computed_gpu_blocks - assert ( - num_external_tokens == num_pending_gpu_blocks * group_config.gpu_block_size - ) + assert len(group_blocks) >= num_gpu_blocks + num_locally_computed_gpu_blocks = num_gpu_blocks + # Skip null placeholder blocks (used for sliding window or mamba padding). + for i, block in enumerate(group_blocks[:num_gpu_blocks]): + if not block.is_null and block.block_hash is None: + num_locally_computed_gpu_blocks = i + break - start_block_idx = num_computed_tokens // group_config.offloaded_block_size - num_blocks = full_block_tokens // group_config.offloaded_block_size + assert ( + num_locally_computed_tokens + <= num_locally_computed_gpu_blocks * gpu_block_size + ) + num_pending_gpu_blocks = num_gpu_blocks - num_locally_computed_gpu_blocks - assert len(request.block_hashes) // self.config.block_size_factor >= num_blocks - offload_keys = group_state.offload_keys[start_block_idx:num_blocks] + num_blocks = cdiv(num_cached_tokens, offloaded_block_size) + assert len(offload_keys) >= num_blocks + if num_pending_gpu_blocks: + start_block_idx = ( + num_locally_computed_gpu_blocks // self.config.block_size_factor + ) + keys_to_load.extend(offload_keys[start_block_idx:num_blocks]) - src_spec = self.manager.prepare_load(offload_keys, req_status.req_context) + dst_block_ids.extend( + block.block_id + for block in group_blocks[ + num_locally_computed_gpu_blocks:num_gpu_blocks + ] + ) + group_sizes.append(num_pending_gpu_blocks) + block_indices.append(num_locally_computed_gpu_blocks) + + group_state.next_stored_block_idx = num_blocks + + src_spec = self.manager.prepare_load(keys_to_load, req_status.req_context) dst_spec = GPULoadStoreSpec( - block_ids[num_computed_gpu_blocks:], - group_sizes=(num_pending_gpu_blocks,), - block_indices=(num_computed_gpu_blocks,), + dst_block_ids, group_sizes=group_sizes, block_indices=block_indices ) self._reqs_to_load[request.request_id] = (src_spec, dst_spec) req_blocks_being_loaded = self._reqs_being_loaded[request.request_id] - req_blocks_being_loaded.update(offload_keys) - group_state.next_stored_block_idx = num_blocks + req_blocks_being_loaded.update(keys_to_load) if self._blocks_being_loaded is not None: self._blocks_being_loaded.update(req_blocks_being_loaded)