Add end-to-end MXFP4 Metal support

build.rs

@@ -11,6 +11,8 @@ fn main() {
         .include("native/metal")
         .file(metal)
         .flag("-fobjc-arc")
+        .flag("-ffast-math")
+        .flag("-mcpu=native")
         .opt_level(3)
         .compile("ds4_metal");
     println!("cargo:rerun-if-changed=native/media");

src/engine.rs

@@ -12,7 +12,9 @@ use crate::model::{ModelChoice, validate_engine_artifacts};
 #[cfg(target_os = "macos")]
 use crate::settings::TurnSettings;
 use crate::settings::{EngineSettings, ReasoningMode};
-use gguf::{F16, F32, Gguf, I32, IQ2_XXS, Q2_K, Q4_0, Q4_K, Q5_K, Q6_K, Q8_0, Tensor, Value};
+use gguf::{
+    F16, F32, Gguf, I32, IQ2_XXS, MXFP4, Q2_K, Q4_0, Q4_K, Q5_K, Q6_K, Q8_0, Tensor, Value,
+};
 #[cfg(target_os = "macos")]
 use kvstore::{KvStore, StoreReason};
 use sha2::{Digest, Sha256};
@@ -53,7 +55,7 @@ pub(crate) use kvstore::sweep_unreachable as sweep_transient_cache;
 pub(crate) use metal::configure_sources as configure_metal_sources;
 
 const DENSE: &[u32] = &[Q8_0, Q4_K, Q4_0];
-const ROUTED: &[u32] = &[Q8_0, IQ2_XXS, Q2_K, Q4_K, Q5_K, Q6_K];
+const ROUTED: &[u32] = &[Q8_0, IQ2_XXS, Q2_K, Q4_K, Q5_K, Q6_K, MXFP4];
 const PLAIN: &[u32] = &[F16, F32];
 const DSPARK_DENSE: &[u32] = &[F16, F32, Q8_0];
 

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((dims.len() as u32).to_le_bytes());
-        bytes.extend(1_u64.to_le_bytes());
+        for dim in dims {
-        bytes.extend(F32.to_le_bytes());
+            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.resize(bytes.len().div_ceil(32) * 32 + payload_bytes, 0);
-        bytes.extend(1_f32.to_le_bytes());
         fs::write(&path, bytes).unwrap();
-
+        path
-        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();
     }
 
     fn push_string(bytes: &mut Vec<u8>, value: &[u8]) {

src/engine/metal.rs

@@ -8,7 +8,7 @@ use glm::GlmExecutor;
 use gpu::*;
 use profile::ExpertProfile;
 
-use super::gguf::{F16, F32, Gguf, IQ2_XXS, Q2_K, Q4_K, Q8_0, Tensor as GgufTensor};
+use super::gguf::{F16, F32, Gguf, IQ2_XXS, MXFP4, Q2_K, Q4_K, Q8_0, Tensor as GgufTensor};
 use super::validation::{DsparkConfig, SupportKind, dspark_config};
 use super::{Model, ModelFamily, Rng, exact_delta_sample};
 use crate::model::ModelChoice;
@@ -20,8 +20,11 @@ use std::env;
 use std::ffi::{CStr, c_char, c_void};
 use std::fs::{self, File};
 use std::io::{Read, Write};
+use std::os::unix::fs::FileExt;
 use std::path::Path;
 use std::ptr::NonNull;
+use std::sync::Arc;
+use std::thread::JoinHandle;
 use std::time::{Duration, Instant, UNIX_EPOCH};
 
 const CHECKPOINT_MAGIC: &[u8; 8] = b"DS4RKV01";
@@ -1664,6 +1667,117 @@ struct DeepSeekModelSpans {
     max_tensor_bytes: u64,
 }
 
+struct PrefillPread {
+    layer: u32,
+    workers: Vec<JoinHandle<Result<(), String>>>,
+}
+
+impl PrefillPread {
+    fn start(model: &Model, layer: u32, spans: &DeepSeekModelSpans) -> Result<Self, String> {
+        const THREADS: usize = 8;
+        const CHUNK: usize = 1024 * 1024;
+
+        let file = Arc::new(File::open(model.main.path()).map_err(|error| {
+            format!(
+                "Cannot open {} for SSD prefill: {error}",
+                model.main.path().display()
+            )
+        })?);
+        let mut ranges = vec![Vec::new(); THREADS];
+        for &(offset, bytes) in &spans.ranges {
+            let part = bytes.div_ceil(THREADS as u64);
+            let mut consumed = 0;
+            for worker_ranges in &mut ranges {
+                if consumed == bytes {
+                    break;
+                }
+                let size = (bytes - consumed).min(part);
+                worker_ranges.push((offset + consumed, size));
+                consumed += size;
+            }
+        }
+        let workers = ranges
+            .into_iter()
+            .map(|ranges| {
+                let file = Arc::clone(&file);
+                std::thread::spawn(move || {
+                    let mut buffer = vec![0_u8; CHUNK];
+                    for (offset, bytes) in ranges {
+                        let mut read = 0;
+                        while read < bytes {
+                            let wanted = (bytes - read).min(CHUNK as u64) as usize;
+                            let count = loop {
+                                match file.read_at(&mut buffer[..wanted], offset + read) {
+                                    Err(error)
+                                        if error.kind() == std::io::ErrorKind::Interrupted => {}
+                                    result => break result.map_err(|error| error.to_string())?,
+                                }
+                            };
+                            if count == 0 {
+                                return Err("unexpected EOF during SSD prefill".into());
+                            }
+                            read += count as u64;
+                        }
+                    }
+                    Ok(())
+                })
+            })
+            .collect();
+        Ok(Self { layer, workers })
+    }
+
+    fn finish(mut self) -> Result<(), String> {
+        for worker in self.workers.drain(..) {
+            worker
+                .join()
+                .map_err(|_| format!("SSD prefill worker panicked for layer {}", self.layer))??;
+        }
+        Ok(())
+    }
+}
+
+impl Drop for PrefillPread {
+    fn drop(&mut self) {
+        for worker in self.workers.drain(..) {
+            let _ = worker.join();
+        }
+    }
+}
+
+fn prefill_pread_enabled() -> bool {
+    !environment_present(c"DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PREAD")
+        && !environment_present(c"DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PREPARE")
+}
+
+fn prefill_selected_addr(rows: u32, shape: super::Shape, weights: &Layer) -> bool {
+    unsafe {
+        ds4_gpu_stream_prefill_batch_selected_addr_enabled(
+            rows,
+            shape.experts as u32,
+            shape.experts_used as u32,
+            weights.expert_gate.kind,
+            weights.expert_down.kind,
+        ) != 0
+    }
+}
+
+fn start_prefill_pread(
+    model: &Model,
+    ssd: &SsdPlan,
+    weights: &Layer,
+    layer: u32,
+    rows: u32,
+) -> Result<PrefillPread, String> {
+    let spans = deepseek_layer_model_spans(
+        model,
+        weights,
+        layer,
+        prefill_selected_addr(rows, model.shape, weights),
+        ssd.per_expert_bytes,
+    )?;
+    PrefillPread::start(model, layer, &spans)
+}
+
 #[derive(Clone, Copy)]
 struct SelectedLoadJob {
     selected: usize,
@@ -1864,6 +1978,24 @@ impl Weights {
     }
 }
 
+fn mxfp4_decode_fast_lookup_allowed(
+    model: &Model,
+    weights: &Weights,
+    quality: bool,
+    ssd_streaming: bool,
+) -> bool {
+    !quality
+        && !ssd_streaming
+        && model.support_kind.is_none()
+        && weights.layers.get(4).is_some_and(|layer| {
+            layer.expert_gate.kind == MXFP4
+                && layer.expert_up.kind == MXFP4
+                && layer.expert_down.kind == MXFP4
+        })
+        && unsafe { ds4_gpu_device_is_pre_m5_apple_silicon() } != 0
+        && !environment_present(c"DS4_METAL_DISABLE_PRE_M5_DECODE_PIPELINE_FAST_LOOKUP")
+}
+
 impl SsdPlan {
     fn new(
         model: &Model,
@@ -2189,7 +2321,7 @@ fn deepseek_model_spans(
         if tensor.bytes == 0 {
             continue;
         }
-        let isolate = tensor.kind == Q4_K && tensor.bytes >= ISOLATED_Q4_BYTES;
+        let isolate = matches!(tensor.kind, Q4_K | MXFP4) && tensor.bytes >= ISOLATED_Q4_BYTES;
         let groups = if isolate
             && tensor.dims.len() == 3
             && tensor.dims[2] == 384
@@ -2819,6 +2951,7 @@ pub(super) struct DeepSeekExecutor {
     checkpoint_tag: [u8; 32],
     model_modified: (u64, u32),
     model_identity: [u8; 32],
+    mxfp4_decode_fast_lookup: bool,
     _context: Context,
     model: Model,
     speculative: EngineSpeculativeSettings,
@@ -2878,15 +3011,19 @@ impl DeepSeekExecutor {
             .transpose()?;
         let initial_ssd_spans = ssd_plan
             .as_ref()
-            .map(|_| deepseek_token_model_spans(&model).map(|spans| spans.ranges))
+            .map(|_| deepseek_token_model_spans(&model))
             .transpose()?;
-        let spans = initial_ssd_spans.as_deref();
+        let spans = initial_ssd_spans
+            .as_ref()
+            .map(|spans| (spans.ranges.as_slice(), spans.max_tensor_bytes));
         let admission = if let Some(plan) = &ssd_plan {
             plan.admission_bytes
         } else {
             resident_deepseek_admission_bytes(&model, context, prefill_chunk)?
         };
         let context_handle = Context::open(&model, quality, ssd.enabled, admission, spans)?;
+        let mxfp4_decode_fast_lookup =
+            mxfp4_decode_fast_lookup_allowed(&model, &weights, quality, ssd.enabled);
         let steering = Steering::load(&model, steering)?;
         let session = Session::new(
             &model,
@@ -2977,6 +3114,7 @@ impl DeepSeekExecutor {
             checkpoint_tag: [0; 32],
             model_modified,
             model_identity,
+            mxfp4_decode_fast_lookup,
             _context: context_handle,
             model,
             speculative,
@@ -3015,13 +3153,24 @@ impl DeepSeekExecutor {
         if let Some(dspark) = &mut self.dspark {
             dspark.begin_capture();
         }
-        if self.ssd.is_some() {
+        let fast_lookup = self.mxfp4_decode_fast_lookup
-            self.encode_streaming_token(token as u32)?;
+            && (self.session.position >= 2048
+                || !environment_present(
+                    c"DS4_METAL_DISABLE_PRE_M5_DECODE_EARLY_PIPELINE_FAST_LOOKUP",
+                ));
+        let previous_fast_lookup =
+            unsafe { ds4_gpu_set_decode_pipeline_fast_lookup(i32::from(fast_lookup)) };
+        let encoded = if self.ssd.is_some() {
+            self.encode_streaming_token(token as u32)
         } else {
-            let commands = Commands::begin()?;
+            (|| {
-            self.encode_token(token as u32)?;
+                let commands = Commands::begin()?;
-            commands.finish()?;
+                self.encode_token(token as u32)?;
-        }
+                commands.finish()
+            })()
+        };
+        unsafe { ds4_gpu_set_decode_pipeline_fast_lookup(previous_fast_lookup) };
+        encoded?;
         self.session.scratch.logits.read_f32(&mut self.logits)?;
         self.session.position += 1;
         self.tokens.push(token);
@@ -3782,9 +3931,9 @@ impl DeepSeekExecutor {
             self.model.shape.model,
             ModelChoice::DeepSeekV4Flash | ModelChoice::DeepSeekV4Flash0731
         ) && self.weights.layers.first().is_some_and(|layer| {
-            layer.expert_gate.kind == Q4_K
+            matches!(layer.expert_gate.kind, Q4_K | MXFP4)
-                && layer.expert_up.kind == Q4_K
+                && layer.expert_up.kind == layer.expert_gate.kind
-                && layer.expert_down.kind == Q4_K
+                && layer.expert_down.kind == layer.expert_gate.kind
         }) {
             64
         } else {
@@ -3899,6 +4048,15 @@ impl DeepSeekExecutor {
                 "DeepSeek prefill token mapping",
             )?;
         }
+        let mut pread = if prefill_pread_enabled() {
+            self.ssd
+                .as_ref()
+                .zip(self.weights.layers.first())
+                .map(|(ssd, weights)| start_prefill_pread(&self.model, ssd, weights, 0, rows))
+                .transpose()?
+        } else {
+            None
+        };
         let commands = Commands::begin()?;
         call(
             unsafe {
@@ -3926,16 +4084,11 @@ impl DeepSeekExecutor {
             .enumerate()
         {
             let started = Instant::now();
-            let layer_selected_addr = self.ssd.is_some()
+            if let Some(job) = pread.take() {
-                && unsafe {
+                job.finish()?;
-                    ds4_gpu_stream_prefill_batch_selected_addr_enabled(
+            }
-                        rows,
+            let layer_selected_addr =
-                        shape.experts as u32,
+                self.ssd.is_some() && prefill_selected_addr(rows, shape, weights);
-                        shape.experts_used as u32,
-                        weights.expert_gate.kind,
-                        weights.expert_down.kind,
-                    ) != 0
-                };
             if let Some(ssd) = &self.ssd {
                 install_deepseek_model_spans(
                     &self.model,
@@ -3948,6 +4101,17 @@ impl DeepSeekExecutor {
                     )?,
                     "DeepSeek prefill layer mapping",
                 )?;
+                if prefill_pread_enabled()
+                    && let Some(next) = self.weights.layers.get(index + 1)
+                {
+                    pread = Some(start_prefill_pread(
+                        &self.model,
+                        ssd,
+                        next,
+                        index as u32 + 1,
+                        rows,
+                    )?);
+                }
             }
             let commands = Commands::begin()?;
             encode_batch_layer(
@@ -5860,49 +6024,6 @@ fn encode_batch_layer(
     )?;
     let gate_row = w.expert_gate.bytes / (w.expert_gate.dims[1] * w.expert_gate.dims[2]);
     let down_row = w.expert_down.bytes / (w.expert_down.dims[1] * w.expert_down.dims[2]);
-    q8_rows(
-        &s.shared_gate,
-        w.shared_gate,
-        shape.embd,
-        shape.ff_expert,
-        &s.norm,
-        rows,
-        map,
-        size,
-    )?;
-    q8_rows(
-        &s.shared_up,
-        w.shared_up,
-        shape.embd,
-        shape.ff_expert,
-        &s.norm,
-        rows,
-        map,
-        size,
-    )?;
-    call(
-        unsafe {
-            ds4_gpu_swiglu_tensor(
-                s.shared_mid.raw(),
-                s.shared_gate.raw(),
-                s.shared_up.raw(),
-                rows * shape.ff_expert as u32,
-                shape.swiglu_clamp,
-                1.0,
-            )
-        },
-        "batch shared expert activation",
-    )?;
-    q8_rows(
-        &s.shared_out,
-        w.shared_down,
-        shape.ff_expert,
-        shape.embd,
-        &s.shared_mid,
-        rows,
-        map,
-        size,
-    )?;
     let mut mid_f16 = false;
     call(
         unsafe {
@@ -5940,7 +6061,66 @@ fn encode_batch_layer(
         },
         "batch routed experts",
     )?;
-    if let Some(steering) = steering.filter(|value| value.ffn_scale != 0.0) {
+    q8_rows(
+        &s.shared_gate,
+        w.shared_gate,
+        shape.embd,
+        shape.ff_expert,
+        &s.norm,
+        rows,
+        map,
+        size,
+    )?;
+    q8_rows(
+        &s.shared_up,
+        w.shared_up,
+        shape.embd,
+        shape.ff_expert,
+        &s.norm,
+        rows,
+        map,
+        size,
+    )?;
+    call(
+        unsafe {
+            ds4_gpu_swiglu_tensor(
+                s.shared_mid.raw(),
+                s.shared_gate.raw(),
+                s.shared_up.raw(),
+                rows * shape.ff_expert as u32,
+                shape.swiglu_clamp,
+                1.0,
+            )
+        },
+        "batch shared expert activation",
+    )?;
+    let ffn_steering = steering.filter(|value| value.ffn_scale != 0.0);
+    let shared_down_f16 = ffn_steering.is_none()
+        && unsafe {
+            ds4_gpu_matmul_q8_0_f16_out_tensor(
+                s.q_half.raw(),
+                map,
+                size,
+                w.shared_down.offset,
+                shape.ff_expert,
+                shape.embd,
+                s.shared_mid.raw(),
+                u64::from(rows),
+            )
+        } != 0;
+    if !shared_down_f16 {
+        q8_rows(
+            &s.shared_out,
+            w.shared_down,
+            shape.ff_expert,
+            shape.embd,
+            &s.shared_mid,
+            rows,
+            map,
+            size,
+        )?;
+    }
+    if let Some(steering) = ffn_steering {
         call(
             unsafe {
                 ds4_gpu_add_tensor(
@@ -5966,6 +6146,21 @@ fn encode_batch_layer(
             },
             "batch steered FFN HC expansion",
         )
+    } else if shared_down_f16 {
+        call(
+            unsafe {
+                ds4_gpu_hc_expand_add_split_half_add_tensor(
+                    s.next_hc.raw(),
+                    s.routed_out.raw(),
+                    s.q_half.raw(),
+                    s.after_attention_hc.raw(),
+                    s.hc_split.raw(),
+                    shape.embd as u32,
+                    shape.hc as u32,
+                )
+            },
+            "batch half shared FFN HC expansion",
+        )
     } else {
         call(
             unsafe {
@@ -8170,10 +8365,10 @@ fn check(result: i32, operation: &str) -> Result<(), String> {
 mod tests {
     use super::{
         compression_ratio, effective_prefill_cap, effective_raw_cap,
-        estimated_deepseek_runtime_bytes, finish_deepseek_model_spans, raw_batch_span,
+        estimated_deepseek_runtime_bytes, finish_deepseek_model_spans,
-        raw_decode_span,
+        gpu::ds4_gpu_print_memory_report, raw_batch_span, raw_decode_span,
     };
-    use crate::engine::{FLASH, PRO};
+    use crate::engine::{FLASH, MXFP4, PRO};
 
     fn installed_artifacts(
         model: crate::model::ModelChoice,
@@ -8281,7 +8476,7 @@ mod tests {
     }
 
     #[test]
-    #[ignore = "requires the installed 0731 Flash GGUF and an Apple M5 device"]
+    #[ignore = "requires a 0731 Flash GGUF and an Apple M5 device"]
     fn flash_0731_m5_decode_performance_gate() {
         use super::{DeepSeekExecutor, Digest, Sha256, configure_sources};
         use crate::engine::{ChatTurn, Model};
@@ -8292,10 +8487,28 @@ mod tests {
         use std::time::Instant;
 
         configure_sources().unwrap();
-        let path = installed_artifacts(ModelChoice::DeepSeekV4Flash0731, false, false).model;
+        let path = std::env::var_os("DS4SERVER_BENCH_MODEL")
+            .map(std::path::PathBuf::from)
+            .unwrap_or_else(|| {
+                installed_artifacts(ModelChoice::DeepSeekV4Flash0731, false, false).model
+            });
+        let prompt_content = std::env::var_os("DS4SERVER_BENCH_CHAT_PROMPT_FILE")
+            .map(|path| std::fs::read_to_string(path).unwrap())
+            .unwrap_or_else(|| "Count from one to two hundred, spelling out every number.".into());
+        let frontier = std::env::var("DS4SERVER_BENCH_FRONTIER")
+            .ok()
+            .map(|value| value.parse::<usize>().unwrap());
         let run = || {
             let model = Model::open_main(&path, ModelChoice::DeepSeekV4Flash0731).unwrap();
-            let prompt = model.render_conversation(
+            let expert_kind = model
+                .main
+                .tensor("blk.4.ffn_gate_exps.weight")
+                .unwrap()
+                .kind;
+            if std::env::var_os("DS4SERVER_BENCH_EXPECT_MXFP4").is_some() {
+                assert_eq!(expert_kind, MXFP4);
+            }
+            let mut prompt = model.render_conversation(
                 "",
                 &[ChatTurn {
                     user: true,
@@ -8304,17 +8517,24 @@ mod tests {
                     skip_previous_eos: false,
                     reasoning: None,
                     reasoning_complete: true,
-                    content: "Count from one to two hundred, spelling out every number.".into(),
+                    content: prompt_content.clone(),
                 }],
                 ReasoningMode::Direct,
             );
+            if let Some(frontier) = frontier {
+                assert!(prompt.len() >= frontier);
+                prompt.truncate(frontier);
+            }
             let eos = model.eos_token();
             let streaming = std::env::var_os("DS4SERVER_BENCH_SSD").is_some();
+            let context = frontier
+                .map(|frontier| u32::try_from(frontier + 129).unwrap())
+                .unwrap_or(4_096);
             let mut executor = DeepSeekExecutor::open(
                 model,
-                4096,
+                context,
                 false,
-                4096,
+                context,
                 100,
                 EngineSpeculativeSettings {
                     mtp_draft_tokens: 1,
@@ -8330,11 +8550,11 @@ mod tests {
                 EngineSsdSettings {
                     enabled: streaming,
                     cold: false,
-                    cache_experts: if streaming { 4_096 } else { 0 },
+                    cache_experts: 0,
                     cache_bytes: 0,
                     full_layers: 0,
                     full_layers_set: false,
-                    preload_experts: if streaming { 4_096 } else { 0 },
+                    preload_experts: 0,
                 },
                 EngineSteeringSettings {
                     file: None,
@@ -8380,9 +8600,10 @@ mod tests {
                 .iter()
                 .map(|byte| format!("{byte:02x}"))
                 .collect::<String>();
+            unsafe { ds4_gpu_print_memory_report(c"benchmark".as_ptr()) };
             let stats = executor.execution_stats();
             eprintln!(
-                "DS4SERVER_METAL_PERF mode={} model=flash-0731 context=4096 prompt={} prefill_tps={:.6} measured={measured} seconds={seconds:.6} tps={tokens_per_second:.6} first_ms={:.6} steady_tps={steady_tokens_per_second:.6} p50_ms={p50:.6} p95_ms={p95:.6} cache_entries={} cache_hits={} cache_misses={} pread_bytes={}",
+                "DS4SERVER_METAL_PERF mode={} model=flash-0731 expert_kind={expert_kind} context={context} prompt={} prefill_tps={:.6} measured={measured} seconds={seconds:.6} tps={tokens_per_second:.6} first_ms={:.6} steady_tps={steady_tokens_per_second:.6} p50_ms={p50:.6} p95_ms={p95:.6} cache_entries={} cache_hits={} cache_misses={} pread_bytes={}",
                 if streaming { "ssd" } else { "resident" },
                 prompt.len(),
                 prompt.len() as f64 / prefill_seconds,
@@ -8396,6 +8617,15 @@ mod tests {
             if std::env::var_os("DS4SERVER_BENCH_TOKENS").is_some() {
                 eprintln!("DS4SERVER_METAL_PROMPT_TOKENS {prompt:?}");
                 eprintln!("DS4SERVER_METAL_GENERATED_TOKENS {generated:?}");
+                let output = generated
+                    .iter()
+                    .filter_map(|&token| executor.model.token_bytes(token))
+                    .flatten()
+                    .collect::<Vec<_>>();
+                eprintln!(
+                    "DS4SERVER_METAL_GENERATED_TEXT {:?}",
+                    String::from_utf8_lossy(&output)
+                );
             }
             assert!(steady_tokens_per_second.is_finite() && steady_tokens_per_second > 0.0);
             steady_tokens_per_second
@@ -8436,7 +8666,7 @@ mod tests {
     }
 
     #[test]
-    #[ignore = "requires the installed 0731 Flash GGUF and an Apple M5 device"]
+    #[ignore = "requires a 0731 Flash GGUF and an Apple M5 device"]
     fn flash_0731_long_context_crosses_indexed_prefill_boundary() {
         use super::{DeepSeekExecutor, argmax, configure_sources};
         use crate::engine::{ChatTurn, Model};
@@ -8446,7 +8676,11 @@ mod tests {
         };
 
         configure_sources().unwrap();
-        let path = installed_artifacts(ModelChoice::DeepSeekV4Flash0731, false, false).model;
+        let path = std::env::var_os("DS4SERVER_BENCH_MODEL")
+            .map(std::path::PathBuf::from)
+            .unwrap_or_else(|| {
+                installed_artifacts(ModelChoice::DeepSeekV4Flash0731, false, false).model
+            });
         let model = Model::open_main(&path, ModelChoice::DeepSeekV4Flash0731).unwrap();
         let prompt = model.render_conversation(
             "",
@@ -8462,6 +8696,7 @@ mod tests {
             ReasoningMode::Direct,
         );
         assert!(prompt.len() > 4_096 && prompt.len() < 8_192);
+        let streaming = std::env::var_os("DS4SERVER_BENCH_SSD").is_some();
         let mut executor = DeepSeekExecutor::open(
             model,
             8_192,
@@ -8480,13 +8715,13 @@ mod tests {
                 dspark_exact_sampling: false,
             },
             EngineSsdSettings {
-                enabled: false,
+                enabled: streaming,
                 cold: false,
-                cache_experts: 0,
+                cache_experts: if streaming { 4_096 } else { 0 },
                 cache_bytes: 0,
                 full_layers: 0,
                 full_layers_set: false,
-                preload_experts: 0,
+                preload_experts: if streaming { 4_096 } else { 0 },
             },
             EngineSteeringSettings {
                 file: None,
@@ -8495,6 +8730,9 @@ mod tests {
             },
         )
         .unwrap();
+        let cancelled = executor.prefill(&prompt, |_| false).unwrap();
+        assert!(cancelled < prompt.len());
+        executor.reset().unwrap();
         assert_eq!(executor.prefill(&prompt, |_| true).unwrap(), prompt.len());
         executor.eval(argmax(executor.logits())).unwrap();
         assert_eq!(executor.position(), prompt.len() as u32 + 1);
@@ -8653,7 +8891,9 @@ mod tests {
 
         configure_sources().unwrap();
         let artifacts = installed_artifacts(ModelChoice::DeepSeekV4Flash0731, false, true);
-        let main_path = artifacts.model;
+        let main_path = std::env::var_os("DS4SERVER_BENCH_MODEL")
+            .map(std::path::PathBuf::from)
+            .unwrap_or(artifacts.model);
         let support_path = artifacts.mtp.unwrap();
         let mut model = Model::open_main(&main_path, ModelChoice::DeepSeekV4Flash0731).unwrap();
         let support = Gguf::open(&support_path).unwrap();
@@ -8672,6 +8912,7 @@ mod tests {
             }],
             crate::settings::ReasoningMode::Direct,
         );
+        let streaming = std::env::var_os("DS4SERVER_BENCH_SSD").is_some();
         let mut executor = DeepSeekExecutor::open(
             model,
             64,
@@ -8690,13 +8931,13 @@ mod tests {
                 dspark_exact_sampling: false,
             },
             EngineSsdSettings {
-                enabled: false,
+                enabled: streaming,
                 cold: false,
-                cache_experts: 0,
+                cache_experts: if streaming { 4_096 } else { 0 },
                 cache_bytes: 0,
                 full_layers: 0,
                 full_layers_set: false,
-                preload_experts: 0,
+                preload_experts: if streaming { 4_096 } else { 0 },
             },
             EngineSteeringSettings {
                 file: None,
@@ -8721,7 +8962,9 @@ mod tests {
         }
         let dspark = executor.dspark.as_ref().unwrap();
         assert!(dspark.drafted > 0);
-        assert!(dspark.accepted > 0);
+        if std::env::var_os("DS4SERVER_BENCH_MODEL").is_none() {
+            assert!(dspark.accepted > 0);
+        }
 
         executor.reset().unwrap();
         executor.dspark.as_mut().unwrap().strict = true;
@@ -8793,10 +9036,19 @@ mod tests {
 
         configure_sources().unwrap();
         let artifacts = installed_artifacts(ModelChoice::DeepSeekV4Flash0731, false, true);
-        validate_engine_artifacts(ModelChoice::DeepSeekV4Flash0731, false, true, &artifacts)
+        let main_path = std::env::var_os("DS4SERVER_BENCH_MODEL")
-            .unwrap();
+            .map(std::path::PathBuf::from)
-        let mut model =
+            .unwrap_or_else(|| {
-            Model::open_main(&artifacts.model, ModelChoice::DeepSeekV4Flash0731).unwrap();
+                validate_engine_artifacts(
+                    ModelChoice::DeepSeekV4Flash0731,
+                    false,
+                    true,
+                    &artifacts,
+                )
+                .unwrap();
+                artifacts.model.clone()
+            });
+        let mut model = Model::open_main(&main_path, ModelChoice::DeepSeekV4Flash0731).unwrap();
         let support = Gguf::open(artifacts.mtp.as_ref().unwrap()).unwrap();
         model.support_kind = Some(validate_support(&support, &model.shape).unwrap());
         model.support = Some(support);
@@ -8813,6 +9065,7 @@ mod tests {
             }],
             ReasoningMode::Direct,
         );
+        let streaming = std::env::var_os("DS4SERVER_BENCH_SSD").is_some();
         let mut executor = DeepSeekExecutor::open(
             model,
             64,
@@ -8831,13 +9084,13 @@ mod tests {
                 dspark_exact_sampling: true,
             },
             EngineSsdSettings {
-                enabled: false,
+                enabled: streaming,
                 cold: false,
-                cache_experts: 0,
+                cache_experts: if streaming { 4_096 } else { 0 },
                 cache_bytes: 0,
                 full_layers: 0,
                 full_layers_set: false,
-                preload_experts: 0,
+                preload_experts: if streaming { 4_096 } else { 0 },
             },
             EngineSteeringSettings {
                 file: None,
@@ -8862,6 +9115,7 @@ mod tests {
             )
             .unwrap();
         assert!(!cycle.is_empty());
+        assert!(executor.dspark.as_ref().unwrap().drafted > 0);
         assert!(executor.logits().iter().all(|logit| logit.is_finite()));
         assert!(executor.session.position >= prompt.len() as u32 + cycle.len() as u32);
     }

src/engine/metal/glm.rs

@@ -434,13 +434,17 @@ impl GlmExecutor {
             .as_ref()
             .map(|plan| glm_streaming_model_spans(&model, &weights, plan))
             .transpose()?;
+        let context_spans = model_spans
+            .as_ref()
+            .map(|(spans, max_tensor_bytes)| (spans.as_slice(), *max_tensor_bytes));
         let context_handle = Context::open(
             &model,
             quality,
             effective_ssd.enabled,
             admission,
-            model_spans.as_deref(),
+            context_spans,
         )?;
+        let model_spans = model_spans.map(|(spans, _)| spans);
         configure_streaming(&model, &weights, streaming.as_ref())?;
         let scratch = GlmScratch::allocate(&model, context)?;
         let caches = (0..weights.layers.len())
@@ -2851,7 +2855,7 @@ fn glm_streaming_model_spans(
     model: &Model,
     weights: &GlmWeights,
     plan: &GlmStreamingPlan,
-) -> Result<Vec<(u64, u64)>, String> {
+) -> Result<(Vec<(u64, u64)>, u64), String> {
     let full_before = model
         .shape
         .leading_dense
@@ -2890,6 +2894,11 @@ fn glm_streaming_model_spans(
         })
         .map(|(_, tensor)| (tensor.offset, tensor.bytes))
         .collect::<Vec<_>>();
+    let max_tensor_bytes = spans
+        .iter()
+        .map(|(_, bytes)| *bytes)
+        .max()
+        .ok_or("GLM SSD streaming found no resident model tensors")?;
     spans.sort_unstable_by_key(|span| span.0);
     let mut merged: Vec<(u64, u64)> = Vec::new();
     for (offset, bytes) in spans {
@@ -2903,10 +2912,7 @@ fn glm_streaming_model_spans(
         }
         merged.push((offset, bytes));
     }
-    if merged.is_empty() {
+    Ok((merged, max_tensor_bytes))
-        return Err("GLM SSD streaming found no resident model tensors".into());
-    }
-    Ok(merged)
 }
 
 fn glm_layer_model_spans(model: &Model, layer: u32) -> Result<Vec<(u64, u64)>, String> {

src/engine/metal/gpu.rs

@@ -95,6 +95,8 @@ unsafe extern "C" {
     pub(super) fn ds4_gpu_flush_commands() -> i32;
     pub(super) fn ds4_gpu_device_is_pre_m5_apple_silicon() -> i32;
     pub(super) fn ds4_gpu_device_is_m5_apple_silicon() -> i32;
+    #[cfg(test)]
+    pub(super) fn ds4_gpu_print_memory_report(label: *const c_char);
     pub(super) fn ds4_gpu_set_decode_pipeline_fast_lookup(enabled: i32) -> i32;
     pub(super) fn ds4_gpu_parallel_ffn_start(
         gate: *mut GpuTensor,
@@ -271,6 +273,16 @@ unsafe extern "C" {
         x: *const GpuTensor,
         rows: u64,
     ) -> i32;
+    pub(super) fn ds4_gpu_matmul_q8_0_f16_out_tensor(
+        out: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        weight: u64,
+        input: u64,
+        output: u64,
+        x: *const GpuTensor,
+        rows: u64,
+    ) -> i32;
     pub(super) fn ds4_gpu_matmul_quant_tensor(
         out: *mut GpuTensor,
         map: *const c_void,
@@ -1356,6 +1368,15 @@ unsafe extern "C" {
         embd: u32,
         hc: u32,
     ) -> i32;
+    pub(super) fn ds4_gpu_hc_expand_add_split_half_add_tensor(
+        out: *mut GpuTensor,
+        block: *const GpuTensor,
+        add_half: *const GpuTensor,
+        residual: *const GpuTensor,
+        split: *const GpuTensor,
+        embd: u32,
+        hc: u32,
+    ) -> i32;
     pub(super) fn ds4_gpu_attention_output_low_q8_tensor(
         low: *mut GpuTensor,
         map: *const c_void,
@@ -1515,15 +1536,14 @@ impl Context {
         quality: bool,
         ssd_streaming: bool,
         admission_bytes: u64,
-        model_spans: Option<&[(u64, u64)]>,
+        model_spans: Option<(&[(u64, u64)], u64)>,
     ) -> Result<Self, String> {
         check(unsafe { ds4_gpu_init() }, "Metal initialization")?;
         unsafe {
             ds4_gpu_set_glm_model(model.shape.family == ModelFamily::Glm);
             ds4_gpu_set_ssd_streaming(ssd_streaming);
-            // DS4 only enables this cache for the pre-M5 MXFP4 decode path.
+            // Decode enables this only around DS4's eligible resident pre-M5
-            // Rust does not accept MXFP4 weights yet, so keep the global
+            // MXFP4 token path; all other work starts from the portable path.
-            // native switch explicitly disabled until that path is admitted.
             ds4_gpu_set_decode_pipeline_fast_lookup(0);
         }
         let recommended = unsafe { ds4_gpu_recommended_working_set_size() };
@@ -1536,7 +1556,7 @@ impl Context {
             ));
         }
         let data_offset = model.main.data_offset();
-        let mapped = if let Some(spans) = model_spans {
+        let mapped = if let Some((spans, max_tensor_bytes)) = model_spans {
             let (offsets, sizes): (Vec<_>, Vec<_>) = spans.iter().copied().unzip();
             unsafe {
                 ds4_gpu_set_model_map_spans(
@@ -1545,7 +1565,7 @@ impl Context {
                     offsets.as_ptr(),
                     sizes.as_ptr(),
                     spans.len() as u32,
-                    model.main.max_tensor_bytes(),
+                    max_tensor_bytes,
                 )
             }
         } else {

src/engine/validation.rs

@@ -1036,6 +1036,21 @@ fn first_u32s<'a>(model: &'a Gguf, keys: &[&str]) -> Result<&'a [u32], String> {
 mod tests {
     use super::*;
 
+    #[test]
+    fn mxfp4_is_limited_to_routed_experts() {
+        let tensor = Tensor {
+            kind: MXFP4,
+            dims: vec![32, 1, 1],
+            offset: 0,
+            bytes: 17,
+        };
+        assert!(validate_tensor("routed", &tensor, ROUTED, &[32, 1, 1]).is_ok());
+        assert_eq!(
+            validate_tensor("dense", &tensor, DENSE, &[32, 1, 1]).unwrap_err(),
+            "tensor dense has unsupported type 39"
+        );
+    }
+
     #[test]
     fn installed_ds4_fixture_opens_and_renders_a_prompt() {
         let path = crate::model::engine_artifacts(