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

fix: ssd streaming and deepseek fixes

Browse Source
This commit is contained in:
Georg Bauer
2026-07-26 21:50:59 +02:00
parent 0d80c217c4
commit 47b405e72b
6 changed files with 457 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
native/metal/ds4_metal.m
View File

@@ -11718,7 +11718,7 @@ static uint32_t ds4_gpu_stream_prefill_batch_selected_addr_auto_min(
return 0;
}
static int ds4_gpu_stream_prefill_batch_selected_addr_enabled(
int ds4_gpu_stream_prefill_batch_selected_addr_enabled(
uint32_t n_tokens,
uint32_t n_total_expert,
uint32_t n_expert,

5
scripts/execution_parity.py
View File

@@ -10,15 +10,16 @@ import sys
ROOT = pathlib.Path(__file__).resolve().parent.parent
HARDWARE_TESTS = (
"flash_resident_and_ssd_streaming_choose_the_same_tokens",
"flash_ssd_streaming_maps_batched_prefill_layers",
"resident_multi_session_switching_preserves_each_kv_frontier",
"legacy_mtp_runs_a_target_owned_greedy_cycle",
"dspark_runs_a_target_owned_greedy_cycle",
"ssd_streaming_supports_legacy_mtp_and_dspark",
"directional_steering_matches_the_ds4_token_oracle",
"resident_and_streamed_glm_match_the_short_code_fixture",
"resident_and_streamed_glm_match_ds4_decode_oracles",
"streamed_glm_uses_ds4_indexed_prefill_for_long_prompts",
"glm_mtp_preserves_target_tokens_and_drafts",
"pro_resident_and_ssd_streaming_choose_the_same_tokens",
"pro_ssd_streaming_full_layer_and_selected_batch_agree",
)
ENDPOINT_SCRIPTS = (
"endpoint_parity.py",

41
src/engine.rs
View File

@@ -376,6 +376,14 @@ struct ResidentSlot {
last_store_tokens: u32,
}
#[cfg(target_os = "macos")]
impl Drop for Generator {
fn drop(&mut self) {
// Match ds4: every resident session graph must die before Metal cleanup.
self.resident_sessions.clear();
}
}
#[cfg(target_os = "macos")]
struct SimulatedMemory {
address: std::ptr::NonNull<std::ffi::c_void>,
@@ -524,9 +532,10 @@ impl Generator {
resident_active: None,
resident_limit: std::env::var("DS4_RESIDENT_SESSIONS")
.ok()
.and_then(|value| value.parse().ok())
.and_then(|value| value.parse::<usize>().ok())
.filter(|limit| *limit > 0)
.unwrap_or(4),
.unwrap_or(1)
.saturating_sub(1),
})
}
@@ -908,21 +917,23 @@ impl Generator {
.map_or((None, 0), |slot| (Some(slot.state), slot.last_store_tokens));
self.executor.swap_resident_state(&mut incoming_state)?;
if let (Some(previous), Some(outgoing)) = (self.resident_active.take(), incoming_state) {
// ponytail: four resident sessions bound memory; raise
// DS4_RESIDENT_SESSIONS when the machine can hold more KV state.
if self.resident_sessions.len() >= self.resident_limit {
let evicted = self.resident_sessions.keys().next().cloned();
if let Some(evicted) = evicted {
self.resident_sessions.remove(&evicted);
// DS4 defaults to one active session; opt into its batched-server
// behavior with DS4_RESIDENT_SESSIONS when memory permits.
if self.resident_limit != 0 {
if self.resident_sessions.len() >= self.resident_limit {
let evicted = self.resident_sessions.keys().next().cloned();
if let Some(evicted) = evicted {
self.resident_sessions.remove(&evicted);
}
}
self.resident_sessions.insert(
previous,
ResidentSlot {
state: outgoing,
last_store_tokens: self.last_store_tokens,
},
);
}
self.resident_sessions.insert(
previous,
ResidentSlot {
state: outgoing,
last_store_tokens: self.last_store_tokens,
},
);
}
self.resident_active = Some(key);
self.checkpoint = None;

463
src/engine/metal.rs
View File

@@ -1655,6 +1655,11 @@ struct SsdPlan {
selected_wait_ns: std::sync::atomic::AtomicU64,
}
struct DeepSeekModelSpans {
ranges: Vec<(u64, u64)>,
max_tensor_bytes: u64,
}
#[derive(Clone, Copy)]
struct SelectedLoadJob {
selected: usize,
@@ -1873,12 +1878,13 @@ impl SsdPlan {
if gate == 0 || down == 0 {
return Err("routed expert tensors have an invalid layout".into());
}
let model_spans = streaming_model_spans(model, weights, per_expert_bytes)?;
let resident_bytes = model_spans.iter().try_fold(0_u64, |total, (_, bytes)| {
let resident_spans = streaming_model_spans(model, weights, per_expert_bytes)?;
let resident_bytes = resident_spans.iter().try_fold(0_u64, |total, (_, bytes)| {
total
.checked_add(*bytes)
.ok_or("resident weight size overflow")
})?;
let model_spans = deepseek_token_model_spans(model)?.ranges;
let runtime_bytes = estimated_deepseek_runtime_bytes(model.shape, context, prefill_chunk);
let recommended = unsafe { ds4_gpu_recommended_working_set_size() };
if recommended == 0 && settings.cache_experts == 0 && settings.cache_bytes == 0 {
@@ -2130,6 +2136,156 @@ fn streaming_model_spans(
Ok(merged)
}
fn deepseek_model_spans(
model: &Model,
purpose: &str,
mut include: impl FnMut(&str) -> bool,
) -> Result<DeepSeekModelSpans, String> {
const ISOLATED_Q4_BYTES: u64 = 2 * 1024 * 1024 * 1024;
let q4_groups = std::env::var("DS4_METAL_Q4_PRO_MAP_GROUPS")
.ok()
.and_then(|value| value.parse::<u32>().ok())
.filter(|groups| *groups > 0 && *groups <= 384 && 384 % *groups == 0)
.unwrap_or(1);
let mut spans = Vec::new();
let mut max_tensor_bytes = 0;
for (_, tensor) in model.main.tensors.iter().filter(|(name, _)| include(name)) {
if tensor.bytes == 0 {
continue;
}
let isolate = tensor.kind == Q4_K && tensor.bytes >= ISOLATED_Q4_BYTES;
let groups = if isolate
&& tensor.dims.len() == 3
&& tensor.dims[2] == 384
&& tensor.bytes.is_multiple_of(u64::from(q4_groups))
{
q4_groups
} else {
1
};
let bytes = tensor.bytes / u64::from(groups);
max_tensor_bytes = max_tensor_bytes.max(bytes);
for group in 0..groups {
spans.push((tensor.offset + u64::from(group) * bytes, bytes, isolate));
}
}
let ranges = finish_deepseek_model_spans(spans, purpose)?;
Ok(DeepSeekModelSpans {
ranges,
max_tensor_bytes,
})
}
fn finish_deepseek_model_spans(
mut spans: Vec<(u64, u64, bool)>,
purpose: &str,
) -> Result<Vec<(u64, u64)>, String> {
spans.sort_unstable_by_key(|span| span.0);
let mut merged: Vec<(u64, u64, bool)> = Vec::new();
for (offset, bytes, isolate) in spans {
let end = offset
.checked_add(bytes)
.ok_or_else(|| format!("{purpose} model span overflow"))?;
if let Some((previous_offset, previous_bytes, previous_isolate)) = merged.last_mut() {
let previous_end = previous_offset.saturating_add(*previous_bytes);
if !isolate && !*previous_isolate && offset <= previous_end {
*previous_bytes = previous_end.max(end) - *previous_offset;
continue;
}
}
merged.push((offset, bytes, isolate));
}
if merged.is_empty() {
return Err(format!("{purpose} has no model tensors"));
}
Ok(merged
.into_iter()
.map(|(offset, bytes, _)| (offset, bytes))
.collect())
}
fn deepseek_token_model_spans(model: &Model) -> Result<DeepSeekModelSpans, String> {
deepseek_model_spans(model, "DeepSeek token embedding", |name| {
name == "token_embd.weight"
})
}
fn deepseek_layer_model_spans(
model: &Model,
weights: &Layer,
layer: u32,
decode_only: bool,
slab_bytes: u64,
) -> Result<DeepSeekModelSpans, String> {
let prefix = format!("blk.{layer}.");
let stream_experts = decode_only
&& (weights.expert_gate.bytes + weights.expert_up.bytes + weights.expert_down.bytes)
/ model.shape.experts
== slab_bytes;
deepseek_model_spans(model, &format!("DeepSeek layer {layer}"), |name| {
(name.starts_with(&prefix)
&& (!stream_experts
|| (!name.ends_with("ffn_gate_exps.weight")
&& !name.ends_with("ffn_up_exps.weight")
&& !name.ends_with("ffn_down_exps.weight"))))
|| (!decode_only && layer == 0 && name == "token_embd.weight")
})
}
fn deepseek_output_model_spans(model: &Model) -> Result<DeepSeekModelSpans, String> {
deepseek_model_spans(model, "DeepSeek output head", |name| {
name.starts_with("output")
})
}
fn deepseek_mtp_base_model_spans(model: &Model) -> Result<DeepSeekModelSpans, String> {
deepseek_model_spans(model, "DeepSeek MTP base weights", |name| {
name == "token_embd.weight" || name.starts_with("output")
})
}
fn install_deepseek_model_spans(
model: &Model,
spans: &DeepSeekModelSpans,
purpose: &str,
) -> Result<(), String> {
let (offsets, sizes): (Vec<_>, Vec<_>) = spans.ranges.iter().copied().unzip();
call(
unsafe {
ds4_gpu_set_model_map_spans(
model.main.map_ptr().cast(),
model.main.len(),
offsets.as_ptr(),
sizes.as_ptr(),
spans.ranges.len() as u32,
spans.max_tensor_bytes,
)
},
purpose,
)
}
fn install_support_model_map(model: &Model, purpose: &str) -> Result<(), String> {
let support = model.support.as_ref().ok_or("support model is missing")?;
call(
unsafe {
ds4_gpu_set_model_map_range(
support.map_ptr().cast(),
support.len(),
support.data_offset(),
support.len() - support.data_offset(),
support.max_tensor_bytes(),
)
},
purpose,
)
}
fn install_speculative_model_maps(model: &Model, purpose: &str) -> Result<(), String> {
install_deepseek_model_spans(model, &deepseek_mtp_base_model_spans(model)?, purpose)?;
install_support_model_map(model, purpose)
}
fn estimated_deepseek_runtime_bytes(shape: super::Shape, context: u32, prefill: u32) -> u64 {
let prefill_cap = effective_prefill_cap(context, prefill);
let raw_cap = effective_raw_cap(shape, context, prefill_cap);
@@ -2743,6 +2899,9 @@ impl DeepSeekExecutor {
pub(super) fn eval(&mut self, token: i32) -> Result<(), String> {
self.eval_target(token)?;
if self.dspark.is_some() && self.ssd.is_some() {
install_speculative_model_maps(&self.model, "DSpark support mapping")?;
}
if let (Some(dspark), Some(support)) = (&mut self.dspark, self.model.support.as_ref()) {
dspark.seed_current_cache(
support,
@@ -2768,9 +2927,13 @@ impl DeepSeekExecutor {
if let Some(dspark) = &mut self.dspark {
dspark.begin_capture();
}
let commands = Commands::begin()?;
self.encode_token(token as u32)?;
commands.finish()?;
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()?;
}
self.session.scratch.logits.read_f32(&mut self.logits)?;
self.session.position += 1;
self.tokens.push(token);
@@ -3029,6 +3192,9 @@ impl DeepSeekExecutor {
return Ok(accepted);
}
if self.dspark.is_some() {
if self.ssd.is_some() {
install_speculative_model_maps(&self.model, "DSpark support mapping")?;
}
let mut dspark = self.dspark.take().expect("DSpark disappeared");
let proposals = dspark.propose(
&self.model,
@@ -3148,6 +3314,9 @@ impl DeepSeekExecutor {
} else {
mtp.state_hc.raw().cast_const()
};
if self.ssd.is_some() {
install_speculative_model_maps(&self.model, "MTP model mapping")?;
}
let commands = Commands::begin()?;
call(
unsafe {
@@ -3388,7 +3557,24 @@ impl DeepSeekExecutor {
let size = self.model.main.len();
let shape = self.model.shape;
let pos = self.session.position;
let batch_selected_addr = self.ssd.is_some()
&& self.weights.layers.first().is_some_and(|layer| unsafe {
ds4_gpu_stream_prefill_batch_selected_addr_enabled(
rows,
shape.experts as u32,
shape.experts_used as u32,
layer.expert_gate.kind,
layer.expert_down.kind,
) != 0
});
if self.ssd.is_some() {
install_deepseek_model_spans(
&self.model,
&deepseek_token_model_spans(&self.model)?,
"DeepSeek prefill token mapping",
)?;
}
let commands = Commands::begin()?;
call(
unsafe {
@@ -3416,6 +3602,19 @@ impl DeepSeekExecutor {
.enumerate()
{
let started = Instant::now();
if let Some(ssd) = &self.ssd {
install_deepseek_model_spans(
&self.model,
&deepseek_layer_model_spans(
&self.model,
weights,
index as u32,
batch_selected_addr,
ssd.per_expert_bytes,
)?,
"DeepSeek prefill layer mapping",
)?;
}
let commands = Commands::begin()?;
encode_batch_layer(
&batch,
@@ -3429,7 +3628,6 @@ impl DeepSeekExecutor {
rows,
self.session.raw_cap,
self.steering.as_ref(),
self.ssd.as_ref(),
)?;
if let Some(profile) = &mut self.profile {
profile.record(
@@ -3459,6 +3657,9 @@ impl DeepSeekExecutor {
std::mem::swap(&mut batch.current_hc, &mut batch.next_hc);
}
if self.dspark.is_some() && self.ssd.is_some() {
install_speculative_model_maps(&self.model, "DSpark prefill support mapping")?;
}
if let (Some(dspark), Some(support)) = (&mut self.dspark, self.model.support.as_ref()) {
dspark.seed_batch_cache(
support,
@@ -3469,6 +3670,13 @@ impl DeepSeekExecutor {
shape,
)?;
}
if self.ssd.is_some() {
install_deepseek_model_spans(
&self.model,
&deepseek_output_model_spans(&self.model)?,
"DeepSeek prefill output mapping",
)?;
}
let output_rows = if collect_tops { rows } else { 1 };
let first_output = rows - output_rows;
@@ -3671,6 +3879,68 @@ impl DeepSeekExecutor {
}
fn encode_token(&mut self, token: u32) -> Result<(), String> {
self.encode_token_embedding(token)?;
for index in 0..self.weights.layers.len() {
self.encode_token_layer(index, token)?;
}
encode_output(
&self.session.scratch,
&self.weights,
self.model.shape,
self.model.main.map_ptr().cast(),
self.model.main.len(),
)
}
fn encode_streaming_token(&mut self, token: u32) -> Result<(), String> {
let slab_bytes = self
.ssd
.as_ref()
.expect("streaming token lost its SSD plan")
.per_expert_bytes;
install_deepseek_model_spans(
&self.model,
&deepseek_token_model_spans(&self.model)?,
"DeepSeek token embedding mapping",
)?;
let commands = Commands::begin()?;
self.encode_token_embedding(token)?;
commands.finish()?;
for index in 0..self.weights.layers.len() {
install_deepseek_model_spans(
&self.model,
&deepseek_layer_model_spans(
&self.model,
&self.weights.layers[index],
index as u32,
true,
slab_bytes,
)?,
"DeepSeek decode layer mapping",
)?;
let commands = Commands::begin()?;
self.encode_token_layer(index, token)?;
commands.finish()?;
}
install_deepseek_model_spans(
&self.model,
&deepseek_output_model_spans(&self.model)?,
"DeepSeek output mapping",
)?;
let commands = Commands::begin()?;
encode_output(
&self.session.scratch,
&self.weights,
self.model.shape,
self.model.main.map_ptr().cast(),
self.model.main.len(),
)?;
commands.finish()
}
fn encode_token_embedding(&mut self, token: u32) -> Result<(), String> {
let shape = self.model.shape;
let map = self.model.main.map_ptr().cast();
let size = self.model.main.len();
@@ -3689,38 +3959,40 @@ impl DeepSeekExecutor {
)
},
"token embedding",
)
}
fn encode_token_layer(&mut self, index: usize, token: u32) -> Result<(), String> {
let shape = self.model.shape;
let scratch = &mut self.session.scratch;
encode_layer(
scratch,
&mut self.session.layers[index],
&self.weights.layers[index],
shape,
self.model.main.map_ptr().cast(),
self.model.main.len(),
index as u32,
self.session.position,
token,
self.session.raw_cap,
self.steering.as_ref(),
self.ssd.as_ref(),
)?;
for (index, layer) in self.weights.layers.iter().enumerate() {
let state = &mut self.session.layers[index];
encode_layer(
scratch,
state,
layer,
shape,
map,
size,
index as u32,
if let Some(profile) = &mut self.profile {
profile.record(
index,
self.session.position,
token,
self.session.raw_cap,
self.steering.as_ref(),
self.ssd.as_ref(),
&scratch.router_selected,
&scratch.router_weights,
1,
index < shape.hash_layers as usize,
)?;
if let Some(profile) = &mut self.profile {
profile.record(
index,
self.session.position,
&scratch.router_selected,
&scratch.router_weights,
1,
index < shape.hash_layers as usize,
)?;
}
if let Some(dspark) = &mut self.dspark {
dspark.capture_decode(index as u32, &scratch.current_hc, shape)?;
}
}
encode_output(scratch, &self.weights, shape, map, size)
if let Some(dspark) = &mut self.dspark {
dspark.capture_decode(index as u32, &scratch.current_hc, shape)?;
}
Ok(())
}
}
@@ -4386,7 +4658,6 @@ fn encode_batch_layer(
rows: u32,
raw_cap: u32,
steering: Option<&Steering>,
ssd: Option<&SsdPlan>,
) -> Result<(), String> {
let hc_dim = shape.hc * shape.embd;
let mix_hc = 2 * shape.hc + shape.hc * shape.hc;
@@ -5160,24 +5431,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]);
if let Some(ssd) = ssd {
let table = StreamExpertTable {
model_map: map,
model_size: size,
layer,
total_experts: shape.experts as u32,
gate_offset: w.expert_gate.offset,
up_offset: w.expert_up.offset,
down_offset: w.expert_down.offset,
gate_expert_bytes: w.expert_gate.dims[1] * gate_row,
down_expert_bytes: w.expert_down.dims[1] * down_row,
};
ssd.begin_selected(
&s.router_selected,
table,
rows.saturating_mul(shape.experts_used as u32),
)?;
}
q8_rows(
&s.shared_gate,
w.shared_gate,
@@ -5221,9 +5474,6 @@ fn encode_batch_layer(
map,
size,
)?;
if let Some(ssd) = ssd {
ssd.finish_selected(false)?;
}
let mut mid_f16 = false;
call(
unsafe {
@@ -6813,7 +7063,8 @@ fn check(result: i32, operation: &str) -> Result<(), String> {
mod tests {
use super::{
compression_ratio, effective_prefill_cap, effective_raw_cap,
estimated_deepseek_runtime_bytes, raw_batch_span, raw_decode_span,
estimated_deepseek_runtime_bytes, finish_deepseek_model_spans, raw_batch_span,
raw_decode_span,
};
use crate::engine::{FLASH, PRO};
@@ -6825,6 +7076,18 @@ mod tests {
assert_eq!(compression_ratio(PRO, 1), 128);
}
#[test]
fn pro_q4_model_spans_remain_isolated() {
assert_eq!(
finish_deepseek_model_spans(
vec![(0, 10, false), (10, 20, true), (30, 10, false)],
"test",
)
.unwrap(),
vec![(0, 10), (10, 20), (30, 10)],
);
}
#[test]
fn decode_raw_span_tracks_the_logical_sliding_window() {
assert_eq!(raw_decode_span(0, 4_352, 128), (1, 0));
@@ -7349,6 +7612,75 @@ mod tests {
assert_eq!(resident, run(path, true));
}
#[test]
#[ignore = "requires the installed 81 GiB Flash GGUF fixture and a Metal device"]
fn flash_ssd_streaming_maps_batched_prefill_layers() {
use super::{DeepSeekExecutor, argmax, configure_sources};
use crate::engine::Model;
use crate::model::ModelChoice;
use crate::settings::{
EngineSpeculativeSettings, EngineSsdSettings, EngineSteeringSettings,
};
use std::path::Path;
configure_sources().unwrap();
let model = Model::open_main(
Path::new(
"../ds4/gguf/DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf",
),
ModelChoice::DeepSeekV4Flash,
)
.unwrap();
let prompt = model.render_conversation(
"",
&[crate::engine::ChatTurn {
user: true,
tool: false,
system: false,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
content: "hi ".repeat(80),
}],
crate::settings::ReasoningMode::Direct,
);
assert!(prompt.len() > 64);
let mut executor = DeepSeekExecutor::open(
model,
512,
false,
256,
100,
EngineSpeculativeSettings {
mtp_draft_tokens: 1,
mtp_margin: 3.0,
glm_mtp: false,
glm_mtp_timing: false,
dspark: false,
dspark_confidence_threshold: 0.9,
dspark_confidence_threshold_set: false,
dspark_strict: false,
},
EngineSsdSettings {
enabled: true,
cold: true,
cache_experts: 16,
cache_bytes: 0,
full_layers: 0,
full_layers_set: false,
preload_experts: 0,
},
EngineSteeringSettings {
file: None,
ffn_scale: 0.0,
attention_scale: 0.0,
},
)
.unwrap();
assert_eq!(executor.prefill(&prompt, |_| true).unwrap(), prompt.len());
executor.eval(argmax(executor.logits())).unwrap();
}
#[test]
#[ignore = "requires the installed 81 GiB Flash GGUF fixture and a Metal device"]
fn resident_multi_session_switching_preserves_each_kv_frontier() {
@@ -7514,7 +7846,7 @@ mod tests {
#[test]
#[ignore = "requires the installed DeepSeek V4 Pro GGUF and Apple Metal"]
fn pro_resident_and_ssd_streaming_choose_the_same_tokens() {
fn pro_ssd_streaming_full_layer_and_selected_batch_agree() {
use super::{DeepSeekExecutor, argmax, configure_sources};
use crate::engine::Model;
use crate::model::ModelChoice;
@@ -7531,7 +7863,7 @@ mod tests {
eprintln!("skipping unavailable Pro fixture: {path}");
return;
}
let run = |streaming| {
let run = |cache_experts| {
let model = Model::open_main(Path::new(&path), ModelChoice::DeepSeekV4Pro).unwrap();
let prompt = model.render_conversation(
"",
@@ -7542,15 +7874,16 @@ mod tests {
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
content: "hi".into(),
content: "hi ".repeat(32),
}],
crate::settings::ReasoningMode::Direct,
);
assert!(prompt.len() > 18);
let mut executor = DeepSeekExecutor::open(
model,
64,
256,
false,
64,
256,
100,
EngineSpeculativeSettings {
mtp_draft_tokens: 1,
@@ -7563,9 +7896,9 @@ mod tests {
dspark_strict: false,
},
EngineSsdSettings {
enabled: streaming,
enabled: true,
cold: true,
cache_experts: if streaming { 32 } else { 0 },
cache_experts,
cache_bytes: 0,
full_layers: 0,
full_layers_set: false,
@@ -7587,6 +7920,6 @@ mod tests {
})
.collect::<Vec<_>>()
};
assert_eq!(run(false), run(true));
assert_eq!(run(32), run(384));
}
}

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

@@ -2725,6 +2725,7 @@ fn glm_streaming_plan(
let max_layer = weights
.layers
.iter()
.chain(weights.nextn.iter())
.filter_map(|layer| layer.sparse)
.map(|weights| {
sparse_expert_bytes(weights, shape.experts)
@@ -3256,7 +3257,7 @@ mod tests {
#[test]
#[ignore = "requires the 197 GiB GLM 5.2 checkpoint and Apple Metal"]
fn resident_and_streamed_glm_match_the_short_code_fixture() {
fn resident_and_streamed_glm_match_ds4_decode_oracles() {
let path = std::env::var("DS4_GLM_MODEL").unwrap_or_else(|_| {
"../ds4/models/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf".into()
});
@@ -3290,6 +3291,9 @@ mod tests {
}
Err(error) => panic!("GLM executor failed to open: {error}"),
};
if streamed {
assert_eq!(executor.ssd_cache_experts, 671);
}
assert_eq!(executor.prefill(&tokens, |_| true).unwrap(), tokens.len());
let token = executor
.logits()
@@ -3302,6 +3306,23 @@ mod tests {
executor.model().token_bytes(token).as_deref(),
Some(b");\n".as_slice())
);
let greeting = executor.model().render_prompt(
"You are a helpful assistant",
"Write one short greeting.",
ReasoningMode::Direct,
);
executor.reset().unwrap();
assert_eq!(
executor.prefill(&greeting, |_| true).unwrap(),
greeting.len()
);
let mut generated = Vec::new();
for _ in 0..4 {
let token = argmax(executor.logits());
generated.push(token);
executor.eval(token).unwrap();
}
assert_eq!(generated, [9703, 0, 2585, 646]);
}
}

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

@@ -74,6 +74,13 @@ unsafe extern "C" {
) -> i32;
pub(super) fn ds4_gpu_stream_expert_cache_note_service_thread();
pub(super) fn ds4_gpu_stream_expert_cache_reset_route_hotness();
pub(super) fn ds4_gpu_stream_prefill_batch_selected_addr_enabled(
n_tokens: u32,
n_total_expert: u32,
n_expert: u32,
gate_type: u32,
down_type: u32,
) -> i32;
pub(super) fn ds4_gpu_glm_stream_expert_cache_begin_selected_load_tensor(
table: *const StreamExpertTable,
selected: *const GpuTensor,