我們提供的 NVIDIA CUDA 深度神經網絡庫(cuDNN) 是一個專門為深度學習應用而設計的 GPU 加速庫,旨在以先進的性能加速深度學習基元。
cuDNN 與 PyTorch、TensorFlow 和 XLA (加速線性代數) 等熱門深度學習框架集成。這些框架抽象化了直接 GPU 編程的復雜性,使您能夠專注于設計和訓練模型,而無需擔心底層硬件。cuDNN 可作為底層性能引擎,確保以更高的效率執行這些框架上的操作。
最近,擴展的點積注意力 (SDPA) 已成為大型語言模型 (LLM) 等重要工作負載中的性能關鍵基元。cuDNN 增加了對該基元的支持,并一直在使用閃存注意力和其他優化來提高其性能釋放,同時擴展功能支持面,以支持一系列注意力用例。
在 NVIDIA H200 Tensor Core GPU 上,cuDNN 可以在 FP8 中實現高達 1.2 PFLOPS 的性能。作為端到端示例,我們的團隊在為 Lama2 70B LoRA 微調啟用 cuDNN FP8 SDPA 后,測得了 1.15 倍的速度提升。該實驗使用了 NVIDIA NeMo 和?NVIDIA Transformer 引擎 在 8-GPU H200 節點上運行。
軟件堆棧 | e2e 加速 |
禁用 cuDNN 的 NeMo | 1x (基準) |
在 BF16 中使用 cuDNN SDPA 的 NeMo | 1.11 倍 |
在 FP8 中使用 cuDNN SDPA 的 NeMo | 1.15 倍 |
在本文中,我將詳細介紹 cuDNN SDPA 可實現的性能,介紹使用方法,并簡要總結 cuDNN 9 中的其他一些值得注意的新功能。
經擴展的點積產品注意力性能
NVIDIA 正式通過在 NVIDIA Omniverse 中開源融合的 Multihead Attention (fMHA) 內核來支持注意力機制。APEX 庫將注意力算法融合到單個內核中。Tri Dao 的創新工作以此內核為起點,實現了閃存注意力的形式,帶來了巨大的性能改進和功能增強。欲了解更多信息,請參閱 Dao-AILab/flash-attention 在 GitHub 上的項目頁面,并查看 FlashAttention-2:通過更好的并行性和作業分區提高注意力 相關論文。
然后, NVIDIA 通過更快、更靈活的實現,在融合注意力方面取得了進展。此實現現在是 NVIDIA 的最新技術成果,現在已經集成到 NVIDIA Hopper 架構 的 GPU 中。
XLA 目前提供了 cuDNN SDPA 的路徑。您可以通過 JAX SDPA API 或依靠 XLA 編譯器將 JAX/PyTorch 中的自定義實現降低到 cuDNN SDPA。
目前,PyTorch 渴望模式 SDPA 不使用 cuDNN,但基于 cuDNN 的實現正在進行中。有關更多信息,請參閱 PyTorch 的相關 Pull Request:Fprop 和 Bprop。
cuDNN SDPA 實施封裝了以下內容:
- 深入了解基礎 NVIDIA 硬件架構
- 實現從非閃存到閃存注意力 v2 的所有先進的 SDPA 算法,以及介于兩者之間的所有算法
- 啟發式算法,可根據問題大小和目標 GPU 架構自動設置性能旋鈕 (例如圖塊大小)
這可在 NVIDIA GPU 上實現出色的可用性能。圖 1 和圖 2 顯示,在各種用例中,cuDNN 9 BF16 的速度比現有最佳 GPU 快 2 倍,而 PyTorch 即時實現 中 cuDNN FP8 的實現速度提升高達 3 倍。
性能提升可實現更長的序列長度以及更短的模型預訓練和微調時間。與其他公開基準(例如 flash-attention/benchmarks)相比,本文僅報告 GPU 時間,不包括主機用度。



SDPA 作為 cuDNN 圖形
cuDNN 中的 SDPA 可指定為張量運算的 cuDNN 圖形。對于給定的圖形,cuDNN 庫擁有一組可以執行該圖形的引擎。
雖然一些圖形可能沒有合適的引擎,但其目的是為在 GPU 上原子化執行的任何 cuDNN 圖形至少提供一個引擎,通常在一個融合內核中,但有時在一小部分協作內核中。您可以將其視為整個框架圖形的“子圖”。
SDPA (各種形式的 SDPA) 是此類子圖的理想示例。支持這些模式的引擎旨在盡可能靈活,而不會影響顯著的性能,而是使用現有的最佳算法 (如 Flash Attention)。這意味著您可以更改注意力計算,但仍然可以在 GPU 上高效運行。本節旨在解釋目前提供的靈活性和支持。
圖 4 顯示了前向傳播 (fprop) 用例的 cuDNN 圖形。這是神經網絡中注意力機制 SDPA 的前向計算所涉及的運算序列。

