Physically Partitioned KVCache Format for MoE Inference
Physically Partitioned KVCache Format for CPU--GPU Load Balancing in MoE Inference
This research provides a method to improve the performance of large language models on GPUs by better managing memory usage, which is relevant for anyone working with or evaluating MoE models.
A new KVCache format for Mixture-of-Experts models allows dynamic load balancing between CPU and GPU memory. It splits the cache into four regions based on device affinity and access pattern, enabling better performance on long-context tasks. On models like DeepSeek-V2-Lite and Mixtral-8x7B, it achieved 1.5x-2.5x speedups on A100 GPUs compared to baselines.
Physically Partitioned KVCache Format for CPU--GPU Load Balancing in MoE Inference
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.
