斯坦福华人天团意外爆冷！AI用纯CUDA-C编内核，竟干翻PyTorch？

新智元2025-05-31

【新智元导读】本想练练手合成点数据，没想到却一不小心干翻了PyTorch专家内核！斯坦福华人团队用纯CUDA-C写出的AI生成内核，瞬间惊艳圈内并登上Hacker News热榜。团队甚至表示：本来不想发这个结果的。就在刚刚，斯坦福HAI华人大神团队又出惊人神作了。他们用纯CUDA-C语言编写的快速AI生成内核，竟然超越了PyTorch！在这个过程中，完全不用借助CUTLASS和Triton等库和...

网页链接

免责声明：本文观点仅代表作者个人观点，不构成本平台的投资建议，本平台不对文章信息准确性、完整性和及时性做出任何保证，亦不对因使用或信赖文章信息引发的任何损失承担责任。

精彩评论

我们需要你的真知灼见来填补这片空白

发表看法

{"i18n":{"language":"zh_CN"},"isChannel":false,"data":{"share":"https://www.laohu8.com/m/news/2539520018?lang=zh_CN&edition=full","thumbnail":"","is_english":false,"pubTime":"2025-05-31 13:46","share_image_url":"https://static.laohu8.com/e9f99090a1c2ed51c021029395664489","id":"2539520018","market":"sh","top_or_hot":-1,"title":"斯坦福华人天团意外爆冷！AI用纯CUDA-C编内核，竟干翻PyTorch？","media":"新智元","content":"<div>\n<p>【新智元导读】本想练练手合成点数据，没想到却一不小心干翻了PyTorch专家内核！斯坦福华人团队用纯CUDA-C写出的AI生成内核，瞬间惊艳圈内并登上Hacker News热榜。团队甚至表示：本来不想发这个结果的。就在刚刚，斯坦福HAI华人大神团队又出惊人神作了。他们用纯CUDA-C语言编写的快速AI生成内核，竟然超越了PyTorch！在这个过程中，完全不用借助CUTLASS和Triton等库和...</p>\n\n<a href=\"https://tech.ifeng.com/c/8jnncyGWJIc\">网页链接</a>\n\n</div>\n","source":"fenghuang_stock","html":"<!DOCTYPE html>\n<html>\n<head>\n<meta http-equiv=\"Content-Type\" content=\"text/html; charset=utf-8\" />\n<meta name=\"viewport\" content=\"width=device-width,initial-scale=1.0,minimum-scale=1.0,maximum-scale=1.0,user-scalable=no\"/>\n<meta name=\"format-detection\" content=\"telephone=no,email=no,address=no\" />\n<title>斯坦福华人天团意外爆冷！AI用纯CUDA-C编内核，竟干翻PyTorch？</title>\n<style type=\"text/css\">\na,abbr,acronym,address,applet,article,aside,audio,b,big,blockquote,body,canvas,caption,center,cite,code,dd,del,details,dfn,div,dl,dt,\nem,embed,fieldset,figcaption,figure,footer,form,h1,h2,h3,h4,h5,h6,header,hgroup,html,i,iframe,img,ins,kbd,label,legend,li,mark,menu,nav,\nobject,ol,output,p,pre,q,ruby,s,samp,section,small,span,strike,strong,sub,summary,sup,table,tbody,td,tfoot,th,thead,time,tr,tt,u,ul,var,video{ font:inherit;margin:0;padding:0;vertical-align:baseline;border:0 }\nbody{ font-size:16px; line-height:1.5; color:#999; background:transparent; }\n.wrapper{ overflow:hidden;word-break:break-all;padding:10px; }\nh1,h2{ font-weight:normal; line-height:1.35; margin-bottom:.6em; }\nh3,h4,h5,h6{ line-height:1.35; margin-bottom:1em; }\nh1{ font-size:24px; }\nh2{ font-size:20px; }\nh3{ font-size:18px; }\nh4{ font-size:16px; }\nh5{ font-size:14px; }\nh6{ font-size:12px; }\np,ul,ol,blockquote,dl,table{ margin:1.2em 0; }\nul,ol{ margin-left:2em; }\nul{ list-style:disc; }\nol{ list-style:decimal; }\nli,li p{ margin:10px 0;}\nimg{ max-width:100%;display:block;margin:0 auto 1em; }\nblockquote{ color:#B5B2B1; border-left:3px solid #aaa; padding:1em; }\nstrong,b{font-weight:bold;}\nem,i{font-style:italic;}\ntable{ width:100%;border-collapse:collapse;border-spacing:1px;margin:1em 0;font-size:.9em; }\nth,td{ padding:5px;text-align:left;border:1px solid #aaa; }\nth{ font-weight:bold;background:#5d5d5d; }\n.symbol-link{font-weight:bold;}\n/* header{ border-bottom:1px solid #494756; } */\n.title{ margin:0 0 8px;line-height:1.3;color:#ddd; }\n.meta {color:#5e5c6d;font-size:13px;margin:0 0 .5em; }\na{text-decoration:none; color:#2a4b87;}\n.meta .head { display: inline-block; overflow: hidden}\n.head .h-thumb { width: 30px; height: 30px; margin: 0; padding: 0; border-radius: 50%; float: left;}\n.head .h-content { margin: 0; padding: 0 0 0 9px; float: left;}\n.head .h-name {font-size: 13px; color: #eee; margin: 0;}\n.head .h-time {font-size: 11px; color: #7E829C; margin: 0;line-height: 11px;}\n.small {font-size: 12.5px; display: inline-block; transform: scale(0.9); -webkit-transform: scale(0.9); transform-origin: left; -webkit-transform-origin: left;}\n.smaller {font-size: 12.5px; display: inline-block; transform: scale(0.8); -webkit-transform: scale(0.8); transform-origin: left; -webkit-transform-origin: left;}\n.bt-text {font-size: 12px;margin: 1.5em 0 0 0}\n.bt-text p {margin: 0}\n</style>\n</head>\n<body>\n<div class=\"wrapper\">\n<header>\n<h2 class=\"title\">\n斯坦福华人天团意外爆冷！AI用纯CUDA-C编内核，竟干翻PyTorch？\n</h2>\n\n<h4 class=\"meta\">\n\n\n2025-05-31 13:46 北京时间&nbsp;&nbsp;&nbsp;<a href=https://tech.ifeng.com/c/8jnncyGWJIc><strong>新智元</strong></a>\n\n\n</h4>\n\n</header>\n<article>\n<div>\n<p>【新智元导读】本想练练手合成点数据，没想到却一不小心干翻了PyTorch专家内核！斯坦福华人团队用纯CUDA-C写出的AI生成内核，瞬间惊艳圈内并登上Hacker News热榜。团队甚至表示：本来不想发这个结果的。就在刚刚，斯坦福HAI华人大神团队又出惊人神作了。他们用纯CUDA-C语言编写的快速AI生成内核，竟然超越了PyTorch！在这个过程中，完全不用借助CUTLASS和Triton等库和...</p>\n\n<a href=\"https://tech.ifeng.com/c/8jnncyGWJIc\">网页链接</a>\n\n</div>\n\n\n</article>\n</div>\n</body>\n</html>\n","isBrief":false,"type":0,"news_type":1,"symbol":"SM","symbol_name":"SM Energy","start_time":0,"source_url":"https://tech.ifeng.com/c/8jnncyGWJIc","article_id":"2539520018","we_media_id":null,"thumbnails":[],"rights":null,"url":"https://stock-news.laohu8.com/highlight/detail?id=2539520018","pubTimestamp":1748670397,"columns":[],"sourceInfo":{"source_id":"fenghuang_stock","name":"凤凰网"},"weMediaInfo":null,"summary":"本想练练手合成点数据，没想到却一不小心干翻了PyTorch专家内核！斯坦福华人团队用纯CUDA-C写出的AI生成内核，瞬间惊艳圈内并登上Hacker News热榜。他们用纯CUDA-C语言编写的快速AI生成内核，竟然超越了PyTorch！然而，意想不到的事情发生了，仅用于测试的合成数据生成本身，竟开始生成非常优秀的内核，甚至超越了人类专家优化的PyTorch基线，而且还利用了高级优化和硬件特性。比如为什么使用FP32内核会比PyTorch更容易实现性能提升，理由就相当有趣。","collect":0,"end_time":0,"defaultTopTitle":"ifeng.com","property":[],"viewcount":null,"language":"zh","relate_stocks":{"SM":"SM Energy","BK4213":"石油与天然气的勘探与生产","DSL":"DoubleLine Income Solutions Fund"},"translate_title":"Stanford Chinese group was unexpectedly upset! AI uses pure CUDA-C to compile the kernel, but it actually defeats PyTorch?","themeId":null,"isJumpTheme":false,"ttsUrl":null,"symbols_score_info":{"DSL":1,"SM":1},"content_text":"【新智元导读】本想练练手合成点数据，没想到却一不小心干翻了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 OuyangAnne Ouyang目前是斯坦福大学计算机科学（CS）博士生，在Scaling Intelligence Lab（可扩展智能实验室）进行研究。她的研究兴趣主要集中在可扩展的自我改进机器学习系统，同时也广泛关注实证机器学习（empirical ML）和性能工程（performance engineering）。此前，她在MIT获得学士和硕士学位，并曾在NVIDIA cuDNN团队工作，负责编写CUDA内核，用于加速GPU上的深度学习工作负载。Azalia MirhoseiniAzalia Mirhoseini是斯坦福大学计算机科学助理教授，也是Scaling Intelligence Lab（可扩展智能实验室）的创始人，并在Google DeepMind兼任高级研究科学家。她的实验室致力于开发可扩展的自主演进人工智能系统与方法论，以期推动通用人工智能的发展。在加入斯坦福大学之前，她曾在Google Brain和Anthropic等业界顶尖的人工智能实验室工作多年。她过往的卓越成就包括：提出混合专家（MoE）神经架构——目前已被前沿的AI模型广泛应用；领导AlphaChip项目——一项将深度强化学习用于布局优化的开创性工作，并成功应用于谷歌AI加速器（TPU）及数据中心CPU等先进芯片的设计中；在测试时计算的Scaling方面有深入的研究Percy LiangPercy 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.htmlhttps://news.ycombinator.com/item?id=44139454","kind":"news","is_publish_news":true,"is_publish_highlight":false,"is_publish_live":false,"is_publish_wemedia":null,"editions":null,"column":"","sentiment":"0","news_tag":"","news_rank":0,"isVideo":false,"video":null,"symbols":[],"gpt_button":1,"need_auth":false,"need_login_tip":false,"code":"91000000","status":"200"},"commentList":[],"isCommentEnd":true,"newsSizeData":{"likeSize":0,"commentSize":0,"repostSize":0,"favoriteSize":0,"likeStatus":false,"favoriteStatus":false},"APP":{"userAgent":"Mozilla/5.0 AppleWebKit/537.36 (KHTML, like Gecko; compatible; ClaudeBot/1.0; +claudebot@anthropic.com)","isDev":false,"isTTM":false,"tenantId":"TBCN","deviceId":"web-server-community-laohu8-v3","version":"4.43.1","shortVersion":"4.43.1","platform":"web","vendor":"web","appName":"laohu8","isIOS":false,"isAndroid":false,"isTiger":false,"isTHS":false,"isWeiXin":false,"isWeiXinMini":false,"isWeiBo":false,"isQQ":false,"isBaiduSwan":false,"isBaiduBox":false,"isDingTalk":false,"isToutiao":false,"isOnePlus":false,"isHuaWei":false,"isXiaomi":false,"isXiaomiWebView":false,"isOppo":false,"isVivo":false,"isSamsung":false,"isMobile":false},"href":"/m/news/2539520018","isCrawlerRequest":true}