因果關系掩膜邏輯降低到單個生成索引、逐點和選擇運算。softmax 邏輯降低到六個 cuDNN 運算 (圖 4)。這種 cuDNN 運算的粒度使您能夠靈活地表達自定義模型。圖 5 顯示了可能的組合。

圖 5 顯示,在頭部維度更大 (256) 且無因遮罩的情況下,cuDNN FP8 正向閃光注意力可實現高達 1.2 PFLOPS 的運算性能。
您還可以使用 BMM1 和 Softmax 之間的任意逐點運算構建自定義圖形。這種靈活性支持尚未發現的新變體。
當創意研究人員調整典型的注意力模式時,他們就不太容易回到優化程度較低的框架實現,從而在性能上達到峰值。圖 2 適用于完整的訓練運行。反向道具注意力圖也具有類似的靈活性。
SDPA 使用演示
有多個 API 切入點可用于創建和運行 cuDNN 圖形:
- 前端 API (包括 C++和 Python 變體)
- 后端 API (僅支持 C 語言)
上一節中描述的所有 cuDNN 圖形概念都適用于這兩個級別的 API.但是,cuDNN 團隊建議您使用 Python 或 C++與 cuDNN 前端 API 進行交互,除非您需要 C 接口。前端 API 明顯更簡潔,并增加了一些便利。
例如,前端 API 擴展了操作節點的概念,使節點能夠封裝多個操作和它們之間的數據流。換言之,它們是方便節點,可以抽象出常見圖形模式 (例如 SDPA) 的細節。節點仍然可以靈活配置眾所周知的變體。此 SDPA 使用演練從最簡單的情況開始,即在 Python 中創建的 SDPA 節點。
我們在 前端庫中的 SDPA Python 示例 中演示了配置選項和基本使用流程:
- 初始化一個具有適當數據類型的
cudnn.pygraph
對象。 - 創建一個包含張量維度、布局和數據類型的
Tensor
對象。 - 創建縮放的點積閃光注意力節點并提供所需的配置。
- 構建圖形并提供設備指針。
- 執行圖形。
# 1. cuDNN graph graph = cudnn.pygraph( io_data_type = cudnn.data_type.BFLOAT16, intermediate_data_type = cudnn.data_type. FLOAT , compute_data_type = cudnn.data_type. FLOAT , ) # 2. tensor attributes q = make_tensor_attr(graph, q_gpu, "q" ) k = make_tensor_attr(graph, k_gpu, "k" ) v = make_tensor_attr(graph, v_gpu, "v" ) attn_scale = make_tensor_attr(graph, attn_scale_cpu, "attn_scale" , is_pass_by_value = True ) # 3. sdpa node with the configurations o, stats = graph.scaled_dot_product_flash_attention( name = "scaled_dot_product_flash_attention" , q = q, k = k, v = v, is_inference = false, attn_scale = attn_scal, bias = None , use_alibi_mask = false, use_padding_mask = false, seq_len_q = None , seq_len_kv = None , use_causal_mask = true, dropout = None , ) # 4. build the graph and provide the device pointers graph.build() variant_pack = { q: q_gpu, k: k_gpu, v: v_gpu, o: o_gpu } # 5. execute the graph graph.execute() |
如果您想獲得比 SDPA 節點默認情況下提供的更高的靈活性,則構建底層 SDPA 圖形的代碼是開源的,因此可以進行自定義。
例如,以下代碼示例可自定義 SDPA 節點內的擴展節點。有關更多信息,請參閱 scaled_dot_product_flash_attention.h 文件。
// Optional scale if (options.inputs.Attn_scale) { // Lower options to scale options auto attn_scale_output = std::make_shared(); attn_scale_output->set_is_virtual( true ); Pointwise_attributes scale_attributes; scale_attributes.set_name( "attn_scale" ); scale_attributes.set_mode(PointwiseMode_t::MUL); scale_attributes.inputs.IN_0 = last_output; scale_attributes.inputs.IN_1 = options.inputs.Attn_scale; last_output = scale_attributes.outputs.OUT_0 = attn_scale_output; auto scale_node = std::make_unique(std::move(scale_attributes), context); sub_nodes.emplace_back(std::move(scale_node)); } |
如果您想詳細了解如何加速您的自定義 Transformer,請參閱 NVIDIA cuDNN 的官方文檔。
cuDNN 9 的其他顯著特性
除了 SDPA 改進之外,cuDNN 9 還引入了其他幾項重要改進:
- 混合輸入精度支持矩陣和卷積
- 改進錯誤報告
- 硬件前向兼容性
- 簡化安裝
混合輸入精度支持矩陣和卷積
要求輸入操作數的數據類型為同一類型 (FP16、FP32) 的矩陣乘法 (matmul) 和卷積 API 不適合 AWQ (激活感知權重量化) 等情況,其中激活在 FP16 中,權重可能在 INT8 中。假設您想以更高的精度計算,則必須投射數據。如果不是在線的,這將需要額外的內存和轉換成本。
cuDNN 現在支持混合輸入精度矩陣和卷積,其中 A 和 B 操作數可以是不同的數據類型,并提供用于性能和內存優化的在線融合類型轉換。您可以選擇將操作數轉換為其他類型。cuDNN 在優化的內核中處理所需的轉換。
圖 6 顯示了 cuDNN 混合輸入精度矩陣與未融合工作流程之間的加速。灰色條表示輸入 A 和 B 分別采用 FP16 和 INT8 精度的情況,其中 A 轉換為 INT8,然后使用 INT32 累加進行 INT8xINT8 矩陣乘法運算。綠色條表示 A 從 INT8 上轉換為 FP16,然后使用 FP32 累加進行 FP16xFP16 矩陣乘法運算的情況。

