一种用于高效张量计算的线程-寄存器解耦GPU执行模型

# 一种线程-寄存器解耦的GPU执行模型，用于高效张量计算  
来源：https://arxiv.org/html/2608.19628  
作者：Zihan Liu†\\dagger, Jingwen Leng∗†\{\}^\{\*\}\\dagger, Yangjie Zhou‡\\ddagger, Yitong Ding†\\dagger, Guanlin Zhu♢\\diamondsuit, Yilu Huang†\\dagger, Chiheng Jin†\\dagger, Chen Zhang†\\dagger, Shixuan Sun†\\dagger, Yu Feng†\\dagger, Anbang Wu†\\dagger, Minyi Guo†\\dagger, Jian Weng♠\\spadesuit, Jiajin Tu♣\\clubsuit, Junsong Wang♣\\clubsuit  
所属机构：  
上海交通大学†\\dagger  
新加坡国立大学‡\\ddagger  
上海科技大学♢\\diamondsuit  
阿卜杜拉国王科技大学♠\\spadesuit  
华为♣\\clubsuit  

联系邮箱：  
\{\\\{altair\.liu, dingyitong, huangyilu, wendy\-hamlet, chenzhang\.sjtu, sunshixuan, y\-feng\}\\\}@sjtu\.edu\.cn  
yj\_zhou@nus\.edu\.sg  
\{\\\{leng\-jw, anbang, guo\-my\}\\\}@cs\.sjtu\.edu\.cn  
jian\.weng@kaust\.edu\.sa  
tjjsh2109@126\.com  
junsongwang@huawei\.com  

---

### 摘要  
现代GPU越来越多地将张量核心集成到执行流水线中。尽管从Ampere架构基于寄存器的操作数供给，到Hopper和Blackwell架构基于内存的无冗余供给，操作数供给方式不断演进，使得聚合张量吞吐量持续增长，但为现代AI工作负载高效编排完整的张量计算流水线仍然充满挑战。我们将根本瓶颈识别为固定并行性和粗粒度调度，这两点在现代AI工作负载中日益凸显——此类工作负载常将多种非GEMM操作与GEMM操作交织执行。为高效编排张量计算，我们提出**FIBER**架构，扩展了GPU的SIMT（单指令多线程）模型。其基本执行实例**纤维（fiber）** 与私有寄存器所有权解耦，仅携带最小的控制状态，通过共享视图访问SM的寄存器。这使得动态并行扩展、细粒度寄存器级数据流调度成为可能，并为矩阵操作数提供了无冗余供给方案。我们扩展了指令集架构（ISA）、微架构和编译器，实现了共享寄存器寻址、无冲突操作数传递以及基于纤维的程序映射。在典型的混合精度LLM服务场景下，**FIBER** 在Ampere上实现了2.25倍的端到端加速（原始FP16计算为1.15倍），在Hopper和Blackwell上分别实现了1.8倍和2.09倍的加速，内核级性能提升最高达2.49倍。  

**注：** †\\dagger 表示通讯作者。  

---

## I 引言  
算法创新与GPU架构之间的良性循环，推动了人工智能模型（如大型语言模型）的快速发展。GPU引入了新的硬件能力，AI研究者可据此开发新的算子乃至新的模型结构。这些算法改进反过来又影响着硬件设计。高效且对程序员友好的SIMT执行模型是这一良性循环的核心。SIMT模型暴露了大规模并行性，并允许程序员从单线程角度进行编程。硬件将线程分组为线程束（warp）以实现高效的锁步执行，并需要大型寄存器文件结构来私有保存线程的执行上下文。  

随着矩阵乘法（GEMM）在现代AI和LLM工作负载中占据主导地位，GPU集成了专用的矩阵加速器，如NVIDIA的张量核心和AMD的矩阵核心。为满足不断增长的TC计算吞吐量及其对操作数的需求，矩阵操作数供给方式不断向基于内存的供给演进，避免了Volta和Ampere中私有寄存器供给方式下的线程束级矩阵片段冗余加载。由于共享内存和张量内存在SM内是共享的，操作数可在不同线程束间复用，这使得操作数供给能够跟上计算吞吐量的快速增长，实现了接近最优的TC利用率。  

