趙天辰，清華大學(xué)電子工程系高能效計(jì)算實(shí)驗(yàn)室研究生，研究方向主要是：面向視覺生成的高效算法，與軟硬件協(xié)同設(shè)計(jì)

近年來，隨著視覺生成模型的發(fā)展，視覺生成任務(wù)的輸入序列長(zhǎng)度逐漸增長(zhǎng)（高分辨率生成，視頻多幀生成，可達(dá)到 10K-100K）。與輸入序列長(zhǎng)度呈平方復(fù)雜度的 Attention 操作，成為主要的性能瓶頸（可占據(jù)全模型的 60-80% 的開銷），有明顯的效率優(yōu)化需求。注意力的稀疏化（Sparse Attention）與低比特量化（Attention Quantization）為常用的 Attention 優(yōu)化技巧，在許多現(xiàn)有應(yīng)用中取得優(yōu)秀的效果。然而，這些方法在視覺生成模型中，在低稠密度（<50%）與低比特（純 INT8/INT4）時(shí)面臨著顯著的性能損失，具有優(yōu)化的需求。

本文圍繞著視覺任務(wù)的 “局部性”（Locality）特點(diǎn)，首先提出了系統(tǒng)的分析框架，識(shí)別出了視覺生成任務(wù) Attention 優(yōu)化的關(guān)鍵挑戰(zhàn)在于 “多樣且分散” 的注意力模式，并且進(jìn)一步探索了該模式的產(chǎn)生原因，并揭示了多樣且分散的注意力模式，可以被統(tǒng)一為代表 “局部聚合” 的塊狀模式。然后，提出了一種簡(jiǎn)單且硬件友好的離線 “Token重排” 方案以實(shí)現(xiàn)注意力模式的統(tǒng)一化，并設(shè)計(jì)了針對(duì)性的稀疏與量化方法，配合高效的 CUDA 系統(tǒng)設(shè)計(jì)，展現(xiàn)了更優(yōu)異的算法性能保持與硬件效率提升。最后，本文討論了該方案更廣泛的應(yīng)用空間，與對(duì)視覺生成算法設(shè)計(jì)的啟發(fā)。

• 論文標(biāo)題：PAROAttention: Pattern-Aware ReOrdering for Efficient Sparse and Quantized Attention in Visual Generation Models
• 論文鏈接：https://arxiv.org/abs/2506.16054
• 項(xiàng)目主頁：https://a-suozhang.xyz/paroattn.github.io/

1. 分析框架：關(guān)鍵問題與如何利用局部性（Locality）

如上文所述，一系列現(xiàn)有的注意力稀疏化與低比特量化方案已取得了進(jìn)展，但是還存在著一定的挑戰(zhàn)與改進(jìn)空間：

對(duì)于稀疏化，一系列現(xiàn)有方案（DiTFastAttn，SparseVideoGen，Sparse-vDiT）嘗試依據(jù)視覺注意力圖的獨(dú)特模式，設(shè)計(jì)針對(duì)性的稀疏掩碼（Sparse Mask，如 “窗口狀的”，“多對(duì)角線”，“垂直線的”），并將其進(jìn)行組合。然而，適配多樣且分散的注意力模式，給稀疏掩碼的設(shè)計(jì)與選擇機(jī)制帶來了嚴(yán)峻的挑戰(zhàn)。本文嘗試采用另一種視角與方法，并不涉及復(fù)雜的掩碼選擇機(jī)制來適配復(fù)雜多樣的注意力模式，而是設(shè)計(jì)方案 “重整注意力模式”。讓多樣且復(fù)雜的注意力模式，統(tǒng)一為硬件友好的塊狀注意力模式，讓稀疏方案設(shè)計(jì)更加簡(jiǎn)單有效。

對(duì)于低比特量化，現(xiàn)有方案（SageAttention 系列）SageAttentionV2 可以將 Attention 中的 QK 計(jì)算（Query 與 Key 的乘法）降低至 INT4，但是 PV（AttentionMap 與 Value 的矩陣乘）計(jì)算仍然需要保持為較高的 FP8。最新版本的 SageAttentionV3 采用了 FP4 量化，但僅在最新的 B 系列 Nvidia GPU 上有支持。本文嘗試分析了更低位寬的定點(diǎn)量化（全流程 INT4）的關(guān)鍵問題，并給出了解決方案。

圖：本文（PAROAttn）的優(yōu)化思路：重整注意力圖以便稀疏與量化處理

