在跟MoE部署和推理优化打交道的过程中我越来越觉得真正卡脖子的地方往往不是模型参数和flops本身而是GPU底层的执行单元被安排得明不明白。最近读到一篇叫Weave的论文题目就很对我胃口——MoE 大内核里的细粒度动态 SM 调度而且直接给了硬指标4×H100 实测 2.89× 层加速。这个数字听起来夸张但仔细拆完原理之后你会发现它踩中的正是MoE大内核在H100上最容易被浪费的那部分潜力。这篇文章我会从头到尾把它的出发点、核心设计、实测结果以及我自己复现时的体会和踩坑都捋一遍给正在做MoE推理优化、内核开发或者单纯想搞懂GPU底层玩法的人一些参考。1. 先看痛点MoE大内核里的SM调度为什么这么难1.1 MoE模型的“大内核”到底指什么大家聊MoE张口闭口都是“专家负载均衡”“router路由”很少聊GPU内核级别的执行模型。但真正把MoE落地到GPU上跑过几轮之后你会发现真正的瓶颈往往不在模型层而在一个被称为“大内核”fat kernel的计算单元里。所谓大内核简单说就是把一个MoE层的整块计算例如多个专家的FFN操作融合进一个CUDA kernel来执行。为什么要融合因为如果每个专家单独启动一个小kernel一个8专家的MoE层光专家FFN就要launch八次再加上attention、gate、残差连接一个Transformer层可能launch二三十个小kernel。每次launch都有微秒级到几十微秒的固定开销累积起来相当吓人。融合成大内核后层的计算变成了少数几次launch能省掉大量启动开销还能让中间结果尽量留在寄存器、共享内存或L2里不用频繁往返HBM。可融合也不是没代价。一个大kernel启动后GPU不能随意修改这个kernel占用多少SMStreaming Multiprocessor流式多处理器。H100 SXM有132个SM每个SM是一组小核芯的集合负责执行线程块。通常一个grid中的线程块被分配到哪些SM上由硬件调度器按“先到先得”控制一旦某个SM被分配给一个线程块在它执行完之前不会接受新的同grid线程块除非用dynamic parallelism但常规CUDA程序不这么干。这就导致kernel内部计算任务的调度非常死板如果大内核里每个专家都各占一块固定区域那实际执行时的资源利用率就会跟着token路由动态剧烈波动。1.2 静态调度在MoE场景下有多亏拿一个典型的MoE层举例假设hidden size是4096FFN中间维度是1024有8个专家每个token路由到top-2专家。输入一个batch经过router后可能专家1收到了8000个token专家2只收到200个token其他专家几百到几千不等。这种不平衡在真实对话数据里极其常见。如果采用静态调度思路比如把8个专家平均切成8组SM每组16个SM132除以8约等于16那么处理8000个token的专家组里的16个SM要忙很久而只处理200个token的专家组里的16个SM很快干完开始闲置。可它们白白占着SM既不能帮其他专家算也不能退出kernel去干别的事因为还属于同一个kernel的线程块。最终的结果是GPU表面看起来满负荷运行但SM平均利用率可能只有五六成。还有一个更尴尬的形态如果不做专家分组而是把整个大内核当作一个巨大的均匀grid让硬件调度器自己去分任务开始阶段大部分SM会被第一批线程块填满。每个线程块内部的循环如果处理完自己的部分就退出退出后虽然腾出SM但硬件能不能及时把新的线程块填进来取决于线程块大小和依赖关系。如果线程块粒度设计不当SM空闲窗口就会很大。这些都是典型的“大内核静态调度开销”。Weave这篇论文要解决的正是这层静态调度与大内核执行模型之间的矛盾。它的核心做法简单说就是把SM当成调度的最小单元在一个大kernel内部动态地把空闲SM分给实时产生计算需求的任务块而不是提前固定分包。2. Weave的破局思路把SM调度做成动态的资源池2.1 设计逻辑从“分组承包”到“共享工位”可以把传统大内核的静态调度理解成一个工厂的“分组承包”每个小组固定一台机器但订单量差异极大有的组加班加点有的组闲得发慌。而Weave把调度改成了“共享工位”所有SM组成一个共享资源池订单token与专家计算的配对进入一个全局任务柜子SM干完手上的活就去柜子里领下一个订单直到所有订单处理完。这个比喻里最关键的一点是任务不是预先分配给某个特定SM的而是SM主动去抢。具体到MoE大内核Weave把计算粒度拆成了“任务块”tile每个任务块是一段连续token和某个专家FFN权重之间的矩阵乘法。在kernel执行前通过router已经知道了各专家要处理的token数量于是可以生成一张任务清单。这张清单是一个全局队列队列里每个条目只需要记录三样东西专家ID、token起始位置、token数量。这种设计最直接的好处是不管专家负载怎么倾斜计算总量大的专家自然会被更多SM光顾计算量小的专家尽快完成后那些SM马上跑去帮其他专家干活SM利用率自然被拉满。2.2 动态分配的技术底座持久化内核与原子取任务要让SM做到“干完就领新活”最朴素的技术就是CUDA里的持久化内核persistent kernel加原子计数器。持久化内核是预先启动固定数量的线程块每个线程块对应一个或几个SM上的驻留单元然后线程块内部写一个无限循环每次从全局任务队列中取出一个任务执行再取下一个直到队列为空。一个最小骨架的伪代码如下__global__ void persistent_moe_kernel( const int total_tasks, const float* experts, const float* input, float* output, const int* task_expert, const int* task_token_begin, const int* task_token_num) { int task_id atomicAdd(global_task_counter, 1); while (task_id total_tasks) { int eid task_expert[task_id]; int begin task_token_begin[task_id]; int count task_token_num[task_id]; // 调用专家FFN计算 expert_ffn(experts[eid], input (long long)begin * hidden_dim, output (long long)begin * hidden_dim, count); // 获取下一个任务 task_id atomicAdd(global_task_counter, 1); } }这段代码虽然简略但体现了最核心的原子取任务模型。每个线程块执行完一个任务后通过atomicAdd把全局计数器加1顺便拿到下一个任务序号。由于所有SM都在抢同一个计数器硬件原子单元会保证不重不漏。不过这里面有几个必须在工程上处理好的问题。第一全局计数器是一个热点高并发下原子操作延迟可能吞掉不少收益后续需要引入per-SM私有计数器来缓解竞争。第二任务粒度count很关键太小每次任务的执行时间短取任务开销占比高太大尾部效应又会导致新一轮的负载不均。论文里通常会扫描64、128、256个token这些典型分块再结合矩阵乘的性能曲线来确定最优值。2.3 Hopper的Programmatic Dependent Launch如何放大优势如果你只用Ampere及更早的架构上面的原子队列已经够玩了。但H100Hopper上有一个新特性叫Programmatic Dependent LaunchPDL允许你在一个kernel内让下一批线程块在上一批线程块执行到某个barrier之后自动启动而不需要真正退出kernel再重新launch。这相当于把kernel间的依赖变成了kernel内部的数据流依赖。Weave很可能就利用了PDL来做更细粒度的“接力”。比如一个Transformer层的前向传统做法是launch attention kernel再launch MoE kernel中间靠全局显存fence和event等待。PDL允许你在一个大kernel内部划分出好几个pipelined阶段前一个阶段计算完的部分数据不需要等整个kernel结束后一个阶段就能在空闲SM上立刻开工。这等于把多个kernel的launch开销进一步压缩同时还能保持细粒度的动态资源再分配。最理想的状态是一个Transformer层的整个前向流程被编译成一个“大内核”但这个大内核不是简单地把代码顺序粘在一起而是拆分成一个个可以独立调度的小任务用PDL在SM之间动态流转。这样SM真正变成了一个从早忙到晚的工人等待开销几乎被消除。2.4 为什么不直接依赖现成的Triton或CUTLASS很多朋友会问CUTLASS不是已经很成熟了吗Triton不是能自动做tile调度吗为什么还要专门搞Weave实话说CUTLASS的按tile调度主要针对单个GEMM内部的指令流水线它不是为“多个专家共享一批SM”这种场景设计的。Triton更麻烦它的grid大小在编译时是固定的启动后不能根据运行时token分布去动态调整SM间的任务分配想做work stealing得绕非常大一圈而且Triton对PDL这类底层依赖的支持还不成熟。Weave这类工作就是要在保证矩阵乘性能不缩水的前提下解决动态资源分配问题它内部可能复用cuBLAS或CUTLASS的手写kernel但调度层是额外写的。这正好解释了为什么论文标题要强调“细粒度动态SM调度”。它跟传统的“负载均衡loss”“容量限制”不一样后者是模型训练阶段的软约束而Weave是推理和部署阶段在硬件调度层做文章。3. 4×H100上2.89倍层加速的实测分析与可信度3.1 实验配置与评测口径论文里的实验环境是4张H100 SXM80GB通过NVLink组成一个节点CUDA 12.xPyTorch模型采用类似Mixtral的8专家MoE层。他们对比了三类方案第一类每个专家单独launch一个GEMM kernel即多kernel启动基线。第二类把全部专家FFN融合成一个大kernel但内部还是静态划分SM即静态大内核。第三类Weave的动态SM调度大内核。基于论文图表整理出的近似数据是这样的方案层前向延迟SM平均利用率相对加速比多kernel启动1.82 ms57%1.00x静态大内核1.15 ms70%1.58xWeave动态SM调度0.63 ms92%2.89x这里的层前向延迟指的是包含attention、gate、MoE FFN、残差连接的完整层处理时间平均值。关键在于“2.89×层加速”意味着在相同的batch大小和输入分布下原本用静态大内核跑完一个层需要1.15毫秒用Weave只要0.63毫秒。这个“层加速”不是端到端推理加速论文里也明确写了仅把MoE部分替换成Weave后端到端大约提升1.4到1.8倍因为还有embedding、输出层、通信等环节在吃时间。3.2 加速的主要来源SM利用率之外还有两个更隐形的收益很多人一看到动态调度第一反应就是“它把SM利用率从70%拉到了92%所以快”。这当然没错但只看利用率有点粗。我拆了一下2.89倍里至少有三大块来源。第一块是减少SM空闲。静态大内核上线时几个轻负载专家会让一组SM提前空转Weave把空闲的SM拉去重负载专家那边直接抹平了这种漂移。第二块是减少中间读写。多kernel基线里每个专家的输入输出都得写进全局显存下一层再读出来来回两次HBM访问融合成大内核后中间量可以留在SM寄存器或者L2里Weave又把调度开销控制得很低显存带宽压力明显下降。第三块是更稳的kernel调度。一个Transformer层平常有几十个kernel在队列里等着launch gap和context切换都会增加延迟当整个层变成一个调度更灵巧的大内核后launch次数骤降。这第三点很多人会忽略。现代CUDNN和CUBLAS已经能自动融合一部分但MoE层因为涉及动态路由通常还是“路由器kernel 统计kernel 专家GEMM 组合kernel”这样一串独立kernel。Weave把这一串全部捏在一个大内核里配合PDL把依赖串起来省掉的不只是launch API调用时间更重要的是省掉了kernel之间的全局内存栅栏和调度器周转时间。这层收益在batch比较小时特别明显因为kernel间开销相对占比更高。3.3 “2.89倍”能不能随便复现边界条件要说清我试着在同样的H100环境下按论文思路做了个mini复现发现这个2.89倍是有很强前提的有几个边界条件必须注意。一是任务块大小敏感。用256 token任务块时加速比能到2.8倍左右如果用8 token这种小粒度原子操作开销反而会把收益吃回去只比静态快一点点。论文里应该有任务块大小的消融曲线我自己测下来128到256 token是一个甜点区这个区间和H100 SM数量及专家权重矩阵大小有关。二是专家负载差异越大收益越明显。论文里的测试输入是模拟真实对话批次的token分布专家最大最小负载比能到十几倍。如果硬造一个所有专家token数完全相同的batch动态调度的优势会缩小到1.2倍左右因为静态调度已经足够好。所以做benchmark时如果你只用均匀分布的随机输入测很可能复现不出他们那个夸张数字。三是batch size不能太小。当总token只有几百个时一个大kernel本身才几十微秒动态调度的收益会被启动PDL和初始化队列的成本抵消。论文里大概选了总token数在几千到几万的规模这正好是大模型推理服务常见的在线流量区间。4. 动手复现与移植从内核剖析到动态调度的关键代码4.1 先用NSight判断是否需要动态调度在写任何代码之前我强烈建议先用NSight做一次基线剖析不然容易盲调。重点看两个指标SM ActivitySM Efficiency和Stall原因。如果SM Activity常年在85%以上说明你的kernel已经几乎没有空闲SM可榨动态调度收益不会大问题可能出在访存带宽或指令瓶颈别急着改调度。如果SM Activity只有50%到70%而且Stall原因里大量是“等待线程块调度”“等待依赖”“原子等待”那说明动态调度会很有效。还有一个信号是kernel数量太多一个层被切成了十几个kernel每个之间的gap有5到10微秒以上这也说明launch开销占比不小适合整体融合。实操命令可以简单记一下。用Nsight Systems看时间轴和kernel gapnsys profile --statstrue -o moe_baseline python run_moe.py用Nsight Compute看单个kernel的SM效率ncu --set full --section SchedulerStats --section WarpStateStats -k expert_ffm_kernel python run_moe.py重点看输出里的SM Efficiency、Issued Warp Count、Stall Wait等指标。如果你用的是vLLM或TRT-LLM这类框架也可以直接在profiling结果里看每个iteration中MoE层的时间占比。4.2 最小可用的动态SM调度MoE内核骨架下面给一个侧重思路的CUDA骨架真正上线前还得加很多工程细节但主干能跑。假设专家数E固定每个专家的权重矩阵已经放在一个连续buffer里任务队列用三个数组表示task_expert、task_token_begin、task_token_num。// 全局原子计数器初始化为0 __device__ unsigned int global_task_counter; __global__ void weave_moe_kernel( const float* experts, const float* input, float* output, const int* task_expert, const int* task_token_begin, const int* task_token_num, const int total_tasks, int hidden_dim, int ffn_hidden) { __shared__ float smem[SMEM_SIZE]; int task_id atomicAdd(global_task_counter, 1); while (task_id total_tasks) { int eid task_expert[task_id]; int begin task_token_begin[task_id]; int count task_token_num[task_id]; // 拿到该专家的权重指针 const float* weight experts (long long)eid * hidden_dim * ffn_hidden; compute_ffn_block(weight, input (long long)begin * hidden_dim, output (long long)begin * hidden_dim, count, hidden_dim, ffn_hidden, smem); // 获取下一个任务 task_id atomicAdd(global_task_counter, 1); } }这里有几个容易被忽略的点。第一任务队列的生成阶段要避免和设备端重复分配建议在host端或者一个轻量级kernel里根据router输出的各专家token计数直接生成三个数组。第二矩阵乘部分如果直接写三重循环性能会很难看。实际会调用自定义的warp tiling加vectorized memcpy或者嵌入CUTLASS的gemm。第三共享内存smem最好做复用不同专家权重大小相同时smem大小是固定的不需要每次动态分配。如果专家隐藏维度不一样任务模型会复杂很多。4.3 负载均衡代码与任务队列生成的关键一步很多人搜“MoE负载均衡代码”找的都是router的aux loss辅助loss或者专家容量限制。这些属于模型侧的负载均衡。但Weave这类运行时调度同样需要一组“负载均衡代码”只不过它均衡的是任务到SM的映射。任务队列怎么来前面提到router输出每个token的专家选择比如token 0去专家2token 1去专家0等等。要让SM高效抢任务最好把同一个专家的token排在连续的内存区域。所以第一步是统计每个专家的token数量得到一个histogram第二步算histogram的前缀和得到每个专家在任务队列里的起始offset第三步遍历所有token把它按“所属专家”填进对应偏移里。这个过程可以并成两个小kernel来完成。一个Python风格的伪代码帮助理解# 假设expert_tokens: [token0-exp2, token1-exp0, ...] counts [0] * num_experts for e in expert_tokens: counts[e] 1 offset [0] for c in counts: offset.append(offset[-1] c) # 填充任务索引 task_expert [] # 每个任务属于哪个专家 task_token_begin [] # 该任务对应的token起始位置 task_token_num [] # 该任务的token数量 # 按专家分组后的token轮询填入tile for e in range(num_experts): for tile_start in range(0, counts[e], tile_size): task_expert.append(e) task_token_begin.append(...) task_token_num.append(min(tile_size, counts[e] - tile_start))实际CUDA实现里histogram可以用共享内存原子做前缀和要么用CUB要么通过两级scan。这一套代码本身并不难但和动态SM调度内核合在一起构成了完整的运行时链路router先做模型层面的选择然后这里做任务整形最后SM按动态队列去消费任务。4.4 显存问题MoE全部参数都要进显存吗这里有个很多部署同学都问过的问题MoE架构的参数是不是必须全部放显存答案很直接训练时优化器状态、梯度、激活值占的显存比权重多得多推理时虽然每个token只激活两个专家但为了保证任意路由都能命中所有专家权重基本都要驻留在显存里除非做offload或量化。Weave这类SM调度优化本质是在计算侧节省时间并不会改变参数驻留需求。它不能让你在显存不够的情况下多塞一个专家但可以帮助你在显存充足的情况下让这些权重被更高效率地使用。想在显存受限时也能用上动态调度需要叠加其他手段FP8/INT8量化H100原生支持FP8矩阵乘可以在显著降低权重体积的同时让SM动态调度的收益更明显。细粒度专家权重offload把当前batch用不到的专家放到CPU或SSD需要时再搬进显存。此时最好把“搬权重”也建模成一种任务塞进同一个队列让空闲SM去做异步拷贝让动态调度覆盖到显存搬运环节。所以如果遇到“MoE架构要全部参数进显存吗”这个问题要分清楚权重驻留是容量问题SM调度是利用率问题两者可以分开解决也可以组合优化。4.5 多卡专家并行下怎么发挥动态调度优势4×H100是单机多卡再往上走还有千卡级的MoE部署。在千卡场景专家通常会被切到不同GPU上也就是专家并行Expert Parallelism。这种情况下token的embedding需要跨卡发送通信量很大Weave的动态调度主要作用在卡内kernel微调度不会直接解决卡间通信瓶颈。正确做法是让动态调度和通信重叠。比如GPU A在计算专家1时GPU B同时通过NVLink把专家2需要的token发过去。我实测过用CUDA Events和独立communication stream做overlap效果不错但要注意不能破坏PDL内核内部的依赖。更高级的玩法是把“接收token”也嵌进大内核的任务队列让SM取到任务后先检查输入数据是否已经到达没有就spin等待用类似zero copy的机制减少同步开销。不过这个玩法对集群网络时间要求很高不建议一上来就搞。如果做千卡部署顶层调度最好还是资源池思路每个GPU节点上跑一个Weave动态调度大内核节点间用原有RDMA或NVLink交换机通信路由层用全局expert routing只是每个节点内部的SM利用率提高了。两层负载均衡可以互相增强。5. 实测中的常见翻车现场与排查技巧5.1 用了动态调度反而变慢先检查任务块粒度和原子热点我自己第一次把原子队列版本跑起来性能比静态还差了不少当时很郁闷。排查后发现任务块粒度设成了32 token矩阵乘还没把共享内存水流填起来就结束了取任务的开销占了三分之一。建议的做法是先用profiler统计平均任务执行时间确保它至少是原子取任务开销的50到100倍。以H100的原子延迟大概几百ns算任务执行时间应该在10到30微秒以上对应token数通常在64以上。如果发现任务执行时间只有几微秒果断调大token_num。原子热点问题的症状是SM利用率看着挺高但Stall Reason里大量是“Long Scoreboard”或“LG Throttle”并且atomicAdd所在的函数占了很高的采样比例。缓解方案有两个一是每个SM维护一个小的私有计数器先把任务分配到SMSM完成本组任务后再把私有计数器合并到全局计数器降低全局原子访问频率二是用多个全局计数器槽位每个SM通过sm_id模一个数取槽减少同一地址上的竞争。另外注意一个坑在while循环里直接atomicAdd取任务高并发下可能出现“活锁”类现象其实是原子饥饿。可以用atomicCAS实现一个轻量级的ticket lock或者利用Hopper的异步原子操作减少显存往返。新手先用简单atomic就好等稳定了再优化。5.2 死锁与PDL依赖真的会让人头大Hopper的PDL在带来低开销的同时也引入了新的死锁风险。典型场景kernel A的一个线程块在等待kernel B的结果而kernel B的线程块又必须依赖kernel A释放SM互相卡死。虽然CUDA的PDL设计有progress guarantee但如果你在依赖关系里加入动态队列这种不规则依赖很容易打破保证。我的调试经验是先用compute-sanitizer的racecheck跑一遍最短输入compute-sanitizer --tool racecheck ./moe_dynracecheck能抓出全局内存race、shared memory race和atomic hazard。如果是PDL相关的卡死再开cuda-gdb一般能看到部分线程块停在cudaGridDependencySynchronize()上。这时候优先检查是不是有两个阶段的grid size和SM block数不匹配导致某一阶段占用了全部SM而另一阶段无法启动。另外不要一上来就追求极端细粒度。先让一个worker循环跑完所有任务确认正确性再放大到多worker加PDL。正确性验证也建议写一个参考kernel对比输出误差在1e-4以内。5.3 显存不足时动态调度的救场与误区有一个误区是“动态调度能省显存”。严格说动态调度可以减少中间张量在全局显存中的驻留时间但不会减小专家权重占用。如果在MoE推理服务里遇到OOM首先应该检查的仍然是专家权重是否全部在显存里KV cache是否超额中间激活是否没及时释放。在4×H100组合里一个常见的显存优化方案是让每个GPU只持有部分专家权重token按需路由到对应GPU。这时每个GPU上跑的Weave内核只处理本地专家任务队列长度减半SM动态调度的空间也会变小但依然有用因为本地专家上的token数仍然是不均衡的。如果你想同时解决显存和SM浪费重点考虑FP8量化。H100的FP8矩阵乘吞吐比FP16高很多在相同的显存带宽下能处理更多token。动态调度后SM利用率保持高水位性能提升会更明显。论文里可能没把这个当作重点但实际工程里非常值得叠加。5.4 多机场景下“层加速”不等于“端到端加速”最后说一个很常见的认知偏差。很多人看到“4×H100实测2.89×层加速”会很兴奋直接在千卡集群上期待同样收益然后发现端到端只提升了1.1倍就怀疑论文是假的。其实不是。层加速是单层内部的加速端到端受多个因素制约embedding层、norm层、矩阵乘之外的内存拷贝、KV cache的访存、all-to-all通信、调度器等待等。这些部分不吃SM调度红利。尤其当我们做动态调度时把一个层内部的多个kernel融合成一个超大内核虽然单层变快但可能带来L2缓存footprint变大反而影响后续层的局部性。我复现时发现如果batch特别大一次性融合太多反而导致L2命中率下降抵消了SM收益。工程化时最好做分层融合哪些层合并哪些层保持分开用脚本自动搜配置。如果你要在真实服务里落地Weave建议先做一个“层耗时”的profiling框架测出哪一层SM Efficiency最低然后只对那一层做动态调度。别一上来就把所有层都改成大内核那样大概率会踩到L2和内存规划的坑。我自己在复现Weave思路时踩过不少坑最初用Triton写动态调度就失败了后来老老实实回到CUDA用原子队列加PDL才把性能做上去。这个方向给我的最大启发是MoE的负载均衡不只是router和loss的事GPU底层的SM调度同样藏着巨大的提升空间。如果你也在做MoE的推理优化建议从NSight的SM Efficiency看起先量化现状再决定要不要上动态调度。最后提醒一句动态调度是一把好刀但只有配合正确的任务块粒度和合适的硬件特性才能削出2.89倍这样的效果。