Paper Notes: DirectKV

Citation

The paper named “No Buffer, No Bottleneck: Efficient Zero-Copy KV Cache Offloading for Long-Context LLMs” is authored from University of Virginia. Published at OSDI 2026.

Motivation

Large Language Models (LLMs) generate a KV cache that grows linearly with sequence length, quickly exceeding GPU HBM (High Bandwidth Memory) capacity. This forces systems to either limit context length or offload the KV cache to CPU memory.

Existing KV cache offloading systems (FlexGen, Neo, Pie, etc.) share a fundamental inefficiency: they follow a “separate kernel” pattern:

1. K/V Projection kernel → writes K/V to CPU memory
2. Later, Attention kernel → re-fetches K/V from CPU memory back to GPU
3. This creates redundant CPU–GPU transfers of the same KV data

This is a “staging” approach: KV blocks are shuttled back and forth between CPU and GPU, wasting precious interconnect bandwidth and increasing latency.

The obvious fix would be zero-copy — let the GPU directly read KV cache from CPU-pinned memory without staging. But naïve zero-copy performs terribly because the CPU–GPU interconnect bandwidth (e.g., NVLink-C2C at ~450 GB/s per direction) is far lower than HBM bandwidth (~4 TB/s per direction). Conventional matrix multiplication tiling assumes all operands are in GPU memory and repeatedly re-fetches operands from wherever they reside — naïvely pointing attention kernels at CPU-resident KV tensors creates a bandwidth bottleneck at the CPU–GPU interconnect, stalling compute.

DirectKV asks: Can we make zero-copy practical by redesigning attention kernels to shift bandwidth pressure away from the CPU–GPU interconnect onto HBM? This motivates their kernel–memory co-design approach: CPU-memory-aware tiling (reuse CPU-fetched KV tiles maximally to shift data movement to HBM), kernel fusion (fuse K/V projection and attention into a single CUDA kernel to eliminate redundant writes and re-fetches), and warp-level pipelining (overlap communication with computation to hide memory latency).

Background

Target hardware is the NVIDIA Grace-Hopper (GH200) superchip: a tightly coupled CPU-GPU architecture where the GPU connects to CPU memory via NVLink-C2C (~450 GB/s per direction). While much faster than PCIe, this bandwidth is still ~10× lower than HBM (~4 TB/s). The chip also provides zero-copy memory access — the GPU can directly read CPU-pinned memory without explicit cudaMemcpyAsync calls.

In standard LLM inference, each transformer layer produces key (K) and value (V) tensors during the attention step. These are stored in a KV cache to avoid recomputing them for every subsequent token. The KV cache size = 2 × num_layers × hidden_dim × sequence_length. For a Llama-8B model at 32K context, this alone can consume tens of GB — pushing GPU memory to its limit.

System Architecture

DirectKV consists of four components that work together across both prefill (processing input prompts) and decode (generating one token at a time) phases:

Kernel Generator (Offline). Compiles a rich set of CUDA kernel candidates ahead of time using C++ template instantiations specialized for different combinations of <data_type, head_dim, tile_size>. This avoids runtime compilation overhead while ensuring each variant is fully optimized. All candidates are stored in an offline repository.

Kernel Adaptor (Runtime). At inference time, selects the best pre-compiled kernel by matching the current request’s precision (FP16/BF16), head dimension, tile size, and execution phase (prefill: multi-token vs. decode: single-token). The selection is lightweight — just substituting template parameters.

Attention Fusion Engine. Fuses K/V projection and attention score computation into a single CUDA kernel launch. This is where the three key ideas (CPU-aware tiling, kernel fusion, warp-level pipelining) execute. It has distinct strategies for prefill vs. decode:

• Prefill phase: iterate over Q (stored in HBM) for each (K,V) tile fetched from CPU memory — since KV cache lives in slow CPU memory, reuse it maximally.
• Decode phase: iterate over all cached (K,V) for the single new query token — since Q is only one token, avoid repeatedly reading intermediate output from HBM.

KV Cache Manager. Allocates pinned host memory (cudaHostAlloc) for KV tensors, making them directly accessible to the GPU via zero-copy pointers. Once written, KV tensors stay in these buffers and are reused across decoding iterations — no staging, no recomputation.

A key constraint driving kernel design is SMEM capacity. Each SM’s shared memory is logically split into projection buffers (for X, Wk, Wv tiles) and attention buffers (for K, V, Q, O tiles). To save space, buffers for Wk and Wv are reused to store newly generated K and V. SMEM allocation must satisfy α·P ≥ 3·m·size(T)·(Dim·N), where P is the L1/SMEM pool size, m is pipeline stages (default 2), and N is the tile size. This constraint determines which <T, Dim, N> configurations are feasible, guiding offline kernel pre-building.

Key Ideas

DirectKV proposes three techniques that work together to make zero-copy KV cache offloading practical:

① CPU-Memory-Aware Tiling. Standard matrix multiplication tiling assumes all tensors are in GPU HBM and re-fetches both A and B tiles repeatedly. DirectKV flips this pattern: when a tile comes from slow CPU memory (like KV cache), it is treated as stationary — fetched once and reused across all inner-loop iterations. The extra reload cost is pushed onto HBM-resident tensors, where bandwidth is 10× higher. This reduces CPU–GPU transfer volume by up to 50%.

