gahow/xserv
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

style: format Rust workspace

Browse Source
This commit is contained in:
Gahow Wang
2026-06-18 18:11:58 +08:00
parent 013465fc06
commit 531cd3fe08
57 changed files with 4045 additions and 1204 deletions
Download Patch File Download Diff File Expand all files Collapse all files

81
crates/xserv-kernels/tests/attention_test.rs
View File

@@ -1,11 +1,21 @@
use xserv_kernels::*;
use xserv_tensor::{Device, Tensor};
fn init() { xserv_cuda::device::set_device(0).unwrap(); }
fn init() {
xserv_cuda::device::set_device(0).unwrap();
}
fn cpu_attention(q: &[f32], k: &[f32], v: &[f32],
batch: usize, heads: usize, q_len: usize, kv_len: usize, head_dim: usize,
causal: bool) -> Vec<f32> {
fn cpu_attention(
q: &[f32],
k: &[f32],
v: &[f32],
batch: usize,
heads: usize,
q_len: usize,
kv_len: usize,
head_dim: usize,
causal: bool,
) -> Vec<f32> {
let mut out = vec![0.0f32; batch * heads * q_len * head_dim];
let scale = 1.0 / (head_dim as f32).sqrt();
@@ -70,8 +80,13 @@ fn check_close(a: &[f32], b: &[f32], atol: f32, name: &str) {
let mut max_err = 0.0f32;
for (i, (x, y)) in a.iter().zip(b).enumerate() {
let err = (x - y).abs();
if err > max_err { max_err = err; }
assert!(err <= atol, "{name}: mismatch at [{i}]: got {x}, expected {y}, err {err}");
if err > max_err {
max_err = err;
}
assert!(
err <= atol,
"{name}: mismatch at [{i}]: got {x}, expected {y}, err {err}"
);
}
println!("{name}: max_err = {max_err:.6e}");
}
@@ -105,7 +120,9 @@ fn test_batched_matmul() {
for i in 0..m {
for j in 0..n {
let mut s = 0.0f32;
for kk in 0..k { s += a_cpu[i * k + kk] * b_cpu[kk * n + j]; }
for kk in 0..k {
s += a_cpu[i * k + kk] * b_cpu[kk * n + j];
}
expected[i * n + j] = s;
}
}
@@ -116,7 +133,10 @@ fn test_batched_matmul() {
#[test]
fn test_attention_no_causal() {
init();
let b = 1; let h = 2; let s = 8; let d = 16;
let b = 1;
let h = 2;
let s = 8;
let d = 16;
let q_data = make_data(b * h * s * d);
let k_data = make_data(b * h * s * d);
let v_data = make_data(b * h * s * d);
@@ -126,13 +146,21 @@ fn test_attention_no_causal() {
let k = Tensor::from_slice(&k_data, &[b, h, s, d]).to_device(Device::Cuda(0));
let v = Tensor::from_slice(&v_data, &[b, h, s, d]).to_device(Device::Cuda(0));
let out = attention(&q, &k, &v, false).to_device(Device::Cpu);
check_close(out.as_slice::<f32>(), &expected, 1e-4, "attention_no_causal");
check_close(
out.as_slice::<f32>(),
&expected,
1e-4,
"attention_no_causal",
);
}
#[test]
fn test_attention_causal() {
init();
let b = 1; let h = 2; let s = 16; let d = 32;
let b = 1;
let h = 2;
let s = 16;
let d = 32;
let q_data = make_data(b * h * s * d);
let k_data = make_data(b * h * s * d);
let v_data = make_data(b * h * s * d);
@@ -148,7 +176,10 @@ fn test_attention_causal() {
#[test]
fn test_attention_causal_larger() {
init();
let b = 2; let h = 4; let s = 64; let d = 64;
let b = 2;
let h = 4;
let s = 64;
let d = 64;
let q_data = make_data(b * h * s * d);
let k_data = make_data(b * h * s * d);
let v_data = make_data(b * h * s * d);
@@ -158,18 +189,28 @@ fn test_attention_causal_larger() {
let k = Tensor::from_slice(&k_data, &[b, h, s, d]).to_device(Device::Cuda(0));
let v = Tensor::from_slice(&v_data, &[b, h, s, d]).to_device(Device::Cuda(0));
let out = attention(&q, &k, &v, true).to_device(Device::Cpu);
check_close(out.as_slice::<f32>(), &expected, 1e-2, "attention_causal_larger");
check_close(
out.as_slice::<f32>(),
&expected,
1e-2,
"attention_causal_larger",
);
}
#[test]
fn test_attention_causal_first_row_sees_only_first_token() {
init();
let b = 1; let h = 1; let s = 4; let d = 8;
let b = 1;
let h = 1;
let s = 4;
let d = 8;
let q_data = make_data(b * h * s * d);
let k_data = make_data(b * h * s * d);
let v_data: Vec<f32> = (0..s * d).map(|i| {
if i < d { 1.0 } else { 0.0 } // only first V row is nonzero
}).collect();
let v_data: Vec<f32> = (0..s * d)
.map(|i| {
if i < d { 1.0 } else { 0.0 } // only first V row is nonzero
})
.collect();
let q = Tensor::from_slice(&q_data, &[b, h, s, d]).to_device(Device::Cuda(0));
let k = Tensor::from_slice(&k_data, &[b, h, s, d]).to_device(Device::Cuda(0));
@@ -181,7 +222,11 @@ fn test_attention_causal_first_row_sees_only_first_token() {
// output[0] should be exactly V[0] = [1, 1, 1, ...1]
let result = out.as_slice::<f32>();
for i in 0..d {
assert!((result[i] - 1.0).abs() < 1e-5,
"first row should equal V[0], got {} at dim {}", result[i], i);
assert!(
(result[i] - 1.0).abs() < 1e-5,
"first row should equal V[0], got {} at dim {}",
result[i],
i
);
}
}

82
crates/xserv-kernels/tests/gemm_test.rs
View File

@@ -1,5 +1,5 @@
use half::bf16;
use xserv_kernels::{matmul, GemmBackend};
use xserv_kernels::{GemmBackend, matmul};
use xserv_tensor::{Device, Tensor};
fn cpu_matmul_f32(a: &[f32], b: &[f32], m: usize, n: usize, k: usize) -> Vec<f32> {
@@ -75,70 +75,110 @@ fn run_gemm_test_bf16(backend: GemmBackend, m: usize, n: usize, k: usize) {
// --- F32 tests ---
#[test]
fn test_gemm_naive_f32_small() { run_gemm_test_f32(GemmBackend::Naive, 4, 4, 4); }
fn test_gemm_naive_f32_small() {
run_gemm_test_f32(GemmBackend::Naive, 4, 4, 4);
}
#[test]
fn test_gemm_naive_f32_medium() { run_gemm_test_f32(GemmBackend::Naive, 64, 64, 64); }
fn test_gemm_naive_f32_medium() {
run_gemm_test_f32(GemmBackend::Naive, 64, 64, 64);
}
#[test]
fn test_gemm_naive_f32_rect() { run_gemm_test_f32(GemmBackend::Naive, 32, 64, 48); }
fn test_gemm_naive_f32_rect() {
run_gemm_test_f32(GemmBackend::Naive, 32, 64, 48);
}
#[test]
fn test_gemm_tiled_f32_small() { run_gemm_test_f32(GemmBackend::Tiled, 4, 4, 4); }
fn test_gemm_tiled_f32_small() {
run_gemm_test_f32(GemmBackend::Tiled, 4, 4, 4);
}
#[test]
fn test_gemm_tiled_f32_medium() { run_gemm_test_f32(GemmBackend::Tiled, 128, 128, 128); }
fn test_gemm_tiled_f32_medium() {
run_gemm_test_f32(GemmBackend::Tiled, 128, 128, 128);
}
#[test]
fn test_gemm_tiled_f32_rect() { run_gemm_test_f32(GemmBackend::Tiled, 65, 33, 97); }
fn test_gemm_tiled_f32_rect() {
run_gemm_test_f32(GemmBackend::Tiled, 65, 33, 97);
}
#[test]
fn test_gemm_cublas_f32_small() { run_gemm_test_f32(GemmBackend::CuBlas, 4, 4, 4); }
fn test_gemm_cublas_f32_small() {
run_gemm_test_f32(GemmBackend::CuBlas, 4, 4, 4);
}
#[test]
fn test_gemm_cublas_f32_medium() { run_gemm_test_f32(GemmBackend::CuBlas, 256, 256, 256); }
fn test_gemm_cublas_f32_medium() {
run_gemm_test_f32(GemmBackend::CuBlas, 256, 256, 256);
}
#[test]
fn test_gemm_cublas_f32_rect() { run_gemm_test_f32(GemmBackend::CuBlas, 65, 33, 97); }
fn test_gemm_cublas_f32_rect() {
run_gemm_test_f32(GemmBackend::CuBlas, 65, 33, 97);
}
// --- BF16 tests ---
#[test]
fn test_gemm_naive_bf16_small() { run_gemm_test_bf16(GemmBackend::Naive, 4, 4, 4); }
fn test_gemm_naive_bf16_small() {
run_gemm_test_bf16(GemmBackend::Naive, 4, 4, 4);
}
#[test]
fn test_gemm_naive_bf16_medium() { run_gemm_test_bf16(GemmBackend::Naive, 64, 64, 64); }
fn test_gemm_naive_bf16_medium() {
run_gemm_test_bf16(GemmBackend::Naive, 64, 64, 64);
}
#[test]
fn test_gemm_tiled_bf16_small() { run_gemm_test_bf16(GemmBackend::Tiled, 4, 4, 4); }
fn test_gemm_tiled_bf16_small() {
run_gemm_test_bf16(GemmBackend::Tiled, 4, 4, 4);
}
#[test]
fn test_gemm_tiled_bf16_medium() { run_gemm_test_bf16(GemmBackend::Tiled, 128, 128, 128); }
fn test_gemm_tiled_bf16_medium() {
run_gemm_test_bf16(GemmBackend::Tiled, 128, 128, 128);
}
#[test]
fn test_gemm_cublas_bf16_small() { run_gemm_test_bf16(GemmBackend::CuBlas, 4, 4, 4); }
fn test_gemm_cublas_bf16_small() {
run_gemm_test_bf16(GemmBackend::CuBlas, 4, 4, 4);
}
#[test]
fn test_gemm_cublas_bf16_medium() { run_gemm_test_bf16(GemmBackend::CuBlas, 256, 256, 256); }
fn test_gemm_cublas_bf16_medium() {
run_gemm_test_bf16(GemmBackend::CuBlas, 256, 256, 256);
}
// --- Custom GEMV tests (M=1, BF16 fast path) ---
#[test]
fn test_gemv_bf16_small() { run_gemm_test_bf16(GemmBackend::CuBlas, 1, 64, 64); }
fn test_gemv_bf16_small() {
run_gemm_test_bf16(GemmBackend::CuBlas, 1, 64, 64);
}
#[test]
fn test_gemv_bf16_medium() { run_gemm_test_bf16(GemmBackend::CuBlas, 1, 256, 256); }
fn test_gemv_bf16_medium() {
run_gemm_test_bf16(GemmBackend::CuBlas, 1, 256, 256);
}
#[test]
fn test_gemv_bf16_4096() { run_gemm_test_bf16(GemmBackend::CuBlas, 1, 4096, 4096); }
fn test_gemv_bf16_4096() {
run_gemm_test_bf16(GemmBackend::CuBlas, 1, 4096, 4096);
}
#[test]
fn test_gemv_bf16_rect() { run_gemm_test_bf16(GemmBackend::CuBlas, 1, 512, 4096); }
fn test_gemv_bf16_rect() {
run_gemm_test_bf16(GemmBackend::CuBlas, 1, 512, 4096);
}
// --- Larger benchmark-style tests ---
#[test]
fn test_gemm_cublas_f32_1024() { run_gemm_test_f32(GemmBackend::CuBlas, 1024, 1024, 1024); }
fn test_gemm_cublas_f32_1024() {
run_gemm_test_f32(GemmBackend::CuBlas, 1024, 1024, 1024);
}
#[test]
fn test_gemm_consistency_all_backends() {

46
crates/xserv-kernels/tests/ops_test.rs
View File

@@ -2,7 +2,9 @@ use half::bf16;
use xserv_kernels::*;
use xserv_tensor::{Device, Tensor};
fn init() { xserv_cuda::device::set_device(0).unwrap(); }
fn init() {
xserv_cuda::device::set_device(0).unwrap();
}
// --- CPU reference implementations ---
@@ -37,10 +39,12 @@ fn cpu_layernorm(x: &[f32], gamma: &[f32], beta: &[f32], eps: f32, hidden: usize
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()
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> {
@@ -88,8 +92,13 @@ fn check_close(result: &[f32], expected: &[f32], atol: f32, name: &str) {
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}");
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}");
}
@@ -208,13 +217,18 @@ 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 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}");
assert!(
(row_sum - 1.0).abs() < 1e-5,
"softmax row {r} sum = {row_sum}"
);
}
}
@@ -247,8 +261,10 @@ fn test_embedding_f32() {
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}");
assert!(
(got - expected).abs() < 1e-6,
"embedding mismatch at [{seq_idx},{i}]: got {got}, expected {expected}"
);
}
}
}
@@ -270,8 +286,8 @@ fn test_rope_f32() {
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 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);
@@ -292,8 +308,8 @@ fn test_rope_position_0_identity() {
.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 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);