Study llmd-fs-backend

定位

vLLM 的 Offloading Connector 插件，负责把 KV Cache 数据在 GPU 显存和共享存储（磁盘/NFS/对象存储）之间高效搬运。

架构分层

vLLM (Python)
  │ 调用 OffloadingConnector → SharedStorageOffloadingSpec
  │
Python 层 (llmd_fs_backend/)
  ├── spec.py       ← 入口：vLLM 加载这个
  ├── manager.py    ← 调度端：lookup, prepare_store/load
  ├── worker.py     ← 工作端：创建 C++ Engine，包装 Handler
  ├── file_mapper.py← 路径映射：block hash → 文件路径
  └── event_publisher.py ← ZMQ 事件通知 KV Cache Manager
  │
  │ pybind11 绑定
  │
C++ 层 (csrc/storage/)
  ├── storage_offload.cpp    ← 核心引擎：异步 I/O 调度
  ├── thread_pool.hpp        ← 线程池：读写优先级队列
  ├── tensor_copier.cu       ← GPU↔CPU 数据搬运
  └── backends/
      ├── fs_io/file_io.cpp  ← 普通文件 I/O (CPU staging)
      └── fs_gds/gds_file_io.cpp ← GDS 直接 I/O (GPU 直通)

Setup

通过 setup.py 把 C++/CUDA 代码编译成 Python 包 storage_offload.so。

用 CUDAExtension 编译，BuildExtension 接管构建。
FS_PORTABLE_WHEEL=1 时静态链接 libstdc++，生成可移植 wheel。
安装后 Python 里直接 import storage_offload。

部署

在 vLLM Pod 里 pip install llmd-fs-connector，然后在 vLLM 启动参数里配置：

--kv-transfer-config '{
  "kv_connector": "OffloadingConnector",
  "kv_connector_extra_config": {
    "spec_name": "SharedStorageOffloadingSpec",
    "spec_module_path": "llmd_fs_backend.spec",
    "shared_storage_path": "/mnt/kv-cache/",
    "block_size": 256
  }
}'

模块详解

spec.py：入口

vLLM 加载 SharedStorageOffloadingSpec，它创建两个核心组件：

get_manager() → 返回 SharedStorageOffloadingManager
get_handlers() → yield 出 GPUToStorageHandler 和 StorageToGPUHandler

manager.py：调度端

运行在 Rank 0，负责决策，不碰实际数据。

lookup(key)：问 FileMapper 文件是否存在
prepare_store(keys)：准备写入，返回存储参数
complete_store(keys, success)：写完后调 event_publisher.publish_blocks_stored() 发 ZMQ 事件

file_mapper.py：路径生成器

把 vLLM 的 block hash 翻译成磁盘文件路径。

配置指纹：SHA256 哈希所有配置字段成 12 位目录名，相同配置共享 Cache。
Rank 隔离：每个 GPU 进程有独立的 _r<rank> 目录。
Hex 分桶：aaa/bb_g0/ 两级目录，防止单目录文件过多。
配置共享：config.json 记录所有字段，多 rank 共用。

event_publisher.py：ZMQ 通知

当 block 成功写到磁盘后，通过 ZMQ PUB 发事件给 Go 端的 KV Cache Manager。

zmq.Context()：I/O 引擎
zmq.PUB socket：广播模式，发完即忘
LINGER=0：关闭时不等待
SNDHWM=100000：发送队列上限
消息格式：[topic, sequence, payload] 三帧，msgpack 编码

worker.py：工作端

StorageOffloadingHandlers

工厂容器，负责搭建 C++ 引擎并生产两个 Handler。

计算 Staging Buffer 大小，必要时降低线程数防内存溢出
创建 C++ StorageOffloadEngine（线程池、CUDA Stream、GDS 配置）
生成 GPUToStorageHandler（PUT）和 StorageToGPUHandler（GET）
两个 Handler 共享同一个 engine 和 pending_jobs 字典

vLLM 通过 spec.get_handlers() 把这两个 Handler 抽走直接调用。

