[Bugfix][KV Offload] Compact canonical MLA rows (#56799)
Signed-off-by: Itay Etelis <92247226+Etelis@users.noreply.github.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>
Itay Etelis committed
Sep 17, 2026 at 06:24 UTC
9854b580dfe518171eed3a6ad76cbd402ca2bdfa
4 files changed
+63
-3
tests/v1/kv_connector/unit/offloading_connector/test_config.py
+40
@@ -17,6 +17,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
17
SchedulerOffloadConfig,
18
)
19
from vllm.platforms import current_platform
20
+from vllm.v1.core.kv_cache_utils import generate_scheduler_kv_cache_config
21
from vllm.v1.kv_cache_interface import (
22
FullAttentionSpec,
23
HiddenStateCacheSpec,
@@ -32,6 +33,8 @@ from vllm.v1.kv_cache_interface import (
33
UniformTypeKVCacheSpecs,
34
)
35
from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
36
+from vllm.v1.kv_offload.file_mapper import FileMapper
37
+from vllm.v1.kv_offload.tiering.spec import TieringOffloadingSpec
38
39
40
def _make_vllm_config(
@@ -876,6 +879,43 @@ def test_canonical_layout_certifies_v2_model_runner():
879
assert _parallelism_agnostic(groups, canonical=True, v2=True)
880
881
882
+@pytest.mark.parametrize("tp_size", [1, 2, 4])
883
+@pytest.mark.parametrize("scheduler", [False, True])
884
+def test_canonical_mla_dsa_rows_are_tp_independent(tp_size, scheduler):
885
+ specs = {
886
+ f"l{i}": _mla_spec(block_size=64, head_size=size, dtype=torch.uint8)
887
+ for i, size in enumerate([656] * 4 + [132] * 3)
888
+ }
889
+ uniform = UniformTypeKVCacheSpecs(block_size=64, kv_cache_specs=specs)
890
+ kv_cache_config = KVCacheConfig(
891
+ num_blocks=4,
892
+ kv_cache_tensors=[
893
+ KVCacheTensor(
894
+ size=4 * uniform.page_size_bytes,
895
+ layers=list(specs),
896
+ layer_stride=0,
897
+ block_stride=uniform.page_size_bytes,
898
+ )
899
+ ],
900
+ kv_cache_groups=[KVCacheGroupSpec(list(specs), uniform)],
901
+ )
902
+ if scheduler:
903
+ kv_cache_config = generate_scheduler_kv_cache_config([kv_cache_config])
904
+ config = _make_vllm_config(
905
+ tensor_parallel_size=tp_size,
906
+ extra_config={"canonical_layout": True, "cpu_bytes_to_use": 196608 * 8},
907
+ )
908
+ config.cache_config.kv_cache_layout = "LBHNC"
909
+ config.parallel_config.distributed_executor_backend = "mp"
910
+ offloading_config = build_offloading_config(config, kv_cache_config)
911
+ spec = TieringOffloadingSpec(offloading_config)
912
+ assert spec.kv_bytes_per_chunk == 196608
913
+ assert spec.num_chunks == 8
914
+ mapper = FileMapper.from_offloading_spec("/cache", spec, parallel_agnostic=True)
915
+ assert mapper.get_run_config()["tp_size"] == 1
916
+ assert mapper.get_run_config()["replicated_layout"] is True
917
+
918
+
919
def test_parallelism_agnostic_disabled_on_v2_model_runner():
920
config = _make_vllm_config()
921
config.use_v2_model_runner = True
tests/v1/kv_connector/unit/offloading_connector/test_worker.py
+6
-2
@@ -97,6 +97,7 @@ def _make_worker(
97
kv_cache_config: KVCacheConfig,
98
replicated_layout: bool = False,
99
rank: int = 0,
100
+ canonical_layout: bool = False,
101
):
102
"""Create an OffloadingConnectorWorker with mocked dependencies."""
103
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.worker import (
@@ -106,6 +107,7 @@ def _make_worker(
107
spec = MagicMock(spec=OffloadingSpec)
108
spec.replicated_layout = replicated_layout
109
spec.config = MagicMock()
110
+ spec.config.canonical_layout = canonical_layout
111
spec.config.parallel.rank = rank
112
spec.get_worker.return_value = MagicMock()
113
@@ -204,11 +206,13 @@ def test_prepare_store_kv_non_writer_marks_completed_without_submit():
206
assert meta.completed_jobs == {7: 1}
207
208
207
-def test_prepare_store_kv_writer_submits_store():
209
+@pytest.mark.parametrize("rank,canonical_layout", [(0, False), (0, True), (1, True)])
210
+def test_prepare_store_kv_writer_submits_store(rank, canonical_layout):
211
worker, _ = _make_worker(
212
KVCacheConfig(num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]),
213
replicated_layout=True,
211
- rank=0,
214
+ rank=rank,
215
+ canonical_layout=canonical_layout,
216
)
217
218
worker.prepare_store_kv(_store_metadata(8))
vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py
+14
@@ -233,6 +233,20 @@ def build_offloading_config(
233
and parallel_config.world_size == tp_size
234
)
235
236
+ if canonical_layout and is_parallelism_agnostic:
237
+ replicated_layout = (
238
+ all(
239
+ type(spec) is MLAAttentionSpec
240
+ for _, group in selected_groups
241
+ for spec in iter_layer_specs(group.kv_cache_spec)
242
+ )
243
+ and parallel_config.nnodes_within_dp == 1
244
+ and (
245
+ parallel_config.world_size == 1
246
+ or parallel_config.distributed_executor_backend == "mp"
247
+ )
248
+ )
249
+
250
kv_events_config = vllm_config.kv_events_config
251
cache_dtype = (
252
vllm_config.model_config.dtype
vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
+3
-1
@@ -53,7 +53,9 @@ class OffloadingConnectorWorker:
53
self.worker: OffloadingWorker | None = None
54
# Non-writers still ack: pending_count waits for world_size per job.
55
self._is_store_writer = (
56
- not self.spec.replicated_layout or self.spec.config.parallel.rank == 0
56
+ not self.spec.replicated_layout
57
+ or self.spec.config.canonical_layout
58
+ or self.spec.config.parallel.rank == 0
59
)
60
61
# job_id -> req_id for in-flight loads.