style: format Rust workspace

crates/xserv-kernels/tests/attention_test.rs

@@ -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
+        );
     }
 }