為尋找 Attention 稀疏與量化的統(tǒng)一解決方案，本文嘗試分析 Attention 效率優(yōu)化中稀疏與低比特量化的關(guān)鍵問題，來自于視覺注意力圖多樣且分散的獨(dú)特?cái)?shù)據(jù)分布 (如下圖左側(cè)所示)：

稀疏注意力方案設(shè)計(jì)需要從 2 方面考慮：保持算法性能，與提升硬件效率。

• 從算法性能角度，需要避免在稀疏過程中錯(cuò)誤的刪除重要值。由于視覺注意力模式存在多樣的結(jié)構(gòu)（對(duì)角線，縱向，塊狀等），且這些特征隨著不同的時(shí)間步，不同的控制信號(hào)而動(dòng)態(tài)變化。因此，注意力模式的多樣性，與動(dòng)態(tài)變化，導(dǎo)致設(shè)計(jì)的掩碼難以完全涵蓋重要值，對(duì)算法性能的保持帶來困難。
• 從硬件效率角度，需要設(shè)計(jì) “結(jié)構(gòu)化” 的稀疏掩碼，以跳過整塊計(jì)算來獲得實(shí)際硬件收益（任意不規(guī)則稀疏需要引入額外的索引操作，且可能導(dǎo)致計(jì)算負(fù)載零散，使得加速收益折損）。特別的，由于 FlashAttention 涉及逐塊進(jìn)行注意力計(jì)算，因此稀疏化的過程中，也應(yīng)考慮如何與其適配。由于視覺注意力圖的模式，往往不與 FlashAttention 中的分塊所對(duì)應(yīng)（對(duì)角線模式中，每個(gè)塊中僅有少量較大值）。因此，注意力圖模式的分散性，使得結(jié)構(gòu)化稀疏難以取得，難以獲得有效的硬件效率提升。

對(duì)于低比特量化算法的設(shè)計(jì)：關(guān)鍵問題為如何盡量減少量化損失。

• 現(xiàn)有工作（如 ViDiT-Q）已經(jīng)分析并指出了低比特量化的關(guān)鍵誤差來源在于 “量化組內(nèi)的數(shù)據(jù)分布差異”，對(duì)于注意力量化，為適配 FlashAttention，需要選擇塊狀的量化分組。然而，“對(duì)角線式” 的視覺注意力模式，導(dǎo)致塊狀的量化分組中，對(duì)角線上的元素成為離群值，帶來了巨大的組內(nèi)數(shù)據(jù)差異，而導(dǎo)致了顯著的量化損失。因此，視覺注意力模式的數(shù)據(jù)分布，導(dǎo)致了顯著的量化損失。

圖：視覺生成稀疏與量化的關(guān)鍵問題來自于多樣分散的注意力模式，與本文的解決方案：采用Token重排以改進(jìn)注意力圖為統(tǒng)一的分塊模式

為解決視覺注意力圖多樣且分散的獨(dú)特?cái)?shù)據(jù)分布給注意力稀疏與量化所帶來的挑戰(zhàn)。本文的技術(shù)路線為：對(duì)注意力圖進(jìn)行 “重整”（Reorganize），以獲得更加統(tǒng)一且易處理的注意力模式。

受到視覺特征提取具有 “局部性” 的先驗(yàn)啟發(fā)（CNN，SwinTransformer 的設(shè)計(jì)理念，與 Hubel 與 Wiesel 的生物學(xué)實(shí)驗(yàn)），本文進(jìn)一步分析了視覺注意力模式多樣性的產(chǎn)生原因，并發(fā)現(xiàn)了 “多樣的視覺注意力模式本質(zhì)上都在描述空間上的局部聚合”。

如下圖所示，在 Transformer 的處理過程中，原本三維空間（F - 幀數(shù)，H，W - 每幀的圖像寬高）會(huì)被轉(zhuǎn)化為一維的標(biāo)記序列（Token Sequence），按照默認(rèn)的 [F,H,W] 的順序排列。這會(huì)導(dǎo)致在除了內(nèi)存上連續(xù)的最后一維（W）之外維度的三維空間相鄰像素，在標(biāo)記序列中呈現(xiàn)為按照一定的間隔排列。

因此，多對(duì)角線的注意力模式，本質(zhì)上是在描述 “其他維度上的局部聚合”，并可以通過Token順序的重排列，轉(zhuǎn)化為代表局部聚合的塊狀模式（將局部聚合的維度轉(zhuǎn)化為內(nèi)存上連續(xù)的維度，如 [F,H,W] -> [F,W,H]）。

