HyperOffload: Graph-Driven Hierarchical Memory Management for Large Language Models on SuperNode Architectures

TL;DR

HyperOffload addresses the memory bottleneck of large language models by elevating remote memory accesses to first-class graph operators within MindSpore on SuperNode hardware. It introduces cache operators (Prefetch, Store, Detach) and a Graph-Driven Execution-Order Optimization to deterministically plan memory movements and hide remote latency behind computation. The framework integrates non-intrusively via MindIR and JIT graph rewriting, enabling compile-time tensor lifetime analysis and memory planning without requiring model changes. Empirical results show substantial peak memory reductions and stable performance across training and long-context inference, demonstrating the practicality of graph-level, compiler-aware memory management for next-generation AI workloads on memory-rich architectures.

Abstract

The rapid evolution of Large Language Models (LLMs) towards long-context reasoning and sparse architectures has pushed memory requirements far beyond the capacity of individual device HBM. While emerging supernode architectures offer terabyte-scale shared memory pools via high-bandwidth interconnects, existing software stacks fail to exploit this hardware effectively. Current runtime-based offloading and swapping techniques operate with a local view, leading to reactive scheduling and exposed communication latency that stall the computation pipeline. In this paper, we propose the SuperNode Memory Management Framework (\textbf{HyperOffload}). It employs a compiler-assisted approach that leverages graph-driven memory management to treat remote memory access as explicit operations in the computation graph, specifically designed for hierarchical SuperNode architectures. Unlike reactive runtime systems, SuperNode represents data movement using cache operators within the compiler's Intermediate Representation (IR). This design enables a global, compile-time analysis of tensor lifetimes and execution dependencies. Leveraging this visibility, we develop a global execution-order refinement algorithm that statically schedules data transfers to hide remote memory latency behind compute-intensive regions. We implement SuperNode within the production deep learning framework MindSpore, adding a remote memory backend and specialized compiler passes. Evaluation on representative LLM workloads shows that SuperNode reduces peak device memory usage by up to 26\% for inference while maintaining end-to-end performance. Our work demonstrates that integrating memory-augmented hardware into the compiler's optimization framework is essential for scaling next-generation AI workloads.

Figures (6)

• Figure 1: Motivation of graph-level remote memory scheduling. Runtime-driven data movement leads to execution bubbles and increased peak device memory usage. By explicitly scheduling communication at the graph level, we enable efficient overlap between computation and remote memory access, allowing large models to run beyond device memory limits transparently.
• Figure 2: SuperNode Hardware Architecture. Compared to (a) traditional SIMT-based GPUs with per-device HBM, (b) SuperNode architectures integrate tile-based NPUs interconnected by ultra-high-bandwidth links and a large shared memory pool. The shared memory pool provides a unified address space, enabling symmetric access to model parameters, KV caches, and activations across all NPUs within the SuperNode.
• Figure 3: Comparison of execution timelines under different compute–communication orchestration strategies. (a) Serial execution without data prefetching, where NPU computation and DMA transfers are fully serialized. (b) Runtime-driven compute–communication overlap, in which the CPU orchestrates DMA prefetching during NPU execution, introducing runtime-induced bubbles. (c) Ideal fully overlapped execution, where DMA prefetching and NPU computation are statically orchestrated ahead of time, achieving bubble-free parallelism.
• Figure 4: Comparison of Communication Overlap Strategies in Heterogeneous Offloading Systems.
• Figure 5: Hierarchical memory interface usage example. The figure demonstrates how model parameters are placed in remote memory and how cache prefetch, store, and detach operations are exposed to users. Both explicit API calls and JIT-based static graph annotations are shown, illustrating flexible and unified interfaces for hierarchical memory management.