【新智元導讀】本想練練手合成點數據，沒想到卻一不小心幹翻了PyTorch專家內核！斯坦福華人團隊用純CUDA-C寫出的AI生成內核，瞬間驚豔圈內並登上Hacker News熱榜。團隊甚至表示：本來不想發這個結果的。

就在剛剛，斯坦福HAI華人大神團隊又出驚人神作了。

他們用純CUDA-C語言編寫的快速AI生成內核，竟然超越了PyTorch！

在這個過程中，完全不用藉助CUTLASS和Triton等庫和領域特定語言（DSL），就能讓性能表現接近PyTorch內置的、經過專家優化的標準生產級內核，甚至在某些情況下還更勝一籌。

作者團隊都是我們熟悉的名字——Anne Ouyang、Azalia Mirhoseini和Percy Liang，有趣的是，他們甚至直言，這個結果其實本不想拿出來發佈。

一經發布，這個發現就引爆了技術圈，現在已經登頂Hacker News總榜第二。

說起來，這個發現還有很多意外的成分。

本來，他們的目標是生成合成數據，來訓練更好的內核生成模型，合成數據生成的設計也十分簡單。

然而，意想不到的事情發生了，僅用於測試的合成數據生成本身，竟開始生成非常優秀的內核，甚至超越了人類專家優化的PyTorch基線，而且還利用了高級優化和硬件特性。

而在此前，這是一項很艱難的挑戰。

由此，研究者們決定提前撰寫博文，把自己的發現分享出來。

總結來說，研究的亮點成果如下：

矩陣乘法（Matmul, FP32）：性能達到PyTorch FP32 torch.matmul的101.3%

二維卷積（Conv2D, FP32）：性能達到PyTorch FP32 torch.nn.Conv2D的179.9%

Softmax（FP32）：性能達到PyTorch FP32 torch.softmax的111.8%

層歸一化（LayerNorm, FP32）：性能達到PyTorch FP32 torch.nn.LayerNorm的484.4%

二維卷積 + ReLU + 最大池化（Conv2D + ReLU + MaxPool, FP32）：性能達到PyTorch FP32參考實現的 290.1%，達到PyTorch FP32 torch.compile()參考實現的189.0%

以上結果在英偉達L40S GPU上進行了基準測試，性能百分比定義爲參考時間除以生成的內核時間。

網友：強制LLM推理，實在太有趣了

在Hacker News上，網友們也對此展開了熱烈討論。

比如爲什麼使用FP32內核會比PyTorch更容易實現性能提升，理由就相當有趣。

如果AI真的能以更低成本，實現更優化的內核，的確潛力巨大。

最令人震撼的就是，無論是最近谷歌的AlphaEvolve，還是o3在Linux內核中發現了零日漏洞，都在提醒我們——

Gemini Pro 2.5和o3已經達到了一個全新的能力水平，那些曾經在其他模型上嘗試失敗的想法，現在突然奏效了。

可以說，我們已經到達了一個節點，LLM能比用人類快得多的速度進行迭代和測試，信息組合、進步和智能應用的蠻力，似乎正在成功！

接下來，我們來看看斯坦福研究者們博客中的具體內容。

博客全文

在博客中，研究者分享了具體方法、五個優化後的內核（包括4個基礎機器學習算子和1個AlexNet模塊的融合內核）、一個優化過程的實例，以及一些思考，關於這些發現對高性能內核生成可能意味着什麼。

可以說，這些內容將是他們後續探索的第一步。

方法

研究者們採用了KernelBench的任務設置（這是他們在2024年12月發佈的一款基於AI的內核生成基準測試）。

具體來說，給定一段torch代碼，LLM會編寫自定義內核來替換原有的torch算子，目標是實現加速。

依照KernelBench最初的設計，參考代碼默認使用FP32精度；在給定的容差閾值（1e-02）下，採用較低精度的解決方案也是被允許的。

此外，由於存在大量針對特定規模的優化手段，KernelBench中的每個問題都設定了具體的輸入大小。

因此，該基準測試旨在找出針對特定問題規模的最快內核，而非一個適用於任意問題規模的高速內核。

而且，研究者會同時運行torch參考代碼和生成的代碼，並通過在多種隨機輸入下比較兩者輸出的數值是否一致，來檢驗其正確性。

當前，在優化內核這個問題上，業界擴展測試時計算資源最常用的方法是順序修訂（sequential revision）。

這是一種多輪迭代的循環：模型首先對內核進行增量式修改，接着檢查其正確性和性能，然後根據結果再次嘗試。

也就是說，要麼修復有問題的內核，要麼進一步提升現有內核的性能。

這個循環過程非常直觀，也容易實現。模型會修復失效的內核，微調可用的內核，一步步優化出性能更佳的版本。

這種方法的主要侷限，在於優化思路缺乏多樣性。

順序循環往往容易陷入局部最優的困境，比如反覆嘗試同類型的轉換，或是在缺乏潛力的優化路徑上無休止地調整。

其結果便是測試時計算資源的低效利用，並且難以促使模型產生具有根本性創新的優化思路。

爲解決這一問題，研究者引入了兩項關鍵改變：

運用自然語言對優化思路進行推理

他們不再於每一步直接生成新的內核，而是以先前嘗試過的思路爲條件，用自然語言生成優化思路，隨後將這些思路具化爲新的代碼變體。

在每個優化步驟進行分支擴展

他們不是每步只改進一個候選方案，而是進行分支擴展，讓每個思路都能派生出多種實現版本，其中性能最佳的內核將作爲下一輪優化的種子。

