kvcache-ai / kvcache-ai/Mooncake

[RFC]: Add GDS offload for mooncake-store based on existing TENT gds transport

Open
#3,786 4 comments 1 reaction 0 assignees View on GitHub
RFC
Dominant language
C++
Stars
6.6k
Forks
1.2k
Avg merge
3d 5h
Merged PRs (30d)
312

Description

### Changes proposed

## Summary

Add opt-in GDS SSD offloading to mooncake store by reusing the existing TENT GDS transport.

The implementation adds:

1. A client-local, fixed-size GDS backing file exposed as a TENT `file://` segment and managed by an aligned extent allocator.
2. A first-class `ReplicaType::GDS` with storage identity, generation, offset, value size, and allocation size in store metadata and wire serialization.
3. Required dual-write behavior for enabled `BatchPut`: memory/NoF writes and GDS writes are submitted in parallel and waited on before finalization.
4. A GDS-only creation fallback when memory replica allocation fails because of capacity or tenant quota.
5. Memory-first `BatchGet` selection, followed by a local matching GDS replica, while retaining existing NoF, local-disk, and disk fallbacks.
6. TENT/Transfer Engine capability checks and GDS buffer-registration tracking needed for sub-slice cuFile requests.

The feature is enabled through environment variables and requires no changes to upstream vLLM, so existing deployments remain unchanged when it is disabled.

## Motivation

Mooncake store currently supports POSIX and `io_uring` file offloading. Its existing data path is as follows:

1. **Write path:** when memory usage in a global segment (CPU memory) reaches the configured eviction watermark, mooncake master selects objects for offloading and places them in the background offload queue. The client then reads the objects from the global segment and writes them to the configured file backend through POSIX or `io_uring`. After the file write completes, the master records the file-backed replica and may release the corresponding memory replica.
2. **Read path:** the client first queries the master for the object's replica metadata and prefers an available memory replica. If the object is no longer present in memory, the client loads it from the file backend through POSIX or `io_uring` and returns the data to the caller.

This data path becomes a bottleneck when host memory is constrained. GPU-to-CPU and CPU-to-SSD bandwidths are not matched: GPU data must first be copied into CPU memory and then written to SSD. **When GPU production outpaces this two-step pipeline, CPU memory and CPU-side file I/O cannot drain the global segment fast enough. Some KV cache is consequently discarded before it can be completely offloaded, preventing more valid KV cache from being offloaded to the SSD even when sufficient SSD capacity remains.**

