hugo/DS4Server
1
0
Fork 0
Code Issues 2 Pull Requests Actions Packages Projects Releases Wiki Activity

Add end-to-end MXFP4 Metal support

Browse Source
This commit is contained in:
Georg Bauer
2026-08-30 20:42:23 +02:00
parent 3977261bd3
commit 28231e1faf
7 changed files with 514 additions and 124 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
src/engine.rs
View File

@@ -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];

119
src/engine/gguf.rs
View File

@@ -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]) {

446
src/engine/metal.rs
View File

@@ -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() {
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::<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)
.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);
}

18
src/engine/metal/glm.rs
View File

@@ -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> {

32
src/engine/metal/gpu.rs
View File

@@ -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 {

15
src/engine/validation.rs
View File

@@ -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(