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>
112 lines
4.2 KiB
Rust
112 lines
4.2 KiB
Rust
//! projection (128 -> 8 + layer-norm) GPU invariants.
|
|
|
|
use approx::assert_relative_eq;
|
|
use ml_alpha::heads::{projection_gpu, ProjectionWeights, HIDDEN_DIM, PROJ_DIM};
|
|
use ml_core::device::MlDevice;
|
|
use rand::{Rng, SeedableRng};
|
|
use rand_chacha::ChaCha8Rng;
|
|
|
|
fn test_device() -> MlDevice {
|
|
MlDevice::cuda(0).expect("CUDA 0 required for ml-alpha tests")
|
|
}
|
|
|
|
#[test]
|
|
fn layer_norm_output_has_near_zero_mean_under_identity_gain() {
|
|
let dev = test_device();
|
|
let mut r = ChaCha8Rng::seed_from_u64(0x90F1);
|
|
let w = ProjectionWeights {
|
|
w: (0..PROJ_DIM * HIDDEN_DIM).map(|_| r.gen_range(-0.1..0.1)).collect(),
|
|
b: vec![0.0; PROJ_DIM],
|
|
ln_gain: vec![1.0; PROJ_DIM],
|
|
ln_bias: vec![0.0; PROJ_DIM],
|
|
};
|
|
let h: Vec<f32> = (0..HIDDEN_DIM).map(|_| r.gen_range(-1.0..1.0)).collect();
|
|
let out = projection_gpu(&dev, &w, &h).expect("gpu");
|
|
let mean: f32 = out.iter().sum::<f32>() / PROJ_DIM as f32;
|
|
assert!(
|
|
mean.abs() < 1e-4,
|
|
"layer-norm should produce zero-mean output, got mean={mean}"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn layer_norm_output_has_near_unit_variance_under_identity_gain() {
|
|
let dev = test_device();
|
|
let mut r = ChaCha8Rng::seed_from_u64(0x91F2);
|
|
let w = ProjectionWeights {
|
|
w: (0..PROJ_DIM * HIDDEN_DIM).map(|_| r.gen_range(-0.5..0.5)).collect(),
|
|
b: vec![0.0; PROJ_DIM],
|
|
ln_gain: vec![1.0; PROJ_DIM],
|
|
ln_bias: vec![0.0; PROJ_DIM],
|
|
};
|
|
let h: Vec<f32> = (0..HIDDEN_DIM).map(|_| r.gen_range(-1.0..1.0)).collect();
|
|
let out = projection_gpu(&dev, &w, &h).expect("gpu");
|
|
let mean: f32 = out.iter().sum::<f32>() / PROJ_DIM as f32;
|
|
let var: f32 = out.iter().map(|v| (v - mean).powi(2)).sum::<f32>() / PROJ_DIM as f32;
|
|
// Variance should be close to 1 (with biased-variance estimator divided
|
|
// by N — the kernel uses the same biased estimator).
|
|
assert!(
|
|
(var - 1.0_f32).abs() < 1e-3,
|
|
"layer-norm output variance should be ≈1, got {var}"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn ln_bias_shifts_output_uniformly() {
|
|
let dev = test_device();
|
|
let mut r = ChaCha8Rng::seed_from_u64(0x92F3);
|
|
let bias_offset = 3.0_f32;
|
|
let w_with = ProjectionWeights {
|
|
w: (0..PROJ_DIM * HIDDEN_DIM).map(|_| r.gen_range(-0.1..0.1)).collect(),
|
|
b: vec![0.0; PROJ_DIM],
|
|
ln_gain: vec![1.0; PROJ_DIM],
|
|
ln_bias: vec![bias_offset; PROJ_DIM],
|
|
};
|
|
let h: Vec<f32> = (0..HIDDEN_DIM).map(|_| r.gen_range(-1.0..1.0)).collect();
|
|
let out = projection_gpu(&dev, &w_with, &h).expect("gpu");
|
|
let mean: f32 = out.iter().sum::<f32>() / PROJ_DIM as f32;
|
|
assert_relative_eq!(mean, bias_offset, epsilon = 1e-4);
|
|
}
|
|
|
|
#[test]
|
|
fn ln_gain_scales_centered_output() {
|
|
let dev = test_device();
|
|
let mut r = ChaCha8Rng::seed_from_u64(0x93F4);
|
|
let gain = 2.5_f32;
|
|
let w = ProjectionWeights {
|
|
w: (0..PROJ_DIM * HIDDEN_DIM).map(|_| r.gen_range(-0.3..0.3)).collect(),
|
|
b: vec![0.0; PROJ_DIM],
|
|
ln_gain: vec![gain; PROJ_DIM],
|
|
ln_bias: vec![0.0; PROJ_DIM],
|
|
};
|
|
let h: Vec<f32> = (0..HIDDEN_DIM).map(|_| r.gen_range(-1.0..1.0)).collect();
|
|
let out = projection_gpu(&dev, &w, &h).expect("gpu");
|
|
let mean: f32 = out.iter().sum::<f32>() / PROJ_DIM as f32;
|
|
let var: f32 = out.iter().map(|v| (v - mean).powi(2)).sum::<f32>() / PROJ_DIM as f32;
|
|
// After layer-norm var ≈ 1; after gain=2.5 var ≈ 6.25.
|
|
assert!(
|
|
(var - gain * gain).abs() < 1e-2,
|
|
"variance should be gain^2 ≈ 6.25, got {var}"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn output_finite_under_zero_input() {
|
|
// Zero hidden -> pre = b -> all equal -> var = 0 -> 1e-6 clamp kicks in.
|
|
let dev = test_device();
|
|
let w = ProjectionWeights {
|
|
w: vec![1.0; PROJ_DIM * HIDDEN_DIM],
|
|
b: vec![3.0; PROJ_DIM],
|
|
ln_gain: vec![1.0; PROJ_DIM],
|
|
ln_bias: vec![0.0; PROJ_DIM],
|
|
};
|
|
let h = vec![0.0; HIDDEN_DIM];
|
|
let out = projection_gpu(&dev, &w, &h).expect("gpu");
|
|
// All pre values are equal (3.0), so after layer-norm everything is 0
|
|
// (zero variance clamped to 1e-6, then (pre - mean) = 0 → 0).
|
|
for v in &out {
|
|
assert!(v.is_finite(), "output must be finite even at zero variance, got {}", v);
|
|
assert!(v.abs() < 1e-3, "output should be ≈0 when all pre values are equal, got {}", v);
|
|
}
|
|
}
|