本文進(jìn)一步驗(yàn)證了，每個(gè)不同的注意力頭（Head），在不同情況下，呈現(xiàn)出一致的在某個(gè)維度上的局部聚合，進(jìn)而可以通過為每個(gè) head 選取合理的Token重排（Token Reorder）方案，將多樣且分散的注意力模式，轉(zhuǎn)化為統(tǒng)一的，硬件友好的塊狀模式，以便于 Attention 的稀疏與量化。該方案利用了算法側(cè)視覺特征提取的局部性（更好的數(shù)值 Locality），并將其與硬件計(jì)算的局部性將對(duì)應(yīng)（更好的內(nèi)存與計(jì)算 Locality），從而獲得了同時(shí)更優(yōu)的算法性能保持，與硬件效率提升。

圖  視覺特征提取 “局部性” 的示意圖

2. 方案設(shè)計(jì)

整體框架

方案流程如下圖所示，對(duì) Attention 計(jì)算的主要瓶頸，兩個(gè)大規(guī)模矩陣乘（QK 與 PV）都進(jìn)行了稀疏與量化優(yōu)化，顯著減少其硬件開銷。本文基于少量矯正數(shù)據(jù)離線決定了每個(gè)注意力頭（Head）的Token重排方案，與對(duì)應(yīng)的稀疏掩碼，幾乎不在推理時(shí)引入額外的開銷。在推理時(shí)，僅需跳過稀疏掩碼所對(duì)應(yīng)的 attention 分塊，并對(duì)剩余的部分逐塊進(jìn)行低比特量化。

圖 PAROAttention 稀疏與量化方案的流程

Token重排方案（PARO：Pattern-Aware Token Reordering）

本文發(fā)現(xiàn)每個(gè)不同的注意力頭（Head），在不同情況下，呈現(xiàn)出一致的在某個(gè)維度上的局部聚合。因此，可以離線地對(duì)每個(gè)注意力頭，選擇恰當(dāng)?shù)腡oken重排方式，將注意力圖轉(zhuǎn)化為展示局部聚合的塊狀（Block-wise）模式。

本文發(fā)現(xiàn)了重排列中的一種特殊方式，維度置換（Permutation），就可以取得不錯(cuò)的效果。對(duì)于視頻生成模型的特征 [F,H,W]，本文為每個(gè)注意力頭 6 種可能的置換方式，離線選取最優(yōu)的置換方式，以獲得需要的數(shù)據(jù)分布方式。由于對(duì)于注意力稀疏與量化，具有不同的數(shù)據(jù)分布需求。因此，本文針對(duì)稀疏和量化分別設(shè)計(jì)了重排方式的選取指標(biāo)，并將兩者組合作為最終指標(biāo)。

• 稀疏角度：為減小結(jié)構(gòu)化稀疏所帶來的損失，要求盡量多的分塊是完全稀疏的（Block Sparse）
• 量化角度：為減少塊內(nèi)數(shù)據(jù)分布差異大而導(dǎo)致的量化損失，要求塊內(nèi)數(shù)據(jù)分布是盡量均勻的（Block Uniform）

如下圖所示，稀疏與量化對(duì)注意力圖的分布需求不同，需要組合兩者需求，才能找到同時(shí)適合兩者的重排方式。經(jīng)過合適的重排處理之后，注意力圖呈現(xiàn)塊狀且較為集中的分布，以適配稀疏與量化處理。

圖：不同重排方式的注意力圖示意

稀疏方案

現(xiàn)有的稀疏注意力方案可分為 2 種方式：（1）動(dòng)態(tài)稀疏方案（如 SpargeAttention）在線依據(jù)注意力值生成稀疏掩碼；（2）靜態(tài)稀疏方案（如 DiTFastAttn）：離線生成稀疏掩碼。兩者各有其優(yōu)劣。盡管本方法設(shè)計(jì)的Token重排（PARO）方案能夠同時(shí)幫助動(dòng)態(tài)與靜態(tài)方案，本文對(duì)兩者優(yōu)劣進(jìn)行的分析，并最終選取了靜態(tài)稀疏方案，作為 PAROAttention 的主要稀疏方案，具體分析如下：

對(duì)于動(dòng)態(tài)稀疏（Dynamic Approach）：

