MoE 推理优化全景指南：从架构原理到 GPU 部署实践

# 简化的 MoE 前向传播
class MoELayer(nn.Module):
    def __init__(self, num_experts=8, top_k=2, hidden_dim=4096):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.gate = nn.Linear(hidden_dim, num_experts)
        self.experts = nn.ModuleList([
            SwiGLU(hidden_dim, hidden_dim * 4) 
            for _ in range(num_experts)
        ])
    
    def forward(self, x):
        # x: [batch, seq_len, hidden_dim]
        gate_logits = self.gate(x)  # [batch, seq_len, num_experts]
        
        # Top-K 路由
        weights, indices = torch.topk(
            F.softmax(gate_logits, dim=-1), 
            k=self.top_k, dim=-1
        )
        weights = weights / weights.sum(dim=-1, keepdim=True)  # 归一化
        
        # 专家计算
        outputs = torch.zeros_like(x)
        for k in range(self.top_k):
            expert_idx = indices[..., k]  # [batch, seq_len]
            w = weights[..., k:k+1]       # [batch, seq_len, 1]
            for e in range(self.num_experts):
                mask = (expert_idx == e)
                if mask.any():
                    token_x = x[mask]
                    expert_out = self.experts[e](token_x)
                    outputs[mask] += w[mask] * expert_out
        
        return outputs

def load_balancing_loss(gate_logits, num_experts, top_k=2):
    """计算负载均衡辅助损失"""
    # gate_logits: [batch, seq_len, num_experts]
    scores = F.softmax(gate_logits, dim=-1)
    
    # 计算每个专家被选中的概率 P_i
    P = scores.mean(dim=[0, 1])  # [num_experts]
    
    # 计算每个专家分配的 token 比例 f_i
    _, indices = torch.topk(scores, k=top_k, dim=-1)
    one_hot = F.one_hot(indices, num_classes=num_experts).float()
    f = one_hot.sum(dim=[0, 1]) / one_hot.sum()
    
    # Auxiliary Loss
    loss = num_experts * (f * P).sum()
    return loss

def dynamic_routing(x, gate_logits, threshold=0.1):
    """基于阈值的动态路由"""
    scores = F.softmax(gate_logits, dim=-1)
    
    # 只激活得分超过阈值的专家
    active_mask = scores > threshold
    num_active = active_mask.sum(dim=-1)
    
    # 对激活的专家重新归一化
    masked_scores = scores * active_mask.float()
    masked_scores = masked_scores / masked_scores.sum(dim=-1, keepdim=True)
    
    return masked_scores, active_mask

# 伪代码：通信与计算重叠
def moe_forward_overlap(x, num_experts_per_gpu):
    # 提前计算下一层需要的 token dispatch
    dispatch_stream = torch.cuda.Stream()
    with torch.cuda.stream(dispatch_stream):
        next_dispatch = all_to_all_async(next_tokens)
    
    # 在当前 stream 上计算
    output = compute_local_experts(x)
    
    # 等待 dispatch 完成
    torch.cuda.current_stream().wait_stream(dispatch_stream)
    return output, next_dispatch.result()

from deep_ep import EPCommunicator

# 初始化通信器
comm = EPCommunicator(
    num_experts=256,
    num_ranks=8,           # 8 张 GPU
    experts_per_rank=32,   # 每张 GPU 32 个专家
    hidden_dim=7168,
    dtype=torch.float8_e4m3fn
)

def moe_layer_forward(x, expert_weights, expert_indices):
    """MoE 层前向传播"""
    # Dispatch: 发送 token 到对应 GPU
    recv_x, recv_weights, recv_indices = comm.dispatch(
        x, expert_weights, expert_indices
    )
    
    # 本地专家计算
    local_output = compute_experts(recv_x, local_experts)
    local_output = local_output * recv_weights
    
    # Combine: 收集结果
    output = comm.combine(local_output, recv_indices)
    return output

