Unified Virtual Memory (UVM)

The foundational characteristic of UVM is the creation of a unified virtual address space that spans both CPU system memory and GPU device memory. This allows pointers to data to be valid and meaningful on both the CPU and GPU, regardless of the data's physical location. This eliminates the need for explicit, manual memory copies (cudaMemcpy) between host and device for data movement, as the hardware and driver can manage migration transparently.

• Simplified Programming: Developers can write code that references data using a single pointer, as if working with a single, large memory pool.
• Key Use Case: Essential for implementing KV cache offloading, where parts of the key-value tensor cache can seamlessly reside in CPU RAM while being accessed by the GPU, without complex manual buffer management.

UVM systems implement on-demand page migration, a mechanism analogous to virtual memory paging in operating systems. When a GPU (or CPU) attempts to access a memory page that resides in the other processor's physical memory, a page fault occurs. The UVM driver then transparently migrates the requested page to the local memory of the accessing processor.

• Efficiency: Data is moved only when needed, optimizing bandwidth usage.
• Granularity: Migration happens at the page level (e.g., 4KB, 64KB), allowing fine-grained movement of data. For KV cache management, this means individual blocks of the cache can be swapped between GPU and CPU memory based on the attention pattern of the decoding process.

A critical capability enabled by UVM is memory oversubscription. The total size of allocated UVM memory can exceed the physical capacity of the GPU's high-bandwidth memory (HBM). The system uses the CPU's RAM and, if necessary, even NVMe storage as a slower, larger backing store.

• Enables Longer Contexts: This directly supports long-context inference by allowing the KV cache for thousands of tokens to be allocated, even if it doesn't all fit in GPU memory at once.
• Performance Trade-off: Accessing data in CPU RAM is significantly slower than GPU HBM, creating a latency/throughput trade-off that must be managed by scheduling and access patterns.

UVM provides coherence guarantees between the CPU and GPU caches. When the CPU writes to a UVM-managed region, subsequent GPU reads will see the updated data, and vice-versa. This is typically achieved through a software-managed coherence model in implementations like CUDA UVM, where explicit synchronization calls (cudaDeviceSynchronize, cudaStreamSynchronize) enforce a consistent view of memory.

• Deterministic Behavior: Ensures correctness when both processors are manipulating shared data structures, such as metadata for managing a partitioned KV cache.
• Developer Responsibility: Unlike fully coherent CPU multi-core systems, the programmer must insert appropriate synchronization points to avoid race conditions.

UVM introduces memory allocation functions (e.g., cudaMallocManaged) that return pointers usable by both CPU and GPU code. These unified pointers simplify code structure. The allocation also comes with hints (e.g., cudaMemAdvise) that allow developers to guide the UVM driver's migration and prefetching policies.

• Common Allocation API: Replaces separate malloc and cudaMalloc calls.
• Performance Optimization: Advising the system that data will be accessed mostly by the GPU (cudaMemAdviseSetPreferredLocation) or that it should be prefetched (cudaMemAdviseSetAccessedBy) can significantly reduce page fault latency, a crucial optimization for managing the streaming access patterns of a KV cache during decoding.

UVM is the enabling hardware/software layer for efficient KV cache offloading strategies. The KV cache can be allocated as a unified memory region. Hot, actively accessed cache entries (e.g., for recent tokens and attention sinks) reside in GPU HBM, while colder, older portions of the context can be transparently paged out to CPU RAM.

• Transparent Management: The attention computation kernels access the KV cache via unified pointers; the UVM system handles the physical location.
• Enables New Optimizations: Allows systems like vLLM to think in terms of virtual cache blocks, which can be non-contiguously mapped to physical memory pages on either CPU or GPU, similar to its PagedAttention concept but across a memory hierarchy.

UVM provides a unified virtual address space across CPU and GPU memory. This allows both the CPU and GPU to access any memory location using the same pointer, managed by the system's memory management unit (MMU) and page tables. The key components enabling this are:

• Page Fault Handling: The GPU can generate a page fault if data is not resident in its local memory, triggering a page migration from CPU RAM or system storage.
• Coherency Protocols: Hardware and software maintain memory coherency, ensuring all processors see a consistent view of data without explicit programmer synchronization.
• On-Demand Paging: Memory pages are moved between CPU and GPU memory only when accessed, enabling efficient use of total system memory capacity.

UVM is critical for KV cache offloading, a technique to handle long-context inference. When a model's key-value cache exceeds GPU memory, UVM allows it to spill over to CPU RAM or NVMe storage. The process involves:

• Transparent Access: The GPU's attention kernels access the KV cache via virtual addresses. If a required block is not in GPU memory, a page fault occurs.
• Background Migration: The UVM driver migrates the required page from CPU memory to GPU memory, allowing the computation to proceed. Less-recently-used pages may be evicted back to the CPU.
• Performance Trade-off: This exchanges higher I/O latency for vastly increased effective cache capacity, enabling context windows of millions of tokens without model architecture changes.

CUDA Unified Virtual Memory is NVIDIA's implementation, a cornerstone for large model inference on their GPUs. Key features include:

• Managed Memory: Allocated via cudaMallocManaged(), this memory is automatically migrated by the CUDA driver.
• Access Counters: The driver uses counters to track page usage, implementing policies like prefetching and data locality optimization.
• GPU Direct RDMA: Supports direct peer-to-peer access and Remote Direct Memory Access (RDMA) over NVLink and InfiniBand, allowing efficient memory sharing across multiple GPUs and nodes.
• Integration with Libraries: Deeply integrated with cuBLAS, cuDNN, and frameworks like TensorRT-LLM, which use UVM for efficient KV cache management and continuous batching.

UVM requires specific hardware and software support to function efficiently:

• GPU Architecture: Requires Pascal (SM 6x) or newer NVIDIA GPUs for full hardware support. Older architectures use slower software-emulated modes.
• CPU-GPU Interconnect: Performance depends on PCIe bandwidth (Gen3/4/5). NVLink provides significantly higher bandwidth for page migrations between CPU and GPU.
• Operating System: Requires driver and OS support for handling GPU page faults. Linux is the primary production environment.
• Page Size: Modern systems support large pages (e.g., 2MB, 1GB) to reduce page table overhead and TLB misses during massive KV cache access.

While UVM simplifies programming, it introduces specific performance characteristics that engineers must design for:

• Latency vs. Capacity: The primary trade-off. Accessing CPU memory can be ~5-10x slower than GPU HBM. NVMe offload can be ~100-1000x slower.
• Page Fault Overhead: The first access to a non-resident page incurs a high latency penalty (tens of microseconds). This makes access patterns critical.
• Prefetching Strategies: Proactive prefetching algorithms are essential to hide migration latency. This involves predicting which KV cache blocks will be needed next based on attention patterns.
• Coherency Overhead: Maintaining coherency across devices adds minor protocol overhead, but is negligible compared to page migration costs.

Modern high-performance inference engines leverage UVM as a foundational layer for advanced memory management:

• vLLM with PagedAttention: While vLLM's PagedAttention manages the KV cache in software blocks, it can use UVM at a lower level to allocate and migrate those blocks between GPU and CPU memory.
• TensorRT-LLM: Uses CUDA UVM for dynamic KV cache allocation and to support its in-flight batching capabilities, allowing seamless handling of variable-length sequences.
• Orchestration with Cache-Aware Scheduling: Schedulers can group requests with similar KV cache footprints to improve locality of reference, minimizing thrashing and page faults within the UVM system.

A memory management technique that moves portions of the key-value cache from GPU memory to cheaper, higher-capacity storage like CPU RAM or NVMe SSDs. This is a primary use case for UVM, as it provides the single, coherent address space required to seamlessly access cached tensors across the CPU-GPU boundary. Offloading trades increased I/O latency for dramatically greater effective cache capacity, enabling longer context windows.

A memory management algorithm, popularized by vLLM, that organizes the KV cache into non-contiguous, fixed-size blocks or pages. This is directly analogous to virtual memory paging in operating systems. UVM provides the hardware/software foundation that makes PagedAttention efficient by handling the address translation and data movement between physical memory locations (GPU, CPU, disk) transparently to the inference engine.

A feature enabled by UVM where data resides in a single physical location accessible by both CPU and GPU, eliminating the need for explicit memcpy operations across the PCIe bus. For KV cache management, this means the CPU can directly prepare or manipulate cache tensors in a buffer that the GPU kernels can immediately access, minimizing latency and overhead for techniques like cache prefetching or partial offloading.

The practice of allocating more memory to processes than the physical GPU memory (VRAM) can hold. UVM is essential for this, as it allows the driver to page less-frequently accessed data (e.g., older parts of a long KV cache) to system RAM. This enables higher batch sizes or longer contexts than the VRAM would normally allow, though performance depends on the page fault handling speed.

An event that occurs when a processor (CPU or GPU) tries to access a memory page that is not resident in its local physical memory. In a UVM system managing a KV cache, a GPU page fault triggers a stall while the required page is fetched from CPU RAM or SSD. The performance of KV cache offloading is heavily dependent on minimizing the frequency and cost of these faults through prefetching algorithms and cache-aware scheduling.