（研究者也會保留一個表現優異的現有內核庫，用於提供種子）。

這種方式解鎖了大規模的並行處理能力，使他們能夠在每一輪探索截然不同的優化方向，避免陷入狹窄的優化路徑。

其結果是，這種測試時循環不再像順序修訂那般，僅僅是與編譯器“對話”，而是更接近一種結構化的探索性搜索。

這種搜索由明確的優化假設指導，並採用大規模並行評估的方式進行。

研究者運行了KernelBench第1級的10個問題，以進行測試。

他們調整了問題規模，以確保內核啓動開銷相對於問題的整體運行時間而言可以忽略不計。

然後，使用OpenAI o3和Gemini 2.5 Pro模型進行了5輪實驗。

下圖展示了首次發現性能最佳內核所在的輪次分佈情況。

可以看到，大多數最優結果出現在靠後的輪次（總共5輪），其中絕大部分出現在第4輪或第5輪。

隨着擴大搜索範圍，研究者還發現：許多高性能內核的優化策略高度相似，集中在少數幾種常見的模式上，這與他們手動編寫內核的經驗也是一致的。

主要的優化類別歸納如下——

內存訪問優化：提升不同內存層級（全局內存、共享內存、寄存器）之間數據遷移的效率，並確保數據訪問方式能夠最大化帶寬、最小化衝突。

異步操作與延遲隱藏：通過將耗時較長的操作（例如全局內存訪問）與計算或其他內存傳輸重疊執行，來隱藏其帶來的延遲。

數據類型與精度優化：在允許的條件下，儘可能使用較低精度的數據類型（如FP16或BF16），以降低內存帶寬需求，提升緩存效率，並有望利用專門的硬件加速單元。

計算與指令優化：提升算術運算本身的效率，削減指令數量，或利用專門的硬件指令。

並行性與佔用率增強：最大化流式多處理器（SM）上活躍線程束（warp）的數量，以便更好地隱藏延遲，提高整體吞吐率。

控制流與循環優化：減少由循環、分支及索引計算等引入的額外開銷。

總結

這次研究者採用的方法，與AI研究中一個日益顯著的趨勢不謀而合——

將強大的推理能力與對多個假設的並行探索相結合，能夠帶來性能的提升。

正如一些近期研究（例如AlphaEvolve、Gemini 2.5 Pro Deep Think）所強調的，我們並不總是需要大規模的重新訓練。

論文地址：https://storage.googleapis.com/deepmind-media/DeepMind.com/Blog/alphaevolve-a-gemini-powered-coding-agent-for-designing-advanced-algorithms/AlphaEvolve.pdf

有時，巧妙的搜索和分支策略便足以催生科學創新、攻克複雜難題，而藉助驗證器進行廣泛搜索，則可能帶來更大的收益。

然而，這並不意味着我們不需要進一步的訓練。

恰恰相反，研究者的這種方法，也有助於生成更優質的合成數據，用以改進未來的模型訓練（這需要更多的問題實例）。

因此，它既是一種強大的測試時擴展方法，也是我們邁向更智能、數據效率更高的模型開發之路的一步。

而且，這次研究者展現的僅僅是初步的成果。這些優化結果的質量看起來相當可觀，但仍有廣闊的提升空間，例如產生更優的優化思路、生成更高質量的最終代碼，以及將此方法應用於日益複雜的內核。

目前，研究者仍在積極改進的兩個具體例子包括：

FP16 Matmul：性能達到torch.matmul的52%

FP16 Flash Attention：性能達到torch.nn.functional.scaled_dot_product_attention的9%

在現代機器學習任務中，FP32的應用不如FP16或BF16普遍，並且在較新的硬件上，針對FP32的優化往往也更少。

這或許能部分解釋，爲何基於FP32的內核更容易在性能上超越PyTorch。

作者介紹

Anne Ouyang

Anne Ouyang目前是斯坦福大學計算機科學（CS）博士生，在Scaling Intelligence Lab（可擴展智能實驗室）進行研究。

她的研究興趣主要集中在可擴展的自我改進機器學習系統，同時也廣泛關注實證機器學習（empirical ML）和性能工程（performance engineering）。

此前，她在MIT獲得學士和碩士學位，並曾在NVIDIA cuDNN團隊工作，負責編寫CUDA內核，用於加速GPU上的深度學習工作負載。

Azalia Mirhoseini

Azalia Mirhoseini是斯坦福大學計算機科學助理教授，也是Scaling Intelligence Lab（可擴展智能實驗室）的創始人，並在Google DeepMind兼任高級研究科學家。

她的實驗室致力於開發可擴展的自主演進人工智能系統與方法論，以期推動通用人工智能的發展。

在加入斯坦福大學之前，她曾在Google Brain和Anthropic等業界頂尖的人工智能實驗室工作多年。

她過往的卓越成就包括：

提出混合專家（MoE）神經架構——目前已被前沿的AI模型廣泛應用；

領導AlphaChip項目——一項將深度強化學習用於佈局優化的開創性工作，併成功應用於谷歌AI加速器（TPU）及數據中心CPU等先進芯片的設計中；

在測試時計算的Scaling方面有深入的研究

Percy Liang

Percy Liang是斯坦福大學計算機科學副教授，兼任基礎模型研究中心（CRFM）主任。同時也是CodaLab Worksheets的創建者，並藉此堅定倡導科研工作的可復現性。

他目前專注於通過開源和嚴格的基準測試，提升基礎模型（特別是大語言模型）的可及性與可理解性。

他曾圍繞機器學習和自然語言處理領域進行了廣泛研究，具體方向包括魯棒性、可解釋性、人機交互、學習理論、知識落地、語義學以及推理等。

