Measuring and Reducing WebGPU Dispatch Overhead for LLM Inference
本文揭示了 WebGPU 的调度开销而非内核质量是浏览器中单批次大语言模型(LLM)推理的主要瓶颈,证明了由于同步混淆导致简单的测量值高估了成本,并得出结论:通过摊销来减少调度次数是最有效的优化策略。
原始论文采用 CC BY 4.0 许可(http://creativecommons.org/licenses/by/4.0/)。 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性,请参阅原始论文。 阅读完整免责声明
想象一下,你正试图在一台电脑上运行一款庞大且复杂的视频游戏,但你必须通过一个非常严格、注重安全的管理者来操作,这个管理者不允许你直接接触硬件。这就是在浏览器中运行人工智能(特别是大语言模型,简称 LLM)的世界。这些模型是聊天机器人背后的“大脑”,它们可以写故事、解数学题并进行对话。为了让这些模型能在你的笔记本电脑或手机上快速运行,而不需要超级计算机,开发者使用了一种名为 WebGPU 的特殊工具。你可以把 WebGPU 想象成一个通用的翻译官,它让你的浏览器能够与你的图形卡(通常用于渲染视频游戏的部件)进行对话,从而处理 AI 所需的大量数学运算。
然而,这里有一个问题。过去,当开发者试图让这些 AI 模型变得更快时,他们专注于让单个数学步骤(称为“内核”,kernels)更加高效,就像是在磨练汽车的引擎。但这篇论文提出了一个不同的问题:如果车本身没问题,但驾驶员在上下车上花费了太多时间呢?在浏览器世界中,每一个数学步骤都需要一次“调度”(dispatch)——即从浏览器发送给图形卡的请求,以开始工作。这个谜团在于:到底有多少时间是浪费在仅仅发送这些请求上,而不是在进行实际的数学运算?理解这一点至关重要,因为如果我们在请求电脑工作上浪费了太多时间,那么无论数学计算多么聪明,聊天机器人的反应都会显得缓慢且迟钝。
“走走停停”的交通堵塞
这篇论文的研究人员发现,大家一直以来测量这些 AI 请求速度的方法都错了。想象一下,你正在计时一名快递员送货所花费的时间。如果你从他们离开仓库开始计时,直到他们开车到达住所、放下包裹,然后再开车回到仓库去取下一个包裹,那么你测量的是整个往返行程。但在 AI 的现实世界中,驱动程序并不会在每送一个包裹后都开回仓库。他们会一次性送下一叠包裹,只在最后才开回仓库。
论文指出,之前的测量方法就像是在为每一个包裹都计时那次完整的往返行程。他们混淆了“发送请求的时间”(调度)与“等待计算机说‘好的,我做完了’的时间”(同步)。这种“等待时间”非常巨大——大约是 450 微秒的停顿。当研究人员将这个等待时间添加到每一个步骤中时,他们认为发送请求的成本比实际情况高出了约 20 倍。
通过使用一种称为“顺序调度”(sequential-dispatch)的新方法,作者弄清楚了如何仅对发送请求这一动作进行计时,而不包含中间漫长的等待。他们发现真实的成本要低得多:在某些系统(Vulkan)上为 24–36 微秒,在其他系统(Metal)上为 32–71 微秒。有趣的是,无论计算机使用的是 "float32" 还是 "float16" 数字(存储小数的两种不同方式),这个成本都是相同的,这证明了延迟来自于浏览器的规则,而非数学运算本身。
真正的瓶颈:过多的停顿
一旦知道了单个请求的真实成本,团队便问道:“这真的重要吗?”为了找出答案,他们进行了一项受控实验。他们采用了一个标准的 AI 模型,并改变了它的封装方式。他们没有向图形卡发送 876 个微小的请求来处理一个单词,而是通过“融合”(fused,即将步骤粘合在一起)了一些步骤,使得显卡只需要接收 564 个请求。
关键点在于:他们并没有让请求内部的数学运算变得更快。他们没有优化代码,也没有减少内存使用。他们只是减少了浏览器需要敲图形卡之门的次数。
结果如何?AI 的速度提升了 53%。生成第一个单词所需的时间从 71.4 毫秒 降至 41.6 毫秒。
这项实验证明,在最常见的设置下(即每次处理一个单词,即“批大小为 1”时),最大的问题并不是数学运算太慢或内存太满。问题仅仅在于有太多的“敲门声”。作者明确排除了“更好的数学代码”或“更少的内存使用”导致提速的可能性。唯一改变的变量就是调度的次数。
这对未来意味着什么
论文的结论是,如果我们希望浏览器中的 AI 聊天机器人运行得更流畅,我们就不应该再试图完善每一个单独的数学步骤,而应该专注于将它们组合在一起。这就像是意识到,为了让快递车更快到达目的地,你不应该只是让司机跑得更快,而应该确保他们一次能搬运更大的箱子,这样就不用频繁往返。
作者建议,解决方案在于“调度摊销”(dispatch amortization)——这是一个高级说法,意思是我们需要将那些“敲门声”的成本分摊到许多任务中,这样延迟才不会造成严重的负面影响。他们指出,这可能不仅需要 AI 运行软件层面的改进,甚至可能需要 WebGPU 规则本身的变革,例如允许浏览器接受一个“命令图”(command graph,即预先规划好的路线),而不是逐一检查每一个步骤。
虽然这些发现是基于特定硬件(如 NVIDIA RTX 5090)和一种特定的 AI 运行方式,但其传达的信息是明确的:对于目前的浏览器 AI 而言,秘诀不在于一个更快的引擎,而在于更少的停顿。
您所在领域的论文太多了?
获取与您研究关键词匹配的最新论文每日摘要——附技术摘要,使用您的语言。