@tduncan / Nvidia-Optimization-Stack-VLLM / commits / d18ed2304a

[KV-offload][FS] : Batch store/load_block in C (#49152)

Signed-off-by: <> Co-authored-by: Varun Sundar Rabindranath <varun-sundar-rabindranath@h100-01.nemg-001.lab.rdu2.dc.redhat.com>

Varun Sundar Rabindranath committed Jul 27, 2026 at 23:21 UTC d18ed2304a2703e3211fc384a58607e754f5b723
4 files changed +452 -58
csrc/fs_io.cpp
+283 -1
@@ -3,19 +3,155 @@
3
4 #include <Python.h>
5
6 +#include <errno.h>
7 +#include <fcntl.h>
8 #include <unistd.h>
9
10 +#include <filesystem>
11 +#include <string>
12 #include <vector>
13
14 +#if defined(O_DIRECT)
15 +constexpr int kODirectFlag = O_DIRECT;
16 +#else
17 +constexpr int kODirectFlag = 0;
18 +#endif
19 +
20 extern "C" {
21
12 -static void _batch_lookup(const std::vector<const char*>& paths,
22 +namespace {
23 +
24 +// Returns 0 on success, or the std::error_code's POSIX-compatible value on
25 +// failure, mirroring the errno convention used by the syscalls below.
26 +inline int ensure_parent_dirs(const std::string& path) {
27 + const auto parent = std::filesystem::path(path).parent_path();
28 + if (parent.empty()) {
29 + return 0;
30 + }
31 + std::error_code ec;
32 + std::filesystem::create_directories(parent, ec);
33 + return ec ? ec.value() : 0;
34 +}
35 +
36 +// Core single-block store: src/size are raw pointer + byte count. Returns 0
37 +// on success, or the errno of the failing step on failure -- captured
38 +// before any subsequent cleanup call can overwrite it. On failure, the temp
39 +// file is removed.
40 +inline int _store_block(const char* tmp_path, const char* dest_path,
41 + const char* src, size_t size, bool use_o_direct) {
42 + if (access(dest_path, F_OK) == 0) {
43 + return 0; // Already present.
44 + }
45 +
46 + if (const int err = ensure_parent_dirs(dest_path); err != 0) {
47 + return err;
48 + }
49 +
50 + const int o_direct_flag = use_o_direct ? kODirectFlag : 0;
51 + const int fd = open(
52 + tmp_path, O_CREAT | O_EXCL | O_WRONLY | O_TRUNC | o_direct_flag, 0644);
53 + if (fd < 0) {
54 + return errno;
55 + }
56 +
57 + const ssize_t written = write(fd, src, size);
58 + if (written < 0 || static_cast<size_t>(written) != size) {
59 + const int err = written < 0 ? errno : EIO;
60 + close(fd); // Best-effort cleanup; the real error is already captured.
61 + unlink(tmp_path);
62 + return err;
63 + }
64 +
65 + if (close(fd) != 0) {
66 + const int err = errno;
67 + unlink(tmp_path);
68 + return err;
69 + }
70 +
71 + if (rename(tmp_path, dest_path) != 0) {
72 + const int err = errno;
73 + unlink(tmp_path);
74 + return err;
75 + }
76 +
77 + return 0;
78 +}
79 +
80 +// Core single-block load: dst/size are raw pointer + byte count. Returns 0
81 +// on success, or the errno of the failing step on failure. On failure,
82 +// the source file is removed since a partially-read block should not be reused.
83 +inline int _load_block(const char* source_path, char* dst, size_t size,
84 + bool use_o_direct) {
85 + const int o_direct_flag = use_o_direct ? kODirectFlag : 0;
86 + const int fd = open(source_path, O_RDONLY | o_direct_flag, 0);
87 + if (fd < 0) {
88 + const int err = errno;
89 + unlink(source_path);
90 + return err;
91 + }
92 +
93 + const ssize_t bytes_read = read(fd, dst, size);
94 + if (bytes_read < 0 || static_cast<size_t>(bytes_read) != size) {
95 + const int err = bytes_read < 0 ? errno : EIO;
96 + close(fd);
97 + unlink(source_path);
98 + return err;
99 + }
100 +
101 + if (close(fd) != 0) {
102 + const int err = errno;
103 + unlink(source_path);
104 + return err;
105 + }
106 +
107 + return 0;
108 +}
109 +
110 +inline void _batch_lookup(const std::vector<const char*>& paths,
111 std::vector<int>& exists_flags) {
112 for (size_t i = 0; i < paths.size(); i++) {
113 exists_flags[i] = (access(paths[i], F_OK) == 0) ? 1 : 0;
114 }
115 }
116
117 +// Helper: extract a list[str] of length n into a vector<const char*>.
118 +// Returns false and sets a Python exception on error.
119 +inline bool extract_str_list(PyObject* list, Py_ssize_t n,
120 + std::vector<const char*>& out) {
121 + for (Py_ssize_t i = 0; i < n; i++) {
122 + out[i] = PyUnicode_AsUTF8AndSize(PyList_GetItem(list, i), nullptr);
123 + if (out[i] == nullptr) {
124 + return false;
125 + }
126 + }
127 + return true;
128 +}
129 +
130 +// Helper: extract a Py_buffer per element of a list[bytes-like] of length n.
131 +// On success, `out` holds n acquired buffers (caller must PyBuffer_Release
132 +// each). On failure, any buffers already acquired are released before
133 +// returning false, and a Python exception is set.
134 +inline bool extract_buffer_list(PyObject* list, Py_ssize_t n, int flags,
135 + std::vector<Py_buffer>& out) {
136 + for (Py_ssize_t i = 0; i < n; i++) {
137 + if (PyObject_GetBuffer(PyList_GetItem(list, i), &out[i], flags) != 0) {
138 + for (Py_ssize_t j = 0; j < i; j++) {
139 + PyBuffer_Release(&out[j]);
140 + }
141 + return false;
142 + }
143 + }
144 + return true;
145 +}
146 +
147 +inline void release_buffer_list(std::vector<Py_buffer>& buffers) {
148 + for (auto& buf : buffers) {
149 + PyBuffer_Release(&buf);
150 + }
151 +}
152 +
153 +} // namespace
154 +
155 /// @brief Check file existence for a batch of paths.
156 /// @param paths list[str] – absolute paths to check.
157 /// @return list[bool] – True if the corresponding path exists, False otherwise.
@@ -51,11 +187,157 @@ static PyObject* batch_lookup(PyObject* /*self*/, PyObject* args) {
187 return result;
188 }
189
190 +/// @brief Store a batch of blocks, each from its own buffer, to disk.
191 +/// @param tmp_paths list[str] – one temp path per block.
192 +/// @param dest_paths list[str] – one destination path per block.
193 +/// @param buffers list[bytes-like] – one source buffer per block.
194 +/// @param use_o_direct bool – whether to open files with O_DIRECT
195 +/// (default True). Ignored where O_DIRECT is unsupported
196 +/// by the platform.
197 +/// @note Releases the GIL for the entire batch. Raises on first error.
198 +static PyObject* batch_store_block(PyObject* /*self*/, PyObject* args) {
199 + PyObject* tmp_paths_obj = nullptr;
200 + PyObject* dest_paths_obj = nullptr;
201 + PyObject* buffers_obj = nullptr;
202 + int use_o_direct = 1;
203 +
204 + if (!PyArg_ParseTuple(args, "O!O!O!|p", &PyList_Type, &tmp_paths_obj,
205 + &PyList_Type, &dest_paths_obj, &PyList_Type,
206 + &buffers_obj, &use_o_direct)) {
207 + return nullptr;
208 + }
209 +
210 + const Py_ssize_t n = PyList_Size(tmp_paths_obj);
211 + if (PyList_Size(dest_paths_obj) != n || PyList_Size(buffers_obj) != n) {
212 + PyErr_SetString(
213 + PyExc_ValueError,
214 + "tmp_paths, dest_paths and buffers must have the same length");
215 + return nullptr;
216 + }
217 +
218 + std::vector<const char*> tmp_paths(n);
219 + std::vector<const char*> dest_paths(n);
220 +
221 + if (!extract_str_list(tmp_paths_obj, n, tmp_paths)) return nullptr;
222 + if (!extract_str_list(dest_paths_obj, n, dest_paths)) return nullptr;
223 +
224 + std::vector<Py_buffer> buffers(n);
225 + if (!extract_buffer_list(buffers_obj, n, PyBUF_SIMPLE, buffers)) {
226 + return nullptr;
227 + }
228 +
229 + Py_ssize_t failed_index = -1;
230 + int failure_errno = 0;
231 +
232 + {
233 + Py_BEGIN_ALLOW_THREADS for (Py_ssize_t i = 0; i < n; i++) {
234 + const char* buf = static_cast<const char*>(buffers[i].buf);
235 + const int err =
236 + _store_block(tmp_paths[i], dest_paths[i], buf,
237 + static_cast<size_t>(buffers[i].len), use_o_direct);
238 + if (err != 0) {
239 + failed_index = i;
240 + failure_errno = err;
241 + break;
242 + }
243 + }
244 + Py_END_ALLOW_THREADS
245 + }
246 +
247 + release_buffer_list(buffers);
248 +
249 + if (failed_index >= 0) {
250 + // PyErr_SetFromErrnoWithFilename() reads the errno to format exception.
251 + errno = failure_errno;
252 + return PyErr_SetFromErrnoWithFilename(PyExc_OSError,
253 + dest_paths[failed_index]);
254 + }
255 +
256 + Py_RETURN_NONE;
257 +}
258 +
259 +/// @brief Load a batch of blocks from disk, each into its own buffer.
260 +/// @param source_paths list[str] – one source path per block.
261 +/// @param buffers list[writable bytes-like] – one destination buffer
262 +/// per block.
263 +/// @param use_o_direct bool – whether to open files with O_DIRECT
264 +/// (default True). Ignored where O_DIRECT is unsupported
265 +/// by the platform.
266 +/// @note Releases the GIL for the entire batch. Raises on first error.
267 +static PyObject* batch_load_block(PyObject* /*self*/, PyObject* args) {
268 + PyObject* source_paths_obj = nullptr;
269 + PyObject* buffers_obj = nullptr;
270 + int use_o_direct = 1;
271 +
272 + if (!PyArg_ParseTuple(args, "O!O!|p", &PyList_Type, &source_paths_obj,
273 + &PyList_Type, &buffers_obj, &use_o_direct)) {
274 + return nullptr;
275 + }
276 +
277 + const Py_ssize_t n = PyList_Size(source_paths_obj);
278 + if (PyList_Size(buffers_obj) != n) {
279 + PyErr_SetString(PyExc_ValueError,
280 + "source_paths and buffers must have the same length");
281 + return nullptr;
282 + }
283 +
284 + std::vector<const char*> source_paths(n);
285 + if (!extract_str_list(source_paths_obj, n, source_paths)) return nullptr;
286 +
287 + std::vector<Py_buffer> buffers(n);
288 + if (!extract_buffer_list(buffers_obj, n, PyBUF_WRITABLE, buffers)) {
289 + return nullptr;
290 + }
291 +
292 + Py_ssize_t failed_index = -1;
293 + int failure_errno = 0;
294 +
295 + {
296 + Py_BEGIN_ALLOW_THREADS for (Py_ssize_t i = 0; i < n; i++) {
297 + char* buf = static_cast<char*>(buffers[i].buf);
298 + const int err =
299 + _load_block(source_paths[i], buf, static_cast<size_t>(buffers[i].len),
300 + use_o_direct);
301 + if (err != 0) {
302 + failed_index = i;
303 + failure_errno = err;
304 + break;
305 + }
306 + }
307 + Py_END_ALLOW_THREADS
308 + }
309 +
310 + release_buffer_list(buffers);
311 +
312 + if (failed_index >= 0) {
313 + // PyErr_SetFromErrnoWithFilename() reads the errno to format exception.
314 + errno = failure_errno;
315 + return PyErr_SetFromErrnoWithFilename(PyExc_OSError,
316 + source_paths[failed_index]);
317 + }
318 +
319 + Py_RETURN_NONE;
320 +}
321 +
322 static PyMethodDef fs_io_C_methods[] = {
323 {"batch_lookup", batch_lookup, METH_VARARGS,
324 "batch_lookup(paths: list[str]) -> list[bool]\n"
325 "\n"
326 "Check file existence for a batch of paths."},
327 + {"batch_store_block", batch_store_block, METH_VARARGS,
328 + "batch_store_block(tmp_paths: list[str], dest_paths: list[str],\n"
329 + " buffers: list[bytes-like],\n"
330 + " use_o_direct: bool = True) -> None\n"
331 + "\n"
332 + "Store a batch of blocks, each from its own buffer, to disk. Raises on "
333 + "first error."},
334 + {"batch_load_block", batch_load_block, METH_VARARGS,
335 + "batch_load_block(source_paths: list[str],\n"
336 + " buffers: list[writable bytes-like],\n"
337 + " use_o_direct: bool = True) -> None\n"
338 + "\n"
339 + "Load a batch of blocks from disk into corresponding buffers. "
340 + "Raises on first error."},
341 {nullptr, nullptr, 0, nullptr},
342 };
343
tests/v1/kv_offload/tiering/test_fs_tier.py
+67 -32
@@ -46,6 +46,7 @@ from vllm.v1.kv_offload.tiering.fs.thread_pool import DualQueueThreadPool
46 # Helpers
47 # ---------------------------------------------------------------------------
48
49 +_NUM_BLOCKS = 8
50 _BLOCK_ELEMENTS = 128 * mmap.PAGESIZE # 2MB per block for pagesize 4096.
51 _DTYPE: torch.dtype = torch.float32
52 _CTX = ReqContext(req_id="test")
@@ -162,7 +163,7 @@ def _page_aligned_rand_tensor(
163
164 @pytest.fixture
165 def fs_tier(tmp_path):
165 - tensor = _page_aligned_zero_tensor(4, _BLOCK_ELEMENTS)
166 + tensor = _page_aligned_zero_tensor(_NUM_BLOCKS, _BLOCK_ELEMENTS)
167 mock_view = memoryview(tensor.numpy())
168 tier = FileSystemTierManager(
169 offloading_spec=_MOCK_OFFLOADING_SPEC,
@@ -178,7 +179,7 @@ def fs_tier(tmp_path):
179
180 @pytest.fixture
181 def fs_tier_with_events(tmp_path):
181 - tensor = _page_aligned_zero_tensor(4, _BLOCK_ELEMENTS)
182 + tensor = _page_aligned_zero_tensor(_NUM_BLOCKS, _BLOCK_ELEMENTS)
183 mock_view = memoryview(tensor.numpy())
184 tier = FileSystemTierManager(
185 offloading_spec=_make_offloading_spec(enable_kv_cache_events=True),
@@ -347,33 +348,43 @@ def test_shutdown_discards_pending_tasks(fs_tier):
348 assert all(not t.is_alive() for t in tier._pool._threads)
349
350
350 -def test_store_load_data_integrity(fs_tier):
351 - """Data written by store must be exactly recovered by load."""
351 +@pytest.mark.parametrize("batch_size", [0, 1, 2, 5])
352 +@pytest.mark.parametrize("use_c_ext", [True, False])
353 +def test_store_load_data_integrity(fs_tier, monkeypatch, use_c_ext, batch_size):
354 + """Data written by store must be exactly recovered by load, for batches
355 + of any size -- including the empty batch."""
356 + import vllm.v1.kv_offload.tiering.fs.io as io_mod
357 +
358 + if use_c_ext and not io_mod._HAS_FSIO_C:
359 + pytest.skip("fs_io_C extension not built")
360 + monkeypatch.setattr(io_mod, "_HAS_FSIO_C", use_c_ext)
361 +
362 tier, tensor = fs_tier
363 # Populate tensor with random data
354 - tensor[:] = _page_aligned_rand_tensor(4, _BLOCK_ELEMENTS)
364 + tensor[:] = _page_aligned_rand_tensor(_NUM_BLOCKS, _BLOCK_ELEMENTS)
365
356 - # Store first 2 blocks
357 - num_store = 2
358 - expected = tensor[:num_store].clone()
366 + keys = [key(i) for i in range(batch_size)]
367 + store_block_ids = list(range(batch_size))
368 + load_block_ids = list(range(_NUM_BLOCKS - batch_size, _NUM_BLOCKS))
369 + expected = tensor[:batch_size].clone()
370
360 - store_ids = list(range(num_store))
361 - keys = [key(i) for i in range(num_store)]
362 -
363 - tier.submit_store(make_job(1, keys, store_ids))
364 - results = drain(tier)
365 - assert all(r.success for r in results)
371 + tier.submit_store(make_job(1, keys, store_block_ids))
372 + store_results = drain(tier)
373 + assert len(store_results) == 1
374 + assert store_results[0].success
375 + assert all(os.path.exists(tier.file_mapper.get_file_name(k)) for k in keys)
376
367 - # Overwrite source blocks to prove data is read from disk
368 - tensor[:num_store] = 0.0
377 + # reset tensor to prove data is read from disk
378 + tensor[:] = 0.0
379
370 - # Load into last 2 blocks
371 - load_ids = [2, 3]
372 - tier.submit_load(make_job(2, keys, load_ids, is_promotion=True))
373 - results = drain(tier)
374 - assert all(r.success for r in results)
380 + # Load into a range disjoint by index from the store ids, to also
381 + # exercise loading a block into a different id than it was stored from.
382 + tier.submit_load(make_job(2, keys, load_block_ids, is_promotion=True))
383 + load_results = drain(tier)
384 + assert len(load_results) == 1
385 + assert load_results[0].success
386
376 - for i, bid in enumerate(load_ids):
387 + for i, bid in enumerate(load_block_ids):
388 assert torch.allclose(tensor[bid], expected[i]), (
389 f"Block {bid} data mismatch after store+load"
390 )
@@ -499,6 +510,30 @@ def test_batch_lookup_dispatch(fs_tier, monkeypatch, use_c_ext):
510 assert results == [LookupResult.HIT, LookupResult.MISS]
511
512
513 +@pytest.mark.parametrize("use_c_ext", [True, False])
514 +def test_out_of_bounds_block_id_smoke(fs_tier, monkeypatch, use_c_ext):
515 + """Smoke test: a block id beyond the primary tensor's block count must
516 + fail the job, for both the C extension and the Python fallback."""
517 + import vllm.v1.kv_offload.tiering.fs.io as io_mod
518 +
519 + if use_c_ext and not io_mod._HAS_FSIO_C:
520 + pytest.skip("fs_io_C extension not built")
521 + monkeypatch.setattr(io_mod, "_HAS_FSIO_C", use_c_ext)
522 +
523 + tier, tensor = fs_tier
524 + out_of_bounds_bid = tensor.shape[0] # one past the last valid block
525 +
526 + tier.submit_store(make_job(1, [key(1)], [out_of_bounds_bid]))
527 + store_results = drain(tier)
528 + assert len(store_results) == 1
529 + assert not store_results[0].success
530 +
531 + tier.submit_load(make_job(2, [key(1)], [out_of_bounds_bid], is_promotion=True))
532 + load_results = drain(tier)
533 + assert len(load_results) == 1
534 + assert not load_results[0].success
535 +
536 +
537 # ---------------------------------------------------------------------------
538 # KV events
539 # ---------------------------------------------------------------------------
@@ -571,14 +606,14 @@ def test_mixed_job_results_emit_event_only_for_successful_job(
606
607 tier = fs_tier_with_events
608 failing_path = tier.file_mapper.get_file_name(key(1))
574 - original_store_block = mgr_mod.store_block
609 + original_batch_store_block = mgr_mod.batch_store_block
610
576 - def flaky_store_block(dest_path, *args, **kwargs):
577 - if dest_path == failing_path:
611 + def flaky_batch_store_block(paths, *args, **kwargs):
612 + if failing_path in paths:
613 raise OSError("injected store failure")
579 - return original_store_block(dest_path, *args, **kwargs)
614 + return original_batch_store_block(paths, *args, **kwargs)
615
581 - monkeypatch.setattr(mgr_mod, "store_block", flaky_store_block)
616 + monkeypatch.setattr(mgr_mod, "batch_store_block", flaky_batch_store_block)
617
618 tier.submit_store(make_job(1, [key(1)], [0]))
619 tier.submit_store(make_job(2, [key(2)], [1]))
@@ -599,14 +634,14 @@ def test_partially_failed_store_emits_no_event(fs_tier_with_events, monkeypatch)
634
635 tier = fs_tier_with_events
636 failing_path = tier.file_mapper.get_file_name(key(2))
602 - original_store_block = mgr_mod.store_block
637 + original_batch_store_block = mgr_mod.batch_store_block
638
604 - def flaky_store_block(dest_path, *args, **kwargs):
605 - if dest_path == failing_path:
639 + def flaky_batch_store_block(paths, *args, **kwargs):
640 + if failing_path in paths:
641 raise OSError("injected store failure")
607 - return original_store_block(dest_path, *args, **kwargs)
642 + return original_batch_store_block(paths, *args, **kwargs)
643
609 - monkeypatch.setattr(mgr_mod, "store_block", flaky_store_block)
644 + monkeypatch.setattr(mgr_mod, "batch_store_block", flaky_batch_store_block)
645
646 tier.submit_store(make_job(1, [key(1), key(2)], [0, 1]))
647 results = drain(tier)
vllm/v1/kv_offload/tiering/fs/io.py
+80 -2
@@ -8,6 +8,18 @@ import os
8 import random
9 import threading
10
11 +try:
12 + from vllm.fs_io_C import ( # pyright: ignore[reportMissingImports]
13 + batch_load_block as batch_load_block_C,
14 + )
15 + from vllm.fs_io_C import (
16 + batch_store_block as batch_store_block_C,
17 + )
18 +
19 + _HAS_FSIO_C = True
20 +except ImportError:
21 + _HAS_FSIO_C = False
22 +
23 logger = logging.getLogger(__name__)
24
25 # O_DIRECT is Linux-specific and not available on macOS
@@ -59,7 +71,23 @@ def _ensure_dirs(path: str) -> None:
71 os.makedirs(os.path.dirname(path), exist_ok=True)
72
73
62 -def store_block(
74 +def _validate_offsets(view: memoryview, offsets: list[int], block_size: int) -> None:
75 + """Raise if any block would read/write past the bounds of `view`.
76 +
77 + Without this, an out-of-range offset silently clips to a shorter (or
78 + empty) slice instead of failing, since memoryview slicing follows
79 + Python's slice-clamping semantics rather than raising.
80 + """
81 + total_len = len(view.cast("B"))
82 + for offset in offsets:
83 + if offset < 0 or offset + block_size > total_len:
84 + raise ValueError(
85 + f"block offset {offset} (block_size {block_size}) is out of "
86 + f"bounds for a buffer of size {total_len}"
87 + )
88 +
89 +
90 +def _store_block(
91 dest_path: str,
92 buffer: memoryview,
93 offset: int,
@@ -104,7 +132,7 @@ def store_block(
132 raise
133
134
107 -def load_block(
135 +def _load_block(
136 source_path: str,
137 view: memoryview,
138 offset: int,
@@ -117,6 +145,7 @@ def load_block(
145 fd: int | None = None
146 view_slice = view.cast("B")[offset : offset + block_size]
147 o_direct = O_DIRECT if use_o_direct else 0
148 +
149 try:
150 fd = os.open(source_path, os.O_RDONLY | o_direct)
151 bytes_read = os.readv(fd, [view_slice])
@@ -133,3 +162,52 @@ def load_block(
162 finally:
163 if fd is not None:
164 os.close(fd)
165 +
166 +
167 +def batch_store_block(
168 + paths: list[str],
169 + view: memoryview,
170 + offsets: list[int],
171 + block_size: int,
172 + use_o_direct: bool = True,
173 +) -> None:
174 + """
175 + Store a batch of KV blocks from a shared buffer to disk in one call.
176 +
177 + Each block buffer[offsets[i] : offsets[i]+block_size] is written atomically
178 + to dest_paths[i] via a temp-file rename. Raises on first error.
179 + """
180 + _validate_offsets(view, offsets, block_size)
181 +
182 + if _HAS_FSIO_C:
183 + view_B = view.cast("B")
184 + view_slices = [view_B[x : x + block_size] for x in offsets]
185 + tmp_paths = [p + _get_tmp_suffix() for p in paths]
186 + return batch_store_block_C(tmp_paths, paths, view_slices, use_o_direct)
187 + else:
188 + for path, offset in zip(paths, offsets):
189 + _store_block(path, view, offset, block_size, use_o_direct)
190 +
191 +
192 +def batch_load_block(
193 + paths: list[str],
194 + view: memoryview,
195 + offsets: list[int],
196 + block_size: int,
197 + use_o_direct: bool = True,
198 +) -> None:
199 + """
200 + Load a batch of KV blocks from disk into a shared buffer in one call.
201 +
202 + Block i is read from source_paths[i] into view[offsets[i] : offsets[i]+block_size].
203 + Raises on first error and removes the offending file.
204 + """
205 + _validate_offsets(view, offsets, block_size)
206 +
207 + if _HAS_FSIO_C:
208 + view_B = view.cast("B")
209 + view_slices = [view_B[x : x + block_size] for x in offsets]
210 + return batch_load_block_C(paths, view_slices, use_o_direct)
211 + else:
212 + for path, offset in zip(paths, offsets):
213 + _load_block(path, view, offset, block_size, use_o_direct)
vllm/v1/kv_offload/tiering/fs/manager.py
+22 -23
@@ -49,7 +49,11 @@ from vllm.v1.kv_offload.tiering.base import (
49 ScheduleEndContext,
50 SecondaryTierManager,
51 )
52 -from vllm.v1.kv_offload.tiering.fs.io import load_block, probe_o_direct, store_block
52 +from vllm.v1.kv_offload.tiering.fs.io import (
53 + batch_load_block,
54 + batch_store_block,
55 + probe_o_direct,
56 +)
57 from vllm.v1.kv_offload.tiering.fs.thread_pool import DualQueueThreadPool
58
59 if TYPE_CHECKING:
@@ -203,33 +207,28 @@ class FileSystemTierManager(SecondaryTierManager):
207 def submit_store(self, job_metadata: JobMetadata) -> None:
208 if self.events is not None:
209 self._store_job_keys[job_metadata.job_id] = list(job_metadata.keys)
206 - tasks = (
207 - functools.partial(
208 - store_block,
209 - self.file_mapper.get_file_name(key),
210 - self._primary_kv_view,
211 - int(bid) * self._block_size,
212 - self._block_size,
213 - self._use_o_direct,
214 - )
215 - for key, bid in zip(job_metadata.keys, job_metadata.block_ids)
210 + task = functools.partial(
211 + batch_store_block,
212 + [self.file_mapper.get_file_name(key) for key in job_metadata.keys],
213 + self._primary_kv_view,
214 + [int(bid) * self._block_size for bid in job_metadata.block_ids],
215 + self._block_size,
216 + self._use_o_direct,
217 )
217 - self._pool.enqueue_store(job_metadata.job_id, len(job_metadata.keys), tasks)
218 + self._pool.enqueue_store(job_metadata.job_id, 1, [task])
219
220 @override
221 def submit_load(self, job_metadata: JobMetadata) -> None:
221 - tasks = (
222 - functools.partial(
223 - load_block,
224 - self.file_mapper.get_file_name(key),
225 - self._primary_kv_view,
226 - int(bid) * self._block_size,
227 - self._block_size,
228 - self._use_o_direct,
229 - )
230 - for key, bid in zip(job_metadata.keys, job_metadata.block_ids)
222 + task = functools.partial(
223 + batch_load_block,
224 + [self.file_mapper.get_file_name(key) for key in job_metadata.keys],
225 + self._primary_kv_view,
226 + [int(bid) * self._block_size for bid in job_metadata.block_ids],
227 + self._block_size,
228 + self._use_o_direct,
229 )
232 - self._pool.enqueue_load(job_metadata.job_id, len(job_metadata.keys), tasks)
230 +
231 + self._pool.enqueue_load(job_metadata.job_id, 1, [task])
232
233 @override
234 def get_finished_jobs(self) -> Iterable[JobResult]: