Tenstorrent芯片架构浅谈

1、背景

        近些年，由于深度学习相关的算法、模型和应用场景均处于高速迭代发展阶段，还未完全收敛，因此对应的计算硬件底座为适配上层应用自然也需要不断更新。其中芯片公司Tenstorrent的芯片架构别具一格，本文尝试一探究竟。

        Tenstorrent成立于2016年，是一家总部位于加拿大多伦多的AI芯片公司，以 AI 加速器起家，现在似乎也瞄准了 RISC-V 架构的处理器。Tenstorrent宣布它正在使用 RISC-V 指令集开发功能强大的处理器，该处理器将拥有自己的架构，称为 Ascalon，并在服务器中与 x86 和ARM 竞争。Ascalon架构将是一个非常宽的乱序核心，承诺更高的IPC。因此，它遵循与 Apple 架构在移动SoC 和 M1/M2 处理器中取得成功的相似路径。

        Tenstorrent实际上基于看似相同的基础开发了几个不同的核心。这些内核在大小、复杂性以及性能方面各不相同，因此它们适用于不同的目的。最弱的版本有 2 宽指令解码，然后是 3 宽、4 宽、6 宽（Alastor架构）版本，最后是 8 宽Ascalon 内核，每个周期能够处理 8条指令。根据PPT，虽然该公司主要谈论在服务器中的使用，但 Ascalon 核心也可用于HPC 和笔记本处理器。

2、产品

        Tenstorrent共设计出3款芯片，其中Jawbridge是一款小型测试芯片，Grayskull和Wormhole则是对外商用芯片，可覆盖训练和推理场景。

Products	Jawbridge	Grayskull	Wormhole
Manufactured	2019	2020	2021
IC Process	GF 14nm	GF 12nm	GF 12nm
Peak Performance	1TFLOPS at FP16, 4TOPS at INT8	92TFLOPS at FP16, 368TOPS at INT8	110TFLOPS at FP16, 430TOPS at INT8
Compute cores	6	120	80
SRAM	6MB total - 1MB/core	120MB total - 1MB/core	120MB total - 1.5MB/core
DRAM Capacity	unknown	8GB	12GB
DRAM Bandwidth	unknown	100GB/s	384GB/s
I/O Interface	1ch. LPDDR4, PCIe Gen4x4	8ch. LPDDR4, PCIe Gen4x16	16 ports of 100G Ethernet, 6ch. GDDR6, PCIe Gen4x16
Scale-out Bandwidth	unknown	192GB/s for NoC	400GB/s for Ethernet
Board Power(TDP)	1.5W	65W / 75W	150W

Wormhole芯片架构示意图（T代表Tensix core，E代表Ethernet接口、G6代表GDDR6）

以最新的Wormhole芯片为例，其主要包含：

• 计算核心Tensix core：
  • 5个标量RISC-V CPU：用于运行时发射调度指令和数据
  • SRAM：作为private memory，用于保存本地计算结果
  • Packet Manager数据包管理器：包含Tensor形状变换、数据传输和数据路由功能
  • Packet Compute Engine数据包计算引擎：INT8算力3TOPS，FP16算力0.75TFLOPS
• DRAM：片上存储，可被Tensix全局访问
• Ethernet接口：用于芯片间互联，便于scale out
• PCIE接口：连接Host，与Host进行数据传输
• Network on Chip(NoC)：数据互联接口，在Ethernet中也含有NoC - System Manager系统管理单元：4个RISC-V CPU，负责中断、时钟、温控等基础功能，及数据路由功能

3、技术特点

        Tenstorrent的芯片架构设计目标是解决模型在训练或推理时无法高效灵活扩展（scale out）的问题，提出2个核心技术点：

3.1 摒弃传统的核间共享式内存架构，采用Multicore Private Memory Model

        这样核间及芯片间的数据通信模式在软件栈上是一致的，极大降低了模型分布式训练或推理的开发难度。在用户角度看使用的硬件是一个Device。由于多核间的数据通信不依赖于共享内存，因此芯片重点在硬件层面和软件层面分别加强了数据通信能力。在硬件层自研片上互联Network on Chip(NoC), NoC是2D双向环路结构，支持片内所有核心和片外芯片间通信。在最新的Wormhole架构芯片中，又新增了Ethernet接口专用于芯片间互联。为了实现芯片内和芯片间软件层面通信接口的一致性，Tenstorrent在软件层设计了基于Pipe的数据路由方式，由Graph Compiler生成当前模型下最优的数据路由策略。对比各大AI芯片公司官宣都会大篇幅介绍计算部件，比如各类DSA架构，Tenstorrent的芯片架构介绍文章中更多是在描述非计算硬件部件。从这里也可以看出，数据通信是一等公民（This communication oriented architecture realizes our vision of converged networking and accelerated AI compute on a single device.），片内和片间高效的数据通信能力是Tenstorrent一直强调的重点。

