Muse Spark Open超快突破:单请求破千tokens/s

SGLang与ML基础设施的目光正投向规模更大的GPU核——megakernel。有人从零开始,用CUDA C++和PTX为Muse Spark Open写出这样一个megakernel,把Blackwell GPU的单请求解码速度从517 tok/s拉到1109 tok/s。

Decode tokens per second, one request at a time: baseline 517, megakernel 1,109, 2.1 times as fast

Muse Spark Open沿用传统MoE结构,每个token会路由到256个专家中的8个,专家权重以4位NVFP4格式存放,解码层共62个。

megakernel是什么?

普通的GPU核一次只干一件事,服务引擎跑完一个解码步大约要发射1450个核。优化不到位时,GPU在核与核之间就会闲下来。megakernel的思路是让一个核常驻所有SM,多种任务同时在里面跑,张量核、TMA、NVLink这些引擎一直保持满速。

CPUs went from batch jobs to multiprogramming to operating systems; GPUs went from one kernel per operation to megakernels

模型与新的张量并行

模型按TP8切分到8张GPU上。传统TP在两个MoE GEMM之间只做一次all-reduce;新方案则在N维上把两个GEMM都切开,加入fp4 all-gather与bf16 all-gather。通信多了一步,但总通信量反而更低,扩展性也更好。

One decoder layer in six steps, with five numbered stops where it waits for the other GPUs; the expert steps are marked as known only after the router runsThe expert GEMMs across eight GPUs: vanilla TP splits GEMM 1 on N and GEMM 2 on K and ends with one bf16 all-reduce; ours splits both on N, with an fp4 all-gather after GEMM 1 and a bf16 all-gather after GEMM 2

megakernel怎么实现

每层只发射一次核,范围从本层attention一直管到下一层的qkv投影和KV写入。每个核在144个SM上铺开12个warp,KV加载、权重预取、张量核计算、SwiGLU以及all-gather/reduce分别由不同warp承担。

Four slots of a receive buffer: all empty before any store; two values landed and the reader keeps polling; all four landed and the reader goes

性能与投机解码

单步解码耗时从12.9 ms压到5.9 ms,核数量由1450降到97。再配上DFlash V2 draft模型,每步验证16 token,平均接受6.6 token,最终跑到1109 tok/s。

One decode step: baseline takes 12.9 ms (draft 1.16, verify 10.9) with 1,450 kernels; Megakernel takes 5.9 ms (draft 0.44, verify 5.3) with 97 kernels

megakernel的内部优化

做法包括权重全缓存预取、指令缓存管理、多级同步(shared memory mbarrier、HBM flag、NVLink直接store),以及不带fence的数据同步,几项合起来让重叠接近完美。

Three places a prefetched weight can wait: straight into the shared-memory ring while the current op runs; into L2 for what does not fit, like the speculatively prefetched experts; or both, with the ring refilling from L2 instead of HBMTwo blocks of code run back to back: 20 KB + 100 KB fit in an SM's instruction cache of about 128 KB, and the 100 KB block runs at its solo 12.4 µs; 100 KB + 100 KB do not fit, and every pass takes 19.5 µsWhere threads meet to synchronize: two warps in the same SM wait on an mbarrier in shared memory; warps on two SMs of one GPU meet at a flag in HBM, polled; a warp on another GPU stores over NVLink straight into this GPU's HBM, where the reader polls for the dataFence-free data synchronization: with a weight stream in flight, the fence before the flag takes 3.6-4.0 µs; when the data is its own flag, the reader starts as soon as the data lands

展望

接下来,Muse Spark Open会支持更大的批次,产品方向则朝GPU操作系统演进。