Physically Partitioned KVCache Format for CPU--GPU Load Balancing in MoE Inference

2026-09-13

AI Summary

为了解决MoE模型在单GPU长上下文推理中CPU-GPU负载不平衡的问题，提出了一种新的KVCache格式InplaceKVCache，并通过WriteScope实现动态调整以优化性能。

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.

Problem

• Mixture-of-Experts
• KVCache
• Load Balancing
• Inference
• CPU-GPU

Innovation

• InplaceKVCache
• CPU-GPU Load Balancing
• Dynamic Workload Sharing
• Long-Context Inference
• Mixture-of-Experts