From 25a0ae51d6121a37aafbc8298d91a09ac26f5920 Mon Sep 17 00:00:00 2001 From: Procr <193802945+procr1337@users.noreply.github.com> Date: Sun, 26 Jul 2026 20:54:34 +0200 Subject: [PATCH 1/2] fix(kv-cache): Support external loads for lockstep MLA groups Signed-off-by: Procr <193802945+procr1337@users.noreply.github.com> --- vllm/v1/core/kv_cache_coordinator.py | 39 +++++++++++++++++++++------- 1 file changed, 30 insertions(+), 9 deletions(-) diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index efc22ad7c9e1..b32419acbe69 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -220,11 +220,6 @@ def allocate_new_computed_blocks( num_local_computed_tokens: The number of local computed tokens. num_external_computed_tokens: The number of external computed tokens. """ - if self.lockstep_mla_allocations and num_external_computed_tokens: - raise NotImplementedError( - "External KV loads are not supported with lockstep MLA allocations." - ) - # A running request is already tracked in num_cached_block and won't # have new prefix-cache hits, so this is a no-op for it. if any( @@ -246,12 +241,38 @@ def allocate_new_computed_blocks( num_external_computed_tokens, ) if num_external_computed_tokens > 0: - for manager in self.single_type_managers: - manager.allocate_external_computed_blocks( + if not self.lockstep_mla_allocations: + for manager in self.single_type_managers: + manager.allocate_external_computed_blocks( + request_id, + num_local_computed_tokens, + num_external_computed_tokens, + ) + else: + managers = self.single_type_managers + assert all( + manager.block_size == managers[0].block_size + for manager in managers[1:] + ), "Lockstep MLA managers must use one effective block size" + + # Every lockstep group is MLA (i.e. full attention), so no group + # skips tokens and every manager would allocate the same blocks. + # Allocate each physical destination block exactly once: the + # connector loads every logical KV group into the same block + # index, but into that group's distinct KV tensors. + old_len = len(managers[0].req_to_blocks[request_id]) + managers[0].allocate_external_computed_blocks( request_id, num_local_computed_tokens, num_external_computed_tokens, ) + shared_external_blocks = managers[0].req_to_blocks[request_id][old_len:] + + # One request reference is required per logical manager. The + # primary allocation owns the first; touch once for each peer. + for manager in managers[1:]: + self.block_pool.touch(shared_external_blocks) + manager.req_to_blocks[request_id].extend(shared_external_blocks) if self.lockstep_mla_allocations: group_blocks = [ @@ -799,10 +820,10 @@ def find_longest_cache_hit( sparse-retention group has not cached yet (0 unless hybrid). """ if self.disable_prefix_cache_for_dcp_hybrid: - blocks: tuple[list[KVCacheBlock], ...] = tuple( + empty_blocks: tuple[list[KVCacheBlock], ...] = tuple( [] for _ in range(len(self.kv_cache_config.kv_cache_groups)) ) - return blocks, 0, 0 + return empty_blocks, 0, 0 num_groups = len(self.kv_cache_config.kv_cache_groups) hit_length = max_cache_hit_length From 7cc9963db08bf1b256adebc7d86d8ae2175b0615 Mon Sep 17 00:00:00 2001 From: Procr <193802945+procr1337@users.noreply.github.com> Date: Mon, 27 Jul 2026 08:07:26 +0000 Subject: [PATCH 2/2] fix(kv-offload): don't CP-scale tokens_per_block for replicated groups build_offloading_config() multiplied every group's block size by the context-parallel factor unconditionally. That was correct when all KV cache groups were CP-sharded, but dcp_replicated groups keep the full cache on every rank, so one of their blocks covers exactly block_size global tokens -- which is why SingleTypeKVCacheManager.__init__ excludes them from the same scaling. The connector therefore saw G * cp where the manager uses G. Since num_gpu_blocks, tokens_per_chunk and hashes_per_chunk all derive from tokens_per_block, the affected group's block and chunk counts came out a factor of cp too small and its offload keys landed on the wrong token boundaries, silently corrupting that group's KV on a load hit. Mirror the manager's rule so both agree on the effective block size. Affects any cp > 1 run combining the offloading connector with a replicated group: DFlash draft groups, the DeepSeek V3.2 sparse indexer, and the DeepSeek V4 compressor. Signed-off-by: Procr <193802945+procr1337@users.noreply.github.com> --- .../kv_connector/v1/offloading/config.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py index 5cf4dffd1af6..f10654c493d1 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py @@ -16,7 +16,7 @@ if TYPE_CHECKING: from vllm.config import VllmConfig - from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheTensor + from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec, KVCacheTensor def is_kv_cache_tensor_packed(kv_cache_tensor: "KVCacheTensor") -> bool: @@ -40,9 +40,23 @@ def build_offloading_config( parallel_config.decode_context_parallel_size * parallel_config.prefill_context_parallel_size ) + + def _tokens_per_block(kv_cache_spec: "KVCacheSpec") -> int: + # Must match `SingleTypeKVCacheManager.__init__`'s effective block size, + # since the offloading key and chunk math indexes the same scheduler + # blocks. Under (D)CP a sharded group's local block covers + # `block_size * cp` tokens of the global sequence, but a + # `dcp_replicated` group keeps the full cache on every rank, so one of + # its blocks covers exactly `block_size` global tokens and must not be + # scaled. Mixed sharded/replicated groups occur with lockstep MLA (an + # MLA cache plus a replicated sparse-indexer cache). + if getattr(kv_cache_spec, "dcp_replicated", False): + return kv_cache_spec.block_size + return kv_cache_spec.block_size * context_parallel_factor + groups = tuple( OffloadingGroupConfig( - tokens_per_block=(group.kv_cache_spec.block_size * context_parallel_factor), + tokens_per_block=_tokens_per_block(group.kv_cache_spec), layer_names=tuple(group.layer_names), ) for group in kv_cache_config.kv_cache_groups