Muse Spark Open解码破千tokens/s：megakernel怎么做到的

SGLang和ML基础设施正在往更大规模的GPU核，也就是megakernel方向走。有人从零开始，用CUDA C++和PTX给Muse Spark Open写了一个megakernel，把Blackwell GPU的单请求解码速度从517 tok/s拉到1109 tok/s。

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

什么是megakernel？

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

模型与新张量并行

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

megakernel实现

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

性能与投机解码

单步解码耗时从12.9 ms降到5.9 ms，核数量从1450减到97。再配上DFlash V2 draft模型，每步验证16 token，平均接受6.6 token，速度达到1109 tok/s。

megakernel内部优化

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

展望

后续计划支持更大的批次，并朝GPU操作系统的方向演进。