GPUDirect Storage (GDS) addresses this mismatch by supporting GPU-to-SSD DMA, so the data can bypass the CPU staging path. The GDS transport has already been integrated into TENT, **but the data path between TENT and mooncake store has not yet been connected.** An earlier draft PR (#2731) attempted to add GDS offload at the mooncake store layer by introducing a separate GDS context. This approach would duplicate cuFile resource management and make mooncake store responsible for transport-specific details that already belong to TENT.

Therefore, the central design goal of this change is to preserve the existing separation of mooncake store and TENT while connecting the two components: TENT owns all GDS cuFile operations, transport selection, and buffer registration, while mooncake store owns only object and replica metadata, allocation/lifecycle decisions, and transfer-request submission. Mooncake store should reuse the existing Transfer Engine and TENT GDS transport instead of introducing another GDS/cuFile implementation. This keeps the GDS data path composable with store's replica semantics and avoids a second, divergent GDS context in mooncake store.

## Non-goals

- Replacing the existing `enable_ssd_offload`/`LOCAL_DISK` POSIX or `io_uring` path.
- Providing remote reads from another client's local GDS file. A remote GDS descriptor falls back to the existing memory, NoF, local-disk, or disk paths.
- Recovering a client's GDS index after restart or promising crash/durability semantics for the local file.

## Architecture

```text
RealClient (env opt-in)
+-- client
+-- GdsStorageManager
| +-- aligned extent allocator
| +-- file:// backing segment handle
| `-- storage_id + generation + read leases
+-- TransferSubmitter::submitFile
| `-- TransferEngine / TENT GDS / cuFile
`-- MasterClient RPCs
`-- MasterService: GDS storage identities and replicas

BatchPut:
reserve local extent
+-- BatchPutStart -> MEMORY/NoF PROCESSING replicas
+-- BatchAddGdsReplicaStart -> GDS PROCESSING replica
`-- submit both write paths -> wait -> commit + BatchPutEnd

BatchGet:
local MEMORY -> any MEMORY -> matching local GDS -> NoF/disk fallbacks
```

### Local GDS storage

`GdsStorageManager` creates/truncates and preallocates the configured backing file, opens it once through `openSegment("file://...")`, and allocates raw value extents aligned to 4 KiB. Reservations are move-only and automatically aborted unless committed or explicitly preserved after an uncertain RPC. A reference-counted read lease keeps an extent alive until all GDS read futures finish.

Each client registers a `storage_id` and `storage_generation` with the master. The generation changes when the file is recreated, preventing an old metadata descriptor from addressing the same offset in a new file.

### Store metadata and RPCs

`ReplicaType::GDS` is an independent replica type rather than an alias for `LOCAL_DISK`:

```text
GdsDescriptor {
owner_client_id
storage_id
storage_generation
value_offset
value_size
allocated_size
}
```

The master exposes storage registration/unregistration, `BatchAddGdsReplicaStart`, and `BatchCreateGdsOnlyObjects`. GDS replicas use the existing `PROCESSING` -> `COMPLETE` lifecycle, participate in `BatchPutEnd` / `BatchPutRevoke`, and are removed during stale-client cleanup. MessagePack serialization uses a stable explicit replica type value and validates the six-field GDS payload.

GDS offload uses a client-local allocator modeled after the existing `offset_allocator_storage_backend`: `GdsStorageManager` preallocates one backing file and creates an `OffsetAllocator` over the byte range `[0, capacity)`, so each object is assigned a reusable file offset rather than requiring one file per object. The requested allocation size is rounded up to a 4 KiB extent, while the descriptor retains both the original `value_size` and the aligned `allocated_size`. The resulting extents are addressed directly by TENT `submitFile` requests, with the same allocator-backed reuse and reclamation model as `OffsetAllocatorStorageBackend`.

Reservations are move-only and use RAII: failed or abandoned operations release their extent, successful commits retain the allocation in the local key index, and uncertain master finalization can preserve it. Reference-counted read leases keep a committed extent allocated until all in-flight GDS reads finish.

Restart recovery is outside the scope of this backend. Initialization recreates the backing file and resets both the allocator and the in-memory GDS index, so objects written by a previous process cannot be discovered or accessed after a restart.

### Write path

When GDS is enabled, each valid key reserves an aligned extent before the normal memory allocation request. The client then submits memory/NoF and GDS transfers without waiting between the two submissions. GDS requests use the cached file segment handle and accumulate the slice offset from the extent base. `submitFile` validates pointer, offset, and length alignment, splits large slices at 16 MiB, and splits batches at 128 entries.

Both paths must succeed for the normal dual-write case. **If the master cannot allocate the requested memory replica due to capacity or tenant quota, the reserved extent can instead become a GDS-only object through the master fallback RPC.** Per-key finalize/revoke groups ensure one key's failure does not abort unrelated keys, and no extent is released before already-submitted futures terminate.

### Read path

Replica selection is explicit and consistent across single-buffer and multi-buffer APIs: local complete memory, any complete memory, a complete GDS replica whose owner/storage generation matches this client, then existing NoF and disk fallbacks. A local GDS read validates the descriptor against the client's extent map, acquires a read lease, submits TENT GDS READ requests directly into the caller's registered destination slices, and releases the lease only after completion.

### TENT integration

The compatibility `TransferEngine` now exposes `hasTransport`. TENT forwards GDS and `io_uring` availability so store can fail fast when the requested path is not actually installed. `GdsTransport` records registered CUDA buffer ranges and resolves an interior source pointer to its registered base and offset before submitting cuFile batch entries. This prevents valid multi-slice requests from being rejected as unregistered buffers.

## Configuration

```bash
export MOONCAKE_GDS_OFFLOAD_ENABLED=1
export MOONCAKE_GDS_OFFLOAD_PATH=/path/to/gds_offload
export MOONCAKE_OFFLOAD_TOTAL_SIZE_LIMIT_BYTES=107374182400

export MC_USE_TENT=1
export MC_STORE_MEMCPY=0
export MC_TENT_CONF='{
"transports": {
"tcp": {"enable": true},
"gds": {"enable": true, "io_batch_depth": 128},
"io_uring": {"enable": false}
}
}'
export CUFILE_ENV_PATH_JSON=/path/to/cufile.json
```

Enablement requires an absolute backing-file path, a non-zero capacity, an active TENT GDS transport, and disabled TENT file fallback transports such as `io_uring`. The capacity currently reuses `MOONCAKE_OFFLOAD_TOTAL_SIZE_LIMIT_BYTES`.

## Validation

The commit adds or updates:

- Client buffer tests for GDS size calculation and slice allocation.
- Master tests for GDS-only object creation without memory segments.
- Optional CUDA/TENT integration coverage for multi-slice GPU -> GDS -> GPU round trips, memory-first reads, memory removal, and GDS-only reads. The integration test is gated by `MOONCAKE_RUN_GDS_TEST=1` and `MOONCAKE_GDS_TEST_PATH`.
- TENT transport selector/hint coverage and GDS registered-buffer handling.
- MessagePack and RPC plumbing for the new descriptor and lifecycle calls.

The integration test requires a CUDA device, a GDS-capable filesystem, and a valid cuFile configuration; ordinary CI can run the non-GDS unit coverage.

## Roadmap

### Delivered in commit `3265a67c`

1. Add store-local aligned extent allocation and file-segment lifecycle.
2. Add master GDS identity registration, replica metadata, cleanup, and RPCs.
3. Integrate parallel memory/GDS `BatchPut` with GDS-only quota/capacity fallback and per-key finalization.
4. Integrate local GDS `BatchGet` fallback and consistent replica selection.
5. Add TENT capability queries, registered-buffer lookup, batching limits, and focused tests.

### Follow-up work

- Add a dedicated GDS capacity/eviction policy and reclaim extents safely.
- Add GDS offload stats in mooncake master
- TBD

### Before submitting a new issue...

- [x] Make sure you already searched for relevant issues and read the [documentation](https://kvcache-ai.github.io/Mooncake/)

Contributor guide

Open the contributing guide

Research direction

The RFC identifies GdsStorageManager, TransferSubmitter::submitFile, MasterService RPCs, and GdsTransport as the main entry points, with client buffer, master, MessagePack/RPC, and TENT transport tests. The roadmap says the core work was delivered in commit 3265a67c; verify ordinary non-GDS coverage first, then use the optional CUDA integration test with MOONCAKE_RUN_GDS_TEST=1 and MOONCAKE_GDS_TEST_PATH. Follow-up scope includes GDS capacity/eviction policy and master statistics.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
backend, distributed-systems, infrastructure
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
20/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.