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

Finish DS4 execution parity

Browse Source
This commit is contained in:
Georg Bauer
2026-07-26 20:44:23 +02:00
parent fd3f8e45dc
commit 0d80c217c4
15 changed files with 16905 additions and 195 deletions
Download Patch File Download Diff File Expand all files Collapse all files

446
src/engine/metal.rs
View File

@@ -2,11 +2,13 @@ mod checkpoint;
mod glm;
mod gpu;
mod hotlist;
mod profile;
use glm::GlmExecutor;
use gpu::*;
use profile::ExpertProfile;
use super::gguf::{F16, F32, Gguf, Q4_K, Q8_0, Tensor as GgufTensor};
use super::gguf::{F16, F32, Gguf, IQ2_XXS, Q4_K, Q8_0, Tensor as GgufTensor};
use super::validation::{DsparkConfig, SupportKind, dspark_config};
use super::{Model, ModelFamily};
use crate::model::ModelChoice;
@@ -1991,7 +1993,12 @@ impl SsdPlan {
}
let mut by_layer = vec![Vec::<(i32, u32)>::new(); model.shape.layers as usize];
let mut loaded = 0_u32;
for &(layer, expert) in hotlist::FLASH {
let hotlist = match model.shape.model {
ModelChoice::DeepSeekV4Flash => hotlist::FLASH,
ModelChoice::DeepSeekV4Pro => hotlist::PRO,
ModelChoice::Glm52 => unreachable!("GLM uses its dedicated executor"),
};
for &(layer, expert) in hotlist {
if loaded == self.preload_experts {
break;
}
@@ -2163,6 +2170,31 @@ fn estimated_deepseek_runtime_bytes(shape: super::Shape, context: u32, prefill:
raw.saturating_add(compressed).saturating_add(scratch)
}
fn resident_deepseek_admission_bytes(
model: &Model,
context: u32,
prefill: u32,
) -> Result<u64, String> {
resident_deepseek_admission_for_weights(
model.main.len() - model.main.data_offset(),
model.shape,
context,
prefill,
)
}
fn resident_deepseek_admission_for_weights(
weight_bytes: u64,
shape: super::Shape,
context: u32,
prefill: u32,
) -> Result<u64, String> {
weight_bytes
.checked_add(estimated_deepseek_runtime_bytes(shape, context, prefill))
.and_then(|bytes| bytes.checked_add(512 * 1024 * 1024))
.ok_or_else(|| "DeepSeek runtime memory size overflow".into())
}
fn effective_prefill_cap(context: u32, requested: u32) -> u32 {
if requested == 0 {
context.clamp(1, DEFAULT_PREFILL_CHUNK)
@@ -2512,7 +2544,18 @@ pub(super) struct ExecutionStats {
pub(super) ssd_resident_bytes: u64,
pub(super) ssd_cache_bytes: u64,
pub(super) ssd_cache_experts: u64,
pub(super) ssd_cache_entries: u64,
pub(super) ssd_preloaded_experts: u64,
pub(super) ssd_cache_hits: u64,
pub(super) ssd_cache_misses: u64,
pub(super) ssd_cache_evictions: u64,
pub(super) ssd_cache_wraps: u64,
pub(super) ssd_buffer_allocs: u64,
pub(super) ssd_buffer_reuses: u64,
pub(super) ssd_pread_bytes: u64,
pub(super) ssd_pread_ms: u64,
pub(super) ssd_evict_advise_bytes: u64,
pub(super) ssd_willneed_advise_bytes: u64,
pub(super) ssd_selected_requests: u64,
pub(super) ssd_requested_bytes: u64,
pub(super) ssd_wait_ms: u64,
@@ -2525,6 +2568,7 @@ pub(super) struct DeepSeekExecutor {
dspark: Option<Dspark>,
steering: Option<Steering>,
ssd: Option<SsdPlan>,
profile: Option<ExpertProfile>,
logits: Vec<f32>,
tokens: Vec<i32>,
quality: bool,
@@ -2539,9 +2583,20 @@ pub(super) struct DeepSeekExecutor {
model_identity: [u8; 32],
_context: Context,
model: Model,
speculative: EngineSpeculativeSettings,
}
pub(super) struct DeepSeekResidentState {
session: Session,
legacy_mtp: Option<LegacyMtp>,
dspark: Option<Dspark>,
logits: Vec<f32>,
tokens: Vec<i32>,
checkpoint_tag: [u8; 32],
}
impl DeepSeekExecutor {
#[allow(dead_code)]
#[allow(clippy::too_many_arguments)]
pub(super) fn open(
model: Model,
@@ -2552,6 +2607,31 @@ impl DeepSeekExecutor {
speculative: EngineSpeculativeSettings,
ssd: EngineSsdSettings,
steering: EngineSteeringSettings,
) -> Result<Self, String> {
Self::open_profile(
model,
context,
quality,
prefill_chunk,
power_percent,
speculative,
ssd,
steering,
None,
)
}
#[allow(clippy::too_many_arguments)]
pub(super) fn open_profile(
model: Model,
context: u32,
quality: bool,
prefill_chunk: u32,
power_percent: u8,
speculative: EngineSpeculativeSettings,
ssd: EngineSsdSettings,
steering: EngineSteeringSettings,
expert_profile_path: Option<&str>,
) -> Result<Self, String> {
let weights = Weights::bind(&model)?;
let mut ssd_plan = ssd
@@ -2559,7 +2639,11 @@ impl DeepSeekExecutor {
.then(|| SsdPlan::new(&model, &weights, ssd, context, prefill_chunk))
.transpose()?;
let spans = ssd_plan.as_ref().map(|plan| plan.model_spans.as_slice());
let admission = ssd_plan.as_ref().map_or(0, |plan| plan.admission_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 steering = Steering::load(&model, steering)?;
let session = Session::new(
@@ -2620,6 +2704,13 @@ impl DeepSeekExecutor {
} else {
model.checkpoint_identity()
};
let profile = ExpertProfile::new(
expert_profile_path,
model.shape.model,
model.shape.layers,
model.shape.experts,
model.shape.experts_used,
)?;
Ok(Self {
weights,
session,
@@ -2627,6 +2718,7 @@ impl DeepSeekExecutor {
dspark,
steering,
ssd: ssd_plan,
profile,
logits: vec![0.0; model.shape.vocab as usize],
tokens: Vec::new(),
quality,
@@ -2645,6 +2737,7 @@ impl DeepSeekExecutor {
model_identity,
_context: context_handle,
model,
speculative,
})
}
@@ -2681,6 +2774,9 @@ impl DeepSeekExecutor {
self.session.scratch.logits.read_f32(&mut self.logits)?;
self.session.position += 1;
self.tokens.push(token);
if let Some(profile) = &self.profile {
profile.write()?;
}
throttle(
&mut self.decode_average,
started.elapsed(),
@@ -3335,6 +3431,16 @@ impl DeepSeekExecutor {
self.steering.as_ref(),
self.ssd.as_ref(),
)?;
if let Some(profile) = &mut self.profile {
profile.record(
index,
pos,
&batch.router_selected,
&batch.router_weights,
rows,
index < shape.hash_layers as usize,
)?;
}
if let Some(dspark) = &mut self.dspark {
dspark.capture_batch(
index as u32,
@@ -3388,6 +3494,9 @@ impl DeepSeekExecutor {
}
self.session.position += rows;
self.tokens.extend_from_slice(tokens);
if let Some(profile) = &self.profile {
profile.write()?;
}
Ok(tops)
}
@@ -3414,16 +3523,29 @@ impl DeepSeekExecutor {
..ExecutionStats::default()
};
if let Some(ssd) = &self.ssd {
let mut native = StreamExpertCacheStats::default();
unsafe { ds4_gpu_stream_expert_cache_get_stats(&mut native) };
let selected = ssd
.selected_experts
.load(std::sync::atomic::Ordering::Relaxed);
stats.ssd_enabled = true;
stats.ssd_resident_bytes = ssd.resident_bytes;
stats.ssd_cache_experts = u64::from(ssd.cache_experts);
stats.ssd_cache_entries = u64::from(native.current_count);
stats.ssd_cache_bytes = ssd
.per_expert_bytes
.saturating_mul(u64::from(ssd.cache_experts));
stats.ssd_preloaded_experts = u64::from(ssd.preload_experts);
stats.ssd_cache_hits = native.hits;
stats.ssd_cache_misses = native.misses;
stats.ssd_cache_evictions = native.evictions;
stats.ssd_cache_wraps = native.wraps;
stats.ssd_buffer_allocs = native.buffer_allocs;
stats.ssd_buffer_reuses = native.buffer_reuses;
stats.ssd_pread_bytes = native.pread_bytes;
stats.ssd_pread_ms = native.pread_ms.max(0.0) as u64;
stats.ssd_evict_advise_bytes = native.evict_advise_bytes;
stats.ssd_willneed_advise_bytes = native.willneed_advise_bytes;
stats.ssd_selected_requests = ssd
.selected_requests
.load(std::sync::atomic::Ordering::Relaxed);
@@ -3466,6 +3588,67 @@ impl DeepSeekExecutor {
Ok(())
}
fn blank_resident_state(&self) -> Result<DeepSeekResidentState, String> {
let session = Session::new(&self.model, self.session.context, self.session.prefill_cap)?;
let legacy_mtp = match (self.model.support_kind, self.model.support.as_ref()) {
(Some(SupportKind::LegacyMtp), Some(support)) => {
let hc_dim = self.model.shape.hc * self.model.shape.embd;
Some(LegacyMtp {
weights: LegacyMtpWeights::bind(support, self.model.shape)?,
layer: LayerState::allocate(&self.model, 1, session.context, session.raw_cap)?,
state_hc: Buffer::floats(hc_dim)?,
next_hc: Buffer::floats(hc_dim)?,
draft_token: None,
raw_rows: 0,
draft_limit: self.speculative.mtp_draft_tokens.max(1) as u32,
margin: self.speculative.mtp_margin,
drafted: 0,
accepted: 0,
})
}
_ => None,
};
let dspark = match (
self.model.support_kind,
self.model.support.as_ref(),
self.speculative.dspark,
) {
(Some(SupportKind::DSpark), Some(support), true) => Some(Dspark::new(
&self.model,
support,
&session,
self.speculative,
self.quality,
)?),
_ => None,
};
Ok(DeepSeekResidentState {
session,
legacy_mtp,
dspark,
logits: vec![0.0; self.model.shape.vocab as usize],
tokens: Vec::new(),
checkpoint_tag: [0; 32],
})
}
pub(super) fn swap_resident_state(
&mut self,
state: &mut Option<DeepSeekResidentState>,
) -> Result<(), String> {
let mut incoming = state
.take()
.map_or_else(|| self.blank_resident_state(), Ok)?;
std::mem::swap(&mut self.session, &mut incoming.session);
std::mem::swap(&mut self.legacy_mtp, &mut incoming.legacy_mtp);
std::mem::swap(&mut self.dspark, &mut incoming.dspark);
std::mem::swap(&mut self.logits, &mut incoming.logits);
std::mem::swap(&mut self.tokens, &mut incoming.tokens);
std::mem::swap(&mut self.checkpoint_tag, &mut incoming.checkpoint_tag);
*state = Some(incoming);
Ok(())
}
pub(super) fn align_prompt(&mut self, tokens: &[i32]) -> Result<usize, String> {
if !tokens.starts_with(&self.tokens) {
self.reset()?;
@@ -3523,6 +3706,16 @@ impl DeepSeekExecutor {
self.steering.as_ref(),
self.ssd.as_ref(),
)?;
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)?;
}
@@ -3537,6 +3730,11 @@ pub(super) enum Executor {
Glm(Box<GlmExecutor>),
}
pub(super) enum ResidentState {
DeepSeek(Box<DeepSeekResidentState>),
Glm(Box<glm::GlmResidentState>),
}
impl Executor {
#[allow(dead_code)]
pub(super) fn open(
@@ -3575,6 +3773,7 @@ impl Executor {
ffn_scale: 0.0,
attention_scale: 0.0,
},
None,
)
}
@@ -3588,9 +3787,10 @@ impl Executor {
speculative: EngineSpeculativeSettings,
ssd: EngineSsdSettings,
steering: EngineSteeringSettings,
expert_profile_path: Option<&str>,
) -> Result<Self, String> {
match model.shape.family {
ModelFamily::DeepSeek => DeepSeekExecutor::open(
ModelFamily::DeepSeek => DeepSeekExecutor::open_profile(
model,
context,
quality,
@@ -3599,12 +3799,20 @@ impl Executor {
speculative,
ssd,
steering,
expert_profile_path,
)
.map(Box::new)
.map(Self::DeepSeek),
ModelFamily::Glm => GlmExecutor::open(model, context, quality, ssd)
.map(Box::new)
.map(Self::Glm),
ModelFamily::Glm => GlmExecutor::open_profile(
model,
context,
quality,
ssd,
speculative,
expert_profile_path,
)
.map(Box::new)
.map(Self::Glm),
}
}
@@ -3627,8 +3835,8 @@ impl Executor {
executor.eval_speculative_greedy(token, max_tokens, reasoning, cancelled)
}
Self::Glm(executor) => {
executor.eval(token)?;
Ok(vec![token])
let _ = reasoning;
executor.eval_speculative_greedy(token, max_tokens, cancelled)
}
}
}
@@ -3654,7 +3862,7 @@ impl Executor {
pub(super) fn execution_stats(&self) -> ExecutionStats {
match self {
Self::DeepSeek(executor) => executor.execution_stats(),
Self::Glm(_) => ExecutionStats::default(),
Self::Glm(executor) => executor.execution_stats(),
}
}
@@ -3686,6 +3894,37 @@ impl Executor {
}
}
pub(super) fn swap_resident_state(
&mut self,
state: &mut Option<ResidentState>,
) -> Result<(), String> {
match self {
Self::DeepSeek(executor) => {
let mut inner = match state.take() {
Some(ResidentState::DeepSeek(state)) => Some(*state),
Some(ResidentState::Glm(_)) => {
return Err("resident session belongs to a different model family".into());
}
None => None,
};
executor.swap_resident_state(&mut inner)?;
*state = inner.map(|state| ResidentState::DeepSeek(Box::new(state)));
}
Self::Glm(executor) => {
let mut inner = match state.take() {
Some(ResidentState::Glm(state)) => Some(*state),
Some(ResidentState::DeepSeek(_)) => {
return Err("resident session belongs to a different model family".into());
}
None => None,
};
executor.swap_resident_state(&mut inner)?;
*state = inner.map(|state| ResidentState::Glm(Box::new(state)));
}
}
Ok(())
}
pub(super) fn align_prompt(&mut self, tokens: &[i32]) -> Result<usize, String> {
match self {
Self::DeepSeek(executor) => executor.align_prompt(tokens),
@@ -6628,6 +6867,30 @@ mod tests {
);
}
#[test]
fn pro_resident_admission_includes_weights_context_and_scratch() {
let weights = 464_627_334_560_u64;
let runtime = estimated_deepseek_runtime_bytes(PRO, 32_768, 4_096);
assert_eq!(
super::resident_deepseek_admission_for_weights(weights, PRO, 32_768, 4_096).unwrap(),
weights + runtime + 512 * 1024 * 1024
);
}
#[test]
fn generated_hotlists_match_each_model_shape() {
for (hotlist, shape) in [
(super::hotlist::FLASH, FLASH),
(super::hotlist::PRO, PRO),
(super::hotlist::GLM52, crate::engine::GLM),
] {
assert!(hotlist.len() >= 4_096);
assert!(hotlist.iter().all(|&(layer, expert)| {
u32::from(layer) < shape.layers && u64::from(expert) < shape.experts
}));
}
}
#[test]
#[ignore = "requires the installed 81 GiB Flash and legacy MTP GGUF fixtures"]
fn legacy_mtp_runs_a_target_owned_greedy_cycle() {
@@ -7086,6 +7349,91 @@ mod tests {
assert_eq!(resident, run(path, true));
}
#[test]
#[ignore = "requires the installed 81 GiB Flash GGUF fixture and a Metal device"]
fn resident_multi_session_switching_preserves_each_kv_frontier() {
use super::{DeepSeekExecutor, argmax, configure_sources};
use crate::engine::Model;
use crate::model::ModelChoice;
use crate::settings::{
EngineSpeculativeSettings, EngineSsdSettings, EngineSteeringSettings, ReasoningMode,
};
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 prompts = ["Reply with A.", "Reply with B."].map(|content| {
model.render_conversation(
"",
&[crate::engine::ChatTurn {
user: true,
tool: false,
system: false,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
content: content.into(),
}],
ReasoningMode::Direct,
)
});
let mut executor = DeepSeekExecutor::open(
model,
64,
false,
64,
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: false,
cold: false,
cache_experts: 0,
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();
executor.prefill(&prompts[0], |_| true).unwrap();
let first = (executor.tokens().to_vec(), argmax(executor.logits()));
let mut inactive = None;
executor.swap_resident_state(&mut inactive).unwrap();
executor.prefill(&prompts[1], |_| true).unwrap();
let second = (executor.tokens().to_vec(), argmax(executor.logits()));
executor.swap_resident_state(&mut inactive).unwrap();
assert_eq!(
(executor.tokens(), argmax(executor.logits())),
(&*first.0, first.1)
);
executor.swap_resident_state(&mut inactive).unwrap();
assert_eq!(
(executor.tokens(), argmax(executor.logits())),
(&*second.0, second.1)
);
}
#[test]
#[ignore = "requires the installed Flash GGUF, DS4 steering fixture, and a Metal device"]
fn directional_steering_matches_the_ds4_token_oracle() {
@@ -7163,4 +7511,82 @@ mod tests {
assert_eq!(tokens, expected);
}
}
#[test]
#[ignore = "requires the installed DeepSeek V4 Pro GGUF and Apple Metal"]
fn pro_resident_and_ssd_streaming_choose_the_same_tokens() {
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 path = std::env::var("DS4_PRO_MODEL").unwrap_or_else(|_| {
"../ds4/models/DeepSeek-V4-Pro-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-Instruct-imatrix.gguf".into()
});
if !Path::new(&path).is_file() {
eprintln!("skipping unavailable Pro fixture: {path}");
return;
}
let run = |streaming| {
let model = Model::open_main(Path::new(&path), ModelChoice::DeepSeekV4Pro).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".into(),
}],
crate::settings::ReasoningMode::Direct,
);
let mut executor = DeepSeekExecutor::open(
model,
64,
false,
64,
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: streaming,
cold: true,
cache_experts: if streaming { 32 } else { 0 },
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();
executor.prefill(&prompt, |_| true).unwrap();
(0..4)
.map(|_| {
let token = argmax(executor.logits());
executor.eval(token).unwrap();
token
})
.collect::<Vec<_>>()
};
assert_eq!(run(false), run(true));
}
}

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

File diff suppressed because it is too large Load Diff

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

@@ -18,6 +18,23 @@ pub(super) struct StreamExpertTable {
pub(super) down_expert_bytes: u64,
}
#[derive(Clone, Copy, Default)]
#[repr(C)]
pub(super) struct StreamExpertCacheStats {
pub(super) configured_count: u32,
pub(super) current_count: u32,
pub(super) hits: u64,
pub(super) misses: u64,
pub(super) evictions: u64,
pub(super) wraps: u64,
pub(super) buffer_allocs: u64,
pub(super) buffer_reuses: u64,
pub(super) evict_advise_bytes: u64,
pub(super) willneed_advise_bytes: u64,
pub(super) pread_bytes: u64,
pub(super) pread_ms: f64,
}
unsafe extern "C" {
pub(super) fn ds4_gpu_init() -> i32;
pub(super) fn ds4_gpu_cleanup();
@@ -43,10 +60,7 @@ unsafe extern "C" {
pub(super) fn ds4_gpu_set_streaming_expert_cache_budget(experts: u32);
pub(super) fn ds4_gpu_set_streaming_expert_cache_expert_bytes(bytes: u64);
pub(super) fn ds4_gpu_recommended_working_set_size() -> u64;
pub(super) fn ds4_gpu_stream_expert_cache_budget_for_expert_size(
gate_expert_bytes: u64,
down_expert_bytes: u64,
) -> u32;
pub(super) fn ds4_gpu_stream_expert_cache_get_stats(stats: *mut StreamExpertCacheStats);
pub(super) fn ds4_gpu_stream_expert_cache_seed_experts(
table: *const StreamExpertTable,
expert_ids: *const i32,
@@ -266,6 +280,17 @@ unsafe extern "C" {
token: u32,
embd: u32,
) -> i32;
pub(super) fn ds4_gpu_embed_tokens_quant_tensor(
out: *mut GpuTensor,
tokens: *const GpuTensor,
map: *const c_void,
size: u64,
weight: u64,
kind: u32,
vocab: u32,
rows: u32,
embd: u32,
) -> i32;
pub(super) fn ds4_gpu_glm_rope_tail_tensor(
x: *mut GpuTensor,
tokens: u32,
@@ -301,6 +326,30 @@ unsafe extern "C" {
cache_f16: bool,
eps: f32,
) -> i32;
pub(super) fn ds4_gpu_glm_kv_lora_rms_norm_tensor(
out: *mut GpuTensor,
kv_raw: *const GpuTensor,
map: *const c_void,
size: u64,
weight: u64,
rows: u32,
kv_raw_dim: u32,
kv_lora: u32,
eps: f32,
) -> i32;
pub(super) fn ds4_gpu_glm_store_compact_kv_tensor(
kv_cache: *mut GpuTensor,
rope_cache: *mut GpuTensor,
kv_norm: *const GpuTensor,
kv_raw: *const GpuTensor,
pos: u32,
rows: u32,
cache_cap: u32,
kv_raw_dim: u32,
kv_lora: u32,
rot: u32,
cache_f16: bool,
) -> i32;
pub(super) fn ds4_gpu_glm_store_indexer_k_tensor(
cache: *mut GpuTensor,
raw: *const GpuTensor,
@@ -327,6 +376,13 @@ unsafe extern "C" {
selected: *mut GpuTensor,
count: u32,
) -> i32;
pub(super) fn ds4_gpu_glm_fill_selected_range_batch_tensor(
selected: *mut GpuTensor,
rows: u32,
pos: u32,
count: u32,
pad_row: u32,
) -> i32;
pub(super) fn ds4_gpu_glm_indexer_rope_tail_tensor(
x: *mut GpuTensor,
tokens: u32,
@@ -353,6 +409,19 @@ unsafe extern "C" {
scale: f32,
cache_f16: bool,
) -> i32;
pub(super) fn ds4_gpu_glm_indexer_scores_batch_tensor(
scores: *mut GpuTensor,
q: *const GpuTensor,
weights: *const GpuTensor,
cache: *const GpuTensor,
visible: u32,
rows: u32,
pos: u32,
heads: u32,
head_dim: u32,
scale: f32,
cache_f16: bool,
) -> i32;
pub(super) fn ds4_gpu_glm_qk_lowrank_typed_tensor(
out: *mut GpuTensor,
q: *const GpuTensor,
@@ -365,6 +434,31 @@ unsafe extern "C" {
q_nope: u32,
q_dim: u32,
) -> i32;
pub(super) fn ds4_gpu_glm_qk_lowrank_typed_batch_tensor(
out: *mut GpuTensor,
q: *const GpuTensor,
map: *const c_void,
size: u64,
weight: u64,
kind: u32,
rows: u32,
heads: u32,
kv_lora: u32,
q_nope: u32,
q_dim: u32,
) -> i32;
pub(super) fn ds4_gpu_glm_value_project_typed_batch_heads_tensor(
heads: *mut GpuTensor,
lora: *const GpuTensor,
map: *const c_void,
size: u64,
weight: u64,
kind: u32,
rows: u32,
n_head: u32,
kv_lora: u32,
value_dim: u32,
) -> i32;
pub(super) fn ds4_gpu_glm_attention_indexed_decode_typed_tensor(
heads_out: *mut GpuTensor,
q: *const GpuTensor,
@@ -392,6 +486,95 @@ unsafe extern "C" {
beta_fast: f32,
beta_slow: f32,
) -> i32;
pub(super) fn ds4_gpu_glm_attention_indexed_batch_lora_causal_tensor(
out: *mut GpuTensor,
q: *const GpuTensor,
qk_low: *const GpuTensor,
kv_cache: *const GpuTensor,
rope_cache: *const GpuTensor,
rows: u32,
pos: u32,
selected: u32,
cache_cap: u32,
cache_f16: bool,
heads: u32,
kv_lora: u32,
q_nope: u32,
rot: u32,
original: u32,
freq_base: f32,
freq_scale: f32,
ext: f32,
attn_factor: f32,
beta_fast: f32,
beta_slow: f32,
) -> i32;
pub(super) fn ds4_gpu_glm_attention_indexed_batch_lora_valid_tensor(
out: *mut GpuTensor,
q: *const GpuTensor,
qk_low: *const GpuTensor,
kv_cache: *const GpuTensor,
rope_cache: *const GpuTensor,
selected: *const GpuTensor,
rows: u32,
selected_count: u32,
cache_cap: u32,
cache_f16: bool,
heads: u32,
kv_lora: u32,
q_nope: u32,
rot: u32,
original: u32,
freq_base: f32,
freq_scale: f32,
ext: f32,
attn_factor: f32,
beta_fast: f32,
beta_slow: f32,
) -> i32;
pub(super) fn ds4_gpu_glm_router_select_batch_tensor(
selected: *mut GpuTensor,
weights: *mut GpuTensor,
probs: *mut GpuTensor,
map: *const c_void,
size: u64,
bias: u64,
logits: *const GpuTensor,
experts: u32,
experts_used: u32,
scale: f32,
rows: u32,
) -> i32;
pub(super) fn ds4_gpu_glm_routed_moe_batch_tensor(
out: *mut GpuTensor,
mid: *mut GpuTensor,
map: *const c_void,
size: u64,
gate_weight: u64,
up_weight: u64,
down_weight: u64,
gate_type: u32,
up_type: u32,
down_type: u32,
gate_expert_bytes: u64,
gate_row_bytes: u64,
up_expert_bytes: u64,
up_row_bytes: u64,
down_expert_bytes: u64,
down_row_bytes: u64,
input: u32,
middle: u32,
output: u32,
selected: *const GpuTensor,
weights: *const GpuTensor,
total_experts: u32,
used_experts: u32,
layer: u32,
x: *const GpuTensor,
rows: u32,
mid_token_stride: u32,
force_resident: bool,
) -> i32;
pub(super) fn ds4_gpu_glm_router_select_tensor(
selected: *mut GpuTensor,
weights: *mut GpuTensor,
@@ -1276,6 +1459,20 @@ impl Buffer {
)
}
pub(super) fn read_i32(&self, values: &mut [i32]) -> Result<(), String> {
check(
unsafe {
ds4_gpu_tensor_read(
self.raw(),
0,
values.as_mut_ptr().cast(),
std::mem::size_of_val(values) as u64,
)
},
"reading Metal integers",
)
}
pub(super) fn read(&self, offset: u64, values: &mut [u8]) -> Result<(), String> {
check(
unsafe {

13391
src/engine/metal/hotlist.rs
View File

File diff suppressed because it is too large Load Diff

307
src/engine/metal/profile.rs Normal file
View File

@@ -0,0 +1,307 @@
use super::gpu::Buffer;
use crate::model::ModelChoice;
use serde_json::{Value, json};
use std::cmp::Ordering;
use std::fs::{self, File};
use std::path::PathBuf;
const CACHE_CAPS: [usize; 10] = [1, 2, 4, 8, 16, 32, 64, 128, 256, 384];
#[derive(Default)]
struct LayerProfile {
records: u64,
counts: Vec<u64>,
weights: Vec<f64>,
caches: Vec<Vec<i32>>,
cache_hits: Vec<u64>,
cache_weight_hits: Vec<f64>,
previous: Option<(u32, Vec<i32>)>,
adjacent_pairs: u64,
adjacent_overlap: f64,
adjacent_jaccard: f64,
hash_router: bool,
}
pub(super) struct ExpertProfile {
path: PathBuf,
model: ModelChoice,
experts: usize,
used: usize,
caps: Vec<usize>,
layers: Vec<LayerProfile>,
}
impl ExpertProfile {
pub(super) fn new(
path: Option<&str>,
model: ModelChoice,
layers: u32,
experts: u64,
used: u64,
) -> Result<Option<Self>, String> {
let Some(path) = path else { return Ok(None) };
if path.trim().is_empty() {
return Err("Expert profile path cannot be empty".into());
}
let experts = usize::try_from(experts).map_err(|_| "expert count is too large")?;
let used = usize::try_from(used).map_err(|_| "selected expert count is too large")?;
let caps = CACHE_CAPS
.into_iter()
.filter(|cap| *cap <= experts)
.collect::<Vec<_>>();
let layers = (0..layers)
.map(|_| LayerProfile {
counts: vec![0; experts],
weights: vec![0.0; experts],
caches: vec![Vec::new(); caps.len()],
cache_hits: vec![0; caps.len()],
cache_weight_hits: vec![0.0; caps.len()],
..LayerProfile::default()
})
.collect();
Ok(Some(Self {
path: path.into(),
model,
experts,
used,
caps,
layers,
}))
}
pub(super) fn record(
&mut self,
layer: usize,
pos: u32,
selected: &Buffer,
weights: &Buffer,
rows: u32,
hash_router: bool,
) -> Result<(), String> {
let count = self
.used
.checked_mul(rows as usize)
.ok_or("expert profile row count overflow")?;
let mut ids = vec![0; count];
let mut route_weights = vec![0.0; count];
super::call(
unsafe { super::gpu::ds4_gpu_end_commands() },
"ending Metal commands for expert profiling",
)?;
let read = selected
.read_i32(&mut ids)
.and_then(|()| weights.read_f32(&mut route_weights));
let resumed = super::call(
unsafe { super::gpu::ds4_gpu_begin_commands() },
"resuming Metal commands after expert profiling",
);
read.and(resumed)?;
for (row, (ids, weights)) in ids
.chunks_exact(self.used)
.zip(route_weights.chunks_exact(self.used))
.enumerate()
{
self.record_row(layer, pos + row as u32, ids, weights, hash_router)?;
}
Ok(())
}
fn record_row(
&mut self,
layer: usize,
pos: u32,
ids: &[i32],
weights: &[f32],
hash_router: bool,
) -> Result<(), String> {
let profile = self
.layers
.get_mut(layer)
.ok_or("expert profile layer is outside the model")?;
profile.records += 1;
profile.hash_router |= hash_router;
if let Some((previous_pos, previous)) = &profile.previous
&& previous_pos.checked_add(1) == Some(pos)
{
let intersection = previous.iter().filter(|id| ids.contains(id)).count();
profile.adjacent_pairs += 1;
profile.adjacent_overlap += intersection as f64 / self.used as f64;
profile.adjacent_jaccard += intersection as f64 / (2 * self.used - intersection) as f64;
}
profile.previous = Some((pos, ids.to_vec()));
for (&id, &weight) in ids.iter().zip(weights) {
let expert = usize::try_from(id)
.ok()
.filter(|expert| *expert < self.experts)
.ok_or_else(|| format!("router selected invalid expert {id} at layer {layer}"))?;
profile.counts[expert] += 1;
profile.weights[expert] += f64::from(weight);
for (index, (&cap, cache)) in self.caps.iter().zip(&mut profile.caches).enumerate() {
if let Some(found) = cache.iter().position(|cached| *cached == id) {
profile.cache_hits[index] += 1;
profile.cache_weight_hits[index] += f64::from(weight);
cache.remove(found);
} else if cache.len() == cap {
cache.pop();
}
cache.insert(0, id);
}
}
Ok(())
}
pub(super) fn write(&self) -> Result<(), String> {
let selections: u64 = self
.layers
.iter()
.map(|layer| layer.counts.iter().sum::<u64>())
.sum();
let weight_total = self
.layers
.iter()
.flat_map(|layer| &layer.weights)
.sum::<f64>();
let cache_summary = self
.caps
.iter()
.enumerate()
.map(|(index, cap)| {
let hits = self
.layers
.iter()
.map(|layer| layer.cache_hits[index])
.sum::<u64>();
let weighted = self
.layers
.iter()
.map(|layer| layer.cache_weight_hits[index])
.sum::<f64>();
json!({
"n": cap,
"hits": hits,
"selections": selections,
"hit_rate": fraction(hits as f64, selections as f64),
"weighted_hit_rate": fraction(weighted, weight_total),
})
})
.collect::<Vec<_>>();
let layers = self
.layers
.iter()
.enumerate()
.map(|(index, layer)| self.layer_json(index, layer))
.collect::<Vec<_>>();
let value = json!({
"source": "ds4 Metal expert locality profile",
"model": self.model.id(),
"layers": self.layers.len(),
"experts": self.experts,
"expert_used": self.used,
"layer_records": self.layers.iter().map(|layer| layer.records).sum::<u64>(),
"selections": selections,
"cache_ns": self.caps,
"cache_summary": cache_summary,
"layers_detail": layers,
});
if let Some(parent) = self
.path
.parent()
.filter(|parent| !parent.as_os_str().is_empty())
{
fs::create_dir_all(parent).map_err(|error| error.to_string())?;
}
let temporary = self.path.with_extension("tmp");
let file = File::create(&temporary).map_err(|error| error.to_string())?;
serde_json::to_writer_pretty(&file, &value).map_err(|error| error.to_string())?;
file.sync_all().map_err(|error| error.to_string())?;
fs::rename(temporary, &self.path).map_err(|error| error.to_string())
}
fn layer_json(&self, index: usize, layer: &LayerProfile) -> Value {
let selections = layer.counts.iter().sum::<u64>();
let total_weight = layer.weights.iter().sum::<f64>();
let mut experts = (0..self.experts)
.filter(|expert| layer.counts[*expert] != 0)
.collect::<Vec<_>>();
experts.sort_by(|a, b| {
layer.counts[*b]
.cmp(&layer.counts[*a])
.then_with(|| {
layer.weights[*b]
.partial_cmp(&layer.weights[*a])
.unwrap_or(Ordering::Equal)
})
.then_with(|| a.cmp(b))
});
let top = experts
.into_iter()
.take(16)
.map(|expert| {
json!({
"id": expert,
"count": layer.counts[expert],
"pct": 100.0 * fraction(layer.counts[expert] as f64, selections as f64),
"weight": layer.weights[expert],
"weight_pct": 100.0 * fraction(layer.weights[expert], total_weight),
})
})
.collect::<Vec<_>>();
let cache = self
.caps
.iter()
.enumerate()
.map(|(cap_index, cap)| {
json!({
"n": cap,
"hits": layer.cache_hits[cap_index],
"hit_rate": fraction(layer.cache_hits[cap_index] as f64, selections as f64),
"weighted_hit_rate": fraction(layer.cache_weight_hits[cap_index], total_weight),
})
})
.collect::<Vec<_>>();
json!({
"layer": index,
"hash_router": layer.hash_router,
"records": layer.records,
"selections": selections,
"unique_experts": layer.counts.iter().filter(|count| **count != 0).count(),
"avg_adjacent_overlap": fraction(layer.adjacent_overlap, layer.adjacent_pairs as f64),
"avg_adjacent_jaccard": fraction(layer.adjacent_jaccard, layer.adjacent_pairs as f64),
"top_experts": top,
"cache": cache,
})
}
}
fn fraction(value: f64, total: f64) -> f64 {
if total > 0.0 { value / total } else { 0.0 }
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn profile_reports_locality_and_lru_hits() {
let path = std::env::temp_dir().join(format!(
"ds4-profile-{}-{}.json",
std::process::id(),
std::thread::current().name().unwrap_or("test")
));
let mut profile = ExpertProfile::new(path.to_str(), ModelChoice::DeepSeekV4Flash, 1, 8, 2)
.unwrap()
.unwrap();
profile
.record_row(0, 10, &[1, 2], &[0.6, 0.4], false)
.unwrap();
profile
.record_row(0, 11, &[1, 3], &[0.7, 0.3], false)
.unwrap();
profile.write().unwrap();
let value: Value = serde_json::from_slice(&fs::read(&path).unwrap()).unwrap();
assert_eq!(value["selections"], 4);
assert_eq!(value["layers_detail"][0]["avg_adjacent_overlap"], 0.5);
assert_eq!(value["cache_summary"][1]["hits"], 1);
let _ = fs::remove_file(path);
}
}