然而，现代AI工作负载日益复杂。注意力机制及其变体、混合精度内核、前馈网络（FFN）模块将多种非GEMM操作与GEMM交织执行。我们的分析表明，尽管TC现已良好集成于GPU，但为这些工作负载编排完整的张量计算流水线仍然次优，其根源在于当前基于线程的SIMT执行模型与私有寄存器所有权绑定。该模型在启动时固定了线程数量及其寄存器预算。但现代工作负载在内存、张量和向量受限阶段之间交替，每个阶段需要不同程度的并行性和寄存器占用率，静态分配无法适应这些阶段依赖的需求。现有的CUDA动态并行依赖于嵌套内核启动，具有显著的运行时开销，而非在内核内动态调整线程数量和寄存器预算。  

该模型还阻碍了高效的线程束级数据流调度。一方面，SM范围的共享内存提供了块内线程束间通信的语义，但需要粗粒度的、基于阶段的调度，缺乏细粒度依赖跟踪。另一方面，具有停顿位的固定延迟计算指令和每线程束依赖计数器可以支持细粒度依赖跟踪，但这些机制仅限于每个线程束的私有寄存器，将细粒度数据流限制在单个线程束内。理想情况下，这些工作负载需要跨执行阶段自适应的动态并行、细粒度数据流调度，并保持无冗余的TC操作数供给。实现这些目标需要重新思考执行实例与寄存器之间的绑定关系。  

我们提出**FIBER**架构，建立在解耦的共享寄存器执行模型上，扩展SIMT并完全兼容。**FIBER**引入了新的并行执行实例**纤维（fiber）**，类似于SIMT线程但不拥有私有寄存器，仅携带最小状态如程序计数器和状态标志。与私有拥有寄存器的SIMT线程不同，纤维通过共享视图访问SM的寄存器文件。这种设计使硬件调度器能够在运行时调整活动纤维的数量，不再受限于预定的私有寄存器分配，支持以微小上下文实现更大的常驻并行，匹配不同计算阶段的并行需求。共享寄存器视图还支持直接的寄存器级数据流和细粒度调度，并由每寄存器依赖跟踪机制支持。共享寄存器还可作为矩阵计算的无冗余操作数供给路径，以保证最优的GEMM性能。这些特性共同使**FIBER**能够高效编排张量计算。  

我们设计了基于**FIBER**的GPU，仅进行最小修改并大部分重用现有硬件架构。首先，我们扩展现有GPU ISA以支持共享寄存器语义、动态并行和细粒度调度。其次，我们引入轻量级微架构增强，保留现有物理寄存器组织，包括用于共享寄存器访问的选择器和交叉开关、解决共享访问端口争用的仲裁器，以及用于细粒度跨纤维依赖跟踪的寄存器忙位图。第三，我们提出**FIBER**编程模型，暴露共享寄存器视图，使纤维可通过寄存器通信并编排细粒度寄存器数据流，同时保持SIMT兼容性。  

我们使用基于Tile-Graph的周期级模拟器GTSim在LLM工作负载上评估**FIBER**。在典型混合精度LLM服务场景下，**FIBER**相比Ampere、Hopper和Blackwell基线分别实现了2.25倍、1.8倍和2.09倍的端到端加速，内核级性能提升最高达2.49倍。面积和功耗开销分别低于0.1%和0.4%。进一步的开销分析表明，即使在极其紧张的硬件预算下，**FIBER**加速的关键使能技术仍然可行。这些结果表明，**FIBER**为张量计算提供了高效编排。  

我们总结贡献如下：  
- 全面分析了GPU上张量计算的瓶颈：静态并行性和粗粒度调度。  
- 设计了**FIBER**架构，将寄存器所有权与并行执行实例解耦，包含ISA扩展、轻量级微架构增强以及兼容CUDA的编程模型。  
- 在典型LLM服务场景下，相比多个GPU基线实现1.8-2.3倍端到端加速，在保持TC性能的同时缩小与理论峰值的差距。  
- 进行了彻底的消融研究，量化了**FIBER**每项关键特性对加速的贡献。  

