如何實現一個高效的Softmax CUDA kernel?

語言: CN / TW / HK

撰文 | 郭冉

Softmax操作是深度學習模型中最常用的操作之一。在深度學習的分類任務中,網路最後的分類器往往是Softmax + CrossEntropy的組合:

如何實現一個高效的Softmax CUDA kernel?

儘管當Softmax和CrossEntropy聯合使用時,其數學推導可以約簡,但還是有很多場景會單獨使用Softmax Op。如BERT的Encoder每一層的attention layer中就單獨使用了Softmax求解attention的概率分佈;GPT-2的attention的multi-head部分也單獨使用了Softmax 等等。

如何實現一個高效的Softmax CUDA kernel?

深度學習框架中的所有計算的運算元都會轉化為GPU上的CUDA kernel function,Softmax操作也不例外。Softmax作為一個被大多數網路廣泛使用的運算元,其CUDA Kernel實現高效性會影響很多網路最終的訓練速度。那麼如何實現一個高效的Softmax CUDA Kernel?本文將會介紹OneFlow中優化的Softmax CUDA Kernel的技巧,並跟cuDNN中的Softmax操作進行實驗對比,結果表明,OneFlow深度優化後的Softmax對視訊記憶體頻寬的利用率可以接近理論上限,遠高於cuDNN的實現。

\

GPU基礎知識與CUDA效能優化原則

\

對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的優化目標就是儘可能充分利用視訊記憶體頻寬資源。

如何評估一個CUDA Kernel是否充分利用了視訊記憶體頻寬資源?

對於視訊記憶體頻寬資源來說,“充分利用“指的是Kernel的有效視訊記憶體讀寫頻寬達到了裝置視訊記憶體頻寬的上限,其中裝置視訊記憶體頻寬可以通過執行 cuda中的的bandwidthTest得到。 Kernel的有效視訊記憶體頻寬通過Kernel讀寫資料量和Kernel執行時間進行評估: $$ 當前Kernel的有效視訊記憶體頻寬 = 讀寫資料量 / 執行時間 $$

Naive的Softmax實現:

在介紹優化技巧之前,我們先看看一個未經優化的Softmax Kernel的最高理論頻寬是多少。如下圖所示,一個最簡單的Softmax計算實現中,分別呼叫幾個基礎的CUDA Kernel function來完成整體的計算:

如何實現一個高效的Softmax CUDA kernel?

假設輸入的資料大小為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倍資料的視訊記憶體,即: $$ Naive Softmax Kernel 有效視訊記憶體頻寬 < 理論頻寬 / 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 與 cuDNN 的對比結果

我們對OneFlow深度優化後的Softmax CUDA Kernel 與 cuDNN中的Softmax Kernel的訪存頻寬進行了對比測試,測試結果如下:

Softmax Kernel Bandwidth:

如何實現一個高效的Softmax CUDA kernel?

Softmax Grad Kernel Bandwidth:

如何實現一個高效的Softmax CUDA kernel?

其中測試環境是 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的實現。

OneFlow深度優化Softmax CUDA Kernel的技巧

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)的情況

下面以前向計算為例分別對三種實現進行介紹:

實現1:每個Warp處理一行或兩行元素。

\

每個 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; }

SoftmaxWarpImpl的實現有以下幾個模板引數:

LOAD、STORE分別代表輸入輸出,使用load.template load<pack_size>(ptr, row_id, col_id);store.template store<pack_size>(ptr, row_id, col_id); 進行讀取和寫入。使用LOAD和STORE有兩個好處:1、可以在CUDA Kernel中只關心計算型別ComputeType,而不用關心具體的資料型別T。2、只需要加幾行程式碼就可以快速支援Softmax和其他Kernel Fuse,減少頻寬需求,提升整體效能。普通的SoftmaxKernel直接使用DirectLoad和DirectStore,FusedSoftmaxKernel如\ FusedScaleSoftmaxDropoutKernel只需要定義一個ScaleLoad結構和一個DropoutStore結構用於對輸入x做Scale預處理和對輸出y做Dropout後處理。

ComputeType代表計算型別。pack_size代表向量化訪存操作的pack元素的個數,我們將幾個元素pack起來讀寫,提升頻寬利用率。cols_per_thread代表每個執行緒處理的元素個數。

thread_group_width代表處理元素的執行緒組的寬度,當cols > pack_size * warp_size時,thread_group_width就是warp_size,即32。當cols < pack_size * warp_size時,就根據cols大小用1/2個warp或1/4個warp來處理每行的元素。採用更小的thread_group_width後,WarpAllReduce需要執行的輪次也相應減少。

