diff --git a/crates/xserv-model/src/bin/sink-attn-check.rs b/crates/xserv-model/src/bin/sink-attn-check.rs new file mode 100644 index 0000000..433611f --- /dev/null +++ b/crates/xserv-model/src/bin/sink-attn-check.rs @@ -0,0 +1,57 @@ +//! Unit-test decode_attention_sink (GPU) against a CPU reference on small random +//! inputs. Isolates the sink/window attention kernel from the rest of gpt-oss. +use half::bf16; +use xserv_kernels::decode_attention_sink; +use xserv_tensor::{Device, Tensor}; + +fn main() { + let (n_heads, n_kv, head_dim, kv_len) = (8usize, 2usize, 4usize, 3usize); + let scale = (head_dim as f32).powf(-0.5); + let n_rep = n_heads / n_kv; + + // deterministic pseudo-random fill + let mut seed = 12345u64; + let mut rnd = || { seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1); ((seed >> 33) as f32 / (1u64 << 31) as f32) - 1.0 }; + + let q: Vec = (0..n_heads * head_dim).map(|_| rnd()).collect(); // [H,1,D] + let k: Vec = (0..n_kv * kv_len * head_dim).map(|_| rnd()).collect(); // [Hkv,kv,D] + let v: Vec = (0..n_kv * kv_len * head_dim).map(|_| rnd()).collect(); + let sinks: Vec = (0..n_heads).map(|_| rnd()).collect(); + + let qb: Vec = q.iter().map(|&x| bf16::from_f32(x)).collect(); + let kb: Vec = k.iter().map(|&x| bf16::from_f32(x)).collect(); + let vb: Vec = v.iter().map(|&x| bf16::from_f32(x)).collect(); + + let qt = Tensor::from_slice(&qb, &[1, n_heads, 1, head_dim]).to_device(Device::Cuda(0)); + let kt = Tensor::from_slice(&kb, &[1, n_kv, kv_len, head_dim]).to_device(Device::Cuda(0)); + let vt = Tensor::from_slice(&vb, &[1, n_kv, kv_len, head_dim]).to_device(Device::Cuda(0)); + let st = Tensor::from_slice(&sinks, &[n_heads]).to_device(Device::Cuda(0)); // f32 + + let out = decode_attention_sink(&qt, &kt, &vt, &st, scale, 0); // [1,H,1,D] + let outh = out.to_device(Device::Cpu); + let og = outh.as_slice::(); + + // CPU reference: for each q head, softmax over [s_0..s_{kv-1}, sink], drop sink, weight V. + let mut max_diff = 0f32; + for h in 0..n_heads { + let kv = h / n_rep; + let mut s = vec![0f32; kv_len]; + let mut m = sinks[h]; + for j in 0..kv_len { + let mut dot = 0f32; + for d in 0..head_dim { dot += q[h * head_dim + d] * k[(kv * kv_len + j) * head_dim + d]; } + s[j] = dot * scale; + if s[j] > m { m = s[j]; } + } + let mut denom = (sinks[h] - m).exp(); + for j in 0..kv_len { denom += (s[j] - m).exp(); } + for d in 0..head_dim { + let mut acc = 0f32; + for j in 0..kv_len { acc += (s[j] - m).exp() / denom * v[(kv * kv_len + j) * head_dim + d]; } + let got = og[h * head_dim + d].to_f32(); + max_diff = max_diff.max((got - acc).abs()); + } + } + println!("decode_attention_sink vs CPU ref: max_abs_diff = {max_diff:.5}"); + println!("{}", if max_diff < 0.05 { "SINK_KERNEL_OK" } else { "SINK_KERNEL_MISMATCH" }); +}