EfficientMoE/MoE-Infinity

📚 Overview

MoE‑Infinity is a library that lets you run huge Mixture‑of‑Experts (MoE) language models on GPUs that don’t have enough VR‑memory to hold the whole model. It does this by keeping only the “hot” expert weights on the GPU and moving the rest to host RAM or an SSD, fetching them just‑in‑time. An activation‑aware cache, pre‑fetching, and CUDA‑graph tricks keep the extra data movement fast enough for real‑time inference.

---

🚀 Key Features

Feature	What it means for you
Expert off‑loading	Large MoE checkpoints (e.g., DeepSeek‑V2, Mixtral, Qwen3‑MoE) can be served on a 16 GB GPU by storing most expert parameters on CPU/SSD.
Activation‑aware caching	Frequently used experts stay resident on‑GPU; the system predicts which experts will be needed next and pre‑loads them.
Speed optimisations	Fused CUDA kernels, CUDA‑graph capture, Marlin INT4 GEMM, FP4/MXFP4 paths, and optional FlashAttention/FlashInfer kernels keep latency low.
HuggingFace‑compatible API	MoE class mimics the usual transformers interface (generate, encode, …). The preferred way is MoE.serve() which starts an async OpenAI‑compatible HTTP server for continuous batching.
Production‑ready serving	OpenAI‑style endpoint with streaming (SSE), paged KV cache, request pre‑emption, hot‑reload, health checks and crash‑recovery logs.
Multi‑GPU support	On a single machine, experts are distributed round‑robin across all visible GPUs; each GPU has its own cache and can use peer‑to‑peer access when possible.
Experimental DFlash	Speculative decoding framework that unifies draft/verify paths and supports rich batching.
Optional accelerators	Plug‑in support for FlashAttention and FlashInfer; FP4‑quantised expert off‑load for DeepSeek‑V4‑Flash.

---

🛠️ Installation Quick‑Start

1. Create a virtual environment (Python 3.10+ recommended).

   conda create -n moe-infinity python=3.12
   conda activate moe-infinity
   

2. Install a released wheel (replace v0.2.1 with the desired tag):

   TAG=v0.2.1
   gh release download "$TAG" --repo EfficientMoE/moe-store --pattern 'moe_store-*.whl' -O /tmp/moe-store.whl
   gh release download "$TAG" --repo EfficientMoE/MoE-Infinity --pattern 'moe_infinity-*.whl' -O /tmp/moe-infinity.whl
   pip install /tmp/moe-store.whl /tmp/moe-infinity.whl
   

3. (Optional) Build from source if you need a custom CUDA compute capability or want the latest code. Follow the “Install from Source” section in the README – it mainly requires CUDA, CUTLASS, and matching PyTorch wheels.
4. Enable optional kernels if you have them installed:

   pip install -e '.[flash_attn]'   # FlashAttention
   pip install -e '.[flashinfer]'   # FlashInfer
   

---

📦 How to Use It

In‑process (simple) example

from transformers import AutoTokenizer
from moe_infinity import MoE
import os, torch

ckpt = "deepseek-ai/DeepSeek-V2-Lite-Chat"
tokenizer = AutoTokenizer.from_pretrained(ckpt, trust_remote_code=True)

config = {
    "offload_path": os.path.expanduser("~/moe-infinity"),
    "device_memory_ratio": 0.75,   # how much GPU memory is reserved for the cache
}
model = MoE(ckpt, config)

inp = tokenizer("translate English to German: How old are you?", return_tensors="pt").input_ids.cuda()
out_ids = model.generate(inp)          # deprecated sync path, kept for compatibility
print(tokenizer.decode(out_ids[0], skip_special_tokens=True))

Production serving (continuous batching)

python -m moe_infinity.entrypoints.openai.api_server_v2 \
    --host 127.0.0.1 \
    --model deepseek-ai/DeepSeek-V2-Lite-Chat \
    --offload-dir ./offload_dir

The server now accepts OpenAI‑style POST /v1/completions requests and streams results via Server‑Sent Events.

---

📊 Supported Models (excerpt)

Family	Example checkpoint
DeepSeek‑V2 / V3	deepseek-ai/DeepSeek-V2-Lite-Chat
DeepSeek‑V4‑Flash (FP4)	deepseek-ai/DeepSeek-V4-Flash
Mixtral	mistralai/Mixtral-8x7B-Instruct-v0.1
Qwen3‑MoE	Qwen/Qwen3-30B-A3B
Qwen3.5‑MoE	Qwen/Qwen3.5-35B-A3B
GLM‑5.2 / 5.3 (FP8)	zai‑org/GLM-5.2‑FP8
GPT‑OSS	openai/gpt-oss‑*
DBRX	databricks/dbrx-instruct
Jamba	ai21labs/Jamba‑*
OLMoE	allenai/OLMoE‑*
Meta NLLB‑MoE	facebook/nllb‑moe‑54b

The library automatically registers a model only when the corresponding class exists in the installed transformers version.

---

🎯 When to Use MoE‑Infinity

• You have a GPU with ≤ 24 GB VRAM but need to serve a 30‑B‑parameter MoE model.
• Latency matters – the cache‑aware off‑loading keeps the extra data‑transfer overhead low enough for interactive chat.
• You want a drop‑in HuggingFace experience while still being able to run a production‑grade HTTP server.
• You run multiple GPUs on a single machine and would like the library to shard experts automatically.

---

📚 Further Reading

• Architecture details – ARCHITECTURE.md
• Model‑compatibility matrix – docs/model-compatibility.md
• Serving guide – docs/serving.md
• DFlash speculative decoding – docs/dflash.md
• Benchmarking best practices – docs/benchmarking.md

---

This summary is based solely on the repository’s README; for the most up‑to‑date information, consult the linked documentation files.