class OffloadedMoE:
    def __init__(self, num_experts=8, gpu_memory_budget_gb=20):
        self.num_experts = num_experts
        self.gpu_experts = {}  # GPU 上的专家缓存
        self.cpu_experts = {}  # CPU 上的完整副本
        self.gpu_budget = gpu_memory_budget_gb
    
    def forward(self, x):
        # 1. Router 选择专家
        expert_indices = self.router(x)
        
        # 2. 检查 GPU 缓存
        for idx in expert_indices:
            if idx not in self.gpu_experts:
                # 从 CPU 加载到 GPU（必要时淘汰冷专家）
                self._load_expert_to_gpu(idx)
        
        # 3. 正常计算
        return self.compute_with_gpu_experts(x, expert_indices)
    
    def _load_expert_to_gpu(self, idx):
        """带 LRU 淘汰的专家加载"""
        if self._gpu_memory_exceeded():
            coldest_expert = self._get_coldest_expert()
            # 异步传输回 CPU
            self._async_offload_to_cpu(coldest_expert)
        
        # 从 CPU 加载到 GPU
        self.gpu_experts[idx] = self.cpu_experts[idx].cuda()
        self.expert_lru.touch(idx)

import torch
from torchao.float8 import Float8Linear

# 将 MoE 层中的专家 FFN 转换为 FP8
def quantize_moe_to_fp8(model):
    for expert in model.moe_experts:
        # W1 和 W2 都转为 FP8
        expert.w1 = Float8Linear.from_float(expert.w1)
        expert.w2 = Float8Linear.from_float(expert.w2)
    return model

# 路由层保持 FP16/BF16 精度（对质量敏感）
# 专家层使用 FP8（对精度不敏感，节省 50% 显存）

from vllm import LLM, SamplingParams

# 加载 Mixtral 8×7B
llm = LLM(
    model="mistralai/Mixtral-8x7B-Instruct-v0.1",
    tensor_parallel_size=4,      # 4 卡 TP
    enforce_eager=False,          # 使用 CUDA Graph 加速
    quantization="fp8",           # FP8 量化
    enable_prefix_caching=True,   # 启用前缀缓存
    max_num_batched_tokens=32768,
    gpu_memory_utilization=0.95,
)

# 推理
prompts = ["解释 Mixture of Experts 的核心思想"]
sampling_params = SamplingParams(
    temperature=0.7,
    max_tokens=512,
    top_p=0.95,
)
outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    print(output.outputs[0].text)

# build.py - TensorRT-LLM 编译 MoE 模型
import tensorrt_llm
from tensorrt_llm import ModelConfig

config = ModelConfig(
    model_name="Mixtral-8x7B",
    precision="fp8",           # FP8 精度
    moe_tp_mode=2,             # MoE Tensor Parallel 模式
    use_fused_mlp=True,        # 融合 FFN kernel
    reduce_fusion=True,        # All-Reduce 融合
    paged_kv_cache=True,       # 分页 KV Cache
    max_batch_size=128,
    max_input_len=4096,
    max_output_len=2048,
)

engine = tensorrt_llm.build(config)
engine.save("mixtral8x7b_fp8.engine")

# 使用 vLLM 的指标接口监控 MoE 运行状态
import requests
import json

# 获取推理指标
metrics = requests.get("http://localhost:8000/metrics").text

# 关键指标
# vllm:gpu_cache_usage_perc: GPU KV Cache 使用率
# vllm:num_requests_running: 正在处理的请求数
# vllm:gpu_num_requests_waiting: 等待队列长度
# vllm:time_to_first_token_seconds: 首 token 延迟

# MoE 特有指标（需要自定义）
def monitor_moe_balance(expert_counts, num_experts=8):
    """监控专家负载均衡"""
    total = sum(expert_counts)
    ideal = total / num_experts
    imbalance = max(abs(c - ideal) for c in expert_counts) / ideal
    print(f"专家负载均衡度: {1 - imbalance:.1%}")
    if imbalance > 0.3:
        print("⚠️ 负载均衡异常，考虑调整路由策略")