Cpu gpu load balancing boosts long context MoE model speed
Physically Partitioned KVCache Format for CPU--GPU Load Balancing in MoE Inference
Distributed, Parallel, and Cluster ComputingMachine Learning
Summary
Large AI models called Mixture-of-Experts (MoE) need to manage memory between the CPU and GPU when processing very long inputs. The authors found that existing methods handle this key memory in a single way, causing inefficiencies. They created a new method, InplaceKVCache, that physically partitions this memory so parts can be optimized separately for CPU or GPU use. This lets them dynamically balance computation between the CPU and GPU without moving data. This results in faster processing speeds on various tested models when handling long input sequences.
What this means in practice
- •For ai infrastructure engineers: Optimize inference speed of large MoE models working with very long sequences on systems with limited GPU memory by dynamically balancing CPU and GPU workloads using InplaceKVCache.
- •For cloud service providers: Improve resource utilization and throughput for large language model services that process long-context queries by integrating CPU-GPU load balancing mechanisms for KVCache management.$Commercial implications: Enables offering more efficient and scalable inference services for large language models, attracting customers demanding long-context capabilities.
Authors
Enda Yu, Dezun Dong, Xiangke Liao
Abstract
Single-GPU long-context inference with Mixture-of-Experts (MoE) models requires spilling the key-value cache (KVCache) to CPU memory. The spilled KV serves two complementary purposes---transferring to the GPU for attention computation, or computing in-place on the CPU---which demand opposing physical states. The optimal split between them varies with workload, yet existing KVCache abstractions offer only storage semantics over a monolithic object of a single physical state, and cannot express dynamic load balancing. We propose InplaceKVCache, the first KVCache abstraction whose format fixes each byte's physical residency at write time, so that the CPU--GPU load balance can be adjusted without moving data after placement. It realizes this as a four-region layout along two dimensions---device affinity and access pattern---turning load balancing into pure scheduling. Built on this abstraction, WriteScope splits CPU--GPU shares along the sequence dimension, and a portable roofline performance model determines the optimal CPU share as sequence length evolves, with online feedback tracking CPU cost drift. On three MoE models (DeepSeek-V2-Lite, Qwen3-30B-A3B, Mixtral-8$\times$7B) with a 32~GB VRAM budget, WriteScope supports end-to-end inference at the 1M-token aggregate scale. In the long-context regime ($\ge$8K), it achieves geometric-mean speedups of $1.5\times$--$2.5\times$ on A100 and $1.4\times$--$1.7\times$ on V100 over four reproduced baselines, while vLLM, SGLang, and KTransformers fail even with a doubled KV budget. A DeepSeek-V4-Flash case study validates composition with native sparse attention.