Prhub

#46713 [FS-Offloading] Batch Lookup in C

原始 PR 作者 varun-sundar-rabindranath 合并时间 2026-06-30 00:28 文件变更 5 提交数 3 评论 15 代码增减 +162 / -2

执行摘要

C 扩展批量文件存在性检查,释放 GIL 提升 KV offload 查找性能

TieredOffloading with FS tier on main yields poor performance. This is mostly due to lookup delays triggered by the FSAsyncLookupManager. We find that the thread backed FSAsyncLookup is severely impacted by the GIL.

该PR适合关注KV offload性能的开发者精读,尤其学习C扩展与Python fallback结合的设计模式、Py_BEGIN_ALLOW_THREADS的使用场景。构建系统团队也可参考CMakeLists中扩展注册的方式。

讨论亮点
  • 命名统一:orozery建议将模块名从fsio_C改为fs_io_C,作者接受并全量修改。
  • 纯C函数提取:orozery建议将核心循环提取为静态函数_batch_lookup,提高清晰度,作者采纳。
  • Docstring添加:orozery建议为batch_lookup添加文档注释,作者补充。
  • 返回类型权衡:orozery建议返回bytes代替list[bool]以节省开销,作者认为list[bool]更符合Python惯用法且性能差异可忽略,决定保留。
  • 注释泛化:orozery建议修正CMakeLists注释避免特指batch_lookup,作者修改为更通用的描述。

实现拆解

  1. 新增C扩展模块:创建csrc/fs_io.cpp,实现batch_lookup函数,内部调用_batch_lookup静态函数使用access系统调用批量检查文件存在,并通过Py_BEGIN_ALLOW_THREADS释放GIL。
  2. 修改FsAsyncLookupManager:在vllm/v1/kv_offload/tiering/fs/manager.py中添加条件导入vllm.fs_io_C,修改batch_lookup方法:先批量收集路径,若C扩展可用则调用batch_lookup_C,否则回退到逐路径os.path.exists生成器。
  3. 集成构建系统:在setup.py中注册fs_io_C扩展(仅Python >=3.11),在CMakeLists.txt中定义扩展目标(define_extension_target),确保在非CUDA设备分支之前构建。
  4. 添加测试覆盖:在test_fs_tier.py中新增test_batch_lookup_c_extension测试C扩展的正确性与输入校验,以及参数化测试test_batch_lookup_dispatch验证C扩展可用/不可用时行为一致。
  5. 预编译扩展打包:在setup.pyexact_members集合中加入vllm/fs_io_C.abi3.so,确保预编译wheel正确包含此扩展。
文件 模块 状态 重要度
csrc/fs_io.cpp C 扩展 added 7.54
vllm/v1/kv_offload/tiering/fs/manager.py 管理器 modified 6.16
tests/v1/kv_offload/tiering/test_fs_tier.py FS 层测试 modified 6.12
setup.py 构建配置 modified 3.92
CMakeLists.txt 构建脚本 modified 2.67

关键符号

batch_lookup _batch_lookup FsAsyncLookupManager.batch_lookup

关键源码片段

csrc/fs_io.cpp dependency-wiring

