化：從2.34ms到0.62ms實戰(zhàn))
1. Softmax算子為什么值得單獨寫一篇優(yōu)化文章做AI算子開發(fā)的朋友應該都有體會Softmax這個算子看著人畜無害實際上特別“刁鉆”。它的數(shù)學形式極其簡單但優(yōu)化空間和踩坑概率在常用算子中絕對排得上號。尤其到了DCU這種國產(chǎn)加速卡上想把Softmax跑出接近理論帶寬的性能需要操心的事情遠比想象中多。先說清楚Softmax在干什么。給定一個輸入向量Softmax把每個元素映射成一個概率值公式長這樣[ \text{Softmax}(x_i) \frac{e^{x_i - \max(x)}}{\sum_{j} e^{x_j - \max(x)}} ]減去最大值那一步不是可有可無是為了數(shù)值穩(wěn)定性。如果不減當輸入里有大數(shù)時(e^{x_i})直接溢出成inf后面全完蛋。這在Transformer的注意力層里尤其致命因為QK^T的點積結(jié)果動輒幾十上百不穩(wěn)定的Softmax會讓訓練直接發(fā)散。但Softmax真正的麻煩不在數(shù)學而在訪存特性。這個算子對每個元素只做幾次浮點運算計算密度極低屬于典型的訪存密集型任務。也就是說性能瓶頸幾乎完全取決于你能多快把數(shù)據(jù)從顯存搬進寄存器、把結(jié)果寫回去。在DCU上做優(yōu)化本質(zhì)上是在跟內(nèi)存帶寬較勁。這篇文章我從一個實際優(yōu)化案例出發(fā)完整走一遍從樸素實現(xiàn)到深度調(diào)優(yōu)的過程。案例背景是某推理模型中的注意力模塊Softmax算子的輸入shape為[batch, heads, seq_len, seq_len]其中batch4heads32seq_len512數(shù)據(jù)類型為FP16。優(yōu)化目標是把單次調(diào)用延遲從2.34ms壓到1ms以內(nèi)。這個場景非常有代表性。它既有行主序二維矩陣的常規(guī)Softmax特征又有大批量小矩陣的并行特性還涉及FP16的精度處理幾乎把DCU上算子優(yōu)化能遇到的典型問題都覆蓋了。適合讀這篇文章的人有三類一是正在做算子移植需要把PyTorch或CUDA代碼遷移到DCU上的工程師二是做推理優(yōu)化被Attention里Softmax耗時困擾的算法工程師三是對并行編程有興趣想看看國產(chǎn)加速卡生態(tài)實際長什么樣的技術(shù)愛好者。先說結(jié)論最終優(yōu)化后的Softmax算子單次調(diào)用耗時從2.34ms降到了0.62ms加速比3.77倍通過了一系列精度對比測試。這個結(jié)果不是靠某個單一技巧拿到的而是一整套策略組合的效果。下面我把每一步的思路、實現(xiàn)細節(jié)和踩坑經(jīng)歷都拆開講。2. DCU硬件結(jié)構(gòu)與并行編程模型的核心認知2.1 DCU的硬件架構(gòu)與計算單元布局要優(yōu)化DCU上的算子首先得搞清楚硬件長什么樣。DCUDeep Computing Unit是基于GPGPU架構(gòu)設(shè)計的加速卡核心計算單元叫計算引擎Compute EngineCE每個CE內(nèi)部包含多個SIMD單元每個SIMD單元又由若干個線程插槽Thread Slot組成。以我使用的DCU型號為例單卡有60個CE每個CE包含4個SIMD單元每個SIMD單元可以同時駐留一定數(shù)量的wavefrontDCU的調(diào)度單位等價于CUDA里的warp通常為64個線程。這意味著一個CE上同時活躍的硬件線程數(shù)非常可觀能夠用大量并行線程來掩蓋訪存延遲。DCU的存儲層次和主流GPGPU類似從快到慢依次是寄存器文件、L1緩存、共享內(nèi)存Local Data ShareLDS、L2緩存、全局顯存。每個SIMD單元有獨立的L1緩存和LDS所有CE共享L2緩存和全局顯存。有個關(guān)鍵區(qū)別需要特別強調(diào)DCU沒有像CPU那樣依賴大緩存來加速隨機訪問而是靠海量線程的并行切換來隱藏訪存延遲。這決定了優(yōu)化思路的核心——你要做的不是減少線程數(shù)而是盡可能多地創(chuàng)建可以并行執(zhí)行的線程讓它們在等待內(nèi)存數(shù)據(jù)時切換執(zhí)行其他計算。2.2 HIP編程模型與DTK工具鏈DCU的編程模型遵循HIPHeterogeneous Interface for Portability規(guī)范。HIP的語法和CUDA高度相似從CUDA代碼遷移到HIP通常只需要做少量替換比如__global__改__global__HIP里也是__global__blockIdx.x變hipBlockIdx_xthreadIdx.x變hipThreadIdx_x。但有些細節(jié)差異會在后面實操部分講清楚。DCU的軟件開發(fā)工具包是DTKDCU Toolkit里面包含HIP編譯器基于LLVM、HIPify工具用于自動把CUDA代碼轉(zhuǎn)換成HIP代碼、性能分析工具等。在使用過程中我用的是DTK 24.04版本編譯命令大致如下hipcc -O3 -stdc17 -DNDEBUG -o softmax_bench softmax_bench.cpp注意-O3在高性能算子編譯中基本是標配但如果你在代碼里用了restrict關(guān)鍵字或者__builtin_assume一定要檢查生成匯編是否真正優(yōu)化到位因為DCU編譯器在某些場景下對復雜指針別名的處理不如預期保守的代碼寫法反而會拖累性能。2.3 Wavefront、線程塊與調(diào)度機制DCU的調(diào)度機制和NVIDIA GPU最直觀的差異就是wavefront大小為64線程而CUDA的warp是32線程。這個差異影響深遠。在Softmax算子優(yōu)化中我們經(jīng)常需要做線程間的數(shù)據(jù)歸約比如求最大值、求和。在CUDA里你習慣用__shfl_down_sync在32個線程內(nèi)做shuffle歸約到了DCU/ROCm平臺對應的是__shfl_down同時因為wavefront有64個線程歸約的步數(shù)多了一級。另一個需要適應的是CE的調(diào)度粒度。DCU的硬件調(diào)度單元以wavefront為單位一個wavefront里的線程執(zhí)行相同的指令如果出現(xiàn)分支分歧會出現(xiàn)串行執(zhí)行不同路徑的情況性能損失明顯。所以在設(shè)計內(nèi)核時盡量保證同一wavefront內(nèi)的線程走相同分支或者干脆避免分支。實際測算下來同樣一份Softmax內(nèi)核我最初從CUDA直接搬過來時性能只有預期的60%左右。排查之后發(fā)現(xiàn)主要問題就在wavefront大小的差異上線程塊維度和歸約邏輯沒有針對64線程重新設(shè)計導致大量線程空轉(zhuǎn)。3. Softmax算子的基礎(chǔ)實現(xiàn)與首輪性能摸底3.1 樸素實現(xiàn)一個線程處理一行在動手優(yōu)化之前先把基礎(chǔ)版本寫出來作為后續(xù)優(yōu)化的基準線。最簡單直觀的思路是讓一個線程處理輸入矩陣的一行。假設(shè)輸入是[M, N]的二維矩陣那就開M個線程每個線程獨立處理一行數(shù)據(jù)。樸素實現(xiàn)的偽代碼邏輯如下__global__ void softmax_naive(const float* input, float* output, int M, int N) { int row hipBlockIdx_x * hipBlockDim_x hipThreadIdx_x; if (row M) return; // 1. 找最大值 float max_val -FLT_MAX; for (int i 0; i N; i) { max_val fmaxf(max_val, input[row * N i]); } // 2. 計算指數(shù)和 float sum 0.0f; for (int i 0; i N; i) { sum expf(input[row * N i] - max_val); } // 3. 歸一化寫回 for (int i 0; i N; i) { output[row * N i] expf(input[row * N i] - max_val) / sum; } }這個實現(xiàn)完全正確但性能慘不忍睹。原因有三第一三次遍歷輸入數(shù)據(jù)。每次都要從全局顯存讀一遍數(shù)據(jù)如果N比較大訪存流量翻了三倍。第二沒有做任何向量化訪存每次只能讀一個float。DCU的全局訪存帶寬是按128字節(jié)為單位的cacheline對齊的標量訪問浪費了大量帶寬。第三大量冗余計算指數(shù)函數(shù)被調(diào)用了兩次。在我測試的shape下這個版本的耗時是2.34ms。注意這個數(shù)字本身就包含了并行執(zhí)行的收益因為M4×32×51265536個線程塊請求已經(jīng)足夠多吞吐量主要被訪存次數(shù)和指令效率限制。3.2 為什么樸素實現(xiàn)會這么慢拿2.34ms來算一下有效帶寬。輸入數(shù)據(jù)量是4×32×512×512×2字節(jié)FP16 64MB加上輸出64MB總共128MB。2.34ms對應的帶寬大約是55GB/s?？碊CU的規(guī)格理論顯存帶寬通常在幾百GB/s到1TB/s以上。55GB/s連理論值的零頭都不到。這中間的差距去哪了訪存模式是罪魁禍首。樸素實現(xiàn)里每個線程訪問的行是連續(xù)的128KB數(shù)據(jù)512×2字節(jié)但線程間訪問的行卻是完全離散的。同一時刻一個wavefront里的64個線程正在訪問64個不同行的首地址這64個地址分布在整個顯存地址空間中每次訪存都觸發(fā)一次完整的cacheline加載無法合并。實際上有效的帶寬利用率非常低大量總線周期花在了等待不連續(xù)地址的數(shù)據(jù)返回上。訪存合并Memory Coalescing是整個優(yōu)化的基石。正確的做法是讓同一wavefront內(nèi)的線程訪問連續(xù)地址也就是把矩陣按列方向拆分給不同線程。后面所有優(yōu)化方案都建立在這個認知之上。3.3 性能分析的基線工具使用在動手改代碼之前先用工具記錄一份基線數(shù)據(jù)方便后續(xù)對照。DCU的生態(tài)里能用到的分析工具主要是dcu_prof和hip_analyze前者做硬件計數(shù)器采樣的時間線分析后者做靜態(tài)代碼檢查。常用的操作是dcu_prof -t softmax_bench ./softmax_bench從prof輸出里可以讀到內(nèi)核執(zhí)行時間、占用率、全局訪存吞吐、L2命中率等關(guān)鍵指標。第一次跑樸素版本時L2命中率只有21%這個數(shù)字直接說明訪存模式有嚴重問題。實操心得每次改動內(nèi)核后先記錄L2命中率和全局訪存吞吐兩個指標如果它們沒有明顯變化說明優(yōu)化方向不對不必糾結(jié)延遲數(shù)字。這兩個指標能幫你快速判斷瓶頸在訪存還是計算。4. 核心優(yōu)化策略從訪存模式到并行劃分的全面調(diào)整4.1 方案一行級并行加向量化訪存樸素實現(xiàn)的問題在于線程塊內(nèi)線程處理不同行導致訪存不合并。第一版優(yōu)化從訪存模式下手——讓一個wavefront協(xié)同處理一行數(shù)據(jù)同時每個線程連續(xù)讀取多個相鄰元素形成向量化訪存。具體做法是每行數(shù)據(jù)由64個線程協(xié)作處理每個線程負責連續(xù)N/64個元素。這樣同一wavefront的線程在某一步訪問的地址是連續(xù)的硬件可以把這些訪問合并成較少的幾次cacheline傳輸。用偽代碼表示__global__ void softmax_v1(const half* input, half* output, int M, int N) { int row hipBlockIdx_x; int tid hipThreadIdx_x; // 0..63 int stride N / 64; const half* row_ptr input row * N; half* out_ptr output row * N; float local_max -FLT_MAX; #pragma unroll 4 for (int i 0; i stride; i) { int idx tid * stride i; float val __half2float(row_ptr[idx]); local_max fmaxf(local_max, val); } // wavefront內(nèi)歸約求全局最大值 for (int offset 32; offset 0; offset 1) { local_max fmaxf(local_max, __shfl_down(local_max, offset)); } // ... 指數(shù)求和、歸一化類似處理 }這里__shfl_down是wavefront內(nèi)部線程間數(shù)據(jù)交換的關(guān)鍵指令它允許線程直接讀取同一wavefront中其他線程的寄存器值不需要經(jīng)過共享內(nèi)存或全局內(nèi)存延遲極低。在DCU上它的實現(xiàn)效率和CUDA的shuffle指令相當是歸約類操作的利器。這一版優(yōu)化后耗時從2.34ms降到1.48ms。提升明顯但還沒達到目標因為每個線程訪問的元素間隔是stride而不是1向量化程度不夠。如果數(shù)據(jù)寬度允許應該用float2或float4類型做顯式向量化訪問。4.2 方案二向量化訪存與循環(huán)展開DCU的編譯器對顯式向量類型的支持非常關(guān)鍵。把數(shù)據(jù)當作float416字節(jié)一次性讀取可以顯著減少訪存指令數(shù)量同時提高cacheline利用率。修改核心循環(huán)const float4* input_v4 reinterpret_castconst float4*(row_ptr); float4 vals[4]; float local_max -FLT_MAX; #pragma unroll 8 for (int i 0; i stride / 4; i) { vals[i % 4] input_v4[tid * (stride / 4) i]; local_max fmaxf(local_max, fmaxf(fmaxf(vals[i % 4].x, vals[i % 4].y), fmaxf(vals[i % 4].z, vals[i % 4].w))); }這里有個前置條件N必須能被4整除N/64也必須能被4整除否則需要處理邊界。好在seq_len512這個場景完全滿足條件。循環(huán)展開的作用是讓編譯器生成更多獨立的訪存指令這些指令可以在等待內(nèi)存返回時并行執(zhí)行提高內(nèi)存級并行MLPMemory Level Parallelism。從實際效果看展開因子8比展開因子4性能更好但繼續(xù)加大收益就不明顯了猜測是寄存器壓力過大導致spill到local memory。這版優(yōu)化跑到了0.98ms終于突破1ms大關(guān)但距離最優(yōu)還有空間。這時瓶頸開始從訪存模式轉(zhuǎn)向計算效率和歸約開銷。4.3 方案三兩遍遍歷合并成一遍在基礎(chǔ)實現(xiàn)中需要三次遍歷數(shù)據(jù)找最大值、算指數(shù)和、算結(jié)果。但實際上第一遍找最大值和第二遍算指數(shù)和是可以合并的現(xiàn)代GPU上常見做法是分塊處理先對每個塊做局部統(tǒng)計再跨塊歸約。這里介紹一種常用技巧online softmax。它允許你在不知道全局最大值的情況下邊讀數(shù)據(jù)邊更新統(tǒng)計量。對于流式數(shù)據(jù)或無法多次訪存的場景非常有用。公式如下維護當前最大值m和累加和sum。每讀到一個新元素x計算[ m \max(m, x) ] [ sum sum \cdot e^{m - m} e^{x - m} ]這樣一來一趟遍歷就能同時完成找最大值和求和。雖然不是所有場景都需要這招多數(shù)情況下兩遍遍歷就夠但理解了它對設(shè)計更復雜的高性能版本會有幫助。實際優(yōu)化中我用的是兩遍遍歷合并到同一內(nèi)核的策略做法是每個線程讀取自己負責的數(shù)據(jù)塊先求出局部最大值然后立即在同一段數(shù)據(jù)上計算部分指數(shù)和。因為局部最大值可能不是全局最大值最后歸一化時需要調(diào)整但這個調(diào)整可以在最終歸約階段一次完成。4.4 方案四LDG緩存策略與__ldg替代DCU的全局內(nèi)存讀取存在L2緩存L2命中與否對性能影響極大。對于Softmax這種數(shù)據(jù)會被多次讀取的操作應該盡量讓數(shù)據(jù)留在L2里。在HIP中可以用__ldg內(nèi)置函數(shù)標記只讀數(shù)據(jù)訪問路徑編譯器會生成非臨時加載指令優(yōu)先命中緩存。實測在部分shape下使用__ldg能額外帶來10%~15%的性能提升。float val __ldg(input[row * N idx]);注意__ldg只對指針指向的只讀數(shù)據(jù)有效。如果你之后對同一內(nèi)存地址做了寫操作編譯器可能優(yōu)化掉__ldg或者更糟緩存一致性反而拖慢速度。這個場景下input和output指針完全分離放心用。4.5 方案五FP16數(shù)據(jù)類型的精度處理我的案例輸入是FP16而計算過程中用FP32做中間累加是必須的。FP16的有效精度只有10位尾數(shù)直接用它累加幾百個指數(shù)值誤差會被放大到不可接受。但這里有個隱蔽的問題在讀取FP16數(shù)據(jù)并轉(zhuǎn)成FP32時轉(zhuǎn)換指令本身有開銷。DCU的__half2float是一條獨立指令大量的轉(zhuǎn)換會拖慢內(nèi)核。優(yōu)化技巧是使用half2向量類型。一個half2寄存器可以同時裝載兩個FP16數(shù)一條指令轉(zhuǎn)換成兩個FP32。如果輸入對齊到4字節(jié)甚至可以用half4。half2 val2 *reinterpret_casthalf2*(row_ptr idx); float2 valf2 __half22float2(val2);這樣訪存指令數(shù)和轉(zhuǎn)換指令數(shù)同時減半整體指令吞吐量大幅提升。這個優(yōu)化比較tricky的地方在于對齊。DCU要求half2指針必須4字節(jié)對齊而輸入張量是64字節(jié)對齊的只要偏移量是2的倍數(shù)就沒問題。在實際代碼里我加了靜態(tài)斷言確保編譯期就能發(fā)現(xiàn)對齊問題。4.6 方案六大矩陣場景的分塊調(diào)度設(shè)計上面的優(yōu)化針對的是seq_len512、行數(shù)很多的情況。但Softmax還有一個典型的性能陷阱當單行數(shù)據(jù)量特別大比如seq_len4096或8192或者矩陣特別小而行的數(shù)量特別多時策略需要完全不同。針對行數(shù)據(jù)很大的場景不能再用“一個wavefront處理一行”的方案因為每行數(shù)據(jù)太大wavefront的寄存器裝不下全部數(shù)據(jù)也沒法一次性歸約。正確的做法是把每行拆分成多個列塊先分別計算局部統(tǒng)計量再用第二個內(nèi)核或第二次pass做全局歸約。這種兩階段處理有個工程細節(jié)要注意第一次pass結(jié)果需要暫存在中間緩沖區(qū)而中間緩沖區(qū)需要按塊對齊分配避免bank conflict。我遇到過的最典型問題就是中間緩沖區(qū)大小沒有按LDS的bank數(shù)對齊導致同時訪問時發(fā)生大量沖突性能比樸素實現(xiàn)還差。針對行數(shù)特別多的場景更應該關(guān)注任務調(diào)度粒度。索引從一維線性化把整個矩陣看成一個大一維數(shù)組按固定大小的tile切分給不同線程塊。然后每個線程塊從tile里提取它需要處理的行的片段。這樣做的好處是線程塊之間的負載更均衡不容易出現(xiàn)某些CE空閑的情況。5. 實戰(zhàn)記錄三版迭代完整性能對比與關(guān)鍵參數(shù)選擇5.1 線程塊尺寸與wavefront對齊的策略選擇線程塊尺寸的設(shè)計直接決定了并行度上限。DCU的一個CE最多可以駐留一定數(shù)量的wavefront超過之后多余線程只能排隊。對于Softmax這種訪存密集型算子并行度要盡量高但也不能無腦加大。在我最終方案里線程塊大小設(shè)為256即4個wavefront256/64。這樣設(shè)置的原因有兩點第一256個線程剛好可以完整覆蓋一行512個FP16數(shù)據(jù)每個線程處理2個half剛好組成一個half2向量不需要額外處理邊界第二256個線程的寄存器占用不會超過CE的資源限制允許足夠多的線程塊同時駐留。關(guān)于每個線程處理多少數(shù)據(jù)有個經(jīng)驗公式理想情況下每個線程處理4~8個FP32元素或者8~16個FP16元素既不會讓訪存指令太少導致延遲掩蓋不足也不會讓寄存器溢出。我的案例中每行512個FP16每個線程處理8個元素即4個half2分配算式是512/(64×4)2個half2每線程。這個組合實測效果最佳。如果行數(shù)是奇數(shù)乘數(shù)例如N768處理方式就會復雜一些需要最后一個wavefront部分線程處理額外元素或者調(diào)整每個線程的數(shù)據(jù)量。大多數(shù)情況下選擇每個線程處理相同數(shù)量的元素然后對越界部分做mask處理性能損失可以控制在幾個百分點內(nèi)。5.2 共享內(nèi)存與寄存器使用的平衡在歸約階段最初版本我用共享內(nèi)存做跨線程數(shù)據(jù)交換。每個線程把自己的局部最大值寫到共享內(nèi)存然后做分階段同步歸約。但很快發(fā)現(xiàn)共享內(nèi)存歸約有明顯的性能代價每次歸約都需要__syncthreads()這個同步指令會讓整個wavefront停下來等待最慢的線程頻繁使用會浪費大量周期。改用warp shuffle歸約后效果立竿見影。shuffle指令直接在寄存器之間交換數(shù)據(jù)不走共享內(nèi)存也不需要同步。因為一個wavefront的64個線程在DCU上是同時調(diào)度的天然保證了一致性shuffle的效率遠高于共享內(nèi)存。在最終代碼里我徹底移除了共享內(nèi)存的使用全部歸約都用shuffle完成。這不僅減少了同步開銷還降低了寄存器壓力因為共享內(nèi)存和寄存器在某些架構(gòu)上是共享資源的。5.3 三版方案的性能測量與放大對比整個調(diào)優(yōu)過程我保留了三個關(guān)鍵版本的測量數(shù)據(jù)整理成表格方便對照。版本核心改進點耗時相對樸素加速比有效帶寬v0樸素線程處理整行三次遍歷2.34ms1x~55GB/sv1向量化wavefront協(xié)作float2訪問1.48ms1.58x~87GB/sv2深度優(yōu)化half2向量shuffle歸約__ldg0.62ms3.77x~207GB/s有效帶寬的計算方法是總數(shù)據(jù)量輸入輸出FP16各64MB除以耗時。v2版本的207GB/s依然沒到理論峰值但考慮到FP16轉(zhuǎn)FP32、指數(shù)計算、shuffle交換這些額外開銷已經(jīng)到了非常合理的水平。這個對比也暴露了一個現(xiàn)象v1到v2的進步不是靠單項優(yōu)化而是simultaneously調(diào)整了訪存寬度、歸約機制、緩存策略和指令混合每個方向一點點累積最后疊加出大提升。單靠某一招想要3.77倍加速基本不可能。5.4 精度驗證與數(shù)據(jù)比對性能再漂亮精度不對就是廢代碼。Softmax的精度驗證我用了三步第一步最大絕對誤差檢查。在同一輸入下用DCU內(nèi)核輸出與PyTorch CPU的float64參考實現(xiàn)對比計算每個元素的最大絕對誤差。我的實現(xiàn)最終最大誤差在2e-4以內(nèi)對FP16輸出來說完全可接受。第二步分布一致性檢查。Softmax的輸出本質(zhì)是一個概率分布檢查每行輸出是否滿足求和接近1。實測最大偏差小于1e-3符合預期。第三步端到端模型驗證。把優(yōu)化后的Softmax接回原推理模型跑一批真實數(shù)據(jù)觀察最終輸出的top-1準確率和原始實現(xiàn)是否一致。這一步最重要因為算子層面的微小誤差在某些場景下會累積放大。特別提醒如果你優(yōu)化的是訓練過程中的Softmax一定要檢查反向傳播的表現(xiàn)因為梯度計算對精度更敏感。我這次只做推理優(yōu)化所以反傳不在范圍內(nèi)如果你需要支持訓練建議在反向Softmax的優(yōu)化上單獨投入時間很多推理場景的優(yōu)化技巧不能直接套用。6. 常見問題與性能排查技巧實錄6.1 為什么我的shuffle歸約結(jié)果不正確在DCU上使用__shfl_down時最容易踩的坑是忘記處理線程數(shù)不等于wavefront大小的情況。如果線程塊大小是128而你只用前64個線程做歸約并shuffle到其他線程就會漏掉數(shù)據(jù)。我調(diào)試時發(fā)現(xiàn)一個更隱蔽的問題__shfl_down的offset如果小于32行為正常但如果offset在32到63之間某些DCU型號的驅(qū)動行為不一致結(jié)果丟數(shù)據(jù)。排查了很久最后改成兩次循環(huán)先做offset32的歸約再做offset16、8、4、2、1的歸約問題消失。還有個常見錯誤是忘記把參與shuffle的變量賦初值。如果某線程的局部最大值初始化為0而不是-FLT_MAX恰好這一行所有輸入都是負數(shù)最終歸約結(jié)果就會錯誤地變成0。這種bug不會崩潰只會讓精度測試不過非常隱蔽。6.2 L2命中率上不去問題出在哪用dcu_prof看到L2命中率低第一反應自然是緩存復用不夠。但對于Softmax這種流式訪問主導的算子L2命中率天生不會太高因為每行數(shù)據(jù)基本只會被讀一次如果用兩遍遍歷會被讀兩次但第二次可能已被擠出L2。如果L2命中率特別低另一個排查方向是線程塊調(diào)度順序。DCU的線程塊調(diào)度順序可以控制L2的空間局部性。嘗試修改hipLaunchKernelGGL的grid維度排列方式把相鄰線程塊映射到共享L2區(qū)域的線程塊能提升命中率。我實測過把grid從[M]改成[M/8, 8]二維grid其中第二維是L2切片索引結(jié)果L2命中率從21%提升到了34%。雖然數(shù)字看著不大但整體耗時降低了約8%。6.3 編譯器自動向量化失敗如何檢查DCU的HIP編譯器對自動向量化的能力有限尤其在循環(huán)內(nèi)有fmaxf這類數(shù)學函數(shù)時常常不會自動生成向量訪存指令。檢查方法很簡單在編譯命令里加-S選項生成匯編文件然后搜索v_load或global_load_dwordx這類指令。如果發(fā)現(xiàn)循環(huán)體里都是global_load_dword4字節(jié)標量加載說明編譯器沒有自動向量化成功。解決辦法是像前面那樣顯式使用float2或half2類型把向量化寫死在代碼邏輯里。編譯器對顯式向量類型的支持很好反匯編能看到global_load_dwordx28字節(jié)或global_load_dwordx416字節(jié)指令。6.4 性能抖動為什么內(nèi)核耗時忽高忽低內(nèi)核耗時不穩(wěn)定可能是其他任務搶占顯存帶寬也可能是時鐘頻率波動。排查方法是連續(xù)運行多次benchmark觀察分布。更有意思的一個原因是DCU在同時跑多個上下文時L2緩存會被共享如果同一個GPU上有其他kernel在跑Softmax的L2命中率會驟降。這是硬件層面的資源競爭代碼層面沒法完全規(guī)避但可以在任務調(diào)度時避免同時運行多個大訪存kernel或使用單獨的GPU實例。6.5 邊界條件處理的最優(yōu)解當N不能被線程數(shù)整除時不要直接開根號強行整除那會讓代碼邏輯臃腫且難以維護。更優(yōu)雅的方案是讓每個線程處理固定數(shù)量的元素最后留一個線程處理尾部數(shù)據(jù)。或者使用grid-stride loop讓每個線程循環(huán)處理多個元素循環(huán)條件里判斷邊界。這個模式對不規(guī)則shape的適配性最好性能損失也很小。6.6 從CUDA代碼遷移到DCU的隱藏問題如果是從CUDA代碼直接遷移hipify工具能完成90%的替換工作但剩下10%會導致性能劇烈下降。最常見的問題第一__syncthreads()的語義在DCU上不如CUDA嚴格某些編譯器優(yōu)化可能導致同步被移除。如果代碼里有復雜的共享內(nèi)存讀寫依賴最好檢查一下反匯編里是否真的存在barrier指令。第二cudaMalloc換成hipMalloc后分配的大內(nèi)存默認屬性可能與CUDA不同訪問延遲更高。建議嘗試hipMallocManaged或調(diào)整對齊屬性。第三__restrict__關(guān)鍵字在DCU編譯器上的優(yōu)化力度不如NVIDIA的nvcc。如果遷移后性能達不到預期嘗試手動把指針加載到局部變量消除每次訪問的指針解引用開銷。7. 關(guān)于DCU算子優(yōu)化生態(tài)的一些補充經(jīng)驗除了Softmax本身這次優(yōu)化過程中接觸到的DCU工具鏈和生態(tài)也值得說幾句。DTK的hipcc編譯器總體質(zhì)量不錯但對某些優(yōu)化模式的支持還不成熟。比如我嘗試過用內(nèi)聯(lián)PTXDCU上對應的是內(nèi)聯(lián)GCN匯編手寫FMA和指數(shù)指令的組合確實能壓掉幾條指令但代碼可維護性急劇下降。除非為了追求極限性能否則不建議在工程代碼里大量使用。DCU的性能分析工具這幾年進步明顯dcu_prof的硬件計數(shù)器覆蓋已經(jīng)比較全。但相比成熟工具還是少了些便利性比如不能直接在時間線上查某條指令的詳細信息。解決方法是自己在代碼里插樁用clock64()記錄關(guān)鍵階段耗時。這個方法土但有效。社區(qū)方面雖然DCU生態(tài)還比不上CUDA但近兩年的文檔和示例代碼質(zhì)量提升很大。遇到問題時先查/opt/dtk目錄下的示例代碼通常能找到對應的內(nèi)核模板再結(jié)合官方性能優(yōu)化指南里的架構(gòu)特性說明大多數(shù)問題都能定位。對于團隊來說如果要做算子遷移建議建立一份自己的性能基線庫把每個算子在不同shape下的樸素實現(xiàn)耗時、優(yōu)化后耗時、理論帶寬和實測帶寬都記錄下來。有了這個基線庫后續(xù)新算子的優(yōu)化進度評估會快很多也能避免把已經(jīng)解決過的坑再踩一遍。8. 最終版本核心代碼參考實現(xiàn)把最終可運行的優(yōu)化內(nèi)核主體貼出來方便參考。這個版本融合了前面所有可行優(yōu)化代碼故意保留了shuffle歸約和half2顯式向量化方便對照理解。#include hip/hip_runtime.h #include hip/hip_fp16.h #include cmath #include cfloat // 假設(shè) M 能被 gridDim.x * blockDim.x 整除此處 M 65536 // 每行 N 512 個 half由 blockDim.x 256 個線程處理 // 每個線程處理 1 個 half2即 2 個 half一個 block 覆蓋 2 行 __global__ void softmax_opt_kernel(const half* __restrict__ input, half* __restrict__ output, int M, int N) { const int tid hipThreadIdx_x; // 0..255 const int wid tid 6; // 所在 wavefront 編號 0..3 const int lane tid 63; // wavefront 內(nèi)線程編號 0..63 // 每個 block 處理 4 行間隔 blockDim.y? 此處簡化為線性處理 int row hipBlockIdx_x * 4 wid; // 一個 wavefront 處理一行 if (row M) return; const int half_per_thread N / 256; // 2 const half2* row_ptr reinterpret_castconst half2*(input row * N); half2* out_ptr reinterpret_casthalf2*(output row * N); float local_max -FLT_MAX; // 每線程負責一個 half2 half2 val __ldg(row_ptr[lane]); float2 valf __half22float2(val); local_max fmaxf(fmaxf(valf.x, valf.y), local_max); // wavefront 內(nèi)歸約最大值 (64 - 1) #pragma unroll for (int offset 32; offset 0; offset 1) { local_max fmaxf(local_max, __shfl_down(local_max, offset)); } float row_max __shfl(local_max, 0); // 廣播到整個 wavefront // 計算指數(shù)與和 float e0 __expf(valf.x - row_max); float e1 __expf(valf.y - row_max); float local_sum e0 e1; #pragma unroll for (int offset 32; offset 0; offset 1) { local_sum __shfl_down(local_sum, offset); } float row_sum __shfl(local_sum, 0); // 歸一化寫回 half2 res; res.x __float2half(e0 / row_sum); res.y __float2half(e1 / row_sum); out_ptr[lane] res; }這段代碼有幾個使用前提行數(shù)必須是block數(shù)×4的整數(shù)倍這個案例滿足N必須等于256的倍數(shù)512滿足。如果shape不滿足需要加上邊界判斷邏輯會復雜一些。關(guān)于__expf和expf的選擇我在優(yōu)化中特意用了快速版本。__expf的精度比expf低一些但速度更快。對Softmax的最終輸出影響在1e-5量級完全可接受。在訓練等需要高精度的場景建議換回標準expf。另一個細節(jié)是__shfl和__shfl_down的配合使用。先通過__shfl_down把最大值歸約到lane 0再用__shfl廣播給全wavefront這個模式比每個線程都保存一份最終值更省指令。9. 后續(xù)擴展方向與實踐建議Softmax的優(yōu)化經(jīng)驗可以直接延伸到其他訪存密集型算子。像LayerNorm、RMSNorm、通道均值方差計算它們的結(jié)構(gòu)都是“讀數(shù)據(jù)-歸約-歸一化-寫回”和Softmax幾乎同構(gòu)。把這次用的shuffle歸約、half2向量化、L2緩存友好的線程塊調(diào)度這幾個技巧遷移過去通常能快速獲得類似幅度的提升。在更復雜的注意力機制中Softmax經(jīng)常和QK^T矩陣乘法、矩陣掩碼操作融合。如果你的融合目標是減少kernel launch次數(shù)可以在Softmax內(nèi)核里加一個參數(shù)判斷是否需要應用attention mask。在DCU上kernel launch的固定開銷比NVIDIA高減少launch次數(shù)對端到端性能的影響更明顯。這也是我后續(xù)在做的工作方向。從業(yè)余時間接觸DCU到現(xiàn)在我最大的體會是硬件不同但底層邏輯共通。訪存合并、歸約效率、指令混合、緩存利用這些在任何GPU架構(gòu)上都是核心命題。只不過在DCU上你需要更主動地管理這些優(yōu)化因為工具鏈和編譯器還沒成熟到自動幫你搞定一切。對于剛開始接觸DCU優(yōu)化的朋友建議從小算子入手別一上來就挑戰(zhàn)大模型端到端調(diào)優(yōu)。Softmax其實是個很好的起點數(shù)學簡單但訪存、歸約、向量化、緩存全涉及一遍。把這個流程走通再看其他算子會輕松很多。這次優(yōu)化到0.62ms之后我并沒有繼續(xù)往下壓。再往下走就要開始碰匯編手寫指數(shù)運算和指令級調(diào)度了收益可能還有20%左右但代碼可讀性和可維護性會急劇惡化。在工程實踐中一個能維護、能debug、能遷移的內(nèi)核比一個快10%但誰也看不懂的內(nèi)核有價值得多。如果你走的是產(chǎn)品化路線記住這個判斷標準。 DCU上Softmax算子深度優(yōu)化：從2.34ms到0.62ms實戰(zhàn)
1. Softmax算子為什么值得單獨寫一篇優(yōu)化文章做AI算子開發(fā)的朋友應該都有體會Softmax這個算子看著人畜無害實際上特別“刁鉆”。它的數(shù)學形式極其簡單但優(yōu)化空間和踩坑概率在常用算子中絕對排得上號。尤其到了DCU這種國產(chǎn)加速卡上想把Softmax跑出接近理論帶寬的性能需要操心的事情遠比想象中多。先說清楚Softmax在干什么。給定一個輸入向量Softmax把每個元素映射成一個概率值公式長這樣[ \text{Softmax}(x_i) \frac{e^{x_i - \max(x)}}{\sum_{j} e^{x_j - \max(x)}} ]減去最大值那一步不是可有可無是為了數(shù)值穩(wěn)定性。如果不減當輸入里有大數(shù)時(e^{x_i})直接溢出成inf后面全完蛋。這在Transformer的注意力層里尤其致命因為QK^T的點積結(jié)果動輒幾十上百不穩(wěn)定的Softmax會讓訓練直接發(fā)散。但Softmax真正的麻煩不在數(shù)學而在訪存特性。這個算子對每個元素只做幾次浮點運算計算密度極低屬于典型的訪存密集型任務。也就是說性能瓶頸幾乎完全取決于你能多快把數(shù)據(jù)從顯存搬進寄存器、把結(jié)果寫回去。在DCU上做優(yōu)化本質(zhì)上是在跟內(nèi)存帶寬較勁。這篇文章我從一個實際優(yōu)化案例出發(fā)完整走一遍從樸素實現(xiàn)到深度調(diào)優(yōu)的過程。案例背景是某推理模型中的注意力模塊Softmax算子的輸入shape為[batch, heads, seq_len, seq_len]其中batch4heads32seq_len512數(shù)據(jù)類型為FP16。優(yōu)化目標是把單次調(diào)用延遲從2.34ms壓到1ms以內(nèi)。這個場景非常有代表性。它既有行主序二維矩陣的常規(guī)Softmax特征又有大批量小矩陣的并行特性還涉及FP16的精度處理幾乎把DCU上算子優(yōu)化能遇到的典型問題都覆蓋了。適合讀這篇文章的人有三類一是正在做算子移植需要把PyTorch或CUDA代碼遷移到DCU上的工程師二是做推理優(yōu)化被Attention里Softmax耗時困擾的算法工程師三是對并行編程有興趣想看看國產(chǎn)加速卡生態(tài)實際長什么樣的技術(shù)愛好者。先說結(jié)論最終優(yōu)化后的Softmax算子單次調(diào)用耗時從2.34ms降到了0.62ms加速比3.77倍通過了一系列精度對比測試。這個結(jié)果不是靠某個單一技巧拿到的而是一整套策略組合的效果。下面我把每一步的思路、實現(xiàn)細節(jié)和踩坑經(jīng)歷都拆開講。2. DCU硬件結(jié)構(gòu)與并行編程模型的核心認知2.1 DCU的硬件架構(gòu)與計算單元布局要優(yōu)化DCU上的算子首先得搞清楚硬件長什么樣。DCUDeep Computing Unit是基于GPGPU架構(gòu)設(shè)計的加速卡核心計算單元叫計算引擎Compute EngineCE每個CE內(nèi)部包含多個SIMD單元每個SIMD單元又由若干個線程插槽Thread Slot組成。以我使用的DCU型號為例單卡有60個CE每個CE包含4個SIMD單元每個SIMD單元可以同時駐留一定數(shù)量的wavefrontDCU的調(diào)度單位等價于CUDA里的warp通常為64個線程。這意味著一個CE上同時活躍的硬件線程數(shù)非常可觀能夠用大量并行線程來掩蓋訪存延遲。DCU的存儲層次和主流GPGPU類似從快到慢依次是寄存器文件、L1緩存、共享內(nèi)存Local Data ShareLDS、L2緩存、全局顯存。每個SIMD單元有獨立的L1緩存和LDS所有CE共享L2緩存和全局顯存。有個關(guān)鍵區(qū)別需要特別強調(diào)DCU沒有像CPU那樣依賴大緩存來加速隨機訪問而是靠海量線程的并行切換來隱藏訪存延遲。這決定了優(yōu)化思路的核心——你要做的不是減少線程數(shù)而是盡可能多地創(chuàng)建可以并行執(zhí)行的線程讓它們在等待內(nèi)存數(shù)據(jù)時切換執(zhí)行其他計算。2.2 HIP編程模型與DTK工具鏈DCU的編程模型遵循HIPHeterogeneous Interface for Portability規(guī)范。HIP的語法和CUDA高度相似從CUDA代碼遷移到HIP通常只需要做少量替換比如__global__改__global__HIP里也是__global__blockIdx.x變hipBlockIdx_xthreadIdx.x變hipThreadIdx_x。但有些細節(jié)差異會在后面實操部分講清楚。DCU的軟件開發(fā)工具包是DTKDCU Toolkit里面包含HIP編譯器基于LLVM、HIPify工具用于自動把CUDA代碼轉(zhuǎn)換成HIP代碼、性能分析工具等。在使用過程中我用的是DTK 24.04版本編譯命令大致如下hipcc -O3 -stdc17 -DNDEBUG -o softmax_bench softmax_bench.cpp注意-O3在高性能算子編譯中基本是標配但如果你在代碼里用了restrict關(guān)鍵字或者__builtin_assume一定要檢查生成匯編是否真正優(yōu)化到位因為DCU編譯器在某些場景下對復雜指針別名的處理不如預期保守的代碼寫法反而會拖累性能。2.3 Wavefront、線程塊與調(diào)度機制DCU的調(diào)度機制和NVIDIA GPU最直觀的差異就是wavefront大小為64線程而CUDA的warp是32線程。這個差異影響深遠。在Softmax算子優(yōu)化中我們經(jīng)常需要做線程間的數(shù)據(jù)歸約比如求最大值、求和。在CUDA里你習慣用__shfl_down_sync在32個線程內(nèi)做shuffle歸約到了DCU/ROCm平臺對應的是__shfl_down同時因為wavefront有64個線程歸約的步數(shù)多了一級。另一個需要適應的是CE的調(diào)度粒度。DCU的硬件調(diào)度單元以wavefront為單位一個wavefront里的線程執(zhí)行相同的指令如果出現(xiàn)分支分歧會出現(xiàn)串行執(zhí)行不同路徑的情況性能損失明顯。所以在設(shè)計內(nèi)核時盡量保證同一wavefront內(nèi)的線程走相同分支或者干脆避免分支。實際測算下來同樣一份Softmax內(nèi)核我最初從CUDA直接搬過來時性能只有預期的60%左右。排查之后發(fā)現(xiàn)主要問題就在wavefront大小的差異上線程塊維度和歸約邏輯沒有針對64線程重新設(shè)計導致大量線程空轉(zhuǎn)。3. Softmax算子的基礎(chǔ)實現(xiàn)與首輪性能摸底3.1 樸素實現(xiàn)一個線程處理一行在動手優(yōu)化之前先把基礎(chǔ)版本寫出來作為后續(xù)優(yōu)化的基準線。最簡單直觀的思路是讓一個線程處理輸入矩陣的一行。假設(shè)輸入是[M, N]的二維矩陣那就開M個線程每個線程獨立處理一行數(shù)據(jù)。樸素實現(xiàn)的偽代碼邏輯如下__global__ void softmax_naive(const float* input, float* output, int M, int N) { int row hipBlockIdx_x * hipBlockDim_x hipThreadIdx_x; if (row M) return; // 1. 找最大值 float max_val -FLT_MAX; for (int i 0; i N; i) { max_val fmaxf(max_val, input[row * N i]); } // 2. 計算指數(shù)和 float sum 0.0f; for (int i 0; i N; i) { sum expf(input[row * N i] - max_val); } // 3. 歸一化寫回 for (int i 0; i N; i) { output[row * N i] expf(input[row * N i] - max_val) / sum; } }這個實現(xiàn)完全正確但性能慘不忍睹。原因有三第一三次遍歷輸入數(shù)據(jù)。每次都要從全局顯存讀一遍數(shù)據(jù)如果N比較大訪存流量翻了三倍。第二沒有做任何向量化訪存每次只能讀一個float。DCU的全局訪存帶寬是按128字節(jié)為單位的cacheline對齊的標量訪問浪費了大量帶寬。第三大量冗余計算指數(shù)函數(shù)被調(diào)用了兩次。在我測試的shape下這個版本的耗時是2.34ms。注意這個數(shù)字本身就包含了并行執(zhí)行的收益因為M4×32×51265536個線程塊請求已經(jīng)足夠多吞吐量主要被訪存次數(shù)和指令效率限制。3.2 為什么樸素實現(xiàn)會這么慢拿2.34ms來算一下有效帶寬。輸入數(shù)據(jù)量是4×32×512×512×2字節(jié)FP16 64MB加上輸出64MB總共128MB。2.34ms對應的帶寬大約是55GB/s?？碊CU的規(guī)格理論顯存帶寬通常在幾百GB/s到1TB/s以上。55GB/s連理論值的零頭都不到。這中間的差距去哪了訪存模式是罪魁禍首。樸素實現(xiàn)里每個線程訪問的行是連續(xù)的128KB數(shù)據(jù)512×2字節(jié)但線程間訪問的行卻是完全離散的。同一時刻一個wavefront里的64個線程正在訪問64個不同行的首地址這64個地址分布在整個顯存地址空間中每次訪存都觸發(fā)一次完整的cacheline加載無法合并。實際上有效的帶寬利用率非常低大量總線周期花在了等待不連續(xù)地址的數(shù)據(jù)返回上。訪存合并Memory Coalescing是整個優(yōu)化的基石。正確的做法是讓同一wavefront內(nèi)的線程訪問連續(xù)地址也就是把矩陣按列方向拆分給不同線程。后面所有優(yōu)化方案都建立在這個認知之上。3.3 性能分析的基線工具使用在動手改代碼之前先用工具記錄一份基線數(shù)據(jù)方便后續(xù)對照。DCU的生態(tài)里能用到的分析工具主要是dcu_prof和hip_analyze前者做硬件計數(shù)器采樣的時間線分析后者做靜態(tài)代碼檢查。常用的操作是dcu_prof -t softmax_bench ./softmax_bench從prof輸出里可以讀到內(nèi)核執(zhí)行時間、占用率、全局訪存吞吐、L2命中率等關(guān)鍵指標。第一次跑樸素版本時L2命中率只有21%這個數(shù)字直接說明訪存模式有嚴重問題。實操心得每次改動內(nèi)核后先記錄L2命中率和全局訪存吞吐兩個指標如果它們沒有明顯變化說明優(yōu)化方向不對不必糾結(jié)延遲數(shù)字。這兩個指標能幫你快速判斷瓶頸在訪存還是計算。4. 核心優(yōu)化策略從訪存模式到并行劃分的全面調(diào)整4.1 方案一行級并行加向量化訪存樸素實現(xiàn)的問題在于線程塊內(nèi)線程處理不同行導致訪存不合并。第一版優(yōu)化從訪存模式下手——讓一個wavefront協(xié)同處理一行數(shù)據(jù)同時每個線程連續(xù)讀取多個相鄰元素形成向量化訪存。具體做法是每行數(shù)據(jù)由64個線程協(xié)作處理每個線程負責連續(xù)N/64個元素。這樣同一wavefront的線程在某一步訪問的地址是連續(xù)的硬件可以把這些訪問合并成較少的幾次cacheline傳輸。用偽代碼表示__global__ void softmax_v1(const half* input, half* output, int M, int N) { int row hipBlockIdx_x; int tid hipThreadIdx_x; // 0..63 int stride N / 64; const half* row_ptr input row * N; half* out_ptr output row * N; float local_max -FLT_MAX; #pragma unroll 4 for (int i 0; i stride; i) { int idx tid * stride i; float val __half2float(row_ptr[idx]); local_max fmaxf(local_max, val); } // wavefront內(nèi)歸約求全局最大值 for (int offset 32; offset 0; offset 1) { local_max fmaxf(local_max, __shfl_down(local_max, offset)); } // ... 指數(shù)求和、歸一化類似處理 }這里__shfl_down是wavefront內(nèi)部線程間數(shù)據(jù)交換的關(guān)鍵指令它允許線程直接讀取同一wavefront中其他線程的寄存器值不需要經(jīng)過共享內(nèi)存或全局內(nèi)存延遲極低。在DCU上它的實現(xiàn)效率和CUDA的shuffle指令相當是歸約類操作的利器。這一版優(yōu)化后耗時從2.34ms降到1.48ms。提升明顯但還沒達到目標因為每個線程訪問的元素間隔是stride而不是1向量化程度不夠。如果數(shù)據(jù)寬度允許應該用float2或float4類型做顯式向量化訪問。4.2 方案二向量化訪存與循環(huán)展開DCU的編譯器對顯式向量類型的支持非常關(guān)鍵。把數(shù)據(jù)當作float416字節(jié)一次性讀取可以顯著減少訪存指令數(shù)量同時提高cacheline利用率。修改核心循環(huán)const float4* input_v4 reinterpret_castconst float4*(row_ptr); float4 vals[4]; float local_max -FLT_MAX; #pragma unroll 8 for (int i 0; i stride / 4; i) { vals[i % 4] input_v4[tid * (stride / 4) i]; local_max fmaxf(local_max, fmaxf(fmaxf(vals[i % 4].x, vals[i % 4].y), fmaxf(vals[i % 4].z, vals[i % 4].w))); }這里有個前置條件N必須能被4整除N/64也必須能被4整除否則需要處理邊界。好在seq_len512這個場景完全滿足條件。循環(huán)展開的作用是讓編譯器生成更多獨立的訪存指令這些指令可以在等待內(nèi)存返回時并行執(zhí)行提高內(nèi)存級并行MLPMemory Level Parallelism。從實際效果看展開因子8比展開因子4性能更好但繼續(xù)加大收益就不明顯了猜測是寄存器壓力過大導致spill到local memory。這版優(yōu)化跑到了0.98ms終于突破1ms大關(guān)但距離最優(yōu)還有空間。這時瓶頸開始從訪存模式轉(zhuǎn)向計算效率和歸約開銷。4.3 方案三兩遍遍歷合并成一遍在基礎(chǔ)實現(xiàn)中需要三次遍歷數(shù)據(jù)找最大值、算指數(shù)和、算結(jié)果。但實際上第一遍找最大值和第二遍算指數(shù)和是可以合并的現(xiàn)代GPU上常見做法是分塊處理先對每個塊做局部統(tǒng)計再跨塊歸約。這里介紹一種常用技巧online softmax。它允許你在不知道全局最大值的情況下邊讀數(shù)據(jù)邊更新統(tǒng)計量。對于流式數(shù)據(jù)或無法多次訪存的場景非常有用。公式如下維護當前最大值m和累加和sum。每讀到一個新元素x計算[ m \max(m, x) ] [ sum sum \cdot e^{m - m} e^{x - m} ]這樣一來一趟遍歷就能同時完成找最大值和求和。雖然不是所有場景都需要這招多數(shù)情況下兩遍遍歷就夠但理解了它對設(shè)計更復雜的高性能版本會有幫助。實際優(yōu)化中我用的是兩遍遍歷合并到同一內(nèi)核的策略做法是每個線程讀取自己負責的數(shù)據(jù)塊先求出局部最大值然后立即在同一段數(shù)據(jù)上計算部分指數(shù)和。因為局部最大值可能不是全局最大值最后歸一化時需要調(diào)整但這個調(diào)整可以在最終歸約階段一次完成。4.4 方案四LDG緩存策略與__ldg替代DCU的全局內(nèi)存讀取存在L2緩存L2命中與否對性能影響極大。對于Softmax這種數(shù)據(jù)會被多次讀取的操作應該盡量讓數(shù)據(jù)留在L2里。在HIP中可以用__ldg內(nèi)置函數(shù)標記只讀數(shù)據(jù)訪問路徑編譯器會生成非臨時加載指令優(yōu)先命中緩存。實測在部分shape下使用__ldg能額外帶來10%~15%的性能提升。float val __ldg(input[row * N idx]);注意__ldg只對指針指向的只讀數(shù)據(jù)有效。如果你之后對同一內(nèi)存地址做了寫操作編譯器可能優(yōu)化掉__ldg或者更糟緩存一致性反而拖慢速度。這個場景下input和output指針完全分離放心用。4.5 方案五FP16數(shù)據(jù)類型的精度處理我的案例輸入是FP16而計算過程中用FP32做中間累加是必須的。FP16的有效精度只有10位尾數(shù)直接用它累加幾百個指數(shù)值誤差會被放大到不可接受。但這里有個隱蔽的問題在讀取FP16數(shù)據(jù)并轉(zhuǎn)成FP32時轉(zhuǎn)換指令本身有開銷。DCU的__half2float是一條獨立指令大量的轉(zhuǎn)換會拖慢內(nèi)核。優(yōu)化技巧是使用half2向量類型。一個half2寄存器可以同時裝載兩個FP16數(shù)一條指令轉(zhuǎn)換成兩個FP32。如果輸入對齊到4字節(jié)甚至可以用half4。half2 val2 *reinterpret_casthalf2*(row_ptr idx); float2 valf2 __half22float2(val2);這樣訪存指令數(shù)和轉(zhuǎn)換指令數(shù)同時減半整體指令吞吐量大幅提升。這個優(yōu)化比較tricky的地方在于對齊。DCU要求half2指針必須4字節(jié)對齊而輸入張量是64字節(jié)對齊的只要偏移量是2的倍數(shù)就沒問題。在實際代碼里我加了靜態(tài)斷言確保編譯期就能發(fā)現(xiàn)對齊問題。4.6 方案六大矩陣場景的分塊調(diào)度設(shè)計上面的優(yōu)化針對的是seq_len512、行數(shù)很多的情況。但Softmax還有一個典型的性能陷阱當單行數(shù)據(jù)量特別大比如seq_len4096或8192或者矩陣特別小而行的數(shù)量特別多時策略需要完全不同。針對行數(shù)據(jù)很大的場景不能再用“一個wavefront處理一行”的方案因為每行數(shù)據(jù)太大wavefront的寄存器裝不下全部數(shù)據(jù)也沒法一次性歸約。正確的做法是把每行拆分成多個列塊先分別計算局部統(tǒng)計量再用第二個內(nèi)核或第二次pass做全局歸約。這種兩階段處理有個工程細節(jié)要注意第一次pass結(jié)果需要暫存在中間緩沖區(qū)而中間緩沖區(qū)需要按塊對齊分配避免bank conflict。我遇到過的最典型問題就是中間緩沖區(qū)大小沒有按LDS的bank數(shù)對齊導致同時訪問時發(fā)生大量沖突性能比樸素實現(xiàn)還差。針對行數(shù)特別多的場景更應該關(guān)注任務調(diào)度粒度。索引從一維線性化把整個矩陣看成一個大一維數(shù)組按固定大小的tile切分給不同線程塊。然后每個線程塊從tile里提取它需要處理的行的片段。這樣做的好處是線程塊之間的負載更均衡不容易出現(xiàn)某些CE空閑的情況。5. 實戰(zhàn)記錄三版迭代完整性能對比與關(guān)鍵參數(shù)選擇5.1 線程塊尺寸與wavefront對齊的策略選擇線程塊尺寸的設(shè)計直接決定了并行度上限。DCU的一個CE最多可以駐留一定數(shù)量的wavefront超過之后多余線程只能排隊。對于Softmax這種訪存密集型算子并行度要盡量高但也不能無腦加大。在我最終方案里線程塊大小設(shè)為256即4個wavefront256/64。這樣設(shè)置的原因有兩點第一256個線程剛好可以完整覆蓋一行512個FP16數(shù)據(jù)每個線程處理2個half剛好組成一個half2向量不需要額外處理邊界第二256個線程的寄存器占用不會超過CE的資源限制允許足夠多的線程塊同時駐留。關(guān)于每個線程處理多少數(shù)據(jù)有個經(jīng)驗公式理想情況下每個線程處理4~8個FP32元素或者8~16個FP16元素既不會讓訪存指令太少導致延遲掩蓋不足也不會讓寄存器溢出。我的案例中每行512個FP16每個線程處理8個元素即4個half2分配算式是512/(64×4)2個half2每線程。這個組合實測效果最佳。如果行數(shù)是奇數(shù)乘數(shù)例如N768處理方式就會復雜一些需要最后一個wavefront部分線程處理額外元素或者調(diào)整每個線程的數(shù)據(jù)量。大多數(shù)情況下選擇每個線程處理相同數(shù)量的元素然后對越界部分做mask處理性能損失可以控制在幾個百分點內(nèi)。5.2 共享內(nèi)存與寄存器使用的平衡在歸約階段最初版本我用共享內(nèi)存做跨線程數(shù)據(jù)交換。每個線程把自己的局部最大值寫到共享內(nèi)存然后做分階段同步歸約。但很快發(fā)現(xiàn)共享內(nèi)存歸約有明顯的性能代價每次歸約都需要__syncthreads()這個同步指令會讓整個wavefront停下來等待最慢的線程頻繁使用會浪費大量周期。改用warp shuffle歸約后效果立竿見影。shuffle指令直接在寄存器之間交換數(shù)據(jù)不走共享內(nèi)存也不需要同步。因為一個wavefront的64個線程在DCU上是同時調(diào)度的天然保證了一致性shuffle的效率遠高于共享內(nèi)存。在最終代碼里我徹底移除了共享內(nèi)存的使用全部歸約都用shuffle完成。這不僅減少了同步開銷還降低了寄存器壓力因為共享內(nèi)存和寄存器在某些架構(gòu)上是共享資源的。5.3 三版方案的性能測量與放大對比整個調(diào)優(yōu)過程我保留了三個關(guān)鍵版本的測量數(shù)據(jù)整理成表格方便對照。版本核心改進點耗時相對樸素加速比有效帶寬v0樸素線程處理整行三次遍歷2.34ms1x~55GB/sv1向量化wavefront協(xié)作float2訪問1.48ms1.58x~87GB/sv2深度優(yōu)化half2向量shuffle歸約__ldg0.62ms3.77x~207GB/s有效帶寬的計算方法是總數(shù)據(jù)量輸入輸出FP16各64MB除以耗時。v2版本的207GB/s依然沒到理論峰值但考慮到FP16轉(zhuǎn)FP32、指數(shù)計算、shuffle交換這些額外開銷已經(jīng)到了非常合理的水平。這個對比也暴露了一個現(xiàn)象v1到v2的進步不是靠單項優(yōu)化而是simultaneously調(diào)整了訪存寬度、歸約機制、緩存策略和指令混合每個方向一點點累積最后疊加出大提升。單靠某一招想要3.77倍加速基本不可能。5.4 精度驗證與數(shù)據(jù)比對性能再漂亮精度不對就是廢代碼。Softmax的精度驗證我用了三步第一步最大絕對誤差檢查。在同一輸入下用DCU內(nèi)核輸出與PyTorch CPU的float64參考實現(xiàn)對比計算每個元素的最大絕對誤差。我的實現(xiàn)最終最大誤差在2e-4以內(nèi)對FP16輸出來說完全可接受。第二步分布一致性檢查。Softmax的輸出本質(zhì)是一個概率分布檢查每行輸出是否滿足求和接近1。實測最大偏差小于1e-3符合預期。第三步端到端模型驗證。把優(yōu)化后的Softmax接回原推理模型跑一批真實數(shù)據(jù)觀察最終輸出的top-1準確率和原始實現(xiàn)是否一致。這一步最重要因為算子層面的微小誤差在某些場景下會累積放大。特別提醒如果你優(yōu)化的是訓練過程中的Softmax一定要檢查反向傳播的表現(xiàn)因為梯度計算對精度更敏感。我這次只做推理優(yōu)化所以反傳不在范圍內(nèi)如果你需要支持訓練建議在反向Softmax的優(yōu)化上單獨投入時間很多推理場景的優(yōu)化技巧不能直接套用。6. 常見問題與性能排查技巧實錄6.1 為什么我的shuffle歸約結(jié)果不正確在DCU上使用__shfl_down時最容易踩的坑是忘記處理線程數(shù)不等于wavefront大小的情況。如果線程塊大小是128而你只用前64個線程做歸約并shuffle到其他線程就會漏掉數(shù)據(jù)。我調(diào)試時發(fā)現(xiàn)一個更隱蔽的問題__shfl_down的offset如果小于32行為正常但如果offset在32到63之間某些DCU型號的驅(qū)動行為不一致結(jié)果丟數(shù)據(jù)。排查了很久最后改成兩次循環(huán)先做offset32的歸約再做offset16、8、4、2、1的歸約問題消失。還有個常見錯誤是忘記把參與shuffle的變量賦初值。如果某線程的局部最大值初始化為0而不是-FLT_MAX恰好這一行所有輸入都是負數(shù)最終歸約結(jié)果就會錯誤地變成0。這種bug不會崩潰只會讓精度測試不過非常隱蔽。6.2 L2命中率上不去問題出在哪用dcu_prof看到L2命中率低第一反應自然是緩存復用不夠。但對于Softmax這種流式訪問主導的算子L2命中率天生不會太高因為每行數(shù)據(jù)基本只會被讀一次如果用兩遍遍歷會被讀兩次但第二次可能已被擠出L2。如果L2命中率特別低另一個排查方向是線程塊調(diào)度順序。DCU的線程塊調(diào)度順序可以控制L2的空間局部性。嘗試修改hipLaunchKernelGGL的grid維度排列方式把相鄰線程塊映射到共享L2區(qū)域的線程塊能提升命中率。我實測過把grid從[M]改成[M/8, 8]二維grid其中第二維是L2切片索引結(jié)果L2命中率從21%提升到了34%。雖然數(shù)字看著不大但整體耗時降低了約8%。6.3 編譯器自動向量化失敗如何檢查DCU的HIP編譯器對自動向量化的能力有限尤其在循環(huán)內(nèi)有fmaxf這類數(shù)學函數(shù)時常常不會自動生成向量訪存指令。檢查方法很簡單在編譯命令里加-S選項生成匯編文件然后搜索v_load或global_load_dwordx這類指令。如果發(fā)現(xiàn)循環(huán)體里都是global_load_dword4字節(jié)標量加載說明編譯器沒有自動向量化成功。解決辦法是像前面那樣顯式使用float2或half2類型把向量化寫死在代碼邏輯里。編譯器對顯式向量類型的支持很好反匯編能看到global_load_dwordx28字節(jié)或global_load_dwordx416字節(jié)指令。6.4 性能抖動為什么內(nèi)核耗時忽高忽低內(nèi)核耗時不穩(wěn)定可能是其他任務搶占顯存帶寬也可能是時鐘頻率波動。排查方法是連續(xù)運行多次benchmark觀察分布。更有意思的一個原因是DCU在同時跑多個上下文時L2緩存會被共享如果同一個GPU上有其他kernel在跑Softmax的L2命中率會驟降。這是硬件層面的資源競爭代碼層面沒法完全規(guī)避但可以在任務調(diào)度時避免同時運行多個大訪存kernel或使用單獨的GPU實例。6.5 邊界條件處理的最優(yōu)解當N不能被線程數(shù)整除時不要直接開根號強行整除那會讓代碼邏輯臃腫且難以維護。更優(yōu)雅的方案是讓每個線程處理固定數(shù)量的元素最后留一個線程處理尾部數(shù)據(jù)。或者使用grid-stride loop讓每個線程循環(huán)處理多個元素循環(huán)條件里判斷邊界。這個模式對不規(guī)則shape的適配性最好性能損失也很小。6.6 從CUDA代碼遷移到DCU的隱藏問題如果是從CUDA代碼直接遷移hipify工具能完成90%的替換工作但剩下10%會導致性能劇烈下降。最常見的問題第一__syncthreads()的語義在DCU上不如CUDA嚴格某些編譯器優(yōu)化可能導致同步被移除。如果代碼里有復雜的共享內(nèi)存讀寫依賴最好檢查一下反匯編里是否真的存在barrier指令。第二cudaMalloc換成hipMalloc后分配的大內(nèi)存默認屬性可能與CUDA不同訪問延遲更高。建議嘗試hipMallocManaged或調(diào)整對齊屬性。第三__restrict__關(guān)鍵字在DCU編譯器上的優(yōu)化力度不如NVIDIA的nvcc。如果遷移后性能達不到預期嘗試手動把指針加載到局部變量消除每次訪問的指針解引用開銷。7. 關(guān)于DCU算子優(yōu)化生態(tài)的一些補充經(jīng)驗除了Softmax本身這次優(yōu)化過程中接觸到的DCU工具鏈和生態(tài)也值得說幾句。DTK的hipcc編譯器總體質(zhì)量不錯但對某些優(yōu)化模式的支持還不成熟。比如我嘗試過用內(nèi)聯(lián)PTXDCU上對應的是內(nèi)聯(lián)GCN匯編手寫FMA和指數(shù)指令的組合確實能壓掉幾條指令但代碼可維護性急劇下降。除非為了追求極限性能否則不建議在工程代碼里大量使用。DCU的性能分析工具這幾年進步明顯dcu_prof的硬件計數(shù)器覆蓋已經(jīng)比較全。但相比成熟工具還是少了些便利性比如不能直接在時間線上查某條指令的詳細信息。解決方法是自己在代碼里插樁用clock64()記錄關(guān)鍵階段耗時。這個方法土但有效。社區(qū)方面雖然DCU生態(tài)還比不上CUDA但近兩年的文檔和示例代碼質(zhì)量提升很大。遇到問題時先查/opt/dtk目錄下的示例代碼通常能找到對應的內(nèi)核模板再結(jié)合官方性能優(yōu)化指南里的架構(gòu)特性說明大多數(shù)問題都能定位。對于團隊來說如果要做算子遷移建議建立一份自己的性能基線庫把每個算子在不同shape下的樸素實現(xiàn)耗時、優(yōu)化后耗時、理論帶寬和實測帶寬都記錄下來。有了這個基線庫后續(xù)新算子的優(yōu)化進度評估會快很多也能避免把已經(jīng)解決過的坑再踩一遍。8. 最終版本核心代碼參考實現(xiàn)把最終可運行的優(yōu)化內(nèi)核主體貼出來方便參考。這個版本融合了前面所有可行優(yōu)化代碼故意保留了shuffle歸約和half2顯式向量化方便對照理解。#include hip/hip_runtime.h #include hip/hip_fp16.h #include cmath #include cfloat // 假設(shè) M 能被 gridDim.x * blockDim.x 整除此處 M 65536 // 每行 N 512 個 half由 blockDim.x 256 個線程處理 // 每個線程處理 1 個 half2即 2 個 half一個 block 覆蓋 2 行 __global__ void softmax_opt_kernel(const half* __restrict__ input, half* __restrict__ output, int M, int N) { const int tid hipThreadIdx_x; // 0..255 const int wid tid 6; // 所在 wavefront 編號 0..3 const int lane tid 63; // wavefront 內(nèi)線程編號 0..63 // 每個 block 處理 4 行間隔 blockDim.y? 此處簡化為線性處理 int row hipBlockIdx_x * 4 wid; // 一個 wavefront 處理一行 if (row M) return; const int half_per_thread N / 256; // 2 const half2* row_ptr reinterpret_castconst half2*(input row * N); half2* out_ptr reinterpret_casthalf2*(output row * N); float local_max -FLT_MAX; // 每線程負責一個 half2 half2 val __ldg(row_ptr[lane]); float2 valf __half22float2(val); local_max fmaxf(fmaxf(valf.x, valf.y), local_max); // wavefront 內(nèi)歸約最大值 (64 - 1) #pragma unroll for (int offset 32; offset 0; offset 1) { local_max fmaxf(local_max, __shfl_down(local_max, offset)); } float row_max __shfl(local_max, 0); // 廣播到整個 wavefront // 計算指數(shù)與和 float e0 __expf(valf.x - row_max); float e1 __expf(valf.y - row_max); float local_sum e0 e1; #pragma unroll for (int offset 32; offset 0; offset 1) { local_sum __shfl_down(local_sum, offset); } float row_sum __shfl(local_sum, 0); // 歸一化寫回 half2 res; res.x __float2half(e0 / row_sum); res.y __float2half(e1 / row_sum); out_ptr[lane] res; }這段代碼有幾個使用前提行數(shù)必須是block數(shù)×4的整數(shù)倍這個案例滿足N必須等于256的倍數(shù)512滿足。如果shape不滿足需要加上邊界判斷邏輯會復雜一些。關(guān)于__expf和expf的選擇我在優(yōu)化中特意用了快速版本。__expf的精度比expf低一些但速度更快。對Softmax的最終輸出影響在1e-5量級完全可接受。在訓練等需要高精度的場景建議換回標準expf。另一個細節(jié)是__shfl和__shfl_down的配合使用。先通過__shfl_down把最大值歸約到lane 0再用__shfl廣播給全wavefront這個模式比每個線程都保存一份最終值更省指令。9. 后續(xù)擴展方向與實踐建議Softmax的優(yōu)化經(jīng)驗可以直接延伸到其他訪存密集型算子。像LayerNorm、RMSNorm、通道均值方差計算它們的結(jié)構(gòu)都是“讀數(shù)據(jù)-歸約-歸一化-寫回”和Softmax幾乎同構(gòu)。把這次用的shuffle歸約、half2向量化、L2緩存友好的線程塊調(diào)度這幾個技巧遷移過去通常能快速獲得類似幅度的提升。在更復雜的注意力機制中Softmax經(jīng)常和QK^T矩陣乘法、矩陣掩碼操作融合。如果你的融合目標是減少kernel launch次數(shù)可以在Softmax內(nèi)核里加一個參數(shù)判斷是否需要應用attention mask。在DCU上kernel launch的固定開銷比NVIDIA高減少launch次數(shù)對端到端性能的影響更明顯。這也是我后續(xù)在做的工作方向。從業(yè)余時間接觸DCU到現(xiàn)在我最大的體會是硬件不同但底層邏輯共通。訪存合并、歸約效率、指令混合、緩存利用這些在任何GPU架構(gòu)上都是核心命題。只不過在DCU上你需要更主動地管理這些優(yōu)化因為工具鏈和編譯器還沒成熟到自動幫你搞定一切。對于剛開始接觸DCU優(yōu)化的朋友建議從小算子入手別一上來就挑戰(zhàn)大模型端到端調(diào)優(yōu)。Softmax其實是個很好的起點數(shù)學簡單但訪存、歸約、向量化、緩存全涉及一遍。把這個流程走通再看其他算子會輕松很多。這次優(yōu)化到0.62ms之后我并沒有繼續(xù)往下壓。再往下走就要開始碰匯編手寫指數(shù)運算和指令級調(diào)度了收益可能還有20%左右但代碼可讀性和可維護性會急劇惡化。在工程實踐中一個能維護、能debug、能遷移的內(nèi)核比一個快10%但誰也看不懂的內(nèi)核有價值得多。如果你走的是產(chǎn)品化路線記住這個判斷標準。