• 在性能保持方面，雖然動(dòng)態(tài)的方案能夠自然適配動(dòng)態(tài)變化的模式。但是由于需要在線產(chǎn)生稀疏掩碼，只能基于 Softmax 之前（Presoftmax）的注意力值，它們的相對(duì)均勻，不呈現(xiàn)明顯模式，難以準(zhǔn)確的識(shí)別出對(duì)應(yīng)模式。
• 在硬件效率方面，動(dòng)態(tài)稀疏方案引入了在線計(jì)算出稀疏掩膜的額外開銷（overhead），該開銷與掩膜預(yù)測(cè)的準(zhǔn)確度互為權(quán)衡，若要獲得準(zhǔn)確的掩膜，則需要引入相對(duì)較大的額外計(jì)算。該額外預(yù)測(cè)過程，一般需要精細(xì)設(shè)計(jì)的 CUDA Kernel 才能夠獲得較高的效率收益。
• 總結(jié)來看，在較低稀疏比下，動(dòng)態(tài)稀疏方式的準(zhǔn)確性與效率提升存在瓶頸，因此本文訴諸靜態(tài)稀疏方案。

對(duì)于靜態(tài)稀疏（Static Approach）：

• 在性能保持方面，由于靜態(tài)確定的注意力圖，難以適配多樣且動(dòng)態(tài)變化的注意力模式，因此靜態(tài)稀疏方案通常會(huì)造成相比動(dòng)態(tài)方案更顯著的性能損失。然而，PAROAttention 的注意力圖重整，已將多樣動(dòng)態(tài)變化的注意力模式，轉(zhuǎn)化為了規(guī)整且統(tǒng)一的模式，解決了這一靜態(tài)稀疏的關(guān)鍵挑戰(zhàn)。因此，通過利用模式更明顯的 Softmax 后注意力圖，能夠獲得比動(dòng)態(tài)方案更優(yōu)的算法性能保持。
• 在硬件效率方面，雖然避免了在線計(jì)算出稀疏掩膜的額外計(jì)算開銷，但是離線稀疏掩碼會(huì)帶來額外的顯存開銷。本文針對(duì)該問題進(jìn)行了對(duì)應(yīng)優(yōu)化（見下文 “CUDA 系統(tǒng)設(shè)計(jì)” 部分）。

經(jīng)過重排列處理之后，注意力圖呈現(xiàn)出統(tǒng)一的集中的分塊模式。因此，本文僅需離線統(tǒng)計(jì)每塊中的 attention 數(shù)據(jù)之和，并設(shè)計(jì)閾值判斷當(dāng)前塊是否需要被跳過（該閾值可以用于調(diào)節(jié)稠密度），就可以離線獲取到稀疏掩碼，在推理時(shí)不引入任何額外開銷（overhead）。如下圖所示，相比其他現(xiàn)有的靜態(tài)注意力稀疏方案，由于預(yù)先對(duì)注意力模式的統(tǒng)一化，PAROAttention 避免了復(fù)雜且受限的掩膜設(shè)計(jì)，而能夠與原圖非常契合的稀疏掩膜。

量化方案：

對(duì)于低比特量化，評(píng)估量化損失的關(guān)鍵指標(biāo)是分組內(nèi)的數(shù)據(jù)差異，現(xiàn)有文獻(xiàn)通常采用不均衡度（Incoherence）進(jìn)行衡量，被定義為當(dāng)前數(shù)據(jù)組中的最大值，除以平均值（x.max () /x.abs ().mean ()）。經(jīng)過合適的Token重排之后，Attention Map 塊內(nèi)的顯著數(shù)據(jù)差異得到明顯緩解，從而可以支持更低位寬的量化。

CUDA 系統(tǒng)設(shè)計(jì)

最小化額外開銷：PAROAttention 所引入的額外開銷主要有以下兩方面，本文在系統(tǒng)層面進(jìn)行了針對(duì)性優(yōu)化以最小化額外開銷。

• 在線的Token重排開銷：雖然重排方式離線確定，但是Token重排的過程（維度置換）需要在線進(jìn)行。為了避免一次顯示的從 GPU Global Memory 到 Shared Memory 的內(nèi)存搬移開銷，本文進(jìn)行了算子融合（Layer Fusion）的操作，僅修改重排前算子寫入地址的順序，所引入的額外開銷可忽略。
• 靜態(tài)稀疏掩碼的顯存開銷：由于注意力圖體量較大，離線決定的稀疏掩碼，可能會(huì)占用 GB 級(jí)別的 GPU 顯存。為減少該開銷，本文采用了預(yù)?。≒refetch）策略，通過新建一個(gè) CUDA Stream，在每次運(yùn)算時(shí)，只讀取當(dāng)前層的稀疏掩碼，可以將額外的顯存開銷降低到若干 MB 級(jí)別。

