Add end-to-end MXFP4 Metal support

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() {
-        return Err("GLM SSD streaming found no resident model tensors".into());
-    }
-    Ok(merged)
+    Ok((merged, max_tensor_bytes))
 }
 
 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.
-            // Rust does not accept MXFP4 weights yet, so keep the global
-            // native switch explicitly disabled until that path is admitted.
+            // Decode enables this only around DS4's eligible resident pre-M5
+            // MXFP4 token path; all other work starts from the portable path.
             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 {