forked from ccf-ai-infra/GPUCodeForces
finish SwAVLoss #124
This commit is contained in:
parent
23cad21ae1
commit
f6e4becd73
|
|
@ -0,0 +1,228 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
cuda_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cmath>
|
||||
|
||||
__inline__ __device__ float warpReduceMax(float val) {
|
||||
for (int offset = 16; offset > 0; offset /= 2)
|
||||
val = fmaxf(val, __shfl_down_sync(0xffffffff, val, offset));
|
||||
return val;
|
||||
}
|
||||
|
||||
__inline__ __device__ float blockReduceMax(float val) {
|
||||
__shared__ float shared[32];
|
||||
int lane = threadIdx.x % 32;
|
||||
int wid = threadIdx.x / 32;
|
||||
int num_warps = blockDim.x / 32;
|
||||
|
||||
val = warpReduceMax(val);
|
||||
if (lane == 0) shared[wid] = val;
|
||||
__syncthreads();
|
||||
|
||||
float res = (threadIdx.x < num_warps) ? shared[threadIdx.x] : -1e38f;
|
||||
if (threadIdx.x < 32) res = warpReduceMax(res);
|
||||
return res;
|
||||
}
|
||||
|
||||
__inline__ __device__ float warpReduceSum(float val) {
|
||||
for (int offset = 16; offset > 0; offset /= 2)
|
||||
val += __shfl_down_sync(0xffffffff, val, offset);
|
||||
return val;
|
||||
}
|
||||
|
||||
__inline__ __device__ float blockReduceSum(float val) {
|
||||
__shared__ float shared[32];
|
||||
int lane = threadIdx.x % 32;
|
||||
int wid = threadIdx.x / 32;
|
||||
int num_warps = blockDim.x / 32;
|
||||
|
||||
val = warpReduceSum(val);
|
||||
if (lane == 0) shared[wid] = val;
|
||||
__syncthreads();
|
||||
|
||||
float res = (threadIdx.x < num_warps) ? shared[threadIdx.x] : 0.0f;
|
||||
if (threadIdx.x < 32) res = warpReduceSum(res);
|
||||
return res;
|
||||
}
|
||||
|
||||
__global__ void normalize_kernel(const float* __restrict__ input,
|
||||
float* output,
|
||||
int rows,
|
||||
int cols) {
|
||||
int bid = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
|
||||
if (bid >= rows) return;
|
||||
|
||||
const float* row_in = input + bid * cols;
|
||||
float* row_out = output + bid * cols;
|
||||
|
||||
float sum_sq = 0.0f;
|
||||
for (int i = tid; i < cols; i += blockDim.x) {
|
||||
float val = row_in[i];
|
||||
sum_sq += val * val;
|
||||
}
|
||||
|
||||
sum_sq = blockReduceSum(sum_sq);
|
||||
|
||||
__shared__ float inv_norm;
|
||||
if (tid == 0) {
|
||||
inv_norm = rsqrtf(sum_sq + 1e-12f);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
float scale = inv_norm;
|
||||
for (int i = tid; i < cols; i += blockDim.x) {
|
||||
row_out[i] = row_in[i] * scale;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void swav_loss_kernel(
|
||||
const float* __restrict__ dots1,
|
||||
const float* __restrict__ dots2,
|
||||
float* loss,
|
||||
int batch_size,
|
||||
int num_proto,
|
||||
float temp,
|
||||
float eps
|
||||
) {
|
||||
int bid = blockIdx.x;
|
||||
if (bid >= batch_size) return;
|
||||
|
||||
int tid = threadIdx.x;
|
||||
|
||||
const float* row1 = dots1 + bid * num_proto;
|
||||
const float* row2 = dots2 + bid * num_proto;
|
||||
|
||||
float local_max1 = -1e38f;
|
||||
float local_max2 = -1e38f;
|
||||
|
||||
for (int i = tid; i < num_proto; i += blockDim.x) {
|
||||
local_max1 = fmaxf(local_max1, row1[i]);
|
||||
local_max2 = fmaxf(local_max2, row2[i]);
|
||||
}
|
||||
|
||||
float max_dot1 = blockReduceMax(local_max1);
|
||||
float max_dot2 = blockReduceMax(local_max2);
|
||||
|
||||
__shared__ float s_max1, s_max2;
|
||||
if (tid == 0) { s_max1 = max_dot1; s_max2 = max_dot2; }
|
||||
__syncthreads();
|
||||
max_dot1 = s_max1;
|
||||
max_dot2 = s_max2;
|
||||
|
||||
float inv_temp = 1.0f / temp;
|
||||
float inv_temp_eps = 1.0f / (temp * eps);
|
||||
|
||||
float l_sum_p1 = 0.0f, l_sum_p2 = 0.0f;
|
||||
float l_sum_q1 = 0.0f, l_sum_q2 = 0.0f;
|
||||
|
||||
for (int i = tid; i < num_proto; i += blockDim.x) {
|
||||
float d1 = row1[i];
|
||||
float d2 = row2[i];
|
||||
|
||||
l_sum_p1 += expf((d1 - max_dot1) * inv_temp);
|
||||
l_sum_p2 += expf((d2 - max_dot2) * inv_temp);
|
||||
|
||||
l_sum_q1 += expf((d1 - max_dot1) * inv_temp_eps);
|
||||
l_sum_q2 += expf((d2 - max_dot2) * inv_temp_eps);
|
||||
}
|
||||
|
||||
float sum_p1 = blockReduceSum(l_sum_p1);
|
||||
float sum_p2 = blockReduceSum(l_sum_p2);
|
||||
float sum_q1 = blockReduceSum(l_sum_q1);
|
||||
float sum_q2 = blockReduceSum(l_sum_q2);
|
||||
|
||||
__shared__ float s_log_sum_p1, s_log_sum_p2, s_sum_q1, s_sum_q2;
|
||||
if (tid == 0) {
|
||||
s_log_sum_p1 = logf(sum_p1);
|
||||
s_log_sum_p2 = logf(sum_p2);
|
||||
s_sum_q1 = sum_q1;
|
||||
s_sum_q2 = sum_q2;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
float l_loss = 0.0f;
|
||||
|
||||
for (int i = tid; i < num_proto; i += blockDim.x) {
|
||||
float d1 = row1[i];
|
||||
float d2 = row2[i];
|
||||
|
||||
float q1 = expf((d1 - max_dot1) * inv_temp_eps) / s_sum_q1;
|
||||
float q2 = expf((d2 - max_dot2) * inv_temp_eps) / s_sum_q2;
|
||||
|
||||
float log_p1 = (d1 - max_dot1) * inv_temp - s_log_sum_p1;
|
||||
float log_p2 = (d2 - max_dot2) * inv_temp - s_log_sum_p2;
|
||||
|
||||
l_loss += q1 * log_p2 + q2 * log_p1;
|
||||
}
|
||||
|
||||
float block_loss = blockReduceSum(l_loss);
|
||||
|
||||
if (tid == 0) {
|
||||
atomicAdd(loss, -0.5f * block_loss / batch_size);
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor swav_forward_cuda(torch::Tensor z1, torch::Tensor z2, torch::Tensor prototypes, float temperature, float epsilon) {
|
||||
auto z1_c = z1.contiguous();
|
||||
auto z2_c = z2.contiguous();
|
||||
auto p_c = prototypes.contiguous();
|
||||
|
||||
int batch_size = z1.size(0);
|
||||
int dim = z1.size(1);
|
||||
int num_proto = prototypes.size(0);
|
||||
|
||||
auto z1_n = torch::empty_like(z1_c);
|
||||
auto z2_n = torch::empty_like(z2_c);
|
||||
auto p_n = torch::empty_like(p_c);
|
||||
|
||||
int norm_block = 128;
|
||||
normalize_kernel<<<batch_size, norm_block>>>(z1_c.data_ptr<float>(), z1_n.data_ptr<float>(), batch_size, dim);
|
||||
normalize_kernel<<<batch_size, norm_block>>>(z2_c.data_ptr<float>(), z2_n.data_ptr<float>(), batch_size, dim);
|
||||
normalize_kernel<<<num_proto, norm_block>>>(p_c.data_ptr<float>(), p_n.data_ptr<float>(), num_proto, dim);
|
||||
|
||||
auto dots1 = torch::matmul(z1_n, p_n.transpose(0, 1));
|
||||
auto dots2 = torch::matmul(z2_n, p_n.transpose(0, 1));
|
||||
|
||||
auto loss = torch::zeros({1}, z1.options());
|
||||
|
||||
swav_loss_kernel<<<batch_size, 256>>>(
|
||||
dots1.data_ptr<float>(),
|
||||
dots2.data_ptr<float>(),
|
||||
loss.data_ptr<float>(),
|
||||
batch_size,
|
||||
num_proto,
|
||||
temperature,
|
||||
epsilon
|
||||
);
|
||||
|
||||
return loss;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
torch::Tensor swav_forward_cuda(torch::Tensor z1, torch::Tensor z2, torch::Tensor prototypes, float temperature, float epsilon);
|
||||
"""
|
||||
|
||||
swav_module = load_inline(
|
||||
name="swav_loss_opt_v2",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=cuda_source,
|
||||
functions=["swav_forward_cuda"],
|
||||
verbose=False
|
||||
)
|
||||
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
def __init__(self, temperature, epsilon):
|
||||
super(ModelNew, self).__init__()
|
||||
self.temperature = temperature
|
||||
self.epsilon = epsilon
|
||||
|
||||
def forward(self, z1, z2, prototypes):
|
||||
return swav_module.swav_forward_cuda(z1, z2, prototypes, self.temperature, self.epsilon)
|
||||
|
|
@ -0,0 +1,50 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, temperature, epsilon):
|
||||
super(Model, self).__init__()
|
||||
self.temperature = temperature
|
||||
self.epsilon = epsilon
|
||||
|
||||
def forward(self, z1: torch.Tensor, z2: torch.Tensor, prototypes: torch.Tensor) -> torch.Tensor:
|
||||
z1 = torch.nn.functional.normalize(z1, dim=1)
|
||||
z2 = torch.nn.functional.normalize(z2, dim=1)
|
||||
prototypes = torch.nn.functional.normalize(prototypes, dim=1)
|
||||
|
||||
scores_1 = torch.matmul(z1, prototypes.T) / self.temperature
|
||||
scores_2 = torch.matmul(z2, prototypes.T) / self.temperature
|
||||
|
||||
with torch.no_grad():
|
||||
q1 = torch.exp(scores_1 / self.epsilon)
|
||||
q1 = q1 / q1.sum(dim=1, keepdim=True)
|
||||
|
||||
q2 = torch.exp(scores_2 / self.epsilon)
|
||||
q2 = q2 / q2.sum(dim=1, keepdim=True)
|
||||
|
||||
p1 = torch.nn.functional.softmax(scores_1, dim=1)
|
||||
p2 = torch.nn.functional.softmax(scores_2, dim=1)
|
||||
|
||||
loss = -0.5 * (torch.mean(torch.sum(q1 * torch.log(p2 + 1e-8), dim=1)) +
|
||||
torch.mean(torch.sum(q2 * torch.log(p1 + 1e-8), dim=1)))
|
||||
|
||||
return loss
|
||||
|
||||
|
||||
batch_size = 16
|
||||
dim = 128
|
||||
num_prototypes = 3000
|
||||
|
||||
|
||||
def get_inputs():
|
||||
z1 = torch.randn(batch_size, dim)
|
||||
z2 = torch.randn(batch_size, dim)
|
||||
prototypes = torch.randn(num_prototypes, dim)
|
||||
return [z1, z2, prototypes]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
temperature = 0.1
|
||||
epsilon = 0.05
|
||||
return [temperature, epsilon]
|
||||
|
|
@ -1,23 +1,27 @@
|
|||
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 code implements fused residual connection + layer normalization using custom CUDA kernels. Key technologies:
|
||||
Custom CUDA kernel extension via torch.utils.cpp_extension.load_inline
|
||||
|
||||
CUDA Warp/Block Reduction: Custom warpReduceSum and blockReduceSum for efficient parallel sum calculations using warp shuffle (__shfl_down_sync)
|
||||
SwAV (Swapping Assignments between Views) loss computation
|
||||
|
||||
Single-Kernel Fusion: Merges residual addition, mean/variance calculation, and normalization into one kernel
|
||||
Multi-kernel design: normalization + loss computation
|
||||
|
||||
Shared Memory Optimization: Uses __shared__ memory for broadcasting mean and inverse std within blocks
|
||||
Warp-level reduction utilities for max/sum operations
|
||||
|
||||
Adaptive Block Size: Dynamically adjusts thread block size based on feature dimension
|
||||
Online clustering with Sinkhorn-Knopp approximation (via epsilon scaling)
|
||||
|
||||
Numerical Stability: Uses rsqrtf for inverse std with epsilon smoothing
|
||||
Prototype-based contrastive learning with temperature scaling
|
||||
|
||||
Memory Coalescing: Thread-strided loops for coalesced memory access
|
||||
Row-wise L2 normalization with shared memory optimization
|
||||
|
||||
PyTorch Integration: Combines custom CUDA kernel with PyTorch's nn.LayerNorm parameters
|
||||
Matrix multiplication for prototype assignment scores
|
||||
|
||||
Use: Transformer architectures, high-performance normalization layers, fused training operations.
|
||||
Numerically stable softmax with max subtraction
|
||||
|
||||
Atomic addition (atomicAdd) for loss accumulation
|
||||
|
||||
Symmetric loss computation between two augmented views
|
||||
|
||||
|
||||
|
||||
|
|
@ -26,22 +30,50 @@ Here's an example to show you the syntax of inline embedding custom CUDA operato
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, normalized_shape):
|
||||
super(Model, self).__init__()
|
||||
self.layer_norm = nn.LayerNorm(normalized_shape)
|
||||
|
||||
def forward(self, x, residual):
|
||||
return self.layer_norm(x + residual)
|
||||
class Model(nn.Module):
|
||||
def __init__(self, temperature, epsilon):
|
||||
super(Model, self).__init__()
|
||||
self.temperature = temperature
|
||||
self.epsilon = epsilon
|
||||
|
||||
def forward(self, z1: torch.Tensor, z2: torch.Tensor, prototypes: torch.Tensor) -> torch.Tensor:
|
||||
z1 = torch.nn.functional.normalize(z1, dim=1)
|
||||
z2 = torch.nn.functional.normalize(z2, dim=1)
|
||||
prototypes = torch.nn.functional.normalize(prototypes, dim=1)
|
||||
|
||||
scores_1 = torch.matmul(z1, prototypes.T) / self.temperature
|
||||
scores_2 = torch.matmul(z2, prototypes.T) / self.temperature
|
||||
|
||||
with torch.no_grad():
|
||||
q1 = torch.exp(scores_1 / self.epsilon)
|
||||
q1 = q1 / q1.sum(dim=1, keepdim=True)
|
||||
|
||||
q2 = torch.exp(scores_2 / self.epsilon)
|
||||
q2 = q2 / q2.sum(dim=1, keepdim=True)
|
||||
|
||||
p1 = torch.nn.functional.softmax(scores_1, dim=1)
|
||||
p2 = torch.nn.functional.softmax(scores_2, dim=1)
|
||||
|
||||
loss = -0.5 * (torch.mean(torch.sum(q1 * torch.log(p2 + 1e-8), dim=1)) +
|
||||
torch.mean(torch.sum(q2 * torch.log(p1 + 1e-8), dim=1)))
|
||||
|
||||
return loss
|
||||
|
||||
|
||||
batch_size = 16
|
||||
seq_len = 64
|
||||
hidden_dim = 256
|
||||
dim = 128
|
||||
num_prototypes = 3000
|
||||
|
||||
|
||||
def get_inputs():
|
||||
x = torch.randn(batch_size, seq_len, hidden_dim)
|
||||
residual = torch.randn(batch_size, seq_len, hidden_dim)
|
||||
return [x, residual]
|
||||
z1 = torch.randn(batch_size, dim)
|
||||
z2 = torch.randn(batch_size, dim)
|
||||
prototypes = torch.randn(num_prototypes, dim)
|
||||
return [z1, z2, prototypes]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
return [hid
|
||||
temperature = 0.1
|
||||
epsilon = 0.05
|
||||
return [temperature, epsilon]
|
||||
|
|
@ -1,152 +0,0 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
cuda_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
__inline__ __device__ float warpReduceSum(float val) {
|
||||
for (int offset = 32 / 2; offset > 0; offset /= 2)
|
||||
val += __shfl_down_sync(0xffffffff, val, offset);
|
||||
return val;
|
||||
}
|
||||
|
||||
__inline__ __device__ float blockReduceSum(float val) {
|
||||
static __shared__ float shared[32];
|
||||
int lane = threadIdx.x % 32;
|
||||
int wid = threadIdx.x / 32;
|
||||
|
||||
val = warpReduceSum(val);
|
||||
|
||||
if (lane == 0) shared[wid] = val;
|
||||
__syncthreads();
|
||||
|
||||
val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : 0.0f;
|
||||
|
||||
if (wid == 0) val = warpReduceSum(val);
|
||||
|
||||
return val;
|
||||
}
|
||||
|
||||
__global__ void residual_add_layernorm_kernel(
|
||||
const float* __restrict__ x,
|
||||
const float* __restrict__ residual,
|
||||
const float* __restrict__ gamma,
|
||||
const float* __restrict__ beta,
|
||||
float* __restrict__ output,
|
||||
int n_rows,
|
||||
int n_cols,
|
||||
float epsilon
|
||||
) {
|
||||
__shared__ float s_mean;
|
||||
__shared__ float s_inv_std;
|
||||
|
||||
int bid = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
|
||||
const float* x_row = x + bid * n_cols;
|
||||
const float* res_row = residual + bid * n_cols;
|
||||
float* out_row = output + bid * n_cols;
|
||||
|
||||
float sum = 0.0f;
|
||||
for (int i = tid; i < n_cols; i += blockDim.x) {
|
||||
sum += x_row[i] + res_row[i];
|
||||
}
|
||||
|
||||
sum = blockReduceSum(sum);
|
||||
|
||||
if (tid == 0) {
|
||||
s_mean = sum / n_cols;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
float mean = s_mean;
|
||||
float sum_sq_diff = 0.0f;
|
||||
|
||||
for (int i = tid; i < n_cols; i += blockDim.x) {
|
||||
float val = x_row[i] + res_row[i];
|
||||
float diff = val - mean;
|
||||
sum_sq_diff += diff * diff;
|
||||
}
|
||||
|
||||
sum_sq_diff = blockReduceSum(sum_sq_diff);
|
||||
|
||||
if (tid == 0) {
|
||||
s_inv_std = rsqrtf(sum_sq_diff / n_cols + epsilon);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
float inv_std = s_inv_std;
|
||||
|
||||
for (int i = tid; i < n_cols; i += blockDim.x) {
|
||||
float val = x_row[i] + res_row[i];
|
||||
out_row[i] = ((val - mean) * inv_std) * gamma[i] + beta[i];
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor residual_add_layernorm_cuda(
|
||||
torch::Tensor x,
|
||||
torch::Tensor residual,
|
||||
torch::Tensor gamma,
|
||||
torch::Tensor beta,
|
||||
float epsilon
|
||||
) {
|
||||
int n_rows = x.size(0) * x.size(1);
|
||||
int n_cols = x.size(2);
|
||||
|
||||
auto output = torch::empty_like(x);
|
||||
|
||||
int block_size = 256;
|
||||
while (block_size < n_cols && block_size < 1024) {
|
||||
block_size *= 2;
|
||||
}
|
||||
|
||||
residual_add_layernorm_kernel<<<n_rows, block_size>>>(
|
||||
x.data_ptr<float>(),
|
||||
residual.data_ptr<float>(),
|
||||
gamma.data_ptr<float>(),
|
||||
beta.data_ptr<float>(),
|
||||
output.data_ptr<float>(),
|
||||
n_rows,
|
||||
n_cols,
|
||||
epsilon
|
||||
);
|
||||
|
||||
return output;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
torch::Tensor residual_add_layernorm_cuda(
|
||||
torch::Tensor x,
|
||||
torch::Tensor residual,
|
||||
torch::Tensor gamma,
|
||||
torch::Tensor beta,
|
||||
float epsilon
|
||||
);
|
||||
"""
|
||||
|
||||
module = load_inline(
|
||||
name="residual_add_layernorm_opt",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=cuda_source,
|
||||
functions=["residual_add_layernorm_cuda"],
|
||||
verbose=True
|
||||
)
|
||||
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
def __init__(self, normalized_shape):
|
||||
super(ModelNew, self).__init__()
|
||||
self.layer_norm = nn.LayerNorm(normalized_shape)
|
||||
self.module = module
|
||||
|
||||
def forward(self, x, residual):
|
||||
return self.module.residual_add_layernorm_cuda(
|
||||
x,
|
||||
residual,
|
||||
self.layer_norm.weight,
|
||||
self.layer_norm.bias,
|
||||
self.layer_norm.eps
|
||||
)
|
||||
|
|
@ -1,22 +0,0 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, normalized_shape):
|
||||
super(Model, self).__init__()
|
||||
self.layer_norm = nn.LayerNorm(normalized_shape)
|
||||
|
||||
def forward(self, x, residual):
|
||||
return self.layer_norm(x + residual)
|
||||
|
||||
batch_size = 16
|
||||
seq_len = 64
|
||||
hidden_dim = 256
|
||||
|
||||
def get_inputs():
|
||||
x = torch.randn(batch_size, seq_len, hidden_dim)
|
||||
residual = torch.randn(batch_size, seq_len, hidden_dim)
|
||||
return [x, residual]
|
||||
|
||||
def get_init_inputs():
|
||||
return [hidden_dim]
|
||||
|
|
@ -4,8 +4,8 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from residuallayernorm_torch import Model, get_inputs, get_init_inputs
|
||||
from residuallayernorm_cuda import ModelNew
|
||||
from SwAVLoss_torch import Model, get_inputs, get_init_inputs
|
||||
from SwAVLoss_cuda import ModelNew
|
||||
|
||||
|
||||
def run_benchmark():
|
||||
|
|
|
|||
Loading…
Reference in New Issue