硬件加速的考量

高效运行大型专家混合 (MoE) 模型需要仔细考量推理阶段的底层硬件。尽管 MoE 的稀疏特性与参数数量相近的密集模型相比，能够减少理论上的浮点运算次数（FLOPs），但在标准硬件上直接运行它们，通常无法转化为实际的加速效果。这主要是因为内存访问模式不规则、分布式设置中的通信开销，以及基于路由决策的动态计算特性。因此，有效地使用硬件加速功能对于实现低延迟和高吞吐量是十分必要的。

MoE 推理的 GPU 加速

现代 GPU，特别是 NVIDIA 的产品（如 Ampere、Hopper），提供了一些功能可用于加速 MoE 推理，但这需要超越标准的密集矩阵乘法库才能达到最佳性能。

自定义核与核融合

标准的深度学习库针对密集操作进行了高度优化。MoE 层涉及一系列操作：计算门控分数、选择 top-k 专家、路由令牌、计算专家函数以及组合结果。通过框架将这些作为独立步骤执行，会引入大量来自核启动以及 GPU 全局内存与其计算单元之间数据移动的开销。

一种主要的优化策略是核融合。通过编写自定义核（例如，使用 CUDA 或 Triton 等库），可以将 MoE 层的多个逻辑步骤组合成一次 GPU 核启动。例如，一个融合核可以：

1. 为一批令牌计算门控分数。
2. 执行 top-k 选择。
3. 收集所选专家的参数或指针。
4. 执行专家计算（通常是专家内部的密集矩阵乘法）。
5. 结合输出，并按门控分数加权。

这最大程度地减少了到全局内存的往返，将中间数据保留在更快的 L1/L2 缓存或流式多处理器（SM）的共享内存中，大幅降低了延迟。

利用 Tensor Cores 和稀疏性功能

GPU 配备了 Tensor Cores 等专用单元，旨在加速矩阵乘法，尤其是在较低精度（FP16、BF16、INT8、FP8）下。尽管专家计算内部通常涉及密集矩阵乘法，直接受益于 Tensor Cores，但 MoE 的整体结构是稀疏的。NVIDIA 的“稀疏性”功能针对结构化稀疏性（例如 2:4 模式），通常不直接适用于 MoE 专家选择的块稀疏性。Tensor Cores 的主要益处在于加速所选专家内部的计算以及门控网络本身的计算，尤其是在与量化结合使用时。

分布式推理的优化通信

如果专家分布在多个 GPU 上（专家并行）以使模型适应内存，推理仍然需要通信。当一批令牌到达时，门控结果确定哪些令牌需要发送到哪些 GPU，以便由相关专家处理。这通常涉及All-to-All通信模式，类似于训练，但根据批处理策略，负载可能更小。NVLink 和 NVSwitch 等高速互连，以及优化的通信库（例如 NCCL），对于最大限度地减少这种数据交换的延迟影响十分重要。在推理过程中，也可以应用诸如通信与计算重叠等技术。

MoE 推理在两个 GPU 上分布的流程。令牌根据门控决策进行路由，这可能需要跨 GPU 通信（由暗示令牌路由的跨集群边界的箭头表示）。

MoE 推理的 TPU 加速

Google 的张量处理单元（TPU）专为加速机器学习工作负载而设计，主要侧重于大规模矩阵运算。

脉动阵列效率

TPU 运用脉动阵列，这种阵列在执行大型密集矩阵乘法方面效率极高。这使得它们对于每个选定专家内部执行的计算非常有效。一旦令牌被路由且相关专家参数加载完毕，TPU 就能非常迅速地处理专家的前向传播。

高带宽内存 (HBM)

TPU 通常配备大量的片上高带宽内存（HBM），与计算单元位于同一封装中。这种高带宽对于 MoE 模型有利，因为它允许更快地将所选专家的参数加载到 TPU 的内存（MEMU）中。最大限度地减少参数获取时间非常重要，特别是考虑到所有专家的总参数量可能非常大。

编译器优化 (XLA)

TPU 的性能高度依赖于 XLA（加速线性代数）编译器。XLA 执行复杂的图优化，包括操作融合、内存布局优化以及根据 TPU 硬件定制的调度。对于 MoE 模型，XLA 可以自动融合门控机制和专家计算的可能部分，从而减少开销，类似于 GPU 上的手动 CUDA 核融合。然而，与自定义 CUDA 核提供的灵活性相比，动态路由逻辑的自动优化程度可能有所不同。

动态路由带来的挑战

虽然 TPU 在静态计算图方面表现出色，但 MoE 固有的动态路由带来了挑战。硬件和编译器针对可预测的数据流进行了优化。高效处理条件执行，即不同令牌激活不同专家（可能需要不同的参数甚至重新编译/分派逻辑），需要仔细实现，并且可能需要针对 TPU 条件计算执行进行优化的特定框架支持。

利用稀疏性和量化

无论使用何种特定加速器（GPU 或 TPU），有两项技术是 MoE 推理硬件加速的基础：

条件参数加载

一个主要瓶颈是加载所选专家的参数。由于每个令牌只有一小部分（例如 top-2）专家是活跃的，理想情况下，只有这些活跃专家的权重应该从主内存（DRAM 或 HBM）加载到加速器更快的本地内存（缓存、SMEM、MEMU）中。高效实现这种“条件加载”需要复杂的内存管理系统，以及路由机制与内存子系统之间的仔细协调。专为分布式 MoE 设计的框架和库通常会包含实现此目的的策略。

量化

量化，即降低模型权重和激活的精度（例如，从 FP32 到 FP16、BF16、INT8 或 FP8），对 MoE 模型特别有效。

1. 减少内存占用： 大幅减少海量专家参数所需的存储空间。
2. 减少内存带宽： 加载活跃专家参数时，需要从主内存传输到计算单元的数据更少。
3. 更快的计算： GPU (Tensor Cores) 和 TPU 都拥有专用硬件单元，为低精度计算提供大幅加速。

有效应用量化通常涉及量化感知训练（QAT）以保持准确性，特别是对于可能对精度变化敏感的路由机制。

MoE 模型在不同硬件和优化级别下的延迟比较。请注意对数刻度。硬件加速（GPU/TPU）相对于 CPU 提供了大幅提升。融合核/优化以及 INT8 量化进一步降低了延迟。

系统级考量

实现 MoE 推理的最佳硬件加速通常是一个系统级问题。它包括：

• 选择合适的硬件： GPU 提供自定义核的灵活性，而 TPU 则擅长专家内部的密集数学计算，并受益于 XLA。选择取决于具体的 MoE 架构、可用的软件栈和性能目标（延迟与吞吐量）。
• 软件栈： 使用对 MoE 特定优化（融合核、高效 All-to-All、条件加载）有内置支持的框架和库（如 DeepSpeed、Tutel、FasterTransformer，或 PyTorch/TensorFlow/JAX 中的专用例程）十分重要。
• 协同设计： 有时，在做出架构选择（例如，专家大小、专家数量、门控机制复杂性）时，需要考虑硬件的限制和能力，以最大限度地提升推理性能。

归根结底，弥合理论上稀疏性带来的计算节省与实际推理速度之间的差距，需要对 MoE 架构和底层硬件加速器的能力都有透彻的理解。采用核融合、优化通信、条件加载和量化等技术进行仔细实现，是充分发挥 MoE 在生产环境中全部潜力的必要条件。