LLM模型如何集成至芯片中

或者说，如何实现每秒生成 17,000 个 token？

一家名为 Taalas 的初创公司最近发布了一款 ASIC 芯片，能够以每秒 17,000 个 token 的推理速度运行 Llama 3.1 8B（3/6 bit 量化）。这相当于在一秒钟内写完大约 30 页 A4 纸的内容。他们声称，与基于 GPU 的推理系统相比，这款芯片的总体拥有成本低 10 倍，耗电量低 10 倍。没错，速度也比当前最先进的推理系统快了大约 10 倍。

我尝试读了他们的博客，发现他们真的把模型权重「硬连线」到了芯片上。起初，这对我来说并不直观。我来自软件行业，对 LLM 只有业余爱好者程度的了解，实在想不明白怎么能把一个 LLM 直接「印」到芯片上。于是，我决定深入研究多篇博客文章、LocalLLaMA 上的讨论，以及相关的硬件概念。这个过程比我原先想象的有趣得多，因此有了这篇文章。

Taalas 成立仅两年半，这是他们推出的第一款芯片。Taalas 的芯片是一款固定功能 ASIC（Application-Specific Integrated Circuit，专用集成电路）。它有点像 CD-ROM、游戏卡带或印刷出来的书，只能容纳一个模型，并且无法重写。

NVIDIA GPU 是如何处理这些工作的？（低效原理 101）

LLM 由一系列按顺序排列的 Layer 构成。例如，Llama 3.1 8B 有 32 个 Layer。每个 Layer 的任务都是对输入做进一步处理和提炼。每一层本质上都是巨大的权重矩阵，也就是模型的「知识」。

当用户输入一个 prompt 时，它会被转换成一个由数字组成的向量，也就是 embeddings。在普通 GPU 上，输入向量首先进入计算核心。随后，GPU 从 VRAM/HBM（GPU 的 RAM）中获取 Layer 1 的权重，执行矩阵乘法，再将中间结果（也就是 activations）存回 VRAM。接着，它获取 Layer 2 的权重和上一层的结果，完成运算后，再次将结果保存到 VRAM。为了生成一个 token，这个循环会一直持续到第 32 层。然后，为了生成下一个 token，GPU 又要把整个 32 层的流程重新走一遍。

因此，数据持续来回传输，不仅会让内存总线引入延迟，还会消耗大量能源。这就是内存带宽瓶颈，有时也会被宽泛地称为 Von Neumann bottleneck，或者「memory wall」。

Taalas 完全绕开了这堵墙。他们直接把 Llama 3.1 的 32 个 Layer 按顺序刻在了一块芯片上。本质上，模型权重变成了蚀刻在硅片中的物理晶体管。

更重要的是，他们还声称发明了一种硬件方案，只需要一个晶体管，就能存储一份 4-bit 数据，并完成与之相关的乘法运算。下文中，我会把它称为他们的「magic multiplier」。

现在，当用户的输入到达后，它会被转换为一个向量，然后流入组成 Layer 1 的物理晶体管。数据通过他们的「magic multiplier」完成乘法运算，而运算结果不会被保存到 VRAM 中。电信号会直接沿着物理导线流入 Layer 2 的晶体管——据我理解，中间会经过 pipeline registers。数据会持续不断地在硅片中流动，直到最终的输出 token 被生成。

所以，他们完全不使用 RAM 吗？

他们不使用外部 DRAM/HBM，但确实会使用少量片上 SRAM。为什么是 SRAM？受成本和复杂度限制，制造商通常不会把 DRAM 和逻辑门混合制造在一起，这也是 GPU 配备独立 VRAM 的原因。（另外，SRAM 目前没有遭遇供应链危机，DRAM 则有。）Taalas 使用片上 SRAM 来存放 KV Cache，也就是当前对话所需的临时记忆或上下文窗口，同时还用它保存用于 fine-tuning 的 LoRA adapters。

但是，为每个模型制造一款定制芯片，不是非常昂贵吗？

从技术上说，确实如此，我也读到过很多提出这一观点的评论。不过，Taalas 设计了一款基础芯片，其中包含规模庞大、通用的逻辑门和晶体管网格。为了把某个特定模型映射到芯片上，他们只需要定制最上面的两层或两道 mask。虽然这个过程依然很慢，但已经比从零开始制造芯片快得多。他们用了两个月时间为 Llama 3.1 8B 开发芯片。在 AI 世界里，一周就像一年，两个月显得慢得离谱；但在定制芯片领域，这个速度据说已经快得惊人。

作为一个没有大型 GPU、只能困在笔记本电脑上运行本地模型的人，我真心期待这类硬件能够尽快实现大规模量产。