[RFC]: Distributed KV Storage Backend for Mooncake Store

## [RFC]: Distributed KV Storage Backend for Mooncake Store

### 1. Introduction

This proposal introduces a **Distributed KV Storage Backend** for Mooncake Store as the fourth storage-backend implementation (alongside `FilePerKey`, `Bucket`, and `OffsetAllocator`). Through the `IDistributedKVClient` abstraction, it interfaces with external distributed KV stores, extending Mooncake's offload storage from the local filesystem to distributed KV storage devices.

The only implementation today is `UbsKVClient`, which provides NVMe KV SSD access by dynamically loading the UBSIO KV Cache library (`libubsio_kvc.so`). `DistributedKVStorageBackend` implements `StorageBackendInterface` and plugs into the existing `FileStorage` pipeline, requiring no changes to upper-layer application code.

**Goal:** give Mooncake Store a **low-cost, large-capacity distributed KV storage tier** that bypasses filesystem overhead and talks to the underlying storage devices directly through a KV interface, suitable for large-scale KV Cache offload scenarios.

### 2. Background and Motivation

As long-context scenarios proliferate (multi-turn conversation, coding assistants, document understanding), KV Cache sizes keep growing. DRAM is expensive and capacity-constrained, making it insufficient for large-scale KV Cache storage.

Mooncake's existing storage backends (`FilePerKeyBackend`, `BucketBackend`, `OffsetAllocatorBackend`) all operate through the filesystem layer — each offload/load maps to file create / open / write-or-read / close syscalls. This per-file overhead becomes a bottleneck under high-frequency, fine-grained I/O workloads.

Advantages of a distributed KV storage backend:

- **Native KV interface, low overhead.** Interacts with underlying storage directly via `Put` / `Get` / `Del` batch APIs, bypassing the filesystem and eliminating per-file syscall overhead.
- **Pluggable client abstraction.** The `IDistributedKVClient` interface allows different KV storage implementations to be wired in; the UBSIO implementation is provided today, and others can be added in the future.
- **Dynamic library loading.** `libubsio_kvc.so` is loaded at runtime via `dlopen` / `dlsym`, with no compile-time linking, keeping deployment flexible.
- **Zero-copy reads.** `UbsKVClient::BatchGet` leverages UBSIO's zero-copy interface (`UbsioBatchGet` returns a device-side pointer), reducing data-movement overhead.

### 3. Technical Design

#### 3.1 Architecture Overview

The integration follows Mooncake's existing storage-backend abstraction. `DistributedKVStorageBackend` implements `StorageBackendInterface` and talks to the underlying KV storage engine through the `IDistributedKVClient` abstraction layer.

```
flowchart LR
A[Application / FileStorage] --> B[StorageBackendInterface]
B --> C[DistributedKVStorageBackend]
C --> D[IDistributedKVClient]
D --> E[UbsKVClient]
E --> F[DlUbsioApi / libubsio_kvc.so]
F --> G[NVMe KV SSD Device]
```

#### 3.2 IDistributedKVClient Interface

`IDistributedKVClient` is a user-facing abstract interface that allows custom KV storage implementations:

| API | Signature | Description |
|-----|-----------|-------------|
| `Init()` | `virtual ErrorCode Init() = 0` | Initialize the connection |
| `BatchPut()` | `virtual tl::expected, ErrorCode> BatchPut(keys, values) = 0` | Batch write; returns a per-key result code |
| `BatchGet()` | `virtual ErrorCode BatchGet(keys, dest_slices) = 0` | Batch read into pre-allocated buffers |
| `Exists()` | `virtual tl::expected Exists(key) = 0` | Check whether a key exists |
| `ScanKeys()` | `virtual ErrorCode ScanKeys(handler) = 0` | Scan all keys; used for restart recovery |

#### 3.3 UbsKVClient Implementation

`UbsKVClient` is the UBSIO implementation of `IDistributedKVClient`:

- **Init:** dynamically loads `libubsio_kvc.so` via `DlUbsioApi::LoadLibrary()`, then calls `UbsioClientInit()` to initialize the client.
- **BatchPut:** converts a `vector` into a C-style pointer array and calls `UbsioBatchPut()` for the batch write.
- **BatchGet:** calls `UbsioBatchGet()` to obtain a device-side zero-copy pointer, `memcpy`s into the target Slice, and then asynchronously calls `UbsioBatchFreeAddress()` (via a thread pool) to release the device-side memory.
- **Exists:** calls `UbsioExist()` to check key existence.
- **ScanKeys:** currently **unsupported** (logs "Ubsio not support scan meta").

