溫馨提示×
您好，登錄后才能下訂單哦！
點擊 登錄注冊 即表示同意《億速云用戶服務條款》
您好，登錄后才能下訂單哦！
這篇文章給大家介紹怎么實現一個高效的Softmax CUDA kernel，內容非常詳細，感興趣的小伙伴們可以參考借鑒，希望對大家能有所幫助。
Softmax操作是深度學習模型中最常用的操作之一。在深度學習的分類任務中，網絡最后的分類器往往是Softmax + CrossEntropy的組合：
盡管當Softmax和CrossEntropy聯合使用時，其數學推導可以約簡，但還是有很多場景會單獨使用Softmax Op。如BERT的Encoder每一層的attention layer中就單獨使用了Softmax求解attention的概率分布；GPT-2的attention的multi-head部分也單獨使用了Softmax 等等。
深度學習框架中的所有計算的算子都會轉化為GPU上的CUDA kernel function，Softmax操作也不例外。Softmax作為一個被大多數網絡廣泛使用的算子，其CUDA Kernel實現高效性會影響很多網絡最終的訓練速度。那么如何實現一個高效的Softmax CUDA Kernel？本文將會介紹OneFlow中優化的Softmax CUDA Kernel的技巧，并跟cuDNN中的Softmax操作進行實驗對比，結果表明，OneFlow深度優化后的Softmax對顯存帶寬的利用率可以接近理論上限，遠高于cuDNN的實現。
對GPU基礎知識的介紹以及CUDA性能優化的原則與目標可以參考之前的文章：
https://zhuanlan.zhihu.com/p/271740706
其中簡單介紹了GPU的硬件結構與執行原理：
Kernel：即CUDA Kernel function，是GPU的基本計算任務描述單元。每個Kernel都會根據配置參數在GPU上由非常多個線程并行執行，GPU計算高效就是因為同時可以由數千個core（thread）同時執行，計算效率遠超CPU。
GPU的線程在邏輯上分為Thread、Block、Grid三級，硬件上分為core、warp兩級；
GPU內存分為Global memory、Shared memory、Local memory三級。
GPU最主要提供的是兩種資源：計算資源 和 顯存帶寬資源。如果我們能將這兩種資源充分利用，且對資源的需求無法再降低，那么性能就優化到了極限，執行時間就會最短。在大多數情況下，深度學習訓練中的GPU計算資源是被充分利用的，而一個GPU CUDA Kernel的優化目標就是盡可能充分利用顯存帶寬資源。
對于顯存帶寬資源來說，“充分利用“指的是Kernel的有效顯存讀寫帶寬達到了設備顯存帶寬的上限，其中設備顯存帶寬可以通過執行 cuda中的的bandwidthTest得到。 Kernel的有效顯存帶寬通過Kernel讀寫數據量和Kernel執行時間進行評估：
當 前 K e r n e l 的 有 效 顯 存 帶 寬 = 讀 寫 數 據 量 / 執 行 時 間 當前Kernel的有效顯存帶寬 = 讀寫數據量 / 執行時間 當前Kernel的有效顯存帶寬=讀寫數據量/執行時間
在介紹優化技巧之前，我們先看看一個未經優化的Softmax Kernel的最高理論帶寬是多少。如下圖所示，一個最簡單的Softmax計算實現中，分別調用幾個基礎的CUDA Kernel function來完成整體的計算：
假設輸入的數據大小為D，shape = (num_rows, num_cols), 即 D = num_rows * num_cols ，最Navie的操作會多次訪問Global memory，其中：
ReduceMax = D + num_rows （read 為 D， write 為 num_rows）
BroadcaseSub = 2 * D + num_rows （read 為 D + num_rows，write 為 D）
Exp = 2 * D （read 、write 均為D）
ReduceSum = D + num_rows（read 為 D， write 為 num_rows）
BroadcastDive = 2 * D + num_rows （read 為 D + num_rows， write 為 D）
總共需要8 * D + 4 * num_rows 的訪存開銷。由于num_rows相比于D可以忽略，則Navie版本的Softmax CUDA Kernel需要訪問至少8倍數據的顯存，即：
N a i v e S o f t m a x K e r n e l 有 效 顯 存 帶 寬 < 理 論 帶 寬 / 8 Naive Softmax Kernel 有效顯存帶寬 < 理論帶寬 / 8 NaiveSoftmaxKernel有效顯存帶寬<理論帶寬/8
對于GeForce RTX? 3090顯卡，其理論帶寬上限為936GB/s，那么Navie版本的Softmax CUDA Kernel利用顯存帶寬的上界就是936 / 8 = 117 GB/s。
在文章 https://zhuanlan.zhihu.com/p/271740706 里，我們在方法：借助shared memory合并帶有Reduce計算的Kernel中提到了對Softmax Kernel的訪存優化到了 2 * D，但這仍然沒有優化到極致。在本文的優化后，大多數場景下OneFlow的Softmax CUDA Kernel的顯存帶寬利用可以逼近理論帶寬。
我們對OneFlow深度優化后的Softmax CUDA Kernel 與 cuDNN中的Softmax Kernel的訪存帶寬進行了對比測試，測試結果如下：
Softmax Kernel Bandwidth：
Softmax Grad Kernel Bandwidth：
其中測試環境是 GeForce RTX? 3090 GPU，數據類型是half，Softmax的Shape = (49152, num_cols)，其中49152 = 32 * 12 * 128，是BERT-base網絡中的attention Tensor的前三維，我們固定了前三維，將最后一維動態變化，測試了從32到32768不同大小的Softmax前向Kernel和反向Kernel的有效顯存帶寬。從上面兩張圖中可以看出OneFlow在多數情況下都可以逼近理論帶寬（800GB/s左右，與cudaMemcpy的訪存帶寬相當。官方公布的理論帶寬為936GB/S）。并且在所有情況下，OneFlow的CUDA Kernel的有效訪存帶寬都優于cuDNN的實現。
Softmax 函數的輸入形狀為:(num_rows, num_cols)，num_cols的變化，會對有效帶寬產生影響；因為，沒有一種 通用 的優化方法可以實現在所有 num_cols的情況下都是傳輸最優的。所以，在 OneFlow 中采用分段函數優化SoftmaxKernel：對于不同 num_cols范圍，選擇不同的實現，以在所有情況下都能達到較高的有效帶寬。見 Optimize softmax cuda kernel
OneFlow 中分三種實現，分段對 softmax 進行優化：
(1) 一個 Warp 處理一行的計算，適用于 num_cols <= 1024 情況
硬件上并行執行的32個線程稱之為一個warp，同一個warp的32個thread執行同一條指令。 warp是GPU調度執行的基本單元
(2) 一個 Block 處理一行的計算，借助 Shared Memory 保存中間結果數據，適用于需要的 Shared Memory 資源滿足 Kernel Launch 的可啟動條件的情況，在本測試環境中是 1024 < num_cols <= 4096
(3) 一個 Block 處理一行的計算，不使用 Shared Memory，重復讀輸入 x，適用于不支持(1)、（2)的情況
下面以前向計算為例分別對三種實現進行介紹：
每個 Warp 處理一行元素，行 Reduce 需要做 Warp 內的 Reduce 操作，這里實現 WarpAllReduce 完成 Warp 內各線程間的求 Global Max 和 Global Sum 操作，WarpAllReduce 是利用Warp級別原語 __shfl_xor_sync 實現的，代碼如下。
template<template<typename> typename ReductionOp, typename T>
__inline__ __device__ T WarpAllReduce(T val) {
for (int mask = kWarpSize / 2; mask > 0; mask /= 2) {
val = ReductionOp<T>()(val, __shfl_xor_sync(0xffffffff, val, mask));
}
return val;
}主體循環邏輯如下，首先根據 num_cols信息算出每個線程要處理的cols_per_thread，每個線程分配cols_per_thread大小的寄存器，將輸入x讀到寄存器中，后續計算均從寄存器中讀取。
理論上來說，每個 Warp 處理一行元素的速度是最快的，但是由于需要使用寄存器將輸入x緩存起來，而寄存器資源是有限的，并且在 num_cols>1024 情況下，使用(2)的 shared memory 方法已經足夠快了，因此僅在 num_cols<=1024 時采用 Warp 的實現。
template<typename T, int pack_size, int cols_per_thread, bool padding>
__global__ void SoftmaxWarpImpl(const int64_t rows, const int64_t cols, const T* x, T* y) {
static_assert(cols_per_thread % pack_size == 0, "");
constexpr int num_packs = cols_per_thread / pack_size;
assert(cols <= cols_per_thread * kWarpSize);
using ComputeType = typename GetComputeType<T>::type;
ComputeType buf[cols_per_thread];
const int global_warp_id = blockIdx.x * blockDim.y + threadIdx.y;
const int num_global_warp = gridDim.x * blockDim.y;
const int lane_id = threadIdx.x;
for (int64_t row = global_warp_id; row < rows; row += num_global_warp) {
const int64_t row_offset = row * cols;
const T* row_x = x + row_offset;
T* row_y = y + row_offset;
ComputeType thread_max = -Inf<ComputeType>();
#pragma unroll
for (int pack_id = 0; pack_id < num_packs; ++pack_id) {
const int col = (pack_id * kWarpSize + lane_id) * pack_size;
if (!padding || col < cols) {
MultiFetch<T, ComputeType, pack_size>()(buf + pack_id * pack_size, row_x + col);
#pragma unroll
for (int i = 0; i < pack_size; ++i) {
thread_max = max(thread_max, buf[pack_id * pack_size + i]);
}
} else {
#pragma unroll
for (int i = 0; i < pack_size; ++i) { buf[pack_id * pack_size + i] = -Inf<ComputeType>(); }
}
}
const ComputeType warp_max = WarpAllReduce<MaxOp, ComputeType>(thread_max);
ComputeType thread_sum = 0;
#pragma unroll
for (int i = 0; i < cols_per_thread; ++i) {
buf[i] = exp(buf[i] - warp_max);
thread_sum += buf[i];
}
const ComputeType warp_sum = WarpAllReduce<SumOp, ComputeType>(thread_sum);
#pragma unroll
for (int i = 0; i < cols_per_thread; ++i) { buf[i] = buf[i] / warp_sum; }
#pragma unroll
for (int i = 0; i < num_packs; ++i) {
const int col = (i * kWarpSize + lane_id) * pack_size;
if (!padding || col < cols) {
MultiStore<ComputeType, T, pack_size>()(row_y + col, buf + i * pack_size);
}
}
}
}一個 Block 處理一行元素，行 Reduce 需要做 Block 內的 Reduce 操作，需要做 Block 內線程同步，利用 BlockAllReduce 完成 Warp 內各線程間的求 Global Max 和 Global Sum 操作。BlockAllReduce 是借助 Cub 的 BlockReduce 方法實現的，代碼如下：
template<template<typename> typename ReductionOp, typename T, int block_size>
__inline__ __device__ T BlockAllReduce(T val) {
typedef cub::BlockReduce<T, block_size> BlockReduce;
__shared__ typename BlockReduce::TempStorage temp_storage;
__shared__ T result_broadcast;
T result = BlockReduce(temp_storage).Reduce(val, ReductionOp<T>());
if (threadIdx.x == 0) { result_broadcast = result; }
__syncthreads();
return result_broadcast;
}主體循環邏輯如下，根據 num_cols算出需要的 Shared Memory 大小作為 Launch Kernel 參數，借助 Shared Memory 保存輸入，后續計算直接從 Shared Memory 讀取。
由于 SM 內的 Shared Memory 資源同樣有限，因此當 num_cols超過一定范圍，kernel 啟動時申請 Shared Memory 超過最大限制，就會出現無法啟動的問題，因此，僅在調用 cudaOccupancyMaxActiveBlocksPerMultiprocessor 返回值大于0時采用 Shared Memory 方案。
此外，需要注意的是，由于 Block 內線程要做同步，當 SM 中正在調度執行的一個 Block 到達同步點時，SM 內可執行 Warp 逐漸減少，若同時執行的 Block 只有一個，則 SM 中可同時執行的 Warp 會在此時逐漸降成0，會導致計算資源空閑，造成浪費，若此時同時有其他 Block 在執行，則在一個 Block 到達同步點時仍然有其他 Block 可以執行。當 block_size 越小時，SM 可同時調度的 Block 越多，因此在這種情況下 block_size 越小越好。但是當在調大 block_size，SM 能同時調度的 Block 數不變的情況下，block_size 應該是越大越好，越大就有越好的并行度。因此代碼中在選擇 block_size 時，對不同 block_size 都計算了 cudaOccupancyMaxActiveBlocksPerMultiprocessor，若結果相同，使用較大的 block_size。
template<typename T, int pack_size, int block_size>
__global__ void SoftmaxBlockSMemImpl(const int64_t rows, const int64_t cols, const T* x, T* y) {
using ComputeType = typename GetComputeType<T>::type;
extern __shared__ __align__(sizeof(ComputeType)) unsigned char shared_buf[];
auto* buf = reinterpret_cast<ComputeType*>(shared_buf);
const int tid = threadIdx.x;
assert(cols % pack_size == 0);
const int num_packs = cols / pack_size;
for (int64_t row = blockIdx.x; row < rows; row += gridDim.x) {
const int64_t row_offset = row * cols;
const T* row_x = x + row_offset;
T* row_y = y + row_offset;
ComputeType thread_max = -Inf<ComputeType>();
for (int pack_id = tid; pack_id < num_packs; pack_id += block_size) {
ComputeType pack[pack_size];
MultiFetch<T, ComputeType, pack_size>()(pack, row_x + pack_id * pack_size);
#pragma unroll
for (int i = 0; i < pack_size; ++i) {
buf[i * num_packs + pack_id] = pack[i];
thread_max = max(thread_max, pack[i]);
}
}
const ComputeType row_max = BlockAllReduce<MaxOp, ComputeType, block_size>(thread_max);
ComputeType thread_sum = 0;
for (int col = tid; col < cols; col += block_size) {
const ComputeType exp_x = exp(buf[col] - row_max);
buf[col] = exp_x;
thread_sum += exp_x;
}
const ComputeType row_sum = BlockAllReduce<SumOp, ComputeType, block_size>(thread_sum);
for (int pack_id = tid; pack_id < num_packs; pack_id += block_size) {
ComputeType pack[pack_size];
#pragma unroll
for (int i = 0; i < pack_size; ++i) {
pack[i] = buf[i * num_packs + pack_id] / row_sum;
thread_max = max(thread_max, pack[i]);
}
MultiStore<ComputeType, T, pack_size>()(row_y + pack_id * pack_size, pack);
}
}
}和實現2一樣，仍然是一個 Block 處理一行元素，不同的是，不再用 Shared Memory 緩存x，而是在每次計算時重新讀輸入 x，這種實現沒有最大 num_cols的限制，可以支持任意大小。
此外，需要注意的是，在這種實現中，block_size 應該設越大越好，block_size 越大，SM 中能同時并行執行的 Block 數就越少，對 cache 的需求就越少，就有更多機會命中 Cache，多次讀x不會多次訪問 Global Memory，因此在實際測試中，在能利用 Cache 情況下，有效帶寬不會因為讀3次x而降低幾倍。
template<typename T, int pack_size, int block_size>
__global__ void SoftmaxBlockUncachedImpl(const int64_t rows, const int64_t cols, const T* x, T* y) {
using ComputeType = typename GetComputeType<T>::type;
const int tid = threadIdx.x;
assert(cols % pack_size == 0);
const int num_packs = cols / pack_size;
for (int64_t row = blockIdx.x; row < rows; row += gridDim.x) {
const int64_t row_offset = row * cols;
const T* row_x = x + row_offset;
T* row_y = y + row_offset;
ComputeType thread_max = -Inf<ComputeType>();
for (int pack_id = tid; pack_id < num_packs; pack_id += block_size) {
ComputeType pack[pack_size];
MultiFetch<T, ComputeType, pack_size>()(pack, row_x + pack_id * pack_size);
#pragma unroll
for (int i = 0; i < pack_size; ++i) { thread_max = max(thread_max, pack[i]); }
}
const ComputeType row_max = BlockAllReduce<MaxOp, ComputeType, block_size>(thread_max);
ComputeType thread_sum = 0;
for (int pack_id = tid; pack_id < num_packs; pack_id += block_size) {
ComputeType pack[pack_size];
MultiFetch<T, ComputeType, pack_size>()(pack, row_x + pack_id * pack_size);
#pragma unroll
for (int i = 0; i < pack_size; ++i) { thread_sum += exp(pack[i] - row_max); }
}
const ComputeType row_sum = BlockAllReduce<SumOp, ComputeType, block_size>(thread_sum);
for (int pack_id = tid; pack_id < num_packs; pack_id += block_size) {
ComputeType pack[pack_size];
MultiFetch<T, ComputeType, pack_size>()(pack, row_x + pack_id * pack_size);
#pragma unroll
for (int i = 0; i < pack_size; ++i) { pack[i] = exp(pack[i] - row_max) / row_sum; }
MultiStore<ComputeType, T, pack_size>()(row_y + pack_id * pack_size, pack);
}
}
}除了以上介紹的針對 softmax 的分段優化技巧外，OneFlow 在 softmax 實現中，還使用了一些通用的公共優化技巧，他們不僅應用于 softmax 中，也應用在其它 kernel 實現中。在此介紹兩種：
1、將 Half 類型 pack 成 Half2 進行存取，在不改變延遲情況下提高指令吞吐，類似 CUDA template for element-wise kernels 的優化
2、Shared Memory 中的 Bank Conflicts
CUDA 將 Shared Memory 按照4字節或8字節大小劃分到32個 bank 中，對于 Volta 架構是4字節，這里以4字節為例，如下圖所示，0-128 bytes地址分別在bank 0-31中，128-256也分別在 bank0-31中。
注：此圖及以下
Bank Conflicts圖片來自VOLTA Architecture and performance optimization
當在一個 Warp 內的不同線程訪問同一個 bank 的不同地址，就會出現 Bank Conflicts。當發生 Bank Conflicts 時，線程間只能順序訪問，增加延遲，下圖是一個 Warp 中 n 個線程同時訪問同一個 bank 中的不同地址時的延遲情況。
注：圖來自Dissecting the NVIDIA Volta GPU Architecture via Microbenchmarking
Bank Conflicts 的幾種情況：
若一個 Warp 內每個線程讀4個字節，順序訪問，則不會有 Bank Conflicts，若一個 Warp 內每個線程讀8個字節，則 Warp 內0號線程訪問的地址在第0和第1個 bank，1號線程訪問的地址在第2和第3個 bank，以此類推，16號線程訪問地址又在第0和第1個 bank內，和0號線程訪問了同一個bank的不同地址，此時即產生了 Bank Conflicts。
在前文的實現(2)中，給 Shared memory 賦值過程中，若采用下面方法，當 pack size=2，每個線程寫連續兩個4 byte 地址，就會產生 Bank Conflicts。
#pragma unroll
for (int j = 0; j < pack_size; ++j) {
buf[pack_id * pack_size * j] = pack[j];
thread_max = max(thread_max, pack[j]);
}因此，在實現(2)中，對Shared memory采用了新的內存布局，避免了同一個Warp訪問相同bank的不同地址，避免了Bank Conflicts。
#pragma unroll
for (int j = 0; j < pack_size; ++j) {
buf[num_packs * j + pack_id] = pack[j];
thread_max = max(thread_max, pack[j]);
}關于怎么實現一個高效的Softmax CUDA kernel就分享到這里了，希望以上內容可以對大家有一定的幫助，可以學到更多知識。如果覺得文章不錯，可以把它分享出去讓更多的人看到。
免責聲明：本站發布的內容（圖片、視頻和文字）以原創、轉載和分享為主，文章觀點不代表本網站立場，如果涉及侵權請聯系站長郵箱：is@yisu.com進行舉報，并提供相關證據，一經查實，將立刻刪除涉嫌侵權內容。
