Optimizing Key-Value Cache Management in LLM Serving with PagedAttention and CUDA IPC — code snippets

snippet-1.py

import torch
from typing import List, Set, Dict

class PhysicalBlockAllocator:
    def __init__(
        self, 
        num_blocks: int, 
        block_size: int, 
        num_layers: int, 
        num_heads: int, 
        head_dim: int, 
        dtype: torch.dtype = torch.bfloat16, 
        device: str = "cuda"
    ):
        self.num_blocks = num_blocks
        self.block_size = block_size  # Number of tokens per block
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.dtype = dtype
        self.device = device
        
        # Pre-allocate global physical KV cache pools
        # Shape: [num_blocks, num_layers, num_heads, block_size, head_dim]
        self.k_cache = torch.empty(
            (num_blocks, num_layers, self.num_heads, block_size, head_dim),
            dtype=self.dtype,
            device=self.device
        )
        self.v_cache = torch.empty(
            (num_blocks, num_layers, self.num_heads, block_size, head_dim),
            dtype=self.dtype,
            device=self.device
        )
        
        self.free_blocks: Set[int] = set(range(num_blocks))
        self.block_table: Dict[str, List[int]] = {}

    def allocate(self, request_id: str, num_tokens: int) -> List[int]:
        num_required_blocks = (num_tokens + self.block_size - 1) // self.block_size
        if len(self.free_blocks) < num_required_blocks:
            raise MemoryError(
                f"OOM in KV Cache Pool. Free blocks: {len(self.free_blocks)}, "
                f"Requested: {num_required_blocks}"
            )
        
        allocated = [self.free_blocks.pop() for _ in range(num_required_blocks)]
        self.block_table[request_id] = allocated
        return allocated

    def free(self, request_id: str) -> None:
        if request_id in self.block_table:
            for block_id in self.block_table[request_id]:
                self.free_blocks.add(block_id)
            del self.block_table[request_id]

snippet-2.txt

#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cstring>

// Exports the raw CUDA memory address of a PyTorch tensor to a CUDA IPC handle
py::bytes export_tensor_ipc(torch::Tensor tensor) {
    TORCH_CHECK(tensor.is_cuda(), "Tensor must reside on a CUDA device");
    void* device_ptr = tensor.data_ptr();
    cudaIpcMemHandle_t handle;
    
    cudaError_t err = cudaIpcGetMemHandle(&handle, device_ptr);
    TORCH_CHECK(err == cudaSuccess, "cudaIpcGetMemHandle failed: ", cudaGetErrorString(err));
    
    return py::bytes(reinterpret_cast<const char*>(&handle), sizeof(handle));
}

// Opens a CUDA IPC handle sent from another process and wraps it in a PyTorch tensor
torch::Tensor import_tensor_ipc(py::bytes raw_handle, torch::IntArrayRef shape, torch::ScalarType dtype, int device_id) {
    cudaIpcMemHandle_t handle;
    std::memcpy(&handle, raw_handle.cast<std::string>().data(), sizeof(handle));
    
    void* device_ptr = nullptr;
    cudaError_t err = cudaIpcOpenMemHandle(&device_ptr, handle, cudaIpcMemLazyEnablePeerAccess);
    TORCH_CHECK(err == cudaSuccess, "cudaIpcOpenMemHandle failed: ", cudaGetErrorString(err));
    
    auto options = torch::TensorOptions().device(torch::kCUDA, device_id).dtype(dtype);
    // Wrap raw GPU pointer. Note: The lifetime of device_ptr must outlive this tensor!
    return torch::from_blob(device_ptr, shape, options);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("export_tensor_ipc", &export_tensor_ipc, "Export CUDA IPC Handle");
    m.def("import_tensor_ipc", &import_tensor_ipc, "Import CUDA IPC Handle");
}

snippet-3.txt

#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cstring>

// Exports a CUDA event to an IPC handle
py::bytes export_cuda_event_ipc(uintptr_t event_ptr) {
    cudaEvent_t event = reinterpret_cast<cudaEvent_t>(event_ptr);
    cudaIpcEventHandle_t handle;
    
    cudaError_t err = cudaIpcGetEventHandle(&handle, event);
    TORCH_CHECK(err == cudaSuccess, "cudaIpcGetEventHandle failed: ", cudaGetErrorString(err));
    
    return py::bytes(reinterpret_cast<const char*>(&handle), sizeof(handle));
}