此前，他於2004年在MIT獲得學士學位，並於2011年在UC伯克利獲得博士學位。

內核優化過程

最後，展示一個Conv2D自動生成思路的優化軌跡示例，torch參考基準時間爲1.41毫秒。

第0輪：7.02毫秒，達到參考性能的20.1%

思路：給定pytorch代碼，用CUDA Kernel替換操作。

第1輪：7.54毫秒，達到參考性能的18.8%

思路：通過使用__ldg加載不變的張量來利用只讀緩存。

第2輪：3.46毫秒，達到參考性能的41.0%

思路：將卷積轉換爲FP16 Tensor-Core GEMM。

作者評論：這是一種算法優化，將卷積轉換爲隱式GEMM，這對於在Tensor Core上高效運行卷積至關重要

第3輪：3.67毫秒，達到參考性能的38.7%

思路：採用雙緩衝cp.async管線，使全局內存加載與Tensor-Core計算重疊。

第4輪：3.46毫秒，達到參考性能的41.0%

思路：給定pytorch代碼，使用隱式矩陣乘法（implicit matmul）的CUDA Kernel替換操作。給定的GEMM內核可能會有幫助。

作者評論：因爲優化涉及到使用GEMM，所以在這一輪開始時，使用了一個之前生成的現有優秀GEMM內核作爲種子，這個思路是手動編寫的

第5輪：1.91毫秒，達到參考性能的74.9%

思路：在每個K-tile循環內部，預計算k_idx分解的內核/輸入索引並在共享內存中重用，以避免冗餘算術運算。

第6輪：1.37毫秒，達到參考性能的103.6%

思路：在共享內存中預計算並緩存N維GEMM索引，以減少輸入數據加載循環內的冗餘算術運算。

第7輪：1.38毫秒，達到參考性能的102.9%

思路：通過使用專用的每warp共享內存緩衝區來並行化CUDA內核輸出，以消除寫入全局內存時的warp級串行化。

第8輪：1.37毫秒，達到參考性能的103.6%

思路：在共享內存中預計算並緩存基礎輸入座標，以減少輸入數據加載過程中的冗餘算術運算。

第9輪：1.36毫秒，達到參考性能的105.1%

思路：對B-fragment加載進行軟件流水線處理，使下一個B-tile的共享內存讀取與當前B-tile的WMMA計算重疊。

第10輪：1.07毫秒，達到參考性能的133.6%

思路：重用共享內存中預計算的N維GEMM分解來進行輸出地址計算，避免冗餘且高成本的除法/模運算。

第11輪：1.21毫秒，達到參考性能的117.4%

思路：在half WMMA運算中移除hi/lo分解，如果最終精度可接受，則依賴標準的FP16累積來提高性能。

第12輪：1.01毫秒，達到參考性能的141.2%

思路：通過在流水線的加載階段動態計算K維索引，使用雙緩衝使Asub（權重）和Bsub（輸入）的K循環全局內存加載與MMA計算重疊。

第13輪：0.795毫秒，達到參考性能的179.9%

思路：通過使用像half2這樣的更寬數據類型，爲加載Asub_pipe和Bsub_pipe實現向量化的共享內存寫入。

最終代碼

最終生成的Conv2D內核代碼，使用了先進的CUDA技術，就是人類自己寫起來都很有挑戰性的那種！