rows_per_access代表每個thread_group一次處理的行數,當cols較小,thread_group_width不是warp_size 32時,若rows能被2整除,我們就讓每個執行緒處理2行來增加指令並行度,從而提升效能。

padding代表當前是否做了padding,若cols不是warp_size的整數倍,我們會把它padding到最近的整數倍處理。

algorithm代表使用的演算法,可選項有Algorithm::kSoftmax或Algorithm::kLogSoftmax。

CUDA Kernel執行的主體迴圈邏輯如下,首先根據 num_cols資訊算出每個執行緒要處理的cols_per_thread,每個執行緒分配rows_per_access * cols_per_thread大小的暫存器,將輸入x讀到暫存器中,後續計算均從暫存器中讀取。

理論上來說,以 Warp為單位處理一行元素的速度是最快的,但是由於需要使用暫存器將輸入x快取起來,而暫存器資源是有限的,並且在 num_cols>1024 情況下,使用(2)的 shared memory 方法已經足夠快了,因此僅在 num_cols<=1024 時採用 Warp 的實現。

``` template global void SoftmaxWarpImpl(LOAD load, STORE store, const int64_t rows, const int64_t cols) { static_assert(cols_per_thread % pack_size == 0, ""); static_assert(thread_group_width <= kWarpSize, ""); static_assert(kWarpSize % thread_group_width == 0, ""); constexpr int num_packs = cols_per_thread / pack_size; assert(cols <= cols_per_thread * thread_group_width); ComputeType buf[rows_per_access][cols_per_thread]; const int global_thread_group_id = blockIdx.x * blockDim.y + threadIdx.y; const int num_global_thread_group = gridDim.x * blockDim.y; const int lane_id = threadIdx.x; for (int64_t row = global_thread_group_id * rows_per_access; row < rows; row += num_global_thread_group * rows_per_access) { ComputeType thread_max[rows_per_access];

pragma unroll

for (int row_id = 0; row_id < rows_per_access; ++row_id) {
  thread_max[row_id] = -Inf<ComputeType>();
  ComputeType* row_buf = buf[row_id];

pragma unroll

  for (int pack_id = 0; pack_id < num_packs; ++pack_id) {
    const int col = (pack_id * thread_group_width + lane_id) * pack_size;
    if (!padding || col < cols) {
      load.template load<pack_size>(row_buf + pack_id * pack_size, row + row_id, col);

pragma unroll

      for (int i = 0; i < pack_size; ++i) {
        thread_max[row_id] = max(thread_max[row_id], row_buf[pack_id * pack_size + i]);
      }
    } else {

pragma unroll

      for (int i = 0; i < pack_size; ++i) {
        row_buf[pack_id * pack_size + i] = -Inf<ComputeType>();
      }
    }
  }
}
ComputeType warp_max[rows_per_access];

pragma unroll

for (int row_id = 0; row_id < rows_per_access; ++row_id) {
  warp_max[row_id] = WarpAllReduce<MaxOp, ComputeType, thread_group_width>(thread_max[row_id]);
}
ComputeType thread_sum[rows_per_access];

pragma unroll

for (int row_id = 0; row_id < rows_per_access; ++row_id) {
  thread_sum[row_id] = 0;
  ComputeType* row_buf = buf[row_id];

pragma unroll

  for (int i = 0; i < cols_per_thread; ++i) {
    if (algorithm == Algorithm::kSoftmax) {
      row_buf[i] = Exp(row_buf[i] - warp_max[row_id]);
      thread_sum[row_id] += row_buf[i];
    } else if (algorithm == Algorithm::kLogSoftmax) {
      row_buf[i] -= warp_max[row_id];
      thread_sum[row_id] += Exp(row_buf[i]);
    } else {
      __trap();
    }
  }
}
ComputeType warp_sum[rows_per_access];

pragma unroll

for (int row_id = 0; row_id < rows_per_access; ++row_id) {
  warp_sum[row_id] = WarpAllReduce<SumOp, ComputeType, thread_group_width>(thread_sum[row_id]);
}

pragma unroll

for (int row_id = 0; row_id < rows_per_access; ++row_id) {
  ComputeType* row_buf = buf[row_id];

pragma unroll

  for (int i = 0; i < cols_per_thread; ++i) {
    if (algorithm == Algorithm::kSoftmax) {
      row_buf[i] = Div(row_buf[i], warp_sum[row_id]);
    } else if (algorithm == Algorithm::kLogSoftmax) {
      row_buf[i] -= Log(warp_sum[row_id]);
    } else {
      __trap();
    }
  }

pragma unroll

  for (int i = 0; i < num_packs; ++i) {
    const int col = (i * thread_group_width + lane_id) * pack_size;
    if (!padding || col < cols) {
      store.template store<pack_size>(row_buf + i * pack_size, row + row_id, col);
    }
  }
}

} } ```

