Compare commits

..

37 Commits
main ... master

Author SHA1 Message Date
xxq250 e0ed2189e5 修改流水线:ci1 2025-10-16 09:27:15 +08:00
zhangliang2021 dd0fc3af99 修改流水线:smoke 2025-10-11 17:21:51 +08:00
zhangliang2021 d2636d28c7 修改流水线:smoke 2025-10-11 17:15:51 +08:00
zhangliang2021 b774ef125d 修改流水线:smoke 2025-09-30 17:28:36 +08:00
zhangliang2021 a2bdb30f98 修改流水线:smoke 2025-09-30 17:24:18 +08:00
zhangliang2021 ee6b5523eb 修改流水线:smoke 2025-09-30 17:21:54 +08:00
zhangliang2021 6eb6783a70 修改流水线:ci1 2025-09-30 17:20:43 +08:00
zhangliang2021 f85485025b 修改流水线:ci1 2025-09-30 17:19:35 +08:00
zhangliang2021 a2d4aa787d 修改流水线:smoke 2025-09-30 16:57:33 +08:00
zhangliang2021 6534a3ec7a 修改流水线:smoke 2025-09-30 16:55:25 +08:00
zhangliang2021 03f10eab3b 修改流水线:smoke 2025-09-30 16:53:55 +08:00
zhangliang2021 dcf6dbf3e6 修改流水线:smoke 2025-09-30 16:53:17 +08:00
zhangliang2021 2e6953baec 修改流水线:smoke 2025-09-30 16:53:17 +08:00
zhangliang2021 e5e698591a 修改流水线:smoke 2025-09-30 16:50:01 +08:00
zhangliang2021 c7407254f6 创建流水线:smoke 2025-09-30 16:50:00 +08:00
zhangliang2021 2ca9ea1e06 修改流水线:ci1 2025-09-30 16:45:27 +08:00
zhangliang2021 5e9088e66e 修改流水线:ci1 2025-09-30 16:44:37 +08:00
zhangliang2021 dca08b99dc 创建流水线:ci1 2025-09-30 16:43:42 +08:00
zhangliang2021 c6cd6ecc88 Delete .gitea/workflows/ci.yml 2025-09-30 16:24:04 +08:00
zhangliang2021 e6a512bb6e Update ci.yml 2025-09-30 16:12:23 +08:00
zhangliang2021 300b5bc46e Update ci.yml 2025-09-30 15:52:35 +08:00
zhangliang2021 6790eccf55 Update ci.yml 2025-09-30 15:51:42 +08:00
zhangliang2021 9ec2371ad3 Update ci.yml 2025-09-28 13:51:28 +08:00
zhangliang2021 02a3562248 Update ci.yml 2025-09-28 13:49:42 +08:00
zhangliang2021 9fd1addada Update ci.yml 2025-09-28 11:34:23 +08:00
zhangliang2021 36e736bdbb Update ci.yml 2025-09-28 11:28:15 +08:00
zhangliang2021 fd197e18bc 修改流水线:ci 2025-09-28 11:21:22 +08:00
zhangliang2021 72e279dee9 修改流水线:ci 2025-09-28 11:12:54 +08:00
zhangliang2021 134baed8e3 修改流水线:ci 2025-09-28 11:09:54 +08:00
zhangliang2021 720e78219c 修改流水线:ci 2025-09-28 11:05:31 +08:00
zhangliang2021 5852c7d17f 修改流水线:ci 2025-09-28 11:04:24 +08:00
zhangliang2021 8eb3e952b4 修改流水线:ci 2025-09-28 11:00:59 +08:00
zhangliang2021 4270cd356b 创建流水线:ci 2025-09-26 17:26:22 +08:00
zhangliang2021 cbf82b7dab Delete .gitea/workflows/ci.yml 2025-09-26 17:25:56 +08:00
zhangliang2021 ec6fff441d 创建流水线:ci 2025-09-26 17:25:15 +08:00
zhangliang2021 97fb51ba74 Delete .gitea/workflows/ci-debug.yml 2025-09-26 17:02:08 +08:00
zhangliang2021 181a0f2faf 创建流水线:ci-debug 2025-09-26 17:01:19 +08:00
2496 changed files with 163 additions and 188604 deletions

13
.gitea/workflows/ci1.yml Normal file
View File

@ -0,0 +1,13 @@
name: "ci1"
on:
pull_request:
types: [opened, reopened]
jobs:
job-18be840c:
name: "任务[job-18be840c]"
# 运行环境,这里就是上面定义的多个 os
runs-on: 'ubuntu-latest'
steps:
- name: "运行Shell脚本[shell-0331a1]"
run: |
curl -X POST https://cicd.metax-tech.com/job/ccf_GPUCodeForces_gitlink_smoke_test --user ccf_bot:1154beee0ae7063df13f10ca76f4869268 --data-urlencode json='{"parameter": [{"name":"pr_number", "value":"${{ github.event.pull_request.number }}"}]}'

View File

@ -0,0 +1,15 @@
name: smoke test
on:
pull_request:
types: [opened, synchronize]
branches: [main]
jobs:
smoke-test:
name: smoke-test
runs-on: ubuntu-latest
steps:
- name: test
run: |
curl -X POST https://cicd.metax-tech.com/job/ccf_GPUCodeForces_gitlink_smoke_test/build --user ccf_bot:1154beee0ae7063df13f10ca76f4869268 --data-urlencode json='{"parameter": [{"name":"pr_number", "value":"${{ github.event.pull_request.number }}"}]}'

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训练营」——让每个人都能轻松了解性能优化的乐趣。

View File

@ -1,4 +1,4 @@
# GPUCodeForces代码解读
# GPUCodeForces赛题解读
## 一、赛题核心定位与整体框架
本赛题属于GPU CUDA 性能优化类任务,要求参赛选手通过自定义 CUDA Kernel 实现各类函数,并与 PyTorch 内置实现进行精度对齐和性能比拼。
这份赛题解读我们以“通过自定义CUDA Kernel实现Swish激活函数”来作为引子让大家从一个具体例子中了解到算子优化的细节。

View File

@ -17,9 +17,6 @@
* 为每个样本提供**标准输出**和**性能指标**,确保结果可复现。
* 最终形成 **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)。
---
@ -203,10 +200,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,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]

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