118 lines
4.0 KiB
Plaintext
118 lines
4.0 KiB
Plaintext
// RMSNorm + LayerNorm forward kernels -- BF16-native.
|
|
//
|
|
// Common header (common_device_functions.cuh) is prepended by build.rs
|
|
// providing: float, (), 0.0f, 1.0f, etc.
|
|
|
|
// ── RMSNorm forward ────────────────────────────────────────────────────
|
|
// Each block processes one sample (row) of the [batch, features] input.
|
|
// Computes: y = x / sqrt(mean(x^2) + eps) * weight
|
|
//
|
|
// Two-pass approach:
|
|
// Pass 1: compute mean(x^2) via parallel reduction in shared memory (f32 accum).
|
|
// Pass 2: normalize and scale each element.
|
|
|
|
extern "C" __global__
|
|
void rmsnorm_forward(const float* __restrict__ x,
|
|
float* __restrict__ y,
|
|
const float* __restrict__ weight,
|
|
int batch,
|
|
int features,
|
|
float eps) {
|
|
int row = blockIdx.x;
|
|
if (row >= batch) return;
|
|
|
|
const float* x_row = x + row * features;
|
|
float* y_row = y + row * features;
|
|
|
|
extern __shared__ float smem[];
|
|
|
|
// Thread-local accumulation of x^2 over features (f32 for precision)
|
|
float local_sq = 0.0f;
|
|
for (int j = threadIdx.x; j < features; j += blockDim.x) {
|
|
float val = (x_row[j]);
|
|
local_sq += val * val;
|
|
}
|
|
|
|
smem[threadIdx.x] = local_sq;
|
|
__syncthreads();
|
|
|
|
// Parallel reduction for sum of squares
|
|
for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
|
|
if (threadIdx.x < stride) {
|
|
smem[threadIdx.x] += smem[threadIdx.x + stride];
|
|
}
|
|
__syncthreads();
|
|
}
|
|
|
|
// rms = sqrt(mean(x^2) + eps)
|
|
float inv_rms = rsqrtf(smem[0] / (float)features + eps);
|
|
|
|
// Normalize and scale
|
|
for (int j = threadIdx.x; j < features; j += blockDim.x) {
|
|
float xf = (x_row[j]);
|
|
float wf = (weight[j]);
|
|
y_row[j] = (xf * inv_rms * wf);
|
|
}
|
|
}
|
|
|
|
// ── LayerNorm forward ──────────────────────────────────────────────────
|
|
// Each block processes one sample (row) of the [batch, features] input.
|
|
// Two-pass approach:
|
|
// Pass 1: compute mean and variance via parallel reduction in shared memory.
|
|
// Pass 2: normalize, scale, and shift each element.
|
|
|
|
extern "C" __global__
|
|
void layernorm_forward(const float* __restrict__ x,
|
|
float* __restrict__ y,
|
|
const float* __restrict__ gamma,
|
|
const float* __restrict__ beta,
|
|
int batch,
|
|
int features,
|
|
float eps) {
|
|
int row = blockIdx.x;
|
|
if (row >= batch) return;
|
|
|
|
const float* x_row = x + row * features;
|
|
float* y_row = y + row * features;
|
|
|
|
// Shared memory: first half for sum, second half for sum of squares
|
|
extern __shared__ float smem[];
|
|
float* s_sum = smem;
|
|
float* s_sq = smem + blockDim.x;
|
|
|
|
// Thread-local accumulation over features (f32 for precision)
|
|
float local_sum = 0.0f;
|
|
float local_sq = 0.0f;
|
|
for (int j = threadIdx.x; j < features; j += blockDim.x) {
|
|
float val = (x_row[j]);
|
|
local_sum += val;
|
|
local_sq += val * val;
|
|
}
|
|
|
|
s_sum[threadIdx.x] = local_sum;
|
|
s_sq[threadIdx.x] = local_sq;
|
|
__syncthreads();
|
|
|
|
// Parallel reduction
|
|
for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
|
|
if (threadIdx.x < stride) {
|
|
s_sum[threadIdx.x] += s_sum[threadIdx.x + stride];
|
|
s_sq[threadIdx.x] += s_sq[threadIdx.x + stride];
|
|
}
|
|
__syncthreads();
|
|
}
|
|
|
|
float mean = s_sum[0] / (float)features;
|
|
float var = s_sq[0] / (float)features - mean * mean;
|
|
float inv_std = rsqrtf(var + eps);
|
|
|
|
// Normalize, scale, shift
|
|
for (int j = threadIdx.x; j < features; j += blockDim.x) {
|
|
float xf = (x_row[j]);
|
|
float gf = (gamma[j]);
|
|
float bf = (beta[j]);
|
|
float normalized = (xf - mean) * inv_std;
|
|
y_row[j] = (normalized * gf + bf);
|
|
}
|
|
}
|