3.2 动态执行

3.2.1 运行时数据压缩

        Tenstorrent芯片内有专门用于数据压缩和解压缩的硬件部件，称为Packet Manager Engine，个人认为Packet Manager Engine的引入是该芯片的重要特色。为最大化利用计算资源且最小化使用内存，在计算单元之外，数据都是以压缩方式存储和流动的，这样就可以设计较小容量的private memory，在芯片内多核间及芯片间的数据通信量也随之降低，进而从整体上可以获得更高的性能功耗比。另外，Packet Manager还可以处理reshape/flatten等tensor形状变换操作，且该操作可以与Compute Engine并行执行，时间上可以overlap。

3.2.2 条件执行

        该特性尤其在稀疏门控专家混合模型（Sparsely-Gated Mixture of Experts）中应用效果显著。这类模型通过门口网络实现只激活模型的部分结构，因此是一种显著增加模型容量和能力而不必成比例增加计算量的方法。复杂的逻辑控制对CPU而言可以轻松应对，但是对AI芯片而言，大量资源都用于计算部件，逻辑控制相关的部件功能较弱，因此在编程时一般是需要尽量避免出现大量逻辑控制语句。在Tenstorrent芯片架构中，片上的逻辑控制单元和计算单元均可以高效运行，避免了CPU fallback的问题，此处也是该芯片的重要特色。

3.2.3 稀疏计算

        除了支持常规的对权重进行稀疏化之外，还支持对激活值进行分块稀疏，进而降低计算量。

3.2.4 动态混合精度

        可以在运行时或AOT阶段设置每个算子的计算精度。该技术在主流深度学习框架中都有实现，因此个人认为这项技术并无特色。

3.3 总结

        Tenstorrent通过软硬协同设计方式，将数据并行和模型并行的部分功能实现下沉到硬件层，有效解决了横向扩展问题，这样就可以替代当前主流深度学习框架在分布式实现方面的大量编码工作，进而降低了深度学习框架的开发和使用门槛。同时，硬件的变化并没有降低软件栈的通用性，其软件栈支持PyTorch等主流框架。另一方面，芯片具有高度模块化，多个芯片可通过标准以太网端口连接在一起，进而扩展成大型AI网络。由于芯片内已集成NoC，因此这种扩展并不需要额外的交换机，因此扩展灵活度很高。

        对比GPU，Tenstorrent官网提到其芯片编程更友好、扩展性更好和更擅长处理稀疏计算和控制流计算场景。

4、芯片介绍

4.1  Tenstorrentt软硬件架构介绍

        Tenstorrentt跟其他架构最大的差别在于MAC核的数量。Tenstorrent有整整120个核，这些核都比我们之前接触的TPU、含光或者Groq要来的小的多，大概架构长这样：

        这个图片里面紫色的CPU不是我们电脑上面的CPU，而是一个很小的RISC的核。小核有一个很大的优势，就是conditional computation。这个芯片相对别的玩家来说TDP要低。

At a peak rate of 368 TOPS, the chip runs on just 65W

        一个小核里面总数大概是一千个int8 的MAC（比如32*32），不过他们也支持fp16跟bf16

Tenstorrent withheld further details, but to achieve the 3-TOPS rating, the tensor engine likely contains about a thousand 8-bit MAC units

        他们也支持类似rowwise的quantization，一组数字共享scale

To save memory space, the design implements a block FP format in which groups of 16 values share the same 8-bit exponent. Tensix defines block FP formats with 8-, 4-, or 2- bit mantissas, trading off throughput for precision. Once the core loads values from memory, it expands them to FP16 before any computation.

        有这些能工作的小核了之后，我们就可以把这些小核串起来。

        四个Synopsys ARC CPU来负责组织120个小核的工作，总共16G DRAM跟16x PCIE。还有注意这些内存是LPDDR，肯定不能跟TPU或者Volta的HBM比较。这些小核之间的通信也有模块的

