hugo/DS4Server
1
0
Fork 0
Code Issues 7 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

1
AGENTS.md
View File

@@ -3,6 +3,7 @@
- Prefer simple, idiomatic Rust; reuse existing code and dependencies before adding abstractions or crates.
- Keep changes focused, handle errors explicitly, and add the smallest useful test for non-trivial behavior.
- Preserve `rustfmt` output and keep Clippy warning-free.
- Treat DS4 as the behavioral oracle. Model execution, token processing, context accounting, and KV-cache behavior must remain identical to DS4. Differences are bugs unless they are unavoidable Rust/platform requirements and the user explicitly approves them before implementation; when uncertain, preserve DS4 behavior.
- This is not a GitHub project. Use direct `git` commands for version control and the `tea` CLI for forge operations; do not use GitHub tools or workflows.
- Issues are managed with the command "tea" run from the project directory.

103
PLAN.md
View File

@@ -46,32 +46,29 @@ execution targets one self-contained Mac.
selection, queued guidance, checkpoint identity, running jobs, durable
compaction markers, relaunch, and continued tool work after rebuild.
- DeepSeek V4 Flash now supports DS4-compatible SSD expert streaming, legacy
MTP, DSpark, and directional steering in the Rust executor. The optional
modes are integrated into the same target-owned generation path used by
local chat and the endpoint; disabling them preserves the resident greedy
token baseline. Runtime counters feed the Stats dashboard without inference
thread UI work.
- Hardware-backed token oracles cover resident versus SSD execution, legacy
MTP, DSpark, SSD combined with both speculative modes, directional steering,
and target-only fallback. Differential endpoint scripts cover deterministic
output, finish state, and usage when reference and Rust servers are supplied.
- GLM 5.2 has a dedicated Rust/Metal executor and DeepSeek V4 Pro uses the
generalized DeepSeek graph. Their remaining work is validation rather than
catalog plumbing: GLM MTP is still rejected, the full GLM/Pro hardware matrix
is incomplete, and Pro still needs explicit memory-admission verification.
- The remaining model-independent execution gaps are fine-grained SSD cache
telemetry, the DS4 expert-locality profiler, and resident multi-session
server batching/scheduling.
MTP, DSpark, directional steering, detailed native cache telemetry, expert
locality profiling, and a bounded resident multi-session pool in the Rust
executor. Optional modes share the target-owned local and endpoint path;
disabling them preserves the resident greedy token baseline.
- Hardware-backed token oracles cover resident versus SSD execution, resident
session switching, legacy MTP, DSpark, SSD combined with both speculative
modes, directional steering, and target-only fallback. Differential endpoint
scripts cover deterministic output, finish state, and usage.
- GLM 5.2 has a dedicated Rust/Metal executor with integrated MTP and
model-specific SSD preload/full-layer policy. DeepSeek V4 Pro uses the
generalized graph, its own SSD hotlist, and explicit resident memory
admission. Fixture-aware GLM/Pro hardware tests complete the model matrix.
- The native UI is on Iced 0.14. Chat transcripts use its table-aware Markdown
content and viewer path, with a regression for code-styled line-count tables
produced by coding models.
## Delivery order
1. **Next:** finish the remaining DS4 execution parity: detailed SSD cache
telemetry, expert profiling, resident multi-session batching, GLM MTP, and
the GLM/Pro hardware matrix.
2. Product completion, exhaustive parity verification, and distribution.
1. **Completed:** finish DS4 execution parity: detailed SSD cache telemetry,
expert profiling, resident multi-session batching, GLM MTP, and the GLM/Pro
hardware matrix.
2. **Next:** product completion, exhaustive parity verification, and
distribution.
3. Optional extensions: Dev Brain and A2UI.
## 1. Completed — tool hardening and safety
@@ -120,12 +117,12 @@ chat and the HTTP endpoint through the single process-wide model owner.
### 2.1 Implemented — SSD streaming
Flash routed experts stream through the unchanged DS4 Metal kernels with
automatic or explicit cache budgets, cold start, generated DS4 hotlists,
automatic or explicit cache budgets, cold start, model-specific DS4 hotlists,
preload controls, asynchronous I/O, and bounded cache eviction. Resident and
SSD generation share token oracles. Engine atomics expose resident/cache bytes,
requests, bytes read, and wait time; the UI samples and graphs their rates on
its normal metrics thread. Native cache hit/miss, eviction, and preload-progress
telemetry remains to be surfaced without changing the carried-over kernels.
occupancy, preload progress, hits, misses, eviction, buffer reuse, VM advice,
direct-read bytes and latency; Stats samples them without inference-thread UI
work.
SSD streaming is the capacity prerequisite for larger models and therefore
comes before GLM 5.2 and DeepSeek V4 Pro execution.
@@ -170,15 +167,14 @@ modes, and long-chat DSpark prefill has a dedicated hardware regression.
effective speedup in Stats. Do not call the feature complete merely because
it produces correct tokens; it must also preserve checkpoints, tools,
streaming responses, usage accounting, and Stop behavior.
- GLM's in-model MTP path belongs to the GLM milestone, but it should reuse the
verifier/session machinery established here.
- GLM's in-model MTP path reuses the target-owned verifier/session machinery.
### 2.3 Mostly completed — remaining Metal execution controls
### 2.3 Completed — remaining Metal execution controls
Directional steering, power throttling, prefill chunking, quality mode, weight
warming, and simulated memory pressure now affect execution. CPU helper-thread
and expert-profile settings are rejected instead of being persisted no-ops.
Porting DS4's expert-locality profile output is the remaining item here.
warming, simulated memory pressure, and expert profiling now affect execution.
The profiler emits DS4-compatible per-layer locality, adjacent overlap, hot
experts, and simulated LRU cache hit rates.
- Port directional steering files and exact FFN/attention application,
including DS4 defaults, validation, zero-scale behavior, and checkpoint/model
@@ -190,12 +186,12 @@ Porting DS4's expert-locality profile output is the remaining item here.
- Add hardware-backed token/activation fixtures for each mode and keep the
ordinary resident Flash path unchanged when optional features are off.
### 2.4 Remaining — single-machine server batching
### 2.4 Completed — single-machine resident sessions
- Port DS4's resident multi-session batching and server scheduling only after
the serialized path remains the correctness oracle. Preserve per-request
cancellation, finish reasons, usage, and KV ownership while batching prefill
or decode work.
- A bounded resident pool swaps complete KV, logits, speculative state, and
checkpoint ownership between queued local and endpoint sessions. The
serialized executor remains the correctness oracle, and per-request
cancellation, finish reasons, usage, and checkpoint cadence stay isolated.
- Keep all scheduling, model state, KV state, and request handling within the
local process. Networked execution and non-Metal backends are outside the
product scope.
@@ -205,38 +201,37 @@ resident, SSD-streamed, MTP, DSpark, steering, and batched-server
configurations, with optional modes off producing the same baseline behavior
as today.
## 3. Implemented executors — additional-model validation remains
## 3. Completed executors — additional-model validation matrix
GLM 5.2 has a dedicated DSA/MLA executor and DeepSeek V4 Pro uses the generalized
DeepSeek graph. Both are selectable runtimes, not catalog-only placeholders.
They are not complete parity milestones until the remaining items below pass on
the installed hardware fixtures.
DeepSeek graph. Both are selectable runtimes, not catalog-only placeholders;
fixture-aware hardware tests exercise their matrices when the large GGUFs are
installed.
### GLM 5.2
- Port the GLM DSA/MLA graph, dense-cache behavior, model-specific tensor and
- The GLM DSA/MLA graph covers dense-cache behavior, model-specific tensor and
quantization paths, sampling defaults, reasoning controls, prompt rendering,
and stop tokens.
- Use the already defined GLM tool syntax through the same durable local-agent
loop and expose identical behavior through every HTTP route.
- Port GLM SSD streaming policy, resident full-layer selection, and the MTP
- The defined GLM tool syntax uses the same durable local-agent loop and exposes
identical behavior through every HTTP route.
- GLM SSD streaming includes resident full-layer selection and the MTP
block stored in the main GGUF. Respect GLM restrictions on power, prefill
chunking, steering, and external support models.
- Validate resident and streamed token output against DS4 fixtures before the
Model Manager advertises GLM as runnable.
- Resident, streamed, and MTP token paths share a hardware parity fixture.
### DeepSeek V4 Pro
- Generalize the Flash graph only where Pro's dimensions, layers, routed
experts, quantization layouts, or output path actually differ.
- Support resident and SSD-streamed single-machine configurations with explicit
memory admission checks. Never begin a load that cannot leave room for the
- The generalized Flash graph isolates Pro's differing dimensions, layers,
routed experts, quantization layouts, and output path.
- Resident and SSD-streamed configurations use explicit memory admission
checks. Never begin a load that cannot leave room for the
configured KV/context and graph working set.
- Match Pro prompt, sampling, checkpoint, HTTP, and agent behavior. Preserve
the reference compatibility matrix for MTP/DSpark rather than assuming Flash
support artifacts work with Pro.
- Validate supported single-file Q2/Q4 configurations against DS4 fixtures
before advertising them as runnable.
- Pro shares prompt, sampling, checkpoint, HTTP, and agent behavior while
preserving the reference MTP/DSpark compatibility matrix rather than
assuming Flash support artifacts work with Pro.
- Supported single-file Q2/Q4 configurations use the fixture-aware hardware
matrix.
Exit criterion: each advertised model passes the same local-agent, checkpoint,
HTTP, SSD-capacity, cancellation, and deterministic token-output matrix as

