为什么内存重排是人工智能计算的隐形税

走进任何现代人工智能实验室、数据中心或自动驾驶汽车开发环境，你都会听到工程师们喋喋不休地谈论浮点运算次数(FLOPS)、运算次数(TOPS)、稀疏性、量化和模型扩展定律。这些指标充斥着新闻头条和产品数据手册。但如果你花时间与那些实际构建或优化这些系统的人交流，你会发现一个截然不同的真相：原始的算术能力并非决定实际性能的关键。

真正重要的是数据移动的效率。而对于当今大多数人工智能加速器来说，数据移动与一个鲜少在编译器和硬件圈子之外被提及的现象——内存重排(memory swizzling)——紧密相连。

内存重排是现代人工智能系统付出的最大隐形代价之一。它既不能提高算法处理效率，也不能提高准确性，更不能降低能耗，更不会带来任何新的洞见。它的存在仅仅是为了弥补几十年前设计选择遗留下来的架构缺陷。随着人工智能模型规模越来越大、结构越来越不规则，这种代价也在不断增加。

本文探讨了内存重排存在的原因、发展历程、代价，以及一种截然不同的架构理念——特别是以寄存器为中心的模型——如何彻底消除对内存重排的需求。

---

一个鲜为人知的问题：数据的存储方式与硬件需求不符

在任何人工智能教程中，张量都被呈现为有序的数学对象，它们以完美的布局整齐地排列在内存中。这些布局对程序员来说直观易懂，并且能够很好地融入PyTorch或TensorFlow等高级框架。

然而，硬件并非如此看待数据。

现代加速器——GPU、TPU和NPU——围绕并行计算单元构建，这些单元期望特定的数据格式：固定大小的图块、严格的对齐边界、具有可预测步长模式的序列，以及能够无冲突地映射到内存库的排列方式。

遗憾的是，现实世界中的张量永远不会以这些格式出现。在处理开始之前，数据必须被重塑、重新分块、重新排序或重新打包成硬件期望的格式。这种重塑过程称为内存重排。

你可以这样理解：算法以矩阵和张量的方式思考；计算硬件则以分块、通道和存储体的方式思考。内存重排就是转换层——一个需要耗费时间和能量的转换过程。

---

为什么分层内存迫使我们进行内存重排？

如今，几乎所有加速器都使用分层内存堆栈，其各层从上到下依次包含：寄存器；共享内存或暂存内存；L1缓存、L2缓存，有时甚至包括L3缓存；高带宽内存(HBM)；以及位于堆栈底部的外部动态随机存取存储器(DRAM)。

每一层的大小、延迟、带宽、访问能耗以及非常重要的对齐约束都各不相同。这是CPU式架构的遗留产物，在这种架构中，缓存隐藏了内存延迟。参见图1和表1。

图1：当前所有硬件处理器中典型分层内存堆栈的容量和带宽属性。(来源：VSORA)

表1：当前所有硬件处理器中典型分层内存堆栈的容量、延迟、带宽和访问能耗。(来源：VSORA)

GPU继承了这种模型，并在此基础上增加了单指令多线程(SIMT)执行。这使得它们拥有惊人的计算能力，但也对数据布局方式极其敏感。如果线程束(warp)中的相邻线程无法访问相邻的内存位置，性能就会急剧下降。如果数据块边界没有对齐，张量核心就会停滞。如果共享内存库发生冲突，所有操作都会被暂停。

TPU也面临着类似的限制，只是机制有所不同。它们的脉动阵列就像精心编排的传送带一样运行。数据必须按正确的顺序和时间到达。如果权重没有按照块优先的格式排列，脉动阵列就无法高效运行。

从智能手机芯片到汽车系统，基于NPU的加速器也面临着同样的问题：多库SRAM、固定的向量宽度以及视觉工作负载的2D局部性要求。如果没有内存重排，数据到达计算引擎时就会“错位”，导致性能急剧下降。

在所有这些情况下，重排都不是一种优化——而是一种生存机制。

---

内存重排的隐性成本

内存重排耗时，有时甚至非常耗时

在实际工作负载中，内存重排通常会消耗总运行时间的20%到60%。这并非笔误。在卷积神经网络中，一半的时间可能都花在了NHWCNCHW转换上；也就是说，这是在内存中布局4D张量的两种不同方式。在Transformer模型中，大量时间被浪费在重塑Q/K/V张量、拆分头、为GEMM重新打包tile以及重组输出上。

内存重排会消耗能量，而能量才是真正的限制因素

一次MAC操作大约消耗四分之一皮焦耳的能量。从DRAM中移动一个值可能需要500皮焦耳的能量。从DRAM中移动数据所消耗的能量大约是执行一次基本乘加运算的1000倍。

内存重排需要读取大量数据块，重新排列后再写回。而且，每个层通常要重复多次这样的操作。当80%的能源预算都用于移动数据而不是进行计算时，内存重排就变得不容忽视。

内存重排会增加内存使用量

大多数内存重排都需要临时缓冲区：打包的图块、暂存缓冲区和重塑的张量。这些额外的内存占用可能会使模型超出L2、L3甚至HBM的极限，从而迫使模型进行更多的数据移动。

内存重排使软件更难编写且更难移植

问问CUDA工程师，是什么让他们夜不能寐。问问TPU编译器工程师，为什么XLA的布局推断代码长达数千页。问问任何为移动设备编写NPU内核的人，为什么他们害怕通道排列。