Although the compute unit can operate only on the local memory, each core can easily access data in other cores using the network-on-a-chip (NoC) interconnect.

        这里的NoC大致业务逻辑如下图

        所以这个NoC类似TPU的ICI，但是我看主要侧重于芯片内部的沟通，而且应该小的多，所以估计没有TPU之类的特别复杂的routing的逻辑，也做不了不同chip之前的沟通。NoC还负责压缩，看描述应该是varint类似的压缩方式。

The packet engine implements hardware data compression. It compresses data before transferring it across the NoC. Depending on the number of zeroes in the data, this compression typically shrinks the data by 50–75%, but the percentage can be even greater on sparse data.

        连接方式是2D torus，之前在TPU那边有介绍，只是这里的2D Torus是小核之间，而之前TPU是卡跟卡之间。

        并行方面，Tenstorrent需要的逻辑相对来说比其他芯片来说复杂的多。这种多核场景里面并行处理的上线肯定比几个大核来的多，但是肯定需要compiler更复杂的配合。视频里给了一个例子

        软件层来说onnx或者pytorch作为前端都行

        可以看到我之前讲的需要compiler复杂的配合。这个现在ppt做的很漂亮，我估计compiler的代码会非常复杂，需要各种不同的计划执行方式来优化不同的模型。

        Tenstorrent跟Grop核Titan相对来说的比较：

        IPS/Watt 大概是这样的，可以看一下含光还是最高，但是Tenstorrent相对来说还是很好的。而且含光为了convolution我记得是有特殊的优化的，不单纯是systolic array。

4.2 总结

        这个多核的架构在Tenstorrent这边说的很好，但是其实之前TPU的paper里面也讨论过了大核跟小核的架构优劣。

        TPUv3的文章里面有提到，

Sixteen 64x64 MXUs would have a little higher utilization (38%–52%) but would need more area. The reason is the MXU area is determined either by the logic for the multipliers or by the wires on its perimeter for the inputs, outputs, and control. In our technology, for 128x128 and larger the MXU's area is limited by the multipliers but area for 64x64 and smaller MXUs is limited by the I/O and control wires.

        所以相对来说并不是core越小越好，要具体问题具体看。对于我来说，太小的core其实优势不一定明显，特别是大公司推荐系统或者图像处理，应为实在batch size不够凑可以不同的访问coalesce起来。当然在汽车上面这个假设不一定成立，但是云端的话还是成立的。

        小core还有一个令人担心的地方是，鉴于并行优化非常的重要，小core是不是能够适用于多种模型，这个是有待市场验证的。

        不过优势还是很明显的，conditional作为一个功能本身绝对是一个特别好的创新。其实conditional 在现在的模型里面不是特别的常见，我觉得很大一个原因是硬件本身支持的不好，所以也省不下来什么，最常见的类似conditional的逻辑还是Mixture of expert，但是这个granularity相对来说就大的多。以后这类支持conditional的硬件出来的越多，越能帮助做模型的人创新。

5、疑惑

        在Tenstorrent芯片构成的网络中，存在4种拓扑结构，比如芯片内多核间、芯片间、服务器间和机架间。这4种拓扑结构对应的通信带宽和网络延迟肯定各不相同。因此需要软件栈中的Graph Compiler生成一个多种约束条件下的cost model，基于该cost model生成一个在当前运行环境下最优的算子与硬件的映射关系图。笔者认为Graph Compiler的实现颇为重要，但鲜有公开资料介绍这方面信息。官方表示会开源其软件栈，但未给出具体时间点。

        AI场景下数据通信延迟时常比计算延迟更高，因此高效协调通信和计算一直是业界努力的方向。为了实现灵活扩展，Tenstorrent强化了数据通信功能，数据通信链路也变得更为复杂。因此通信和计算协同问题更为凸显，Tenstorrent如何有效解决这个问题，目前我们尚不清楚。

        关于开发者生态，当前深度学习相关软件大都是将NV GPU作为第一优先级进行适配支持。Tenstorrent芯片的软件栈相对独立，虽然可以支持PyTorch、TensorFlow等框架，但如何推动这些深度学习框架在Tenstorrent芯片上适配开发，并针对性优化，路途漫漫。