zizi05
/
GPUCodeForces
forked from ccf-ai-infra/GPUCodeForces
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fixes softmax #3

This commit is contained in:
ZZZJ 2025-10-30 17:54:54 +08:00
parent 8c0422b9d9
commit a2b285d27b
4 changed files with 297 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

29
S1/3/prompt.txt Normal file
View File

@ -0,0 +1,29 @@
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
# softmax_torch.py
import torch
import torch.nn as nn
class Model(nn.Module):
"""使用 PyTorch 内置 nn.Softmax 的基准实现。"""
def __init__(self):
super().__init__()
self.softmax = nn.Softmax(dim=-1)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.softmax(x)
batch_size = 256
feature_dim = 4096
def get_inputs():
x = torch.randn(batch_size, feature_dim) * 5
return [x]
def get_init_inputs():
return []

88
S1/3/run_code.py Normal file
View File

@ -0,0 +1,88 @@
###########################################################
# 性能和精度验证程序
###########################################################
import torch
import torch.nn as nn
import time
from softmax_torch import Model, get_inputs, get_init_inputs
from softmax_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()

158
S1/3/softmax_cuda.py Normal file
View File

@ -0,0 +1,158 @@
# softmax_cuda.py
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
from softmax_torch import feature_dim
assert feature_dim % 4 == 0, "Feature dimension must be a multiple of 4 for float4 vectorization"
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
self._compile_cuda_kernel()
def _compile_cuda_kernel(self):
cpp_source = """
#include <torch/extension.h>
torch::Tensor softmax_forward_cuda(torch::Tensor input);
"""
cuda_source = f"""
#include <cuda_runtime.h>
#include <float.h>
#define BLOCK_SIZE 512
#define WARP_SIZE 32
__device__ __forceinline__ float warp_reduce_max(float val) {{
for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2)
val = fmaxf(val, __shfl_down_sync(0xffffffff, val, offset));
return val;
}}
__device__ __forceinline__ float warp_reduce_sum(float val) {{
for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2)
val += __shfl_down_sync(0xffffffff, val, offset);
return val;
}}
// Optimized single-pass fused Softmax kernel with float4 vectorization
__global__ void softmax_fused_vectorized_kernel(
const float* __restrict__ input,
float* __restrict__ output,
int batch_size,
int feature_dim
) {{
extern __shared__ float sdata[];
float* s_reducers = sdata;
float* s_x_cache = &sdata[BLOCK_SIZE / WARP_SIZE];
const int feature_dim_div4 = feature_dim / 4;
int row = blockIdx.x;
if (row >= batch_size) return;
const float4* x4 = reinterpret_cast<const float4*>(input + row * feature_dim);
float4* y4 = reinterpret_cast<float4*>(output + row * feature_dim);
for (int i_vec = threadIdx.x; i_vec < feature_dim_div4; i_vec += BLOCK_SIZE) {{
float4 val4 = x4[i_vec];
float* cache_ptr = s_x_cache + i_vec * 4;
cache_ptr[0] = val4.x;
cache_ptr[1] = val4.y;
cache_ptr[2] = val4.z;
cache_ptr[3] = val4.w;
}}
__syncthreads();
float thread_max = -FLT_MAX;
for (int i = threadIdx.x; i < feature_dim; i += BLOCK_SIZE) {{
thread_max = fmaxf(thread_max, s_x_cache[i]);
}}
float warp_max = warp_reduce_max(thread_max);
int warp_id = threadIdx.x / WARP_SIZE;
int lane_id = threadIdx.x % WARP_SIZE;
if (lane_id == 0) s_reducers[warp_id] = warp_max;
__syncthreads();
thread_max = (threadIdx.x < BLOCK_SIZE / WARP_SIZE) ? s_reducers[lane_id] : -FLT_MAX;
if (warp_id == 0) warp_max = warp_reduce_max(thread_max);
if (threadIdx.x == 0) s_reducers[0] = warp_max;
__syncthreads();
float row_max = s_reducers[0];
float thread_sum = 0.0f;
for (int i = threadIdx.x; i < feature_dim; i += BLOCK_SIZE) {{
thread_sum += expf(s_x_cache[i] - row_max);
}}
float warp_sum = warp_reduce_sum(thread_sum);
if (lane_id == 0) s_reducers[warp_id] = warp_sum;
__syncthreads();
thread_sum = (threadIdx.x < BLOCK_SIZE / WARP_SIZE) ? s_reducers[lane_id] : 0.0f;
if (warp_id == 0) warp_sum = warp_reduce_sum(thread_sum);
if (threadIdx.x == 0) s_reducers[0] = warp_sum;
__syncthreads();
float row_sum = s_reducers[0];
float inv_row_sum = 1.0f / row_sum;
for (int i_vec = threadIdx.x; i_vec < feature_dim_div4; i_vec += BLOCK_SIZE) {{
float* cache_ptr = s_x_cache + i_vec * 4;
float4 val4;
// Read 4 scalars from shared memory and calculate
val4.x = expf(cache_ptr[0] - row_max) * inv_row_sum;
val4.y = expf(cache_ptr[1] - row_max) * inv_row_sum;
val4.z = expf(cache_ptr[2] - row_max) * inv_row_sum;
val4.w = expf(cache_ptr[3] - row_max) * inv_row_sum;
y4[i_vec] = val4;
}}
}}
torch::Tensor softmax_forward_cuda(torch::Tensor input) {{
input = input.contiguous();
int batch_size = input.size(0);
int feature_dim = input.size(1);
if (feature_dim % 4 != 0) {{
AT_ERROR("Feature dimension must be a multiple of 4 for this kernel.");
}}
auto output = torch::empty_like(input);
const int threads = BLOCK_SIZE;
const int blocks = batch_size;
size_t shared_mem_size = (BLOCK_SIZE / WARP_SIZE + feature_dim) * sizeof(float);
softmax_fused_vectorized_kernel<<<blocks, threads, shared_mem_size>>>(
input.data_ptr<float>(),
output.data_ptr<float>(),
batch_size,
feature_dim
);
return output;
}}
"""
self.softmax_op = load_inline(
name="softmax_fused_vectorized_op",
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=["softmax_forward_cuda"],
extra_cuda_cflags=["-O3", "--use_fast_math"],
verbose=False
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.softmax_op.softmax_forward_cuda(x.contiguous())

22
S1/3/softmax_torch.py Normal file
View File

@ -0,0 +1,22 @@
# softmax_torch.py
import torch
import torch.nn as nn
class Model(nn.Module):
"""使用 PyTorch 内置 nn.Softmax 的基准实现。"""
def __init__(self):
super().__init__()
self.softmax = nn.Softmax(dim=-1)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.softmax(x)
batch_size = 256
feature_dim = 4096
def get_inputs():
x = torch.randn(batch_size, feature_dim) * 5
return [x]
def get_init_inputs():
return []