Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 27 additions & 0 deletions csrc/cumem_ipc_interposer/Makefile
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
CUDA_HOME ?= /usr/local/cuda
CC ?= cc
CFLAGS ?= -O2
MIN_CUDA_VERSION ?= 13000
CPPFLAGS += -I$(CUDA_HOME)/include
CFLAGS += -fPIC -Wall -Wextra -Werror
LDFLAGS += -shared
LDLIBS += -ldl -pthread

TARGET := liblmcache_cumem_shareable.so
SOURCE := cumem_shareable_interposer.c

.PHONY: all clean check-cuda-version

all: check-cuda-version $(TARGET)

check-cuda-version:
@test -f "$(CUDA_HOME)/include/cuda.h"
@$(CC) $(CPPFLAGS) -dM -E -x c -include cuda.h /dev/null | \
awk -v minimum="$(MIN_CUDA_VERSION)" \
'/CUDA_VERSION/ { if ($$3 < minimum) exit 1; found=1 } END { exit !found }'

$(TARGET): $(SOURCE)
$(CC) $(CPPFLAGS) $(CFLAGS) $(LDFLAGS) -o $@ $< $(LDLIBS)

clean:
rm -f $(TARGET)
47 changes: 47 additions & 0 deletions csrc/cumem_ipc_interposer/cumem_shareable_interposer.c
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
// SPDX-License-Identifier: Apache-2.0
#define _GNU_SOURCE

#include <cuda.h>
#include <dlfcn.h>
#include <pthread.h>

/*
* vLLM owns allocation, mapping, sleep/wake, and release. This interposer only
* changes cuMemCreate's requested handle type so CUDA 13.3 allocations can be
* exported as POSIX FDs. In particular, replace driver-selected FABRIC handles:
* FABRIC handles cannot be transported through same-UID SCM_RIGHTS.
*/
typedef CUresult (*cuMemCreate_fn)(CUmemGenericAllocationHandle *, size_t,
const CUmemAllocationProp *,
unsigned long long);

static cuMemCreate_fn real_cuMemCreate;
static pthread_once_t resolve_once = PTHREAD_ONCE_INIT;

static void resolve_driver_symbol(void) {
/*
* CUDA may be loaded in an extension's local dependency scope, where
* RTLD_NEXT cannot see it. Resolve against the driver DSO explicitly.
*/
void *driver = dlopen("libcuda.so.1", RTLD_NOW | RTLD_LOCAL);
if (driver != NULL) {
real_cuMemCreate = (cuMemCreate_fn)dlsym(driver, "cuMemCreate");
}
}

CUresult cuMemCreate(CUmemGenericAllocationHandle *handle, size_t size,
const CUmemAllocationProp *prop,
unsigned long long flags) {
pthread_once(&resolve_once, resolve_driver_symbol);
if (real_cuMemCreate == NULL) {
return CUDA_ERROR_NOT_INITIALIZED;
}
if (prop == NULL ||
prop->requestedHandleTypes == CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR) {
return real_cuMemCreate(handle, size, prop, flags);
}

CUmemAllocationProp shareable = *prop;
shareable.requestedHandleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR;
return real_cuMemCreate(handle, size, &shareable, flags);
}
20 changes: 20 additions & 0 deletions docs/design/integration/vllm/hybrid-kv-cache-groups.md
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,26 @@ store is skipped and nothing is committed — a later retrieve simply misses and
the engine recomputes. The non-GPU transfer path rejects multi-group transfers
outright.

### Exact recurrent boundaries

Align-mode recurrent block tables are sparse and mutable, so their positional
entries are not authoritative store sources. vLLM core hands the connector the
exact committed block for each LMCache chunk boundary. The scheduler records
those handoffs by request, engine group, and boundary token count, then replaces
the recurrent group’s positional IDs before emitting a store.

The raw transfer shape remains unchanged: a recurrent group with 512-token
blocks and a 4096-token LMCache chunk receives eight IDs per chunk. The first
seven are null-block placeholders and the final ID is the exact boundary block.
This lets grouped atomic-store validation retain its full-coverage invariant
while the all-null-chunk logic continues to reject invalid recurrent objects.
If any boundary needed by a store has no exact handoff, the connector emits no
store for that boundary rather than caching a stale positional state. Handoffs
can trail scheduler token accounting by one step, so the connector emits the
longest contiguous prefix for which every recurrent group has an exact
handoff and defers the remaining suffix. This prevents each newly scheduled
suffix from suppressing an already-ready prefix throughout a chunked prefill.

## Example

vLLM exposes two engine groups — group 0: layers [0,2,4], group 1: [1,3]. If
Expand Down
66 changes: 66 additions & 0 deletions docs/design/v1/platform/cuda/cumem_ipc.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
# cuMem CUDA IPC

## Purpose

vLLM sleep-mode KV caches use CUDA virtual-memory-management allocations.
Ordinary `cudaIpcGetMemHandle` and PyTorch storage sharing do not safely export
those allocations. LMCache therefore uses a POSIX-FD-shareable cuMem path for
LMCache-driven transfers while preserving vLLM's ownership of the allocation.

## Allocation contract

The serving worker must preload
`csrc/cumem_ipc_interposer/liblmcache_cumem_shareable.so`. The interposer changes
only `CUmemAllocationProp.requestedHandleTypes` passed to `cuMemCreate`, forcing
`CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR` even when the driver selected a
FABRIC handle. Allocation, mapping, sleep/wake, and release remain owned by
vLLM and CUDA.

Build against the CUDA 13.3 toolkit:

```bash
make -C csrc/cumem_ipc_interposer CUDA_HOME=/usr/local/cuda-13.3
```

Launch each vLLM worker with the resulting library in `LD_PRELOAD`.

## Descriptor transport

`LMCACHE_CUMEM_BROKER_DIR` must point to a pre-existing directory that:

- is mounted at the same absolute path in the vLLM and LMCache containers;
- is owned by the UID shared by both processes;
- is writable/searchable by that UID; and
- is a real directory, not a symlink.

The exporter creates a mode-`0600` AF_UNIX socket directly below that root.
Pickled registration metadata carries the socket path, an allocation identity,
and a random capability token. The allocation FD itself crosses the socket with
`SCM_RIGHTS`; it is never serialized. A private `/tmp` path is rejected because
container mount namespaces can resolve the same text to different directories.

## Mapping and view contract

The sidecar imports one mapping per `(device UUID, allocation ID)`. Multiple HMA
tensor aliases increment references to that mapping and the final alias release
performs, in order:

1. stream/device synchronization;
2. `cuMemUnmap`;
3. `cuMemAddressFree`; and
4. `cuMemRelease`.

Tensor descriptors preserve allocation extent, allocation-relative storage
offset, physical storage extent, shape, stride, dtype, and device UUID. All view
metadata is validated before mapping and again before reconstruction.

Unregister drops tensor aliases and mappings, then runs allocator collection.
It deliberately does not call `cudaDeviceReset`; small CUDA contexts remain
available so the same sidecar PID can register fresh worker allocations.

## Transfer-mode fallback

In `auto` mode, an exporter failure before `REGISTER_KV_CACHE` is sent may
select engine-driven transfer. Explicit `lmcache_driven` mode fails closed.
After a registration exists, the context never changes transfer mode during a
request.
Loading
Loading