[fs_connector][feat]: Add Python backend implementation for fs backend of offloading connector (fs_connector)

Signed-off-by: Kfir Toledo <[email protected]>

Author: kfirtoledo

.gitignore

@@ -22,6 +22,25 @@ __pycache__/
 *.pyd
 *.python-version
 
+# Python build artifacts
+*.egg-info/
+build/
+dist/
+
+# C++/CUDA build outputs
+*.o
+*.so
+*.d
+*.a
+
+# Ninja build files
+build.ninja
+.ninja_log
+.ninja_deps
+
+# Temporary pip build directories
+*.egg-info/
+
 # Go workspace file
 go.work
 go.work.sum

kv_connectors/llmd_fs_backend/README.md

@@ -0,0 +1,98 @@
+# llmd-fs-backend README
+
+## Overview
+The llmd-fs-backend extends the native [vLLM Offloading Connector](#offloading-connector-docs) to support a file system backend.
+This backend provides a shared-storage offloading layer for vLLM. It moves KV-cache blocks between GPU and shared storage efficiently using:
+
+- Async CUDA copies or GPU kernels
+- Pinned memory pools
+- Multi-threaded I/O workers
+- NUMA-aware CPU affinity
+- Atomic file writes and zero-copy reads
+
+The fs connector (llmd_fs_backend) is used for shared storage but it can also work with local disk.
+
+For architectural clarity, the fs connector is not responsible for cleanup. Storage systems should manage this.
+For simple setups, see the **Storage Cleanup** section.
+
+<img src="./docs/images/fs_connector.png" width="400" />
+
+## System Requirements
+- vLLM version 0.11.0 or above, which includes the Offloading Connector
+
+## Installation
+
+```bash
+apt-get update && apt-get install -y libnuma-dev
+pip install git+https://github.com/llm-d-kv-cache-manager.git#subdirectory=kv_connectors/llmd_fs_backend
+```
+
+This installs:
+- Python module `llmd_fs_backend`
+- CUDA extension `storage_offload.so`
+
+## Configuration Flags
+
+### Connector parameters
+
+- `shared_storage_path`: filesystem path for store and load the KV files.
+- `block_size`: number of GPU blocks grouped into each file (must be in granulaity of GPU block size that)
+- `threads_per_gpu`: number of I/O threads per GPU
+- `max_pinned_memory_gb`: total pinned memory limit
+
+### Environment variables
+- `STORAGE_CONNECTOR_DEBUG`: enable debug logs
+- `USE_KERNEL_COPY_WRITE`: enable GPU-kernel writes (default 0)
+- `USE_KERNEL_COPY_READ`: enable GPU-kernel reads (default 1)
+
+## Example vLLM YAML
+
+To load the fs connector:
+
+```yaml
+--kv-transfer-config '{
+  "kv_connector": "OffloadingConnector",
+  "kv_role": "kv_both",
+  "kv_connector_extra_config": {
+    "spec_name": "SharedStorageOffloadingSpec",
+    "spec_module_path": "llmd_fs_backend.spec",
+    "shared_storage_path": "/mnt/files-storage/kv-cache/",
+    "block_size": 256,
+    "threads_per_gpu": "64"
+  }
+}'
+--distributed_executor_backend "mp"
+```
+
+A full deployment example can be found in the [`docs`](./docs/deployment) folder.
+
+It is recommended to use multiprocess mode by setting:
+`--distributed_executor_backend "mp"`
+
+To configure environment variables:
+
+```yaml
+env:
+- name: STORAGE_CONNECTOR_DEBUG
+  value: 1
+```
+
+## Storage Cleanup
+TBD
+
+## Troubleshooting
+
+### Missing `numa.h`
+Install the required package:
+
+```bash
+apt-get install -y libnuma-dev
+```
+
+---
+
+## Link Aliases
+
+- **Offloading Connector Docs**
+  <a name="offloading-connector-docs"></a>
+  https://docs.vllm.ai/en/stable/features/disagg_prefill/#usage-example:~:text=backends%22%3A%5B%22UCX%22%2C%20%22GDS%22%5D%7D%7D%27-,OffloadingConnector,-%3A%20enable%20offloading%20of

kv_connectors/llmd_fs_backend/docs/images/fs_connector.png

@@ -0,0 +1,39 @@
+[build-system]
+requires = [
+    "setuptools>=65",
+    "wheel",
+    "torch",
+    "ninja"
+]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "llmd_fs_connector"
+version = "0.1.0"
+description = "Standalone llm-d fs storage connector"
+readme = "README.md"
+authors = [
+    { name = "Kfir", email = "[email protected]" }
+]
+maintainers = [
+    { name = "llm-d community" }
+]
+requires-python = ">=3.9"
+dependencies = [
+    "torch>=2.1",
+]
+
+[tool.setuptools]
+packages = ["llmd_fs_backend"]
+package-dir = {"" = "src"}
+
+[tool.setuptools.package-data]
+llmd_fs_backend = ["*.so"]
+
+[project.optional-dependencies]
+dev = [
+    "vllm",
+    "pytest",
+    "black",
+    "ruff",
+]

kv_connectors/llmd_fs_backend/pyproject.toml

@@ -0,0 +1,9 @@
+from vllm.logger import init_logger
+from vllm.v1.kv_offload.factory import OffloadingSpecFactory
+
+logger = init_logger(__name__)
+
+# Register SharedStorageOffloadingSpec to offloading connector
+OffloadingSpecFactory.register_spec("SharedStorageOffloadingSpec",
+                                    "vllm.v1.kv_offload.shared_storage",
+                                    "SharedStorageOffloadingSpec")

kv_connectors/llmd_fs_backend/src/llmd_fs_backend/__init__.py

@@ -0,0 +1,109 @@
+import os
+import torch
+from pathlib import Path
+from collections.abc import Iterable
+from typing import Optional
+
+from vllm.v1.core.kv_cache_utils import BlockHash
+from llmd_fs_backend.mediums import SharedStorageLoadStoreSpec
+from vllm.v1.kv_offload.abstract import (
+    LoadStoreSpec,
+    OffloadingManager,
+    PrepareStoreOutput,
+)
+from llmd_fs_backend.worker import StorageOffloadingHandler
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class SharedStorageOffloadingManager(OffloadingManager):
+    """
+    SharedStorageOffloadingManager manages KV offloading to a shared storage medium.
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        tp_size: int,
+        tp_rank: int,
+        dtype: torch.dtype,
+        root_dir: str = "/tmp/shared-kv",
+    ) -> None:
+
+        # Basic metadata about the model and tensor parallelism
+        self.model_name = model_name
+        self.tp_size = tp_size
+        self.tp_rank = tp_rank
+        self.dtype = dtype
+
+        # Resolve base directory where KV files for this model and tp rank are stored
+        self.base_path: Path = StorageOffloadingHandler.get_kv_cache_base_path(
+            dtype=dtype,
+            model_name=model_name,
+            tp_size=tp_size,
+            tp_rank=tp_rank,
+            root_dir=root_dir,
+        )
+
+    # ----------------------------------------------------------------------
+    # Lookup
+    # ----------------------------------------------------------------------
+    def lookup(self, block_hashes: Iterable[BlockHash]) -> int:
+        """
+        Return how many consecutive blocks from the start are already offloaded.
+        """
+        hit_count = 0
+        for block_hash in block_hashes:
+            file_path = StorageOffloadingHandler.get_file_name(self.base_path, block_hash)
+            if not os.path.exists(file_path):
+                break
+            hit_count += 1
+        return hit_count
+
+    # ----------------------------------------------------------------------
+    # Load
+    # ----------------------------------------------------------------------
+    def prepare_load(self, block_hashes: Iterable[BlockHash]) -> LoadStoreSpec:
+        """
+        For shared storage, loading is stateless - return specs that point to files.
+        """
+        return SharedStorageLoadStoreSpec(block_hashes)
+
+    def touch(self, block_hashes: Iterable[BlockHash]):
+        """
+        Update access times if desired.
+        Shared storage version does nothing here because updates are handled
+        by the file thread for performance reasons.
+        """
+        pass
+
+    def complete_load(self, block_hashes: Iterable[BlockHash]):
+        """Stateless load - no post-load action needed."""
+        pass
+
+    # ----------------------------------------------------------------------
+    # Store
+    # ----------------------------------------------------------------------
+    def prepare_store(self, block_hashes: Iterable[BlockHash]) -> Optional[PrepareStoreOutput]:
+        """
+        Prepare storing new blocks.
+        Shared storage always accepts new blocks. Eviction is not needed.
+        If a file already exists, the file thread handles it.
+        """
+        block_hashes_to_store = list(block_hashes)
+
+        # Set up store spec
+        store_spec = SharedStorageLoadStoreSpec(block_hashes_to_store)
+
+        return PrepareStoreOutput(
+            block_hashes_to_store=block_hashes_to_store,
+            store_spec=store_spec,
+            block_hashes_evicted=[],  # no eviction needed
+        )
+
+    def complete_store(self, block_hashes: Iterable[BlockHash], success: bool = True):
+        """
+        For shared storage, storing is stateless - no action needed.
+        """
+        pass

kv_connectors/llmd_fs_backend/src/llmd_fs_backend/factory.py

@@ -0,0 +1,27 @@
+import numpy as np
+from typing import Iterable
+from vllm.v1.core.kv_cache_utils import BlockHash
+from vllm.v1.kv_offload.abstract import LoadStoreSpec
+
+class SharedStorageLoadStoreSpec(LoadStoreSpec):
+    """
+    Spec for loading and storing KV blocks on shared storage.
+    Stores block hashes internally as a numpy array.
+    """
+
+    def __init__(self, block_hashes: Iterable[BlockHash]):
+        # Validate all items are bytes (BlockHash)
+        block_hashes = list(block_hashes)
+        for h in block_hashes:
+            if not isinstance(h, (bytes, bytearray)):
+                raise TypeError(f"Expected BlockHash (bytes-like), got {type(h).__name__}")
+
+        # Store directly as object array of bytes
+        self.block_hashes = np.array(block_hashes, dtype=object)
+
+    def __repr__(self) -> str:
+        return repr(self.block_hashes)
+
+    @staticmethod
+    def medium() -> str:
+        return "SHARED_STORAGE"