@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7cc302dd87
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 5, 22:15 UTC

[kv_offload+HMA][7/N]: Support register_kv_caches for hybrid models (#37853)

Signed-off-by: Or Ozeri <oro@il.ibm.com>

Or Ozeri committed Mar 27, 2026 at 08:38 UTC 7cc302dd87cba546e1cf9b9967317437142438cb
13 files changed +1477 -759
tests/v1/kv_connector/unit/offloading_connector/__init__.py
tests/v1/kv_connector/unit/offloading_connector/conftest.py new
+7
@@ -0,0 +1,7 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from tests.v1.kv_connector.unit.offloading_connector.utils import (
4 + request_runner,
5 +)
6 +
7 +__all__ = ["request_runner"]
tests/v1/kv_connector/unit/offloading_connector/test_metrics.py new
+151
@@ -0,0 +1,151 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
4 + OffloadingConnectorStats,
5 +)
6 +from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
7 + OffloadingConnector,
8 +)
9 +
10 +
11 +def test_build_kv_connector_stats_with_none():
12 + """Test that build_kv_connector_stats returns empty stats when given None."""
13 + stats = OffloadingConnector.build_kv_connector_stats(data=None)
14 +
15 + assert stats is not None
16 + assert isinstance(stats, OffloadingConnectorStats)
17 + assert len(stats.data) == 0
18 + assert stats.is_empty()
19 +
20 +
21 +def test_build_kv_connector_stats_with_empty_dict():
22 + """Test that build_kv_connector_stats returns empty stats with empty dict."""
23 + stats = OffloadingConnector.build_kv_connector_stats(data={})
24 +
25 + assert stats is not None
26 + assert isinstance(stats, OffloadingConnectorStats)
27 + assert len(stats.data) == 0
28 + assert stats.is_empty()
29 +
30 +
31 +def test_build_kv_connector_stats_reconstructs_offload_stats():
32 + """Test that OffloadingConnector stats are properly reconstructed with
33 + correct data."""
34 + serialized_data = {
35 + "CPU_to_GPU": [
36 + {"op_size": 16, "op_time": 1.0},
37 + {"op_size": 8, "op_time": 0.5},
38 + ],
39 + "GPU_to_CPU": [
40 + {"op_size": 1, "op_time": 0.1},
41 + {"op_size": 2, "op_time": 0.2},
42 + ],
43 + }
44 +
45 + stats = OffloadingConnector.build_kv_connector_stats(data=serialized_data)
46 +
47 + offload_connector_stats = stats
48 + assert isinstance(offload_connector_stats, OffloadingConnectorStats)
49 + assert offload_connector_stats.data["CPU_to_GPU"] == [
50 + {"op_size": 16, "op_time": 1.0},
51 + {"op_size": 8, "op_time": 0.5},
52 + ]
53 + assert offload_connector_stats.data["GPU_to_CPU"] == [
54 + {"op_size": 1, "op_time": 0.1},
55 + {"op_size": 2, "op_time": 0.2},
56 + ]
57 +
58 +
59 +def test_aggregate_same_connector():
60 + """Test aggregating stats from the same connector type."""
61 + stats1 = OffloadingConnectorStats(
62 + data={
63 + "CPU_to_GPU": [
64 + {"op_size": 16, "op_time": 1.0},
65 + {"op_size": 8, "op_time": 0.5},
66 + ],
67 + "GPU_to_CPU": [
68 + {"op_size": 1, "op_time": 0.1},
69 + {"op_size": 2, "op_time": 0.2},
70 + ],
71 + }
72 + )
73 +
74 + stats2 = OffloadingConnectorStats(
75 + data={
76 + "CPU_to_GPU": [
77 + {"op_size": 3, "op_time": 0.2},
78 + {"op_size": 7, "op_time": 0.9},
79 + ],
80 + "GPU_to_CPU": [{"op_size": 16, "op_time": 2}],
81 + }
82 + )
83 +
84 + result = stats1.aggregate(stats2)
85 +
86 + assert result is stats1 # Should return self
87 + offload_connector_stats = result
88 + assert offload_connector_stats.data["CPU_to_GPU"] == [
89 + {"op_size": 16, "op_time": 1.0},
90 + {"op_size": 8, "op_time": 0.5},
91 + {"op_size": 3, "op_time": 0.2},
92 + {"op_size": 7, "op_time": 0.9},
93 + ]
94 + assert offload_connector_stats.data["GPU_to_CPU"] == [
95 + {"op_size": 1, "op_time": 0.1},
96 + {"op_size": 2, "op_time": 0.2},
97 + {"op_size": 16, "op_time": 2},
98 + ]
99 +
100 +
101 +def test_reduce():
102 + """Test that reduce() correctly reduces all nested connector stats."""
103 + stats = OffloadingConnectorStats(
104 + data={
105 + "CPU_to_GPU": [
106 + {"op_size": 16, "op_time": 1.0},
107 + {"op_size": 8, "op_time": 0.5},
108 + {"op_size": 3, "op_time": 0.2},
109 + {"op_size": 7, "op_time": 0.9},
110 + ],
111 + "GPU_to_CPU": [
112 + {"op_size": 1, "op_time": 0.1},
113 + {"op_size": 2, "op_time": 0.2},
114 + {"op_size": 16, "op_time": 2},
115 + ],
116 + }
117 + )
118 +
119 + reduced = stats.reduce()
120 +
121 + assert isinstance(reduced, dict)
122 + # Check that the stats were reduced (should have aggregated values)
123 + assert "CPU_to_GPU_total_bytes" in reduced
124 + assert "CPU_to_GPU_total_time" in reduced
125 + assert "GPU_to_CPU_total_bytes" in reduced
126 + assert "GPU_to_CPU_total_time" in reduced
127 + assert reduced["CPU_to_GPU_total_bytes"] == 34
128 + assert reduced["CPU_to_GPU_total_time"] == 2.6
129 + assert reduced["GPU_to_CPU_total_time"] == 2.3
130 + assert reduced["GPU_to_CPU_total_bytes"] == 19
131 +
132 +
133 +def test_reset():
134 + """Test that reset() resets all nested connector stats."""
135 + offload_connector_stats = OffloadingConnectorStats(
136 + data={
137 + "CPU_to_GPU": [
138 + {"op_size": 3, "op_time": 0.2},
139 + {"op_size": 7, "op_time": 0.9},
140 + ],
141 + "GPU_to_CPU": [{"op_size": 16, "op_time": 2}],
142 + }
143 + )
144 +
145 + assert not offload_connector_stats.is_empty()
146 +
147 + offload_connector_stats.reset()
148 +
149 + # After reset, stats should be empty
150 + assert offload_connector_stats.is_empty()
151 + assert len(offload_connector_stats.data) == 0
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py new
+341
@@ -0,0 +1,341 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from collections.abc import Iterable
4 +
5 +import pytest
6 +
7 +from tests.v1.kv_connector.unit.offloading_connector.utils import (
8 + generate_store_output,
9 +)
10 +from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID
11 +from vllm.distributed.kv_events import BlockRemoved, BlockStored
12 +from vllm.v1.core.kv_cache_utils import BlockHash
13 +from vllm.v1.kv_offload.abstract import OffloadingEvent
14 +from vllm.v1.request import RequestStatus
15 +
16 +
17 +@pytest.mark.parametrize("async_scheduling", [True, False])
18 +def test_offloading_connector(request_runner, async_scheduling: bool):
19 + offloaded_block_size = 12
20 + gpu_block_size = 4
21 + num_gpu_blocks = 100
22 + block_size_factor = offloaded_block_size // gpu_block_size
23 +
24 + runner = request_runner(
25 + offloaded_block_size=offloaded_block_size,
26 + gpu_block_size=gpu_block_size,
27 + num_gpu_blocks=num_gpu_blocks,
28 + async_scheduling=async_scheduling,
29 + )
30 +
31 + # 3 blocks, store just the middle block (skip first and last)
32 + # blocks = [0, 1, 2], [3, 4, 5], [6, 7, 8]
33 + runner.new_request(token_ids=[0] * offloaded_block_size * 3)
34 + runner.manager.prepare_store.side_effect = (
35 + lambda block_hashes: generate_store_output(list(block_hashes)[1:2])
36 + )
37 + runner.run(decoded_tokens=[0])
38 +
39 + # add block missing 1 token -> no offload
40 + runner.run(
41 + decoded_tokens=[0] * (offloaded_block_size - 1),
42 + expected_stored_gpu_block_indexes=(3, 4, 5),
43 + )
44 + runner.manager.prepare_store.assert_not_called()
45 +
46 + # +1 token -> single block, fail prepare_store
47 + runner.manager.prepare_store.side_effect = lambda block_hashes: None
48 + runner.run(decoded_tokens=[0])
49 + runner.manager.prepare_store.assert_called()
50 +
51 + # 1 more block (+ token for async scheduling)
52 + # now set block_hashes_to_store = []
53 + runner.manager.prepare_store.side_effect = (
54 + lambda block_hashes: generate_store_output([])
55 + )
56 + runner.run(decoded_tokens=[0] * (offloaded_block_size + 1))
57 +
58 + # 1 more block (+ token for kicking off offloading)
59 + # now check touch was called with all 6 blocks
60 + runner.manager.prepare_store.side_effect = (
61 + lambda block_hashes: generate_store_output(block_hashes)
62 + )
63 + runner.run(
64 + decoded_tokens=[0] * (offloaded_block_size + 1),
65 + expected_stored_gpu_block_indexes=(15, 16, 17),
66 + )
67 + runner.manager.touch.assert_called()
68 + block_hashes1 = list(runner.manager.touch.call_args.args[0])
69 + assert len(block_hashes1) == 6
70 +
71 + # terminate request
72 + runner.run(decoded_tokens=[EOS_TOKEN_ID])
73 +
74 + # create a new request differing only on the last token
75 + runner.new_request(token_ids=[0] * (offloaded_block_size * 6 - 1) + [1])
76 + runner.run(decoded_tokens=[0])
77 + runner.manager.touch.assert_called()
78 + block_hashes2 = list(runner.manager.touch.call_args.args[0])
79 + assert len(block_hashes2) == 6
80 +
81 + # verify hashes are the same, except for the last block
82 + assert block_hashes1[:5] == block_hashes2[:5]
83 + assert block_hashes1[5] != block_hashes2[5]
84 +
85 + # terminate request
86 + runner.run(
87 + decoded_tokens=[EOS_TOKEN_ID],
88 + expected_stored_gpu_block_indexes=tuple(range(6 * block_size_factor)),
89 + )
90 +
91 + # full_block_tokens - num_computed_tokens < offloaded_block_size
92 + runner.new_request(
93 + token_ids=[0] * gpu_block_size + [1] * (offloaded_block_size - gpu_block_size)
94 + )
95 + runner.manager.prepare_store.side_effect = (
96 + lambda block_hashes: generate_store_output([])
97 + )
98 + runner.run(decoded_tokens=[EOS_TOKEN_ID])
99 + runner.manager.lookup.assert_not_called()
100 +
101 + # single block lookup with no hits
102 + runner.new_request(token_ids=[1] * offloaded_block_size)
103 + runner.manager.prepare_store.side_effect = (
104 + lambda block_hashes: generate_store_output([])
105 + )
106 + runner.run(decoded_tokens=[EOS_TOKEN_ID])
107 + runner.manager.lookup.assert_called()
108 + assert len(list(runner.manager.lookup.call_args.args[0])) == 1
109 +
110 + # single block lookup with a hit
111 + runner.scheduler.reset_prefix_cache()
112 + runner.new_request(token_ids=[0] * offloaded_block_size)
113 + runner.manager.prepare_store.side_effect = (
114 + lambda block_hashes: generate_store_output([])
115 + )
116 + runner.manager.lookup.return_value = 1
117 + runner.run(
118 + decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(0, 1, 2)
119 + )
120 +
121 + # single block lookup with a hit in a middle block
122 + runner.new_request(
123 + token_ids=[0] * offloaded_block_size * 2 + [1] * offloaded_block_size
124 + )
125 + runner.manager.prepare_store.side_effect = (
126 + lambda block_hashes: generate_store_output([])
127 + )
128 + runner.manager.lookup.return_value = 1
129 + runner.run(
130 + decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(3, 4, 5)
131 + )
132 +
133 + # test take_events
134 + def to_hashes(int_hashes: list[int]) -> list[BlockHash]:
135 + return [BlockHash(str(i).encode()) for i in int_hashes]
136 +
137 + def take_events() -> Iterable[OffloadingEvent]:
138 + yield OffloadingEvent(
139 + block_hashes=to_hashes([1, 2, 3]), block_size=16, medium="A", removed=False
140 + )
141 + yield OffloadingEvent(
142 + block_hashes=to_hashes([4, 5, 6]), block_size=32, medium="B", removed=True
143 + )
144 +
145 + runner.manager.take_events.side_effect = take_events
146 + events = list(runner.scheduler_connector.take_events())
147 + assert len(events) == 2
148 + event = events[0]
149 + assert isinstance(event, BlockStored)
150 + assert event.block_hashes == to_hashes([1, 2, 3])
151 + assert event.block_size == 16
152 + assert event.medium == "A"
153 + assert event.token_ids == []
154 + assert event.parent_block_hash is None
155 + assert event.lora_id is None
156 + assert event.lora_name is None
157 + event = events[1]
158 + assert isinstance(event, BlockRemoved)
159 + assert event.block_hashes == to_hashes([4, 5, 6])
160 + assert event.medium == "B"
161 +
162 +
163 +@pytest.mark.parametrize("async_scheduling", [True, False])
164 +def test_request_preemption(request_runner, async_scheduling: bool):
165 + offloaded_block_size = 12
166 + gpu_block_size = 4
167 + num_gpu_blocks = 100
168 +
169 + runner = request_runner(
170 + offloaded_block_size=offloaded_block_size,
171 + gpu_block_size=gpu_block_size,
172 + num_gpu_blocks=num_gpu_blocks,
173 + async_scheduling=async_scheduling,
174 + )
175 +
176 + free_block_queue = runner.scheduler.kv_cache_manager.block_pool.free_block_queue
177 + num_free_blocks_empty = free_block_queue.num_free_blocks
178 +
179 + # 2 blocks, store all, without flushing
180 + # blocks = [0, 1, 2], [3, 4, 5]
181 + runner.new_request(token_ids=[0] * offloaded_block_size * 2)
182 + runner.manager.prepare_store.side_effect = (
183 + lambda block_hashes: generate_store_output(block_hashes)
184 + )
185 + runner.run(
186 + decoded_tokens=[0],
187 + complete_transfers=False,
188 + )
189 +
190 + # decode 2 more blocks - 1 gpu block, storing [6, 7, 8] (no flush)
191 + runner.manager.prepare_store.side_effect = (
192 + lambda block_hashes: generate_store_output(block_hashes)
193 + )
194 + runner.run(
195 + decoded_tokens=[0] * (2 * offloaded_block_size - gpu_block_size),
196 + complete_transfers=False,
197 + )
198 +
199 + # simulate KV cache running out of space
200 + free_block_queue.num_free_blocks = 0
201 +
202 + # request should be preempted now
203 + runner.run(
204 + decoded_tokens=[],
205 + complete_transfers=False,
206 + expected_flushed_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
207 + expected_stored_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
208 + )
209 +
210 + # restore KV cache space and reset GPU prefix cache
211 + free_block_queue.num_free_blocks = num_free_blocks_empty
212 + runner.scheduler.reset_prefix_cache()
213 +
214 + # request should now return from preemption
215 + # re-load [0, ..., 8] from the CPU and store [9, 10, 11]
216 + runner.manager.lookup.return_value = 3
217 + runner.manager.prepare_store.side_effect = (
218 + lambda block_hashes: generate_store_output(block_hashes)
219 + )
220 + runner.run(
221 + decoded_tokens=[0] * gpu_block_size,
222 + expected_loaded_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
223 + )
224 +
225 + runner.run(
226 + decoded_tokens=[EOS_TOKEN_ID],
227 + expected_stored_gpu_block_indexes=(9, 10, 11),
228 + )
229 +
230 +
231 +@pytest.mark.parametrize("async_scheduling", [True, False])
232 +def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: bool):
233 + offloaded_block_size = 12
234 + gpu_block_size = 4
235 + num_gpu_blocks = 100
236 +
237 + runner = request_runner(
238 + offloaded_block_size=offloaded_block_size,
239 + gpu_block_size=gpu_block_size,
240 + num_gpu_blocks=num_gpu_blocks,
241 + async_scheduling=async_scheduling,
242 + )
243 +
244 + # store 1 blocks
245 + runner.new_request(token_ids=[0] * offloaded_block_size)
246 + runner.manager.prepare_store.side_effect = (
247 + lambda block_hashes: generate_store_output(block_hashes)
248 + )
249 + runner.run(
250 + decoded_tokens=[EOS_TOKEN_ID],
251 + expected_stored_gpu_block_indexes=(0, 1, 2),
252 + )
253 +
254 + # start a request to load the first block, but don't complete
255 + runner.scheduler.reset_prefix_cache()
256 + runner.new_request(token_ids=[0] * offloaded_block_size)
257 + runner.manager.lookup.return_value = 1
258 + runner.run(
259 + decoded_tokens=[],
260 + complete_transfers=False,
261 + )
262 +
263 + # request triggered a load
264 + transfer_jobs = list(runner.offloading_spec.handler.transfer_specs)
265 + assert transfer_jobs
266 +
267 + # start a new request to load the same first block
268 + runner.new_request(token_ids=[0] * offloaded_block_size)
269 + runner.manager.lookup.return_value = 1
270 + runner.run(
271 + decoded_tokens=[],
272 + complete_transfers=False,
273 + )
274 +
275 + # request did not trigger a load
276 + assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs)
277 +
278 + # complete transfers
279 + runner.manager.prepare_store.side_effect = (
280 + lambda block_hashes: generate_store_output([])
281 + )
282 + runner.run(
283 + decoded_tokens=[EOS_TOKEN_ID],
284 + expected_loaded_gpu_block_indexes=(0, 1, 2),
285 + )
286 +
287 + # second request will use the GPU prefix cache
288 + assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs)
289 +
290 +
291 +@pytest.mark.parametrize("async_scheduling", [True, False])
292 +def test_abort_loading_requests(request_runner, async_scheduling: bool):
293 + offloaded_block_size = 12
294 + gpu_block_size = 4
295 + num_gpu_blocks = 100
296 +
297 + runner = request_runner(
298 + offloaded_block_size=offloaded_block_size,
299 + gpu_block_size=gpu_block_size,
300 + num_gpu_blocks=num_gpu_blocks,
301 + async_scheduling=async_scheduling,
302 + )
303 +
304 + # store 1 blocks
305 + runner.new_request(token_ids=[0] * offloaded_block_size)
306 + runner.manager.prepare_store.side_effect = (
307 + lambda block_hashes: generate_store_output(block_hashes)
308 + )
309 + runner.run(
310 + decoded_tokens=[EOS_TOKEN_ID],
311 + expected_stored_gpu_block_indexes=(0, 1, 2),
312 + )
313 +
314 + # start a request to load the first block, but don't complete
315 + runner.scheduler.reset_prefix_cache()
316 + runner.new_request(token_ids=[0] * offloaded_block_size)
317 + runner.manager.lookup.return_value = 1
318 + runner.run(
319 + decoded_tokens=[],
320 + complete_transfers=False,
321 + )
322 +
323 + # request triggered a load
324 + transfer_jobs = list(runner.offloading_spec.handler.transfer_specs)
325 + assert transfer_jobs
326 +
327 + # abort request
328 + req_id = str(runner.req_id)
329 + runner.scheduler.finish_requests((req_id,), RequestStatus.FINISHED_ABORTED)
330 +
331 + # verify request is not deleted
332 + assert req_id in runner.scheduler.requests
333 +
334 + # complete loading request
335 + runner.run(
336 + decoded_tokens=[],
337 + expected_loaded_gpu_block_indexes=(0, 1, 2),
338 + )
339 +
340 + # assert request is deleted
341 + assert req_id not in runner.scheduler.requests
tests/v1/kv_connector/unit/offloading_connector/test_worker.py new
+504
@@ -0,0 +1,504 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from collections import defaultdict
4 +from unittest.mock import MagicMock, patch
5 +
6 +import pytest
7 +import torch
8 +
9 +from vllm.platforms import current_platform
10 +from vllm.utils.torch_utils import get_dtype_size
11 +from vllm.v1.attention.backend import AttentionBackend
12 +from vllm.v1.attention.backends.registry import AttentionBackendEnum
13 +from vllm.v1.attention.backends.utils import set_kv_cache_layout
14 +from vllm.v1.kv_cache_interface import (
15 + FullAttentionSpec,
16 + KVCacheConfig,
17 + KVCacheGroupSpec,
18 + KVCacheTensor,
19 + MambaSpec,
20 + MLAAttentionSpec,
21 + UniformTypeKVCacheSpecs,
22 +)
23 +from vllm.v1.kv_offload.spec import (
24 + CanonicalKVCacheRef,
25 + CanonicalKVCaches,
26 + OffloadingSpec,
27 +)
28 +
29 +NUM_BLOCKS = 10
30 +BLOCK_SIZE = 16
31 +NUM_KV_HEADS = 4
32 +HEAD_SIZE = 64
33 +DTYPE = torch.float16
34 +
35 +# Attention backends to test
36 +ATTN_BACKENDS: list[str] = []
37 +if current_platform.is_cuda():
38 + ATTN_BACKENDS = [
39 + "FLASH_ATTN",
40 + "FLEX_ATTENTION",
41 + "FLASHINFER",
42 + "TRITON_ATTN",
43 + ]
44 +elif current_platform.is_rocm():
45 + ATTN_BACKENDS = ["TRITON_ATTN"]
46 +
47 +# ---------------------------------------------------------------------------
48 +# Helpers
49 +# ---------------------------------------------------------------------------
50 +
51 +
52 +def _allocate_and_reshape_kv_caches(
53 + kv_cache_config: KVCacheConfig,
54 + attn_groups: list[list],
55 + device: torch.device,
56 +):
57 + """
58 + Use the real GPUModelRunner allocation and reshape methods to produce
59 + kv_caches, just like the model runner does during initialization.
60 + """
61 + from vllm.v1.worker.gpu_model_runner import GPUModelRunner
62 +
63 + # Some backends (e.g. FlashAttention) query the KV cache layout during
64 + # reshape, which ultimately calls get_current_vllm_config(). Setting
65 + # the layout override avoids needing a full VllmConfig context.
66 + set_kv_cache_layout("NHD")
67 + try:
68 + runner = object.__new__(GPUModelRunner)
69 + runner.device = device
70 + runner.runner_only_attn_layers = set()
71 + runner.attn_groups = attn_groups
72 + runner.kv_cache_config = kv_cache_config
73 + runner.cache_config = MagicMock(cache_dtype="auto")
74 + runner.shared_kv_cache_layers = {}
75 + runner.model_config = MagicMock()
76 + runner.model_config.hf_config.model_type = ""
77 + runner.compilation_config = MagicMock(
78 + static_forward_context=defaultdict(MagicMock)
79 + )
80 + runner.kv_caches = []
81 +
82 + kernel_block_sizes = [BLOCK_SIZE] * len(kv_cache_config.kv_cache_groups)
83 + return runner.initialize_kv_cache_tensors(kv_cache_config, kernel_block_sizes)
84 + finally:
85 + set_kv_cache_layout(None)
86 +
87 +
88 +def _make_mock_layer(backend_cls: type[AttentionBackend]):
89 + """
90 + Create a mock AttentionLayerBase whose get_attn_backend returns backend_cls.
91 + """
92 + layer = MagicMock()
93 + layer.get_attn_backend.return_value = backend_cls
94 + return layer
95 +
96 +
97 +def _make_worker(kv_cache_config: KVCacheConfig):
98 + """
99 + Create an OffloadingConnectorWorker with mocked dependencies.
100 + """
101 + from vllm.distributed.kv_transfer.kv_connector.v1.offloading.worker import (
102 + OffloadingConnectorWorker,
103 + )
104 +
105 + spec = MagicMock(spec=OffloadingSpec)
106 + spec.kv_cache_config = kv_cache_config
107 + spec.vllm_config = MagicMock()
108 + spec.get_handlers.return_value = iter([])
109 +
110 + worker = OffloadingConnectorWorker(spec=spec)
111 + worker.worker = MagicMock()
112 +
113 + return worker, spec
114 +
115 +
116 +# ---------------------------------------------------------------------------
117 +# Tests
118 +# ---------------------------------------------------------------------------
119 +
120 +
121 +@pytest.mark.parametrize("backend", ATTN_BACKENDS)
122 +@patch(
123 + "vllm.distributed.kv_transfer.kv_connector.v1.offloading"
124 + ".worker.get_layers_from_vllm_config"
125 +)
126 +def test_register_kv_caches(mock_get_layers, backend):
127 + """Test register_kv_caches with multiple groups covering all layer types.
128 +
129 + Creates one FullAttention group, one MLA group, one Mamba group, and
130 + one Mamba-padded group. Each group has GROUP_SIZE layers.
131 +
132 + KVCacheTensors are shared across all groups mirroring the real allocation
133 + in kv_cache_utils.py: tensor i is shared by layer i from every group.
134 + The padded-mamba group has a different page size so its layers get their
135 + own dedicated tensors.
136 +
137 + Uses the real GPUModelRunner.initialize_kv_cache_tensors to produce
138 + kv_caches, which automatically applies
139 + _update_hybrid_attention_mamba_layout for hybrid models.
140 +
141 + Verifies that the canonicalized CanonicalKVCaches has the correct
142 + block tensors, tensor_idx references, and page sizes across all groups.
143 + """
144 + from vllm.v1.attention.backends.mla.indexer import (
145 + DeepseekV32IndexerBackend,
146 + )
147 + from vllm.v1.worker.utils import AttentionGroup
148 +
149 + MLA_HEAD_SIZE = NUM_KV_HEADS * HEAD_SIZE * 2
150 +
151 + # padded mamba (missing HEAD_SIZE)
152 + CONV_STATE_SHAPE = (BLOCK_SIZE * NUM_KV_HEADS, HEAD_SIZE)
153 + UNALIGNED_SSM_STATE_SHAPE = (BLOCK_SIZE * NUM_KV_HEADS - 1, HEAD_SIZE)
154 +
155 + PAGE_SIZE_BYTES = 2 * BLOCK_SIZE * NUM_KV_HEADS * HEAD_SIZE * get_dtype_size(DTYPE)
156 + unaligned_mamba_page_size = PAGE_SIZE_BYTES - HEAD_SIZE * get_dtype_size(DTYPE)
157 +
158 + # unpadded mamba (fills page exactly)
159 + ALIGNED_SSM_STATE_SHAPE = (BLOCK_SIZE * NUM_KV_HEADS, HEAD_SIZE)
160 +
161 + backend_cls = AttentionBackendEnum[backend].get_class()
162 +
163 + attn_spec = FullAttentionSpec(
164 + block_size=BLOCK_SIZE,
165 + num_kv_heads=NUM_KV_HEADS,
166 + head_size=HEAD_SIZE,
167 + dtype=DTYPE,
168 + )
169 + mla_spec = MLAAttentionSpec(
170 + block_size=BLOCK_SIZE,
171 + num_kv_heads=1,
172 + head_size=MLA_HEAD_SIZE,
173 + dtype=DTYPE,
174 + )
175 + unaligned_mamba_spec = MambaSpec(
176 + block_size=BLOCK_SIZE,
177 + shapes=(CONV_STATE_SHAPE, UNALIGNED_SSM_STATE_SHAPE),
178 + dtypes=(DTYPE, DTYPE),
179 + page_size_padded=PAGE_SIZE_BYTES,
180 + )
181 + aligned_mamba_spec = MambaSpec(
182 + block_size=BLOCK_SIZE,
183 + shapes=(CONV_STATE_SHAPE, ALIGNED_SSM_STATE_SHAPE),
184 + dtypes=(DTYPE, DTYPE),
185 + page_size_padded=PAGE_SIZE_BYTES,
186 + )
187 +
188 + assert attn_spec.page_size_bytes == PAGE_SIZE_BYTES
189 + assert mla_spec.page_size_bytes == PAGE_SIZE_BYTES
190 + assert unaligned_mamba_spec.page_size_bytes == PAGE_SIZE_BYTES
191 + assert aligned_mamba_spec.page_size_bytes == PAGE_SIZE_BYTES
192 +
193 + GROUP_SIZE = 3
194 +
195 + # -- Build per-group layer info ----------------------------------------
196 + layer_idx = 0
197 +
198 + attn_layer_names = []
199 + for _ in range(GROUP_SIZE):
200 + attn_layer_names.append(f"model.layers.{layer_idx}.self_attn")
201 + layer_idx += 1
202 +
203 + mla_layer_names = []
204 + for _ in range(GROUP_SIZE):
205 + mla_layer_names.append(f"model.layers.{layer_idx}.self_attn")
206 + layer_idx += 1
207 +
208 + unaligned_mamba_layer_names = []
209 + for _ in range(GROUP_SIZE):
210 + unaligned_mamba_layer_names.append(f"model.layers.{layer_idx}.mamba_unpadded")
211 + layer_idx += 1
212 +
213 + aligned_mamba_layer_names = []
214 + for _ in range(GROUP_SIZE - 1):
215 + aligned_mamba_layer_names.append(f"model.layers.{layer_idx}.mamba_padded")
216 + layer_idx += 1
217 +
218 + layer_groups = [
219 + attn_layer_names,
220 + mla_layer_names,
221 + unaligned_mamba_layer_names,
222 + aligned_mamba_layer_names,
223 + ]
224 +
225 + kv_cache_tensors: list[KVCacheTensor] = []
226 + for i in range(GROUP_SIZE):
227 + shared_by: list[str] = []
228 + for group_layer_names in layer_groups:
229 + if len(group_layer_names) > i:
230 + shared_by.append(group_layer_names[i])
231 + kv_cache_tensors.append(
232 + KVCacheTensor(
233 + size=PAGE_SIZE_BYTES * NUM_BLOCKS,
234 + shared_by=shared_by,
235 + )
236 + )
237 +
238 + kv_cache_groups = [
239 + KVCacheGroupSpec(layer_names=attn_layer_names, kv_cache_spec=attn_spec),
240 + KVCacheGroupSpec(layer_names=mla_layer_names, kv_cache_spec=mla_spec),
241 + KVCacheGroupSpec(
242 + layer_names=unaligned_mamba_layer_names, kv_cache_spec=unaligned_mamba_spec
243 + ),
244 + KVCacheGroupSpec(
245 + layer_names=aligned_mamba_layer_names, kv_cache_spec=aligned_mamba_spec
246 + ),
247 + ]
248 +
249 + attn_groups = [
250 + [
251 + AttentionGroup(
252 + backend=backend_cls,
253 + layer_names=attn_layer_names,
254 + kv_cache_spec=attn_spec,
255 + kv_cache_group_id=0,
256 + ),
257 + AttentionGroup(
258 + backend=DeepseekV32IndexerBackend,
259 + layer_names=mla_layer_names,
260 + kv_cache_spec=mla_spec,
261 + kv_cache_group_id=1,
262 + ),
263 + AttentionGroup(
264 + backend=DeepseekV32IndexerBackend, # unused for mamba
265 + layer_names=unaligned_mamba_layer_names,
266 + kv_cache_spec=unaligned_mamba_spec,
267 + kv_cache_group_id=2,
268 + ),
269 + AttentionGroup(
270 + backend=DeepseekV32IndexerBackend, # unused for mamba
271 + layer_names=aligned_mamba_layer_names,
272 + kv_cache_spec=aligned_mamba_spec,
273 + kv_cache_group_id=3,
274 + ),
275 + ]
276 + ]
277 +
278 + kv_cache_config = KVCacheConfig(
279 + num_blocks=NUM_BLOCKS,
280 + kv_cache_tensors=kv_cache_tensors,
281 + kv_cache_groups=kv_cache_groups,
282 + )
283 +
284 + kv_caches = _allocate_and_reshape_kv_caches(
285 + kv_cache_config,
286 + attn_groups,
287 + device=torch.device("cuda:0"),
288 + )
289 +
290 + mock_layers: dict[str, MagicMock] = {}
291 + for layer_name in attn_layer_names:
292 + mock_layers[layer_name] = _make_mock_layer(backend_cls)
293 + for layer_name in mla_layer_names:
294 + mock_layers[layer_name] = _make_mock_layer(DeepseekV32IndexerBackend)
295 + mock_get_layers.return_value = mock_layers
296 +
297 + worker, spec = _make_worker(kv_cache_config)
298 + worker.register_kv_caches(kv_caches)
299 +
300 + canonical = spec.get_handlers.call_args[0][0]
301 + assert isinstance(canonical, CanonicalKVCaches)
302 +
303 + # -- Expected block tensors ----------------------------------------------
304 + # All tensors have the same padded page size (PAGE_SIZE_BYTES).
305 + # Tensor 0: shared by attn[0], mla[0], mamba_unaligned[0], mamba_aligned[0]
306 + # Tensor 1: shared by attn[1], mla[1], mamba_unaligned[1], mamba_aligned[1]
307 + # Tensor 2: shared by attn[2], mla[2], mamba_unaligned[2]
308 + # (mamba_aligned has only GROUP_SIZE-1 = 2 layers)
309 + expected_tensors = [
310 + (NUM_BLOCKS, PAGE_SIZE_BYTES),
311 + (NUM_BLOCKS, PAGE_SIZE_BYTES),
312 + (NUM_BLOCKS, PAGE_SIZE_BYTES),
313 + ]
314 +
315 + # -- Expected group data refs (order matches kv_cache_groups) -------------
316 + ref = CanonicalKVCacheRef
317 + expected_group_refs = [
318 + # attn group: layers attn[0..2] → tensors 0,1,2 with full page size
319 + [
320 + ref(tensor_idx=0, page_size_bytes=PAGE_SIZE_BYTES),
321 + ref(tensor_idx=1, page_size_bytes=PAGE_SIZE_BYTES),
322 + ref(tensor_idx=2, page_size_bytes=PAGE_SIZE_BYTES),
323 + ],
324 + # mla group: layers mla[0..2] → tensors 0,1,2 with full page size
325 + [
326 + ref(tensor_idx=0, page_size_bytes=PAGE_SIZE_BYTES),
327 + ref(tensor_idx=1, page_size_bytes=PAGE_SIZE_BYTES),
328 + ref(tensor_idx=2, page_size_bytes=PAGE_SIZE_BYTES),
329 + ],
330 + # unaligned mamba group: layers [0..2] → tensors 0,1,2 with unaligned page
331 + [
332 + ref(tensor_idx=0, page_size_bytes=unaligned_mamba_page_size),
333 + ref(tensor_idx=1, page_size_bytes=unaligned_mamba_page_size),
334 + ref(tensor_idx=2, page_size_bytes=unaligned_mamba_page_size),
335 + ],
336 + # aligned mamba group: layers [0..1] → tensors 0,1 with full page size
337 + [
338 + ref(tensor_idx=0, page_size_bytes=PAGE_SIZE_BYTES),
339 + ref(tensor_idx=1, page_size_bytes=PAGE_SIZE_BYTES),
340 + ],
341 + ]
342 +
343 + # Verify block tensors
344 + assert len(canonical.tensors) == len(expected_tensors)
345 + for block_tensor, (exp_num_blocks, exp_page_size) in zip(
346 + canonical.tensors, expected_tensors
347 + ):
348 + tensor = block_tensor.tensor
349 + assert tensor.dtype == torch.int8
350 + assert tensor.shape == (exp_num_blocks, exp_page_size)
351 + assert block_tensor.page_size_bytes == exp_page_size
352 +
353 + # Verify group data refs
354 + assert len(canonical.group_data_refs) == len(expected_group_refs)
355 + for actual_refs, exp_refs in zip(canonical.group_data_refs, expected_group_refs):
356 + assert len(actual_refs) == len(exp_refs)
357 + for actual, expected in zip(actual_refs, exp_refs):
358 + assert actual.tensor_idx == expected.tensor_idx
359 + assert actual.page_size_bytes == expected.page_size_bytes
360 +
361 +
362 +@pytest.mark.parametrize("backend", ATTN_BACKENDS)
363 +@patch(
364 + "vllm.distributed.kv_transfer.kv_connector.v1.offloading"
365 + ".worker.get_layers_from_vllm_config"
366 +)
367 +def test_register_kv_caches_uniform_type(mock_get_layers, backend):
368 + """Test register_kv_caches with UniformTypeKVCacheSpecs.
369 +
370 + Two attention layers use the same backend but different num_kv_heads,
371 + giving them different per-layer page sizes. Each has its own
372 + KVCacheTensor and are wrapped in a UniformTypeKVCacheSpecs group.
373 + Verifies that each layer gets the correct tensor_idx and
374 + page_size_bytes in its block data ref.
375 + """
376 + from vllm.v1.worker.utils import AttentionGroup
377 +
378 + backend_cls = AttentionBackendEnum[backend].get_class()
379 +
380 + layer_a = "model.layers.0.self_attn"
381 + layer_b = "model.layers.1.self_attn"
382 + spec_a = FullAttentionSpec(
383 + block_size=BLOCK_SIZE,
384 + num_kv_heads=NUM_KV_HEADS,
385 + head_size=HEAD_SIZE,
386 + dtype=DTYPE,
387 + )
388 + spec_b = FullAttentionSpec(
389 + block_size=BLOCK_SIZE,
390 + num_kv_heads=NUM_KV_HEADS * 2,
391 + head_size=HEAD_SIZE,
392 + dtype=DTYPE,
393 + )
394 + assert spec_a.page_size_bytes != spec_b.page_size_bytes
395 +
396 + uniform_spec = UniformTypeKVCacheSpecs(
397 + block_size=BLOCK_SIZE,
398 + kv_cache_specs={layer_a: spec_a, layer_b: spec_b},
399 + )
400 +
401 + kv_cache_config = KVCacheConfig(
402 + num_blocks=NUM_BLOCKS,
403 + kv_cache_tensors=[
404 + KVCacheTensor(
405 + size=spec_a.page_size_bytes * NUM_BLOCKS,
406 + shared_by=[layer_a],
407 + ),
408 + KVCacheTensor(
409 + size=spec_b.page_size_bytes * NUM_BLOCKS,
410 + shared_by=[layer_b],
411 + ),
412 + ],
413 + kv_cache_groups=[
414 + KVCacheGroupSpec(
415 + layer_names=[layer_a, layer_b],
416 + kv_cache_spec=uniform_spec,
417 + )
418 + ],
419 + )
420 +
421 + attn_groups = [
422 + [
423 + AttentionGroup(
424 + backend=backend_cls,
425 + layer_names=[layer_a],
426 + kv_cache_spec=spec_a,
427 + kv_cache_group_id=0,
428 + ),
429 + AttentionGroup(
430 + backend=backend_cls,
431 + layer_names=[layer_b],
432 + kv_cache_spec=spec_b,
433 + kv_cache_group_id=0,
434 + ),
435 + ]
436 + ]
437 +
438 + kv_caches = _allocate_and_reshape_kv_caches(
439 + kv_cache_config,
440 + attn_groups,
441 + device=torch.device("cuda:0"),
442 + )
443 +
444 + mock_get_layers.return_value = {
445 + layer_a: _make_mock_layer(backend_cls),
446 + layer_b: _make_mock_layer(backend_cls),
447 + }
448 +
449 + worker, spec = _make_worker(kv_cache_config)
450 + worker.register_kv_caches(kv_caches)
451 +
452 + canonical = spec.get_handlers.call_args[0][0]
453 + assert isinstance(canonical, CanonicalKVCaches)
454 +
455 + unbinds = backend_cls.get_name() in ("FLASH_ATTN", "FLEX_ATTENTION")
456 + tensors_per_layer = 2 if unbinds else 1
457 +
458 + for block_tensor in canonical.tensors:
459 + assert block_tensor.tensor.dtype == torch.int8
460 +
461 + # Single group with refs from both layers
462 + assert len(canonical.group_data_refs) == 1
463 + group_refs = canonical.group_data_refs[0]
464 + assert len(group_refs) == 2 * tensors_per_layer
465 +
466 + if unbinds:
467 + half_a = spec_a.page_size_bytes // 2
468 + half_b = spec_b.page_size_bytes // 2
469 +
470 + assert len(canonical.tensors) == 4
471 + assert canonical.tensors[0].page_size_bytes == half_a
472 + assert canonical.tensors[1].page_size_bytes == half_a
473 + assert canonical.tensors[2].page_size_bytes == half_b
474 + assert canonical.tensors[3].page_size_bytes == half_b
475 + assert canonical.tensors[0].tensor.shape == (NUM_BLOCKS, half_a)
476 + assert canonical.tensors[1].tensor.shape == (NUM_BLOCKS, half_a)
477 + assert canonical.tensors[2].tensor.shape == (NUM_BLOCKS, half_b)
478 + assert canonical.tensors[3].tensor.shape == (NUM_BLOCKS, half_b)
479 +
480 + assert group_refs[0] == CanonicalKVCacheRef(
481 + tensor_idx=0, page_size_bytes=half_a
482 + )
483 + assert group_refs[1] == CanonicalKVCacheRef(
484 + tensor_idx=1, page_size_bytes=half_a
485 + )
486 + assert group_refs[2] == CanonicalKVCacheRef(
487 + tensor_idx=2, page_size_bytes=half_b
488 + )
489 + assert group_refs[3] == CanonicalKVCacheRef(
490 + tensor_idx=3, page_size_bytes=half_b
491 + )
492 + else:
493 + assert len(canonical.tensors) == 2
494 + assert canonical.tensors[0].page_size_bytes == spec_a.page_size_bytes
495 + assert canonical.tensors[1].page_size_bytes == spec_b.page_size_bytes
496 + assert canonical.tensors[0].tensor.shape == (NUM_BLOCKS, spec_a.page_size_bytes)
497 + assert canonical.tensors[1].tensor.shape == (NUM_BLOCKS, spec_b.page_size_bytes)
498 +
499 + assert group_refs[0] == CanonicalKVCacheRef(
500 + tensor_idx=0, page_size_bytes=spec_a.page_size_bytes
501 + )
502 + assert group_refs[1] == CanonicalKVCacheRef(
503 + tensor_idx=1, page_size_bytes=spec_b.page_size_bytes
504 + )
tests/v1/kv_connector/unit/offloading_connector/utils.py renamed
+15 -487
@@ -9,16 +9,17 @@ from unittest.mock import MagicMock
9 import pytest
10 import torch
11
12 +from tests.v1.kv_connector.unit.utils import (
13 + EOS_TOKEN_ID,
14 + create_model_runner_output,
15 + create_vllm_config,
16 +)
17 from vllm import SamplingParams
13 -from vllm.config import KVTransferConfig, VllmConfig
14 -from vllm.distributed.kv_events import BlockRemoved, BlockStored
18 +from vllm.config import KVTransferConfig, VllmConfig, set_current_vllm_config
19 from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole
20 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
21 OffloadingConnectorMetadata,
22 )
19 -from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
20 - OffloadingConnectorStats,
21 -)
23 from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
24 OffloadingConnector,
25 )
@@ -39,7 +40,6 @@ from vllm.v1.kv_cache_interface import (
40 )
41 from vllm.v1.kv_offload.abstract import (
42 LoadStoreSpec,
42 - OffloadingEvent,
43 OffloadingManager,
44 PrepareStoreOutput,
45 )
@@ -51,15 +51,9 @@ from vllm.v1.kv_offload.worker.worker import (
51 TransferSpec,
52 )
53 from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, KVConnectorOutput
54 -from vllm.v1.request import Request, RequestStatus
54 +from vllm.v1.request import Request
55 from vllm.v1.structured_output import StructuredOutputManager
56
57 -from .utils import (
58 - EOS_TOKEN_ID,
59 - create_model_runner_output,
60 - create_vllm_config,
61 -)
62 -
57
58 class MockLoadStoreSpec(LoadStoreSpec):
59 def __init__(self, block_hashes: Iterable[BlockHash]):
@@ -125,7 +119,7 @@ class MockOffloadingSpec(OffloadingSpec):
119 return self.manager
120
121 def get_handlers(
128 - self, _, __
122 + self, _
123 ) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], OffloadingHandler]]:
124 yield GPULoadStoreSpec, MockLoadStoreSpec, self.handler
125 yield MockLoadStoreSpec, GPULoadStoreSpec, self.handler
@@ -179,7 +173,7 @@ class RequestRunner:
173 kv_role="kv_both",
174 kv_connector_extra_config={
175 "spec_name": "MockOffloadingSpec",
182 - "spec_module_path": "tests.v1.kv_connector.unit.test_offloading_connector", # noqa: E501
176 + "spec_module_path": "tests.v1.kv_connector.unit.offloading_connector.utils", # noqa: E501
177 "block_size": offloaded_block_size,
178 },
179 )
@@ -217,10 +211,12 @@ class RequestRunner:
211 )
212
213 # register worker kv_caches to enable OffloadingWorker creations
220 - self.worker_connector.register_cross_layers_kv_cache(
221 - kv_cache=torch.empty(0),
222 - attn_backend=FlashAttentionBackend,
223 - )
214 + # set_current_vllm_config is needed for get_kv_cache_layout() to work
215 + with set_current_vllm_config(vllm_config):
216 + self.worker_connector.register_cross_layers_kv_cache(
217 + kv_cache=torch.empty(0),
218 + attn_backend=FlashAttentionBackend,
219 + )
220
221 # extract connector of scheduler
222 scheduler_connector = self.scheduler.connector
@@ -521,471 +517,3 @@ def generate_store_output(block_hashes: Iterable[BlockHash]):
517 store_spec=MockLoadStoreSpec(block_hashes),
518 block_hashes_evicted=[],
519 )
524 -
525 -
526 -@pytest.mark.parametrize("async_scheduling", [True, False])
527 -def test_offloading_connector(request_runner, async_scheduling: bool):
528 - offloaded_block_size = 12
529 - gpu_block_size = 4
530 - num_gpu_blocks = 100
531 - block_size_factor = offloaded_block_size // gpu_block_size
532 -
533 - runner = request_runner(
534 - offloaded_block_size=offloaded_block_size,
535 - gpu_block_size=gpu_block_size,
536 - num_gpu_blocks=num_gpu_blocks,
537 - async_scheduling=async_scheduling,
538 - )
539 -
540 - # 3 blocks, store just the middle block (skip first and last)
541 - # blocks = [0, 1, 2], [3, 4, 5], [6, 7, 8]
542 - runner.new_request(token_ids=[0] * offloaded_block_size * 3)
543 - runner.manager.prepare_store.side_effect = (
544 - lambda block_hashes: generate_store_output(list(block_hashes)[1:2])
545 - )
546 - runner.run(decoded_tokens=[0])
547 -
548 - # add block missing 1 token -> no offload
549 - runner.run(
550 - decoded_tokens=[0] * (offloaded_block_size - 1),
551 - expected_stored_gpu_block_indexes=(3, 4, 5),
552 - )
553 - runner.manager.prepare_store.assert_not_called()
554 -
555 - # +1 token -> single block, fail prepare_store
556 - runner.manager.prepare_store.side_effect = lambda block_hashes: None
557 - runner.run(decoded_tokens=[0])
558 - runner.manager.prepare_store.assert_called()
559 -
560 - # 1 more block (+ token for async scheduling)
561 - # now set block_hashes_to_store = []
562 - runner.manager.prepare_store.side_effect = (
563 - lambda block_hashes: generate_store_output([])
564 - )
565 - runner.run(decoded_tokens=[0] * (offloaded_block_size + 1))
566 -
567 - # 1 more block (+ token for kicking off offloading)
568 - # now check touch was called with all 6 blocks
569 - runner.manager.prepare_store.side_effect = (
570 - lambda block_hashes: generate_store_output(block_hashes)
571 - )
572 - runner.run(
573 - decoded_tokens=[0] * (offloaded_block_size + 1),
574 - expected_stored_gpu_block_indexes=(15, 16, 17),
575 - )
576 - runner.manager.touch.assert_called()
577 - block_hashes1 = list(runner.manager.touch.call_args.args[0])
578 - assert len(block_hashes1) == 6
579 -
580 - # terminate request
581 - runner.run(decoded_tokens=[EOS_TOKEN_ID])
582 -
583 - # create a new request differing only on the last token
584 - runner.new_request(token_ids=[0] * (offloaded_block_size * 6 - 1) + [1])
585 - runner.run(decoded_tokens=[0])
586 - runner.manager.touch.assert_called()
587 - block_hashes2 = list(runner.manager.touch.call_args.args[0])
588 - assert len(block_hashes2) == 6
589 -
590 - # verify hashes are the same, except for the last block
591 - assert block_hashes1[:5] == block_hashes2[:5]
592 - assert block_hashes1[5] != block_hashes2[5]
593 -
594 - # terminate request
595 - runner.run(
596 - decoded_tokens=[EOS_TOKEN_ID],
597 - expected_stored_gpu_block_indexes=tuple(range(6 * block_size_factor)),
598 - )
599 -
600 - # full_block_tokens - num_computed_tokens < offloaded_block_size
601 - runner.new_request(
602 - token_ids=[0] * gpu_block_size + [1] * (offloaded_block_size - gpu_block_size)
603 - )
604 - runner.manager.prepare_store.side_effect = (
605 - lambda block_hashes: generate_store_output([])
606 - )
607 - runner.run(decoded_tokens=[EOS_TOKEN_ID])
608 - runner.manager.lookup.assert_not_called()
609 -
610 - # single block lookup with no hits
611 - runner.new_request(token_ids=[1] * offloaded_block_size)
612 - runner.manager.prepare_store.side_effect = (
613 - lambda block_hashes: generate_store_output([])
614 - )
615 - runner.run(decoded_tokens=[EOS_TOKEN_ID])
616 - runner.manager.lookup.assert_called()
617 - assert len(list(runner.manager.lookup.call_args.args[0])) == 1
618 -
619 - # single block lookup with a hit
620 - runner.scheduler.reset_prefix_cache()
621 - runner.new_request(token_ids=[0] * offloaded_block_size)
622 - runner.manager.prepare_store.side_effect = (
623 - lambda block_hashes: generate_store_output([])
624 - )
625 - runner.manager.lookup.return_value = 1
626 - runner.run(
627 - decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(0, 1, 2)
628 - )
629 -
630 - # single block lookup with a hit in a middle block
631 - runner.new_request(
632 - token_ids=[0] * offloaded_block_size * 2 + [1] * offloaded_block_size
633 - )
634 - runner.manager.prepare_store.side_effect = (
635 - lambda block_hashes: generate_store_output([])
636 - )
637 - runner.manager.lookup.return_value = 1
638 - runner.run(
639 - decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(3, 4, 5)
640 - )
641 -
642 - # test take_events
643 - def to_hashes(int_hashes: list[int]) -> list[BlockHash]:
644 - return [BlockHash(str(i).encode()) for i in int_hashes]
645 -
646 - def take_events() -> Iterable[OffloadingEvent]:
647 - yield OffloadingEvent(
648 - block_hashes=to_hashes([1, 2, 3]), block_size=16, medium="A", removed=False
649 - )
650 - yield OffloadingEvent(
651 - block_hashes=to_hashes([4, 5, 6]), block_size=32, medium="B", removed=True
652 - )
653 -
654 - runner.manager.take_events.side_effect = take_events
655 - events = list(runner.scheduler_connector.take_events())
656 - assert len(events) == 2
657 - event = events[0]
658 - assert isinstance(event, BlockStored)
659 - assert event.block_hashes == to_hashes([1, 2, 3])
660 - assert event.block_size == 16
661 - assert event.medium == "A"
662 - assert event.token_ids == []
663 - assert event.parent_block_hash is None
664 - assert event.lora_id is None
665 - assert event.lora_name is None
666 - event = events[1]
667 - assert isinstance(event, BlockRemoved)
668 - assert event.block_hashes == to_hashes([4, 5, 6])
669 - assert event.medium == "B"
670 -
671 -
672 -@pytest.mark.parametrize("async_scheduling", [True, False])
673 -def test_request_preemption(request_runner, async_scheduling: bool):
674 - offloaded_block_size = 12
675 - gpu_block_size = 4
676 - num_gpu_blocks = 100
677 -
678 - runner = request_runner(
679 - offloaded_block_size=offloaded_block_size,
680 - gpu_block_size=gpu_block_size,
681 - num_gpu_blocks=num_gpu_blocks,
682 - async_scheduling=async_scheduling,
683 - )
684 -
685 - free_block_queue = runner.scheduler.kv_cache_manager.block_pool.free_block_queue
686 - num_free_blocks_empty = free_block_queue.num_free_blocks
687 -
688 - # 2 blocks, store all, without flushing
689 - # blocks = [0, 1, 2], [3, 4, 5]
690 - runner.new_request(token_ids=[0] * offloaded_block_size * 2)
691 - runner.manager.prepare_store.side_effect = (
692 - lambda block_hashes: generate_store_output(block_hashes)
693 - )
694 - runner.run(
695 - decoded_tokens=[0],
696 - complete_transfers=False,
697 - )
698 -
699 - # decode 2 more blocks - 1 gpu block, storing [6, 7, 8] (no flush)
700 - runner.manager.prepare_store.side_effect = (
701 - lambda block_hashes: generate_store_output(block_hashes)
702 - )
703 - runner.run(
704 - decoded_tokens=[0] * (2 * offloaded_block_size - gpu_block_size),
705 - complete_transfers=False,
706 - )
707 -
708 - # simulate KV cache running out of space
709 - free_block_queue.num_free_blocks = 0
710 -
711 - # request should be preempted now
712 - runner.run(
713 - decoded_tokens=[],
714 - complete_transfers=False,
715 - expected_flushed_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
716 - expected_stored_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
717 - )
718 -
719 - # restore KV cache space and reset GPU prefix cache
720 - free_block_queue.num_free_blocks = num_free_blocks_empty
721 - runner.scheduler.reset_prefix_cache()
722 -
723 - # request should now return from preemption
724 - # re-load [0, ..., 8] from the CPU and store [9, 10, 11]
725 - runner.manager.lookup.return_value = 3
726 - runner.manager.prepare_store.side_effect = (
727 - lambda block_hashes: generate_store_output(block_hashes)
728 - )
729 - runner.run(
730 - decoded_tokens=[0] * gpu_block_size,
731 - expected_loaded_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
732 - )
733 -
734 - runner.run(
735 - decoded_tokens=[EOS_TOKEN_ID],
736 - expected_stored_gpu_block_indexes=(9, 10, 11),
737 - )
738 -
739 -
740 -@pytest.mark.parametrize("async_scheduling", [True, False])
741 -def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: bool):
742 - offloaded_block_size = 12
743 - gpu_block_size = 4
744 - num_gpu_blocks = 100
745 -
746 - runner = request_runner(
747 - offloaded_block_size=offloaded_block_size,
748 - gpu_block_size=gpu_block_size,
749 - num_gpu_blocks=num_gpu_blocks,
750 - async_scheduling=async_scheduling,
751 - )
752 -
753 - # store 1 blocks
754 - runner.new_request(token_ids=[0] * offloaded_block_size)
755 - runner.manager.prepare_store.side_effect = (
756 - lambda block_hashes: generate_store_output(block_hashes)
757 - )
758 - runner.run(
759 - decoded_tokens=[EOS_TOKEN_ID],
760 - expected_stored_gpu_block_indexes=(0, 1, 2),
761 - )
762 -
763 - # start a request to load the first block, but don't complete
764 - runner.scheduler.reset_prefix_cache()
765 - runner.new_request(token_ids=[0] * offloaded_block_size)
766 - runner.manager.lookup.return_value = 1
767 - runner.run(
768 - decoded_tokens=[],
769 - complete_transfers=False,
770 - )
771 -
772 - # request triggered a load
773 - transfer_jobs = list(runner.offloading_spec.handler.transfer_specs)
774 - assert transfer_jobs
775 -
776 - # start a new request to load the same first block
777 - runner.new_request(token_ids=[0] * offloaded_block_size)
778 - runner.manager.lookup.return_value = 1
779 - runner.run(
780 - decoded_tokens=[],
781 - complete_transfers=False,
782 - )
783 -
784 - # request did not trigger a load
785 - assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs)
786 -
787 - # complete transfers
788 - runner.manager.prepare_store.side_effect = (
789 - lambda block_hashes: generate_store_output([])
790 - )
791 - runner.run(
792 - decoded_tokens=[EOS_TOKEN_ID],
793 - expected_loaded_gpu_block_indexes=(0, 1, 2),
794 - )
795 -
796 - # second request will use the GPU prefix cache
797 - assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs)
798 -
799 -
800 -@pytest.mark.parametrize("async_scheduling", [True, False])
801 -def test_abort_loading_requests(request_runner, async_scheduling: bool):
802 - offloaded_block_size = 12
803 - gpu_block_size = 4
804 - num_gpu_blocks = 100
805 -
806 - runner = request_runner(
807 - offloaded_block_size=offloaded_block_size,
808 - gpu_block_size=gpu_block_size,
809 - num_gpu_blocks=num_gpu_blocks,
810 - async_scheduling=async_scheduling,
811 - )
812 -
813 - # store 1 blocks
814 - runner.new_request(token_ids=[0] * offloaded_block_size)
815 - runner.manager.prepare_store.side_effect = (
816 - lambda block_hashes: generate_store_output(block_hashes)
817 - )
818 - runner.run(
819 - decoded_tokens=[EOS_TOKEN_ID],
820 - expected_stored_gpu_block_indexes=(0, 1, 2),
821 - )
822 -
823 - # start a request to load the first block, but don't complete
824 - runner.scheduler.reset_prefix_cache()
825 - runner.new_request(token_ids=[0] * offloaded_block_size)
826 - runner.manager.lookup.return_value = 1
827 - runner.run(
828 - decoded_tokens=[],
829 - complete_transfers=False,
830 - )
831 -
832 - # request triggered a load
833 - transfer_jobs = list(runner.offloading_spec.handler.transfer_specs)
834 - assert transfer_jobs
835 -
836 - # abort request
837 - req_id = str(runner.req_id)
838 - runner.scheduler.finish_requests((req_id,), RequestStatus.FINISHED_ABORTED)
839 -
840 - # verify request is not deleted
841 - assert req_id in runner.scheduler.requests
842 -
843 - # complete loading request
844 - runner.run(
845 - decoded_tokens=[],
846 - expected_loaded_gpu_block_indexes=(0, 1, 2),
847 - )
848 -
849 - # assert request is deleted
850 - assert req_id not in runner.scheduler.requests
851 -
852 -
853 -class TestOffloadingConnectorStats:
854 - """Tests for OffloadingConnector stats reconstruction and operations."""
855 -
856 - def test_build_kv_connector_stats_with_none(self):
857 - """Test that build_kv_connector_stats returns empty stats when given None."""
858 - stats = OffloadingConnector.build_kv_connector_stats(data=None)
859 -
860 - assert stats is not None
861 - assert isinstance(stats, OffloadingConnectorStats)
862 - assert len(stats.data) == 0
863 - assert stats.is_empty()
864 -
865 - def test_build_kv_connector_stats_with_empty_dict(self):
866 - """Test that build_kv_connector_stats returns empty stats with empty dict."""
867 - stats = OffloadingConnector.build_kv_connector_stats(data={})
868 -
869 - assert stats is not None
870 - assert isinstance(stats, OffloadingConnectorStats)
871 - assert len(stats.data) == 0
872 - assert stats.is_empty()
873 -
874 - def test_build_kv_connector_stats_reconstructs_offload_stats(self):
875 - """Test that OffloadingConnector stats are properly reconstructed with
876 - correct data."""
877 - serialized_data = {
878 - "CPU_to_GPU": [
879 - {"op_size": 16, "op_time": 1.0},
880 - {"op_size": 8, "op_time": 0.5},
881 - ],
882 - "GPU_to_CPU": [
883 - {"op_size": 1, "op_time": 0.1},
884 - {"op_size": 2, "op_time": 0.2},
885 - ],
886 - }
887 -
888 - stats = OffloadingConnector.build_kv_connector_stats(data=serialized_data)
889 -
890 - offload_connector_stats = stats
891 - assert isinstance(offload_connector_stats, OffloadingConnectorStats)
892 - assert offload_connector_stats.data["CPU_to_GPU"] == [
893 - {"op_size": 16, "op_time": 1.0},
894 - {"op_size": 8, "op_time": 0.5},
895 - ]
896 - assert offload_connector_stats.data["GPU_to_CPU"] == [
897 - {"op_size": 1, "op_time": 0.1},
898 - {"op_size": 2, "op_time": 0.2},
899 - ]
900 -
901 - def test_aggregate_same_connector(self):
902 - """Test aggregating stats from the same connector type."""
903 - stats1 = OffloadingConnectorStats(
904 - data={
905 - "CPU_to_GPU": [
906 - {"op_size": 16, "op_time": 1.0},
907 - {"op_size": 8, "op_time": 0.5},
908 - ],
909 - "GPU_to_CPU": [
910 - {"op_size": 1, "op_time": 0.1},
911 - {"op_size": 2, "op_time": 0.2},
912 - ],
913 - }
914 - )
915 -
916 - stats2 = OffloadingConnectorStats(
917 - data={
918 - "CPU_to_GPU": [
919 - {"op_size": 3, "op_time": 0.2},
920 - {"op_size": 7, "op_time": 0.9},
921 - ],
922 - "GPU_to_CPU": [{"op_size": 16, "op_time": 2}],
923 - }
924 - )
925 -
926 - result = stats1.aggregate(stats2)
927 -
928 - assert result is stats1 # Should return self
929 - offload_connector_stats = result
930 - assert offload_connector_stats.data["CPU_to_GPU"] == [
931 - {"op_size": 16, "op_time": 1.0},
932 - {"op_size": 8, "op_time": 0.5},
933 - {"op_size": 3, "op_time": 0.2},
934 - {"op_size": 7, "op_time": 0.9},
935 - ]
936 - assert offload_connector_stats.data["GPU_to_CPU"] == [
937 - {"op_size": 1, "op_time": 0.1},
938 - {"op_size": 2, "op_time": 0.2},
939 - {"op_size": 16, "op_time": 2},
940 - ]
941 -
942 - def test_reduce(self):
943 - """Test that reduce() correctly reduces all nested connector stats."""
944 - stats = OffloadingConnectorStats(
945 - data={
946 - "CPU_to_GPU": [
947 - {"op_size": 16, "op_time": 1.0},
948 - {"op_size": 8, "op_time": 0.5},
949 - {"op_size": 3, "op_time": 0.2},
950 - {"op_size": 7, "op_time": 0.9},
951 - ],
952 - "GPU_to_CPU": [
953 - {"op_size": 1, "op_time": 0.1},
954 - {"op_size": 2, "op_time": 0.2},
955 - {"op_size": 16, "op_time": 2},
956 - ],
957 - }
958 - )
959 -
960 - reduced = stats.reduce()
961 -
962 - assert isinstance(reduced, dict)
963 - # Check that the stats were reduced (should have aggregated values)
964 - assert "CPU_to_GPU_total_bytes" in reduced
965 - assert "CPU_to_GPU_total_time" in reduced
966 - assert "GPU_to_CPU_total_bytes" in reduced
967 - assert "GPU_to_CPU_total_time" in reduced
968 - assert reduced["CPU_to_GPU_total_bytes"] == 34
969 - assert reduced["CPU_to_GPU_total_time"] == 2.6
970 - assert reduced["GPU_to_CPU_total_time"] == 2.3
971 - assert reduced["GPU_to_CPU_total_bytes"] == 19
972 -
973 - def test_reset(self):
974 - """Test that reset() resets all nested connector stats."""
975 - offload_connector_stats = OffloadingConnectorStats(
976 - data={
977 - "CPU_to_GPU": [
978 - {"op_size": 3, "op_time": 0.2},
979 - {"op_size": 7, "op_time": 0.9},
980 - ],
981 - "GPU_to_CPU": [{"op_size": 16, "op_time": 2}],
982 - }
983 - )
984 -
985 - assert not offload_connector_stats.is_empty()
986 -
987 - offload_connector_stats.reset()
988 -
989 - # After reset, stats should be empty
990 - assert offload_connector_stats.is_empty()
991 - assert len(offload_connector_stats.data) == 0
tests/v1/kv_connector/unit/test_nixl_connector.py
+3
@@ -91,6 +91,9 @@ def clear_kv_transfer():
91 yield
92 if has_kv_transfer_group():
93 ensure_kv_transfer_shutdown()
94 + # Reset any KV cache layout override set during tests so it doesn't
95 + # leak into tests in other modules.
96 + set_kv_cache_layout(None)
97
98
99 def get_default_xfer_telemetry(
tests/v1/kv_offload/test_cpu_gpu.py
+88 -127
@@ -6,32 +6,20 @@ import time
6 import pytest
7 import torch
8
9 -from vllm.platforms import current_platform
9 from vllm.utils.torch_utils import set_random_seed
11 -from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
10 from vllm.v1.kv_offload.mediums import CPULoadStoreSpec, GPULoadStoreSpec
11 +from vllm.v1.kv_offload.spec import (
12 + CanonicalKVCacheRef,
13 + CanonicalKVCaches,
14 + CanonicalKVCacheTensor,
15 +)
16 from vllm.v1.kv_offload.worker.cpu_gpu import CpuGpuOffloadingHandlers
17
15 -BACKENDS_TO_TEST = [FlashAttentionBackend]
16 -
17 -if not current_platform.is_rocm():
18 - from vllm.v1.attention.backends.flashinfer import FlashInferBackend
19 -
20 - BACKENDS_TO_TEST.append(FlashInferBackend)
21 -
22 - from vllm.v1.attention.backends.mla.flashattn_mla import FlashAttnMLABackend
23 -
24 - BACKENDS_TO_TEST.append(FlashAttnMLABackend)
25 -
18 NUM_GPU_BLOCKS = [64]
19 NUM_CPU_BLOCKS = [256]
28 -KERNEL_BLOCK_SIZES = [16]
29 -LOGICAL_BLOCK_SIZES = [16, 32]
30 -LOGICAL_BLOCKS_PER_CPU_BLOCK = [1, 3]
31 -HEAD_SIZES = [64]
32 -NUM_HEADS = [8]
33 -NUM_LAYERS = [4]
34 -DTYPES = [torch.bfloat16]
20 +GPU_PAGE_SIZES = [512, 1024]
21 +BLOCK_SIZE_FACTORS = [1, 3]
22 +NUM_TENSORS = [4]
23 SEEDS = [0]
24 CUDA_DEVICES = ["cuda:0"]
25 NUM_MAPPINGS = [3]
@@ -39,15 +27,11 @@ NUM_MAPPINGS = [3]
27
28 @pytest.mark.parametrize("gpu_to_cpu", [True, False])
29 @pytest.mark.parametrize("num_mappings", NUM_MAPPINGS)
42 -@pytest.mark.parametrize("head_size", HEAD_SIZES)
43 -@pytest.mark.parametrize("num_heads", NUM_HEADS)
44 -@pytest.mark.parametrize("kernel_block_size", KERNEL_BLOCK_SIZES)
45 -@pytest.mark.parametrize("logical_block_size", LOGICAL_BLOCK_SIZES)
46 -@pytest.mark.parametrize("logical_blocks_per_cpu_block", LOGICAL_BLOCKS_PER_CPU_BLOCK)
30 +@pytest.mark.parametrize("gpu_page_size_bytes", GPU_PAGE_SIZES)
31 +@pytest.mark.parametrize("block_size_factor", BLOCK_SIZE_FACTORS)
32 @pytest.mark.parametrize("num_gpu_blocks", NUM_GPU_BLOCKS)
33 @pytest.mark.parametrize("num_cpu_blocks", NUM_CPU_BLOCKS)
49 -@pytest.mark.parametrize("num_layers", NUM_LAYERS)
50 -@pytest.mark.parametrize("dtype", DTYPES)
34 +@pytest.mark.parametrize("num_tensors", NUM_TENSORS)
35 @pytest.mark.parametrize("seed", SEEDS)
36 @pytest.mark.parametrize("device", CUDA_DEVICES)
37 @torch.inference_mode()
@@ -55,113 +39,89 @@ def test_transfer(
39 default_vllm_config,
40 gpu_to_cpu: bool,
41 num_mappings: int,
58 - head_size: int,
59 - num_heads: int,
60 - kernel_block_size: int,
61 - logical_block_size: int,
62 - logical_blocks_per_cpu_block: int,
42 + gpu_page_size_bytes: int,
43 + block_size_factor: int,
44 num_gpu_blocks: int,
45 num_cpu_blocks: int,
65 - num_layers: int,
66 - dtype: torch.dtype,
46 + num_tensors: int,
47 seed: int,
48 device: str,
49 ) -> None:
50 set_random_seed(seed)
51
72 - # create per-layer GPU KV caches based on available attn_backends
73 - attn_backends_list = BACKENDS_TO_TEST
74 -
75 - assert logical_block_size % kernel_block_size == 0
76 - kernel_blocks_per_gpu_block = logical_block_size // kernel_block_size
77 - num_gpu_kernel_blocks = num_gpu_blocks * kernel_blocks_per_gpu_block
78 -
79 - gpu_caches = {}
80 - attn_backends = {}
81 - for i in range(num_layers):
82 - layer_name = f"layer {i}"
83 -
84 - attn_backend = attn_backends_list[i % len(attn_backends_list)]
85 - attn_backends[layer_name] = attn_backend
86 -
87 - gpu_cache_shape = attn_backend.get_kv_cache_shape(
88 - num_gpu_kernel_blocks, kernel_block_size, num_heads, head_size
52 + # build CanonicalKVCacheTensor list: one per tensor
53 + kv_cache_tensors: list[CanonicalKVCacheTensor] = []
54 + for i in range(num_tensors):
55 + gpu_tensor = torch.randint(
56 + -128,
57 + 127,
58 + (num_gpu_blocks, gpu_page_size_bytes),
59 + dtype=torch.int8,
60 + device=device,
61 + )
62 + kv_cache_tensors.append(
63 + CanonicalKVCacheTensor(
64 + tensor=gpu_tensor,
65 + page_size_bytes=gpu_page_size_bytes,
66 + )
67 )
90 - gpu_caches[layer_name] = torch.rand(gpu_cache_shape, dtype=dtype, device=device)
68
92 - # create handler
93 - cpu_block_size = logical_blocks_per_cpu_block * logical_block_size
94 - kernel_blocks_per_cpu_block = cpu_block_size // kernel_block_size
69 + # one group containing all tensors, one data ref per tensor
70 + kv_cache_groups_data_refs: list[list[CanonicalKVCacheRef]] = [
71 + [
72 + CanonicalKVCacheRef(
73 + tensor_idx=i,
74 + page_size_bytes=gpu_page_size_bytes,
75 + )
76 + for i in range(num_tensors)
77 + ]
78 + ]
79 +
80 + kv_caches = CanonicalKVCaches(
81 + tensors=kv_cache_tensors,
82 + group_data_refs=kv_cache_groups_data_refs,
83 + )
84 handlers = CpuGpuOffloadingHandlers(
96 - attn_backends=attn_backends,
97 - gpu_block_size=logical_block_size,
98 - cpu_block_size=cpu_block_size,
85 + kv_caches=kv_caches,
86 + block_size_factor=block_size_factor,
87 num_cpu_blocks=num_cpu_blocks,
100 - gpu_caches=gpu_caches,
88 )
89
90 # select block mappings
104 - gpu_blocks = random.sample(
105 - range(num_gpu_blocks), num_mappings * logical_blocks_per_cpu_block
106 - )
91 + gpu_blocks = random.sample(range(num_gpu_blocks), num_mappings * block_size_factor)
92 cpu_blocks = random.sample(range(num_cpu_blocks), num_mappings)
93
109 - # convert gpu blocks to kernel block size
110 - gpu_blocks_in_kernel_block_size = []
111 - for gpu_block in gpu_blocks:
112 - base_block_id = gpu_block * kernel_blocks_per_gpu_block
113 - for i in range(kernel_blocks_per_gpu_block):
114 - gpu_blocks_in_kernel_block_size.append(i + base_block_id)
115 -
116 - # convert cpu blocks to gpu block size
117 - cpu_blocks_in_kernel_block_size = []
118 - for cpu_block in cpu_blocks:
119 - base_block_id = cpu_block * kernel_blocks_per_cpu_block
120 - for i in range(kernel_blocks_per_cpu_block):
121 - cpu_blocks_in_kernel_block_size.append(i + base_block_id)
122 -
123 - # maybe skip some GPU block to test reading from the middle of a CPU block
94 + # expand cpu blocks to gpu-page granularity for uniform comparison:
95 + # each cpu block maps to block_size_factor consecutive sub-blocks
96 + cpu_blocks_expanded = [
97 + cpu_block * block_size_factor + j
98 + for cpu_block in cpu_blocks
99 + for j in range(block_size_factor)
100 + ]
101 +
102 + # maybe skip some GPU blocks to test reading from the middle of a CPU block
103 if not gpu_to_cpu:
125 - gpu_blocks_to_skip = logical_blocks_per_cpu_block - 1
126 - gpu_blocks = gpu_blocks[gpu_blocks_to_skip:]
127 - kernel_blocks_to_skip = gpu_blocks_to_skip * kernel_blocks_per_gpu_block
128 - gpu_blocks_in_kernel_block_size = gpu_blocks_in_kernel_block_size[
129 - kernel_blocks_to_skip:
130 - ]
131 - cpu_blocks_in_kernel_block_size = cpu_blocks_in_kernel_block_size[
132 - kernel_blocks_to_skip:
133 - ]
104 + blocks_to_skip = block_size_factor - 1
105 + gpu_blocks = gpu_blocks[blocks_to_skip:]
106 + cpu_blocks_expanded = cpu_blocks_expanded[blocks_to_skip:]
107
108 # set transfer direction
109 if gpu_to_cpu:
110 handler = handlers.gpu_to_cpu_handler
138 - src_blocks = gpu_blocks
139 - dst_blocks = cpu_blocks
140 - src_spec = GPULoadStoreSpec(src_blocks, group_sizes=(len(src_blocks),))
141 - dst_spec = CPULoadStoreSpec(dst_blocks)
142 - src_blocks_in_kernel_block_size = gpu_blocks_in_kernel_block_size
143 - dst_blocks_in_kernel_block_size = cpu_blocks_in_kernel_block_size
144 - dst_size_in_kernel_blocks = num_cpu_blocks * kernel_blocks_per_cpu_block
111 + src_spec = GPULoadStoreSpec(gpu_blocks, group_sizes=(len(gpu_blocks),))
112 + dst_spec = CPULoadStoreSpec(cpu_blocks)
113 + dst_to_src = dict(zip(cpu_blocks_expanded, gpu_blocks))
114 + num_dst_sub_blocks = num_cpu_blocks * block_size_factor
115 else:
116 handler = handlers.cpu_to_gpu_handler
147 - src_blocks = cpu_blocks
148 - dst_blocks = gpu_blocks
149 - src_spec = CPULoadStoreSpec(src_blocks)
150 - dst_spec = GPULoadStoreSpec(dst_blocks, group_sizes=(len(dst_blocks),))
151 - src_blocks_in_kernel_block_size = cpu_blocks_in_kernel_block_size
152 - dst_blocks_in_kernel_block_size = gpu_blocks_in_kernel_block_size
153 - dst_size_in_kernel_blocks = num_gpu_blocks * kernel_blocks_per_gpu_block
154 -
155 - # build dst -> src mapping
156 - dst_to_src = {}
157 - for src_block, dst_block in zip(
158 - src_blocks_in_kernel_block_size, dst_blocks_in_kernel_block_size
159 - ):
160 - dst_to_src[dst_block] = src_block
117 + src_spec = CPULoadStoreSpec(cpu_blocks)
118 + dst_spec = GPULoadStoreSpec(gpu_blocks, group_sizes=(len(gpu_blocks),))
119 + dst_to_src = dict(zip(gpu_blocks, cpu_blocks_expanded))
120 + num_dst_sub_blocks = num_gpu_blocks
121
122 # clone src and dst tensors before transfer
163 - orig_src_caches = [x.clone() for x in handler.src_tensors]
164 - orig_dst_caches = [x.clone() for x in handler.dst_tensors]
123 + orig_src_tensors = [x.clone() for x in handler.src_tensors]
124 + orig_dst_tensors = [x.clone() for x in handler.dst_tensors]
125
126 # call transfer function
127 start_time = time.time()
@@ -180,11 +140,8 @@ def test_transfer(
140 if gpu_to_cpu
141 else ("CPU", "GPU")
142 )
183 - assert (
184 - finished[0].transfer_size
185 - == handler.total_block_size_in_bytes
186 - * handler.dst_block_size_factor
187 - * len(dst_blocks)
143 + assert finished[0].transfer_size == (
144 + len(gpu_blocks) * handler.group_block_size_in_bytes[0]
145 )
146 assert finished[0].transfer_time > 0
147 assert finished[0].transfer_time < (time.time() - start_time)
@@ -192,19 +149,23 @@ def test_transfer(
149 time.sleep(0.1)
150
151 # verify src tensors did not change
195 - for orig_tensor, tensor in zip(orig_src_caches, handler.src_tensors):
152 + for orig_tensor, tensor in zip(orig_src_tensors, handler.src_tensors):
153 assert torch.equal(orig_tensor, tensor)
154
198 - # verify dst tensors
199 - for dst_block in range(dst_size_in_kernel_blocks):
200 - src_block_candidate = dst_to_src.get(dst_block)
201 - for src_cache, dst_cache, orig_dst_cache in zip(
202 - handler.src_tensors,
203 - handler.dst_tensors,
204 - orig_dst_caches,
205 - ):
206 - if src_block_candidate is not None:
207 - expected_value = src_cache[src_block_candidate]
155 + # verify dst tensors at gpu-page granularity.
156 + for src_tensor, dst_tensor, orig_dst_tensor in zip(
157 + handler.src_tensors,
158 + handler.dst_tensors,
159 + orig_dst_tensors,
160 + ):
161 + # view both GPU and CPU tensors as (n, gpu_page_size_bytes) for comparison.
162 + src_view = src_tensor.view(-1, gpu_page_size_bytes)
163 + dst_view = dst_tensor.view(-1, gpu_page_size_bytes)
164 + orig_dst_view = orig_dst_tensor.view(-1, gpu_page_size_bytes)
165 + for dst_sub_block in range(num_dst_sub_blocks):
166 + src_sub_block = dst_to_src.get(dst_sub_block)
167 + if src_sub_block is not None:
168 + expected = src_view[src_sub_block]
169 else:
209 - expected_value = orig_dst_cache[dst_block]
210 - torch.testing.assert_close(dst_cache[dst_block].cpu(), expected_value.cpu())
170 + expected = orig_dst_view[dst_sub_block]
171 + torch.testing.assert_close(dst_view[dst_sub_block].cpu(), expected.cpu())
vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
+227 -15
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from collections import defaultdict
4 +from dataclasses import replace
5
6 import torch
7
@@ -18,7 +19,17 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
19 from vllm.logger import init_logger
20 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
21 from vllm.v1.attention.backend import AttentionBackend
21 -from vllm.v1.kv_offload.spec import OffloadingSpec
22 +from vllm.v1.kv_cache_interface import (
23 + AttentionSpec,
24 + MambaSpec,
25 + UniformTypeKVCacheSpecs,
26 +)
27 +from vllm.v1.kv_offload.spec import (
28 + CanonicalKVCacheRef,
29 + CanonicalKVCaches,
30 + CanonicalKVCacheTensor,
31 + OffloadingSpec,
32 +)
33 from vllm.v1.kv_offload.worker.worker import (
34 OffloadingWorker,
35 TransferSpec,
@@ -53,17 +64,13 @@ class OffloadingConnectorWorker:
64 self._job_counter = job_id + 1
65 return job_id
66
56 - def _register_handlers(
57 - self,
58 - kv_caches: dict[str, torch.Tensor],
59 - attn_backends: dict[str, type[AttentionBackend]],
60 - ):
61 - for src_cls, dst_cls, handler in self.spec.get_handlers(
62 - kv_caches, attn_backends
63 - ):
67 + def _register_handlers(self, kv_caches: CanonicalKVCaches):
68 + for src_cls, dst_cls, handler in self.spec.get_handlers(kv_caches):
69 self.worker.register_handler(src_cls, dst_cls, handler)
70
66 - def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
71 + def register_kv_caches(
72 + self, kv_caches: dict[str, torch.Tensor | list[torch.Tensor]]
73 + ):
74 layer_names = list(kv_caches.keys())
75 layers = get_layers_from_vllm_config(
76 self.spec.vllm_config,
@@ -73,16 +80,221 @@ class OffloadingConnectorWorker:
80 attn_backends = {
81 layer_name: layers[layer_name].get_attn_backend()
82 for layer_name in layer_names
83 + if layer_name in layers
84 }
77 - self._register_handlers(kv_caches, attn_backends)
85 +
86 + # layer_name -> list of matching KV cache tensors
87 + # such that each tensor starts with the num_blocks dimension.
88 + # FlashAttention layers which use the (2, num_blocks, ...) layout
89 + # will possibly map to 2 tensors, one per K and one per V.
90 + # All other layers will probably map to a single tensor.
91 + tensors_per_block: dict[str, tuple[torch.Tensor, ...]] = {}
92 + # layer_name -> size of (un-padded) page in bytes
93 + unpadded_page_size_bytes: dict[str, int] = {}
94 + # layer_name -> size of page in bytes
95 + page_size_bytes: dict[str, int] = {}
96 + for kv_cache_group in self.spec.kv_cache_config.kv_cache_groups:
97 + group_layer_names = kv_cache_group.layer_names
98 + group_kv_cache_spec = kv_cache_group.kv_cache_spec
99 + if isinstance(group_kv_cache_spec, UniformTypeKVCacheSpecs):
100 + per_layer_specs = group_kv_cache_spec.kv_cache_specs
101 + else:
102 + per_layer_specs = {}
103 + for layer_name in group_layer_names:
104 + layer_kv_cache_spec = per_layer_specs.get(
105 + layer_name, group_kv_cache_spec
106 + )
107 + if isinstance(layer_kv_cache_spec, AttentionSpec):
108 + layer_kv_cache = kv_caches[layer_name]
109 + assert isinstance(layer_kv_cache, torch.Tensor)
110 + assert layer_kv_cache.storage_offset() == 0
111 +
112 + # get the logical dimension for num_blocks
113 + test_shape = attn_backends[layer_name].get_kv_cache_shape(
114 + num_blocks=1234,
115 + block_size=16,
116 + num_kv_heads=1,
117 + head_size=256,
118 + )
119 + num_blocks_logical_dim = test_shape.index(1234)
120 +
121 + # sort the logical dimensions by stride (high to low)
122 + # to get a physical-to-logical mapping:
123 + # physical_to_logical[physical_pos] = logical_dim
124 + logical_strides = layer_kv_cache.stride()
125 + physical_to_logical = sorted(
126 + range(len(logical_strides)),
127 + key=lambda idx: logical_strides[idx],
128 + reverse=True,
129 + )
130 +
131 + num_blocks_physical_dim = physical_to_logical.index(
132 + num_blocks_logical_dim
133 + )
134 + if num_blocks_physical_dim == 0:
135 + num_blocks = layer_kv_cache.shape[num_blocks_logical_dim]
136 + storage = layer_kv_cache.untyped_storage()
137 + page = layer_kv_cache_spec.page_size_bytes
138 + tensors_per_block[layer_name] = (
139 + torch.tensor(
140 + [],
141 + dtype=torch.int8,
142 + device=layer_kv_cache.device,
143 + )
144 + .set_(storage)
145 + .view(num_blocks, page),
146 + )
147 + page_size_bytes[layer_name] = (
148 + layer_kv_cache_spec.page_size_bytes
149 + )
150 + else:
151 + # Flash Attention case: (2, num_blocks, ...)
152 + assert test_shape[0] == 2
153 + assert physical_to_logical[0] == 0
154 + assert num_blocks_physical_dim == 1
155 +
156 + # unbind the tensor to separate K and V tensors
157 + num_blocks = layer_kv_cache.shape[num_blocks_logical_dim]
158 + half_page_size = layer_kv_cache_spec.page_size_bytes // 2
159 + storage = layer_kv_cache.untyped_storage()
160 + raw = (
161 + torch.tensor(
162 + [],
163 + dtype=torch.int8,
164 + device=layer_kv_cache.device,
165 + )
166 + .set_(storage)
167 + .view(2, num_blocks, half_page_size)
168 + )
169 + tensors_per_block[layer_name] = tuple(raw.unbind(0))
170 +
171 + page_size_bytes[layer_name] = half_page_size
172 +
173 + unpadded_page_size_bytes[layer_name] = page_size_bytes[layer_name]
174 +
175 + elif isinstance(layer_kv_cache_spec, MambaSpec):
176 + state_tensors = kv_caches[layer_name]
177 + assert isinstance(state_tensors, list)
178 +
179 + # re-construct the raw (num_blocks, page_size) tensor
180 + # from the first state tensor
181 + assert len(state_tensors) > 0
182 + first_state_tensor = state_tensors[0]
183 + assert first_state_tensor.storage_offset() == 0
184 + num_blocks = first_state_tensor.shape[0]
185 + tensor = (
186 + torch.tensor(
187 + [],
188 + dtype=torch.int8,
189 + device=first_state_tensor.device,
190 + )
191 + .set_(first_state_tensor.untyped_storage())
192 + .view((num_blocks, layer_kv_cache_spec.page_size_bytes))
193 + )
194 + tensors_per_block[layer_name] = (tensor,)
195 +
196 + page_size_bytes[layer_name] = layer_kv_cache_spec.page_size_bytes
197 + unpadded_page_size_bytes[layer_name] = replace(
198 + layer_kv_cache_spec, page_size_padded=None
199 + ).page_size_bytes
200 +
201 + else:
202 + raise NotImplementedError
203 +
204 + block_tensors: list[CanonicalKVCacheTensor] = []
205 + block_data_refs: dict[str, list[CanonicalKVCacheRef]] = defaultdict(list)
206 + for kv_cache_tensor in self.spec.kv_cache_config.kv_cache_tensors:
207 + tensor_layer_names = kv_cache_tensor.shared_by
208 +
209 + # verify all layers in the group reference the exact same tensors
210 + assert len({len(tensors_per_block[n]) for n in tensor_layer_names}) == 1
211 + assert (
212 + len({tensors_per_block[n][0].data_ptr() for n in tensor_layer_names})
213 + == 1
214 + )
215 + assert (
216 + len({tensors_per_block[n][0].stride() for n in tensor_layer_names}) == 1
217 + )
218 +
219 + # pick the first layer to represent the group
220 + first_layer_name = tensor_layer_names[0]
221 + for tensor in tensors_per_block[first_layer_name]:
222 + block_tensors.append(
223 + CanonicalKVCacheTensor(
224 + tensor=tensor,
225 + page_size_bytes=page_size_bytes[first_layer_name],
226 + )
227 + )
228 +
229 + curr_tensor_idx = len(block_tensors) - 1
230 + for layer_name in tensor_layer_names:
231 + block_data_refs[layer_name].append(
232 + CanonicalKVCacheRef(
233 + tensor_idx=curr_tensor_idx,
234 + page_size_bytes=(unpadded_page_size_bytes[layer_name]),
235 + )
236 + )
237 +
238 + group_data_refs: list[list[CanonicalKVCacheRef]] = []
239 + for kv_cache_group in self.spec.kv_cache_config.kv_cache_groups:
240 + group_refs: list[CanonicalKVCacheRef] = []
241 + for layer_name in kv_cache_group.layer_names:
242 + group_refs += block_data_refs[layer_name]
243 + group_data_refs.append(group_refs)
244 +
245 + canonical_kv_caches = CanonicalKVCaches(
246 + tensors=block_tensors,
247 + group_data_refs=group_data_refs,
248 + )
249 +
250 + self._register_handlers(canonical_kv_caches)
251
252 def register_cross_layers_kv_cache(
253 self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend]
254 ):
82 - cross_layer_name = "ALL_LAYERS"
83 - kv_caches = {cross_layer_name: kv_cache}
84 - attn_backends = {cross_layer_name: attn_backend}
85 - self._register_handlers(kv_caches, attn_backends)
255 + # verify that num_blocks is at physical position 0 in the cross-layers
256 + # tensor layout.
257 + test_shape = attn_backend.get_kv_cache_shape(
258 + num_blocks=1234, block_size=16, num_kv_heads=1, head_size=256
259 + )
260 + num_blocks_logical_dim = test_shape.index(1234) + 1
261 + physical_to_logical = attn_backend.get_kv_cache_stride_order(
262 + include_num_layers_dimension=True
263 + )
264 + num_blocks_physical_dim = physical_to_logical.index(num_blocks_logical_dim)
265 + assert num_blocks_physical_dim == 0
266 +
267 + kv_cache_groups = self.spec.kv_cache_config.kv_cache_groups
268 + assert len(kv_cache_groups) == 1
269 + kv_cache_spec = kv_cache_groups[0].kv_cache_spec
270 + num_layers = len(kv_cache_groups[0].layer_names)
271 + page_size_bytes = kv_cache_spec.page_size_bytes * num_layers
272 +
273 + assert kv_cache.storage_offset() == 0
274 + storage = kv_cache.untyped_storage()
275 + assert len(storage) % page_size_bytes == 0
276 + num_blocks = len(storage) // page_size_bytes
277 + tensor = (
278 + torch.tensor(
279 + [],
280 + dtype=torch.int8,
281 + device=kv_cache.device,
282 + )
283 + .set_(storage)
284 + .view(num_blocks, page_size_bytes)
285 + )
286 + kv_cache_tensor = CanonicalKVCacheTensor(
287 + tensor=tensor, page_size_bytes=page_size_bytes
288 + )
289 + # in cross layers layout, there's currently only a single group
290 + kv_cache_data_ref = CanonicalKVCacheRef(
291 + tensor_idx=0, page_size_bytes=page_size_bytes
292 + )
293 + canonical_kv_caches = CanonicalKVCaches(
294 + tensors=[kv_cache_tensor], group_data_refs=[[kv_cache_data_ref]]
295 + )
296 +
297 + self._register_handlers(canonical_kv_caches)
298
299 def handle_preemptions(self, kv_connector_metadata: OffloadingConnectorMetadata):
300 for job_id, transfer_spec in self._unsubmitted_store_jobs:
vllm/v1/attention/backends/utils.py
+2 -1
@@ -78,9 +78,10 @@ def get_kv_cache_layout():
78 return cache_layout
79
80
81 -def set_kv_cache_layout(cache_layout: KVCacheLayoutType):
81 +def set_kv_cache_layout(cache_layout: KVCacheLayoutType | None):
82 global _KV_CACHE_LAYOUT_OVERRIDE
83 _KV_CACHE_LAYOUT_OVERRIDE = cache_layout
84 + get_kv_cache_layout.cache_clear()
85
86
87 @dataclass
vllm/v1/kv_offload/cpu/spec.py
+4 -14
@@ -2,17 +2,14 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from collections.abc import Iterator
4
5 -import torch
6 -
5 from vllm.config import VllmConfig
6 from vllm.platforms import current_platform
9 -from vllm.v1.attention.backend import AttentionBackend
7 from vllm.v1.kv_cache_interface import KVCacheConfig
8 from vllm.v1.kv_offload.abstract import LoadStoreSpec, OffloadingManager
9 from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
10 from vllm.v1.kv_offload.mediums import CPULoadStoreSpec, GPULoadStoreSpec
11 from vllm.v1.kv_offload.reuse_manager import FilterReusedOffloadingManager
15 -from vllm.v1.kv_offload.spec import OffloadingSpec
12 +from vllm.v1.kv_offload.spec import CanonicalKVCaches, OffloadingSpec
13 from vllm.v1.kv_offload.worker.cpu_gpu import CpuGpuOffloadingHandlers
14 from vllm.v1.kv_offload.worker.worker import OffloadingHandler
15
@@ -90,9 +87,7 @@ class CPUOffloadingSpec(OffloadingSpec):
87 return self._manager
88
89 def get_handlers(
93 - self,
94 - kv_caches: dict[str, torch.Tensor],
95 - attn_backends: dict[str, type[AttentionBackend]],
90 + self, kv_caches: CanonicalKVCaches
91 ) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], OffloadingHandler]]:
92 if not self._handlers:
93 if not current_platform.is_cuda_alike():
@@ -100,15 +95,10 @@ class CPUOffloadingSpec(OffloadingSpec):
95 "CPU Offloading is currently only supported on CUDA-alike GPUs"
96 )
97
103 - assert len(self.gpu_block_size) == 1
104 - gpu_block_size = self.gpu_block_size[0]
105 -
98 self._handlers = CpuGpuOffloadingHandlers(
107 - attn_backends=attn_backends,
108 - gpu_block_size=gpu_block_size,
109 - cpu_block_size=gpu_block_size * self.block_size_factor,
99 + kv_caches=kv_caches,
100 + block_size_factor=self.block_size_factor,
101 num_cpu_blocks=self.num_blocks,
111 - gpu_caches=kv_caches,
102 )
103
104 assert self._handlers is not None
vllm/v1/kv_offload/spec.py
+53 -6
@@ -2,12 +2,12 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from abc import ABC, abstractmethod
4 from collections.abc import Iterator
5 +from dataclasses import dataclass
6 from typing import TYPE_CHECKING
7
8 import torch
9
10 from vllm.logger import init_logger
10 -from vllm.v1.attention.backend import AttentionBackend
11 from vllm.v1.kv_offload.abstract import LoadStoreSpec, OffloadingManager
12 from vllm.v1.kv_offload.worker.worker import OffloadingHandler
13
@@ -18,6 +18,56 @@ if TYPE_CHECKING:
18 logger = init_logger(__name__)
19
20
21 +@dataclass
22 +class CanonicalKVCacheTensor:
23 + """
24 + A canonicalized KV cache tensor whose first dimension is num_blocks.
25 +
26 + For attention backends where the raw tensor has num_blocks at a
27 + non-leading physical dimension (e.g. FlashAttention's
28 + (2, num_blocks, ...) layout), the tensor is split so that each
29 + resulting CanonicalKVCacheTensor starts with (num_blocks, ...).
30 + """
31 +
32 + # The KV cache tensor with shape (num_blocks, ...)
33 + tensor: torch.Tensor
34 + # The (possibly padded) page size per block in bytes
35 + page_size_bytes: int
36 +
37 +
38 +@dataclass
39 +class CanonicalKVCacheRef:
40 + """
41 + Per-layer (or group of layers) reference to a specific (by index)
42 + CanonicalKVCacheTensor and records the un-padded page size used by that layer.
43 + """
44 +
45 + # Index into the list of CanonicalKVCacheTensor objects
46 + tensor_idx: int
47 + # The un-padded page size per block in bytes
48 + page_size_bytes: int
49 +
50 +
51 +@dataclass
52 +class CanonicalKVCaches:
53 + """
54 + Canonicalized block-level representation of the KV caches.
55 +
56 + Composed of:
57 + - Unique list of KV cache data tensors,
58 + each with shape (num_blocks, page_size_in_bytes) and int8 dtype.
59 + - Per-group data references of the tensors.
60 + i.e. how each KV cache group maps to the tensors.
61 + """
62 +
63 + # Ordered list of unique block tensors, each with shape
64 + # (num_blocks, ...).
65 + tensors: list[CanonicalKVCacheTensor]
66 + # Per-KV-cache-group list of data references that map each layer
67 + # in the group to the appropriate entry in the tensors list.
68 + group_data_refs: list[list[CanonicalKVCacheRef]]
69 +
70 +
71 class OffloadingSpec(ABC):
72 """Spec for an offloading connector"""
73
@@ -73,16 +123,13 @@ class OffloadingSpec(ABC):
123
124 @abstractmethod
125 def get_handlers(
76 - self,
77 - kv_caches: dict[str, torch.Tensor],
78 - attn_backends: dict[str, type[AttentionBackend]],
126 + self, kv_caches: CanonicalKVCaches
127 ) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], OffloadingHandler]]:
128 """
129 Get offloading handlers along with their respective src and dst types.
130
131 Args:
84 - kv_caches: A dictionary of layer_name -> gpu_kv_cache tensor.
85 - attn_backends: A dictionary of layer_name -> AttentionBackend.
132 + kv_caches: Canonicalized KV caches.
133
134 Yields:
135 Tuples of (src_type, dst_type, offloading_handler).
vllm/v1/kv_offload/worker/cpu_gpu.py
+82 -109
@@ -9,8 +9,8 @@ import torch
9 from vllm import _custom_ops as ops
10 from vllm.logger import init_logger
11 from vllm.utils.platform_utils import is_pin_memory_available
12 -from vllm.v1.attention.backend import AttentionBackend
12 from vllm.v1.kv_offload.mediums import BlockIDsLoadStoreSpec
13 +from vllm.v1.kv_offload.spec import CanonicalKVCacheRef, CanonicalKVCaches
14 from vllm.v1.kv_offload.worker.worker import (
15 OffloadingHandler,
16 TransferResult,
@@ -73,39 +73,72 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
73
74 def __init__(
75 self,
76 - src_tensors: list[torch.Tensor],
77 - dst_tensors: list[torch.Tensor],
78 - src_block_size_factor: int,
79 - dst_block_size_factor: int,
76 + gpu_tensors: list[torch.Tensor],
77 + cpu_tensors: list[torch.Tensor],
78 + block_size_factor: int,
79 + kv_cache_groups_data_refs: list[list[CanonicalKVCacheRef]],
80 + gpu_to_cpu: bool,
81 ):
82 """
83 Initialize a SingleDirectionOffloadingHandler.
84
85 Args:
85 - src_tensors: list of KV cache tensors to copy from.
86 - dst_tensors: list of KV cache tensors to copy to.
87 - Order should match src_tensors.
88 - src_block_size_factor: The number of kernel blocks
89 - per KV block in a source tensor.
90 - dst_block_size_factor: The number of kernel blocks
91 - per KV block in a destination tensor.
86 + gpu_tensors: list of GPU KV cache tensors.
87 + Each of shape (num_gpu_blocks, gpu_page_size_bytes) with dtype int8.
88 + cpu_tensors: list of CPU KV cache tensors.
89 + Each of shape (num_cpu_blocks, cpu_page_size_bytes) with dtype int8.
90 + Order should match gpu_tensors.
91 + kv_cache_groups_data_refs: list of CanonicalKVCacheRef per group.
92 + gpu_to_cpu: if True, transfer from GPU to CPU; otherwise CPU to GPU.
93 """
93 - assert len(src_tensors) == len(dst_tensors)
94 + assert len(gpu_tensors) == len(cpu_tensors)
95 + assert len(gpu_tensors) > 0
96 +
97 + # assert a single KV group until transfer_async supports multiple groups
98 + assert len(kv_cache_groups_data_refs) == 1
99 +
100 + # assert input tensors are as expected
101 + for gpu_tensor, cpu_tensor in zip(gpu_tensors, cpu_tensors):
102 + assert gpu_tensor.dtype == torch.int8
103 + assert gpu_tensor.ndim == 2
104 + assert gpu_tensor.is_cuda
105 + assert cpu_tensor.dtype == torch.int8
106 + assert cpu_tensor.ndim == 2
107 + assert cpu_tensor.device.type == "cpu"
108 + _, gpu_page_size = gpu_tensor.shape
109 + _, cpu_page_size = cpu_tensor.shape
110 + assert cpu_page_size == gpu_page_size * block_size_factor
111 +
112 + self.src_tensors: list[torch.Tensor] = (
113 + gpu_tensors if gpu_to_cpu else cpu_tensors
114 + )
115 + self.dst_tensors: list[torch.Tensor] = (
116 + cpu_tensors if gpu_to_cpu else gpu_tensors
117 + )
118 + self.gpu_to_cpu: bool = gpu_to_cpu
119
95 - self.src_tensors: list[torch.Tensor] = src_tensors
96 - self.dst_tensors: list[torch.Tensor] = dst_tensors
97 - min_block_size_factor = min(src_block_size_factor, dst_block_size_factor)
98 - self.src_block_size_factor: int = src_block_size_factor // min_block_size_factor
99 - self.dst_block_size_factor: int = dst_block_size_factor // min_block_size_factor
120 + # GPU blocks may be smaller
121 + # cpu_page_size = gpu_page_size * block_size_factor.
122 + self.src_block_size_factor = 1 if self.gpu_to_cpu else block_size_factor
123 + self.dst_block_size_factor = block_size_factor if self.gpu_to_cpu else 1
124
101 - self.block_size_in_bytes = [
102 - tensor.element_size() * tensor.stride(0) * min_block_size_factor
103 - for tensor in src_tensors
125 + # per-tensor block size in byte
126 + self.tensor_block_size_in_bytes = [
127 + gpu_tensor.shape[1] for gpu_tensor in gpu_tensors
128 ]
105 - self.total_block_size_in_bytes = sum(self.block_size_in_bytes)
129
107 - assert len(src_tensors) > 0
108 - self.gpu_to_cpu: bool = self.src_tensors[0].is_cuda
130 + # per-group block size in bytes
131 + self.group_block_size_in_bytes = []
132 + for kv_cache_group_data_refs in kv_cache_groups_data_refs:
133 + group_block_size_in_bytes = 0
134 + for kv_cache_data_ref in kv_cache_group_data_refs:
135 + # TODO(orozery): use kv_cache_data_ref.page_size_bytes
136 + # once swap_blocks support it
137 + group_block_size_in_bytes += self.tensor_block_size_in_bytes[
138 + kv_cache_data_ref.tensor_idx
139 + ]
140 + self.group_block_size_in_bytes.append(group_block_size_in_bytes)
141 +
142 self.transfer_type = ("GPU", "CPU") if self.gpu_to_cpu else ("CPU", "GPU")
143 # job_id -> event
144 self._transfer_events: dict[int, torch.Event] = {}
@@ -167,7 +200,7 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
200 for src_tensor, dst_tensor, block_size_in_bytes in zip(
201 self.src_tensors,
202 self.dst_tensors,
170 - self.block_size_in_bytes,
203 + self.tensor_block_size_in_bytes,
204 ):
205 ops.swap_blocks(
206 src_tensor,
@@ -184,7 +217,7 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
217 stream=stream,
218 start_event=start_event,
219 end_event=end_event,
187 - num_bytes=dst_sub_block_count * self.total_block_size_in_bytes,
220 + num_bytes=dst_sub_block_count * self.group_block_size_in_bytes[0],
221 )
222 )
223
@@ -223,102 +256,42 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
256 class CpuGpuOffloadingHandlers:
257 def __init__(
258 self,
226 - gpu_block_size: int,
227 - cpu_block_size: int,
259 + kv_caches: CanonicalKVCaches,
260 + block_size_factor: int,
261 num_cpu_blocks: int,
229 - gpu_caches: dict[str, torch.Tensor],
230 - attn_backends: dict[str, type[AttentionBackend]],
262 ):
232 - assert gpu_caches
233 - assert cpu_block_size % gpu_block_size == 0
234 -
235 - # find kernel block size and determine layout per each gpu tensor
236 - kernel_block_size: int | None = None
237 - # list of (gpu_tensor, split_k_and_v)
238 - parsed_gpu_tensors: list[tuple[torch.Tensor, bool]] = []
239 - for layer_name, gpu_tensor in gpu_caches.items():
240 - gpu_shape = gpu_tensor.shape
241 - attn_backend = attn_backends[layer_name]
242 - test_shape = attn_backend.get_kv_cache_shape(
243 - num_blocks=1234, block_size=16, num_kv_heads=1, head_size=256
244 - )
245 -
246 - has_layers_dim = False
247 - split_k_and_v = False
248 - if len(gpu_shape) != len(test_shape):
249 - # cross-layers tensor
250 - # shape is (num_blocks, ...)
251 - assert len(gpu_shape) == len(test_shape) + 1
252 - has_layers_dim = True
253 - # prepend a dummy num_layers=80 to test_shape
254 - test_shape = (80,) + test_shape
255 - elif test_shape[0] != 1234:
256 - # shape should be (2, num_blocks, ...)
257 - assert test_shape[0] == 2
258 - assert test_shape[1] == 1234
259 - assert gpu_shape[0] == 2
260 - split_k_and_v = True
261 -
262 - if has_layers_dim:
263 - # in the cross layers case, the registered kv cache tensor
264 - # shape matches the physical layout, whereas test_shape
265 - # is the logical layout.
266 - # To match them, we need to permute test_shape
267 - try:
268 - kv_cache_stride_order = attn_backend.get_kv_cache_stride_order(
269 - include_num_layers_dimension=has_layers_dim
270 - )
271 - assert len(kv_cache_stride_order) == len(gpu_shape)
272 - except (AttributeError, NotImplementedError):
273 - kv_cache_stride_order = tuple(range(len(gpu_shape)))
274 -
275 - test_shape = tuple(test_shape[i] for i in kv_cache_stride_order)
276 -
277 - # find block_size (16) dimension index
278 - block_size_idx = test_shape.index(16)
279 - if kernel_block_size is not None:
280 - assert kernel_block_size == gpu_shape[block_size_idx]
281 - else:
282 - kernel_block_size = gpu_shape[block_size_idx]
283 - assert gpu_block_size % kernel_block_size == 0
284 -
285 - parsed_gpu_tensors.append((gpu_tensor, split_k_and_v))
286 -
287 - assert kernel_block_size is not None
288 - cpu_block_size_factor = cpu_block_size // kernel_block_size
289 - gpu_block_size_factor = gpu_block_size // kernel_block_size
290 - num_cpu_kernel_blocks = num_cpu_blocks * cpu_block_size_factor
291 -
292 - # allocate cpu tensors
263 pin_memory = is_pin_memory_available()
294 - logger.info("Allocating %d CPU tensors...", len(parsed_gpu_tensors))
264 + logger.info("Allocating %d CPU tensors...", len(kv_caches.tensors))
265 gpu_tensors: list[torch.Tensor] = []
266 cpu_tensors: list[torch.Tensor] = []
297 - for gpu_tensor, split_k_and_v in parsed_gpu_tensors:
298 - cpu_shape = list(gpu_tensor.shape)
299 - cpu_shape[1 if split_k_and_v else 0] = num_cpu_kernel_blocks
300 -
301 - logger.debug("Allocating CPU tensor of shape %r", cpu_shape)
267 + for kv_cache_tensor in kv_caches.tensors:
268 + gpu_page_size_bytes = kv_cache_tensor.page_size_bytes
269 + gpu_tensor = kv_cache_tensor.tensor.view(torch.int8).view(
270 + (-1, gpu_page_size_bytes)
271 + )
272 + cpu_page_size_bytes = gpu_page_size_bytes * block_size_factor
273 cpu_tensor = torch.zeros(
303 - cpu_shape,
304 - dtype=gpu_tensor.dtype,
274 + (num_cpu_blocks, cpu_page_size_bytes),
275 + dtype=torch.int8,
276 device="cpu",
277 pin_memory=pin_memory,
278 )
279
309 - gpu_tensors.extend(gpu_tensor.unbind(0) if split_k_and_v else [gpu_tensor])
310 - cpu_tensors.extend(cpu_tensor.unbind(0) if split_k_and_v else [cpu_tensor])
280 + gpu_tensors.append(gpu_tensor)
281 + cpu_tensors.append(cpu_tensor)
282
283 self.gpu_to_cpu_handler = SingleDirectionOffloadingHandler(
313 - src_tensors=gpu_tensors,
314 - dst_tensors=cpu_tensors,
315 - src_block_size_factor=gpu_block_size_factor,
316 - dst_block_size_factor=cpu_block_size_factor,
284 + gpu_tensors=gpu_tensors,
285 + cpu_tensors=cpu_tensors,
286 + block_size_factor=block_size_factor,
287 + kv_cache_groups_data_refs=kv_caches.group_data_refs,
288 + gpu_to_cpu=True,
289 )
290
291 self.cpu_to_gpu_handler = SingleDirectionOffloadingHandler(
320 - src_tensors=cpu_tensors,
321 - dst_tensors=gpu_tensors,
322 - src_block_size_factor=cpu_block_size_factor,
323 - dst_block_size_factor=gpu_block_size_factor,
292 + gpu_tensors=gpu_tensors,
293 + cpu_tensors=cpu_tensors,
294 + block_size_factor=block_size_factor,
295 + kv_cache_groups_data_refs=kv_caches.group_data_refs,
296 + gpu_to_cpu=False,
297 )