Add end-to-end MXFP4 Metal support

src/engine/gguf.rs

@@ -19,6 +19,7 @@ pub(super) const Q5_K: u32 = 13;
 pub(super) const Q6_K: u32 = 14;
 pub(super) const IQ2_XXS: u32 = 16;
 pub(super) const I32: u32 = 26;
+pub(super) const MXFP4: u32 = 39;
 
 #[derive(Clone, Debug)]
 pub(super) enum Value {
@@ -120,6 +121,12 @@ impl Gguf {
             let relative_offset = cursor.u64()?;
             let (block_elements, block_bytes) = tensor_type(kind)
                 .ok_or_else(|| format!("tensor {name} uses unsupported GGUF type {kind}"))?;
+            if kind == MXFP4 && !dims[0].is_multiple_of(block_elements) {
+                return Err(format!(
+                    "tensor {name} MXFP4 row length {} is not aligned to {block_elements} values",
+                    dims[0]
+                ));
+            }
             let blocks = elements
                 .checked_add(block_elements - 1)
                 .ok_or_else(|| format!("tensor {name} size overflows"))?
@@ -357,6 +364,7 @@ fn tensor_type(kind: u32) -> Option<(u64, u64)> {
         28 => (1, 8),
         29 => (256, 56),
         30 => (1, 2),
+        39 => (32, 17),
         _ => return None,
     })
 }
@@ -529,16 +537,104 @@ impl<'a> Cursor<'a> {
 mod tests {
     use super::*;
     use std::fs;
+    use std::sync::atomic::{AtomicU64, Ordering};
     use std::time::{SystemTime, UNIX_EPOCH};
 
+    static FIXTURE_ID: AtomicU64 = AtomicU64::new(0);
+
     #[test]
     fn maps_metadata_and_tensor_payload_without_copying_weights() {
+        let path = write_gguf(&[1], F32, 4);
+
+        let model = Gguf::open(&path).unwrap();
+        assert_eq!(model.bytes("general.architecture").unwrap(), b"deepseek4");
+        assert_eq!(model.tensor("weight").unwrap().dims, [1]);
+        assert_eq!(model.tensor_data("weight").unwrap(), [0; 4]);
+        model.warm().unwrap();
+        fs::remove_file(path).unwrap();
+    }
+
+    #[test]
+    fn accepts_exact_mxfp4_block_layout() {
+        let path = write_gguf(&[32, 2], MXFP4, 34);
+        let model = Gguf::open(&path).unwrap();
+        let tensor = model.tensor("weight").unwrap();
+        assert_eq!(tensor.bytes, 34);
+        assert_eq!(model.tensor_data("weight").unwrap().len(), 34);
+        fs::remove_file(path).unwrap();
+    }
+
+    #[test]
+    fn rejects_mxfp4_rows_that_are_not_block_aligned() {
+        let path = write_gguf(&[33, 2], MXFP4, 34);
+        let error = Gguf::open(&path).err().unwrap();
+        assert!(error.contains("MXFP4 row length 33 is not aligned to 32 values"));
+        fs::remove_file(path).unwrap();
+    }
+
+    #[test]
+    fn rejects_truncated_mxfp4_payload() {
+        let path = write_gguf(&[32, 2], MXFP4, 33);
+        let error = Gguf::open(&path).err().unwrap();
+        assert!(error.contains("points outside the GGUF file"));
+        fs::remove_file(path).unwrap();
+    }
+
+    #[test]
+    fn rejects_mxfp4_dimension_overflow() {
+        let path = write_gguf(&[32, u64::MAX], MXFP4, 0);
+        let error = Gguf::open(&path).err().unwrap();
+        assert!(error.contains("tensor weight element count overflows"));
+        fs::remove_file(path).unwrap();
+    }
+
+    #[test]
+    fn keeps_unknown_tensor_kinds_rejected() {
+        let path = write_gguf(&[32], 38, 0);
+        let error = Gguf::open(&path).err().unwrap();
+        assert!(error.contains("unsupported GGUF type 38"));
+        fs::remove_file(path).unwrap();
+    }
+
+    #[test]
+    fn mxfp4_scalar_dot_matches_independent_values() {
+        const VALUES: [f32; 16] = [
+            0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0,
+        ];
+        let mut block = [0xf6_u8; 17];
+        block[0] = 127;
+        let activation = [1.0; 32];
+        let scale = f32::from_bits(u32::from(block[0]) << 23);
+        let mut dot = 0.0;
+        for (index, packed) in block[1..].iter().copied().enumerate() {
+            dot += scale * VALUES[usize::from(packed & 0x0f)] * activation[index];
+            dot += scale * VALUES[usize::from(packed >> 4)] * activation[index + 16];
+        }
+        assert_eq!(dot, -32.0);
+
+        for exponent in [0_u8, 1, 126, 127, 128, 254] {
+            let bits = if exponent == 0 {
+                0x0040_0000
+            } else {
+                u32::from(exponent) << 23
+            };
+            let expected = 2.0_f32.powi(if exponent == 0 {
+                -127
+            } else {
+                i32::from(exponent) - 127
+            });
+            assert_eq!(f32::from_bits(bits), expected);
+        }
+    }
+
+    fn write_gguf(dims: &[u64], kind: u32, payload_bytes: usize) -> PathBuf {
         let path = std::env::temp_dir().join(format!(
-            "ds4-server-gguf-{}",
+            "ds4-server-gguf-{}-{}",
             SystemTime::now()
                 .duration_since(UNIX_EPOCH)
                 .unwrap()
-                .as_nanos()
+                .as_nanos(),
+            FIXTURE_ID.fetch_add(1, Ordering::Relaxed),
         ));
         let mut bytes = Vec::new();
         bytes.extend(MAGIC.to_le_bytes());
@@ -549,20 +645,15 @@ mod tests {
         bytes.extend(8_u32.to_le_bytes());
         push_string(&mut bytes, b"deepseek4");
         push_string(&mut bytes, b"weight");
-        bytes.extend(1_u32.to_le_bytes());
-        bytes.extend(1_u64.to_le_bytes());
-        bytes.extend(F32.to_le_bytes());
+        bytes.extend((dims.len() as u32).to_le_bytes());
+        for dim in dims {
+            bytes.extend(dim.to_le_bytes());
+        }
+        bytes.extend(kind.to_le_bytes());
         bytes.extend(0_u64.to_le_bytes());
-        bytes.resize(bytes.len().div_ceil(32) * 32, 0);
-        bytes.extend(1_f32.to_le_bytes());
+        bytes.resize(bytes.len().div_ceil(32) * 32 + payload_bytes, 0);
         fs::write(&path, bytes).unwrap();
-
-        let model = Gguf::open(&path).unwrap();
-        assert_eq!(model.bytes("general.architecture").unwrap(), b"deepseek4");
-        assert_eq!(model.tensor("weight").unwrap().dims, [1]);
-        assert_eq!(model.tensor_data("weight").unwrap(), 1_f32.to_le_bytes());
-        model.warm().unwrap();
-        fs::remove_file(path).unwrap();
+        path
     }
 
     fn push_string(bytes: &mut Vec<u8>, value: &[u8]) {