[Feature]: MoE Active Expert Management --moe-gpu-prefetch <num>

[leoustc]

🚀 The feature, motivation and pitch

🚀 The feature, motivation and pitch

I would like to propose a new feature for vLLM to improve memory efficiency and scalability for sparse Mixture-of-Experts (MoE) models:

--moe-gpu-prefetch <num>

This feature introduces expert-level GPU memory management, allowing only a limited number of experts to reside on GPU while the rest are offloaded to CPU memory.

---

Motivation

Modern MoE models (e.g., Gemma, Mixtral-style architectures) contain a large number of experts per layer (e.g., 128), but only a small subset (e.g., top-8) are activated per token.

Current approaches to handle GPU memory constraints include:

• Full model residency on GPU (high memory requirement)
• Layer-level CPU offload (coarse-grained and inefficient)

These approaches do not align well with the fundamental property of MoE:

Activation is sparse at the expert level, not at the layer level.

As a result:

• GPU memory is wasted on inactive experts
• Data movement is excessive when offloading entire layers
• Large MoE models cannot efficiently run on limited GPU setups

---

Proposed feature

--moe-gpu-prefetch <num>

Where <num> defines the number of expert models that can reside ("hot") on GPU at any time.

Execution model:

• All expert weights are stored in CPU memory

• GPU memory contains:

  • shared layers
  • router
  • a fixed number of <num> expert slots

• During inference:

  • If a routed expert is already on GPU → execute directly
  • If not → load it from CPU into GPU (evicting another expert if needed)

This effectively turns GPU into a hot expert cache.

---

🔧 Core Design: GPU Expert Slot Mapping

The core of this feature is to introduce a fixed set of physical GPU expert slots, and maintain a dynamic mapping between these slots and logical expert IDs.

Key abstraction

GPU Expert Slots (size = <num>)

Slot 0 → Expert 12
Slot 1 → Expert 45
Slot 2 → Expert 7
...
Slot N-1 → Expert 63

Runtime structures:

slot_to_expert:  slot_id → expert_id
expert_to_slot:  expert_id → slot_id (or -1 if not resident)

Where:

• A slot is a fixed GPU memory region capable of storing one expert’s weights
• An expert ID is the logical expert index in the MoE model

---

Execution flow

1. Router selects top-k experts per token

2. For each expert:

   • Check expert_to_slot

     • Hit → directly execute using the corresponding GPU slot
     • Miss → trigger load

3. On miss:

   • Select a victim slot (e.g., LRU or frequency-based)
   • Evict the current expert from that slot
   • Load the required expert from CPU to GPU
   • Update mappings

4. Continue execution

---

Design philosophy

Instead of managing memory at the layer level:

This feature manages memory at the expert level.

This shifts the model execution from:

• Layer-centric memory management

to:

• Expert-centric memory management

The GPU is no longer treated as a place to hold the entire model, but as:

a bounded, dynamically managed cache of active experts.

---

Benefits

• Enables large MoE models to run on limited GPU memory
• Reduces memory waste from inactive experts
• Avoids coarse layer-level data movement
• Aligns with sparse MoE execution semantics
• Maintains correctness (no change to outputs)
• Provides deterministic GPU memory usage

---

Relation to existing vLLM design

This design is conceptually similar to vLLM’s KV cache block management:

• KV cache → manages token-level memory blocks
• Expert slots → manage expert-level memory blocks

This allows:

• O(1) expert residency lookup
• clean integration into the MoE execution path
• minimal impact on the scheduler

---

Alternatives

1. Full GPU residency

   • Requires large GPU memory
   • Not scalable for large MoE models

2. Layer-level CPU offload

   • Moves entire layers between CPU and GPU
   • Ignores MoE sparsity
   • Causes excessive data movement

Compared to these:

Expert-level offload is more fine-grained and better aligned with MoE routing behavior.

---

Summary

This proposal introduces a simple but powerful abstraction:

a fixed number of GPU expert slots mapped dynamically to active expert IDs.

This enables vLLM to:

• scale MoE inference under constrained GPU memory
• better utilize sparse activation patterns
• move toward a more flexible, expert-aware execution model