#### 3.4 DlUbsioApi — Dynamic Library Loading

`DlUbsioApi` wraps the loading and invocation of the UBSIO dynamic library:

- Uses `dlopen` / `dlsym` to load `libubsio_kvc.so` at runtime.
- Wraps 13 C function pointers: `UbsioClientInit` / `Put` / `Get` / `Exist` / `Delete` / `GetLength` / `BatchPut` / `BatchGet` / `BatchGetDirect` / `BatchExist` / `BatchDelete` / `BatchGetLength` / `BatchFree`.
- Thread-safe (load/unload guarded by a `std::mutex`).

#### 3.5 StorageBackendInterface Integration

`DistributedKVStorageBackend` implements `StorageBackendInterface`:

- **`Init()`:** creates and initializes an `UbsKVClient`.
- **`BatchOffload()`:** concatenates multiple Slices into a contiguous byte string, calls `BatchPut` to write to the KV backend, and on success notifies the master via `complete_handler` to add a `LOCAL_DISK` replica.
- **`BatchLoad()`:** reads data from the KV backend into an RDMA-registered CPU ClientBuffer.
- **`IsExist()`:** delegates to `kv_client_->Exists()`.
- **`IsEnableOffloading()`:** checks `total_size_ < total_size_limit && total_keys_ < total_keys_limit`.
- **`ScanMeta()`:** scans existing data, invoking the handler in batches to restore master metadata.

**Backend registration:**

- A new enum value `StorageBackendType::kDistributedKV`.
- The `CreateStorageBackend()` factory instantiates `DistributedKVStorageBackend` on a type match.
- The backend is selected via an environment variable: `MOONCAKE_OFFLOAD_STORAGE_BACKEND_DESCRIPTOR=distributed_kv_storage_backend`.

#### 3.6 Data Flow

```
Write path (BatchOffload):
Inference engine → FileStorage::Heartbeat()
→ FileStorage::OffloadObjects()
→ GPU→CPU D2H copy (pinned_buffer_pool_)
→ DistributedKVStorageBackend::BatchOffload()
→ concatenate multiple Slices into a contiguous string
→ UbsKVClient::BatchPut() → DlUbsioApi::UbsioBatchPut() → libubsio_kvc.so
→ complete_handler → client_->NotifyOffloadSuccess() (notify master)

Read path (BatchLoad):
Inference engine → FileStorage::BatchGet()
→ DistributedKVStorageBackend::BatchLoad()
→ UbsKVClient::BatchGet() → DlUbsioApi::UbsioBatchGet() (zero-copy returns device pointer)
→ memcpy into the RDMA-registered CPU ClientBuffer
→ asynchronous UbsioBatchFreeAddress (thread pool)
→ TransferEngine RDMA pull to the requesting GPU
```

#### 3.7 Configuration

Activated via an environment variable:

```bash
export MOONCAKE_OFFLOAD_STORAGE_BACKEND_DESCRIPTOR=distributed_kv_storage_backend
```

`FileStorageConfig::Validate()` skips the `storage_filepath` path validation for the `kDistributedKV` type (the KV backend needs no local file path).

### 4. Key Source Files

| File | Description |
|------|-------------|
| `mooncake-store/include/storage_backend.h` | Core interface definitions: `IDistributedKVClient`, `UbsKVClient`, `DistributedKVStorageBackend`, `StorageBackendType::kDistributedKV` |
| `mooncake-store/include/dl_ubsio_api.h` | UBSIO dynamic-library loading wrapper `DlUbsioApi`; declares all ubsio C function-pointer types |
| `mooncake-store/src/storage_backend.cpp` | Full implementation of `UbsKVClient` and `DistributedKVStorageBackend` |
| `mooncake-store/src/dl_ubsio_api.cpp` | `DlUbsioApi::LoadLibrary()` / `CleanupLibrary()` implementation |
| `mooncake-store/src/file_storage.cpp` | Integration point: env-var parsing, `Validate()` path-skip, `Heartbeat()` empty-list optimization |
| `mooncake-store/src/CMakeLists.txt` | Added `dl_ubsio_api.cpp` to the build list |
| `mooncake-store/tests/storage_backend_test.cpp` | `MockDistributedKVClient` + 18 `DistributedKVStorageBackend` unit tests |

