319 lines
10 KiB
Rust
319 lines
10 KiB
Rust
use half::bf16;
|
|
use xserv_kernels::*;
|
|
use xserv_tensor::{Device, Tensor};
|
|
|
|
fn init() {
|
|
xserv_cuda::device::set_device(0).unwrap();
|
|
}
|
|
|
|
// --- CPU reference implementations ---
|
|
|
|
fn cpu_rmsnorm(x: &[f32], gamma: &[f32], eps: f32, hidden: usize) -> Vec<f32> {
|
|
let rows = x.len() / hidden;
|
|
let mut out = vec![0.0f32; x.len()];
|
|
for r in 0..rows {
|
|
let row = &x[r * hidden..(r + 1) * hidden];
|
|
let sum_sq: f32 = row.iter().map(|v| v * v).sum();
|
|
let rms_inv = 1.0 / (sum_sq / hidden as f32 + eps).sqrt();
|
|
for i in 0..hidden {
|
|
out[r * hidden + i] = row[i] * rms_inv * gamma[i];
|
|
}
|
|
}
|
|
out
|
|
}
|
|
|
|
fn cpu_layernorm(x: &[f32], gamma: &[f32], beta: &[f32], eps: f32, hidden: usize) -> Vec<f32> {
|
|
let rows = x.len() / hidden;
|
|
let mut out = vec![0.0f32; x.len()];
|
|
for r in 0..rows {
|
|
let row = &x[r * hidden..(r + 1) * hidden];
|
|
let mean: f32 = row.iter().sum::<f32>() / hidden as f32;
|
|
let var: f32 = row.iter().map(|v| (v - mean) * (v - mean)).sum::<f32>() / hidden as f32;
|
|
let inv_std = 1.0 / (var + eps).sqrt();
|
|
for i in 0..hidden {
|
|
out[r * hidden + i] = gamma[i] * (row[i] - mean) * inv_std + beta[i];
|
|
}
|
|
}
|
|
out
|
|
}
|
|
|
|
fn cpu_gelu(x: &[f32]) -> Vec<f32> {
|
|
let sqrt_2_over_pi = 0.7978845608f32;
|
|
x.iter()
|
|
.map(|&v| {
|
|
let inner = sqrt_2_over_pi * (v + 0.044715 * v * v * v);
|
|
0.5 * v * (1.0 + inner.tanh())
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
fn cpu_silu(x: &[f32]) -> Vec<f32> {
|
|
x.iter().map(|&v| v / (1.0 + (-v).exp())).collect()
|
|
}
|
|
|
|
fn cpu_softmax(x: &[f32], cols: usize) -> Vec<f32> {
|
|
let rows = x.len() / cols;
|
|
let mut out = vec![0.0f32; x.len()];
|
|
for r in 0..rows {
|
|
let row = &x[r * cols..(r + 1) * cols];
|
|
let max = row.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
|
|
let exps: Vec<f32> = row.iter().map(|v| (v - max).exp()).collect();
|
|
let sum: f32 = exps.iter().sum();
|
|
for i in 0..cols {
|
|
out[r * cols + i] = exps[i] / sum;
|
|
}
|
|
}
|
|
out
|
|
}
|
|
|
|
fn cpu_rope(x: &mut [f32], positions: &[u32], num_heads: usize, head_dim: usize, theta: f32) {
|
|
let half_dim = head_dim / 2;
|
|
let num_tokens = positions.len();
|
|
for t in 0..num_tokens {
|
|
let pos = positions[t] as f32;
|
|
for h in 0..num_heads {
|
|
for i in 0..half_dim {
|
|
let freq = 1.0 / theta.powf(2.0 * i as f32 / head_dim as f32);
|
|
let angle = pos * freq;
|
|
let cos_val = angle.cos();
|
|
let sin_val = angle.sin();
|
|
let base = (t * num_heads + h) * head_dim;
|
|
let x0 = x[base + i];
|
|
let x1 = x[base + i + half_dim];
|
|
x[base + i] = x0 * cos_val - x1 * sin_val;
|
|
x[base + i + half_dim] = x1 * cos_val + x0 * sin_val;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
fn check_close(result: &[f32], expected: &[f32], atol: f32, name: &str) {
|
|
assert_eq!(result.len(), expected.len(), "{name}: length mismatch");
|
|
let mut max_err = 0.0f32;
|
|
for (i, (r, e)) in result.iter().zip(expected).enumerate() {
|
|
let err = (r - e).abs();
|
|
if err > max_err {
|
|
max_err = err;
|
|
}
|
|
assert!(
|
|
err <= atol,
|
|
"{name}: mismatch at [{i}]: got {r}, expected {e}, err {err}"
|
|
);
|
|
}
|
|
println!("{name}: max_err = {max_err:.6e}");
|
|
}
|
|
|
|
fn make_data(n: usize) -> Vec<f32> {
|
|
(0..n).map(|i| ((i % 17) as f32 - 8.0) * 0.1).collect()
|
|
}
|
|
|
|
// === RMSNorm ===
|
|
|
|
#[test]
|
|
fn test_rmsnorm_f32() {
|
|
init();
|
|
let hidden = 768;
|
|
let rows = 4;
|
|
let x_data = make_data(rows * hidden);
|
|
let gamma_data: Vec<f32> = (0..hidden).map(|i| 0.5 + (i % 3) as f32 * 0.2).collect();
|
|
let expected = cpu_rmsnorm(&x_data, &gamma_data, 1e-5, hidden);
|
|
|
|
let x = Tensor::from_slice(&x_data, &[rows, hidden]).to_device(Device::Cuda(0));
|
|
let gamma = Tensor::from_slice(&gamma_data, &[hidden]).to_device(Device::Cuda(0));
|
|
let out = rmsnorm(&x, &gamma, 1e-5).to_device(Device::Cpu);
|
|
check_close(out.as_slice::<f32>(), &expected, 1e-4, "rmsnorm_f32");
|
|
}
|
|
|
|
#[test]
|
|
fn test_rmsnorm_bf16() {
|
|
init();
|
|
let hidden = 768;
|
|
let rows = 4;
|
|
let x_f32 = make_data(rows * hidden);
|
|
let gamma_f32: Vec<f32> = (0..hidden).map(|i| 0.5 + (i % 3) as f32 * 0.2).collect();
|
|
let expected = cpu_rmsnorm(&x_f32, &gamma_f32, 1e-5, hidden);
|
|
|
|
let x_bf16: Vec<bf16> = x_f32.iter().map(|&v| bf16::from_f32(v)).collect();
|
|
let gamma_bf16: Vec<bf16> = gamma_f32.iter().map(|&v| bf16::from_f32(v)).collect();
|
|
let x = Tensor::from_slice(&x_bf16, &[rows, hidden]).to_device(Device::Cuda(0));
|
|
let gamma = Tensor::from_slice(&gamma_bf16, &[hidden]).to_device(Device::Cuda(0));
|
|
let out = rmsnorm(&x, &gamma, 1e-5).to_device(Device::Cpu);
|
|
|
|
let result: Vec<f32> = out.as_slice::<bf16>().iter().map(|v| v.to_f32()).collect();
|
|
check_close(&result, &expected, 0.05, "rmsnorm_bf16");
|
|
}
|
|
|
|
// === LayerNorm ===
|
|
|
|
#[test]
|
|
fn test_layernorm_f32() {
|
|
init();
|
|
let hidden = 768;
|
|
let rows = 4;
|
|
let x_data = make_data(rows * hidden);
|
|
let gamma_data: Vec<f32> = (0..hidden).map(|i| 0.8 + (i % 5) as f32 * 0.1).collect();
|
|
let beta_data: Vec<f32> = (0..hidden).map(|i| ((i % 7) as f32 - 3.0) * 0.01).collect();
|
|
let expected = cpu_layernorm(&x_data, &gamma_data, &beta_data, 1e-5, hidden);
|
|
|
|
let x = Tensor::from_slice(&x_data, &[rows, hidden]).to_device(Device::Cuda(0));
|
|
let gamma = Tensor::from_slice(&gamma_data, &[hidden]).to_device(Device::Cuda(0));
|
|
let beta = Tensor::from_slice(&beta_data, &[hidden]).to_device(Device::Cuda(0));
|
|
let out = layernorm(&x, &gamma, &beta, 1e-5).to_device(Device::Cpu);
|
|
check_close(out.as_slice::<f32>(), &expected, 1e-4, "layernorm_f32");
|
|
}
|
|
|
|
// === GELU ===
|
|
|
|
#[test]
|
|
fn test_gelu_f32() {
|
|
init();
|
|
let data = make_data(10000);
|
|
let expected = cpu_gelu(&data);
|
|
let x = Tensor::from_slice(&data, &[10000]).to_device(Device::Cuda(0));
|
|
let out = gelu(&x).to_device(Device::Cpu);
|
|
check_close(out.as_slice::<f32>(), &expected, 1e-5, "gelu_f32");
|
|
}
|
|
|
|
#[test]
|
|
fn test_gelu_bf16() {
|
|
init();
|
|
let data_f32 = make_data(10000);
|
|
let expected = cpu_gelu(&data_f32);
|
|
let data_bf16: Vec<bf16> = data_f32.iter().map(|&v| bf16::from_f32(v)).collect();
|
|
let x = Tensor::from_slice(&data_bf16, &[10000]).to_device(Device::Cuda(0));
|
|
let out = gelu(&x).to_device(Device::Cpu);
|
|
let result: Vec<f32> = out.as_slice::<bf16>().iter().map(|v| v.to_f32()).collect();
|
|
check_close(&result, &expected, 0.02, "gelu_bf16");
|
|
}
|
|
|
|
// === SiLU ===
|
|
|
|
#[test]
|
|
fn test_silu_f32() {
|
|
init();
|
|
let data = make_data(10000);
|
|
let expected = cpu_silu(&data);
|
|
let x = Tensor::from_slice(&data, &[10000]).to_device(Device::Cuda(0));
|
|
let out = silu(&x).to_device(Device::Cpu);
|
|
check_close(out.as_slice::<f32>(), &expected, 1e-5, "silu_f32");
|
|
}
|
|
|
|
// === Softmax ===
|
|
|
|
#[test]
|
|
fn test_softmax_f32() {
|
|
init();
|
|
let rows = 8;
|
|
let cols = 256;
|
|
let data = make_data(rows * cols);
|
|
let expected = cpu_softmax(&data, cols);
|
|
let x = Tensor::from_slice(&data, &[rows, cols]).to_device(Device::Cuda(0));
|
|
let out = softmax(&x).to_device(Device::Cpu);
|
|
check_close(out.as_slice::<f32>(), &expected, 1e-5, "softmax_f32");
|
|
}
|
|
|
|
#[test]
|
|
fn test_softmax_sum_to_one() {
|
|
init();
|
|
let rows = 4;
|
|
let cols = 2048;
|
|
let data: Vec<f32> = (0..rows * cols)
|
|
.map(|i| ((i % 31) as f32 - 15.0) * 0.5)
|
|
.collect();
|
|
let x = Tensor::from_slice(&data, &[rows, cols]).to_device(Device::Cuda(0));
|
|
let out = softmax(&x).to_device(Device::Cpu);
|
|
let result = out.as_slice::<f32>();
|
|
for r in 0..rows {
|
|
let row_sum: f32 = result[r * cols..(r + 1) * cols].iter().sum();
|
|
assert!(
|
|
(row_sum - 1.0).abs() < 1e-5,
|
|
"softmax row {r} sum = {row_sum}"
|
|
);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_softmax_large_values() {
|
|
init();
|
|
let data = vec![1000.0f32, 1001.0, 999.0, 1000.5];
|
|
let expected = cpu_softmax(&data, 4);
|
|
let x = Tensor::from_slice(&data, &[1, 4]).to_device(Device::Cuda(0));
|
|
let out = softmax(&x).to_device(Device::Cpu);
|
|
check_close(out.as_slice::<f32>(), &expected, 1e-5, "softmax_large");
|
|
}
|
|
|
|
// === Embedding ===
|
|
|
|
#[test]
|
|
fn test_embedding_f32() {
|
|
init();
|
|
let vocab_size = 100;
|
|
let hidden = 64;
|
|
let table_data: Vec<f32> = (0..vocab_size * hidden).map(|i| i as f32 * 0.01).collect();
|
|
let token_ids: Vec<u32> = vec![0, 5, 99, 42, 1];
|
|
|
|
let table = Tensor::from_slice(&table_data, &[vocab_size, hidden]).to_device(Device::Cuda(0));
|
|
let out = embedding(&table, &token_ids).to_device(Device::Cpu);
|
|
|
|
assert_eq!(out.shape(), &[5, hidden]);
|
|
let result = out.as_slice::<f32>();
|
|
for (seq_idx, &tid) in token_ids.iter().enumerate() {
|
|
for i in 0..hidden {
|
|
let expected = table_data[tid as usize * hidden + i];
|
|
let got = result[seq_idx * hidden + i];
|
|
assert!(
|
|
(got - expected).abs() < 1e-6,
|
|
"embedding mismatch at [{seq_idx},{i}]: got {got}, expected {expected}"
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
// === RoPE ===
|
|
|
|
#[test]
|
|
fn test_rope_f32() {
|
|
init();
|
|
let num_tokens = 4;
|
|
let num_heads = 2;
|
|
let head_dim = 8;
|
|
let theta = 10000.0f32;
|
|
let positions: Vec<u32> = vec![0, 1, 2, 3];
|
|
|
|
let x_data: Vec<f32> = (0..num_tokens * num_heads * head_dim)
|
|
.map(|i| ((i % 13) as f32 - 6.0) * 0.1)
|
|
.collect();
|
|
let mut expected = x_data.clone();
|
|
cpu_rope(&mut expected, &positions, num_heads, head_dim, theta);
|
|
|
|
let x =
|
|
Tensor::from_slice(&x_data, &[num_tokens, num_heads, head_dim]).to_device(Device::Cuda(0));
|
|
let cache = RopeCache::new(64, head_dim, theta);
|
|
rope_inplace(&x, &cache, &positions);
|
|
|
|
let out = x.to_device(Device::Cpu);
|
|
check_close(out.as_slice::<f32>(), &expected, 1e-4, "rope_f32");
|
|
}
|
|
|
|
#[test]
|
|
fn test_rope_position_0_identity() {
|
|
init();
|
|
// At position 0, all angles are 0, so cos=1, sin=0 → identity transform
|
|
let num_tokens = 1;
|
|
let num_heads = 2;
|
|
let head_dim = 8;
|
|
let positions: Vec<u32> = vec![0];
|
|
|
|
let x_data: Vec<f32> = (0..num_tokens * num_heads * head_dim)
|
|
.map(|i| (i as f32 + 1.0) * 0.1)
|
|
.collect();
|
|
|
|
let x =
|
|
Tensor::from_slice(&x_data, &[num_tokens, num_heads, head_dim]).to_device(Device::Cuda(0));
|
|
let cache = RopeCache::new(64, head_dim, 10000.0);
|
|
rope_inplace(&x, &cache, &positions);
|
|
|
|
let out = x.to_device(Device::Cpu);
|
|
check_close(out.as_slice::<f32>(), &x_data, 1e-6, "rope_pos0");
|
|
}
|