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

Align DeepSeek and GLM execution with DS4

Browse Source
This commit is contained in:
Georg Bauer committed 2026-09-11 17:18:45 +02:00
1 parent 02db0968ae
commit 48c2f751b4
27 files changed
+4517 -400
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/engine/metal/checkpoint.rs
+1 -3
View File
@@ -302,9 +302,7 @@ impl DeepSeekExecutor {
self.logits = logits;
self.checkpoint_tag = checkpoint_tag;
if let Some(dspark) = &mut self.dspark {
dspark.capture_mask = 0;
dspark.cache_start = 0;
dspark.cache_len = 0;
dspark.reset_cache();
}
Ok(())
}
src/engine/metal/glm.rs
+437 -41
View File
@@ -1008,7 +1008,8 @@ impl GlmExecutor {
)?;
}
std::mem::swap(&mut self.scratch.current, &mut self.scratch.next);
if layer_index + 1 == DECODE_FLUSH_LAYERS {
if glm_decode_flush_layer(layer_index + 1, self.weights.layers.len(), self.ssd.enabled)
{
call(
unsafe { ds4_gpu_flush_commands() },
"flushing the GLM decode graph",
@@ -1419,7 +1420,15 @@ impl GlmExecutor {
causal_range: bool,
) -> Result<(), String> {
let dsa = layer.dsa()?;
self.encode_batch_attention_raw(batch, dsa, cache, pos, rows, selected_count, causal_range)
self.encode_batch_attention_raw(
batch,
dsa,
cache,
pos,
rows,
selected_count,
if causal_range { pos + rows } else { 0 },
)
}
#[allow(clippy::too_many_arguments)]
@@ -1431,7 +1440,7 @@ impl GlmExecutor {
pos: u32,
rows: u32,
selected_count: u32,
causal_range: bool,
dense_limit: u32,
) -> Result<(), String> {
let shape = self.model.shape;
let map = self.model.main.map_ptr().cast();
@@ -1446,7 +1455,12 @@ impl GlmExecutor {
)?;
let mut start = 0;
while start < rows {
let slice = (rows - start).min(INDEXED_PREFILL_ATTN_SLICE);
let (slice, causal_range) = glm_attention_slice(pos + start, rows - start, dense_limit);
let attention_selected_count = if causal_range {
(pos + rows).min(dense_limit)
} else {
selected_count
};
let q = batch
.q
.view(u64::from(start) * q_dim * 4, u64::from(slice) * q_dim * 4)?;
@@ -1471,7 +1485,7 @@ impl GlmExecutor {
cache.kv.raw(),
pos + start,
slice,
selected_count,
attention_selected_count,
self.context,
CACHE_F16,
shape.heads as u32,
@@ -1492,7 +1506,7 @@ impl GlmExecutor {
cache.rope.raw(),
slice,
pos + start,
selected_count,
attention_selected_count,
self.context,
CACHE_F16,
shape.heads as u32,
@@ -2007,6 +2021,29 @@ impl GlmExecutor {
result
}
pub(crate) fn rewind_speculative_output(&mut self, pos: usize) -> Result<(), String> {
if pos > self.tokens.len() {
return Err("GLM emitted token frontier exceeds evaluated tokens".into());
}
if pos == self.tokens.len() {
return Ok(());
}
if self.model.shape.model == ModelChoice::Glm53Flash {
return if self.rewind_glm53_mtp(pos)? {
Ok(())
} else {
Err("GLM could not restore the emitted token frontier".into())
};
}
self.tokens.truncate(pos);
if let Some(mtp) = &mut self.mtp {
mtp.pending = None;
mtp.parent = None;
mtp.rollback = None;
}
Ok(())
}
fn rewind_glm53_mtp(&mut self, pos: usize) -> Result<bool, String> {
let Some((start, first)) = self.mtp.as_ref().and_then(|m| m.rollback) else {
return Ok(false);
@@ -3137,7 +3174,7 @@ impl GlmExecutor {
)?;
let visible = pos + rows;
let causal = pos < dense_limit;
let causal = visible <= dense_limit;
let selected_count = if causal {
visible.min(dense_limit)
} else {
@@ -3250,7 +3287,15 @@ impl GlmExecutor {
},
"projecting batched GLM 5.3 low-rank queries",
)?;
self.encode_batch_attention_raw(batch, weights, cache, pos, rows, selected_count, causal)?;
self.encode_batch_attention_raw(
batch,
weights,
cache,
pos,
rows,
selected_count,
dense_limit,
)?;
glm53_project_rows(
&batch.attn_out,
weights.output,
@@ -3467,7 +3512,7 @@ impl GlmExecutor {
"expanding GLM 5.3 FFN hyperconnections",
)?;
std::mem::swap(&mut self.scratch.hc_current, &mut self.scratch.hc_next);
if ordinal + 1 == DECODE_FLUSH_LAYERS {
if glm_decode_flush_layer(ordinal + 1, self.weights.layers.len(), self.ssd.enabled) {
call(
unsafe { ds4_gpu_flush_commands() },
"flushing GLM 5.3 decode",
@@ -4555,14 +4600,12 @@ impl GlmExecutor {
}
let pos = self.position();
let mut rows =
indexed_prefill_rows(pos, remaining.len(), self.model.shape.indexer_top_k as u32);
if self.model.shape.model == ModelChoice::Glm53Flash {
let dense = glm53_dense_limit(self.context, self.ssd.enabled);
if pos < dense {
rows = rows.min((dense - pos) as usize);
}
}
let rows = indexed_prefill_rows(
self.model.shape.model,
pos,
remaining.len(),
self.model.shape.indexer_top_k as u32,
);
let cancelled = self.eval_batch(
&remaining[..rows],
&mut progress,
@@ -4760,13 +4803,7 @@ impl GlmExecutor {
let shape = self.model.shape;
let map = self.model.main.map_ptr().cast();
let size = self.model.main.len();
let mut dense_limit = glm53_dense_limit(self.context, self.ssd.enabled);
// DS4 routes the entire verification pair through indexed attention
// when it no longer fits the full-attention span. Normal prefill
// already splits at this boundary before entering the batch function.
if pos < dense_limit && pos + rows > dense_limit {
dense_limit = 0;
}
let dense_limit = glm53_dense_limit(self.context, self.ssd.enabled);
// Both HC split/sum and expansion derive their row count from the
// output view, not a rows argument. Match DS4's token-sized views even
// when a short prefill reuses the full 2048-row workspace.
@@ -4930,6 +4967,16 @@ impl GlmExecutor {
},
"expanding batched GLM 5.3 attention hyperconnections",
)?;
#[cfg(test)]
{
commands = compare_glm53_reference_stages(
commands,
index,
pos,
rows as usize * shape.embd as usize,
&[("attn_out", &batch.attn_out)],
)?;
}
glm53_hc_pre(
after_attn,
&batch.ffn_norm,
@@ -4947,6 +4994,16 @@ impl GlmExecutor {
size,
)?;
self.encode_glm53_ffn_batch(batch, layer, index as u32, rows)?;
#[cfg(test)]
{
commands = compare_glm53_reference_stages(
commands,
index,
pos,
rows as usize * shape.embd as usize,
&[("ffn_out", &batch.next)],
)?;
}
if let Some(steering) = &self.steering {
steering.apply(&batch.next, index as u32, rows, false)?;
}
@@ -4963,10 +5020,11 @@ impl GlmExecutor {
},
"expanding batched GLM 5.3 FFN hyperconnections",
)?;
let drain = glm_prefill_flush_layers(rows, read_logits)
&& (index + 1).is_multiple_of(PREFILL_DRAIN_LAYERS)
&& index + 1 < self.weights.layers.len();
if glm_prefill_flush_layers(rows, read_logits) {
if (index + 1).is_multiple_of(PREFILL_DRAIN_LAYERS)
&& index + 1 < self.weights.layers.len()
{
if drain {
commands.finish()?;
commands = Commands::begin()?;
} else {
@@ -4989,11 +5047,15 @@ impl GlmExecutor {
if let Some(views) = &mut hc_views {
views.swap(2, 3);
}
let done = u32::try_from(
u64::from(rows) * (index as u64 + 1) / self.weights.layers.len() as u64,
)
.expect("GLM 5.3 prefill progress is bounded by the batch");
cancelled |= !progress(pos + done);
// DS4 reports completed work at drains, not every encoded layer.
// An asynchronous flush alone does not advance GPU completion.
if drain {
let done = u32::try_from(
u64::from(rows) * (index as u64 + 1) / self.weights.layers.len() as u64,
)
.expect("GLM 5.3 prefill progress is bounded by the batch");
cancelled |= !progress(pos + done);
}
}
commands.finish()?;
@@ -5024,6 +5086,7 @@ impl GlmExecutor {
self.logits = logits;
}
self.tokens.extend_from_slice(tokens);
cancelled |= !progress(self.position());
if let Some(profile) = &self.profile {
profile.write()?;
}
@@ -5513,14 +5576,35 @@ fn glm_prefill_flush_layers(rows: u32, logits_requested: bool) -> bool {
logits_requested && rows > 8
}
fn indexed_prefill_rows(pos: u32, remaining: usize, top_k: u32) -> usize {
fn glm_decode_flush_layer(completed: usize, layers: usize, ssd_streaming: bool) -> bool {
// The resident indexed DS4 graph flushes every four layers, except the
// final layer. Static SSD mappings use the expert loader's own boundaries.
!ssd_streaming && completed < layers && completed.is_multiple_of(DECODE_FLUSH_LAYERS)
}
fn indexed_prefill_rows(model: ModelChoice, pos: u32, remaining: usize, top_k: u32) -> usize {
let mut rows = remaining.min(INDEXED_PREFILL_CHUNK as usize);
if pos < top_k {
// GLM 5.3's active DS4 graph keeps full chunks across both the old
// top-k boundary and the dense/sparse boundary. Only attention is sliced.
if model != ModelChoice::Glm53Flash && pos < top_k {
rows = rows.min((top_k - pos) as usize);
}
rows
}
fn glm_attention_slice(pos: u32, remaining: u32, dense_limit: u32) -> (u32, bool) {
let causal = pos < dense_limit;
let rows = remaining.min(INDEXED_PREFILL_ATTN_SLICE);
(
if causal {
rows.min(dense_limit - pos)
} else {
rows
},
causal,
)
}
fn full_indexer_layer(shape: super::super::Shape, layer: usize) -> bool {
if shape.model == ModelChoice::Glm53Flash {
return layer < (shape.layers - shape.nextn) as usize && !glm53_kda_layer(shape, layer);
@@ -6432,12 +6516,72 @@ fn f32_project_rows(
)
}
// Diagnostic only: no stage reads, extra drains or environment checks in the
// release app. Match original DS4's existing layer0/bootstrap tensor dumps.
#[cfg(test)]
fn compare_glm53_reference_stages(
commands: Commands,
layer: usize,
pos: u32,
values: usize,
tensors: &[(&str, &Buffer)],
) -> Result<Commands, String> {
let Some(prefix) = std::env::var_os("DS4SERVER_GLM_STAGE_REFERENCE") else {
return Ok(commands);
};
let selected_pos = std::env::var("DS4SERVER_GLM_STAGE_POS")
.map_or(Ok(0), |v| v.parse::<u32>())
.map_err(|e| e.to_string())?;
if layer != 0 || pos != selected_pos {
return Ok(commands);
}
commands.finish()?;
for &(name, tensor) in tensors {
let path = format!("{}_{name}-{layer}_pos{pos}.bin", prefix.to_string_lossy());
let bytes = std::fs::read(&path).map_err(|e| format!("{path}: {e}"))?;
if values == 0 || bytes.len() != values * 4 {
return Err(format!(
"reference tensor geometry mismatch: {path}: {} bytes, expected {}",
bytes.len(),
values * 4
));
}
let reference: Vec<f32> = bytes
.chunks_exact(4)
.map(|b| f32::from_le_bytes(b.try_into().unwrap()))
.collect();
let mut actual = vec![0.0; reference.len()];
tensor.read_f32(&mut actual)?;
if !actual.iter().chain(&reference).all(|v| v.is_finite()) {
return Err(format!("nonfinite stage {name}"));
}
let max_abs = actual
.iter()
.zip(&reference)
.map(|(a, b)| (a - b).abs())
.fold(0.0_f32, f32::max);
let rms = (actual
.iter()
.zip(&reference)
.map(|(a, b)| f64::from(a - b).powi(2))
.sum::<f64>()
/ actual.len() as f64)
.sqrt();
eprintln!(
"{}",
serde_json::json!({"event":"glm_stage_compare", "name":name, "layer":layer, "pos":pos, "values":actual.len(), "max_abs":max_abs, "rms":rms})
);
}
Commands::begin()
}
#[cfg(test)]
mod tests {
use super::{
GlmExecutor, argmax, dynamic_expert_budget, full_indexer_layer, glm_prefill_flush_layers,
glm53_dense_compact_prefill, glm53_dense_limit, indexed_prefill_rows,
live_prefix_rewind_target, streaming_token_prefill_eligible,
GlmExecutor, argmax, dynamic_expert_budget, full_indexer_layer, glm_attention_slice,
glm_decode_flush_layer, glm_prefill_flush_layers, glm53_dense_compact_prefill,
glm53_dense_limit, indexed_prefill_rows, live_prefix_rewind_target,
streaming_token_prefill_eligible,
};
use crate::engine::{GLM, Model, ReasoningMode};
use crate::model::ModelChoice;
@@ -6464,9 +6608,30 @@ mod tests {
assert!(!streaming_token_prefill_eligible(8192, 0, 65, 64));
assert!(!streaming_token_prefill_eligible(8192, 8180, 13, 64));
assert!(!streaming_token_prefill_eligible(65_536, 4096, 1, 64));
assert_eq!(indexed_prefill_rows(0, 5000, 2048), 2048);
assert_eq!(indexed_prefill_rows(2048, 5000, 2048), 2048);
assert_eq!(indexed_prefill_rows(0, 17, 2048), 17);
assert_eq!(
indexed_prefill_rows(ModelChoice::Glm52, 0, 5000, 2048),
2048
);
assert_eq!(
indexed_prefill_rows(ModelChoice::Glm52, 2048, 5000, 2048),
2048
);
assert_eq!(indexed_prefill_rows(ModelChoice::Glm52, 0, 17, 2048), 17);
assert_eq!(
indexed_prefill_rows(ModelChoice::Glm52, 9, 2626, 2048),
2039
);
for pos in [9, 2047, 4095, 8191] {
assert_eq!(
indexed_prefill_rows(ModelChoice::Glm53Flash, pos, 2626, 2048),
2048
);
}
for boundary in [4096, 8192] {
assert_eq!(glm_attention_slice(boundary - 1, 2, boundary), (1, true));
assert_eq!(glm_attention_slice(boundary, 1, boundary), (1, false));
assert_eq!(glm_attention_slice(9, 2048, boundary), (2048, true));
}
assert_eq!(glm53_dense_limit(32_768, false), 4096);
assert_eq!(glm53_dense_limit(32_768, true), 8192);
assert_eq!(glm53_dense_limit(65_536, true), 4096);
@@ -6477,6 +6642,212 @@ mod tests {
assert!(!glm_prefill_flush_layers(2048, false));
}
#[test]
fn glm_resident_decode_flushes_periodically_without_a_final_or_ssd_flush() {
for layers in [1, 4, 5, 76, 78] {
let actual = (1..=layers)
.filter(|&completed| glm_decode_flush_layer(completed, layers, false))
.collect::<Vec<_>>();
assert_eq!(actual, (4..layers).step_by(4).collect::<Vec<_>>());
assert!((1..=layers).all(|completed| !glm_decode_flush_layer(completed, layers, true)));
}
}
#[test]
#[ignore = "installed GLM and DS4SERVER_GLM_REFERENCE_LOG from the standalone reference"]
fn glm53_reference_prompt_tokens_match_shared_runtime() {
let path = std::env::var("DS4SERVER_GLM_REFERENCE_LOG").unwrap();
let events = std::fs::read_to_string(path)
.unwrap()
.lines()
.map(|line| serde_json::from_str::<serde_json::Value>(line).unwrap())
.collect::<Vec<_>>();
let start = events
.iter()
.find(|e| e["event"] == "reference_start")
.unwrap();
assert_eq!(start["family"], "glm");
let model = Model::open_main(
std::path::Path::new(start["model"].as_str().unwrap()),
ModelChoice::Glm53Flash,
)
.unwrap();
let readme = std::fs::read_to_string(start["readme"].as_str().unwrap()).unwrap();
let prompts = [
"Reply with exactly: OK".to_owned(),
format!("Give a summary of the following text:\n\n{readme}"),
"Tell me a complete short story about a lighthouse keeper. Do not ask questions.".into(),
"Write a Python function is_prime(n: int) -> bool, followed by five assert examples. No tools.".into(),
];
let ids = |v: &serde_json::Value| {
v.as_array()
.unwrap()
.iter()
.map(|v| v.as_i64().unwrap() as i32)
.collect::<Vec<_>>()
};
let mut frontier = Vec::new();
for (index, prompt) in prompts.iter().enumerate() {
let expected = if index <= 1 {
model.render_prompt("", prompt, ReasoningMode::Low)
} else {
let mut expected = frontier.clone();
expected.extend(model.tokenizer.encode_continuation(
prompt,
ReasoningMode::Low,
false,
));
expected
};
let receipt = events
.iter()
.find(|e| e["event"] == "reference_prefill" && e["turn"] == index)
.unwrap();
assert_eq!(
ids(&receipt["tokens"]),
expected,
"reference turn {index} prompt differs"
);
let cached = if index <= 1 { 9 } else { frontier.len() };
assert_eq!(
receipt["cached_tokens"], cached,
"bootstrap/continuation differs"
);
let result = events
.iter()
.find(|e| e["event"] == "reference_result" && e["turn"] == index)
.unwrap();
if index != 0 {
assert_eq!(result["finish_reason"], "stop");
}
frontier = expected;
frontier.extend(ids(&result["token_ids"]));
}
}
#[test]
#[ignore = "requires the installed GLM 5.3 Flash Q2 checkpoint and Apple Metal"]
fn glm53_reference_logits_replay_separates_sampling_from_execution() {
use crate::engine::{Rng, sample};
let events: Vec<serde_json::Value> =
std::fs::read_to_string(std::env::var("DS4SERVER_GLM_REFERENCE_LOG").unwrap())
.unwrap()
.lines()
.map(|l| serde_json::from_str(l).unwrap())
.collect();
let event = |name: &str| {
events
.iter()
.find(|e| e["event"] == name && (name == "reference_start" || e["turn"] == 1))
.unwrap()
};
let start = event("reference_start");
assert_eq!(start["family"], "glm");
assert_eq!(start["acceleration"], false);
assert_eq!(start["seed"], 42);
let trace = event("reference_logits_trace");
let path: std::ffi::OsString = serde_json::from_value(trace["path"].clone()).unwrap();
let bytes = std::fs::read(std::path::PathBuf::from(path)).unwrap();
let n = trace["vocab"].as_u64().unwrap() as usize;
assert_eq!(bytes.len(), 32 * n * 4);
let values: Vec<f32> = bytes
.chunks_exact(4)
.map(|b| f32::from_le_bytes(b.try_into().unwrap()))
.collect();
let token_ids = |v: &serde_json::Value| {
v.as_array()
.unwrap()
.iter()
.map(|v| v.as_i64().unwrap() as i32)
.collect::<Vec<_>>()
};
let prompt = token_ids(&event("reference_prefill")["tokens"]);
let expected = token_ids(&event("reference_result")["token_ids"]);
assert!(expected.len() >= 32);
// First feed the exact C-produced rows through the production Rust
// sampler, so a model-graph difference cannot masquerade as RNG drift.
let mut host_rng = Rng::new(42);
let host: Vec<i32> = values
.chunks_exact(n)
.map(|row| sample(row, 0.6, 0.95, 0.0, 0, &mut host_rng))
.collect();
assert_eq!(
host,
expected[..32],
"Rust sampler differs on original DS4 logits"
);
super::super::configure_sources().unwrap();
let model = Model::open_main(
std::path::Path::new(start["model"].as_str().unwrap()),
ModelChoice::Glm53Flash,
)
.unwrap();
assert_eq!(n, model.shape.vocab as usize);
let mut executor = GlmExecutor::open(
model,
32768,
false,
EngineSsdSettings {
enabled: false,
cold: false,
cache_experts: 0,
cache_bytes: 0,
full_layers: 0,
full_layers_set: false,
preload_experts: 0,
},
)
.unwrap();
assert_eq!(event("reference_prefill")["cached_tokens"], 9);
for part in [&prompt[..9], &prompt[9..]] {
assert_eq!(
executor
.prefill(part, |pos| {
eprintln!("reference replay prefill {pos}");
true
})
.unwrap(),
part.len()
);
}
let mut gpu_rng = Rng::new(42);
let mut different = Vec::new();
for (step, row) in values.chunks_exact(n).enumerate() {
let actual = executor.logits();
let chosen = sample(actual, 0.6, 0.95, 0.0, 0, &mut gpu_rng);
let max_abs = actual
.iter()
.zip(row)
.map(|(a, b)| (a - b).abs())
.fold(0.0_f32, f32::max);
let rms = (actual
.iter()
.zip(row)
.map(|(a, b)| f64::from(a - b).powi(2))
.sum::<f64>()
/ n as f64)
.sqrt();
eprintln!(
"{}",
serde_json::json!({"event":"glm_logits_compare", "step":step, "expected":expected[step], "actual":chosen, "max_abs":max_abs, "rms":rms})
);
if chosen != expected[step]
|| !actual
.iter()
.zip(row)
.all(|(a, b)| a.to_bits() == b.to_bits())
{
different.push(step);
}
// Never diverge the context: both backends see the reference token.
executor.eval(expected[step]).unwrap();
}
assert!(
different.is_empty(),
"logits or sampled outputs differ on identical history at steps {different:?}"
);
}
#[test]
#[ignore = "requires the installed GLM 5.3 Flash Q2 checkpoint and Apple Metal"]
fn glm53_m5_text_performance_gate() {
@@ -7171,12 +7542,30 @@ mod tests {
.unwrap();
let prepare = |executor: &mut GlmExecutor| {
executor.reset().unwrap();
let mut reported = Vec::new();
executor
.prefill(&prompt, |pos| {
reported.push(pos);
eprintln!("verify2 prefill {pos}");
true
})
.unwrap();
if prompt.len() <= super::INDEXED_PREFILL_CHUNK as usize {
let layers = executor.weights.layers.len();
let mut expected = vec![0];
if prompt.len() > 8 {
expected.extend(
(super::PREFILL_DRAIN_LAYERS..layers)
.step_by(super::PREFILL_DRAIN_LAYERS)
.map(|done| (prompt.len() * done / layers) as u32),
);
}
expected.push(prompt.len() as u32);
assert_eq!(
reported, expected,
"progress must follow completed GPU work"
);
}
};
let recurrent = |executor: &GlmExecutor| {
executor
@@ -7230,6 +7619,9 @@ mod tests {
eprintln!("partial prefill HC workspace guards passed");
}
let before = recurrent(&executor);
assert!(executor.rewind_speculative_output(frontier + 1).is_err());
executor.rewind_speculative_output(frontier).unwrap();
assert_eq!(recurrent(&executor), before);
let first = argmax(executor.logits());
executor.eval(first).unwrap();
let after_first = recurrent(&executor);
@@ -7259,7 +7651,11 @@ mod tests {
executor.mtp.as_ref().unwrap().rollback,
Some((start, first))
);
assert!(executor.rewind_glm53_mtp(start as usize + keep).unwrap());
// The UI/headless consumer uses this same path to remove an
// MTP-returned stop token before retaining the ongoing chat.
executor
.rewind_speculative_output(start as usize + keep)
.unwrap();
assert_eq!(executor.position(), start + keep as u32);
assert!(
recurrent(&executor) == *if keep == 0 { &before } else { &after_first },
src/engine/metal/gpu.rs
+3
View File
@@ -45,6 +45,9 @@ pub(super) struct QwenKernelArgs {
pub(super) struct GpuCanarySample {
pub(super) scheduled_seconds: f64,
pub(super) completed_seconds: f64,
pub(super) gpu_wait_seconds: f64,
pub(super) gpu_interval_seconds: f64,
pub(super) host_return_seconds: f64,
}
#[derive(Clone, Copy, Default)]
src/engine/metal/markov.rs
+327
View File
@@ -0,0 +1,327 @@
//! DSpark's CPU Markov head: persistent workers, DS4 row partition/reduction.
use super::{F16, F32, Gguf, Q8_0, Weight, dense_dot_bytes, dense_dot_q8, quantize_q8_activation};
use memmap2::Mmap;
use std::sync::{Arc, mpsc};
use std::thread::{self, JoinHandle};
type Best = (usize, f32);
struct Input {
values: Vec<f32>,
quantized: Option<(Vec<i8>, Vec<f32>)>,
logits: Vec<f32>,
}
struct Matrix {
map: Arc<Mmap>,
weight: Weight,
width: usize,
rows: usize,
row_bytes: usize,
}
impl Matrix {
fn best(&self, input: &Input, start: usize, end: usize) -> Best {
let bytes = &self.map[self.weight.offset as usize..];
let mut best = (start, -f32::MAX);
for token in start..end {
let row = &bytes[token * self.row_bytes..(token + 1) * self.row_bytes];
let dot = input.quantized.as_ref().map_or_else(
|| dense_dot_bytes(self.weight.kind, self.width, row, &input.values),
|(values, scales)| dense_dot_q8(row, values, scales, self.width),
);
let score = input.logits[token] + dot;
if score > best.1 {
best = (token, score);
}
}
best
}
}
struct Worker {
input: Option<mpsc::SyncSender<Arc<Input>>>,
result: mpsc::Receiver<Best>,
thread: Option<JoinHandle<()>>,
}
impl Drop for Worker {
fn drop(&mut self) {
self.input.take();
if let Some(thread) = self.thread.take() {
let _ = thread.join();
}
}
}
pub(super) struct MarkovPool {
matrix: Arc<Matrix>,
workers: Vec<Worker>,
chunk: usize,
}
fn worker_count(online: usize, requested: Option<&str>) -> usize {
requested
.and_then(|value| value.parse::<usize>().ok())
.filter(|&value| value > 0)
.unwrap_or(online.min(12))
.clamp(1, 32)
}
impl MarkovPool {
pub(super) fn new(model: &Gguf, weight: Weight) -> Result<Self, String> {
let threads = worker_count(
thread::available_parallelism().map_or(1, std::num::NonZero::get),
std::env::var("DS4_THREADS").ok().as_deref(),
);
Self::from_map(model.shared_map(), weight, threads)
}
fn from_map(map: Arc<Mmap>, weight: Weight, threads: usize) -> Result<Self, String> {
let width = usize::try_from(weight.dims[0]).map_err(|_| "Markov width overflow")?;
let rows = usize::try_from(weight.dims[1]).map_err(|_| "Markov rows overflow")?;
if width == 0 || rows == 0 || rows > i32::MAX as usize {
return Err("DSpark dense argmax has invalid dimensions".into());
}
let row_bytes = match weight.kind {
F32 => width.checked_mul(4),
F16 => width.checked_mul(2),
Q8_0 => width.div_ceil(32).checked_mul(34),
_ => None,
}
.ok_or("unsupported DSpark dense tensor layout")?;
let bytes = row_bytes
.checked_mul(rows)
.ok_or("DSpark dense argmax size overflow")?;
if weight.offset > map.len() as u64 || bytes as u64 > map.len() as u64 - weight.offset {
return Err("DSpark dense argmax is outside the GGUF mapping".into());
}
let matrix = Arc::new(Matrix {
map,
weight,
width,
rows,
row_bytes,
});
let mut pool = Self {
matrix,
workers: Vec::new(),
chunk: rows.div_ceil(threads),
};
for slot in 1..threads {
let matrix = Arc::clone(&pool.matrix);
let start = slot * pool.chunk;
let end = (start + pool.chunk).min(rows);
let (input, jobs) = mpsc::sync_channel::<Arc<Input>>(1);
let (results, result) = mpsc::channel();
let thread = thread::Builder::new()
.name(format!("dspark-markov-{slot}"))
.spawn(move || {
while let Ok(input) = jobs.recv() {
let best = matrix.best(&input, start, end);
// Return buffer ownership before signalling completion.
drop(input);
if results.send(best).is_err() {
break;
}
}
})
.map_err(|error| format!("Cannot start DSpark Markov worker: {error}"))?;
pool.workers.push(Worker {
input: Some(input),
result,
thread: Some(thread),
});
}
Ok(pool)
}
pub(super) fn argmax(&mut self, values: &[f32], logits: &mut Vec<f32>) -> Result<i32, String> {
if values.len() != self.matrix.width || logits.len() != self.matrix.rows {
return Err("DSpark dense argmax has mismatched dimensions".into());
}
let input = Arc::new(Input {
values: values.to_vec(),
quantized: (self.matrix.weight.kind == Q8_0).then(|| quantize_q8_activation(values)),
logits: std::mem::take(logits),
});
let parallel = !self.workers.is_empty() && self.matrix.rows >= 512;
let mut failed = false;
if parallel {
for worker in &self.workers {
failed |= worker
.input
.as_ref()
.expect("worker input open")
.send(Arc::clone(&input))
.is_err();
}
}
let mut best = self.matrix.best(
&input,
0,
if parallel {
self.chunk
} else {
self.matrix.rows
},
);
if parallel {
// Join results in slot order: equal scores keep the first token.
// Drain every dispatched job even when another worker failed.
for worker in &self.workers {
match worker.result.recv() {
Ok(candidate) if candidate.1 > best.1 => best = candidate,
Ok(_) => {}
Err(_) => failed = true,
}
}
}
*logits = Arc::try_unwrap(input)
.map_err(|_| "DSpark worker retained input")?
.logits;
if failed {
return Err("DSpark Markov worker failed".into());
}
Ok(best.0 as i32)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
#[ignore = "requires installed 0731 DSpark support; CPU dispatch diagnostic, not parity acceptance"]
fn installed_markov_worker_dispatch() {
use std::time::Instant;
let artifacts = crate::model::engine_artifacts(
crate::model::ModelChoice::DeepSeekV4Flash0731,
true,
&crate::app::models_path(),
);
let model = Gguf::open(artifacts.support.as_ref().unwrap()).unwrap();
let weight = Weight::bind(&model, "mtp.2.markov_head.markov_w2.weight").unwrap();
let state_weight = Weight::bind(&model, "mtp.2.markov_head.markov_w1.weight").unwrap();
let values = super::super::dense_row(&model, state_weight, 671).unwrap();
let mut pool12 = MarkovPool::from_map(model.shared_map(), weight, 12).unwrap();
let mut pool18 = MarkovPool::from_map(model.shared_map(), weight, 18).unwrap();
let mut logits = vec![0.0; weight.dims[1] as usize];
let expected = pool12.argmax(&values, &mut logits).unwrap();
for round in 0..4 {
let order = if round % 2 == 0 { [0, 1, 2] } else { [2, 1, 0] };
for mode in order {
let started = Instant::now();
for _ in 0..128 {
let token = if mode == 0 {
// Same row arithmetic/input ownership, old per-call
// dispatch pattern: isolates worker lifetime/count.
let matrix = &pool18.matrix;
let input = Input {
values: values.clone(),
quantized: Some(quantize_q8_activation(&values)),
logits: std::mem::take(&mut logits),
};
let chunk = matrix.rows.div_ceil(18);
let best = thread::scope(|scope| {
let handles = (0..matrix.rows)
.step_by(chunk)
.map(|start| {
let input = &input;
scope.spawn(move || {
matrix.best(input, start, (start + chunk).min(matrix.rows))
})
})
.collect::<Vec<_>>();
handles.into_iter().map(|h| h.join().unwrap()).fold(
(0, -f32::MAX),
|best, candidate| {
if candidate.1 > best.1 {
candidate
} else {
best
}
},
)
});
logits = input.logits;
best.0 as i32
} else if mode == 1 {
pool12.argmax(&values, &mut logits).unwrap()
} else {
pool18.argmax(&values, &mut logits).unwrap()
};
assert_eq!(token, expected);
}
println!(
"markov_dispatch round={round} mode={} calls=128 elapsed_ms={:.3}",
["scoped18", "persistent12", "persistent18"][mode],
started.elapsed().as_secs_f64() * 1000.0
);
}
}
}
#[test]
fn markov_workers_reuse_threads_and_buffers_with_ordered_ties() {
assert_eq!(worker_count(18, None), 12);
assert_eq!(worker_count(4, Some("0")), 4);
assert_eq!(worker_count(18, Some("99")), 32);
assert_eq!(worker_count(18, Some("1")), 1);
let rows = 1025;
let mut map = memmap2::MmapMut::map_anon(rows * 4).unwrap();
for (row, bytes) in map.chunks_exact_mut(4).enumerate() {
bytes.copy_from_slice(&((row % 17) as f32).to_le_bytes());
}
let map = Arc::new(map.make_read_only().unwrap());
let weight = Weight {
offset: 0,
kind: F32,
bytes: (rows * 4) as u64,
dims: [1, rows as u64, 1],
};
for threads in [1, 4, 12] {
let mut pool = MarkovPool::from_map(Arc::clone(&map), weight, threads).unwrap();
let ids = pool
.workers
.iter()
.map(|w| w.thread.as_ref().unwrap().thread().id())
.collect::<Vec<_>>();
let mut logits = vec![0.0; rows];
let pointer = logits.as_ptr();
assert_eq!(pool.argmax(&[2.0], &mut logits).unwrap(), 16);
logits[rows - 1] = 1000.0;
assert_eq!(pool.argmax(&[2.0], &mut logits).unwrap(), (rows - 1) as i32);
assert_eq!(logits.as_ptr(), pointer);
assert!(pool.argmax(&[], &mut logits).is_err());
assert_eq!(logits.len(), rows);
assert_eq!(
ids,
pool.workers
.iter()
.map(|w| w.thread.as_ref().unwrap().thread().id())
.collect::<Vec<_>>()
);
if let Some(worker) = pool.workers.first_mut() {
let (closed, receiver) = mpsc::sync_channel(1);
drop(receiver);
worker.input.replace(closed);
worker.thread.take().unwrap().join().unwrap();
assert!(pool.argmax(&[2.0], &mut logits).is_err());
assert_eq!(logits.as_ptr(), pointer);
assert_eq!(logits[rows - 1], 1000.0);
}
}
assert!(
MarkovPool::from_map(
map,
Weight {
offset: 1,
..weight
},
1
)
.is_err()
);
}
}
src/engine/metal/qwen_mtplx/request.rs
+3 -3
View File
@@ -357,8 +357,8 @@ fn qwen_request_progress_uses_ui_frontiers_and_decode_timer() {
let timing = PromptTiming {
evaluated_tokens: 20,
eval_seconds: 0.125,
mtp_history_seconds: 0.025,
restore_seconds: 0.010,
mtp_history_seconds: Some(0.025),
restore_seconds: Some(0.010),
};
assert_eq!(metrics.snapshot().prompt_timing, None);
metrics.set_prompt_timing(Some(timing));
@@ -399,7 +399,7 @@ fn qwen_request_progress_uses_ui_frontiers_and_decode_timer() {
let exact_hit = PromptTiming {
evaluated_tokens: 0,
eval_seconds: 0.0,
mtp_history_seconds: 0.0,
mtp_history_seconds: Some(0.0),
..timing
};
metrics.set_prompt_timing(Some(exact_hit));
src/engine/metal/qwen_mtplx/session_cache.rs
+2 -2
View File
@@ -1163,8 +1163,8 @@ impl Execution {
observe(TurnProgress::PromptReady(crate::metrics::PromptTiming {
evaluated_tokens: prompt.len() - cached,
eval_seconds: prepared.eval_seconds + history_seconds,
mtp_history_seconds: history_seconds,
restore_seconds,
mtp_history_seconds: Some(history_seconds),
restore_seconds: Some(restore_seconds),
}))?;
// The bank owns this boundary before decode starts, including when a
// subsequent callback aborts the request. Keeping it in this stack
src/engine/metal/qwen_mtplx_tests.rs
+2 -1
View File
@@ -78,7 +78,8 @@ fn mtplx_canonical_source_bodies_preserve_pinned_hashes() {
.split("// Runtime unit SHA256: ")
.skip(1)
.collect::<Vec<_>>();
assert_eq!(runtime_units.len(), 22);
// tools/mtplx-kernel-source.py --check verifies this complete pinned export.
assert_eq!(runtime_units.len(), 26);
for unit in runtime_units {
let expected = unit.lines().next().unwrap();
let body = unit