xlite-dev/LeetCUDA

解決的問題

LeetCUDA 為初學者提供了掌握現代 CUDA 編程的全面、結構化學習路徑，專門針對 AI 工作負載的內核優化。它彌補了基礎 CUDA 知識與關鍵 AI 算子（如矩陣乘法 (GEMM) 與注意力機制）的高性能實現之間的鴻溝。

工作原理

該專案透過從「Easy」到「Hard++」難度的 200 多個 CUDA 內核的進階過程來組織學習。每個主題都遵循一致的工作流程：實現自定義 CUDA 內核、建立 PyTorch Python 綁定並執行測試。它特別強調使用 Tensor Cores（透過 WMMA、MMA 與 CuTe），並支援包括 FP32、FP16、BF16 與 FP8 在內的各種數據類型。該儲存庫包含 HGEMM 與 FlashAttention-2 的高性能實現，展示了諸如 block swizzling、multi-staging 與 fine-grained tiling 等高級技術。

適用對象

專為想要學習如何為深度學習編寫高性能 CUDA 內核的初學者與開發人員，以及準備參加 AI 基礎設施領域技術面試的人士設計。

亮點

• 廣泛的內核庫：涵蓋逐元素操作、歸約與複雜 AI 算子的 200 多個內核。
• 高性能基準測試：包括實現達到 cuBLAS 性能 98-100% 的 HGEMM 實現。
• 高級注意力實現：採用純 MMA PTX 指令並具有 Split-Q 與 shared QKV SMEM 等優化的 FlashAttention-2。
• 現代硬體支援：建置指令碼針對 Ada Lovelace (sm_89)、Hopper (sm_90a) 與 Blackwell (sm_120a) 架構。
• PyTorch 整合：所有內核都設計為可與 PyTorch 綁定一起使用，以便於測試與整合。