style: cargo fmt --all — align the tree with the CI fmt gate

hanzo-3d/src/geometry.rs

@@ -14,7 +14,11 @@ pub struct Vec3 {
 }
 
 impl Vec3 {
-    pub const ZERO: Vec3 = Vec3 { x: 0.0, y: 0.0, z: 0.0 };
+    pub const ZERO: Vec3 = Vec3 {
+        x: 0.0,
+        y: 0.0,
+        z: 0.0,
+    };
 
     pub const fn new(x: f32, y: f32, z: f32) -> Self {
         Vec3 { x, y, z }
@@ -110,7 +114,12 @@ pub struct Quat {
 }
 
 impl Quat {
-    pub const IDENTITY: Quat = Quat { w: 1.0, x: 0.0, y: 0.0, z: 0.0 };
+    pub const IDENTITY: Quat = Quat {
+        w: 1.0,
+        x: 0.0,
+        y: 0.0,
+        z: 0.0,
+    };
 
     pub const fn new(w: f32, x: f32, y: f32, z: f32) -> Self {
         Quat { w, x, y, z }
@@ -159,7 +168,14 @@ pub struct Camera {
 impl Camera {
     /// A camera at the origin looking down `+z` with the given intrinsics.
     pub fn look_forward(fx: f32, fy: f32, cx: f32, cy: f32) -> Self {
-        Camera { fx, fy, cx, cy, rot: Mat3::IDENTITY, t: Vec3::ZERO }
+        Camera {
+            fx,
+            fy,
+            cx,
+            cy,
+            rot: Mat3::IDENTITY,
+            t: Vec3::ZERO,
+        }
     }
 
     /// World point → `(u, v, depth)`. `depth` is the camera-space `z`.

hanzo-3d/src/io.rs

@@ -60,24 +60,31 @@ pub fn obj_to_mesh(src: &str) -> Result<Mesh> {
         let mut it = line.split_whitespace();
         match it.next() {
             Some("v") => {
-                let v = read_vec3(&mut it).map_err(|e| Error(format!("line {}: v: {}", lineno + 1, e.0)))?;
+                let v = read_vec3(&mut it)
+                    .map_err(|e| Error(format!("line {}: v: {}", lineno + 1, e.0)))?;
                 m.vertices.push(v);
             }
             Some("vn") => {
-                let n = read_vec3(&mut it).map_err(|e| Error(format!("line {}: vn: {}", lineno + 1, e.0)))?;
+                let n = read_vec3(&mut it)
+                    .map_err(|e| Error(format!("line {}: vn: {}", lineno + 1, e.0)))?;
                 m.normals.push(n);
             }
             Some("f") => {
                 let mut idx = [0u32; 3];
                 for (k, slot) in idx.iter_mut().enumerate() {
-                    let tok = it.next().ok_or_else(|| Error(format!("line {}: f needs 3 verts", lineno + 1)))?;
+                    let tok = it
+                        .next()
+                        .ok_or_else(|| Error(format!("line {}: f needs 3 verts", lineno + 1)))?;
                     // "a", "a/b", "a//c" — the vertex index is the first field, 1-based.
                     let first = tok.split('/').next().unwrap_or(tok);
-                    let one_based: u32 = first
-                        .parse()
-                        .map_err(|_| Error(format!("line {}: bad face index {:?}", lineno + 1, tok)))?;
+                    let one_based: u32 = first.parse().map_err(|_| {
+                        Error(format!("line {}: bad face index {:?}", lineno + 1, tok))
+                    })?;
                     if one_based == 0 {
-                        return err(format!("line {}: face vertex 0 (OBJ is 1-indexed)", lineno + 1));
+                        return err(format!(
+                            "line {}: face vertex 0 (OBJ is 1-indexed)",
+                            lineno + 1
+                        ));
                     }
                     let _ = k;
                     *slot = one_based - 1;
@@ -122,14 +129,21 @@ pub fn ply_to_mesh(src: &str) -> Result<Mesh> {
     let mut lines = body.lines().filter(|l| !l.trim().is_empty());
     let mut m = Mesh::default();
     for _ in 0..n_vert {
-        let line = lines.next().ok_or_else(|| Error("ply: truncated vertices".into()))?;
+        let line = lines
+            .next()
+            .ok_or_else(|| Error("ply: truncated vertices".into()))?;
         let mut it = line.split_whitespace();
         m.vertices.push(read_vec3(&mut it)?);
     }
     for _ in 0..n_face {
-        let line = lines.next().ok_or_else(|| Error("ply: truncated faces".into()))?;
+        let line = lines
+            .next()
+            .ok_or_else(|| Error("ply: truncated faces".into()))?;
         let mut it = line.split_whitespace();
-        let k: usize = it.next().and_then(|t| t.parse().ok()).ok_or_else(|| Error("ply: bad face count".into()))?;
+        let k: usize = it
+            .next()
+            .and_then(|t| t.parse().ok())
+            .ok_or_else(|| Error("ply: bad face count".into()))?;
         if k != 3 {
             return err(format!("ply: only triangles supported, got a {k}-gon"));
         }
@@ -190,7 +204,9 @@ pub fn ply_to_splat(src: &str) -> Result<GaussianSplat> {
     let mut out = GaussianSplat::default();
     let mut lines = body.lines().filter(|l| !l.trim().is_empty());
     for _ in 0..n {
-        let line = lines.next().ok_or_else(|| Error("ply: truncated splats".into()))?;
+        let line = lines
+            .next()
+            .ok_or_else(|| Error("ply: truncated splats".into()))?;
         let mut it = line.split_whitespace();
         let mut f = [0f32; 14];
         for (i, slot) in f.iter_mut().enumerate() {
@@ -239,7 +255,9 @@ fn split_ply_header(src: &str) -> Result<(PlyHeader<'_>, &str)> {
         return err("not a PLY (missing 'ply' magic)");
     }
     let marker = "end_header";
-    let pos = src.find(marker).ok_or_else(|| Error("ply: no end_header".into()))?;
+    let pos = src
+        .find(marker)
+        .ok_or_else(|| Error("ply: no end_header".into()))?;
     let header = &src[..pos];
     let body = &src[pos + marker.len()..];
     Ok((PlyHeader { text: header }, body))

hanzo-3d/src/lib.rs

@@ -42,12 +42,18 @@ pub fn unit_cube() -> Mesh {
         v(0.0, 1.0, 1.0),
     ];
     let faces = vec![
-        [0, 1, 2], [0, 2, 3], // back
-        [4, 6, 5], [4, 7, 6], // front
-        [0, 4, 5], [0, 5, 1], // bottom
-        [3, 2, 6], [3, 6, 7], // top
-        [0, 3, 7], [0, 7, 4], // left
-        [1, 5, 6], [1, 6, 2], // right
+        [0, 1, 2],
+        [0, 2, 3], // back
+        [4, 6, 5],
+        [4, 7, 6], // front
+        [0, 4, 5],
+        [0, 5, 1], // bottom
+        [3, 2, 6],
+        [3, 6, 7], // top
+        [0, 3, 7],
+        [0, 7, 4], // left
+        [1, 5, 6],
+        [1, 6, 2], // right
     ];
     Mesh::new(vertices, faces)
 }

hanzo-3d/src/representations.rs

@@ -34,13 +34,21 @@ pub struct Material {
 
 impl Default for Material {
     fn default() -> Self {
-        Material { base_color: [0.8, 0.8, 0.8, 1.0], roughness: 1.0, metallic: 0.0 }
+        Material {
+            base_color: [0.8, 0.8, 0.8, 1.0],
+            roughness: 1.0,
+            metallic: 0.0,
+        }
     }
 }
 
 impl Mesh {
     pub fn new(vertices: Vec<Vec3>, faces: Vec<[u32; 3]>) -> Self {
-        Mesh { vertices, faces, ..Default::default() }
+        Mesh {
+            vertices,
+            faces,
+            ..Default::default()
+        }
     }
 
     pub fn vertex_count(&self) -> usize {
@@ -95,7 +103,10 @@ pub struct Voxel {
 
 impl Voxel {
     pub fn new(resolution: u32) -> Self {
-        Voxel { resolution, occupied: Vec::new() }
+        Voxel {
+            resolution,
+            occupied: Vec::new(),
+        }
     }
 }
 

hanzo-flash-attn-build/src/lib.rs

@@ -108,7 +108,10 @@ pub fn cutlass_include_arg(cutlass_dir: &Path) -> String {
 /// A directory is a usable CUTLASS checkout iff its primary public header is
 /// present. This is cheap and robust against partial downloads.
 fn is_valid_cutlass(dir: &Path) -> bool {
-    dir.join("include").join("cutlass").join("cutlass.h").is_file()
+    dir.join("include")
+        .join("cutlass")
+        .join("cutlass.h")
+        .is_file()
 }
 
 /// Download the pinned CUTLASS source archive from GitHub and extract it into
@@ -136,9 +139,12 @@ fn download_and_extract(out_dir: &Path, commit: &str) -> Result<()> {
 
     // Extract underneath out_dir; the archive's own top-level directory is
     // `cutlass-<commit>/`, matching the path `fetch_cutlass` expects.
-    archive
-        .unpack(out_dir)
-        .with_context(|| format!("failed to extract CUTLASS archive into {}", out_dir.display()))?;
+    archive.unpack(out_dir).with_context(|| {
+        format!(
+            "failed to extract CUTLASS archive into {}",
+            out_dir.display()
+        )
+    })?;
 
     Ok(())
 }
@@ -148,7 +154,8 @@ fn download_and_extract(out_dir: &Path, commit: &str) -> Result<()> {
 #[allow(dead_code)]
 fn read_all(mut r: impl Read) -> Result<Vec<u8>> {
     let mut buf = Vec::new();
-    r.read_to_end(&mut buf).context("failed to read response body")?;
+    r.read_to_end(&mut buf)
+        .context("failed to read response body")?;
     Ok(buf)
 }
 
@@ -185,7 +192,11 @@ mod tests {
         fs::create_dir_all(tmp.join("include")).unwrap();
         assert!(!is_valid_cutlass(&tmp));
         fs::create_dir_all(tmp.join("include").join("cutlass")).unwrap();
-        fs::write(tmp.join("include").join("cutlass").join("cutlass.h"), b"// hdr").unwrap();
+        fs::write(
+            tmp.join("include").join("cutlass").join("cutlass.h"),
+            b"// hdr",
+        )
+        .unwrap();
         assert!(is_valid_cutlass(&tmp));
         let _ = fs::remove_dir_all(&tmp);
     }

hanzo-flash-attn/src/lib.rs

@@ -257,7 +257,8 @@ impl hanzo_ml::CustomOp3 for FlashAttn {
                     let qcl = Layout::contiguous(q_l.shape());
                     let kcl = Layout::contiguous(k_l.shape());
                     let vcl = Layout::contiguous(v_l.shape());
-                    let (out, shape) = self.cuda_fwd_t::<f16>(&qf, &qcl, &kf, &kcl, &vf, &vcl, false)?;
+                    let (out, shape) =
+                        self.cuda_fwd_t::<f16>(&qf, &qcl, &kf, &kcl, &vf, &vcl, false)?;
                     let out_l = Layout::contiguous(&shape);
                     let out = out.to_dtype(&out_l, hanzo_ml::DType::BF16)?;
                     Ok((out, shape))
@@ -744,7 +745,8 @@ impl hanzo_ml::CustomOp3 for FlashAttnVarLen {
                     let qcl = Layout::contiguous(q_l.shape());
                     let kcl = Layout::contiguous(k_l.shape());
                     let vcl = Layout::contiguous(v_l.shape());
-                    let (out, shape) = self.cuda_fwd_t::<f16>(&qf, &qcl, &kf, &kcl, &vf, &vcl, false)?;
+                    let (out, shape) =
+                        self.cuda_fwd_t::<f16>(&qf, &qcl, &kf, &kcl, &vf, &vcl, false)?;
                     let out_l = Layout::contiguous(&shape);
                     let out = out.to_dtype(&out_l, hanzo_ml::DType::BF16)?;
                     Ok((out, shape))

hanzo-kernels/src/ffi.rs

@@ -653,40 +653,117 @@ extern "C" {
     // i-quant codebook dense MMQ prefill launchers (int8 WMMA via load_tiles_iq*). Same signature as
     // the K-quant launchers above; defined by the weak launch_mmq_gguf_<iq*> in each mmq_instance_iq*.cu.
     pub fn launch_mmq_gguf_iq2_xxs(
-        tmp_fixup: *mut c_void, x: *const c_void, y: *const c_void, dst: *mut c_void,
-        ncols_x: i64, nrows_x: i64, ncols_y: i64, stride_row_x: i64, stride_col_dst: i64,
-        cc: i32, nsm: i32, smpbo: i64, warp_size: i32, stream: *mut c_void,
+        tmp_fixup: *mut c_void,
+        x: *const c_void,
+        y: *const c_void,
+        dst: *mut c_void,
+        ncols_x: i64,
+        nrows_x: i64,
+        ncols_y: i64,
+        stride_row_x: i64,
+        stride_col_dst: i64,
+        cc: i32,
+        nsm: i32,
+        smpbo: i64,
+        warp_size: i32,
+        stream: *mut c_void,
     );
     pub fn launch_mmq_gguf_iq2_xs(
-        tmp_fixup: *mut c_void, x: *const c_void, y: *const c_void, dst: *mut c_void,
-        ncols_x: i64, nrows_x: i64, ncols_y: i64, stride_row_x: i64, stride_col_dst: i64,
-        cc: i32, nsm: i32, smpbo: i64, warp_size: i32, stream: *mut c_void,
+        tmp_fixup: *mut c_void,
+        x: *const c_void,
+        y: *const c_void,
+        dst: *mut c_void,
+        ncols_x: i64,
+        nrows_x: i64,
+        ncols_y: i64,
+        stride_row_x: i64,
+        stride_col_dst: i64,
+        cc: i32,
+        nsm: i32,
+        smpbo: i64,
+        warp_size: i32,
+        stream: *mut c_void,
     );
     pub fn launch_mmq_gguf_iq2_s(
-        tmp_fixup: *mut c_void, x: *const c_void, y: *const c_void, dst: *mut c_void,
-        ncols_x: i64, nrows_x: i64, ncols_y: i64, stride_row_x: i64, stride_col_dst: i64,
-        cc: i32, nsm: i32, smpbo: i64, warp_size: i32, stream: *mut c_void,
+        tmp_fixup: *mut c_void,
+        x: *const c_void,
+        y: *const c_void,
+        dst: *mut c_void,
+        ncols_x: i64,
+        nrows_x: i64,
+        ncols_y: i64,
+        stride_row_x: i64,
+        stride_col_dst: i64,
+        cc: i32,
+        nsm: i32,
+        smpbo: i64,
+        warp_size: i32,
+        stream: *mut c_void,
     );
     pub fn launch_mmq_gguf_iq3_xxs(
-        tmp_fixup: *mut c_void, x: *const c_void, y: *const c_void, dst: *mut c_void,
-        ncols_x: i64, nrows_x: i64, ncols_y: i64, stride_row_x: i64, stride_col_dst: i64,
-        cc: i32, nsm: i32, smpbo: i64, warp_size: i32, stream: *mut c_void,
+        tmp_fixup: *mut c_void,
+        x: *const c_void,
+        y: *const c_void,
+        dst: *mut c_void,
+        ncols_x: i64,
+        nrows_x: i64,
+        ncols_y: i64,
+        stride_row_x: i64,
+        stride_col_dst: i64,
+        cc: i32,
+        nsm: i32,
+        smpbo: i64,
+        warp_size: i32,
+        stream: *mut c_void,
     );
     pub fn launch_mmq_gguf_iq3_s(
-        tmp_fixup: *mut c_void, x: *const c_void, y: *const c_void, dst: *mut c_void,
-        ncols_x: i64, nrows_x: i64, ncols_y: i64, stride_row_x: i64, stride_col_dst: i64,
-        cc: i32, nsm: i32, smpbo: i64, warp_size: i32, stream: *mut c_void,
+        tmp_fixup: *mut c_void,
+        x: *const c_void,
+        y: *const c_void,
+        dst: *mut c_void,
+        ncols_x: i64,
+        nrows_x: i64,
+        ncols_y: i64,
+        stride_row_x: i64,
+        stride_col_dst: i64,
+        cc: i32,
+        nsm: i32,
+        smpbo: i64,
+        warp_size: i32,
+        stream: *mut c_void,
     );
     pub fn launch_mmq_gguf_iq4_xs(
-        tmp_fixup: *mut c_void, x: *const c_void, y: *const c_void, dst: *mut c_void,
-        ncols_x: i64, nrows_x: i64, ncols_y: i64, stride_row_x: i64, stride_col_dst: i64,
-        cc: i32, nsm: i32, smpbo: i64, warp_size: i32, stream: *mut c_void,
+        tmp_fixup: *mut c_void,
+        x: *const c_void,
+        y: *const c_void,
+        dst: *mut c_void,
+        ncols_x: i64,
+        nrows_x: i64,
+        ncols_y: i64,
+        stride_row_x: i64,
+        stride_col_dst: i64,
+        cc: i32,
+        nsm: i32,
+        smpbo: i64,
+        warp_size: i32,
+        stream: *mut c_void,
     );
     // IQ1_S (1-bit grid + delta): DS4 ds-layout (d + delta*sum) -- distinct GPU-packed iq1s_grid_gpu.
     pub fn launch_mmq_gguf_iq1_s(
-        tmp_fixup: *mut c_void, x: *const c_void, y: *const c_void, dst: *mut c_void,
-        ncols_x: i64, nrows_x: i64, ncols_y: i64, stride_row_x: i64, stride_col_dst: i64,
-        cc: i32, nsm: i32, smpbo: i64, warp_size: i32, stream: *mut c_void,
+        tmp_fixup: *mut c_void,
+        x: *const c_void,
+        y: *const c_void,
+        dst: *mut c_void,
+        ncols_x: i64,
+        nrows_x: i64,
+        ncols_y: i64,
+        stride_row_x: i64,
+        stride_col_dst: i64,
+        cc: i32,
+        nsm: i32,
+        smpbo: i64,
+        warp_size: i32,
+        stream: *mut c_void,
     );
 
     // Expert-grouped MoE prefill launchers (llama mul_mat_id). Same int8 MMQ core as the dense
@@ -695,123 +772,344 @@ extern "C" {
     // `stride_channel_x` steps one expert's weight bank. Output dst is [ncols_dst, nrows_x]. Defined by
     // DEFINE_MMQ_GGUF_MOE_LAUNCHER in each mmq_instance_*.cu.
     pub fn launch_mmq_gguf_moe_q4_0(
-        tmp_fixup: *mut c_void, x: *const c_void, y: *const c_void,
-        ids_dst: *const c_void, expert_bounds: *const c_void, dst: *mut c_void,
-        ncols_x: i64, nrows_x: i64, ncols_dst: i64, ncols_max: i64,
-        stride_row_x: i64, stride_channel_x: i64, n_experts: i64,
-        cc: i32, nsm: i32, smpbo: i64, warp_size: i32, stream: *mut c_void,
+        tmp_fixup: *mut c_void,
+        x: *const c_void,
+        y: *const c_void,
+        ids_dst: *const c_void,
+        expert_bounds: *const c_void,
+        dst: *mut c_void,
+        ncols_x: i64,
+        nrows_x: i64,
+        ncols_dst: i64,
+        ncols_max: i64,
+        stride_row_x: i64,
+        stride_channel_x: i64,
+        n_experts: i64,
+        cc: i32,
+        nsm: i32,
+        smpbo: i64,
+        warp_size: i32,
+        stream: *mut c_void,
     );
     pub fn launch_mmq_gguf_moe_q4_1(
-        tmp_fixup: *mut c_void, x: *const c_void, y: *const c_void,
-        ids_dst: *const c_void, expert_bounds: *const c_void, dst: *mut c_void,
-        ncols_x: i64, nrows_x: i64, ncols_dst: i64, ncols_max: i64,
-        stride_row_x: i64, stride_channel_x: i64, n_experts: i64,
-        cc: i32, nsm: i32, smpbo: i64, warp_size: i32, stream: *mut c_void,
+        tmp_fixup: *mut c_void,
+        x: *const c_void,
+        y: *const c_void,
+        ids_dst: *const c_void,
+        expert_bounds: *const c_void,
+        dst: *mut c_void,
+        ncols_x: i64,
+        nrows_x: i64,
+        ncols_dst: i64,
+        ncols_max: i64,
+        stride_row_x: i64,
+        stride_channel_x: i64,
+        n_experts: i64,
+        cc: i32,
+        nsm: i32,
+        smpbo: i64,
+        warp_size: i32,
+        stream: *mut c_void,
     );
     pub fn launch_mmq_gguf_moe_q5_0(
-        tmp_fixup: *mut c_void, x: *const c_void, y: *const c_void,
-        ids_dst: *const c_void, expert_bounds: *const c_void, dst: *mut c_void,
-        ncols_x: i64, nrows_x: i64, ncols_dst: i64, ncols_max: i64,
-        stride_row_x: i64, stride_channel_x: i64, n_experts: i64,
-        cc: i32, nsm: i32, smpbo: i64, warp_size: i32, stream: *mut c_void,
+        tmp_fixup: *mut c_void,
+        x: *const c_void,
+        y: *const c_void,
+        ids_dst: *const c_void,
+        expert_bounds: *const c_void,
+        dst: *mut c_void,
+        ncols_x: i64,
+        nrows_x: i64,
+        ncols_dst: i64,
+        ncols_max: i64,
+        stride_row_x: i64,
+        stride_channel_x: i64,
+        n_experts: i64,
+        cc: i32,
+        nsm: i32,
+        smpbo: i64,
+        warp_size: i32,
+        stream: *mut c_void,
     );
     pub fn launch_mmq_gguf_moe_q5_1(
-        tmp_fixup: *mut c_void, x: *const c_void, y: *const c_void,
-        ids_dst: *const c_void, expert_bounds: *const c_void, dst: *mut c_void,
-        ncols_x: i64, nrows_x: i64, ncols_dst: i64, ncols_max: i64,
-        stride_row_x: i64, stride_channel_x: i64, n_experts: i64,
-        cc: i32, nsm: i32, smpbo: i64, warp_size: i32, stream: *mut c_void,
+        tmp_fixup: *mut c_void,
+        x: *const c_void,
+        y: *const c_void,
+        ids_dst: *const c_void,
+        expert_bounds: *const c_void,
+        dst: *mut c_void,
+        ncols_x: i64,
+        nrows_x: i64,
+        ncols_dst: i64,
+        ncols_max: i64,
+        stride_row_x: i64,
+        stride_channel_x: i64,
+        n_experts: i64,
+        cc: i32,
+        nsm: i32,
+        smpbo: i64,
+        warp_size: i32,
+        stream: *mut c_void,
     );
     pub fn launch_mmq_gguf_moe_q8_0(
-        tmp_fixup: *mut c_void, x: *const c_void, y: *const c_void,
-        ids_dst: *const c_void, expert_bounds: *const c_void, dst: *mut c_void,
-        ncols_x: i64, nrows_x: i64, ncols_dst: i64, ncols_max: i64,
-        stride_row_x: i64, stride_channel_x: i64, n_experts: i64,
-        cc: i32, nsm: i32, smpbo: i64, warp_size: i32, stream: *mut c_void,
+        tmp_fixup: *mut c_void,
+        x: *const c_void,
+        y: *const c_void,
+        ids_dst: *const c_void,
+        expert_bounds: *const c_void,
+        dst: *mut c_void,
+        ncols_x: i64,
+        nrows_x: i64,
+        ncols_dst: i64,
+        ncols_max: i64,
+        stride_row_x: i64,
+        stride_channel_x: i64,
+        n_experts: i64,
+        cc: i32,
+        nsm: i32,
+        smpbo: i64,
+        warp_size: i32,
+        stream: *mut c_void,
     );
     pub fn launch_mmq_gguf_moe_q2_k(
-        tmp_fixup: *mut c_void, x: *const c_void, y: *const c_void,
-        ids_dst: *const c_void, expert_bounds: *const c_void, dst: *mut c_void,
-        ncols_x: i64, nrows_x: i64, ncols_dst: i64, ncols_max: i64,
-        stride_row_x: i64, stride_channel_x: i64, n_experts: i64,
-        cc: i32, nsm: i32, smpbo: i64, warp_size: i32, stream: *mut c_void,
+        tmp_fixup: *mut c_void,
+        x: *const c_void,
+        y: *const c_void,
+        ids_dst: *const c_void,
+        expert_bounds: *const c_void,
+        dst: *mut c_void,
+        ncols_x: i64,
+        nrows_x: i64,
+        ncols_dst: i64,
+        ncols_max: i64,
+        stride_row_x: i64,
+        stride_channel_x: i64,
+        n_experts: i64,
+        cc: i32,
+        nsm: i32,
+        smpbo: i64,
+        warp_size: i32,
+        stream: *mut c_void,
     );
     pub fn launch_mmq_gguf_moe_q3_k(
-        tmp_fixup: *mut c_void, x: *const c_void, y: *const c_void,
-        ids_dst: *const c_void, expert_bounds: *const c_void, dst: *mut c_void,
-        ncols_x: i64, nrows_x: i64, ncols_dst: i64, ncols_max: i64,
-        stride_row_x: i64, stride_channel_x: i64, n_experts: i64,
-        cc: i32, nsm: i32, smpbo: i64, warp_size: i32, stream: *mut c_void,
+        tmp_fixup: *mut c_void,
+        x: *const c_void,
+        y: *const c_void,
+        ids_dst: *const c_void,
+        expert_bounds: *const c_void,
+        dst: *mut c_void,
+        ncols_x: i64,
+        nrows_x: i64,
+        ncols_dst: i64,
+        ncols_max: i64,
+        stride_row_x: i64,
+        stride_channel_x: i64,
+        n_experts: i64,
+        cc: i32,
+        nsm: i32,
+        smpbo: i64,
+        warp_size: i32,
+        stream: *mut c_void,
     );
     pub fn launch_mmq_gguf_moe_q4_k(
-        tmp_fixup: *mut c_void, x: *const c_void, y: *const c_void,
-        ids_dst: *const c_void, expert_bounds: *const c_void, dst: *mut c_void,
-        ncols_x: i64, nrows_x: i64, ncols_dst: i64, ncols_max: i64,
-        stride_row_x: i64, stride_channel_x: i64, n_experts: i64,
-        cc: i32, nsm: i32, smpbo: i64, warp_size: i32, stream: *mut c_void,
+        tmp_fixup: *mut c_void,
+        x: *const c_void,
+        y: *const c_void,
+        ids_dst: *const c_void,
+        expert_bounds: *const c_void,
+        dst: *mut c_void,
+        ncols_x: i64,
+        nrows_x: i64,
+        ncols_dst: i64,
+        ncols_max: i64,
+        stride_row_x: i64,
+        stride_channel_x: i64,
+        n_experts: i64,
+        cc: i32,
+        nsm: i32,
+        smpbo: i64,
+        warp_size: i32,
+        stream: *mut c_void,
     );
     pub fn launch_mmq_gguf_moe_q5_k(
-        tmp_fixup: *mut c_void, x: *const c_void, y: *const c_void,
-        ids_dst: *const c_void, expert_bounds: *const c_void, dst: *mut c_void,
-        ncols_x: i64, nrows_x: i64, ncols_dst: i64, ncols_max: i64,
-        stride_row_x: i64, stride_channel_x: i64, n_experts: i64,
-        cc: i32, nsm: i32, smpbo: i64, warp_size: i32, stream: *mut c_void,
+        tmp_fixup: *mut c_void,
+        x: *const c_void,
+        y: *const c_void,
+        ids_dst: *const c_void,
+        expert_bounds: *const c_void,
+        dst: *mut c_void,
+        ncols_x: i64,
+        nrows_x: i64,
+        ncols_dst: i64,
+        ncols_max: i64,
+        stride_row_x: i64,
+        stride_channel_x: i64,
+        n_experts: i64,
+        cc: i32,
+        nsm: i32,
+        smpbo: i64,
+        warp_size: i32,
+        stream: *mut c_void,
     );
     pub fn launch_mmq_gguf_moe_q6_k(
-        tmp_fixup: *mut c_void, x: *const c_void, y: *const c_void,
-        ids_dst: *const c_void, expert_bounds: *const c_void, dst: *mut c_void,
-        ncols_x: i64, nrows_x: i64, ncols_dst: i64, ncols_max: i64,
-        stride_row_x: i64, stride_channel_x: i64, n_experts: i64,
-        cc: i32, nsm: i32, smpbo: i64, warp_size: i32, stream: *mut c_void,
+        tmp_fixup: *mut c_void,
+        x: *const c_void,
+        y: *const c_void,
+        ids_dst: *const c_void,
+        expert_bounds: *const c_void,
+        dst: *mut c_void,
+        ncols_x: i64,
+        nrows_x: i64,
+        ncols_dst: i64,
+        ncols_max: i64,
+        stride_row_x: i64,
+        stride_channel_x: i64,
+        n_experts: i64,
+        cc: i32,
+        nsm: i32,
+        smpbo: i64,
+        warp_size: i32,
+        stream: *mut c_void,
     );
     // i-quant codebook MoE MMQ prefill launchers (DEFINE_MMQ_GGUF_MOE_LAUNCHER in mmq_instance_iq*.cu).
     pub fn launch_mmq_gguf_moe_iq2_xxs(
-        tmp_fixup: *mut c_void, x: *const c_void, y: *const c_void,
-        ids_dst: *const c_void, expert_bounds: *const c_void, dst: *mut c_void,
-        ncols_x: i64, nrows_x: i64, ncols_dst: i64, ncols_max: i64,
-        stride_row_x: i64, stride_channel_x: i64, n_experts: i64,
-        cc: i32, nsm: i32, smpbo: i64, warp_size: i32, stream: *mut c_void,
+        tmp_fixup: *mut c_void,
+        x: *const c_void,
+        y: *const c_void,
+        ids_dst: *const c_void,
+        expert_bounds: *const c_void,
+        dst: *mut c_void,
+        ncols_x: i64,
+        nrows_x: i64,
+        ncols_dst: i64,
+        ncols_max: i64,
+        stride_row_x: i64,
+        stride_channel_x: i64,
+        n_experts: i64,
+        cc: i32,
+        nsm: i32,
+        smpbo: i64,
+        warp_size: i32,
+        stream: *mut c_void,
     );
     pub fn launch_mmq_gguf_moe_iq2_xs(
-        tmp_fixup: *mut c_void, x: *const c_void, y: *const c_void,
-        ids_dst: *const c_void, expert_bounds: *const c_void, dst: *mut c_void,
-        ncols_x: i64, nrows_x: i64, ncols_dst: i64, ncols_max: i64,
-        stride_row_x: i64, stride_channel_x: i64, n_experts: i64,
-        cc: i32, nsm: i32, smpbo: i64, warp_size: i32, stream: *mut c_void,
+        tmp_fixup: *mut c_void,
+        x: *const c_void,
+        y: *const c_void,
+        ids_dst: *const c_void,
+        expert_bounds: *const c_void,
+        dst: *mut c_void,
+        ncols_x: i64,
+        nrows_x: i64,
+        ncols_dst: i64,
+        ncols_max: i64,
+        stride_row_x: i64,
+        stride_channel_x: i64,
+        n_experts: i64,
+        cc: i32,
+        nsm: i32,
+        smpbo: i64,
+        warp_size: i32,
+        stream: *mut c_void,
     );
     pub fn launch_mmq_gguf_moe_iq2_s(
-        tmp_fixup: *mut c_void, x: *const c_void, y: *const c_void,
-        ids_dst: *const c_void, expert_bounds: *const c_void, dst: *mut c_void,
-        ncols_x: i64, nrows_x: i64, ncols_dst: i64, ncols_max: i64,
-        stride_row_x: i64, stride_channel_x: i64, n_experts: i64,
-        cc: i32, nsm: i32, smpbo: i64, warp_size: i32, stream: *mut c_void,
+        tmp_fixup: *mut c_void,
+        x: *const c_void,
+        y: *const c_void,
+        ids_dst: *const c_void,
+        expert_bounds: *const c_void,
+        dst: *mut c_void,
+        ncols_x: i64,
+        nrows_x: i64,
+        ncols_dst: i64,
+        ncols_max: i64,
+        stride_row_x: i64,
+        stride_channel_x: i64,
+        n_experts: i64,
+        cc: i32,
+        nsm: i32,
+        smpbo: i64,
+        warp_size: i32,
+        stream: *mut c_void,
     );
     pub fn launch_mmq_gguf_moe_iq3_xxs(
-        tmp_fixup: *mut c_void, x: *const c_void, y: *const c_void,
-        ids_dst: *const c_void, expert_bounds: *const c_void, dst: *mut c_void,
-        ncols_x: i64, nrows_x: i64, ncols_dst: i64, ncols_max: i64,
-        stride_row_x: i64, stride_channel_x: i64, n_experts: i64,
-        cc: i32, nsm: i32, smpbo: i64, warp_size: i32, stream: *mut c_void,
+        tmp_fixup: *mut c_void,
+        x: *const c_void,
+        y: *const c_void,
+        ids_dst: *const c_void,
+        expert_bounds: *const c_void,
+        dst: *mut c_void,
+        ncols_x: i64,
+        nrows_x: i64,
+        ncols_dst: i64,
+        ncols_max: i64,
+        stride_row_x: i64,
+        stride_channel_x: i64,
+        n_experts: i64,
+        cc: i32,
+        nsm: i32,
+        smpbo: i64,
+        warp_size: i32,
+        stream: *mut c_void,
     );
     pub fn launch_mmq_gguf_moe_iq3_s(
-        tmp_fixup: *mut c_void, x: *const c_void, y: *const c_void,
-        ids_dst: *const c_void, expert_bounds: *const c_void, dst: *mut c_void,
-        ncols_x: i64, nrows_x: i64, ncols_dst: i64, ncols_max: i64,
-        stride_row_x: i64, stride_channel_x: i64, n_experts: i64,
-        cc: i32, nsm: i32, smpbo: i64, warp_size: i32, stream: *mut c_void,
+        tmp_fixup: *mut c_void,
+        x: *const c_void,
+        y: *const c_void,
+        ids_dst: *const c_void,
+        expert_bounds: *const c_void,
+        dst: *mut c_void,
+        ncols_x: i64,
+        nrows_x: i64,
+        ncols_dst: i64,
+        ncols_max: i64,
+        stride_row_x: i64,
+        stride_channel_x: i64,
+        n_experts: i64,
+        cc: i32,
+        nsm: i32,
+        smpbo: i64,
+        warp_size: i32,
+        stream: *mut c_void,
     );
     pub fn launch_mmq_gguf_moe_iq4_xs(
-        tmp_fixup: *mut c_void, x: *const c_void, y: *const c_void,
-        ids_dst: *const c_void, expert_bounds: *const c_void, dst: *mut c_void,
-        ncols_x: i64, nrows_x: i64, ncols_dst: i64, ncols_max: i64,
-        stride_row_x: i64, stride_channel_x: i64, n_experts: i64,
-        cc: i32, nsm: i32, smpbo: i64, warp_size: i32, stream: *mut c_void,
+        tmp_fixup: *mut c_void,
+        x: *const c_void,
+        y: *const c_void,
+        ids_dst: *const c_void,
+        expert_bounds: *const c_void,
+        dst: *mut c_void,
+        ncols_x: i64,
+        nrows_x: i64,
+        ncols_dst: i64,
+        ncols_max: i64,
+        stride_row_x: i64,
+        stride_channel_x: i64,
+        n_experts: i64,
+        cc: i32,
+        nsm: i32,
+        smpbo: i64,
+        warp_size: i32,
+        stream: *mut c_void,
     );
     pub fn launch_mmq_gguf_moe_iq1_s(
-        tmp_fixup: *mut c_void, x: *const c_void, y: *const c_void,
-        ids_dst: *const c_void, expert_bounds: *const c_void, dst: *mut c_void,
-        ncols_x: i64, nrows_x: i64, ncols_dst: i64, ncols_max: i64,
-        stride_row_x: i64, stride_channel_x: i64, n_experts: i64,
-        cc: i32, nsm: i32, smpbo: i64, warp_size: i32, stream: *mut c_void,
+        tmp_fixup: *mut c_void,
+        x: *const c_void,
+        y: *const c_void,
+        ids_dst: *const c_void,
+        expert_bounds: *const c_void,
+        dst: *mut c_void,
+        ncols_x: i64,
+        nrows_x: i64,
+        ncols_dst: i64,
+        ncols_max: i64,
+        stride_row_x: i64,
+        stride_channel_x: i64,
+        n_experts: i64,
+        cc: i32,
+        nsm: i32,
+        smpbo: i64,
+        warp_size: i32,
+        stream: *mut c_void,
     );
 }

hanzo-metal-kernels/src/kernels/quantized.rs

@@ -217,8 +217,21 @@ pub fn call_quantized_matmul_mm_t(
     dst: &Buffer,
 ) -> Result<(), MetalKernelError> {
     call_quantized_matmul_mm_t_offset(
-        device, ep, kernels, dtype, src0_shape, src0_stride, src0, 0, src1_shape,
-        src1_stride, src1, src1_offset, dst_shape, dst_offset, dst,
+        device,
+        ep,
+        kernels,
+        dtype,
+        src0_shape,
+        src0_stride,
+        src0,
+        0,
+        src1_shape,
+        src1_stride,
+        src1,
+        src1_offset,
+        dst_shape,
+        dst_offset,
+        dst,
     )
 }
 
@@ -331,4 +344,3 @@ pub fn call_quantized_matmul_mm_t_offset(
 fn divide(m: usize, b: usize) -> usize {
     m.div_ceil(b)
 }
-

hanzo-metal-kernels/src/metal/commands.rs

@@ -128,9 +128,7 @@ impl Commands {
         residency_set: &ResidencySet,
     ) -> Result<Self, MetalKernelError> {
         let compute_per_buffer = match std::env::var("METAL_COMPUTE_PER_BUFFER") {
-            Ok(val) => val
-                .parse()
-                .unwrap_or(DEFAULT_METAL_COMPUTE_PER_BUFFER),
+            Ok(val) => val.parse().unwrap_or(DEFAULT_METAL_COMPUTE_PER_BUFFER),
             _ => DEFAULT_METAL_COMPUTE_PER_BUFFER,
         };
 

hanzo-ml-examples/examples/bert_single_file_binary/main.rs

@@ -6,9 +6,9 @@ extern crate accelerate_src;
 use hanzo_transformers::models::bert::{BertModel, Config as BertConfig, DTYPE};
 
 use anyhow::{Error as E, Result};
+use clap::Parser;
 use hanzo_ml::{Device, Tensor};
 use hanzo_nn::VarBuilder;
-use clap::Parser;
 use tokenizers::{PaddingParams, Tokenizer};
 
 #[derive(Parser, Debug)]

hanzo-ml-examples/examples/gguf-tokenizer.rs

@@ -5,9 +5,9 @@ use std::{
 };
 
 use anyhow::{Context, Result};
+use clap::Parser;
 use hanzo_ml::quantized::gguf_file;
 use hanzo_ml::quantized::tokenizer::TokenizerFromGguf;
-use clap::Parser;
 use hf_hub::{api::sync::Api, Repo, RepoType};
 use tokenizers::Tokenizer;
 

hanzo-ml-examples/examples/nomic-bert/main.rs

@@ -7,9 +7,9 @@ extern crate accelerate_src;
 use hanzo_transformers::models::nomic_bert::{self, Config, NomicBertModel};
 
 use anyhow::{bail, Error as E, Result};
+use clap::Parser;
 use hanzo_ml::{DType, Tensor};
 use hanzo_nn::VarBuilder;
-use clap::Parser;
 use hf_hub::{api::sync::Api, Repo, RepoType};
 use tokenizers::{PaddingParams, Tokenizer};
 

hanzo-ml-examples/examples/paddleocr-vl/main.rs

@@ -57,10 +57,10 @@ extern crate intel_mkl_src;
 extern crate accelerate_src;
 
 use anyhow::{Error as E, Result};
+use clap::{Parser, ValueEnum};
 use hanzo_ml::{DType, Device, Tensor};
 use hanzo_nn::VarBuilder;
 use hanzo_transformers::models::paddleocr_vl::{Config, PaddleOCRVLModel};
-use clap::{Parser, ValueEnum};
 use tokenizers::Tokenizer;
 
 const DEFAULT_MODEL_ID: &str = "PaddlePaddle/PaddleOCR-VL";

hanzo-ml-examples/examples/z_image/main.rs

@@ -33,6 +33,7 @@
 //! <https://huggingface.co/Tongyi-MAI/Z-Image-Turbo>
 
 use anyhow::{Error as E, Result};
+use clap::Parser;
 use hanzo_ml::{DType, IndexOp, Tensor};
 use hanzo_nn::VarBuilder;
 use hanzo_transformers::models::z_image::{
@@ -40,7 +41,6 @@ use hanzo_transformers::models::z_image::{
     FlowMatchEulerDiscreteScheduler, SchedulerConfig, TextEncoderConfig, VaeConfig,
     ZImageTextEncoder, ZImageTransformer2DModel,
 };
-use clap::Parser;
 use hf_hub::api::sync::Api;
 use tokenizers::Tokenizer;
 

hanzo-ml/benches/benchmarks/affine.rs

@@ -1,6 +1,6 @@
 use crate::benchmarks::{BenchDevice, BenchDeviceHandler};
-use hanzo_ml::{DType, Device, Tensor};
 use criterion::{criterion_group, Criterion, Throughput};
+use hanzo_ml::{DType, Device, Tensor};
 use std::hint::black_box;
 use std::time::Instant;
 

hanzo-ml/benches/benchmarks/binary.rs

@@ -1,6 +1,6 @@
 use crate::benchmarks::{BenchDevice, BenchDeviceHandler};
-use hanzo_ml::{DType, Device, Tensor};
 use criterion::{criterion_group, Criterion, Throughput};
+use hanzo_ml::{DType, Device, Tensor};
 use std::hint::black_box;
 use std::time::Instant;
 

hanzo-ml/benches/benchmarks/broadcast.rs

@@ -1,6 +1,6 @@
 use crate::benchmarks::{BenchDevice, BenchDeviceHandler};
-use hanzo_ml::{DType, Device, Tensor};
 use criterion::{criterion_group, Criterion, Throughput};
+use hanzo_ml::{DType, Device, Tensor};
 use std::hint::black_box;
 use std::time::Instant;
 

hanzo-ml/benches/benchmarks/cat.rs

@@ -1,6 +1,6 @@
 use crate::benchmarks::{BenchDevice, BenchDeviceHandler};
-use hanzo_ml::{DType, Device, Tensor};
 use criterion::{criterion_group, Criterion, Throughput};
+use hanzo_ml::{DType, Device, Tensor};
 use std::hint::black_box;
 use std::time::Instant;
 

hanzo-ml/benches/benchmarks/contiguous.rs

@@ -1,6 +1,6 @@
 use crate::benchmarks::{BenchDevice, BenchDeviceHandler};
-use hanzo_ml::{DType, Device, Tensor};
 use criterion::{criterion_group, Criterion, Throughput};
+use hanzo_ml::{DType, Device, Tensor};
 use std::hint::black_box;
 use std::time::Instant;
 

hanzo-ml/benches/benchmarks/conv_transpose2d.rs

@@ -1,6 +1,6 @@
 use crate::benchmarks::{BenchDevice, BenchDeviceHandler};
-use hanzo_ml::{DType, Device, Tensor};
 use criterion::{criterion_group, Criterion, Throughput};
+use hanzo_ml::{DType, Device, Tensor};
 use std::hint::black_box;
 use std::time::Instant;
 

hanzo-ml/benches/benchmarks/copy.rs

@@ -1,6 +1,6 @@
 use crate::benchmarks::{BenchDevice, BenchDeviceHandler};
-use hanzo_ml::{Device, Tensor, WithDType};
 use criterion::{criterion_group, Criterion, Throughput};
+use hanzo_ml::{Device, Tensor, WithDType};
 use std::hint::black_box;
 use std::time::Instant;
 

hanzo-ml/benches/benchmarks/matmul.rs

@@ -1,6 +1,6 @@
 use crate::benchmarks::{BenchDevice, BenchDeviceHandler};
-use hanzo_ml::{DType, Device, Tensor};
 use criterion::{criterion_group, Criterion, Throughput};
+use hanzo_ml::{DType, Device, Tensor};
 use std::hint::black_box;
 use std::time::Instant;
 

hanzo-ml/benches/benchmarks/qmatmul.rs

@@ -1,9 +1,9 @@
 use crate::benchmarks::{BenchDevice, BenchDeviceHandler};
+use criterion::{criterion_group, Criterion, Throughput};
 use hanzo_ml::{
     quantized::{self, GgmlDType, QMatMul},
     Device, Module, Tensor,
 };
-use criterion::{criterion_group, Criterion, Throughput};
 use std::hint::black_box;
 use std::time::Instant;
 

hanzo-ml/benches/benchmarks/random.rs

@@ -1,6 +1,6 @@
 use crate::benchmarks::{BenchDevice, BenchDeviceHandler};
-use hanzo_ml::{DType, Device, Tensor};
 use criterion::{criterion_group, Criterion, Throughput};
+use hanzo_ml::{DType, Device, Tensor};
 use std::hint::black_box;
 use std::time::Instant;
 

hanzo-ml/benches/benchmarks/reduce.rs

@@ -1,7 +1,7 @@
 use crate::benchmarks::{BenchDevice, BenchDeviceHandler};
-use hanzo_ml::{DType, Device, Tensor};
 use criterion::{criterion_group, Criterion, Throughput};
 use half::{bf16, f16};
+use hanzo_ml::{DType, Device, Tensor};
 use std::hint::black_box;
 use std::time::Instant;
 

hanzo-ml/benches/benchmarks/vec_dot.rs

@@ -1,6 +1,6 @@
-use hanzo_ml::cpu::kernels::VecOps;
 use criterion::{criterion_group, BatchSize, Criterion, Throughput};
 use half::{bf16, f16};
+use hanzo_ml::cpu::kernels::VecOps;
 
 fn bench_vec_dot_f32(c: &mut Criterion, k: usize) {
     let a: Vec<f32> = (0..k).map(|i| i as f32).collect();

hanzo-ml/benches/benchmarks/where_cond.rs

@@ -1,6 +1,6 @@
 use crate::benchmarks::{BenchDevice, BenchDeviceHandler};
-use hanzo_ml::{DType, Device, Tensor};
 use criterion::{criterion_group, Criterion, Throughput};
+use hanzo_ml::{DType, Device, Tensor};
 use std::hint::black_box;
 use std::time::Instant;
 

hanzo-ml/src/compute.rs

@@ -305,8 +305,14 @@ mod tests {
     fn policies_set_expected_dtypes() {
         let dev = Device::Cpu;
         let s = ComputeCtx::serving(dev.clone());
-        assert_eq!((s.compute, s.kv, s.accum), (DType::BF16, DType::F32, DType::F32));
+        assert_eq!(
+            (s.compute, s.kv, s.accum),
+            (DType::BF16, DType::F32, DType::F32)
+        );
         let p = ComputeCtx::ds4_parity(dev);
-        assert_eq!((p.compute, p.kv, p.accum), (DType::F32, DType::F32, DType::F32));
+        assert_eq!(
+            (p.compute, p.kv, p.accum),
+            (DType::F32, DType::F32, DType::F32)
+        );
     }
 }

hanzo-ml/src/cuda_backend/cudnn.rs

@@ -108,7 +108,10 @@ pub(crate) fn launch_conv2d<
     };
     let workspace_size = match conv2d.get_workspace_size(alg) {
         Ok(ws) => ws,
-        Err(_) => { alg = conv2d.pick_algorithm()?; conv2d.get_workspace_size(alg)? }
+        Err(_) => {
+            alg = conv2d.pick_algorithm()?;
+            conv2d.get_workspace_size(alg)?
+        }
     };
     let mut workspace = dev.cuda_stream().alloc_zeros::<u8>(workspace_size)?;
     unsafe {

hanzo-ml/src/cuda_backend/device.rs

@@ -1,7 +1,9 @@
 use crate::backend::{BackendDevice, BackendStorage};
 use crate::{CpuStorage, CpuStorageRef, DType, Layout, Result, Shape};
 pub use cudarc;
-use cudarc::driver::{result as cuda_result, sys as cuda_sys, CudaFunction, DevicePtr, UnifiedSlice};
+use cudarc::driver::{
+    result as cuda_result, sys as cuda_sys, CudaFunction, DevicePtr, UnifiedSlice,
+};
 use float8::F8E4M3;
 use half::{bf16, f16};
 pub use hanzo_kernels as kernels;

hanzo-ml/src/custom_op.rs

@@ -17,7 +17,10 @@ use std::sync::Arc;
 /// `#[cfg(feature = "vulkan")]`), so it never touches other backends.
 #[cfg(feature = "vulkan")]
 fn log_vulkan_custom_op_bail(name: &str, l: &Layout) {
-    if std::env::var("VK_PROFILE").map(|v| v != "0").unwrap_or(false) {
+    if std::env::var("VK_PROFILE")
+        .map(|v| v != "0")
+        .unwrap_or(false)
+    {
         eprintln!(
             "[VK_PROFILE] custom-op bail op={name} shape={:?} (no vulkan_fwd; would round-trip/err)",
             l.shape().dims()

hanzo-ml/src/dummy_wgpu_backend.rs

@@ -21,7 +21,9 @@ impl From<String> for WgpuError {
 }
 
 fn err<T>() -> Result<T> {
-    Err(Error::Msg("hanzo was not compiled with wgpu support".into()))
+    Err(Error::Msg(
+        "hanzo was not compiled with wgpu support".into(),
+    ))
 }
 
 macro_rules! fail {

hanzo-ml/src/fattn_decode.rs

@@ -46,7 +46,11 @@ fn fattn_decode_cpu_f32(
             // end = qpos + 1 (causal), start = max(0, qpos+1-window) (its own window).
             let qpos = kv_len - q_len + s;
             let end = qpos + 1;
-            let start = if window != 0 && end > window { end - window } else { 0 };
+            let start = if window != 0 && end > window {
+                end - window
+            } else {
+                0
+            };
             let qbase = (h * q_len + s) * HEAD_DIM;
             let qh = &q[qbase..qbase + HEAD_DIM];
 
@@ -93,16 +97,25 @@ fn fattn_decode_cpu_f32(
 /// Derive `(n_head, n_kv_head, kv_len, q_len)` from the tensor shapes and validate the contract.
 fn dims(q: &Tensor, k: &Tensor, v: &Tensor) -> Result<(usize, usize, usize, usize)> {
     if k.dims().last().copied() != Some(HEAD_DIM) {
-        crate::bail!("fattn_decode: k last dim must be {HEAD_DIM}, got {:?}", k.dims());
+        crate::bail!(
+            "fattn_decode: k last dim must be {HEAD_DIM}, got {:?}",
+            k.dims()
+        );
     }
     if k.dims() != v.dims() {
-        crate::bail!("fattn_decode: k {:?} and v {:?} must have equal shape", k.dims(), v.dims());
+        crate::bail!(
+            "fattn_decode: k {:?} and v {:?} must have equal shape",
+            k.dims(),
+            v.dims()
+        );
     }
     // q is [n_head, q_len, 512] (canonical) or [n_head, 512] (q_len == 1).
     let (n_head, q_len) = match q.dims() {
         [n_head, hd] if *hd == HEAD_DIM => (*n_head, 1usize),
         [n_head, q_len, hd] if *hd == HEAD_DIM => (*n_head, *q_len),
-        other => crate::bail!("fattn_decode: q must be [n_head, 512] or [n_head, q_len, 512], got {other:?}"),
+        other => crate::bail!(
+            "fattn_decode: q must be [n_head, 512] or [n_head, q_len, 512], got {other:?}"
+        ),
     };
     let (n_kv_head, kv_len) = match k.dims() {
         [kv_len, _] => (1usize, *kv_len),
@@ -110,7 +123,9 @@ fn dims(q: &Tensor, k: &Tensor, v: &Tensor) -> Result<(usize, usize, usize, usiz
         other => crate::bail!("fattn_decode: k rank must be 2 or 3, got {other:?}"),
     };
     if n_head == 0 || n_kv_head == 0 || n_head % n_kv_head != 0 {
-        crate::bail!("fattn_decode: n_head {n_head} must be a nonzero multiple of n_kv_head {n_kv_head}");
+        crate::bail!(
+            "fattn_decode: n_head {n_head} must be a nonzero multiple of n_kv_head {n_kv_head}"
+        );
     }
     if q_len == 0 || !(1..=8).contains(&q_len) {
         crate::bail!("fattn_decode: q_len {q_len} must be in 1..=8");
@@ -141,7 +156,10 @@ pub fn fattn_decode_f32_hd512(
     let sinks = match sinks {
         Some(s) => {
             if s.elem_count() != n_head {
-                crate::bail!("fattn_decode: sinks must have {n_head} elems, got {}", s.elem_count());
+                crate::bail!(
+                    "fattn_decode: sinks must have {n_head} elems, got {}",
+                    s.elem_count()
+                );
             }
             Some(s.to_dtype(crate::DType::F32)?.contiguous()?)
         }
@@ -173,7 +191,17 @@ pub fn fattn_decode_f32_hd512(
         }
         #[cfg(feature = "cuda")]
         Device::Cuda(dev) => cuda_impl(
-            dev, &q, &k, &v, sinks.as_ref(), n_head, n_kv_head, kv_len, q_len, window, scale,
+            dev,
+            &q,
+            &k,
+            &v,
+            sinks.as_ref(),
+            n_head,
+            n_kv_head,
+            kv_len,
+            q_len,
+            window,
+            scale,
             out_shape,
         ),
         _ => crate::bail!("fattn_decode_f32_hd512: unsupported device"),
@@ -273,9 +301,15 @@ mod tests {
         let mut rng = StdRng::seed_from_u64(
             0x00d5_4f10 ^ (kv_len as u64) ^ ((n_kv_head as u64) << 20) ^ ((q_len as u64) << 40),
         );
-        let q: Vec<f32> = (0..n_head * q_len * HEAD_DIM).map(|_| rng.random::<f32>() - 0.5).collect();
-        let k: Vec<f32> = (0..n_kv_head * kv_len * HEAD_DIM).map(|_| rng.random::<f32>() - 0.5).collect();
-        let v: Vec<f32> = (0..n_kv_head * kv_len * HEAD_DIM).map(|_| rng.random::<f32>() - 0.5).collect();
+        let q: Vec<f32> = (0..n_head * q_len * HEAD_DIM)
+            .map(|_| rng.random::<f32>() - 0.5)
+            .collect();
+        let k: Vec<f32> = (0..n_kv_head * kv_len * HEAD_DIM)
+            .map(|_| rng.random::<f32>() - 0.5)
+            .collect();
+        let v: Vec<f32> = (0..n_kv_head * kv_len * HEAD_DIM)
+            .map(|_| rng.random::<f32>() - 0.5)
+            .collect();
         let sinks: Option<Vec<f32>> = if with_sink {
             Some((0..n_head).map(|_| rng.random::<f32>() - 0.5).collect())
         } else {
@@ -303,7 +337,16 @@ mod tests {
         assert_eq!(out.dims(), &[n_head, q_len, HEAD_DIM]);
         let got = out.flatten_all()?.to_vec1::<f32>()?;
         let want = fattn_decode_cpu_f32(
-            &q, &k, &v, sinks.as_deref(), n_head, n_kv_head, kv_len, q_len, window, scale,
+            &q,
+            &k,
+            &v,
+            sinks.as_deref(),
+            n_head,
+            n_kv_head,
+            kv_len,
+            q_len,
+            window,
+            scale,
         );
 
         assert_eq!(got.len(), want.len());

hanzo-ml/src/lib.rs

@@ -51,6 +51,7 @@
 mod accelerate;
 pub mod backend;
 pub mod backprop;
+pub mod compute;
 pub mod conv;
 mod convert;
 pub mod cpu;
@@ -67,6 +68,8 @@ mod dummy_metal_backend;
 pub mod dummy_vulkan_backend;
 pub mod dummy_wgpu_backend;
 pub mod error;
+pub mod fattn_decode;
+mod hc_sinkhorn;
 mod indexer;
 pub mod layout;
 #[cfg(feature = "metal")]
@@ -77,15 +80,12 @@ pub mod model_delta;
 pub mod npy;
 pub mod op;
 pub mod pickle;
-pub mod compute;
 pub mod quantized;
 #[cfg(feature = "rocm")]
 pub mod rocm_backend;
 pub mod safetensors;
 pub mod scalar;
 pub mod shape;
-pub mod fattn_decode;
-mod hc_sinkhorn;
 mod sort;
 mod storage;
 pub mod streaming;

hanzo-ml/src/metal_backend/mod.rs

@@ -2198,7 +2198,9 @@ mod dsl_metal_dispatch {
         };
         let (rows, n) = (37usize, 128usize);
         let (nt, eps) = (n, 1e-5f32);
-        let x: Vec<f32> = (0..rows * n).map(|i| ((i % 17) as f32 - 8.0) * 0.1).collect();
+        let x: Vec<f32> = (0..rows * n)
+            .map(|i| ((i % 17) as f32 - 8.0) * 0.1)
+            .collect();
         let w: Vec<f32> = (0..n).map(|i| 1.0 + (i % 5) as f32 * 0.1).collect();
         let want = rms_ref(&x, &w, rows, n, eps);
 
@@ -2224,8 +2226,16 @@ mod dsl_metal_dispatch {
             e.set_input_buffer(3, Some(&*eb), 0);
             e.set_input_buffer(4, Some(&*ndb), 0);
             e.set_input_buffer(5, Some(&*ib), 0);
-            let grid = objc2_metal::MTLSize { width: rows, height: 1, depth: 1 };
-            let tg = objc2_metal::MTLSize { width: nt, height: 1, depth: 1 };
+            let grid = objc2_metal::MTLSize {
+                width: rows,
+                height: 1,
+                depth: 1,
+            };
+            let tg = objc2_metal::MTLSize {
+                width: nt,
+                height: 1,
+                depth: 1,
+            };
             e.dispatch_thread_groups(grid, tg);
         }
         md.wait_until_completed().unwrap();

hanzo-ml/src/model_delta/bitdelta.rs

@@ -343,8 +343,10 @@ mod tests {
     }
 
     fn write_ckpt(path: &Path, tensors: &[(&str, Tensor)]) {
-        let map: HashMap<String, Tensor> =
-            tensors.iter().map(|(k, v)| (k.to_string(), v.clone())).collect();
+        let map: HashMap<String, Tensor> = tensors
+            .iter()
+            .map(|(k, v)| (k.to_string(), v.clone()))
+            .collect();
         save_checkpoint(&map, path).unwrap();
     }
 
@@ -404,11 +406,17 @@ mod tests {
 
         write_ckpt(
             &base,
-            &[("w", Tensor::from_slice(&base_v, &shape, &Device::Cpu).unwrap())],
+            &[(
+                "w",
+                Tensor::from_slice(&base_v, &shape, &Device::Cpu).unwrap(),
+            )],
         );
         write_ckpt(
             &ft,
-            &[("w", Tensor::from_slice(&ft_v, &shape, &Device::Cpu).unwrap())],
+            &[(
+                "w",
+                Tensor::from_slice(&ft_v, &shape, &Device::Cpu).unwrap(),
+            )],
         );
 
         encode(&base, &ft, &bdfile).unwrap();
@@ -470,11 +478,17 @@ mod tests {
         let recon = tmp("ub_recon.safetensors");
         write_ckpt(
             &base,
-            &[("w", Tensor::from_slice(&base_v, &shape, &Device::Cpu).unwrap())],
+            &[(
+                "w",
+                Tensor::from_slice(&base_v, &shape, &Device::Cpu).unwrap(),
+            )],
         );
         write_ckpt(
             &ft,
-            &[("w", Tensor::from_slice(&ft_v, &shape, &Device::Cpu).unwrap())],
+            &[(
+                "w",
+                Tensor::from_slice(&ft_v, &shape, &Device::Cpu).unwrap(),
+            )],
         );
 
         encode(&base, &ft, &bdfile).unwrap();
@@ -508,7 +522,10 @@ mod tests {
         write_ckpt(
             &base,
             &[
-                ("a", Tensor::zeros((2, 3), DType::F32, &Device::Cpu).unwrap()),
+                (
+                    "a",
+                    Tensor::zeros((2, 3), DType::F32, &Device::Cpu).unwrap(),
+                ),
                 ("b", Tensor::zeros((5,), DType::F32, &Device::Cpu).unwrap()),
             ],
         );
@@ -564,7 +581,11 @@ mod tests {
         // dtype preserved on the reconstructed tensor.
         assert_eq!(recon_map["w"].dtype(), DType::BF16);
         // untouched tensor copied through unchanged.
-        let frozen_v: Vec<f32> = recon_map["frozen"].flatten_all().unwrap().to_vec1().unwrap();
+        let frozen_v: Vec<f32> = recon_map["frozen"]
+            .flatten_all()
+            .unwrap()
+            .to_vec1()
+            .unwrap();
         assert_eq!(frozen_v, vec![9.0, 8.0]);
 
         for p in [base, ft, bdfile, recon] {
@@ -579,11 +600,17 @@ mod tests {
         let bdfile = tmp("em.bitdelta");
         write_ckpt(
             &base,
-            &[("w", Tensor::zeros((2, 2), DType::F32, &Device::Cpu).unwrap())],
+            &[(
+                "w",
+                Tensor::zeros((2, 2), DType::F32, &Device::Cpu).unwrap(),
+            )],
         );
         write_ckpt(
             &ft,
-            &[("w", Tensor::zeros((3, 3), DType::F32, &Device::Cpu).unwrap())],
+            &[(
+                "w",
+                Tensor::zeros((3, 3), DType::F32, &Device::Cpu).unwrap(),
+            )],
         );
 
         let err = encode(&base, &ft, &bdfile).unwrap_err();

hanzo-ml/src/model_delta/soup.rs

@@ -212,8 +212,10 @@ mod tests {
     }
 
     fn write(path: &Path, tensors: &[(&str, Tensor)]) {
-        let map: HashMap<String, Tensor> =
-            tensors.iter().map(|(k, v)| (k.to_string(), v.clone())).collect();
+        let map: HashMap<String, Tensor> = tensors
+            .iter()
+            .map(|(k, v)| (k.to_string(), v.clone()))
+            .collect();
         save_checkpoint(&map, path).unwrap();
     }
 
@@ -375,7 +377,10 @@ mod tests {
         let o = tmp("dp_o.safetensors");
         write(
             &base,
-            &[("w", t(&[1.0, 1.0], &[2])), ("frozen", t(&[7.0, 8.0], &[2]))],
+            &[
+                ("w", t(&[1.0, 1.0], &[2])),
+                ("frozen", t(&[7.0, 8.0], &[2])),
+            ],
         );
         write(&ft, &[("w", t(&[3.0, 3.0], &[2]))]);
 

hanzo-ml/src/quantized/cuda.rs

@@ -32,7 +32,8 @@ pub fn set_force_dmmv(f: bool) {
 // dp4a path for *correctness*, e.g. DeepSeek-V4's W8A8 QAT, turns it on at load). When on,
 // QCudaStorage::fwd tries it first and falls back to the legacy on-GPU q8_1 kernels for any
 // dtype/shape it doesn't support. Mirrors `FORCE_DMMV`: an AtomicBool seeded once from the env.
-pub(crate) static FAST_MMQ: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
+pub(crate) static FAST_MMQ: std::sync::atomic::AtomicBool =
+    std::sync::atomic::AtomicBool::new(false);
 static FAST_MMQ_INIT: std::sync::Once = std::sync::Once::new();
 
 /// Programmatically enable/disable the fast int8 mmq/mmvq path (overrides the env default,
@@ -467,7 +468,10 @@ fn mul_mat_vec_iquant_dp4a(
         crate::bail!("only bsize between 1 and 8 are supported, got {b_size}");
     }
     if y.len() != ncols * b_size {
-        crate::bail!("unexpected y size {}, ncols {ncols} bsize {b_size}", y.len());
+        crate::bail!(
+            "unexpected y size {}, ncols {ncols} bsize {b_size}",
+            y.len()
+        );
     }
     let nblk32 = ncols / 32;
 
@@ -771,13 +775,17 @@ impl QCudaStorage {
         n: usize,
         k: usize,
     ) -> Result<CudaStorage> {
-        let suffix = iquant_dp4a_suffix(self.dtype)
-            .ok_or_else(|| crate::Error::Msg(format!("no i-quant MoE kernel for {:?}", self.dtype)).bt())?;
+        let suffix = iquant_dp4a_suffix(self.dtype).ok_or_else(|| {
+            crate::Error::Msg(format!("no i-quant MoE kernel for {:?}", self.dtype)).bt()
+        })?;
         if k % 256 != 0 {
             crate::bail!("i-quant MoE ncols {k} must be a multiple of 256");
         }
         if x_flat.len() != nrows * k {
-            crate::bail!("unexpected x_flat size {}, nrows {nrows} k {k}", x_flat.len());
+            crate::bail!(
+                "unexpected x_flat size {}, nrows {nrows} k {k}",
+                x_flat.len()
+            );
         }
         let dev = self.device();
         let nblk32 = k / 32;
@@ -1220,8 +1228,15 @@ impl QCudaStorage {
         if ncols != k {
             crate::bail!("mismatch on matmul dim {self_shape:?} {:?}", rhs_l.shape())
         }
-        let out =
-            mul_mat_vec_iquant_dp4a(&self.data, &rhs, self.dtype, ncols, nrows, b_size, self.device())?;
+        let out = mul_mat_vec_iquant_dp4a(
+            &self.data,
+            &rhs,
+            self.dtype,
+            ncols,
+            nrows,
+            b_size,
+            self.device(),
+        )?;
         let mut out_shape = rhs_l.shape().dims().to_vec();
         out_shape.pop();
         out_shape.push(nrows);
@@ -1584,7 +1599,8 @@ mod test {
                 .collect();
 
             // LOAD onto CUDA (this is the call that used to bail for these dtypes).
-            let qcuda = match QStorage::from_data(std::borrow::Cow::Owned(bytes), &cuda_dev, dtype)? {
+            let qcuda = match QStorage::from_data(std::borrow::Cow::Owned(bytes), &cuda_dev, dtype)?
+            {
                 QStorage::Cuda(s) => s,
                 _ => unreachable!("from_data on a CUDA device must yield CUDA storage"),
             };
@@ -1654,7 +1670,8 @@ mod test {
                 _ => unreachable!(),
             };
             let bytes = qcpu.data()?.into_owned();
-            let qcuda = match QStorage::from_data(std::borrow::Cow::Owned(bytes), &cuda_dev, dtype)? {
+            let qcuda = match QStorage::from_data(std::borrow::Cow::Owned(bytes), &cuda_dev, dtype)?
+            {
                 QStorage::Cuda(s) => s,
                 _ => unreachable!(),
             };
@@ -1666,8 +1683,9 @@ mod test {
             let input_l = crate::Layout::contiguous((m, k));
             let self_shape: crate::Shape = (n, k).into();
 
-            let (q_out, q_shape) = crate::quantized::fast_mmq::try_fwd(&qcuda, &self_shape, &input, &input_l)?
-                .expect("fast_mmq::try_fwd must support q8_1 dense prefill");
+            let (q_out, q_shape) =
+                crate::quantized::fast_mmq::try_fwd(&qcuda, &self_shape, &input, &input_l)?
+                    .expect("fast_mmq::try_fwd must support q8_1 dense prefill");
             assert_eq!(q_shape.dims().to_vec(), vec![m, n]);
             let q_got = dev.clone_dtoh(&q_out.as_cuda_slice::<f32>()?.as_view())?;
 
@@ -1687,7 +1705,10 @@ mod test {
                     }
                     let gq = q_got[r * n + j];
                     let gd = d_got[r * n + j];
-                    assert!(gq.is_finite() && gd.is_finite(), "{dtype:?}: non-finite output");
+                    assert!(
+                        gq.is_finite() && gd.is_finite(),
+                        "{dtype:?}: non-finite output"
+                    );
                     max_abs = max_abs.max(acc.abs());
                     err_q = err_q.max((gq - acc).abs());
                     err_d = err_d.max((gd - acc).abs());
@@ -1696,9 +1717,18 @@ mod test {
             }
             let tol_q = 0.05 * max_abs + 1e-3;
             let tol_d = 1e-3 * max_abs + 1e-4;
-            assert!(err_q <= tol_q, "{dtype:?}: DENSE QMMQ vs oracle err {err_q} > tol {tol_q} (max_abs {max_abs})");
-            assert!(err_d <= tol_d, "{dtype:?}: DEQUANT-DENSE vs oracle err {err_d} > tol {tol_d} (max_abs {max_abs})");
-            assert!(err_qd <= tol_q, "{dtype:?}: QMMQ vs DEQUANT-DENSE err {err_qd} > tol {tol_q} (max_abs {max_abs})");
+            assert!(
+                err_q <= tol_q,
+                "{dtype:?}: DENSE QMMQ vs oracle err {err_q} > tol {tol_q} (max_abs {max_abs})"
+            );
+            assert!(
+                err_d <= tol_d,
+                "{dtype:?}: DEQUANT-DENSE vs oracle err {err_d} > tol {tol_d} (max_abs {max_abs})"
+            );
+            assert!(
+                err_qd <= tol_q,
+                "{dtype:?}: QMMQ vs DEQUANT-DENSE err {err_qd} > tol {tol_q} (max_abs {max_abs})"
+            );
             eprintln!("[DENSE_QMMQ] {dtype:?}: max_abs={max_abs:.4} err_qmmq={err_q:.5} err_dequant={err_d:.6} err_qmmq_vs_dequant={err_qd:.5}");
         }
         Ok(())
@@ -1745,7 +1775,10 @@ mod test {
 
         for dtype in [GgmlDType::Q4K, GgmlDType::Q6K, GgmlDType::Q8_0] {
             let mut bank = QCudaStorage::zeros(&dev, e * n * k, dtype)?;
-            bank.quantize(&CudaStorage::wrap_cuda_slice(dev.clone_htod(&w)?, dev.clone()))?;
+            bank.quantize(&CudaStorage::wrap_cuda_slice(
+                dev.clone_htod(&w)?,
+                dev.clone(),
+            ))?;
             let deq = bank.dequantize(e * n * k)?;
             let w_deq = dev.clone_dtoh(&deq.as_cuda_slice::<f32>()?.as_view())?;
             let self_shape: crate::Shape = (e, n, k).into();

hanzo-ml/src/quantized/fast_mmq.rs

@@ -556,7 +556,11 @@ pub(crate) fn indexed_moe_grouped(
         let e = (e as usize).min(e_cnt - 1);
         let p = cursor[e] as usize;
         ids_dst[p] = s as i32;
-        quantize_ids[p] = if input_dim1 == 1 { (s / topk) as i32 } else { s as i32 };
+        quantize_ids[p] = if input_dim1 == 1 {
+            (s / topk) as i32
+        } else {
+            s as i32
+        };
         cursor[e] += 1;
     }
     let ncols_max = counts.iter().copied().max().unwrap_or(0) as i64;
@@ -609,8 +613,8 @@ pub(crate) fn indexed_moe_grouped(
             qids_ptr,
             scratch_ptr,
             0,
-            k as i64,        // ne00 (real row length)
-            k as i64,        // s01  (input row stride, elements)
+            k as i64, // ne00 (real row length)
+            k as i64, // s01  (input row stride, elements)
             0,
             0,
             k_padded as i64, // ne0  (padded loop bound)

hanzo-ml/src/quantized/gguf_file.rs

@@ -172,9 +172,11 @@ impl TensorInfo {
                 // (page-aligned base + 32-aligned GGUF offsets satisfy this for every real model).
                 let addr = mmap.as_ptr() as usize + start;
                 if addr.is_multiple_of(self.ggml_dtype.type_align()) {
-                    let storage = super::QStorage::Cpu(
-                        self.ggml_dtype.from_mmap(mmap.clone(), start, n_blocks),
-                    );
+                    let storage = super::QStorage::Cpu(self.ggml_dtype.from_mmap(
+                        mmap.clone(),
+                        start,
+                        n_blocks,
+                    ));
                     QTensor::new(storage, dims)
                 } else {
                     super::ggml_file::qtensor_from_ggml(

hanzo-ml/src/quantized/k_quants.rs

@@ -2595,7 +2595,10 @@ pub fn matmul<T: GgmlType>(
     // already has plenty of work to spread.
     //
     // `MIN_LEN` overrides the heuristic entirely (must be >= 1).
-    let min_len = match std::env::var("MIN_LEN").ok().and_then(|s| s.parse::<usize>().ok()) {
+    let min_len = match std::env::var("MIN_LEN")
+        .ok()
+        .and_then(|s| s.parse::<usize>().ok())
+    {
         Some(v) if v >= 1 => v,
         _ => {
             if m == 1 {
@@ -2616,8 +2619,8 @@ pub fn matmul<T: GgmlType>(
             .with_min_len(min_len)
             .with_max_len(512)
             .for_each(|(col_idx, dst)| {
-                let rhs_col = &rhs_t
-                    [col_idx * weight_blocks_per_row..(col_idx + 1) * weight_blocks_per_row];
+                let rhs_col =
+                    &rhs_t[col_idx * weight_blocks_per_row..(col_idx + 1) * weight_blocks_per_row];
                 *dst = T::vec_dot(k, rhs_col, lhs_row);
             });
     }

hanzo-ml/src/quantized/metal.rs

@@ -534,8 +534,12 @@ impl QMetalStorage {
                 };
                 self.moe_expert_matmul(xs, eid as usize * expert_bytes, m, n, k)?
             };
-            let y_e =
-                crate::tensor::from_storage(crate::Storage::Metal(y_e), (m, n), none.clone(), false);
+            let y_e = crate::tensor::from_storage(
+                crate::Storage::Metal(y_e),
+                (m, n),
+                none.clone(),
+                false,
+            );
             out_flat = out_flat.index_add(&idx, &y_e, 0)?;
         }
 

hanzo-ml/src/quantized/quant_format.rs

@@ -32,7 +32,7 @@
 #![allow(unused_macros)]
 #![allow(unused_imports)]
 
-use super::k_quants::{GgmlType, BlockQ8_0, QK8_0};
+use super::k_quants::{BlockQ8_0, GgmlType, QK8_0};
 use half::f16;
 
 /// The one shared matmul building block behind every decode-only type's `vec_dot`.

hanzo-ml/src/rocm_backend/device.rs

@@ -4,12 +4,11 @@ use half::{bf16, f16};
 use hanzo_rocm_kernels::compile::KernelCache;
 use std::sync::{Arc, Mutex, RwLock};
 
-use super::wrappers::{
-    DevicePool, SendSyncDeviceMemory, SendSyncPseudoRng,
-    SendSyncRocblasHandle, SendSyncStream,
-};
 #[cfg(feature = "rocm-miopen")]
 use super::wrappers::SendSyncMIOpenHandle;
+use super::wrappers::{
+    DevicePool, SendSyncDeviceMemory, SendSyncPseudoRng, SendSyncRocblasHandle, SendSyncStream,
+};
 use super::{Affine, RocmError, RocmStorage, RocmStorageSlice};
 use rocm_rs::hip::Device as HipDevice;
 
@@ -133,7 +132,9 @@ impl RocmDevice {
         let mut guard = self.rocrand.lock().unwrap();
         if guard.is_none() {
             let mut rng = SendSyncPseudoRng::new(rocm_rs::rocrand::rng_type::PSEUDO_DEFAULT)
-                .map_err(|e| crate::Error::Msg(format!("Failed to create rocrand generator: {}", e)))?;
+                .map_err(|e| {
+                    crate::Error::Msg(format!("Failed to create rocrand generator: {}", e))
+                })?;
             rng.set_seed(*self.seed_value.read().unwrap())
                 .map_err(|e| crate::Error::Msg(format!("Failed to set rocrand seed: {}", e)))?;
             *guard = Some(rng);

hanzo-ml/src/rocm_backend/mod.rs

@@ -505,7 +505,8 @@ pub fn launch_config(num_elems: usize) -> (rocm_rs::hip::Dim3, rocm_rs::hip::Dim
 /// `qdw_traits<WTYPE>` in the .hip = a fully wired type, for BOTH decode and MoE. No per-quant
 /// kernel, no per-quant launcher: `matvec_quant` reads the (type, activation) pair off this enum
 /// and dispatches the single core's f16/bf16 entry point. Mirrors the CPU `for_each_quant!` table.
-static FORCE_SCALAR_MATVEC: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
+static FORCE_SCALAR_MATVEC: std::sync::atomic::AtomicBool =
+    std::sync::atomic::AtomicBool::new(false);
 
 /// Test-only: force `dp4a_active` to false so the scalar `qmatvec_core` runs, letting the bit-exact
 /// numeric oracles gate the scalar decode against the exact reference. Never set in production.
@@ -586,11 +587,29 @@ impl RocmQuantType {
     /// Elements per block (must divide `k`). Matches `qdw_traits<WTYPE>::ELEMS` in quant.hip.
     pub fn block_elems(self) -> usize {
         match self {
-            Self::Q8_0 | Self::Q4_0 | Self::Q4_1 | Self::Q5_0 | Self::Q5_1 | Self::Q8_1
-            | Self::IQ4_NL | Self::MXFP4 => 32,
-            Self::Q4K | Self::Q6K | Self::IQ4_XS | Self::TQ2_0 | Self::Q2K | Self::Q3K | Self::Q5K
-            | Self::IQ2_XXS | Self::IQ2_XS | Self::IQ2_S | Self::IQ3_XXS | Self::IQ3_S
-            | Self::TQ1_0 | Self::IQ1_S | Self::IQ1_M => 256,
+            Self::Q8_0
+            | Self::Q4_0
+            | Self::Q4_1
+            | Self::Q5_0
+            | Self::Q5_1
+            | Self::Q8_1
+            | Self::IQ4_NL
+            | Self::MXFP4 => 32,
+            Self::Q4K
+            | Self::Q6K
+            | Self::IQ4_XS
+            | Self::TQ2_0
+            | Self::Q2K
+            | Self::Q3K
+            | Self::Q5K
+            | Self::IQ2_XXS
+            | Self::IQ2_XS
+            | Self::IQ2_S
+            | Self::IQ3_XXS
+            | Self::IQ3_S
+            | Self::TQ1_0
+            | Self::IQ1_S
+            | Self::IQ1_M => 256,
         }
     }
 
@@ -682,8 +701,17 @@ impl RocmQuantType {
             Self::Q5_0 => "qmmq_q5_0_f16",
             Self::Q5_1 => "qmmq_q5_1_f16",
             Self::Q8_1 => "qmmq_q8_1_f16",
-            Self::Q2K | Self::Q3K | Self::IQ2_XXS | Self::IQ2_XS | Self::IQ2_S | Self::IQ3_XXS
-            | Self::IQ3_S | Self::IQ4_NL | Self::TQ1_0 | Self::IQ1_S | Self::IQ1_M
+            Self::Q2K
+            | Self::Q3K
+            | Self::IQ2_XXS
+            | Self::IQ2_XS
+            | Self::IQ2_S
+            | Self::IQ3_XXS
+            | Self::IQ3_S
+            | Self::IQ4_NL
+            | Self::TQ1_0
+            | Self::IQ1_S
+            | Self::IQ1_M
             | Self::MXFP4 => {
                 unreachable!("prefill_kernel: {self:?} is decode-only (gated by qmmq_capable)")
             }
@@ -729,11 +757,22 @@ impl RocmQuantType {
             (Self::Q8_1, 64) => "moe_qmmq_q8_1_tm64_f16",
             (Self::Q8_1, _) => "moe_qmmq_q8_1_tm32_f16",
             (
-                Self::Q2K | Self::Q3K | Self::IQ2_XXS | Self::IQ2_XS | Self::IQ2_S | Self::IQ3_XXS
-                | Self::IQ3_S | Self::IQ4_NL | Self::TQ1_0 | Self::IQ1_S | Self::IQ1_M
+                Self::Q2K
+                | Self::Q3K
+                | Self::IQ2_XXS
+                | Self::IQ2_XS
+                | Self::IQ2_S
+                | Self::IQ3_XXS
+                | Self::IQ3_S
+                | Self::IQ4_NL
+                | Self::TQ1_0
+                | Self::IQ1_S
+                | Self::IQ1_M
                 | Self::MXFP4,
                 _,
-            ) => unreachable!("moe_prefill_kernel: {self:?} is decode-only (gated by qmmq_capable)"),
+            ) => {
+                unreachable!("moe_prefill_kernel: {self:?} is decode-only (gated by qmmq_capable)")
+            }
         }
     }
 
@@ -985,7 +1024,10 @@ impl Act {
             RocmStorageSlice::F16(_) => Ok(Act::F16),
             RocmStorageSlice::BF16(_) => Ok(Act::Bf16),
             RocmStorageSlice::F32(_) => Ok(Act::F32),
-            other => crate::bail!("dp4a activation must be f16/bf16/f32, got {:?}", other.dtype()),
+            other => crate::bail!(
+                "dp4a activation must be f16/bf16/f32, got {:?}",
+                other.dtype()
+            ),
         }
     }
 
@@ -1483,7 +1525,10 @@ impl RocmDevice {
         match &w.slice {
             RocmStorageSlice::F16(_) => launch_gemv!(F16, f16, "dense_gemv_f16"),
             RocmStorageSlice::F32(_) => launch_gemv!(F32, f32, "dense_gemv_f32"),
-            other => crate::bail!("dense_gemv: weight dtype {:?} unsupported (f16/f32 only)", other.dtype()),
+            other => crate::bail!(
+                "dense_gemv: weight dtype {:?} unsupported (f16/f32 only)",
+                other.dtype()
+            ),
         }
     }
 
@@ -1609,8 +1654,14 @@ impl RocmDevice {
             )?;
         }
         Ok((
-            RocmStorage { slice: RocmStorageSlice::U8(xq), device: self.clone() },
-            RocmStorage { slice: RocmStorageSlice::F16(xd), device: self.clone() },
+            RocmStorage {
+                slice: RocmStorageSlice::U8(xq),
+                device: self.clone(),
+            },
+            RocmStorage {
+                slice: RocmStorageSlice::F16(xd),
+                device: self.clone(),
+            },
         ))
     }
 
@@ -1664,7 +1715,10 @@ impl RocmDevice {
             RocmStorageSlice::F16(s) => ("quantize_q8_1", s.as_ptr()),
             RocmStorageSlice::F32(s) => ("quantize_q8_1_f32", s.as_ptr()),
             RocmStorageSlice::BF16(s) => ("quantize_q8_1_bf16", s.as_ptr()),
-            other => crate::bail!("quantize_q8_1: x must be f16/f32/bf16, got {:?}", other.dtype()),
+            other => crate::bail!(
+                "quantize_q8_1: x must be f16/f32/bf16, got {:?}",
+                other.dtype()
+            ),
         };
         let nblk = k / 32;
         let xq = self.alloc::<u8>(m * k)?;
@@ -1697,9 +1751,18 @@ impl RocmDevice {
             )?;
         }
         Ok((
-            RocmStorage { slice: RocmStorageSlice::U8(xq), device: self.clone() },
-            RocmStorage { slice: RocmStorageSlice::F16(xd), device: self.clone() },
-            RocmStorage { slice: RocmStorageSlice::I32(xs), device: self.clone() },
+            RocmStorage {
+                slice: RocmStorageSlice::U8(xq),
+                device: self.clone(),
+            },
+            RocmStorage {
+                slice: RocmStorageSlice::F16(xd),
+                device: self.clone(),
+            },
+            RocmStorage {
+                slice: RocmStorageSlice::I32(xs),
+                device: self.clone(),
+            },
         ))
     }
 
@@ -1720,7 +1783,10 @@ impl RocmDevice {
     ) -> Result<RocmStorage> {
         use hanzo_rocm_kernels::kernel::{KernelSource, QuantKernel};
         if k % qt.block_elems() != 0 {
-            crate::bail!("qmmq_quant({qt:?}): k={k} not a multiple of block size {}", qt.block_elems());
+            crate::bail!(
+                "qmmq_quant({qt:?}): k={k} not a multiple of block size {}",
+                qt.block_elems()
+            );
         }
         // Pre-quantize activations once. SYMMETRIC weights need only xq/xd; ASYMMETRIC (Q4_K) also
         // needs the per-block int8 sum xs (the q8_1 bias term). xq/xd/xs own their VRAM until the end
@@ -1735,7 +1801,11 @@ impl RocmDevice {
         };
         // Dummy 1-elem xs for the symmetric path (the kernel's asym branch elides, so it is never
         // dereferenced; we just need a non-null device pointer to bind to the kernel parameter).
-        let xs_dummy = if xs_opt.is_none() { Some(self.alloc::<i32>(1)?) } else { None };
+        let xs_dummy = if xs_opt.is_none() {
+            Some(self.alloc::<i32>(1)?)
+        } else {
+            None
+        };
         let xq_ptr = match &xq.slice {
             RocmStorageSlice::U8(s) => s.as_ptr(),
             _ => crate::bail!("qmmq_quant: xq must be u8"),
@@ -1789,7 +1859,10 @@ impl RocmDevice {
                 ],
             )?;
         }
-        Ok(RocmStorage { slice: RocmStorageSlice::F16(out), device: self.clone() })
+        Ok(RocmStorage {
+            slice: RocmStorageSlice::F16(out),
+            device: self.clone(),
+        })
     }
 
     /// FUSED indexed-MoE PREFILL GEMM. Groups the routed slots by expert (host-side counting sort over
@@ -1812,7 +1885,10 @@ impl RocmDevice {
         use hanzo_rocm_kernels::kernel::{KernelSource, QuantKernel};
         use std::ffi::c_void;
         if k % qt.block_elems() != 0 {
-            crate::bail!("moe_qmmq_quant({qt:?}): k={k} not a multiple of block size {}", qt.block_elems());
+            crate::bail!(
+                "moe_qmmq_quant({qt:?}): k={k} not a multiple of block size {}",
+                qt.block_elems()
+            );
         }
         let wbank_mem = match &wbank.slice {
             RocmStorageSlice::U8(m) => m,
@@ -1860,7 +1936,10 @@ impl RocmDevice {
         // beats 128 across the whole prefill range; TILE_M=32 wins the few-token decode micro-batches
         // (~1-8 tok/expert, where a 128-tile is ~1-6% full); 128 only pays off for very dense routing
         // (avg >> 64), which needs a very long prompt. MOE_TILE_M overrides for the A/B.
-        let tile_m: i32 = match std::env::var("MOE_TILE_M").ok().and_then(|s| s.parse().ok()) {
+        let tile_m: i32 = match std::env::var("MOE_TILE_M")
+            .ok()
+            .and_then(|s| s.parse().ok())
+        {
             Some(v @ (32 | 64 | 128)) => v,
             _ if avg_tok >= 160 => 128,
             _ if avg_tok >= 24 => 64,
@@ -1882,7 +1961,10 @@ impl RocmDevice {
         let num_row_tiles = tile_expert.len();
         if num_row_tiles == 0 {
             let out = self.alloc::<f16>(nslots * n)?;
-            return Ok(RocmStorage { slice: RocmStorageSlice::F16(out), device: self.clone() });
+            return Ok(RocmStorage {
+                slice: RocmStorageSlice::F16(out),
+                device: self.clone(),
+            });
         }
         // M-tile fill telemetry (LEVER-1 measurement): valid rows vs WMMA tile capacity. The wasted
         // fraction = empty WMMA M-rows = the directly-recoverable matrix-core work. Env-gated.
@@ -1902,7 +1984,11 @@ impl RocmDevice {
             let (xq, xd, xs) = self.quantize_q8_1(x, nslots, k)?;
             (xq, xd, Some(xs))
         };
-        let xs_dummy = if xs_opt.is_none() { Some(self.alloc::<i32>(1)?) } else { None };
+        let xs_dummy = if xs_opt.is_none() {
+            Some(self.alloc::<i32>(1)?)
+        } else {
+            None
+        };
         let xq_ptr = match &xq.slice {
             RocmStorageSlice::U8(s) => s.as_ptr(),
             _ => crate::bail!("moe_qmmq_quant: xq must be u8"),
@@ -1965,7 +2051,10 @@ impl RocmDevice {
                 ],
             )?;
         }
-        Ok(RocmStorage { slice: RocmStorageSlice::F16(out), device: self.clone() })
+        Ok(RocmStorage {
+            slice: RocmStorageSlice::F16(out),
+            device: self.clone(),
+        })
     }
 
     /// FUSED MoE expert-combine: `out[i,j] = sum_e scores[i,e] * ys[i,e,j]` in ONE launch. Replaces
@@ -2017,7 +2106,10 @@ impl RocmDevice {
                         ],
                     )?;
                 }
-                RocmStorage { slice: RocmStorageSlice::$variant(out), device: self.clone() }
+                RocmStorage {
+                    slice: RocmStorageSlice::$variant(out),
+                    device: self.clone(),
+                }
             }};
         }
         let out = match &ys.slice {
@@ -2081,8 +2173,14 @@ impl RocmDevice {
             )?;
         }
         Ok((
-            RocmStorage { slice: RocmStorageSlice::U32(ids), device: self.clone() },
-            RocmStorage { slice: RocmStorageSlice::F32(w), device: self.clone() },
+            RocmStorage {
+                slice: RocmStorageSlice::U32(ids),
+                device: self.clone(),
+            },
+            RocmStorage {
+                slice: RocmStorageSlice::F32(w),
+                device: self.clone(),
+            },
         ))
     }
 }
@@ -2315,11 +2413,9 @@ impl RocmStorage {
         let r_off = residual_layout.start_offset();
         let alpha_off = alpha_layout.start_offset();
         let (x_mem, r_mem, alpha_mem) = match (&x.slice, &residual.slice, &alpha.slice) {
-            (
-                RocmStorageSlice::F32(xm),
-                RocmStorageSlice::F32(rm),
-                RocmStorageSlice::F32(am),
-            ) => (xm, rm, am),
+            (RocmStorageSlice::F32(xm), RocmStorageSlice::F32(rm), RocmStorageSlice::F32(am)) => {
+                (xm, rm, am)
+            }
             _ => crate::bail!(
                 "add_rms_norm: f32 inputs required (x={:?} residual={:?} alpha={:?})",
                 x.slice.dtype(),
@@ -2511,7 +2607,10 @@ impl RocmStorage {
         }
         let pos_mem = match &positions.slice {
             RocmStorageSlice::U32(m) => m,
-            other => crate::bail!("rope_positions positions must be u32, got {:?}", other.dtype()),
+            other => crate::bail!(
+                "rope_positions positions must be u32, got {:?}",
+                other.dtype()
+            ),
         };
         let device = self.device.clone();
         let q_el = b * h * t * d;
@@ -2614,7 +2713,10 @@ impl RocmStorage {
             RocmStorageSlice::F32(_) => launch_rope_pos!(F32, f32, "rope_positions_f32"),
             RocmStorageSlice::F16(_) => launch_rope_pos!(F16, f16, "rope_positions_f16"),
             RocmStorageSlice::BF16(_) => launch_rope_pos!(BF16, bf16, "rope_positions_bf16"),
-            other => crate::bail!("rope_positions on rocm unsupported dtype {:?}", other.dtype()),
+            other => crate::bail!(
+                "rope_positions on rocm unsupported dtype {:?}",
+                other.dtype()
+            ),
         }
     }
 
@@ -2660,7 +2762,10 @@ impl RocmStorage {
         }
         let pos_mem = match &positions.slice {
             RocmStorageSlice::U32(m) => m,
-            other => crate::bail!("rope_norm_positions positions must be u32, got {:?}", other.dtype()),
+            other => crate::bail!(
+                "rope_norm_positions positions must be u32, got {:?}",
+                other.dtype()
+            ),
         };
         let device = self.device.clone();
         let q_el = b * h * t * d;
@@ -2682,7 +2787,12 @@ impl RocmStorage {
         macro_rules! launch_rope_norm {
             ($variant:ident, $ty:ty, $func:literal) => {{
                 let (q_mem, k_mem, qw_mem, kw_mem, cos_mem, sin_mem) = match (
-                    &q.slice, &k.slice, &q_weight.slice, &k_weight.slice, &cos.slice, &sin.slice,
+                    &q.slice,
+                    &k.slice,
+                    &q_weight.slice,
+                    &k_weight.slice,
+                    &cos.slice,
+                    &sin.slice,
                 ) {
                     (
                         RocmStorageSlice::$variant(a),
@@ -2692,7 +2802,9 @@ impl RocmStorage {
                         RocmStorageSlice::$variant(e),
                         RocmStorageSlice::$variant(f),
                     ) => (a, b, c, d, e, f),
-                    _ => crate::bail!("rope_norm_positions: q/k/weights/cos/sin must share q's dtype"),
+                    _ => crate::bail!(
+                        "rope_norm_positions: q/k/weights/cos/sin must share q's dtype"
+                    ),
                 };
                 let q_out = device.alloc::<$ty>(q_el)?;
                 let k_out = device.alloc::<$ty>(k_el)?;
@@ -2736,8 +2848,14 @@ impl RocmStorage {
                     )?;
                 }
                 Ok((
-                    Self { slice: RocmStorageSlice::$variant(q_out), device: device.clone() },
-                    Self { slice: RocmStorageSlice::$variant(k_out), device: device.clone() },
+                    Self {
+                        slice: RocmStorageSlice::$variant(q_out),
+                        device: device.clone(),
+                    },
+                    Self {
+                        slice: RocmStorageSlice::$variant(k_out),
+                        device: device.clone(),
+                    },
                 ))
             }};
         }
@@ -2746,7 +2864,10 @@ impl RocmStorage {
             RocmStorageSlice::F32(_) => launch_rope_norm!(F32, f32, "rope_norm_positions_f32"),
             RocmStorageSlice::F16(_) => launch_rope_norm!(F16, f16, "rope_norm_positions_f16"),
             RocmStorageSlice::BF16(_) => launch_rope_norm!(BF16, bf16, "rope_norm_positions_bf16"),
-            other => crate::bail!("rope_norm_positions on rocm unsupported dtype {:?}", other.dtype()),
+            other => crate::bail!(
+                "rope_norm_positions on rocm unsupported dtype {:?}",
+                other.dtype()
+            ),
         }
     }
 
@@ -2774,7 +2895,9 @@ impl RocmStorage {
         let (kb, hkv, lk, kd) = k_layout.shape().dims4()?;
         let (vb, vhkv, vlk, vd) = v_layout.shape().dims4()?;
         if d != FA_DH || kd != FA_DH || vd != FA_DH {
-            crate::bail!("flash_attn on rocm requires head_dim == {FA_DH}, got q={d} k={kd} v={vd}");
+            crate::bail!(
+                "flash_attn on rocm requires head_dim == {FA_DH}, got q={d} k={kd} v={vd}"
+            );
         }
         if (kb, vb, vhkv, vlk) != (b, b, hkv, lk) {
             crate::bail!("flash_attn q/k/v batch/kv-head/kv-len mismatch {q_layout:?} {k_layout:?} {v_layout:?}");
@@ -2862,7 +2985,10 @@ impl RocmStorage {
         match &q.slice {
             RocmStorageSlice::F16(_) => launch_flash!(F16, f16, "flash_attn_f16"),
             RocmStorageSlice::BF16(_) => launch_flash!(BF16, bf16, "flash_attn_bf16"),
-            other => crate::bail!("flash_attn on rocm unsupported dtype {:?} (f16/bf16 only)", other.dtype()),
+            other => crate::bail!(
+                "flash_attn on rocm unsupported dtype {:?} (f16/bf16 only)",
+                other.dtype()
+            ),
         }
     }
 
@@ -2923,7 +3049,10 @@ impl RocmStorage {
             RocmStorageSlice::F32(_) => launch_softmax!(F32, f32, "softmax_f32"),
             RocmStorageSlice::F16(_) => launch_softmax!(F16, f16, "softmax_f16"),
             RocmStorageSlice::BF16(_) => launch_softmax!(BF16, bf16, "softmax_bf16"),
-            other => crate::bail!("softmax_last_dim on rocm unsupported dtype {:?}", other.dtype()),
+            other => crate::bail!(
+                "softmax_last_dim on rocm unsupported dtype {:?}",
+                other.dtype()
+            ),
         }
     }
 
@@ -3827,90 +3956,27 @@ impl BackendStorage for RocmStorage {
 
         let src_dtype = self.slice.dtype();
         let slice = match dtype {
-            DType::U8 => cast_launch!(
-                dev,
-                grid,
-                block,
-                el,
-                num_dims,
-                ds,
-                src_ptr,
-                src_dtype,
-                u8,
-                U8
-            ),
-            DType::U32 => cast_launch!(
-                dev,
-                grid,
-                block,
-                el,
-                num_dims,
-                ds,
-                src_ptr,
-                src_dtype,
-                u32,
-                U32
-            ),
-            DType::I64 => cast_launch!(
-                dev,
-                grid,
-                block,
-                el,
-                num_dims,
-                ds,
-                src_ptr,
-                src_dtype,
-                i64,
-                I64
-            ),
-            DType::BF16 => cast_launch!(
-                dev,
-                grid,
-                block,
-                el,
-                num_dims,
-                ds,
-                src_ptr,
-                src_dtype,
-                bf16,
-                BF16
-            ),
-            DType::F16 => cast_launch!(
-                dev,
-                grid,
-                block,
-                el,
-                num_dims,
-                ds,
-                src_ptr,
-                src_dtype,
-                f16,
-                F16
-            ),
-            DType::F32 => cast_launch!(
-                dev,
-                grid,
-                block,
-                el,
-                num_dims,
-                ds,
-                src_ptr,
-                src_dtype,
-                f32,
-                F32
-            ),
-            DType::F64 => cast_launch!(
-                dev,
-                grid,
-                block,
-                el,
-                num_dims,
-                ds,
-                src_ptr,
-                src_dtype,
-                f64,
-                F64
-            ),
+            DType::U8 => {
+                cast_launch!(dev, grid, block, el, num_dims, ds, src_ptr, src_dtype, u8, U8)
+            }
+            DType::U32 => {
+                cast_launch!(dev, grid, block, el, num_dims, ds, src_ptr, src_dtype, u32, U32)
+            }
+            DType::I64 => {
+                cast_launch!(dev, grid, block, el, num_dims, ds, src_ptr, src_dtype, i64, I64)
+            }
+            DType::BF16 => {
+                cast_launch!(dev, grid, block, el, num_dims, ds, src_ptr, src_dtype, bf16, BF16)
+            }
+            DType::F16 => {
+                cast_launch!(dev, grid, block, el, num_dims, ds, src_ptr, src_dtype, f16, F16)
+            }
+            DType::F32 => {
+                cast_launch!(dev, grid, block, el, num_dims, ds, src_ptr, src_dtype, f32, F32)
+            }
+            DType::F64 => {
+                cast_launch!(dev, grid, block, el, num_dims, ds, src_ptr, src_dtype, f64, F64)
+            }
             DType::I16 | DType::I32 => {
                 return Err(crate::Error::Msg(
                     "i16/i32 dtypes are not supported for to_dtype on ROCm".to_string(),
@@ -3942,8 +4008,7 @@ impl BackendStorage for RocmStorage {
     fn where_cond(&self, l: &Layout, a: &Self, la: &Layout, b: &Self, lb: &Layout) -> Result<Self> {
         let device = self.device.clone();
         let numel = l.shape().elem_count();
-        let (ds, num_dims) =
-            DimsStrides::build(l.dims(), &[l.stride(), la.stride(), lb.stride()])?;
+        let (ds, num_dims) = DimsStrides::build(l.dims(), &[l.stride(), la.stride(), lb.stride()])?;
         let (cond_prefix, cond_ptr) = match &self.slice {
             RocmStorageSlice::U8(s) => ("where_u8", unsafe { s.offset_ptr(l.start_offset()) }
                 as *mut std::ffi::c_void),
@@ -4313,7 +4378,12 @@ impl BackendStorage for RocmStorage {
         // the filled inline payload. Both are byte-identical math.
         let ids_contig = ids_l.is_contiguous();
         let (ds, num_dims) = if ids_contig {
-            (DimsStrides { v: [0; ROCM_DS_MAX * ROCM_DS_SETS] }, 0)
+            (
+                DimsStrides {
+                    v: [0; ROCM_DS_MAX * ROCM_DS_SETS],
+                },
+                0,
+            )
         } else {
             DimsStrides::build(ids_l.shape().dims(), &[ids_l.stride()])?
         };
@@ -4325,15 +4395,24 @@ impl BackendStorage for RocmStorage {
 
         let prefix_base = if ids_contig { "isc" } else { "is" };
         let (ids_prefix, ids_ptr): (String, *mut std::ffi::c_void) = match &idx.slice {
-            RocmStorageSlice::U32(s) => (format!("{prefix_base}_u32"), unsafe {
-                s.offset_ptr(ids_l.start_offset())
-            } as *mut std::ffi::c_void),
-            RocmStorageSlice::U8(s) => (format!("{prefix_base}_u8"), unsafe {
-                s.offset_ptr(ids_l.start_offset())
-            } as *mut std::ffi::c_void),
-            RocmStorageSlice::I64(s) => (format!("{prefix_base}_i64"), unsafe {
-                s.offset_ptr(ids_l.start_offset())
-            } as *mut std::ffi::c_void),
+            RocmStorageSlice::U32(s) => {
+                (
+                    format!("{prefix_base}_u32"),
+                    unsafe { s.offset_ptr(ids_l.start_offset()) } as *mut std::ffi::c_void,
+                )
+            }
+            RocmStorageSlice::U8(s) => {
+                (
+                    format!("{prefix_base}_u8"),
+                    unsafe { s.offset_ptr(ids_l.start_offset()) } as *mut std::ffi::c_void,
+                )
+            }
+            RocmStorageSlice::I64(s) => {
+                (
+                    format!("{prefix_base}_i64"),
+                    unsafe { s.offset_ptr(ids_l.start_offset()) } as *mut std::ffi::c_void,
+                )
+            }
             _ => crate::bail!("index_select ids should be u8, u32, or i64"),
         };
         let ids_prefix = ids_prefix.as_str();

hanzo-ml/src/storage.rs

@@ -5,9 +5,9 @@ use crate::scalar::Scalar;
 use crate::RocmStorage;
 #[cfg(feature = "vulkan")]
 use crate::VulkanStorage;
-use crate::{CpuStorage, CudaStorage, DType, Device, Error, Layout, MetalStorage, Result, Shape};
 #[cfg(feature = "wgpu")]
 use crate::WgpuStorage;
+use crate::{CpuStorage, CudaStorage, DType, Device, Error, Layout, MetalStorage, Result, Shape};
 use crate::{CustomOp1, CustomOp2, CustomOp3, InplaceOp1, InplaceOp2, InplaceOp3};
 
 // We do not want to implement Clone on Storage as cloning may fail because of
@@ -512,9 +512,7 @@ impl Storage {
                 c.vulkan_fwd(s1, l1, s2, l2, s3, l3)
             }
             #[cfg(feature = "wgpu")]
-            (Self::Wgpu(s1), Self::Wgpu(s2), Self::Wgpu(s3)) => {
-                c.wgpu_fwd(s1, l1, s2, l2, s3, l3)
-            }
+            (Self::Wgpu(s1), Self::Wgpu(s2), Self::Wgpu(s3)) => c.wgpu_fwd(s1, l1, s2, l2, s3, l3),
             _ => unreachable!(),
         }
     }
@@ -1286,9 +1284,7 @@ impl Storage {
                 Ok(src.copy_strided_src(dst, dst_offset, src_l)?)
             }
             #[cfg(feature = "wgpu")]
-            (Self::Wgpu(src), Self::Wgpu(dst)) => {
-                Ok(src.copy_strided_src(dst, dst_offset, src_l)?)
-            }
+            (Self::Wgpu(src), Self::Wgpu(dst)) => Ok(src.copy_strided_src(dst, dst_offset, src_l)?),
             (lhs, rhs) => Err(Error::DeviceMismatchBinaryOp {
                 lhs: lhs.device().location(),
                 rhs: rhs.device().location(),

hanzo-ml/src/vulkan_backend.rs

@@ -511,7 +511,8 @@ impl VulkanDevice {
             packed[o] = u16::from_le_bytes([src[0], src[1]]) as u32;
             for j in 0..8 {
                 let b = 2 + j * 4;
-                packed[o + 1 + j] = u32::from_le_bytes([src[b], src[b + 1], src[b + 2], src[b + 3]]);
+                packed[o + 1 + j] =
+                    u32::from_le_bytes([src[b], src[b + 1], src[b + 2], src[b + 3]]);
             }
         }
         self.upload_u32(&packed)
@@ -689,10 +690,17 @@ impl VulkanDevice {
             crate::bail!("flash_attn_gpu: head_dim {d} > 256 (kernel limit)");
         }
         if q.count < bh * lq * d {
-            crate::bail!("flash_attn_gpu: q count {} < bh*lq*d {}", q.count, bh * lq * d);
+            crate::bail!(
+                "flash_attn_gpu: q count {} < bh*lq*d {}",
+                q.count,
+                bh * lq * d
+            );
         }
         if k.count < bh * lk * d || v.count < bh * lk * d {
-            crate::bail!("flash_attn_gpu: k/v count too small for bh*lk*d {}", bh * lk * d);
+            crate::bail!(
+                "flash_attn_gpu: k/v count too small for bh*lk*d {}",
+                bh * lk * d
+            );
         }
         let out = self.alloc_f32(bh * lq * d)?;
         // Push block matches flash_attn.comp: {u32 bh, lq, lk, d; f32 scale; u32 causal}.
@@ -786,7 +794,13 @@ impl VulkanDevice {
     /// Q4_K decode matvec via int8 dp4a (default where int_dot8; VK_DP4A_DECODE_OFF forces scalar):
     /// quantize x to q8_1, then dp4a the Q4_K codes (column dp4a at mcount=1). ~1.8x faster than the scalar subgroup matvec
     /// on gfx1151; the q8_1 activation quant adds ~0.5-1% vs the scalar reference (gated < 2e-2).
-    pub fn matvec_q4k_dp4a(&self, wq: &VulkanStorage, x: &[f32], nout: usize, k: usize) -> Result<Vec<f32>> {
+    pub fn matvec_q4k_dp4a(
+        &self,
+        wq: &VulkanStorage,
+        x: &[f32],
+        nout: usize,
+        k: usize,
+    ) -> Result<Vec<f32>> {
         let xin = self.upload_f32(x)?;
         let (xq, xs, xsum) = self.quantize_act_q8(&xin, 1, k)?;
         let out = self.alloc_f32(nout)?;
@@ -802,7 +816,13 @@ impl VulkanDevice {
 
     /// Q4_K decode matvec via the SCALAR float path (forces the non-dp4a kernel regardless of the
     /// VK_DP4A_DECODE_OFF default): the exact CPU-faithful reference for the dp4a gates.
-    pub fn matvec_q4k_scalar(&self, wq: &VulkanStorage, x: &[f32], nout: usize, k: usize) -> Result<Vec<f32>> {
+    pub fn matvec_q4k_scalar(
+        &self,
+        wq: &VulkanStorage,
+        x: &[f32],
+        nout: usize,
+        k: usize,
+    ) -> Result<Vec<f32>> {
         let xin = self.upload_f32(x)?;
         let out = self.alloc_f32(nout)?;
         let push = push_u32(&[nout as u32, k as u32, 0u32]);
@@ -926,7 +946,12 @@ impl VulkanDevice {
             let (xq, xs, xsum) = self.quantize_act_q8(x, 1, k)?;
             let bufs = [wq.buffer, xq.buffer, xs.buffer, xsum.buffer, out.buffer];
             let pushd = push_u32(&[0u32, 1u32, nout as u32, k as u32, woff as u32]);
-            self.dispatch("mul_mat_q4k_dp4a", &bufs, &pushd, ((nout as u32).div_ceil(WG1D), 1, 1))?;
+            self.dispatch(
+                "mul_mat_q4k_dp4a",
+                &bufs,
+                &pushd,
+                ((nout as u32).div_ceil(WG1D), 1, 1),
+            )?;
             return Ok(out);
         }
         let push = push_u32(&[nout as u32, k as u32, woff as u32]);
@@ -978,7 +1003,11 @@ impl VulkanDevice {
         k: usize,
     ) -> Result<VulkanStorage> {
         if x.count < nrows * k {
-            crate::bail!("moe_matvec_gpu: x count {} < nrows*k {}", x.count, nrows * k);
+            crate::bail!(
+                "moe_matvec_gpu: x count {} < nrows*k {}",
+                x.count,
+                nrows * k
+            );
         }
         if ids.count < nrows {
             crate::bail!("moe_matvec_gpu: ids count {} < nrows {nrows}", ids.count);
@@ -1387,7 +1416,13 @@ impl VulkanDevice {
         let out = self.alloc_f32(bh * seq_len * v_dim)?;
         let push = push_u32(&[seq_len as u32, k_dim as u32, v_dim as u32]);
         let bufs = [
-            q.buffer, k.buffer, v.buffer, g.buffer, beta.buffer, state.buffer, out.buffer,
+            q.buffer,
+            k.buffer,
+            v.buffer,
+            g.buffer,
+            beta.buffer,
+            state.buffer,
+            out.buffer,
         ];
         let v_tiles = (v_dim as u32).div_ceil(WG1D);
         // Chunked prefill kernel bounds its shared key array at K<=128; fall back to the sequential
@@ -1518,7 +1553,10 @@ impl VulkanDevice {
     #[allow(clippy::too_many_arguments)]
     pub fn paged_attention_vk(&self, p: &PagedAttnArgs<'_>) -> Result<VulkanStorage> {
         if p.head_size as u64 > 256 {
-            crate::bail!("vulkan paged_attn: head_size {} > 256 (shared mem bound)", p.head_size);
+            crate::bail!(
+                "vulkan paged_attn: head_size {} > 256 (shared mem bound)",
+                p.head_size
+            );
         }
         let out = self.alloc_f32(p.num_seqs * p.num_heads * p.head_size)?;
         // push: 10 u32 then 1 f32 (scale). std430 scalar packing, tightly packed.
@@ -1736,9 +1774,8 @@ impl VulkanDevice {
             let src = &data[blk * 210..blk * 210 + 210];
             let dst = &mut words[blk * 53..blk * 53 + 53];
             // Copy the 210 bytes into the low 210 bytes of the 212-byte (53 u32) padded block.
-            let dst_bytes: &mut [u8] = unsafe {
-                std::slice::from_raw_parts_mut(dst.as_mut_ptr() as *mut u8, 53 * 4)
-            };
+            let dst_bytes: &mut [u8] =
+                unsafe { std::slice::from_raw_parts_mut(dst.as_mut_ptr() as *mut u8, 53 * 4) };
             dst_bytes[..210].copy_from_slice(src);
         }
         self.upload_u32(&words)
@@ -1986,7 +2023,10 @@ impl VulkanDevice {
         let total = nout * nblocks;
         let want = total * 66;
         if data.len() != want {
-            crate::bail!("quantize_iq2xxs: data len {} != {nout}*{nblocks}*66 = {want}", data.len());
+            crate::bail!(
+                "quantize_iq2xxs: data len {} != {nout}*{nblocks}*66 = {want}",
+                data.len()
+            );
         }
         let mut words = vec![0u32; total * 17];
         for blk in 0..total {
@@ -2050,7 +2090,10 @@ impl VulkanDevice {
         let total = nout * nblocks;
         let want = total * 74;
         if data.len() != want {
-            crate::bail!("quantize_iq2xs: data len {} != {nout}*{nblocks}*74 = {want}", data.len());
+            crate::bail!(
+                "quantize_iq2xs: data len {} != {nout}*{nblocks}*74 = {want}",
+                data.len()
+            );
         }
         let mut words = vec![0u32; total * 19];
         for blk in 0..total {
@@ -2113,7 +2156,10 @@ impl VulkanDevice {
         let total = nout * nblocks;
         let want = total * 56;
         if data.len() != want {
-            crate::bail!("quantize_iq1m: data len {} != {nout}*{nblocks}*56 = {want}", data.len());
+            crate::bail!(
+                "quantize_iq1m: data len {} != {nout}*{nblocks}*56 = {want}",
+                data.len()
+            );
         }
         let mut words = vec![0u32; total * 14];
         for blk in 0..total {
@@ -2176,7 +2222,10 @@ impl VulkanDevice {
         let total = nout * nblocks;
         let want = total * 50;
         if data.len() != want {
-            crate::bail!("quantize_iq1s: data len {} != {nout}*{nblocks}*50 = {want}", data.len());
+            crate::bail!(
+                "quantize_iq1s: data len {} != {nout}*{nblocks}*50 = {want}",
+                data.len()
+            );
         }
         let mut words = vec![0u32; total * 13];
         for blk in 0..total {
@@ -2239,7 +2288,10 @@ impl VulkanDevice {
         let total = nout * nblocks;
         let want = total * 110;
         if data.len() != want {
-            crate::bail!("quantize_iq3s: data len {} != {nout}*{nblocks}*110 = {want}", data.len());
+            crate::bail!(
+                "quantize_iq3s: data len {} != {nout}*{nblocks}*110 = {want}",
+                data.len()
+            );
         }
         let mut words = vec![0u32; total * 28];
         for blk in 0..total {
@@ -2302,7 +2354,10 @@ impl VulkanDevice {
         let total = nout * nblocks;
         let want = total * 98;
         if data.len() != want {
-            crate::bail!("quantize_iq3xxs: data len {} != {nout}*{nblocks}*98 = {want}", data.len());
+            crate::bail!(
+                "quantize_iq3xxs: data len {} != {nout}*{nblocks}*98 = {want}",
+                data.len()
+            );
         }
         let mut words = vec![0u32; total * 25];
         for blk in 0..total {
@@ -2365,7 +2420,10 @@ impl VulkanDevice {
         let total = nout * nblocks;
         let want = total * 82;
         if data.len() != want {
-            crate::bail!("quantize_iq2s: data len {} != {nout}*{nblocks}*82 = {want}", data.len());
+            crate::bail!(
+                "quantize_iq2s: data len {} != {nout}*{nblocks}*82 = {want}",
+                data.len()
+            );
         }
         let mut words = vec![0u32; total * 21];
         for blk in 0..total {
@@ -2712,10 +2770,7 @@ impl VulkanDevice {
     // Allocate a buffer that is guaranteed HOST_VISIBLE (for staging). Forces the host-visible
     // policy regardless of the device strategy; staging buffers are transient and always small
     // relative to the host-visible heap (one tensor's worth at a time). Not pooled.
-    unsafe fn raw_buffer_host_visible(
-        &self,
-        bytes: u64,
-    ) -> Result<(vk::Buffer, vk::DeviceMemory)> {
+    unsafe fn raw_buffer_host_visible(&self, bytes: u64) -> Result<(vk::Buffer, vk::DeviceMemory)> {
         let bytes = bytes.max(4);
         let dev = self.dev();
         let info = vk::BufferCreateInfo::default()
@@ -3154,9 +3209,7 @@ impl VulkanDevice {
             if reads_inflight_write {
                 let bar = [vk::MemoryBarrier::default()
                     .src_access_mask(vk::AccessFlags::SHADER_WRITE)
-                    .dst_access_mask(
-                        vk::AccessFlags::SHADER_READ | vk::AccessFlags::SHADER_WRITE,
-                    )];
+                    .dst_access_mask(vk::AccessFlags::SHADER_READ | vk::AccessFlags::SHADER_WRITE)];
                 dev.cmd_pipeline_barrier(
                     s.cmd,
                     vk::PipelineStageFlags::COMPUTE_SHADER,
@@ -3448,7 +3501,12 @@ const BATCH_CAP: u32 = 4096;
 // when nothing is recorded. Caller must hold the submitter lock. When `profile` is set, prints the
 // per-batch phase breakdown (recording / submit / fence-wait time, dispatch + barrier counts) so
 // the real GPU shows where per-token milliseconds go; the timers are only sampled when profiling.
-fn flush_locked(dev: &ash::Device, queue: vk::Queue, s: &mut Submitter, profile: bool) -> Result<()> {
+fn flush_locked(
+    dev: &ash::Device,
+    queue: vk::Queue,
+    s: &mut Submitter,
+    profile: bool,
+) -> Result<()> {
     if !s.recording {
         return Ok(());
     }
@@ -3681,7 +3739,9 @@ impl BackendDevice for VulkanDevice {
                 instance.get_physical_device_features2(pdev, &mut f2);
             }
             let int_dot8 = idot_feat.shader_integer_dot_product != 0
-                && std::env::var("VK_INT_DOT").map(|v| v != "0").unwrap_or(true);
+                && std::env::var("VK_INT_DOT")
+                    .map(|v| v != "0")
+                    .unwrap_or(true);
 
             // Build the enabled-extension list dynamically: coopmat and push_descriptor are
             // independent and either may be present. push_descriptor needs no extra device feature
@@ -3719,8 +3779,8 @@ impl BackendDevice for VulkanDevice {
             }
             let device = instance.create_device(pdev, &dci, None).map_err(vkerr)?;
             // Load push_descriptor device fns now that the device exists with the extension enabled.
-            let push_descriptor = use_pd
-                .then(|| ash::khr::push_descriptor::Device::new(&instance, &device));
+            let push_descriptor =
+                use_pd.then(|| ash::khr::push_descriptor::Device::new(&instance, &device));
             let queue = device.get_device_queue(qfi, 0);
             let mem_props = instance.get_physical_device_memory_properties(pdev);
 
@@ -4277,10 +4337,14 @@ impl VulkanStorage {
         v_dim: usize,
     ) -> Result<VulkanStorage> {
         if k_dim > GDN_STEP_MAX_K {
-            crate::bail!("vulkan: gdn_step head_k_dim {k_dim} exceeds GDN_STEP_MAX_K {GDN_STEP_MAX_K}");
+            crate::bail!(
+                "vulkan: gdn_step head_k_dim {k_dim} exceeds GDN_STEP_MAX_K {GDN_STEP_MAX_K}"
+            );
         }
         if !(state_l.is_contiguous() && state_l.start_offset() == 0) {
-            crate::bail!("vulkan: gdn_step state must be contiguous and offset 0 (it is updated in place)");
+            crate::bail!(
+                "vulkan: gdn_step state must be contiguous and offset 0 (it is updated in place)"
+            );
         }
         let mut qk = None;
         let mut kk = None;
@@ -4323,7 +4387,9 @@ impl VulkanStorage {
         k_size: usize,
     ) -> Result<VulkanStorage> {
         if k_size > GDN_CONV_MAX_K {
-            crate::bail!("vulkan: gdn_conv1d_step kernel {k_size} exceeds GDN_CONV_MAX_K {GDN_CONV_MAX_K}");
+            crate::bail!(
+                "vulkan: gdn_conv1d_step kernel {k_size} exceeds GDN_CONV_MAX_K {GDN_CONV_MAX_K}"
+            );
         }
         if !(state_l.is_contiguous() && state_l.start_offset() == 0) {
             crate::bail!("vulkan: gdn_conv1d_step conv_state must be contiguous and offset 0 (updated in place)");
@@ -4498,7 +4564,9 @@ impl BackendStorage for VulkanStorage {
             perm.push(d);
             (self.contiguous(&layout.permute(&perm)?)?, dims[d])
         } else {
-            crate::bail!("vulkan: reduce over multiple axes at once not supported (got {sum_dims:?})");
+            crate::bail!(
+                "vulkan: reduce over multiple axes at once not supported (got {sum_dims:?})"
+            );
         };
         let rows: usize = layout.shape().elem_count() / cols;
         // arg-reductions return u32 indices; value reductions return f32.
@@ -5129,7 +5197,8 @@ impl BackendStorage for VulkanStorage {
         // it: if the two scratch buffers won't fit in the largest usable heap's *free* bytes (plus a
         // margin), fall through to the fp32 tiled GEMM, which needs no extra f16 buffers — correct
         // result, just slower, instead of an OOM abort.
-        let scratch_bytes = ((b * m * k * 2).max(4) as u64).saturating_add((b * k * n * 2).max(4) as u64);
+        let scratch_bytes =
+            ((b * m * k * 2).max(4) as u64).saturating_add((b * k * n * 2).max(4) as u64);
         if self.device.inner.cm_use
             && self.device.scratch_fits(scratch_bytes)
             && matches!(self.device.coopmat_info(), Some((16, 16, 16)))
@@ -5195,7 +5264,12 @@ impl BackendStorage for VulkanStorage {
             crate::bail!("vulkan: copy_strided_src supports rank <= 6, got {rank}");
         }
         let strides = src_l.stride();
-        let mut p = vec![n as u32, rank as u32, src_l.start_offset() as u32, dst_offset as u32];
+        let mut p = vec![
+            n as u32,
+            rank as u32,
+            src_l.start_offset() as u32,
+            dst_offset as u32,
+        ];
         let mut shape6 = [0u32; 6];
         let mut stride6 = [0u32; 6];
         for d in 0..rank {
@@ -5343,8 +5417,13 @@ mod dsl_dispatch_proof {
         let ws = dev.upload_f32(&w).unwrap();
         let out = dev.alloc_f32(N).unwrap();
         let groups = (N as u32).div_ceil(WG);
-        dev.dispatch("dsl_mul", &[xs.buffer, ws.buffer, out.buffer], &[], (groups, 1, 1))
-            .unwrap();
+        dev.dispatch(
+            "dsl_mul",
+            &[xs.buffer, ws.buffer, out.buffer],
+            &[],
+            (groups, 1, 1),
+        )
+        .unwrap();
         dev.flush().unwrap();
         let got = out.to_vec_f32().unwrap();
 
@@ -5354,8 +5433,15 @@ mod dsl_dispatch_proof {
             .map(|(a, b)| (a - b).abs())
             .fold(0.0f32, f32::max);
         eprintln!("[dsl-proof] DSL kernel via ml::VulkanDevice::dispatch  N={N} groups={groups}  maxerr={maxerr:.2e}");
-        eprintln!("[dsl-proof] first4 got={:?} ref={:?}", &got[..4], &refout[..4]);
-        assert_eq!(maxerr, 0.0, "DSL-generated kernel not bit-exact through ml dispatch");
+        eprintln!(
+            "[dsl-proof] first4 got={:?} ref={:?}",
+            &got[..4],
+            &refout[..4]
+        );
+        assert_eq!(
+            maxerr, 0.0,
+            "DSL-generated kernel not bit-exact through ml dispatch"
+        );
     }
 
     // Generalization: a REDUCTION kernel (matvec, comptime k=32/rows=64) auto-processed by the
@@ -5374,7 +5460,9 @@ mod dsl_dispatch_proof {
                 return;
             }
         };
-        let w: Vec<f32> = (0..ROWS * K).map(|i| (i as f32 % 7.0) * 0.25 - 1.0).collect();
+        let w: Vec<f32> = (0..ROWS * K)
+            .map(|i| (i as f32 % 7.0) * 0.25 - 1.0)
+            .collect();
         let x: Vec<f32> = (0..K).map(|i| (i as f32 % 5.0) * 0.5 - 1.0).collect();
         let refout: Vec<f32> = (0..ROWS)
             .map(|r| {
@@ -5389,15 +5477,31 @@ mod dsl_dispatch_proof {
         let ws = dev.upload_f32(&w).unwrap();
         let xs = dev.upload_f32(&x).unwrap();
         let out = dev.alloc_f32(ROWS).unwrap();
-        dev.dispatch("dsl_matvec", &[ws.buffer, xs.buffer, out.buffer], &[], (1, 1, 1))
-            .unwrap();
+        dev.dispatch(
+            "dsl_matvec",
+            &[ws.buffer, xs.buffer, out.buffer],
+            &[],
+            (1, 1, 1),
+        )
+        .unwrap();
         dev.flush().unwrap();
         let got = out.to_vec_f32().unwrap();
 
-        let maxerr = got.iter().zip(&refout).map(|(a, b)| (a - b).abs()).fold(0.0f32, f32::max);
+        let maxerr = got
+            .iter()
+            .zip(&refout)
+            .map(|(a, b)| (a - b).abs())
+            .fold(0.0f32, f32::max);
         eprintln!("[dsl-proof] DSL matvec via ml::VulkanDevice::dispatch  rows={ROWS} k={K}  maxerr={maxerr:.2e}");
-        eprintln!("[dsl-proof] first4 got={:?} ref={:?}", &got[..4], &refout[..4]);
-        assert!(maxerr < 1e-4, "DSL matvec through ml diverged: maxerr={maxerr:.3e}");
+        eprintln!(
+            "[dsl-proof] first4 got={:?} ref={:?}",
+            &got[..4],
+            &refout[..4]
+        );
+        assert!(
+            maxerr < 1e-4,
+            "DSL matvec through ml diverged: maxerr={maxerr:.3e}"
+        );
     }
 
     // The production `rms_norm` method now dispatches the DSL block-per-row kernel (rms_norm_blk.spv):
@@ -5445,11 +5549,18 @@ mod dsl_dispatch_proof {
         let mut worst = 0f32;
         for out in &outs {
             let got = out.to_vec_f32().unwrap();
-            let err = got.iter().zip(&cpu).map(|(a, b)| (a - b).abs()).fold(0f32, f32::max);
+            let err = got
+                .iter()
+                .zip(&cpu)
+                .map(|(a, b)| (a - b).abs())
+                .fold(0f32, f32::max);
             worst = worst.max(err);
         }
         eprintln!("[rms_norm-prod] {ROWS}x{N} x8-batched  maxerr={worst:.2e}  (DSL kernel via production path, pooled scalars, no sync crutch)");
-        assert!(worst < 1e-3, "production rms_norm (DSL kernel) diverged from CPU: {worst:.3e}");
+        assert!(
+            worst < 1e-3,
+            "production rms_norm (DSL kernel) diverged from CPU: {worst:.3e}"
+        );
 
         // Clean kernel-only bench: FIXED buffers (no per-iter 64MB alloc/upload) dispatched in a loop
         // with one fence -- the true kernel cost, not the method's allocation overhead. Same 5-SSBO /
@@ -5512,11 +5623,18 @@ mod dsl_dispatch_proof {
         let mut worst = 0f32;
         for out in &outs {
             let got = out.to_vec_f32().unwrap();
-            let err = got.iter().zip(&cpu).map(|(a, b)| (a - b).abs()).fold(0f32, f32::max);
+            let err = got
+                .iter()
+                .zip(&cpu)
+                .map(|(a, b)| (a - b).abs())
+                .fold(0f32, f32::max);
             worst = worst.max(err);
         }
         eprintln!("[softmax-prod] {ROWS}x{M} x8-batched  maxerr={worst:.2e}  (DSL kernel via production path, pooled n, no sync crutch)");
-        assert!(worst < 1e-4, "production softmax (DSL kernel) diverged from CPU: {worst:.3e}");
+        assert!(
+            worst < 1e-4,
+            "production softmax (DSL kernel) diverged from CPU: {worst:.3e}"
+        );
 
         let iters = 50;
         let bytes = (2 * ROWS * M * 4) as f64; // read x + write out
@@ -5524,11 +5642,13 @@ mod dsl_dispatch_proof {
         let ndb = dev.upload_u32(&[M as u32]).unwrap();
         let bufs = [xs.buffer, out.buffer, ndb.buffer];
         let grid = (ROWS as u32, 1, 1);
-        dev.dispatch_out("softmax_rows", &bufs, 1, &[], grid).unwrap();
+        dev.dispatch_out("softmax_rows", &bufs, 1, &[], grid)
+            .unwrap();
         dev.synchronize().unwrap();
         let t = std::time::Instant::now();
         for _ in 0..iters {
-            dev.dispatch_out("softmax_rows", &bufs, 1, &[], grid).unwrap();
+            dev.dispatch_out("softmax_rows", &bufs, 1, &[], grid)
+                .unwrap();
         }
         dev.synchronize().unwrap();
         let ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;

hanzo-ml/src/wgpu_backend.rs

@@ -140,7 +140,11 @@ pub struct WgpuDevice {
 
 impl std::fmt::Debug for WgpuDevice {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "WgpuDevice({}, {})", self.inner.gpu_id, self.inner.adapter_desc)
+        write!(
+            f,
+            "WgpuDevice({}, {})",
+            self.inner.gpu_id, self.inner.adapter_desc
+        )
     }
 }
 
@@ -155,7 +159,11 @@ pub struct WgpuStorage {
 
 impl std::fmt::Debug for WgpuStorage {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "WgpuStorage(count={}, dtype={:?})", self.count, self.dtype)
+        write!(
+            f,
+            "WgpuStorage(count={}, dtype={:?})",
+            self.count, self.dtype
+        )
     }
 }
 
@@ -202,14 +210,16 @@ impl WgpuDevice {
             label: Some(name),
             source: wgpu::ShaderSource::Wgsl(src.into()),
         });
-        let pipeline = Arc::new(dev.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
-            label: Some(name),
-            layout: None,
-            module: &module,
-            entry_point: "main",
-            compilation_options: wgpu::PipelineCompilationOptions::default(),
-            cache: None,
-        }));
+        let pipeline = Arc::new(
+            dev.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
+                label: Some(name),
+                layout: None,
+                module: &module,
+                entry_point: "main",
+                compilation_options: wgpu::PipelineCompilationOptions::default(),
+                cache: None,
+            }),
+        );
         self.inner
             .pipelines
             .lock()
@@ -256,7 +266,8 @@ impl WgpuDevice {
             layout: &bgl,
             entries: &entries,
         });
-        let mut enc = dev.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some(name) });
+        let mut enc =
+            dev.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some(name) });
         {
             let mut cpass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor {
                 label: Some(name),
@@ -292,13 +303,15 @@ impl WgpuDevice {
         let nwords = init.len().div_ceil(4).max(1);
         let mut words = vec![0u8; nwords * 4];
         words[..init.len()].copy_from_slice(init);
-        let buffer = self.dev().create_buffer_init(&wgpu::util::BufferInitDescriptor {
-            label: Some("storage_init"),
-            contents: &words,
-            usage: wgpu::BufferUsages::STORAGE
-                | wgpu::BufferUsages::COPY_SRC
-                | wgpu::BufferUsages::COPY_DST,
-        });
+        let buffer = self
+            .dev()
+            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
+                label: Some("storage_init"),
+                contents: &words,
+                usage: wgpu::BufferUsages::STORAGE
+                    | wgpu::BufferUsages::COPY_SRC
+                    | wgpu::BufferUsages::COPY_DST,
+            });
         Arc::new(buffer)
     }
 
@@ -405,7 +418,13 @@ impl WgpuDevice {
     }
 
     /// Host-input convenience: uploads `x` then runs [`matvec_q4_0_gpu`]. Returns the f32 result.
-    pub fn matvec_q4_0(&self, wq: &WgpuStorage, x: &[f32], nout: usize, k: usize) -> Result<Vec<f32>> {
+    pub fn matvec_q4_0(
+        &self,
+        wq: &WgpuStorage,
+        x: &[f32],
+        nout: usize,
+        k: usize,
+    ) -> Result<Vec<f32>> {
         if x.len() != k {
             crate::bail!("matvec_q4_0: x len {} != k {k}", x.len());
         }
@@ -414,7 +433,13 @@ impl WgpuDevice {
     }
 
     /// Host-input convenience: uploads `x` then runs [`matvec_q8_0_gpu`]. Returns the f32 result.
-    pub fn matvec_q8_0(&self, wq: &WgpuStorage, x: &[f32], nout: usize, k: usize) -> Result<Vec<f32>> {
+    pub fn matvec_q8_0(
+        &self,
+        wq: &WgpuStorage,
+        x: &[f32],
+        nout: usize,
+        k: usize,
+    ) -> Result<Vec<f32>> {
         if x.len() != k {
             crate::bail!("matvec_q8_0: x len {} != k {k}", x.len());
         }
@@ -448,7 +473,13 @@ impl WgpuDevice {
     }
 
     /// Host-input convenience: uploads `x` then runs [`matvec_q4k_gpu`]. Returns the f32 result.
-    pub fn matvec_q4k(&self, wq: &WgpuStorage, x: &[f32], nout: usize, k: usize) -> Result<Vec<f32>> {
+    pub fn matvec_q4k(
+        &self,
+        wq: &WgpuStorage,
+        x: &[f32],
+        nout: usize,
+        k: usize,
+    ) -> Result<Vec<f32>> {
         if x.len() != k {
             crate::bail!("matvec_q4k: x len {} != k {k}", x.len());
         }
@@ -478,7 +509,11 @@ impl WgpuDevice {
         k: usize,
     ) -> Result<WgpuStorage> {
         if x.count < nrows * k {
-            crate::bail!("moe_matvec_gpu: x count {} < nrows*k {}", x.count, nrows * k);
+            crate::bail!(
+                "moe_matvec_gpu: x count {} < nrows*k {}",
+                x.count,
+                nrows * k
+            );
         }
         if ids.count < nrows {
             crate::bail!("moe_matvec_gpu: ids count {} < nrows {nrows}", ids.count);
@@ -515,10 +550,17 @@ impl WgpuDevice {
             crate::bail!("flash_attn_gpu: head_dim {d} > 256 (kernel limit)");
         }
         if q.count < bh * lq * d {
-            crate::bail!("flash_attn_gpu: q count {} < bh*lq*d {}", q.count, bh * lq * d);
+            crate::bail!(
+                "flash_attn_gpu: q count {} < bh*lq*d {}",
+                q.count,
+                bh * lq * d
+            );
         }
         if k.count < bh * lk * d || v.count < bh * lk * d {
-            crate::bail!("flash_attn_gpu: k/v count too small for bh*lk*d {}", bh * lk * d);
+            crate::bail!(
+                "flash_attn_gpu: k/v count too small for bh*lk*d {}",
+                bh * lk * d
+            );
         }
         let out = self.alloc_f32(bh * lq * d)?;
         // Uniform block matches flash_attn.wgsl: {u32 bh, lq, lk, d; f32 scale; u32 causal}.
@@ -572,8 +614,9 @@ impl WgpuDevice {
             usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
             mapped_at_creation: false,
         });
-        let mut enc =
-            dev.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("readback") });
+        let mut enc = dev.create_command_encoder(&wgpu::CommandEncoderDescriptor {
+            label: Some("readback"),
+        });
         enc.copy_buffer_to_buffer(src, 0, &staging, 0, bytes);
         self.queue().submit(std::iter::once(enc.finish()));
 
@@ -682,7 +725,11 @@ impl WgpuStorage {
     // Buffer for a contiguous, offset-0 f32 view of `layout`: no copy when the storage already is
     // one (returns its own buffer), otherwise materializes a packed copy returned via `keep` (held
     // alive by the caller until the dispatch is recorded). Mirrors VulkanStorage::contig_buf.
-    fn contig_buf(&self, layout: &Layout, keep: &mut Option<WgpuStorage>) -> Result<Arc<wgpu::Buffer>> {
+    fn contig_buf(
+        &self,
+        layout: &Layout,
+        keep: &mut Option<WgpuStorage>,
+    ) -> Result<Arc<wgpu::Buffer>> {
         if layout.is_contiguous() && layout.start_offset() == 0 {
             Ok(self.buffer.clone())
         } else {
@@ -941,10 +988,17 @@ impl BackendDevice for WgpuDevice {
         match s {
             CpuStorage::F32(v) => self.upload_f32(v),
             CpuStorage::U32(v) => self.upload_u32(v),
-            CpuStorage::F16(v) => self.upload_f32(&v.iter().map(|x| x.to_f32()).collect::<Vec<_>>()),
-            CpuStorage::BF16(v) => self.upload_f32(&v.iter().map(|x| x.to_f32()).collect::<Vec<_>>()),
+            CpuStorage::F16(v) => {
+                self.upload_f32(&v.iter().map(|x| x.to_f32()).collect::<Vec<_>>())
+            }
+            CpuStorage::BF16(v) => {
+                self.upload_f32(&v.iter().map(|x| x.to_f32()).collect::<Vec<_>>())
+            }
             CpuStorage::U8(v) => self.upload_u32(&v.iter().map(|&x| x as u32).collect::<Vec<_>>()),
-            _ => crate::bail!("wgpu: only f32/u32/f16/bf16/u8 supported, got {:?}", s.dtype()),
+            _ => crate::bail!(
+                "wgpu: only f32/u32/f16/bf16/u8 supported, got {:?}",
+                s.dtype()
+            ),
         }
     }
 
@@ -952,7 +1006,13 @@ impl BackendDevice for WgpuDevice {
         self.storage_from_cpu_storage(&s)
     }
 
-    fn rand_uniform(&self, shape: &Shape, dtype: DType, min: f64, max: f64) -> Result<Self::Storage> {
+    fn rand_uniform(
+        &self,
+        shape: &Shape,
+        dtype: DType,
+        min: f64,
+        max: f64,
+    ) -> Result<Self::Storage> {
         if dtype != DType::F32 {
             crate::bail!("wgpu: rand_uniform only f32, got {dtype:?}");
         }
@@ -967,7 +1027,13 @@ impl BackendDevice for WgpuDevice {
         self.upload_f32(&data)
     }
 
-    fn rand_normal(&self, shape: &Shape, dtype: DType, mean: f64, std: f64) -> Result<Self::Storage> {
+    fn rand_normal(
+        &self,
+        shape: &Shape,
+        dtype: DType,
+        mean: f64,
+        std: f64,
+    ) -> Result<Self::Storage> {
         if dtype != DType::F32 {
             crate::bail!("wgpu: rand_normal only f32, got {dtype:?}");
         }
@@ -1195,12 +1261,21 @@ impl BackendStorage for WgpuStorage {
         let cb = self.contig_buf(layout, &mut ck)?;
         let n = layout.shape().elem_count();
         let out = self.device.alloc_f32(n)?;
-        self.device
-            .dispatch(kernel, &[&cb, &out.buffer], &pack_u32(&[n as u32]), Self::groups_1d(n))?;
+        self.device.dispatch(
+            kernel,
+            &[&cb, &out.buffer],
+            &pack_u32(&[n as u32]),
+            Self::groups_1d(n),
+        )?;
         Ok(out)
     }
 
-    fn binary_impl<B: BinaryOpT>(&self, rhs: &Self, lhs_l: &Layout, rhs_l: &Layout) -> Result<Self> {
+    fn binary_impl<B: BinaryOpT>(
+        &self,
+        rhs: &Self,
+        lhs_l: &Layout,
+        rhs_l: &Layout,
+    ) -> Result<Self> {
         let kernel: &'static str = match B::NAME {
             "add" => "add",
             "sub" => "sub",
@@ -1224,8 +1299,12 @@ impl BackendStorage for WgpuStorage {
         let rb = rhs.contig_buf(rhs_l, &mut rk)?;
         let n = lhs_l.shape().elem_count();
         let out = self.device.alloc_f32(n)?;
-        self.device
-            .dispatch(kernel, &[&lb, &rb, &out.buffer], &pack_u32(&[n as u32]), Self::groups_1d(n))?;
+        self.device.dispatch(
+            kernel,
+            &[&lb, &rb, &out.buffer],
+            &pack_u32(&[n as u32]),
+            Self::groups_1d(n),
+        )?;
         Ok(out)
     }
 
@@ -1239,7 +1318,10 @@ impl BackendStorage for WgpuStorage {
     ) -> Result<Self> {
         // self is the condition mask (u32). The kernel reads it linearly.
         if self.dtype != DType::U32 {
-            crate::bail!("wgpu: where_cond requires u32 condition, got {:?}", self.dtype);
+            crate::bail!(
+                "wgpu: where_cond requires u32 condition, got {:?}",
+                self.dtype
+            );
         }
         let mut condk = None;
         let condb = if l.is_contiguous() && l.start_offset() == 0 {
@@ -1463,7 +1545,12 @@ impl BackendStorage for WgpuStorage {
             crate::bail!("wgpu: copy_strided_src supports rank <= 6, got {rank}");
         }
         let strides = src_l.stride();
-        let mut p = vec![n as u32, rank as u32, src_l.start_offset() as u32, dst_offset as u32];
+        let mut p = vec![
+            n as u32,
+            rank as u32,
+            src_l.start_offset() as u32,
+            dst_offset as u32,
+        ];
         let mut shape6 = [0u32; 6];
         let mut stride6 = [0u32; 6];
         for d in 0..rank {
@@ -1502,8 +1589,12 @@ impl BackendStorage for WgpuStorage {
             src_offset as u32,
             dst_offset as u32,
         ]);
-        self.device
-            .dispatch("copy2d", &[&self.buffer, &dst.buffer], &p, Self::groups_1d(total))
+        self.device.dispatch(
+            "copy2d",
+            &[&self.buffer, &dst.buffer],
+            &p,
+            Self::groups_1d(total),
+        )
     }
 
     fn const_set(&mut self, s: crate::scalar::Scalar, layout: &Layout) -> Result<()> {

hanzo-ml/tests/cuda_iquant_bench.rs

@@ -37,14 +37,22 @@ fn synth(dtype: GgmlDType, nout: usize, k: usize) -> Vec<u8> {
 }
 
 // m == 1 -> decode (mmvq dp4a); m > 8 -> prefill (qmmq int8-WMMA GEMM).
-fn bench_shape(dev: &Device, dtype: GgmlDType, nout: usize, k: usize, m: usize) -> hanzo_ml::Result<()> {
+fn bench_shape(
+    dev: &Device,
+    dtype: GgmlDType,
+    nout: usize,
+    k: usize,
+    m: usize,
+) -> hanzo_ml::Result<()> {
     let raw = synth(dtype, nout, k);
     let bytes = raw.len();
     let qs = QStorage::from_data(Cow::Owned(raw), dev, dtype)?;
     let q = QTensor::new(qs, (nout, k))?;
     let matmul = QMatMul::from_qtensor(q)?;
     let x = Tensor::from_vec(
-        (0..m * k).map(|i| (pbyte(i) as f32 / 128.0) - 1.0).collect(),
+        (0..m * k)
+            .map(|i| (pbyte(i) as f32 / 128.0) - 1.0)
+            .collect(),
         (m, k),
         dev,
     )?;
@@ -64,9 +72,7 @@ fn bench_shape(dev: &Device, dtype: GgmlDType, nout: usize, k: usize, m: usize)
     let us = t0.elapsed().as_secs_f64() * 1e6 / iters as f64;
     let path = if m == 1 { "native-dp4a" } else { "native-qmmq" };
     let kind = if m == 1 { "decode" } else { "prefill" };
-    println!(
-        "[{path:>16}] {kind} {dtype:?} [n={nout:>5} k={k:>5} m={m:>4}]  {us:9.2} us/call",
-    );
+    println!("[{path:>16}] {kind} {dtype:?} [n={nout:>5} k={k:>5} m={m:>4}]  {us:9.2} us/call",);
     Ok(())
 }
 
@@ -83,7 +89,11 @@ fn bench_iq2xxs_decode() {
     // Decode (m=1, mmvq) + prefill (m=128, qmmq) on realistic proj shapes (attn k=2048; FFN k=4096),
     // for IQ2_XXS (grid codebook) and IQ4_XS (LUT codebook -- the dominant type in UD-IQ2_M MoE quants).
     for dtype in [GgmlDType::IQ2_XXS, GgmlDType::IQ4_XS] {
-        for &(n, k, m) in &[(4096usize, 4096usize, 1usize), (8192, 4096, 1), (4096, 4096, 128)] {
+        for &(n, k, m) in &[
+            (4096usize, 4096usize, 1usize),
+            (8192, 4096, 1),
+            (4096, 4096, 128),
+        ] {
             bench_shape(&dev, dtype, n, k, m).unwrap();
         }
     }

hanzo-ml/tests/cuda_iquant_tests.rs

@@ -199,7 +199,8 @@ fn run_case(dev: &Device, dtype: GgmlDType, nout: usize, k: usize) -> hanzo_ml::
     let qs_cuda = QStorage::from_data(Cow::Owned(raw), dev, dtype)?;
     let q_cuda = QTensor::new(qs_cuda, (nout, k))?;
     let matmul = QMatMul::from_qtensor(q_cuda)?;
-    let x_host: Vec<f32> = q8_1_roundtrip(&(0..k).map(|i| pseudo(i + 1_000_003)).collect::<Vec<_>>());
+    let x_host: Vec<f32> =
+        q8_1_roundtrip(&(0..k).map(|i| pseudo(i + 1_000_003)).collect::<Vec<_>>());
     let x = Tensor::from_vec(x_host.clone(), (1, k), dev)?;
     let y_cuda: Vec<f32> = matmul.forward(&x)?.flatten_all()?.to_vec1::<f32>()?;
     assert_eq!(y_cuda.len(), nout);
@@ -240,7 +241,10 @@ fn check(dtype: GgmlDType) {
     // A couple of decode shapes (attn-proj k=2048, ffn k=4096; non-square nout).
     for &(nout, k) in &[(512usize, 2048usize), (768, 4096)] {
         let s = run_case(&dev, dtype, nout, k).unwrap();
-        println!("{dtype:?} [{nout}x{k}]: max_abs={:.3e} max_rel={:.3e}", s.max_abs, s.max_rel);
+        println!(
+            "{dtype:?} [{nout}x{k}]: max_abs={:.3e} max_rel={:.3e}",
+            s.max_abs, s.max_rel
+        );
         // max_rel (RMS-normalized) is the correctness metric. The dp4a path sums the integer
         // grid x activation products EXACTLY in int32 and applies the f32 scale once, so it is actually
         // MORE precise than the reference, which materializes each weight as f32 (`to_float` rounds
@@ -327,8 +331,11 @@ fn cuda_matvec_iq2xxs_batched_matches_cpu() {
     .to_vec1::<f32>()
     .unwrap();
     let matmul = QMatMul::from_qtensor(
-        QTensor::new(QStorage::from_data(Cow::Owned(raw), &dev, GgmlDType::IQ2_XXS).unwrap(), (nout, k))
-            .unwrap(),
+        QTensor::new(
+            QStorage::from_data(Cow::Owned(raw), &dev, GgmlDType::IQ2_XXS).unwrap(),
+            (nout, k),
+        )
+        .unwrap(),
     )
     .unwrap();
     for &b in &[1usize, 5, 8] {
@@ -354,9 +361,13 @@ fn cuda_matvec_iq2xxs_batched_matches_cpu() {
                 rs += r * r;
             }
         }
-        let max_rel = ((sse / (b * nout) as f64).sqrt() / (rs / (b * nout) as f64).sqrt().max(1e-9)) as f32;
+        let max_rel =
+            ((sse / (b * nout) as f64).sqrt() / (rs / (b * nout) as f64).sqrt().max(1e-9)) as f32;
         println!("IQ2_XXS batched b={b}: max_rel={max_rel:.3e}");
-        assert!(max_rel < 1e-3, "batched b={b} diverged: max_rel={max_rel:.3e}");
+        assert!(
+            max_rel < 1e-3,
+            "batched b={b} diverged: max_rel={max_rel:.3e}"
+        );
     }
 }
 
@@ -378,7 +389,11 @@ fn run_prefill(dev: &Device, dtype: GgmlDType, nout: usize, k: usize, m: usize)
     .to_vec1::<f32>()
     .unwrap();
     let matmul = QMatMul::from_qtensor(
-        QTensor::new(QStorage::from_data(Cow::Owned(raw), dev, dtype).unwrap(), (nout, k)).unwrap(),
+        QTensor::new(
+            QStorage::from_data(Cow::Owned(raw), dev, dtype).unwrap(),
+            (nout, k),
+        )
+        .unwrap(),
     )
     .unwrap();
     // m>8 activation, pre-quantized to q8_1 levels per row (idempotent re-quant -> bit-identical inputs).
@@ -389,7 +404,13 @@ fn run_prefill(dev: &Device, dtype: GgmlDType, nout: usize, k: usize, m: usize)
         .flat_map(q8_1_roundtrip)
         .collect();
     let x = Tensor::from_vec(x_host.clone(), (m, k), dev).unwrap();
-    let y: Vec<f32> = matmul.forward(&x).unwrap().flatten_all().unwrap().to_vec1::<f32>().unwrap();
+    let y: Vec<f32> = matmul
+        .forward(&x)
+        .unwrap()
+        .flatten_all()
+        .unwrap()
+        .to_vec1::<f32>()
+        .unwrap();
     // Reference: GPU dense matmul of the exact f32 codebook weights x the same activation.
     let w_gpu = Tensor::from_vec(w_deq, (nout, k), dev).unwrap();
     let y_ref: Vec<f32> = x
@@ -435,7 +456,10 @@ fn cuda_qmmq_iquant_prefill_matches_dequant() {
     ] {
         for &m in &[16usize, 32] {
             let s = run_prefill(&dev, dtype, 512, 2048, m);
-            println!("qmmq {dtype:?} prefill m={m}: max_abs={:.3e} max_rel={:.3e}", s.max_abs, s.max_rel);
+            println!(
+                "qmmq {dtype:?} prefill m={m}: max_abs={:.3e} max_rel={:.3e}",
+                s.max_abs, s.max_rel
+            );
             assert!(
                 s.max_rel < 1e-3 && s.max_abs < 1e-1,
                 "{dtype:?} qmmq prefill m={m} diverged: max_abs={:.3e} max_rel={:.3e}",
@@ -484,10 +508,15 @@ fn run_moe(
     // Activation: [t, topk|1, k], pre-quantized to q8_1 levels (so dp4a's re-quant is idempotent and
     // the reference sees identical activations).
     let in1 = if shared { 1 } else { topk };
-    let x_host: Vec<f32> =
-        q8_1_roundtrip(&(0..t * in1 * k).map(|i| pseudo(i + 1_234_567)).collect::<Vec<_>>());
+    let x_host: Vec<f32> = q8_1_roundtrip(
+        &(0..t * in1 * k)
+            .map(|i| pseudo(i + 1_234_567))
+            .collect::<Vec<_>>(),
+    );
     let x_cuda = Tensor::from_vec(x_host.clone(), (t, in1, k), dev).unwrap();
-    let ids_host: Vec<u32> = (0..t * topk).map(|i| (pbyte(i * 7 + 3) as u32) % e_cnt as u32).collect();
+    let ids_host: Vec<u32> = (0..t * topk)
+        .map(|i| (pbyte(i * 7 + 3) as u32) % e_cnt as u32)
+        .collect();
     let ids_cuda = Tensor::from_vec(ids_host.clone(), (t, topk), dev).unwrap();
 
     let y_cuda: Vec<f32> = q_cuda
@@ -589,8 +618,9 @@ fn cuda_moe_route_numeric() {
             for &norm in &[false, true] {
                 // logits in [-4, 4) -> a peaked softmax (real dynamic range, not near-uniform).
                 let seed = ntok * 131 + n_experts * 17 + topk * 7 + norm as usize;
-                let logits: Vec<f32> =
-                    (0..ntok * n_experts).map(|i| pseudo(i + seed) * 4.0).collect();
+                let logits: Vec<f32> = (0..ntok * n_experts)
+                    .map(|i| pseudo(i + seed) * 4.0)
+                    .collect();
                 let lg_cuda = Tensor::from_vec(logits.clone(), (ntok, n_experts), &dev).unwrap();
                 let lg_cpu = Tensor::from_vec(logits, (ntok, n_experts), &cpu).unwrap();
 

hanzo-ml/tests/gguf_i32_skip_test.rs

@@ -21,7 +21,10 @@ fn gguf_header_loads_skipping_integer_tensors() {
     // routing tensors are skipped and the header parses.
     let content = Content::read(&mut f).expect("gguf header must parse (I32 tensors skipped)");
     let n = content.tensor_infos.len();
-    assert!(n > 0, "expected weight tensors after skipping integer auxiliaries");
+    assert!(
+        n > 0,
+        "expected weight tensors after skipping integer auxiliaries"
+    );
     // None of the kept tensors are the skipped integer routing maps.
     let leaked_int = content
         .tensor_infos

hanzo-ml/tests/gguf_mmap_test.rs

@@ -31,7 +31,9 @@ fn tmp_path() -> PathBuf {
 /// A deterministic f32 tensor with enough spread that quantization is non-trivial.
 fn sample(rows: usize, cols: usize, dev: &Device) -> Result<Tensor> {
     let n = rows * cols;
-    let data: Vec<f32> = (0..n).map(|i| ((i as f32) * 0.017).sin() * 1.3 - 0.2).collect();
+    let data: Vec<f32> = (0..n)
+        .map(|i| ((i as f32) * 0.017).sin() * 1.3 - 0.2)
+        .collect();
     Tensor::from_vec(data, (rows, cols), dev)
 }
 
@@ -47,11 +49,7 @@ fn write_fixture(path: &Path) -> Result<Vec<&'static str>> {
     gguf_file::write(
         &mut file,
         &[],
-        &[
-            ("blk.q4_0", &q4_0),
-            ("blk.q8_0", &q8_0),
-            ("blk.f32", &f32t),
-        ],
+        &[("blk.q4_0", &q4_0), ("blk.q8_0", &q8_0), ("blk.f32", &f32t)],
     )?;
     file.sync_all()?;
     Ok(vec!["blk.q4_0", "blk.q8_0", "blk.f32"])
@@ -138,9 +136,7 @@ fn gguf_mmap_tensors_reference_mapping_not_heap() {
 
     let mut total_mapped_bytes = 0usize;
     for name in names {
-        let t = mapped
-            .tensor_mmap(&mmap, name, &dev)
-            .expect("tensor_mmap");
+        let t = mapped.tensor_mmap(&mmap, name, &dev).expect("tensor_mmap");
         let d = t.data().expect("data");
         let p = d.as_ptr() as usize;
         assert!(

hanzo-ml/tests/pth_tests.rs

@@ -15,8 +15,7 @@ fn test_pth_with_key() {
 
 #[test]
 fn test_pth_fortran_contiguous() {
-    let tensors =
-        hanzo_ml::pickle::PthTensors::new("tests/fortran_tensor_3d.pth", None).unwrap();
+    let tensors = hanzo_ml::pickle::PthTensors::new("tests/fortran_tensor_3d.pth", None).unwrap();
     let tensor = tensors.get("tensor_fortran").unwrap().unwrap();
 
     assert_eq!(tensor.dims3().unwrap(), (2, 3, 4));

hanzo-ml/tests/quantized_tests.rs

@@ -1481,7 +1481,9 @@ fn iquant_indexed_moe_cpu() -> Result<()> {
     let (e, n, k) = (4usize, 128usize, 256usize);
     let (t, topk) = (3usize, 2usize);
 
-    let x_vec: Vec<f32> = (0..t * topk * k).map(|i| (i as f32 * 0.013).sin() * 0.5).collect();
+    let x_vec: Vec<f32> = (0..t * topk * k)
+        .map(|i| (i as f32 * 0.013).sin() * 0.5)
+        .collect();
     let x = Tensor::from_vec(x_vec, (t, topk, k), dev)?;
     let ids_vec: Vec<u32> = (0..t * topk).map(|i| (i * 3 % e) as u32).collect();
     let ids = Tensor::from_vec(ids_vec.clone(), (t, topk), dev)?;

hanzo-ml/tests/spark_bench.rs

@@ -163,7 +163,11 @@ fn bench_matmul(dev: &Device) -> hanzo_ml::Result<Row> {
         s.push(t0.elapsed().as_secs_f64() * 1e6);
     }
     // matmul FLOPs = 2*m*n*k.
-    Ok(report("f32_matmul_4096", s, 2.0 * m as f64 * n as f64 * k as f64))
+    Ok(report(
+        "f32_matmul_4096",
+        s,
+        2.0 * m as f64 * n as f64 * k as f64,
+    ))
 }
 
 // ---- quant matvec: y[nout] = Wq * x[k], nout = k = 4096 (single decode row) ---------------------
@@ -217,7 +221,9 @@ fn bench_moe(dev: &Device) -> hanzo_ml::Result<Row> {
     let qs = QStorage::from_data(std::borrow::Cow::Owned(bytes), dev, dtype)?;
     let qv = QTensor::new(qs, (e_cnt, n, k))?;
     let x: Vec<f32> = (0..t * topk * k).map(|i| pseudo(i + 11) * 0.7).collect();
-    let ids: Vec<u32> = (0..t * topk).map(|i| ((i * 7 + 3) % e_cnt) as u32).collect();
+    let ids: Vec<u32> = (0..t * topk)
+        .map(|i| ((i * 7 + 3) % e_cnt) as u32)
+        .collect();
     let x_t = Tensor::from_vec(x, (t, topk, k), dev)?;
     let ids_t = Tensor::from_vec(ids, (t, topk), dev)?;
 
@@ -271,7 +277,9 @@ fn spark_bench() -> hanzo_ml::Result<()> {
     } else {
         "wgpu"
     };
-    println!("=== spark_bench backend={backend} adapter=\"{adapter}\" warmup={WARMUP} iters={ITERS} ===");
+    println!(
+        "=== spark_bench backend={backend} adapter=\"{adapter}\" warmup={WARMUP} iters={ITERS} ==="
+    );
 
     let mut rows = Vec::new();
 
@@ -299,7 +307,10 @@ fn spark_bench() -> hanzo_ml::Result<()> {
         Err(e) => rows.push(skip_row("flash_attn_d128_s512", &format!("ERR: {e}"))),
     }
 
-    println!("{:<26}  {:>12}  {:>12}  {}", "kernel", "us/call", "GFLOP/s", "note");
+    println!(
+        "{:<26}  {:>12}  {:>12}  {}",
+        "kernel", "us/call", "GFLOP/s", "note"
+    );
     println!("{}", "-".repeat(72));
     for r in &rows {
         if r.us.is_nan() {

hanzo-ml/tests/tensor_tests.rs

@@ -1,5 +1,5 @@
-use hanzo_ml::{test_device, test_utils, DType, Device, IndexOp, Result, Tensor, D};
 use float8::F8E4M3;
+use hanzo_ml::{test_device, test_utils, DType, Device, IndexOp, Result, Tensor, D};
 
 fn zeros(device: &Device) -> Result<()> {
     let tensor = Tensor::zeros((5, 2), DType::F32, device)?;

hanzo-ml/tests/vulkan_quant_tests.rs

@@ -547,7 +547,6 @@ fn vulkan_matvec_tq2_0_matches_cpu() -> hanzo_ml::Result<()> {
     Ok(())
 }
 
-
 #[test]
 fn vulkan_matvec_iq2s_matches_cpu() -> hanzo_ml::Result<()> {
     let Some(dev) = gpu() else { return Ok(()) };
@@ -559,7 +558,9 @@ fn vulkan_matvec_iq2s_matches_cpu() -> hanzo_ml::Result<()> {
         );
         assert!(
             s.max_rel < 1e-3 && s.max_abs < 1e-3,
-            "IQ2_S GPU/CPU mismatch: max_abs={} max_rel={}", s.max_abs, s.max_rel
+            "IQ2_S GPU/CPU mismatch: max_abs={} max_rel={}",
+            s.max_abs,
+            s.max_rel
         );
     }
     Ok(())
@@ -576,7 +577,9 @@ fn vulkan_matvec_iq3xxs_matches_cpu() -> hanzo_ml::Result<()> {
         );
         assert!(
             s.max_rel < 1e-3 && s.max_abs < 1e-3,
-            "IQ3_XXS GPU/CPU mismatch: max_abs={} max_rel={}", s.max_abs, s.max_rel
+            "IQ3_XXS GPU/CPU mismatch: max_abs={} max_rel={}",
+            s.max_abs,
+            s.max_rel
         );
     }
     Ok(())
@@ -593,7 +596,9 @@ fn vulkan_matvec_iq3s_matches_cpu() -> hanzo_ml::Result<()> {
         );
         assert!(
             s.max_rel < 1e-3 && s.max_abs < 1e-3,
-            "IQ3_S GPU/CPU mismatch: max_abs={} max_rel={}", s.max_abs, s.max_rel
+            "IQ3_S GPU/CPU mismatch: max_abs={} max_rel={}",
+            s.max_abs,
+            s.max_rel
         );
     }
     Ok(())
@@ -610,7 +615,9 @@ fn vulkan_matvec_iq1s_matches_cpu() -> hanzo_ml::Result<()> {
         );
         assert!(
             s.max_rel < 1e-3 && s.max_abs < 1e-3,
-            "IQ1_S GPU/CPU mismatch: max_abs={} max_rel={}", s.max_abs, s.max_rel
+            "IQ1_S GPU/CPU mismatch: max_abs={} max_rel={}",
+            s.max_abs,
+            s.max_rel
         );
     }
     Ok(())
@@ -627,7 +634,9 @@ fn vulkan_matvec_iq1m_matches_cpu() -> hanzo_ml::Result<()> {
         );
         assert!(
             s.max_rel < 1e-3 && s.max_abs < 1e-3,
-            "IQ1_M GPU/CPU mismatch: max_abs={} max_rel={}", s.max_abs, s.max_rel
+            "IQ1_M GPU/CPU mismatch: max_abs={} max_rel={}",
+            s.max_abs,
+            s.max_rel
         );
     }
     Ok(())
@@ -829,7 +838,9 @@ fn moe_case(
     // Activations [t, topk, k] (per-slot inputs) and router ids [t, topk].
     let x_host: Vec<f32> = (0..t * topk * k).map(|i| pseudo(i + 11) * 0.7).collect();
     // Deterministic expert assignment spread across all experts.
-    let ids_host: Vec<u32> = (0..t * topk).map(|i| ((i * 7 + 3) % e_cnt) as u32).collect();
+    let ids_host: Vec<u32> = (0..t * topk)
+        .map(|i| ((i * 7 + 3) % e_cnt) as u32)
+        .collect();
 
     // Build the Vulkan QTensor bank the loader way and run the fused MoE forward.
     let qs_vk = QStorage::from_data(std::borrow::Cow::Owned(bytes), dev, dtype)?;
@@ -920,9 +931,9 @@ fn vulkan_flash_attn_matches_cpu() -> hanzo_ml::Result<()> {
     // (bh, lq, lk, d, causal). d = 128 is Qwen3's head dim.
     let cases = [
         (8usize, 16usize, 16usize, 128usize, false), // prefill, non-causal
-        (8, 16, 16, 128, true),                       // prefill, causal
-        (8, 1, 64, 128, true),                        // decode: single query over a 64-key cache
-        (4, 7, 13, 64, false),                        // ragged Lq != Lk, smaller d
+        (8, 16, 16, 128, true),                      // prefill, causal
+        (8, 1, 64, 128, true),                       // decode: single query over a 64-key cache
+        (4, 7, 13, 64, false),                       // ragged Lq != Lk, smaller d
     ];
     for &(bh, lq, lk, d, causal) in &cases {
         let max_abs = flash_case(&dev, bh, lq, lk, d, causal)?;
@@ -991,7 +1002,11 @@ fn vulkan_q4k_tiled2d_matches_default() -> hanzo_ml::Result<()> {
                 max_abs = max_abs.max((a - b).abs());
                 max_ref = max_ref.max(a.abs());
             }
-            let rel = if max_ref > 0.0 { max_abs / max_ref } else { max_abs };
+            let rel = if max_ref > 0.0 {
+                max_abs / max_ref
+            } else {
+                max_abs
+            };
             assert!(
                 rel < 2e-3,
                 "2d-tiled != default (m={m}, nout={nout}, k={k}): max_abs={max_abs}, rel={rel}"
@@ -1067,7 +1082,11 @@ fn vulkan_q4k_dp4a2d_matches_default() -> hanzo_ml::Result<()> {
                 max_abs = max_abs.max((a - b).abs());
                 max_ref = max_ref.max(a.abs());
             }
-            let rel = if max_ref > 0.0 { max_abs / max_ref } else { max_abs };
+            let rel = if max_ref > 0.0 {
+                max_abs / max_ref
+            } else {
+                max_abs
+            };
             assert!(
                 rel < 1.5e-2,
                 "dp4a-2d != default (m={m}, nout={nout}, k={k}): max_abs={max_abs}, max_ref={max_ref}, rel={rel}"
@@ -1150,7 +1169,11 @@ fn vulkan_q4k_coopmat_matches_default() -> hanzo_ml::Result<()> {
                 max_abs = max_abs.max((a - b).abs());
                 max_ref = max_ref.max(a.abs());
             }
-            let rel = if max_ref > 0.0 { max_abs / max_ref } else { max_abs };
+            let rel = if max_ref > 0.0 {
+                max_abs / max_ref
+            } else {
+                max_abs
+            };
             assert!(
                 rel < 1e-2,
                 "coopmat != legacy (m={m}, nout={nout}, k={k}): max_abs={max_abs}, rel={rel}"
@@ -1258,10 +1281,16 @@ fn vulkan_q4k_kernel_bench() -> hanzo_ml::Result<()> {
 // it removes the inter-op global barrier that serializes decode (see LLM.md / paper 06).
 #[test]
 fn vulkan_rope_norm_matches_unfused() {
-    let Some(dev) = gpu() else { return; };
+    let Some(dev) = gpu() else {
+        return;
+    };
     let vk = dev.as_vulkan_device().unwrap();
     // (b, h, t, d): decode (t=1) + short prefill; q/k head counts; head dims 64 and 128.
-    for &(b, h, t, d) in &[(1usize, 4usize, 1usize, 128usize), (1, 8, 5, 64), (2, 4, 3, 128)] {
+    for &(b, h, t, d) in &[
+        (1usize, 4usize, 1usize, 128usize),
+        (1, 8, 5, 64),
+        (2, 4, 3, 128),
+    ] {
         let hd = d / 2;
         let n = b * h * t * d;
         let x: Vec<f32> = (0..n).map(|i| pseudo(i + 7)).collect();
@@ -1274,7 +1303,10 @@ fn vulkan_rope_norm_matches_unfused() {
         for row in 0..b * h * t {
             let base = row * d;
             let mut ss = 0f32;
-            for i in 0..d { let v = x[base + i]; ss += v * v; }
+            for i in 0..d {
+                let v = x[base + i];
+                ss += v * v;
+            }
             let denom = (ss / d as f32 + eps).sqrt();
             let i_t = row % t;
             for i_d in 0..hd {
@@ -1286,12 +1318,19 @@ fn vulkan_rope_norm_matches_unfused() {
                 refv[i2] = x1 * s + x2 * c;
             }
         }
-        let got = vk.rope_norm_f32(&x, &weight, &cs, &sn, b, h, t, d, eps).unwrap();
+        let got = vk
+            .rope_norm_f32(&x, &weight, &cs, &sn, b, h, t, d, eps)
+            .unwrap();
         assert_eq!(got.len(), n);
         let mut maxabs = 0f32;
-        for i in 0..n { maxabs = maxabs.max((got[i] - refv[i]).abs()); }
+        for i in 0..n {
+            maxabs = maxabs.max((got[i] - refv[i]).abs());
+        }
         eprintln!("rope_norm b{b} h{h} t{t} d{d}: maxabs={maxabs:.3e}");
-        assert!(maxabs < 1e-4, "rope_norm (b{b} h{h} t{t} d{d}) mismatch maxabs={maxabs}");
+        assert!(
+            maxabs < 1e-4,
+            "rope_norm (b{b} h{h} t{t} d{d}) mismatch maxabs={maxabs}"
+        );
     }
 }
 
@@ -1299,7 +1338,9 @@ fn vulkan_rope_norm_matches_unfused() {
 // add-then-rms_norm chain bit-exactly. Rung 2 of the Vulkan decode op-fusion campaign.
 #[test]
 fn vulkan_add_rmsnorm_matches_unfused() {
-    let Some(dev) = gpu() else { return; };
+    let Some(dev) = gpu() else {
+        return;
+    };
     let vk = dev.as_vulkan_device().unwrap();
     for &(nrows, m) in &[(1usize, 2048usize), (5, 64), (3, 4096)] {
         let n = nrows * m;
@@ -1313,15 +1354,27 @@ fn vulkan_add_rmsnorm_matches_unfused() {
         for row in 0..nrows {
             let base = row * m;
             let mut ss = 0f32;
-            for i in 0..m { let v = x[base + i] + res[base + i]; s_ref[base + i] = v; ss += v * v; }
+            for i in 0..m {
+                let v = x[base + i] + res[base + i];
+                s_ref[base + i] = v;
+                ss += v * v;
+            }
             let denom = (ss / m as f32 + eps).sqrt();
-            for i in 0..m { y_ref[base + i] = s_ref[base + i] / denom * alpha[i]; }
+            for i in 0..m {
+                y_ref[base + i] = s_ref[base + i] / denom * alpha[i];
+            }
         }
         let (s_gpu, y_gpu) = vk.add_rmsnorm_f32(&x, &res, &alpha, nrows, m, eps).unwrap();
         let (mut ms, mut my) = (0f32, 0f32);
-        for i in 0..n { ms = ms.max((s_gpu[i] - s_ref[i]).abs()); my = my.max((y_gpu[i] - y_ref[i]).abs()); }
+        for i in 0..n {
+            ms = ms.max((s_gpu[i] - s_ref[i]).abs());
+            my = my.max((y_gpu[i] - y_ref[i]).abs());
+        }
         eprintln!("add_rmsnorm r{nrows} m{m}: s_maxabs={ms:.3e} y_maxabs={my:.3e}");
-        assert!(ms < 1e-5 && my < 1e-4, "add_rmsnorm r{nrows} m{m} mismatch s={ms} y={my}");
+        assert!(
+            ms < 1e-5 && my < 1e-4,
+            "add_rmsnorm r{nrows} m{m} mismatch s={ms} y={my}"
+        );
     }
 }
 
@@ -1330,7 +1383,9 @@ fn vulkan_add_rmsnorm_matches_unfused() {
 // reference) regardless of the dp4a default. Validates the Vulkan decode lever (1.8x).
 #[test]
 fn vulkan_q4k_dp4a_decode_matches_scalar() {
-    let Some(dev) = gpu() else { return; };
+    let Some(dev) = gpu() else {
+        return;
+    };
     let vk = dev.as_vulkan_device().unwrap();
     for &(nout, k) in SHAPES {
         let x: Vec<f32> = (0..k).map(|i| pseudo(i + 5)).collect();
@@ -1347,7 +1402,10 @@ fn vulkan_q4k_dp4a_decode_matches_scalar() {
             }
             let rel = (sse / refsq.max(1e-9)).sqrt();
             eprintln!("dp4a-{name} vs scalar nout={nout} k={k}: rel={rel:.3e}");
-            assert!(rel < 2e-2, "dp4a {name} rel err {rel} too large (nout={nout} k={k})");
+            assert!(
+                rel < 2e-2,
+                "dp4a {name} rel err {rel} too large (nout={nout} k={k})"
+            );
         }
     }
 }

hanzo-ml/tests/wgpu_quant_tests.rs

@@ -98,7 +98,10 @@ fn gpu() -> Option<Device> {
     match Device::new_wgpu(0) {
         Ok(d) => {
             if let Ok(g) = d.as_wgpu_device() {
-                eprintln!("[wgpu_quant_tests] selected adapter: {}", g.adapter_description());
+                eprintln!(
+                    "[wgpu_quant_tests] selected adapter: {}",
+                    g.adapter_description()
+                );
             }
             Some(d)
         }
@@ -319,7 +322,9 @@ fn moe_case(
     let bytes = q_bank.data()?.into_owned();
 
     let x_host: Vec<f32> = (0..t * topk * k).map(|i| pseudo(i + 11) * 0.7).collect();
-    let ids_host: Vec<u32> = (0..t * topk).map(|i| ((i * 7 + 3) % e_cnt) as u32).collect();
+    let ids_host: Vec<u32> = (0..t * topk)
+        .map(|i| ((i * 7 + 3) % e_cnt) as u32)
+        .collect();
 
     let qs = QStorage::from_data(std::borrow::Cow::Owned(bytes), dev, dtype)?;
     let q = QTensor::new(qs, (e_cnt, n, k))?;

hanzo-nn/src/attention/cpu_flash/varlen.rs

@@ -7,8 +7,8 @@
 
 #![allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)]
 
-use hanzo_ml::{DType, Device, Result, Storage, Tensor};
 use half::f16;
+use hanzo_ml::{DType, Device, Result, Storage, Tensor};
 use rayon::prelude::*;
 
 use super::standard::FLASH_ATTN_POOL;

hanzo-train/src/cache.rs

@@ -204,7 +204,10 @@ mod tests {
             return Ok(());
         }
         let cache = Cache::open(dir)?;
-        assert_eq!(cache.n_fused, 5, "target_layer_ids [1,11,21,31,41] => 5 fused");
+        assert_eq!(
+            cache.n_fused, 5,
+            "target_layer_ids [1,11,21,31,41] => 5 fused"
+        );
         assert_eq!(cache.manifest.hidden_size, 4096);
 
         // Record 1 has seq_len 274 (verified against the idx layout).

hanzo-train/src/model.rs

@@ -44,7 +44,11 @@ impl EmbedReader {
         }
         let shape: Vec<usize> = meta["shape"]
             .as_array()
-            .map(|a| a.iter().filter_map(|v| v.as_u64().map(|x| x as usize)).collect())
+            .map(|a| {
+                a.iter()
+                    .filter_map(|v| v.as_u64().map(|x| x as usize))
+                    .collect()
+            })
             .unwrap_or_default();
         if shape != [vocab, hidden] {
             hanzo_ml::bail!("embed_tokens.weight shape {shape:?} != [{vocab}, {hidden}]");
@@ -75,7 +79,11 @@ impl EmbedReader {
             hanzo_ml::bail!("embed row {token} out of vocab {}", self.vocab);
         }
         let off = self.data_start + (token * self.hidden * 2) as u64;
-        Ok(self.read_f16_at(off, self.hidden)?.into_iter().map(|h| h.to_f32()).collect())
+        Ok(self
+            .read_f16_at(off, self.hidden)?
+            .into_iter()
+            .map(|h| h.to_f32())
+            .collect())
     }
 
     /// The full `[vocab, hidden]` f16 table — materialized only at checkpoint-save time.
@@ -137,19 +145,25 @@ pub struct Dspark {
     hidden_norm: Tensor,
     layers: Vec<Layer>,
     norm: Tensor,
-    markov_w1: Tensor, // [vocab, rank] (nn.Embedding weight)
-    markov_w2: Tensor, // [vocab, rank] (nn.Linear rank->vocab weight)
-    embed: EmbedReader, // FROZEN [vocab, hidden], streamed on demand
+    markov_w1: Tensor,    // [vocab, rank] (nn.Embedding weight)
+    markov_w2: Tensor,    // [vocab, rank] (nn.Linear rank->vocab weight)
+    embed: EmbedReader,   // FROZEN [vocab, hidden], streamed on demand
     mask_embed: Vec<f32>, // FROZEN embedding row for `mask_token_id` (block slots 1..)
-    lm_head: Linear,   // FROZEN weight [vocab, hidden] (F16)
-    conf_w: Tensor,    // zero-init [1, hidden] (load-compat, untrained)
-    conf_b: Tensor,    // zero-init [1]
-    cos: Tensor,       // [max_pos, head_dim]
-    sin: Tensor,       // [max_pos, head_dim]
+    lm_head: Linear,      // FROZEN weight [vocab, hidden] (F16)
+    conf_w: Tensor,       // zero-init [1, hidden] (load-compat, untrained)
+    conf_b: Tensor,       // zero-init [1]
+    cos: Tensor,          // [max_pos, head_dim]
+    sin: Tensor,          // [max_pos, head_dim]
     train_markov: bool, // when false, the Markov head is frozen (excluded from the optimizer + detached)
 }
 
-fn tvar(vm: &VarMap, dev: &Device, shape: impl Into<Shape>, name: &str, init: hanzo_nn::Init) -> Result<Tensor> {
+fn tvar(
+    vm: &VarMap,
+    dev: &Device,
+    shape: impl Into<Shape>,
+    name: &str,
+    init: hanzo_nn::Init,
+) -> Result<Tensor> {
     vm.get(shape, name, init, DType::F32, dev)
 }
 
@@ -160,7 +174,10 @@ impl Dspark {
     /// which removes its AdamW state — a memory lever for constrained hosts.
     pub fn new(cfg: DsparkCfg, init_path: &Path, dev: &Device, train_markov: bool) -> Result<Self> {
         use hanzo_nn::Init;
-        let normal = Init::Randn { mean: 0.0, stdev: 0.02 };
+        let normal = Init::Randn {
+            mean: 0.0,
+            stdev: 0.02,
+        };
         let ones = Init::Const(1.0);
         let vm = VarMap::new();
 
@@ -177,23 +194,74 @@ impl Dspark {
         for i in 0..cfg.layers {
             let p = |s: &str| format!("layers.{i}.{s}");
             layers.push(Layer {
-                q_proj: Linear::new(tvar(&vm, dev, (qd, h), &p("self_attn.q_proj.weight"), normal)?, None),
-                k_proj: Linear::new(tvar(&vm, dev, (kvd, h), &p("self_attn.k_proj.weight"), normal)?, None),
-                v_proj: Linear::new(tvar(&vm, dev, (kvd, h), &p("self_attn.v_proj.weight"), normal)?, None),
-                o_proj: Linear::new(tvar(&vm, dev, (h, qd), &p("self_attn.o_proj.weight"), normal)?, None),
+                q_proj: Linear::new(
+                    tvar(&vm, dev, (qd, h), &p("self_attn.q_proj.weight"), normal)?,
+                    None,
+                ),
+                k_proj: Linear::new(
+                    tvar(&vm, dev, (kvd, h), &p("self_attn.k_proj.weight"), normal)?,
+                    None,
+                ),
+                v_proj: Linear::new(
+                    tvar(&vm, dev, (kvd, h), &p("self_attn.v_proj.weight"), normal)?,
+                    None,
+                ),
+                o_proj: Linear::new(
+                    tvar(&vm, dev, (h, qd), &p("self_attn.o_proj.weight"), normal)?,
+                    None,
+                ),
                 q_norm: tvar(&vm, dev, hd, &p("self_attn.q_norm.weight"), ones)?,
                 k_norm: tvar(&vm, dev, hd, &p("self_attn.k_norm.weight"), ones)?,
                 input_ln: tvar(&vm, dev, h, &p("input_layernorm.weight"), ones)?,
                 post_ln: tvar(&vm, dev, h, &p("post_attention_layernorm.weight"), ones)?,
-                gate: Linear::new(tvar(&vm, dev, (cfg.intermediate, h), &p("mlp.gate_proj.weight"), normal)?, None),
-                up: Linear::new(tvar(&vm, dev, (cfg.intermediate, h), &p("mlp.up_proj.weight"), normal)?, None),
-                down: Linear::new(tvar(&vm, dev, (h, cfg.intermediate), &p("mlp.down_proj.weight"), normal)?, None),
+                gate: Linear::new(
+                    tvar(
+                        &vm,
+                        dev,
+                        (cfg.intermediate, h),
+                        &p("mlp.gate_proj.weight"),
+                        normal,
+                    )?,
+                    None,
+                ),
+                up: Linear::new(
+                    tvar(
+                        &vm,
+                        dev,
+                        (cfg.intermediate, h),
+                        &p("mlp.up_proj.weight"),
+                        normal,
+                    )?,
+                    None,
+                ),
+                down: Linear::new(
+                    tvar(
+                        &vm,
+                        dev,
+                        (h, cfg.intermediate),
+                        &p("mlp.down_proj.weight"),
+                        normal,
+                    )?,
+                    None,
+                ),
             });
         }
 
         let norm = tvar(&vm, dev, h, "norm.weight", Init::Const(cfg.final_norm_init))?;
-        let markov_w1 = tvar(&vm, dev, (cfg.vocab, cfg.markov_rank), "markov_head.markov_w1.weight", normal)?;
-        let markov_w2 = tvar(&vm, dev, (cfg.vocab, cfg.markov_rank), "markov_head.markov_w2.weight", normal)?;
+        let markov_w1 = tvar(
+            &vm,
+            dev,
+            (cfg.vocab, cfg.markov_rank),
+            "markov_head.markov_w1.weight",
+            normal,
+        )?;
+        let markov_w2 = tvar(
+            &vm,
+            dev,
+            (cfg.vocab, cfg.markov_rank),
+            "markov_head.markov_w2.weight",
+            normal,
+        )?;
 
         // FROZEN embed streamed on demand (not resident); FROZEN lm_head kept in F16 (1.05GB, vs
         // 2.1GB f32). The CPU backend supports F16 matmul (not BF16). `lm_head`'s [hidden, vocab]
@@ -241,7 +309,9 @@ impl Dspark {
             return all;
         }
         let skip = [self.markov_w1.id(), self.markov_w2.id()];
-        all.into_iter().filter(|v| !skip.contains(&v.id())).collect()
+        all.into_iter()
+            .filter(|v| !skip.contains(&v.id()))
+            .collect()
     }
 
     /// Fuse the target hidden states into the DSpark context memory, over the whole sequence at once
@@ -287,7 +357,8 @@ impl Dspark {
 
         // create_noise_embed: slot 0 = anchor token, slots 1.. = mask_token_id. Streamed from the
         // FROZEN embedding; the block state is a plain leaf (never optimized, no grad to the table).
-        let anchor_tok = tok(a).ok_or_else(|| hanzo_ml::Error::msg("anchor token out of vocab"))? as usize;
+        let anchor_tok =
+            tok(a).ok_or_else(|| hanzo_ml::Error::msg("anchor token out of vocab"))? as usize;
         let mut hv = Vec::with_capacity(bs * h);
         hv.extend_from_slice(&self.embed.row_f32(anchor_tok)?);
         for _ in 1..bs {
@@ -304,7 +375,9 @@ impl Dspark {
 
         for layer in &self.layers {
             let normed = ops::rms_norm(&hstate.contiguous()?, &layer.input_ln, eps)?;
-            let attn = self.attention(layer, &normed, &ctx, &cos_draft, &sin_draft, &cos_full, &sin_full, &mask)?;
+            let attn = self.attention(
+                layer, &normed, &ctx, &cos_draft, &sin_draft, &cos_full, &sin_full, &mask,
+            )?;
             hstate = hstate.add(&attn)?;
             let normed2 = ops::rms_norm(&hstate.contiguous()?, &layer.post_ln, eps)?;
             let gate = ops::silu(&layer.gate.forward(&normed2)?)?;
@@ -318,7 +391,9 @@ impl Dspark {
         // later, once, over the whole batch — see `frozen_head_ce`.)
         let mut prevs = Vec::with_capacity(bs);
         for k in 0..bs {
-            prevs.push(tok(a + k).ok_or_else(|| hanzo_ml::Error::msg("markov prev token out of vocab"))?);
+            prevs.push(
+                tok(a + k).ok_or_else(|| hanzo_ml::Error::msg("markov prev token out of vocab"))?,
+            );
         }
         let prev_t = Tensor::from_vec(prevs, (bs,), &self.dev)?;
         // Frozen Markov ⇒ detach so no vocab×rank grad is formed and it stays at init.
@@ -344,12 +419,21 @@ impl Dspark {
             return Ok(None);
         }
         let nv = targets.len();
-        Ok(Some((block_out.narrow(0, 0, nv)?, bias.narrow(0, 0, nv)?, targets)))
+        Ok(Some((
+            block_out.narrow(0, 0, nv)?,
+            bias.narrow(0, 0, nv)?,
+            targets,
+        )))
     }
 
     /// Cross-entropy through the FROZEN `lm_head` for a whole batch, returning
     /// `(loss_value, surrogate)`. See [`frozen_head_ce`]; this just supplies the head weight.
-    pub fn head_ce(&self, block_cat: &Tensor, bias_cat: &Tensor, targets: &Tensor) -> Result<(Tensor, Tensor)> {
+    pub fn head_ce(
+        &self,
+        block_cat: &Tensor,
+        bias_cat: &Tensor,
+        targets: &Tensor,
+    ) -> Result<(Tensor, Tensor)> {
         frozen_head_ce(block_cat, self.lm_head.weight(), bias_cat, targets)
     }
 
@@ -381,7 +465,8 @@ impl Dspark {
         let q = apply_rope(&q, cos_draft, sin_draft)?;
 
         // Keys/values: [fused context ‖ block].
-        let k = Tensor::cat(&[&layer.k_proj.forward(ctx)?, &layer.k_proj.forward(x)?], 0)?.reshape((kl, nkv, hd))?;
+        let k = Tensor::cat(&[&layer.k_proj.forward(ctx)?, &layer.k_proj.forward(x)?], 0)?
+            .reshape((kl, nkv, hd))?;
         let k = ops::rms_norm(&k.contiguous()?, &layer.k_norm, eps)?;
         let k = k.transpose(0, 1)?.contiguous()?; // [nkv, kl, hd]
         let k = apply_rope(&k, cos_full, sin_full)?;
@@ -411,7 +496,10 @@ impl Dspark {
         for (k, v) in self.varmap.data().lock().unwrap().iter() {
             map.insert(k.clone(), v.as_tensor().clone());
         }
-        map.insert("embed_tokens.weight".into(), self.embed.full_f16(&self.dev)?);
+        map.insert(
+            "embed_tokens.weight".into(),
+            self.embed.full_f16(&self.dev)?,
+        );
         map.insert("lm_head.weight".into(), self.lm_head.weight().clone());
         map.insert("confidence_head.proj.weight".into(), self.conf_w.clone());
         map.insert("confidence_head.proj.bias".into(), self.conf_b.clone());
@@ -484,7 +572,12 @@ fn repeat_kv(x: &Tensor, nrep: usize) -> Result<Tensor> {
 }
 
 /// Precompute neox RoPE cos/sin tables `[max_pos, head_dim]` (each half duplicated).
-fn rope_tables(max_pos: usize, head_dim: usize, theta: f64, dev: &Device) -> Result<(Tensor, Tensor)> {
+fn rope_tables(
+    max_pos: usize,
+    head_dim: usize,
+    theta: f64,
+    dev: &Device,
+) -> Result<(Tensor, Tensor)> {
     let half = head_dim / 2;
     let mut cos = vec![0f32; max_pos * head_dim];
     let mut sin = vec![0f32; max_pos * head_dim];
@@ -507,7 +600,12 @@ fn rope_tables(max_pos: usize, head_dim: usize, theta: f64, dev: &Device) -> Res
 /// Additive DSpark attention mask `[block, ctx_len+block]`: `0` where attended, `-inf` where masked.
 /// A context column `j` is attended iff `j < anchor_pos`; every block column is attended
 /// (bidirectional block self-attention). Mirrors `create_dspark_attention_mask`.
-pub fn dspark_mask(ctx_len: usize, block: usize, anchor_pos: usize, dev: &Device) -> Result<Tensor> {
+pub fn dspark_mask(
+    ctx_len: usize,
+    block: usize,
+    anchor_pos: usize,
+    dev: &Device,
+) -> Result<Tensor> {
     let kl = ctx_len + block;
     let neg = f32::NEG_INFINITY;
     let mut row = vec![0f32; kl];
@@ -599,8 +697,14 @@ pub fn verify_checkpoint(path: &Path, cfg: &DsparkCfg) -> Result<()> {
         ("hidden_norm.weight".into(), vec![h]),
         ("norm.weight".into(), vec![h]),
         ("lm_head.weight".into(), vec![cfg.vocab, h]),
-        ("markov_head.markov_w1.weight".into(), vec![cfg.vocab, cfg.markov_rank]),
-        ("markov_head.markov_w2.weight".into(), vec![cfg.vocab, cfg.markov_rank]),
+        (
+            "markov_head.markov_w1.weight".into(),
+            vec![cfg.vocab, cfg.markov_rank],
+        ),
+        (
+            "markov_head.markov_w2.weight".into(),
+            vec![cfg.vocab, cfg.markov_rank],
+        ),
         ("confidence_head.proj.weight".into(), vec![1, h]),
         ("confidence_head.proj.bias".into(), vec![1]),
     ];
@@ -649,13 +753,16 @@ mod tests {
         // vocab=2, rank=2. w1 rows = prev-token latents; w2 rows = per-vocab projection.
         let w1 = Tensor::from_vec(vec![1f32, 0., 0., 1.], (2, 2), &dev)?; // [[1,0],[0,1]]
         let w2 = Tensor::from_vec(vec![1f32, 2., 3., 4.], (2, 2), &dev)?; // [[1,2],[3,4]]
-        // prev token 0 -> latent [1,0]; bias = latent · w2^T = [1*1+0*2, 1*3+0*4] = [1,3].
+                                                                          // prev token 0 -> latent [1,0]; bias = latent · w2^T = [1*1+0*2, 1*3+0*4] = [1,3].
         let prev = Tensor::from_vec(vec![0u32], (1,), &dev)?;
         let bias = markov_bias(&w1, &w2, &prev)?;
         assert_eq!(bias.to_vec2::<f32>()?, vec![vec![1.0, 3.0]]);
         // prev token 1 -> latent [0,1]; bias = [2,4].
         let prev1 = Tensor::from_vec(vec![1u32], (1,), &dev)?;
-        assert_eq!(markov_bias(&w1, &w2, &prev1)?.to_vec2::<f32>()?, vec![vec![2.0, 4.0]]);
+        assert_eq!(
+            markov_bias(&w1, &w2, &prev1)?.to_vec2::<f32>()?,
+            vec![vec![2.0, 4.0]]
+        );
         Ok(())
     }
 
@@ -670,26 +777,34 @@ mod tests {
             &dev,
         )?)?;
         let lm_w = Tensor::from_vec(
-            (0..v * h).map(|i| (i % 7) as f32 * 0.13 - 0.4).collect::<Vec<_>>(),
+            (0..v * h)
+                .map(|i| (i % 7) as f32 * 0.13 - 0.4)
+                .collect::<Vec<_>>(),
             (v, h),
             &dev,
         )?;
         let bias = Var::from_tensor(&Tensor::from_vec(
-            (0..n * v).map(|i| (i % 5) as f32 * 0.05).collect::<Vec<_>>(),
+            (0..n * v)
+                .map(|i| (i % 5) as f32 * 0.05)
+                .collect::<Vec<_>>(),
             (n, v),
             &dev,
         )?)?;
         let targets = Tensor::from_vec(vec![1u32, 4, 2], (n,), &dev)?;
 
         // Reference: full autograd through the head (lm_w kept f32 so no rounding).
-        let z_ref = block.as_tensor().matmul(&lm_w.t()?)?.add(bias.as_tensor())?;
+        let z_ref = block
+            .as_tensor()
+            .matmul(&lm_w.t()?)?
+            .add(bias.as_tensor())?;
         let loss_ref = hanzo_nn::loss::cross_entropy(&z_ref, &targets)?;
         let g_ref = loss_ref.backward()?;
         let gb_ref = g_ref.get(&block).unwrap().flatten_all()?.to_vec1::<f32>()?;
         let gbias_ref = g_ref.get(&bias).unwrap().flatten_all()?.to_vec1::<f32>()?;
 
         // Surrogate: same grads, no [hidden, vocab] head grad formed.
-        let (loss_val, surrogate) = frozen_head_ce(block.as_tensor(), &lm_w, bias.as_tensor(), &targets)?;
+        let (loss_val, surrogate) =
+            frozen_head_ce(block.as_tensor(), &lm_w, bias.as_tensor(), &targets)?;
         let g = surrogate.backward()?;
         let gb = g.get(&block).unwrap().flatten_all()?.to_vec1::<f32>()?;
         let gbias = g.get(&bias).unwrap().flatten_all()?.to_vec1::<f32>()?;
@@ -715,21 +830,37 @@ mod tests {
         let dev = Device::Cpu;
         let (n, d, h, v) = (8usize, 6usize, 5usize, 11usize);
         let x = Tensor::from_vec(
-            (0..n * d).map(|i| ((i * 7) % 13) as f32 * 0.1 - 0.6).collect::<Vec<_>>(),
+            (0..n * d)
+                .map(|i| ((i * 7) % 13) as f32 * 0.1 - 0.6)
+                .collect::<Vec<_>>(),
             (n, d),
             &dev,
         )?;
         let w = Tensor::from_vec(
-            (0..v * h).map(|i| ((i * 5) % 17) as f32 * 0.1 - 0.8).collect::<Vec<_>>(),
+            (0..v * h)
+                .map(|i| ((i * 5) % 17) as f32 * 0.1 - 0.8)
+                .collect::<Vec<_>>(),
             (v, h),
             &dev,
         )?;
         let bias = Tensor::zeros((n, v), DType::F32, &dev)?;
-        let targets = Tensor::from_vec((0..n as u32).map(|i| (i * 3) % v as u32).collect::<Vec<_>>(), (n,), &dev)?;
+        let targets = Tensor::from_vec(
+            (0..n as u32)
+                .map(|i| (i * 3) % v as u32)
+                .collect::<Vec<_>>(),
+            (n,),
+            &dev,
+        )?;
         // theta = 0 ⇒ block = 0 ⇒ uniform logits ⇒ CE ≈ ln(v) at step 0.
         let theta = Var::from_tensor(&Tensor::zeros((d, h), DType::F32, &dev)?)?;
 
-        let mut opt = AdamW::from_slice(&[&theta], ParamsAdamW { lr: 0.2, ..Default::default() })?;
+        let mut opt = AdamW::from_slice(
+            &[&theta],
+            ParamsAdamW {
+                lr: 0.2,
+                ..Default::default()
+            },
+        )?;
         let mut first = 0f32;
         let mut last = 0f32;
         for step in 0..60 {
@@ -746,8 +877,14 @@ mod tests {
             }
         }
         let ln_v = (v as f32).ln();
-        assert!((first - ln_v).abs() < 0.2, "start CE {first} should be ~ln(v) {ln_v}");
-        assert!(last < first * 0.6, "CE should drop substantially: {first} -> {last}");
+        assert!(
+            (first - ln_v).abs() < 0.2,
+            "start CE {first} should be ~ln(v) {ln_v}"
+        );
+        assert!(
+            last < first * 0.6,
+            "CE should drop substantially: {first} -> {last}"
+        );
         Ok(())
     }
 

hanzo-training/src/grpo/math.rs

@@ -91,7 +91,10 @@ pub fn policy_gradient_loss(
 ) -> Result<Tensor> {
     assert_eq!(advantages.len(), seq_logprobs.len());
     assert_eq!(advantages.len(), completion_lengths.len());
-    assert!(!advantages.is_empty(), "empty batch in policy_gradient_loss");
+    assert!(
+        !advantages.is_empty(),
+        "empty batch in policy_gradient_loss"
+    );
 
     // Per-completion contribution: A_i * (logprob_i / len_i). We build it as a
     // sum of differentiable scalars then divide by N for the mean.
@@ -185,8 +188,16 @@ mod tests {
         let r = vec![0.0f32, 10.0];
         let a = group_advantages(&r, true, 1e-4);
         // With ddof=1 std of {0,10} is ~7.071; (10-5)/7.071 ~= 0.707
-        assert!((a[1] - std::f32::consts::FRAC_1_SQRT_2).abs() < 1e-2, "got {}", a[1]);
-        assert!((a[0] + std::f32::consts::FRAC_1_SQRT_2).abs() < 1e-2, "got {}", a[0]);
+        assert!(
+            (a[1] - std::f32::consts::FRAC_1_SQRT_2).abs() < 1e-2,
+            "got {}",
+            a[1]
+        );
+        assert!(
+            (a[0] + std::f32::consts::FRAC_1_SQRT_2).abs() < 1e-2,
+            "got {}",
+            a[0]
+        );
     }
 
     #[test]

hanzo-training/src/grpo/mod.rs

@@ -82,11 +82,8 @@ mod toy_tests {
             let varmap = VarMap::new();
             let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device);
             // Initialize logits to zero -> uniform policy at start.
-            let logits = vb.get_with_hints(
-                (max_len, vocab),
-                "logits",
-                hanzo_nn::Init::Const(0.0),
-            )?;
+            let logits =
+                vb.get_with_hints((max_len, vocab), "logits", hanzo_nn::Init::Const(0.0))?;
             Ok(Self {
                 logits,
                 varmap,
@@ -135,7 +132,7 @@ mod toy_tests {
             for (pos, &tok) in completion_tokens.iter().enumerate() {
                 debug_assert!((tok as usize) < self.vocab);
                 let idx = Tensor::new(&[tok], &self.device)?; // [1]
-                // row `pos`, then gather the token's logprob -> scalar tensor
+                                                              // row `pos`, then gather the token's logprob -> scalar tensor
                 let row = logp.narrow(0, pos, 1)?.squeeze(0)?; // [vocab]
                 let val = row.index_select(&idx, 0)?.squeeze(0)?; // scalar
                 total = Some(match total {
@@ -234,8 +231,7 @@ mod toy_tests {
             seed: 7,
         };
         let rewards: Vec<Box<dyn Reward>> = vec![Box::new(TokenMatchReward::new(2))];
-        let mut trainer =
-            GrpoTrainer::with_reference(cfg, policy, rewards, Some(reference))?;
+        let mut trainer = GrpoTrainer::with_reference(cfg, policy, rewards, Some(reference))?;
 
         let prompts = vec![vec![1u32], vec![2u32]];
         let mut last = trainer.step(&prompts)?;

hanzo-training/src/grpo/policy.rs

@@ -55,11 +55,7 @@ pub trait Policy {
     /// (a single scalar per completion) is the sequence-level objective; a
     /// per-token vector is also valid for token-level KL but the GRPO PG term
     /// only needs the sum (see `math::policy_gradient_loss`).
-    fn sequence_logprob(
-        &self,
-        prompt_tokens: &[u32],
-        completion_tokens: &[u32],
-    ) -> Result<Tensor>;
+    fn sequence_logprob(&self, prompt_tokens: &[u32], completion_tokens: &[u32]) -> Result<Tensor>;
 
     /// Optional decoder from token ids to text, used to populate
     /// [`Completion::completion_text`] for text-based rewards. Default: `None`.

hanzo-training/src/grpo/trainer.rs

@@ -146,9 +146,12 @@ impl<P: Policy, R: Policy> GrpoTrainer<P, R> {
                 .wrapping_add((pi as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15));
             let mut sampler = LogitsSampler::with_seed(&self.config, seed);
 
-            let mut group =
-                self.policy
-                    .sample_group(prompt, g, self.config.max_completion_len, &mut sampler)?;
+            let mut group = self.policy.sample_group(
+                prompt,
+                g,
+                self.config.max_completion_len,
+                &mut sampler,
+            )?;
             debug_assert_eq!(group.completions.len(), g);
 
             // Populate decoded text for text-based rewards if the policy can.
@@ -163,7 +166,11 @@ impl<P: Policy, R: Policy> GrpoTrainer<P, R> {
 
             let mean = group_rewards.iter().copied().sum::<f32>() / g as f32;
             let var = if g > 1 {
-                group_rewards.iter().map(|&r| (r - mean).powi(2)).sum::<f32>() / (g as f32 - 1.0)
+                group_rewards
+                    .iter()
+                    .map(|&r| (r - mean).powi(2))
+                    .sum::<f32>()
+                    / (g as f32 - 1.0)
             } else {
                 0.0
             };