import torchimport torch.nn as nnimport torch.nn.functional as Ffrom torch.utils.cpp_extension import load_inlineconv2d_implicit_gemm_cuda_source = r"""#include <torch/extension.h>#include <ATen/cuda/CUDAContext.h> // For at::cuda::getCurrentCUDAStream()#include <mma.h>#include <cuda_fp16.h>using namespace nvcuda;// WMMA tile dimensions#define WMMA_M 16#define WMMA_N 16#define WMMA_K 16// Skew padding for shared memory to avoid bank conflicts#define SKEW_HALF 8 // 8 half elements (16 bytes)// CUDA built-in warpSize is 32 for supported architectures (sm_70+)// This constant is used for host-side configuration (e.g. blockDim)#define CUDA_WARP_SIZE_CONST 32// Threadblock configuration#define WARPS_PER_BLOCK 8// THREADS_PER_BLOCK must be evaluatable by host compiler for blockDim configuration#define THREADS_PER_BLOCK (WARPS_PER_BLOCK * CUDA_WARP_SIZE_CONST)// Macro-tile dimensions computed by a threadblock// BLOCK_M_TILES_WMMA * WMMA_M = output channels processed by a block// BLOCK_N_TILES_WMMA * WMMA_N = output spatial elements processed by a block#define BLOCK_M_TILES_WMMA 8#define BLOCK_N_TILES_WMMA 8#define TILE_M_PER_BLOCK (BLOCK_M_TILES_WMMA * WMMA_M) // e.g., 8 * 16 = 128 (for C_out dimension)#define TILE_N_PER_BLOCK (BLOCK_N_TILES_WMMA * WMMA_N) // e.g., 8 * 16 = 128 (for N_batch * H_out * W_out dimension)// Vector size for shared memory writes (half2)#define VECTOR_SIZE_H2 2// Struct to hold precomputed N-dimension GEMM indicesstruct NDecomposed {int ow_eff;int oh_eff;int n_batch_idx;bool isValidPixel; // True if this pixel_idx is within N_gemm boundsint h_in_base;int w_in_base;};__global__ void conv2d_implicit_gemm_wmma_kernel(const float* __restrict__ input_ptr, // Input: (N, Cin, Hin, Win)const float* __restrict__ weight_ptr, // Weights: (Cout, Cin, Kh, Kw)const float* __restrict__ bias_ptr, // Bias: (Cout) or nullptrfloat* __restrict__ output_ptr, // Output: (N, Cout, Hout, Wout)const int N_batch, const int C_in, const int H_in, const int W_in,const int C_out, const int K_h, const int K_w,const int stride_h, const int stride_w,const int pad_h, const int pad_w,const int H_out, const int W_out,const int M_gemm, // C_outconst int N_gemm, // N_batch * H_out * W_outconst int K_gemm // C_in * K_h * K_w) {// Thread identificationconst int warp_id = threadIdx.x / warpSize; // 0 .. WARPS_PER_BLOCK-1const int lane_id = threadIdx.x % warpSize; // 0 .. 31 (or warpSize-1)// Top-left corner of the macro-tile this block is responsible for in GEMM termsconst int block_row_gemm_start = TILE_M_PER_BLOCK * blockIdx.y;const int block_col_gemm_start = TILE_N_PER_BLOCK * blockIdx.x;// Shared memory for tiles of A (weights) and B (input/im2col) - Double Buffered for K-loop pipelining__shared__ half Asub_pipe[2][TILE_M_PER_BLOCK][WMMA_K + SKEW_HALF];__shared__ half Bsub_pipe[2][TILE_N_PER_BLOCK][WMMA_K + SKEW_HALF];// Shared memory for precomputed N-indices__shared__ NDecomposed n_params_sh[TILE_N_PER_BLOCK];// Shared memory for output stage (per-warp buffers)__shared__ float C_shmem_output_buffers[WARPS_PER_BLOCK][WMMA_M][WMMA_N];// Accumulator fragments per warp.wmma::fragment<wmma::accumulator, WMMA_M, WMMA_N, WMMA_K, float> acc_frag[BLOCK_N_TILES_WMMA];#pragma unrollfor (int i = 0; i < BLOCK_N_TILES_WMMA; ++i) {wmma::fill_fragment(acc_frag[i], 0.0f);}// Populate n_params_sh once at the beginning of the kernelif (threadIdx.x < TILE_N_PER_BLOCK) {int r_b_tile_idx = threadIdx.x;int current_pixel_idx = block_col_gemm_start + r_b_tile_idx;if (current_pixel_idx < N_gemm) {n_params_sh[r_b_tile_idx].ow_eff = current_pixel_idx % W_out;int temp_div_wout = current_pixel_idx / W_out;n_params_sh[r_b_tile_idx].oh_eff = temp_div_wout % H_out;n_params_sh[r_b_tile_idx].n_batch_idx = temp_div_wout / H_out;n_params_sh[r_b_tile_idx].isValidPixel = true;n_params_sh[r_b_tile_idx].h_in_base = n_params_sh[r_b_tile_idx].oh_eff * stride_h - pad_h;n_params_sh[r_b_tile_idx].w_in_base = n_params_sh[r_b_tile_idx].ow_eff * stride_w - pad_w;} else {n_params_sh[r_b_tile_idx].isValidPixel = false;n_params_sh[r_b_tile_idx].ow_eff = 0;n_params_sh[r_b_tile_idx].oh_eff = 0;n_params_sh[r_b_tile_idx].n_batch_idx = 0;n_params_sh[r_b_tile_idx].h_in_base = 0;n_params_sh[r_b_tile_idx].w_in_base = 0;}}__syncthreads();// Constants for vectorized shared memory loading// Number of half2 elements along K-dim for a shared memory tile rowconst int NUM_H2_ELEMENTS_IN_K_DIM = WMMA_K / VECTOR_SIZE_H2;// Number of thread groups, where each group has NUM_H2_ELEMENTS_IN_K_DIM threads.// Each group is responsible for loading the K-dimension for one M-row (for A) or N-row (for B) at a time,// iterating over M-rows or N-rows with this step size.const int NUM_ROW_PROCESSING_GROUPS = THREADS_PER_BLOCK / NUM_H2_ELEMENTS_IN_K_DIM;// --- K-Loop Pipelining ---int num_k_tiles = (K_gemm + WMMA_K - 1) / WMMA_K;

// --- Prologue: Load first k-tile (k_tile_iter = 0) into pipe_idx = 0 ---if (num_k_tiles > 0) {int k_tile_start_prologue = 0;int current_pipe_idx_prologue = 0;// Load Asub_pipe[0] for k_tile_iter = 0{// This thread is responsible for the 'h2_idx_in_k_dim_A'-th half2 element// in the K-dimension of the shared memory tile.int h2_idx_in_k_dim_A = threadIdx.x % NUM_H2_ELEMENTS_IN_K_DIM;// Starting 'half' index in shared memory for this half2 write.int shmem_k_start_for_h2_A = h2_idx_in_k_dim_A * VECTOR_SIZE_H2;// Global k-indices for the two half elements.int k_global_A_0 = k_tile_start_prologue + shmem_k_start_for_h2_A;int k_global_A_1 = k_tile_start_prologue + shmem_k_start_for_h2_A + 1;// Decompose k_global_A_0int kw_eff_reg_A_0 = 0, kh_eff_reg_A_0 = 0, ic_eff_reg_A_0 = 0;bool is_valid_k_A_0 = (k_global_A_0 < K_gemm);if (is_valid_k_A_0) {kw_eff_reg_A_0 = k_global_A_0 % K_w;int temp_div_kw_A_0 = k_global_A_0 / K_w;kh_eff_reg_A_0 = temp_div_kw_A_0 % K_h;ic_eff_reg_A_0 = temp_div_kw_A_0 / K_h;}// Decompose k_global_A_1int kw_eff_reg_A_1 = 0, kh_eff_reg_A_1 = 0, ic_eff_reg_A_1 = 0;bool is_valid_k_A_1 = (k_global_A_1 < K_gemm);if (is_valid_k_A_1) {kw_eff_reg_A_1 = k_global_A_1 % K_w;int temp_div_kw_A_1 = k_global_A_1 / K_w;kh_eff_reg_A_1 = temp_div_kw_A_1 % K_h;ic_eff_reg_A_1 = temp_div_kw_A_1 / K_h;}

// This thread belongs to 'm_row_group_id_A'-th group of threads.// This group iterates over M-rows of the Asub_pipe tile.int m_row_group_id_A = threadIdx.x / NUM_H2_ELEMENTS_IN_K_DIM;for (int r_a_tile_base = m_row_group_id_A; r_a_tile_base < TILE_M_PER_BLOCK; r_a_tile_base += NUM_ROW_PROCESSING_GROUPS) {int oc_idx = block_row_gemm_start + r_a_tile_base;float weight_val_0 = 0.0f;if (oc_idx < C_out && is_valid_k_A_0) {weight_val_0 = weight_ptr[oc_idx * C_in * K_h * K_w +ic_eff_reg_A_0 * K_h * K_w +kh_eff_reg_A_0 * K_w +kw_eff_reg_A_0];}float weight_val_1 = 0.0f;if (oc_idx < C_out && is_valid_k_A_1) {weight_val_1 = weight_ptr[oc_idx * C_in * K_h * K_w +ic_eff_reg_A_1 * K_h * K_w +kh_eff_reg_A_1 * K_w +kw_eff_reg_A_1];}half2* smem_ptr_h2_A = reinterpret_cast<half2*>(&Asub_pipe[current_pipe_idx_prologue][r_a_tile_base][shmem_k_start_for_h2_A]);*smem_ptr_h2_A = make_half2(__float2half(weight_val_0), __float2half(weight_val_1));}}// Load Bsub_pipe[0] for k_tile_iter = 0{int h2_idx_in_k_dim_B = threadIdx.x % NUM_H2_ELEMENTS_IN_K_DIM;int shmem_k_start_for_h2_B = h2_idx_in_k_dim_B * VECTOR_SIZE_H2;int k_global_B_0 = k_tile_start_prologue + shmem_k_start_for_h2_B;int k_global_B_1 = k_tile_start_prologue + shmem_k_start_for_h2_B + 1;int kw_eff_reg_B_0 = 0, kh_eff_reg_B_0 = 0, ic_eff_reg_B_0 = 0;bool is_valid_k_B_0 = (k_global_B_0 < K_gemm);if (is_valid_k_B_0) {kw_eff_reg_B_0 = k_global_B_0 % K_w;int temp_div_kw_B_0 = k_global_B_0 / K_w;kh_eff_reg_B_0 = temp_div_kw_B_0 % K_h;ic_eff_reg_B_0 = temp_div_kw_B_0 / K_h;}int kw_eff_reg_B_1 = 0, kh_eff_reg_B_1 = 0, ic_eff_reg_B_1 = 0;bool is_valid_k_B_1 = (k_global_B_1 < K_gemm);if (is_valid_k_B_1) {kw_eff_reg_B_1 = k_global_B_1 % K_w;int temp_div_kw_B_1 = k_global_B_1 / K_w;kh_eff_reg_B_1 = temp_div_kw_B_1 % K_h;ic_eff_reg_B_1 = temp_div_kw_B_1 / K_h;}int n_row_group_id_B = threadIdx.x / NUM_H2_ELEMENTS_IN_K_DIM;for (int r_b_tile_base = n_row_group_id_B; r_b_tile_base < TILE_N_PER_BLOCK; r_b_tile_base += NUM_ROW_PROCESSING_GROUPS) {float input_val_0 = 0.0f;if (n_params_sh[r_b_tile_base].isValidPixel && is_valid_k_B_0) {const NDecomposed& current_n_params = n_params_sh[r_b_tile_base];int h_in_eff_0 = current_n_params.h_in_base + kh_eff_reg_B_0;int w_in_eff_0 = current_n_params.w_in_base + kw_eff_reg_B_0;if (h_in_eff_0 >= 0 && h_in_eff_0 < H_in && w_in_eff_0 >= 0 && w_in_eff_0 < W_in) {input_val_0 = input_ptr[current_n_params.n_batch_idx * C_in * H_in * W_in +ic_eff_reg_B_0 * H_in * W_in +h_in_eff_0 * W_in +w_in_eff_0];}}float input_val_1 = 0.0f;if (n_params_sh[r_b_tile_base].isValidPixel && is_valid_k_B_1) {const NDecomposed& current_n_params = n_params_sh[r_b_tile_base];int h_in_eff_1 = current_n_params.h_in_base + kh_eff_reg_B_1;int w_in_eff_1 = current_n_params.w_in_base + kw_eff_reg_B_1;if (h_in_eff_1 >= 0 && h_in_eff_1 < H_in && w_in_eff_1 >= 0 && w_in_eff_1 < W_in) {input_val_1 = input_ptr[current_n_params.n_batch_idx * C_in * H_in * W_in +ic_eff_reg_B_1 * H_in * W_in +h_in_eff_1 * W_in +w_in_eff_1];}}half2* smem_ptr_h2_B = reinterpret_cast<half2*>(&Bsub_pipe[current_pipe_idx_prologue][r_b_tile_base][shmem_k_start_for_h2_B]);*smem_ptr_h2_B = make_half2(__float2half(input_val_0), __float2half(input_val_1));}}}// Loop over the K_gemm dimension in tiles of WMMA_Kfor (int k_tile_iter = 0; k_tile_iter < num_k_tiles; ++k_tile_iter) {__syncthreads(); // Sync point for pipeliningint compute_pipe_idx = k_tile_iter % 2;int load_pipe_idx = (k_tile_iter + 1) % 2;// --- Load Stage for next k-tile (k_tile_iter + 1) into load_pipe_idx ---int k_tile_start_for_load = (k_tile_iter + 1) * WMMA_K;if (k_tile_start_for_load < K_gemm) {// Load Asub_pipe[load_pipe_idx]{int h2_idx_in_k_dim_A = threadIdx.x % NUM_H2_ELEMENTS_IN_K_DIM;int shmem_k_start_for_h2_A = h2_idx_in_k_dim_A * VECTOR_SIZE_H2;int k_global_A_0 = k_tile_start_for_load + shmem_k_start_for_h2_A;int k_global_A_1 = k_tile_start_for_load + shmem_k_start_for_h2_A + 1;int kw_eff_reg_A_0 = 0, kh_eff_reg_A_0 = 0, ic_eff_reg_A_0 = 0;bool is_valid_k_A_0 = (k_global_A_0 < K_gemm);if (is_valid_k_A_0) {kw_eff_reg_A_0 = k_global_A_0 % K_w;int temp_div_kw_A_0 = k_global_A_0 / K_w;kh_eff_reg_A_0 = temp_div_kw_A_0 % K_h;ic_eff_reg_A_0 = temp_div_kw_A_0 / K_h;}int kw_eff_reg_A_1 = 0, kh_eff_reg_A_1 = 0, ic_eff_reg_A_1 = 0;bool is_valid_k_A_1 = (k_global_A_1 < K_gemm);if (is_valid_k_A_1) {kw_eff_reg_A_1 = k_global_A_1 % K_w;int temp_div_kw_A_1 = k_global_A_1 / K_w;kh_eff_reg_A_1 = temp_div_kw_A_1 % K_h;ic_eff_reg_A_1 = temp_div_kw_A_1 / K_h;}

int m_row_group_id_A = threadIdx.x / NUM_H2_ELEMENTS_IN_K_DIM;for (int r_a_tile_base = m_row_group_id_A; r_a_tile_base < TILE_M_PER_BLOCK; r_a_tile_base += NUM_ROW_PROCESSING_GROUPS) {int oc_idx = block_row_gemm_start + r_a_tile_base;float weight_val_0 = 0.0f;if (oc_idx < C_out && is_valid_k_A_0) {weight_val_0 = weight_ptr[oc_idx * C_in * K_h * K_w +ic_eff_reg_A_0 * K_h * K_w +kh_eff_reg_A_0 * K_w +kw_eff_reg_A_0];}float weight_val_1 = 0.0f;if (oc_idx < C_out && is_valid_k_A_1) {weight_val_1 = weight_ptr[oc_idx * C_in * K_h * K_w +ic_eff_reg_A_1 * K_h * K_w +kh_eff_reg_A_1 * K_w +kw_eff_reg_A_1];}half2* smem_ptr_h2_A = reinterpret_cast<half2*>(&Asub_pipe[load_pipe_idx][r_a_tile_base][shmem_k_start_for_h2_A]);*smem_ptr_h2_A = make_half2(__float2half(weight_val_0), __float2half(weight_val_1));}}// Load Bsub_pipe[load_pipe_idx]{int h2_idx_in_k_dim_B = threadIdx.x % NUM_H2_ELEMENTS_IN_K_DIM;int shmem_k_start_for_h2_B = h2_idx_in_k_dim_B * VECTOR_SIZE_H2;int k_global_B_0 = k_tile_start_for_load + shmem_k_start_for_h2_B;int k_global_B_1 = k_tile_start_for_load + shmem_k_start_for_h2_B + 1;int kw_eff_reg_B_0 = 0, kh_eff_reg_B_0 = 0, ic_eff_reg_B_0 = 0;bool is_valid_k_B_0 = (k_global_B_0 < K_gemm);if (is_valid_k_B_0) {kw_eff_reg_B_0 = k_global_B_0 % K_w;int temp_div_kw_B_0 = k_global_B_0 / K_w;kh_eff_reg_B_0 = temp_div_kw_B_0 % K_h;ic_eff_reg_B_0 = temp_div_kw_B_0 / K_h;}int kw_eff_reg_B_1 = 0, kh_eff_reg_B_1 = 0, ic_eff_reg_B_1 = 0;bool is_valid_k_B_1 = (k_global_B_1 < K_gemm);if (is_valid_k_B_1) {kw_eff_reg_B_1 = k_global_B_1 % K_w;int temp_div_kw_B_1 = k_global_B_1 / K_w;kh_eff_reg_B_1 = temp_div_kw_B_1 % K_h;ic_eff_reg_B_1 = temp_div_kw_B_1 / K_h;}int n_row_group_id_B = threadIdx.x / NUM_H2_ELEMENTS_IN_K_DIM;for (int r_b_tile_base = n_row_group_id_B; r_b_tile_base < TILE_N_PER_BLOCK; r_b_tile_base += NUM_ROW_PROCESSING_GROUPS) {float input_val_0 = 0.0f;if (n_params_sh[r_b_tile_base].isValidPixel && is_valid_k_B_0) {const NDecomposed& current_n_params = n_params_sh[r_b_tile_base];int h_in_eff_0 = current_n_params.h_in_base + kh_eff_reg_B_0;int w_in_eff_0 = current_n_params.w_in_base + kw_eff_reg_B_0;if (h_in_eff_0 >= 0 && h_in_eff_0 < H_in && w_in_eff_0 >= 0 && w_in_eff_0 < W_in) {input_val_0 = input_ptr[current_n_params.n_batch_idx * C_in * H_in * W_in +ic_eff_reg_B_0 * H_in * W_in +h_in_eff_0 * W_in +w_in_eff_0];}}float input_val_1 = 0.0f;if (n_params_sh[r_b_tile_base].isValidPixel && is_valid_k_B_1) {const NDecomposed& current_n_params = n_params_sh[r_b_tile_base];int h_in_eff_1 = current_n_params.h_in_base + kh_eff_reg_B_1;int w_in_eff_1 = current_n_params.w_in_base + kw_eff_reg_B_1;if (h_in_eff_1 >= 0 && h_in_eff_1 < H_in && w_in_eff_1 >= 0 && w_in_eff_1 < W_in) {input_val_1 = input_ptr[current_n_params.n_batch_idx * C_in * H_in * W_in +ic_eff_reg_B_1 * H_in * W_in +h_in_eff_1 * W_in +w_in_eff_1];}}half2* smem_ptr_h2_B = reinterpret_cast<half2*>(&Bsub_pipe[load_pipe_idx][r_b_tile_base][shmem_k_start_for_h2_B]);*smem_ptr_h2_B = make_half2(__float2half(input_val_0), __float2half(input_val_1));}}}// --- Compute Stage for current k-tile (k_tile_iter) using compute_pipe_idx ---int a_row_start_in_tile = warp_id * WMMA_M;wmma::fragment<wmma::matrix_a, WMMA_M, WMMA_N, WMMA_K, half, wmma::row_major> a_frag;wmma::load_matrix_sync(a_frag, &Asub_pipe[compute_pipe_idx][a_row_start_in_tile][0], WMMA_K + SKEW_HALF);wmma::fragment<wmma::matrix_b, WMMA_M, WMMA_N, WMMA_K, half, wmma::col_major> b_frag_inner_pipe[2];if (BLOCK_N_TILES_WMMA > 0) {int b_col_start_in_tile_current = 0 * WMMA_N;wmma::load_matrix_sync(b_frag_inner_pipe[0], &Bsub_pipe[compute_pipe_idx][b_col_start_in_tile_current][0], WMMA_K + SKEW_HALF);}

int current_inner_pipe_idx = 0;#pragma unrollfor (int n_tile = 0; n_tile < BLOCK_N_TILES_WMMA; ++n_tile) {int next_inner_pipe_idx = 1 - current_inner_pipe_idx;if (n_tile < BLOCK_N_TILES_WMMA - 1) {int b_col_start_in_tile_next = (n_tile + 1) * WMMA_N;wmma::load_matrix_sync(b_frag_inner_pipe[next_inner_pipe_idx], &Bsub_pipe[compute_pipe_idx][b_col_start_in_tile_next][0], WMMA_K + SKEW_HALF);}wmma::mma_sync(acc_frag[n_tile], a_frag, b_frag_inner_pipe[current_inner_pipe_idx], acc_frag[n_tile]);

current_inner_pipe_idx = next_inner_pipe_idx;}}__syncthreads();// Store results from accumulator fragments to global memory#pragma unrollfor (int n_tile = 0; n_tile < BLOCK_N_TILES_WMMA; ++n_tile) {wmma::store_matrix_sync(&C_shmem_output_buffers[warp_id][0][0], acc_frag[n_tile], WMMA_N, wmma::mem_row_major);for (int elem_idx_in_frag = lane_id; elem_idx_in_frag < WMMA_M * WMMA_N; elem_idx_in_frag += warpSize) {int r_frag = elem_idx_in_frag / WMMA_N;int c_frag = elem_idx_in_frag % WMMA_N;int oc_idx = block_row_gemm_start + (warp_id * WMMA_M) + r_frag;

int offset_in_block_N_processing = (n_tile * WMMA_N) + c_frag;if (oc_idx < C_out && offset_in_block_N_processing < TILE_N_PER_BLOCK &&n_params_sh[offset_in_block_N_processing].isValidPixel) {const NDecomposed& current_n_params = n_params_sh[offset_in_block_N_processing];int ow_eff = current_n_params.ow_eff;int oh_eff = current_n_params.oh_eff;int n_batch_idx = current_n_params.n_batch_idx;float val = C_shmem_output_buffers[warp_id][r_frag][c_frag];if (bias_ptr != nullptr) {val += bias_ptr[oc_idx];}output_ptr[n_batch_idx * C_out * H_out * W_out +oc_idx * H_out * W_out +oh_eff * W_out +ow_eff] = val;}}}}torch::Tensor conv2d_implicit_gemm_cuda(torch::Tensor input, torch::Tensor weight, torch::Tensor bias,int N_batch, int C_in, int H_in, int W_in,int C_out, int K_h, int K_w,int stride_h, int stride_w, int pad_h, int pad_w,int H_out, int W_out) {TORCH_CHECK(input.device().is_cuda(), "Input must be a CUDA tensor");TORCH_CHECK(weight.device().is_cuda(), "Weight must be a CUDA tensor");TORCH_CHECK(input.dtype() == torch::kFloat32, "Input must be float32");TORCH_CHECK(weight.dtype() == torch::kFloat32, "Weight must be float32");if (bias.defined()) {TORCH_CHECK(bias.device().is_cuda(), "Bias must be a CUDA tensor");TORCH_CHECK(bias.dtype() == torch::kFloat32, "Bias must be float32");TORCH_CHECK(bias.dim() == 1 && bias.size(0) == C_out, "Bias has wrong shape");}TORCH_CHECK(input.dim() == 4, "Input must be 4D");TORCH_CHECK(weight.dim() == 4, "Weight must be 4D");TORCH_CHECK(input.size(0) == N_batch, "Input N_batch mismatch");TORCH_CHECK(input.size(1) == C_in, "Input C_in mismatch");TORCH_CHECK(input.size(2) == H_in, "Input H_in mismatch");TORCH_CHECK(input.size(3) == W_in, "Input W_in mismatch");TORCH_CHECK(weight.size(0) == C_out, "Weight C_out mismatch");TORCH_CHECK(weight.size(1) == C_in, "Weight C_in mismatch");TORCH_CHECK(weight.size(2) == K_h, "Weight K_h mismatch");TORCH_CHECK(weight.size(3) == K_w, "Weight K_w mismatch");auto output = torch::zeros({N_batch, C_out, H_out, W_out}, input.options());const int M_gemm = C_out;const int N_gemm = N_batch * H_out * W_out;const int K_gemm = C_in * K_h * K_w;if (M_gemm == 0 || N_gemm == 0) {return output;}if (K_gemm == 0) {if (bias.defined()) {output = output + bias.reshape({1, C_out, 1, 1});}return output;}dim3 block_dim(THREADS_PER_BLOCK);dim3 grid_dim((N_gemm + TILE_N_PER_BLOCK - 1) / TILE_N_PER_BLOCK,(M_gemm + TILE_M_PER_BLOCK - 1) / TILE_M_PER_BLOCK);const float* bias_ptr_data = bias.defined() ? bias.data_ptr<float>() : nullptr;cudaStream_t stream = at::cuda::getCurrentCUDAStream();conv2d_implicit_gemm_wmma_kernel<<<grid_dim, block_dim, 0, stream>>>(input.data_ptr<float>(),weight.data_ptr<float>(),bias_ptr_data,output.data_ptr<float>(),N_batch, C_in, H_in, W_in,C_out, K_h, K_w,stride_h, stride_w, pad_h, pad_w,H_out, W_out,M_gemm, N_gemm, K_gemm);

AT_CUDA_CHECK(cudaGetLastError());return output;}"""conv2d_implicit_gemm_cuda_declaration = r"""torch::Tensor conv2d_implicit_gemm_cuda(torch::Tensor input, torch::Tensor weight, torch::Tensor bias,int N_batch, int C_in, int H_in, int W_in,int C_out, int K_h, int K_w,int stride_h, int stride_w, int pad_h, int pad_w,int H_out, int W_out);"""# JIT compile the CUDA kernelcustom_conv2d_wmma_ops = load_inline(name="custom_conv2d_wmma_ops_optimized_k_pipe_vec_smem", # Changed name to avoid collisioncpp_sources=conv2d_implicit_gemm_cuda_declaration,cuda_sources=conv2d_implicit_gemm_cuda_source,functions=["conv2d_implicit_gemm_cuda"],verbose=True,extra_cuda_cflags=["-arch=sm_70", "--use_fast_math", "-std=c++17"])class ModelNew(nn.Module):def __init__(self, num_classes=1000): # num_classes is part of original signature, kept for consistencysuper(ModelNew, self).__init__()

