斯坦福华人天团意外爆冷！AI用纯CUDA-C编内核，竟干翻PyTorch？
新智元报道
编辑：Aeneas 好困
【新智元导读】本想练练手合成点数据，没想到却一不小心干翻了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）所强调的，我们并不总是需要大规模的重新训练。
有时，巧妙的搜索和分支策略便足以催生科学创新、攻克复杂难题，而借助验证器进行广泛搜索，则可能带来更大的收益。
然而，这并不意味着我们不需要进一步的训练。
恰恰相反，研究者的这种方法，也有助于生成更优质的合成数据，用以改进未来的模型训练（这需要更多的问题实例）。
因此，它既是一种强大的测试时扩展方法，也是我们迈向更智能、数据效率更高的模型开发之路的一步。
而且，这次研究者展现的仅仅是初步的成果。这些优化结果的质量看起来相当可观，但仍有广阔的提升空间，例如产生更优的优化思路、生成更高质量的最终代码，以及将此方法应用于日益复杂的内核。
目前，研究者仍在积极改进的两个具体例子包括：
- 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 torch import torch.nn as nn import torch.nn.functional as F from torch.utils.cpp_extension import load_inline conv2d_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 indices struct NDecomposed { int ow_eff; int oh_eff; int n_batch_idx; bool isValidPixel; // True if this pixel_idx is within N_gemm bounds int 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 nullptr float* __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_out const int N_gemm, // N_batch * H_out * W_out const int K_gemm // C_in * K_h * K_w ) { // Thread identification const int warp_id = threadIdx.x / warpSize; // 0 .. WARPS_PER_BLOCK-1 const 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 terms const 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 unroll for (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 kernel if (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 row const 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_0 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; } // Decompose k_global_A_1 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; } // 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_K for (int k_tile_iter = 0; k_tile_iter < num_k_tiles; ++k_tile_iter) { __syncthreads(); // Sync point for pipelining int 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 unroll for (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 unroll for (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 kernel custom_conv2d_wmma_ops = load_inline( name="custom_conv2d_wmma_ops_optimized_k_pipe_vec_smem", # Changed name to avoid collision cpp_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 consistency super(ModelNew, self).__init__() # Define Conv1 parameters (matching the original model) self.in_channels = 3 self.out_channels = 96 self.kernel_size_val = 11 # Assuming square kernel self.stride_val = 4 # Assuming square stride self.padding_val = 2 # Assuming square padding # Create a temporary Conv2d layer to initialize weights and bias temp_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 present self.register_parameter('conv1_bias', None) self.custom_conv_op = custom_conv2d_wmma_ops.conv2d_implicit_gemm_cuda def forward(self, x): N_batch = x.size(0) # C_in_runtime = x.size(1) # Should match self.in_channels H_in = x.size(2) W_in = x.size(3) # Calculate output dimensions H_out = (H_in + 2 * self.padding_val - self.kernel_size_val) // self.stride_val + 1 W_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_w self.stride_val, self.stride_val, # stride_h, stride_w self.padding_val, self.padding_val, # pad_h, pad_w H_out, W_out ) return x
（文：新智元）