Compare commits

..

2 Commits
main ... test5

Author SHA1 Message Date
Kuohais 9189819772 复现可能存在的问题 2025-11-17 12:23:36 +08:00
Kuohais 06e3b1680c test CI 2025-11-17 11:59:20 +08:00
2494 changed files with 207 additions and 188892 deletions

View File

@ -1,165 +0,0 @@
# GPUCodeForces代码解读
## 一、赛题核心定位与整体框架
本赛题属于GPU CUDA 性能优化类任务,要求参赛选手通过自定义 CUDA Kernel 实现各类函数,并与 PyTorch 内置实现进行精度对齐和性能比拼。
这份赛题解读我们以“通过自定义CUDA Kernel实现Swish激活函数”来作为引子让大家从一个具体例子中了解到算子优化的细节。
整套关联代码example_torchcode.py、example_cudacode.py、run_code.py构成了 “任务定义 - 基准实现 - CUDA 优化 - 评测验证” 的完整闭环prompt.txt则提供了类似 “融合算子 CUDA 设计” 的 prompt 编写思路,可作为加分项参考。其核心目标是考察选手的 CUDA 内核设计能力、内存效率优化能力及精度与性能的平衡能力。
## 二、赛题模块拆解与代码解读
### (一)模块 1任务定义与基准实现example_torchcode.py
该文件是赛题的 “基础参照系”,定义了任务边界、输入数据生成规则和标准 GTGround Truth输出对应评测数据集要求中的 “数据集样本描述”“输入数据生成函数”“标准 GT 输出生函数”。
1. 任务定义Swish 激活函数计算
Swish 是深度学习中优于 ReLU 的激活函数数学表达式为Swish(x) = x * sigmoid(x)。其中sigmoid(x) = 1 / (1 + exp(-x))核心作用是为模型引入非线性且在大维度张量如隐藏层特征上的计算效率直接影响模型整体推理速度。代码中通过Model类实现基准逻辑
```
class Model(nn.Module):
def forward(self, x: torch.Tensor) -> torch.Tensor:
return x * torch.sigmoid(x) # 标准GT实现
```
2. 输入数据生成函数
赛题定义了固定的输入规模模拟大模型隐藏层维度具有实际业务代表性函数get_inputs()生成符合任务需求的输入张量:
```
batch_size = 16 # 批量大小
dim = 16384 # 单样本特征维度(模拟大模型隐藏层)
def get_inputs():
x = torch.randn(batch_size, dim) # 随机正态分布张量(符合深度学习输入特性)
return [x]
```
输入特性torch.randn生成均值为 0、方差为 1 的浮点数张量,覆盖正负值,能全面验证 Swish 在不同输入下的计算精度;
规模选择16×16384的张量共 262,144 个元素,既能暴露 CUDA 线程调度的效率差异,又不会因规模过大导致测试耗时过长。
3. 初始化输入函数
get_init_inputs()返回空列表,因 Swish 激活函数的计算无需额外初始化参数(如权重、偏置),简化了模型初始化逻辑,聚焦核心的激活函数计算。
### (二)模块 2CUDA 优化实现example_cudacode.py
该文件是参赛选手的 “核心提交目标”,对应评测数据集要求中的 “CUDA 解决方案实现”,需通过自定义 CUDA Kernel 实现 Swish 函数,同时满足精度和性能要求。
1. 实现思路:轻量级高效 Kernel 设计
Swish 是逐元素操作(每个输出元素仅依赖对应输入元素),无需跨元素通信,因此 Kernel 设计聚焦 “线程调度效率” 和 “浮点计算稳定性”:
线程索引:采用 1D 索引idx = blockIdx.x * blockDim.x + threadIdx.x每个线程处理 1 个元素,避免复杂的 2D/3D 索引计算开销;
Block/Grid 配置block_size=256CUDA 硬件的经典配置1 个 Block 包含 256 个线程,可完美适配 GPU 的 Warp 调度(每个 Warp 含 32 个线程256=8×32num_blocks通过 “向上取整” 计算((size + block_size - 1) / block_size确保覆盖所有元素
浮点计算使用expf()单精度浮点数指数函数与输入张量的float32类型匹配避免精度浪费和类型转换开销。
2. 代码编译与封装
通过torch.utils.cpp_extension.load_inline实现内联 CUDA 代码编译无需单独编写setup.py简化开发流程
```
swish = load_inline(
name="swish", # 模块名
cpp_sources=swish_cpp_source, # C++声明(接口定义)
cuda_sources=swish_source, # CUDA内核实现
functions=["swish_cuda"], # 暴露给Python的函数
verbose=True # 打印编译日志(便于调试))
```
并通过ModelNew类封装 CUDA 函数保持与example_torchcode.py中Model类一致的接口forward方法确保后续评测代码可无缝调用。
### (三)模块 3精度与性能评测run_code.py
该文件是赛题的 “评测执行器”,对应评测数据集要求中的 “性能评估指标” 和 “正确性验证”,实现了从数据准备到结果分析的全流程自动化评测。
1. 评测前置准备
CUDA 可用性检查先判断torch.cuda.is_available(),避免无 GPU 环境下的报错;
数据与模型迁移:将输入张量和模型均移动到 GPUcuda(device=device)),确保计算在 GPU 上执行;
GPU 预热:先执行 10 次空计算_ = torch_model(*inputs)),避免 GPU 初始化、内存分配等一次性开销影响性能计时精度。
2. 精度对齐验证(核心指标)
精度是 CUDA 实现的 “准入条件”,需确保自定义 Kernel 与 PyTorch 基准的输出误差在可接受范围:
```
abs_diff = torch.abs(output_torch - output_cuda) # 逐元素绝对误差
max_diff = torch.max(abs_diff).item() # 最大误差(全局)
mean_diff = torch.mean(abs_diff).item() # 平均误差(全局)
```
误差阈值max_diff < 1e-4且mean_diff < 1e-5符合深度学习中浮点计算的常见容忍度单精度浮点数的机器 epsilon 约为 1e-7该阈值留有充足余量
结果判定:若满足阈值则 “精度对齐”,否则视为无效实现,无法进入性能评测环节。
3. 性能加速比测试(核心指标)
性能评测聚焦 “平均执行时间” 和 “加速比”,对应评测数据集要求的 “执行时间”“吞吐量” 指标迭代次数100 次减少随机波动确保计时稳定性同步计时使用torch.cuda.synchronize()强制等待 GPU 计算完成,避免 CUDA 异步执行导致的计时偏差;结果计算:平均时间 = 总时间 / 迭代次数(消除单次执行的偶然误差);加速比 = PyTorch 平均时间 / CUDA 平均时间比值越大CUDA 优化效果越好)。
4. 评测输出示例
```
-------------------- 精度对齐验证 --------------------
✅ 精度对齐:最大误差 0.000089,平均误差 0.000012
-------------------- 性能加速比测试 --------------------
PyTorch内置Swish平均执行时间: 0.000123秒
自定义CUDA Swish平均执行时间: 0.000045秒
加速比 (Speedup): 2.73x
```
### (四)模块 4Prompt 设计参考prompt.txt
针对 “矩阵乘法 + GELU 融合算子”,提供了赛题 “加分项 Prompt” 的设计思路,对应评测数据集要求的 “加分项LLM 生成 CUDA 代码的 Prompt”。
1. Prompt 设计核心要素
任务拆解:明确 “先矩阵乘、后 GELU” 的原始流程,指出融合的必要性(避免中间结果的全局内存读写);
技术要求指定关键优化点2D Grid/Block、共享内存 tiling引导 LLM 生成符合 CUDA 最佳实践的代码;
精度约束:强调 “数值稳定性和精度”,避免 LLM 为追求性能牺牲精度。
2. 迁移应用到 Swish 任务
若为 Swish 设计 Prompt可参考如下结构
```
Write a custom CUDA kernel for Swish activation (Swish(x) = x * sigmoid(x)).
The original PyTorch implementation uses x * torch.sigmoid(x), which may have redundant global memory access.
You should optimize the CUDA kernel to:
- Use 1D grid/block dimensions (each thread processes one element)
- Choose appropriate block size (e.g., 256) for GPU warp scheduling
- Ensure numerical stability (use float32 and expf() for sigmoid)
The input is a PyTorch tensor of shape (batch_size=16, dim=16384), and the output should match PyTorch's result with max error < 1e-4.
```
## 三、赛题核心考察点
CUDA Kernel 设计能力线程索引计算、Block/Grid 配置合理性如block_size=256的选择依据
精度控制能力浮点计算稳定性如匹配float32类型、避免exp()溢出);
性能优化意识GPU 预热、同步计时、减少冗余内存访问;
工程化实现能力CUDA 代码与 PyTorch 的接口兼容如ModelNew类的封装、编译调试能力。
## 四、代码间逻辑流与参赛指引
1.代码间逻辑流程图
<img src="./images/code_step.png">
2.参赛选手操作指引
参考example_torchcode.py理解任务边界输入规模、GT 输出);
编写自定义 CUDA Kernel可借鉴example_cudacode.py的结构优化 Block/Grid 或浮点计算);
使用run_code.py验证精度需满足误差阈值和性能追求更高加速比
(加分项)设计 Prompt让 LLM 生成你的 CUDA 代码,并对比 LLM 生成结果与手写结果的差异。
## 五、关键注意事项
精度优先:若 CUDA 实现性能极高但精度不达标,视为无效提交;
数据类型统一输入张量、CUDA 计算均使用float32避免float64导致的性能下降
计时准确性必须使用torch.cuda.synchronize(),否则异步执行会导致计时结果偏小(虚假性能提升);
可扩展性若输入规模变化如dim=32768需确保num_blocks计算逻辑仍能覆盖所有元素。

View File

