fix(kv-cache): Support external loads for lockstep MLA groups

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

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