No Buffer, No Bottleneck: Efficient Zero-Copy KV Cache Offloading for Long-Context LLMs

Let zero-copy CPU offloading actually work with tiling

Featured image

Topic

Let zero-copy CPU offloading actually work with tiling Carve SMEM of GPU into buffer to save the data from CPU. Use tiling -> leverage CPU side’s data locality -> less fetch from CPU

Background

Prior work: CPU offloading

[offload KV cache to CPUs: need to stage KV cache into GPU memory] -> wasting GPU memory and doubling CPU–GPU transfers -> creates capacity and bandwidth bottlenecks.

  1. Swap-based approaches
    • place KV blocks in CPU memory (or SSD)
    • use large GPU staging buffers to overlap transfers with computation.
  2. store KV in remote or disaggregated memory and fetch it on demand
    • Still swap-like model: KV must be staged into GPU memory before attention.

[offloads computation & KV cache to CPUs]

  1. partially offloads attention computation and KV cache to CPUs, or moves the entire attention computation onto CPUs to avoid repeated PCIe transfers -> sacrifices the full compute capability of the GPU / due to the bandwidth gap vs. CPU-GPU interconnects <-> HBM

[Zero copy offloading / GPU kernels directly access CPU-resident KV tensors] Zerocopy eliminates staging buffers and reduces HBM usage. Zero-Copy eliminates the need for a large GPU buffer, reducing memory usage by 50% compared to Baseline and by 33% compared to Swap.

However, The effectiveness of zero-copy primarily depends on CPU-GPU bandwidth. CPU–GPU bandwidth becomes the dominant bottleneck.

Problem 1 : Prior CPU offloading leaves GPU compute units underutilized

Prior works remain fundamentally limited by the large bandwidth gap /between CPU–GPU interconnects (40–60 GB/s over PCIe) and HBM (3–4 TB/s) -> leaves GPU compute units underutilized and exacerbates the GPU memory wall

Even with NVLink-C2C, a significant bandwidth gap remains between the CPU–GPU interconnect (900 GB/s bidirectional) <-> HBM bandwidth (4 TB/s).

Problem 2: but naive application performs poorly with zero copy.

Solution: Kernel-memory codesign

DirectKV: first zerocopy KV cache offloading system for modern heterogeneous CPU–GPU platforms

root cause of the reason why naive zerocopy underperforms : repeated CPU-memory fetches & poor L2 locality -> introduce three optimizations:

1. CPU-memory-aware tiling to shift the bottleneck to HBM

using shared memory to stream CPU-resident tensors efficiently and reuse them across tiles./ smart tiling, considering access pattern

shift bandwidth pressure from the CPU–GPU interconnect to HBM tiling through SMEM shared memory (SMEM) within each streaming multiprocessor (SM) can be leveraged to make zero-copy practical. Acting as a fast on-chip buffer, SMEM enables flexible tiling strategies that restructure data access in matrix multiplication.

SMEM is carved from the unified L1/SMEM pool (e.g., 256 KB on GH200), we conservatively reserve a fraction α (80% by default) for SMEM and leave the remainder to L1 cache. @ isn’t it reduce L1 cache size then?

CPU-memory-aware Zero copy -> less CPU-GPU Transfer -> lower latency

But the efficiency of zero-copy offloading is fundamentally limited by the bandwidth of the CPU–GPU interconnect.

2. warp-level pipelining to overlap transfer and compute

while one set of warps fetches the next tiles from memory, others concurrently compute on the current tiles.

-> warp-level pipelining sustains higher HBM throughput & reduce latency

SM threads can be divided into multiple warp groups, with each group independently handling tasks such as computation or communication.

3. kernel fusion to eliminate redundant KV refetching

fuse the projection and attention score computation into a single CUDA kernel to improve the throughput of data transfer redundant writes of K back to CPU memory and the subsequent re-fetching in the next kernel -> higher HBM throughput & lower latency by improving memory locality and bandwidth utilization.

kernel fusion of KV projection and attention /that keeps K/V tensors in SMEM rather than round-tripping them through CPU memory.

In conventional implementations, projection and attention score computation are launched as separate CUDA kernels. -> This separation forces the tensors generated by the projection kernel to be written back to CPU memory (to update the KV cache) and then re-read by the attention kernel creating redundant CPU–GPU transfers exacerbating the bandwidth gap between NVLink-C2C(900GB/s aggregate,450GB/s per direction) and HBM (4 TB/s per direction).

the generated K and V tiles are retained in SMEM and immediately consumed by the subsequent attention computation, eliminating redundant writes to and reads from CPU memory.

The efficiency of zero-copy offloading is fundamentally limited by the bandwidth of the CPU–GPU interconnect. -> leverage NVLink-C2C instead of PCIe -> faster CPU-GPU interconnect

Evaluation

DirectKV

  1. reduces CPU–GPU transfer volume by up to 50% : thanks to cpu memory aware zero copy
  2. cuts GPU memory usage by 43% : zero copy
  3. improves end-to-end performance by up to 1.2× compared to existing solutions.

Thoughts & Questions