From 28231e1faf055c487f0d2ec6408fc48f66f33ee8 Mon Sep 17 00:00:00 2001 From: Georg Bauer Date: Sun, 30 Aug 2026 20:42:23 +0200 Subject: [PATCH] Add end-to-end MXFP4 Metal support --- build.rs | 2 + src/engine.rs | 6 +- src/engine/gguf.rs | 119 +++++++++-- src/engine/metal.rs | 446 ++++++++++++++++++++++++++++++--------- src/engine/metal/glm.rs | 18 +- src/engine/metal/gpu.rs | 32 ++- src/engine/validation.rs | 15 ++ 7 files changed, 514 insertions(+), 124 deletions(-) diff --git a/build.rs b/build.rs index b546f52..17c7949 100644 --- a/build.rs +++ b/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"); diff --git a/src/engine.rs b/src/engine.rs index c325479..62177a7 100644 --- a/src/engine.rs +++ b/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]; diff --git a/src/engine/gguf.rs b/src/engine/gguf.rs index d52d314..cb86ef0 100644 --- a/src/engine/gguf.rs +++ b/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, value: &[u8]) { diff --git a/src/engine/metal.rs b/src/engine/metal.rs index 777434c..4ca18db 100644 --- a/src/engine/metal.rs +++ b/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>>, +} + +impl PrefillPread { + fn start(model: &Model, layer: u32, spans: &DeepSeekModelSpans) -> Result { + 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 { + 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() { - self.encode_streaming_token(token as u32)?; + let fast_lookup = self.mxfp4_decode_fast_lookup + && (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)?; - commands.finish()?; - } + (|| { + let commands = Commands::begin()?; + 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 - && layer.expert_up.kind == Q4_K - && layer.expert_down.kind == Q4_K + matches!(layer.expert_gate.kind, Q4_K | MXFP4) + && layer.expert_up.kind == layer.expert_gate.kind + && 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() - && 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 - }; + if let Some(job) = pread.take() { + job.finish()?; + } + let layer_selected_addr = + self.ssd.is_some() && prefill_selected_addr(rows, shape, weights); 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, - raw_decode_span, + estimated_deepseek_runtime_bytes, finish_deepseek_model_spans, + 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::().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::(); + 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::>(); + 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) - .unwrap(); - let mut model = - Model::open_main(&artifacts.model, ModelChoice::DeepSeekV4Flash0731).unwrap(); + let main_path = std::env::var_os("DS4SERVER_BENCH_MODEL") + .map(std::path::PathBuf::from) + .unwrap_or_else(|| { + 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); } diff --git a/src/engine/metal/glm.rs b/src/engine/metal/glm.rs index b1fd27b..88b62ff 100644 --- a/src/engine/metal/glm.rs +++ b/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, 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::>(); + 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, String> { diff --git a/src/engine/metal/gpu.rs b/src/engine/metal/gpu.rs index 0b4da0f..18fd1ea 100644 --- a/src/engine/metal/gpu.rs +++ b/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 { 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 { diff --git a/src/engine/validation.rs b/src/engine/validation.rs index 0a339ef..ae79e4d 100644 --- a/src/engine/validation.rs +++ b/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(