### 5. Test Coverage

`storage_backend_test.cpp` contains **18** `DistributedKVStorageBackend` test cases:

- **Init:** success / failure
- **BatchOffload:** basic / multiple keys / empty batch / uninitialized / over capacity limit / over key-count limit
- **BatchLoad:** basic / key not found / uninitialized
- **IsExist:** exists / not exists / uninitialized
- **IsEnableOffloading:** within limits / over capacity / over key-count
- **ScanMeta:** basic / empty / uninitialized / failure

Tests use `MockDistributedKVClient` (an in-memory map simulating the KV store) and inject `kv_client_` directly via `friend class StorageBackendTest`.

### 6. Current Limitations & Future Work

| # | Limitation | Impact | Priority |
|---|-----------|--------|----------|
| 1 | **ScanKeys unsupported**: `UbsKVClient::ScanKeys()` only logs an error and returns `OK` | After a process restart, master metadata cannot be recovered via `ScanMeta` — a critical missing capability | **P0** |
| 2 | **No data deletion / eviction mechanism**: `eviction_handler` is ignored; data in the KV backend only grows | Storage space cannot be reclaimed; long-running deployments will exhaust capacity | **P0** |
| 3 | **Hardcoded UbsKVClient**: `Init()` directly does `make_unique()`; no config injection of other `IDistributedKVClient` implementations | The abstraction exists but the factory is incomplete; KV clients cannot be switched at runtime | **P1** |
| 4 | **Extra memcpy in BatchGet**: after obtaining a device pointer via zero-copy, a `memcpy` into the target Slice is still needed | Adds one memory copy, hurting read performance | **P1** |
| 5 | **HBM direct-read path not implemented**: `DlUbsioApi` wraps `UbsioBatchGetWithHBM`, but `UbsKVClient::BatchGet` does not use it | Cannot read directly into GPU memory; missing a GPU direct-read optimization path | **P2** |
| 6 | **Incomplete BatchPut partial-success handling**: data for failed keys may already have been written but not cleaned up | Orphaned data may remain in the KV backend | **P2** |
| 7 | **No Delete interface invoked**: `DlUbsioApi` wraps `UbsioDelete` / `UbsioBatchDelete`, but the backend never calls them | Data cannot be proactively deleted | **P1** |
| 8 | **Incomplete config validation**: `Validate()` only skips path validation; no KV-backend-specific config validation | Missing validation of key configs such as the UBSIO library path | **P2** |
| 9 | **Init thread safety**: uses `atomic` double-check, but two threads may still create an `UbsKVClient` concurrently | Concurrent init may leak resources | **P2** |

### 7. Compatibility

- **Interface compatible:** `DistributedKVStorageBackend` fully implements `StorageBackendInterface`; all callers using `FileStorage` need no changes.
- **Non-intrusive:** existing backends (`BucketBackend`, `FilePerKeyBackend`, `OffsetAllocatorBackend`) are entirely unaffected.
- **Runtime switching:** the backend is selected via a single environment variable, supporting A/B comparison at deployment time without recompilation.
- **Dynamic-library decoupling:** `libubsio_kvc.so` is loaded at runtime via `dlopen`, with no compile-time linking; its absence does not affect other backends.

### 8. Alternatives Considered

**SSD KV Engine (#1990):** that proposal compiles SSD KV Engine as a static library integrated into Mooncake Store, providing self-contained single-node SSD offload. By contrast, this proposal's `DistributedKVStorageBackend` supports distributed KV storage through the `IDistributedKVClient` abstraction layer, focusing more on storage expansion in multi-node scenarios. The two can coexist, covering single-node and distributed cases respectively.

**NVMe KV Backend (#1957):** that proposal introduces a four-layer Backend → Connector → Executor separation — more complete in architecture but also more complex. The current implementation can be viewed as a simplified version of that architecture (`DistributedKVStorageBackend` ≈ Backend, `UbsKVClient` ≈ Connector, `DlUbsioApi` ≈ Executor), and can gradually evolve toward the full architecture later.

Start with mooncake-store/include/storage_backend.h and src/storage_backend.cpp, then inspect the integration points in src/file_storage.cpp and the 18 cases in tests/storage_backend_test.cpp. Compare the listed limitations with the desired backend behavior and clarify which capability is in scope; done requires an agreed design, implementation criteria, and corresponding tests.

Written by the indexing model from the issue text.