style: format Rust workspace

crates/xserv-kernels/src/argmax.rs

@@ -2,8 +2,13 @@ use std::ffi::c_void;
 use xserv_tensor::{DType, Device, Tensor};
 
 unsafe extern "C" {
-    fn launch_argmax_bf16(logits: *const c_void, out_idx: *mut c_void,
-                          rows: i32, cols: i32, stream: *mut c_void);
+    fn launch_argmax_bf16(
+        logits: *const c_void,
+        out_idx: *mut c_void,
+        rows: i32,
+        cols: i32,
+        stream: *mut c_void,
+    );
 }
 
 /// GPU argmax over the last dim of a [rows, cols] BF16 tensor.
@@ -19,7 +24,10 @@ pub fn argmax_bf16_to_host(logits: &Tensor) -> Vec<u32> {
     assert_eq!(logits.ndim(), 2, "argmax expects a 2D [rows, cols] tensor");
     assert_eq!(logits.dtype(), DType::BF16, "argmax kernel is BF16-only");
     assert!(logits.is_contiguous(), "argmax requires contiguous input");
-    assert!(matches!(logits.device(), Device::Cuda(_)), "argmax requires GPU input");
+    assert!(
+        matches!(logits.device(), Device::Cuda(_)),
+        "argmax requires GPU input"
+    );
 
     let rows = logits.shape()[0];
     let cols = logits.shape()[1];
@@ -35,7 +43,8 @@ pub fn argmax_bf16_to_host(logits: &Tensor) -> Vec<u32> {
         launch_argmax_bf16(
             logits.data_ptr() as *const c_void,
             out.as_mut_ptr() as *mut c_void,
-            rows as i32, cols as i32,
+            rows as i32,
+            cols as i32,
             xserv_cuda::current_stream_raw(),
         );
     }
@@ -44,9 +53,8 @@ pub fn argmax_bf16_to_host(logits: &Tensor) -> Vec<u32> {
     out.copy_to_host(&mut host_bytes).expect("argmax D2H");
     drop(out); // returned to pool
 
-    let host_i32: &[i32] = unsafe {
-        std::slice::from_raw_parts(host_bytes.as_ptr() as *const i32, rows)
-    };
+    let host_i32: &[i32] =
+        unsafe { std::slice::from_raw_parts(host_bytes.as_ptr() as *const i32, rows) };
     host_i32.iter().map(|&v| v as u32).collect()
 }
 
@@ -62,4 +70,3 @@ pub fn argmax_bf16_single(logits: &Tensor) -> u32 {
     };
     argmax_bf16_to_host(&view)[0]
 }
-