本文的主要作者包括 Dennis Abts、Garrin Kimmell、Andrew Ling 等，全部由 Groq Inc. 公司以及 KAIST (John Kim) 的研究人员组成。文章发表在 ISCA ‘22（2022年6月18日至22日，纽约，美国）会议上，题为《A Software-Defined Tensor Streaming Multiprocessor for Large-Scale Machine Learning》。文章的 DOI 为 10.1145⁄3470496.3527405。
该研究的背景聚焦于高性能计算（HPC）和大规模机器学习的发展趋势。随着自然语言处理（NLP）和深度学习模型（如 GPT-3）的规模不断扩大，模型训练和推理的计算需求显著增长。现代机器学习模型可能包含数十亿甚至数千亿个参数，需要效率更高的大规模计算系统以满足存储能力、传输带宽以及处理能力的需求。然而，传统的多处理器系统（包括 CPU 和 GPU）在网络资源的利用以及运行时的确定性方面仍有局限。
为了解决这些挑战，本文提出了一种基于软件定义、规模可扩展的 Tensor Streaming Processor (TSP) 多处理架构。研究的主要目标如下： 1. 提供具有高吞吐量和低延迟的可扩展通信架构。 2. 实现一种新颖的软件调度（software-scheduled）网络，消除动态争用导致的不确定性。 3. 推动适用于各种机器学习训练和推理任务的基础架构，从而实现更强的灵活性和计算确定性。
本文提出了一个名为 TSP（Tensor Streaming Processor）的处理器架构，该架构通过多个模块部分细致描述了其设计、实现及相关实验流程：
1. 整体架构的描述
TSP 的架构设计基于一种“软件定义”的 Dragonfly 拓扑结构，用于连接多个处理单元（Processing Elements，PEs）。其通信网络由高带宽、低延迟的全局互联支撑，可以支持最多 10,440 个 TSP 节点，且在 3 微秒以内完成端到端延迟的通信。此外，每个 TSP 节点额外贡献了 220 MiB 的本地全局共享设置存储（SRAM），这被实现为逻辑共享但物理分布的存储器。
2. 单个 TSP 的运行模式与扩展 每个 TSP 内部采用了生产者-消费者流模型（Producer-Consumer Stream Model），其功能单元通过 320 个 SIMD（单指令多数据）元素执行指令。为了在多个 TSP 节点间保留编程模式的连贯性，整个系统需要具备确定性网络架构，避免动态路由的复杂性。
3. 软件调度的实现
研究中最突出的创新是软件调度的网络组织形式：
- 路由方法： 定义了基于软件的源路由机制，所有通信路径在编译阶段通过静态调度方式确定；避免了传统硬件网络中因动态争用而引入的随机时延。
- 流控方式： 消除了基于硬件的流控机制，与之相对，所有的链路通信均以软件显式调度，确保传输数据不会因硬件拥塞问题而丢失。
- 编码协议： 通信协议极简化，将数据封装开销降到最低，仅使用 2.5% 的额外头尾信息实现消息传输。
4. 网络拓扑设计
TSP 基于 Dragonfly 网络拓扑实施分级、全局互联，其特点为：
- 节点内完善连接： 在单个 4U 节点内，使用 7 条本地链接全联通 8 个 TSP 单元，以实现低成本的同步通信。
- 全局路由： 每个 TSP 至多分配 4 条全局链接，支持高效的多机架扩展，最大化减少了网络直径及带宽瓶颈问题。
5. 硬件与时间同步
通过硬件校准计数器（Hardware-Aligned Counters, HAC）和软件延迟校准指令（Runtime Deskew），多 TSP 系统实现了逻辑上全局同步的通信模型。以下是主要机制：
- HAC 跨 chip-to-chip (C2C) 链路同步，实现单跳网络延迟的稳定；
- 软件层面提供指令如 notify、sync，确保分布式任务按周期边界运行。
在实验评估中，本文对 TSP 系统在以下多个任务上的性能进行测试：
1. 分布式矩阵乘法 矩阵操作是深度学习中的基础操作。研究通过列分割和行分割两种分布方式，将大规模矩阵计算分发到多个 TSP 单元上。在多 TSP 系统中，测得最大延迟随 TSP 数目增加显著减少。同时，系统吞吐量（TFlops）能够近似线性扩展。例如，100 TSP 系统在执行 650,000×650,000 的矩阵乘法时，比传统 GPU（如 NVIDIA V100 集群）提升了 100 倍以上的速度。
2. All-Reduce 集合通信 研究通过软件调度的方式优化了 All-Reduce 操作的吞吐量和延迟。例如，在 8 台 TSP 的实验中，网络延迟仅为 2.1 微秒，比同参数环境中的 NVIDIA A100 系统表现更优。系统无需流控标志变量（如显存信号灯），直接按照编译时确定的数据流顺序实施计算，依赖于全同步的数据视图实现高效通信。
3. Transformer 模型 对 Bert-Large 模型的推理任务进行了 24,240 多次运行记录，结果显示推理输出平均延迟为 1225 微秒，且无显著波动。此外，Transformer 模型的吞吐量可随着 TSP 数目的增加而近线性增长，这表明其拓展性极强。
4. Cholesky 引理分解 对复杂 Cholesky 分解任务进行评估，通过多阶段向量矩阵乘法流水线实现，多 TSP 系统取得约 1.5 倍的性能提升，单节点吞吐量高达 22.4 FP16 TFLOPS，远超过文献中已知的最佳实现。
本文的贡献主要体现在以下方面： 1. 技术创新： 首次提出软件调度的 Dragonfly 网络，具有高度确定性，同时消除了动态调度网络的不稳定性。 2. 高效扩展性： 通过软件定义的多 TSP 系统架构，打破传统处理器（包括 GPU）在大规模并行计算中的性能瓶颈。 3. 学术与实践意义并存： 本研究对现有的高性能机器学习与 HPC 系统实现了实践指导，例如在 NLP 和矩阵计算等领域的性能突破。
本文是对现有高性能计算和机器学习基础架构的一次重要创新，极大增强了大型并行系统的确定性和可扩展性，为解决日益增长的深度学习模型需求提供了突破性解决方案。