From 6b6a1b52ab6be33f71c5b5c0a1d0cf3349f11898 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] 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