② Kernel Fusion (Projection + Attention). Instead of launching separate kernels for K/V projection and attention, DirectKV fuses them into a single CUDA kernel. Generated K/V tiles stay in shared memory (SMEM) and are immediately consumed by the attention computation — no redundant writes to CPU memory, no re-fetches. The fused design sustains up to 3.5× higher HBM throughput and 2.5–3.0× lower latency compared to separate kernels.

③ Warp-Level Pipelining. Within each fused kernel, warp groups are specialized into producer, consumer, and storer roles. While one warp group computes on the current tile, another prefetches the next tile from memory, and a third writes results to CPU. This overlap hides memory latency and keeps compute units busy. The Hopper GPU’s Tensor Memory Accelerator (TMA) further accelerates the async data movement.

These three ideas share a common philosophy: shift bandwidth pressure from the CPU–GPU interconnect onto HBM, where the extra traffic can be absorbed without becoming a bottleneck.

Experiments & Results

Setup. NVIDIA GH200 (Hopper GPU, 96 GB HBM3, NVLink-C2C), CUDA 12.4, PyTorch 2.3. Models: Llama-3.1-8B, OPT-13B, OPT-30B. Datasets: ShareGPT and Alpaca, Poisson arrivals up to 30 req/s, context lengths 1K–32K tokens. Baselines: SGLang (HBM-resident), Pie (swap-based offloading), Neo (CPU-GPU pipelining), FlexGen (multi-tier offloading with compression).

End-to-end latency (Fig. 10). DirectKV achieves the lowest latency among all offloading systems. On Llama-8B at 30 req/s: 0.75s vs. 1.55–2.95s for baselines. When SGLang runs out of memory (OPT-13B and OPT-30B at high load), DirectKV continues to serve while maintaining low latency.

Context length scaling (Fig. 11a). DirectKV maintains the lowest latency across all sequence lengths. At 32K tokens, Neo, Pie, and SGLang OOM; DirectKV remains efficient while FlexGen incurs much higher latency due to multi-tier offloading overhead.

GPU memory (Fig. 11b). DirectKV uses 47 GB GPU memory on average — 43% less than Neo (86 GB), Pie (88 GB), and FlexGen (74 GB). SGLang consumes 92 GB and OOMs at long contexts. Savings come from storing KV cache in CPU-pinned memory accessed via zero-copy, eliminating staging buffers.

CPU-aware tiling (Fig. 12). Naïve zero-copy repeatedly fetches KV from CPU memory. CPU-aware tiling reorganizes access to maximize reuse of each fetched KV block, reducing CPU–GPU transfer volume by up to 50% and latency by up to 70%.

Fused kernel (Fig. 13). Fusing projection and attention into one kernel sustains up to 3.5× higher HBM throughput and 2.5–3.0× lower latency vs. separate kernels.

NVLink-C2C vs. PCIe (Fig. 14). NVLink-C2C reduces attention latency by up to 4.2× compared to PCIe. DirectKV benefits most from the higher bandwidth; competing systems remain bottlenecked by staging overhead.

Fundamental Underlying Technologies

Kernel Fusion (Projection + Attention)

Kernel fusion is a classic GPU optimization: instead of launching separate kernels that each write their results to global memory (HBM) and have the next kernel read them back, multiple operations are merged into a single kernel that passes data through on-chip shared memory (SMEM).

In the context of LLM inference, the standard path is:

1. Projection kernel: K = X·Wk, V = X·Wv → write K, V to HBM
2. Attention kernel: read K, V from HBM → compute softmax(Q·Kᵀ/√d)·V

This means K and V make a round-trip through HBM: written by kernel 1, read back by kernel 2. HBM bandwidth (hundreds of GB/s to a few TB/s) is already the scarcest resource on a GPU, and this redundant traffic consumes it for no computational benefit.

How fusion helps. By merging projection and attention into one kernel, K and V tiles are generated in registers/SMEM and immediately consumed by the attention computation — they never leave the on-chip memory hierarchy. This eliminates the HBM write-read round-trip for K/V, substantially reducing memory traffic and latency. The fused design typically sustains higher HBM throughput (since freed bandwidth can be used for useful data movement) and delivers lower latency (no kernel launch overhead between projection and attention, no redundant data movement).

Prerequisite: shared memory (SMEM). Fusion is only feasible when the working set of both fused operations fits within SMEM. Each SM’s SMEM is limited (e.g., ~228 KB on Hopper architectures), so tile sizes must be chosen so that projection buffers (X, Wk, Wv) and attention buffers (K, V, Q, O) together fit. This is a key constraint that guides kernel design — fatter tiles improve compute efficiency but risk SMEM overflow. Buffer reuse (e.g., overwriting Wk/Wv buffers with freshly computed K/V) helps reduce footprint.

Broader context. Kernel fusion is widely used in high-performance attention implementations (FlashAttention, xFormers, etc.) and is not specific to any one paper. It is a general technique that applies whenever multiple GPU kernels touch the same intermediate data, and is especially impactful in memory-bound workloads like LLM inference.