Weight-Context-Switch Primitive for Local LLM Memory Management

Overview

This research analyzes the failure of traditional operating system (OS) process-state management when applied to the gigabyte-scale memory footprints of local LLMs. Traditional OS paging is designed for general-purpose workloads, but LLM inference is dominated by the Key-Value (KV) cache—a dynamic state that grows linearly with sequence length and batch size [C000, C003]. For a model like Llama-3-70B in FP16, the KV cache consumes approximately 160 KB per token ^[C009]. In long-context deployments (100K to 10M tokens), this cache becomes the primary driver of GPU-hour waste and system instability ^[C004].

The objective is to implement a "Weight-Context-Switch" primitive. This primitive treats model weights and KV caches as a single, kernel-swappable resource on Unified Memory Architecture (uma) hardware, such as NVIDIA Grace Hopper, which utilizes a 900 GB/s NVLink-C2C interconnect to create a shared address space between CPU and GPU ^[C008]. This approach addresses the structural failures of current memory orchestration:

This shift is critical because the "memory wall" has moved from static model weights—which are now mitigated by quantization ^[C001]—to the dynamic KV cache ^[C000]. Without a first-principles primitive for switching these resources, the system remains trapped between the binary choice of maximizing "intelligence" (larger context windows) and maintaining system availability (preventing OOM crashes) [C002, C006].

Landscape

Current approaches to LLM memory orchestration focus on mitigating the "memory wall" created by the Key-Value (KV) cache, which grows linearly with sequence length and batch size ^[C004]. Because of the high memory requirements for models like Llama-3-70B ^[C008], the industry has split into three primary architectural directions: virtualized paging, unified memory hardware, and stateful system services.

Virtualized Memory Management

The dominant software approach, led by vLLM, utilizes pagedattention to treat GPU memory like operating system RAM ^[C000]. By paging KV cache blocks on demand, it eliminates the fragmentation caused by pre-allocating contiguous memory for the maximum sequence length ^[C009]. However, this introduces a "virtual memory tax" because the KV cache becomes non-contiguous in both physical and virtual memory, increasing pointer-chasing overhead ^[C007]. To resolve this, vAttention decouples virtual and physical allocation using CUDA APIs, maintaining virtual contiguity to allow the use of standard attention kernels while still preventing physical fragmentation ^[C007].

Unified Memory and hardware Acceleration

hardware-level solutions aim to eliminate the latency of explicit data transfers between CPU and GPU. NVIDIA Grace Hopper and Grace Blackwell architectures use NVLink C2C (900 GB/s) to create a single unified memory address space ^[C008]. This allows the GPU to access CPU memory directly for KV cache offloading without redundant copies, bypassing the bottleneck of PCIe Gen 5 ^[C008]. Parallel to this, heterogeneous prototypes like TFLOP utilize CPU-FPGA systems with 4-bit product quantization on both weights and KV caches to reduce the memory footprint during the decoding stage ^[C001].

Stateful Context Switching

For edge and mobile deployments, the challenge shifts from throughput to system availability. The LLMS framework treats LLMs as a stateful system service, decoupling app and LLM contexts to prevent the Android Low Memory Killer (LMK) from terminating processes ^[C002]. It implements Tolerance-Aware Compression and IO-Recompute Pipelined Loading to reduce context switching latency by up to two orders of magnitude ^[C002].

Comparative Memory Orchestration Strategies

The shift toward uma hardware and stateful swapping suggests that treating the KV cache as a kernel-swappable resource—rather than a transient application buffer—is required to support gigabyte-scale local footprints without triggering system-level crashes [C002, C004].

Key Findings

The primary driver of system instability and GPU-hour waste in LLM serving has shifted from model weight fitting to KV-cache dynamics [C000, C004]. At scale, this renders traditional OS process-state management ineffective, as KV cache overflows trigger significant context-switching latency when swapping from HBM to DDR ^[C004].

Research reveals a critical tension between memory flexibility and kernel throughput. While PagedAttention eliminates fragmentation, it forces a non-contiguous virtual memory layout that introduces a "virtual memory tax" through pointer-chasing and TLB bottlenecks [C005, C007]. In contrast, vAttention leverages CUDA virtual memory management APIs to maintain virtual memory contiguity, resulting in up to a 1.23x throughput increase over PagedAttention-based kernels ^[C007].

On mobile devices, where the Android Low Memory Killer (LMK) presents a binary failure mode, LLMS mitigates switching overhead by treating the KV cache as a stateful system service ^[C002]. It employs Tolerance-Aware Compression and an LCTRU (Least Compression-Tolerable and Recently-Used) queue to prioritize which cache chunks to evict or compress, reducing latency by two orders of magnitude compared to baseline solutions ^[C002].

At the hardware level, the "Weight-Context-Switch" primitive is most viable on Unified Memory Architecture (uma) hardware. The NVIDIA Grace Hopper and Blackwell architectures utilize NVLink C2C to provide a 900 GB/s memory-coherent interconnect ^[C008]. This allows the GPU to access CPU memory without explicit data transfers or redundant copies, effectively treating the system's total memory as a single pool for both weights and KV caches ^[C008].

Finally, evidence suggests that "worst-case" provisioning for context length leads to 4-8$\times$ throughput waste because 80-95% of production requests are short ^[C006]. Implementing dual-pool token-budget routing—partitioning fleets into specialized short-context and long-context pools—reduces GPU-hour consumption by 31-42% ^[C006].

Tensions and Tradeoffs

