RFC: CPU Offload for Mixture-of-Experts (MoE) Inference in vLLM
PR：#31938
Summary
This PR proposes a CPU Offload Module for MoE inference in vLLM, enabling a large portion of expert weights and computation to be dynamically offloaded to the CPU while keeping only a small, hot subset of experts cached on GPU.
The design supports:
- Hybrid GPU–CPU execution
- Pinned-memory–based weight streaming
- Asynchronous GPU ↔ CPU interaction via callback
- AVX / AMX–optimized CPU MoE kernels
- Two execution modes:
- DBO (Dual Batch Overlap)
- Prefetch-based miss expert execution
This approach significantly reduces GPU memory footprint and improves scalability for large MoE models with thousands of experts.
Motivation
Problem Statement
Large MoE models (e.g., Deepseek-like architectures) suffer from GPU memory pressure during inference:
- Expert weights dominate memory usage
- Only a small subset of experts is active per token
Existing approaches typically:
- Require all expert weights on GPU
- Perform naive CPU fallback with high synchronization overhead
Key Observations
- Expert access is sparse and skewed: Only a small fraction of experts are accessed per token, with certain experts being accessed more frequently than others
- CPU memory capacity is abundant: CPU RAM can easily accommodate the full set of expert weights that would otherwise overwhelm GPU memory
- Modern CPUs (AVX512 / AMX) can efficiently compute MoE FFNs: CPU-based computation can be competitive for expert forward passes when properly optimized
- GPU kernels can overlap compute with CPU execution if callbacks are used: Asynchronous execution patterns enable efficient resource utilization
Goals of This RFC
- Reduce GPU memory usage without sacrificing throughput
- Enable fine-grained expert caching
- Overlap GPU compute, CPU compute, and data movement
- Integrate cleanly with vLLM's execution model
Proposed Change.
Proposed Change
Architecture Overview
+-------------------+ callback +----------------------+
| GPU | <-----------------> | CPU |
| | update cache experts | |
| Cached Experts | <----------------- | Offloaded Experts |
| (hot experts) | | (cold experts) |
| | | |
| fused_experts() | | AVX / AMX MoE FFN |
+-------------------+ +----------------------+
Core Design Philosophy
The core design principle is that the GPU no longer stores all expert weights for each layer, but instead caches only a limited number of hot experts. The CPU maintains the complete set of experts and dynamically determines which experts need to be copied to the GPU and which should be computed directly on the CPU based on actual token routing behavior.
The entire mechanism revolves around:
- Expert cache management
- Miss buffer handling
- Copy policy decisions
- CPU/GPU computation overlap
Key Components
- Python Offload Manager (CpuOffloadInfer): Orchestrates the offload process, manages expert cache state, and coordinates GPU-CPU interactions
- GPU Expert Cache: Limited-capacity cache storing hot experts on GPU
- Miss Expert Buffer (double-buffered): Temporary buffer for experts that miss the cache during forward passes
- CPU MoE Execution Engine: AVX/AMX-optimized kernels for computing expert forward passes on CPU
- GPU↔CPU Callback-based Synchronization: Asynchronous communication mechanism for coordinating GPU and CPU execution
Initialization Phase
During model initialization:
- All MoE expert weights for each layer are fully loaded and permanently resident in CPU pinned memory
- The GPU allocates an Expert Cache with capacity
cache_expert_num for each layer, storing the most frequently accessed experts
- The GPU cache is not static; experts are dynamically managed based on runtime token routing behavior
To track the state of experts in the GPU cache, the system maintains per-layer metadata:
- cache_map: Maps expert IDs to their positions in the GPU cache
- miss_map: Tracks which experts are currently in the miss buffer
- policy_sort: Maintains priority ordering for expert replacement decisions
Forward Pass Execution Flow
Step 1: Expert Cache Policy Matching
At the start of a forward pass, the model has already obtained topk_ids for each token from the router. The system calls expert_cache_policy to match these topk_ids against the current layer's cache state.
This process outputs two key pieces of information:
- cpu_topk_ids: Which tokens' experts require CPU computation
- copy_map: The set of experts that need to be copied from CPU to GPU in this forward pass
Important: copy_map does not directly correspond to "experts copied to GPU cache". It is simply a list of experts that need to be copied in this pass, and their final destination depends on the execution mode.
Step 2: Execution Mode Selection
The system operates in two primary execution modes:
DBO Mode (Dual Batch Overlap)
When the system is in DBO mode or in decode/small batch scenarios, the forward pass enters a fully parallel CPU-GPU execution path:
- Experts in
copy_map are asynchronously copied to the GPU Expert Cache for subsequent fused_experts computation
- CPU immediately begins computing miss experts
- CPU computation, GPU computation, and expert copying are deliberately placed in different execution threads
- Overlap is achieved through vLLM's DBO scheduling mechanism: while the GPU computes fused experts for the current batch, the CPU is already working on miss experts for the next step or the same step, maximizing resource utilization and reducing decode latency
Prefetch Mode
In Prefetch mode (typically for larger prefill batches), system behavior adjusts based on the number of tokens in the batch:
-
As token count increases, more experts are triggered in the forward pass
-
The system dynamically calculates n_copy to limit the maximum number of experts copied in this pass
-
If n_copy is less than the total number of experts:
- CPU still participates in computation
- Experts in
copy_map are not placed in the GPU cache
- Instead, they are copied to a dedicated Miss Expert Buffer (
temp_layer)
- GPU uses this temp buffer to execute
fused_experts
- CPU computes the remaining experts that were not copied
- Results from both paths are merged at the output stage
-
When batch size is extremely large and n_copy covers all or nearly all experts:
- The system automatically degrades to "full GPU mode"
- CPU no longer participates in computation
- All experts are copied and
fused_experts computation is completed on the GPU side
- This is not an additional branch logic, but a natural consequence of the Prefetch strategy when copy count reaches the threshold
Double-Buffered Miss Expert Buffer Management: To prevent miss experts from being overwritten during cross-layer execution, the system globally maintains only two Miss Expert Buffers, using layer_id % 2 for double-buffering:
- Even-numbered layers use buffer 0
- Odd-numbered layers use buffer 1
By coordinating with independent CUDA streams and events:
- Copy and computation on the same buffer are strictly serialized
- Different buffers can form a natural pipeline
- Expert copying and computation for adjacent layers can interleave, enabling efficient pipelining
Feedback Period.
No response
CC List.
No response
Any Other Things.
No response
Before submitting a new issue...
RFC: CPU Offload for Mixture-of-Experts (MoE) Inference in vLLM
PR：#31938
Summary
This PR proposes a CPU Offload Module for MoE inference in vLLM, enabling a large portion of expert weights and computation to be dynamically offloaded to the CPU while keeping only a small, hot subset of experts cached on GPU.
The design supports:
This approach significantly reduces GPU memory footprint and improves scalability for large MoE models with thousands of experts.
Motivation
Problem Statement
Large MoE models (e.g., Deepseek-like architectures) suffer from GPU memory pressure during inference:
Existing approaches typically:
Key Observations
Goals of This RFC
Proposed Change.
Proposed Change
Architecture Overview
Core Design Philosophy
The core design principle is that the GPU no longer stores all expert weights for each layer, but instead caches only a limited number of hot experts. The CPU maintains the complete set of experts and dynamically determines which experts need to be copied to the GPU and which should be computed directly on the CPU based on actual token routing behavior.
The entire mechanism revolves around:
Key Components
Initialization Phase
During model initialization:
cache_expert_numfor each layer, storing the most frequently accessed expertsTo track the state of experts in the GPU cache, the system maintains per-layer metadata:
Forward Pass Execution Flow
Step 1: Expert Cache Policy Matching
At the start of a forward pass, the model has already obtained
topk_idsfor each token from the router. The system callsexpert_cache_policyto match thesetopk_idsagainst the current layer's cache state.This process outputs two key pieces of information:
Important:
copy_mapdoes not directly correspond to "experts copied to GPU cache". It is simply a list of experts that need to be copied in this pass, and their final destination depends on the execution mode.Step 2: Execution Mode Selection
The system operates in two primary execution modes:
DBO Mode (Dual Batch Overlap)
When the system is in DBO mode or in decode/small batch scenarios, the forward pass enters a fully parallel CPU-GPU execution path:
copy_mapare asynchronously copied to the GPU Expert Cache for subsequentfused_expertscomputationPrefetch Mode
In Prefetch mode (typically for larger prefill batches), system behavior adjusts based on the number of tokens in the batch:
As token count increases, more experts are triggered in the forward pass
The system dynamically calculates
n_copyto limit the maximum number of experts copied in this passIf
n_copyis less than the total number of experts:copy_mapare not placed in the GPU cachetemp_layer)fused_expertsWhen batch size is extremely large and
n_copycovers all or nearly all experts:fused_expertscomputation is completed on the GPU sideDouble-Buffered Miss Expert Buffer Management: To prevent miss experts from being overwritten during cross-layer execution, the system globally maintains only two Miss Expert Buffers, using
layer_id % 2for double-buffering:By coordinating with independent CUDA streams and events:
Feedback Period.
No response
CC List.
No response
Any Other Things.
No response
Before submitting a new issue...