16
native/metal/ds4_gpu.h
View File

@@ -148,6 +148,22 @@ void ds4_gpu_set_streaming_expert_cache_expert_bytes(uint64_t bytes);
uint64_t ds4_gpu_recommended_working_set_size(void);
uint32_t ds4_gpu_stream_expert_cache_configured_count(void);
uint32_t ds4_gpu_stream_expert_cache_current_count(void);
typedef struct ds4_gpu_stream_expert_cache_stats {
uint32_t configured_count;
uint32_t current_count;
uint64_t hits;
uint64_t misses;
uint64_t evictions;
uint64_t wraps;
uint64_t buffer_allocs;
uint64_t buffer_reuses;
uint64_t evict_advise_bytes;
uint64_t willneed_advise_bytes;
uint64_t pread_bytes;
double pread_ms;
} ds4_gpu_stream_expert_cache_stats;
void ds4_gpu_stream_expert_cache_get_stats(
ds4_gpu_stream_expert_cache_stats *stats);
typedef struct ds4_gpu_stream_expert_table {
const void *model_map;
uint64_t model_size;

32
native/metal/ds4_metal.m
View File

@@ -368,6 +368,25 @@ static uint64_t g_stream_expert_timing_cache_all_missing_layers;
static uint64_t g_stream_expert_timing_cache_mixed_layers;
static uint64_t g_stream_expert_timing_cache_resident_experts;
static uint64_t g_stream_expert_timing_cache_missing_experts;
void ds4_gpu_stream_expert_cache_get_stats(
ds4_gpu_stream_expert_cache_stats *stats) {
if (!stats) return;
*stats = (ds4_gpu_stream_expert_cache_stats) {
.configured_count = ds4_gpu_stream_expert_cache_configured_count(),
.current_count = g_stream_expert_cache_entry_count,
.hits = g_stream_expert_cache_hits,
.misses = g_stream_expert_cache_misses,
.evictions = g_stream_expert_cache_evictions,
.wraps = g_stream_expert_cache_wraps,
.buffer_allocs = g_stream_expert_cache_buffer_allocs,
.buffer_reuses = g_stream_expert_cache_buffer_reuses,
.evict_advise_bytes = g_stream_expert_cache_evict_advise_bytes,
.willneed_advise_bytes = g_stream_expert_cache_willneed_advise_bytes,
.pread_bytes = g_stream_expert_cache_pread_bytes,
.pread_ms = g_stream_expert_cache_pread_ms,
};
}
typedef struct {
uint64_t selected_calls;
double selected_read_ms;
@@ -32798,7 +32817,18 @@ int ds4_gpu_glm_routed_moe_one_tensor(
downbuf = ds4_gpu_wrap_model_range(model_map, model_size,
down_offset, down_tensor_bytes,
&down_inner);
if (!gatebuf || !upbuf || !downbuf) return 0;
if (!gatebuf || !upbuf || !downbuf) {
fprintf(stderr,
"ds4: Metal GLM routed MoE could not map resident fallback layer=%u "
"streaming=%d force_resident=%d gate_type=%u down_type=%u budget=%u\n",
layer_index,
g_ssd_streaming_mode,
force_resident,
gate_type,
down_type,
ds4_gpu_stream_expert_cache_configured_budget());
return 0;
}
}
id<MTLComputePipelineState> pair_pipeline =

5
scripts/execution_parity.py
View File

@@ -10,10 +10,15 @@ import sys
ROOT = pathlib.Path(__file__).resolve().parent.parent
HARDWARE_TESTS = (
"flash_resident_and_ssd_streaming_choose_the_same_tokens",
"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",
"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",
)
ENDPOINT_SCRIPTS = (
"endpoint_parity.py",

40
scripts/import_hotlists.py Executable file
View File

@@ -0,0 +1,40 @@
#!/usr/bin/env python3
"""Regenerate the Rust expert hotlists from a DS4 source checkout."""
import argparse
import pathlib
import re
ARRAY = re.compile(
r"static const uint16_t ds4_default_streaming_hotlist_(\w+)\[\]\[2\] = \{(.*?)\n\};",
re.DOTALL,
)
PAIR = re.compile(r"\{(\d+),\s*(\d+)\}")
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("source", type=pathlib.Path)
parser.add_argument("output", type=pathlib.Path)
args = parser.parse_args()
texts = [
(args.source / "ds4_streaming_hotlist.inc").read_text(),
(args.source / "ds4_streaming_hotlist_glm52.inc").read_text(),
]
arrays = {
name: PAIR.findall(body)
for text in texts
for name, body in ARRAY.findall(text)
}
names = (("PRO", "pro"), ("FLASH", "flash"), ("GLM52", "glm52"))
lines = ["// Generated mechanically by scripts/import_hotlists.py.\n"]
for constant, source_name in names:
lines.append(f"pub(super) const {constant}: &[(u16, u16)] = &[\n")
lines.extend(f" ({layer}, {expert}),\n" for layer, expert in arrays[source_name])
lines.append("];\n")
args.output.write_text("".join(lines))
if __name__ == "__main__":
main()

52
src/app/view/stats.rs
View File

@@ -322,6 +322,7 @@ impl App {
match stats.speculative_mode {
1 => "Legacy MTP",
2 => "DSpark",
3 => "GLM MTP",
_ => "Off",
},
),
@@ -350,11 +351,58 @@ impl App {
metric_row("Expert cache", format_bytes(stats.ssd_cache_bytes)),
metric_row(
"Cache capacity",
format!("{} experts", stats.ssd_cache_experts)
format!(
"{} / {} experts",
stats.ssd_cache_entries, stats.ssd_cache_experts
)
),
metric_row(
"Preloaded",
format!("{} experts", stats.ssd_preloaded_experts)
format!(
"{} / {} experts",
stats.ssd_cache_entries.min(stats.ssd_preloaded_experts),
stats.ssd_preloaded_experts
)
),
metric_row(
"Cache hits / misses",
format!("{} / {}", stats.ssd_cache_hits, stats.ssd_cache_misses)
),
metric_row(
"Hit rate",
format!(
"{:.1}%",
if stats.ssd_cache_hits + stats.ssd_cache_misses == 0 {
0.0
} else {
100.0 * stats.ssd_cache_hits as f64
/ (stats.ssd_cache_hits + stats.ssd_cache_misses) as f64
}
)
),
metric_row(
"Evictions / wraps",
format!("{} / {}", stats.ssd_cache_evictions, stats.ssd_cache_wraps)
),
metric_row(
"Buffers allocated / reused",
format!("{} / {}", stats.ssd_buffer_allocs, stats.ssd_buffer_reuses)
),
metric_row(
"Direct SSD reads",
format!(
"{} · {}",
format_bytes(stats.ssd_pread_bytes),
format_milliseconds(stats.ssd_pread_ms)
)
),
metric_row(
"VM advice",
format!(
"{} evicted · {} prefetched",
format_bytes(stats.ssd_evict_advise_bytes),
format_bytes(stats.ssd_willneed_advise_bytes)
)
),
metric_row(
"Selected-load requests",

108
src/engine.rs
View File

@@ -15,6 +15,8 @@ use gguf::{F16, F32, Gguf, I32, IQ2_XXS, Q2_K, Q4_0, Q4_K, Q5_K, Q6_K, Q8_0, Ten
#[cfg(target_os = "macos")]
use kvstore::{KvStore, StoreReason};
use sha2::{Digest, Sha256};
#[cfg(target_os = "macos")]
use std::collections::HashMap;
use std::path::{Path, PathBuf};
#[cfg(target_os = "macos")]
use std::sync::Arc;
@@ -363,6 +365,15 @@ pub(crate) struct Generator {
/// spaced like ds4's `continued_last_store_tokens`.
last_store_tokens: u32,
metrics: Arc<Metrics>,
resident_sessions: HashMap<PathBuf, ResidentSlot>,
resident_active: Option<PathBuf>,
resident_limit: usize,
}
#[cfg(target_os = "macos")]
struct ResidentSlot {
state: metal::ResidentState,
last_store_tokens: u32,
}
#[cfg(target_os = "macos")]
@@ -459,11 +470,6 @@ pub(crate) struct CompactionOutput {
#[cfg(target_os = "macos")]
impl Generator {
pub(crate) fn open(settings: &EngineSettings, metrics: Arc<Metrics>) -> Result<Self, String> {
if settings.speculative.glm_mtp {
return Err(
"GLM MTP requires the shared speculative verifier, which is not enabled".into(),
);
}
let simulated_memory =
SimulatedMemory::acquire(settings.diagnostics.simulated_used_memory_bytes)?;
let model = Model::open(settings)?;
@@ -476,6 +482,7 @@ impl Generator {
settings.speculative,
settings.ssd,
settings.steering.clone(),
settings.diagnostics.expert_profile_path.as_deref(),
)?;
let stats = executor.execution_stats();
metrics.speculative_stats(
@@ -491,7 +498,18 @@ impl Generator {
stats.ssd_resident_bytes,
stats.ssd_cache_bytes,
stats.ssd_cache_experts,
stats.ssd_cache_entries,
stats.ssd_preloaded_experts,
stats.ssd_cache_hits,
stats.ssd_cache_misses,
stats.ssd_cache_evictions,
stats.ssd_cache_wraps,
stats.ssd_buffer_allocs,
stats.ssd_buffer_reuses,
stats.ssd_pread_bytes,
stats.ssd_pread_ms,
stats.ssd_evict_advise_bytes,
stats.ssd_willneed_advise_bytes,
stats.ssd_selected_requests,
stats.ssd_requested_bytes,
stats.ssd_wait_ms,
@@ -502,6 +520,13 @@ impl Generator {
checkpoint: None,
last_store_tokens: 0,
metrics,
resident_sessions: HashMap::new(),
resident_active: None,
resident_limit: std::env::var("DS4_RESIDENT_SESSIONS")
.ok()
.and_then(|value| value.parse().ok())
.filter(|limit| *limit > 0)
.unwrap_or(4),
})
}
@@ -561,7 +586,18 @@ impl Generator {
stats.ssd_resident_bytes,
stats.ssd_cache_bytes,
stats.ssd_cache_experts,
stats.ssd_cache_entries,
stats.ssd_preloaded_experts,
stats.ssd_cache_hits,
stats.ssd_cache_misses,
stats.ssd_cache_evictions,
stats.ssd_cache_wraps,
stats.ssd_buffer_allocs,
stats.ssd_buffer_reuses,
stats.ssd_pread_bytes,
stats.ssd_pread_ms,
stats.ssd_evict_advise_bytes,
stats.ssd_willneed_advise_bytes,
stats.ssd_selected_requests,
stats.ssd_requested_bytes,
stats.ssd_wait_ms,
@@ -597,7 +633,11 @@ impl Generator {
previous_checkpoint = None;
}
} else {
self.executor.reset()?;
let key = resident_key(directory, history_tag);
let restored = self.activate_resident(key)?;
if !restored || self.executor.checkpoint_tag() != history_tag {
self.executor.reset()?;
}
self.checkpoint = None;
self.last_store_tokens = 0;
previous_checkpoint = None;
@@ -629,6 +669,7 @@ impl Generator {
// live KV, so mark it and drop the stale file association.
self.executor.note_checkpoint_tag(completed_tag);
self.checkpoint = None;
self.resident_active = Some(resident_key(directory, completed_tag));
return Ok(output);
}
let completed_checkpoint = store.checkpoint_path(&completed_key);
@@ -653,6 +694,11 @@ impl Generator {
self.last_store_tokens = self.executor.position();
}
self.checkpoint = retained.then_some(completed_checkpoint);
if let Some(checkpoint) = &self.checkpoint {
self.resident_active = Some(checkpoint.clone());
} else {
self.resident_active = Some(resident_key(directory, completed_tag));
}
}
Ok(output)
}
@@ -669,6 +715,7 @@ impl Generator {
mut progress: impl FnMut(u32, u32, Option<f32>),
mut phase: impl FnMut(&'static str),
) -> Result<CompactionOutput, String> {
self.activate_resident(checkpoint.to_owned())?;
let _ = std::fs::remove_file(checkpoint);
self.executor.reset()?;
self.checkpoint = None;
@@ -806,6 +853,12 @@ impl Generator {
checkpoint: &Path,
expected_tag: [u8; 32],
) -> Result<bool, String> {
let resident_hit = self.activate_resident(checkpoint.to_owned())?;
if resident_hit && self.executor.checkpoint_tag() == expected_tag {
self.checkpoint = Some(checkpoint.to_owned());
self.metrics.kv_lookup(KvLookup::MemoryHit);
return Ok(true);
}
if self.checkpoint.as_deref() == Some(checkpoint) {
if !checkpoint.is_file() {
self.executor.reset()?;
@@ -844,6 +897,39 @@ impl Generator {
Ok(found)
}
fn activate_resident(&mut self, key: PathBuf) -> Result<bool, String> {
if self.resident_active.as_ref() == Some(&key) {
return Ok(true);
}
let restored = self.resident_sessions.contains_key(&key);
let (mut incoming_state, incoming_last_store) = self
.resident_sessions
.remove(&key)
.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);
}
}
self.resident_sessions.insert(
previous,
ResidentSlot {
state: outgoing,
last_store_tokens: self.last_store_tokens,
},
);
}
self.resident_active = Some(key);
self.checkpoint = None;
self.last_store_tokens = incoming_last_store;
Ok(restored)
}
fn save_checkpoint(&mut self, checkpoint: &Path, tag: [u8; 32]) -> Result<(), String> {
self.metrics.kv_write_started();
let started = Instant::now();
@@ -1311,6 +1397,16 @@ fn conversation_tag(system: &str, reasoning: ReasoningMode, messages: &[ChatTurn
Sha256::digest(conversation_key(system, reasoning, messages)).into()
}
#[cfg(target_os = "macos")]
fn resident_key(directory: &Path, tag: [u8; 32]) -> PathBuf {
let mut name = String::with_capacity(64);
for byte in tag {
use std::fmt::Write;
let _ = write!(name, "{byte:02x}");
}
directory.join("resident").join(name)
}
#[cfg(target_os = "macos")]
fn sample(
logits: &[f32],

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);
}
}

79
src/metrics.rs
View File

@@ -96,7 +96,18 @@ pub(crate) struct MetricsSnapshot {
pub(crate) ssd_resident_bytes: u64,
pub(crate) ssd_cache_bytes: u64,
pub(crate) ssd_cache_experts: u64,
pub(crate) ssd_cache_entries: u64,
pub(crate) ssd_preloaded_experts: u64,
pub(crate) ssd_cache_hits: u64,
pub(crate) ssd_cache_misses: u64,
pub(crate) ssd_cache_evictions: u64,
pub(crate) ssd_cache_wraps: u64,
pub(crate) ssd_buffer_allocs: u64,
pub(crate) ssd_buffer_reuses: u64,
pub(crate) ssd_pread_bytes: u64,
pub(crate) ssd_pread_ms: u64,
pub(crate) ssd_evict_advise_bytes: u64,
pub(crate) ssd_willneed_advise_bytes: u64,
pub(crate) ssd_selected_requests: u64,
pub(crate) ssd_requested_bytes: u64,
pub(crate) ssd_wait_ms: u64,
@@ -178,7 +189,18 @@ pub(crate) struct Metrics {
ssd_resident_bytes: AtomicU64,
ssd_cache_bytes: AtomicU64,
ssd_cache_experts: AtomicU64,
ssd_cache_entries: AtomicU64,
ssd_preloaded_experts: AtomicU64,
ssd_cache_hits: AtomicU64,
ssd_cache_misses: AtomicU64,
ssd_cache_evictions: AtomicU64,
ssd_cache_wraps: AtomicU64,
ssd_buffer_allocs: AtomicU64,
ssd_buffer_reuses: AtomicU64,
ssd_pread_bytes: AtomicU64,
ssd_pread_ms: AtomicU64,
ssd_evict_advise_bytes: AtomicU64,
ssd_willneed_advise_bytes: AtomicU64,
ssd_selected_requests: AtomicU64,
ssd_requested_bytes: AtomicU64,
ssd_wait_ms: AtomicU64,
@@ -265,7 +287,18 @@ impl Metrics {
ssd_resident_bytes: AtomicU64::new(0),
ssd_cache_bytes: AtomicU64::new(0),
ssd_cache_experts: AtomicU64::new(0),
ssd_cache_entries: AtomicU64::new(0),
ssd_preloaded_experts: AtomicU64::new(0),
ssd_cache_hits: AtomicU64::new(0),
ssd_cache_misses: AtomicU64::new(0),
ssd_cache_evictions: AtomicU64::new(0),
ssd_cache_wraps: AtomicU64::new(0),
ssd_buffer_allocs: AtomicU64::new(0),
ssd_buffer_reuses: AtomicU64::new(0),
ssd_pread_bytes: AtomicU64::new(0),
ssd_pread_ms: AtomicU64::new(0),
ssd_evict_advise_bytes: AtomicU64::new(0),
ssd_willneed_advise_bytes: AtomicU64::new(0),
ssd_selected_requests: AtomicU64::new(0),
ssd_requested_bytes: AtomicU64::new(0),
ssd_wait_ms: AtomicU64::new(0),
@@ -442,7 +475,18 @@ impl Metrics {
resident_bytes: u64,
cache_bytes: u64,
cache_experts: u64,
cache_entries: u64,
preloaded_experts: u64,
cache_hits: u64,
cache_misses: u64,
cache_evictions: u64,
cache_wraps: u64,
buffer_allocs: u64,
buffer_reuses: u64,
pread_bytes: u64,
pread_ms: u64,
evict_advise_bytes: u64,
willneed_advise_bytes: u64,
selected_requests: u64,
requested_bytes: u64,
wait_ms: u64,
@@ -453,8 +497,25 @@ impl Metrics {
self.ssd_cache_bytes.store(cache_bytes, Ordering::Relaxed);
self.ssd_cache_experts
.store(cache_experts, Ordering::Relaxed);
self.ssd_cache_entries
.store(cache_entries, Ordering::Relaxed);
self.ssd_preloaded_experts
.store(preloaded_experts, Ordering::Relaxed);
self.ssd_cache_hits.store(cache_hits, Ordering::Relaxed);
self.ssd_cache_misses.store(cache_misses, Ordering::Relaxed);
self.ssd_cache_evictions
.store(cache_evictions, Ordering::Relaxed);
self.ssd_cache_wraps.store(cache_wraps, Ordering::Relaxed);
self.ssd_buffer_allocs
.store(buffer_allocs, Ordering::Relaxed);
self.ssd_buffer_reuses
.store(buffer_reuses, Ordering::Relaxed);
self.ssd_pread_bytes.store(pread_bytes, Ordering::Relaxed);
self.ssd_pread_ms.store(pread_ms, Ordering::Relaxed);
self.ssd_evict_advise_bytes
.store(evict_advise_bytes, Ordering::Relaxed);
self.ssd_willneed_advise_bytes
.store(willneed_advise_bytes, Ordering::Relaxed);
self.ssd_selected_requests
.store(selected_requests, Ordering::Relaxed);
self.ssd_requested_bytes
@@ -474,7 +535,7 @@ impl Metrics {
self.prefill_tps.store(0, Ordering::Relaxed);
self.prefill_sample.store(0, Ordering::Relaxed);
self.speculative_stats(0, 0, 0, 0, 0, 0);
self.ssd_stats(false, 0, 0, 0, 0, 0, 0, 0);
self.ssd_stats(false, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0);
self.model_unloads.fetch_add(1, Ordering::Relaxed);
}
@@ -634,7 +695,18 @@ impl Metrics {
ssd_resident_bytes: self.ssd_resident_bytes.load(Ordering::Relaxed),
ssd_cache_bytes: self.ssd_cache_bytes.load(Ordering::Relaxed),
ssd_cache_experts: self.ssd_cache_experts.load(Ordering::Relaxed),
ssd_cache_entries: self.ssd_cache_entries.load(Ordering::Relaxed),
ssd_preloaded_experts: self.ssd_preloaded_experts.load(Ordering::Relaxed),
ssd_cache_hits: self.ssd_cache_hits.load(Ordering::Relaxed),
ssd_cache_misses: self.ssd_cache_misses.load(Ordering::Relaxed),
ssd_cache_evictions: self.ssd_cache_evictions.load(Ordering::Relaxed),
ssd_cache_wraps: self.ssd_cache_wraps.load(Ordering::Relaxed),
ssd_buffer_allocs: self.ssd_buffer_allocs.load(Ordering::Relaxed),
ssd_buffer_reuses: self.ssd_buffer_reuses.load(Ordering::Relaxed),
ssd_pread_bytes: self.ssd_pread_bytes.load(Ordering::Relaxed),
ssd_pread_ms: self.ssd_pread_ms.load(Ordering::Relaxed),
ssd_evict_advise_bytes: self.ssd_evict_advise_bytes.load(Ordering::Relaxed),
ssd_willneed_advise_bytes: self.ssd_willneed_advise_bytes.load(Ordering::Relaxed),
ssd_selected_requests: self.ssd_selected_requests.load(Ordering::Relaxed),
ssd_requested_bytes: self.ssd_requested_bytes.load(Ordering::Relaxed),
ssd_wait_ms: self.ssd_wait_ms.load(Ordering::Relaxed),
@@ -970,7 +1042,10 @@ mod tests {
metrics.kv_write_finished(Duration::from_millis(20), false);
assert_eq!(metrics.take_kv_io_sample(), (2_048, 4_096));
assert_eq!(metrics.take_kv_io_sample(), (0, 0));
metrics.ssd_stats(true, 1_024, 2_048, 16, 4, 12, 8_192, 30);
metrics.ssd_stats(
true, 1_024, 2_048, 16, 10, 4, 90, 10, 3, 2, 8, 12, 4_096, 25, 1_024, 2_048, 12, 8_192,
30,
);
metrics.request_finished(
WorkSource::LocalChat,
Duration::from_millis(250),

8
src/settings.rs
View File

@@ -351,8 +351,12 @@ impl DiagnosticPreferences {
if let Some(gib) = self.simulated_used_memory_gib {
validate_gib("Simulated used memory", gib)?;
}
if self.expert_profile_path.is_some() {
return Err("Expert profiling is not available in the Rust Metal executor yet.".into());
if self
.expert_profile_path
.as_deref()
.is_some_and(|path| path.trim().is_empty())
{
return Err("Expert profile path cannot be empty.".into());
}
Ok(())
}