Rohan亲自测完Kog AI的3000 token/s,把单用户推理速度拉高了10-30倍,这套monokernel设计可能改写低延迟推理的玩法,做实时AI产品的团队必须盯紧。
Kog团队在标准数据中心GPU上实现了极高的单用户推理速度,在8× AMD MI300X GPUs上达到3,000 tokens/s,在8× NVIDIA H200上达到2,100 tokens/s。相比常规推理速度(约100-300 tokens/s),实现了10-30倍提升。其核心思路是将LLM解码视为内存流问题,通过协同设计monokernel、重建同步机制、针对性内存访问映射及采用延迟张量并行的Laneformer模型架构,消除了传统流程的阻塞点。
我必须亲自测试才能相信这不可思议的推理速度。
在标准数据中心GPU上,为1位用户提供3,000 tokens/s的速度。
他们利用了GPU生成模型token时一个隐藏的效率差距。
@Kog__AI 刚刚在8× AMD MI300X GPU上实现了3,000 tokens/s,在8× NVIDIA H200上实现了2,100 tokens/s(FP16,不使用推测性解码)。
他们的技术预览基于一个2B模型,并展示了其技术将如何以类似速度扩展到前沿大型MoE模型。
这个数字非常巨大,因为对于2B到8B模型,正常的低批次GPU解码通常每请求只有100到300 tokens/s,因此Kog声称用户的实际感受速度提升了10倍到30倍。
他们的诀窍是:将大语言模型解码视为内存流问题,而不是主要数学问题,从而获得这种速度。
对于1位用户、批次大小为1的情况,GPU并不会像训练或大批次服务那样执行大规模高效的矩阵乘法运算;而是为每个新模型token反复从高带宽内存中拉取模型的活跃权重,因此速度取决于这些权重能否持续流畅地流动。
常规推理栈不断打断这种流动。
它们为模型的不同部分运行许多独立的GPU程序,在内存中搬运中间结果,在同步点等待,与CPU通信以进行调度或采样,然后逐个token重复这个过程。
Kog的解决方案是协同设计三个通常被分别调优的部分:运行时、底层GPU代码和模型架构。
最大的工程举措是monokernel,整个解码过程作为一个常驻GPU的单一程序运行,包括采样,因此系统不需要不断停下来进行内核启动、CPU调度和中间内存往返。
他们还重新构建了同步机制,因为其自身的测量表明,网格同步消耗了约35%的模型token生成时间;他们不再让每个计算单元等待一个宽泛的屏障,而是让每个单元只等待其确切需要的数据。
在AMD MI300X上,他们还根据小芯片布局映射内存访问,因为内存延迟会因发出请求的晶片不同而变化。
此外,他们的Laneformer模型使用了Delayed Tensor Parallelism,让跨GPU通信在后台进行,而不是阻塞每一层。