BaseStorageOffloadingHandler

真正的搬运工，负责把逻辑 Block 请求翻译成物理文件路径。

transfer_async(job_id, spec)：发起异步传输，提交给 C++ 引擎
_build_file_block_mapping：拆解任务，把 Block 按文件归属分组，算出文件路径和偏移量
_build_transfer：遍历所有 KV Cache Group，拼成完整传输列表
get_finished()：收割已完成任务，计算传输耗时和吞吐量
wait(job_ids)：阻塞等待任务完成

StorageEngine (C++ 接口)

通过 pybind11 暴露的 C++ 引擎，提供：

async_store_gpu_blocks()：GPU → 磁盘
async_load_gpu_blocks()：磁盘 → GPU
get_finished()：返回完成的 Job ID 列表
wait_job()：阻塞等待

GDS (GPUDirect Storage)

NVIDIA 技术，让 GPU 直接和 NVMe 磁盘对话，跳过 CPU 内存中转。

普通模式：磁盘 → CPU 内存 → GPU（两次拷贝）
GDS 模式：磁盘 → PCIe → GPU（一次拷贝，更低延迟）
通过 libcufile 实现，运行时 dlopen 动态加载

数据流

写操作 (GPU → 磁盘)

vLLM 显存满了，调 manager.prepare_store()
vLLM 调 GPUToStorageHandler.transfer_async()
Handler 用 FileMapper 算出文件路径，调 C++ async_store_gpu_blocks()
C++ 线程池执行 GPU → 磁盘传输
完成后 vLLM 调 manager.complete_store()
Manager 调 event_publisher.publish_blocks_stored() 发 ZMQ 事件

读操作 (磁盘 → GPU)

vLLM 发现 Cache Miss，调 manager.lookup()
Manager 查文件是否存在，返回 True
vLLM 调 manager.prepare_load() 获取文件路径
vLLM 调 StorageToGPUHandler.transfer_async()
C++ 线程池执行磁盘 → GPU 传输
vLLM 拿到数据，继续推理

事件发布机制

有两条并行的事件发布路线，把 KV Cache 的变化通知给 Go 端的 KV Cache Manager。

路线 1：vLLM 原生路线（GPU KV Events）

vLLM Scheduler 每个调度步骤调用 connector.take_events() → OffloadingConnectorScheduler.take_events() → Manager.take_events()，收集 Manager 内部缓存的事件，打包成 KVEventBatch，通过 vLLM 自带的 ZmqEventPublisher 发布到端口 5557。

路线 2：StorageEventPublisher 路线（FS Backend 自带）

当 C++ 写完磁盘，vLLM 调用 manager.complete_store(keys, success=True)，Manager 内部直接调用 StorageEventPublisher.publish_blocks_stored()，通过独立的 ZMQ PUB 发送到端口 5559。

对比

	路线 1 (vLLM 原生)	路线 2 (StorageEventPublisher)
触发点	vLLM Scheduler 每个步骤调用 `take_events()`	`complete_store()` 被调用时
发布者	vLLM 的 `ZmqEventPublisher`	`StorageEventPublisher`
端口	5557（vLLM 配置）	5559（FS Backend 配置）
Topic 格式	`kv@{ip}@{model}`	`kv@SHARED_STORAGE@{model}`
可控开关	`enable_kv_cache_events`	`enable_events`

为什么需要路线 2？

来自 PR #571（Alberto Perdomo），原因：

更低延迟：不等待 vLLM 调度步骤，I/O 完成后立即发送。
独立 Topic 标记：Topic 里包含 SHARED_STORAGE 这个 medium 标记，vLLM 原生 GPU 事件没有这个元数据。这让 Go 端可以区分事件来自存储层还是 GPU 层。
架构解耦：FS Backend 作为独立插件，自带发布器确保即使 vLLM 没配置事件发布，存储层事件也能可靠发送。

实际效果：Go 端的 KV Cache Manager 可以订阅任意一个端口（或两个都订阅），收到的都是 BlockStored/BlockRemoved 事件。