核心新增 C 扩展文件,实现 GIL 释放的批量文件存在性检查。

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project#include <Python.h>
#include <unistd.h>
#include <vector>extern "C" {// 批量文件存在性检查(不持有 GIL,避免多线程竞争)
static void _batch_lookup(const std::vector<const char*>& paths,
                          std::vector<int>& exists_flags) {
  for (size_t i = 0; i < paths.size(); i++) {
    // 使用 access(2) 检查文件是否存在,F_OK 模式
    exists_flags[i] = (access(paths[i], F_OK) == 0) ? 1 : 0;
  }
}/// @brief Check file existence for a batch of paths.
/// @param paths list[str] – absolute paths to check.
/// @return list[bool] – True if the corresponding path exists, False otherwise.
/// @note Releases the GIL for the entire batch. File existence via access(2).
static PyObject* batch_lookup(PyObject* /*self*/, PyObject* args) {
  PyObject* path_list;
  if (!PyArg_ParseTuple(args, "O!", &PyList_Type, &path_list)) {
    return nullptr;
  }  const Py_ssize_t n = PyList_Size(path_list);
  std::vector<const char*> paths(n);
  // 提取所有路径字符串,若转换失败立即返回 nullptr
  for (Py_ssize_t i = 0; i < n; i++) {
    paths[i] = PyUnicode_AsUTF8AndSize(PyList_GetItem(path_list, i), nullptr);
    if (paths[i] == nullptr) {
      return nullptr;
    }
  }  std::vector<int> exists_flags(n);
  {
    // 释放 GIL,允许其他线程并行执行
    Py_BEGIN_ALLOW_THREADS
    _batch_lookup(paths, exists_flags);
    Py_END_ALLOW_THREADS
  }  // 将 int 数组转换为 Python list[bool]
  PyObject* result = PyList_New(n);
  if (result == nullptr) {
    return nullptr;
  }
  for (Py_ssize_t i = 0; i < n; i++) {
    PyList_SetItem(result, i, PyBool_FromLong(exists_flags[i]));
  }
  return result;
}static PyMethodDef fs_io_C_methods[] = {
    {"batch_lookup", batch_lookup, METH_VARARGS,
     "batch_lookup(paths: list[str]) -> list[bool]\n"
     "\n"
     "Check file existence for a batch of paths."},
    {nullptr, nullptr, 0, nullptr},
};static struct PyModuleDef fs_io_C_module = {
    PyModuleDef_HEAD_INIT, "fs_io_C", "Filesystem helpers for KV offload", -1,
    fs_io_C_methods,
};PyMODINIT_FUNC PyInit_fs_io_C(void) { return PyModule_Create(&fs_io_C_module); }} // extern "C"
vllm/v1/kv_offload/tiering/fs/manager.py dependency-wiring

修改 FsAsyncLookupManager,优先使用 C 扩展,否则 fallback 到 Python 实现。

try:
    from vllm.fs_io_C import batch_lookup as batch_lookup_C
    _HAS_BATCH_LOOKUP_C = True
except ImportError:
    _HAS_BATCH_LOOKUP_C = Falseclass FsAsyncLookupManager(AsyncLookupManager):
    """Async lookup manager for FileSystemTierManager."""
​
    def __init__(
        self,
        tier: "FileSystemTierManager",
        tier_type: str,
    ) -> None:
        super().__init__(tier_type=tier_type)
        self._tier = tier
​
    def batch_lookup(
        self, keys: list[OffloadKey], req_context: ReqContext
    ) -> Iterable[bool]:
        # 先收集所有路径,避免生成器反复调用 get_file_name
        paths = [self._tier.file_mapper.get_file_name(k) for k in keys]
        if _HAS_BATCH_LOOKUP_C:
            # C 扩展:整个 batch 通过一次 GIL 释放完成,大幅减少 GIL 竞争
            return batch_lookup_C(paths)
        # Fallback: 纯 Python 路径,每次 os.path.exists 都持有 GIL
        return (os.path.exists(p) for p in paths)

评论区精华

将 C 扩展从 fsio_C 重命名为 fs_io_C style

orozery 建议将模块名统一为 fs_io_C,避免缩写不一致。

结论:作者接受并修改所有引用。 · 已解决

提取纯 C 逻辑函数 _batch_lookup 设计

orozery 建议将文件检查逻辑从 batch_lookup 中抽出为独立静态函数,提高代码清晰性和测试性。

结论:作者采纳并实现。 · 已解决

batch_lookup 添加 docstring documentation

orozery 建议为函数添加 docstring 说明参数和返回值。

结论:作者添加。 · 已解决

返回类型保留 bool list 而非 bytes 设计

orozery 建议返回 bytes 对象代替 bool 列表以节省空间,作者认为 bool 列表更清晰,性能差异可忽略,决定保留。

结论:保持 bool list。 · 已解决

风险与影响

1) C扩展仅支持Python >=3.11,低版本自动回退到Python实现,性能收益丧失;
2) 使用POSIX access系统调用,不支持Windows平台(但vLLM主要部署在Linux,风险可控);
3) C扩展构建失败时_HAS_BATCH_LOOKUP_C为False自动fallback,服务不中断;
4) 输入参数校验严格,非字符串元素抛出TypeError,符合预期。

影响范围:所有使用FileSystemTierManager作为KV offload二级存储的配置(即secondary_tiers类型为fs的TieredOffloadingSpec)。通过释放GIL,批量路径检查不再阻塞Python线程,显著降低TTFT和ITL延迟,提升吞吐。无API或使用方式变更,对非FS tier用户无影响。

Python 3.11+ 依赖 C 扩展构建失败 fallback 仅 POSIX 平台

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论