@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / d0a100c87a
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 2, 22:00 UTC

File system secondary tier implemented in python (#41735)

Signed-off-by: Rotem Shavitt <rshavitt@gmail.com> Signed-off-by: Or Ozeri <oro@il.ibm.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>

Rotem Shavitt committed May 24, 2026 at 21:14 UTC d0a100c87af832ad97ade60b8ec7610018a08427
9 files changed +980 -4
tests/v1/kv_connector/unit/test_offloading_connector.py
+56 -4
@@ -107,7 +107,7 @@ class MockSubscriber:
107 self.sub.close()
108
109
110 -def _wait_for_prefix_cache_reset(llm: LLM) -> None:
110 +def _wait_for_prefix_cache_reset(llm: LLM, reset_connector: bool = False) -> None:
111 """Wait for async offload transfers to finish so prefix cache can reset.
112
113 The GPU-to-CPU offload runs on a CUDA stream asynchronously. While blocks
@@ -115,10 +115,14 @@ def _wait_for_prefix_cache_reset(llm: LLM) -> None:
115 ``False``. Between retries we send a dummy single-token prefill to force
116 the engine to step, which polls the worker for completed transfers and
117 frees GPU blocks.
118 +
119 + Args:
120 + llm: The LLM instance to reset.
121 + reset_connector: If True, also reset the KV connector state.
122 """
123 _dummy_params = SamplingParams(max_tokens=1)
124 deadline = time.monotonic() + _RESET_CACHE_TIMEOUT
121 - while not llm.reset_prefix_cache():
125 + while not llm.reset_prefix_cache(reset_connector=reset_connector):
126 if time.monotonic() > deadline:
127 raise TimeoutError(
128 "reset_prefix_cache did not succeed within "
@@ -133,7 +137,9 @@ def _wait_for_prefix_cache_reset(llm: LLM) -> None:
137 )
138
139
136 -def _latency_test(llm: LLM, subscriber: MockSubscriber | None):
140 +def _latency_test(
141 + llm: LLM, subscriber: MockSubscriber | None, reset_connector: bool = False
142 +):
143 sampling_params = SamplingParams(max_tokens=1)
144
145 num_times_cpu_better_than_cold = 0
@@ -163,7 +169,7 @@ def _latency_test(llm: LLM, subscriber: MockSubscriber | None):
169
170 # Wait for the async CPU offload to finish, then reset prefix cache
171 # so the next generate() must reload from CPU rather than GPU.
166 - _wait_for_prefix_cache_reset(llm)
172 + _wait_for_prefix_cache_reset(llm, reset_connector=reset_connector)
173
174 # Verify CPU stored events arrived (offload is done before we
175 # attempt to load from CPU).
@@ -337,3 +343,49 @@ def test_tiering_offloading() -> None:
343 finally:
344 subscriber.close()
345 del llm
346 +
347 +
348 +def test_fs_tiering_offloading(tmp_path) -> None:
349 + """Tests OffloadingConnector with TieringOffloadingSpec
350 + + fs_python secondary tier."""
351 + extra_config: dict = {
352 + "cpu_bytes_to_use": 1 << 30,
353 + "block_size": 48,
354 + "spec_name": "TieringOffloadingSpec",
355 + "secondary_tiers": [{"type": "fs_python", "root_dir": str(tmp_path)}],
356 + }
357 + kv_transfer_config = KVTransferConfig(
358 + kv_connector="OffloadingConnector",
359 + kv_role="kv_both",
360 + kv_connector_extra_config=extra_config,
361 + )
362 +
363 + port: int
364 + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
365 + s.bind(("0.0.0.0", 0))
366 + port = s.getsockname()[1]
367 + events_endpoint = f"tcp://*:{port}"
368 + kv_events_config = KVEventsConfig(
369 + enable_kv_cache_events=True,
370 + publisher="zmq",
371 + endpoint=events_endpoint,
372 + topic="test",
373 + )
374 +
375 + llm = LLM(
376 + model="meta-llama/Llama-3.2-1B-Instruct",
377 + max_model_len=512,
378 + gpu_memory_utilization=0.5,
379 + kv_events_config=kv_events_config,
380 + kv_transfer_config=kv_transfer_config,
381 + )
382 + subscriber = MockSubscriber(
383 + events_endpoint.replace("*", "127.0.0.1"),
384 + topic=kv_events_config.topic,
385 + )
386 + try:
387 + _latency_test(llm, subscriber, reset_connector=True)
388 + _accuracy_test(llm, subscriber)
389 + finally:
390 + subscriber.close()
391 + del llm
tests/v1/kv_offload/test_file_mapper.py new
+127
@@ -0,0 +1,127 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for FileMapper."""
4 +
5 +from unittest.mock import MagicMock
6 +
7 +from vllm.v1.kv_offload.base import (
8 + OffloadingSpec,
9 + make_offload_key,
10 +)
11 +from vllm.v1.kv_offload.file_mapper import FileMapper
12 +
13 +# ---------------------------------------------------------------------------
14 +# Shared mocks (mirrors test_fs_tier.py pattern)
15 +# ---------------------------------------------------------------------------
16 +
17 +_MOCK_VLLM_CONFIG = MagicMock()
18 +_MOCK_VLLM_CONFIG.model_config.model = "test-model"
19 +_MOCK_VLLM_CONFIG.cache_config.block_size = 16
20 +_MOCK_VLLM_CONFIG.cache_config.cache_dtype = "torch.float32"
21 +_MOCK_VLLM_CONFIG.parallel_config.tensor_parallel_size = 1
22 +_MOCK_VLLM_CONFIG.parallel_config.pipeline_parallel_size = 1
23 +_MOCK_VLLM_CONFIG.parallel_config.prefill_context_parallel_size = 1
24 +_MOCK_VLLM_CONFIG.parallel_config.decode_context_parallel_size = 1
25 +_MOCK_VLLM_CONFIG.parallel_config.rank = 0
26 +
27 +_MOCK_KV_CACHE_CONFIG = MagicMock()
28 +_MOCK_KV_CACHE_CONFIG.kv_cache_groups = []
29 +
30 +_MOCK_OFFLOADING_SPEC = MagicMock(spec=OffloadingSpec)
31 +_MOCK_OFFLOADING_SPEC.vllm_config = _MOCK_VLLM_CONFIG
32 +_MOCK_OFFLOADING_SPEC.kv_cache_config = _MOCK_KV_CACHE_CONFIG
33 +_MOCK_OFFLOADING_SPEC.block_size_factor = 1
34 +
35 +
36 +# ---------------------------------------------------------------------------
37 +# Helper
38 +# ---------------------------------------------------------------------------
39 +
40 +
41 +def make_mapper_from_offloading_spec(**kwargs) -> FileMapper:
42 + """Helper to create FileMapper with customizable mock config."""
43 + # Create a copy of the mock config to avoid modifying the global one
44 + mock_vllm_config = MagicMock()
45 + mock_vllm_config.model_config.model = kwargs.get("model_name", "test-model")
46 + mock_vllm_config.cache_config.block_size = kwargs.get("hash_block_size", 16)
47 + mock_vllm_config.cache_config.cache_dtype = (
48 + f"torch.{kwargs.get('dtype', 'float16')}"
49 + )
50 + mock_vllm_config.parallel_config.tensor_parallel_size = kwargs.get("tp_size", 1)
51 + mock_vllm_config.parallel_config.pipeline_parallel_size = kwargs.get("pp_size", 1)
52 + mock_vllm_config.parallel_config.prefill_context_parallel_size = kwargs.get(
53 + "pcp_size", 1
54 + )
55 + mock_vllm_config.parallel_config.decode_context_parallel_size = kwargs.get(
56 + "dcp_size", 1
57 + )
58 + mock_vllm_config.parallel_config.rank = kwargs.get("rank", 0)
59 +
60 + mock_kv_cache_config = MagicMock()
61 + mock_kv_cache_config.kv_cache_groups = []
62 +
63 + mock_offloading_spec = MagicMock(spec=OffloadingSpec)
64 + mock_offloading_spec.vllm_config = mock_vllm_config
65 + mock_offloading_spec.kv_cache_config = mock_kv_cache_config
66 + mock_offloading_spec.block_size_factor = kwargs.get("block_size_factor", 1)
67 +
68 + return FileMapper.from_offloading_spec(
69 + root_dir=kwargs.get("root_dir", "/tmp/cache"),
70 + offloading_spec=mock_offloading_spec,
71 + gpu_blocks_per_file=mock_offloading_spec.block_size_factor,
72 + )
73 +
74 +
75 +# ---------------------------------------------------------------------------
76 +# Tests
77 +# ---------------------------------------------------------------------------
78 +
79 +
80 +def test_get_file_name_full_structure():
81 + """
82 + Path must match: <base_path>_r<rank>/<hhh>/<hh>_g<group_idx>/<hash_hex>.bin
83 +
84 + Concretely:
85 + - The segment immediately after base_path must end with `_r0`
86 + - The next segment is the first 3 hex chars of the block hash
87 + - The next segment is <2 hex chars>_g<group_idx>
88 + - The final segment is <full hash hex>.bin
89 + """
90 + rank = 3
91 + group_idx = 2
92 + block_hash = bytes(range(8)) # deterministic, non-zero bytes
93 + fm = make_mapper_from_offloading_spec(rank=rank)
94 + key = make_offload_key(block_hash, group_idx)
95 + path = fm.get_file_name(key)
96 +
97 + expected_path = (
98 + "/tmp/cache/test-model_588656ebcc66_r3/000/10_g2/0001020304050607.bin"
99 + )
100 + assert path == expected_path
101 +
102 +
103 +def test_get_run_config_fields():
104 + fm = make_mapper_from_offloading_spec(
105 + model_name="my-model",
106 + dtype="bfloat16",
107 + tp_size=2,
108 + )
109 + cfg = fm.get_run_config()
110 + assert cfg == {
111 + "model_name": "my-model",
112 + "hash_block_size": 16,
113 + "gpu_blocks_per_file": 1,
114 + "tp_size": 2,
115 + "pp_size": 1,
116 + "pcp_size": 1,
117 + "dcp_size": 1,
118 + "dtype": "bfloat16",
119 + "kv_cache_groups": [],
120 + "inference_engine": "vllm",
121 + }
122 +
123 +
124 +def test_get_config_file_path():
125 + fm = make_mapper_from_offloading_spec()
126 + config_path = fm.get_config_file_path()
127 + assert config_path == f"{fm.base_path}/config.json"
tests/v1/kv_offload/test_fs_tier.py new
+256
@@ -0,0 +1,256 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Unit tests for FileSystemTierManager.
5 +
6 +These tests use real disk I/O to verify the Python filesystem tier implementation.
7 +The tier manager writes KV cache blocks to disk and reads them back, verifying
8 +data integrity throughout the process.
9 +"""
10 +
11 +import os
12 +import time
13 +from unittest.mock import MagicMock
14 +
15 +import numpy as np
16 +import pytest
17 +import torch
18 +
19 +from vllm.v1.kv_offload.base import OffloadKey, ReqContext, make_offload_key
20 +from vllm.v1.kv_offload.tiering.base import JobMetadata
21 +from vllm.v1.kv_offload.tiering.fs.manager import (
22 + FileSystemTierManager,
23 +)
24 +
25 +# ---------------------------------------------------------------------------
26 +# Helpers
27 +# ---------------------------------------------------------------------------
28 +
29 +_BLOCK_ELEMENTS = 512 * 1024 # 2 MB per block (float32 × 512K = 2MB)
30 +_DTYPE = torch.float32
31 +_CTX = ReqContext(req_id="test")
32 +
33 +_MOCK_VLLM_CONFIG = MagicMock()
34 +_MOCK_VLLM_CONFIG.model_config.model = "test-model"
35 +_MOCK_VLLM_CONFIG.cache_config.block_size = 16
36 +_MOCK_VLLM_CONFIG.cache_config.cache_dtype = "torch.float32"
37 +_MOCK_VLLM_CONFIG.parallel_config.tensor_parallel_size = 1
38 +_MOCK_VLLM_CONFIG.parallel_config.pipeline_parallel_size = 1
39 +_MOCK_VLLM_CONFIG.parallel_config.prefill_context_parallel_size = 1
40 +_MOCK_VLLM_CONFIG.parallel_config.decode_context_parallel_size = 1
41 +_MOCK_VLLM_CONFIG.parallel_config.rank = 0
42 +
43 +_MOCK_KV_CACHE_CONFIG = MagicMock()
44 +_MOCK_KV_CACHE_CONFIG.kv_cache_groups = []
45 +
46 +_MOCK_OFFLOADING_SPEC = MagicMock()
47 +_MOCK_OFFLOADING_SPEC.vllm_config = _MOCK_VLLM_CONFIG
48 +_MOCK_OFFLOADING_SPEC.kv_cache_config = _MOCK_KV_CACHE_CONFIG
49 +_MOCK_OFFLOADING_SPEC.block_size_factor = 1
50 +
51 +
52 +def key(n: int) -> OffloadKey:
53 + return make_offload_key(n.to_bytes(8, "big"), 0)
54 +
55 +
56 +def make_job(
57 + job_id: int,
58 + keys: list[OffloadKey],
59 + block_ids: list[int] | None = None,
60 + is_promotion: bool = False,
61 +) -> JobMetadata:
62 + if block_ids is None:
63 + block_ids = list(range(len(keys)))
64 + return JobMetadata(
65 + job_id=job_id,
66 + keys=keys,
67 + block_ids=np.array(block_ids, dtype=np.int64),
68 + is_promotion=is_promotion,
69 + req_context=_CTX,
70 + )
71 +
72 +
73 +def drain(tier: FileSystemTierManager, max_rounds: int = 40) -> list:
74 + """
75 + Call get_finished() repeatedly until no new results arrive for 5
76 + consecutive rounds or max_rounds is reached.
77 + """
78 + results = []
79 + idle = 0
80 + for _ in range(max_rounds):
81 + time.sleep(0.01)
82 + new = list(tier.get_finished())
83 + results.extend(new)
84 + if new:
85 + idle = 0
86 + else:
87 + idle += 1
88 + if idle >= 5:
89 + break
90 + return results
91 +
92 +
93 +# ---------------------------------------------------------------------------
94 +# Fixtures
95 +# ---------------------------------------------------------------------------
96 +
97 +
98 +@pytest.fixture
99 +def fs_tier(tmp_path):
100 + tensor = torch.zeros((4, _BLOCK_ELEMENTS), dtype=_DTYPE)
101 + mock_view = memoryview(tensor.numpy())
102 + tier = FileSystemTierManager(
103 + offloading_spec=_MOCK_OFFLOADING_SPEC,
104 + primary_kv_view=mock_view,
105 + tier_type="fs_python",
106 + root_dir=str(tmp_path),
107 + n_read_threads=4,
108 + n_write_threads=4,
109 + )
110 + yield tier, tensor
111 + tier.shutdown()
112 +
113 +
114 +# ---------------------------------------------------------------------------
115 +# Tests
116 +# ---------------------------------------------------------------------------
117 +
118 +
119 +def test_lookup_empty_tier(fs_tier):
120 + tier, _ = fs_tier
121 + assert tier.lookup(key(1), _CTX) is False
122 + assert tier.lookup(key(2), _CTX) is False
123 +
124 +
125 +def test_store_creates_file_and_lookup_succeeds(fs_tier):
126 + tier, _ = fs_tier
127 + job = make_job(1, [key(1)], [0])
128 + tier.submit_store(job)
129 + results = drain(tier)
130 + assert len(results) == 1
131 + assert results[0].success
132 + assert tier.lookup(key(1), _CTX) is True
133 + dest = tier.file_mapper.get_file_name(key(1))
134 + assert os.path.exists(dest), f"Expected file at {dest}"
135 +
136 +
137 +def test_store_then_load_roundtrip(fs_tier):
138 + tier, _ = fs_tier
139 + job_s = make_job(1, [key(1), key(2)], [0, 1])
140 + tier.submit_store(job_s)
141 + store_results = drain(tier)
142 + assert all(r.success for r in store_results)
143 +
144 + assert tier.lookup(key(1), _CTX) is True
145 + assert tier.lookup(key(2), _CTX) is True
146 +
147 + job_l = make_job(2, [key(1), key(2)], [2, 3], is_promotion=True)
148 + tier.submit_load(job_l)
149 + load_results = drain(tier)
150 + assert all(r.success for r in load_results)
151 + # Blocks stay on disk after load
152 + assert tier.lookup(key(1), _CTX) is True
153 + assert tier.lookup(key(2), _CTX) is True
154 +
155 +
156 +def test_invalid_path_raises_at_construction():
157 + """Construction must fail immediately when the config file cannot be written."""
158 + tensor = torch.zeros((32, _BLOCK_ELEMENTS), dtype=_DTYPE)
159 + mock_view = memoryview(tensor.numpy())
160 +
161 + with pytest.raises(OSError):
162 + FileSystemTierManager(
163 + offloading_spec=_MOCK_OFFLOADING_SPEC,
164 + primary_kv_view=mock_view,
165 + tier_type="fs_python",
166 + root_dir="/dev/null/invalid_path",
167 + )
168 +
169 +
170 +def test_failed_load_missing_file(fs_tier):
171 + """Test that loading a block whose file does not exist results in a failed job."""
172 + tier, _ = fs_tier
173 + job = make_job(1, [key(99)], [0], is_promotion=True)
174 + tier.submit_load(job)
175 + results = drain(tier)
176 + assert len(results) == 1
177 + assert not results[0].success
178 +
179 +
180 +def test_multiple_jobs_tracked_independently(fs_tier):
181 + tier, _ = fs_tier
182 + job1 = make_job(1, [key(1)], [0])
183 + job2 = make_job(2, [key(2)], [1])
184 + tier.submit_store(job1)
185 + tier.submit_store(job2)
186 + results = drain(tier)
187 + job_ids = {r.job_id for r in results}
188 + assert job_ids == {1, 2}
189 + assert tier.lookup(key(1), _CTX) is True
190 + assert tier.lookup(key(2), _CTX) is True
191 +
192 +
193 +def test_multi_block_job_partial_failure(fs_tier):
194 + """A load job where one block file is missing yields a single failed JobResult."""
195 + tier, _ = fs_tier
196 + # Store two of three keys
197 + tier.submit_store(make_job(1, [key(10), key(11)], [0, 1]))
198 + assert all(r.success for r in drain(tier))
199 +
200 + # Load all three — key(99) was never stored
201 + tier.submit_load(
202 + make_job(2, [key(10), key(11), key(99)], [0, 1, 2], is_promotion=True)
203 + )
204 + results = drain(tier)
205 +
206 + assert len(results) == 1
207 + assert results[0].job_id == 2
208 + assert not results[0].success
209 +
210 +
211 +def test_shutdown_discards_pending_tasks(fs_tier):
212 + """Shutdown clears both queues and stops all worker threads without draining."""
213 + tier, _ = fs_tier
214 + # Submit many tasks to ensure some remain pending
215 + for i in range(10):
216 + tier.submit_store(make_job(i, [key(i)], [i % 4]))
217 +
218 + # Shutdown immediately without draining
219 + tier.shutdown()
220 +
221 + # Verify queues are cleared and threads stopped
222 + assert len(tier._pool._load_q) == 0
223 + assert len(tier._pool._store_q) == 0
224 + assert all(not t.is_alive() for t in tier._pool._threads)
225 +
226 +
227 +def test_store_load_data_integrity(fs_tier):
228 + """Data written by store must be exactly recovered by load."""
229 + tier, tensor = fs_tier
230 + # Populate tensor with random data
231 + tensor[:] = torch.rand((4, _BLOCK_ELEMENTS), dtype=_DTYPE)
232 +
233 + # Store first 2 blocks
234 + num_store = 2
235 + expected = tensor[:num_store].clone()
236 +
237 + store_ids = list(range(num_store))
238 + keys = [key(i) for i in range(num_store)]
239 +
240 + tier.submit_store(make_job(1, keys, store_ids))
241 + results = drain(tier)
242 + assert all(r.success for r in results)
243 +
244 + # Overwrite source blocks to prove data is read from disk
245 + tensor[:num_store] = 0.0
246 +
247 + # Load into last 2 blocks
248 + load_ids = [2, 3]
249 + tier.submit_load(make_job(2, keys, load_ids, is_promotion=True))
250 + results = drain(tier)
251 + assert all(r.success for r in results)
252 +
253 + for i, bid in enumerate(load_ids):
254 + assert torch.allclose(tensor[bid], expected[i]), (
255 + f"Block {bid} data mismatch after store+load"
256 + )
vllm/v1/kv_offload/file_mapper.py new
+126
@@ -0,0 +1,126 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import hashlib
5 +import json
6 +
7 +from vllm.v1.kv_offload.base import (
8 + OffloadingSpec,
9 + OffloadKey,
10 + get_offload_block_hash,
11 + get_offload_group_idx,
12 +)
13 +
14 +_BASE_PATH_HASH_LEN = 12
15 +_CONFIG_FILENAME = "config.json"
16 +
17 +
18 +class FileMapper:
19 + """
20 + FileMapper maps KV blocks (given by their hash) to file names.
21 + """
22 +
23 + def __init__(
24 + self,
25 + root_dir: str,
26 + model_name: str,
27 + hash_block_size: int,
28 + gpu_blocks_per_file: int,
29 + tp_size: int,
30 + pp_size: int,
31 + pcp_size: int,
32 + dcp_size: int,
33 + rank: int,
34 + dtype: str,
35 + kv_cache_groups: list[dict] | None = None,
36 + inference_engine: str = "vllm",
37 + parallel_agnostic: bool = False,
38 + ):
39 + """
40 + Initialize the file mapper. Each worker constructs its own, but
41 + `config.json` is shared across workers since rank lives outside the hash.
42 + When `parallel_agnostic=True`, tp/pp/pcp/dcp are forced to 1 and rank
43 + to 0 so multiple parallelism layouts collapse into the same folder.
44 + """
45 + if parallel_agnostic:
46 + tp_size = pp_size = pcp_size = dcp_size = 1
47 + rank = 0
48 + self.rank: int = rank
49 + self.fields: dict = {
50 + "model_name": model_name,
51 + "hash_block_size": hash_block_size,
52 + "gpu_blocks_per_file": gpu_blocks_per_file,
53 + "tp_size": tp_size,
54 + "pp_size": pp_size,
55 + "pcp_size": pcp_size,
56 + "dcp_size": dcp_size,
57 + "dtype": str(dtype),
58 + "kv_cache_groups": kv_cache_groups or [],
59 + "inference_engine": inference_engine,
60 + }
61 + self.base_path: str = self._compute_base_path(root_dir, self.fields)
62 +
63 + @classmethod
64 + def from_offloading_spec(
65 + cls,
66 + root_dir: str,
67 + offloading_spec: OffloadingSpec,
68 + gpu_blocks_per_file: int = 1,
69 + parallel_agnostic: bool = False,
70 + ) -> "FileMapper":
71 + """Build a FileMapper from an OffloadingSpec."""
72 + vllm_config = offloading_spec.vllm_config
73 + kv_cache_config = offloading_spec.kv_cache_config
74 +
75 + parallel_config = vllm_config.parallel_config
76 + dtype = str(vllm_config.cache_config.cache_dtype).replace("torch.", "")
77 + kv_cache_groups = [
78 + {
79 + "block_size": group.kv_cache_spec.block_size,
80 + "layer_names": list(group.layer_names),
81 + }
82 + for group in kv_cache_config.kv_cache_groups
83 + ]
84 + return cls(
85 + root_dir=root_dir,
86 + model_name=vllm_config.model_config.model,
87 + hash_block_size=vllm_config.cache_config.block_size,
88 + gpu_blocks_per_file=gpu_blocks_per_file,
89 + tp_size=parallel_config.tensor_parallel_size,
90 + pp_size=parallel_config.pipeline_parallel_size,
91 + pcp_size=parallel_config.prefill_context_parallel_size,
92 + dcp_size=parallel_config.decode_context_parallel_size,
93 + rank=parallel_config.rank,
94 + dtype=dtype,
95 + kv_cache_groups=kv_cache_groups,
96 + parallel_agnostic=parallel_agnostic,
97 + )
98 +
99 + def get_file_name(self, key: OffloadKey) -> str:
100 + """Map an OffloadKey to <base>_r<rank>/<hhh>/<hh>_g<group_idx>/<hash>.bin."""
101 + hash_hex = get_offload_block_hash(key).hex()
102 + group_idx = get_offload_group_idx(key)
103 + subfolder1, subfolder2 = hash_hex[:3], hash_hex[3:5]
104 + return (
105 + f"{self.base_path}_r{self.rank}"
106 + f"/{subfolder1}/{subfolder2}_g{group_idx}/{hash_hex}.bin"
107 + )
108 +
109 + def get_run_config(self) -> dict:
110 + return dict(self.fields)
111 +
112 + def get_config_file_path(self) -> str:
113 + return f"{self.base_path}/{_CONFIG_FILENAME}"
114 +
115 + @staticmethod
116 + def _compute_base_path(root_dir: str, fields: dict) -> str:
117 + """
118 + Layout: <root_dir>/<safe_model_name>_<sha256-prefix>/.
119 + safe_model_name replaces '/' with '_' so HuggingFace IDs don't nest.
120 + """
121 + canonical = json.dumps(fields, sort_keys=True, separators=(",", ":"))
122 + digest = hashlib.sha256(canonical.encode("utf-8")).hexdigest()[
123 + :_BASE_PATH_HASH_LEN
124 + ]
125 + safe_model_name = fields["model_name"].replace("/", "_")
126 + return f"{root_dir}/{safe_model_name}_{digest}"
vllm/v1/kv_offload/tiering/factory.py
+6
@@ -57,3 +57,9 @@ SecondaryTierFactory.register_tier(
57 "vllm.v1.kv_offload.tiering.example.manager",
58 "ExampleSecondaryTierManager",
59 )
60 +
61 +SecondaryTierFactory.register_tier(
62 + "fs_python",
63 + "vllm.v1.kv_offload.tiering.fs.manager",
64 + "FileSystemTierManager",
65 +)
vllm/v1/kv_offload/tiering/fs/__init__.py
vllm/v1/kv_offload/tiering/fs/io.py new
+101
@@ -0,0 +1,101 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import logging
5 +import os
6 +import random
7 +import threading
8 +
9 +logger = logging.getLogger(__name__)
10 +
11 +# O_DIRECT is Linux-specific and not available on macOS
12 +O_DIRECT = getattr(os, "O_DIRECT", 0)
13 +
14 +# Thread-local storage for unique temporary file suffixes
15 +_thread_local = threading.local()
16 +
17 +
18 +def _get_tmp_suffix() -> str:
19 + """Generate a thread-local unique suffix for temporary files."""
20 + try:
21 + return _thread_local.tmp_suffix
22 + except AttributeError:
23 + _thread_local.tmp_suffix = f"_{random.randint(0, 2**63 - 1)}.tmp"
24 + return _thread_local.tmp_suffix
25 +
26 +
27 +def _ensure_dirs(path: str) -> None:
28 + """Create parent directories of *path* if they don't exist."""
29 + os.makedirs(os.path.dirname(path), exist_ok=True)
30 +
31 +
32 +def store_block(
33 + dest_path: str,
34 + buffer: memoryview,
35 + offset: int,
36 + block_size: int,
37 +) -> None:
38 + """
39 + Store callback: Writes to a temp file then atomically replaces the destination.
40 + """
41 + # Check if block already exists to avoid redundant writes
42 + if os.path.exists(dest_path):
43 + return
44 +
45 + tmp_path = dest_path + _get_tmp_suffix()
46 + # Ensure parent directories exist
47 + _ensure_dirs(dest_path)
48 +
49 + # Write block atomically. Cast to a flat byte view so the slice uses byte
50 + # indices; the raw memoryview may be multi-dimensional with itemsize > 1.
51 + view_slice = buffer.cast("B")[offset : offset + block_size]
52 + try:
53 + fd = os.open(
54 + tmp_path,
55 + os.O_CREAT | os.O_EXCL | os.O_WRONLY | os.O_TRUNC | O_DIRECT,
56 + 0o644,
57 + )
58 + try:
59 + written = os.write(fd, view_slice)
60 + if written < len(view_slice):
61 + raise OSError(
62 + f"Short write: expected {len(view_slice)} bytes, wrote {written}"
63 + )
64 + finally:
65 + os.close(fd)
66 + os.replace(tmp_path, dest_path)
67 + except Exception:
68 + try:
69 + os.remove(tmp_path)
70 + except OSError as cleanup_exc:
71 + logger.warning("Failed to remove temp file %s: %s", tmp_path, cleanup_exc)
72 + raise
73 +
74 +
75 +def load_block(
76 + source_path: str,
77 + view: memoryview,
78 + offset: int,
79 + block_size: int,
80 +) -> None:
81 + """
82 + Load callback: read one KV block from disk. Remove the file on failure.
83 + """
84 + fd: int | None = None
85 + view_slice = view.cast("B")[offset : offset + block_size]
86 + try:
87 + fd = os.open(source_path, os.O_RDONLY | O_DIRECT)
88 + bytes_read = os.readv(fd, [view_slice])
89 + if bytes_read < block_size:
90 + raise OSError(f"Short read: expected {block_size} bytes, read {bytes_read}")
91 + except Exception:
92 + try:
93 + os.remove(source_path)
94 + except OSError as cleanup_exc:
95 + logger.warning(
96 + "Failed to remove unreadable file %s: %s", source_path, cleanup_exc
97 + )
98 + raise
99 + finally:
100 + if fd is not None:
101 + os.close(fd)
vllm/v1/kv_offload/tiering/fs/manager.py new
+150
@@ -0,0 +1,150 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +FileSystemTierManager: Pure-Python file system secondary tier for KV cache offloading.
5 +
6 +Store path:
7 + Data is written to a temp file (<dest_path.tmp>) via os.write,
8 + then os.replace'd to the final path (without .tmp).
9 +
10 +Load path:
11 + Data is read from the block file directly via os.readv into the
12 + provided memoryview slice.
13 +
14 +File naming: <base_path>_r<rank>/<hhh>/<hh>_g<group_idx>/<hash_hex>.bin
15 + (hash-based subdirectories to limit directory fan-out)
16 +"""
17 +
18 +import functools
19 +import json
20 +import os
21 +from collections.abc import Iterable
22 +from typing import TYPE_CHECKING
23 +
24 +from vllm.logger import init_logger
25 +from vllm.v1.kv_offload.base import OffloadKey, ReqContext
26 +from vllm.v1.kv_offload.file_mapper import FileMapper
27 +from vllm.v1.kv_offload.tiering.base import (
28 + JobMetadata,
29 + JobResult,
30 + SecondaryTierManager,
31 +)
32 +from vllm.v1.kv_offload.tiering.fs.io import load_block, store_block
33 +from vllm.v1.kv_offload.tiering.fs.thread_pool import DualQueueThreadPool
34 +
35 +if TYPE_CHECKING:
36 + from vllm.v1.kv_offload.base import OffloadingSpec
37 +
38 +logger = init_logger(__name__)
39 +
40 +
41 +class FileSystemTierManager(SecondaryTierManager):
42 + """
43 + Pure-Python disk-backed secondary tier.
44 +
45 + Read-priority threads service load jobs preferentially; write-priority
46 + threads service store jobs preferentially. Both groups can drain either
47 + queue, so neither starves.
48 +
49 + submit_store / submit_load are non-blocking: they enqueue tasks and return.
50 + get_finished() polls job completion and returns completed JobResults.
51 +
52 + """
53 +
54 + def __init__(
55 + self,
56 + offloading_spec: "OffloadingSpec",
57 + primary_kv_view: memoryview,
58 + tier_type: str,
59 + root_dir: str,
60 + n_read_threads: int = 16,
61 + n_write_threads: int = 16,
62 + ):
63 + """
64 + Args:
65 + offloading_spec: contains the vllm_config, kv_cache_config
66 + and block_size_factor.
67 + primary_kv_view: Memoryview of the primary tier's CPU KV cache.
68 + tier_type: Tier type identifier, set by SecondaryTierFactory.
69 + root_dir: Root directory for block files.
70 + n_read_threads: Number of read-priority I/O threads.
71 + n_write_threads: Number of write-priority I/O threads.
72 + """
73 + super().__init__(offloading_spec, primary_kv_view, tier_type)
74 +
75 + # Extract block size from primary view
76 + assert primary_kv_view.strides is not None, (
77 + "primary_kv_view.strides cannot be None"
78 + )
79 + self._block_size: int = primary_kv_view.strides[0]
80 +
81 + # Create file mapper
82 + self.file_mapper = FileMapper.from_offloading_spec(
83 + root_dir=root_dir,
84 + offloading_spec=offloading_spec,
85 + gpu_blocks_per_file=offloading_spec.block_size_factor,
86 + )
87 +
88 + # Write config file
89 + config_path = self.file_mapper.get_config_file_path()
90 + os.makedirs(os.path.dirname(config_path), exist_ok=True)
91 + if not os.path.exists(config_path):
92 + with open(config_path, "w") as f:
93 + json.dump(
94 + self.file_mapper.get_run_config(), f, indent=2, sort_keys=True
95 + )
96 +
97 + self._pool = DualQueueThreadPool(
98 + n_read_threads,
99 + n_write_threads,
100 + thread_name_prefix="vllm_kv_py_fs",
101 + )
102 +
103 + def lookup(
104 + self, key: OffloadKey, req_context: ReqContext | None = None
105 + ) -> bool | None:
106 + return os.path.exists(self.file_mapper.get_file_name(key))
107 +
108 + def submit_store(self, job_metadata: JobMetadata) -> None:
109 + tasks = (
110 + functools.partial(
111 + store_block,
112 + self.file_mapper.get_file_name(key),
113 + self._primary_kv_view,
114 + int(bid) * self._block_size,
115 + self._block_size,
116 + )
117 + for key, bid in zip(job_metadata.keys, job_metadata.block_ids)
118 + )
119 + self._pool.enqueue_store(job_metadata.job_id, len(job_metadata.keys), tasks)
120 +
121 + def submit_load(self, job_metadata: JobMetadata) -> None:
122 + tasks = (
123 + functools.partial(
124 + load_block,
125 + self.file_mapper.get_file_name(key),
126 + self._primary_kv_view,
127 + int(bid) * self._block_size,
128 + self._block_size,
129 + )
130 + for key, bid in zip(job_metadata.keys, job_metadata.block_ids)
131 + )
132 + self._pool.enqueue_load(job_metadata.job_id, len(job_metadata.keys), tasks)
133 +
134 + def get_finished(self) -> Iterable[JobResult]:
135 + """
136 + Collect completed jobs from the finished-jobs queue.
137 + """
138 + return (
139 + JobResult(job_id=job_id, success=success)
140 + for job_id, success in self._pool.get_finished()
141 + )
142 +
143 + def shutdown(self) -> None:
144 + """
145 + Release resources held by this tier.
146 +
147 + Shuts down the thread pool, clearing pending tasks and waiting for
148 + active threads to complete.
149 + """
150 + self._pool.shutdown(wait=True)
vllm/v1/kv_offload/tiering/fs/thread_pool.py new
+158
@@ -0,0 +1,158 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Thread pool:
5 + Two queues (load, store) and two sets of threads:
6 + - Load-priority threads: drain the load queue first, then the store queue.
7 + - Store-priority threads: drain the store queue first, then the load queue.
8 + Load jobs are enqueued to the load queue; store jobs to the store queue.
9 +"""
10 +
11 +import threading
12 +from collections import deque
13 +from collections.abc import Callable, Iterable
14 +
15 +from vllm.logger import init_logger
16 +from vllm.v1.kv_offload.tiering.base import JobId
17 +
18 +logger = init_logger(__name__)
19 +
20 +
21 +class JobState:
22 + """
23 + Thread-safe completion tracker for a set of per-block I/O tasks.
24 +
25 + Each task calls task_done(success) when it finishes.
26 + """
27 +
28 + __slots__ = ("_job_id", "_n_tasks", "_completed", "_success", "_lock")
29 +
30 + def __init__(self, job_id: JobId, n_tasks: int) -> None:
31 + self._job_id: JobId = job_id
32 + self._n_tasks = n_tasks
33 + self._completed = 0
34 + self._success = True
35 + self._lock = threading.Lock()
36 +
37 + @property
38 + def job_id(self) -> JobId:
39 + return self._job_id
40 +
41 + def task_done(self, success: bool) -> tuple[bool, bool]:
42 + """Returns if job completed and success flag"""
43 + with self._lock:
44 + self._completed += 1
45 + if not success:
46 + self._success = False
47 + return self._completed == self._n_tasks, self._success
48 +
49 +
50 +class DualQueueThreadPool:
51 + """
52 + Thread pool with two task queues (load and store) and two thread groups.
53 +
54 + Load-priority threads drain the load queue first, then fall back to the
55 + store queue. Store-priority threads do the reverse. Both queues share
56 + a single condition variable.
57 + """
58 +
59 + def __init__(
60 + self,
61 + n_read_threads: int,
62 + n_write_threads: int,
63 + thread_name_prefix: str = "fs_secondary_tier",
64 + ) -> None:
65 + self._load_q: deque = deque()
66 + self._store_q: deque = deque()
67 + self._condition = threading.Condition(threading.Lock())
68 + self._stop = False
69 + self._threads: list[threading.Thread] = []
70 + self._finished_q: deque[tuple[JobId, bool]] = deque()
71 +
72 + for i in range(n_read_threads):
73 + t = threading.Thread(
74 + target=self._worker,
75 + args=(True,),
76 + name=f"{thread_name_prefix}_l{i}",
77 + daemon=True,
78 + )
79 + t.start()
80 + self._threads.append(t)
81 +
82 + for i in range(n_write_threads):
83 + t = threading.Thread(
84 + target=self._worker,
85 + args=(False,),
86 + name=f"{thread_name_prefix}_s{i}",
87 + daemon=True,
88 + )
89 + t.start()
90 + self._threads.append(t)
91 +
92 + def enqueue_load(
93 + self,
94 + job_id: JobId,
95 + n_tasks: int,
96 + tasks: Iterable[Callable],
97 + ) -> None:
98 + """Enqueue load tasks for a job (high-priority for load-priority threads)."""
99 + state = JobState(job_id, n_tasks)
100 + with self._condition:
101 + for fn in tasks:
102 + self._load_q.append((fn, state))
103 + self._condition.notify(n_tasks)
104 +
105 + def enqueue_store(
106 + self,
107 + job_id: JobId,
108 + n_tasks: int,
109 + tasks: Iterable[Callable],
110 + ) -> None:
111 + """Enqueue store tasks for a job (high-priority for store-priority threads)."""
112 + state = JobState(job_id, n_tasks)
113 + with self._condition:
114 + for fn in tasks:
115 + self._store_q.append((fn, state))
116 + self._condition.notify(n_tasks)
117 +
118 + def get_finished(self) -> list[tuple[JobId, bool]]:
119 + jobs = []
120 + while self._finished_q:
121 + jobs.append(self._finished_q.popleft())
122 + return jobs
123 +
124 + def shutdown(self, wait: bool = True) -> None:
125 + with self._condition:
126 + self._stop = True
127 + self._load_q.clear()
128 + self._store_q.clear()
129 + self._condition.notify_all()
130 + if wait:
131 + for t in self._threads:
132 + t.join()
133 +
134 + def _worker(self, load_priority: bool) -> None:
135 + # Wait for tasks, process from primary queue first, fall back to secondary.
136 + while True:
137 + with self._condition:
138 + self._condition.wait_for(
139 + lambda: self._stop or self._load_q or self._store_q
140 + )
141 + if self._stop:
142 + return
143 + primary = self._load_q if load_priority else self._store_q
144 + secondary = self._store_q if load_priority else self._load_q
145 + task, state = primary.popleft() if primary else secondary.popleft()
146 + try:
147 + task()
148 + job_finished, success = state.task_done(True)
149 + except Exception as exc:
150 + logger.error(
151 + "FileSystemTierManagerPython: job %s block I/O failed: %s",
152 + state.job_id,
153 + exc,
154 + )
155 + job_finished, success = state.task_done(False)
156 +
157 + if job_finished:
158 + self._finished_q.append((state.job_id, success))