[FS-Offloading] Batch Lookup in C (#46713)
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
Jun 29, 2026 at 12:28 UTC
c8fb2963bd1baebbdd28062097096b59b2ba3189
5 files changed
+162
-2
CMakeLists.txt
+15
@@ -140,6 +140,21 @@ if(Python_VERSION VERSION_GREATER_EQUAL "3.11")
140
WITH_SOABI)
141
endif()
142
143
+#
144
+# fs_io extension (pure CXX; must stay above the non-CUDA device branch
145
+# so CPU builds define the target before the early return).
146
+# GIL-releasing filesystem helpers for FileSystemTierManager.
147
+#
148
+if(Python_VERSION VERSION_GREATER_EQUAL "3.11")
149
+ define_extension_target(
150
+ fs_io_C
151
+ DESTINATION vllm
152
+ LANGUAGE CXX
153
+ SOURCES csrc/fs_io.cpp
154
+ USE_SABI 3.11
155
+ WITH_SOABI)
156
+endif()
157
+
158
#
159
# Forward the non-CUDA device extensions to external CMake scripts.
160
#
csrc/fs_io.cpp
new
+69
@@ -0,0 +1,69 @@
1
+// SPDX-License-Identifier: Apache-2.0
2
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+#include <Python.h>
5
+
6
+#include <unistd.h>
7
+
8
+#include <vector>
9
+
10
+extern "C" {
11
+
12
+static void _batch_lookup(const std::vector<const char*>& paths,
13
+ std::vector<int>& exists_flags) {
14
+ for (size_t i = 0; i < paths.size(); i++) {
15
+ exists_flags[i] = (access(paths[i], F_OK) == 0) ? 1 : 0;
16
+ }
17
+}
18
+
19
+/// @brief Check file existence for a batch of paths.
20
+/// @param paths list[str] – absolute paths to check.
21
+/// @return list[bool] – True if the corresponding path exists, False otherwise.
22
+/// @note Releases the GIL for the entire batch. File existence via access(2).
23
+static PyObject* batch_lookup(PyObject* /*self*/, PyObject* args) {
24
+ PyObject* path_list;
25
+ if (!PyArg_ParseTuple(args, "O!", &PyList_Type, &path_list)) {
26
+ return nullptr;
27
+ }
28
+
29
+ const Py_ssize_t n = PyList_Size(path_list);
30
+ std::vector<const char*> paths(n);
31
+ for (Py_ssize_t i = 0; i < n; i++) {
32
+ paths[i] = PyUnicode_AsUTF8AndSize(PyList_GetItem(path_list, i), nullptr);
33
+ if (paths[i] == nullptr) {
34
+ return nullptr;
35
+ }
36
+ }
37
+
38
+ std::vector<int> exists_flags(n);
39
+ {
40
+ Py_BEGIN_ALLOW_THREADS _batch_lookup(paths, exists_flags);
41
+ Py_END_ALLOW_THREADS
42
+ }
43
+
44
+ PyObject* result = PyList_New(n);
45
+ if (result == nullptr) {
46
+ return nullptr;
47
+ }
48
+ for (Py_ssize_t i = 0; i < n; i++) {
49
+ PyList_SetItem(result, i, PyBool_FromLong(exists_flags[i]));
50
+ }
51
+ return result;
52
+}
53
+
54
+static PyMethodDef fs_io_C_methods[] = {
55
+ {"batch_lookup", batch_lookup, METH_VARARGS,
56
+ "batch_lookup(paths: list[str]) -> list[bool]\n"
57
+ "\n"
58
+ "Check file existence for a batch of paths."},
59
+ {nullptr, nullptr, 0, nullptr},
60
+};
61
+
62
+static struct PyModuleDef fs_io_C_module = {
63
+ PyModuleDef_HEAD_INIT, "fs_io_C", "Filesystem helpers for KV offload", -1,
64
+ fs_io_C_methods,
65
+};
66
+
67
+PyMODINIT_FUNC PyInit_fs_io_C(void) { return PyModule_Create(&fs_io_C_module); }
68
+
69
+} // extern "C"
setup.py
+2
@@ -777,6 +777,7 @@ class precompiled_wheel_utils:
777
"vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so",
778
"vllm/cumem_allocator.abi3.so",
779
"vllm/spinloop.abi3.so",
780
+ "vllm/fs_io_C.abi3.so",
781
# ROCm-specific libraries
782
"vllm/_rocm_C.abi3.so",
783
}
@@ -1104,6 +1105,7 @@ if _is_cuda() or _is_hip():
1105
1106
if sys.version_info >= (3, 11):
1107
ext_modules.append(CMakeExtension(name="vllm.spinloop"))
1108
+ ext_modules.append(CMakeExtension(name="vllm.fs_io_C"))
1109
1110
if _is_hip():
1111
ext_modules.append(CMakeExtension(name="vllm._rocm_C"))
tests/v1/kv_offload/tiering/test_fs_tier.py
+64
-1
@@ -103,7 +103,7 @@ def lookup_and_wait(
103
keys: list[OffloadKey],
104
ctx: ReqContext = _CTX,
105
timeout: float = 1.0,
106
-) -> list[bool]:
106
+) -> list[LookupResult]:
107
"""Perform a full async lookup cycle and return resolved results."""
108
for k in keys:
109
tier.lookup(k, ctx)
@@ -332,3 +332,66 @@ def test_wait_idle_blocks_until_tasks_complete():
332
gate.set()
333
pool.shutdown(wait=True)
334
waiter.join(timeout=5.0)
335
+
336
+
337
+def test_batch_lookup_c_extension(tmp_path):
338
+ """Validates batch_lookup_C: empty, single, all-existing, all-missing,
339
+ mixed ordering, and input type validation."""
340
+ try:
341
+ from vllm.fs_io_C import batch_lookup as batch_lookup_C
342
+ except ImportError:
343
+ pytest.skip("fs_io_C extension not built")
344
+
345
+ # Setup
346
+ all_exist = [str(tmp_path / f"e{i}.bin") for i in range(3)]
347
+ for p in all_exist:
348
+ open(p, "w").close()
349
+ all_missing = [str(tmp_path / f"m{i}.bin") for i in range(3)]
350
+
351
+ # Empty list
352
+ assert batch_lookup_C([]) == []
353
+
354
+ # Single existing / missing
355
+ assert batch_lookup_C([all_exist[0]]) == [True]
356
+ assert batch_lookup_C([all_missing[0]]) == [False]
357
+
358
+ # All existing / all missing
359
+ assert batch_lookup_C(all_exist) == [True, True, True]
360
+ assert batch_lookup_C(all_missing) == [False, False, False]
361
+
362
+ # Mixed — verifies index ordering is preserved
363
+ paths = [val for pair in zip(all_exist, all_missing) for val in pair]
364
+ assert batch_lookup_C(paths) == [True, False, True, False, True, False]
365
+
366
+ # Input validation: non-list argument
367
+ with pytest.raises(TypeError):
368
+ batch_lookup_C(("/tmp/foo",))
369
+ with pytest.raises(TypeError):
370
+ batch_lookup_C(None)
371
+
372
+ # Input validation: non-str elements in list
373
+ with pytest.raises(TypeError):
374
+ batch_lookup_C([None])
375
+ with pytest.raises(TypeError):
376
+ batch_lookup_C([b"/tmp/foo"])
377
+ with pytest.raises(TypeError):
378
+ batch_lookup_C([42])
379
+ with pytest.raises(TypeError):
380
+ batch_lookup_C([all_exist[0], None]) # valid first, invalid mid-list
381
+
382
+
383
+@pytest.mark.parametrize("use_c_ext", [True, False])
384
+def test_batch_lookup_dispatch(fs_tier, monkeypatch, use_c_ext):
385
+ import vllm.v1.kv_offload.tiering.fs.manager as mgr_mod
386
+
387
+ if use_c_ext and not mgr_mod._HAS_BATCH_LOOKUP_C:
388
+ pytest.skip("fs_io_C extension not built")
389
+
390
+ monkeypatch.setattr(mgr_mod, "_HAS_BATCH_LOOKUP_C", use_c_ext)
391
+
392
+ tier, _ = fs_tier
393
+ tier.submit_store(make_job(1, [key(1)], [0]))
394
+ assert all(r.success for r in drain(tier))
395
+
396
+ results = lookup_and_wait(tier, [key(1), key(2)])
397
+ assert results == [LookupResult.HIT, LookupResult.MISS]
vllm/v1/kv_offload/tiering/fs/manager.py
+12
-1
@@ -21,6 +21,13 @@ import os
21
from collections.abc import Iterable
22
from typing import TYPE_CHECKING
23
24
+try:
25
+ from vllm.fs_io_C import batch_lookup as batch_lookup_C
26
+
27
+ _HAS_BATCH_LOOKUP_C = True
28
+except ImportError:
29
+ _HAS_BATCH_LOOKUP_C = False
30
+
31
from typing_extensions import override
32
33
from vllm.logger import init_logger
@@ -56,7 +63,11 @@ class FsAsyncLookupManager(AsyncLookupManager):
63
def batch_lookup(
64
self, keys: list[OffloadKey], req_context: ReqContext
65
) -> Iterable[bool]:
59
- return (os.path.exists(self._tier.file_mapper.get_file_name(k)) for k in keys)
66
+ paths = [self._tier.file_mapper.get_file_name(k) for k in keys]
67
+ if _HAS_BATCH_LOOKUP_C:
68
+ # C extension: GIL released for the entire faccessat() batch.
69
+ return batch_lookup_C(paths)
70
+ return (os.path.exists(p) for p in paths)
71
72
73
class FileSystemTierManager(SecondaryTierManager):