A high-throughput and memory-efficient inference and serving engine for LLMs
[KV offload] Parallel-agnostic fs-tier cache for single full-attention group (#44733)
Signed-off-by: Itay Etelis <itay.etelis@ibm.com> Co-authored-by: Itay Etelis <itay.etelis@ibm.com>
Itay Etelis committed
Jun 11, 2026 at 15:48 UTC
c3662b36ea768da448722accd108f8968eeef586
4 files changed
+102
-3
tests/v1/kv_offload/test_file_mapper.py
+86
-1
@@ -4,6 +4,14 @@
4
5
from unittest.mock import MagicMock
6
7
+import torch
8
+
9
+from vllm.v1.kv_cache_interface import (
10
+ FullAttentionSpec,
11
+ KVCacheGroupSpec,
12
+ MLAAttentionSpec,
13
+ SlidingWindowSpec,
14
+)
15
from vllm.v1.kv_offload.base import (
16
OffloadingSpec,
17
make_offload_key,
@@ -58,7 +66,7 @@ def make_mapper_from_offloading_spec(**kwargs) -> FileMapper:
66
mock_vllm_config.parallel_config.rank = kwargs.get("rank", 0)
67
68
mock_kv_cache_config = MagicMock()
61
- mock_kv_cache_config.kv_cache_groups = []
69
+ mock_kv_cache_config.kv_cache_groups = kwargs.get("kv_cache_groups", [])
70
71
mock_offloading_spec = MagicMock(spec=OffloadingSpec)
72
mock_offloading_spec.vllm_config = mock_vllm_config
@@ -69,6 +77,7 @@ def make_mapper_from_offloading_spec(**kwargs) -> FileMapper:
77
root_dir=kwargs.get("root_dir", "/tmp/cache"),
78
offloading_spec=mock_offloading_spec,
79
gpu_blocks_per_file=mock_offloading_spec.block_size_factor,
80
+ parallel_agnostic=kwargs.get("parallel_agnostic", False),
81
)
82
83
@@ -125,3 +134,79 @@ def test_get_config_file_path():
134
fm = make_mapper_from_offloading_spec()
135
config_path = fm.get_config_file_path()
136
assert config_path == f"{fm.base_path}/config.json"
137
+
138
+
139
+# ---------------------------------------------------------------------------
140
+# parallel_agnostic: honored only for a single non-MLA full-attention group
141
+# ---------------------------------------------------------------------------
142
+
143
+
144
+def _full_attention_group() -> KVCacheGroupSpec:
145
+ return KVCacheGroupSpec(
146
+ layer_names=["layer0"],
147
+ kv_cache_spec=FullAttentionSpec(
148
+ block_size=16, num_kv_heads=4, head_size=128, dtype=torch.float32
149
+ ),
150
+ )
151
+
152
+
153
+def _sliding_window_group() -> KVCacheGroupSpec:
154
+ return KVCacheGroupSpec(
155
+ layer_names=["layer0"],
156
+ kv_cache_spec=SlidingWindowSpec(
157
+ block_size=16,
158
+ num_kv_heads=4,
159
+ head_size=128,
160
+ dtype=torch.float32,
161
+ sliding_window=128,
162
+ ),
163
+ )
164
+
165
+
166
+def test_parallel_agnostic_enabled_for_single_full_attention():
167
+ # tp/rank are collapsed out of the namespace so the cache is shared
168
+ # across tensor-parallel sizes.
169
+ fm = make_mapper_from_offloading_spec(
170
+ tp_size=2,
171
+ rank=1,
172
+ kv_cache_groups=[_full_attention_group()],
173
+ parallel_agnostic=True,
174
+ )
175
+ assert fm.fields["tp_size"] == 1
176
+ assert fm.rank == 0
177
+
178
+
179
+def test_parallel_agnostic_disabled_for_multiple_groups():
180
+ # More than one KV-cache group (hybrid model) => keep per-layout namespacing.
181
+ fm = make_mapper_from_offloading_spec(
182
+ tp_size=2,
183
+ kv_cache_groups=[_full_attention_group(), _full_attention_group()],
184
+ parallel_agnostic=True,
185
+ )
186
+ assert fm.fields["tp_size"] == 2
187
+
188
+
189
+def test_parallel_agnostic_disabled_for_non_full_attention():
190
+ # Single group but not full attention (sliding window) => keep namespacing.
191
+ fm = make_mapper_from_offloading_spec(
192
+ tp_size=2,
193
+ kv_cache_groups=[_sliding_window_group()],
194
+ parallel_agnostic=True,
195
+ )
196
+ assert fm.fields["tp_size"] == 2
197
+
198
+
199
+def test_parallel_agnostic_excludes_mla():
200
+ # MLA latent KV is replicated per rank, so its offloaded blocks are not
201
+ # parallelism-invariant: the opt-in must not collapse tp/rank.
202
+ group = KVCacheGroupSpec(
203
+ layer_names=["layer0"],
204
+ kv_cache_spec=MLAAttentionSpec(
205
+ block_size=16, num_kv_heads=1, head_size=576, dtype=torch.float32
206
+ ),
207
+ )
208
+ fm = make_mapper_from_offloading_spec(
209
+ tp_size=2, rank=1, kv_cache_groups=[group], parallel_agnostic=True
210
+ )
211
+ assert fm.fields["tp_size"] == 2
212
+ assert fm.rank == 1
vllm/v1/kv_offload/file_mapper.py
+10
@@ -4,6 +4,7 @@
4
import hashlib
5
import json
6
7
+from vllm.v1.kv_cache_interface import FullAttentionSpec, MLAAttentionSpec
8
from vllm.v1.kv_offload.base import (
9
OffloadingSpec,
10
OffloadKey,
@@ -81,6 +82,15 @@ class FileMapper:
82
}
83
for group in kv_cache_config.kv_cache_groups
84
]
85
+ # Only a single full-attention group is parallelism-invariant. MLA is
86
+ # excluded: its latent KV is replicated per rank, never head-sharded.
87
+ groups = kv_cache_config.kv_cache_groups
88
+ spec = groups[0].kv_cache_spec if len(groups) == 1 else None
89
+ parallel_agnostic = (
90
+ parallel_agnostic
91
+ and isinstance(spec, FullAttentionSpec)
92
+ and not isinstance(spec, MLAAttentionSpec)
93
+ )
94
return cls(
95
root_dir=root_dir,
96
model_name=vllm_config.model_config.model,
vllm/v1/kv_offload/tiering/fs/manager.py
+2
-1
@@ -107,11 +107,12 @@ class FileSystemTierManager(SecondaryTierManager):
107
)
108
self._block_size: int = primary_kv_view.strides[0]
109
110
- # Create file mapper
110
+ # Opt in; FileMapper enables it only for a parallelism-invariant block.
111
self.file_mapper = FileMapper.from_offloading_spec(
112
root_dir=root_dir,
113
offloading_spec=offloading_spec,
114
gpu_blocks_per_file=offloading_spec.block_size_factor,
115
+ parallel_agnostic=True,
116
)
117
118
# Write config file
vllm/v1/kv_offload/tiering/obj/manager.py
+4
-1
@@ -108,7 +108,10 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):
108
self._primary_reg = None
109
self._block_size_bytes: int = 0
110
root_dir = f"{prefix}/" if prefix else ""
111
- self._file_mapper = FileMapper.from_offloading_spec(root_dir, offloading_spec)
111
+ # Opt in; FileMapper enables it only for a parallelism-invariant block.
112
+ self._file_mapper = FileMapper.from_offloading_spec(
113
+ root_dir, offloading_spec, parallel_agnostic=True
114
+ )
115
self._next_obj_dev_id: int = 1 # dev_id=0 is reserved for _exists() probes
116
117
self._probe_connectivity()