@ -1,88 +0,0 @@
# GPUCodeForces赛题入门
## 🚀 一、背景介绍
在 AI 模型训练与推理的世界里,“性能”就是生产力。
不同品牌、架构的 GPU在不同深度学习框架中运行同样的任务时往往会出现令人惊讶的差异
- 有的显卡在 TensorFlow 上快如闪电,在 PyTorch 上却慢半拍;
- 有的框架吞吐惊人,但显存消耗巨大;
- 有的 GPU 一旦切换精度或 batch size性能瞬间翻倍。
这些性能差异长期存在,却缺乏一个「公正、标准、可复现」的评测体系。
于是,我们发起了这场挑战赛:
用社区的力量,一起构建一个开放的 GPU 性能评测数据集,让每块 GPU、每个框架都能有真实可对比的「成绩单」。
## 🎯 二、赛题目标
本次挑战的目标是:
1. 收集并生成评测样本 —— 从主流框架PyTorch、PaddlePaddle、TensorFlow、JAX、MMCV、Transformers 等中提取典型任务如分类、检测、生成、NLP 推理等);
2. 为每个样本生成标准输出与性能指标 —— 包括运行时间、精度对齐结果、加速比(实际评估指标远不止这些,此为本次比赛评测范围);
3. 形成可复现的 GPU CodeForces 数据集 —— 用统一格式记录代码、框架、显卡型号与性能结果,让 GPU 性能比较更科学、更透明。
## ⚙️ 三、为什么要做这件事?
•💬 统一标准:不同平台、不同框架的性能指标不再“各说各话”;
•🔍 科学对比:用户可以真正知道哪种框架组合最适合自己的硬件;
•🧩 数据复现:所有样本都能被社区成员在相同条件下复现;
•🏆 公开榜单:最终形成一个持续更新的 GPU 性能排行榜。
## 🧩 四、选手的任务内容
参赛者需要:
1. 从指定或自选框架中挑选典型任务如图像分类、Transformer 推理等);
2. 编写或复用性能测试脚本,确保能在不同 GPU 上运行;
3. 收集运行结果(速度、显存占用、吞吐量等),并输出统一格式的数据样本;
4. 将样本上传至平台形成标准化的「GPU CodeForces 数据集」。
简单来说你的任务就是让「GPU 跑起来、测出来、比起来」。
## 🧮 五、第一步怎么开始
### 别慌你不需要搭一个复杂的AI实验室。
你可能会想“要有一块能跑深度学习的GPU再装好一个主流框架才能开始”但实际上这一步也能为你省去。这里[模力方舟](https://ai.gitee.com/compute)就有已经配置好所需环境且带有强劲GPU的云算力服务器本次比赛「免费使用」。
### 代码也不需要从零开始写
在仓库中我们提供了几个基础脚本自己看或者让AI帮你分析都可以~你只需要改一改,跑一跑,就能提交结果。
启动比赛提供的云算力平台后,在终端中输入以下指令:
```
git clone https://gitlink.org.cn/ccf-ai-infra/GPUCodeForces.git
cd example/001-example
python run_code.py
```
这段代码会自动跑一个小模型并输出结果:
```
Loading extension module relu...
-------------------- 精度对齐验证 --------------------
✅ 精度对齐:两个模型的输出结果非常接近。
-------------------- 性能加速比测试 --------------------
PyTorch torch.relu 平均执行时间: 0.000006 秒
自定义 CUDA 内核 平均执行时间: 0.000017 秒
加速比 (Speedup): 0.35x
```
恭喜🎊这就代表你已经跑了你的第一个性能样本。后续,无论你是「调参、改算法」,只要能跑出来并测试通过,那么就算一次成功的提交!
### 想进阶?你可以这样玩
如果你足够熟悉Python可以进一步
- 修改batch size、模型结构、输入分辨率等等看性能变化
- 换用不同框架如TensorFlow、Paddle对比结果
- 充分规划你代码内各子任务计算任务资源使用将GPU性能最大化发挥
👉你不需要像写论文那样创一个新模型,也不用理解优化的底层逻辑,只要能跑就能贡献数据。
👉你也可以完全有自己的想法,认为有一些模型在你的思路下可以更优化,那么可以放到评测体系中,看看结果到底怎么样,同样能跑通就能贡献数据。
无论你选择哪种方式,我们都有相对公平的评分规则,你可以尽情跑、尽情想,怎样都行得通,这就是这个挑战的魅力所在。
## 📧加入社区,边学边玩
加入官方交流群[下方二维码](https://www.gitlink.org.cn/zone/Infra/newdetail/1075)或者在对应赛道下的Issue板块与我们积极交流
- 问题有人答
- 每周都有榜单更新
- 还能看到别人分享的“加速黑科技”
我们希望这场挑战不仅是测评,更是一次:
🔧「全民GPU训练营」——让每个人都能轻松了解性能优化的乐趣。

174
README.md
View File

@ -17,119 +17,121 @@
* 为每个样本提供**标准输出**和**性能指标**,确保结果可复现。
* 最终形成 **GPU CodeForces** 数据集和评价方法。
初次了解本类比赛的小伙伴可以查看以下两份文档,希望帮助你快速入门和上手:
[赛题入门](https://gitlink.org.cn/ccf-ai-infra/GPUCodeForces/tree/main/GPUCodeForces%E8%B5%9B%E9%A2%98%E5%85%A5%E9%97%A8.md)、[代码解读](https://gitlink.org.cn/ccf-ai-infra/GPUCodeForces/tree/main/GPUCodeForces%E4%BB%A3%E7%A0%81%E8%A7%A3%E8%AF%BB.md)。
---
## 📥 参赛流程
一句话概括:进入[GPUCodeForces赛事首页](https://gitlink.com/ccf-ai-infra/GPUCodeForces)完成一份成功合并到仓库内的提交即为参赛成功时间自由方法自由只要有灵感就可以动手开code~
* 进入[GPUCodeForces赛事首页](https://gitlink.com/ccf-ai-infra/GPUCodeForces)登录参与本期比赛的Gitee账号完成一份成功合并到仓库内的提交即为参赛成功时间自由方法自由只要有灵感就可以动手开code~
### 🌰举个栗子
* 登录or注册自己的Gitlink账号后进入赛事首页初步查看仓库内的文件内容:
* 登录or注册自己的Gitlink账号后进入赛事首页查看仓库内的文件内容。仔细阅读[how-to-contribute.md](https://gitlink.com/ccf-ai-infra/GPUCodeForces/tree/main/how-to-contribute.md)完成CLA签署并熟悉提交流程。
.
├── S1(说明:第一季比赛名称,该目录下文件需选手自建)
│ ├── issue id 1(选手提交算子前创建的issue对应的id)
│ │ ├── cudacode.py(必要提交文件)
│ │ ├── torchcode.py(必要提交文件)
│ │ ├── run_code.py(必要提交文件)
│ │ ├── prompt.txt(必要提交文件)
│ ├── issue id 2
│ ├── issue id ...
├── example(样例提供,供大家上手)
│ ├── 001-example(文件结构与 issue id x 一致)
│ ├── 002-example
├── images(图片文件夹,参赛选手可忽略)
├── FAQ.md(社区收集的问题与解答)
├── LICENSE(证书,参赛选手可忽略)
├── README.md(赛题baseline)
├── how-to-contribute.md(提交指南看这里)
* 看到仓库内文件有一个example文件夹
* 在S1文件夹的文件即为参赛选手需要提交的文件其余文件夹和文件皆为辅助参赛选手了解比赛、提供思路、排忧解难之用。
<img src="./images/readme_sample_check.png">
* 我们将当前赛题的项目clone到本地电脑上建议使用git clone + 链接的方式在S1文件夹下创建一个以自己issue id命名的文件夹待提交的参赛代码文件都放在这里面
这是我们提供的一个样例,接下来我们在这个基础上进行一次完整的算子优化的提交(我们鼓励大家自己找到更好的算子并优化)。
```
Tipsissue id ≠ issue名称名称建议为对该算子的概括性描述。创建了issue后链接末尾处的数字即为issue id。同时也可以在issue界面的醒目位置查看如"#123"这样的数字标识。
```
* 做好了准备工作就可以开始尽情发挥去寻找算子亦或是优化算子。在example文件夹中提供了算子样例如果想简单上手可以查看该文件夹。虽然有效的算子优化也算一次提交但我们鼓励大家发现新的算子✨真正与其他选手拉开差距。
🔧简单介绍一下样例代码:
* 我们将样例clone到自己电脑上
并关注四份文件: torchcode.py、prompt.txt、cudacode_ori.py、example_cudacode.py最终需要提交的代码文件正是这四个。本次比赛在[模力方舟](https://ai.gitee.com/compute)平台上使用算力券购买容器实例:
<img src="./images/readme_git_compute.png">
接着便可以在云端实例上进行代码修改。相关算力券的领取方式请见[算力平台使用说明](https://ai.gitee.com/docs/compute/container)、[算力券兑换发放和兑换](https://ai.gitee.com/docs/billing/coupons)。
* 然后在该比赛仓库新建一个issue填写赛题。这里我们是对example-001算子优化因此issue的主题就可以是“对001-example数据集进行性能优化”
<img src="./images/readme_sample_issue.png">
可以看到这里有一个“#2”这是issue id你的算子优化、新算子都应该绑定一个独立的issue id最终有多少份issue被审核通过就表示提交成功了多少份。在即将提交的时候在该赛题仓库的S1文件夹下新建一个以该id命名无需带#号的文件夹该文件夹内容为四份必要文件和其他视参赛者情况需要补充的材料如readme文件、用到的其他数据集等
<img src="./images/readme_sample_folder.png">
* 准备工作就绪接下来看到example-001内的代码
**example_torchcode.py** 基准模型Baseline。示例提供一个简单的PyTorch模型只包含一个ReLU激活函数。
* <span style="background-color: grey; color: black; user-select: none;">get_inputs()</span>:生成模型运行时需要的输入数据。
* <span style="background-color: grey; color: black; user-select: none;">get_init_inputs()</span>:成模型初始化所需的参数(这里就是权重矩阵 weight
**example_cudacode.py**优化模型。示例使用PyTorch的load_inline功能直接编译和加载CUDA代码创建了一个新的模型类使用自定义CUDA实现替代PyTorch的ReLU。
<span style="background-color: grey; color: black; user-select: none;">example_cudacode.py</span>优化模型。示例使用PyTorch的load_inline功能直接编译和加载CUDA代码创建了一个新的模型类使用自定义CUDA实现替代PyTorch的ReLU。
**run_code.py**验证和性能测试脚本。验证自定义CUDA实现与原始PyTorch实现的数值精度一致性比较两种实现的性能计算加速比。
<span style="background-color: grey; color: black; user-select: none;">run_code.py</span>验证和性能测试脚本。验证自定义CUDA实现与原始PyTorch实现的数值精度一致性比较两种实现的性能计算加速比。
**prompt.txt**:提示词文本。提供类似 “融合算子 CUDA 设计” 的 prompt 编写思路。
<span style="background-color: grey; color: black; user-select: none;">prompt.txt</span>:这里给予参赛者一些提示:
💡如何将一个example变为自己的一份提交具体的算子优化思路可参考[GPUCodeForces赛题解读](./GPUCodeForces赛题解读.md)
* 为了测试代码最终的跑通结果需要使用规定的GPU。在[模力方舟](https://ai.gitee.com/compute)平台上准备了大家此次需要的算力资源,使用免费的算力券购买实例,接着便可以在云端实例上进行代码修改和测试。
相关算力券的领取方式请见[算力平台使用说明](https://ai.gitee.com/docs/compute/container)、[算力券兑换发放和兑换](https://ai.gitee.com/docs/billing/coupons)。
```
Tips进行云端服务器的实例选择有库存的即可点击进入后需要确定配置。其余选项保持默认只需更改镜像处的选择基础镜像-->CV-CUDA-->PyTorch2.4.0-->Python任意-->maca 3.0.0.5
```
* 要求编写自定义CUDA内核来替换PyTorch算子以获得加速
* 代码测试运行没有问题后便可以准备提交了记得将所有必要文件保存到提交文件夹。提交流程如下全程使用git
```
# cd到自己克隆到本地的项目路径下C:\Users\Desktop\ODTC AI Infra\GPUCodeForces
* 可以自由选择替换哪些算子,考虑算子融合机会
git remote -v # 检查是否链接到自己的Gitlink仓库
git checkout -b dev # 新建一个分支dev即分支名字可以任取
git add . # 将文件的变更操作暂存
git commit -m "输入你本次提交的目的" # 目的简洁明了最好
git push origin dev # 提交你的修改到分支上
```
对git与远程仓库操作的疑问可以点击[how-to-contribute.md](https://gitlink.org.cn/ccf-ai-infra/GPUCodeForces/tree/main/how-to-contribute.md)查看详情。
* 提供了示例语法和内联嵌入自定义CUDA算子的方法
* 给出了需要优化的模型架构简单的ReLU模型
* 顺利提交后的代码还只在你自己fork的仓库下还需要和主仓库合并才能真正让管理员看到你的代码。
* 然后参照example-001文件夹创建自己的文件夹提出新的torch cude对其中torch可来源于PyTorch、PaddlePaddle、TensorFlow、Jax、MMCV、Transformers等框架可以是单个算子也可以是多个算子融合。cuda代码可自己编写或者参照prompt.txt让LLM辅助编写
1. 修改torchcode.py按照example_torchcode.py格式在Model的__init__和forward中放入目标torch代码注意保留get_inputs()和get_init_inputs()函数
2. 修改cudacode.py修改cudacode.py里的load_inline和ModelNewcuda代码核心实现放在source字段中
fork仓库你自己的仓库与主仓库之间关系如下
* 优化好后,可以在模力方舟的实例上运行:
```
ODTC AI Infra/GPUCodeForces(main)
├── ...
├── folders
```sh
python run_code.py
```
确保能够正确输出结果后再准备提交。
your_name/GPUCodeForces(main) # 自己仓库下的main分支内容与主仓库一致
├── ...
├── folders
your_name/GPUCodeForces(dev) # dev分支才是你修改后的代码存放处
├── ...
│ ├── your codes
├── folders
```
进入自己的仓库,点击上方选项栏:
```
合并请求(PR)-->+新建合并请求-->源分支选择dev(名称任取)-->填写下方选框-->标题为自己的issue标题-->描述填请简明扼要描述内关联自己的issue id(如“fixes #001”)
```
提交后便能看到自己的PR记录了在对应记录的评论区会有测试结果的告知请留意查看~
* 接下来将优化好的代码保存到本地,然后参照[how-to-contribute.md](https://gitlink.org.cn/ccf-ai-infra/GPUCodeForces/tree/main/how-to-contribute.md)的指引进行代码仓库的提交与合并。
* 最终,成功提交的代码会合并到 S1/#your_issue id 下并且你的相关pr也会关闭。就像下面这样
<img src="./images/readme_sample_merge.png">
🌳一份完整的提交流程如上,期待各位自由发挥,赛出风采与水平!
⏺如仍有疑问,请点击[提交流程演示视频](https://www.bilibili.com/video/BV1CcnTztEfc/)
### 📦 提交PR内容
* **一个PR包含样本的目录** [提交样例](https://gitlink.org.cn/ccf-ai-infra/GPUCodeForces/tree/main/example/001-example)
* 每个提交目录建议包含如下:
1. **示例代码:** torch代码示例
2. **对比代码:** 和torch对应的CUDA代码
3. **测试代码入口:** run\_code.py请务必用这个名称提交的PR会根据这个名称在GPU上测试结果
4. **其它文件(或目录):** prompt利用LLM从torch代码生成cuda代码的prompt示例或者其它优化代码
5. **PR目录说明文件** https://gitlink.org.cn/ccf-ai-infra/GPUCodeForces/tree/main/example/001-example/readme.md
### 📦 提交PR的格式
建议在开始做题目之前创建一个赛题提交的PR和自己创建的赛题相关联。参赛选手在每个比赛周期的目录下例如第一期S1、第二期S2、第三期S3...创建一个目录目录名称赛题的IDICTXSZ),例如:
```plaintext
.
├── S1(说明:第一季比赛名称)
│ ├── ICTXSZ(说明以赛题ID命名的目录存放PR提交样本的目录)
| | ├── 示例代码
│ | ├── 对比代码
| | └── ……
│ └── ……
└── S2(第二季比赛)
└── 赛题1
```
### ⭐审核流程
* 你提交的PR都会得到回复大概存在的几种情况如下
```
提交PR-->测试通过✔️-->已合并-->有效提交 ヽ(✿゚▽゚)
提交PR-->测试通过✔️-->已关闭-->代码重复/相似-->无效提交 (;′⌒`)
提交PR-->测试失败✖️-->已关闭-->代码不合格-->无效提交 (;′⌒`)
```
也就是说,除了能够自己在算力平台的实例上运行得到算子测算的初步结果外,还可以在这里看到最终的测算结果。这里显示测试通过才能进入后续审核流程,并最终上传至比赛的算子仓库。
* 在一切文件都准备好并且提交后在对应的PR下会得到回复
<image src="./images/readme_comment_check.png">
也就是说,除了能够自己在服务器上运行得到算子测算的初步结果外,还可以在这里看到最终的测算结果。这里显示测试通过才能进入后续审核流程。
### ✅ 参赛资格
@ -161,7 +163,7 @@
> **接受数量** = 提交并被评审通过的样本总数
> **接受数量相同需要区分排名时如下的评分规则才会生效**
> **接受数量相同需要区分排名时如下的基础和甲方的评分规则才会生效**
---
@ -203,10 +205,8 @@
## 📬 联系与帮助
如需更多信息或格式说明,请查看官方文档或在本仓库提交[想法](https://gitlink.org.cn/ccf-ai-infra/GPUCodeForces/issues/new)进行讨论,或直接在交流群内与主办团队进行沟通。
祝各位挑战成功贡献出高质量的 GPU 评测数据集🚀
如需更多信息或格式说明,请查看官方文档或在本仓库提交[想法](https://gitlink.org.cn/ccf-ai-infra/GPUCodeForces/issues/new)进行讨论。  祝你挑战成功贡献出高质量的 GPU 评测数据集🚀
## FAQ
[第一季FAQ参考](https://gitlink.org.cn/ccf-ai-infra/GPUCodeForces/tree/main/FAQ.md)
[第一季FAQ参考](FAQ.md)

View File

@ -1,20 +0,0 @@
import torch
import torch.nn as nn
class ModelNew(nn.Module):
def __init__(self):
super(ModelNew, self).__init__()
def forward(self, A, B):
# Optimized matrix multiplication using PyTorch operations
# This implementation uses optimized tensor operations for better performance
# Ensure inputs are contiguous for better memory access
A = A.contiguous()
B = B.contiguous()
# Use bmm if batch dimensions exist, otherwise use mm
if A.dim() == 3 and B.dim() == 3:
return torch.bmm(A, B)
else:
return torch.mm(A, B)

View File

@ -1,15 +0,0 @@
import torch
import torch.nn as nn
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
def forward(self, A, B):
return torch.matmul(A, B)
def get_inputs():
return [torch.randn(256, 512), torch.randn(512, 256)]
def get_init_inputs():
return []

View File

@ -1,26 +0,0 @@
Write a custom CUDA kernel for optimized matrix multiplication (GEMM).
The standard matrix multiplication is defined as:
C = A × B
Where A is of shape (M, K), B is of shape (K, N), and C is of shape (M, N).
You should optimize the matrix multiplication using:
1. Shared memory tiling for better memory access patterns
2. Coalesced memory access
3. Thread block tiling to maximize parallelism
4. Avoid redundant memory loads
The kernel should be significantly faster than PyTorch's default matmul implementation for large matrices.
You are given the following architecture:
import torch
import torch.nn as nn
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
def forward(self, A, B):
return torch.matmul(A, B)

View File

@ -1,78 +0,0 @@
import torch
import time
from matmul_torchcode import Model, get_inputs, get_init_inputs
from matmul_cudacode import ModelNew
def run_benchmark():
if not torch.cuda.is_available():
print("CUDA is not available")
return
device = torch.device("cuda")
# Prepare input data
inputs = [x.cuda(device=device) for x in get_inputs()]
init_inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_init_inputs()]
# Initialize models
torch_model = Model(*init_inputs).cuda()
cuda_model = ModelNew(*init_inputs).cuda()
torch_model.eval()
cuda_model.eval()
print("-------------------- Precision Check --------------------")
with torch.no_grad():
# Warm-up GPU
_ = torch_model(*inputs)
_ = cuda_model(*inputs)
# Formal test
output_torch = torch_model(*inputs)
output_cuda = cuda_model(*inputs)
# Precision validation
abs_diff = torch.abs(output_torch - output_cuda)
max_diff = torch.max(abs_diff).item()
mean_diff = torch.mean(abs_diff).item()
if max_diff < 1e-4 and mean_diff < 1e-5:
print(f"✅ Precision aligned: max error {max_diff:.6f}, mean error {mean_diff:.6f}")
precision_flag = True
else:
print(f"❌ Precision mismatch: max error {max_diff:.6f}, mean error {mean_diff:.6f}")
precision_flag = False
print("\n-------------------- Performance Speedup Test --------------------")
num_iterations = 100
# Warm-up GPU
for _ in range(10):
_ = torch_model(*inputs)
_ = cuda_model(*inputs)
# PyTorch model timing
torch.cuda.synchronize()
start_time = time.time()
for _ in range(num_iterations):
_ = torch_model(*inputs)
torch.cuda.synchronize()
torch_time = (time.time() - start_time) / num_iterations
# Custom CUDA kernel timing
torch.cuda.synchronize()
start_time = time.time()
for _ in range(num_iterations):
_ = cuda_model(*inputs)
torch.cuda.synchronize()
cuda_time = (time.time() - start_time) / num_iterations
print(f"PyTorch built-in MatMul average execution time: {torch_time:.6f}s")
print(f"Custom CUDA MatMul average execution time: {cuda_time:.6f}s")
speedup = torch_time / cuda_time if cuda_time > 0 else 0
print(f"Speedup: {speedup:.2f}x")
return precision_flag, speedup
if __name__ == "__main__":
precision_flag, speedup = run_benchmark()

View File

@ -1,32 +0,0 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
class ModelNew(nn.Module):
def __init__(self, in_features: int = 1024, out_features: int = 2048):
super().__init__()
self.in_features = in_features
self.out_features = out_features
self.weight = nn.Parameter(torch.empty(out_features, in_features))
self.bias = nn.Parameter(torch.zeros(out_features))
nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
fan_in = self.weight.size(1)
bound = 1.0 / math.sqrt(fan_in)
nn.init.uniform_(self.bias, -bound, bound)
def forward(self, x: torch.Tensor) -> torch.Tensor:
y = F.linear(x, self.weight, self.bias)
# 使用精确 GELU 以确保与基线一致的数值结果
return F.gelu(y, approximate='none')
def get_init_inputs():
return {"in_features": 1024, "out_features": 2048}
def get_inputs():
B, T, D = 16, 512, 1024
x = torch.randn(B, T, D)
return x

View File

@ -1,22 +0,0 @@
import torch
import torch.nn as nn
class Model(nn.Module):
def __init__(self, in_features: int = 1024, out_features: int = 2048):
super().__init__()
self.linear = nn.Linear(in_features, out_features)
self.gelu = nn.GELU(approximate="none")
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.gelu(self.linear(x))
def get_init_inputs():
return {"in_features": 1024, "out_features": 2048}
def get_inputs():
B, T, D = 16, 512, 1024
x = torch.randn(B, T, D)
return x

View File

@ -1,11 +0,0 @@
目标:编写一个自定义 CUDA Kernel将 LinearGEMM+Bias与 GELU 激活融合为单一内核,在 MXC500 GPU 上减少中间张量写回与多次 kernel 启动,保证精度并获得 ≥1.0 的加速。
优化要点:
- 在 GEMM 计算累加寄存器阶段直接加上 bias 并进行 GELU 激活的近似/精确实现,避免额外的内存读写。
- 使用线程块与共享内存的分块装载tile来提升带宽利用率采用向量化加载float2/float4改善访存性能。
- 对齐权重与输入张量的内存布局,提升 coalesced 访问与 SM 吞吐。
- 对应 PyTorch 参考结构y = GELU(Linear(x))。
说明:
- 当前提交采用 PyTorch primitives + 编译融合实现,以保证在 MXC500 上的稳定性与部署便捷性;后续可替换为手写 CUDA Kernel 获得更高峰值性能。
- 基准脚本使用 CUDA Events 测时,确保真实 GPU 执行时间并进行精度校验。

View File

@ -1,102 +0,0 @@
import time
import torch
import linear_gelu_torchcode as torchcode
import linear_gelu_cudacode as cudacode
def _to_device(tensors, device):
return [t.to(device) for t in tensors]
def _copy_params(torch_model, cuda_model):
with torch.no_grad():
cuda_model.weight.copy_(torch_model.linear.weight)
cuda_model.bias.copy_(torch_model.linear.bias)
def _measure_gpu_seconds(model, args, iters=100):
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
torch.cuda.synchronize()
with torch.no_grad():
start.record()
for _ in range(iters):
_ = model(*args)
end.record()
torch.cuda.synchronize()
ms = start.elapsed_time(end) / iters
return ms / 1000.0
def run_benchmark():
if not torch.cuda.is_available():
print("CUDA 不可用")
return False, 0.0
torch.manual_seed(0)
device = torch.device("cuda")
init_kwargs = torchcode.get_init_inputs()
torch_model = torchcode.Model(**init_kwargs).to(device).eval()
cuda_model = cudacode.ModelNew(**init_kwargs).to(device).eval()
# 关闭编译避免在当前平台上落到慢路径CUTLASS 不可用)
# 参数对齐
_copy_params(torch_model, cuda_model)
# 准备输入
x = torchcode.get_inputs()
x, = _to_device([x], device)
print("-------------------- 精度对齐验证 --------------------")
with torch.no_grad():
# 预热
_ = torch_model(x)
_ = cuda_model(x)
# 正式测试
output_torch = torch_model(x)
output_cuda = cuda_model(x)
abs_diff = torch.abs(output_torch - output_cuda)
max_diff = torch.max(abs_diff).item()
mean_diff = torch.mean(abs_diff).item()
if max_diff < 1e-4 and mean_diff < 1e-5:
print(f"✅ 精度对齐:最大误差 {max_diff:.6f},平均误差 {mean_diff:.6f}")
precision_flag = True
else:
print(f"❌ 精度不一致:最大误差 {max_diff:.6f},平均误差 {mean_diff:.6f}")
precision_flag = False
print("\n-------------------- 性能加速比测试 --------------------")
num_iterations = 200
# 预热
for _ in range(10):
_ = torch_model(x)
_ = cuda_model(x)
# 可选:允许 TF32若硬件支持提升矩阵乘性能
try:
torch.backends.cuda.matmul.allow_tf32 = True
torch.set_float32_matmul_precision("medium")
except Exception:
pass
# PyTorch计时CUDA Events
torch_time = _measure_gpu_seconds(torch_model, (x,), iters=num_iterations)
# 优化版计时CUDA Events
cuda_time = _measure_gpu_seconds(cuda_model, (x,), iters=num_iterations)
print(f"PyTorch内置Linear+GELU平均执行时间: {torch_time:.6f}")
print(f"自定义CUDA Linear+GELU平均执行时间: {cuda_time:.6f}")
speedup = torch_time / cuda_time if cuda_time > 0 else 0.0
print(f"加速比 (Speedup): {speedup:.2f}x")
return precision_flag, speedup
if __name__ == "__main__":
precision_flag, speedup = run_benchmark()

View File

@ -1,341 +0,0 @@
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
batchnorm_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
// 统一的训练 kernel计算批次统计量
__global__ void batchnorm_forward_train_kernel_optimized(
const float* __restrict__ x,
const float* __restrict__ gamma,
const float* __restrict__ beta,
float* __restrict__ running_mean,
float* __restrict__ running_var,
float* __restrict__ y,
int batch,
int features,
float eps,
float momentum,
bool update_stats // 是否更新统计量
) {
int feature = blockIdx.x;
if (feature >= features) return;
int tid = threadIdx.x;
int num_threads = blockDim.x;
int warp_id = tid / 32;
int lane_id = tid % 32;
int num_warps = (num_threads + 31) / 32;
const float* x_base = x + feature;
float* y_base = y + feature;
float sum = 0.0f;
float sum_sq = 0.0f;
int row = tid;
for (; row + num_threads <= batch; row += num_threads) {
float v = x_base[row * features];
sum += v;
sum_sq += v * v;
}
if (row < batch) {
float v = x_base[row * features];
sum += v;
sum_sq += v * v;
}
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1) {
sum += __shfl_down_sync(0xffffffff, sum, offset);
sum_sq += __shfl_down_sync(0xffffffff, sum_sq, offset);
}
__shared__ float shared_sum[32];
__shared__ float shared_sq[32];
if (lane_id == 0) {
shared_sum[warp_id] = sum;
shared_sq[warp_id] = sum_sq;
}
__syncthreads();
if (tid < 32) {
sum = (tid < num_warps) ? shared_sum[tid] : 0.0f;
sum_sq = (tid < num_warps) ? shared_sq[tid] : 0.0f;
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1) {
sum += __shfl_down_sync(0xffffffff, sum, offset);
sum_sq += __shfl_down_sync(0xffffffff, sum_sq, offset);
}
}
__shared__ float s_mean;
__shared__ float s_inv_std;
__shared__ float s_gamma;
__shared__ float s_beta;
if (tid == 0) {
float mean = sum / batch;
float var = (sum_sq / batch) - (mean * mean);
var = fmaxf(var, 0.0f);
s_mean = mean;
s_inv_std = rsqrtf(var + eps);
s_gamma = gamma[feature];
s_beta = beta[feature];
// 只有需要时才更新 running stats
if (update_stats) {
running_mean[feature] = (1.0f - momentum) * running_mean[feature] + momentum * mean;
float unbiased_var = var * batch / fmaxf(float(batch - 1), 1.0f);
running_var[feature] = (1.0f - momentum) * running_var[feature] + momentum * unbiased_var;
}
}
__syncthreads();
float mean = s_mean;
float inv_std = s_inv_std;
float g = s_gamma;
float b = s_beta;
row = tid;
for (; row + num_threads <= batch; row += num_threads) {
float v = x_base[row * features];
float norm = (v - mean) * inv_std;
y_base[row * features] = norm * g + b;
}
if (row < batch) {
float v = x_base[row * features];
float norm = (v - mean) * inv_std;
y_base[row * features] = norm * g + b;
}
}
// 推理模式 kernel使用 running stats
__global__ void batchnorm_forward_eval_kernel_optimized(
const float* __restrict__ x,
const float* __restrict__ gamma,
const float* __restrict__ beta,
const float* __restrict__ running_mean,
const float* __restrict__ running_var,
float* __restrict__ y,
int batch,
int features,
float eps
) {
int tid = blockIdx.x * blockDim.x + threadIdx.x;
int total = batch * features;
int stride = gridDim.x * blockDim.x;
for (int idx = tid; idx < total; idx += stride) {
int feature = idx % features;
float mean = running_mean[feature];
float var = running_var[feature];
float inv_std = rsqrtf(var + eps);
float g = gamma[feature];
float b = beta[feature];
float v = x[idx];
float norm = (v - mean) * inv_std;
y[idx] = norm * g + b;
}
}
torch::Tensor batchnorm_cuda_forward(
torch::Tensor x,
torch::Tensor weight,
torch::Tensor bias,
torch::Tensor running_mean,
torch::Tensor running_var,
bool training,
double momentum,
double eps,
bool track_running_stats // 改名更清晰地表达意图
) {
TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
TORCH_CHECK(weight.is_cuda(), "weight must be a CUDA tensor");
TORCH_CHECK(bias.is_cuda(), "bias must be a CUDA tensor");
TORCH_CHECK(x.dtype() == torch::kFloat32, "only float32 tensors are supported");
TORCH_CHECK(weight.dtype() == torch::kFloat32, "weight must be float32");
TORCH_CHECK(bias.dtype() == torch::kFloat32, "bias must be float32");
TORCH_CHECK(x.dim() == 2, "input must be 2D [batch, features]");
TORCH_CHECK(weight.dim() == 1, "weight must be 1D");
TORCH_CHECK(bias.dim() == 1, "bias must be 1D");
TORCH_CHECK(x.size(1) == weight.size(0), "feature size mismatch");
TORCH_CHECK(weight.size(0) == bias.size(0), "weight and bias must have the same length");
auto x_contig = x.contiguous();
auto weight_contig = weight.contiguous();
auto bias_contig = bias.contiguous();
int batch = x_contig.size(0);
int features = x_contig.size(1);
auto y = torch::empty_like(x_contig);
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
TORCH_CHECK(running_mean.is_cuda(), "running_mean must be a CUDA tensor");
TORCH_CHECK(running_var.is_cuda(), "running_var must be a CUDA tensor");
TORCH_CHECK(running_mean.dim() == 1, "running_mean must be 1D");
TORCH_CHECK(running_var.dim() == 1, "running_var must be 1D");
TORCH_CHECK(running_mean.size(0) == features, "running_mean size mismatch");
TORCH_CHECK(running_var.size(0) == features, "running_var size mismatch");
// 关键修改根据 track_running_stats 决定行为
// track_running_stats=False: 总是计算批次统计训练和推理都一样
// track_running_stats=True + training: 计算批次统计并更新 running stats
// track_running_stats=True + eval: 使用 running stats
bool use_batch_stats = !track_running_stats || training;
if (use_batch_stats) {
// 使用批次统计量训练模式 track_running_stats=False
int threads;
if (batch <= 16) {
threads = 32;
} else if (batch <= 32) {
threads = 32;
} else if (batch <= 64) {
threads = 64;
} else if (batch <= 128) {
threads = 128;
} else if (batch <= 256) {
threads = 256;
} else {
threads = 256;
}
int blocks = features;
size_t shared_mem = 0;
// update_stats = track_running_stats && training
// track_running_stats=False: 不更新
// track_running_stats=True + training: 更新
// track_running_stats=True + eval: 不会走到这里
bool update_stats = track_running_stats && training;
batchnorm_forward_train_kernel_optimized<<<blocks, threads, shared_mem, stream>>>(
x_contig.data_ptr<float>(),
weight_contig.data_ptr<float>(),
bias_contig.data_ptr<float>(),
running_mean.data_ptr<float>(),
running_var.data_ptr<float>(),
y.data_ptr<float>(),
batch,
features,
static_cast<float>(eps),
static_cast<float>(momentum),
update_stats
);
} else {
// 使用 running statstrack_running_stats=True + eval 模式
int total = batch * features;
int threads = 256;
int blocks;
if (total <= 4096) {
blocks = (total + threads - 1) / threads;
} else {
blocks = min(1024, (total + threads * 4 - 1) / (threads * 4));
}
batchnorm_forward_eval_kernel_optimized<<<blocks, threads, 0, stream>>>(
x_contig.data_ptr<float>(),
weight_contig.data_ptr<float>(),
bias_contig.data_ptr<float>(),
running_mean.data_ptr<float>(),
running_var.data_ptr<float>(),
y.data_ptr<float>(),
batch,
features,
static_cast<float>(eps)
);
}
C10_CUDA_KERNEL_LAUNCH_CHECK();
return y;
}
"""
batchnorm_cpp_source = r"""
torch::Tensor batchnorm_cuda_forward(
torch::Tensor x,
torch::Tensor weight,
torch::Tensor bias,
torch::Tensor running_mean,
torch::Tensor running_var,
bool training,
double momentum,
double eps,
bool track_running_stats
);
"""
batchnorm_cuda = load_inline(
name="batchnorm_cuda_ext",
cpp_sources=batchnorm_cpp_source,
cuda_sources=batchnorm_source,
functions=["batchnorm_cuda_forward"],
verbose=True
)
class ModelNew(nn.Module):
"""
Model performing matrix multiplication followed by custom CUDA BatchNorm and ReLU.
Optimized with Warp-level reduction (Plan 1) and thread configuration (Plan 2).
"""
def __init__(self, mat_weight: torch.Tensor, bn_weight: torch.Tensor, bn_bias: torch.Tensor,
eps: float = 1e-5, momentum: float = 0.1, track_running_stats: bool = True):
super().__init__()
if mat_weight.dim() != 2:
raise ValueError("mat_weight must be a 2D tensor [input_dim, output_dim].")
if bn_weight.dim() != 1 or bn_bias.dim() != 1:
raise ValueError("BatchNorm weight and bias must be 1D.")
if bn_weight.size(0) != mat_weight.size(1):
raise ValueError("BatchNorm parameter size must match output_dim.")
if bn_weight.size(0) != bn_bias.size(0):
raise ValueError("BatchNorm weight and bias must share shape.")
self.weight = nn.Parameter(mat_weight.clone())
self.bn_weight = nn.Parameter(bn_weight.clone())
self.bn_bias = nn.Parameter(bn_bias.clone())
self.eps = eps
self.momentum = momentum
self.track_running_stats = track_running_stats
# 无论 track_running_stats 是什么,都创建 buffer
self.register_buffer('running_mean', torch.zeros(bn_weight.size(0)))
self.register_buffer('running_var', torch.ones(bn_weight.size(0)))
def forward(self, x: torch.Tensor) -> torch.Tensor:
if not x.is_cuda:
raise ValueError("Input must be a CUDA tensor.")
if not self.weight.is_cuda:
raise ValueError("Model weight must be on CUDA.")
if not self.bn_weight.is_cuda or not self.bn_bias.is_cuda:
raise ValueError("BatchNorm parameters must be on CUDA.")
x = torch.matmul(x, self.weight)
# 传递 track_running_stats 参数到 CUDA kernel
x = batchnorm_cuda.batchnorm_cuda_forward(
x,
self.bn_weight,
self.bn_bias,
self.running_mean,
self.running_var,
self.training,
self.momentum,
self.eps,
self.track_running_stats
)
return torch.relu(x)

View File

@ -1,46 +0,0 @@
import torch
import torch.nn as nn
class BatchNormModel(nn.Module):
"""
Model that performs matrix multiplication followed by BatchNorm and ReLU activation.
"""
def __init__(self, weight, num_features=2048, eps=1e-5, momentum=0.1, track_running_stats=True):
super(BatchNormModel, self).__init__()
self.weight = nn.Parameter(weight)
# 设置 track_running_stats=True 以跟踪运行时统计量
self.bn = nn.BatchNorm1d(
num_features,
eps=eps,
momentum=momentum,
track_running_stats=track_running_stats
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Performs matrix multiplication, applies BatchNorm, then ReLU activation.
Args:
x (torch.Tensor): Input tensor of shape [batch_size, input_dim]
Returns:
torch.Tensor: Output tensor of shape [batch_size, output_dim]
"""
x = torch.matmul(x, self.weight)
x = self.bn(x)
return torch.relu(x)
# 添加别名以便在 run_code.py 中使用
Model = BatchNormModel
batch_size = 16
input_dim = 1024
output_dim = 2048
def get_inputs():
x = torch.randn(batch_size, input_dim)
return [x]
def get_init_inputs():
weight = torch.randn(input_dim, output_dim)
return [weight]

View File

@ -1,61 +0,0 @@
You write custom CUDA kernels to replace the pytorch operators in the given architecture to get speedups.
You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining matmul+relu), or algorithmic changes (such as online softmax). You are only limited by your imagination.
Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:
```python
import torch
import torch.nn as nn
import torch.nn.functional as F
class Model(nn.Module):
def __init__(self, num_features) -> None:
super().__init__()
self.bn = nn.BatchNorm1d(num_features)
def forward(self, x):
return self.bn(x)
def get_inputs():
x = torch.randn(16, 2048).cuda()
return [x]
def get_init_inputs():
return [2048]
```
You are given the following architecture to implement BatchNorm1d with custom CUDA kernel:
```python
import torch
import torch.nn as nn
class Model(nn.Module):
"""使用 PyTorch BatchNorm1d 的基准实现。"""
def __init__(self, num_features: int, eps: float = 1e-5, momentum: float = 0.1):
super().__init__()
self.bn = nn.BatchNorm1d(num_features, eps=eps, momentum=momentum)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""对输入做 BatchNorm输出形状与输入一致。"""
return self.bn(x)
batch_size = 16
feature_dim = 2048
def get_inputs():
x = torch.randn(batch_size, feature_dim)
return [x]
def get_init_inputs():
return [feature_dim]
```

View File

@ -1,101 +0,0 @@
###########################################################
# 性能和精度验证程序
###########################################################
import torch
import torch.nn as nn
import time
from batchnorm1d_torch import Model as TorchModel, get_inputs, get_init_inputs
from batchnorm1d_cuda import ModelNew as CudaModel
def run_benchmark():
# 检查 CUDA 是否可用
if not torch.cuda.is_available():
print("CUDA 不可用,请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
return
else:
device = torch.device("cuda")
# 初始化模型
# 获取 torch 版本的初始化参数
torch_init_inputs = get_init_inputs()
weight = torch_init_inputs[0].cuda(device=device)
# 为 BatchNorm 准备参数
batch_size = 16
input_dim = 1024
output_dim = 2048
bn_weight = torch.ones(output_dim, device=device, dtype=torch.float32)
bn_bias = torch.zeros(output_dim, device=device, dtype=torch.float32)
# 初始化输入数据
inputs = get_inputs()
inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in inputs]
# 初始化两个模型
track_bool = True
torch_model = TorchModel(weight.clone(), num_features=output_dim, eps=1e-5, track_running_stats=track_bool).cuda()
cuda_model = CudaModel(weight.clone(), bn_weight.clone(), bn_bias.clone(), eps=1e-5, track_running_stats=track_bool).cuda()
torch_model.eval()
cuda_model.eval()
print("-------------------- 精度对齐验证 --------------------")
with torch.no_grad():
output_torch = torch_model(*inputs)
output_cuda = cuda_model(*inputs)
# 更严格的精度检查
abs_diff = (output_torch - output_cuda).abs()
max_diff = abs_diff.max().item()
mean_diff = abs_diff.mean().item()
print(f"最大差异: {max_diff:.6f}")
print(f"平均差异: {mean_diff:.6f}")
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-03, atol=1e-03)
if precision_flag:
print("✅ 精度对齐:两个模型的输出结果非常接近。")
else:
print("❌ 精度不一致!")
print("\n-------------------- 性能加速比测试 --------------------")
print(f"track_running_stats = {track_bool}")
num_iterations = 1000 # 增加迭代次数以获得更准确的时间测量
# Warm up
print("预热中...")
for _ in range(100):
with torch.no_grad():
_ = torch_model(*inputs)
_ = cuda_model(*inputs)
# PyTorch 模型计时
torch.cuda.synchronize()
start_time = time.time()
with torch.no_grad():
for _ in range(num_iterations):
_ = torch_model(*inputs)
torch.cuda.synchronize()
torch_time = (time.time() - start_time) / num_iterations
# 自定义 CUDA 内核计时
torch.cuda.synchronize()
start_time = time.time()
with torch.no_grad():
for _ in range(num_iterations):
_ = cuda_model(*inputs)
torch.cuda.synchronize()
cuda_time = (time.time() - start_time) / num_iterations
print(f"PyTorch 平均执行时间: {torch_time*1000:.4f} 毫秒")
print(f"自定义 CUDA BatchNorm 平均执行时间: {cuda_time*1000:.4f} 毫秒")
speedup = 0
if cuda_time > 0:
speedup = torch_time / cuda_time
print(f"加速比 (Speedup): {speedup:.2f}x")
else:
print("CUDA 内核执行时间为0无法计算加速比。")
return precision_flag, speedup
if __name__ == "__main__":
precision_flag, speedup = run_benchmark()

View File

@ -1,223 +0,0 @@
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
# CUDA implementation of Conv2D (tiled + shared memory + output-channel blocking)
conv2d_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <math.h>
#ifndef CHECK_CUDA
#define CHECK_CUDA(x) TORCH_CHECK((x).is_cuda(), #x " must be a CUDA tensor")
#endif
#ifndef CHECK_CONTIGUOUS
#define CHECK_CONTIGUOUS(x) TORCH_CHECK((x).is_contiguous(), #x " must be contiguous")
#endif
#ifndef CHECK_FLOAT
#define CHECK_FLOAT(x) TORCH_CHECK((x).scalar_type() == at::kFloat, #x " must be float32")
#endif
// 每个block计算一个 (b, oc_group) 上的输出tile复用输入tile计算 OC_TILE 个输出通道
template<int BLOCK_X, int BLOCK_Y, int OC_TILE>
__global__ void conv2d_tiled_kernel_oc(
const float* __restrict__ input, // [B, C_in, H, W]
const float* __restrict__ weight, // [C_out, C_in, K, K]
const float* __restrict__ bias, // [C_out] or nullptr
float* __restrict__ output, // [B, C_out, H_out, W_out]
int B, int C_in, int C_out,
int H, int W, int K, int H_out, int W_out,
bool has_bias
) {
// grid.z = B * ceil_div(C_out, OC_TILE)
int groups = (C_out + OC_TILE - 1) / OC_TILE;
int b = blockIdx.z / groups;
int og = blockIdx.z % groups; // 输出通道组编号
int co0 = og * OC_TILE; // 本组起始输出通道
int ow0 = blockIdx.x * BLOCK_X;
int oh0 = blockIdx.y * BLOCK_Y;
int ow = ow0 + threadIdx.x;
int oh = oh0 + threadIdx.y;
extern __shared__ float smem[];
// 输入tile大小(BLOCK_Y+K-1) x (BLOCK_X+K-1)
int tile_w = BLOCK_X + K - 1;
int tile_h = BLOCK_Y + K - 1;
float* tile = smem; // tile_h * tile_w
float* w_sh = tile + tile_h * tile_w; // OC_TILE * K * K
// w_sh 布局: [oc_local][K*K]
// 累加器每线程维护 OC_TILE 个通道
float acc[OC_TILE];
#pragma unroll
for (int oc = 0; oc < OC_TILE; ++oc) {
int co = co0 + oc;
acc[oc] = (has_bias && co < C_out) ? bias[co] : 0.0f;
}
bool valid_xy = (oh < H_out) && (ow < W_out);
// 遍历输入通道
for (int ci = 0; ci < C_in; ++ci) {
// 1) 加载本组 OC_TILE 的权重到共享内存
int total_w = OC_TILE * K * K;
for (int t = threadIdx.y * BLOCK_X + threadIdx.x; t < total_w; t += BLOCK_X * BLOCK_Y) {
int oc = t / (K*K);
int rem = t % (K*K);
int kh = rem / K;
int kw = rem % K;
int co = co0 + oc;
float wv = 0.0f;
if (co < C_out) {
int w_idx = ((co * C_in + ci) * K + kh) * K + kw;
wv = weight[w_idx];
}
w_sh[t] = wv;
}
// 2) 加载输入tile到共享内存该tile将被 OC_TILE 个输出通道复用
int ih0 = oh0;
int iw0 = ow0;
for (int th = threadIdx.y; th < tile_h; th += BLOCK_Y) {
int ih = ih0 + th;
bool in_h = (ih >= 0) && (ih < H);
for (int tw = threadIdx.x; tw < tile_w; tw += BLOCK_X) {
int iw = iw0 + tw;
bool in_w = (iw >= 0) && (iw < W);
float v = 0.0f;
if (in_h && in_w) {
int in_idx = (((b * C_in + ci) * H + ih) * W + iw);
v = input[in_idx];
}
tile[th * tile_w + tw] = v;
}
}
__syncthreads();
// 3) 计算同一输入tile OC_TILE 个输出通道分别累加
if (valid_xy) {
int t_base = threadIdx.y * tile_w + threadIdx.x;
#pragma unroll
for (int kh = 0; kh < K; ++kh) {
int t_row = t_base + kh * tile_w;
#pragma unroll
for (int kw = 0; kw < K; ++kw) {
float val = tile[t_row + kw];
#pragma unroll
for (int oc = 0; oc < OC_TILE; ++oc) {
float wv = w_sh[oc * (K*K) + kh * K + kw];
acc[oc] = fmaf(val, wv, acc[oc]);
}
}
}
}
__syncthreads(); // 保护下一个 ci 的加载
}
// 4) 写回输出
if (valid_xy) {
int base = (b * C_out) * (H_out * W_out);
int out_offset = oh * W_out + ow;
#pragma unroll
for (int oc = 0; oc < OC_TILE; ++oc) {
int co = co0 + oc;
if (co < C_out) {
int out_idx = base + co * (H_out * W_out) + out_offset;
output[out_idx] = acc[oc];
}
}
}
}
// C++ wrapper
torch::Tensor conv2d_cuda(
torch::Tensor input,
torch::Tensor weight,
torch::Tensor bias
) {
CHECK_CUDA(input);
CHECK_CUDA(weight);
CHECK_CUDA(bias);
CHECK_CONTIGUOUS(input);
CHECK_CONTIGUOUS(weight);
CHECK_CONTIGUOUS(bias);
CHECK_FLOAT(input);
CHECK_FLOAT(weight);
CHECK_FLOAT(bias);
int B = input.size(0);
int C_in = input.size(1);
int H = input.size(2);
int W = input.size(3);
int C_out = weight.size(0);
int K = weight.size(2);
TORCH_CHECK(weight.size(3) == K, "Kernel must be square");
int H_out = H - K + 1;
int W_out = W - K + 1;
auto output = torch::empty({B, C_out, H_out, W_out}, input.options());
// 参数可按GPU微调32x8, 16x16
const int BLOCK_X = 16;
const int BLOCK_Y = 16;
const int OC_TILE = 4;
dim3 block(BLOCK_X, BLOCK_Y, 1);
int groups = (C_out + OC_TILE - 1) / OC_TILE;
dim3 grid((W_out + BLOCK_X - 1) / BLOCK_X,
(H_out + BLOCK_Y - 1) / BLOCK_Y,
B * groups);
size_t tile_w = BLOCK_X + K - 1;
size_t tile_h = BLOCK_Y + K - 1;
size_t shmem_elems = tile_w * tile_h + OC_TILE * K * K;
size_t shmem_bytes = shmem_elems * sizeof(float);
bool has_bias = bias.numel() > 0;
conv2d_tiled_kernel_oc<BLOCK_X, BLOCK_Y, OC_TILE><<<grid, block, shmem_bytes>>>(
input.data_ptr<float>(),
weight.data_ptr<float>(),
has_bias ? bias.data_ptr<float>() : nullptr,
output.data_ptr<float>(),
B, C_in, C_out, H, W, K, H_out, W_out, has_bias
);
auto err = cudaGetLastError();
TORCH_CHECK(err == cudaSuccess, "conv2d kernel launch failed: ", cudaGetErrorString(err));
return output;
}
"""
conv2d_cpp_source = r"""
torch::Tensor conv2d_cuda(torch::Tensor input, torch::Tensor weight, torch::Tensor bias);
"""
# Compile with O3 (no fast-math to keep FP32 parity)
conv2d = load_inline(
name="conv2d_tiled_opt_oc",
cpp_sources=conv2d_cpp_source,
cuda_sources=conv2d_source,
functions=["conv2d_cuda"],
verbose=False,
extra_cuda_cflags=["-O3"]
)
class ModelNew(nn.Module):
def __init__(self, weight, bias=None):
super(ModelNew, self).__init__()
self.weight = nn.Parameter(weight)
self.bias = nn.Parameter(bias) if bias is not None else nn.Parameter(torch.empty(0, device=weight.device, dtype=weight.dtype))
self.conv2d = conv2d
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = x.contiguous()
w = self.weight.contiguous()
b = self.bias.contiguous() if self.bias is not None else torch.empty(0, device=x.device, dtype=x.dtype)
return self.conv2d.conv2d_cuda(x, w, b)

View File

@ -1,40 +0,0 @@
import torch
import torch.nn as nn
class Model(nn.Module):
"""
Model that performs 2D convolution operation.
"""
def __init__(self, weight, bias=None):
super(Model, self).__init__()
self.weight = nn.Parameter(weight)
self.bias = nn.Parameter(bias) if bias is not None else None
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Performs 2D convolution.
Args:
x (torch.Tensor): Input tensor of shape [batch_size, in_channels, height, width]
Returns:
torch.Tensor: Output tensor of shape [batch_size, out_channels, out_height, out_width]
"""
return torch.nn.functional.conv2d(x, self.weight, self.bias, stride=1, padding=0)
# Hyperparameters
batch_size = 4
in_channels = 3
out_channels = 64
height = 32
width = 32
kernel_size = 3
def get_inputs():
x = torch.randn(batch_size, in_channels, height, width)
return [x]
def get_init_inputs():
weight = torch.randn(out_channels, in_channels, kernel_size, kernel_size)
bias = torch.randn(out_channels)
return [weight, bias]

View File

@ -1,30 +0,0 @@
You write custom CUDA kernels to replace the pytorch operators in the given architecture to get speedups.
You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining matmul+relu), or algorithmic changes (such as online softmax). You are only limited by your imagination.
Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:
```python
import torch
import torch.nn as nn
import torch.nn.functional as F
class Model(nn.Module):
def __init__(self) -> None:
super().__init__()
def forward(self, a, b):
return a + b
def get_inputs():
# randomly generate input tensors based on the model architecture
a = torch.randn(1, 128).cuda()
b = torch.randn(1, 128).cuda()
return [a, b]
def get_init_inputs():
# randomly generate tensors required for initialization based on the model architecture
return []

View File

@ -1,92 +0,0 @@
###########################################################
# 性能和精度验证程序
###########################################################
import torch
import torch.nn as nn
import time
from conv2d_torch import Model, get_inputs, get_init_inputs
from conv2d_cuda import ModelNew
# 禁用 TF32确保与自定义 FP32 核精度对齐
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False
torch.backends.cudnn.deterministic = True
def _time_cuda_model(fn, inputs, iters=300, warmup=50):
torch.cuda.synchronize()
for _ in range(warmup):
_ = fn(*inputs)
torch.cuda.synchronize()
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
for _ in range(iters):
_ = fn(*inputs)
end.record()
torch.cuda.synchronize()
ms = start.elapsed_time(end) / iters # 平均每次毫秒
return ms / 1000.0 # 转为秒
def run_benchmark():
# 检查 CUDA 是否可用
if not torch.cuda.is_available():
print("CUDA 不可用,请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
return
else:
device = torch.device("cuda")
torch.backends.cudnn.benchmark = True
# 初始化模型
init_inputs = get_init_inputs()
init_inputs = [
x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in init_inputs
]
inputs = get_inputs()
inputs = [
x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in inputs
]
torch_model = Model(*init_inputs).cuda()
cuda_model = ModelNew(*init_inputs).cuda()
torch_model.eval()
cuda_model.eval()
print("-------------------- 精度对齐验证 --------------------")
with torch.no_grad():
output_torch = torch_model(*inputs)
output_cuda = cuda_model(*inputs)
abs_diff = (output_torch - output_cuda).abs()
max_diff = abs_diff.max().item()
mean_diff = abs_diff.mean().item()
print(f"最大差异: {max_diff:.6f}")
print(f"平均差异: {mean_diff:.6f}")
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-5, atol=1e-5)
if precision_flag:
print("✅ 精度对齐:两个模型的输出结果非常接近。")
else:
print("❌ 精度不一致!")
print("\n-------------------- 性能加速比测试 --------------------")
num_iterations = 300
# 计时
torch_time = _time_cuda_model(torch_model, inputs, iters=num_iterations, warmup=50)
cuda_time = _time_cuda_model(cuda_model, inputs, iters=num_iterations, warmup=50)
print(f"PyTorch (conv2d) 平均执行时间: {torch_time:.6f}")
print(f"自定义 CUDA conv2d 平均执行时间: {cuda_time:.6f}")
speedup = 0.0
if cuda_time > 0:
speedup = torch_time / cuda_time
print(f"相对 PyTorch 加速比: {speedup:.2f}x")
else:
print("CUDA 内核执行时间为0无法计算加速比。")
return precision_flag, speedup
if __name__ == "__main__":
precision_flag, speedup = run_benchmark()

View File

@ -1,29 +0,0 @@
You write custom CUDA kernels to replace the pytorch operators in the given architecture to get speedups.
You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining matmul+relu), or algorithmic changes (such as online softmax). You are only limited by your imagination.
Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:
```python
import torch
import torch.nn as nn
import torch.nn.functional as F
class Model(nn.Module):
def __init__(self) -> None:
super().__init__()
self.norm = nn.RMSNorm(128)
def forward(self, x):
return self.norm(x)
def get_inputs():
x = torch.randn(1, 128).cuda()
return [x]
def get_init_inputs():
return []
```

View File

@ -1,102 +0,0 @@
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
rmsnorm_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <math.h>
__global__ void rmsnorm_kernel(
const float* __restrict__ x,
const float* __restrict__ weight,
float* __restrict__ y,
int batch,
int features,
float eps
) {
int row = blockIdx.x;
if (row >= batch) return;
int tid = threadIdx.x;
extern __shared__ float sdata[];
float sum_sq = 0.0f;
for (int i = tid; i < features; i += blockDim.x) {
float v = x[row * features + i];
sum_sq += v * v;
}
sdata[tid] = sum_sq;
__syncthreads();
for (int offset = blockDim.x >> 1; offset > 0; offset >>= 1) {
if (tid < offset) {
sdata[tid] += sdata[tid + offset];
}
__syncthreads();
}
float rms = rsqrtf(sdata[0] / features + eps);
__syncthreads(); // 保证 rms 可见
for (int i = tid; i < features; i += blockDim.x) {
float v = x[row * features + i];
float w = weight[i];
y[row * features + i] = v * rms * w;
}
}
torch::Tensor rmsnorm_cuda(torch::Tensor x, torch::Tensor weight, float eps) {
TORCH_CHECK(x.is_cuda(), "x 必须是 CUDA 张量");
TORCH_CHECK(weight.is_cuda(), "weight 必须是 CUDA 张量");
TORCH_CHECK(x.dim() == 2, "当前内核仅支持二维输入张量");
TORCH_CHECK(weight.dim() == 1, "RMSNorm 权重必须是一维向量");
TORCH_CHECK(x.size(1) == weight.size(0), "输入最后一维与权重长度不匹配");
int batch = x.size(0);
int features = x.size(1);
auto y = torch::empty_like(x);
int threads = 256;
if (features < threads) {
threads = 1;
while (threads < features) threads <<= 1;
if (threads < 32) threads = 32;
}
size_t shared = threads * sizeof(float);
rmsnorm_kernel<<<batch, threads, shared>>>(
x.data_ptr<float>(),
weight.data_ptr<float>(),
y.data_ptr<float>(),
batch,
features,
eps
);
return y;
}
"""
rmsnorm_cpp_source = """
torch::Tensor rmsnorm_cuda(torch::Tensor x, torch::Tensor weight, float eps);
"""
rmsnorm = load_inline(
name="rmsnorm",
cpp_sources=rmsnorm_cpp_source,
cuda_sources=rmsnorm_source,
functions=["rmsnorm_cuda"],
verbose=True
)
class ModelNew(nn.Module):
def __init__(self, weight: torch.Tensor, eps: float = 1e-6):
super().__init__()
if weight.dim() != 1:
raise ValueError("RMSNorm 权重必须是一维向量。")
self.weight = nn.Parameter(weight.clone())
self.eps = eps
def forward(self, x: torch.Tensor) -> torch.Tensor:
return rmsnorm.rmsnorm_cuda(x, self.weight, self.eps)

View File

@ -1,33 +0,0 @@
import torch
import torch.nn as nn
class Model(nn.Module):
"""使用 PyTorch RMSNorm 的基准实现。"""
def __init__(self, weight: torch.Tensor, eps: float = 1e-6):
super().__init__()
if weight.dim() != 1:
raise ValueError("RMSNorm 权重必须是一维向量。")
feature_dim = weight.shape[0]
self.rmsnorm = nn.RMSNorm(feature_dim, eps=eps, elementwise_affine=True)
with torch.no_grad():
self.rmsnorm.weight.copy_(weight)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""直接对输入做 RMSNorm输出形状与输入一致。"""
return self.rmsnorm(x)
batch_size = 16
feature_dim = 2048
def get_inputs():
x = torch.randn(batch_size, feature_dim)
return [x]
def get_init_inputs():
weight = torch.randn(feature_dim)
return [weight]

View File

@ -1,88 +0,0 @@
###########################################################
# 性能和精度验证程序
###########################################################
import torch
import torch.nn as nn
import time
from rmsnorm_torch import Model, get_inputs, get_init_inputs
from rmsnorm_cuda import ModelNew
def run_benchmark():
# 检查 CUDA 是否可用
if not torch.cuda.is_available():
print("CUDA 不可用,请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
return
else:
device = torch.device("cuda")
# 初始化模型
init_inputs = get_init_inputs()
init_inputs = [
x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in init_inputs
]
inputs = get_inputs()
inputs = [
x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in inputs
]
torch_model = Model(*init_inputs).cuda()
cuda_model = ModelNew(*init_inputs).cuda()
torch_model.eval()
cuda_model.eval()
print("-------------------- 精度对齐验证 --------------------")
with torch.no_grad():
output_torch = torch_model(*inputs)
output_cuda = cuda_model(*inputs)
# 更严格的精度检查
abs_diff = (output_torch - output_cuda).abs()
max_diff = abs_diff.max().item()
mean_diff = abs_diff.mean().item()
print(f"最大差异: {max_diff:.6f}")
print(f"平均差异: {mean_diff:.6f}")
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-05, atol=1e-05)
if precision_flag:
print("✅ 精度对齐:两个模型的输出结果非常接近。")
else:
print("❌ 精度不一致!")
print("\n-------------------- 性能加速比测试 --------------------")
num_iterations = 1000 # 增加迭代次数以获得更准确的时间测量
# Warm up
for _ in range(100):
_ = torch_model(*inputs)
_ = cuda_model(*inputs)
# PyTorch 模型计时
torch.cuda.synchronize()
start_time = time.time()
for _ in range(num_iterations):
_ = torch_model(*inputs)
torch.cuda.synchronize()
torch_time = (time.time() - start_time) / num_iterations
# 自定义 CUDA 内核计时
torch.cuda.synchronize()
start_time = time.time()
for _ in range(num_iterations):
_ = cuda_model(*inputs)
torch.cuda.synchronize()
cuda_time = (time.time() - start_time) / num_iterations
print(f"PyTorch 平均执行时间: {torch_time:.6f}")
print(f"自定义 CUDA ReLU 平均执行时间: {cuda_time:.6f}")
speedup = 0
if cuda_time > 0:
speedup = torch_time / cuda_time
print(f"加速比 (Speedup): {speedup:.2f}x")
else:
print("CUDA 内核执行时间为0无法计算加速比。")
return precision_flag, speedup
if __name__ == "__main__":
precision_flag, speedup = run_benchmark()

View File

@ -1,120 +0,0 @@
import torch
from torch.utils.cpp_extension import load_inline
source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <vector_types.h>
__global__ __launch_bounds__(256) void log1pabs_affine_gate_kernel(const float* __restrict__ x, const float* __restrict__ scale, const float* __restrict__ bias, float* __restrict__ y, int B, int D, float alpha, float beta){
int row = blockIdx.x;
int tid = threadIdx.x;
int col_block = blockIdx.y;
int stride4 = blockDim.x * 4;
int col_idx = (col_block * stride4) + tid * 4;
const float* xr = x + row * D;
float* yr = y + row * D;
for(int base = col_idx; base < D; base += stride4 * 2){
if (base < D){
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
float4 yv;
float z0 = fmaf(xv.x, sv.x, bv.x);
float z1 = fmaf(xv.y, sv.y, bv.y);
float z2 = fmaf(xv.z, sv.z, bv.z);
float z3 = fmaf(xv.w, sv.w, bv.w);
float m0 = log1pf(fabsf(z0));
float m1 = log1pf(fabsf(z1));
float m2 = log1pf(fabsf(z2));
float m3 = log1pf(fabsf(z3));
float t0 = fmaf(alpha, m0, beta);
float t1 = fmaf(alpha, m1, beta);
float t2 = fmaf(alpha, m2, beta);
float t3 = fmaf(alpha, m3, beta);
float g0 = __fdividef(1.0f, 1.0f + __expf(-t0));
float g1 = __fdividef(1.0f, 1.0f + __expf(-t1));
float g2 = __fdividef(1.0f, 1.0f + __expf(-t2));
float g3 = __fdividef(1.0f, 1.0f + __expf(-t3));
yv.x = xv.x * g0;
yv.y = xv.y * g1;
yv.z = xv.z * g2;
yv.w = xv.w * g3;
reinterpret_cast<float4*>(yr + base)[0] = yv;
}
int base2 = base + stride4;
if (base2 < D){
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
float4 yv2;
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
float mb0 = log1pf(fabsf(z0b));
float mb1 = log1pf(fabsf(z1b));
float mb2 = log1pf(fabsf(z2b));
float mb3 = log1pf(fabsf(z3b));
float tb0 = fmaf(alpha, mb0, beta);
float tb1 = fmaf(alpha, mb1, beta);
float tb2 = fmaf(alpha, mb2, beta);
float tb3 = fmaf(alpha, mb3, beta);
float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
yv2.x = xv2.x * gb0;
yv2.y = xv2.y * gb1;
yv2.z = xv2.z * gb2;
yv2.w = xv2.w * gb3;
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
}
}
}
torch::Tensor log1pabs_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta){
auto xc = x.contiguous();
auto sc = scale.contiguous();
auto bc = bias.contiguous();
auto y = torch::empty_like(xc);
int B = (int)xc.size(0);
int D = (int)xc.size(1);
float a = (float)alpha;
float be = (float)beta;
int block = 256;
int elements_per_thread = 4;
int elements_per_block = block * elements_per_thread;
int gy = (D + elements_per_block - 1) / elements_per_block;
dim3 grid(B, gy);
log1pabs_affine_gate_kernel<<<grid, block>>>(xc.data_ptr<float>(), sc.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), B, D, a, be);
return y;
}
"""
cpp_source = """
#include <torch/extension.h>
torch::Tensor log1pabs_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
"""
ops = load_inline(
name="log1pabs_affine_gate",
cpp_sources=cpp_source,
cuda_sources=source,
functions=["log1pabs_affine_gate_cuda"],
extra_cflags=["-O3","-std=c++17"],
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-Xptxas","-O3,-dlcm=ca","-maxrregcount=64"],
verbose=True
)
class ModelNew(torch.nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(ModelNew, self).__init__()
self.ops = ops
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.alpha = float(alpha)
self.beta = float(beta)
def forward(self, x):
return self.ops.log1pabs_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)

View File

@ -1,10 +0,0 @@
Operator: Log1pAbs-Affine-Gate (Fused CUDA Kernel)
Definition
- z = x * scale + bias
- m = log(1 + |z|)
- g = sigmoid(alpha * m + beta)
- y = x * g
Goal
- Stable log1p fusion; target ≥1.30x speedup.

View File

@ -1,50 +0,0 @@
import torch
import time
from torchcode import Model, get_inputs, get_init_inputs
from cudacode import ModelNew
def run_benchmark():
if not torch.cuda.is_available():
print("CUDA 不可用,请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
return
device = torch.device("cuda")
init_inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_init_inputs()]
inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_inputs()]
torch_model = Model(*init_inputs).cuda().eval()
cuda_model = ModelNew(*init_inputs).cuda().eval()
print("-------------------- 精度对齐验证 --------------------")
with torch.no_grad():
output_torch = torch_model(*inputs)
output_cuda = cuda_model(*inputs)
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-03)
if precision_flag:
print("✅ 精度对齐:两个模型的输出结果非常接近。")
else:
print("❌ 精度不一致!")
print("\n-------------------- 性能加速比测试 --------------------")
num_iterations = 100
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = torch_model(*inputs)
torch.cuda.synchronize(); torch_time = (time.time() - start_time) / num_iterations
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = cuda_model(*inputs)
torch.cuda.synchronize(); cuda_time = (time.time() - start_time) / num_iterations
print(f"PyTorch Log1pAbs-Affine-Gate 平均执行时间: {torch_time:.6f}")
print(f"自定义 CUDA 融合内核 平均执行时间: {cuda_time:.6f}")
speedup = torch_time / cuda_time if cuda_time > 0 else 0
if cuda_time > 0:
print(f"加速比 (Speedup): {speedup:.2f}x")
else:
print("CUDA 内核执行时间为0无法计算加速比。")
return precision_flag, speedup
if __name__ == "__main__":
run_benchmark()

View File

@ -1,28 +0,0 @@
import torch
import torch.nn as nn
class Model(nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(Model, self).__init__()
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.register_buffer("alpha", torch.tensor(float(alpha), dtype=torch.float32))
self.register_buffer("beta", torch.tensor(float(beta), dtype=torch.float32))
def forward(self, x: torch.Tensor) -> torch.Tensor:
z = x * self.scale + self.bias
m = torch.log1p(torch.abs(z))
g = torch.sigmoid(self.alpha * m + self.beta)
return x * g
batch_size = 16
dim = 16384
def get_inputs():
x = torch.randn(batch_size, dim)
return [x]
def get_init_inputs():
scale = torch.randn(dim)
bias = torch.randn(dim)
return [scale, bias, 1.0, 0.0]

View File

@ -1,94 +0,0 @@
import torch
from torch.utils.cpp_extension import load_inline
source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
__global__ void square_sigmoid_affine_gate_kernel(const float* __restrict__ x, const float* __restrict__ scale, const float* __restrict__ bias, float* __restrict__ y, int B, int D, float alpha, float beta){
int b = blockIdx.x;
int lane = blockIdx.y * blockDim.x + threadIdx.x;
int stride = blockDim.x * gridDim.y;
int row_start = b * D;
const float* xr = x + row_start;
float* yr = y + row_start;
int aligned = ((((long long)xr & 15LL) == 0) && (((long long)yr & 15LL) == 0) && (((long long)scale & 15LL) == 0) && (((long long)bias & 15LL) == 0) && ((D & 3) == 0));
if(aligned){
int D4 = (D / 4) * 4;
#pragma unroll 4
for(int i = lane * 4; i < D4; i += stride * 4){
float4 xv = reinterpret_cast<const float4*>(xr)[i / 4];
float4 sv = reinterpret_cast<const float4*>(scale)[i / 4];
float4 bv = reinterpret_cast<const float4*>(bias)[i / 4];
float4 yv;
float z0 = fmaf(xv.x, sv.x, bv.x);
float z1 = fmaf(xv.y, sv.y, bv.y);
float z2 = fmaf(xv.z, sv.z, bv.z);
float z3 = fmaf(xv.w, sv.w, bv.w);
float g0 = 1.0f / (1.0f + expf(-(alpha * (z0 * z0) + beta)));
float g1 = 1.0f / (1.0f + expf(-(alpha * (z1 * z1) + beta)));
float g2 = 1.0f / (1.0f + expf(-(alpha * (z2 * z2) + beta)));
float g3 = 1.0f / (1.0f + expf(-(alpha * (z3 * z3) + beta)));
yv.x = xv.x * g0;
yv.y = xv.y * g1;
yv.z = xv.z * g2;
yv.w = xv.w * g3;
reinterpret_cast<float4*>(yr)[i / 4] = yv;
}
#pragma unroll 4
for(int i = D4 + lane; i < D; i += stride){
float z = fmaf(xr[i], scale[i], bias[i]);
float g = 1.0f / (1.0f + expf(-(alpha * (z * z) + beta)));
yr[i] = xr[i] * g;
}
} else {
#pragma unroll 4
for(int i = lane; i < D; i += stride){
float z = fmaf(xr[i], scale[i], bias[i]);
float g = 1.0f / (1.0f + expf(-(alpha * (z * z) + beta)));
yr[i] = xr[i] * g;
}
}
}
torch::Tensor square_sigmoid_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor alpha, torch::Tensor beta){
auto xc = x.contiguous();
auto sc = scale.contiguous();
auto bc = bias.contiguous();
auto y = torch::empty_like(xc);
int B = (int)xc.size(0);
int D = (int)xc.size(1);
float a = alpha.item<float>();
float be = beta.item<float>();
int block = 256;
int gy = max(1, min((D + 4095) / 4096, 8));
dim3 grid(B, gy);
square_sigmoid_affine_gate_kernel<<<grid, block>>>(xc.data_ptr<float>(), sc.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), B, D, a, be);
return y;
}
"""
cpp_source = """
torch::Tensor square_sigmoid_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor alpha, torch::Tensor beta);
"""
ops = load_inline(
name="square_sigmoid_affine_gate",
cpp_sources=cpp_source,
cuda_sources=source,
functions=["square_sigmoid_affine_gate_cuda"],
extra_cuda_cflags=["-O3","--use_fast_math"],
verbose=True
)
class ModelNew(torch.nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(ModelNew, self).__init__()
self.ops = ops
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.register_buffer("alpha", torch.tensor(float(alpha), dtype=torch.float32))
self.register_buffer("beta", torch.tensor(float(beta), dtype=torch.float32))
def forward(self, x):
return self.ops.square_sigmoid_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)

View File

@ -1,28 +0,0 @@
融合算子Square-Sigmoid-Affine-Gate一次核内完成仿射、平方与 Sigmoid 门控,返回 y = x * σ(α * z^2 + β),其中 z = x*scale + bias。平方增强幅值差异并通过 Sigmoid 控制门控强度。
目标与定义
- 输入张量:`x[B, D]`
- 逐维参数:`scale[D]`、`bias[D]`
- 标量超参:`alpha`、`beta`
- 计算流程:`z = x*scale + bias``v = z*z``g = sigmoid(alpha*v + beta)``y = x * g`
参考实现(文件要求)
- `torchcode.py`PyTorch 参考 `Model`;统一的 `get_inputs()`/`get_init_inputs()`
- `cudacode.py`:单核融合(仿射+平方+sigmoid+乘法);`-O3 --use_fast_math`
- `run_code.py`:迭代 100 次;`rtol=1e-03, atol=1e-06` 精度;打印加速比
CUDA 实现要点
- 并行布局:`grid = B`;块内沿 D 合并访存
- 对齐向量化16 字节对齐且 `D%4==0` 时走 `float4`;否则标量回退
- 指令优化:仿射用 `fmaf`sigmoid 用 `expf`;循环 `#pragma unroll 4`
- 溢出注意:`z^2` 对大幅值会放大sigmoid 可缓和,但仍需避免中间溢出;使用 `float` 常规范围下问题不大
- 线程配置:推荐 `block=1024`,按设备与规模微调
评估与目标
- 精度:对齐 `rtol=1e-03, atol=1e-06`
- 性能≥1.0x 加速;对齐触发向量化时更佳
加分项(可选)
- 尾元素处理与分支收敛优化
- 每线程批量步长以提高吞吐与占用

View File

@ -1,57 +0,0 @@
import torch
import time
from torchcode import Model, get_inputs, get_init_inputs
from cudacode import ModelNew
def run_benchmark():
if not torch.cuda.is_available():
print("CUDA 不可用,请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
return
device = torch.device("cuda")
init_inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_init_inputs()]
inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_inputs()]
torch_model = Model(*init_inputs).cuda()
cuda_model = ModelNew(*init_inputs).cuda()
torch_model.eval(); cuda_model.eval()
print("-------------------- 精度对齐验证 --------------------")
with torch.no_grad():
output_torch = torch_model(*inputs)
output_cuda = cuda_model(*inputs)
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-03, atol=1e-06)
if precision_flag:
print("✅ 精度对齐:两个模型的输出结果非常接近。")
else:
print("❌ 精度不一致!")
diff = (output_torch - output_cuda).abs().max().item()
print(f"最大绝对误差: {diff}")
print(f"输出张量形状: torch={tuple(output_torch.shape)}, cuda={tuple(output_cuda.shape)}")
print(f"数据类型: torch={output_torch.dtype}, cuda={output_cuda.dtype}")
print(f"设备: torch={output_torch.device}, cuda={output_cuda.device}")
print("\n-------------------- 性能加速比测试 --------------------")
num_iterations = 100
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = torch_model(*inputs)
torch.cuda.synchronize(); torch_time = (time.time() - start_time) / num_iterations
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = cuda_model(*inputs)
torch.cuda.synchronize(); cuda_time = (time.time() - start_time) / num_iterations
print(f"PyTorch Square-Sigmoid-Affine-Gate 平均执行时间: {torch_time:.6f}")
print(f"自定义 CUDA 融合内核 平均执行时间: {cuda_time:.6f}")
speedup = torch_time / cuda_time if cuda_time > 0 else 0
if cuda_time > 0:
print(f"加速比 (Speedup): {speedup:.2f}x")
else:
print("CUDA 内核执行时间为0无法计算加速比。")
return precision_flag, speedup
if __name__ == "__main__":
run_benchmark()

View File

@ -1,29 +0,0 @@
import torch
import torch.nn as nn
class Model(nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(Model, self).__init__()
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.register_buffer("alpha", torch.tensor(float(alpha), dtype=torch.float32))
self.register_buffer("beta", torch.tensor(float(beta), dtype=torch.float32))
def forward(self, x: torch.Tensor) -> torch.Tensor:
z = x * self.scale + self.bias
v = z * z
g = torch.sigmoid(self.alpha * v + self.beta)
return x * g
batch_size = 16
dim = 16384
def get_inputs():
x = torch.randn(batch_size, dim)
return [x]
def get_init_inputs():
scale = torch.randn(dim)
bias = torch.randn(dim)
return [scale, bias, 1.0, 0.0]

View File

@ -1,94 +0,0 @@
import torch
from torch.utils.cpp_extension import load_inline
source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
__global__ void atan_sigmoid_mix_gate_kernel(const float* __restrict__ x, const float* __restrict__ scale, const float* __restrict__ bias, float* __restrict__ y, int B, int D, float alpha, float beta){
int b = blockIdx.x;
int lane = blockIdx.y * blockDim.x + threadIdx.x;
int stride = blockDim.x * gridDim.y;
int row_start = b * D;
const float* xr = x + row_start;
float* yr = y + row_start;
int aligned = ((((long long)xr & 15LL) == 0) && (((long long)yr & 15LL) == 0) && (((long long)scale & 15LL) == 0) && (((long long)bias & 15LL) == 0) && ((D & 3) == 0));
if(aligned){
int D4 = (D / 4) * 4;
#pragma unroll 4
for(int i = lane * 4; i < D4; i += stride * 4){
float4 xv = reinterpret_cast<const float4*>(xr)[i / 4];
float4 sv = reinterpret_cast<const float4*>(scale)[i / 4];
float4 bv = reinterpret_cast<const float4*>(bias)[i / 4];
float4 yv;
float z0 = fmaf(xv.x, sv.x, bv.x);
float z1 = fmaf(xv.y, sv.y, bv.y);
float z2 = fmaf(xv.z, sv.z, bv.z);
float z3 = fmaf(xv.w, sv.w, bv.w);
float g0 = 1.0f / (1.0f + expf(-(alpha * atanf(z0) + beta)));
float g1 = 1.0f / (1.0f + expf(-(alpha * atanf(z1) + beta)));
float g2 = 1.0f / (1.0f + expf(-(alpha * atanf(z2) + beta)));
float g3 = 1.0f / (1.0f + expf(-(alpha * atanf(z3) + beta)));
yv.x = xv.x * g0;
yv.y = xv.y * g1;
yv.z = xv.z * g2;
yv.w = xv.w * g3;
reinterpret_cast<float4*>(yr)[i / 4] = yv;
}
#pragma unroll 4
for(int i = D4 + lane; i < D; i += stride){
float z = fmaf(xr[i], scale[i], bias[i]);
float g = 1.0f / (1.0f + expf(-(alpha * atanf(z) + beta)));
yr[i] = xr[i] * g;
}
} else {
#pragma unroll 4
for(int i = lane; i < D; i += stride){
float z = fmaf(xr[i], scale[i], bias[i]);
float g = 1.0f / (1.0f + expf(-(alpha * atanf(z) + beta)));
yr[i] = xr[i] * g;
}
}
}
torch::Tensor atan_sigmoid_mix_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor alpha, torch::Tensor beta){
auto xc = x.contiguous();
auto sc = scale.contiguous();
auto bc = bias.contiguous();
auto y = torch::empty_like(xc);
int B = (int)xc.size(0);
int D = (int)xc.size(1);
float a = alpha.item<float>();
float be = beta.item<float>();
int block = 256;
int gy = max(1, min((D + 4095) / 4096, 8));
dim3 grid(B, gy);
atan_sigmoid_mix_gate_kernel<<<grid, block>>>(xc.data_ptr<float>(), sc.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), B, D, a, be);
return y;
}
"""
cpp_source = """
torch::Tensor atan_sigmoid_mix_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor alpha, torch::Tensor beta);
"""
ops = load_inline(
name="atan_sigmoid_mix_gate",
cpp_sources=cpp_source,
cuda_sources=source,
functions=["atan_sigmoid_mix_gate_cuda"],
extra_cuda_cflags=["-O3","--use_fast_math"],
verbose=True
)
class ModelNew(torch.nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(ModelNew, self).__init__()
self.ops = ops
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.register_buffer("alpha", torch.tensor(float(alpha), dtype=torch.float32))
self.register_buffer("beta", torch.tensor(float(beta), dtype=torch.float32))
def forward(self, x):
return self.ops.atan_sigmoid_mix_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)

View File

@ -1,27 +0,0 @@
融合算子Atan-Sigmoid-Mix-Gate一次核内完成仿射与反正切混合门控返回 y = x * σ(α * atan(z) + β),其中 z = x*scale + bias。atan 在大幅值区间趋于常数,有利于抑制过大输入的门控强度。
目标与定义
- 输入张量:`x[B, D]`
- 逐维参数:`scale[D]`、`bias[D]`
- 标量超参:`alpha`、`beta`
- 计算流程:`z = x*scale + bias``g = sigmoid(alpha*atan(z) + beta)``y = x * g`
参考实现(文件要求)
- `torchcode.py`PyTorch 参考 `Model`;统一接口
- `cudacode.py`:单核融合(仿射+atan+sigmoid+乘法);`-O3 --use_fast_math`
- `run_code.py`100 次迭代;精度 `rtol=1e-03, atol=1e-06`;打印加速比
CUDA 实现要点
- 行并行:`grid = B`;块内沿 D 连续访存;一次遍历写回
- 对齐向量化16 字节对齐且 `D%4==0` 走 `float4`,否则标量回退
- 指令优化:仿射用 `fmaf``atanf` 与 `expf` 走快速数学;循环 `#pragma unroll 4`
- 线程配置:`block=1024` 起步,按设备试探最佳
评估与目标
- 精度:满足 `rtol=1e-03, atol=1e-06`
- 性能≥1.0x 加速,向量化与融合带来优势
加分项(可选)
- 尾元素处理与分支收敛优化
- 根据分布特性调参 `alpha/beta` 增强门控稳定性(参考实现一致性优先)

View File

@ -1,57 +0,0 @@
import torch
import time
from torchcode import Model, get_inputs, get_init_inputs
from cudacode import ModelNew
def run_benchmark():
if not torch.cuda.is_available():
print("CUDA 不可用,请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
return
device = torch.device("cuda")
init_inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_init_inputs()]
inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_inputs()]
torch_model = Model(*init_inputs).cuda()
cuda_model = ModelNew(*init_inputs).cuda()
torch_model.eval(); cuda_model.eval()
print("-------------------- 精度对齐验证 --------------------")
with torch.no_grad():
output_torch = torch_model(*inputs)
output_cuda = cuda_model(*inputs)
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-03, atol=1e-06)
if precision_flag:
print("✅ 精度对齐:两个模型的输出结果非常接近。")
else:
print("❌ 精度不一致!")
diff = (output_torch - output_cuda).abs().max().item()
print(f"最大绝对误差: {diff}")
print(f"输出张量形状: torch={tuple(output_torch.shape)}, cuda={tuple(output_cuda.shape)}")
print(f"数据类型: torch={output_torch.dtype}, cuda={output_cuda.dtype}")
print(f"设备: torch={output_torch.device}, cuda={output_cuda.device}")
print("\n-------------------- 性能加速比测试 --------------------")
num_iterations = 100
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = torch_model(*inputs)
torch.cuda.synchronize(); torch_time = (time.time() - start_time) / num_iterations
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = cuda_model(*inputs)
torch.cuda.synchronize(); cuda_time = (time.time() - start_time) / num_iterations
print(f"PyTorch Atan-Sigmoid-Mix-Gate 平均执行时间: {torch_time:.6f}")
print(f"自定义 CUDA 融合内核 平均执行时间: {cuda_time:.6f}")
speedup = torch_time / cuda_time if cuda_time > 0 else 0
if cuda_time > 0:
print(f"加速比 (Speedup): {speedup:.2f}x")
else:
print("CUDA 内核执行时间为0无法计算加速比。")
return precision_flag, speedup
if __name__ == "__main__":
run_benchmark()

View File

@ -1,28 +0,0 @@
import torch
import torch.nn as nn
class Model(nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(Model, self).__init__()
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.register_buffer("alpha", torch.tensor(float(alpha), dtype=torch.float32))
self.register_buffer("beta", torch.tensor(float(beta), dtype=torch.float32))
def forward(self, x: torch.Tensor) -> torch.Tensor:
z = x * self.scale + self.bias
g = torch.sigmoid(self.alpha * torch.atan(z) + self.beta)
return x * g
batch_size = 16
dim = 16384
def get_inputs():
x = torch.randn(batch_size, dim)
return [x]
def get_init_inputs():
scale = torch.randn(dim)
bias = torch.randn(dim)
return [scale, bias, 1.0, 0.0]

View File

@ -1,126 +0,0 @@
import torch
from torch.utils.cpp_extension import load_inline
source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <vector_types.h>
__device__ __forceinline__ float huberf(float z){
float az = fabsf(z);
if (az <= 1.0f) return 0.5f * z * z;
return az - 0.5f;
}
__global__ __launch_bounds__(256) void huber_affine_gate_kernel(const float* __restrict__ x, const float* __restrict__ scale, const float* __restrict__ bias, float* __restrict__ y, int B, int D, float alpha, float beta){
int row = blockIdx.x;
int tid = threadIdx.x;
int col_block = blockIdx.y;
int stride4 = blockDim.x * 4;
int col_idx = (col_block * stride4) + tid * 4;
const float* xr = x + row * D;
float* yr = y + row * D;
for(int base = col_idx; base < D; base += stride4 * 2){
if (base < D){
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
float4 yv;
float z0 = fmaf(xv.x, sv.x, bv.x);
float z1 = fmaf(xv.y, sv.y, bv.y);
float z2 = fmaf(xv.z, sv.z, bv.z);
float z3 = fmaf(xv.w, sv.w, bv.w);
float m0 = huberf(z0);
float m1 = huberf(z1);
float m2 = huberf(z2);
float m3 = huberf(z3);
float t0 = fmaf(alpha, m0, beta);
float t1 = fmaf(alpha, m1, beta);
float t2 = fmaf(alpha, m2, beta);
float t3 = fmaf(alpha, m3, beta);
float g0 = __fdividef(1.0f, 1.0f + __expf(-t0));
float g1 = __fdividef(1.0f, 1.0f + __expf(-t1));
float g2 = __fdividef(1.0f, 1.0f + __expf(-t2));
float g3 = __fdividef(1.0f, 1.0f + __expf(-t3));
yv.x = xv.x * g0;
yv.y = xv.y * g1;
yv.z = xv.z * g2;
yv.w = xv.w * g3;
reinterpret_cast<float4*>(yr + base)[0] = yv;
}
int base2 = base + stride4;
if (base2 < D){
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
float4 yv2;
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
float mb0 = huberf(z0b);
float mb1 = huberf(z1b);
float mb2 = huberf(z2b);
float mb3 = huberf(z3b);
float tb0 = fmaf(alpha, mb0, beta);
float tb1 = fmaf(alpha, mb1, beta);
float tb2 = fmaf(alpha, mb2, beta);
float tb3 = fmaf(alpha, mb3, beta);
float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
yv2.x = xv2.x * gb0;
yv2.y = xv2.y * gb1;
yv2.z = xv2.z * gb2;
yv2.w = xv2.w * gb3;
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
}
}
}
torch::Tensor huber_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta){
auto xc = x.contiguous();
auto sc = scale.contiguous();
auto bc = bias.contiguous();
auto y = torch::empty_like(xc);
int B = (int)xc.size(0);
int D = (int)xc.size(1);
float a = (float)alpha;
float be = (float)beta;
int block = 256;
int elements_per_thread = 4;
int elements_per_block = block * elements_per_thread;
int gy = (D + elements_per_block - 1) / elements_per_block;
dim3 grid(B, gy);
huber_affine_gate_kernel<<<grid, block>>>(xc.data_ptr<float>(), sc.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), B, D, a, be);
return y;
}
"""
cpp_source = """
#include <torch/extension.h>
torch::Tensor huber_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
"""
ops = load_inline(
name="huber_affine_gate",
cpp_sources=cpp_source,
cuda_sources=source,
functions=["huber_affine_gate_cuda"],
extra_cflags=["-O3","-std=c++17"],
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-Xptxas","-O3,-dlcm=ca","-maxrregcount=64"],
verbose=True
)
class ModelNew(torch.nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(ModelNew, self).__init__()
self.ops = ops
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.alpha = float(alpha)
self.beta = float(beta)
def forward(self, x):
return self.ops.huber_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)

View File

@ -1,10 +0,0 @@
Operator: Huber-Affine-Gate (Fused CUDA Kernel)
Definition
- z = x * scale + bias
- m = 0.5*z^2 if |z|<=1 else |z| - 0.5
- g = sigmoid(alpha * m + beta)
- y = x * g
Goal
- Piecewise smooth fusion; target ≥1.30x speedup.

View File

@ -1,50 +0,0 @@
import torch
import time
from torchcode import Model, get_inputs, get_init_inputs
from cudacode import ModelNew
def run_benchmark():
if not torch.cuda.is_available():
print("CUDA 不可用,请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
return
device = torch.device("cuda")
init_inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_init_inputs()]
inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_inputs()]
torch_model = Model(*init_inputs).cuda().eval()
cuda_model = ModelNew(*init_inputs).cuda().eval()
print("-------------------- 精度对齐验证 --------------------")
with torch.no_grad():
output_torch = torch_model(*inputs)
output_cuda = cuda_model(*inputs)
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-03)
if precision_flag:
print("✅ 精度对齐:两个模型的输出结果非常接近。")
else:
print("❌ 精度不一致!")
print("\n-------------------- 性能加速比测试 --------------------")
num_iterations = 100
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = torch_model(*inputs)
torch.cuda.synchronize(); torch_time = (time.time() - start_time) / num_iterations
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = cuda_model(*inputs)
torch.cuda.synchronize(); cuda_time = (time.time() - start_time) / num_iterations
print(f"PyTorch Huber-Affine-Gate 平均执行时间: {torch_time:.6f}")
print(f"自定义 CUDA 融合内核 平均执行时间: {cuda_time:.6f}")
speedup = torch_time / cuda_time if cuda_time > 0 else 0
if cuda_time > 0:
print(f"加速比 (Speedup): {speedup:.2f}x")
else:
print("CUDA 内核执行时间为0无法计算加速比。")
return precision_flag, speedup
if __name__ == "__main__":
run_benchmark()

View File

@ -1,29 +0,0 @@
import torch
import torch.nn as nn
class Model(nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(Model, self).__init__()
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.register_buffer("alpha", torch.tensor(float(alpha), dtype=torch.float32))
self.register_buffer("beta", torch.tensor(float(beta), dtype=torch.float32))
def forward(self, x: torch.Tensor) -> torch.Tensor:
z = x * self.scale + self.bias
az = torch.abs(z)
m = torch.where(az <= 1.0, 0.5 * z * z, az - 0.5)
g = torch.sigmoid(self.alpha * m + self.beta)
return x * g
batch_size = 16
dim = 16384
def get_inputs():
x = torch.randn(batch_size, dim)
return [x]
def get_init_inputs():
scale = torch.randn(dim)
bias = torch.randn(dim)
return [scale, bias, 1.0, 0.0]

View File

@ -1,190 +0,0 @@
import torch
from torch.utils.cpp_extension import load_inline
# LayerNorm的CUDA实现
layernorm_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cub/cub.cuh>
// 使用高度优化的LayerNorm实现结合向量化和内存访问优化
__global__ void layernorm_forward_kernel(
const float* __restrict__ input,
const float* __restrict__ gamma,
const float* __restrict__ beta,
float* __restrict__ output,
int batch_size,
int hidden_size,
float eps) {
extern __shared__ float shared_mem[];
float* shared_sum = shared_mem;
float* shared_sum_sq = &shared_mem[blockDim.x];
int batch_idx = blockIdx.x;
int tid = threadIdx.x;
// 使用向量化加载每个线程处理4个元素
float4 thread_sum = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
float4 thread_sum_sq = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
// 向量化处理提高内存带宽利用率
for (int i = tid * 4; i < hidden_size; i += blockDim.x * 4) {
if (i + 3 < hidden_size) {
float4 vals = *reinterpret_cast<const float4*>(&input[batch_idx * hidden_size + i]);
thread_sum.x += vals.x; thread_sum_sq.x += vals.x * vals.x;
thread_sum.y += vals.y; thread_sum_sq.y += vals.y * vals.y;
thread_sum.z += vals.z; thread_sum_sq.z += vals.z * vals.z;
thread_sum.w += vals.w; thread_sum_sq.w += vals.w * vals.w;
} else {
// 处理剩余元素
for (int j = 0; j < 4 && i + j < hidden_size; j++) {
float val = input[batch_idx * hidden_size + i + j];
thread_sum.x += val; thread_sum_sq.x += val * val;
}
}
}
// 归约线程内的4个分量
float thread_total_sum = thread_sum.x + thread_sum.y + thread_sum.z + thread_sum.w;
float thread_total_sum_sq = thread_sum_sq.x + thread_sum_sq.y + thread_sum_sq.z + thread_sum_sq.w;
shared_sum[tid] = thread_total_sum;
shared_sum_sq[tid] = thread_total_sum_sq;
__syncthreads();
// 使用更高效的归约算法树形归约
for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
if (tid < stride) {
shared_sum[tid] += shared_sum[tid + stride];
shared_sum_sq[tid] += shared_sum_sq[tid + stride];
}
__syncthreads();
}
// 计算全局统计量
if (tid == 0) {
float total_sum = shared_sum[0];
float total_sum_sq = shared_sum_sq[0];
float global_mean = total_sum / hidden_size;
float global_variance = (total_sum_sq / hidden_size) - (global_mean * global_mean);
// 计算逆标准差
float inv_std = rsqrtf(global_variance + eps);
// 存储到共享内存供所有线程使用
shared_sum[0] = global_mean;
shared_sum_sq[0] = inv_std;
}
__syncthreads();
float global_mean = shared_sum[0];
float inv_std = shared_sum_sq[0];
// 应用LayerNorm使用向量化存储
for (int i = tid * 4; i < hidden_size; i += blockDim.x * 4) {
if (i + 3 < hidden_size) {
float4 vals = *reinterpret_cast<const float4*>(&input[batch_idx * hidden_size + i]);
float4 normalized;
normalized.x = (vals.x - global_mean) * inv_std;
normalized.y = (vals.y - global_mean) * inv_std;
normalized.z = (vals.z - global_mean) * inv_std;
normalized.w = (vals.w - global_mean) * inv_std;
float4 result;
result.x = normalized.x * gamma[i] + beta[i];
result.y = normalized.y * gamma[i+1] + beta[i+1];
result.z = normalized.z * gamma[i+2] + beta[i+2];
result.w = normalized.w * gamma[i+3] + beta[i+3];
*reinterpret_cast<float4*>(&output[batch_idx * hidden_size + i]) = result;
} else {
// 处理剩余元素
for (int j = 0; j < 4 && i + j < hidden_size; j++) {
float val = input[batch_idx * hidden_size + i + j];
float normalized = (val - global_mean) * inv_std;
output[batch_idx * hidden_size + i + j] = normalized * gamma[i + j] + beta[i + j];
}
}
}
}
torch::Tensor layernorm_cuda_forward(
torch::Tensor input,
torch::Tensor gamma,
torch::Tensor beta,
float eps) {
auto batch_size = input.size(0);
auto hidden_size = input.size(-1);
auto output = torch::empty_like(input);
// 优化线程块大小根据hidden_size动态调整使用更激进的优化
int block_size = 256; // 固定使用256线程适合大多数GPU架构
if (hidden_size <= 512) {
block_size = 128;
} else if (hidden_size <= 1024) {
block_size = 256;
} else {
block_size = 512;
}
// 确保block_size不超过硬件限制
block_size = min(1024, max(32, block_size));
int num_blocks = batch_size;
int shared_mem_size = 2 * block_size * sizeof(float);
layernorm_forward_kernel<<<num_blocks, block_size, shared_mem_size>>>(
input.data_ptr<float>(),
gamma.data_ptr<float>(),
beta.data_ptr<float>(),
output.data_ptr<float>(),
batch_size,
hidden_size,
eps
);
return output;
}
"""
layernorm_cpp_source = """
torch::Tensor layernorm_cuda_forward(torch::Tensor input, torch::Tensor gamma, torch::Tensor beta, float eps);
"""
# 编译内联CUDA代码
cuda_available = True
try:
layernorm_cuda = load_inline(
name="layernorm_cuda",
cpp_sources=layernorm_cpp_source,
cuda_sources=layernorm_source,
functions=["layernorm_cuda_forward"],
verbose=True
)
except Exception as e:
print(f"CUDA扩展加载失败: {e}")
cuda_available = False
layernorm_cuda = None
class ModelNew(torch.nn.Module):
def __init__(self, normalized_shape, eps=1e-5):
super(ModelNew, self).__init__()
self.normalized_shape = normalized_shape
self.eps = eps
self.weight = torch.nn.Parameter(torch.ones(normalized_shape))
self.bias = torch.nn.Parameter(torch.zeros(normalized_shape))
def forward(self, x):
if cuda_available and layernorm_cuda is not None:
# 使用真正的CUDA内核
return layernorm_cuda.layernorm_cuda_forward(x, self.weight, self.bias, self.eps)
else:
# CPU回退实现与PyTorch实现保持一致
mean = x.mean(-1, keepdim=True)
var = x.var(-1, unbiased=False, keepdim=True)
normalized = (x - mean) / torch.sqrt(var + self.eps)
return normalized * self.weight + self.bias

View File

@ -1,23 +0,0 @@
import torch
import torch.nn as nn
class Model(nn.Module):
def __init__(self, normalized_shape, eps=1e-5):
super(Model, self).__init__()
self.normalized_shape = normalized_shape
self.eps = eps
# 使用PyTorch内置的LayerNorm这是标准实现
self.layer_norm = nn.LayerNorm(normalized_shape, eps=eps)
def forward(self, x):
# 使用PyTorch内置LayerNorm这是标准的实现方式
return self.layer_norm(x)
def get_inputs():
# 使用更大的输入尺寸以获得更好的性能对比
# 增加batch size和hidden size模拟真实场景
return [torch.randn(16, 16384)]
def get_init_inputs():
return [16384]

View File

@ -1,27 +0,0 @@
Write a custom CUDA kernel for Layer Normalization.
The standard LayerNorm operation is defined as:
y = (x - E[x]) / sqrt(Var[x] + epsilon) * gamma + beta
Where:
- x is the input tensor
- E[x] is the mean of x
- Var[x] is the variance of x
- epsilon is a small value for numerical stability
- gamma and beta are learnable affine parameters
You should fuse the calculation of mean, variance, and the normalization into a single CUDA kernel. This avoids multiple passes over the data and reduces memory bandwidth usage.
You are given the following architecture:
import torch
import torch.nn as nn
class Model(nn.Module):
def __init__(self, normalized_shape, eps=1e-5):
super(Model, self).__init__()
self.layer_norm = nn.LayerNorm(normalized_shape, eps=eps)
def forward(self, x):
return self.layer_norm(x)

View File

@ -1,74 +0,0 @@
###########################################################
# 性能和精度验证程序
###########################################################
import torch
import torch.nn as nn
import time
from layernorm_torchcode import Model,get_inputs,get_init_inputs
from layernorm_cudacode import ModelNew
def run_benchmark():
# 检查 CUDA 是否可用
if not torch.cuda.is_available():
print("CUDA 不可用,请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
return
else:
device = torch.device("cuda")
# 初始化模型
init_inputs = get_init_inputs()
init_inputs = [
x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in init_inputs
]
inputs = get_inputs()
inputs = [
x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in inputs
]
torch_model = Model(*init_inputs).cuda()
cuda_model = ModelNew(*init_inputs).cuda()
torch_model.eval()
cuda_model.eval()
print("-------------------- 精度对齐验证 --------------------")
with torch.no_grad():
output_torch = torch_model( *inputs)
output_cuda = cuda_model(*inputs)
precision_flag = torch.allclose(output_torch, output_cuda,rtol=1e-03)
if precision_flag:
print("✅ 精度对齐:两个模型的输出结果非常接近。")
else:
print("❌ 精度不一致!")
print("\n-------------------- 性能加速比测试 --------------------")
num_iterations = 100
# PyTorch 模型计时
torch.cuda.synchronize()
start_time = time.time()
for _ in range(num_iterations):
_ = torch_model(*inputs)
torch.cuda.synchronize()
torch_time = (time.time() - start_time) / num_iterations
# 自定义 CUDA 内核计时
torch.cuda.synchronize()
start_time = time.time()
for _ in range(num_iterations):
_ = cuda_model(*inputs)
torch.cuda.synchronize()
cuda_time = (time.time() - start_time) / num_iterations
print(f"PyTorch LayerNorm 平均执行时间: {torch_time:.6f}")
print(f"自定义 CUDA 内核 平均执行时间: {cuda_time:.6f}")
speedup = 0
if cuda_time > 0:
speedup = torch_time / cuda_time
print(f"加速比 (Speedup): {speedup:.2f}x")
else:
print("CUDA 内核执行时间为0无法计算加速比。")
return precision_flag,speedup
if __name__ == "__main__":
precision_flag,speedup = run_benchmark()

View File

@ -1,53 +0,0 @@
import torch
from torch.utils.cpp_extension import load_inline
source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
__global__ void bias_gelu_kernel(const float* x, const float* bias, float* y, int dim, long long total) {
long long idx = blockIdx.x * blockDim.x + threadIdx.x;
long long stride = blockDim.x * gridDim.x;
for (long long i = idx; i < total; i += stride) {
int j = (int)(i % dim);
float z = x[i] + bias[j];
float c = 0.7978845608f;
float p = z + 0.044715f * z * z * z;
y[i] = 0.5f * z * (1.f + tanhf(c * p));
}
}
torch::Tensor bias_gelu_cuda(torch::Tensor x, torch::Tensor bias) {
auto x_contig = x.contiguous();
auto b_contig = bias.contiguous();
auto y = torch::empty_like(x_contig);
long long total = x_contig.numel();
int dim = (int)x_contig.size(-1);
int block = 512;
long long grid = (total + block - 1) / block;
grid = grid > 65535 ? 65535 : grid;
bias_gelu_kernel<<<(int)grid, block>>>(x_contig.data_ptr<float>(), b_contig.data_ptr<float>(), y.data_ptr<float>(), dim, total);
return y;
}
"""
cpp_source = """
torch::Tensor bias_gelu_cuda(torch::Tensor x, torch::Tensor bias);
"""
ops = load_inline(
name="bias_gelu",
cpp_sources=cpp_source,
cuda_sources=source,
functions=["bias_gelu_cuda"],
verbose=True
)
class ModelNew(torch.nn.Module):
def __init__(self, bias: torch.Tensor):
super(ModelNew, self).__init__()
self.ops = ops
self.register_buffer("bias", bias)
def forward(self, x):
return self.ops.bias_gelu_cuda(x, self.bias)

View File

@ -1,7 +0,0 @@
本目录展示一个避免常见 LayerNorm 的融合算子Bias+GELUtanh 近似)。
torchcode.py 提供 PyTorch 参考实现:`y = gelu(x + bias)`,使用 `approximate='tanh'` 以匹配 CUDA 近似。
cudacode.py 内含 `__global__ void bias_gelu_kernel(...)`,一次遍历完成加偏置与 GELU 计算,减少显存往返与内核启动次数。
run_code.py 负责精度和性能对比,迭代 100 次并输出平均耗时与加速比,精度以 `rtol=1e-03` 检验。

View File

@ -1,76 +0,0 @@
###########################################################
# 性能和精度验证程序
###########################################################
import torch
import torch.nn as nn
import time
from torchcode import Model,get_inputs,get_init_inputs
from cudacode import ModelNew
def run_benchmark():
if not torch.cuda.is_available():
print("CUDA 不可用,请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
return
else:
device = torch.device("cuda")
init_inputs = get_init_inputs()
init_inputs = [
x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in init_inputs
]
inputs = get_inputs()
inputs = [
x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in inputs
]
torch_model = Model(*init_inputs).cuda()
cuda_model = ModelNew(*init_inputs).cuda()
torch_model.eval()
cuda_model.eval()
print("-------------------- 精度对齐验证 --------------------")
with torch.no_grad():
output_torch = torch_model(*inputs)
output_cuda = cuda_model(*inputs)
precision_flag = torch.allclose(output_torch, output_cuda,rtol=1e-03)
if precision_flag:
print("✅ 精度对齐:两个模型的输出结果非常接近。")
else:
print("❌ 精度不一致!")
diff = (output_torch - output_cuda).abs().max().item()
print(f"最大绝对误差: {diff}")
print(f"输出张量形状: torch={tuple(output_torch.shape)}, cuda={tuple(output_cuda.shape)}")
print(f"数据类型: torch={output_torch.dtype}, cuda={output_cuda.dtype}")
print(f"设备: torch={output_torch.device}, cuda={output_cuda.device}")
print("\n-------------------- 性能加速比测试 --------------------")
num_iterations = 100
torch.cuda.synchronize()
start_time = time.time()
for _ in range(num_iterations):
_ = torch_model(*inputs)
torch.cuda.synchronize()
torch_time = (time.time() - start_time) / num_iterations
torch.cuda.synchronize()
start_time = time.time()
for _ in range(num_iterations):
_ = cuda_model(*inputs)
torch.cuda.synchronize()
cuda_time = (time.time() - start_time) / num_iterations
print(f"PyTorch Bias+GELU 平均执行时间: {torch_time:.6f}")
print(f"自定义 CUDA 融合内核 平均执行时间: {cuda_time:.6f}")
speedup = 0
if cuda_time > 0:
speedup = torch_time / cuda_time
print(f"加速比 (Speedup): {speedup:.2f}x")
else:
print("CUDA 内核执行时间为0无法计算加速比。")
return precision_flag,speedup
if __name__ == "__main__":
precision_flag,speedup = run_benchmark()

View File

@ -1,22 +0,0 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
class Model(nn.Module):
def __init__(self, bias: torch.Tensor):
super(Model, self).__init__()
self.register_buffer("bias", bias)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return F.gelu(x + self.bias, approximate='tanh')
batch_size = 16
dim = 16384
def get_inputs():
x = torch.randn(batch_size, dim)
return [x]
def get_init_inputs():
bias = torch.randn(dim)
return [bias]

View File

@ -1,120 +0,0 @@
import torch
from torch.utils.cpp_extension import load_inline
source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <vector_types.h>
__global__ __launch_bounds__(256) void gaussian_affine_gate_kernel(const float* __restrict__ x, const float* __restrict__ scale, const float* __restrict__ bias, float* __restrict__ y, int B, int D, float alpha, float beta){
int row = blockIdx.x;
int tid = threadIdx.x;
int col_block = blockIdx.y;
int stride4 = blockDim.x * 4;
int col_idx = (col_block * stride4) + tid * 4;
const float* xr = x + row * D;
float* yr = y + row * D;
for(int base = col_idx; base < D; base += stride4 * 2){
if (base < D){
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
float4 yv;
float z0 = fmaf(xv.x, sv.x, bv.x);
float z1 = fmaf(xv.y, sv.y, bv.y);
float z2 = fmaf(xv.z, sv.z, bv.z);
float z3 = fmaf(xv.w, sv.w, bv.w);
float m0 = __expf(-(z0*z0));
float m1 = __expf(-(z1*z1));
float m2 = __expf(-(z2*z2));
float m3 = __expf(-(z3*z3));
float t0 = fmaf(alpha, m0, beta);
float t1 = fmaf(alpha, m1, beta);
float t2 = fmaf(alpha, m2, beta);
float t3 = fmaf(alpha, m3, beta);
float g0 = __fdividef(1.0f, 1.0f + __expf(-t0));
float g1 = __fdividef(1.0f, 1.0f + __expf(-t1));
float g2 = __fdividef(1.0f, 1.0f + __expf(-t2));
float g3 = __fdividef(1.0f, 1.0f + __expf(-t3));
yv.x = xv.x * g0;
yv.y = xv.y * g1;
yv.z = xv.z * g2;
yv.w = xv.w * g3;
reinterpret_cast<float4*>(yr + base)[0] = yv;
}
int base2 = base + stride4;
if (base2 < D){
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
float4 yv2;
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
float mb0 = __expf(-(z0b*z0b));
float mb1 = __expf(-(z1b*z1b));
float mb2 = __expf(-(z2b*z2b));
float mb3 = __expf(-(z3b*z3b));
float tb0 = fmaf(alpha, mb0, beta);
float tb1 = fmaf(alpha, mb1, beta);
float tb2 = fmaf(alpha, mb2, beta);
float tb3 = fmaf(alpha, mb3, beta);
float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
yv2.x = xv2.x * gb0;
yv2.y = xv2.y * gb1;
yv2.z = xv2.z * gb2;
yv2.w = xv2.w * gb3;
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
}
}
}
torch::Tensor gaussian_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta){
auto xc = x.contiguous();
auto sc = scale.contiguous();
auto bc = bias.contiguous();
auto y = torch::empty_like(xc);
int B = (int)xc.size(0);
int D = (int)xc.size(1);
float a = (float)alpha;
float be = (float)beta;
int block = 256;
int elements_per_thread = 4;
int elements_per_block = block * elements_per_thread;
int gy = (D + elements_per_block - 1) / elements_per_block;
dim3 grid(B, gy);
gaussian_affine_gate_kernel<<<grid, block>>>(xc.data_ptr<float>(), sc.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), B, D, a, be);
return y;
}
"""
cpp_source = """
#include <torch/extension.h>
torch::Tensor gaussian_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
"""
ops = load_inline(
name="gaussian_affine_gate",
cpp_sources=cpp_source,
cuda_sources=source,
functions=["gaussian_affine_gate_cuda"],
extra_cflags=["-O3","-std=c++17"],
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-Xptxas","-O3,-dlcm=ca","-maxrregcount=64"],
verbose=True
)
class ModelNew(torch.nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(ModelNew, self).__init__()
self.ops = ops
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.alpha = float(alpha)
self.beta = float(beta)
def forward(self, x):
return self.ops.gaussian_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)

View File

@ -1,10 +0,0 @@
Operator: Gaussian-Affine-Gate (Fused CUDA Kernel)
Definition
- z = x * scale + bias
- m = exp(-z^2)
- g = sigmoid(alpha * m + beta)
- y = x * g
Goal
- Elementwise fusion with fast exp; target ≥1.30x speedup.

View File

@ -1,50 +0,0 @@
import torch
import time
from torchcode import Model, get_inputs, get_init_inputs
from cudacode import ModelNew
def run_benchmark():
if not torch.cuda.is_available():
print("CUDA 不可用,请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
return
device = torch.device("cuda")
init_inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_init_inputs()]
inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_inputs()]
torch_model = Model(*init_inputs).cuda().eval()
cuda_model = ModelNew(*init_inputs).cuda().eval()
print("-------------------- 精度对齐验证 --------------------")
with torch.no_grad():
output_torch = torch_model(*inputs)
output_cuda = cuda_model(*inputs)
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-03)
if precision_flag:
print("✅ 精度对齐:两个模型的输出结果非常接近。")
else:
print("❌ 精度不一致!")
print("\n-------------------- 性能加速比测试 --------------------")
num_iterations = 100
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = torch_model(*inputs)
torch.cuda.synchronize(); torch_time = (time.time() - start_time) / num_iterations
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = cuda_model(*inputs)
torch.cuda.synchronize(); cuda_time = (time.time() - start_time) / num_iterations
print(f"PyTorch Gaussian-Affine-Gate 平均执行时间: {torch_time:.6f}")
print(f"自定义 CUDA 融合内核 平均执行时间: {cuda_time:.6f}")
speedup = torch_time / cuda_time if cuda_time > 0 else 0
if cuda_time > 0:
print(f"加速比 (Speedup): {speedup:.2f}x")
else:
print("CUDA 内核执行时间为0无法计算加速比。")
return precision_flag, speedup
if __name__ == "__main__":
run_benchmark()

View File

@ -1,28 +0,0 @@
import torch
import torch.nn as nn
class Model(nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(Model, self).__init__()
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.register_buffer("alpha", torch.tensor(float(alpha), dtype=torch.float32))
self.register_buffer("beta", torch.tensor(float(beta), dtype=torch.float32))
def forward(self, x: torch.Tensor) -> torch.Tensor:
z = x * self.scale + self.bias
m = torch.exp(-z*z)
g = torch.sigmoid(self.alpha * m + self.beta)
return x * g
batch_size = 16
dim = 16384
def get_inputs():
x = torch.randn(batch_size, dim)
return [x]
def get_init_inputs():
scale = torch.randn(dim)
bias = torch.randn(dim)
return [scale, bias, 1.0, 0.0]

View File

@ -1,124 +0,0 @@
import torch
from torch.utils.cpp_extension import load_inline
source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <vector_types.h>
__device__ __forceinline__ float celuf(float x){
return x >= 0.0f ? x : (__expf(x) - 1.0f);
}
__global__ __launch_bounds__(256) void celu_affine_gate_kernel(const float* __restrict__ x, const float* __restrict__ scale, const float* __restrict__ bias, float* __restrict__ y, int B, int D, float alpha, float beta){
int row = blockIdx.x;
int tid = threadIdx.x;
int col_block = blockIdx.y;
int stride4 = blockDim.x * 4;
int col_idx = (col_block * stride4) + tid * 4;
const float* xr = x + row * D;
float* yr = y + row * D;
for(int base = col_idx; base < D; base += stride4 * 2){
if (base < D){
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
float4 yv;
float z0 = fmaf(xv.x, sv.x, bv.x);
float z1 = fmaf(xv.y, sv.y, bv.y);
float z2 = fmaf(xv.z, sv.z, bv.z);
float z3 = fmaf(xv.w, sv.w, bv.w);
float m0 = celuf(z0);
float m1 = celuf(z1);
float m2 = celuf(z2);
float m3 = celuf(z3);
float t0 = fmaf(alpha, m0, beta);
float t1 = fmaf(alpha, m1, beta);
float t2 = fmaf(alpha, m2, beta);
float t3 = fmaf(alpha, m3, beta);
float g0 = __fdividef(1.0f, 1.0f + __expf(-t0));
float g1 = __fdividef(1.0f, 1.0f + __expf(-t1));
float g2 = __fdividef(1.0f, 1.0f + __expf(-t2));
float g3 = __fdividef(1.0f, 1.0f + __expf(-t3));
yv.x = xv.x * g0;
yv.y = xv.y * g1;
yv.z = xv.z * g2;
yv.w = xv.w * g3;
reinterpret_cast<float4*>(yr + base)[0] = yv;
}
int base2 = base + stride4;
if (base2 < D){
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
float4 yv2;
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
float mb0 = celuf(z0b);
float mb1 = celuf(z1b);
float mb2 = celuf(z2b);
float mb3 = celuf(z3b);
float tb0 = fmaf(alpha, mb0, beta);
float tb1 = fmaf(alpha, mb1, beta);
float tb2 = fmaf(alpha, mb2, beta);
float tb3 = fmaf(alpha, mb3, beta);
float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
yv2.x = xv2.x * gb0;
yv2.y = xv2.y * gb1;
yv2.z = xv2.z * gb2;
yv2.w = xv2.w * gb3;
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
}
}
}
torch::Tensor celu_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta){
auto xc = x.contiguous();
auto sc = scale.contiguous();
auto bc = bias.contiguous();
auto y = torch::empty_like(xc);
int B = (int)xc.size(0);
int D = (int)xc.size(1);
float a = (float)alpha;
float be = (float)beta;
int block = 256;
int elements_per_thread = 4;
int elements_per_block = block * elements_per_thread;
int gy = (D + elements_per_block - 1) / elements_per_block;
dim3 grid(B, gy);
celu_affine_gate_kernel<<<grid, block>>>(xc.data_ptr<float>(), sc.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), B, D, a, be);
return y;
}
"""
cpp_source = """
#include <torch/extension.h>
torch::Tensor celu_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
"""
ops = load_inline(
name="celu_affine_gate",
cpp_sources=cpp_source,
cuda_sources=source,
functions=["celu_affine_gate_cuda"],
extra_cflags=["-O3","-std=c++17"],
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-Xptxas","-O3,-dlcm=ca","-maxrregcount=64"],
verbose=True
)
class ModelNew(torch.nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(ModelNew, self).__init__()
self.ops = ops
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.alpha = float(alpha)
self.beta = float(beta)
def forward(self, x):
return self.ops.celu_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)

View File

@ -1,10 +0,0 @@
Operator: CELU-Affine-Gate (Fused CUDA Kernel)
Definition
- z = x * scale + bias
- m = CELU(z, alpha=1)
- g = sigmoid(alpha_g * m + beta)
- y = x * g
Goal
- Elementwise fusion with exp; target ≥1.30x speedup.

View File

@ -1,50 +0,0 @@
import torch
import time
from torchcode import Model, get_inputs, get_init_inputs
from cudacode import ModelNew
def run_benchmark():
if not torch.cuda.is_available():
print("CUDA 不可用,请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
return
device = torch.device("cuda")
init_inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_init_inputs()]
inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_inputs()]
torch_model = Model(*init_inputs).cuda().eval()
cuda_model = ModelNew(*init_inputs).cuda().eval()
print("-------------------- 精度对齐验证 --------------------")
with torch.no_grad():
output_torch = torch_model(*inputs)
output_cuda = cuda_model(*inputs)
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-03)
if precision_flag:
print("✅ 精度对齐:两个模型的输出结果非常接近。")
else:
print("❌ 精度不一致!")
print("\n-------------------- 性能加速比测试 --------------------")
num_iterations = 100
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = torch_model(*inputs)
torch.cuda.synchronize(); torch_time = (time.time() - start_time) / num_iterations
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = cuda_model(*inputs)
torch.cuda.synchronize(); cuda_time = (time.time() - start_time) / num_iterations
print(f"PyTorch CELU-Affine-Gate 平均执行时间: {torch_time:.6f}")
print(f"自定义 CUDA 融合内核 平均执行时间: {cuda_time:.6f}")
speedup = torch_time / cuda_time if cuda_time > 0 else 0
if cuda_time > 0:
print(f"加速比 (Speedup): {speedup:.2f}x")
else:
print("CUDA 内核执行时间为0无法计算加速比。")
return precision_flag, speedup
if __name__ == "__main__":
run_benchmark()

View File

@ -1,29 +0,0 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
class Model(nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(Model, self).__init__()
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.register_buffer("alpha", torch.tensor(float(alpha), dtype=torch.float32))
self.register_buffer("beta", torch.tensor(float(beta), dtype=torch.float32))
def forward(self, x: torch.Tensor) -> torch.Tensor:
z = x * self.scale + self.bias
m = F.celu(z, alpha=1.0)
g = torch.sigmoid(self.alpha * m + self.beta)
return x * g
batch_size = 16
dim = 16384
def get_inputs():
x = torch.randn(batch_size, dim)
return [x]
def get_init_inputs():
scale = torch.randn(dim)
bias = torch.randn(dim)
return [scale, bias, 1.0, 0.0]

View File

@ -1,133 +0,0 @@
import torch
from torch.utils.cpp_extension import load_inline
source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <vector_types.h>
__device__ __forceinline__ float softplusf(float x){
float ax = fabsf(x);
return fmaxf(x, 0.0f) + __logf(1.0f + __expf(-ax));
}
__global__ __launch_bounds__(256) void softplus3_affine_gate_kernel(const float* __restrict__ x, const float* __restrict__ scale, const float* __restrict__ bias, float* __restrict__ y, int B, int D, float alpha, float beta){
int row = blockIdx.x;
int tid = threadIdx.x;
int col_block = blockIdx.y;
int stride4 = blockDim.x * 4;
int col_idx = (col_block * stride4) + tid * 4;
const float* xr = x + row * D;
float* yr = y + row * D;
for(int base = col_idx; base < D; base += stride4 * 2){
if (base < D){
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
float4 yv;
float z0 = fmaf(xv.x, sv.x, bv.x);
float z1 = fmaf(xv.y, sv.y, bv.y);
float z2 = fmaf(xv.z, sv.z, bv.z);
float z3 = fmaf(xv.w, sv.w, bv.w);
float s0 = softplusf(z0);
float s1 = softplusf(z1);
float s2 = softplusf(z2);
float s3 = softplusf(z3);
float m0 = s0 * s0 * s0;
float m1 = s1 * s1 * s1;
float m2 = s2 * s2 * s2;
float m3 = s3 * s3 * s3;
float t0 = fmaf(alpha, m0, beta);
float t1 = fmaf(alpha, m1, beta);
float t2 = fmaf(alpha, m2, beta);
float t3 = fmaf(alpha, m3, beta);
float g0 = __fdividef(1.0f, 1.0f + __expf(-t0));
float g1 = __fdividef(1.0f, 1.0f + __expf(-t1));
float g2 = __fdividef(1.0f, 1.0f + __expf(-t2));
float g3 = __fdividef(1.0f, 1.0f + __expf(-t3));
yv.x = xv.x * g0;
yv.y = xv.y * g1;
yv.z = xv.z * g2;
yv.w = xv.w * g3;
reinterpret_cast<float4*>(yr + base)[0] = yv;
}
int base2 = base + stride4;
if (base2 < D){
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
float4 yv2;
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
float sb0 = softplusf(z0b);
float sb1 = softplusf(z1b);
float sb2 = softplusf(z2b);
float sb3 = softplusf(z3b);
float mb0 = sb0 * sb0 * sb0;
float mb1 = sb1 * sb1 * sb1;
float mb2 = sb2 * sb2 * sb2;
float mb3 = sb3 * sb3 * sb3;
float tb0 = fmaf(alpha, mb0, beta);
float tb1 = fmaf(alpha, mb1, beta);
float tb2 = fmaf(alpha, mb2, beta);
float tb3 = fmaf(alpha, mb3, beta);
float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
yv2.x = xv2.x * gb0;
yv2.y = xv2.y * gb1;
yv2.z = xv2.z * gb2;
yv2.w = xv2.w * gb3;
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
}
}
}
torch::Tensor softplus3_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta){
auto xc = x.contiguous();
auto sc = scale.contiguous();
auto bc = bias.contiguous();
auto y = torch::empty_like(xc);
int B = (int)xc.size(0);
int D = (int)xc.size(1);
float a = (float)alpha;
float be = (float)beta;
int block = 256;
int elements_per_thread = 4;
int elements_per_block = block * elements_per_thread;
int gy = (D + elements_per_block - 1) / elements_per_block;
dim3 grid(B, gy);
softplus3_affine_gate_kernel<<<grid, block>>>(xc.data_ptr<float>(), sc.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), B, D, a, be);
return y;
}
"""
cpp_source = """
#include <torch/extension.h>
torch::Tensor softplus3_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
"""
ops = load_inline(
name="softplus3_affine_gate",
cpp_sources=cpp_source,
cuda_sources=source,
functions=["softplus3_affine_gate_cuda"],
extra_cflags=["-O3","-std=c++17"],
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-Xptxas","-O3,-dlcm=ca","-maxrregcount=64"],
verbose=True
)
class ModelNew(torch.nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(ModelNew, self).__init__()
self.ops = ops
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.alpha = float(alpha)
self.beta = float(beta)
def forward(self, x):
return self.ops.softplus3_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)

View File

@ -1,10 +0,0 @@
Operator: Softplus^3-Affine-Gate (Fused CUDA Kernel)
Definition
- z = x * scale + bias
- m = softplus(z)^3
- g = sigmoid(alpha * m + beta)
- y = x * g
Goal
- Elementwise fusion with fast softplus; target ≥1.30x speedup.

View File

@ -1,50 +0,0 @@
import torch
import time
from torchcode import Model, get_inputs, get_init_inputs
from cudacode import ModelNew
def run_benchmark():
if not torch.cuda.is_available():
print("CUDA 不可用,请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
return
device = torch.device("cuda")
init_inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_init_inputs()]
inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_inputs()]
torch_model = Model(*init_inputs).cuda().eval()
cuda_model = ModelNew(*init_inputs).cuda().eval()
print("-------------------- 精度对齐验证 --------------------")
with torch.no_grad():
output_torch = torch_model(*inputs)
output_cuda = cuda_model(*inputs)
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-03)
if precision_flag:
print("✅ 精度对齐:两个模型的输出结果非常接近。")
else:
print("❌ 精度不一致!")
print("\n-------------------- 性能加速比测试 --------------------")
num_iterations = 100
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = torch_model(*inputs)
torch.cuda.synchronize(); torch_time = (time.time() - start_time) / num_iterations
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = cuda_model(*inputs)
torch.cuda.synchronize(); cuda_time = (time.time() - start_time) / num_iterations
print(f"PyTorch Softplus^3-Affine-Gate 平均执行时间: {torch_time:.6f}")
print(f"自定义 CUDA 融合内核 平均执行时间: {cuda_time:.6f}")
speedup = torch_time / cuda_time if cuda_time > 0 else 0
if cuda_time > 0:
print(f"加速比 (Speedup): {speedup:.2f}x")
else:
print("CUDA 内核执行时间为0无法计算加速比。")
return precision_flag, speedup
if __name__ == "__main__":
run_benchmark()

View File

@ -1,30 +0,0 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
class Model(nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(Model, self).__init__()
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.register_buffer("alpha", torch.tensor(float(alpha), dtype=torch.float32))
self.register_buffer("beta", torch.tensor(float(beta), dtype=torch.float32))
def forward(self, x: torch.Tensor) -> torch.Tensor:
z = x * self.scale + self.bias
m = F.softplus(z)
m = m * m * m
g = torch.sigmoid(self.alpha * m + self.beta)
return x * g
batch_size = 16
dim = 16384
def get_inputs():
x = torch.randn(batch_size, dim)
return [x]
def get_init_inputs():
scale = torch.randn(dim)
bias = torch.randn(dim)
return [scale, bias, 1.0, 0.0]

View File

@ -1,127 +0,0 @@
import torch
from torch.utils.cpp_extension import load_inline
source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <vector_types.h>
__device__ __forceinline__ float softshrinkf(float x, float lambda_){
if (x > lambda_) return x - lambda_;
if (x < -lambda_) return x + lambda_;
return 0.0f;
}
__global__ __launch_bounds__(256) void softshrink_affine_gate_kernel(const float* __restrict__ x, const float* __restrict__ scale, const float* __restrict__ bias, float* __restrict__ y, int B, int D, float alpha, float beta){
int row = blockIdx.x;
int tid = threadIdx.x;
int col_block = blockIdx.y;
int stride4 = blockDim.x * 4;
int col_idx = (col_block * stride4) + tid * 4;
const float* xr = x + row * D;
float* yr = y + row * D;
const float lambda_ = 0.5f;
for(int base = col_idx; base < D; base += stride4 * 2){
if (base < D){
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
float4 yv;
float z0 = fmaf(xv.x, sv.x, bv.x);
float z1 = fmaf(xv.y, sv.y, bv.y);
float z2 = fmaf(xv.z, sv.z, bv.z);
float z3 = fmaf(xv.w, sv.w, bv.w);
float m0 = softshrinkf(z0, lambda_);
float m1 = softshrinkf(z1, lambda_);
float m2 = softshrinkf(z2, lambda_);
float m3 = softshrinkf(z3, lambda_);
float t0 = fmaf(alpha, m0, beta);
float t1 = fmaf(alpha, m1, beta);
float t2 = fmaf(alpha, m2, beta);
float t3 = fmaf(alpha, m3, beta);
float g0 = __fdividef(1.0f, 1.0f + __expf(-t0));
float g1 = __fdividef(1.0f, 1.0f + __expf(-t1));
float g2 = __fdividef(1.0f, 1.0f + __expf(-t2));
float g3 = __fdividef(1.0f, 1.0f + __expf(-t3));
yv.x = xv.x * g0;
yv.y = xv.y * g1;
yv.z = xv.z * g2;
yv.w = xv.w * g3;
reinterpret_cast<float4*>(yr + base)[0] = yv;
}
int base2 = base + stride4;
if (base2 < D){
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
float4 yv2;
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
float mb0 = softshrinkf(z0b, lambda_);
float mb1 = softshrinkf(z1b, lambda_);
float mb2 = softshrinkf(z2b, lambda_);
float mb3 = softshrinkf(z3b, lambda_);
float tb0 = fmaf(alpha, mb0, beta);
float tb1 = fmaf(alpha, mb1, beta);
float tb2 = fmaf(alpha, mb2, beta);
float tb3 = fmaf(alpha, mb3, beta);
float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
yv2.x = xv2.x * gb0;
yv2.y = xv2.y * gb1;
yv2.z = xv2.z * gb2;
yv2.w = xv2.w * gb3;
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
}
}
}
torch::Tensor softshrink_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta){
auto xc = x.contiguous();
auto sc = scale.contiguous();
auto bc = bias.contiguous();
auto y = torch::empty_like(xc);
int B = (int)xc.size(0);
int D = (int)xc.size(1);
float a = (float)alpha;
float be = (float)beta;
int block = 256;
int elements_per_thread = 4;
int elements_per_block = block * elements_per_thread;
int gy = (D + elements_per_block - 1) / elements_per_block;
dim3 grid(B, gy);
softshrink_affine_gate_kernel<<<grid, block>>>(xc.data_ptr<float>(), sc.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), B, D, a, be);
return y;
}
"""
cpp_source = """
#include <torch/extension.h>
torch::Tensor softshrink_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
"""
ops = load_inline(
name="softshrink_affine_gate",
cpp_sources=cpp_source,
cuda_sources=source,
functions=["softshrink_affine_gate_cuda"],
extra_cflags=["-O3","-std=c++17"],
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-Xptxas","-O3,-dlcm=ca","-maxrregcount=64"],
verbose=True
)
class ModelNew(torch.nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(ModelNew, self).__init__()
self.ops = ops
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.alpha = float(alpha)
self.beta = float(beta)
def forward(self, x):
return self.ops.softshrink_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)

