@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 91e4521f9f
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, 17:15 UTC

[Feat][v1] Simple yet General CPU KV Cache Offloading (#37160)

Signed-off-by: Yifan Qiao <yifanqiao@berkeley.edu> Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>

Yifan Qiao committed Mar 31, 2026 at 17:58 UTC 91e4521f9f46f24aff9c77f275ab581af6776445
15 files changed +2964 -3
tests/v1/simple_kv_offload/__init__.py
tests/v1/simple_kv_offload/test_integration.py new
+193
@@ -0,0 +1,193 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Integration tests for SimpleCPUOffloadConnector with real models."""
4 +
5 +import time
6 +
7 +import pytest
8 +
9 +from vllm import LLM, SamplingParams, TokensPrompt
10 +from vllm.config import KVTransferConfig
11 +from vllm.platforms import current_platform
12 +
13 +if not current_platform.is_cuda():
14 + pytest.skip("Requires CUDA", allow_module_level=True)
15 +
16 +# Small models for default CI / local runs (accuracy only).
17 +SMALL_MODELS = [
18 + "meta-llama/Llama-3.2-1B-Instruct",
19 + "google/gemma-3-1b-it",
20 +]
21 +
22 +# Large models for optional perf runs only (slow to load and execute).
23 +PERF_MODELS = [
24 + "meta-llama/Llama-3.1-8B",
25 + "openai/gpt-oss-20b",
26 +]
27 +
28 +
29 +def _make_llm(model: str, lazy: bool, cpu_bytes_to_use: int) -> LLM:
30 + kv_transfer_config = KVTransferConfig(
31 + kv_connector="SimpleCPUOffloadConnector",
32 + kv_role="kv_both",
33 + kv_connector_extra_config={
34 + "cpu_bytes_to_use": cpu_bytes_to_use,
35 + "lazy_offload": lazy,
36 + },
37 + )
38 + return LLM(
39 + model=model,
40 + kv_cache_memory_bytes=40 << 30, # 40 GiB
41 + disable_hybrid_kv_cache_manager=False,
42 + enable_prefix_caching=True,
43 + kv_transfer_config=kv_transfer_config,
44 + )
45 +
46 +
47 +def _flush_gpu_cache(llm: LLM, sampling_params: SamplingParams, seed: int = 0):
48 + """Generate enough filler requests to allocate the entire GPU KV cache.
49 +
50 + This pushes all prior blocks through the free queue so that the lazy
51 + cursor offloads them to CPU before they are evicted.
52 + """
53 + cache_config = llm.llm_engine.vllm_config.cache_config
54 + num_gpu_blocks = cache_config.num_gpu_blocks
55 + block_size = cache_config.block_size
56 + # Use 1.2x GPU capacity to give the lazy cursor enough scheduling steps
57 + # to walk past all target blocks near the tail of the free queue.
58 + total_tokens_needed = int(num_gpu_blocks * block_size * 1.5)
59 +
60 + # Use token-id prompts so each filler is unique (no prefix sharing).
61 + # Split into multiple requests to stay under max_model_len.
62 + max_tokens_per_req = 4096
63 + num_fillers = (total_tokens_needed + max_tokens_per_req - 1) // max_tokens_per_req
64 + batch_size = 10
65 + for i in range(0, num_fillers, batch_size):
66 + batch_end = min(i + batch_size, num_fillers)
67 + filler_prompts = []
68 + for j in range(i, batch_end):
69 + ids = [seed * num_fillers + j + 1] * max_tokens_per_req
70 + filler_prompts.append(TokensPrompt(prompt_token_ids=ids))
71 + llm.generate(filler_prompts, sampling_params, use_tqdm=False)
72 +
73 +
74 +def _accuracy_test(llm: LLM, lazy: bool = False):
75 + """Verify that CPU-loaded KV produces correct output."""
76 + sampling_params = SamplingParams(max_tokens=1, temperature=0)
77 + prompt = "hi " * 2000 + "Let's count to ten. One, two, three, "
78 +
79 + # Cold run — populate GPU cache and trigger CPU offload
80 + cold_output = llm.generate(prompt, sampling_params, use_tqdm=False)[0]
81 +
82 + # CPU hit runs
83 + test_count = 10
84 + success_count = 0
85 + expected = cold_output.outputs[0].text
86 + for i in range(test_count):
87 + if lazy:
88 + _flush_gpu_cache(llm, sampling_params, seed=i)
89 + time.sleep(2) # let engine core drain pending transfers
90 +
91 + # Reset GPU prefix cache so next run must load from CPU
92 + if not llm.reset_prefix_cache():
93 + print(f"GPU prefix cache reset failed for iteration {i}")
94 +
95 + output = llm.generate(prompt, sampling_params, use_tqdm=False)[0]
96 + if output.outputs[0].text == expected:
97 + success_count += 1
98 +
99 + assert success_count >= 0.5 * test_count, (
100 + f"Accuracy too low: {success_count}/{test_count} matched '{expected}'"
101 + )
102 +
103 +
104 +def _latency_test(llm: LLM, lazy: bool = False):
105 + """Verify CPU cache hit is faster than cold compute."""
106 + sampling_params = SamplingParams(max_tokens=1, seed=42)
107 + prompt_token_ids = [0] * 10001
108 +
109 + num_times_cpu_better = 0
110 + num_tests = 10
111 + for i in range(num_tests):
112 + prompt_token_ids[0] = i
113 + prompts = [TokensPrompt(prompt_token_ids=prompt_token_ids)]
114 +
115 + # Cold
116 + time.sleep(2) # let engine core drain pending transfers
117 + if not llm.reset_prefix_cache():
118 + print(f"GPU prefix cache reset failed for iteration {i}")
119 + start = time.time()
120 + llm.generate(prompts, sampling_params, use_tqdm=False)
121 + cold_time = time.time() - start
122 +
123 + if lazy:
124 + _flush_gpu_cache(llm, sampling_params, seed=i)
125 + else:
126 + # Eager mode: GPU hit ensures store completion is processed.
127 + llm.generate(prompts, sampling_params, use_tqdm=False)
128 +
129 + time.sleep(2) # let engine core drain pending transfers
130 + if not llm.reset_prefix_cache():
131 + print(f"GPU prefix cache reset failed for iteration {i}")
132 +
133 + # CPU hit
134 + start = time.time()
135 + llm.generate(prompts, sampling_params, use_tqdm=False)
136 + cpu_time = time.time() - start
137 +
138 + if cpu_time < cold_time:
139 + num_times_cpu_better += 1
140 +
141 + assert num_times_cpu_better >= 0.8 * num_tests, (
142 + f"CPU hit only faster {num_times_cpu_better}/{num_tests} times"
143 + )
144 +
145 +
146 +@pytest.mark.optional
147 +@pytest.mark.slow_test
148 +@pytest.mark.parametrize("model", SMALL_MODELS)
149 +def test_simple_cpu_offload_accuracy(model: str):
150 + """Store to CPU, reset GPU, load from CPU; verify output matches baseline."""
151 + llm = _make_llm(model, False, 1 << 30) # 1GB
152 + try:
153 + _accuracy_test(llm, lazy=False)
154 + finally:
155 + del llm
156 +
157 +
158 +@pytest.mark.optional
159 +@pytest.mark.slow_test
160 +@pytest.mark.parametrize("model", PERF_MODELS)
161 +def test_simple_cpu_offload_perf_latency(model: str):
162 + """CPU KV hit should beat cold prefill on long context (large models only)."""
163 + llm = _make_llm(model, False, 10 << 30) # 10GB
164 + try:
165 + _latency_test(llm, lazy=False)
166 + finally:
167 + del llm
168 +
169 +
170 +@pytest.mark.optional
171 +@pytest.mark.slow_test
172 +@pytest.mark.parametrize("model", SMALL_MODELS)
173 +def test_simple_cpu_offload_accuracy_lazy(model: str):
174 + """Lazy mode: flush GPU cache to trigger CPU offload, then verify hit."""
175 + # CPU must be larger than GPU KV cache to avoid evicting offloaded blocks.
176 + llm = _make_llm(model, True, 80 << 30) # 80GB
177 + try:
178 + _accuracy_test(llm, lazy=True)
179 + finally:
180 + del llm
181 +
182 +
183 +@pytest.mark.optional
184 +@pytest.mark.slow_test
185 +@pytest.mark.parametrize("model", PERF_MODELS)
186 +def test_simple_cpu_offload_perf_latency_lazy(model: str):
187 + """Lazy mode: CPU KV hit should beat cold prefill (large models only)."""
188 + # CPU must be larger than GPU KV cache to avoid evicting offloaded blocks.
189 + llm = _make_llm(model, True, 80 << 30) # 80GB
190 + try:
191 + _latency_test(llm, lazy=True)
192 + finally:
193 + del llm
tests/v1/simple_kv_offload/test_scheduler.py new
+1137
@@ -0,0 +1,1137 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for SimpleCPUOffloadScheduler."""
4 +
5 +from __future__ import annotations
6 +
7 +from dataclasses import dataclass
8 +
9 +import torch
10 +
11 +from vllm import SamplingParams
12 +from vllm.config import (
13 + CacheConfig,
14 + DeviceConfig,
15 + KVTransferConfig,
16 + ModelConfig,
17 + SchedulerConfig,
18 + VllmConfig,
19 +)
20 +from vllm.utils.hashing import sha256
21 +from vllm.v1.core.block_pool import BlockPool
22 +from vllm.v1.core.kv_cache_manager import KVCacheBlocks
23 +from vllm.v1.core.kv_cache_utils import (
24 + get_request_block_hasher,
25 + init_none_hash,
26 + make_block_hash_with_group_id,
27 +)
28 +from vllm.v1.core.sched.output import (
29 + CachedRequestData,
30 + NewRequestData,
31 + SchedulerOutput,
32 +)
33 +from vllm.v1.kv_cache_interface import (
34 + FullAttentionSpec,
35 + KVCacheConfig,
36 + KVCacheGroupSpec,
37 + KVCacheTensor,
38 +)
39 +from vllm.v1.outputs import KVConnectorOutput
40 +from vllm.v1.request import Request
41 +from vllm.v1.simple_kv_offload.manager import SimpleCPUOffloadScheduler
42 +from vllm.v1.simple_kv_offload.metadata import SimpleCPUOffloadWorkerMetadata
43 +
44 +# ---------------------------------------------------------------------------
45 +# Constants
46 +# ---------------------------------------------------------------------------
47 +BLOCK_SIZE = 16
48 +HEAD_SIZE = 16
49 +NUM_KV_HEADS = 1
50 +DTYPE = torch.float16
51 +# bytes per block per tensor:
52 +# block_size * num_kv_heads * head_size * 2 (K+V) * element_size
53 +_BYTES_PER_BLOCK = BLOCK_SIZE * NUM_KV_HEADS * HEAD_SIZE * 2 * DTYPE.itemsize
54 +
55 +# Ensure none_hash is initialized once
56 +init_none_hash(sha256)
57 +
58 +
59 +# ---------------------------------------------------------------------------
60 +# Helpers
61 +# ---------------------------------------------------------------------------
62 +
63 +
64 +def _make_kv_cache_config(
65 + num_blocks: int,
66 + num_groups: int = 1,
67 +) -> KVCacheConfig:
68 + """Build a KVCacheConfig with non-empty kv_cache_tensors."""
69 + groups = []
70 + tensors = []
71 + for g in range(num_groups):
72 + layer_names = [f"layer_{g}"]
73 + groups.append(
74 + KVCacheGroupSpec(
75 + layer_names,
76 + FullAttentionSpec(
77 + block_size=BLOCK_SIZE,
78 + num_kv_heads=NUM_KV_HEADS,
79 + head_size=HEAD_SIZE,
80 + dtype=DTYPE,
81 + ),
82 + )
83 + )
84 + tensors.append(
85 + KVCacheTensor(
86 + size=_BYTES_PER_BLOCK * num_blocks,
87 + shared_by=layer_names,
88 + )
89 + )
90 + return KVCacheConfig(
91 + num_blocks=num_blocks,
92 + kv_cache_tensors=tensors,
93 + kv_cache_groups=groups,
94 + )
95 +
96 +
97 +def _make_vllm_config(block_size: int = BLOCK_SIZE) -> VllmConfig:
98 + """Minimal VllmConfig for scheduler tests (no GPU)."""
99 + model_config = ModelConfig(
100 + model="facebook/opt-125m",
101 + trust_remote_code=True,
102 + dtype="float16",
103 + seed=42,
104 + )
105 + scheduler_config = SchedulerConfig(
106 + max_num_seqs=16,
107 + max_num_batched_tokens=64,
108 + max_model_len=10000,
109 + enable_chunked_prefill=True,
110 + is_encoder_decoder=False,
111 + )
112 + cache_config = CacheConfig(
113 + block_size=block_size,
114 + gpu_memory_utilization=0.9,
115 + enable_prefix_caching=True,
116 + )
117 + kv_transfer_config = KVTransferConfig(
118 + kv_connector="SimpleCPUOffloadConnector",
119 + kv_role="kv_both",
120 + )
121 + return VllmConfig(
122 + scheduler_config=scheduler_config,
123 + model_config=model_config,
124 + cache_config=cache_config,
125 + kv_transfer_config=kv_transfer_config,
126 + device_config=DeviceConfig("cpu"),
127 + )
128 +
129 +
130 +@dataclass
131 +class SchedulerFixture:
132 + """Bundle returned by make_scheduler for convenient access."""
133 +
134 + scheduler: SimpleCPUOffloadScheduler
135 + gpu_block_pool: BlockPool
136 + vllm_config: VllmConfig
137 + kv_cache_config: KVCacheConfig
138 + num_groups: int = 1
139 +
140 +
141 +def make_scheduler(
142 + num_cpu_blocks: int = 8,
143 + num_gpu_blocks: int = 16,
144 + num_groups: int = 1,
145 + lazy: bool = False,
146 +) -> SchedulerFixture:
147 + """Build a SimpleCPUOffloadScheduler with small block pools."""
148 + kv_cache_config = _make_kv_cache_config(num_gpu_blocks, num_groups)
149 + vllm_config = _make_vllm_config()
150 + cpu_capacity_bytes = _BYTES_PER_BLOCK * num_cpu_blocks * num_groups
151 +
152 + sched = SimpleCPUOffloadScheduler(
153 + vllm_config=vllm_config,
154 + kv_cache_config=kv_cache_config,
155 + cpu_capacity_bytes=cpu_capacity_bytes,
156 + lazy_offload=lazy,
157 + )
158 +
159 + # Build a real GPU block pool and bind it
160 + gpu_block_pool = BlockPool(
161 + num_gpu_blocks=num_gpu_blocks,
162 + enable_caching=True,
163 + hash_block_size=BLOCK_SIZE,
164 + )
165 + sched.bind_gpu_block_pool(gpu_block_pool)
166 +
167 + return SchedulerFixture(
168 + scheduler=sched,
169 + gpu_block_pool=gpu_block_pool,
170 + vllm_config=vllm_config,
171 + kv_cache_config=kv_cache_config,
172 + num_groups=num_groups,
173 + )
174 +
175 +
176 +_req_counter = 0
177 +
178 +
179 +def make_request(
180 + num_blocks: int = 2,
181 + request_id: str | None = None,
182 +) -> Request:
183 + """Create a Request with deterministic block hashes."""
184 + global _req_counter
185 + _req_counter += 1
186 + if request_id is None:
187 + request_id = f"req-{_req_counter}"
188 +
189 + num_tokens = num_blocks * BLOCK_SIZE
190 + start = _req_counter * 10000
191 + prompt_token_ids = list(range(start, start + num_tokens))
192 + sampling_params = SamplingParams(max_tokens=1)
193 +
194 + req = Request(
195 + request_id=request_id,
196 + prompt_token_ids=prompt_token_ids,
197 + sampling_params=sampling_params,
198 + pooling_params=None,
199 + mm_features=None,
200 + block_hasher=get_request_block_hasher(BLOCK_SIZE, sha256),
201 + )
202 + return req
203 +
204 +
205 +def make_scheduler_output(
206 + req_id_to_num_tokens: dict[str, int],
207 + *,
208 + new_reqs: dict[str, tuple[list[int], ...]] | None = None,
209 + cached_req_new_blocks: dict[str, tuple[list[int], ...] | None] | None = None,
210 +) -> SchedulerOutput:
211 + """Build a minimal SchedulerOutput with num_scheduled_tokens.
212 +
213 + Args:
214 + new_reqs: For first-time requests, maps req_id -> block_ids tuple.
215 + These are placed into ``scheduled_new_reqs`` as ``NewRequestData``.
216 + cached_req_new_blocks: For returning (cached) requests, maps
217 + req_id -> new_block_ids (incremental) or None.
218 + These are placed into ``scheduled_cached_reqs``.
219 + """
220 + scheduled_new_reqs: list[NewRequestData] = []
221 + if new_reqs:
222 + for req_id, block_ids in new_reqs.items():
223 + scheduled_new_reqs.append(
224 + NewRequestData(
225 + req_id=req_id,
226 + prompt_token_ids=None,
227 + mm_features=[],
228 + sampling_params=None,
229 + pooling_params=None,
230 + block_ids=block_ids,
231 + num_computed_tokens=0,
232 + lora_request=None,
233 + )
234 + )
235 +
236 + if cached_req_new_blocks:
237 + cached_req_ids = list(cached_req_new_blocks.keys())
238 + cached_new_block_ids = [cached_req_new_blocks[rid] for rid in cached_req_ids]
239 + cached_reqs = CachedRequestData(
240 + req_ids=cached_req_ids,
241 + resumed_req_ids=set(),
242 + new_token_ids=[[] for _ in cached_req_ids],
243 + all_token_ids={},
244 + new_block_ids=cached_new_block_ids,
245 + num_computed_tokens=[0] * len(cached_req_ids),
246 + num_output_tokens=[0] * len(cached_req_ids),
247 + )
248 + else:
249 + cached_reqs = CachedRequestData.make_empty()
250 +
251 + return SchedulerOutput(
252 + scheduled_new_reqs=scheduled_new_reqs,
253 + scheduled_cached_reqs=cached_reqs,
254 + num_scheduled_tokens=req_id_to_num_tokens,
255 + total_num_scheduled_tokens=sum(req_id_to_num_tokens.values()),
256 + scheduled_spec_decode_tokens={},
257 + scheduled_encoder_inputs={},
258 + num_common_prefix_blocks=[],
259 + preempted_req_ids=set(),
260 + finished_req_ids=set(),
261 + free_encoder_mm_hashes=[],
262 + )
263 +
264 +
265 +def simulate_store_completion(
266 + scheduler: SimpleCPUOffloadScheduler,
267 + event_idx: int,
268 +) -> None:
269 + """Simulate worker reporting a store event completion."""
270 + output = KVConnectorOutput(
271 + finished_recving=set(),
272 + kv_connector_worker_meta=SimpleCPUOffloadWorkerMetadata(
273 + completed_store_events={event_idx: scheduler._expected_worker_count},
274 + ),
275 + )
276 + scheduler.update_connector_output(output)
277 +
278 +
279 +def simulate_load_completion(
280 + scheduler: SimpleCPUOffloadScheduler,
281 + req_ids: set[str],
282 +) -> None:
283 + """Simulate worker reporting load completions for requests."""
284 + output = KVConnectorOutput(
285 + finished_sending=set(),
286 + finished_recving=req_ids,
287 + )
288 + scheduler.update_connector_output(output)
289 +
290 +
291 +def get_cpu_free_blocks(scheduler: SimpleCPUOffloadScheduler) -> int:
292 + """Return number of free CPU blocks."""
293 + return scheduler.cpu_block_pool.get_num_free_blocks()
294 +
295 +
296 +def _allocate_gpu_blocks(
297 + gpu_block_pool: BlockPool,
298 + request: Request,
299 + num_blocks: int,
300 + group_id: int = 0,
301 +) -> list:
302 + """Allocate GPU blocks, cache them with hashes, return block list.
303 +
304 + Mimics what KVCacheManager does: allocate blocks from pool, then
305 + register them in the prefix cache via cache_full_blocks so that
306 + re-allocation properly evicts stale hashes.
307 + """
308 + blocks = gpu_block_pool.get_new_blocks(num_blocks)
309 + num_full = min(num_blocks, len(request.block_hashes))
310 + if num_full > 0:
311 + gpu_block_pool.cache_full_blocks(
312 + request=request,
313 + blocks=blocks,
314 + num_cached_blocks=0,
315 + num_full_blocks=num_full,
316 + block_size=BLOCK_SIZE,
317 + kv_cache_group_id=group_id,
318 + )
319 + return blocks
320 +
321 +
322 +def _alloc_and_register(
323 + fix: SchedulerFixture,
324 + request: Request,
325 + num_blocks: int,
326 + *,
327 + confirmed: bool = True,
328 + group_id: int = 0,
329 +) -> KVCacheBlocks:
330 + """Allocate GPU blocks and return KVCacheBlocks.
331 +
332 + Block IDs are no longer registered in a mock KVCacheManager; instead
333 + tests pass them through ``make_scheduler_output`` so that
334 + ``yield_req_data`` can pick them up.
335 +
336 + If ``confirmed`` is True, advance ``request.num_computed_tokens`` to simulate
337 + the scheduler's ``_update_after_schedule`` from a prior step.
338 + """
339 + gpu_blocks = _allocate_gpu_blocks(
340 + fix.gpu_block_pool, request, num_blocks, group_id=group_id
341 + )
342 + kv_blocks = KVCacheBlocks(blocks=(gpu_blocks,))
343 + if confirmed:
344 + request.num_computed_tokens = num_blocks * BLOCK_SIZE
345 + return kv_blocks
346 +
347 +
348 +# ---------------------------------------------------------------------------
349 +# Test 1a: Eager store-and-load roundtrip
350 +# ---------------------------------------------------------------------------
351 +def test_eager_store_and_load_roundtrip() -> None:
352 + """Eager mode: store blocks on compute, complete store, verify cache hit."""
353 + fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False)
354 + sched = fix.scheduler
355 +
356 + num_blocks = 2
357 + req = make_request(num_blocks=num_blocks)
358 +
359 + kv_blocks = _alloc_and_register(fix, req, num_blocks)
360 + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
361 + block_ids = kv_blocks.get_block_ids()
362 + sched_out = make_scheduler_output(
363 + {req.request_id: num_blocks * BLOCK_SIZE},
364 + new_reqs={req.request_id: block_ids},
365 + )
366 +
367 + meta = sched.build_connector_meta(sched_out)
368 + assert meta.store_event >= 0, "Expected a store event to be scheduled"
369 + assert len(meta.store_gpu_blocks) > 0
370 + assert len(meta.store_cpu_blocks) == len(meta.store_gpu_blocks)
371 + simulate_store_completion(sched, meta.store_event)
372 +
373 + # New request with same tokens should get CPU cache hit
374 + req2 = Request(
375 + request_id="req-eager-load",
376 + prompt_token_ids=req.prompt_token_ids,
377 + sampling_params=req.sampling_params,
378 + pooling_params=None,
379 + mm_features=None,
380 + block_hasher=req._block_hasher,
381 + )
382 + hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)
383 + assert hit_tokens == num_blocks * BLOCK_SIZE
384 + assert is_async is True
385 +
386 + gpu_blocks2 = fix.gpu_block_pool.get_new_blocks(num_blocks)
387 + kv_blocks2 = KVCacheBlocks(blocks=(gpu_blocks2,))
388 + sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens)
389 +
390 + block_ids2 = kv_blocks2.get_block_ids()
391 + sched_out2 = make_scheduler_output(
392 + {req2.request_id: 1},
393 + new_reqs={req2.request_id: block_ids2},
394 + )
395 + meta2 = sched.build_connector_meta(sched_out2)
396 + assert meta2.load_event >= 0, "Expected a load event to be assigned"
397 + assert len(meta2.load_gpu_blocks) > 0
398 + assert len(meta2.load_cpu_blocks) == len(meta2.load_gpu_blocks)
399 +
400 +
401 +# ---------------------------------------------------------------------------
402 +# Test 1b: Lazy store-and-load roundtrip
403 +# ---------------------------------------------------------------------------
404 +def _flush_old_blocks_to_lru_head(
405 + gpu_pool: BlockPool,
406 + num_filler_blocks: int,
407 +) -> list:
408 + """Allocate filler blocks so that previously-freed (hashed) blocks migrate
409 + to the LRU head of the free queue. Returns the filler blocks (caller must
410 + free them later to restore pool capacity).
411 +
412 + In a real engine the same thing happens naturally: after one request
413 + finishes and frees its blocks, subsequent requests allocate from the LRU
414 + head, consuming the unhashed blocks and leaving the old hashed blocks at
415 + the front of the queue.
416 + """
417 + fillers = gpu_pool.get_new_blocks(num_filler_blocks)
418 + return fillers
419 +
420 +
421 +def test_lazy_store_and_load_roundtrip() -> None:
422 + """Lazy mode: schedule a request, finish it so its hashed blocks are freed,
423 + then schedule new requests so the old blocks migrate to the LRU head.
424 + The lazy scanner offloads them to CPU. Re-scheduling the old request
425 + triggers a CPU cache hit + load.
426 +
427 + GPU pool: 8 blocks (7 usable). _target_free = ceil(64/16) = 4.
428 + """
429 + fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=8, lazy=True)
430 + sched = fix.scheduler
431 + gpu_pool = fix.gpu_block_pool
432 +
433 + num_blocks = 2
434 +
435 + # --- Step 1: Schedule req_old, compute, and finish ---
436 + req_old = make_request(num_blocks=num_blocks)
437 + gpu_blocks_old = _allocate_gpu_blocks(gpu_pool, req_old, num_blocks, group_id=0)
438 + gpu_pool.free_blocks(gpu_blocks_old)
439 +
440 + # Allocate filler blocks so req_old's hashed blocks move to LRU head.
441 + # 7 usable - 2 (req_old freed) = 5 other free blocks to consume.
442 + fillers = _flush_old_blocks_to_lru_head(gpu_pool, num_filler_blocks=5)
443 +
444 + # --- Step 2: Lazy scanner should offload req_old's blocks ---
445 + sched_out = make_scheduler_output({})
446 + meta = sched.build_connector_meta(sched_out)
447 + assert meta.store_event >= 0, "Expected lazy store to offload old blocks"
448 + assert len(meta.store_gpu_blocks) == num_blocks
449 + simulate_store_completion(sched, meta.store_event)
450 +
451 + # Free fillers to restore pool capacity.
452 + gpu_pool.free_blocks(fillers)
453 +
454 + # --- Step 3: Re-schedule req_old — should get CPU cache hit ---
455 + req_old2 = Request(
456 + request_id="req-old-reload",
457 + prompt_token_ids=req_old.prompt_token_ids,
458 + sampling_params=req_old.sampling_params,
459 + pooling_params=None,
460 + mm_features=None,
461 + block_hasher=req_old._block_hasher,
462 + )
463 + hit_tokens, is_async = sched.get_num_new_matched_tokens(
464 + req_old2, num_computed_tokens=0
465 + )
466 + assert hit_tokens == num_blocks * BLOCK_SIZE, (
467 + f"Expected {num_blocks * BLOCK_SIZE} hit tokens, got {hit_tokens}"
468 + )
469 + assert is_async is True
470 +
471 + # Allocate fresh GPU blocks for the load.
472 + gpu_blocks_load = gpu_pool.get_new_blocks(num_blocks)
473 + kv_blocks_load = KVCacheBlocks(blocks=(gpu_blocks_load,))
474 + sched.update_state_after_alloc(
475 + req_old2, kv_blocks_load, num_external_tokens=hit_tokens
476 + )
477 +
478 + sched_out2 = make_scheduler_output({req_old2.request_id: 1})
479 + meta2 = sched.build_connector_meta(sched_out2)
480 + assert meta2.load_event >= 0, "Expected a load event to be assigned"
481 + assert len(meta2.load_gpu_blocks) > 0
482 +
483 +
484 +# ---------------------------------------------------------------------------
485 +# Test 2a: Eager duplicate store is skipped
486 +# ---------------------------------------------------------------------------
487 +def test_eager_duplicate_store_skipped() -> None:
488 + """Eager: storing the same block hashes twice should not allocate new CPU blocks."""
489 + fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False)
490 + sched = fix.scheduler
491 +
492 + num_blocks = 2
493 + req = make_request(num_blocks=num_blocks)
494 +
495 + kv_blocks = _alloc_and_register(fix, req, num_blocks)
496 + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
497 + block_ids = kv_blocks.get_block_ids()
498 + sched_out = make_scheduler_output(
499 + {req.request_id: num_blocks * BLOCK_SIZE},
500 + new_reqs={req.request_id: block_ids},
501 + )
502 +
503 + meta1 = sched.build_connector_meta(sched_out)
504 + assert meta1.store_event >= 0
505 + simulate_store_completion(sched, meta1.store_event)
506 + cpu_free_after_first = get_cpu_free_blocks(sched)
507 +
508 + # Second request with identical hashes — should skip store
509 + req2 = Request(
510 + request_id="req-dup-eager",
511 + prompt_token_ids=req.prompt_token_ids,
512 + sampling_params=req.sampling_params,
513 + pooling_params=None,
514 + mm_features=None,
515 + block_hasher=req._block_hasher,
516 + )
517 + kv_blocks2 = _alloc_and_register(fix, req2, num_blocks)
518 + sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=0)
519 + block_ids2 = kv_blocks2.get_block_ids()
520 + sched_out2 = make_scheduler_output(
521 + {req2.request_id: num_blocks * BLOCK_SIZE},
522 + new_reqs={req2.request_id: block_ids2},
523 + )
524 +
525 + meta2 = sched.build_connector_meta(sched_out2)
526 + if meta2.store_event >= 0:
527 + assert len(meta2.store_cpu_blocks) == 0, (
528 + "Expected no new CPU blocks for duplicate hashes"
529 + )
530 + assert get_cpu_free_blocks(sched) == cpu_free_after_first
531 +
532 +
533 +# ---------------------------------------------------------------------------
534 +# Test 2b: Lazy duplicate store is skipped
535 +# ---------------------------------------------------------------------------
536 +def test_lazy_duplicate_store_skipped() -> None:
537 + """Lazy: blocks already offloaded to CPU should not be offloaded again.
538 +
539 + Same pattern as the lazy roundtrip: flush old blocks to LRU head, offload,
540 + then repeat with the same hashes and verify no new CPU allocation.
541 + """
542 + fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=8, lazy=True)
543 + sched = fix.scheduler
544 + gpu_pool = fix.gpu_block_pool
545 +
546 + num_blocks = 2
547 + req = make_request(num_blocks=num_blocks)
548 +
549 + # Schedule + finish → hashed blocks in free queue
550 + gpu_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0)
551 + gpu_pool.free_blocks(gpu_blocks)
552 +
553 + # Flush old blocks to LRU head, then trigger lazy offload.
554 + fillers = _flush_old_blocks_to_lru_head(gpu_pool, num_filler_blocks=5)
555 + meta1 = sched.build_connector_meta(make_scheduler_output({}))
556 + assert meta1.store_event >= 0
557 + simulate_store_completion(sched, meta1.store_event)
558 + gpu_pool.free_blocks(fillers)
559 + cpu_free_after_first = get_cpu_free_blocks(sched)
560 +
561 + # Allocate blocks with the same hashes and free them again.
562 + # The scanner should see they are already in CPU cache and skip them.
563 + req2 = Request(
564 + request_id="req-dup-lazy",
565 + prompt_token_ids=req.prompt_token_ids,
566 + sampling_params=req.sampling_params,
567 + pooling_params=None,
568 + mm_features=None,
569 + block_hasher=req._block_hasher,
570 + )
571 + gpu_blocks2 = _allocate_gpu_blocks(gpu_pool, req2, num_blocks, group_id=0)
572 + gpu_pool.free_blocks(gpu_blocks2)
573 +
574 + # Flush again so the hashed blocks are at LRU head for the scanner.
575 + fillers2 = _flush_old_blocks_to_lru_head(gpu_pool, num_filler_blocks=5)
576 + meta2 = sched.build_connector_meta(make_scheduler_output({}))
577 + gpu_pool.free_blocks(fillers2)
578 +
579 + # Either no store event, or zero new CPU blocks (already cached).
580 + if meta2.store_event >= 0:
581 + assert len(meta2.store_cpu_blocks) == 0, (
582 + "Expected no new CPU blocks for duplicate hashes"
583 + )
584 + assert get_cpu_free_blocks(sched) == cpu_free_after_first
585 +
586 +
587 +# ---------------------------------------------------------------------------
588 +# Test 3: LRU eviction order
589 +# ---------------------------------------------------------------------------
590 +def test_lru_eviction_order() -> None:
591 + """With limited CPU space, oldest blocks should be evicted first.
592 +
593 + CPU block pool: num_cpu_blocks=5 -> 4 free usable blocks (1 taken by null_block).
594 + After storing 4 blocks (2 req_a + 2 req_b), all free slots are occupied by
595 + cached blocks (ref_cnt=0, in hash map). When 2 more are stored (req_c),
596 + 2 LRU blocks from req_a get evicted from the cache to make room.
597 + """
598 + # 5 total = 4 usable (null_block takes 1), filling exactly with 4 blocks
599 + fix = make_scheduler(num_cpu_blocks=5, num_gpu_blocks=16, lazy=False)
600 + sched = fix.scheduler
601 +
602 + # Fill CPU with 4 blocks: 2 requests x 2 blocks (in LRU insertion order)
603 + req_a = make_request(num_blocks=2)
604 + req_b = make_request(num_blocks=2)
605 +
606 + kv_a = _alloc_and_register(fix, req_a, 2)
607 + kv_b = _alloc_and_register(fix, req_b, 2)
608 + sched.update_state_after_alloc(req_a, kv_a, num_external_tokens=0)
609 + sched.update_state_after_alloc(req_b, kv_b, num_external_tokens=0)
610 +
611 + ids_a = kv_a.get_block_ids()
612 + ids_b = kv_b.get_block_ids()
613 + sched_out = make_scheduler_output(
614 + {
615 + req_a.request_id: 2 * BLOCK_SIZE,
616 + req_b.request_id: 2 * BLOCK_SIZE,
617 + },
618 + new_reqs={
619 + req_a.request_id: ids_a,
620 + req_b.request_id: ids_b,
621 + },
622 + )
623 + meta = sched.build_connector_meta(sched_out)
624 + assert meta.store_event >= 0
625 + simulate_store_completion(sched, meta.store_event)
626 +
627 + # Verify all 4 blocks are cached in CPU hash map
628 + for i, bhash in enumerate(req_a.block_hashes[:2]):
629 + bhash_with_group = make_block_hash_with_group_id(bhash, 0)
630 + assert (
631 + sched.cpu_block_pool.cached_block_hash_to_block.get_one_block(
632 + bhash_with_group
633 + )
634 + is not None
635 + ), f"req_a block {i} should be cached after store"
636 + for i, bhash in enumerate(req_b.block_hashes[:2]):
637 + bhash_with_group = make_block_hash_with_group_id(bhash, 0)
638 + assert (
639 + sched.cpu_block_pool.cached_block_hash_to_block.get_one_block(
640 + bhash_with_group
641 + )
642 + is not None
643 + ), f"req_b block {i} should be cached after store"
644 +
645 + # Store 2 more blocks from a new request - must evict 2 LRU blocks (req_a)
646 + req_c = make_request(num_blocks=2)
647 + kv_c = _alloc_and_register(fix, req_c, 2)
648 + sched.update_state_after_alloc(req_c, kv_c, num_external_tokens=0)
649 +
650 + ids_c = kv_c.get_block_ids()
651 + sched_out2 = make_scheduler_output(
652 + {req_c.request_id: 2 * BLOCK_SIZE},
653 + new_reqs={req_c.request_id: ids_c},
654 + )
655 + meta2 = sched.build_connector_meta(sched_out2)
656 + assert meta2.store_event >= 0
657 + simulate_store_completion(sched, meta2.store_event)
658 +
659 + # req_a hashes should be evicted from CPU (they were LRU)
660 + for i, bhash in enumerate(req_a.block_hashes[:2]):
661 + bhash_with_group = make_block_hash_with_group_id(bhash, 0)
662 + cache_map = sched.cpu_block_pool.cached_block_hash_to_block
663 + cached = cache_map.get_one_block(bhash_with_group)
664 + assert cached is None, f"req_a block {i} should have been evicted"
665 +
666 + # req_b and req_c hashes should be present
667 + for i, bhash in enumerate(req_b.block_hashes[:2]):
668 + bhash_with_group = make_block_hash_with_group_id(bhash, 0)
669 + cache_map = sched.cpu_block_pool.cached_block_hash_to_block
670 + cached = cache_map.get_one_block(bhash_with_group)
671 + assert cached is not None, f"req_b block {i} should still be cached"
672 +
673 + for i, bhash in enumerate(req_c.block_hashes[:2]):
674 + bhash_with_group = make_block_hash_with_group_id(bhash, 0)
675 + cache_map = sched.cpu_block_pool.cached_block_hash_to_block
676 + cached = cache_map.get_one_block(bhash_with_group)
677 + assert cached is not None, f"req_c block {i} should still be cached"
678 +
679 +
680 +# ---------------------------------------------------------------------------
681 +# Test 4: Touched blocks survive eviction
682 +# ---------------------------------------------------------------------------
683 +def test_touched_blocks_survive_eviction() -> None:
684 + """Touching CPU blocks updates their LRU position, protecting them from eviction."""
685 + # 5 total = 4 usable (null_block takes 1)
686 + fix = make_scheduler(num_cpu_blocks=5, num_gpu_blocks=16, lazy=False)
687 + sched = fix.scheduler
688 +
689 + # Fill CPU with 4 blocks (req_a: 2, req_b: 2) in LRU order
690 + req_a = make_request(num_blocks=2)
691 + req_b = make_request(num_blocks=2)
692 +
693 + kv_a = _alloc_and_register(fix, req_a, 2)
694 + kv_b = _alloc_and_register(fix, req_b, 2)
695 + sched.update_state_after_alloc(req_a, kv_a, num_external_tokens=0)
696 + sched.update_state_after_alloc(req_b, kv_b, num_external_tokens=0)
697 +
698 + ids_a = kv_a.get_block_ids()
699 + ids_b = kv_b.get_block_ids()
700 + sched_out = make_scheduler_output(
701 + {
702 + req_a.request_id: 2 * BLOCK_SIZE,
703 + req_b.request_id: 2 * BLOCK_SIZE,
704 + },
705 + new_reqs={
706 + req_a.request_id: ids_a,
707 + req_b.request_id: ids_b,
708 + },
709 + )
710 + meta = sched.build_connector_meta(sched_out)
711 + simulate_store_completion(sched, meta.store_event)
712 +
713 + # Touch req_a's CPU blocks to make them most-recently-used
714 + cpu_pool = sched.cpu_block_pool
715 + for bhash in req_a.block_hashes[:2]:
716 + bhash_with_group = make_block_hash_with_group_id(bhash, 0)
717 + cached_blk = cpu_pool.cached_block_hash_to_block.get_one_block(bhash_with_group)
718 + assert cached_blk is not None
719 + cpu_pool.touch([cached_blk])
720 + # Undo touch to return ref_cnt to 0
721 + # (so it's a free candidate but at MRU position)
722 + cpu_pool.free_blocks([cached_blk])
723 +
724 + # Now store 2 more blocks; req_b (LRU front) should be evicted, not req_a
725 + req_c = make_request(num_blocks=2)
726 + kv_c = _alloc_and_register(fix, req_c, 2)
727 + sched.update_state_after_alloc(req_c, kv_c, num_external_tokens=0)
728 +
729 + ids_c = kv_c.get_block_ids()
730 + sched_out2 = make_scheduler_output(
731 + {req_c.request_id: 2 * BLOCK_SIZE},
732 + new_reqs={req_c.request_id: ids_c},
733 + )
734 + meta2 = sched.build_connector_meta(sched_out2)
735 + simulate_store_completion(sched, meta2.store_event)
736 +
737 + # req_b should be evicted (LRU), req_a and req_c should survive
738 + for i, bhash in enumerate(req_b.block_hashes[:2]):
739 + bhash_with_group = make_block_hash_with_group_id(bhash, 0)
740 + cached = cpu_pool.cached_block_hash_to_block.get_one_block(bhash_with_group)
741 + assert cached is None, f"req_b block {i} should have been evicted (it was LRU)"
742 +
743 + for i, bhash in enumerate(req_a.block_hashes[:2]):
744 + bhash_with_group = make_block_hash_with_group_id(bhash, 0)
745 + cached = cpu_pool.cached_block_hash_to_block.get_one_block(bhash_with_group)
746 + assert cached is not None, f"req_a block {i} should survive (was touched/MRU)"
747 +
748 +
749 +# ---------------------------------------------------------------------------
750 +# Test 5: Preemption no CPU block leak
751 +# ---------------------------------------------------------------------------
752 +def test_preemption_no_cpu_block_leak() -> None:
753 + """request_finished during in-flight load defers cleanup;
754 + completes after load done."""
755 + fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False)
756 + sched = fix.scheduler
757 +
758 + num_blocks = 2
759 +
760 + # First: store blocks to CPU
761 + req = make_request(num_blocks=num_blocks)
762 + kv_blocks = _alloc_and_register(fix, req, num_blocks)
763 + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
764 + block_ids = kv_blocks.get_block_ids()
765 + sched_out = make_scheduler_output(
766 + {req.request_id: num_blocks * BLOCK_SIZE},
767 + new_reqs={req.request_id: block_ids},
768 + )
769 + meta = sched.build_connector_meta(sched_out)
770 + simulate_store_completion(sched, meta.store_event)
771 +
772 + # Create new request with same tokens, check hit
773 + req2 = Request(
774 + request_id="req-preempt-load",
775 + prompt_token_ids=req.prompt_token_ids,
776 + sampling_params=req.sampling_params,
777 + pooling_params=None,
778 + mm_features=None,
779 + block_hasher=req._block_hasher,
780 + )
781 + hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)
782 + assert hit_tokens > 0
783 +
784 + gpu_blocks2 = fix.gpu_block_pool.get_new_blocks(num_blocks)
785 + kv_blocks2 = KVCacheBlocks(blocks=(gpu_blocks2,))
786 + sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens)
787 +
788 + # Assign load_event via build_connector_meta
789 + block_ids2 = kv_blocks2.get_block_ids()
790 + sched_out2 = make_scheduler_output(
791 + {req2.request_id: 1},
792 + new_reqs={req2.request_id: block_ids2},
793 + )
794 + meta2 = sched.build_connector_meta(sched_out2)
795 + assert meta2.load_event >= 0
796 +
797 + # Request finishes BEFORE load completes -> deferred
798 + sched.request_finished(req2, block_ids=[])
799 + assert req2.request_id in sched._reqs_to_load
800 + assert sched._reqs_to_load[req2.request_id].finished is True
801 +
802 + # Now simulate load completion -> cleanup fires
803 + simulate_load_completion(sched, {req2.request_id})
804 + assert req2.request_id not in sched._reqs_to_load
805 +
806 +
807 +# ---------------------------------------------------------------------------
808 +# Test 6: Eager store preemption cleanup
809 +# ---------------------------------------------------------------------------
810 +def test_eager_store_preemption_cleanup() -> None:
811 + """In eager mode, finishing a request during in-flight store defers cleanup."""
812 + fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False)
813 + sched = fix.scheduler
814 +
815 + num_blocks = 2
816 + req = make_request(num_blocks=num_blocks)
817 + kv_blocks = _alloc_and_register(fix, req, num_blocks)
818 + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
819 +
820 + block_ids = kv_blocks.get_block_ids()
821 + sched_out = make_scheduler_output(
822 + {req.request_id: num_blocks * BLOCK_SIZE},
823 + new_reqs={req.request_id: block_ids},
824 + )
825 + meta = sched.build_connector_meta(sched_out)
826 + store_event = meta.store_event
827 + assert store_event >= 0
828 +
829 + # The request gets store_events populated
830 + assert req.request_id in sched._reqs_to_store
831 + store_state = sched._reqs_to_store[req.request_id]
832 + assert store_event in store_state.store_events
833 +
834 + # Finish request while store still in-flight -> deferred
835 + sched.request_finished(req, block_ids=[])
836 + assert req.request_id in sched._reqs_to_store
837 + assert sched._reqs_to_store[req.request_id].finished is True
838 +
839 + # Simulate store completion -> deferred cleanup fires
840 + simulate_store_completion(sched, store_event)
841 + assert req.request_id not in sched._reqs_to_store
842 +
843 +
844 +# ---------------------------------------------------------------------------
845 +# Test 7: In-flight finish deferred cleanup (load variant)
846 +# ---------------------------------------------------------------------------
847 +def test_inflight_finish_deferred_cleanup() -> None:
848 + """Store, then start a load, request_finished defers,
849 + load completion fires cleanup."""
850 + fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False)
851 + sched = fix.scheduler
852 +
853 + num_blocks = 2
854 +
855 + # Store
856 + req = make_request(num_blocks=num_blocks)
857 + kv_blocks = _alloc_and_register(fix, req, num_blocks)
858 + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
859 + block_ids = kv_blocks.get_block_ids()
860 + sched_out = make_scheduler_output(
861 + {req.request_id: num_blocks * BLOCK_SIZE},
862 + new_reqs={req.request_id: block_ids},
863 + )
864 + meta = sched.build_connector_meta(sched_out)
865 + simulate_store_completion(sched, meta.store_event)
866 +
867 + # Load
868 + req2 = Request(
869 + request_id="req-inflight-load",
870 + prompt_token_ids=req.prompt_token_ids,
871 + sampling_params=req.sampling_params,
872 + pooling_params=None,
873 + mm_features=None,
874 + block_hasher=req._block_hasher,
875 + )
876 + hit_tokens, _ = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)
877 + assert hit_tokens > 0
878 +
879 + gpu_blocks2 = fix.gpu_block_pool.get_new_blocks(num_blocks)
880 + kv_blocks2 = KVCacheBlocks(blocks=(gpu_blocks2,))
881 + sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens)
882 +
883 + block_ids2 = kv_blocks2.get_block_ids()
884 + sched_out2 = make_scheduler_output(
885 + {req2.request_id: 1},
886 + new_reqs={req2.request_id: block_ids2},
887 + )
888 + meta2 = sched.build_connector_meta(sched_out2)
889 + assert meta2.load_event >= 0
890 +
891 + # Finish before load completes
892 + sched.request_finished(req2, block_ids=[])
893 + assert req2.request_id in sched._reqs_to_load
894 +
895 + # Simulate load completion -> request removed
896 + simulate_load_completion(sched, {req2.request_id})
897 + assert req2.request_id not in sched._reqs_to_load
898 +
899 +
900 +# ---------------------------------------------------------------------------
901 +# Test 8: Null GPU blocks are skipped in store and load transfer pairs
902 +# ---------------------------------------------------------------------------
903 +def test_multi_group_null_blocks_skipped() -> None:
904 + """Null GPU blocks (no block_hash) must not appear in store or load pairs.
905 +
906 + In eager store mode, _prepare_eager_store_specs skips blocks whose
907 + block_hash is None (null blocks have no hash). We verify this by mixing
908 + real hashed blocks with unhashed (null-like) blocks in a single group and
909 + checking that only real blocks appear in the store list.
910 + """
911 + fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, num_groups=1, lazy=False)
912 + sched = fix.scheduler
913 + gpu_pool = fix.gpu_block_pool
914 +
915 + num_blocks = 2
916 + req = make_request(num_blocks=num_blocks)
917 +
918 + # Allocate real blocks (with hashes) and use the null_block as a placeholder
919 + gpu_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0)
920 + null_block = gpu_pool.null_block
921 +
922 + # Mix: [real_block, null_block] — null_block has no hash, should be skipped
923 + mixed_blocks = [gpu_blocks[0], null_block]
924 + kv_blocks = KVCacheBlocks(blocks=(mixed_blocks,))
925 + req.num_computed_tokens = num_blocks * BLOCK_SIZE
926 + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
927 +
928 + block_ids = kv_blocks.get_block_ids()
929 + sched_out = make_scheduler_output(
930 + {req.request_id: num_blocks * BLOCK_SIZE},
931 + new_reqs={req.request_id: block_ids},
932 + )
933 + meta = sched.build_connector_meta(sched_out)
934 +
935 + # Null block's ID should NOT appear in store_gpu_blocks
936 + null_block_id = null_block.block_id
937 + assert null_block_id not in meta.store_gpu_blocks, (
938 + f"Null block id {null_block_id} should not appear in store transfer pairs"
939 + )
940 +
941 + # Only real block should be scheduled for store
942 + assert len(meta.store_gpu_blocks) == 1
943 + assert gpu_blocks[0].block_id in meta.store_gpu_blocks
944 +
945 + # Complete the store
946 + assert meta.store_event >= 0
947 + simulate_store_completion(sched, meta.store_event)
948 +
949 + # Create matching request and get load hit
950 + req2 = Request(
951 + request_id="req-null-load",
952 + prompt_token_ids=req.prompt_token_ids,
953 + sampling_params=req.sampling_params,
954 + pooling_params=None,
955 + mm_features=None,
956 + block_hasher=req._block_hasher,
957 + )
958 + hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)
959 + # Only 1 block was stored (the real one)
960 + assert hit_tokens == BLOCK_SIZE
961 + assert is_async is True
962 +
963 + # Allocate new GPU blocks for the load
964 + gpu_blocks2 = gpu_pool.get_new_blocks(1)
965 + kv_blocks2 = KVCacheBlocks(blocks=([gpu_blocks2[0], null_block],))
966 + sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens)
967 +
968 + sched_out2 = make_scheduler_output({req2.request_id: 1})
969 + meta2 = sched.build_connector_meta(sched_out2)
970 +
971 + # Null block's ID should NOT appear in load_gpu_blocks
972 + assert null_block_id not in meta2.load_gpu_blocks, (
973 + f"Null block id {null_block_id} should not appear in load transfer pairs"
974 + )
975 +
976 +
977 +# ---------------------------------------------------------------------------
978 +# Test 9: Chunked prefill accumulates block_ids across steps
979 +# ---------------------------------------------------------------------------
980 +def test_chunked_prefill_reads_live_block_ids() -> None:
981 + """With chunked prefill, block IDs accumulate across scheduler steps.
982 + _prepare_eager_store_specs reads block IDs from scheduler_output via
983 + yield_req_data, so the store should reflect the updated (larger) block
984 + list, not a stale snapshot."""
985 + fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False)
986 + sched = fix.scheduler
987 +
988 + num_blocks = 4
989 + req = make_request(num_blocks=num_blocks)
990 +
991 + # First chunk: allocate 2 blocks
992 + kv_blocks_first = _alloc_and_register(fix, req, 2)
993 + sched.update_state_after_alloc(req, kv_blocks_first, num_external_tokens=0)
994 +
995 + assert req.request_id in sched._reqs_to_store
996 + # Should still be exactly 1 entry in _reqs_to_store
997 + assert list(sched._reqs_to_store.keys()).count(req.request_id) == 1
998 +
999 + # Build connector meta with 2 blocks — stores the first 2
1000 + ids_first = kv_blocks_first.get_block_ids()
1001 + sched_out1 = make_scheduler_output(
1002 + {req.request_id: 2 * BLOCK_SIZE},
1003 + new_reqs={req.request_id: ids_first},
1004 + )
1005 + meta1 = sched.build_connector_meta(sched_out1)
1006 + assert meta1.store_event >= 0
1007 + assert len(meta1.store_gpu_blocks) == 2
1008 + simulate_store_completion(sched, meta1.store_event)
1009 +
1010 + # Second chunk: allocate 4 blocks total (2 new ones)
1011 + kv_blocks_second = _alloc_and_register(fix, req, num_blocks)
1012 + # update_state_after_alloc is idempotent for store registration
1013 + sched.update_state_after_alloc(req, kv_blocks_second, num_external_tokens=0)
1014 +
1015 + # Still exactly 1 entry
1016 + assert list(sched._reqs_to_store.keys()).count(req.request_id) == 1
1017 +
1018 + # The second chunk's NEW block IDs (positions 2,3) are passed as
1019 + # cached_req_new_blocks. The full block_ids include both old and new,
1020 + # but yield_req_data only appends the new_block_ids for cached reqs.
1021 + ids_second_full = kv_blocks_second.get_block_ids()
1022 + # New blocks are those beyond the first chunk
1023 + new_block_ids = tuple(ids_second_full[g][2:] for g in range(len(ids_second_full)))
1024 + sched_out2 = make_scheduler_output(
1025 + {req.request_id: 2 * BLOCK_SIZE},
1026 + cached_req_new_blocks={req.request_id: new_block_ids},
1027 + )
1028 + meta2 = sched.build_connector_meta(sched_out2)
1029 + assert meta2.store_event >= 0
1030 + # Only the 2 NEW blocks should be stored (first 2 already done)
1031 + assert len(meta2.store_gpu_blocks) == 2
1032 +
1033 +
1034 +# ---------------------------------------------------------------------------
1035 +# Test 10: Partial GPU prefix hit + CPU load + new compute blocks
1036 +# ---------------------------------------------------------------------------
1037 +def test_partial_gpu_prefix_plus_cpu_load() -> None:
1038 + """When GPU has a prefix cache hit for the first N blocks, CPU has a
1039 + hit for the next M blocks, and there are P new blocks needing fresh
1040 + compute, the block layout is:
1041 +
1042 + | comp (N) | ext_comp (M) | new (P) |
1043 +
1044 + External blocks sit in the middle — not at the beginning or end.
1045 + The load path must target hashes at positions [N, N+M).
1046 +
1047 + Request: 6 blocks (0..5).
1048 + - Store all 6 to CPU.
1049 + - New request: GPU prefix cache hits blocks 0,1 (hashed).
1050 + CPU hits blocks 2,3. Blocks 4,5 are new (need compute).
1051 + - update_state_after_alloc receives 6 GPU blocks:
1052 + [0,1] hashed (comp), [2,3] unhashed (ext_comp), [4,5] unhashed (new).
1053 + - Load must target hash positions 2,3.
1054 + """
1055 + fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False)
1056 + sched = fix.scheduler
1057 + gpu_pool = fix.gpu_block_pool
1058 +
1059 + num_blocks = 6
1060 + req = make_request(num_blocks=num_blocks)
1061 +
1062 + # Store all 6 blocks to CPU via eager store.
1063 + kv_blocks = _alloc_and_register(fix, req, num_blocks)
1064 + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
1065 + block_ids = kv_blocks.get_block_ids()
1066 + sched_out = make_scheduler_output(
1067 + {req.request_id: num_blocks * BLOCK_SIZE},
1068 + new_reqs={req.request_id: block_ids},
1069 + )
1070 + meta = sched.build_connector_meta(sched_out)
1071 + assert meta.store_event >= 0
1072 + simulate_store_completion(sched, meta.store_event)
1073 +
1074 + # New request with same tokens — but only partial GPU prefix hit.
1075 + req2 = Request(
1076 + request_id="req-partial-gpu",
1077 + prompt_token_ids=req.prompt_token_ids,
1078 + sampling_params=req.sampling_params,
1079 + pooling_params=None,
1080 + mm_features=None,
1081 + block_hasher=req._block_hasher,
1082 + )
1083 +
1084 + # GPU prefix cache hits the first 2 blocks.
1085 + gpu_local_computed = 2 * BLOCK_SIZE
1086 + hit_tokens, is_async = sched.get_num_new_matched_tokens(
1087 + req2, num_computed_tokens=gpu_local_computed
1088 + )
1089 + # CPU should hit blocks 2,3 (not 4,5 — those are beyond the CPU range).
1090 + num_cpu_hit_blocks = 2
1091 + # Actually CPU has all 6 stored; it returns hits starting from position 2.
1092 + # The number of CPU hit blocks = min(remaining request blocks, CPU cached).
1093 + # Here remaining = 6 - 2 = 4 blocks are in CPU, so hit = 4 * BLOCK_SIZE.
1094 + num_cpu_hit_blocks = 4
1095 + assert hit_tokens == num_cpu_hit_blocks * BLOCK_SIZE, (
1096 + f"Expected {num_cpu_hit_blocks * BLOCK_SIZE} CPU hit tokens, got {hit_tokens}"
1097 + )
1098 + assert is_async is True
1099 +
1100 + # Simulate what the real scheduler does: only accept 2 of the 4 CPU hit
1101 + # blocks as external (e.g. due to budget constraints), leaving 2 new
1102 + # blocks for fresh compute.
1103 + num_ext_blocks = 2
1104 + num_new_blocks = 2
1105 + external_tokens = num_ext_blocks * BLOCK_SIZE
1106 +
1107 + # Build block list matching real layout: | comp(2) | ext_comp(2) | new(2) |
1108 + # comp: GPU prefix cache hit — blocks with hashes
1109 + gpu_comp = _allocate_gpu_blocks(gpu_pool, req2, 2, group_id=0)
1110 + # ext_comp + new: freshly allocated, no hashes
1111 + gpu_ext_and_new = gpu_pool.get_new_blocks(num_ext_blocks + num_new_blocks)
1112 + all_gpu_blocks = gpu_comp + gpu_ext_and_new
1113 + kv_blocks2 = KVCacheBlocks(blocks=(all_gpu_blocks,))
1114 +
1115 + # Critical call: with 2 hashed comp blocks and 2 external tokens worth
1116 + # of blocks, the manager must derive skipped=2 and load hashes [2,3].
1117 + sched.update_state_after_alloc(
1118 + req2, kv_blocks2, num_external_tokens=external_tokens
1119 + )
1120 +
1121 + block_ids2 = kv_blocks2.get_block_ids()
1122 + sched_out2 = make_scheduler_output(
1123 + {req2.request_id: num_new_blocks * BLOCK_SIZE},
1124 + new_reqs={req2.request_id: block_ids2},
1125 + )
1126 + meta2 = sched.build_connector_meta(sched_out2)
1127 + assert meta2.load_event >= 0, "Expected a load event for partial GPU + CPU hit"
1128 + assert len(meta2.load_gpu_blocks) == num_ext_blocks
1129 + assert len(meta2.load_cpu_blocks) == num_ext_blocks
1130 +
1131 + # Verify the load targets the ext_comp GPU blocks (positions 2,3),
1132 + # not the comp blocks (0,1) or new blocks (4,5).
1133 + ext_block_ids = [b.block_id for b in gpu_ext_and_new[:num_ext_blocks]]
1134 + for bid in meta2.load_gpu_blocks:
1135 + assert bid in ext_block_ids, (
1136 + f"Load GPU block {bid} should be an ext_comp block, not a comp or new block"
1137 + )
vllm/config/vllm.py
+5 -1
@@ -660,7 +660,11 @@ class VllmConfig:
660 )
661
662 if kv_offloading_backend == "native":
663 - self.kv_transfer_config.kv_connector = "OffloadingConnector"
663 + if envs.VLLM_USE_SIMPLE_KV_OFFLOAD:
664 + config_connector = "SimpleCPUOffloadConnector"
665 + else:
666 + config_connector = "OffloadingConnector"
667 + self.kv_transfer_config.kv_connector = config_connector
668 self.kv_transfer_config.kv_connector_extra_config.update(
669 {"cpu_bytes_to_use": kv_offloading_size * (1 << 30)}
670 )
vllm/distributed/kv_transfer/kv_connector/factory.py
+7
@@ -202,6 +202,7 @@ KVConnectorFactory.register_connector(
202 "vllm.distributed.kv_transfer.kv_connector.v1.decode_bench_connector",
203 "DecodeBenchConnector",
204 )
205 +
206 KVConnectorFactory.register_connector(
207 "MooncakeConnector",
208 "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector",
@@ -213,3 +214,9 @@ KVConnectorFactory.register_connector(
214 "vllm.distributed.kv_transfer.kv_connector.v1.flexkv_connector",
215 "FlexKVConnectorV1",
216 )
217 +
218 +KVConnectorFactory.register_connector(
219 + "SimpleCPUOffloadConnector",
220 + "vllm.distributed.kv_transfer.kv_connector.v1.simple_cpu_offload_connector",
221 + "SimpleCPUOffloadConnector",
222 +)
vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py new
+247
@@ -0,0 +1,247 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""SimpleCPUOffloadConnector: minimal CPU KV cache offloading."""
4 +
5 +from collections.abc import Iterable
6 +from typing import TYPE_CHECKING, Any
7 +
8 +import torch
9 +
10 +from vllm.config import VllmConfig
11 +from vllm.distributed.kv_events import KVCacheEvent
12 +from vllm.distributed.kv_transfer.kv_connector.v1.base import (
13 + KVConnectorBase_V1,
14 + KVConnectorMetadata,
15 + KVConnectorRole,
16 + SupportsHMA,
17 +)
18 +from vllm.logger import init_logger
19 +from vllm.v1.core.sched.output import SchedulerOutput
20 +from vllm.v1.outputs import KVConnectorOutput
21 +from vllm.v1.simple_kv_offload.manager import (
22 + SimpleCPUOffloadScheduler,
23 +)
24 +from vllm.v1.simple_kv_offload.metadata import (
25 + SimpleCPUOffloadMetadata,
26 +)
27 +from vllm.v1.simple_kv_offload.worker import (
28 + SimpleCPUOffloadWorker,
29 +)
30 +
31 +if TYPE_CHECKING:
32 + from vllm.forward_context import ForwardContext
33 + from vllm.v1.attention.backend import AttentionMetadata
34 + from vllm.v1.core.block_pool import BlockPool
35 + from vllm.v1.core.kv_cache_manager import KVCacheBlocks
36 + from vllm.v1.kv_cache_interface import KVCacheConfig
37 + from vllm.v1.request import Request
38 +
39 +logger = init_logger(__name__)
40 +
41 +# Default CPU capacity: 8 GB
42 +DEFAULT_CPU_CAPACITY_BYTES = 8 * (1024**3)
43 +
44 +
45 +class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA):
46 + """CPU KV cache offloading with custom kernel transfers and BlockPool LRU."""
47 +
48 + def __init__(
49 + self,
50 + vllm_config: VllmConfig,
51 + role: KVConnectorRole,
52 + kv_cache_config: "KVCacheConfig | None" = None,
53 + ):
54 + super().__init__(vllm_config, role, kv_cache_config)
55 +
56 + enable_prefix_caching = vllm_config.cache_config.enable_prefix_caching
57 + extra_config = self._kv_transfer_config.kv_connector_extra_config or {}
58 +
59 + cpu_capacity_bytes = int(
60 + extra_config.get("cpu_bytes_to_use", DEFAULT_CPU_CAPACITY_BYTES)
61 + )
62 + # cpu_bytes_to_use is server-wide for compatibility;
63 + # cpu_bytes_to_use_per_rank overrides for per-rank capacity.
64 + world_size = vllm_config.parallel_config.world_size
65 + cpu_capacity_per_rank = cpu_capacity_bytes // world_size
66 + if "cpu_bytes_to_use_per_rank" in extra_config:
67 + explicit = int(extra_config["cpu_bytes_to_use_per_rank"])
68 + if explicit != cpu_capacity_per_rank:
69 + logger.warning(
70 + "cpu_bytes_to_use_per_rank (%.2f GB) != "
71 + "cpu_bytes_to_use/world_size (%.2f GB). Using per-rank value.",
72 + explicit / (1024**3),
73 + cpu_capacity_per_rank / (1024**3),
74 + )
75 + cpu_capacity_per_rank = explicit
76 +
77 + lazy_offload = bool(extra_config.get("lazy_offload", False))
78 +
79 + self.scheduler_manager: SimpleCPUOffloadScheduler | None = None
80 + self.worker_handler: SimpleCPUOffloadWorker | None = None
81 +
82 + if not enable_prefix_caching:
83 + logger.warning(
84 + "Detected prefix caching disabled, disabling CPU offload "
85 + "since it requires prefix caching."
86 + )
87 + return
88 +
89 + logger.info(
90 + "SimpleCPUOffloadConnector: role=%s, "
91 + "per_rank=%.2f GB, world_size=%d, mode=%s",
92 + role.name,
93 + cpu_capacity_per_rank / (1024**3),
94 + world_size,
95 + "lazy" if lazy_offload else "eager",
96 + )
97 +
98 + if role == KVConnectorRole.SCHEDULER:
99 + self.scheduler_manager = SimpleCPUOffloadScheduler(
100 + vllm_config,
101 + kv_cache_config,
102 + cpu_capacity_per_rank,
103 + lazy_offload=lazy_offload,
104 + )
105 + elif role == KVConnectorRole.WORKER:
106 + self.worker_handler = SimpleCPUOffloadWorker(
107 + vllm_config, kv_cache_config, cpu_capacity_per_rank
108 + )
109 +
110 + # --- Worker-side methods ---
111 +
112 + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None:
113 + if self.worker_handler is not None:
114 + self.worker_handler.register_kv_caches(kv_caches)
115 +
116 + def bind_connector_metadata(
117 + self,
118 + connector_metadata: KVConnectorMetadata,
119 + ) -> None:
120 + super().bind_connector_metadata(connector_metadata)
121 + if self.worker_handler is not None:
122 + assert isinstance(connector_metadata, SimpleCPUOffloadMetadata)
123 + self.worker_handler.bind_connector_metadata(connector_metadata)
124 +
125 + def clear_connector_metadata(self) -> None:
126 + super().clear_connector_metadata()
127 + if self.worker_handler is not None:
128 + self.worker_handler.clear_connector_metadata()
129 +
130 + def handle_preemptions(self, kv_connector_metadata: KVConnectorMetadata) -> None:
131 + if self.worker_handler is not None:
132 + assert isinstance(kv_connector_metadata, SimpleCPUOffloadMetadata)
133 + self.worker_handler.handle_preemptions(kv_connector_metadata)
134 +
135 + def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None:
136 + pass # Launch loads ops in get_finished() after launching model execution
137 +
138 + def wait_for_layer_load(self, layer_name: str) -> None:
139 + pass # Always load asynchronously and deferred to get_finished()
140 +
141 + def save_kv_layer(
142 + self,
143 + layer_name: str,
144 + kv_layer: torch.Tensor,
145 + attn_metadata: "AttentionMetadata",
146 + **kwargs: Any,
147 + ) -> None:
148 + pass # Always save asynchronously and deferred to get_finished()
149 +
150 + def wait_for_save(self) -> None:
151 + pass # All stores are driven by get_finished() and no wait needed
152 +
153 + def get_finished(
154 + self,
155 + finished_req_ids: set[str],
156 + ) -> tuple[set[str] | None, set[str] | None]:
157 + if self.worker_handler is not None:
158 + return self.worker_handler.get_finished(finished_req_ids)
159 + return None, None
160 +
161 + def build_connector_worker_meta(self):
162 + if self.worker_handler is not None:
163 + return self.worker_handler.build_connector_worker_meta()
164 + return None
165 +
166 + # --- Scheduler-side methods ---
167 +
168 + # NOTE: New API only for SimpleCPUOffloadConnector.
169 + def bind_gpu_block_pool(self, gpu_block_pool: "BlockPool") -> None:
170 + if self.scheduler_manager is not None:
171 + self.scheduler_manager.bind_gpu_block_pool(gpu_block_pool)
172 +
173 + def get_num_new_matched_tokens(
174 + self,
175 + request: "Request",
176 + num_computed_tokens: int,
177 + ) -> tuple[int | None, bool]:
178 + if self.scheduler_manager is not None:
179 + return self.scheduler_manager.get_num_new_matched_tokens(
180 + request, num_computed_tokens
181 + )
182 + return 0, False
183 +
184 + def update_state_after_alloc(
185 + self,
186 + request: "Request",
187 + blocks: "KVCacheBlocks",
188 + num_external_tokens: int,
189 + ) -> None:
190 + if self.scheduler_manager is not None:
191 + self.scheduler_manager.update_state_after_alloc(
192 + request, blocks, num_external_tokens
193 + )
194 +
195 + def build_connector_meta(
196 + self,
197 + scheduler_output: SchedulerOutput,
198 + ) -> KVConnectorMetadata:
199 + if self.scheduler_manager is not None:
200 + return self.scheduler_manager.build_connector_meta(scheduler_output)
201 + return SimpleCPUOffloadMetadata()
202 +
203 + def update_connector_output(
204 + self,
205 + connector_output: KVConnectorOutput,
206 + ) -> None:
207 + if self.scheduler_manager is not None:
208 + self.scheduler_manager.update_connector_output(connector_output)
209 +
210 + def request_finished(
211 + self,
212 + request: "Request",
213 + block_ids: list[int],
214 + ) -> tuple[bool, dict[str, Any] | None]:
215 + if self.scheduler_manager is not None:
216 + return self.scheduler_manager.request_finished(request, block_ids)
217 + return False, None
218 +
219 + def request_finished_all_groups(
220 + self,
221 + request: "Request",
222 + block_ids: tuple[list[int], ...],
223 + ) -> tuple[bool, dict[str, Any] | None]:
224 + if self.scheduler_manager is not None:
225 + return self.scheduler_manager.request_finished_all_groups(
226 + request, block_ids
227 + )
228 + return False, None
229 +
230 + # NOTE: New API only for SimpleCPUOffloadConnector.
231 + def has_pending_transfers(self) -> bool:
232 + if self.scheduler_manager is not None:
233 + return self.scheduler_manager.has_pending_stores()
234 + return False
235 +
236 + def take_events(self) -> Iterable[KVCacheEvent]:
237 + if self.scheduler_manager is not None:
238 + return self.scheduler_manager.take_events()
239 + return []
240 +
241 + def reset_cache(self) -> bool | None:
242 + raise NotImplementedError(
243 + "SimpleCPUOffloadConnector does not support reset_cache(). "
244 + "reset_prefix_cache() requires synchronizing all pending "
245 + "CPU offload transfers before clearing GPU prefix cache blocks, "
246 + "which is not yet implemented."
247 + )
vllm/envs.py
+4
@@ -1674,6 +1674,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
1674 "VLLM_XPU_ENABLE_XPU_GRAPH": lambda: bool(
1675 int(os.getenv("VLLM_XPU_ENABLE_XPU_GRAPH", "0"))
1676 ),
1677 + # Enable simple KV offload.
1678 + "VLLM_USE_SIMPLE_KV_OFFLOAD": lambda: bool(
1679 + int(os.getenv("VLLM_USE_SIMPLE_KV_OFFLOAD", "0"))
1680 + ),
1681 }
1682
1683
vllm/v1/core/sched/scheduler.py
+7
@@ -234,6 +234,13 @@ class Scheduler(SchedulerInterface):
234 hash_block_size=self.block_size,
235 metrics_collector=self.kv_metrics_collector,
236 )
237 + # Bind GPU block pool to the KV connector. This must happen after
238 + # kv_cache_manager is constructed so block_pool is available.
239 + if self.connector is not None and hasattr(
240 + self.connector, "bind_gpu_block_pool"
241 + ):
242 + self.connector.bind_gpu_block_pool(self.kv_cache_manager.block_pool)
243 +
244 self.use_pp = self.parallel_config.pipeline_parallel_size > 1
245 self.use_v2_model_runner = envs.VLLM_USE_V2_MODEL_RUNNER
246 self.scheduler_reserve_full_isl = (
vllm/v1/metrics/stats.py
+10 -2
@@ -281,8 +281,16 @@ class PromptTokenStats:
281
282 self.computed += prompt_len - num_cached_tokens
283 self.external_kv_transfer += num_external_computed_tokens
284 - self.local_cache_hit += (
285 - num_cached_tokens + recomputed - num_external_computed_tokens
284 + # FIXME(yifan): local_cache_hit can go negative after preemption.
285 + # num_cached_tokens is a one-time snapshot from first scheduling and
286 + # is never reset on preemption, while num_external_computed_tokens is
287 + # overwritten on re-scheduling. If CPU offload finds more tokens on
288 + # the second pass than the original total, the subtraction underflows.
289 + # A fundamental fix is to track the first-time num_external_computed_tokens
290 + # as a separate metric rather than reusing num_external_computed_tokens
291 + # for metric directly.
292 + self.local_cache_hit += max(
293 + 0, (num_cached_tokens + recomputed - num_external_computed_tokens)
294 )
295 self.cached_tokens += num_cached_tokens
296 self.recomputed_tokens += recomputed
vllm/v1/simple_kv_offload/__init__.py
vllm/v1/simple_kv_offload/copy_backend.py new
+97
@@ -0,0 +1,97 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""DMA copy backend for GPU<->CPU block transfers."""
4 +
5 +from __future__ import annotations
6 +
7 +import queue
8 +import threading
9 +
10 +import torch
11 +
12 +from vllm.logger import init_logger
13 +from vllm.platforms import current_platform
14 +from vllm.v1.simple_kv_offload.cuda_mem_ops import (
15 + BatchMemcpyParams,
16 + build_params,
17 + copy_blocks,
18 +)
19 +
20 +logger = init_logger(__name__)
21 +
22 +
23 +class DmaCopyBackend:
24 + """cuMemcpyBatchAsync copy backend (background thread)."""
25 +
26 + def __init__(self) -> None:
27 + self._store_params: BatchMemcpyParams | None = None
28 + self._load_params: BatchMemcpyParams | None = None
29 + self._load_stream: torch.cuda.Stream | None = None
30 + self._store_stream: torch.cuda.Stream | None = None
31 + self._queue: queue.SimpleQueue | None = None
32 + self._thread: threading.Thread | None = None
33 + self._shutdown: bool = False
34 +
35 + def init(
36 + self,
37 + gpu_caches: dict[str, torch.Tensor],
38 + cpu_caches: dict[str, torch.Tensor],
39 + device: torch.device,
40 + load_stream: torch.cuda.Stream,
41 + store_stream: torch.cuda.Stream,
42 + ) -> None:
43 + self._load_stream = load_stream
44 + self._store_stream = store_stream
45 +
46 + self._store_params = build_params(gpu_caches, cpu_caches, store_stream)
47 + self._load_params = build_params(cpu_caches, gpu_caches, load_stream)
48 +
49 + self._queue = queue.SimpleQueue()
50 + self._thread = threading.Thread(
51 + target=self._copy_loop,
52 + args=(self._queue, device, load_stream, store_stream),
53 + daemon=True,
54 + )
55 + self._thread.start()
56 +
57 + def launch_copy(
58 + self,
59 + src_blocks: list[int],
60 + dst_blocks: list[int],
61 + is_store: bool,
62 + event_idx: int,
63 + events_list: list[tuple[int, torch.Event]],
64 + ) -> None:
65 + params = self._store_params if is_store else self._load_params
66 + assert params is not None and self._queue is not None
67 + self._queue.put(
68 + (src_blocks, dst_blocks, params, is_store, event_idx, events_list)
69 + )
70 +
71 + def shutdown(self) -> None:
72 + if self._shutdown:
73 + return
74 + self._shutdown = True
75 + if self._queue is not None:
76 + self._queue.put(None)
77 + if self._thread is not None:
78 + self._thread.join(timeout=5.0)
79 +
80 + @staticmethod
81 + def _copy_loop(
82 + q: queue.SimpleQueue,
83 + device: torch.device,
84 + load_stream: torch.cuda.Stream,
85 + store_stream: torch.cuda.Stream,
86 + ) -> None:
87 + current_platform.set_device(device)
88 + while True:
89 + item = q.get()
90 + if item is None:
91 + return
92 + src_blocks, dst_blocks, params, is_store, event_idx, events_list = item
93 + copy_blocks(src_blocks, dst_blocks, params)
94 + stream = store_stream if is_store else load_stream
95 + event = torch.Event()
96 + event.record(stream)
97 + events_list.append((event_idx, event))
vllm/v1/simple_kv_offload/cuda_mem_ops.py new
+153
@@ -0,0 +1,153 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Low-level CUDA memory helpers: pinning and batch DMA transfers."""
4 +
5 +import ctypes
6 +from typing import Any, NamedTuple
7 +
8 +import numpy as np
9 +import torch
10 +
11 +from vllm.logger import init_logger
12 +
13 +logger = init_logger(__name__)
14 +
15 +
16 +def pin_tensor(tensor: torch.Tensor) -> None:
17 + """Pin a CPU tensor via cudaHostRegister.
18 +
19 + This bypasses PyTorch's CUDACachingHostAllocator which rounds
20 + every ``pin_memory=True`` allocation up to the next power of 2
21 + (e.g. 100 GB becomes 128 GB).
22 + """
23 + err = torch.cuda.cudart().cudaHostRegister(tensor.data_ptr(), tensor.nbytes, 0)
24 + if err.value != 0:
25 + raise RuntimeError(f"cudaHostRegister failed: {err}")
26 +
27 +
28 +class _CUmemLocation(ctypes.Structure):
29 + _fields_ = [("type", ctypes.c_uint), ("id", ctypes.c_int)]
30 +
31 +
32 +class _CUmemcpyAttributes(ctypes.Structure):
33 + _fields_ = [
34 + ("srcAccessOrder", ctypes.c_uint),
35 + ("srcLocHint", _CUmemLocation),
36 + ("dstLocHint", _CUmemLocation),
37 + ("flags", ctypes.c_uint),
38 + ]
39 +
40 +
41 +_BATCH_MEMCPY_FUNC_TYPE = ctypes.CFUNCTYPE(
42 + ctypes.c_uint, # CUresult
43 + ctypes.c_void_p,
44 + ctypes.c_void_p,
45 + ctypes.c_void_p,
46 + ctypes.c_size_t,
47 + ctypes.c_void_p,
48 + ctypes.c_void_p,
49 + ctypes.c_size_t,
50 + ctypes.c_void_p,
51 + ctypes.c_void_p,
52 +)
53 +
54 +# Resolved lazily on first use.
55 +_batch_memcpy_fn: Any = None
56 +
57 +
58 +def _resolve_batch_memcpy():
59 + """Resolve cuMemcpyBatchAsync via cuGetProcAddress (one-time)."""
60 + from cuda.bindings import driver as drv
61 +
62 + err, ptr, _ = drv.cuGetProcAddress(b"cuMemcpyBatchAsync", 12080, 0)
63 + if err != drv.CUresult.CUDA_SUCCESS:
64 + raise RuntimeError(f"cuGetProcAddress(cuMemcpyBatchAsync) failed: {err}")
65 + return _BATCH_MEMCPY_FUNC_TYPE(ptr)
66 +
67 +
68 +class BatchMemcpyParams(NamedTuple):
69 + src_bases: np.ndarray # [num_layers] uint64 — data_ptr per layer
70 + dst_bases: np.ndarray # [num_layers] uint64
71 + bpb: np.ndarray # [num_layers] uint64 — bytes per block
72 + num_layers: int
73 + attrs: _CUmemcpyAttributes
74 + attrs_idx: ctypes.c_size_t
75 + # NOTE: cuMemcpyBatchAsync_v2() removed fail_idx field, but we use
76 + # cuMemcpyBatchAsync() with fail_idx for backward compatibility
77 + fail_idx: ctypes.c_size_t
78 + stream_handle: int # raw cudaStream_t / CUstream
79 +
80 +
81 +def build_params(
82 + src_caches: dict[str, torch.Tensor],
83 + dst_caches: dict[str, torch.Tensor],
84 + stream: torch.cuda.Stream,
85 +) -> BatchMemcpyParams:
86 + global _batch_memcpy_fn
87 + if _batch_memcpy_fn is None:
88 + _batch_memcpy_fn = _resolve_batch_memcpy()
89 +
90 + assert list(src_caches.keys()) == list(dst_caches.keys())
91 + src_tensors = list(src_caches.values())
92 + dst_tensors = list(dst_caches.values())
93 +
94 + src_bases, dst_bases, bpb = [], [], []
95 + for s, d in zip(src_tensors, dst_tensors):
96 + s_bpb = s.stride(0) * s.element_size()
97 + assert s_bpb == d.stride(0) * d.element_size()
98 + src_bases.append(s.data_ptr())
99 + dst_bases.append(d.data_ptr())
100 + bpb.append(s_bpb)
101 +
102 + # Refer to https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__MEM.html#group__CUDA__MEM_1g6f1ff58e3065df3eb4b573dba77ad31f for details. # noqa: E501
103 + attrs = _CUmemcpyAttributes(srcAccessOrder=3) # ANY
104 +
105 + return BatchMemcpyParams(
106 + src_bases=np.array(src_bases, dtype=np.uint64),
107 + dst_bases=np.array(dst_bases, dtype=np.uint64),
108 + bpb=np.array(bpb, dtype=np.uint64),
109 + num_layers=len(src_tensors),
110 + attrs=attrs,
111 + attrs_idx=ctypes.c_size_t(0),
112 + fail_idx=ctypes.c_size_t(0),
113 + stream_handle=stream.cuda_stream,
114 + )
115 +
116 +
117 +def copy_blocks(
118 + src_block_ids: list[int],
119 + dst_block_ids: list[int],
120 + params: BatchMemcpyParams,
121 +) -> None:
122 + """Copy blocks via cuMemcpyBatchAsync."""
123 + n = len(src_block_ids)
124 + if n == 0:
125 + return
126 +
127 + src_ids = np.array(src_block_ids, dtype=np.uint64)
128 + dst_ids = np.array(dst_block_ids, dtype=np.uint64)
129 +
130 + src_all = (
131 + params.src_bases[:, None] + src_ids[None, :] * params.bpb[:, None]
132 + ).ravel()
133 + dst_all = (
134 + params.dst_bases[:, None] + dst_ids[None, :] * params.bpb[:, None]
135 + ).ravel()
136 + sz_all = np.repeat(params.bpb, n)
137 +
138 + total = n * params.num_layers
139 + err = _batch_memcpy_fn(
140 + dst_all.ctypes.data,
141 + src_all.ctypes.data,
142 + sz_all.ctypes.data,
143 + total,
144 + ctypes.addressof(params.attrs),
145 + ctypes.byref(params.attrs_idx),
146 + 1,
147 + ctypes.byref(params.fail_idx),
148 + params.stream_handle,
149 + )
150 + if err != 0:
151 + raise RuntimeError(
152 + f"cuMemcpyBatchAsync failed: err={err} failIdx={params.fail_idx.value}"
153 + )
vllm/v1/simple_kv_offload/manager.py new
+739
@@ -0,0 +1,739 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Scheduler-side manager for SimpleCPUOffloadConnector."""
4 +
5 +import contextlib
6 +from collections.abc import Iterable
7 +from dataclasses import dataclass, field
8 +from typing import TYPE_CHECKING, Any
9 +
10 +from vllm.config import VllmConfig
11 +from vllm.distributed.kv_events import KVCacheEvent
12 +from vllm.distributed.kv_transfer.kv_connector.utils import yield_req_data
13 +from vllm.logger import init_logger
14 +from vllm.utils.math_utils import cdiv
15 +from vllm.v1.core.block_pool import BlockPool
16 +from vllm.v1.core.kv_cache_coordinator import (
17 + KVCacheCoordinator,
18 + get_kv_cache_coordinator,
19 +)
20 +from vllm.v1.core.sched.output import SchedulerOutput
21 +from vllm.v1.kv_cache_interface import (
22 + FullAttentionSpec,
23 + MambaSpec,
24 + SlidingWindowSpec,
25 +)
26 +from vllm.v1.outputs import KVConnectorOutput
27 +from vllm.v1.simple_kv_offload.metadata import (
28 + SimpleCPUOffloadMetadata,
29 + SimpleCPUOffloadWorkerMetadata,
30 +)
31 +
32 +if TYPE_CHECKING:
33 + from vllm.v1.core.kv_cache_manager import KVCacheBlocks
34 + from vllm.v1.core.kv_cache_utils import KVCacheBlock
35 + from vllm.v1.kv_cache_interface import KVCacheConfig
36 + from vllm.v1.request import Request
37 +
38 +logger = init_logger(__name__)
39 +
40 +
41 +@dataclass
42 +class TransferMeta:
43 + gpu_block_ids: list[int]
44 + cpu_block_ids: list[int]
45 +
46 +
47 +@dataclass
48 +class LoadRequestState:
49 + request: "Request"
50 + transfer_meta: TransferMeta
51 + load_event: int | None = None
52 + finished: bool = False
53 +
54 +
55 +# NOTE: This per-request state is only used in eager mode.
56 +@dataclass
57 +class StoreRequestState:
58 + request: "Request"
59 + # Accumulated block IDs from scheduler_output via yield_req_data.
60 + block_ids: tuple[list[int], ...]
61 + # Per-group cursors tracking how many blocks have been stored/skipped.
62 + num_stored_blocks: list[int]
63 + store_events: set[int] = field(default_factory=set)
64 + finished: bool = False
65 +
66 +
67 +class SimpleCPUOffloadScheduler:
68 + """Scheduler-side manager for CPU offloading."""
69 +
70 + def __init__(
71 + self,
72 + vllm_config: VllmConfig,
73 + kv_cache_config: "KVCacheConfig | None",
74 + cpu_capacity_bytes: int,
75 + lazy_offload: bool = False,
76 + ):
77 + self.vllm_config = vllm_config
78 + self.kv_cache_config = kv_cache_config
79 + self.enable_kv_cache_events = (
80 + vllm_config.kv_events_config is not None
81 + and vllm_config.kv_events_config.enable_kv_cache_events
82 + )
83 + # NOTE: We use the same block size for both GPU and CPU.
84 + self.block_size = vllm_config.cache_config.block_size
85 + # Derive a CPU KVCacheConfig from the GPU config and build a coordinator
86 + assert kv_cache_config is not None
87 + self.cpu_kv_cache_config = self._derive_cpu_config(
88 + kv_cache_config, cpu_capacity_bytes
89 + )
90 + self.num_cpu_blocks = self.cpu_kv_cache_config.num_blocks
91 + # Find the full attention kv group for prefix cache matching.
92 + self.fa_gidx = -1
93 + for g_idx, g in enumerate(self.cpu_kv_cache_config.kv_cache_groups):
94 + if isinstance(g.kv_cache_spec, FullAttentionSpec):
95 + self.fa_gidx = g_idx
96 + break
97 + assert 0 <= self.fa_gidx < len(self.cpu_kv_cache_config.kv_cache_groups)
98 +
99 + logger.info(
100 + "SimpleCPUOffloadScheduler: Allocating %d CPU blocks (%.2f GB, mode=%s)",
101 + self.num_cpu_blocks,
102 + cpu_capacity_bytes / (1024**3),
103 + "lazy" if lazy_offload else "eager",
104 + )
105 +
106 + # TODO (yifan): maybe need to enable kv_cache_events and metrics_collector here.
107 + dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
108 + pcp_world_size = vllm_config.parallel_config.prefill_context_parallel_size
109 + assert dcp_world_size == 1 and pcp_world_size == 1
110 + self.cpu_coordinator: KVCacheCoordinator = get_kv_cache_coordinator(
111 + kv_cache_config=self.cpu_kv_cache_config,
112 + max_model_len=vllm_config.model_config.max_model_len,
113 + use_eagle=False,
114 + enable_caching=True,
115 + enable_kv_cache_events=self.enable_kv_cache_events,
116 + dcp_world_size=dcp_world_size,
117 + pcp_world_size=pcp_world_size,
118 + hash_block_size=self.block_size,
119 + )
120 + self.cpu_block_pool: BlockPool = self.cpu_coordinator.block_pool
121 +
122 + # GPU block pool reference - bound after scheduler builds kv_cache_manager
123 + self._gpu_block_pool: BlockPool | None = None
124 +
125 + # Load metadata
126 + self._reqs_to_load: dict[str, LoadRequestState] = {}
127 + # Inverse map: load_event_idx -> req_ids. Keyed by load_event_idx because
128 + # the worker reports completions by event index, not request id.
129 + self._load_event_to_reqs: dict[int, list[str]] = {}
130 +
131 + # Store metadata
132 + self._lazy_mode = lazy_offload
133 + # Lazy mode: use a cursor to track the last scanned block in the GPU free queue.
134 + self._cursor: KVCacheBlock | None = None
135 + if self._lazy_mode:
136 + self._target_free = self._estimate_lazy_target_blocks(
137 + kv_cache_config,
138 + vllm_config.scheduler_config.max_num_batched_tokens,
139 + )
140 + else:
141 + self._target_free = 0
142 + self._store_event_to_blocks: dict[int, TransferMeta] = {}
143 + # Eager mode only
144 + self._reqs_to_store: dict[str, StoreRequestState] = {}
145 + self._store_event_to_reqs: dict[int, list[str]] = {}
146 +
147 + # Event counters
148 + self._load_event_counter: int = 0
149 + self._store_event_counter: int = 0
150 +
151 + # For TP/PP: track partial store completions across steps.
152 + # Events must be reported by all world_size workers before considered complete.
153 + self._expected_worker_count = vllm_config.parallel_config.world_size
154 + self._store_event_pending_counts: dict[int, int] = {}
155 +
156 + @staticmethod
157 + def _derive_cpu_config(
158 + gpu_config: "KVCacheConfig", cpu_capacity_bytes: int
159 + ) -> "KVCacheConfig":
160 + """Derive a CPU KVCacheConfig from the GPU config.
161 + Same kv_cache_groups, num_blocks scaled by CPU/GPU memory ratio."""
162 + # Import here to avoid potential circular imports
163 + from vllm.v1.kv_cache_interface import KVCacheConfig as KVCacheConfigCls
164 + from vllm.v1.kv_cache_interface import KVCacheTensor
165 +
166 + assert len(gpu_config.kv_cache_tensors) > 0
167 +
168 + gpu_total_bytes = sum(t.size for t in gpu_config.kv_cache_tensors)
169 + num_gpu_blocks = gpu_config.num_blocks
170 + num_cpu_blocks = max(1, num_gpu_blocks * cpu_capacity_bytes // gpu_total_bytes)
171 + # Create CPU kv_cache_tensors mirroring GPU by scaling size proportionally.
172 + cpu_tensors = [
173 + KVCacheTensor(
174 + size=t.size // num_gpu_blocks * num_cpu_blocks,
175 + shared_by=list(t.shared_by),
176 + )
177 + for t in gpu_config.kv_cache_tensors
178 + ]
179 +
180 + return KVCacheConfigCls(
181 + num_blocks=num_cpu_blocks,
182 + kv_cache_tensors=cpu_tensors,
183 + kv_cache_groups=gpu_config.kv_cache_groups,
184 + )
185 +
186 + @staticmethod
187 + def _estimate_lazy_target_blocks(
188 + kv_cache_config: "KVCacheConfig", max_num_batched_tokens: int
189 + ) -> int:
190 + """GPU blocks to keep available (free/offloaded) per step in lazy mode."""
191 + WATERMARK_RATIO = 1.0 # Reserve larger space to avoid running out of GPU blocks
192 + target = 0
193 + for g in kv_cache_config.kv_cache_groups:
194 + spec = g.kv_cache_spec
195 + if isinstance(spec, MambaSpec):
196 + target += 2
197 + elif isinstance(spec, SlidingWindowSpec):
198 + target += cdiv(spec.sliding_window, spec.block_size) + 1
199 + else:
200 + target += cdiv(max_num_batched_tokens, spec.block_size)
201 + return int(target * (1 + WATERMARK_RATIO))
202 +
203 + def bind_gpu_block_pool(self, gpu_block_pool: BlockPool) -> None:
204 + """Bind GPU block pool so that we can touch blocks during stores.
205 + Called by Scheduler after kv_cache_manager is ready."""
206 + self._gpu_block_pool = gpu_block_pool
207 +
208 + def get_num_new_matched_tokens(
209 + self, request: "Request", num_computed_tokens: int
210 + ) -> tuple[int | None, bool]:
211 + """Return (num_new_tokens, is_async) from consecutive CPU cache hits."""
212 + skipped = num_computed_tokens // self.block_size
213 + remaining_hashes = request.block_hashes[skipped:]
214 +
215 + if not remaining_hashes:
216 + return 0, False
217 + # Must recompute at least the last token, matching the logic in
218 + # kv_cache_manager.get_computed_blocks().
219 + max_hit_len = request.num_tokens - 1 - num_computed_tokens
220 + if max_hit_len <= 0:
221 + return 0, False
222 + _, hit_length = self.cpu_coordinator.find_longest_cache_hit(
223 + remaining_hashes, max_hit_len
224 + )
225 +
226 + if hit_length > 0:
227 + return hit_length, True
228 + return 0, False
229 +
230 + # TODO(yifan): this API now only matches the suffix part of the prefix cache. A more
231 + # general API should scan blocks in both GPU and CPU block pool in a single pass.
232 + def update_state_after_alloc(
233 + self,
234 + request: "Request",
235 + blocks: "KVCacheBlocks",
236 + num_external_tokens: int,
237 + ) -> None:
238 + req_id = request.request_id
239 + block_ids_by_group = blocks.get_block_ids()
240 + num_groups = len(block_ids_by_group)
241 +
242 + # Store tracking (eager mode only). Register the request;
243 + # block IDs are accumulated from scheduler_output in
244 + # _prepare_eager_store_specs via yield_req_data.
245 + if not self._lazy_mode and req_id not in self._reqs_to_store:
246 + self._reqs_to_store[req_id] = StoreRequestState(
247 + request=request,
248 + block_ids=tuple([] for _ in range(num_groups)),
249 + num_stored_blocks=[0] * num_groups,
250 + )
251 +
252 + if num_external_tokens == 0:
253 + return
254 +
255 + num_blocks_to_load = num_external_tokens // self.block_size
256 + assert num_blocks_to_load > 0
257 +
258 + skipped = sum(blk.block_hash is not None for blk in blocks.blocks[self.fa_gidx])
259 + num_computed_tokens = skipped * self.block_size
260 + hashes_to_load = request.block_hashes[skipped : skipped + num_blocks_to_load]
261 +
262 + # Find CPU cached blocks across all groups.
263 + max_hit_len = len(hashes_to_load) * self.block_size
264 + cpu_hit_blocks, hit_length = self.cpu_coordinator.find_longest_cache_hit(
265 + hashes_to_load, max_hit_len
266 + )
267 + assert hit_length == num_external_tokens, (
268 + f"Expected {num_external_tokens} hit tokens, got {hit_length}"
269 + )
270 +
271 + # Build transfer pairs across all groups.
272 + total_computed_tokens = num_computed_tokens + num_external_tokens
273 + kv_cache_groups = self.cpu_kv_cache_config.kv_cache_groups
274 +
275 + gpu_block_ids: list[int] = []
276 + cpu_block_ids: list[int] = []
277 + cpu_blocks_to_touch: list[KVCacheBlock] = []
278 +
279 + for g in range(num_groups):
280 + cpu_blocks_g = cpu_hit_blocks[g]
281 + n_ext_g = len(cpu_blocks_g)
282 + if n_ext_g == 0:
283 + continue
284 +
285 + # Number of blocks in the computed range for this group.
286 + g_block_size = kv_cache_groups[g].kv_cache_spec.block_size
287 + n_computed_g = cdiv(total_computed_tokens, g_block_size)
288 +
289 + # Back-trace: ext blocks sit at the tail of the computed range.
290 + gpu_ext_start = n_computed_g - n_ext_g
291 + group_gpu_ids = block_ids_by_group[g]
292 +
293 + for i, cpu_blk in enumerate(cpu_blocks_g):
294 + # Skip null blocks (e.g. sliding window or mamba padding).
295 + if cpu_blk.is_null:
296 + continue
297 + gpu_block_ids.append(group_gpu_ids[gpu_ext_start + i])
298 + cpu_block_ids.append(cpu_blk.block_id)
299 + cpu_blocks_to_touch.append(cpu_blk)
300 +
301 + # Touch CPU blocks to prevent eviction during async load.
302 + self.cpu_block_pool.touch(cpu_blocks_to_touch)
303 +
304 + # Touch GPU blocks to prevent freeing during async load
305 + assert self._gpu_block_pool is not None
306 + self._gpu_block_pool.touch(
307 + [self._gpu_block_pool.blocks[bid] for bid in gpu_block_ids]
308 + )
309 +
310 + assert self._reqs_to_load.get(req_id) is None
311 + self._reqs_to_load[req_id] = LoadRequestState(
312 + request=request, transfer_meta=TransferMeta(gpu_block_ids, cpu_block_ids)
313 + )
314 +
315 + def build_connector_meta(
316 + self,
317 + scheduler_output: SchedulerOutput,
318 + ) -> SimpleCPUOffloadMetadata:
319 + # --- Stores ---
320 + store_event = -1
321 + store_gpu, store_cpu, store_req_ids = self.prepare_store_specs(scheduler_output)
322 + if store_gpu:
323 + store_event = self._store_event_counter
324 + self._store_event_counter += 1
325 + self._store_event_to_blocks[store_event] = TransferMeta(
326 + store_gpu, store_cpu
327 + )
328 + if store_req_ids: # For eager mode only, track req->blocks mapping
329 + self._store_event_to_reqs[store_event] = store_req_ids
330 + for req_id in store_req_ids:
331 + store_state = self._reqs_to_store.get(req_id)
332 + if store_state is not None:
333 + store_state.store_events.add(store_event)
334 +
335 + # --- Loads ---
336 + load_event = -1
337 + load_gpu: list[int] = []
338 + load_cpu: list[int] = []
339 + load_req_ids: list[str] = []
340 + for req_id, load_state in self._reqs_to_load.items():
341 + if load_state.load_event is not None:
342 + continue
343 + assert load_state.transfer_meta is not None
344 + load_gpu.extend(load_state.transfer_meta.gpu_block_ids)
345 + load_cpu.extend(load_state.transfer_meta.cpu_block_ids)
346 + load_req_ids.append(req_id)
347 + if load_req_ids:
348 + load_event = self._load_event_counter
349 + self._load_event_counter += 1
350 + for req_id in load_req_ids:
351 + self._reqs_to_load[req_id].load_event = load_event
352 + self._load_event_to_reqs[load_event] = load_req_ids
353 +
354 + result = SimpleCPUOffloadMetadata(
355 + load_event=load_event,
356 + load_gpu_blocks=load_gpu,
357 + load_cpu_blocks=load_cpu,
358 + load_event_to_reqs=self._load_event_to_reqs,
359 + store_event=store_event,
360 + store_gpu_blocks=store_gpu,
361 + store_cpu_blocks=store_cpu,
362 + need_flush=bool(scheduler_output.preempted_req_ids),
363 + )
364 + return result
365 +
366 + def prepare_store_specs(
367 + self, scheduler_output: SchedulerOutput
368 + ) -> tuple[list[int], list[int], list[str]]:
369 + """Prepare store specs for the store event."""
370 + if self._lazy_mode:
371 + return self._prepare_lazy_store_specs()
372 + else:
373 + return self._prepare_eager_store_specs(scheduler_output)
374 +
375 + def _prepare_lazy_store_specs(
376 + self,
377 + ) -> tuple[list[int], list[int], list[str]]:
378 + """Single-pass cursor walk: offload cached GPU blocks near eviction.
379 +
380 + Walks the GPU free queue from the cursor, counting blocks that are
381 + free-or-offloaded (safe for the allocator to evict). Stops when
382 + target_free blocks are covered or CPU capacity is reached.
383 + """
384 + gpu_pool = self._gpu_block_pool
385 + if gpu_pool is None or self._target_free <= 0:
386 + return [], [], []
387 +
388 + free_queue = gpu_pool.free_block_queue
389 + cpu_pool = self.cpu_block_pool
390 + num_cpu_free = cpu_pool.get_num_free_blocks()
391 +
392 + # Validate cursor: stale if block was removed from free queue.
393 + if self._cursor is not None and self._cursor.ref_cnt > 0:
394 + self._cursor = None
395 +
396 + # Determine start node.
397 + if self._cursor is None:
398 + node = free_queue.fake_free_list_head.next_free_block
399 + else:
400 + node = self._cursor.next_free_block
401 +
402 + tail = free_queue.fake_free_list_tail
403 + gpu_ids: list[int] = []
404 + block_hashes: list[bytes] = []
405 + covered = 0
406 + last_visited = self._cursor
407 +
408 + while (
409 + node is not None
410 + and node is not tail
411 + and covered < self._target_free
412 + and len(gpu_ids) < num_cpu_free
413 + ):
414 + last_visited = node
415 + bhash = node.block_hash
416 +
417 + if (
418 + bhash is not None
419 + and not node.is_null
420 + and cpu_pool.cached_block_hash_to_block.get_one_block(bhash) is None
421 + ):
422 + gpu_ids.append(node.block_id)
423 + block_hashes.append(bhash)
424 +
425 + covered += 1
426 + node = node.next_free_block
427 +
428 + self._cursor = last_visited
429 +
430 + # Batch-allocate CPU blocks and stamp hashes.
431 + if gpu_ids:
432 + cpu_blocks = cpu_pool.get_new_blocks(len(gpu_ids))
433 + cpu_ids = [blk.block_id for blk in cpu_blocks]
434 + for cpu_blk, bhash in zip(cpu_blocks, block_hashes): # type: ignore[assignment]
435 + cpu_blk._block_hash = bhash # type: ignore[assignment]
436 + # Touch GPU blocks to prevent eviction during async copy.
437 + gpu_pool.touch([gpu_pool.blocks[bid] for bid in gpu_ids])
438 + else:
439 + cpu_ids = []
440 +
441 + return gpu_ids, cpu_ids, []
442 +
443 + def _prepare_eager_store_specs(
444 + self, scheduler_output: SchedulerOutput
445 + ) -> tuple[list[int], list[int], list[str]]:
446 + """Identify newly computed blocks to offload from scheduler requests.
447 +
448 + Only considers blocks whose KV data has been **confirmed computed** by
449 + the GPU. This means blocks from the current step are NOT stored until the
450 + next step. If a request finishes in the same step as its last full block,
451 + that block may be missed. (TODO: flush on finish.)
452 +
453 + Returns:
454 + (gpu_block_ids, cpu_block_ids, req_ids) for the store event.
455 + """
456 +
457 + merged_gpu_block_ids: list[int] = []
458 + merged_cpu_block_ids: list[int] = []
459 + req_ids: list[str] = []
460 +
461 + gpu_block_pool = self._gpu_block_pool
462 + if gpu_block_pool is None:
463 + return [], [], []
464 + cpu_block_pool = self.cpu_block_pool
465 + num_free = cpu_block_pool.get_num_free_blocks()
466 + kv_cache_groups = self.cpu_kv_cache_config.kv_cache_groups
467 + num_groups = len(kv_cache_groups)
468 + gpu_blocks_this_step: set[int] = set()
469 +
470 + for req_id, new_block_id_groups, preempted in yield_req_data(scheduler_output):
471 + state = self._reqs_to_store.get(req_id)
472 + if state is None or state.finished:
473 + continue
474 +
475 + # Accumulate new block IDs.
476 + if preempted:
477 + state.block_ids = tuple([] for _ in range(num_groups))
478 + state.num_stored_blocks = [0] * num_groups
479 + if new_block_id_groups:
480 + for g in range(min(num_groups, len(new_block_id_groups))):
481 + if new_block_id_groups[g] is not None:
482 + state.block_ids[g].extend(new_block_id_groups[g])
483 +
484 + num_new_tokens = scheduler_output.num_scheduled_tokens.get(req_id, 0)
485 + if num_new_tokens == 0:
486 + continue
487 +
488 + block_ids_by_group = state.block_ids
489 + if not block_ids_by_group:
490 + continue
491 +
492 + # --- Phase 1: Scan blocks, classify as cached vs to-store ---
493 + gpu_block_ids: list[int] = []
494 + block_hashes_to_store: list[bytes] = []
495 + advanced_per_group: list[int] = [0] * num_groups
496 + out_of_space = False
497 + # Confirmed tokens: KV data written and visible to all streams.
498 + req = state.request
499 + confirmed_tokens = req.num_computed_tokens - req.num_output_placeholders
500 +
501 + for g in range(num_groups):
502 + # FIXME (yifan): handle CPU cache eviction, where
503 + # num_stored_blocks can be stale and omit evicted blocks in
504 + # the middle of the request.
505 + already_stored_g = state.num_stored_blocks[g]
506 + group_gpu_ids = block_ids_by_group[g]
507 +
508 + # Cap to blocks with confirmed KV data.
509 + g_block_size = kv_cache_groups[g].kv_cache_spec.block_size
510 + ready_blocks_g = confirmed_tokens // g_block_size
511 + scannable = group_gpu_ids[already_stored_g:ready_blocks_g]
512 +
513 + for gpu_block_id in scannable:
514 + gpu_block = gpu_block_pool.blocks[gpu_block_id]
515 + if gpu_block.is_null:
516 + advanced_per_group[g] += 1
517 + continue
518 +
519 + bhash_with_group = gpu_block.block_hash
520 + if bhash_with_group is None:
521 + break
522 +
523 + # Check if this group's data is already scheduled for store
524 + # in this step or already cached in CPU.
525 + if (
526 + gpu_block_id in gpu_blocks_this_step
527 + or cpu_block_pool.cached_block_hash_to_block.get_one_block(
528 + bhash_with_group
529 + )
530 + is not None
531 + ):
532 + advanced_per_group[g] += 1
533 + continue
534 +
535 + if num_free <= 0:
536 + out_of_space = True
537 + break
538 + num_free -= 1
539 +
540 + gpu_block_ids.append(gpu_block_id)
541 + block_hashes_to_store.append(bhash_with_group)
542 + advanced_per_group[g] += 1
543 +
544 + if out_of_space:
545 + break
546 +
547 + # --- Phase 2: Batch allocate CPU blocks and stamp hashes ---
548 + n_to_alloc = len(gpu_block_ids)
549 + if n_to_alloc > 0:
550 + cpu_blocks_alloc = cpu_block_pool.get_new_blocks(n_to_alloc)
551 + cpu_block_ids = [blk.block_id for blk in cpu_blocks_alloc]
552 + for cpu_blk, bhash in zip(cpu_blocks_alloc, block_hashes_to_store):
553 + cpu_blk._block_hash = bhash # type: ignore[assignment]
554 + else:
555 + cpu_block_ids = []
556 +
557 + if cpu_block_ids:
558 + req_ids.append(req_id)
559 + merged_gpu_block_ids.extend(gpu_block_ids)
560 + merged_cpu_block_ids.extend(cpu_block_ids)
561 + gpu_blocks_this_step.update(gpu_block_ids)
562 +
563 + # Touch GPU blocks to prevent freeing during async copy
564 + gpu_block_pool.touch(
565 + [gpu_block_pool.blocks[bid] for bid in gpu_block_ids]
566 + )
567 +
568 + logger.debug(
569 + "Request %s: Scheduling store of %d blocks to CPU (%d groups)",
570 + req_id,
571 + len(cpu_block_ids),
572 + num_groups,
573 + )
574 +
575 + # Advance per-group cursors (includes cached hits + newly stored)
576 + for g in range(num_groups):
577 + state.num_stored_blocks[g] += advanced_per_group[g]
578 +
579 + return merged_gpu_block_ids, merged_cpu_block_ids, req_ids
580 +
581 + def update_connector_output(self, connector_output: KVConnectorOutput) -> None:
582 + """Handle async transfer completions from worker.
583 +
584 + Load completions arrive via finished_recving (real req_ids).
585 + Store completions arrive via kv_connector_worker_meta as
586 + per-event worker counts. We accumulate across steps and process
587 + a store event only when all workers have reported completion.
588 + """
589 + # --- Load completions ---
590 + for req_id in list(connector_output.finished_recving or []):
591 + self._cleanup_load_request(req_id)
592 +
593 + # --- Store completions ---
594 + meta = connector_output.kv_connector_worker_meta
595 + if not isinstance(meta, SimpleCPUOffloadWorkerMetadata):
596 + return
597 + for event_idx, count in meta.completed_store_events.items():
598 + total = self._store_event_pending_counts.get(event_idx, 0) + count
599 + if total >= self._expected_worker_count:
600 + self._store_event_pending_counts.pop(event_idx, None)
601 + self._process_store_event(event_idx)
602 + else:
603 + self._store_event_pending_counts[event_idx] = total
604 +
605 + def _process_store_event(self, event_idx: int) -> None:
606 + """Process a fully-completed store event."""
607 + transfer = self._store_event_to_blocks.pop(event_idx)
608 + self._process_store_completion(transfer.gpu_block_ids, transfer.cpu_block_ids)
609 + logger.debug(
610 + "Store event %d completed: cached %d blocks to CPU",
611 + event_idx,
612 + len(transfer.cpu_block_ids),
613 + )
614 +
615 + # Eager only: update per-req state
616 + if not self._lazy_mode:
617 + for req_id in self._store_event_to_reqs.pop(event_idx, []):
618 + state = self._reqs_to_store.get(req_id)
619 + if state is None:
620 + continue
621 + state.store_events.discard(event_idx)
622 + if state.finished and not state.store_events:
623 + self._cleanup_store_request(req_id)
624 +
625 + def _process_store_completion(
626 + self, gpu_block_ids: list[int], cpu_block_ids: list[int]
627 + ) -> None:
628 + """Cache CPU blocks per-group and release GPU refs.
629 +
630 + Block hashes were stamped on CPU blocks at allocation time (in
631 + ``_prepare_*_store_specs``). Here we just register them in the
632 + cache map so they become discoverable by the load path.
633 + """
634 + assert len(cpu_block_ids) == len(gpu_block_ids)
635 +
636 + cpu_blocks = [self.cpu_block_pool.blocks[bid] for bid in cpu_block_ids]
637 +
638 + for cpu_block in cpu_blocks:
639 + bhash = cpu_block.block_hash
640 + assert bhash is not None
641 + self.cpu_block_pool.cached_block_hash_to_block.insert(bhash, cpu_block)
642 +
643 + # Free CPU and GPU blocks' ref counts to turn them into prefix cache
644 + self.cpu_block_pool.free_blocks(cpu_blocks)
645 + assert self._gpu_block_pool is not None
646 + self._gpu_block_pool.free_blocks(
647 + self._gpu_block_pool.blocks[bid] for bid in gpu_block_ids
648 + )
649 +
650 + def has_pending_stores(self) -> bool:
651 + """Return True if there are in-flight store transfers."""
652 + return bool(self._store_event_to_blocks)
653 +
654 + def request_finished(
655 + self,
656 + request: "Request",
657 + block_ids: list[int],
658 + ) -> tuple[bool, dict[str, Any] | None]:
659 + """Always returns (False, None). GPU blocks are protected by ref_cnt,
660 + so the scheduler can free blocks immediately."""
661 + req_id = request.request_id
662 +
663 + # Handle load: defer cleanup if load is in-flight
664 + load_state = self._reqs_to_load.get(req_id)
665 + if load_state is not None:
666 + if load_state.load_event is not None:
667 + load_state.finished = True # Defer: load in-flight
668 + else:
669 + self._cleanup_load_request(req_id)
670 +
671 + # Handle store (eager mode only): defer cleanup if stores in-flight
672 + if not self._lazy_mode:
673 + store_state = self._reqs_to_store.get(req_id)
674 + if store_state is not None:
675 + if store_state.store_events:
676 + store_state.finished = True # Defer: stores in-flight
677 + else:
678 + self._cleanup_store_request(req_id)
679 +
680 + return False, None
681 +
682 + def request_finished_all_groups(
683 + self,
684 + request: "Request",
685 + block_ids: tuple[list[int], ...],
686 + ) -> tuple[bool, dict[str, Any] | None]:
687 + return self.request_finished(request, block_ids=[])
688 +
689 + def _cleanup_load_request(self, req_id: str) -> None:
690 + """Release all load resources for a request.
691 +
692 + Shared between request_finished() and update_connector_output() paths.
693 + Removes the request from _reqs_to_load, cleans up event mappings,
694 + and frees CPU/GPU touch refs.
695 + """
696 + state = self._reqs_to_load.pop(req_id, None)
697 + if state is None:
698 + return
699 + # Remove from load event mapping (only this req, not whole event)
700 + if state.load_event is not None:
701 + reqs = self._load_event_to_reqs.get(state.load_event)
702 + if reqs is not None:
703 + with contextlib.suppress(ValueError):
704 + reqs.remove(req_id)
705 + if not reqs:
706 + self._load_event_to_reqs.pop(state.load_event, None)
707 +
708 + if state.transfer_meta is not None:
709 + # Free CPU touch refs
710 + self.cpu_block_pool.free_blocks(
711 + self.cpu_block_pool.blocks[bid]
712 + for bid in state.transfer_meta.cpu_block_ids
713 + )
714 + # Free GPU touch refs
715 + assert self._gpu_block_pool is not None
716 + self._gpu_block_pool.free_blocks(
717 + self._gpu_block_pool.blocks[bid]
718 + for bid in state.transfer_meta.gpu_block_ids
719 + )
720 +
721 + def _cleanup_store_request(self, req_id: str) -> None:
722 + """Release store metadata for a request.
723 +
724 + Metadata-only cleanup but no block freeing. Job completion handles
725 + block caching and GPU ref freeing via _process_store_completion().
726 + """
727 + state = self._reqs_to_store.pop(req_id, None)
728 + if state is None:
729 + return
730 + for event_idx in list(state.store_events):
731 + if (reqs := self._store_event_to_reqs.get(event_idx)) is not None:
732 + with contextlib.suppress(ValueError):
733 + reqs.remove(req_id)
734 + if not reqs:
735 + self._store_event_to_reqs.pop(event_idx, None)
736 + state.store_events.clear()
737 +
738 + def take_events(self) -> Iterable[KVCacheEvent]:
739 + return self.cpu_block_pool.take_events()
vllm/v1/simple_kv_offload/metadata.py new
+60
@@ -0,0 +1,60 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Metadata for SimpleCPUOffloadConnector."""
4 +
5 +from dataclasses import dataclass, field
6 +
7 +from vllm.distributed.kv_transfer.kv_connector.v1.base import (
8 + KVConnectorMetadata,
9 + KVConnectorWorkerMetadata,
10 +)
11 +
12 +INVALID_JOB_ID = -1
13 +
14 +
15 +@dataclass
16 +class SimpleCPUOffloadMetadata(KVConnectorMetadata):
17 + """
18 + Metadata passed from scheduler to worker for CPU offload operations.
19 +
20 + The worker receives flat block lists keyed by a monotonic event_idx.
21 + Job->req_id translation is handled by the scheduler-side manager
22 + (via inverse maps), so the worker never knows about request identities.
23 + """
24 +
25 + # Load event per step. INVALID_JOB_ID means no blocks to load this step.
26 + load_event: int = INVALID_JOB_ID
27 + load_gpu_blocks: list[int] = field(default_factory=list)
28 + load_cpu_blocks: list[int] = field(default_factory=list)
29 + # Reverse map: load_event->req_ids, for tracking requests with finished load events
30 + load_event_to_reqs: dict[int, list[str]] = field(default_factory=dict)
31 +
32 + # Store event per step. INVALID_JOB_ID means no blocks to store this step.
33 + store_event: int = INVALID_JOB_ID
34 + store_gpu_blocks: list[int] = field(default_factory=list)
35 + store_cpu_blocks: list[int] = field(default_factory=list)
36 +
37 + # Whether any requests were preempted this step and need flush pending transfers.
38 + need_flush: bool = False
39 +
40 +
41 +@dataclass
42 +class SimpleCPUOffloadWorkerMetadata(KVConnectorWorkerMetadata):
43 + """Worker -> Scheduler metadata for completed store events.
44 +
45 + Each worker reports {event_idx: 1} for newly completed stores.
46 + ``aggregate()`` sums counts across workers within a step.
47 + The scheduler-side manager accumulates across steps and processes
48 + a store completion only when count reaches ``world_size``.
49 + """
50 +
51 + completed_store_events: dict[int, int]
52 +
53 + def aggregate(
54 + self, other: "KVConnectorWorkerMetadata"
55 + ) -> "KVConnectorWorkerMetadata":
56 + assert isinstance(other, SimpleCPUOffloadWorkerMetadata)
57 + merged = dict(self.completed_store_events)
58 + for k, v in other.completed_store_events.items():
59 + merged[k] = merged.get(k, 0) + v
60 + return SimpleCPUOffloadWorkerMetadata(completed_store_events=merged)
vllm/v1/simple_kv_offload/worker.py new
+305
@@ -0,0 +1,305 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Worker-side handler for SimpleCPUOffloadConnector."""
4 +
5 +from typing import TYPE_CHECKING
6 +
7 +import torch
8 +
9 +from vllm.config import VllmConfig
10 +from vllm.logger import init_logger
11 +from vllm.utils.platform_utils import is_pin_memory_available
12 +from vllm.v1.simple_kv_offload.copy_backend import DmaCopyBackend
13 +from vllm.v1.simple_kv_offload.cuda_mem_ops import pin_tensor
14 +from vllm.v1.simple_kv_offload.metadata import (
15 + SimpleCPUOffloadMetadata,
16 + SimpleCPUOffloadWorkerMetadata,
17 +)
18 +
19 +if TYPE_CHECKING:
20 + from vllm.v1.kv_cache_interface import KVCacheConfig
21 +
22 +logger = init_logger(__name__)
23 +
24 +
25 +class SimpleCPUOffloadWorker:
26 + """Worker-side handler for CPU offloading transfers."""
27 +
28 + def __init__(
29 + self,
30 + vllm_config: VllmConfig,
31 + kv_cache_config: "KVCacheConfig | None",
32 + cpu_capacity_bytes: int,
33 + ):
34 + self.vllm_config = vllm_config
35 + self.kv_cache_config = kv_cache_config
36 + self.cpu_capacity_bytes = cpu_capacity_bytes
37 +
38 + self.gpu_kv_caches: dict[str, torch.Tensor] | None = None
39 + self.cpu_kv_caches: dict[str, torch.Tensor] | None = None
40 + self.device: torch.device | None = None
41 + self.num_cpu_blocks: int = 0
42 +
43 + # CUDA streams for the async transfers
44 + self.load_stream: torch.cuda.Stream | None = None
45 + self.store_stream: torch.cuda.Stream | None = None
46 +
47 + self._backend = DmaCopyBackend()
48 +
49 + # Ordered (event_idx, Event). Events pre-allocated on main thread.
50 + self._load_events: list[tuple[int, torch.Event]] = []
51 + self._store_events: list[tuple[int, torch.Event]] = []
52 + # High-water marks: highest event_idx completed per stream.
53 + # When the event list is empty, the hwm covers all prior events.
54 + self._load_hwm: int = -1
55 + self._store_hwm: int = -1
56 +
57 + # Metadata for the current step
58 + self._connector_metadata: SimpleCPUOffloadMetadata | None = None
59 +
60 + # Pending event index sets, populated in bind_connector_metadata
61 + self._pending_load_event_indices: set[int] = set()
62 + self._pending_store_event_indices: set[int] = set()
63 + # Completed store events to report via build_connector_worker_meta
64 + self._completed_store_events: dict[int, int] = {}
65 +
66 + def register_kv_caches(
67 + self,
68 + kv_caches: dict[str, torch.Tensor],
69 + ) -> None:
70 + """Register GPU KV caches and allocate pinned CPU tensors.
71 + The worker will infer the underlying raw storage from the kv_caches.
72 +
73 + Args:
74 + kv_caches: Per-layer GPU KV caches. Values are either a single
75 + tensor (attention layers) or a list of tensors (Mamba layers
76 + in hybrid models). All values are included for offloading
77 + by resolving to their underlying raw storage.
78 + """
79 + if not kv_caches:
80 + logger.warning("No KV caches to offload.")
81 + return
82 +
83 + # Resolve each entry to a representative tensor for storage
84 + # deduplication. For attention layers the value is already a tensor;
85 + # for Mamba layers it is a list of tensors that all share the same
86 + # underlying raw storage, so we take the first one.
87 + def _repr_tensor(v: torch.Tensor | list[torch.Tensor]) -> torch.Tensor:
88 + assert isinstance(v, torch.Tensor | list)
89 + return v if isinstance(v, torch.Tensor) else v[0]
90 +
91 + any_tensor = _repr_tensor(next(iter(kv_caches.values())))
92 + self.device = any_tensor.device
93 +
94 + assert self.kv_cache_config is not None
95 + num_blocks = self.kv_cache_config.num_blocks
96 +
97 + # Deduplicate: multiple layers may share the same backing storage.
98 + seen_ptrs: dict[int, tuple[str, torch.Tensor]] = {}
99 + for name, value in kv_caches.items():
100 + tensor = _repr_tensor(value)
101 + ptr = tensor.untyped_storage().data_ptr()
102 + if ptr not in seen_ptrs:
103 + seen_ptrs[ptr] = (name, tensor)
104 +
105 + # Build [num_blocks, block_bytes] int8 views from each unique
106 + # storage so that stride(0) gives block_bytes for the copy op.
107 + #
108 + # The physical layout varies across attention backends:
109 + # FlashAttn/ROCm: (2, num_blocks, ...) -> K/V outermost, 2 segments
110 + # FlashInfer/MLA: (num_blocks, ...) -> blocks outermost, 1 segment
111 + # We derive page_size_bytes = storage.nbytes() // num_blocks, then
112 + # classify dims: any dim whose byte-stride exceeds page_size_bytes
113 + # must be an outer segment dim (e.g. the K/V dim of size 2). A less
114 + # hacky way is to update the interface with the layout.
115 + unique_gpu_caches: dict[str, torch.Tensor] = {}
116 + for name, tensor in seen_ptrs.values():
117 + storage = tensor.untyped_storage()
118 + raw = torch.empty(0, dtype=torch.int8, device=self.device).set_(
119 + storage, 0, (storage.nbytes(),)
120 + )
121 + el = tensor.element_size()
122 + page_size_bytes = storage.nbytes() // num_blocks
123 + outer_dims = [
124 + d for d in range(tensor.ndim) if tensor.stride(d) * el > page_size_bytes
125 + ]
126 + if not outer_dims:
127 + unique_gpu_caches[name] = raw.view(num_blocks, -1)
128 + else:
129 + seg_stride = tensor.stride(outer_dims[0]) * el
130 + for idx in range(tensor.shape[outer_dims[0]]):
131 + offset = idx * seg_stride
132 + chunk = raw[offset : offset + seg_stride]
133 + unique_gpu_caches[f"{name}.{idx}"] = chunk.view(num_blocks, -1)
134 +
135 + # Compute per-tensor bytes_per_block. Tensors may have different
136 + # page_size_bytes (e.g., UniformTypeKVCacheSpecs with varying head_size).
137 + per_tensor_bpb = [
138 + t.stride(0) * t.element_size() for t in unique_gpu_caches.values()
139 + ]
140 + total_bytes_per_block = sum(per_tensor_bpb)
141 +
142 + self.num_cpu_blocks = max(1, self.cpu_capacity_bytes // total_bytes_per_block)
143 +
144 + logger.info(
145 + "SimpleCPUOffloadWorker: %d unique GPU KV tensors, "
146 + "allocating %d CPU blocks (%.2f GB)",
147 + len(unique_gpu_caches),
148 + self.num_cpu_blocks,
149 + (self.num_cpu_blocks * total_bytes_per_block) / (1024**3),
150 + )
151 +
152 + pin_memory = is_pin_memory_available()
153 + if not pin_memory:
154 + logger.warning(
155 + "Pinned memory not available. CPU offload performance may be degraded."
156 + )
157 +
158 + self.gpu_kv_caches = unique_gpu_caches
159 + self.cpu_kv_caches = {}
160 + for name, gpu_tensor in unique_gpu_caches.items():
161 + cpu_shape = (self.num_cpu_blocks,) + gpu_tensor.shape[1:]
162 + # Allocate non-pinned first, then pin via cudaHostRegister to
163 + # bypass PyTorch's CUDACachingHostAllocator which rounds up to
164 + # the next power of 2 (e.g. 100 GB -> 128 GB).
165 + tensor = torch.zeros(cpu_shape, dtype=gpu_tensor.dtype, device="cpu")
166 + if pin_memory:
167 + pin_tensor(tensor)
168 + self.cpu_kv_caches[name] = tensor
169 +
170 + # Use lowest priority so KV cache I/O yields to compute streams.
171 + low_pri, _ = torch.cuda.Stream.priority_range()
172 + self.load_stream = torch.cuda.Stream(priority=low_pri)
173 + self.store_stream = torch.cuda.Stream(priority=low_pri)
174 +
175 + # Initialize copy backend with caches and streams.
176 + self._backend.init(
177 + self.gpu_kv_caches,
178 + self.cpu_kv_caches,
179 + self.device,
180 + self.load_stream,
181 + self.store_stream,
182 + )
183 +
184 + def bind_connector_metadata(self, metadata: SimpleCPUOffloadMetadata) -> None:
185 + self._connector_metadata = metadata
186 + if metadata.load_event >= 0:
187 + self._pending_load_event_indices.add(metadata.load_event)
188 + if metadata.store_event >= 0:
189 + self._pending_store_event_indices.add(metadata.store_event)
190 +
191 + def clear_connector_metadata(self) -> None:
192 + self._connector_metadata = None
193 +
194 + def start_load_kv(self) -> None:
195 + # NOTE: we defer launching both load and store to get_finished(),
196 + # which runs after model execution. This hides the CPU-side
197 + # block copy op overhead (~5ms) behind GPU compute.
198 + pass
199 +
200 + def wait_for_save(self) -> None:
201 + pass
202 +
203 + def get_finished(
204 + self,
205 + finished_req_ids: set[str],
206 + ) -> tuple[set[str] | None, set[str] | None]:
207 + """Submit transfers and report completed events to the scheduler.
208 +
209 + Called after model execution. The manager only schedules stores for
210 + blocks whose KV data is confirmed computed, so we launch both loads
211 + and stores immediately — no deferral or cross-stream sync needed.
212 +
213 + Returns:
214 + tuple of (finished_sending, finished_recving).
215 + - finished_sending: always None (stores use worker metadata).
216 + - finished_recving: req_ids whose loads have completed.
217 + """
218 + # (1) Submit transfers
219 + metadata = self._connector_metadata
220 + if metadata is not None:
221 + # Launch loads (CPU->GPU).
222 + if metadata.load_cpu_blocks:
223 + self._backend.launch_copy(
224 + metadata.load_cpu_blocks,
225 + metadata.load_gpu_blocks,
226 + is_store=False,
227 + event_idx=metadata.load_event,
228 + events_list=self._load_events,
229 + )
230 + # Launch stores (GPU->CPU).
231 + if metadata.store_gpu_blocks:
232 + self._backend.launch_copy(
233 + metadata.store_gpu_blocks,
234 + metadata.store_cpu_blocks,
235 + is_store=True,
236 + event_idx=metadata.store_event,
237 + events_list=self._store_events,
238 + )
239 +
240 + # (2) Track completed transfer events
241 + finished_recving: set[str] = set()
242 +
243 + if self._pending_load_event_indices:
244 + load_wm = self._poll_stream_events(is_store=False)
245 + for j in [j for j in self._pending_load_event_indices if j <= load_wm]:
246 + self._pending_load_event_indices.discard(j)
247 + req_ids = (
248 + metadata.load_event_to_reqs.get(j) if metadata is not None else None
249 + )
250 + if req_ids:
251 + finished_recving.update(req_ids)
252 +
253 + if self._pending_store_event_indices:
254 + store_wm = self._poll_stream_events(is_store=True)
255 + for j in [j for j in self._pending_store_event_indices if j <= store_wm]:
256 + self._pending_store_event_indices.discard(j)
257 + self._completed_store_events[j] = 1
258 +
259 + return None, finished_recving or None
260 +
261 + def build_connector_worker_meta(self) -> SimpleCPUOffloadWorkerMetadata | None:
262 + """Return completed store events since the last call."""
263 + if not self._completed_store_events:
264 + return None
265 + meta = SimpleCPUOffloadWorkerMetadata(
266 + completed_store_events=self._completed_store_events,
267 + )
268 + self._completed_store_events = {}
269 + return meta
270 +
271 + def handle_preemptions(
272 + self, kv_connector_metadata: SimpleCPUOffloadMetadata
273 + ) -> None:
274 + """Sync all in-flight transfers before preempted blocks are reused."""
275 + if not kv_connector_metadata.need_flush:
276 + return
277 + self._flush_and_sync_all()
278 +
279 + def _flush_and_sync_all(self) -> None:
280 + """Synchronize all in-flight transfer events."""
281 + for event_idx, event in self._load_events:
282 + event.synchronize()
283 + self._load_hwm = event_idx
284 + self._load_events.clear()
285 +
286 + for event_idx, event in self._store_events:
287 + event.synchronize()
288 + self._store_hwm = event_idx
289 + self._store_events.clear()
290 +
291 + def _poll_stream_events(self, is_store: bool) -> int:
292 + """Non-blocking poll for completed events and return the high-water mark."""
293 + events = self._store_events if is_store else self._load_events
294 + hwm = self._store_hwm if is_store else self._load_hwm
295 + while events:
296 + event_idx, event = events[0]
297 + if not event.query():
298 + break
299 + hwm = event_idx
300 + events.pop(0)
301 + if is_store:
302 + self._store_hwm = hwm
303 + else:
304 + self._load_hwm = hwm
305 + return hwm