實現2: 一個Block處理一行元素

一個 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; }

CUDA Kernel的各個模板引數在(1)中做了詳細介紹,執行的主體迴圈邏輯如下,根據 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 global void SoftmaxBlockSMemImpl(LOAD load, STORE store, const int64_t rows, const int64_t cols) { extern shared align(sizeof(double)) unsigned char shared_buf[]; auto* buf = reinterpret_cast(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) { ComputeType thread_max = -Inf(); for (int pack_id = tid; pack_id < num_packs; pack_id += block_size) { ComputeType pack[pack_size]; load.template load(pack, row, 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) {
  if (algorithm == Algorithm::kSoftmax) {
    const ComputeType exp_x = Exp(buf[col] - row_max);
    buf[col] = exp_x;
    thread_sum += exp_x;
  } else {
    const ComputeType x = buf[col] - row_max;
    buf[col] = 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) {
    if (algorithm == Algorithm::kSoftmax) {
      pack[i] = Div(buf[i * num_packs + pack_id], row_sum);
    } else if (algorithm == Algorithm::kLogSoftmax) {
      pack[i] = buf[i * num_packs + pack_id] - Log(row_sum);
    } else {
      __trap();
    }
    thread_max = max(thread_max, pack[i]);
  }
  store.template store<pack_size>(pack, row, pack_id * pack_size);
}

} } ```

實現3:一個Block處理一行的元素,不使用Shared Memory,重複讀輸入x

和實現2一樣,仍然是一個 Block 處理一行元素,不同的是,不再用 Shared Memory 快取輸入x,而是在每次計算時重新讀輸入 x,這種實現沒有最大 num_cols的限制,可以支援任意大小。

此外,需要注意的是,在這種實現中,block_size 應該設越大越好,block_size 越大,SM 中能同時並行執行的 Block 數就越少,對 cache 的需求就越少,就有更多機會命中 Cache,多次讀x不會多次訪問 Global Memory,因此在實際測試中,在能利用 Cache 情況下,有效頻寬不會因為讀3次x而降低幾倍。

``` template global void SoftmaxBlockUncachedImpl(LOAD load, STORE store, const int64_t rows, const int64_t cols) { 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) { ComputeType thread_max = -Inf(); for (int pack_id = tid; pack_id < num_packs; pack_id += block_size) { ComputeType pack[pack_size]; load.template load(pack, row, 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];
  load.template load<pack_size>(pack, row, 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];
  load.template load<pack_size>(pack, row, pack_id * pack_size);

pragma unroll

  for (int i = 0; i < pack_size; ++i) {
    if (algorithm == Algorithm::kSoftmax) {
      pack[i] = Div(Exp(pack[i] - row_max), row_sum);
    } else if (algorithm == Algorithm::kLogSoftmax) {
      pack[i] = (pack[i] - row_max) - Log(row_sum);
    } else {
      __trap();
    }
  }
  store.template store<pack_size>(pack, row, pack_id * pack_size);
}

} } ```

公共優化技巧

除了以上介紹的針對 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中。

如何實現一個高效的Softmax CUDA kernel?

:此圖及以下 Bank Conflicts 圖片來自VOLTA Architecture and performance optimization

當在一個 Warp 內的不同執行緒訪問同一個 bank 的不同地址,就會出現 Bank Conflicts。當發生 Bank Conflicts 時,執行緒間只能順序訪問,增加延遲,下圖是一個 Warp 中 n 個執行緒同時訪問同一個 bank 中的不同地址時的延遲情況。

如何實現一個高效的Softmax CUDA kernel?

:圖來自Dissecting the NVIDIA Volta GPU Architecture via Microbenchmarking

Bank Conflicts 的幾種情況:

如何實現一個高效的Softmax CUDA kernel?

如何實現一個高效的Softmax CUDA kernel?

若一個 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]);
  }

```

參考資料:

Using CUDA Warp-Level Primitives | NVIDIA Developer Blog

CUDA Pro Tip: Increase Performance with Vectorized Memory Access

Dissecting the NVIDIA Volta GPU Architecture via Microbenchmarking

VOLTA Architecture and performance optimization

歡迎下載體驗OneFlow深度學習框架:https://github.com/Oneflow-Inc/oneflow/