9ddde632b2
Single-block 8-thread kernel; thread j computes its own 128-dim dot product, then thread 0 computes block-wide mean/var, then each thread applies the per-output affine layer-norm. No atomicAdd; reductions are single-thread (8 elements — negligible cost). Tests (5/5 on sm_86) assert: - layer-norm zero-mean output under identity gain - layer-norm unit-variance output under identity gain - ln_bias shifts mean uniformly - ln_gain scales variance (var = gain^2) - finite output under zero input (variance clamp activates) Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
52 lines
1.6 KiB
Plaintext
52 lines
1.6 KiB
Plaintext
// projection.cu — h_128 -> z_8 with affine + layer-norm.
|
|
//
|
|
// pre[j] = sum_i w[j, i] * h[i] + b[j] for j in 0..8
|
|
// mean = (1/8) * sum_j pre[j]
|
|
// var = (1/8) * sum_j (pre[j] - mean)^2
|
|
// out[j] = ln_gain[j] * (pre[j] - mean) * rsqrt(var + 1e-6) + ln_bias[j]
|
|
//
|
|
// Single block, 8 threads (one per output unit). Mean and variance
|
|
// computed by a single thread (block-coordinated) to keep the kernel
|
|
// simple — the cost is negligible for 8 outputs. No atomicAdd anywhere.
|
|
|
|
extern "C" __global__ void projection_kernel(
|
|
const float* __restrict__ w, // [8, 128]
|
|
const float* __restrict__ b, // [8]
|
|
const float* __restrict__ ln_gain, // [8]
|
|
const float* __restrict__ ln_bias, // [8]
|
|
const float* __restrict__ h, // [128]
|
|
float* __restrict__ out // [8]
|
|
) {
|
|
__shared__ float pre[8];
|
|
__shared__ float mean_var[2]; // [mean, var]
|
|
|
|
int j = threadIdx.x;
|
|
if (j < 8) {
|
|
float v = b[j];
|
|
for (int i = 0; i < 128; ++i) {
|
|
v += w[j * 128 + i] * h[i];
|
|
}
|
|
pre[j] = v;
|
|
}
|
|
__syncthreads();
|
|
|
|
if (j == 0) {
|
|
float s = 0.0f, ss = 0.0f;
|
|
for (int k = 0; k < 8; ++k) {
|
|
s += pre[k];
|
|
ss += pre[k] * pre[k];
|
|
}
|
|
const float mean = s * 0.125f; // /8
|
|
const float var = ss * 0.125f - mean * mean;
|
|
mean_var[0] = mean;
|
|
mean_var[1] = fmaxf(var, 1e-6f);
|
|
}
|
|
__syncthreads();
|
|
|
|
if (j < 8) {
|
|
const float mean = mean_var[0];
|
|
const float invstd = rsqrtf(mean_var[1]);
|
|
out[j] = ln_gain[j] * (pre[j] - mean) * invstd + ln_bias[j];
|
|
}
|
|
}
|