Practitioners implementing local LLM memory orchestration face a primary conflict between maximizing the available memory budget for the KV-cache and maintaining raw kernel execution throughput. This creates several critical architectural tensions:

Memory Flexibility vs. Kernel Throughput
The industry standard, PagedAttention, increases batch sizes by treating GPU memory like OS RAM ^[C000], but introduces the aforementioned "virtual memory tax" ^[C007]. Alternative approaches like vAttention attempt to retain virtual contiguity to support standard attention kernels out-of-the-box, improving throughput ^[C007].

Context Intelligence vs. System Availability
On mobile devices, there is a binary tradeoff between the "intelligence" of a model (determined by its context window) and system stability ^[C002]. Increasing the KV cache to support longer reasoning chains raises the probability that the Android Low Memory Killer (LMK) will terminate the process ^[C002]. This necessitates fine-grained, chunk-wise compression and swapping to reduce context-switching latency ^[C002].

Offloading Latency vs. Memory Capacity
When KV cache requirements exceed GPU HBM, practitioners must offload data to DDR or SSDs ^[C004]. While this prevents out-of-Memory (OOM) crashes, it introduces significant context-switching latency ^[C004]. Unified Memory Architecture (uma) hardware, such as NVIDIA Grace Hopper, attempts to resolve this via NVLink C2C, providing 900 GB/s of memory-coherent bandwidth to treat CPU memory as a seamless extension of GPU memory ^[C008].

Worst-Case Provisioning vs. Resource Efficiency
Provisioning fleets for worst-case context lengths often leads to 4-8$\times$ throughput waste because the majority of actual requests are short ^[C006]. This mismatch triggers reliability issues, including preemption and request rejections, forcing a tradeoff between a simplified "one-size-fits-all" configuration and the complexity of dual-pool routing based on estimated token budgets ^[C006].

Opportunities

Priority Build Targets

To implement a first-principles 'Weight-Context-Switch' primitive, development should focus on three architectural components:

1. uma-Aware state Orchestrator: Build a kernel-level manager for NVIDIA Grace Hopper or Blackwell architectures that leverages the 900 GB/s NVLink C2C memory-coherent interconnect to treat CPU and GPU memory as a single address space ^[C008]. This eliminates explicit data transfers and redundant copies when swapping KV caches between HBM and DDR ^[C008].
2. Hybrid Virtual Memory Layer: Integrate vAttention's approach of decoupling virtual and physical memory allocation using CUDA APIs ^[C007] with TLB subregion contiguity ^[C005]. This would allow the system to retain virtual contiguity for the KV cache—improving throughput ^[C007]—while reducing TLB miss stalls by coalescing up to 512 pages into a single entry ^[C005].
3. Tolerance-Aware Eviction Engine: Implement a kernel-level LCTRU (Least Compression-Tolerable and Recently-Used) queue ^[C002]. This engine should prioritize swapping or compressing KV cache chunks based on measured accuracy tolerance, reducing context-switching latency by up to two orders of magnitude on memory-constrained devices ^[C002].

Critical Research Questions

The transition to a unified swappable resource requires empirical answers to the following:

• The Virtualization Tax: Can the performance overhead of non-contiguous memory access in PagedAttention ^[C000] be fully negated by combining kernel fusion with TLB coalescing ^[C005], or is there an inherent throughput ceiling when moving away from contiguous HBM layouts?
• Quantization Drift: When applying 4-bit product quantization to both model weights and the KV cache—as seen in the TFLOP prototype ^[C001]—how does the resulting representation drift impact the fidelity of long-context reasoning (100K+ tokens) compared to FP16 [C001, C004]?
• Offload Thresholds: At what specific context length does the latency of swapping KV caches to DDR ^[C004] exceed the throughput gains provided by continuous batching and paged memory [C000, C009]?

Memory Management Trade-offs

References

• [C000] Efficient Memory Management for Large Language Model Serving with PagedAttention — https://doi.org/10.1145/3600006.3613165
• [C001] TFLOP: Towards Energy-Efficient LLM Inference An FPGA-Affinity Accelerator with Unified LUT-based OPtimization — https://www.semanticscholar.org/paper/f6efc61b2656d856b22c51e6564f091d8ae730d3
• [C002] LLM as a System Service on Mobile Devices — https://doi.org/10.48550/arxiv.2403.11805
• [C003] Characterizing the Behavior and Impact of KV Caching on Transformer Inferences Under Concurrency — https://doi.org/10.1109/ipdps64566.2025.00108
• [C004] Challenges in Deploying Long-Context Transformers: A Theoretical Peak Performance Analysis — https://doi.org/10.48550/arxiv.2405.08944
• [C005] Enabling Large-Reach TLBs for High-Throughput Processors by Exploiting Memory Subregion Contiguity — https://arxiv.org/abs/2110.08613
• [C006] Dual-Pool Token-Budget Routing for Cost-Efficient and Reliable LLM Serving — https://openalex.org/W7153670721
• [C007] vAttention: Dynamic Memory Management for Serving LLMs without PagedAttention — https://arxiv.org/abs/2405.04437
• [C008] Accelerate Large-Scale LLM Inference and KV Cache Offload — https://developer.nvidia.com/blog/accelerate-large-scale-llm-inference-and-kv-cache-offload-with-cpu-gpu-memory-sharing/
• [C009] vLLM Explained: PagedAttention, Continuous Batching, and Deploying High ... — https://www.runpod.io/articles/guides/vllm-pagedattention-continuous-batching