第5章:脉动阵列原理与设计
本章深入探讨脉动阵列(Systolic Array)的基本原理和设计思想,以Google TPU为例分析其架构演进。我们将从数据流动模式、计算同步机制入手,理解Weight-stationary设计的优势,并通过TPU四代产品的架构对比,掌握脉动阵列在实际应用中的优化策略。通过本章学习,读者将能够分析脉动阵列的利用率,设计高效的矩阵运算映射方案,并理解不同控制流架构的权衡。
5.1 脉动阵列基本概念
5.1.1 数据流动与计算同步
脉动阵列是一种规则的处理单元(PE)阵列结构,数据像心脏泵血一样有节奏地在PE间流动。每个PE在每个时钟周期执行相同类型的操作,通过精心设计的数据流动模式实现高效的并行计算。这种架构最早由H.T. Kung和Charles Leiserson在1978年提出,其设计理念是让数据在处理单元间有规律地"脉动"流过,每个数据项在流经系统时被多次使用,从而最大化数据重用率。
脉动阵列的核心特征:
- 局部通信:每个PE只与相邻PE通信,避免了长距离布线,这在VLSI设计中极为重要,因为长距离布线会带来更大的延迟、功耗和面积开销。在7nm工艺下,跨越1mm的金属线延迟可达100ps,而相邻PE间的短距离连接仅需10ps,这种10倍的延迟差异在高频设计中至关重要
- 规则结构:PE排列规则,便于VLSI实现和扩展,规则的结构使得物理设计更容易,时序收敛更可预测。规则的结构还便于采用层次化的时钟树综合,减少时钟偏斜,提高时序裕量
- 同步计算:所有PE在统一时钟下工作,简化控制逻辑,避免了复杂的握手协议和异步接口设计。同步设计使得功能验证更加直接,可以使用传统的静态时序分析工具进行时序收敛
- 数据重用:通过数据在PE间传递实现高效重用,每个数据项被多个PE使用,大幅降低了内存带宽需求。例如,在矩阵乘法中,每个输入元素被重用N次(N为阵列维度),相比直接从内存读取,可以减少N倍的带宽需求
脉动阵列的数学基础建立在矩阵运算的规律性上。考虑矩阵乘法 $C = A \times B$,其中 $C_{ij} = \sum_{k} A_{ik} \times B_{kj}$。这个计算具有高度的规律性和并行性:每个输出元素的计算相互独立,且都需要执行相同的运算序列(乘累加)。脉动阵列正是利用了这种规律性,将计算在空间(多个PE)和时间(多个周期)上展开。
深入理解脉动阵列的理论基础需要从计算的数据依赖图(Data Dependence Graph, DDG)入手。对于矩阵乘法,DDG具有高度规则的结构:每个 $C_{ij}$ 的计算依赖于第i行的所有A元素和第j列的所有B元素。这种规则性使得可以通过系统化的方法将算法映射到硬件架构上。
空间-时间映射理论: 脉动阵列的设计本质上是将算法的迭代空间(iteration space)映射到处理器的空间-时间域。对于三重循环的矩阵乘法:
for i = 0 to M-1:
for j = 0 to N-1:
for k = 0 to K-1:
C[i][j] += A[i][k] * B[k][j]
这个迭代空间是一个三维空间 $(i,j,k)$,脉动阵列将其映射为:
- 空间维度:$(i,j)$ 映射到PE阵列的物理位置
- 时间维度:$k$ 映射到时间步,每个时间步处理一个k值
这种映射的数学表达为: $$\begin{bmatrix} PE_x \\ PE_y \\ t \end{bmatrix} = \begin{bmatrix} 1 & 0 & 0 \\ 0 & 1 & 0 \\ 0 & 0 & 1 \end{bmatrix} \begin{bmatrix} i \\ j \\ k \end{bmatrix}$$ 这是最简单的恒等映射,实际设计中可能使用更复杂的仿射变换来优化数据重用和减少通信。
从数据流的角度看,脉动阵列实现了三种关键的数据移动模式:
- 单播到多播的转换:一个输入数据通过PE间的传递被多个计算单元使用。这种机制将原本的一对多通信转换为多个一对一通信,极大地简化了互连网络的设计。具体而言,如果使用总线实现一对多通信,需要驱动N个负载,这会导致延迟和功耗都随N线性增长;而脉动阵列的邻近通信只需要驱动单个负载
- 流水线化的累加:部分和在PE间传递并逐步累加,避免了全局归约的开销。传统的树状归约需要log(N)级的加法器和复杂的控制逻辑,而脉动阵列通过空间展开实现了线性时间的累加。更重要的是,流水线化的累加可以与新的乘法运算重叠,实现计算的连续性
- 时空映射:将算法的循环嵌套映射到硬件的空间维度和时间维度。这种映射将串行算法的循环迭代转换为并行硬件的空间分布和时间步进,实现了算法到硬件的自然映射
数据复用的定量分析: 脉动阵列的一个核心优势是数据复用。我们可以定量分析不同级别的数据复用:
-
时间复用(Temporal Reuse):同一个PE在不同时间使用相同数据 - 在Weight-stationary设计中,权重在PE中驻留整个计算过程 - 复用因子:等于处理的批次数B
-
空间复用(Spatial Reuse):不同PE使用相同数据 - 输入激活在水平方向传播,被每行的所有PE使用 - 复用因子:等于阵列的列数N
-
总复用因子:$R_{total} = R_{temporal} \times R_{spatial} = B \times N$
这意味着每个数据项从内存读取一次后,可以被使用 $B \times N$ 次,极大地提高了运算强度(Operational Intensity): $$OI = \frac{计算量}{数据移动量} = \frac{2MNK}{M \times K + K \times N + M \times N} \approx \frac{2K}{1 + \frac{1}{M} + \frac{1}{N}}$$ 当M和N足够大时,$OI \approx 2K$,这解释了为什么脉动阵列特别适合深度(K维度大)的矩阵乘法。
基本的一维脉动阵列可以表示为:
时刻 t=0: [a0] → PE0 → PE1 → PE2 → PE3
时刻 t=1: [a1] → PE0 → PE1 → PE2 → PE3
↓ ↓ ↓ ↓
[b0] [b1] [b2] [b3]
对于矩阵乘法 $C = A \times B$,二维脉动阵列的数据流动模式:
b00 b01 b02 b03
↓ ↓ ↓ ↓
a00 → [PE00][PE01][PE02][PE03]
a10 → [PE10][PE11][PE12][PE13]
a20 → [PE20][PE21][PE22][PE23]
a30 → [PE30][PE31][PE32][PE33]
每个PE执行乘累加操作:$c_{ij} = c_{ij} + a_{ik} \times b_{kj}$
脉动执行的时序分析对理解其效率至关重要。对于一个 $m \times n$ 的脉动阵列计算 $M \times K$ 和 $K \times N$ 的矩阵乘法:
启动延迟(Startup Latency):第一个结果出现前需要等待的周期数
- 水平方向数据传播:$n-1$ 个周期(数据需要逐级传递到最右侧PE)
- 垂直方向数据传播:$m-1$ 个周期(数据需要逐级传递到最底部PE)
- 总启动延迟:$L_{startup} = m + n - 2$ 个周期
- 这个延迟是不可避免的,因为第一个输出需要等待对应的输入数据到达右下角的PE
稳态吞吐量(Steady-State Throughput):
- 每个周期产生的运算数:$m \times n$ 次MAC操作
- 吞吐量:$Throughput = m \times n \times f_{clock}$ ops/s
- 其中 $f_{clock}$ 是时钟频率
总执行时间: $$T_{total} = L_{startup} + \lceil \frac{M}{m} \rceil \times \lceil \frac{K}{1} \rceil \times \lceil \frac{N}{n} \rceil$$ 这个公式揭示了脉动阵列设计的一个关键权衡:更大的阵列有更高的并行度,但也有更长的启动延迟。对于小矩阵,启动延迟可能占主导地位,导致利用率下降。
效率分析的临界点: 当矩阵维度满足 $M \times N > \alpha \times (m + n)^2$ 时(其中$\alpha$为经验系数,通常为2-4),脉动阵列才能达到较高的效率。这解释了为什么脉动阵列特别适合大规模矩阵运算,而对于小矩阵可能不如简单的向量处理器高效。
在实际系统中,常用的优化技术包括:
- 流水线重叠:在处理前一个矩阵块的同时预取下一个块的数据
- 批处理聚合:将多个小矩阵运算聚合成大矩阵运算
- 动态配置:根据矩阵大小动态调整使用的PE数量
5.1.2 Weight-stationary设计选择
脉动阵列有三种主要的数据流模式:
- Weight-stationary (WS):权重固定在PE中,输入和输出流动
- Output-stationary (OS):输出固定在PE中,权重和输入流动
- Row-stationary (RS):每行PE保持特定数据,其他数据流动
TPU选择Weight-stationary设计的关键原因深植于深度学习工作负载的特性。在推理场景中,模型权重是固定的,而输入数据不断变化。Weight-stationary设计完美契合了这一特点。
从能量效率的角度看,Weight-stationary的优势源于数据移动能耗的基本物理定律。根据电容充放电理论,移动n位数据的能耗为 $E = \frac{1}{2} C \times V^2 \times f \times \alpha$,其中$C$为线路电容,$V$为电压,$f$为频率,$\alpha$为翻转因子。在现代工艺下,电容随着距离增加而线性增长:$C = C_0 + c_w \times L$,其中$C_0$是驱动器和接收器的固定电容,$c_w$是单位长度的线电容(典型值0.2fF/μm),$L$是连线长度。
能耗层次分析: 现代芯片中,数据移动的能耗呈现明显的层次结构:
- 寄存器访问:~0.1 pJ/bit
- L1 Cache访问:~0.5 pJ/bit
- L2 Cache访问:~5 pJ/bit
- 片上SRAM访问:~10 pJ/bit
- DRAM访问:~100 pJ/bit
这个能耗金字塔清楚地表明,将数据保持在尽可能接近计算单元的地方至关重要。Weight-stationary通过将权重锁定在PE的本地寄存器中(最低能耗层次),实现了能效最优化。
优势分析:
- 能效最优:权重静止避免了大量数据移动,减少动态功耗。在现代工艺节点下,数据移动的能耗往往超过计算本身。45nm工艺下,32位浮点运算约需50pJ,而从DRAM读取32位数据需要640pJ,是计算的12.8倍
- 带宽效率:权重预加载后可重复使用,适合批处理。当处理批量请求时,同一组权重可以处理多个输入样本,极大地提高了权重的时间局部性
- 控制简单:权重地址生成逻辑简化,减少控制开销。权重按照固定模式加载到PE中,不需要复杂的动态调度
让我们通过一个具体例子深入理解Weight-stationary的优势。考虑一个全连接层,批大小B=128,输入维度I=1024,输出维度O=1024:
这个例子代表了典型的Transformer模型中的FFN层。在这种规模下,权重矩阵占据2MB存储(FP16),而每个批次的激活数据也为256KB。如果采用传统的计算方式,每次计算都需要从内存加载权重,这将成为性能瓶颈。
对于典型的神经网络层计算,设批大小为 $B$,输入特征为 $I$,输出特征为 $O$:
- 权重数据量:$W = I \times O$
- 输入激活量:$A_{in} = B \times I$
- 输出激活量:$A_{out} = B \times O$
在Weight-stationary模式下:
- 权重读取次数:1(预加载)
- 输入读取次数:$O$(每个输出通道一次)
- 输出写入次数:1
总数据移动量:$D_{WS} = W + O \times A_{in} + A_{out}$
相比之下,Output-stationary的数据移动量: $D_{OS} = B \times W + A_{in} + A_{out}$
当批大小 $B$ 较大时,$D_{WS} < D_{OS}$,Weight-stationary更有优势。
实际硬件实现中,Weight-stationary还带来了其他系统级优势:
- 预取优化:权重的访问模式完全可预测,可以实现完美的预取,隐藏内存延迟。现代DDR4/5的延迟高达60-100ns,通过预取可以完全隐藏这些延迟
- 功耗门控:权重寄存器在加载后可以关闭写使能,降低动态功耗。在28nm工艺下,关闭写使能可以节省约15%的寄存器功耗
- 容错设计:权重存储可以使用ECC保护,因为写入频率低,校验开销可接受。SECDED ECC只增加12.5%的存储开销,但能显著提高可靠性
- 压缩友好:权重可以压缩存储,在加载到PE时解压,因为解压只需要做一次。典型的权重压缩率可达2-4倍,节省大量存储带宽
然而,Weight-stationary也有其局限性:
- 灵活性受限:不同层的权重尺寸差异大,可能导致PE利用率不均。例如,ResNet的第一层卧7×7卷积,而后续层多为1×1或3×3,这种差异使Weight-stationary难以适应
- 批大小敏感:小批量时效率下降,因为权重重用机会减少。当批大小为1时,每个权重只被使用一次,Weight-stationary的优势完全消失
- 动态网络挑战:对于动态图或条件执行的网络,权重预加载策略复杂。需要预测所有可能的执行路径并预加载相应权重,这可能导致大量的内存浪费
这些局限性解释了为什么一些NPU(如Nvidia的Tensor Core)选择了更灵活的Output-stationary或Row-stationary设计。不同的设计选择反映了对不同应用场景的优化重点。
5.1.3 阵列维度与利用率分析
脉动阵列的维度选择直接影响硬件利用率和性能。设阵列维度为 $M \times N$,分析不同场景下的利用率是一个多目标优化问题,需要在面积、功耗、性能之间寻找平衡。
矩阵乘法映射效率:
对于矩阵乘法 $C_{P \times Q} = A_{P \times K} \times B_{K \times Q}$,映射到 $M \times N$ 的脉动阵列涉及三个维度的分块:
-
完美匹配:当 $P = kM, Q = lN$($k,l$ 为整数)时,利用率接近100% - 这种情况下,矩阵恰好可以分成整数个tiles - 每个tile完全利用脉动阵列的所有PE - 实际应用中很少出现完美匹配
-
维度不匹配:实际利用率 $U = \frac{\lceil P/M \rceil \times \lceil Q/N \rceil}{P \times Q} \times \frac{P \times Q}{M \times N}$ - 第一项表示tiling的效率损失 - 第二项表示最后一个tile的利用率 - 综合利用率通常在60-85%之间
-
边界效应:需要padding处理 - Padding开销:$(M - P \mod M) \times (N - Q \mod N)$ - 有效计算比例:$\eta = \frac{P \times Q}{M \times N \times \lceil P/M \rceil \times \lceil Q/N \rceil}$ - Padding不仅增加计算量,还增加内存访问
维度选择的设计空间探索:
选择最优的阵列维度需要考虑多个因素,这是一个复杂的多目标优化问题。
- 目标工作负载分析: - 统计常见层的维度分布:需要收集目标应用的所有层参数,建立维度分布直方图 - 识别性能关键层:通常遵循帕累托原则,20%的层占据80%的计算量 - 考虑未来模型趋势:Transformer模型倾向于使用更大的隐藏维度(如4096、8192)
以BERT-Large为例,其关键维度包括:
- 注意力矩阵乘法:[batch×seq, 1024] × [1024, 1024]
- FFN第一层:[batch×seq, 1024] × [1024, 4096]
- FFN第二层:[batch×seq, 4096] × [4096, 1024]
对于batch=32, seq=512的典型配置,第一个维度达到16384,这就需要考虑如何高效地将其映射到物理阵列上。
-
硬件约束: - 芯片面积:$Area \propto M \times N$(PE数量)。在现代工艺下,每个INT8 MAC单元约需100-200平方微米,256×256阵列需要约13平方毫米 - 布线复杂度:$Routing \propto M + N$(边界PE的连接)。随着阵列增大,全局布线变得困难,可能需要多级布线网络 - 时钟频率:更大的阵列可能需要降低频率以满足时序。经验公式:$f_{max} \approx \frac{c}{\sqrt{M \times N}}$,其中$c$为工艺相关常数
-
功耗考虑: - 静态功耗:$P_{static} \propto M \times N$ - 动态功耗:$P_{dynamic} = \alpha \times C \times V^2 \times f \times U$ - 其中U是利用率,低利用率意味着功耗效率低
批处理优化:
对于批量矩阵乘法,可以通过合并批次提高利用率:
- 将批次维度展开到矩阵维度
- 批次 $B$ 的 GEMM 可转换为:$(B \times P) \times K$ 和 $K \times Q$
稀疏处理策略:
对于2:4结构化稀疏:
- 理论计算减少50%
- 实际加速比:$S = \frac{2}{1 + \alpha}$,其中 $\alpha$ 为索引开销系数
- 典型情况下 $\alpha \approx 0.2-0.3$,实际加速约1.5-1.6倍
5.2 TPU v1-v4i架构演进
5.2.1 TPU v1:推理专用架构
TPU v1是Google的第一代张量处理器,专为数据中心推理工作负载设计。这个项目始于2013年,当时Google意识到如果所有用户每天使用3分钟的语音搜索,需要的计算资源将是当时数据中心容量的两倍。这促使Google开发专用硬件来加速神经网络推理。
项目背景与动机:
2013年,Google面临着一个关键挑战:深度学习模型在语音识别、图像分类等任务上取得了突破性进展,但部署成本极高。具体数据显示:
- 语音搜索的计算需求以每年10倍的速度增长
- 如果使用CPU处理所有语音请求,需要将数据中心规模扩大100倍
- GPU虽然提供了加速,但功耗过高(250W),不适合大规模数据中心部署
这促使Google做出了一个大胆的决定:设计专用的推理芯片。项目代号"Manhattan",目标是在15个月内完成从架构设计到批量部署。
设计理念与目标:
TPU v1的设计哲学是"做一件事,并把它做到极致"。与GPU试图兼顾图形渲染和通用计算不同,TPU v1专注于神经网络推理这一个任务。这种专注带来了几个关键优势:
- 去除了不必要的硬件:没有纹理单元、光栅化器、几何处理器等GPU特有组件
- 简化了编程模型:不需要CUDA那样复杂的编程接口,只需要简单的矩阵运算原语
- 优化了数据通路:专门为张量运算设计,去除了通用处理器的分支预测、乱序执行等开销
架构决策的深层考虑:
-
为什么选择INT8而非FP16/FP32: - Google的大规模实验表明,对于已训练模型,INT8量化的精度损失可控(<1%) - INT8乘法器面积仅为FP32的1/6,功耗为1/10 - 内存带宽需求减少4倍,这在内存墙问题日益严重的今天尤为重要
-
为什么选择256×256的阵列规模: - 太小(如64×64):控制开销比例过高,利用率受限 - 太大(如512×512):面积和功耗呈平方增长,时序收敛困难 - 256×256是sweet spot:在28nm工艺下,可以达到700MHz主频,同时保持合理的面积(~80mm²)
-
为什么采用脉动阵列而非其他架构: - SIMD向量处理器:带宽需求过高,每个周期需要读取所有操作数 - 多核MIMD:同步开销大,编程模型复杂 - 脉动阵列:数据重用率高,编程模型简单,硬件效率最高
核心规格:
- 脉动阵列:256×256 (65,536个MAC单元)
- 峰值性能:92 TOPS (INT8)
- 片上存储:28MB (24MB统一缓冲区 + 4MB累加器)
- 内存带宽:34GB/s (DDR3)
- 功耗:40W (典型)
- 工艺节点:28nm
- 芯片面积:331mm²(其中脉动阵列占约24%)
MXU (Matrix Multiply Unit) 设计:
统一缓冲区 (24MB)
↓
┌──────────────┐
│ 权重FIFO │ ← 权重加载
└──────────────┘
↓
┌──────────────────────┐
│ │
│ 256×256 脉动阵列 │ ← 激活输入
│ │
└──────────────────────┘
↓
┌──────────────┐
│ 累加器(4MB) │
└──────────────┘
↓
激活单元
关键设计决策:
- INT8量化:8位整数运算,功耗效率最优。Google的实验表明,对于已训练好的模型,INT8量化仅带来0.1-0.5%的精度损失,但能将计算密度提升4倍(相比FP32)
- 大阵列尺寸:256×256提供高并行度。这个尺寸是经过仔细权衡的结果:既要足够大以摊销控制开销,又要避免过大导致利用率下降
- 统一缓冲区:灵活分配权重和激活存储。24MB的容量可以存储整个ResNet-50的权重,避免频繁的DRAM访问
量化策略深入分析:
TPU v1采用的INT8量化包含几个关键技术,这些技术后来成为业界标准:
- 非对称量化:$q = \text{round}(r/s) + z$ - r:实数值 - s:缩放因子(scale) - z:零点(zero point)
非对称量化相比对称量化的优势在于可以更好地处理激活值的非对称分布。例如,ReLU后的激活值都是非负的,使用非对称量化可以充分利用INT8的表示范围。
- 逐通道量化: 每个输出通道使用独立的缩放因子,这是因为不同通道的权重分布差异很大。实验表明:
- 逐层量化:精度损失2-5%
- 逐通道量化:精度损失<0.5%
- 开销:需要存储额外的缩放因子(每通道4字节)
- 激活函数融合: ReLU6($y = \min(\max(x, 0), 6)$)的使用有多重考虑:
- 自然限制输出范围到[0, 6],简化量化
- 防止激活值爆炸,提高数值稳定性
- 硬件实现简单,只需要比较器和多路选择器
- 量化感知的编译优化: 编译器在量化时会进行以下优化:
- 常量折叠:预计算可以离线完成的量化操作
- 重排序:将量化/反量化操作移动到计算密集区域的边界
- 融合:将量化与其他操作(如激活函数)融合,减少内存访问
性能分析:
- 理论峰值:$256 \times 256 \times 2 \times 700\text{MHz} = 91.75$ TOPS
- 实际利用率:
- 大矩阵(>1024×1024):75-85%
- 中等矩阵(256-1024):60-75%
- 小矩阵(<256):30-50%
- 瓶颈分析:
- 计算瓶颈:当矩阵足够大且在片上缓存中时
- 内存瓶颈:当工作集超过24MB统一缓冲区时
- 同步瓶颈:当有大量小矩阵运算时
实际应用案例:
- InceptionV3:相比Haswell CPU提速71倍,相比K80 GPU提速3.5倍
- LSTM:在RankBrain应用中,延迟降低至7ms(满足10ms的SLA要求)
- 能效比:相比K80 GPU提升29倍(TOPS/W)
5.2.2 TPU v2/v3:训练能力扩展
TPU v2标志着Google从推理专用向训练推理一体化的战略转变。这个转变背后有深刻的技术和商业考量。
从v1到v2的架构革新:
TPU v2不是v1的简单升级,而是一次彻底的架构重新设计。主要驱动因素包括:
- 训练需求激增:2016-2017年,Google内部的模型训练需求每3.5个月翻倍
- 模型规模爆炸:BERT、GPT等大模型需要数千GPU天的训练时间
- 软硬件协同:TensorFlow的发展需要配套的训练硬件
TPU v2关键升级:
- 脉动阵列:128×128 (每个核心),相比v1的256×256看似"缩水",实际是为了更好的频率扩展
- 核心数量:2个核心/芯片,实现芯片级并行
- 浮点支持:bfloat16训练,这是Google的独创格式
- HBM容量:16GB HBM2,带宽密度是DDR3的20倍
- 内存带宽:600GB/s,解决了v1的内存瓶颈
- 峰值性能:45 TFLOPS/核心(bfloat16)
- 工艺节点:升级到16nm FinFET
bfloat16的设计智慧:
bfloat16(Brain Floating Point)是Google专门为深度学习设计的数据格式:
- 格式:1位符号 + 8位指数 + 7位尾数
- 动态范围:与FP32相同(指数位数相同)
- 精度:低于FP16,但对深度学习够用
为什么选择bfloat16而非FP16?
- 动态范围优先:深度学习中,动态范围比精度更重要
- 转换简单:FP32到bfloat16只需截断尾数,无需调整指数
- 硬件友好:可以复用FP32的指数处理逻辑
TPU v3增强:
- HBM容量:32GB(2倍提升)
- 内存带宽:900GB/s
- 峰值性能:123 TFLOPS(2.7倍提升)
- 液冷支持:功耗密度增加
架构创新:
-
向量处理单元(VPU): - 处理非矩阵运算(激活函数、归一化) - SIMD架构,128宽向量 - 支持超越函数(exp、log、sqrt)
-
跨核心互连: - 2D mesh拓扑 - 支持模型并行和数据并行 - ICI (Inter-Core Interconnect) 带宽:50GB/s/方向
-
混合精度训练:
前向传播:bfloat16 计算
反向传播:bfloat16 计算
权重更新:float32 累加
误差分析: $\epsilon_{total} = \epsilon_{round} + \epsilon_{trunc} + \epsilon_{accum}$
其中:
- $\epsilon_{round}$:舍入误差 ≈ $2^{-8}$ (bfloat16)
- $\epsilon_{trunc}$:截断误差 ≈ $2^{-7}$
- $\epsilon_{accum}$:累积误差 ≈ $\sqrt{n} \times \epsilon_{round}$
5.2.3 TPU v4i:推理优化版本
TPU v4i是v4的推理优化变体,针对大规模推理部署。
优化特性:
- 功耗优化:相比v4降低约30%
- 成本优化:去除不必要的训练特性
- 密度提升:机架密度提高2倍
性能规格:
- INT8性能:275 TOPS
- bfloat16性能:137 TFLOPS
- 内存容量:32GB HBM
- 内存带宽:1.2TB/s
推理优化技术:
-
动态批处理: - 请求合并减少延迟 - 批大小自适应:$B_{opt} = \min(B_{max}, \lceil \frac{L_{target}}{L_{single}} \rceil)$ - 其中 $L_{target}$ 为目标延迟,$L_{single}$ 为单样本延迟
-
算子融合优化:
Conv → BN → ReLU 融合为单个kernel
GEMM → Bias → Activation 融合执行
内存访问减少:$M_{fused} = M_{input} + M_{output}$ 相比未融合:$M_{unfused} = M_{input} + 2 \times M_{intermediate} + M_{output}$
- 稀疏加速支持: - 2:4结构化稀疏硬件支持 - 稀疏索引压缩:4位索引编码 - 理论加速:2×,实际加速:1.6-1.8×
5.2.4 架构演进趋势分析
通过四代TPU的演进,可以总结出以下趋势:
-
专用化程度提升: - v1:纯推理 - v2/v3:训练+推理 - v4i:推理专用优化
-
内存系统演进: - 带宽增长:34GB/s → 1.2TB/s (35×) - 容量增长:28MB → 32GB (1000×) - 带宽/计算比:0.37 → 4.36 bytes/op
-
精度支持扩展: - v1:INT8 - v2+:bfloat16、INT8、混合精度 - v4:FP32累加器
-
系统级优化: - 单芯片 → 多核心 → Pod级系统 - 软件栈成熟度提升 - 编译器优化深化
5.3 控制流与指令集
5.3.1 VLIW vs MIMD vs SIMD架构对比
脉动阵列的控制方式直接影响其灵活性和效率。TPU采用了CISC风格的指令集,结合了多种控制模式的优点。理解这些控制模式的权衡对于设计高效的NPU至关重要。
控制架构的历史演进:
计算机体系结构的发展历程为我们提供了宝贵的经验。从Flynn分类法(1966年)开始,并行架构被分为SISD、SIMD、MISD和MIMD四类。NPU设计主要关注后三种,每种都有其独特的优势和适用场景。
Flynn分类法的现代解读:
- SISD(Single Instruction Single Data):传统的冯·诺依曼架构,不适合大规模并行计算
- SIMD(Single Instruction Multiple Data):一条指令控制多个数据通道,GPU和向量处理器的基础
- MISD(Multiple Instruction Single Data):理论概念,实际应用罕见
- MIMD(Multiple Instruction Multiple Data):多核CPU的基础,每个核心独立执行
在NPU设计中,纯粹的分类已经不够,现代架构往往是混合式的:
SIMD (Single Instruction Multiple Data):
- 特点:一条指令控制多个数据通道,所有PE执行相同操作
- 优势:控制简单,硬件开销小,每个PE不需要独立的取指和译码单元
- 劣势:灵活性受限,分支效率低,遇到条件执行时需要屏蔽部分PE
在脉动阵列中,SIMD是最自然的选择,因为矩阵乘法本质上是规则的、无分支的计算。每个PE执行相同的MAC操作,只是操作数不同。这种规律性使得SIMD控制极其高效。
SIMD在TPU中的具体实现:
- 统一控制器:单个控制器向所有PE广播控制信号
- 锁步执行:所有PE在同一时钟周期执行相同操作
- 条件执行处理:通过predicate mask实现部分PE的条件执行
控制开销的定量分析:
- 传统MIMD:每个PE需要:
- 指令缓存:~2KB(64条32位指令)
- 译码器:~1000个晶体管
- 分支预测器:~2000个晶体管
- 总开销:~5000个晶体管/PE
- SIMD脉动阵列:
- 共享控制器:~50000个晶体管(服务256×256个PE)
- 每PE平均:<1个晶体管
- 面积节省:>99%的控制逻辑开销
这种巨大的面积节省使得可以将更多的晶体管用于实际计算,这是脉动阵列能够达到高计算密度的关键。
脉动阵列中的SIMD体现:
MATRIX_MULTIPLY:
所有PE执行: c += a × b
数据不同,操作相同
VLIW (Very Long Instruction Word):
- 特点:单条指令包含多个操作
- 优势:编译时调度,无需硬件调度
- 劣势:代码密度低,编译器复杂
TPU的VLIW设计:
指令格式:
[MXU_OP | VPU_OP | DMA_OP | SCALAR_OP]
8bit 8bit 8bit 8bit
MIMD (Multiple Instruction Multiple Data):
- 特点:多个独立指令流
- 优势:最大灵活性
- 劣势:硬件复杂,同步开销大
5.3.2 指令调度与依赖管理
TPU的指令调度采用静态调度为主、动态调整为辅的策略。这种设计选择反映了对编译时优化和运行时效率的平衡考虑。
静态调度的优势:
- 确定性延迟:每条指令的执行时间在编译时确定,便于性能预测
- 硬件简化:不需要复杂的乱序执行单元、重排序缓冲区等
- 功耗优化:避免了动态调度的功耗开销(约占CPU功耗的20-30%)
- 调试友好:执行顺序与程序顺序一致,便于调试和验证
指令级并行(ILP)分析:
设指令序列为 $I = \{i_1, i_2, ..., i_n\}$,依赖图 $G = (V, E)$:
- 节点 $V$:指令
- 边 $E$:数据依赖,边权重为依赖距离(周期数)
关键路径长度:$L_{critical} = \max_{path} \sum_{i \in path} latency(i)$
并行度量化:
- 理论并行度:$P_{theory} = \frac{\sum_{i} latency(i)}{L_{critical}}$
- 实际并行度:$P_{actual} = \frac{\sum_{i} latency(i)}{T_{execution}}$
- 并行效率:$\eta = \frac{P_{actual}}{P_{theory}}$
在TPU中,典型的并行效率可达80-90%,这得益于:
- 规则的计算模式(矩阵运算)
- 可预测的内存访问模式
- 编译器的深度优化
依赖类型:
- RAW (Read After Write):真依赖
i1: C = MATRIX_MUL(A, B)
i2: D = ACTIVATE(C) // 必须等待i1完成
- WAR (Write After Read):反依赖
i1: C = LOAD(addr1)
i2: STORE(addr1, D) // 必须等待i1读取完成
- WAW (Write After Write):输出依赖
i1: C = COMPUTE1()
i2: C = COMPUTE2() // 必须保序
调度算法:
TPU编译器使用改进的List Scheduling算法:
def schedule_instructions(dag, resources):
ready_list = get_ready_nodes(dag)
schedule = []
time = 0
while ready_list:
# 优先级函数
priority = lambda n: (
-critical_path_length(n), # 关键路径优先
-resource_usage(n), # 资源密集优先
node_id(n) # 确定性tie-breaking
)
node = max(ready_list, key=priority)
if can_schedule(node, resources, time):
schedule.append((node, time))
update_resources(resources, node, time)
ready_list.remove(node)
ready_list.extend(get_newly_ready(dag, node))
else:
time += 1
return schedule
流水线设计:
TPU采用5级流水线:
- IF (Instruction Fetch):取指令
- ID (Instruction Decode):解码和操作数读取
- EX1 (Execute 1):地址计算/标量运算
- EX2 (Execute 2):矩阵运算启动
- WB (Write Back):结果写回
流水线冒险处理:
- 数据冒险:通过forwarding解决
- 控制冒险:分支预测(准确率>95%)
- 结构冒险:资源仲裁和stall
5.3.3 TPU指令集架构
TPU指令集包含以下主要类别:
1. 矩阵运算指令:
MatrixMultiply(src1_addr, src2_addr, dst_addr, M, K, N)
- src1: M×K矩阵地址
- src2: K×N矩阵地址
- dst: M×N结果地址
- 执行周期: ceil(M/256) × ceil(N/256) × ceil(K/256)
2. 向量运算指令:
VectorOp(op_type, src_addr, dst_addr, length)
- op_type: ADD/MUL/ReLU/Sigmoid等
- 支持广播和element-wise操作
- SIMD宽度: 128
3. 数据传输指令:
DMA_Load(dram_addr, sram_addr, size, stride)
DMA_Store(sram_addr, dram_addr, size, stride)
- 支持2D/3D传输模式
- 最大传输: 16MB
- 双缓冲支持
4. 同步与控制指令:
Barrier(barrier_id) // 同步点
Branch(condition, target) // 条件跳转
Loop(count, body_addr) // 硬件循环
指令编码示例:
31 28 27 24 23 20 19 16 15 12 11 8 7 0
┌────┬──────┬──────┬──────┬───────┬───────┬───────┐
│ Op │ Dst │ Src1 │ Src2 │ Imm12 │ Flags │ Rsvd │
└────┴──────┴──────┴──────┴───────┴───────┴───────┘
Op: 4位操作码
Dst/Src: 4位寄存器编号
Imm12: 12位立即数
Flags: 4位标志位(条件码等)
5.3.4 性能计数器与调试支持
TPU提供丰富的性能计数器用于性能分析和调试:
硬件性能计数器:
-
计算相关: - MXU利用率:
active_cycles / total_cycles- VPU利用率 - 整数运算单元利用率 -
内存相关: - HBM带宽利用率 - Cache命中率 - DMA传输效率
-
同步开销: - Barrier等待周期 - Pipeline stall周期 - 依赖等待时间
性能分析公式:
实际性能 = 理论峰值 × 利用率 × 效率
其中:
- 利用率 = $\frac{有效计算周期}{总周期}$
- 效率 = $\frac{有效运算数}{总运算数}$
对于矩阵乘法: $$Performance = \frac{2MNK}{T_{execution}} = \frac{2MNK}{T_{compute} + T_{memory} + T_{sync}}$$
其中:
- $T_{compute}$:计算时间
- $T_{memory}$:数据传输时间
- $T_{sync}$:同步开销
调试支持特性:
- 断点支持:硬件断点和watchpoint
- 单步执行:指令级单步
- 寄存器和内存dump
- Trace buffer:记录执行历史
本章小结
本章深入探讨了脉动阵列的基本原理和Google TPU的架构演进。我们学习了:
-
脉动阵列基础: - 数据流动的三种模式(WS/OS/RS)及其权衡 - Weight-stationary设计在能效和带宽效率上的优势 - 阵列维度选择对利用率的影响,边界效应的处理
-
TPU架构演进: - 从推理专用(v1)到训练支持(v2/v3)再到推理优化(v4i) - 内存系统的大幅提升:带宽35倍增长,容量1000倍增长 - 混合精度训练和稀疏加速的硬件支持
-
控制流设计: - SIMD/VLIW/MIMD架构的权衡 - 静态调度为主的指令调度策略 - 5级流水线设计和冒险处理机制
关键公式回顾:
- 脉动阵列利用率:$U = \frac{实际计算量}{阵列规模 \times 执行周期}$
- Weight-stationary数据移动:$D_{WS} = W + O \times A_{in} + A_{out}$
- 性能计算:$Performance = \frac{2MNK}{T_{compute} + T_{memory} + T_{sync}}$
练习题
基础题
习题5.1:矩阵乘法映射分析 给定一个128×128的脉动阵列,需要计算矩阵乘法 $C_{384 \times 512} = A_{384 \times 256} \times B_{256 \times 512}$。请计算:
- 需要多少个tiling块?
- 理论利用率是多少?
- 如果采用padding到最近的128倍数,padding开销是多少?
Hint:考虑如何将大矩阵分块映射到脉动阵列上。
参考答案
-
Tiling分块: - M维度:$\lceil 384/128 \rceil = 3$ 块 - N维度:$\lceil 512/128 \rceil = 4$ 块 - K维度:$\lceil 256/128 \rceil = 2$ 块 - 总块数:$3 \times 4 = 12$ 个输出块
-
理论利用率: - 不考虑K维度的利用率:$U = \frac{384 \times 512}{3 \times 128 \times 4 \times 128} = \frac{196608}{196608} = 100\%$ - 考虑所有维度:需要 $3 \times 4 \times 2 = 24$ 次脉动阵列调用
-
Padding开销: - M维度无需padding(384 = 3×128) - N维度无需padding(512 = 4×128)
- K维度无需padding(256 = 2×128) - 总padding开销:0
习题5.2:数据移动量计算 对于批大小B=32,输入特征I=1024,输出特征O=2048的全连接层,分别计算Weight-stationary和Output-stationary模式下的数据移动量。假设权重为FP16(2字节),激活为INT8(1字节)。
Hint:分别统计权重、输入激活和输出激活的读写次数。
参考答案
Weight-stationary模式:
- 权重加载:$W = 1024 \times 2048 \times 2 = 4MB$(一次性加载)
- 输入读取:$A_{in} = 32 \times 1024 \times 1 \times 2048 = 64MB$(每个输出通道读一次)
- 输出写入:$A_{out} = 32 \times 2048 \times 1 = 64KB$
- 总数据移动:$D_{WS} = 4MB + 64MB + 64KB ≈ 68.06MB$
Output-stationary模式:
- 权重读取:$W = 1024 \times 2048 \times 2 \times 32 = 128MB$(每个batch读一次)
- 输入读取:$A_{in} = 32 \times 1024 \times 1 = 32KB$
- 输出写入:$A_{out} = 32 \times 2048 \times 1 = 64KB$
- 总数据移动:$D_{OS} = 128MB + 32KB + 64KB ≈ 128.19MB$
Weight-stationary节省:$(128.19 - 68.06) / 128.19 = 46.9\%$
习题5.3:稀疏加速计算 一个256×256的脉动阵列支持2:4结构化稀疏。对于一个包含75%稀疏度的权重矩阵(随机稀疏),转换为2:4结构化稀疏后,请计算:
- 结构化稀疏的实际稀疏度
- 理论加速比
- 考虑20%索引开销后的实际加速比
Hint:2:4稀疏要求每4个元素中恰好2个为零。
参考答案
-
结构化稀疏的实际稀疏度: - 2:4稀疏固定为50%稀疏度 - 原始75%随机稀疏转换为2:4后,稀疏度降为50%
-
理论加速比: - 相对于密集计算:$S_{theory} = \frac{1}{1-0.5} = 2×$ - 相对于75%随机稀疏:无法直接在硬件上加速
-
实际加速比(考虑索引开销): - $S_{actual} = \frac{2}{1 + 0.2} = \frac{2}{1.2} = 1.67×$ - 有效计算减少:$1 - \frac{1}{1.67} = 40\%$
挑战题
习题5.4:流水线性能分析 TPU的5级流水线中,矩阵乘法指令需要100个周期完成,向量运算需要10个周期,DMA传输需要50个周期。给定以下指令序列:
I1: DMA_Load(weights) // 50 cycles
I2: DMA_Load(inputs) // 50 cycles
I3: MatrixMultiply() // 100 cycles, 依赖I1,I2
I4: VectorAdd() // 10 cycles, 依赖I3
I5: VectorReLU() // 10 cycles, 依赖I4
I6: DMA_Store(outputs) // 50 cycles, 依赖I5
请计算:
- 不考虑流水线的总执行时间
- 理想流水线执行时间(无stall)
- 如果I3必须等待I1和I2都完成才能开始,实际执行时间是多少?
Hint:画出流水线执行时序图。
参考答案
-
顺序执行总时间: $T_{seq} = 50 + 50 + 100 + 10 + 10 + 50 = 270$ cycles
-
理想流水线(假设可以完美重叠): - I1和I2可以并行(DMA双通道):50 cycles - I3开始:100 cycles - I4-I6顺序执行:10 + 10 + 50 = 70 cycles - 理想时间:$\max(50, 50) + 100 + 70 = 220$ cycles
-
实际执行时间(考虑依赖):
时间 0-50: I1执行,I2并行执行
时间 50-150: I3执行(等待I1,I2完成)
时间 150-160: I4执行
时间 160-170: I5执行
时间 170-220: I6执行
总时间:220 cycles
流水线效率:$\eta = \frac{270}{220} = 1.23×$ 加速
习题5.5:TPU Pod系统性能建模 一个TPU v3 Pod包含1024个TPU芯片,采用2D torus互连拓扑。每个芯片有2个核心,每核心123 TFLOPS。对于训练一个GPT-3规模的模型(175B参数),请分析:
- 理论峰值算力
- 如果采用数据并行,每个芯片需要存储多少模型参数?
- 如果采用模型并行,通信开销如何估算?
Hint:考虑参数、梯度、优化器状态的存储需求。
参考答案
-
理论峰值算力: - 单芯片:$2 \times 123 = 246$ TFLOPS - Pod总算力:$1024 \times 246 = 251,904$ TFLOPS ≈ 252 PFLOPS
-
数据并行存储需求: - 模型参数:175B × 2 bytes (bfloat16) = 350GB - 梯度:175B × 2 bytes = 350GB - Adam优化器状态(动量+二阶矩):175B × 4 bytes = 700GB - 总需求:350 + 350 + 700 = 1400GB per replica - 单芯片HBM:32GB - 需要:$\lceil 1400/32 \rceil = 44$ 个芯片才能存储完整模型
-
模型并行通信开销: - 采用1024路模型并行 - 每个芯片存储:175B/1024 ≈ 171M参数 - 前向传播:每层需要All-Gather激活 - 反向传播:需要Reduce-Scatter梯度 - 通信量/step:$2 \times batch_size \times hidden_dim \times 2$ bytes - 对于hidden_dim=12288,batch=1:通信量≈49KB/层 - 96层transformer:总通信≈4.7MB/sample
通信时间估算(50GB/s互连): $T_{comm} = \frac{4.7MB}{50GB/s} = 94μs/sample$
习题5.6:编译器Tiling策略优化 给定一个卷积层:输入[1,224,224,3],卷积核[7,7,3,64],步长2,padding 3。目标是映射到128×128的脉动阵列上。请设计:
- Im2col转换后的矩阵维度
- 最优的tiling策略(最小化数据移动)
- 估算总的数据移动量
Hint:考虑空间维度和通道维度的不同tiling方式。
参考答案
-
Im2col转换: - 输出尺寸:$\lfloor (224 + 2×3 - 7)/2 \rfloor + 1 = 112$ - Im2col矩阵:
- 输入矩阵A:[112×112, 7×7×3] = [12544, 147]
- 权重矩阵B:[7×7×3, 64] = [147, 64]
- 输出矩阵C:[12544, 64]
-
Tiling策略: - 策略1:空间维度tiling
- M维度(空间):$\lceil 12544/128 \rceil = 98$ 块
- N维度(输出通道):$\lceil 64/128 \rceil = 1$ 块
- K维度(卷积核):$\lceil 147/128 \rceil = 2$ 块
- 策略2:优化tiling(利用64<128)
- 将2个空间位置打包:[256, 147] × [147, 64]
- M维度:$\lceil 12544/256 \rceil = 49$ 块
- 仍需要98次计算,但可以更好地利用阵列
- 数据移动量估算: - 输入数据(Im2col后):12544 × 147 × 1 byte = 1.84MB - 权重数据:147 × 64 × 1 byte = 9.4KB - 输出数据:12544 × 64 × 1 byte = 803KB
采用Weight-stationary:
- 权重加载1次:9.4KB
- 输入读取98次(每个M块):1.84MB × 98 = 180MB
- 输出写入1次:803KB
- 总数据移动:≈181MB
优化:通过输入数据重用减少读取
- 相邻输出位置共享输入数据
- 实际数据移动可减少约30-40%
常见陷阱与错误
1. 阵列利用率陷阱
错误:假设脉动阵列总是能达到100%利用率
原因:
- 矩阵维度与阵列尺寸不匹配
- Pipeline启动和排空开销
- 数据依赖导致的bubble
正确做法:
- 选择合适的批大小和tiling参数
- 考虑padding开销
- 使用双缓冲隐藏数据传输
2. 内存带宽估算错误
错误:只考虑计算时间,忽略数据传输
示例:
错误估算:256×256阵列,700MHz
计算时间 = M×N×K / (256×256×700M) = ...
正确估算:
总时间 = max(计算时间, 数据传输时间)
如果数据传输时间 > 计算时间,则memory-bound
3. 精度损失累积
错误:在量化时不考虑误差累积
问题场景:
- 长序列的累加操作
- 深度网络的梯度消失
- 大batch size的归约操作
解决方案:
- 使用更高精度的累加器(如INT32累加INT8乘法)
- 关键层保持高精度
- 定期重新量化
4. 稀疏模式不匹配
错误:假设任何稀疏模式都能加速
实际情况:
- 硬件只支持特定稀疏模式(如2:4)
- 非结构化稀疏需要额外索引开销
- 稀疏度太低时overhead大于收益
5. 编译器优化假设
错误:手动优化一定比编译器好
事实:
- 现代编译器(如XLA)有复杂的优化pass
- 手动tiling可能破坏编译器优化机会
- 需要profile-guided optimization
最佳实践检查清单
设计阶段
- [ ] 确定目标工作负载的算子分布
- [ ] 分析典型矩阵维度和batch size
- [ ] 评估内存带宽需求 vs 计算需求
- [ ] 选择合适的数值精度(INT8/BF16/FP32)
- [ ] 确定脉动阵列维度(功耗、面积、性能权衡)
实现阶段
- [ ] 实现高效的地址生成单元
- [ ] 设计灵活的数据通路(支持不同精度)
- [ ] 实现双缓冲减少等待
- [ ] 添加必要的调试和性能计数器
- [ ] 考虑时钟域交叉(CDC)设计
验证阶段
- [ ] 覆盖所有支持的数据类型
- [ ] 测试边界条件(小矩阵、非对齐尺寸)
- [ ] 验证数值精度(与参考模型对比)
- [ ] 压力测试(最大带宽、最大计算)
- [ ] 验证错误处理机制
优化阶段
- [ ] Profile实际工作负载
- [ ] 识别性能瓶颈(计算/内存/同步)
- [ ] 优化数据布局减少bank conflict
- [ ] 调整tiling参数
- [ ] 实现算子融合减少内存访问
部署阶段
- [ ] 制定功耗预算和散热方案
- [ ] 准备性能调优指南
- [ ] 提供调试工具和文档
- [ ] 建立性能回归测试
- [ ] 规划未来升级路径