兼容性：PAROAttention 的稀疏與量化方案都逐塊處理，可直接與兼容 FlashAttention。由于重排與稀疏掩碼均離線完成，無需精細(xì)的 CUDA Kernel 優(yōu)化，僅需基于 FlashAttention 進(jìn)行跳過整塊計(jì)算的支持，能夠廣泛適配各種場(chǎng)景。

3. 軟硬件實(shí)驗(yàn)結(jié)果

算法性能保持效果

本文在主流視頻（CogVideo）與圖片生成模型（Flux）上測(cè)試了多方面指標(biāo)，包括了：

• 視頻質(zhì)量指標(biāo)：CLIPSIM 衡量語義一致性；VQA 衡量視頻質(zhì)量；FlowScore 衡量時(shí)間一致性；
• 與浮點(diǎn)生成差異：如 FVD-FP16 衡量特征空間差異，PSNR/CosSim 衡量像素空間差異，SSIM 衡量結(jié)構(gòu)相似性。

典型的實(shí)驗(yàn)結(jié)論概括如下：

（1）其他基線的稀疏方案在相對(duì)較高稀疏比（50%）時(shí)，仍會(huì)造成可觀的質(zhì)量損失，包括內(nèi)容變化，圖像模糊等；而 PAROAttention 的稀疏化方案，可以在 20% 的較高稀疏比情況下，依然生成和浮點(diǎn)結(jié)果非常相似的結(jié)果，獲得比基線方案 50% 更好的多方面指標(biāo)。

（2）Token重排方案 PARO，并不局限于靜態(tài)稀疏方案。其與動(dòng)態(tài)稀疏方案 SpargeAttention 能夠直接適配，并提升生成效果。將 30% 稠密度的 SpargeAttention 組合 PARO，可以獲得與 50% 稠密度 SpargeAttention 同等的生成質(zhì)量。將加速比從 1.67x 提升至 2.22x。

（3）相比于 SageAttentionV2（QK INT4，PV FP8），PAROAttention 的量化方案可以在無精度損失的情況下，進(jìn)一步將 PV 量化到 INT4。

（4）PAROAttention 的稀疏與量化方案可以并行使用，最激進(jìn)的優(yōu)化方案（50%+INT4）相比浮點(diǎn)能取得近 10 倍的 Attention 部分延遲優(yōu)化，同時(shí)獲得與僅能取得 2x 左右延遲優(yōu)化的基線方法類似的算法性能保持。

硬件加速效果

本文進(jìn)一步對(duì)系統(tǒng)層面優(yōu)化技巧進(jìn)行了分析，關(guān)鍵實(shí)驗(yàn)結(jié)論如下：

（1）PAROAttention 的稀疏方案，同時(shí)取得了更優(yōu)的算法性能保持與效率提升。以 50% 稠密度為例，PAROAttention 取得了 1.73x 的 attention 加速，超過同等情況下的 SpargeAttention（1.67x）與 SparseVideoGen（1.42x），由于靜態(tài)稀疏方案幾乎不會(huì)引入額外開銷，而基線方案的在線稀疏掩膜生成 / 選擇會(huì)造成 6% 到 10% 左右的額外開銷，該開銷在更低稠密度下顯得給更為明顯。

（2）PAROAttention 的加速比與理論上限較為接近（50% 稠密度，理論 2 倍，實(shí)際 1.73 倍），凸顯了方案的硬件友好性。

（3）PAROAttention 的各方面額外開銷 overhead 得到了有效減少，控制在整體的 1% 之內(nèi)。

總結(jié)與未來指引

總結(jié)來看，本文關(guān)注了視覺生成任務(wù)的 “局部性” 特性。通過一個(gè)簡(jiǎn)單且有效的Token重排操作，可以同時(shí)實(shí)現(xiàn)算法側(cè)視覺特征提取的局部性（更好的數(shù)值 Locality），并將其與硬件計(jì)算的局部性相對(duì)應(yīng)（更好的內(nèi)存與計(jì)算 Locality），從而獲得了同時(shí)更優(yōu)的算法性能保持，與硬件效率提升。PAROAttention 的方案主要圍繞推理效率優(yōu)化設(shè)計(jì)，但是采用Token重排來更好利用特征提取局部性的思想并不局限于推理優(yōu)化中。不同的注意力頭自主的學(xué)習(xí)到在不同維度上的局部聚合，可以啟發(fā)優(yōu)化訓(xùn)練方法，與圖像的參數(shù)化方式，三維空間的位置編碼設(shè)計(jì)，并進(jìn)一步推動(dòng)具有合理歸納偏置（Inductive Bias）的視覺基座模型的構(gòu)建。