改進錯誤報告
日志記錄在軟件開發中至關重要,對于使用 cuDNN 的深度學習框架等復雜系統尤其如此。過去,cuDNN 的一個常見痛點是難以調試錯誤和警告。我們一直在不斷改進錯誤報告以解決這一問題。
cuDNN 9 增加了以下功能:
- 更具體的錯誤代碼
- 進行分類以幫助組織數量增加的錯誤碼
- 符合日志記錄規范的嵌入式日志記錄級別
cuDNNGetLastErrorString
是一個新的函數,用于以編程方式檢索最新的錯誤消息。
欲了解更多信息,請參閱 錯誤報告和 API 日志記錄 部分內容,詳見開發者指南。
硬件前向兼容性
在版本 9.0.0 之前,cuDNN 庫在庫發布之日支持最新的公開可用 GPU 架構。例如,cuDNN 8.9.x 通過 NVIDIA Hopper (即計算能力 9.0) 提供支持。不支持在未來的 GPU 架構上運行 8.9.x。
但是,cuDNN 9 對 API 的一大子集具有硬件前向兼容性。這意味著,僅將此 API 子集與 cuDNN v9 結合使用的程序將在未來的 GPU 上正常運行,并且這些程序的用戶不會被強制升級其 cuDNN 安裝以使用未來的 GPU。
在計算能力大于 9.0 的 GPU 上運行時,前向兼容性不是糾錯,而是意味著庫可以找到功能等效的實現,并使用 PTX JIT 針對新架構進行定位。
欲了解更多關于支持限制和最佳實踐的信息,請參閱 硬件前向兼容性 部分內容,在開發者指南中。
簡化安裝
在 Python 環境中,除了庫之外,您現在還可以使用 pip 安裝新的 Python 前端:
# cuDNN library for CUDA 12 pip install nvidia-cudnn-cu12 # cuDNN frontend pip install git+https: //github .com /NVIDIA/cudnn-frontend .git |
cuDNN 9 還簡化了 RPM 和 Debian 元軟件包的安裝流程。例如,以下命令可在 Ubuntu 22 上安裝 cuDNN:
wget https: //developer .download.nvidia.com /compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1 .1-1_all.deb sudo dpkg -i cuda-keyring_1.1-1_all.deb sudo apt-get update sudo apt-get install -y cudnn |
這個新流程簡化了安裝 Debian 軟件包時的密鑰環設置,并使用 cudnn
元軟件包,以提高安裝效率。
如果您需要了解更多信息和完整說明,請參閱 cuDNN 安裝指南,以獲取詳細的安裝步驟和配置說明。
后續步驟
如果您有反饋、問題或意見,可以在 cuDNN NVIDIA 開發者論壇(適用于 cuDNN 庫主題) 或 NVIDIA/cudnn-frontend 在 GitHub 上發布 (適用于前端主題)。另外,您可以 下載 cuDNN 開始使用。
請持續關注,了解未來更多新的 cuDNN 功能。本文中概述的進展是一個重要的里程碑,但未來還會有更多進展。
隨著 AI 不斷將行業推向軟硬件集成的極限, NVIDIA 不斷優化性能并改善用戶體驗,以便 cuDNN 能夠在深度學習框架和圖形編譯器中得到更有效、更廣泛的使用。
致謝
NVIDIA cuDNN 團隊與公司內的許多其他團隊緊密合作,為這篇博文提供了技術內容。
?