新智元报道
编辑:alan上月底,创业公司Groq的产品一夜爆火。
凭借自研的硬件加速器LPU,达成了500个token/s的神级推理速度,当场秒杀了ChatGPT。
Groq提供的响应速度刷新了人们的认知,而这要归功于背后的语言处理单元硬件LPU(language processing unit hardware)。
Groq的研发团队在LPU上应用了创新的硬件架构设计,并配套了强大的编译器。
下面让我们跟随Substack的专栏作家Abhinav Upadhyay一起,一步步揭开Groq LPU底层架构的神秘面纱。
Groq LPU的神秘面纱
论文地址:https://dl.acm.org/doi/10.1109/ISCA45697.2020.00023
论文地址:https://dl.acm.org/doi/abs/10.1145/3470496.3527405
两篇工作分别在2020年和2022年发表在计算机体系结构顶会ISCA上,后一篇还是获奖论文。
这两篇文章解释了Groq的张量流处理器(TSP)的设计和实现,以及他们如何使用TSP构建分布式推理引擎。
尽管没有正式声明,但LPU很可能是基于这个分布式系统来进行设计和扩展的。
那么,我们就首先详细分解一下TSP及其编译器的架构,然后以此为基础来分析Groq如何使用这些TSP,构建可靠且高吞吐量的分布式AI推理引擎。
TSP的架构与传统的CPU或GPU芯片有很大不同,主要目的是为了让TSP硬件更具确定性。
这里就先要提一嘴CPU或GPU的不确定性。
基于微架构的设计,在CPU和GPU上执行指令是不确定的,——即无法保证特定指令何时执行、完成需要多长时间以及何时提供结果。
举个例子,现代CPU一般具有如下设计:
所有这些都使得CPU中指令执行的顺序和时间不确定且难以推理。而GPU还有其他一些非确定性因素,包括缓存、共享和全局内存、动态资源分区等。非确定性带来的问题是,我们很难推理程序的性能,也很难保证最坏情况下的性能限制。因此,Groq为TSP提出了一个全新的设计,高度并行,且没有不确定行为。这消除了硬件的复杂性,使编译器能够获得更大的权力,精确调度和控制指令的执行,保证对程序性能的限制。下面,让我们从内部了解TSP的架构是什么样子的。- 超标量(Super scalar architecture):每个周期能够发出多条指令;
- 乱序执行(Out-of-order execution):以任意顺序执行指令;
- 预测执行(Speculative execution):对于分支,它会猜测分支条件是真是假,并提前预测执行该分支以提高吞吐量(当然如果猜错了,就需要放弃并返回另一条分支);
- 指令流水线(Instruction pipelining):将指令分为多个阶段,以流水线的方式执行,再次提高了指令吞吐量;
- 多级缓存(Multiple levels of caches):CPU有2到3级缓存,可以减少从内存加载数据带来的延迟。
在了解了TSP的架构之后,让我们将注意力转移到它的核心:指令执行。TSP中的指令执行MXM:用于执行矩阵运算
SXM:用于对矢量进行移位和旋转操作
MEM:内存读/写运算
VXM:向量上的算术运算
ICU:指令控制单元,这个有点特殊,就是上图(b)底部那一条水平的蓝色条,它负责获取和调度指令并在其他切片上执行。
- 320个通道的编程抽象:TSP芯片中的每个tile都能够以SIMD方式在矢量的16个单元(16个通道)上运行。垂直切片由20个这样的tile组成,因此总共有320个SIMD通道可供执行;由于TSP硬件中没有非确定性行为,因此编译器可以准确了解每条指令的延迟,以及程序中的数据流(DNN的计算图等)。编译器识别计算任务之间的依赖关系,并分配到TSP的可用功能单元上并行执行。TSP编程模型依赖于两个关键要素:
- 144个独立指令队列:芯片上有144个指令队列,每个周期能够发出一条或多条指令。编译器可以完全控制每个队列中的程序顺序;
- 每个通道64个逻辑流:每个通道可以访问64个逻辑流,可用于移动操作数或结果,其中32个可用于向东移动数据,而另外32个用于向西移动数据;
- 220M全局共享SRAM。
编译器的后端可以跟踪片上任何流的位置和使用时间,称为软件定义硬件。硬件中的确定性数据路径
通过ISA获得的有关指令延迟的信息
从TSP扩展到LPU
TSP是LPU的基础单元。许多TSP以机架的形式组合在一起,形成一个能够提供大量吞吐量的分布式系统。
- 当两个TSP互连时,其中一个TSP将其HAC值传输给对方。然后,对方将该值返回发送方。发送方观察当前HAC值与返回值之间的差值。
- 这个差值就代表了两个设备之间链路的延迟。此过程重复多次,得到两个TSP之间的平均链路延迟。
- 之后,两个设备以父子关系排列。父级定期将当前HAC值发送给子级。子级将平均链路延迟与自己的HAC值相加,并与自己的HAC值进行比较。
- 两个值之间的差值表示由于连续时钟漂移而导致的初始未对准。然后子级调整其HAC值以减小此差异。在多次重复此过程后,两个TSP的HAC值会收敛在一个小邻域内,表示链路延迟的抖动。
- 协议允许两个TSP相互同步,并且可以通过在网络中建立生成树来扩展TSP多跳网络。
- 在单个TSP级别,有几个独立的功能单元和144个独立的指令队列。为了同步它们,TSP支持SYNC和NOTIFY指令。SYNC指令将所有指令队列置于停放状态,其中一个队列充当通知程序。当通知器发出 NOTIFY指令时,该指令被广播到芯片上的所有队列,此时它们被同步并恢复操作。
- 对于多TSP系统,两个TSP使用HAC相互同步,另外每个TSP都支持DESKEW指令,用于停止处理任何后续指令,直到TSP的HAC溢出。
- 要扩展多跳系统,可以在生成树的每个hop上重复执行以上方案。
编译器在软件计划网络中的作用
到目前为止,编译器能够对TSP内以及整个网络中的数据移动进行周期准确的了解。编译器知道在源TSP上注入向量的确切时间以及它到达目标TSP的确切时间,称为软件计划网络。
编译器不是动态管理数据流,而是在编译时静态解析所有内容。