Meta 的 816B 稀疏模型搬上单台四芯片 TPU:整步合成一个内核,单请求每 token 快 1.94 倍
Meta 的开源混合专家模型 Muse Spark Open 总参数 816B、每 token 约激活 42B,此前在 TPU v7 上服务时,一次验证步的耗时是单纯从 HBM 流式读取权重所需时间的 2.9 倍。SGLang-TorchTPU 把一次解码或验证步的全部 62 层写成一个 Pallas 程序,让权重、KV 与跨设备传输都排在计算之前发出。单请求下每输出 token 从 3.78 ms 降到 1.95 ms,两个并发请求时为 1.85 倍,并发继续增大时收益收窄。