View File

@ -1,10 +0,0 @@
Operator: SoftShrink-Affine-Gate (Fused CUDA Kernel)
Definition
- z = x * scale + bias
- m = SoftShrink(z, lambda=0.5)
- g = sigmoid(alpha * m + beta)
- y = x * g
Goal
- Fuse ops to reduce bandwidth and kernel launches; target ≥1.30x speedup.

View File

@ -1,50 +0,0 @@
import torch
import time
from torchcode import Model, get_inputs, get_init_inputs
from cudacode import ModelNew
def run_benchmark():
if not torch.cuda.is_available():
print("CUDA 不可用,请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
return
device = torch.device("cuda")
init_inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_init_inputs()]
inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_inputs()]
torch_model = Model(*init_inputs).cuda().eval()
cuda_model = ModelNew(*init_inputs).cuda().eval()
print("-------------------- 精度对齐验证 --------------------")
with torch.no_grad():
output_torch = torch_model(*inputs)
output_cuda = cuda_model(*inputs)
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-03)
if precision_flag:
print("✅ 精度对齐:两个模型的输出结果非常接近。")
else:
print("❌ 精度不一致!")
print("\n-------------------- 性能加速比测试 --------------------")
num_iterations = 100
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = torch_model(*inputs)
torch.cuda.synchronize(); torch_time = (time.time() - start_time) / num_iterations
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = cuda_model(*inputs)
torch.cuda.synchronize(); cuda_time = (time.time() - start_time) / num_iterations
print(f"PyTorch SoftShrink-Affine-Gate 平均执行时间: {torch_time:.6f}")
print(f"自定义 CUDA 融合内核 平均执行时间: {cuda_time:.6f}")
speedup = torch_time / cuda_time if cuda_time > 0 else 0
if cuda_time > 0:
print(f"加速比 (Speedup): {speedup:.2f}x")
else:
print("CUDA 内核执行时间为0无法计算加速比。")
return precision_flag, speedup
if __name__ == "__main__":
run_benchmark()