---

## II 背景  
本节首先介绍LLM工作负载及其日益增长的复杂性背景。然后解释基本的GPU硬件架构及其编程方法。  

### II-A LLM工作负载  
大型语言模型大多采用基于Transformer的架构。Transformer模块的核心组件是注意力机制，其基本多头形式如公式(1)所示。注意力是最高度优化的内核之一，FlashAttention系列即为例证，它利用GPU特性执行高效融合计算。  

\[
\displaystyle\text{Attention}(\text{Q}, \text{K}, \text{V}) = \text{Softmax}\left(\text{Q} \cdot \text{K}^{\top} / \sqrt{\text{d}}\right) \cdot \text{V} \quad (1)
\]  

在状态注意力之前，线性投影将前一层的隐藏状态转换为查询（Q）、键（K）和值（V）状态，每个状态由与相应权重矩阵的GEMM生成。注意力之后，生成的上下文向量通过输出投影（再次以GEMM实现）。然后激活值流入前馈网络（FFN），该网络由GEMM层与轻量级非线性激活交织主导；FFN的输出成为下一个Transformer层的输入。  

尽管许多架构中出现了额外组件如RMSNorm、RoPE和Alibi，但GEMM操作仍是现代LLM中占主导的计算瓶颈。为应对快速增长的模型规模，现代LLM采用多种优化，包括高效注意力和**混合精度计算**，通常表示为WxAy，指示x位权重和y位激活。本研究考虑了LLM架构的典型计算阶段和这些广泛采用的优化。  

### II-B 基本GPU架构与编程  
GPU利用大规模线程级并行性提供高计算吞吐量。GPU由多个处理块组成，称为流式多处理器（SM），通过片外全局内存通信。本工作以NVIDIA GPU为基线。自Volta起，每个SM被组织为四个子分区（SP），共享统一的L1缓存和共享内存空间。在每个子分区中，指令先由调度器和分发器处理，然后发布到多个流水线之一，包括CUDA核心（SIMD单元）、张量核心、特殊功能单元（SFU）和加载/存储单元。  

编程NVIDIA GPU时，开发者遵循SIMT模型，所有线程概念上执行相同的指令序列。程序被分解为独立的线程块，每个块调度到一个SM上；多个块可在同一SM上并发运行以隐藏延迟。块内的线程被分组为32个线程的**线程束（warp）**（基本调度单位），以锁步方式执行。  

**图1：** TC操作数供给的两种方式。（左）寄存器操作数供给。（右）共享/张量内存操作数供给。  

---

## III 高效张量计算的挑战  
现代GPU越来越多地集成张量核心或矩阵核心以获得高矩阵计算吞吐量。为满足这种增长的吞吐量对操作数供给的需求，NVIDIA从Hopper架构开始采用共享/张量内存操作数供给，消除了Volta和Ampere中使用的寄存器供给方式下的重复矩阵操作数加载。如图1所示，在寄存器供给下，线程束3重新加载已由线程束1和2加载的片段A1, B1，因为线程束无法通过寄存器交换数据；消除这种重复为TC扩展扫除了关键障碍。内存供给使得GEMM能够实现最优的TC吞吐量。  

然而，现代工作负载（尤其是LLM）并非纯粹的GEMM：SIMD/SFU流水线上的多种非GEMM操作与GEMM交织，如注意力、FFN及更复杂的变体。因此，仅高效集成TC是不够的；GEMM还必须与其他非GEMM操作高效编排。本节我们识别这种编排中的两个主要瓶颈。  

### III-A 静态并行性  
两种操作数供给方式都保留了SIMT执行模型：每个线程私有拥有其寄存器，且每线程的寄存器分配和活动线程数在启动时固定。CUDA动态并行可以在运行时生成额外工作，但仅通过...