# Define Conv1 parameters (matching the original model)self.in_channels = 3self.out_channels = 96self.kernel_size_val = 11 # Assuming square kernelself.stride_val = 4 # Assuming square strideself.padding_val = 2 # Assuming square padding# Create a temporary Conv2d layer to initialize weights and biastemp_conv = nn.Conv2d(in_channels=self.in_channels,out_channels=self.out_channels,kernel_size=self.kernel_size_val,stride=self.stride_val,padding=self.padding_val,bias=True # nn.Conv2d has bias=True by default)self.conv1_weight = nn.Parameter(temp_conv.weight.detach().clone())if temp_conv.bias is not None:self.conv1_bias = nn.Parameter(temp_conv.bias.detach().clone())else:# Correctly register 'conv1_bias' as None if not presentself.register_parameter('conv1_bias', None)self.custom_conv_op = custom_conv2d_wmma_ops.conv2d_implicit_gemm_cudadef forward(self, x):N_batch = x.size(0)# C_in_runtime = x.size(1) # Should match self.in_channelsH_in = x.size(2)W_in = x.size(3)# Calculate output dimensionsH_out = (H_in + 2 * self.padding_val - self.kernel_size_val) // self.stride_val + 1W_out = (W_in + 2 * self.padding_val - self.kernel_size_val) // self.stride_val + 1

# Bias tensor handling: pass an undefined tensor if bias is None.# The C++ TORCH_CHECK(bias.defined()) handles this by providing nullptr to kernel.bias_tensor = self.conv1_bias if self.conv1_bias is not None else torch.Tensor()x = self.custom_conv_op(x, self.conv1_weight, bias_tensor,N_batch, self.in_channels, H_in, W_in,self.out_channels, self.kernel_size_val, self.kernel_size_val, # K_h, K_wself.stride_val, self.stride_val, # stride_h, stride_wself.padding_val, self.padding_val, # pad_h, pad_wH_out, W_out)return x

參考資料：

https://crfm.stanford.edu/2025/05/28/fast-kernels.html

https://news.ycombinator.com/item?id=44139454