View File

@ -1,29 +0,0 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
class Model(nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(Model, self).__init__()
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.register_buffer("alpha", torch.tensor(float(alpha), dtype=torch.float32))
self.register_buffer("beta", torch.tensor(float(beta), dtype=torch.float32))
def forward(self, x: torch.Tensor) -> torch.Tensor:
z = x * self.scale + self.bias
m = F.softshrink(z, lambd=0.5)
g = torch.sigmoid(self.alpha * m + self.beta)
return x * g
batch_size = 16
dim = 16384
def get_inputs():
x = torch.randn(batch_size, dim)
return [x]
def get_init_inputs():
scale = torch.randn(dim)
bias = torch.randn(dim)
return [scale, bias, 1.0, 0.0]

View File

@ -1,125 +0,0 @@
import torch
from torch.utils.cpp_extension import load_inline
source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <vector_types.h>
__device__ __forceinline__ float hardshrinkf(float x, float lambda_){
return (x > lambda_ || x < -lambda_) ? x : 0.0f;
}
__global__ __launch_bounds__(256) void hardshrink_affine_gate_kernel(const float* __restrict__ x, const float* __restrict__ scale, const float* __restrict__ bias, float* __restrict__ y, int B, int D, float alpha, float beta){
int row = blockIdx.x;
int tid = threadIdx.x;
int col_block = blockIdx.y;
int stride4 = blockDim.x * 4;
int col_idx = (col_block * stride4) + tid * 4;
const float* xr = x + row * D;
float* yr = y + row * D;
const float lambda_ = 0.5f;
for(int base = col_idx; base < D; base += stride4 * 2){
if (base < D){
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
float4 yv;
float z0 = fmaf(xv.x, sv.x, bv.x);
float z1 = fmaf(xv.y, sv.y, bv.y);
float z2 = fmaf(xv.z, sv.z, bv.z);
float z3 = fmaf(xv.w, sv.w, bv.w);
float m0 = hardshrinkf(z0, lambda_);
float m1 = hardshrinkf(z1, lambda_);
float m2 = hardshrinkf(z2, lambda_);
float m3 = hardshrinkf(z3, lambda_);
float t0 = fmaf(alpha, m0, beta);
float t1 = fmaf(alpha, m1, beta);
float t2 = fmaf(alpha, m2, beta);
float t3 = fmaf(alpha, m3, beta);
float g0 = __fdividef(1.0f, 1.0f + __expf(-t0));
float g1 = __fdividef(1.0f, 1.0f + __expf(-t1));
float g2 = __fdividef(1.0f, 1.0f + __expf(-t2));
float g3 = __fdividef(1.0f, 1.0f + __expf(-t3));
yv.x = xv.x * g0;
yv.y = xv.y * g1;
yv.z = xv.z * g2;
yv.w = xv.w * g3;
reinterpret_cast<float4*>(yr + base)[0] = yv;
}
int base2 = base + stride4;
if (base2 < D){
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
float4 yv2;
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
float mb0 = hardshrinkf(z0b, lambda_);
float mb1 = hardshrinkf(z1b, lambda_);
float mb2 = hardshrinkf(z2b, lambda_);
float mb3 = hardshrinkf(z3b, lambda_);
float tb0 = fmaf(alpha, mb0, beta);
float tb1 = fmaf(alpha, mb1, beta);
float tb2 = fmaf(alpha, mb2, beta);
float tb3 = fmaf(alpha, mb3, beta);
float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
yv2.x = xv2.x * gb0;
yv2.y = xv2.y * gb1;
yv2.z = xv2.z * gb2;
yv2.w = xv2.w * gb3;
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
}
}
}
torch::Tensor hardshrink_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta){
auto xc = x.contiguous();
auto sc = scale.contiguous();
auto bc = bias.contiguous();
auto y = torch::empty_like(xc);
int B = (int)xc.size(0);
int D = (int)xc.size(1);
float a = (float)alpha;
float be = (float)beta;
int block = 256;
int elements_per_thread = 4;
int elements_per_block = block * elements_per_thread;
int gy = (D + elements_per_block - 1) / elements_per_block;
dim3 grid(B, gy);
hardshrink_affine_gate_kernel<<<grid, block>>>(xc.data_ptr<float>(), sc.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), B, D, a, be);
return y;
}
"""
cpp_source = """
#include <torch/extension.h>
torch::Tensor hardshrink_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
"""
ops = load_inline(
name="hardshrink_affine_gate",
cpp_sources=cpp_source,
cuda_sources=source,
functions=["hardshrink_affine_gate_cuda"],
extra_cflags=["-O3","-std=c++17"],
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-Xptxas","-O3,-dlcm=ca","-maxrregcount=64"],
verbose=True
)
class ModelNew(torch.nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(ModelNew, self).__init__()
self.ops = ops
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.alpha = float(alpha)
self.beta = float(beta)
def forward(self, x):
return self.ops.hardshrink_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)

View File

@ -1,10 +0,0 @@
Operator: HardShrink-Affine-Gate (Fused CUDA Kernel)
Definition
- z = x * scale + bias
- m = HardShrink(z, lambda=0.5)
- g = sigmoid(alpha * m + beta)
- y = x * g
Goal
- Fuse ops to reduce memory passes and launches; target ≥1.30x speedup.

View File

@ -1,50 +0,0 @@
import torch
import time
from torchcode import Model, get_inputs, get_init_inputs
from cudacode import ModelNew
def run_benchmark():
if not torch.cuda.is_available():
print("CUDA 不可用,请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
return
device = torch.device("cuda")
init_inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_init_inputs()]
inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_inputs()]
torch_model = Model(*init_inputs).cuda().eval()
cuda_model = ModelNew(*init_inputs).cuda().eval()
print("-------------------- 精度对齐验证 --------------------")
with torch.no_grad():
output_torch = torch_model(*inputs)
output_cuda = cuda_model(*inputs)
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-03)
if precision_flag:
print("✅ 精度对齐:两个模型的输出结果非常接近。")
else:
print("❌ 精度不一致!")
print("\n-------------------- 性能加速比测试 --------------------")
num_iterations = 100
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = torch_model(*inputs)
torch.cuda.synchronize(); torch_time = (time.time() - start_time) / num_iterations
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = cuda_model(*inputs)
torch.cuda.synchronize(); cuda_time = (time.time() - start_time) / num_iterations
print(f"PyTorch HardShrink-Affine-Gate 平均执行时间: {torch_time:.6f}")
print(f"自定义 CUDA 融合内核 平均执行时间: {cuda_time:.6f}")
speedup = torch_time / cuda_time if cuda_time > 0 else 0
if cuda_time > 0:
print(f"加速比 (Speedup): {speedup:.2f}x")
else:
print("CUDA 内核执行时间为0无法计算加速比。")
return precision_flag, speedup
if __name__ == "__main__":
run_benchmark()

View File

@ -1,29 +0,0 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
class Model(nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(Model, self).__init__()
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.register_buffer("alpha", torch.tensor(float(alpha), dtype=torch.float32))
self.register_buffer("beta", torch.tensor(float(beta), dtype=torch.float32))
def forward(self, x: torch.Tensor) -> torch.Tensor:
z = x * self.scale + self.bias
m = F.hardshrink(z, lambd=0.5)
g = torch.sigmoid(self.alpha * m + self.beta)
return x * g
batch_size = 16
dim = 16384
def get_inputs():
x = torch.randn(batch_size, dim)
return [x]
def get_init_inputs():
scale = torch.randn(dim)
bias = torch.randn(dim)
return [scale, bias, 1.0, 0.0]

View File

@ -1,120 +0,0 @@
import torch
from torch.utils.cpp_extension import load_inline
source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <vector_types.h>
__global__ __launch_bounds__(256) void laplacian_affine_gate_kernel(const float* __restrict__ x, const float* __restrict__ scale, const float* __restrict__ bias, float* __restrict__ y, int B, int D, float alpha, float beta){
int row = blockIdx.x;
int tid = threadIdx.x;
int col_block = blockIdx.y;
int stride4 = blockDim.x * 4;
int col_idx = (col_block * stride4) + tid * 4;
const float* xr = x + row * D;
float* yr = y + row * D;
for(int base = col_idx; base < D; base += stride4 * 2){
if (base < D){
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
float4 yv;
float z0 = fmaf(xv.x, sv.x, bv.x);
float z1 = fmaf(xv.y, sv.y, bv.y);
float z2 = fmaf(xv.z, sv.z, bv.z);
float z3 = fmaf(xv.w, sv.w, bv.w);
float m0 = __expf(-fabsf(z0));
float m1 = __expf(-fabsf(z1));
float m2 = __expf(-fabsf(z2));
float m3 = __expf(-fabsf(z3));
float t0 = fmaf(alpha, m0, beta);
float t1 = fmaf(alpha, m1, beta);
float t2 = fmaf(alpha, m2, beta);
float t3 = fmaf(alpha, m3, beta);
float g0 = __fdividef(1.0f, 1.0f + __expf(-t0));
float g1 = __fdividef(1.0f, 1.0f + __expf(-t1));
float g2 = __fdividef(1.0f, 1.0f + __expf(-t2));
float g3 = __fdividef(1.0f, 1.0f + __expf(-t3));
yv.x = xv.x * g0;
yv.y = xv.y * g1;
yv.z = xv.z * g2;
yv.w = xv.w * g3;
reinterpret_cast<float4*>(yr + base)[0] = yv;
}
int base2 = base + stride4;
if (base2 < D){
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
float4 yv2;
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
float mb0 = __expf(-fabsf(z0b));
float mb1 = __expf(-fabsf(z1b));
float mb2 = __expf(-fabsf(z2b));
float mb3 = __expf(-fabsf(z3b));
float tb0 = fmaf(alpha, mb0, beta);
float tb1 = fmaf(alpha, mb1, beta);
float tb2 = fmaf(alpha, mb2, beta);
float tb3 = fmaf(alpha, mb3, beta);
float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
yv2.x = xv2.x * gb0;
yv2.y = xv2.y * gb1;
yv2.z = xv2.z * gb2;
yv2.w = xv2.w * gb3;
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
}
}
}
torch::Tensor laplacian_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta){
auto xc = x.contiguous();
auto sc = scale.contiguous();
auto bc = bias.contiguous();
auto y = torch::empty_like(xc);
int B = (int)xc.size(0);
int D = (int)xc.size(1);
float a = (float)alpha;
float be = (float)beta;
int block = 256;
int elements_per_thread = 4;
int elements_per_block = block * elements_per_thread;
int gy = (D + elements_per_block - 1) / elements_per_block;
dim3 grid(B, gy);
laplacian_affine_gate_kernel<<<grid, block>>>(xc.data_ptr<float>(), sc.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), B, D, a, be);
return y;
}
"""
cpp_source = """
#include <torch/extension.h>
torch::Tensor laplacian_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
"""
ops = load_inline(
name="laplacian_affine_gate",
cpp_sources=cpp_source,
cuda_sources=source,
functions=["laplacian_affine_gate_cuda"],
extra_cflags=["-O3","-std=c++17"],
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-Xptxas","-O3,-dlcm=ca","-maxrregcount=64"],
verbose=True
)
class ModelNew(torch.nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(ModelNew, self).__init__()
self.ops = ops
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.alpha = float(alpha)
self.beta = float(beta)
def forward(self, x):
return self.ops.laplacian_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)

