arXiv:2605.17889cs.LG2026-05中稿 · DAC '26

通过CPU-GPU协同优化,大幅提升MoE模型推理吞吐量。

CoX-MoE: Coalesced Expert Execution for High-Throughput MoE Inference with AMX-Enabled CPU-GPU Co-Execution

论文配图:CoX-MoE: Coalesced Expert Execution for High-Throughput MoE Inference with AMX-Enabled CPU-GPU Co-Execution
图 1 · 摘自论文原文
  • 采用大批次+专家聚类策略,减少碎片化计算负担。
  • 预分配高频专家至GPU,降低跨设备传输开销。
  • 在真实硬件上实现7.1倍于现有方案的吞吐提升。

Mixture-of-Experts(MoE)架构通过稀疏激活提升计算效率,但面向高吞吐的推理面临严重显存压力,源于庞大的参数量和中间数据。现有方法尝试通过微批次专家卸载或向CPU卸载计算来缓解,但微批次导致工作负载碎片化,降低操作强度,使专家执行变为内存瓶颈;而CPU卸载受限于慢速PCIe传输,且不适用于解码阶段的注意力计算。因此系统利用率低下,严重制约端到端吞吐。本文提出CoX-MoE,一种基于AMX的CPU-GPU协同系统,通过合并专家执行与智能调度策略,全面优化MoE推理。CoX-MoE引入两项关键机制:(i) 聚合感知调度策略,采用普通批处理而非微批次进行专家计算,并选择性卸载注意力计算;(ii) 静态专家感知分层方案,预先将高频激活专家分配至GPU,减轻PCIe传输开销,平衡CPU与GPU负载。相比最优框架,CoX-MoE在吞吐上分别达到FlexGen的7.1倍和MoE-Lightning的2.4倍。

原文摘要 · Abstract (English)

The Mixture-of-Experts (MoE) architecture improves computational efficiency via sparse expert activation, but throughput-oriented inference faces substantial GPU memory pressure due to a significant parameter size and intermediate data. Prior works attempt to mitigate this using expert offloading with micro-batching or by offloading computation to the CPU. However, the fragmented workload resulting from micro-batching degrades operational intensity, causing expert execution to become memory-bound. Meanwhile, CPU offloading is constrained by slow PCIe transfers and its limited applicability to attention computation in the decode stage. Consequently, these inefficiencies prevent effective system utilization, severely restricting the end-to-end throughput of MoE inference. To address these challenges, this paper proposes CoX-MoE, an Advanced Matrix Extensions (AMX)-enabled CPU-GPU collaborative system that comprehensively optimizes MoE inference by combining coalesced expert execution with strategic workload orchestration for higher throughput. CoX-MoE introduces (i) a coalescing-aware orchestration policy to jointly optimize resource allocation by adopting ordinary batch, instead of micro-batch, for expert computation and selective attention offloading, and (ii) a static expert-aware stratification scheme that pre-assigns frequently activated experts to the GPU, mitigating PCIe transfer overhead and balancing workload for the CPU and GPU during inference. Compared to state-of-the-art frameworks, CoX-MoE delivers significant gains, achieving up to 7.1x and 2.4x higher throughput than FlexGen and MoE-Lightning, respectively.

MoE推理异构计算高效部署硬件协同

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。