GC-MoE：用CUDA Green Contexts解耦MoE通信与计算

为什么MoE解码阶段的专家并行效率低

在专家并行（Expert Parallelism）下，MoE层每一层都要经历三件事：all-to-all的token分发、分组GEMM、以及combine回收。到了decode阶段，每个rank手里只有1到256个token，单卡计算结果寥寥，而权重搬运和通信延迟却一点没少。

这意味着decode场景下的MoE是一个典型的"通信与计算互相拖累"的结构。以DeepGEMM的Mega MoE为代表的融合算子试图用重叠（overlap）来掩盖延迟，但它的做法是让通信warp和计算warp共享同一批SM：寄存器、共享内存、issue slot全部混在一起竞争，同步域覆盖整个SM阵列，每次replay还要收尾清空workspace并做一次跨rank barrier。重叠是有了，但两边的资源预算无法独立调优。

All-to-all 通信是专家并行的核心开销：每个 token 需根据门控路由结果被发送到持有对应专家的 GPU，计算完成后再收回。在 decode 阶段，单次迭代每卡仅有极少 token，all-to-all 传输的数据量虽小，但固定延迟（kernel 启动、NVLink/IB 传输握手）无法被摊薄，通信时间占比急剧上升。

**分组 GEMM（Grouped GEMM）**指将多个专家的小矩阵乘批处理为一个 kernel 调用，以避免逐专家启动 kernel 的开销。但每个专家分到的 token 数高度不均，导致负载不均衡和 SM 利用率下降。

Mega MoE / 融合算子的思路是把 dispatch、GEMM、combine 塞进单个 persistent kernel，用 warp specialization 让通信与计算重叠，省去 kernel 启动和中间 workspace 往返，但代价是资源竞争与同步域耦合。

GC-MoE的核心思路：把设备切成两个独立分区

GC-MoE是一个面向NVIDIA Hopper和Blackwell GPU的开源FP8/FP4 MoE推理算子，它借助CUDA Green Contexts把设备划分成两个分区：一个小的通信分区，一个大的计算分区。

这个划分带来的关键在于：每个角色拥有自己的kernel、自己的资源预算、自己的同步域。分发和GEMM可以各自独立调参，不再互相抢资源。在B300上，GC-MoE只从148个SM中划出20个给分发使用，剩下的全部留给计算。

三节点CUDA Graph与token-ring握手

调度层面，GC-MoE用一个三节点的CUDA Graph把dispatch kernel和grouped-GEMM kernel作为并发的根节点启动（Hopper上因为有融合的shared experts，图是四节点）。token-ring握手让计算在第一批ring block被填充时就开始，而不是等整个分发完成；双缓冲的接收计数器则消除了replay之间尾部的跨rank barrier。

CUDA Graph 将一系列 kernel 启动及其依赖关系捕获为一张静态图，一次提交即可由驱动按依赖调度，消除了逐 kernel 的 CPU 启动开销与 launch gap。"三节点"指图中包含 dispatch、grouped-GEMM 以及 combine/收尾三个 kernel 节点，其中前两者作为并发的根节点同时启动，由 token-ring 握手在 device 侧建立细粒度的生产者-消费者同步，而非依赖图级别的串行边。

token-ring 握手是一种环形缓冲区式的跨 kernel 同步：dispatch 每填满一批 ring block 就递增计数器，计算 kernel 轮询该计数器，一旦首批数据就绪即开始 GEMM，实现计算与分发的流水重叠。双缓冲的接收计数器则让下一次 replay 可以复用上一轮的缓冲区，无需在每轮结束时插入跨 rank barrier。

计算分区的两套MMA实现

计算分区针对两代硬件走了不同路线：Blackwell上跑operand-swapped的tcgen05 block-scaled MMA，Hopper上跑register-A的WGMMA并在寄存器内完成FP4解码。两者都融合了SwiGLU、requantization以及可选的shared experts。

CUDA Green Contexts 是 CUDA 12.4 引入的硬件级资源分区机制，允许将 GPU 的 SM 划分为若干互不重叠的集合，每个 context 拥有独立的 SM 配额。与 MPS 或流优先级不同，Green Contexts 的分区是空间隔离的：不同分区上的 kernel 不会竞争同一批 SM，各自的占用率、共享内存和寄存器预算互不影响。这使得通信 kernel 和计算 kernel 可以分别按各自最优的 block size、occupancy 和同步策略独立调优，而不必为对方的资源需求让步。该特性目前仅在 Hopper 及更新架构上完整支持。

WGMMA（Warpgroup Matrix Multiply-Accumulate）是 Hopper 架构引入的异步张量核心指令，以 128 个线程组成的 warpgroup 为单位执行，操作数 A 可来自寄存器或共享内存。register-A 指 A 操作数常驻寄存器，避免共享内存往返，适合 decode 阶段的小 M 尺寸。

tcgen05 MMA 是 Blackwell 架构的第五代张量核心指令集，引入 Tensor Memory（TMEM）作为累加器与操作数的专用存储层级，支持 block-scaled 格式——即 FP4/FP8 微缩放（MX）格式，每个 32 元素块共享一个指数缩放因子，在保持精度的同时将权重位宽压到 4 bit。operand-swapped 指交换 A/B 操作数的角色，使权重以更有利的布局进入 MMA，减少布局转换开销。

实测结果：稳定加速且全面领先

在8张H20 GPU上，GC-MoE对DeepSeek^广告-V4.1-Flash的decode在routed和shared-expert两个套件上平均加速1.188倍，峰值分别为1.296倍和1.353倍。在8张B300 GPU上，它对GLM-5.2和DeepSeek-V4.1-Flash的加速为1.106倍，带融合shared experts时为1.123倍，峰值1.209倍，并且在全部54个测量点上以相同精度击败Mega MoE。

这些数字的量级看起来不夸张，但考虑到Mega MoE本身已经是高度优化的融合算子，能在所有测量点上保持一致领先，说明资源隔离带来的独立调优空间是真实有效的，而不是某个配置下的偶然胜出。

对MoE推理工程的一点观察

GC-MoE的价值不只在于加速比，更在于它给出了一种工程范式：与其把所有阶段塞进一个kernel里靠共享SM做重叠，不如用硬件级的资源分区把不同职责彻底隔离，让各自的调度策略互不干扰。Green Contexts原本更多被用于延迟敏感型场景的资源预留，这里被拿来做通信与计算的解耦，思路值得借鉴。

当然，这种方案也依赖Hopper/Blackwell这一代硬件特性，且分区比例（比如B300上的20/148）需要针对具体模型和拓扑调优。对于decode批量很小的推理服务，通信分区占比是一个需要实测扫参的关键参数。