// Imports a CUDA event IPC handle
uintptr_t import_cuda_event_ipc(py::bytes raw_handle) {
    cudaIpcEventHandle_t handle;
    std::memcpy(&handle, raw_handle.cast<std::string>().data(), sizeof(handle));
    
    cudaEvent_t event;
    cudaError_t err = cudaIpcOpenEventHandle(&event, handle);
    TORCH_CHECK(err == cudaSuccess, "cudaIpcOpenEventHandle failed: ", cudaGetErrorString(err));
    
    return reinterpret_cast<uintptr_t>(event);
}

PYBIND11_MODULE(TORCH_EVENT_EXTENSION_NAME, m) {
    m.def("export_cuda_event_ipc", &export_cuda_event_ipc, "Export CUDA Event IPC Handle");
    m.def("import_cuda_event_ipc", &import_cuda_event_ipc, "Import CUDA Event IPC Handle");
}

snippet-4.py

import socket
import struct
import json
from typing import Tuple, List

class MetadataIPCClient:
    def __init__(self, path: str):
        self.path = path
        self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)

    def connect(self):
        self.sock.connect(self.path)

    def send_kv_metadata_with_event(
        self, 
        request_id: str, 
        block_ids: List[int], 
        k_handle: bytes, 
        v_handle: bytes, 
        event_handle: bytes,
        shape: Tuple[int, ...]
    ):
        meta_dict = {
            "request_id": request_id,
            "block_ids": block_ids,
            "shape": shape
        }
        meta_bytes = json.dumps(meta_dict).encode('utf-8')
        
        # Header: [JSON meta length (4B), K length (4B), V length (4B), Event length (4B)]
        header = struct.pack("!IIII", len(meta_bytes), len(k_handle), len(v_handle), len(event_handle))
        
        # Send all data
        self.sock.sendall(header + meta_bytes + k_handle + v_handle + event_handle)

snippet-5.py

import torch
import cuda_ipc_extension
import cuda_event_extension
from typing import Tuple, List

