forked from ccf-ai-infra/GPUCodeForces
Merge pull request 'finish DoubleGLU #44' (#592) from uucoco/GPUCodeForces:uucoco44 into main
This commit is contained in:
commit
5f40b2e9a9
|
|
@ -0,0 +1,117 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
|
||||
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 double_glu_cuda(torch::Tensor input);
|
||||
"""
|
||||
|
||||
cuda_source = """
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
__device__ __forceinline__ float sigmoid_f(float x) {
|
||||
return 1.0f / (1.0f + expf(-x));
|
||||
}
|
||||
|
||||
__global__ void double_glu_vec4_kernel(
|
||||
const float4* __restrict__ x,
|
||||
float4* __restrict__ y,
|
||||
int chunk_vec_dim,
|
||||
int total_chunk_vecs)
|
||||
{
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int stride = gridDim.x * blockDim.x;
|
||||
|
||||
for (int i = idx; i < total_chunk_vecs; i += stride) {
|
||||
int row = i / chunk_vec_dim;
|
||||
int col = i % chunk_vec_dim;
|
||||
|
||||
// Input width = 4 * chunk
|
||||
// Output width = 2 * chunk
|
||||
|
||||
int row_offset_in = row * 4 * chunk_vec_dim;
|
||||
int row_offset_out = row * 2 * chunk_vec_dim;
|
||||
|
||||
// Process Pair 1 (G1, X1)
|
||||
int g1_idx = row_offset_in + col;
|
||||
int x1_idx = row_offset_in + chunk_vec_dim + col;
|
||||
|
||||
float4 g1 = x[g1_idx];
|
||||
float4 x1 = x[x1_idx];
|
||||
float4 out1;
|
||||
|
||||
out1.x = sigmoid_f(g1.x) * x1.x;
|
||||
out1.y = sigmoid_f(g1.y) * x1.y;
|
||||
out1.z = sigmoid_f(g1.z) * x1.z;
|
||||
out1.w = sigmoid_f(g1.w) * x1.w;
|
||||
|
||||
y[row_offset_out + col] = out1;
|
||||
|
||||
// Process Pair 2 (G2, X2)
|
||||
int g2_idx = row_offset_in + 2 * chunk_vec_dim + col;
|
||||
int x2_idx = row_offset_in + 3 * chunk_vec_dim + col;
|
||||
|
||||
float4 g2 = x[g2_idx];
|
||||
float4 x2 = x[x2_idx];
|
||||
float4 out2;
|
||||
|
||||
out2.x = sigmoid_f(g2.x) * x2.x;
|
||||
out2.y = sigmoid_f(g2.y) * x2.y;
|
||||
out2.z = sigmoid_f(g2.z) * x2.z;
|
||||
out2.w = sigmoid_f(g2.w) * x2.w;
|
||||
|
||||
y[row_offset_out + chunk_vec_dim + col] = out2;
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor double_glu_cuda(torch::Tensor input) {
|
||||
auto x_c = input.contiguous();
|
||||
|
||||
int last_dim = x_c.size(-1);
|
||||
TORCH_CHECK(last_dim % 16 == 0, "Feature dim must be divisible by 16 (4 chunks * float4) for optimization");
|
||||
|
||||
auto out_sizes = x_c.sizes().vec();
|
||||
out_sizes.back() /= 2;
|
||||
auto output = torch::empty(out_sizes, x_c.options());
|
||||
|
||||
int chunk_dim = last_dim / 4;
|
||||
int chunk_vec_dim = chunk_dim / 4;
|
||||
|
||||
int batch_size = x_c.numel() / last_dim;
|
||||
int total_chunk_vecs = batch_size * chunk_vec_dim;
|
||||
|
||||
int threads = 256;
|
||||
int blocks = (total_chunk_vecs + threads - 1) / threads;
|
||||
if (blocks > 65535) blocks = 65535;
|
||||
if (blocks == 0) blocks = 1;
|
||||
|
||||
double_glu_vec4_kernel<<<blocks, threads>>>(
|
||||
reinterpret_cast<const float4*>(x_c.data_ptr<float>()),
|
||||
reinterpret_cast<float4*>(output.data_ptr<float>()),
|
||||
chunk_vec_dim,
|
||||
total_chunk_vecs
|
||||
);
|
||||
|
||||
return output;
|
||||
}
|
||||
"""
|
||||
|
||||
self.op = load_inline(
|
||||
name="double_glu_opt_vec4",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=cuda_source,
|
||||
functions=["double_glu_cuda"],
|
||||
extra_cuda_cflags=["-O3", "--use_fast_math"],
|
||||
verbose=False
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.op.double_glu_cuda(x)
|
||||
|
|
@ -0,0 +1,21 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
g1, x1, g2, x2 = x.chunk(4, dim=-1)
|
||||
return torch.cat([torch.sigmoid(g1) * x1, torch.sigmoid(g2) * x2], dim=-1)
|
||||
|
||||
batch_size = 128
|
||||
feature_dim = 4096
|
||||
|
||||
def get_inputs():
|
||||
x = torch.randn(batch_size, feature_dim, dtype=torch.float32)
|
||||
return [x]
|
||||
|
||||
def get_init_inputs():
|
||||
return []
|
||||
|
|
@ -0,0 +1,66 @@
|
|||
You write custom CUDA kernels to replace the pytorch operators in the given GeGLU 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 chunk+gelu+elementwise_mul), or algorithmic changes (such as optimized memory access patterns). You are only limited by your imagination.
|
||||
|
||||
This CUDA kernel implements optimized Double Gated Linear Unit (Double GLU) with:
|
||||
|
||||
Memory Optimization:
|
||||
|
||||
Vectorized memory access using float4 for 4x bandwidth
|
||||
|
||||
Contiguous tensor inputs for coalesced memory access
|
||||
|
||||
Processes two GLU pairs simultaneously per thread
|
||||
|
||||
Parallelization Strategy:
|
||||
|
||||
Grid-stride loop for efficient workload distribution
|
||||
|
||||
256 threads per block optimal configuration
|
||||
|
||||
Automatic grid size calculation with 65535 block limit
|
||||
|
||||
Computational Optimization:
|
||||
|
||||
Double GLU: Two parallel GLU operations sigmoid(gate) * activation
|
||||
|
||||
Optimized sigmoid: 1.0f / (1.0f + expf(-x))
|
||||
|
||||
Fast math compilation flags for optimized exponential
|
||||
|
||||
Efficient indexing for four input chunks (G1, X1, G2, X2)
|
||||
|
||||
Work Distribution:
|
||||
|
||||
Each thread processes 8 total elements (4 per GLU pair) via float4
|
||||
|
||||
Processes two independent GLU operations simultaneously
|
||||
|
||||
Input divided into four equal chunks, output into two chunks
|
||||
|
||||
Requires input feature dimension divisible by 16 for optimal performance
|
||||
|
||||
The implementation maximizes throughput by processing two GLU operations in parallel through vectorization and efficient memory access patterns.
|
||||
|
||||
Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
g1, x1, g2, x2 = x.chunk(4, dim=-1)
|
||||
return torch.cat([torch.sigmoid(g1) * x1, torch.sigmoid(g2) * x2], dim=-1)
|
||||
|
||||
batch_size = 128
|
||||
feature_dim = 4096
|
||||
|
||||
def get_inputs():
|
||||
x = torch.randn(batch_size, feature_dim, dtype=torch.float32)
|
||||
return [x]
|
||||
|
||||
def get_init_inputs():
|
||||
return []
|
||||
|
|
@ -0,0 +1,77 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from DoubleGLU_torch import Model, get_inputs, get_init_inputs
|
||||
from DoubleGLU_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)
|
||||
|
||||
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 torch.relu 平均执行时间: {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()
|
||||
Loading…
Reference in New Issue