final submitted kernel (150 lines, 6,482 chars) — full source, untruncated
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
_cpp_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cublas_v2.h>
void linear_epilogue_launch(torch::Tensor y, torch::Tensor bias, double bias_scale, double negative_slope);
torch::Tensor fused_linear_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
double multiplier, double negative_slope) {
TORCH_CHECK(x.is_cuda(), "x must be CUDA");
TORCH_CHECK(weight.is_cuda(), "weight must be CUDA");
TORCH_CHECK(bias.is_cuda(), "bias must be CUDA");
TORCH_CHECK(x.scalar_type() == at::kFloat, "only fp32 supported");
TORCH_CHECK(weight.scalar_type() == at::kFloat, "only fp32 supported");
TORCH_CHECK(bias.scalar_type() == at::kFloat, "only fp32 supported");
TORCH_CHECK(x.dim() == 2 && weight.dim() == 2 && bias.dim() == 1, "bad dimensions");
TORCH_CHECK(x.is_contiguous() && weight.is_contiguous() && bias.is_contiguous(), "inputs must be contiguous");
const int64_t M64 = x.size(0);
const int64_t K64 = x.size(1);
const int64_t N64 = weight.size(0);
TORCH_CHECK(weight.size(1) == K64 && bias.size(0) == N64, "incompatible dimensions");
TORCH_CHECK(M64 <= INT_MAX && N64 <= INT_MAX && K64 <= INT_MAX, "dimensions too large");
const int M = (int)M64;
const int N = (int)N64;
const int K = (int)K64;
c10::cuda::CUDAGuard device_guard(x.device());
auto y = torch::empty({M64, N64}, x.options());
// Row-major Y = X * W^T. Interpreting the same storage as column-major,
// compute Y^T (N x M) = W (N x K) * X^T (K x M), where W is obtained as
// transpose of the column-major K x N view of the row-major weight tensor.
cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
const float alpha = static_cast<float>(multiplier);
const float beta = 0.0f;
cublasStatus_t stat = cublasSgemm(handle,
CUBLAS_OP_T, CUBLAS_OP_N,
N, M, K,
&alpha,
weight.data_ptr<float>(), K,
x.data_ptr<float>(), K,
&beta,
y.data_ptr<float>(), N);
TORCH_CHECK(stat == CUBLAS_STATUS_SUCCESS, "cublasSgemm failed");
linear_epilogue_launch(y, bias, multiplier, negative_slope);
return y;
}
'''
_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
__global__ void epilogue_scalar_kernel(float* __restrict__ y,
const float* __restrict__ bias,
long long total,
int n,
float bias_scale,
float negative_slope) {
long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
long long stride = (long long)blockDim.x * gridDim.x;
for (; idx < total; idx += stride) {
int col = (int)(idx % n);
float v = y[idx] + bias[col] * bias_scale;
y[idx] = (v >= 0.0f) ? v : v * negative_slope;
}
}
__global__ void epilogue_vec4_kernel(float* __restrict__ y,
const float* __restrict__ bias,
long long vec_total,
int n,
float bias_scale,
float negative_slope) {
long long vidx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
long long stride = (long long)blockDim.x * gridDim.x;
float4* __restrict__ y4 = reinterpret_cast<float4*>(y);
for (; vidx < vec_total; vidx += stride) {
long long base = vidx * 4LL;
int col = (int)(base % n);
float4 v = y4[vidx];
// Fast path used by the benchmark dimensions (n is a multiple of 4).
float b0 = bias[col + 0] * bias_scale;
float b1 = bias[col + 1] * bias_scale;
float b2 = bias[col + 2] * bias_scale;
float b3 = bias[col + 3] * bias_scale;
v.x += b0; v.y += b1; v.z += b2; v.w += b3;
v.x = (v.x >= 0.0f) ? v.x : v.x * negative_slope;
v.y = (v.y >= 0.0f) ? v.y : v.y * negative_slope;
v.z = (v.z >= 0.0f) ? v.z : v.z * negative_slope;
v.w = (v.w >= 0.0f) ? v.w : v.w * negative_slope;
y4[vidx] = v;
}
}
void linear_epilogue_launch(torch::Tensor y, torch::Tensor bias, double bias_scale, double negative_slope) {
const long long total = y.numel();
const int n = (int)y.size(1);
const int threads = 256;
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
float bs = static_cast<float>(bias_scale);
float ns = static_cast<float>(negative_slope);
// Vectorized path for aligned, four-divisible rows. The supplied problem has
// N=8192 and CUDA allocations are suitably aligned.
if ((total % 4LL) == 0 && (n % 4) == 0) {
long long vt = total / 4LL;
int blocks = (int)((vt + threads - 1) / threads);
if (blocks > 65535) blocks = 65535;
epilogue_vec4_kernel<<<blocks, threads, 0, stream>>>(y.data_ptr<float>(), bias.data_ptr<float>(), vt, n, bs, ns);
} else {
int blocks = (int)((total + threads - 1) / threads);
if (blocks > 65535) blocks = 65535;
epilogue_scalar_kernel<<<blocks, threads, 0, stream>>>(y.data_ptr<float>(), bias.data_ptr<float>(), total, n, bs, ns);
}
}
'''
_ext = load_inline(
name="fused_linear_leakyrelu_ext_v1",
cpp_sources=_cpp_src,
cuda_sources=_cuda_src,
functions=["fused_linear_forward"],
extra_cflags=["-O3"],
extra_cuda_cflags=["-O3", "--use_fast_math"],
extra_ldflags=["-lcublas"],
with_cuda=True,
verbose=False,
)
class ModelNew(nn.Module):
def __init__(self, in_features, out_features, multiplier, negative_slope):
super(ModelNew, self).__init__()
self.gemm = nn.Linear(in_features, out_features)
self.multiplier = multiplier
self.negative_slope = negative_slope
def forward(self, x):
return _ext.fused_linear_forward(x, self.gemm.weight, self.gemm.bias,
float(self.multiplier), float(self.negative_slope))