从苹果神经引擎中恢复 50 GB/S 的性能

# 从苹果神经引擎中重新获得 50 GB/s 的速度 来源：https://eiln.github.io/posts/ane-dma.html ## 引言 苹果 M3 神经引擎中的一个 RTL 性能问题，会导致 DRAM 权重流传输吞吐量从标称的 45-60 GB/s 降低至 17-19 GB/s。这发生在总权重大小为 1 MiB 的整数倍时，目前影响了 ANEMLL (https://github.com/anemll/anemll) 15 个模型中的 7 个。通过在内核 DMA 引擎的推测预取环中避开有问题的路径，Llama 3.2 1B 的每秒 token 吞吐量从 10.0 提升至 24.3（DRAM 使用率从 24.7 提升至 60.0 GB/s），Qwen3-8B 从 1.36 提升至 2.97（DRAM 使用率从 22.4 提升至 48.7 GB/s）。 ## 发现过程 在对神经引擎的 DRAM 权重流传输吞吐量（单位：GB/s）进行单 token 解码分析时： \[ X[1,D] \times W[D,N] = Y[1,N] \] 当 \(N=4096\) 时，我注意到 \(D=1536\) 的运行速度比 \(D=2048\) 快了近 **3 倍**，后者是 Llama 3.2 中使用的默认值。 STATIC（纯 KernelDMA）每个副本的中位数微秒，N=4096： | D | 576 | 768 | 1024 | 1280 | 1536 | 2048 | |------|-----|-----|------|------|------|------| | rep a | 150.4 | 196.8 | 238.1 | 293.8 | 310.9 | 997.6 | | rep b | 157.8 | 190.2 | 250.1 | 288.3 | 326.8 | 995.0 | | rep c | 148.7 | 189.6 | 249.4 | 275.2 | 316.5 | 995.4 | 在 D = 2048 附近区域扫描 D 值： [图表链接：带宽在 D 接近 2048 时的变化] 嗯？ 当 D=2048 时，吞吐量为 16.93 GB/s。当 D=2016 时，吞吐量为 44.5 GB/s，这意味着 ``` 44.505062 − 16.930761 = 27.574301 GB/s (降低了 61.96%)。 ``` 一次 **27.57 GB/s 的下降**，从 44.5 降到 16.93 GB/s。请注意，扫描数据是在 M3 Air 上收集的，在单次运行中相同热/负载条件下重复了 40 次。我还确保 ANE 寄存器文件的 DMA 大小和地址是唯一改变的变量： ``` D=2044 D=2048 D=2052 TD+0x004 预估周期数 0x000001ea 0x000001eb 0x000001ec TD+0x078 核心 1 基地址 0x000ff800 0x00100000 0x00100800 TD+0x07c 核心 2 基地址 0x001ff000 0x00200000 0x00201000 ... TD+0x0b0 核心 15 基地址 0x00ef8800 0x00f00000 0x00f07800 TD+0x0b4–0x0f0 核心大小 ×16 0x000ff800 0x00100000 0x00100800 TD+0x134 Common.Cin 0x000007fc 0x00000800 0x00000804 TD+0x1f0 L2 源步长 0x00007fc0 0x00008000 0x00008040 TD+0x1f4 未知步长镜像 0x00007fc0 0x00008000 0x00008040 TD+0x214 L2 结果基地址 0x00008fc0 0x00009000 0x00009050 ``` 然后我对整个 D 值范围进行了扫描： 这是个好主意，因为我发现在 D = 2048 处出现了一个共振峰。从没想过会对吞吐量（GB/s）与张量维度（D）进行快速傅里叶变换（FFT），但结果如下： [图表链接：吞吐量随 D 变化的 FFT 结果] 显然，内存控制器的吞吐量在张量维度空间中有一个波长为 2048 的主导谐波。遗憾的是，这是一个波谷 :(。 所有 D = 2048 的倍数都被限制在固定的带宽下限 17-19 GB/s： [图表链接：在 2048 的倍数处的带宽凹陷] 在 D = 2048 的倍数处，吞吐量从标称的 45-60 GB/s 急剧下降至 17-19 GB/s，并在大约 256 行之后恢复到标称值。这不是一个 RTL 正确性错误，因为内核 DMA 仍然正确地完成了传输。但是，2048 附近的请求被强制进入一个单独的、信用不足的问题状态，不合理地将吞吐量限制在 28-43 GB/s（最坏情况从 60 降至 17），而这些传输大小不幸非常常见。 ## 假设 1 - DRAM 空间相关性 > 16 个核心是否以 2 的幂次步长映射到同一个 DRAM bank？ DRAM 是并行数据接口：DRAM 带宽是 DQ（数据）引脚数量乘以每个引脚的数据速率， \[ \text{DRAM BW} = N \times R = 128\ \text{bit} \times 6.4\ \text{GT/s} = 102.4\ \text{GB/s} \] M3 的 LPDDR-6400 的 102.4 GB/s 与其标称的 100 GB/s 相符。持续的 DRAM 带宽严格受限于 DQ 利用率，每一 GB/s 低于 102.4 GB/s DRAM 上限的情况，都意味着 DQ 线路空闲了额外的周期。DRAM 简述：DRAM 内存控制器使用并行访问通过高速 DQ 引脚流式传输比特；大型 DRAM 阵列被划分为 bank，而带宽（大致）取决于将并行请求分散到不同 bank 中。 如果资源是独立的，并行性可以提高吞吐量。如果并行请求者访问相同的资源，它们的请求将背靠背串行化，并有效地被限制在单一速率。在 ~17-19 GB/s 的限制下（而其相邻的请求运行在 45-60 GB/s），这可能是由 2 的幂次边界处的崩溃引起的。从底层开始分析也是好的：如果额外的 AXI 请求请求的是相同的物理 bank，那么它们无济于事。 --- ### 核心争用 神经引擎有多种并行途径，第一类是核心级并行。ANE 有 16 个并行核心。核心通过将缓冲区均匀分区到 \(N\) 个核心来分配工作，并相互约定处理不同的切片。我们知道核心被分配来获取权重缓冲区的不同切片，但 ANE 仍然有 16 个核心在同一周期内并行地从 DRAM 请求它们的切片。 如果每个核心从 DRAM 获取自己的切片，那么流式延迟应该花费相同的时间，无论是一个核心还是所有 16 个核心启用，因为它们的请求应该被并行服务。然而，如果由于任何核心争用导致带宽减少，那么对于受限的 D=2048 情况，减少核心数量可能会反直觉地增加吞吐量。对 D=2048 和 D=2016 扫描活跃核心数量： [图表链接：延迟随活跃核心数的变化] 对于 D=2016 和 D=2048，从 1 到 16 个活跃核心，延迟都是恒定的，这意味着即使在 core=1 时也存在限流。问题存在于每个核心的级别，并且问题在多个核心上被复制。 --- ### 地址争用 即使排除了核心级争用，我仍然怀疑由于 2 的幂次周期性而导致了一些 DRAM 争用。像 2048 这样的 2 的幂次步长在每次轮换时增加 \(2^k\)，这意味着低位比特 \([0..k-1]\) 是恒定的。DRAM 对物理地址进行哈希，使得步进访问模式在不同的 bank 之间得到空间解相关，因此哈希折叠低位比特或映射高位 \(2^k\) 比特可以解释 2 的幂次周期性。 为了测试 DRAM 空间相关性是否是问题所在，我们扰乱了权重获取的地址。地址被随机扰乱并分布在整个约 64 MiB 的 IOVA 区域（59.90 MiB 跨度）内，因此它在页面内和页面外都被扰乱了。为了排除无风扇 M3 Air 上的温度漂移，基线和扰乱样本在运行间交替进行，因此任何温度变化都同等影响两种条件。 基线的中位数吞吐量为 31.37 GB/s，随机扰乱地址的中位数吞吐量为 32.29 GB/s。随机扰乱的持续吞吐量平均高出 1 GB/s，表明我们在这次运行中可能通过扰乱解决了一些空间相关性，但是 (1) 这尚未在所有情况下得到证实 (2) 扰乱无法恢复解释崩溃所需的约 +200% 吞吐量下降。 ## 假设 2 - RTL 整数回绕 回忆一下，崩溃在 D = 2048 的每个整数倍处重复： [图表链接：在 2048 倍数处的崩溃] 问：在精确的 2 的幂次整数边界处什么会重复？ 答：固定宽度数字逻辑中的整数溢出。 ``` module line_counter ( input wire clk, input wire reset, input wire advance, output reg [13:0] line_count ); always @(posedge clk) begin if (reset) line_count <= 14'h0000; else if (advance) line_count <= line_count + 1'b1; // 在 0x3fff + 1 -> 0x0000 时回绕 end endmodule ``` --- ### 核心维度 \[ X[1,D] \times W[D,N] = Y[1,N] \] 其中 - \(D\) (Cin)：每个卷积核的长度：\(D\) 个 FP16（2 字节）权重，即 \(2D\) 字节。 - \(N\) (Cout)：卷积核的数量。每个核心处理 \(N/16\) 个卷积核。 由于原始图表是在固定 N = 4096 的情况下扫描 D，我们实际上从未确定凹陷是由 \(D\) 引起的，还是由 \(D\) 和 \(N\) 的乘积（决定了每个核心在整个任务中必须处理的总卷积核字节数）引起的。 \[ \text{字节/核心} = \underbrace{\frac{N}{16}}_{\text{卷积核/核心}} \times \underbrace{D}_{\text{权重/卷积核}} \times \underbrace{2}_{\text{字节/权重}} \] 为了区分未知变量，我们反向扫描 \(D\) 和 \(N\)，使编译的任务每个核心都有相同的 1 MiB 静态卷积核数据。执行的寄存器文件的十六进制差异表明只有相关字段（地址、大小）发生了变化： [图表链接：反向扫描 D 和 N 得到的 1 MiB 每核心] 任何 D 和 N 的组合构成每个核心 1 MiB 的总卷积核字节数，都会使核心吞吐量崩溃至观察到的 17 GB/s。考虑到常驻的“L1”KMem 每核心是 64 KiB，我们现在知道有某种推测预取/信用机制在操作这个 1 MiB。 反过来，我们现在知道可以通过不传输 1 MiB 的倍数来避免崩溃；编译器可以通过分割任何编译后每核心恰好产生 1 MiB 内核 DMA 的任务来绕过它。 --- ### 推测预取 高带宽内存控制器有很多理由操作最小传输*行*粒度，而不是单个字节（https://www.goodreads.com/quotes/11711388-of-course-i-d-also-suggest-that-whoever-was-the-genius）。 逆向工程是一门艺术。如果每次传输都以某种行粒度发生，比如内核 DMA 的 64 字节行粒度，那么内核 DMA 控制器和任何预取逻辑也会以行的逻辑单位（而不是字节）来编写。现在转换思路以行为单位思考： 在 N=4096 时，每个 D += 2048 会为总请求字节数增加 1 MiB： \[ 256\ \text{内核/核心} \times 4\ \text{KiB/内核} = 1\ \text{MiB/核心} \] 如果内核 DMA 行粒度是 64 字节（我们知道从 \(2^6\) 字节对齐地址），那么 1 MiB 传输总共请求 \(\texttt{0x4000}\) 个 64 字节行： \[ 1\ \text{MiB/核心} \div 64\ \text{B/行} = 16,384\ \text{行/核心} = \texttt{0x4000}\ \text{行/核心} \] 我们现在推测在 \(\texttt{0x4000}\) 行处有某个计数器发生了回绕。 ``` always @(posedge clk) begin if (reset) line_count <= 14'h0000; else if (advance) line_count <= line_count + 1'b1; // 在 0x3fff + 1 -> 0x0000 时回绕 end ``` 在 \(\texttt{0x4000}\) 或 \(2^{14}\) 处发生回绕，对应于 14 位的存储。还有什么是 \(2^{14}\)？苹果硅芯片上使用的 16 KiB 虚拟内存页面大小。使用 16 KiB 页面，地址位 (\([13:0]\)) 是页面偏移量，不会因虚拟到物理转换而改变。对低位地址位进行运算的预取算法 (`addr & 0x3fff`) 将具有 14 位回绕。 --- 定义 \(k\) 为传输所跨越的 \(\texttt{0x4000}\) 行周期的次数，我称之为一圈： \[ k \equiv \frac{D}{\texttt{0x4000}} \] 定义 \(x\) 为距离第 \(k\) 个凹陷的行数： \[ \text{行/核心} = k \cdot \texttt{0x4000} + x \] 归一化到 \(x\) 后，V 形凹陷在围绕凹陷的 \(x = \pm256\) 行处恢复到正常值，对于围绕 \(\texttt{0x4000}\) 的所有 \(k\) 圈都成立。 [图表链接：凹陷恢复图] ``` 64 B/行 * 256 行 = 16 KiB = 一个页面。 ``` 凹陷恰好发生在相当于一个虚拟内存页面的 DMA 行范围内。这看起来像是一个以页面深度为大小的前瞻预取窗口。 手动叠加 \(D=2048\)（\(k=1\)）和 \(D=4096\)（\(k=2\)）的带宽曲线，当重新居中到凹陷处时，产生几乎相同的带宽曲线： [图表链接：k=1 和 k=2 凹陷处的叠加] 现在绘制每个 \(k\) 的带宽在每个样本 \(x = 0, 32, 64, 128, 256\)（左图）的值；注意每个 \(k\) 曲线随着 \(x \to 0\) 向内收敛。 [图表链接：k 圈带宽扇形图] 一个重要的发现是，每个圈 \(k\) 的时间曲线实际上就是圈 1 曲线按 \(k\) 倍垂直缩放；圈 6 比圈 1 陡峭约 6 倍。通过 \(k\cdot\texttt{0x4000}\) 重新居中每个凹陷后，每个 \(k\) 的带宽曲线都坍缩到同一条线上。每圈斜率在测量数据中是真正的 \(k\) 线性（每个的 \(R^2 = 0.96–0.99\)），斜率 \(= 3.18 \cdot k\) μs/行（右图）。 这些观察结果共同表明： - (1) 受限的传输“剖面”在每个 \(k\) 处重复。如果每个周期经历相同的受限带宽剖面，那么 \(k\) 个周期通过相同剖面传输 \(k \times\) 更多的字节，给出一个 \(k\) 倍更陡的时间曲线。 - (2) 每个周期内的带宽剖面主要由相对中心位移 \(x\) 决定。因此相同的 \(x\) 在每个 \(\texttt{0x4000}\) 行处重现相同的带宽状态。因此内部状态只知道当前 1 MiB 圈内的位置，而不知道传输处于哪个圈 \(k\) 或还剩多少圈。一个 \(k=6\) 的传输包含大约六倍数量的字节经历相同的 \(x\) 相关的速率，因此其额外延迟大约是 \(k=1\) 的 6 倍。 - (3) 在 \(x=0\) 时，零乘以零等于零，每个周期都恰好落在相同的病态状态上。因此无论 \(k\) 如何，传输速率都坍缩到相同的 \(B(0)\)。具体来说，随 \(k\) 缩放的不是带宽坍缩本身，而是在该坍缩速率下传输的数据量。 因此 \(x\) 选择带宽状态；\(k\) 决定该状态重复多少次。 --- 预取环前瞻一次请求一个 1 MiB 环。该环将总传输大小视为 \(k\) 个重复的 1 MiB 池需要获取： \[ S(k,x) = 64\,(k \cdot \texttt{0x4000} + x) = k \cdot 1\ \text{MiB} + 64x \quad \text{字节/核心} \] 如果每个 1 MiB 预取具有某个带宽曲线 \(B(x)\)，那么总传输时间为： \[ t(k,x) = \frac{S(k,x)}{B(x)} = \frac{k \cdot 1\ \text{MiB} + 64x}{B(x)} \] 对 \(x\) 的偏导数为 \[ \boxed{ \frac{\partial t(k,x)}{\partial x} \approx -k \cdot 1\ \text{MiB} \, \frac{B'(x)}{B(x)^2} } \] 对于小的（<256）x 值，这是 \(k\) 主导的： \[ \boxed{ \frac{\partial t}{\partial x} \propto k } \] 下限 = 18 GB/s，每圈 16 MiB，得到 900 μs/圈（与 median_us(x=0)/k = 900 相符）。从下限（~18）恢复到肩部（~45）在 256 行上大致将每圈时间减半，平均 = (900−380)/256 = 2 μs/行。如果靠近凹陷处更陡峭...