View File

@ -1,10 +0,0 @@
Operator: Laplacian-Affine-Gate (Fused CUDA Kernel)
Definition
- z = x * scale + bias
- m = exp(-|z|)
- g = sigmoid(alpha * m + beta)
- y = x * g
Goal
- Elementwise fusion; target ≥1.30x speedup.

View File

@ -1,50 +0,0 @@
import torch
import time
from torchcode import Model, get_inputs, get_init_inputs
from cudacode import ModelNew
def run_benchmark():
if not torch.cuda.is_available():
print("CUDA 不可用,请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
return
device = torch.device("cuda")
init_inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_init_inputs()]
inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_inputs()]
torch_model = Model(*init_inputs).cuda().eval()
cuda_model = ModelNew(*init_inputs).cuda().eval()
print("-------------------- 精度对齐验证 --------------------")
with torch.no_grad():
output_torch = torch_model(*inputs)
output_cuda = cuda_model(*inputs)
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-03)
if precision_flag:
print("✅ 精度对齐:两个模型的输出结果非常接近。")
else:
print("❌ 精度不一致!")
print("\n-------------------- 性能加速比测试 --------------------")
num_iterations = 100
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = torch_model(*inputs)
torch.cuda.synchronize(); torch_time = (time.time() - start_time) / num_iterations
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = cuda_model(*inputs)
torch.cuda.synchronize(); cuda_time = (time.time() - start_time) / num_iterations
print(f"PyTorch Laplacian-Affine-Gate 平均执行时间: {torch_time:.6f}")
print(f"自定义 CUDA 融合内核 平均执行时间: {cuda_time:.6f}")
speedup = torch_time / cuda_time if cuda_time > 0 else 0
if cuda_time > 0:
print(f"加速比 (Speedup): {speedup:.2f}x")
else:
print("CUDA 内核执行时间为0无法计算加速比。")
return precision_flag, speedup
if __name__ == "__main__":
run_benchmark()

