DAK: Direct-Access-Enabled GPU Memory Offloading with Optimal Efficiency for LLM Inference

2026-04-28

AI Summary

Large language model inference is constrained by GPU memory capacity and bandwidth, and existing offloading approaches rely on prefetching, which incurs HBM contention, memory waste, and pipeline bubbles. This work proposes the first end-to-end offloading framework that enables GPUs to directly access remote memory. It leverages the Tensor Memory Accelerator (TMA) to asynchronously fetch weights and KV caches directly into shared memory, combined with a greedy algorithm to optimize offloading ratios, proactive congestion control, and TMA multicast. These techniques collectively eliminate interconnect bottlenecks and read amplification. The framework achieves up to 3× speedup on NVLink-C2C systems and 1.8× on PCIe systems, approaching the theoretical aggregate bandwidth limit.

Abstract

LLM inference is constrained by GPU memory capacity and bandwidth. Tiered memory architectures mitigate this by allowing the GPU to offload memory to the remote tier. However, existing memory offloading frameworks rely on prefetching data into local GPU HBM. This approach underutilizes system resources by introducing HBM contention, squandering memory capacity, and creating pipeline bubbles. We show that enabling direct GPU access to remote memory significantly outperforms prefetching, achieving optimal aggregate system bandwidth. We propose DAK, an end-to-end direct-access memory offloading framework that repurposes the Tensor Memory Accelerator (TMA) to asynchronously fetch offloaded weights and KV caches directly from remote memory into GPU shared memory (SMEM). To maximize remote access performance, DAK introduces a greedy algorithm to determine optimal per-operation offloading ratios, alongside active congestion control and TMA multicast to eliminate interconnect bottlenecks and read amplification. Evaluations across diverse architectures show that DAK achieves near-optimal bandwidth aggregation, with up to 3$\times$ performance gains on NVLink-C2C and 1.8$\times$ on PCIe systems compared to state-of-the-art memory offloading baselines.

Problem

• LLM inference
• GPU memory offloading
• memory bandwidth
• tiered memory
• resource underutilization

Innovation

• direct-access memory offloading
• Tensor Memory Accelerator (TMA)
• LLM inference
• remote memory
• bandwidth aggregation