[Paper] Janus：拆分注意力和专家以实现可扩展的 MoE 推理

发布: 3天前 (2025年12月16日 GMT+8 00:53)

7 min read

原文: arXiv

Source: arXiv - 2512.13525v1

Overview

该论文介绍了 Janus，一个将大型混合专家（Mixture‑of‑Experts，MoE）模型的注意力和专家组件拆分并在不同 GPU 子集群上运行的系统。通过将这两部分视为独立工作负载，Janus 能够为每个部分分配恰当的资源，显著提升推理可扩展性并降低实时应用的延迟。

系统划分 – 作者将 MoE 模型拆分为两个逻辑阶段：
- (a) 注意力阶段（标准 Transformer 自注意力）
- (b) 专家阶段（稀疏前馈层）。
  每个阶段在各自的 GPU 子集群上运行。
两阶段通信
- 阶段 1：在同一节点内，高速 NVLink 或 PCIe 在注意力 GPU 与专家 GPU 之间交换 token 表征。
- 阶段 2：跨节点时，使用带宽感知的集合通信（例如基于 NCCL 的 all‑gather）仅传输被激活的专家输出，从而保持跨节点流量低。
GPU‑内核调度器 – 自定义内核监控每个 token 选中的专家数量，并重新分配工作，使任何 GPU 都不会成为瓶颈。由于 MoE 推理受内存限制，调度器直接在 GPU 上运行，以避免 CPU‑GPU 同步延迟。
动态资源管理 – 控制器观察运行时指标（吞吐量、延迟、GPU 内存使用），并可以：
- (i) 为注意力或专家子集群添加或移除 GPU，和
- (ii) 在 GPU 之间迁移专家权重，以保持负载均衡。

所有组件均基于现有深度学习框架（如 PyTorch）构建，并使用标准 CUDA/NCCL 原语，使该方法能够在典型数据中心 GPU 集群中迁移。

指标	Janus	基准 MoE 系统	加速比
每 GPU 吞吐量（tokens/s）	高出 3.9 倍	–	3.9×
95 百分位每标记延迟（BERT‑large‑MoE）	≤ 12 ms	20 ms	降低 1.7 倍
GPU 内存利用率（attention vs. expert）	45 % / 78 %（平衡）	70 % / 70 %（过度配置）	–
扩展效率（最多 64 GPU）	85 %	55 %	–

关键要点

Cost‑effective inference – 云服务提供商可以在需要时启动较小的注意力集群和较大的专家集群，从而降低 GPU 租赁费用。
Latency‑critical services – 实时推荐、搜索或对话式 AI 能够在不对整个 MoE 模型进行过度配置的情况下，实现每个 token 小于 10 ms 的目标。
Simplified deployment pipelines – 由于 Janus 使用标准的 CUDA/NCCL 和 PyTorch，现有的针对 transformer 模型的 CI/CD 工作流可以在最少代码修改的情况下进行扩展。
Future‑proofing – 随着 MoE 模型规模增长至数十亿参数，独立扩展稀疏专家部分的能力对于在当前 GPU 硬件上保持可行的推理至关重要。

总体而言，Janus 证明了将注意力和专家视为独立且可调的服务，是在生产环境中实现可扩展、低延迟 MoE 推理的实用路径。