View File

@ -1,28 +0,0 @@
import torch
import torch.nn as nn
class Model(nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(Model, self).__init__()
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.register_buffer("alpha", torch.tensor(float(alpha), dtype=torch.float32))
self.register_buffer("beta", torch.tensor(float(beta), dtype=torch.float32))
def forward(self, x: torch.Tensor) -> torch.Tensor:
z = x * self.scale + self.bias
m = torch.exp(-torch.abs(z))
g = torch.sigmoid(self.alpha * m + self.beta)
return x * g
batch_size = 16
dim = 16384
def get_inputs():
x = torch.randn(batch_size, dim)
return [x]
def get_init_inputs():
scale = torch.randn(dim)
bias = torch.randn(dim)
return [scale, bias, 1.0, 0.0]

View File

@ -1,129 +0,0 @@
import torch
from torch.utils.cpp_extension import load_inline
source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <vector_types.h>
__global__ __launch_bounds__(256) void logcosh_affine_gate_kernel(const float* __restrict__ x, const float* __restrict__ scale, const float* __restrict__ bias, float* __restrict__ y, int B, int D, float alpha, float beta){
int row = blockIdx.x;
int tid = threadIdx.x;
int col_block = blockIdx.y;
int stride4 = blockDim.x * 4;
int col_idx = (col_block * stride4) + tid * 4;
const float* xr = x + row * D;
float* yr = y + row * D;
const float LOG2 = 0.6931471805599453094f;
for(int base = col_idx; base < D; base += stride4 * 2){
if (base < D){
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
float4 yv;
float z0 = fmaf(xv.x, sv.x, bv.x);
float z1 = fmaf(xv.y, sv.y, bv.y);
float z2 = fmaf(xv.z, sv.z, bv.z);
float z3 = fmaf(xv.w, sv.w, bv.w);
float a0 = fabsf(z0);
float a1 = fabsf(z1);
float a2 = fabsf(z2);
float a3 = fabsf(z3);
float m0 = a0 + __logf(1.0f + __expf(-2.0f * a0)) - LOG2;
float m1 = a1 + __logf(1.0f + __expf(-2.0f * a1)) - LOG2;
float m2 = a2 + __logf(1.0f + __expf(-2.0f * a2)) - LOG2;
float m3 = a3 + __logf(1.0f + __expf(-2.0f * a3)) - LOG2;
float t0 = fmaf(alpha, m0, beta);
float t1 = fmaf(alpha, m1, beta);
float t2 = fmaf(alpha, m2, beta);
float t3 = fmaf(alpha, m3, beta);
float g0 = __fdividef(1.0f, 1.0f + __expf(-t0));
float g1 = __fdividef(1.0f, 1.0f + __expf(-t1));
float g2 = __fdividef(1.0f, 1.0f + __expf(-t2));
float g3 = __fdividef(1.0f, 1.0f + __expf(-t3));
yv.x = xv.x * g0;
yv.y = xv.y * g1;
yv.z = xv.z * g2;
yv.w = xv.w * g3;
reinterpret_cast<float4*>(yr + base)[0] = yv;
}
int base2 = base + stride4;
if (base2 < D){
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
float4 yv2;
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
float ab0 = fabsf(z0b);
float ab1 = fabsf(z1b);
float ab2 = fabsf(z2b);
float ab3 = fabsf(z3b);
float mb0 = ab0 + __logf(1.0f + __expf(-2.0f * ab0)) - LOG2;
float mb1 = ab1 + __logf(1.0f + __expf(-2.0f * ab1)) - LOG2;
float mb2 = ab2 + __logf(1.0f + __expf(-2.0f * ab2)) - LOG2;
float mb3 = ab3 + __logf(1.0f + __expf(-2.0f * ab3)) - LOG2;
float tb0 = fmaf(alpha, mb0, beta);
float tb1 = fmaf(alpha, mb1, beta);
float tb2 = fmaf(alpha, mb2, beta);
float tb3 = fmaf(alpha, mb3, beta);
float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
yv2.x = xv2.x * gb0;
yv2.y = xv2.y * gb1;
yv2.z = xv2.z * gb2;
yv2.w = xv2.w * gb3;
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
}
}
}
torch::Tensor logcosh_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta){
auto xc = x.contiguous();
auto sc = scale.contiguous();
auto bc = bias.contiguous();
auto y = torch::empty_like(xc);
int B = (int)xc.size(0);
int D = (int)xc.size(1);
float a = (float)alpha;
float be = (float)beta;
int block = 256;
int elements_per_thread = 4;
int elements_per_block = block * elements_per_thread;
int gy = (D + elements_per_block - 1) / elements_per_block;
dim3 grid(B, gy);
logcosh_affine_gate_kernel<<<grid, block>>>(xc.data_ptr<float>(), sc.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), B, D, a, be);
return y;
}
"""
cpp_source = """
#include <torch/extension.h>
torch::Tensor logcosh_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
"""
ops = load_inline(
name="logcosh_affine_gate",
cpp_sources=cpp_source,
cuda_sources=source,
functions=["logcosh_affine_gate_cuda"],
extra_cflags=["-O3","-std=c++17"],
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-Xptxas","-O3,-dlcm=ca","-maxrregcount=64"],
verbose=True
)
class ModelNew(torch.nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(ModelNew, self).__init__()
self.ops = ops
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.alpha = float(alpha)
self.beta = float(beta)
def forward(self, x):
return self.ops.logcosh_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)

View File

@ -1,10 +0,0 @@
Operator: LogCosh-Affine-Gate (Fused CUDA Kernel)
Definition
- z = x * scale + bias
- m = log(cosh(z))
- g = sigmoid(alpha * m + beta)
- y = x * g
Goal
- Numerically stable logcosh fusion; target ≥1.30x speedup.

View File

@ -1,50 +0,0 @@
import torch
import time
from torchcode import Model, get_inputs, get_init_inputs
from cudacode import ModelNew
def run_benchmark():
if not torch.cuda.is_available():
print("CUDA 不可用,请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
return
device = torch.device("cuda")
init_inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_init_inputs()]
inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_inputs()]
torch_model = Model(*init_inputs).cuda().eval()
cuda_model = ModelNew(*init_inputs).cuda().eval()
print("-------------------- 精度对齐验证 --------------------")
with torch.no_grad():
output_torch = torch_model(*inputs)
output_cuda = cuda_model(*inputs)
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-03)
if precision_flag:
print("✅ 精度对齐:两个模型的输出结果非常接近。")
else:
print("❌ 精度不一致!")
print("\n-------------------- 性能加速比测试 --------------------")
num_iterations = 100
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = torch_model(*inputs)
torch.cuda.synchronize(); torch_time = (time.time() - start_time) / num_iterations
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = cuda_model(*inputs)
torch.cuda.synchronize(); cuda_time = (time.time() - start_time) / num_iterations
print(f"PyTorch LogCosh-Affine-Gate 平均执行时间: {torch_time:.6f}")
print(f"自定义 CUDA 融合内核 平均执行时间: {cuda_time:.6f}")
speedup = torch_time / cuda_time if cuda_time > 0 else 0
if cuda_time > 0:
print(f"加速比 (Speedup): {speedup:.2f}x")
else:
print("CUDA 内核执行时间为0无法计算加速比。")
return precision_flag, speedup
if __name__ == "__main__":
run_benchmark()

View File

@ -1,29 +0,0 @@
import torch
import torch.nn as nn
class Model(nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(Model, self).__init__()
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.register_buffer("alpha", torch.tensor(float(alpha), dtype=torch.float32))
self.register_buffer("beta", torch.tensor(float(beta), dtype=torch.float32))
def forward(self, x: torch.Tensor) -> torch.Tensor:
z = x * self.scale + self.bias
az = torch.abs(z)
m = az + torch.log1p(torch.exp(-2.0 * az)) - torch.log(torch.tensor(2.0))
g = torch.sigmoid(self.alpha * m + self.beta)
return x * g
batch_size = 16
dim = 16384
def get_inputs():
x = torch.randn(batch_size, dim)
return [x]
def get_init_inputs():
scale = torch.randn(dim)
bias = torch.randn(dim)
return [scale, bias, 1.0, 0.0]

View File

@ -1,133 +0,0 @@
import torch
from torch.utils.cpp_extension import load_inline
source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <vector_types.h>
__device__ __forceinline__ float softplusf(float x){
float ax = fabsf(x);
return fmaxf(x, 0.0f) + __logf(1.0f + __expf(-ax));
}
__global__ __launch_bounds__(256) void softplus_sqrt_affine_gate_kernel(const float* __restrict__ x, const float* __restrict__ scale, const float* __restrict__ bias, float* __restrict__ y, int B, int D, float alpha, float beta){
int row = blockIdx.x;
int tid = threadIdx.x;
int col_block = blockIdx.y;
int stride4 = blockDim.x * 4;
int col_idx = (col_block * stride4) + tid * 4;
const float* xr = x + row * D;
float* yr = y + row * D;
for(int base = col_idx; base < D; base += stride4 * 2){
if (base < D){
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
float4 yv;
float z0 = fmaf(xv.x, sv.x, bv.x);
float z1 = fmaf(xv.y, sv.y, bv.y);
float z2 = fmaf(xv.z, sv.z, bv.z);
float z3 = fmaf(xv.w, sv.w, bv.w);
float s0 = softplusf(z0);
float s1 = softplusf(z1);
float s2 = softplusf(z2);
float s3 = softplusf(z3);
float m0 = sqrtf(s0);
float m1 = sqrtf(s1);
float m2 = sqrtf(s2);
float m3 = sqrtf(s3);
float t0 = fmaf(alpha, m0, beta);
float t1 = fmaf(alpha, m1, beta);
float t2 = fmaf(alpha, m2, beta);
float t3 = fmaf(alpha, m3, beta);
float g0 = __fdividef(1.0f, 1.0f + __expf(-t0));
float g1 = __fdividef(1.0f, 1.0f + __expf(-t1));
float g2 = __fdividef(1.0f, 1.0f + __expf(-t2));
float g3 = __fdividef(1.0f, 1.0f + __expf(-t3));
yv.x = xv.x * g0;
yv.y = xv.y * g1;
yv.z = xv.z * g2;
yv.w = xv.w * g3;
reinterpret_cast<float4*>(yr + base)[0] = yv;
}
int base2 = base + stride4;
if (base2 < D){
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
float4 yv2;
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
float sb0 = softplusf(z0b);
float sb1 = softplusf(z1b);
float sb2 = softplusf(z2b);
float sb3 = softplusf(z3b);
float mb0 = sqrtf(sb0);
float mb1 = sqrtf(sb1);
float mb2 = sqrtf(sb2);
float mb3 = sqrtf(sb3);
float tb0 = fmaf(alpha, mb0, beta);
float tb1 = fmaf(alpha, mb1, beta);
float tb2 = fmaf(alpha, mb2, beta);
float tb3 = fmaf(alpha, mb3, beta);
float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
yv2.x = xv2.x * gb0;
yv2.y = xv2.y * gb1;
yv2.z = xv2.z * gb2;
yv2.w = xv2.w * gb3;
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
}
}
}
torch::Tensor softplus_sqrt_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta){
auto xc = x.contiguous();
auto sc = scale.contiguous();
auto bc = bias.contiguous();
auto y = torch::empty_like(xc);
int B = (int)xc.size(0);
int D = (int)xc.size(1);
float a = (float)alpha;
float be = (float)beta;
int block = 256;
int elements_per_thread = 4;
int elements_per_block = block * elements_per_thread;
int gy = (D + elements_per_block - 1) / elements_per_block;
dim3 grid(B, gy);
softplus_sqrt_affine_gate_kernel<<<grid, block>>>(xc.data_ptr<float>(), sc.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), B, D, a, be);
return y;
}
"""
cpp_source = """
#include <torch/extension.h>
torch::Tensor softplus_sqrt_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
"""
ops = load_inline(
name="softplus_sqrt_affine_gate",
cpp_sources=cpp_source,
cuda_sources=source,
functions=["softplus_sqrt_affine_gate_cuda"],
extra_cflags=["-O3","-std=c++17"],
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-Xptxas","-O3,-dlcm=ca","-maxrregcount=64"],
verbose=True
)
class ModelNew(torch.nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(ModelNew, self).__init__()
self.ops = ops
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.alpha = float(alpha)
self.beta = float(beta)
def forward(self, x):
return self.ops.softplus_sqrt_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)