class PrefillEngine:
    def __init__(self, allocator: PhysicalBlockAllocator, ipc_client: MetadataIPCClient):
        self.allocator = allocator
        self.ipc_client = ipc_client
        self.stream = torch.cuda.Stream()

    def process_request(self, request_id: str, prompt_tokens: torch.Tensor):
        with torch.cuda.stream(self.stream):
            # 1. Compute prompt KV cache
            prompt_k, prompt_v = self._run_forward_pass(prompt_tokens)
            
            num_tokens = prompt_tokens.shape[0]
            block_ids = self.allocator.allocate(request_id, num_tokens)
            
            # 2. Scatter into physical cache pool
            self._write_to_physical_cache(prompt_k, prompt_v, block_ids)
            
            # 3. Record CUDA Event to mark compute completion on this stream
            event = torch.cuda.Event(enable_timing=False, interprocess=True)
            event.record(self.stream)
            
            # 4. Generate CUDA IPC handles
            k_ipc_handle = cuda_ipc_extension.export_tensor_ipc(self.allocator.k_cache)
            v_ipc_handle = cuda_ipc_extension.export_tensor_ipc(self.allocator.v_cache)
            event_ipc_handle = cuda_event_extension.export_cuda_event_ipc(event.cuda_event)
            
            # 5. Broadcast to Decode Engine
            self.ipc_client.send_kv_metadata_with_event(
                request_id=request_id,
                block_ids=block_ids,
                k_handle=k_ipc_handle,
                v_handle=v_ipc_handle,
                event_handle=event_ipc_handle,
                shape=tuple(self.allocator.k_cache.shape)
            )
            
    def _run_forward_pass(self, tokens: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        seq_len = tokens.shape[0]
        device = tokens.device
        k = torch.randn((seq_len, self.allocator.num_layers, self.allocator.num_heads, self.allocator.head_dim), dtype=self.allocator.dtype, device=device)
        v = torch.randn((seq_len, self.allocator.num_layers, self.allocator.num_heads, self.allocator.head_dim), dtype=self.allocator.dtype, device=device)
        return k, v

    def _write_to_physical_cache(self, k: torch.Tensor, v: torch.Tensor, block_ids: List[int]):
        block_size = self.allocator.block_size
        num_tokens = k.shape[0]
        for idx, block_id in enumerate(block_ids):
            start = idx * block_size
            end = min(start + block_size, num_tokens)
            size = end - start
            self.allocator.k_cache[block_id, :, :, :size, :] = k[start:end].permute(1, 2, 0, 3)
            self.allocator.v_cache[block_id, :, :, :size, :] = v[start:end].permute(1, 2, 0, 3)

snippet-6.py

import socket
import struct
import json
import torch
import cuda_ipc_extension
import cuda_event_extension
from typing import Dict, List, Tuple

class DecodeEngine:
    def __init__(self, device_id: int):
        self.device_id = device_id
        self.request_block_tables: Dict[str, List[int]] = {}
        self.imported_caches: Dict[str, Tuple[torch.Tensor, torch.Tensor]] = {}
        self.stream = torch.cuda.Stream()

    def handle_new_request(self, client_sock: socket.socket):
        # 1. Parse header: [JSON len, K len, V len, Event len]
        header = client_sock.recv(16)
        if not header:
            return
        meta_len, k_len, v_len, ev_len = struct.unpack("!IIII", header)
        
        # 2. Read payloads
        meta_bytes = client_sock.recv(meta_len)
        k_handle = client_sock.recv(k_len)
        v_handle = client_sock.recv(v_len)
        ev_handle = client_sock.recv(ev_len)
        
        meta = json.loads(meta_bytes.decode('utf-8'))
        request_id = meta["request_id"]
        block_ids = meta["block_ids"]
        shape = meta["shape"]
        
        # 3. Import IPC handles
        k_cache = cuda_ipc_extension.import_tensor_ipc(
            k_handle, shape, torch.bfloat16, self.device_id
        )
        v_cache = cuda_ipc_extension.import_tensor_ipc(
            v_handle, shape, torch.bfloat16, self.device_id
        )
        
        # 4. Import the CUDA Event and synchronize the stream
        raw_event_ptr = cuda_event_extension.import_cuda_event_ipc(ev_handle)
        py_event = torch.cuda.Event(_cuda_event=raw_event_ptr)
        
        with torch.cuda.stream(self.stream):
            # Instruct decode stream to block until the prefill event triggers
            self.stream.wait_event(py_event)
            
            # Save mapping references
            self.imported_caches[request_id] = (k_cache, v_cache)
            self.request_block_tables[request_id] = block_ids
            
    def release_request(self, request_id: str):
        if request_id in self.imported_caches:
            del self.imported_caches[request_id]
        if request_id in self.request_block_tables:
            del self.request_block_tables[request_id]

snippet-7.yaml

apiVersion: apps/v1
kind: Deployment
metadata:
  name: decoupled-llm-serving
  labels:
    app: decoupled-llm
spec:
  replicas: 1
  selector:
    matchLabels:
      app: decoupled-llm
  template:
    metadata:
      labels:
        app: decoupled-llm
    spec:
      hostIPC: true  # CRITICAL: Shares host IPC namespace so CUDA drivers can resolve memory pointers
      containers:
        - name: prefill-engine
          image: llm-serving:latest
          command: ["python", "prefill_server.py"]
          resources:
            limits:
              nvidia.com/gpu: "1"
          volumeMounts:
            - name: dshm
              mountPath: /dev/shm
            - name: shared-uds
              mountPath: /var/run/llm
        - name: decode-engine
          image: llm-serving:latest
          command: ["python", "decode_server.py"]
          resources:
            limits:
              nvidia.com/gpu: "1"
          volumeMounts:
            - name: dshm
              mountPath: /dev/shm
            - name: shared-uds
              mountPath: /var/run/llm
      volumes:
        - name: dshm
          emptyDir:
            medium: Memory
            sizeLimit: 16Gi
        - name: shared-uds
          emptyDir: {}