Alternatives

No response

Additional context

Demo repo: leoustc/vllm-moe

• Hardware: A10-40GB
• Model: /models/gemma-4-26B-A4B-it

GPU limit	Prefetch num	Output tok/s	Status
0.50	16	20.43	OK
0.50	32	NA	KV cache startup failure
0.50	64	NA	KV cache startup failure
0.75	16	20.64	OK
0.75	32	34.78	OK
0.75	64	NA	KV cache startup failure
0.95	16	21.71	OK
0.95	32	33.86	OK
0.95	64	56.51	OK
0.95	72	60.31	OK
0.95	96	NA	active expert cache startup failure

with this feature, we can fine control the GPU resource in MoE like model and still get a good performance on large model on a small GPU.

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[leoustc]

https://github.com/leoustc/vllm-moe/blob/main/BENCHMARK.md

[sbates130272]

Can this include a mode to store the other experts' parameters on NVMe and load from there? Saving on CPU memory footprint?

[wsierek]

Isn't this duplicate of the following PR:
#37190

[leoustc]

Isn't this duplicate of the following PR: #37190

Good question. There is overlap, but I do not think this is an exact duplicate of #37190.

PR #37190 implements a persistent GPU expert cache via --moe-expert-cache-size, with CPU-pinned expert weights and LFRU eviction. That is closest to the cache-residency part of this proposal.

The difference in this issue is that --moe-gpu-prefetch <num> is intended as an active expert management / prefetch mode:

• CPU remains the source of truth for all expert weights.
• GPU keeps a bounded active expert cache.
• A background pager thread manages missing expert loads. (key difference)
• Runtime separates resident experts and missing experts into waves. (key difference)
• Missing experts are prioritized using routed token counts, so hot/more-needed experts are loaded first.
• The goal is not only an LRU/LFRU cache, but overlapping/organizing expert movement around routing demand. (key difference)

So I see #37190 as related and probably a useful base/reference, especially for the provider abstraction and cache policy. This proposal is more focused on prefetch scheduling and wave-based active expert execution. I’m happy to align terminology/design with #37190 if maintainers prefer one shared interface.

[leoustc]

Isn't this duplicate of the following PR: #37190

Good question. There is overlap, but I do not think this is an exact duplicate of #37190.

PR #37190 implements a persistent GPU expert cache via --moe-expert-cache-size, with CPU-pinned expert weights and LFRU eviction. That is closest to the cache-residency part of this proposal.

The difference in this issue is that --moe-gpu-prefetch <num> is intended as an active expert management / prefetch mode:

• CPU remains the source of truth for all expert weights.
• GPU keeps a bounded active expert cache.
• A background pager thread manages missing expert loads. (key difference)
• Runtime separates resident experts and missing experts into waves. (key difference)
• Missing experts are prioritized using routed token counts, so hot/more-needed experts are loaded first.
• The goal is not only an LRU/LFRU cache, but overlapping/organizing expert movement around routing demand. (key difference)

So I see #37190 as related and probably a useful base/reference, especially for the provider abstraction and cache policy. This proposal is more focused on prefetch scheduling and wave-based active expert execution. I’m happy to align terminology/design with #37190 if maintainers prefer one shared interface.

One more clarification on the design goal:

1. The core abstraction I want to introduce is a GPU expert slot layer between the logical expert ID and the physical GPU-resident expert slot ID.

   In other words, routing still produces real/logical expert IDs, but execution uses a dynamic mapping:

   logical expert id -> GPU expert slot id
   

   This is similar in spirit to vLLM’s KV page/block management: KV cache maps token/block-level logical state to physical memory blocks; this feature maps MoE expert-level logical IDs to physical GPU expert slots.

   Because of that, I think this should be treated as a general MoE memory-management feature in vLLM, not only as a model-specific cache optimization.

2. The dedicated cache-management/pager thread is also an important part of the proposal.

   The first version can be conservative, but in the long term this thread can manage, prefetch, and eventually predict expert residency independently of the main compute path. That is the key direction for improving performance beyond a synchronous cache miss/load path.

[github-actions]

This issue has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this issue should remain open. Thank you!