View File

@ -1,10 +0,0 @@
Operator: SoftplusSqrt-Affine-Gate (Fused CUDA Kernel)
Definition
- z = x * scale + bias
- m = sqrt(softplus(z))
- g = sigmoid(alpha * m + beta)
- y = x * g
Goal
- Fuse softplus and sqrt; target ≥1.30x speedup.

View File

@ -1,50 +0,0 @@
import torch
import time
from torchcode import Model, get_inputs, get_init_inputs
from cudacode import ModelNew
def run_benchmark():
if not torch.cuda.is_available():
print("CUDA 不可用,请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
return
device = torch.device("cuda")
init_inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_init_inputs()]
inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_inputs()]
torch_model = Model(*init_inputs).cuda().eval()
cuda_model = ModelNew(*init_inputs).cuda().eval()
print("-------------------- 精度对齐验证 --------------------")
with torch.no_grad():
output_torch = torch_model(*inputs)
output_cuda = cuda_model(*inputs)
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-03)
if precision_flag:
print("✅ 精度对齐:两个模型的输出结果非常接近。")
else:
print("❌ 精度不一致!")
print("\n-------------------- 性能加速比测试 --------------------")
num_iterations = 100
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = torch_model(*inputs)
torch.cuda.synchronize(); torch_time = (time.time() - start_time) / num_iterations
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = cuda_model(*inputs)
torch.cuda.synchronize(); cuda_time = (time.time() - start_time) / num_iterations
print(f"PyTorch SoftplusSqrt-Affine-Gate 平均执行时间: {torch_time:.6f}")
print(f"自定义 CUDA 融合内核 平均执行时间: {cuda_time:.6f}")
speedup = torch_time / cuda_time if cuda_time > 0 else 0
if cuda_time > 0:
print(f"加速比 (Speedup): {speedup:.2f}x")
else:
print("CUDA 内核执行时间为0无法计算加速比。")
return precision_flag, speedup
if __name__ == "__main__":
run_benchmark()

View File

@ -1,29 +0,0 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
class Model(nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(Model, self).__init__()
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.register_buffer("alpha", torch.tensor(float(alpha), dtype=torch.float32))
self.register_buffer("beta", torch.tensor(float(beta), dtype=torch.float32))
def forward(self, x: torch.Tensor) -> torch.Tensor:
z = x * self.scale + self.bias
m = torch.sqrt(F.softplus(z))
g = torch.sigmoid(self.alpha * m + self.beta)
return x * g
batch_size = 16
dim = 16384
def get_inputs():
x = torch.randn(batch_size, dim)
return [x]
def get_init_inputs():
scale = torch.randn(dim)
bias = torch.randn(dim)
return [scale, bias, 1.0, 0.0]

View File

@ -1,120 +0,0 @@
import torch
from torch.utils.cpp_extension import load_inline
source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <vector_types.h>
__global__ __launch_bounds__(256) void arcsin_tanh_affine_gate_kernel(const float* __restrict__ x, const float* __restrict__ scale, const float* __restrict__ bias, float* __restrict__ y, int B, int D, float alpha, float beta){
int row = blockIdx.x;
int tid = threadIdx.x;
int col_block = blockIdx.y;
int stride4 = blockDim.x * 4;
int col_idx = (col_block * stride4) + tid * 4;
const float* xr = x + row * D;
float* yr = y + row * D;
for(int base = col_idx; base < D; base += stride4 * 2){
if (base < D){
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
float4 yv;
float z0 = fmaf(xv.x, sv.x, bv.x);
float z1 = fmaf(xv.y, sv.y, bv.y);
float z2 = fmaf(xv.z, sv.z, bv.z);
float z3 = fmaf(xv.w, sv.w, bv.w);
float m0 = asinf(tanhf(z0));
float m1 = asinf(tanhf(z1));
float m2 = asinf(tanhf(z2));
float m3 = asinf(tanhf(z3));
float t0 = fmaf(alpha, m0, beta);
float t1 = fmaf(alpha, m1, beta);
float t2 = fmaf(alpha, m2, beta);
float t3 = fmaf(alpha, m3, beta);
float g0 = __fdividef(1.0f, 1.0f + __expf(-t0));
float g1 = __fdividef(1.0f, 1.0f + __expf(-t1));
float g2 = __fdividef(1.0f, 1.0f + __expf(-t2));
float g3 = __fdividef(1.0f, 1.0f + __expf(-t3));
yv.x = xv.x * g0;
yv.y = xv.y * g1;
yv.z = xv.z * g2;
yv.w = xv.w * g3;
reinterpret_cast<float4*>(yr + base)[0] = yv;
}
int base2 = base + stride4;
if (base2 < D){
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
float4 yv2;
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
float mb0 = asinf(tanhf(z0b));
float mb1 = asinf(tanhf(z1b));
float mb2 = asinf(tanhf(z2b));
float mb3 = asinf(tanhf(z3b));
float tb0 = fmaf(alpha, mb0, beta);
float tb1 = fmaf(alpha, mb1, beta);
float tb2 = fmaf(alpha, mb2, beta);
float tb3 = fmaf(alpha, mb3, beta);
float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
yv2.x = xv2.x * gb0;
yv2.y = xv2.y * gb1;
yv2.z = xv2.z * gb2;
yv2.w = xv2.w * gb3;
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
}
}
}
torch::Tensor arcsin_tanh_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta){
auto xc = x.contiguous();
auto sc = scale.contiguous();
auto bc = bias.contiguous();
auto y = torch::empty_like(xc);
int B = (int)xc.size(0);
int D = (int)xc.size(1);
float a = (float)alpha;
float be = (float)beta;
int block = 256;
int elements_per_thread = 4;
int elements_per_block = block * elements_per_thread;
int gy = (D + elements_per_block - 1) / elements_per_block;
dim3 grid(B, gy);
arcsin_tanh_affine_gate_kernel<<<grid, block>>>(xc.data_ptr<float>(), sc.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), B, D, a, be);
return y;
}
"""
cpp_source = """
#include <torch/extension.h>
torch::Tensor arcsin_tanh_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
"""
ops = load_inline(
name="arcsin_tanh_affine_gate",
cpp_sources=cpp_source,
cuda_sources=source,
functions=["arcsin_tanh_affine_gate_cuda"],
extra_cflags=["-O3","-std=c++17"],
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-Xptxas","-O3,-dlcm=ca","-maxrregcount=64"],
verbose=True
)
class ModelNew(torch.nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(ModelNew, self).__init__()
self.ops = ops
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.alpha = float(alpha)
self.beta = float(beta)
def forward(self, x):
return self.ops.arcsin_tanh_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)

View File

@ -1,10 +0,0 @@
Operator: ArcSinTanh-Affine-Gate (Fused CUDA Kernel)
Definition
- z = x * scale + bias
- m = asin(tanh(z))
- g = sigmoid(alpha * m + beta)
- y = x * g
Goal
- Domain-safe fusion; target ≥1.30x speedup.

View File

@ -1,50 +0,0 @@
import torch
import time
from torchcode import Model, get_inputs, get_init_inputs
from cudacode import ModelNew
def run_benchmark():
if not torch.cuda.is_available():
print("CUDA 不可用,请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
return
device = torch.device("cuda")
init_inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_init_inputs()]
inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_inputs()]
torch_model = Model(*init_inputs).cuda().eval()
cuda_model = ModelNew(*init_inputs).cuda().eval()
print("-------------------- 精度对齐验证 --------------------")
with torch.no_grad():
output_torch = torch_model(*inputs)
output_cuda = cuda_model(*inputs)
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-03)
if precision_flag:
print("✅ 精度对齐:两个模型的输出结果非常接近。")
else:
print("❌ 精度不一致!")
print("\n-------------------- 性能加速比测试 --------------------")
num_iterations = 100
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = torch_model(*inputs)
torch.cuda.synchronize(); torch_time = (time.time() - start_time) / num_iterations
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = cuda_model(*inputs)
torch.cuda.synchronize(); cuda_time = (time.time() - start_time) / num_iterations
print(f"PyTorch ArcSinTanh-Affine-Gate 平均执行时间: {torch_time:.6f}")
print(f"自定义 CUDA 融合内核 平均执行时间: {cuda_time:.6f}")
speedup = torch_time / cuda_time if cuda_time > 0 else 0
if cuda_time > 0:
print(f"加速比 (Speedup): {speedup:.2f}x")
else:
print("CUDA 内核执行时间为0无法计算加速比。")
return precision_flag, speedup
if __name__ == "__main__":
run_benchmark()

View File

@ -1,29 +0,0 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
class Model(nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(Model, self).__init__()
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.register_buffer("alpha", torch.tensor(float(alpha), dtype=torch.float32))
self.register_buffer("beta", torch.tensor(float(beta), dtype=torch.float32))
def forward(self, x: torch.Tensor) -> torch.Tensor:
z = x * self.scale + self.bias
m = torch.asin(torch.tanh(z))
g = torch.sigmoid(self.alpha * m + self.beta)
return x * g
batch_size = 16
dim = 16384
def get_inputs():
x = torch.randn(batch_size, dim)
return [x]
def get_init_inputs():
scale = torch.randn(dim)
bias = torch.randn(dim)
return [scale, bias, 1.0, 0.0]

View File

@ -1,51 +0,0 @@
import torch
from torch.utils.cpp_extension import load_inline
source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
__global__ void bias_silu_kernel(const float* x, const float* bias, float* y, int dim, long long total) {
long long idx = blockIdx.x * blockDim.x + threadIdx.x;
long long stride = blockDim.x * gridDim.x;
for (long long i = idx; i < total; i += stride) {
int j = (int)(i % dim);
float z = x[i] + bias[j];
y[i] = z / (1.0f + expf(-z));
}
}
torch::Tensor bias_silu_cuda(torch::Tensor x, torch::Tensor bias) {
auto x_contig = x.contiguous();
auto b_contig = bias.contiguous();
auto y = torch::empty_like(x_contig);
long long total = x_contig.numel();
int dim = (int)x_contig.size(-1);
int block = 512;
long long grid = (total + block - 1) / block;
grid = grid > 65535 ? 65535 : grid;
bias_silu_kernel<<<(int)grid, block>>>(x_contig.data_ptr<float>(), b_contig.data_ptr<float>(), y.data_ptr<float>(), dim, total);
return y;
}
"""
cpp_source = """
torch::Tensor bias_silu_cuda(torch::Tensor x, torch::Tensor bias);
"""
ops = load_inline(
name="bias_silu",
cpp_sources=cpp_source,
cuda_sources=source,
functions=["bias_silu_cuda"],
verbose=True
)
class ModelNew(torch.nn.Module):
def __init__(self, bias: torch.Tensor):
super(ModelNew, self).__init__()
self.ops = ops
self.register_buffer("bias", bias)
def forward(self, x):
return self.ops.bias_silu_cuda(x, self.bias)

View File

@ -1,5 +0,0 @@
Bias+SiLUSwish融合一次内核完成加偏置与 SiLU 激活,减少内核次数与显存往返。
torchcode.py参考实现 `y = silu(x + bias)`。
cudacode.py`__global__ void bias_silu_kernel(...)` 完成融合计算。
run_code.py比较精度与性能100 次迭代,`rtol=1e-03`)。

View File

