LLM System: 通信计算融合算子实现 01 - sm80GEMM + ReduceScatter

Sun, 05 Jul 2026 00:00:00 +0800

flux实现：sm80 gemm+rs

看flux的论文来说，原理并不难，修改的是ffn的第二个gemm的epi阶段，让epi阶段的写回操作变成写入reducescatter的通信buffer，相当于做了一次零拷贝。几个关键的点：1. 如何用cutlass自定义epi阶段。2. 用了什么ptx。3. 清零？

cutlass自定义epilogue–sm80

常规的epi，写回gmem，这里改成写在scatter-aware-memory（我自己起的名字）。cutlass封装自定义epi的逻辑抽象是：输出tile存到哪里，以及做什么运算，这分别是两个模版类：

using EVT_D = decltype(this->evt_d(kparams));
using StoreD = decltype(this->evt_store_d(kparams));
using EVT = cutlass::epilogue::threadblock::Sm80EVT<StoreD, EVT_D>;

StoreD就是规定了怎么存，EVTD就是规定了怎么算。 EVT全称epilogue visitor tree就是定义了一个epi阶段的计算-访存操作流，用树状图的形式保存epi阶段的一堆操作（做成图应该是为了方便接入nvcc？不懂为啥非得在这强调是tree）一个树节点（操作）是一个visitor。

custom_evt_d()

EVT_Compute0 = alpha * accumulator

代码里对应 VisitorCompute<cutlass::multiplies, …>，两个输入是 VisitorScalarBroadcast 和 VisitorAccFetch。第一个负责广播 alpha，第二个v从 GEMM accumulator 里拿当前 tile 的结果，然后 multiply

EVT_Compute1 = beta * C + EVT_Compute0 = beta * C + alpha * accumulator

代码里对应 VisitorAuxLoadGemmk 先把 C/bias 读进 epilogue，然后 VisitorCompute<cutlass::multiply_add, …> 做 beta * C + alpha * acc。因为CUTLASS 2.x 没有 SM90 那种 SrcFetch 替代物，所以这里用 VisitorAuxLoadGemmk 来读 C。

ReduceScatter on Echo的技术博客

LLM System: 通信计算融合算子实现 01 - sm80GEMM + ReduceScatter

flux实现：sm80 gemm+rs

cutlass自定义epilogue–sm80

custom_evt_d()