234 lines
6.6 KiB
Plaintext
234 lines
6.6 KiB
Plaintext
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* normalization.cu — CUDA normalization kernels for Foundry
|
|
*
|
|
* Target: sm_37 (K80, Kepler)
|
|
* Compile: nvcc -ptx -arch=sm_37 normalization.cu -o normalization.ptx
|
|
*
|
|
* RMS norm: single-block reduction with shared memory
|
|
* Softmax: one block per row, two-pass (max + exp-sum) with shared memory
|
|
* QK-RMSNorm: one block per head, per-head RMS over head_dim, shared weight
|
|
*
|
|
* PTX symbols emitted (for embedding into ../cuda/fc_cuda_ptx.h):
|
|
* FC_PTX_NORMALIZATION -> kernels fc_rms_norm, fc_softmax, fc_qk_rms_norm
|
|
*/
|
|
|
|
extern "C" {
|
|
|
|
/*
|
|
* RMS normalization: out[i] = (x[i] / rms) * weight[i]
|
|
* where rms = sqrt(mean(x^2) + eps)
|
|
*
|
|
* Launched with 1 block, blockDim.x threads.
|
|
* Shared memory: blockDim.x * sizeof(float)
|
|
*
|
|
* weight may be NULL (passed as 0 device pointer) — if so, skip scaling.
|
|
*/
|
|
__global__ void fc_rms_norm(const float *x, const float *weight, float *out,
|
|
int n, float eps) {
|
|
extern __shared__ float sdata[];
|
|
|
|
int tid = threadIdx.x;
|
|
int threads = blockDim.x;
|
|
|
|
/* Phase 1: compute partial sum of squares */
|
|
float sum_sq = 0.0f;
|
|
for (int i = tid; i < n; i += threads) {
|
|
float v = x[i];
|
|
sum_sq += v * v;
|
|
}
|
|
sdata[tid] = sum_sq;
|
|
__syncthreads();
|
|
|
|
/* Shared memory reduction */
|
|
for (int s = threads / 2; s > 0; s >>= 1) {
|
|
if (tid < s) {
|
|
sdata[tid] += sdata[tid + s];
|
|
}
|
|
__syncthreads();
|
|
}
|
|
|
|
/* Phase 2: compute RMS and normalize */
|
|
float rms = sqrtf(sdata[0] / (float)n + eps);
|
|
float inv_rms = 1.0f / rms;
|
|
|
|
for (int i = tid; i < n; i += threads) {
|
|
float v = x[i] * inv_rms;
|
|
if (weight) {
|
|
v *= weight[i];
|
|
}
|
|
out[i] = v;
|
|
}
|
|
}
|
|
|
|
/*
|
|
* Softmax over last dimension.
|
|
* x: [rows, cols], out: [rows, cols]
|
|
*
|
|
* Launched with rows blocks, blockDim.x threads per block.
|
|
* Shared memory: blockDim.x * sizeof(float) * 2
|
|
* First half for max reduction, second half for sum reduction.
|
|
*/
|
|
__global__ void fc_softmax(const float *x, float *out, int rows, int cols) {
|
|
extern __shared__ float sdata[];
|
|
|
|
int row = blockIdx.x;
|
|
if (row >= rows) return;
|
|
|
|
int tid = threadIdx.x;
|
|
int threads = blockDim.x;
|
|
|
|
const float *x_row = x + row * cols;
|
|
float *o_row = out + row * cols;
|
|
|
|
float *smax = sdata; /* [threads] for max reduction */
|
|
float *ssum = sdata + threads; /* [threads] for sum reduction */
|
|
|
|
/* Pass 1: find max */
|
|
float local_max = -1e30f;
|
|
for (int i = tid; i < cols; i += threads) {
|
|
float v = x_row[i];
|
|
if (v > local_max) local_max = v;
|
|
}
|
|
smax[tid] = local_max;
|
|
__syncthreads();
|
|
|
|
for (int s = threads / 2; s > 0; s >>= 1) {
|
|
if (tid < s) {
|
|
if (smax[tid + s] > smax[tid])
|
|
smax[tid] = smax[tid + s];
|
|
}
|
|
__syncthreads();
|
|
}
|
|
float row_max = smax[0];
|
|
|
|
/* Pass 2: compute exp(x - max) and sum */
|
|
float local_sum = 0.0f;
|
|
for (int i = tid; i < cols; i += threads) {
|
|
float v = expf(x_row[i] - row_max);
|
|
o_row[i] = v; /* store intermediate exp values */
|
|
local_sum += v;
|
|
}
|
|
ssum[tid] = local_sum;
|
|
__syncthreads();
|
|
|
|
for (int s = threads / 2; s > 0; s >>= 1) {
|
|
if (tid < s) {
|
|
ssum[tid] += ssum[tid + s];
|
|
}
|
|
__syncthreads();
|
|
}
|
|
float inv_sum = 1.0f / ssum[0];
|
|
|
|
/* Pass 3: normalize */
|
|
for (int i = tid; i < cols; i += threads) {
|
|
o_row[i] *= inv_sum;
|
|
}
|
|
}
|
|
|
|
/*
|
|
* Per-head QK-RMSNorm, in place.
|
|
* data: [n_heads * head_dim] fp32; each head is RMS-normalized over its own
|
|
* head_dim and scaled by the shared `weight` [head_dim]. The norm weight is
|
|
* shared across all heads — the GGUF stores it pre-baked (incl. any (1+w)
|
|
* offset), so this is just one RMS-norm per head.
|
|
*
|
|
* Mirrors the CPU reference Pantheon/src/models/qwen35moe_attn.c
|
|
* qwen35moe_qk_norm_heads (lines 39-46), which calls pn_vrms_norm_f32 once per
|
|
* head: out[i] = (x[i] / sqrt(mean(x^2) + eps)) * weight[i].
|
|
*
|
|
* Launched with `n_heads` blocks, blockDim.x threads (use 256).
|
|
* Shared memory: blockDim.x * sizeof(float) (sum-of-squares reduction).
|
|
*
|
|
* weight may be NULL (passed as 0 device pointer) — if so, skip scaling.
|
|
*/
|
|
__global__ void fc_qk_rms_norm(float *data, const float *weight,
|
|
int n_heads, int head_dim, float eps) {
|
|
extern __shared__ float sdata[];
|
|
|
|
int head = blockIdx.x;
|
|
if (head >= n_heads) return;
|
|
|
|
int tid = threadIdx.x;
|
|
int threads = blockDim.x;
|
|
float *h = data + (size_t)head * head_dim;
|
|
|
|
/* Phase 1: partial sum of squares over this head's dims */
|
|
float sum_sq = 0.0f;
|
|
for (int i = tid; i < head_dim; i += threads) {
|
|
float v = h[i];
|
|
sum_sq += v * v;
|
|
}
|
|
sdata[tid] = sum_sq;
|
|
__syncthreads();
|
|
|
|
for (int s = threads / 2; s > 0; s >>= 1) {
|
|
if (tid < s) {
|
|
sdata[tid] += sdata[tid + s];
|
|
}
|
|
__syncthreads();
|
|
}
|
|
|
|
/* Phase 2: normalize and scale by the shared weight */
|
|
float rms = sqrtf(sdata[0] / (float)head_dim + eps);
|
|
float inv_rms = 1.0f / rms;
|
|
|
|
for (int i = tid; i < head_dim; i += threads) {
|
|
float v = h[i] * inv_rms;
|
|
if (weight) {
|
|
v *= weight[i];
|
|
}
|
|
h[i] = v;
|
|
}
|
|
}
|
|
|
|
/*
|
|
* Batched (prefill) RMS norm: one block per ROW of x [rows][n], each row
|
|
* normalized independently with the exact fc_rms_norm op order (same
|
|
* blockDim -> same reduction tree -> bit-exact per row).
|
|
*
|
|
* Launched with `rows` blocks, blockDim.x threads (the single-row
|
|
* launcher's choice: min(256, n)). Shared: blockDim.x * sizeof(float).
|
|
*/
|
|
__global__ void fc_rms_norm_nx(const float *x, const float *weight, float *out,
|
|
int n, float eps, int rows) {
|
|
extern __shared__ float sdata[];
|
|
|
|
int row = blockIdx.x;
|
|
if (row >= rows) return;
|
|
|
|
int tid = threadIdx.x;
|
|
int threads = blockDim.x;
|
|
const float *xr = x + (size_t)row * n;
|
|
float *outr = out + (size_t)row * n;
|
|
|
|
float sum_sq = 0.0f;
|
|
for (int i = tid; i < n; i += threads) {
|
|
float v = xr[i];
|
|
sum_sq += v * v;
|
|
}
|
|
sdata[tid] = sum_sq;
|
|
__syncthreads();
|
|
|
|
for (int s = threads / 2; s > 0; s >>= 1) {
|
|
if (tid < s) {
|
|
sdata[tid] += sdata[tid + s];
|
|
}
|
|
__syncthreads();
|
|
}
|
|
|
|
float rms = sqrtf(sdata[0] / (float)n + eps);
|
|
float inv_rms = 1.0f / rms;
|
|
|
|
for (int i = tid; i < n; i += threads) {
|
|
float v = xr[i] * inv_rms;
|
|
if (weight) {
|
|
v *= weight[i];
|
|
}
|
|
outr[i] = v;
|
|
}
|
|
}
|
|
|
|
} /* extern "C" */
|