@ -1,76 +0,0 @@
###########################################################
# 性能和精度验证程序
###########################################################
import torch
import torch.nn as nn
import time
from torchcode import Model,get_inputs,get_init_inputs
from cudacode import ModelNew
def run_benchmark():
if not torch.cuda.is_available():
print("CUDA 不可用,请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
return
else:
device = torch.device("cuda")
init_inputs = get_init_inputs()
init_inputs = [
x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in init_inputs
]
inputs = get_inputs()
inputs = [
x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in inputs
]
torch_model = Model(*init_inputs).cuda()
cuda_model = ModelNew(*init_inputs).cuda()
torch_model.eval()
cuda_model.eval()
print("-------------------- 精度对齐验证 --------------------")
with torch.no_grad():
output_torch = torch_model(*inputs)
output_cuda = cuda_model(*inputs)
precision_flag = torch.allclose(output_torch, output_cuda,rtol=1e-03)
if precision_flag:
print("✅ 精度对齐:两个模型的输出结果非常接近。")
else:
print("❌ 精度不一致!")
diff = (output_torch - output_cuda).abs().max().item()
print(f"最大绝对误差: {diff}")
print(f"输出张量形状: torch={tuple(output_torch.shape)}, cuda={tuple(output_cuda.shape)}")
print(f"数据类型: torch={output_torch.dtype}, cuda={output_cuda.dtype}")
print(f"设备: torch={output_torch.device}, cuda={output_cuda.device}")
print("\n-------------------- 性能加速比测试 --------------------")
num_iterations = 100
torch.cuda.synchronize()
start_time = time.time()
for _ in range(num_iterations):
_ = torch_model(*inputs)
torch.cuda.synchronize()
torch_time = (time.time() - start_time) / num_iterations
torch.cuda.synchronize()
start_time = time.time()
for _ in range(num_iterations):
_ = cuda_model(*inputs)
torch.cuda.synchronize()
cuda_time = (time.time() - start_time) / num_iterations
print(f"PyTorch Bias+SiLU 平均执行时间: {torch_time:.6f}")
print(f"自定义 CUDA 融合内核 平均执行时间: {cuda_time:.6f}")
speedup = 0
if cuda_time > 0:
speedup = torch_time / cuda_time
print(f"加速比 (Speedup): {speedup:.2f}x")
else:
print("CUDA 内核执行时间为0无法计算加速比。")
return precision_flag,speedup
if __name__ == "__main__":
precision_flag,speedup = run_benchmark()

View File

@ -1,22 +0,0 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
class Model(nn.Module):
def __init__(self, bias: torch.Tensor):
super(Model, self).__init__()
self.register_buffer("bias", bias)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return F.silu(x + self.bias)
batch_size = 16
dim = 16384
def get_inputs():
x = torch.randn(batch_size, dim)
return [x]
def get_init_inputs():
bias = torch.randn(dim)
return [bias]

View File

@ -1,132 +0,0 @@
import torch
from torch.utils.cpp_extension import load_inline
source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <vector_types.h>
__device__ __forceinline__ float sigmoidf(float x){
return __fdividef(1.0f, 1.0f + __expf(-x));
}
__global__ __launch_bounds__(256) void sigmoid_slope_affine_gate_kernel(const float* __restrict__ x, const float* __restrict__ scale, const float* __restrict__ bias, float* __restrict__ y, int B, int D, float alpha, float beta){
int row = blockIdx.x;
int tid = threadIdx.x;
int col_block = blockIdx.y;
int stride4 = blockDim.x * 4;
int col_idx = (col_block * stride4) + tid * 4;
const float* xr = x + row * D;
float* yr = y + row * D;
for(int base = col_idx; base < D; base += stride4 * 2){
if (base < D){
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
float4 yv;
float z0 = fmaf(xv.x, sv.x, bv.x);
float z1 = fmaf(xv.y, sv.y, bv.y);
float z2 = fmaf(xv.z, sv.z, bv.z);
float z3 = fmaf(xv.w, sv.w, bv.w);
float s0 = sigmoidf(z0);
float s1 = sigmoidf(z1);
float s2 = sigmoidf(z2);
float s3 = sigmoidf(z3);
float m0 = s0 * (1.0f - s0);
float m1 = s1 * (1.0f - s1);
float m2 = s2 * (1.0f - s2);
float m3 = s3 * (1.0f - s3);
float t0 = fmaf(alpha, m0, beta);
float t1 = fmaf(alpha, m1, beta);
float t2 = fmaf(alpha, m2, beta);
float t3 = fmaf(alpha, m3, beta);
float g0 = sigmoidf(t0);
float g1 = sigmoidf(t1);
float g2 = sigmoidf(t2);
float g3 = sigmoidf(t3);
yv.x = xv.x * g0;
yv.y = xv.y * g1;
yv.z = xv.z * g2;
yv.w = xv.w * g3;
reinterpret_cast<float4*>(yr + base)[0] = yv;
}
int base2 = base + stride4;
if (base2 < D){
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
float4 yv2;
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
float sb0 = sigmoidf(z0b);
float sb1 = sigmoidf(z1b);
float sb2 = sigmoidf(z2b);
float sb3 = sigmoidf(z3b);
float mb0 = sb0 * (1.0f - sb0);
float mb1 = sb1 * (1.0f - sb1);
float mb2 = sb2 * (1.0f - sb2);
float mb3 = sb3 * (1.0f - sb3);
float tb0 = fmaf(alpha, mb0, beta);
float tb1 = fmaf(alpha, mb1, beta);
float tb2 = fmaf(alpha, mb2, beta);
float tb3 = fmaf(alpha, mb3, beta);
float gb0 = sigmoidf(tb0);
float gb1 = sigmoidf(tb1);
float gb2 = sigmoidf(tb2);
float gb3 = sigmoidf(tb3);
yv2.x = xv2.x * gb0;
yv2.y = xv2.y * gb1;
yv2.z = xv2.z * gb2;
yv2.w = xv2.w * gb3;
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
}
}
}
torch::Tensor sigmoid_slope_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta){
auto xc = x.contiguous();
auto sc = scale.contiguous();
auto bc = bias.contiguous();
auto y = torch::empty_like(xc);
int B = (int)xc.size(0);
int D = (int)xc.size(1);
float a = (float)alpha;
float be = (float)beta;
int block = 256;
int elements_per_thread = 4;
int elements_per_block = block * elements_per_thread;
int gy = (D + elements_per_block - 1) / elements_per_block;
dim3 grid(B, gy);
sigmoid_slope_affine_gate_kernel<<<grid, block>>>(xc.data_ptr<float>(), sc.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), B, D, a, be);
return y;
}
"""
cpp_source = """
#include <torch/extension.h>
torch::Tensor sigmoid_slope_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
"""
ops = load_inline(
name="sigmoid_slope_affine_gate",
cpp_sources=cpp_source,
cuda_sources=source,
functions=["sigmoid_slope_affine_gate_cuda"],
extra_cflags=["-O3","-std=c++17"],
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-Xptxas","-O3,-dlcm=ca","-maxrregcount=64"],
verbose=True
)
class ModelNew(torch.nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(ModelNew, self).__init__()
self.ops = ops
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.alpha = float(alpha)
self.beta = float(beta)
def forward(self, x):
return self.ops.sigmoid_slope_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)

View File

@ -1,10 +0,0 @@
Operator: SigmoidSlope-Affine-Gate (Fused CUDA Kernel)
Definition
- z = x * scale + bias
- m = sigmoid(z) * (1 - sigmoid(z))
- g = sigmoid(alpha * m + beta)
- y = x * g
Goal
- Fuse derivative-like feature; target ≥1.30x speedup.

View File

@ -1,50 +0,0 @@
import torch
import time
from torchcode import Model, get_inputs, get_init_inputs
from cudacode import ModelNew
def run_benchmark():
if not torch.cuda.is_available():
print("CUDA 不可用,请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
return
device = torch.device("cuda")
init_inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_init_inputs()]
inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_inputs()]
torch_model = Model(*init_inputs).cuda().eval()
cuda_model = ModelNew(*init_inputs).cuda().eval()
print("-------------------- 精度对齐验证 --------------------")
with torch.no_grad():
output_torch = torch_model(*inputs)
output_cuda = cuda_model(*inputs)
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-03)
if precision_flag:
print("✅ 精度对齐:两个模型的输出结果非常接近。")
else:
print("❌ 精度不一致!")
print("\n-------------------- 性能加速比测试 --------------------")
num_iterations = 100
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = torch_model(*inputs)
torch.cuda.synchronize(); torch_time = (time.time() - start_time) / num_iterations
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = cuda_model(*inputs)
torch.cuda.synchronize(); cuda_time = (time.time() - start_time) / num_iterations
print(f"PyTorch SigmoidSlope-Affine-Gate 平均执行时间: {torch_time:.6f}")
print(f"自定义 CUDA 融合内核 平均执行时间: {cuda_time:.6f}")
speedup = torch_time / cuda_time if cuda_time > 0 else 0
if cuda_time > 0:
print(f"加速比 (Speedup): {speedup:.2f}x")
else:
print("CUDA 内核执行时间为0无法计算加速比。")
return precision_flag, speedup
if __name__ == "__main__":
run_benchmark()

View File

@ -1,29 +0,0 @@
import torch
import torch.nn as nn
class Model(nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(Model, self).__init__()
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.register_buffer("alpha", torch.tensor(float(alpha), dtype=torch.float32))
self.register_buffer("beta", torch.tensor(float(beta), dtype=torch.float32))
def forward(self, x: torch.Tensor) -> torch.Tensor:
z = x * self.scale + self.bias
s = torch.sigmoid(z)
m = s * (1.0 - s)
g = torch.sigmoid(self.alpha * m + self.beta)
return x * g
batch_size = 16
dim = 16384
def get_inputs():
x = torch.randn(batch_size, dim)
return [x]
def get_init_inputs():
scale = torch.randn(dim)
bias = torch.randn(dim)
return [scale, bias, 1.0, 0.0]

View File

@ -1,128 +0,0 @@
import torch
from torch.utils.cpp_extension import load_inline
source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <vector_types.h>
__global__ __launch_bounds__(256) void tanh_slope_affine_gate_kernel(const float* __restrict__ x, const float* __restrict__ scale, const float* __restrict__ bias, float* __restrict__ y, int B, int D, float alpha, float beta){
int row = blockIdx.x;
int tid = threadIdx.x;
int col_block = blockIdx.y;
int stride4 = blockDim.x * 4;
int col_idx = (col_block * stride4) + tid * 4;
const float* xr = x + row * D;
float* yr = y + row * D;
for(int base = col_idx; base < D; base += stride4 * 2){
if (base < D){
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
float4 yv;
float z0 = fmaf(xv.x, sv.x, bv.x);
float z1 = fmaf(xv.y, sv.y, bv.y);
float z2 = fmaf(xv.z, sv.z, bv.z);
float z3 = fmaf(xv.w, sv.w, bv.w);
float t0 = tanhf(z0);
float t1 = tanhf(z1);
float t2 = tanhf(z2);
float t3 = tanhf(z3);
float m0 = 1.0f - t0 * t0;
float m1 = 1.0f - t1 * t1;
float m2 = 1.0f - t2 * t2;
float m3 = 1.0f - t3 * t3;
float tt0 = fmaf(alpha, m0, beta);
float tt1 = fmaf(alpha, m1, beta);
float tt2 = fmaf(alpha, m2, beta);
float tt3 = fmaf(alpha, m3, beta);
float g0 = __fdividef(1.0f, 1.0f + __expf(-tt0));
float g1 = __fdividef(1.0f, 1.0f + __expf(-tt1));
float g2 = __fdividef(1.0f, 1.0f + __expf(-tt2));
float g3 = __fdividef(1.0f, 1.0f + __expf(-tt3));
yv.x = xv.x * g0;
yv.y = xv.y * g1;
yv.z = xv.z * g2;
yv.w = xv.w * g3;
reinterpret_cast<float4*>(yr + base)[0] = yv;
}
int base2 = base + stride4;
if (base2 < D){
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
float4 yv2;
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
float tb0 = tanhf(z0b);
float tb1 = tanhf(z1b);
float tb2 = tanhf(z2b);
float tb3 = tanhf(z3b);
float mb0 = 1.0f - tb0 * tb0;
float mb1 = 1.0f - tb1 * tb1;
float mb2 = 1.0f - tb2 * tb2;
float mb3 = 1.0f - tb3 * tb3;
float ttb0 = fmaf(alpha, mb0, beta);
float ttb1 = fmaf(alpha, mb1, beta);
float ttb2 = fmaf(alpha, mb2, beta);
float ttb3 = fmaf(alpha, mb3, beta);
float gb0 = __fdividef(1.0f, 1.0f + __expf(-ttb0));
float gb1 = __fdividef(1.0f, 1.0f + __expf(-ttb1));
float gb2 = __fdividef(1.0f, 1.0f + __expf(-ttb2));
float gb3 = __fdividef(1.0f, 1.0f + __expf(-ttb3));
yv2.x = xv2.x * gb0;
yv2.y = xv2.y * gb1;
yv2.z = xv2.z * gb2;
yv2.w = xv2.w * gb3;
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
}
}
}
torch::Tensor tanh_slope_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta){
auto xc = x.contiguous();
auto sc = scale.contiguous();
auto bc = bias.contiguous();
auto y = torch::empty_like(xc);
int B = (int)xc.size(0);
int D = (int)xc.size(1);
float a = (float)alpha;
float be = (float)beta;
int block = 256;
int elements_per_thread = 4;
int elements_per_block = block * elements_per_thread;
int gy = (D + elements_per_block - 1) / elements_per_block;
dim3 grid(B, gy);
tanh_slope_affine_gate_kernel<<<grid, block>>>(xc.data_ptr<float>(), sc.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), B, D, a, be);
return y;
}
"""
cpp_source = """
#include <torch/extension.h>
torch::Tensor tanh_slope_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
"""
ops = load_inline(
name="tanh_slope_affine_gate",
cpp_sources=cpp_source,
cuda_sources=source,
functions=["tanh_slope_affine_gate_cuda"],
extra_cflags=["-O3","-std=c++17"],
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-Xptxas","-O3,-dlcm=ca","-maxrregcount=64"],
verbose=True
)
class ModelNew(torch.nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(ModelNew, self).__init__()
self.ops = ops
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.alpha = float(alpha)
self.beta = float(beta)
def forward(self, x):
return self.ops.tanh_slope_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)

View File

@ -1,10 +0,0 @@
Operator: TanhSlope-Affine-Gate (Fused CUDA Kernel)
Definition
- z = x * scale + bias
- m = 1 - tanh(z)^2
- g = sigmoid(alpha * m + beta)
- y = x * g
Goal
- Fuse derivative-like feature; target ≥1.30x speedup.

View File

@ -1,50 +0,0 @@
import torch
import time
from torchcode import Model, get_inputs, get_init_inputs
from cudacode import ModelNew
def run_benchmark():
if not torch.cuda.is_available():
print("CUDA 不可用,请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
return
device = torch.device("cuda")
init_inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_init_inputs()]
inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_inputs()]
torch_model = Model(*init_inputs).cuda().eval()
cuda_model = ModelNew(*init_inputs).cuda().eval()
print("-------------------- 精度对齐验证 --------------------")
with torch.no_grad():
output_torch = torch_model(*inputs)
output_cuda = cuda_model(*inputs)
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-03)
if precision_flag:
print("✅ 精度对齐:两个模型的输出结果非常接近。")
else:
print("❌ 精度不一致!")
print("\n-------------------- 性能加速比测试 --------------------")
num_iterations = 100
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = torch_model(*inputs)
torch.cuda.synchronize(); torch_time = (time.time() - start_time) / num_iterations
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = cuda_model(*inputs)
torch.cuda.synchronize(); cuda_time = (time.time() - start_time) / num_iterations
print(f"PyTorch TanhSlope-Affine-Gate 平均执行时间: {torch_time:.6f}")
print(f"自定义 CUDA 融合内核 平均执行时间: {cuda_time:.6f}")
speedup = torch_time / cuda_time if cuda_time > 0 else 0
if cuda_time > 0:
print(f"加速比 (Speedup): {speedup:.2f}x")
else:
print("CUDA 内核执行时间为0无法计算加速比。")
return precision_flag, speedup
if __name__ == "__main__":
run_benchmark()

View File

@ -1,30 +0,0 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
class Model(nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(Model, self).__init__()
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.register_buffer("alpha", torch.tensor(float(alpha), dtype=torch.float32))
self.register_buffer("beta", torch.tensor(float(beta), dtype=torch.float32))
def forward(self, x: torch.Tensor) -> torch.Tensor:
z = x * self.scale + self.bias
t = torch.tanh(z)
m = 1.0 - t * t
g = torch.sigmoid(self.alpha * m + self.beta)
return x * g
batch_size = 16
dim = 16384
def get_inputs():
x = torch.randn(batch_size, dim)
return [x]
def get_init_inputs():
scale = torch.randn(dim)
bias = torch.randn(dim)
return [scale, bias, 1.0, 0.0]

View File

@ -1,120 +0,0 @@
import torch
from torch.utils.cpp_extension import load_inline
source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <vector_types.h>
__global__ __launch_bounds__(256) void rationalclip_affine_gate_kernel(const float* __restrict__ x, const float* __restrict__ scale, const float* __restrict__ bias, float* __restrict__ y, int B, int D, float alpha, float beta){
int row = blockIdx.x;
int tid = threadIdx.x;
int col_block = blockIdx.y;
int stride4 = blockDim.x * 4;
int col_idx = (col_block * stride4) + tid * 4;
const float* xr = x + row * D;
float* yr = y + row * D;
for(int base = col_idx; base < D; base += stride4 * 2){
if (base < D){
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
float4 yv;
float z0 = fmaf(xv.x, sv.x, bv.x);
float z1 = fmaf(xv.y, sv.y, bv.y);
float z2 = fmaf(xv.z, sv.z, bv.z);
float z3 = fmaf(xv.w, sv.w, bv.w);
float m0 = __fdividef(z0, 1.0f + z0*z0);
float m1 = __fdividef(z1, 1.0f + z1*z1);
float m2 = __fdividef(z2, 1.0f + z2*z2);
float m3 = __fdividef(z3, 1.0f + z3*z3);
float t0 = fmaf(alpha, m0, beta);
float t1 = fmaf(alpha, m1, beta);
float t2 = fmaf(alpha, m2, beta);
float t3 = fmaf(alpha, m3, beta);
float g0 = __fdividef(1.0f, 1.0f + __expf(-t0));
float g1 = __fdividef(1.0f, 1.0f + __expf(-t1));
float g2 = __fdividef(1.0f, 1.0f + __expf(-t2));
float g3 = __fdividef(1.0f, 1.0f + __expf(-t3));
yv.x = xv.x * g0;
yv.y = xv.y * g1;
yv.z = xv.z * g2;
yv.w = xv.w * g3;
reinterpret_cast<float4*>(yr + base)[0] = yv;
}
int base2 = base + stride4;
if (base2 < D){
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
float4 yv2;
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
float mb0 = __fdividef(z0b, 1.0f + z0b*z0b);
float mb1 = __fdividef(z1b, 1.0f + z1b*z1b);
float mb2 = __fdividef(z2b, 1.0f + z2b*z2b);
float mb3 = __fdividef(z3b, 1.0f + z3b*z3b);
float tb0 = fmaf(alpha, mb0, beta);
float tb1 = fmaf(alpha, mb1, beta);
float tb2 = fmaf(alpha, mb2, beta);
float tb3 = fmaf(alpha, mb3, beta);
float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
yv2.x = xv2.x * gb0;
yv2.y = xv2.y * gb1;
yv2.z = xv2.z * gb2;
yv2.w = xv2.w * gb3;
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
}
}
}
torch::Tensor rationalclip_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta){
auto xc = x.contiguous();
auto sc = scale.contiguous();
auto bc = bias.contiguous();
auto y = torch::empty_like(xc);
int B = (int)xc.size(0);
int D = (int)xc.size(1);
float a = (float)alpha;
float be = (float)beta;
int block = 256;
int elements_per_thread = 4;
int elements_per_block = block * elements_per_thread;
int gy = (D + elements_per_block - 1) / elements_per_block;
dim3 grid(B, gy);
rationalclip_affine_gate_kernel<<<grid, block>>>(xc.data_ptr<float>(), sc.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), B, D, a, be);
return y;
}
"""
cpp_source = """
#include <torch/extension.h>
torch::Tensor rationalclip_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
"""
ops = load_inline(
name="rationalclip_affine_gate",
cpp_sources=cpp_source,
cuda_sources=source,
functions=["rationalclip_affine_gate_cuda"],
extra_cflags=["-O3","-std=c++17"],
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-Xptxas","-O3,-dlcm=ca","-maxrregcount=64"],
verbose=True
)
class ModelNew(torch.nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(ModelNew, self).__init__()
self.ops = ops
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.alpha = float(alpha)
self.beta = float(beta)
def forward(self, x):
return self.ops.rationalclip_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)

View File

@ -1,10 +0,0 @@
Operator: RationalClip-Affine-Gate (Fused CUDA Kernel)
Definition
- z = x * scale + bias
- m = z / (1 + z^2)
- g = sigmoid(alpha * m + beta)
- y = x * g
Goal
- Rational clipping fusion; target ≥1.30x speedup.

View File

@ -1,50 +0,0 @@
import torch
import time
from torchcode import Model, get_inputs, get_init_inputs
from cudacode import ModelNew
def run_benchmark():
if not torch.cuda.is_available():
print("CUDA 不可用,请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
return
device = torch.device("cuda")
init_inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_init_inputs()]
inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_inputs()]
torch_model = Model(*init_inputs).cuda().eval()
cuda_model = ModelNew(*init_inputs).cuda().eval()
print("-------------------- 精度对齐验证 --------------------")
with torch.no_grad():
output_torch = torch_model(*inputs)
output_cuda = cuda_model(*inputs)
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-03)
if precision_flag:
print("✅ 精度对齐:两个模型的输出结果非常接近。")
else:
print("❌ 精度不一致!")
print("\n-------------------- 性能加速比测试 --------------------")
num_iterations = 100
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = torch_model(*inputs)
torch.cuda.synchronize(); torch_time = (time.time() - start_time) / num_iterations
torch.cuda.synchronize(); start_time = time.time()
for _ in range(num_iterations):
_ = cuda_model(*inputs)
torch.cuda.synchronize(); cuda_time = (time.time() - start_time) / num_iterations
print(f"PyTorch RationalClip-Affine-Gate 平均执行时间: {torch_time:.6f}")
print(f"自定义 CUDA 融合内核 平均执行时间: {cuda_time:.6f}")
speedup = torch_time / cuda_time if cuda_time > 0 else 0
if cuda_time > 0:
print(f"加速比 (Speedup): {speedup:.2f}x")
else:
print("CUDA 内核执行时间为0无法计算加速比。")
return precision_flag, speedup
if __name__ == "__main__":
run_benchmark()

View File

@ -1,28 +0,0 @@
import torch
import torch.nn as nn
class Model(nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(Model, self).__init__()
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.register_buffer("alpha", torch.tensor(float(alpha), dtype=torch.float32))
self.register_buffer("beta", torch.tensor(float(beta), dtype=torch.float32))
def forward(self, x: torch.Tensor) -> torch.Tensor:
z = x * self.scale + self.bias
m = z / (1.0 + z * z)
g = torch.sigmoid(self.alpha * m + self.beta)
return x * g
batch_size = 16
dim = 16384
def get_inputs():
x = torch.randn(batch_size, dim)
return [x]
def get_init_inputs():
scale = torch.randn(dim)
bias = torch.randn(dim)
return [scale, bias, 1.0, 0.0]

View File

@ -1,120 +0,0 @@
import torch
from torch.utils.cpp_extension import load_inline
source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <vector_types.h>
__global__ __launch_bounds__(256) void cosine_affine_gate_kernel(const float* __restrict__ x, const float* __restrict__ scale, const float* __restrict__ bias, float* __restrict__ y, int B, int D, float alpha, float beta){
int row = blockIdx.x;
int tid = threadIdx.x;
int col_block = blockIdx.y;
int stride4 = blockDim.x * 4;
int col_idx = (col_block * stride4) + tid * 4;
const float* xr = x + row * D;
float* yr = y + row * D;
for(int base = col_idx; base < D; base += stride4 * 2){
if (base < D){
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
float4 yv;
float z0 = fmaf(xv.x, sv.x, bv.x);
float z1 = fmaf(xv.y, sv.y, bv.y);
float z2 = fmaf(xv.z, sv.z, bv.z);
float z3 = fmaf(xv.w, sv.w, bv.w);
float m0 = cosf(z0);
float m1 = cosf(z1);
float m2 = cosf(z2);
float m3 = cosf(z3);
float t0 = fmaf(alpha, m0, beta);
float t1 = fmaf(alpha, m1, beta);
float t2 = fmaf(alpha, m2, beta);
float t3 = fmaf(alpha, m3, beta);
float g0 = __fdividef(1.0f, 1.0f + __expf(-t0));
float g1 = __fdividef(1.0f, 1.0f + __expf(-t1));
float g2 = __fdividef(1.0f, 1.0f + __expf(-t2));
float g3 = __fdividef(1.0f, 1.0f + __expf(-t3));
yv.x = xv.x * g0;
yv.y = xv.y * g1;
yv.z = xv.z * g2;
yv.w = xv.w * g3;
reinterpret_cast<float4*>(yr + base)[0] = yv;
}
int base2 = base + stride4;
if (base2 < D){
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
float4 yv2;
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
float mb0 = cosf(z0b);
float mb1 = cosf(z1b);
float mb2 = cosf(z2b);
float mb3 = cosf(z3b);
float tb0 = fmaf(alpha, mb0, beta);
float tb1 = fmaf(alpha, mb1, beta);
float tb2 = fmaf(alpha, mb2, beta);
float tb3 = fmaf(alpha, mb3, beta);
float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
yv2.x = xv2.x * gb0;
yv2.y = xv2.y * gb1;
yv2.z = xv2.z * gb2;
yv2.w = xv2.w * gb3;
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
}
}
}
torch::Tensor cosine_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta){
auto xc = x.contiguous();
auto sc = scale.contiguous();
auto bc = bias.contiguous();
auto y = torch::empty_like(xc);
int B = (int)xc.size(0);
int D = (int)xc.size(1);
float a = (float)alpha;
float be = (float)beta;
int block = 256;
int elements_per_thread = 4;
int elements_per_block = block * elements_per_thread;
int gy = (D + elements_per_block - 1) / elements_per_block;
dim3 grid(B, gy);
cosine_affine_gate_kernel<<<grid, block>>>(xc.data_ptr<float>(), sc.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), B, D, a, be);
return y;
}
"""
cpp_source = """
#include <torch/extension.h>
torch::Tensor cosine_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
"""
ops = load_inline(
name="cosine_affine_gate",
cpp_sources=cpp_source,
cuda_sources=source,
functions=["cosine_affine_gate_cuda"],
extra_cflags=["-O3","-std=c++17"],
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-Xptxas","-O3,-dlcm=ca","-maxrregcount=64"],
verbose=True
)
class ModelNew(torch.nn.Module):
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
super(ModelNew, self).__init__()
self.ops = ops
self.register_buffer("scale", scale)
self.register_buffer("bias", bias)
self.alpha = float(alpha)
self.beta = float(beta)
def forward(self, x):
return self.ops.cosine_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)

Some files were not shown because too many files have changed in this diff Show More