走進任何現代人工智能實驗室、數據中心或自動駕駛汽車開發環境，你都會聽到工程師們喋喋不休地談論浮點運算次數(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)