答案是：内存重排。由于硬件需要非常具体的布局，软件必须承载巨大的复杂性。每一种新的模型架构——无论是卷积神经网络(CNN)、长短期记忆网络(LSTM)、Transformer模型还是扩散模型——都会增加新的布局模式，而这些模式必须得到支持。

最终形成了一个由布局启发式算法、张量变换和对性能要求极高的内存管理机制维系的生态系统。

---

主流架构如何依赖内存重排

1.英伟达GPU

Tensor核心需要特定的tile-main布局。共享内存被分块存储，避免冲突需要内存重排。线程束必须合并内存访问；否则，效率会大幅下降。即使是地球上优化程度最高的GPU库cuBLAS和cuDNN，也充满了内部内存重排内核。

2.谷歌TPU

TPU依赖于脉动阵列。数据流经这些阵列必须完全有序。权重和激活值会不断重新排列，以与脉动结构保持一致。XLA的大部分功能都是为了管理数据布局而存在的。

3.AMD CDNA、ARM Ethos、苹果ANE和高通AI引擎

所有这些架构都执行内存重排。Morton分块、交错、通道堆叠等等。这是一种通用模式。每个使用分层内存的架构都继承了对内存重排的需求。

---

另一种理念：从根本上消除内存重排

现在，想象一下，退后一步，从根本上重新思考人工智能硬件。与其接受如今复杂的内存层次结构——缓存层、共享内存块、分块式SRAM和对齐规则——不如设想一种基于更简单前提的架构。

如果根本没有内存层次结构呢？如果整个系统围绕着一个庞大而扁平的寄存器阵列运转呢？如果编译器而非硬件，以确定性的精度协调每一次数据移动呢？如果所有常见的担忧——对齐、存储体冲突、分块策略和合并规则——都因为不再重要而消失呢？

这就是以寄存器为中心的架构背后的理念。数据不再沿着内存层级上下移动，而是直接驻留在执行计算的寄存器中。这种架构的组织方式不是围绕数据的移动，而是围绕数据的可用性。

这意味着：

无需预热或担心缓存未命中

无需调度线程束

无需避免内存冲突

无需匹配数据块大小

无需考虑张量布局

对形状或步长不敏感，因此完全无需进行内存重排

在这样的系统中，编译器始终精确地知道每个值所在的位置，以及它接下来应该去往何处。它不会进行推测、预取、分块或依赖启发式算法。它不会祈祷硬件运行正常。相反，数据放置变成了一个可解决、可预测的问题。

最终，机器的吞吐量保持稳定，延迟变得可预测，能耗大幅下降，因为不必要的数据移动已被排除在外。在这个系统中，性能不再受制于内存操作——真正的计算，也就是数学运算，终于成为了核心。

---

人工智能的未来：为何以寄存器为中心的架构至关重要

随着人工智能系统的演进，整齐划一的张量和完美的矩形计算单元的世界正在逐渐消逝。现代模型不再是可预测的、步调一致的密集层堆叠。相反，它们向各个方向扩展：它们能够接收多模态输入，整合稀疏和不规则结构，进行自适应推理，并处理越来越长的序列。它们还必须实时响应安全关键型应用，并且必须在严格的能耗预算内完成这些操作——从汽车到边缘设备，无一例外。

换句话说，塑造GPU和TPU架构的那些假设——对规则性、密集网格和整齐计算单元的预期——正在瓦解。未来的工作负载不再符合硬件的预期。

以寄存器为中心的架构提供了一条截然不同的道路。因为它直接在数据所在的位置进行操作，而不是强制将数据转换为适合计算单元的格式，所以它绕过了内存重排的整个机制。它不依赖于固定的张量形状。

当访问模式变得不规则或动态时，它不会出现故障。它避免了为了满足计算单元的需求而进行代价高昂的内存重排。随着模型变得越来越异构和复杂，这种架构能够随着其复杂性的增加而扩展，而不是与之对抗。

这不仅仅是渐进式的改进，它代表了我们对人工智能计算思维方式的转变。通过消除不必要的数据移动——现代加速器中最大的瓶颈和能耗来源——以寄存器为中心的架构使硬件能够适应人工智能本身复杂多变的现实。

内存重排是每个分层内存加速器都必须付出的隐形代价。它是GPU、TPU、NPU以及几乎所有人工智能芯片运行的基础。同时，它也日益成为一个负担。内存重排会引入延迟、消耗能源、增加内存使用量并使软件复杂化——而对实际的数学运算却没有任何贡献。

一种以寄存器为中心的架构通过移除导致内存重排的层级结构，从根本上消除了内存重排。它用确定性的数据流取代了猜测和启发式方法。它优先考虑局部性，而无需重新排列数据。它让算法驱动硬件，而不是反过来。

随着人工智能工作负载变得更加不规则、动态和对功耗敏感，能够保持数据稳定和可预测(而不是无休止地重新排列数据)的架构将定义下一代计算。

内存重排是上一个硬件时代的必要补救措施，它不应该定义下一个时代。

Lauro Rizzatti是VSORA的商业顾问，VSORA是一家提供重新定义性能的硅半导体解决方案的技术公司。他是一位知名的芯片设计验证顾问，也是硬件仿真领域的行业专家。

(责编：Franklin)