From 05753fcc646ca4c729b11fd7bb0499d24ea9c145 Mon Sep 17 00:00:00 2001
From: Lucas Wilkinson <lwilkinson@neuralmagic.com>
Date: Wed, 26 Aug 2026 08:19:13 -0400
Subject: [PATCH] [KV offload] Scale uniform-type attention groups by DCP in
 build_offloading_config

`build_offloading_config` scales a group's `tokens_per_block` by
`decode_context_parallel_size` only when the group spec is a bare
`AttentionSpec`. Models whose attention layers are merged into a single
`UniformTypeKVCacheSpecs` group (GLM-5.2 / DeepSeek-V3.2: MLA layers plus
the sparse-indexer K cache) are block-sharded across DCP ranks exactly like
a plain attention group, but skipped the multiplier. Under DCP the hash
block size from `resolve_kv_cache_block_sizes` is `block_size * dcp`, so the
OffloadingConnector failed at startup with

    AssertionError: tokens_per_block=64 not divisible by tokens_per_hash=512.

Treat a group as DCP-sharded when every layer spec it covers is an
`AttentionSpec` (via `iter_layer_specs`), which keeps the Mamba behaviour
unchanged and matches how the scheduler block size is resolved.

Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>
---
 .../unit/offloading_connector/test_config.py  | 33 +++++++++++++++++++
 .../kv_connector/v1/offloading/config.py      | 15 ++++++++-
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py
index afb0a9bda54e..be567a74f0fd 100644
--- a/tests/v1/kv_connector/unit/offloading_connector/test_config.py
+++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py
@@ -358,6 +358,39 @@ def test_dcp_scales_attention_but_not_mamba_group_blocks():
     ] == [1, 3]
 
 
+def test_dcp_scales_uniform_type_attention_group_blocks():
+    # A single UniformTypeKVCacheSpecs group of attention-only layers (e.g.
+    # GLM-5.2 / DeepSeek-V3.2 MLA + sparse-indexer layers) is DCP-sharded
+    # like a bare attention group, so its logical block must scale by DCP
+    # to stay divisible by the DCP-scaled hash block size.
+    config = _make_vllm_config(tensor_parallel_size=2, decode_context_parallel_size=2)
+    config.speculative_config = None
+    layer_specs = {
+        "mla_layer": _mla_spec(head_size=512),
+        "indexer_layer": _mla_spec(head_size=128),
+    }
+    uniform_spec = UniformTypeKVCacheSpecs(block_size=16, kv_cache_specs=layer_specs)
+    num_blocks = 4
+    kv_cache_config = KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=spec.page_size_bytes * num_blocks,
+                layers=[name],
+                layer_stride=spec.page_size_bytes,
+                block_stride=spec.page_size_bytes,
+            )
+            for name, spec in layer_specs.items()
+        ],
+        kv_cache_groups=[KVCacheGroupSpec(list(layer_specs), uniform_spec)],
+    )
+
+    offloading_config = build_offloading_config(config, kv_cache_config)
+
+    assert tuple(group.tokens_per_block for group in offloading_config.groups) == (32,)
+    assert offloading_config.cache.tokens_per_hash == 32
+
+
 def test_preserves_data_parallel_config():
     config = _make_vllm_config()
     config.parallel_config.data_parallel_index = 2
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 70919ce24e75..f348d31735ef 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py
@@ -8,7 +8,9 @@
 from vllm.v1.kv_cache_interface import (
     AttentionSpec,
     FullAttentionSpec,
+    KVCacheSpec,
     MLAAttentionSpec,
+    iter_layer_specs,
 )
 from vllm.v1.kv_offload.config import (
     OffloadingCacheConfig,
@@ -23,6 +25,17 @@
     from vllm.v1.kv_cache_interface import KVCacheConfig
 
 
+def _is_dcp_sharded(spec: "KVCacheSpec") -> bool:
+    """Attention groups are block-sharded across DCP ranks, so a logical
+    (scheduler/hash) block spans ``block_size * dcp`` tokens. A uniform-type
+    group made only of attention specs (e.g. MLA + sparse indexer layers) is
+    sharded the same way; anything else keeps its full per-rank block."""
+    layer_specs = iter_layer_specs(spec)
+    return len(layer_specs) > 0 and all(
+        isinstance(member, AttentionSpec) for member in layer_specs
+    )
+
+
 def build_offloading_config(
     vllm_config: "VllmConfig",
     kv_cache_config: "KVCacheConfig",
@@ -41,7 +54,7 @@ def build_offloading_config(
                 group.kv_cache_spec.block_size
                 * (
                     parallel_config.decode_context_parallel_size
-                    if isinstance(group.kv_cache_spec, AttentionSpec)
+                    if _is_dcp_sharded(group.kv_cache_spec)
                     else 1
                 )
             ),