From 47b405e72bba5048479824bca9fd13bd4dc7e388 Mon Sep 17 00:00:00 2001 From: Georg Bauer Date: Sun, 26 Jul 2026 21:50:59 +0200 Subject: [PATCH] fix: ssd streaming and deepseek fixes --- native/metal/ds4_metal.m | 2 +- scripts/execution_parity.py | 5 +- src/engine.rs | 41 ++-- src/engine/metal.rs | 463 +++++++++++++++++++++++++++++++----- src/engine/metal/glm.rs | 23 +- src/engine/metal/gpu.rs | 7 + 6 files changed, 457 insertions(+), 84 deletions(-) diff --git a/native/metal/ds4_metal.m b/native/metal/ds4_metal.m index 6ed87e5..bcb4480 100644 --- a/native/metal/ds4_metal.m +++ b/native/metal/ds4_metal.m @@ -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, diff --git a/scripts/execution_parity.py b/scripts/execution_parity.py index e7eb0d8..5afebae 100755 --- a/scripts/execution_parity.py +++ b/scripts/execution_parity.py @@ -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", diff --git a/src/engine.rs b/src/engine.rs index c62a0ac..a64d22d 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -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, @@ -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::().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; diff --git a/src/engine/metal.rs b/src/engine/metal.rs index 62c8f49..688803e 100644 --- a/src/engine/metal.rs +++ b/src/engine/metal.rs @@ -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 { + 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::().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, 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 { + 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 { + 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 { + deepseek_model_spans(model, "DeepSeek output head", |name| { + name.starts_with("output") + }) +} + +fn deepseek_mtp_base_model_spans(model: &Model) -> Result { + 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::>() }; - assert_eq!(run(false), run(true)); + assert_eq!(run(32), run(384)); } } diff --git a/src/engine/metal/glm.rs b/src/engine/metal/glm.rs index 32660bf..474e601 100644 --- a/src/engine/metal/glm.rs +++ b/src/engine/metal/glm.rs @@ -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]); } } diff --git a/src/engine/metal/gpu.rs b/src/engine/metal/gpu.rs index 1604e4d..7000241 100644 --- a/src/engine/metal/gpu.rs +++ b/src/engine/metal/gpu.rs @@ -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,