Integrate DS4 execution parity in Rust

This commit is contained in:
Georg Bauer
2026-07-26 17:58:05 +02:00
parent c9f0c3661c
commit 4420b81117
20 changed files with 11643 additions and 358 deletions

70
PLAN.md
View File

@@ -45,18 +45,31 @@ execution targets one self-contained Mac.
- Focused coverage exercises triggers, summary bounds and sanitizing, tail - Focused coverage exercises triggers, summary bounds and sanitizing, tail
selection, queued guidance, checkpoint identity, running jobs, durable selection, queued guidance, checkpoint identity, running jobs, durable
compaction markers, relaunch, and continued tool work after rebuild. compaction markers, relaunch, and continued tool work after rebuild.
- The next baseline gap is SSD streaming. Speculative decoding, steering, GLM - DeepSeek V4 Flash now supports DS4-compatible SSD expert streaming, legacy
5.2 execution, and DeepSeek V4 Pro execution are not implemented in the Rust MTP, DSpark, and directional steering in the Rust executor. The optional
executor. Related catalog, validation, and preference plumbing must not be modes are integrated into the same target-owned generation path used by
treated as runtime support. 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.
## Delivery order ## Delivery order
1. **Next:** remaining DS4 execution technology, starting with SSD streaming, then 1. **Next:** finish the remaining DS4 execution parity: detailed SSD cache
speculative decoding and the other Metal/runtime parity work. telemetry, expert profiling, resident multi-session batching, GLM MTP, and
2. Additional model execution: GLM 5.2 and DeepSeek V4 Pro. the GLM/Pro hardware matrix.
3. Product completion, exhaustive parity verification, and distribution. 2. Product completion, exhaustive parity verification, and distribution.
4. Optional extensions: Dev Brain and A2UI. 3. Optional extensions: Dev Brain and A2UI.
## 1. Completed — tool hardening and safety ## 1. Completed — tool hardening and safety
@@ -95,13 +108,21 @@ Exit criterion: use the agent for a real inspect/edit/test cycle while every
side effect is visible, risky actions require consent, Stop works at every side effect is visible, risky actions require consent, Stop works at every
stage, and no file tool can escape the selected project. stage, and no file tool can escape the selected project.
## 2. DS4 execution technology parity ## 2. In progress — DS4 execution technology parity
Goal: finish the model-independent Metal/runtime capabilities in `ds4.c` Goal: finish the model-independent Metal/runtime capabilities in `ds4.c`
before adding larger model families. Every capability must be shared by local before adding larger model families. Every capability must be shared by local
chat and the HTTP endpoint through the single process-wide model owner. chat and the HTTP endpoint through the single process-wide model owner.
### 2.1 SSD streaming — first runtime priority ### 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,
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.
SSD streaming is the capacity prerequisite for larger models and therefore SSD streaming is the capacity prerequisite for larger models and therefore
comes before GLM 5.2 and DeepSeek V4 Pro execution. comes before GLM 5.2 and DeepSeek V4 Pro execution.
@@ -122,7 +143,14 @@ comes before GLM 5.2 and DeepSeek V4 Pro execution.
the cache and I/O layer model-aware so later GLM/Pro milestones add policy and the cache and I/O layer model-aware so later GLM/Pro milestones add policy and
graph support rather than a second streaming subsystem. graph support rather than a second streaming subsystem.
### 2.2 Speculative decoding: legacy MTP and DSpark ### 2.2 Completed for Flash — legacy MTP and DSpark
Both support GGUF formats are validated and executed by Rust. Draft generation,
confidence/margin gates, target verification, accepted-prefix commit, rollback,
strict/quality fallback, cancellation, checkpoint identity, streaming output,
usage accounting, and Stats counters use the existing generation lifecycle.
Sampled decoding remains target-only as in DS4. SSD streaming works with both
modes, and long-chat DSpark prefill has a dedicated hardware regression.
- Load and validate the optional Flash legacy-MTP or DSpark support GGUF without - Load and validate the optional Flash legacy-MTP or DSpark support GGUF without
treating either as a standalone model. Preserve exact support-kind and target treating either as a standalone model. Preserve exact support-kind and target
@@ -142,7 +170,12 @@ comes before GLM 5.2 and DeepSeek V4 Pro execution.
- GLM's in-model MTP path belongs to the GLM milestone, but it should reuse the - GLM's in-model MTP path belongs to the GLM milestone, but it should reuse the
verifier/session machinery established here. verifier/session machinery established here.
### 2.3 Remaining Metal execution controls ### 2.3 Mostly 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.
- Port directional steering files and exact FFN/attention application, - Port directional steering files and exact FFN/attention application,
including DS4 defaults, validation, zero-scale behavior, and checkpoint/model including DS4 defaults, validation, zero-scale behavior, and checkpoint/model
@@ -154,7 +187,7 @@ comes before GLM 5.2 and DeepSeek V4 Pro execution.
- Add hardware-backed token/activation fixtures for each mode and keep the - Add hardware-backed token/activation fixtures for each mode and keep the
ordinary resident Flash path unchanged when optional features are off. ordinary resident Flash path unchanged when optional features are off.
### 2.4 Single-machine server batching ### 2.4 Remaining — single-machine server batching
- Port DS4's resident multi-session batching and server scheduling only after - Port DS4's resident multi-session batching and server scheduling only after
the serialized path remains the correctness oracle. Preserve per-request the serialized path remains the correctness oracle. Preserve per-request
@@ -169,11 +202,12 @@ resident, SSD-streamed, MTP, DSpark, steering, and batched-server
configurations, with optional modes off producing the same baseline behavior configurations, with optional modes off producing the same baseline behavior
as today. as today.
## 3. Additional model execution ## 3. Implemented executors — additional-model validation remains
Start these only after the shared capacity and execution technology above is GLM 5.2 has a dedicated DSA/MLA executor and DeepSeek V4 Pro uses the generalized
stable. Catalog entries, settings, tokenizer work, or GGUF validation alone do DeepSeek graph. Both are selectable runtimes, not catalog-only placeholders.
not make a model selectable for inference. They are not complete parity milestones until the remaining items below pass on
the installed hardware fixtures.
### GLM 5.2 ### GLM 5.2

62
scripts/execution_parity.py Executable file
View File

@@ -0,0 +1,62 @@
#!/usr/bin/env python3
"""Run the DS4 token oracles, then the existing endpoint smoke corpus."""
import argparse
import pathlib
import subprocess
import sys
ROOT = pathlib.Path(__file__).resolve().parent.parent
HARDWARE_TESTS = (
"flash_resident_and_ssd_streaming_choose_the_same_tokens",
"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",
)
ENDPOINT_SCRIPTS = (
"endpoint_parity.py",
"endpoint_reasoning.py",
"endpoint_continuation.py",
)
def run(command):
print("+", " ".join(map(str, command)), flush=True)
subprocess.run(command, cwd=ROOT, check=True)
def main():
parser = argparse.ArgumentParser()
parser.add_argument(
"--endpoint",
help="also run the existing endpoint parity scripts against this URL",
)
parser.add_argument(
"--skip-hardware",
action="store_true",
help="skip tests requiring the local Flash/MTP/DSpark GGUF fixtures",
)
args = parser.parse_args()
if not args.skip_hardware:
for test in HARDWARE_TESTS:
run(
[
"cargo",
"test",
"--all-features",
test,
"--",
"--ignored",
"--nocapture",
]
)
if args.endpoint:
for script in ENDPOINT_SCRIPTS:
run([sys.executable, ROOT / "scripts" / script, args.endpoint])
if __name__ == "__main__":
main()

View File

@@ -0,0 +1,23 @@
Reproduce the following C code EXACTLY, character for character, inside a single code block and output nothing else:
```c
static uint32_t clamp_u32(uint32_t v, uint32_t lo, uint32_t hi) {
if (v < lo) return lo;
if (v > hi) return hi;
return v;
}
static uint32_t ring_advance(uint32_t pos, uint32_t cap) {
uint32_t next = pos + 1u;
return next >= cap ? 0u : next;
}
static int scratch_init(scratch *s, uint32_t ctx_size) {
if (ctx_size == 0u) ctx_size = 1u;
s->ctx_size = ctx_size;
s->comp_cap = ctx_size / 4u + 2u;
s->rows = clamp_u32(s->comp_cap, 1u, 4096u);
s->head = 0u;
return s->rows > 0u ? 0 : -1;
}
```

131
scripts/speculative_parity.py Executable file
View File

@@ -0,0 +1,131 @@
#!/usr/bin/env python3
"""Compare DS4 and Rust endpoints configured for the same decoding mode."""
import argparse
import json
import urllib.request
CASES = [
"hi",
"Reply with exactly three words describing a calm sea.",
"Write the first eight positive odd numbers separated by commas.",
]
def post(base_url, path, payload):
request = urllib.request.Request(
base_url.rstrip("/") + path,
data=json.dumps(payload).encode(),
headers={"Content-Type": "application/json"},
)
with urllib.request.urlopen(request, timeout=600) as response:
return json.load(response)
def chat(base_url, prompt):
result = post(
base_url,
"/v1/chat/completions",
{
"model": "deepseek-v4-flash",
"messages": [{"role": "user", "content": prompt}],
"reasoning_effort": "none",
"temperature": 0,
"max_tokens": 64,
},
)
choice = result["choices"][0]
message = choice["message"]
return {
"content": message.get("content"),
"reasoning": message.get("reasoning_content"),
"finish_reason": choice.get("finish_reason"),
"usage": result.get("usage"),
}
def anthropic(base_url, prompt):
result = post(
base_url,
"/v1/messages",
{
"model": "deepseek-v4-flash",
"messages": [{"role": "user", "content": prompt}],
"thinking": {"type": "disabled"},
"temperature": 0,
"max_tokens": 64,
},
)
return {
"content": result.get("content"),
"stop_reason": result.get("stop_reason"),
"usage": result.get("usage"),
}
def responses(base_url, prompt):
result = post(
base_url,
"/v1/responses",
{
"model": "deepseek-v4-flash",
"input": prompt,
"reasoning": {"effort": "none"},
"temperature": 0,
"max_output_tokens": 64,
},
)
return {
"output": normalize(result.get("output")),
"status": result.get("status"),
"usage": result.get("usage"),
}
def normalize(value):
if isinstance(value, list):
return [normalize(item) for item in value]
if isinstance(value, dict):
return {
key: normalize(item)
for key, item in value.items()
if key not in {"id", "created_at"}
}
return value
def main():
parser = argparse.ArgumentParser()
parser.add_argument("reference_url", help="DS4 reference server URL")
parser.add_argument("rust_url", help="Rust DS4Server URL in the same mode")
args = parser.parse_args()
failures = []
for prompt in CASES:
for name, request in (
("chat", chat),
("anthropic", anthropic),
("responses", responses),
):
reference = request(args.reference_url, prompt)
rust = request(args.rust_url, prompt)
if reference != rust:
failures.append(
{
"case": name,
"prompt": prompt,
"reference": reference,
"rust": rust,
}
)
else:
print(f"ok {name}: {prompt}")
if failures:
print(json.dumps(failures, indent=2, ensure_ascii=False))
raise SystemExit(1)
print("all deterministic DS4/Rust endpoint outputs and usage records match")
if __name__ == "__main__":
main()

View File

@@ -90,6 +90,9 @@ pub(crate) struct App {
pub(super) kv_cache_report: KvCacheReport, pub(super) kv_cache_report: KvCacheReport,
last_cache_scan: Instant, last_cache_scan: Instant,
last_http_requests: u64, last_http_requests: u64,
last_ssd_selected_requests: u64,
last_ssd_requested_bytes: u64,
last_ssd_wait_ms: u64,
#[cfg(target_os = "macos")] #[cfg(target_os = "macos")]
generation_service: Option<GenerationService>, generation_service: Option<GenerationService>,
#[cfg(target_os = "macos")] #[cfg(target_os = "macos")]
@@ -136,6 +139,9 @@ pub(super) struct MetricsPoint {
pub(super) http_requests_per_second: f32, pub(super) http_requests_per_second: f32,
pub(super) kv_read_bytes_per_second: f32, pub(super) kv_read_bytes_per_second: f32,
pub(super) kv_write_bytes_per_second: f32, pub(super) kv_write_bytes_per_second: f32,
pub(super) ssd_requests_per_second: f32,
pub(super) ssd_bytes_per_second: f32,
pub(super) ssd_wait_ms_per_second: f32,
} }
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
@@ -154,6 +160,7 @@ pub(crate) enum Message {
FocusNext, FocusNext,
FocusPrevious, FocusPrevious,
PreferenceModelChanged(ModelChoice), PreferenceModelChanged(ModelChoice),
PreferenceLegacyMtpChanged(bool),
PreferenceDsparkChanged(bool), PreferenceDsparkChanged(bool),
PreferenceTimeoutChanged(String), PreferenceTimeoutChanged(String),
PreferenceEndpointPortChanged(String), PreferenceEndpointPortChanged(String),
@@ -313,6 +320,9 @@ impl App {
kv_cache_report: KvCacheReport::default(), kv_cache_report: KvCacheReport::default(),
last_cache_scan: Instant::now() - CACHE_SCAN_INTERVAL, last_cache_scan: Instant::now() - CACHE_SCAN_INTERVAL,
last_http_requests: 0, last_http_requests: 0,
last_ssd_selected_requests: 0,
last_ssd_requested_bytes: 0,
last_ssd_wait_ms: 0,
#[cfg(target_os = "macos")] #[cfg(target_os = "macos")]
generation_service, generation_service,
#[cfg(target_os = "macos")] #[cfg(target_os = "macos")]
@@ -416,6 +426,9 @@ impl App {
kv_cache_report: KvCacheReport::default(), kv_cache_report: KvCacheReport::default(),
last_cache_scan: Instant::now() - CACHE_SCAN_INTERVAL, last_cache_scan: Instant::now() - CACHE_SCAN_INTERVAL,
last_http_requests: 0, last_http_requests: 0,
last_ssd_selected_requests: 0,
last_ssd_requested_bytes: 0,
last_ssd_wait_ms: 0,
#[cfg(target_os = "macos")] #[cfg(target_os = "macos")]
generation_service, generation_service,
#[cfg(target_os = "macos")] #[cfg(target_os = "macos")]
@@ -535,6 +548,7 @@ impl App {
Message::PreferenceModelChanged(model) => { Message::PreferenceModelChanged(model) => {
self.preference_draft.model = model; self.preference_draft.model = model;
if !model.supports_dspark() { if !model.supports_dspark() {
self.preference_draft.legacy_mtp_enabled = false;
self.preference_draft.dspark_enabled = false; self.preference_draft.dspark_enabled = false;
self.preference_draft.dspark_confidence_threshold.clear(); self.preference_draft.dspark_confidence_threshold.clear();
self.preference_draft.dspark_strict = false; self.preference_draft.dspark_strict = false;
@@ -552,6 +566,16 @@ impl App {
} }
self.preference_error = None; self.preference_error = None;
} }
Message::PreferenceLegacyMtpChanged(enabled) => {
self.preference_draft.legacy_mtp_enabled =
self.preference_draft.model.supports_dspark() && enabled;
if self.preference_draft.legacy_mtp_enabled {
self.preference_draft.dspark_enabled = false;
self.preference_draft.dspark_confidence_threshold.clear();
self.preference_draft.dspark_strict = false;
}
self.preference_error = None;
}
Message::PreferenceDsparkChanged(enabled) => { Message::PreferenceDsparkChanged(enabled) => {
self.preference_draft.dspark_enabled = self.preference_draft.dspark_enabled =
self.preference_draft.model.supports_dspark() && enabled; self.preference_draft.model.supports_dspark() && enabled;
@@ -559,7 +583,7 @@ impl App {
self.preference_draft.dspark_confidence_threshold.clear(); self.preference_draft.dspark_confidence_threshold.clear();
self.preference_draft.dspark_strict = false; self.preference_draft.dspark_strict = false;
} else { } else {
self.preference_draft.ssd_streaming = false; self.preference_draft.legacy_mtp_enabled = false;
} }
self.preference_error = None; self.preference_error = None;
} }
@@ -665,7 +689,7 @@ impl App {
.is_empty() .is_empty()
{ {
self.preference_draft.dspark_enabled = true; self.preference_draft.dspark_enabled = true;
self.preference_draft.ssd_streaming = false; self.preference_draft.legacy_mtp_enabled = false;
} }
self.preference_error = None; self.preference_error = None;
} }
@@ -674,17 +698,12 @@ impl App {
self.preference_draft.model.supports_dspark() && value; self.preference_draft.model.supports_dspark() && value;
if self.preference_draft.dspark_strict { if self.preference_draft.dspark_strict {
self.preference_draft.dspark_enabled = true; self.preference_draft.dspark_enabled = true;
self.preference_draft.ssd_streaming = false; self.preference_draft.legacy_mtp_enabled = false;
} }
self.preference_error = None; self.preference_error = None;
} }
Message::PreferenceSsdChanged(value) => { Message::PreferenceSsdChanged(value) => {
self.preference_draft.ssd_streaming = value; self.preference_draft.ssd_streaming = value;
if value {
self.preference_draft.dspark_enabled = false;
self.preference_draft.dspark_confidence_threshold.clear();
self.preference_draft.dspark_strict = false;
}
self.preference_error = None; self.preference_error = None;
} }
Message::PreferenceSsdColdChanged(value) => { Message::PreferenceSsdColdChanged(value) => {
@@ -1226,7 +1245,23 @@ impl App {
let kv_read_bytes_per_second = kv_read_bytes as f32 / METRICS_SAMPLE_INTERVAL.as_secs_f32(); let kv_read_bytes_per_second = kv_read_bytes as f32 / METRICS_SAMPLE_INTERVAL.as_secs_f32();
let kv_write_bytes_per_second = let kv_write_bytes_per_second =
kv_write_bytes as f32 / METRICS_SAMPLE_INTERVAL.as_secs_f32(); kv_write_bytes as f32 / METRICS_SAMPLE_INTERVAL.as_secs_f32();
let sample_seconds = METRICS_SAMPLE_INTERVAL.as_secs_f32();
let ssd_requests_per_second = snapshot
.ssd_selected_requests
.saturating_sub(self.last_ssd_selected_requests)
as f32
/ sample_seconds;
let ssd_bytes_per_second = snapshot
.ssd_requested_bytes
.saturating_sub(self.last_ssd_requested_bytes)
as f32
/ sample_seconds;
let ssd_wait_ms_per_second =
snapshot.ssd_wait_ms.saturating_sub(self.last_ssd_wait_ms) as f32 / sample_seconds;
self.last_http_requests = snapshot.http_requests; self.last_http_requests = snapshot.http_requests;
self.last_ssd_selected_requests = snapshot.ssd_selected_requests;
self.last_ssd_requested_bytes = snapshot.ssd_requested_bytes;
self.last_ssd_wait_ms = snapshot.ssd_wait_ms;
self.metrics_history.push_back(MetricsPoint { self.metrics_history.push_back(MetricsPoint {
decode_tokens_per_second: if snapshot.phase == crate::metrics::RuntimePhase::Generating decode_tokens_per_second: if snapshot.phase == crate::metrics::RuntimePhase::Generating
{ {
@@ -1238,6 +1273,9 @@ impl App {
http_requests_per_second, http_requests_per_second,
kv_read_bytes_per_second, kv_read_bytes_per_second,
kv_write_bytes_per_second, kv_write_bytes_per_second,
ssd_requests_per_second,
ssd_bytes_per_second,
ssd_wait_ms_per_second,
}); });
if self.metrics_history.len() > 120 { if self.metrics_history.len() > 120 {
self.metrics_history.pop_front(); self.metrics_history.pop_front();

View File

@@ -3,6 +3,7 @@ use super::*;
#[derive(Clone)] #[derive(Clone)]
pub(super) struct PreferenceDraft { pub(super) struct PreferenceDraft {
pub(super) model: ModelChoice, pub(super) model: ModelChoice,
pub(super) legacy_mtp_enabled: bool,
pub(super) dspark_enabled: bool, pub(super) dspark_enabled: bool,
pub(super) idle_timeout_minutes: String, pub(super) idle_timeout_minutes: String,
pub(super) endpoint_port: String, pub(super) endpoint_port: String,
@@ -51,6 +52,7 @@ impl PreferenceDraft {
let speculative = &runtime.speculative; let speculative = &runtime.speculative;
Self { Self {
model: config.model, model: config.model,
legacy_mtp_enabled: speculative.legacy_mtp_enabled,
dspark_enabled: speculative.dspark_enabled, dspark_enabled: speculative.dspark_enabled,
idle_timeout_minutes: config.idle_timeout_minutes.to_string(), idle_timeout_minutes: config.idle_timeout_minutes.to_string(),
endpoint_port: config.endpoint.port.to_string(), endpoint_port: config.endpoint.port.to_string(),
@@ -136,6 +138,7 @@ impl PreferenceDraft {
Ok(SpeculativePreferences { Ok(SpeculativePreferences {
mtp_draft_tokens: parse_positive_i32("MTP draft tokens", &self.mtp_draft_tokens)?, mtp_draft_tokens: parse_positive_i32("MTP draft tokens", &self.mtp_draft_tokens)?,
mtp_margin: parse_f32("MTP margin", &self.mtp_margin)?, mtp_margin: parse_f32("MTP margin", &self.mtp_margin)?,
legacy_mtp_enabled: self.legacy_mtp_enabled,
glm_mtp: self.glm_mtp, glm_mtp: self.glm_mtp,
glm_mtp_timing: self.glm_mtp_timing, glm_mtp_timing: self.glm_mtp_timing,
dspark_enabled: self.dspark_enabled, dspark_enabled: self.dspark_enabled,

View File

@@ -3,6 +3,19 @@ use iced::widget::column;
impl App { impl App {
pub(super) fn preferences_panel(&self) -> Element<'_, Message> { pub(super) fn preferences_panel(&self) -> Element<'_, Message> {
let legacy_mtp_toggle: Option<fn(bool) -> Message> = self
.preference_draft
.model
.supports_dspark()
.then_some(Message::PreferenceLegacyMtpChanged);
let legacy_mtp = hint(
checkbox(
"Enable legacy MTP for this model",
self.preference_draft.legacy_mtp_enabled,
)
.on_toggle_maybe(legacy_mtp_toggle),
"Uses the managed one-stage MTP support GGUF. The target model verifies every drafted token; it is mutually exclusive with DSpark.",
);
let dspark_toggle: Option<fn(bool) -> Message> = self let dspark_toggle: Option<fn(bool) -> Message> = self
.preference_draft .preference_draft
.model .model
@@ -14,7 +27,7 @@ impl App {
self.preference_draft.dspark_enabled, self.preference_draft.dspark_enabled,
) )
.on_toggle_maybe(dspark_toggle), .on_toggle_maybe(dspark_toggle),
"Speculative decoding with the managed DSpark draft artifact: a small model proposes tokens that the main model verifies in one pass. Usually a large speedup, and it cannot be combined with SSD streaming.", "Speculative decoding with the managed DSpark draft artifact: a small model proposes tokens that the main model verifies in one pass. Usually a large speedup; the target model may also stream routed experts from SSD.",
); );
let glm_mtp_toggle: Option<fn(bool) -> Message> = (self.preference_draft.model let glm_mtp_toggle: Option<fn(bool) -> Message> = (self.preference_draft.model
== ModelChoice::Glm52) == ModelChoice::Glm52)
@@ -311,6 +324,7 @@ impl App {
.on_toggle_maybe(glm_mtp_timing_toggle), .on_toggle_maybe(glm_mtp_timing_toggle),
"Records per-stage timings of the speculative path to the log, to show where the acceleration actually goes. A diagnostic aid that costs a little throughput.", "Records per-stage timings of the speculative path to the log, to show where the acceleration actually goes. A diagnostic aid that costs a little throughput.",
), ),
legacy_mtp,
dspark, dspark,
preference_input_row( preference_input_row(
"DSpark confidence threshold", "DSpark confidence threshold",
@@ -326,7 +340,7 @@ impl App {
"Lets the draft model only propose, never decide: every token is sampled by the full model. Gives up some of the speedup in exchange for output identical to non-speculative decoding.", "Lets the draft model only propose, never decide: every token is sampled by the full model. Gives up some of the speedup in exchange for output identical to non-speculative decoding.",
), ),
text(if self.preference_draft.model.supports_dspark() { text(if self.preference_draft.model.supports_dspark() {
"DSpark uses the managed support artifact; entering a threshold or enabling strict mode also enables DSpark." "Legacy MTP and DSpark use separate managed support artifacts; entering a DSpark threshold or enabling strict mode selects DSpark."
} else if self.preference_draft.model == ModelChoice::Glm52 { } else if self.preference_draft.model == ModelChoice::Glm52 {
"GLM MTP is integrated; DSpark is unavailable for this model." "GLM MTP is integrated; DSpark is unavailable for this model."
} else { } else {
@@ -339,9 +353,10 @@ impl App {
|engine| { |engine| {
let settings = engine.speculative; let settings = engine.speculative;
format!( format!(
"Engine: MTP draft {} • margin {} • GLM MTP {} • timing {} • DSpark {} • confidence {}{} • target-only {}", "Engine: MTP draft {} • margin {} legacy MTP {} GLM MTP {} • timing {} • DSpark {} • confidence {}{} • target-only {}",
settings.mtp_draft_tokens, settings.mtp_draft_tokens,
settings.mtp_margin, settings.mtp_margin,
if self.preference_draft.legacy_mtp_enabled { "on" } else { "off" },
if settings.glm_mtp { "on" } else { "off" }, if settings.glm_mtp { "on" } else { "off" },
if settings.glm_mtp_timing { "on" } else { "off" }, if settings.glm_mtp_timing { "on" } else { "off" },
if settings.dspark { "on" } else { "off" }, if settings.dspark { "on" } else { "off" },
@@ -357,7 +372,7 @@ impl App {
hint( hint(
checkbox("Enable SSD-backed model streaming", self.preference_draft.ssd_streaming) checkbox("Enable SSD-backed model streaming", self.preference_draft.ssd_streaming)
.on_toggle(Message::PreferenceSsdChanged), .on_toggle(Message::PreferenceSsdChanged),
"Leaves the routed expert weights on disk and pages them in as they are needed, so a model larger than this machine's memory still runs. Every cache miss waits for the SSD, and DSpark cannot run alongside it.", "Leaves the routed expert weights on disk and pages them in as they are needed, so a model larger than this machine's memory still runs. Every cache miss waits for the SSD; speculative support weights remain resident while target experts stream.",
), ),
hint( hint(
checkbox("Skip automatic expert preload", self.preference_draft.ssd_streaming_cold) checkbox("Skip automatic expert preload", self.preference_draft.ssd_streaming_cold)
@@ -381,7 +396,7 @@ impl App {
text_input("Automatic", &self.preference_draft.ssd_preload_experts) text_input("Automatic", &self.preference_draft.ssd_preload_experts)
.on_input(Message::PreferenceSsdPreloadChanged), .on_input(Message::PreferenceSsdPreloadChanged),
), ),
text("A blank full-layer value is automatic; an explicit 0 disables fully resident GLM layers. SSD streaming and DSpark are mutually exclusive.") text("A blank full-layer value is automatic; an explicit 0 disables fully resident GLM layers. Flash legacy MTP and DSpark support weights remain resident when target experts stream.")
.size(12), .size(12),
text(engine.as_ref().map_or_else( text(engine.as_ref().map_or_else(
|| "Effective SSD settings will appear after valid values are entered." || "Effective SSD settings will appear after valid values are entered."

View File

@@ -198,6 +198,55 @@ impl App {
.spacing(7) .spacing(7)
.into(), .into(),
); );
let ssd_activity = stats_panel(
"SSD STREAMING ACTIVITY · LAST 24 SECONDS",
column![
mini_chart(
&self.metrics_history,
|point| point.ssd_requests_per_second,
Color::from_rgb8(240, 180, 70),
),
row![
text("Selected loads")
.size(12)
.color(Color::from_rgb8(240, 180, 70)),
Space::with_width(Length::Fill),
text(format!("{:.1}/s", latest.ssd_requests_per_second))
.size(12)
.color(muted_text()),
],
mini_chart(
&self.metrics_history,
|point| point.ssd_bytes_per_second,
Color::from_rgb8(67, 194, 203),
),
row![
text("Requested expert data")
.size(12)
.color(Color::from_rgb8(67, 194, 203)),
Space::with_width(Length::Fill),
text(format_rate(latest.ssd_bytes_per_second))
.size(12)
.color(muted_text()),
],
mini_chart(
&self.metrics_history,
|point| point.ssd_wait_ms_per_second,
Color::from_rgb8(220, 80, 86),
),
row![
text("Inference wait")
.size(12)
.color(Color::from_rgb8(220, 80, 86)),
Space::with_width(Length::Fill),
text(format!("{:.0} ms/s", latest.ssd_wait_ms_per_second))
.size(12)
.color(muted_text()),
],
]
.spacing(7)
.into(),
);
let model = stats_panel( let model = stats_panel(
"MODEL CORE", "MODEL CORE",
@@ -249,6 +298,86 @@ impl App {
.spacing(9) .spacing(9)
.into(), .into(),
); );
let acceptance = if stats.drafted_tokens == 0 {
0.0
} else {
stats.accepted_draft_tokens as f64 / stats.drafted_tokens as f64
};
let target_passes = stats
.speculative_cycles
.saturating_add(stats.verifier_passes);
let effective_speedup = if target_passes == 0 {
1.0
} else {
stats
.speculative_cycles
.saturating_add(stats.accepted_draft_tokens) as f64
/ target_passes as f64
};
let speculative = stats_panel(
"SPECULATIVE DECODING",
column![
metric_row(
"Mode",
match stats.speculative_mode {
1 => "Legacy MTP",
2 => "DSpark",
_ => "Off",
},
),
metric_row("Cycles", format_count(stats.speculative_cycles)),
metric_row("Drafted", format_count(stats.drafted_tokens)),
metric_row("Accepted", format_count(stats.accepted_draft_tokens)),
metric_row("Acceptance", format!("{:.1}%", acceptance * 100.0)),
metric_row(
"Target verifier passes",
format_count(stats.verifier_passes)
),
metric_row("Verifier wall time", format_milliseconds(stats.verifier_ms)),
metric_row(
"Effective target-pass speedup",
format!("{effective_speedup:.2}×")
),
]
.spacing(9)
.into(),
);
let ssd = stats_panel(
"SSD EXPERT STREAMING",
column![
metric_row("State", if stats.ssd_enabled { "Enabled" } else { "Off" }),
metric_row("Resident weights", format_bytes(stats.ssd_resident_bytes)),
metric_row("Expert cache", format_bytes(stats.ssd_cache_bytes)),
metric_row(
"Cache capacity",
format!("{} experts", stats.ssd_cache_experts)
),
metric_row(
"Preloaded",
format!("{} experts", stats.ssd_preloaded_experts)
),
metric_row(
"Selected-load requests",
format_count(stats.ssd_selected_requests)
),
metric_row(
"Requested expert bytes",
format_bytes(stats.ssd_requested_bytes)
),
metric_row("Selected-load wait", format_milliseconds(stats.ssd_wait_ms)),
metric_row(
"Average load wait",
format_milliseconds(
stats
.ssd_wait_ms
.checked_div(stats.ssd_selected_requests)
.unwrap_or(0)
)
),
]
.spacing(9)
.into(),
);
let cache = stats_panel( let cache = stats_panel(
"KV CACHE", "KV CACHE",
column![ column![
@@ -352,10 +481,12 @@ impl App {
heading, heading,
headline, headline,
throughput, throughput,
ssd_activity,
kv_io, kv_io,
requests, requests,
disc, disc,
row![model, runtime].spacing(10), row![model, runtime].spacing(10),
row![speculative, ssd].spacing(10),
row![cache, server].spacing(10), row![cache, server].spacing(10),
text("Counters are published by the runtime with relaxed atomics and sampled by the UI every 200 ms.") text("Counters are published by the runtime with relaxed atomics and sampled by the UI every 200 ms.")
.size(11) .size(11)

View File

@@ -22,7 +22,22 @@ use std::sync::atomic::{AtomicBool, Ordering};
#[cfg(target_os = "macos")] #[cfg(target_os = "macos")]
use std::time::Instant; use std::time::Instant;
use tokenizer::Tokenizer; use tokenizer::Tokenizer;
use validation::{validate_dspark, validate_main}; use validation::{SupportKind, validate_main, validate_support};
#[cfg(target_os = "macos")]
unsafe extern "C" {
fn mmap(
address: *mut std::ffi::c_void,
length: usize,
protection: i32,
flags: i32,
fd: i32,
offset: i64,
) -> *mut std::ffi::c_void;
fn mlock(address: *const std::ffi::c_void, length: usize) -> i32;
fn munlock(address: *const std::ffi::c_void, length: usize) -> i32;
fn munmap(address: *mut std::ffi::c_void, length: usize) -> i32;
}
pub(crate) use validation::validate_model_artifact; pub(crate) use validation::validate_model_artifact;
@@ -189,6 +204,7 @@ const GLM: Shape = Shape {
pub(crate) struct Model { pub(crate) struct Model {
main: Gguf, main: Gguf,
support: Option<Gguf>, support: Option<Gguf>,
support_kind: Option<SupportKind>,
shape: Shape, shape: Shape,
tokenizer: Tokenizer, tokenizer: Tokenizer,
} }
@@ -206,10 +222,17 @@ impl Model {
#[allow(dead_code)] #[allow(dead_code)]
pub(crate) fn open(settings: &EngineSettings) -> Result<Self, String> { pub(crate) fn open(settings: &EngineSettings) -> Result<Self, String> {
let mut model = Self::open_main(&settings.artifacts.model, settings.model)?; let mut model = Self::open_main(&settings.artifacts.model, settings.model)?;
if settings.execution.warm_weights {
model.main.warm()?;
}
if let Some(path) = &settings.artifacts.mtp { if let Some(path) = &settings.artifacts.mtp {
let support = Gguf::open(path)?; let support = Gguf::open(path)?;
validate_dspark(&support, &model.shape)?; let kind = validate_support(&support, &model.shape)?;
if settings.execution.warm_weights {
support.warm()?;
}
model.support = Some(support); model.support = Some(support);
model.support_kind = Some(kind);
} }
Ok(model) Ok(model)
} }
@@ -228,6 +251,7 @@ impl Model {
Ok(Self { Ok(Self {
main, main,
support: None, support: None,
support_kind: None,
shape, shape,
tokenizer, tokenizer,
}) })
@@ -333,6 +357,7 @@ impl Model {
#[cfg(target_os = "macos")] #[cfg(target_os = "macos")]
pub(crate) struct Generator { pub(crate) struct Generator {
executor: metal::Executor, executor: metal::Executor,
_simulated_memory: Option<SimulatedMemory>,
checkpoint: Option<PathBuf>, checkpoint: Option<PathBuf>,
/// Token frontier of the last transient store, so continued checkpoints are /// Token frontier of the last transient store, so continued checkpoints are
/// spaced like ds4's `continued_last_store_tokens`. /// spaced like ds4's `continued_last_store_tokens`.
@@ -340,6 +365,68 @@ pub(crate) struct Generator {
metrics: Arc<Metrics>, metrics: Arc<Metrics>,
} }
#[cfg(target_os = "macos")]
struct SimulatedMemory {
address: std::ptr::NonNull<std::ffi::c_void>,
bytes: usize,
}
#[cfg(target_os = "macos")]
impl SimulatedMemory {
fn acquire(bytes: u64) -> Result<Option<Self>, String> {
if bytes == 0 {
return Ok(None);
}
let bytes = usize::try_from(bytes).map_err(|_| "simulated memory size is too large")?;
let address = unsafe { mmap(std::ptr::null_mut(), bytes, 1 | 2, 2 | 0x1000, -1, 0) };
let Some(address) = std::ptr::NonNull::new(address) else {
return Err(format!(
"cannot reserve simulated used memory: {}",
std::io::Error::last_os_error()
));
};
if address.as_ptr() as isize == -1 {
return Err(format!(
"cannot reserve simulated used memory: {}",
std::io::Error::last_os_error()
));
}
let chunk = 256 * 1024 * 1024;
let mut locked = 0;
while locked < bytes {
let length = (bytes - locked).min(chunk);
let start = unsafe { address.as_ptr().cast::<u8>().add(locked) };
for page in (0..length).step_by(16 * 1024) {
unsafe { start.add(page).write((page / (16 * 1024)) as u8) };
}
unsafe { start.add(length - 1).write(1) };
if unsafe { mlock(start.cast(), length) } != 0 {
if locked != 0 {
unsafe { munlock(address.as_ptr(), locked) };
}
unsafe { munmap(address.as_ptr(), bytes) };
return Err(format!(
"cannot lock simulated used memory after {:.2} GiB: {}",
locked as f64 / 1_073_741_824.0,
std::io::Error::last_os_error()
));
}
locked += length;
}
Ok(Some(Self { address, bytes }))
}
}
#[cfg(target_os = "macos")]
impl Drop for SimulatedMemory {
fn drop(&mut self) {
unsafe {
munlock(self.address.as_ptr(), self.bytes);
munmap(self.address.as_ptr(), self.bytes);
}
}
}
#[derive(Clone)] #[derive(Clone)]
pub(crate) struct ChatTurn { pub(crate) struct ChatTurn {
pub(crate) user: bool, pub(crate) user: bool,
@@ -377,25 +464,41 @@ impl Generator {
"GLM MTP requires the shared speculative verifier, which is not enabled".into(), "GLM MTP requires the shared speculative verifier, which is not enabled".into(),
); );
} }
if settings.speculative.dspark let simulated_memory =
|| (settings.ssd.enabled && settings.model != ModelChoice::Glm52) SimulatedMemory::acquire(settings.diagnostics.simulated_used_memory_bytes)?;
|| settings.steering.file.is_some()
{
return Err(
"DSpark, SSD streaming, and steering are not yet available in the Rust executor"
.into(),
);
}
let model = Model::open(settings)?; let model = Model::open(settings)?;
let executor = metal::Executor::open_configured( let executor = metal::Executor::open_configured(
model, model,
settings.context_tokens.max(1) as u32, settings.context_tokens.max(1) as u32,
settings.execution.quality, settings.execution.quality,
settings.execution.prefill_chunk, settings.execution.prefill_chunk,
settings.execution.power_percent,
settings.speculative,
settings.ssd, settings.ssd,
settings.steering.clone(),
)?; )?;
let stats = executor.execution_stats();
metrics.speculative_stats(
stats.speculative_mode,
stats.speculative_cycles,
stats.drafted_tokens,
stats.accepted_draft_tokens,
stats.verifier_passes,
stats.verifier_ms,
);
metrics.ssd_stats(
stats.ssd_enabled,
stats.ssd_resident_bytes,
stats.ssd_cache_bytes,
stats.ssd_cache_experts,
stats.ssd_preloaded_experts,
stats.ssd_selected_requests,
stats.ssd_requested_bytes,
stats.ssd_wait_ms,
);
Ok(Self { Ok(Self {
executor, executor,
_simulated_memory: simulated_memory,
checkpoint: None, checkpoint: None,
last_store_tokens: 0, last_store_tokens: 0,
metrics, metrics,
@@ -422,8 +525,9 @@ impl Generator {
checkpoint, checkpoint,
conversation_tag(&settings.system_prompt, settings.reasoning_mode, history), conversation_tag(&settings.system_prompt, settings.reasoning_mode, history),
)?; )?;
let (mut output, prompt_complete) = let result = self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress);
self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress)?; self.publish_execution_stats();
let (mut output, prompt_complete) = result?;
let mut completed = messages.to_vec(); let mut completed = messages.to_vec();
completed.push(output.message.clone()); completed.push(output.message.clone());
output.previous_checkpoint_bytes = output.previous_checkpoint_bytes =
@@ -442,6 +546,28 @@ impl Generator {
Ok(output) Ok(output)
} }
fn publish_execution_stats(&self) {
let stats = self.executor.execution_stats();
self.metrics.speculative_stats(
stats.speculative_mode,
stats.speculative_cycles,
stats.drafted_tokens,
stats.accepted_draft_tokens,
stats.verifier_passes,
stats.verifier_ms,
);
self.metrics.ssd_stats(
stats.ssd_enabled,
stats.ssd_resident_bytes,
stats.ssd_cache_bytes,
stats.ssd_cache_experts,
stats.ssd_preloaded_experts,
stats.ssd_selected_requests,
stats.ssd_requested_bytes,
stats.ssd_wait_ms,
);
}
pub(crate) fn generate_transient( pub(crate) fn generate_transient(
&mut self, &mut self,
directory: &Path, directory: &Path,
@@ -481,8 +607,9 @@ impl Generator {
self.metrics.kv_lookup(KvLookup::MemoryHit); self.metrics.kv_lookup(KvLookup::MemoryHit);
} }
let (mut output, prompt_complete) = let result = self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress);
self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress)?; self.publish_execution_stats();
let (mut output, prompt_complete) = result?;
let mut completed = messages.to_vec(); let mut completed = messages.to_vec();
completed.push(output.message.clone()); completed.push(output.message.clone());
output.previous_checkpoint_bytes = previous_checkpoint output.previous_checkpoint_bytes = previous_checkpoint
@@ -809,6 +936,7 @@ impl Generator {
} }
completed completed
}; };
self.publish_execution_stats();
if completed != suffix.len() { if completed != suffix.len() {
return Ok(( return Ok((
GenerationOutput { GenerationOutput {
@@ -826,11 +954,12 @@ impl Generator {
progress(self.executor.position(), self.executor.context(), Some(0.0)); progress(self.executor.position(), self.executor.context(), Some(0.0));
let generation_started = Instant::now(); let generation_started = Instant::now();
let mut generated_tokens = 0_u32; let mut generated_tokens = 0_u32;
for _ in 0..settings let generation_limit = settings
.max_generated_tokens .max_generated_tokens
.max(0) .max(0)
.min((max_context - self.executor.position() as usize) as i32) .min((max_context - self.executor.position() as usize) as i32)
{ as u32;
while generated_tokens < generation_limit {
if cancelled.load(Ordering::Relaxed) { if cancelled.load(Ordering::Relaxed) {
append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &[], true); append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &[], true);
flush_generated( flush_generated(
@@ -887,86 +1016,132 @@ impl Generator {
true, true,
)); ));
} }
if self.executor.model().is_think_start_token(token) { let cycle = if settings.temperature <= 0.0 {
append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &[], true); self.executor.eval_speculative_greedy(
reasoning = true; token,
generated.reasoning.get_or_insert_default(); generation_limit - generated_tokens,
} else if self.executor.model().is_think_end_token(token) { settings.reasoning_mode,
append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &[], true); cancelled,
if let Some(reasoning_text) = &mut generated.reasoning )?
&& emit_safe_text( } else {
reasoning_text, self.executor.eval(token)?;
&mut emitted_reasoning, vec![token]
&settings.stops, };
true, self.publish_execution_stats();
true, for token in cycle {
emit, if generated_tokens >= generation_limit
) || self
.executor
.model()
.is_stop_token_for_reasoning(token, settings.reasoning_mode)
{ {
return Ok(( append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &[], true);
GenerationOutput { flush_generated(
message: generated, &mut generated,
prompt_tokens: prompt_tokens as u32,
cached_tokens: reused as u32,
completion_tokens: generated_tokens + 1,
finish_reason: "stop",
previous_checkpoint_bytes: None,
checkpoint_bytes: 0,
},
false,
));
}
reasoning = false;
generated.reasoning_complete = true;
emit(false, String::new());
} else if let Some(bytes) = self.executor.model().token_bytes(token) {
append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &bytes, false);
let stopped = if reasoning {
let text = generated.reasoning.get_or_insert_default();
emit_safe_text(
text,
&mut emitted_reasoning, &mut emitted_reasoning,
&settings.stops,
false,
true,
emit,
)
} else {
generated.reasoning_complete = true;
emit_safe_text(
&mut generated.content,
&mut emitted_content, &mut emitted_content,
&settings.stops, &settings.stops,
false,
false,
emit, emit,
) );
};
if stopped {
return Ok(( return Ok((
GenerationOutput { GenerationOutput {
message: generated, message: generated,
prompt_tokens: prompt_tokens as u32, prompt_tokens: prompt_tokens as u32,
cached_tokens: reused as u32, cached_tokens: reused as u32,
completion_tokens: generated_tokens + 1, completion_tokens: generated_tokens,
finish_reason: "stop", finish_reason: "stop",
previous_checkpoint_bytes: None, previous_checkpoint_bytes: None,
checkpoint_bytes: 0, checkpoint_bytes: 0,
}, },
false, true,
)); ));
} }
if self.executor.model().is_think_start_token(token) {
append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &[], true);
reasoning = true;
generated.reasoning.get_or_insert_default();
} else if self.executor.model().is_think_end_token(token) {
append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &[], true);
if let Some(reasoning_text) = &mut generated.reasoning
&& emit_safe_text(
reasoning_text,
&mut emitted_reasoning,
&settings.stops,
true,
true,
emit,
)
{
return Ok((
GenerationOutput {
message: generated,
prompt_tokens: prompt_tokens as u32,
cached_tokens: reused as u32,
completion_tokens: generated_tokens + 1,
finish_reason: "stop",
previous_checkpoint_bytes: None,
checkpoint_bytes: 0,
},
false,
));
}
reasoning = false;
generated.reasoning_complete = true;
emit(false, String::new());
} else if let Some(bytes) = self.executor.model().token_bytes(token) {
append_generated_bytes(
&mut generated,
reasoning,
&mut pending_utf8,
&bytes,
false,
);
let stopped = if reasoning {
let text = generated.reasoning.get_or_insert_default();
emit_safe_text(
text,
&mut emitted_reasoning,
&settings.stops,
false,
true,
emit,
)
} else {
generated.reasoning_complete = true;
emit_safe_text(
&mut generated.content,
&mut emitted_content,
&settings.stops,
false,
false,
emit,
)
};
if stopped {
return Ok((
GenerationOutput {
message: generated,
prompt_tokens: prompt_tokens as u32,
cached_tokens: reused as u32,
completion_tokens: generated_tokens + 1,
finish_reason: "stop",
previous_checkpoint_bytes: None,
checkpoint_bytes: 0,
},
false,
));
}
}
generated_tokens += 1;
progress(
self.executor.position(),
self.executor.context(),
Some(
generated_tokens as f32
/ generation_started.elapsed().as_secs_f32().max(1.0e-6),
),
);
} }
self.executor.eval(token)?;
generated_tokens += 1;
progress(
self.executor.position(),
self.executor.context(),
Some(
generated_tokens as f32
/ generation_started.elapsed().as_secs_f32().max(1.0e-6),
),
);
} }
append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &[], true); append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &[], true);
flush_generated( flush_generated(

View File

@@ -1,4 +1,4 @@
use memmap2::{Mmap, MmapOptions}; use memmap2::{Advice, Mmap, MmapOptions};
use sha2::{Digest, Sha256}; use sha2::{Digest, Sha256};
use std::collections::HashMap; use std::collections::HashMap;
use std::fs::File; use std::fs::File;
@@ -215,6 +215,23 @@ impl Gguf {
self.max_tensor_bytes self.max_tensor_bytes
} }
pub(super) fn warm(&self) -> Result<(), String> {
let start = self.data_offset as usize;
if start >= self.map.len() {
return Ok(());
}
self.map
.advise_range(Advice::WillNeed, start, self.map.len() - start)
.map_err(|error| format!("Cannot warm {}: {error}", self.path.display()))?;
let mut checksum = 0_u64;
for offset in (start..self.map.len()).step_by(16 * 1024) {
checksum = checksum.wrapping_add(u64::from(self.map[offset]));
}
checksum = checksum.wrapping_add(u64::from(self.map[self.map.len() - 1]));
std::hint::black_box(checksum);
Ok(())
}
pub(super) fn tensor(&self, name: &str) -> Result<&Tensor, String> { pub(super) fn tensor(&self, name: &str) -> Result<&Tensor, String> {
self.tensors self.tensors
.get(name) .get(name)
@@ -544,6 +561,7 @@ mod tests {
assert_eq!(model.bytes("general.architecture").unwrap(), b"deepseek4"); assert_eq!(model.bytes("general.architecture").unwrap(), b"deepseek4");
assert_eq!(model.tensor("weight").unwrap().dims, [1]); assert_eq!(model.tensor("weight").unwrap().dims, [1]);
assert_eq!(model.tensor_data("weight").unwrap(), 1_f32.to_le_bytes()); assert_eq!(model.tensor_data("weight").unwrap(), 1_f32.to_le_bytes());
model.warm().unwrap();
fs::remove_file(path).unwrap(); fs::remove_file(path).unwrap();
} }

File diff suppressed because it is too large Load Diff

View File

@@ -301,6 +301,15 @@ impl DeepSeekExecutor {
self.tokens = tokens; self.tokens = tokens;
self.logits = logits; self.logits = logits;
self.checkpoint_tag = checkpoint_tag; self.checkpoint_tag = checkpoint_tag;
if let Some(mtp) = &mut self.legacy_mtp {
mtp.draft_token = None;
mtp.raw_rows = 0;
}
if let Some(dspark) = &mut self.dspark {
dspark.capture_mask = 0;
dspark.cache_start = 0;
dspark.cache_len = 0;
}
Ok(()) Ok(())
} }
} }

View File

@@ -231,7 +231,7 @@ impl GlmExecutor {
} }
let weights = GlmWeights::bind(&model)?; let weights = GlmWeights::bind(&model)?;
let admission = admission_bytes(&model, &weights, context, ssd)?; let admission = admission_bytes(&model, &weights, context, ssd)?;
let context_handle = Context::open(&model, quality, ssd.enabled, admission)?; let context_handle = Context::open(&model, quality, ssd.enabled, admission, None)?;
configure_streaming(&model, &weights, ssd)?; configure_streaming(&model, &weights, ssd)?;
let scratch = GlmScratch::allocate(&model, context)?; let scratch = GlmScratch::allocate(&model, context)?;
let caches = (0..weights.layers.len()) let caches = (0..weights.layers.len())

View File

@@ -28,6 +28,14 @@ unsafe extern "C" {
map_size: u64, map_size: u64,
max_tensor_bytes: u64, max_tensor_bytes: u64,
) -> i32; ) -> i32;
pub(super) fn ds4_gpu_set_model_map_spans(
model_map: *const c_void,
model_size: u64,
offsets: *const u64,
sizes: *const u64,
count: u32,
max_tensor_bytes: u64,
) -> i32;
pub(super) fn ds4_gpu_set_quality(quality: bool); pub(super) fn ds4_gpu_set_quality(quality: bool);
pub(super) fn ds4_gpu_set_glm_model(enabled: bool); pub(super) fn ds4_gpu_set_glm_model(enabled: bool);
pub(super) fn ds4_gpu_set_ssd_streaming(enabled: bool); pub(super) fn ds4_gpu_set_ssd_streaming(enabled: bool);
@@ -39,6 +47,19 @@ unsafe extern "C" {
gate_expert_bytes: u64, gate_expert_bytes: u64,
down_expert_bytes: u64, down_expert_bytes: u64,
) -> u32; ) -> u32;
pub(super) fn ds4_gpu_stream_expert_cache_seed_experts(
table: *const StreamExpertTable,
expert_ids: *const i32,
expert_priorities: *const u32,
experts: u32,
) -> i32;
pub(super) fn ds4_gpu_stream_expert_cache_begin_selected_load(
table: *const StreamExpertTable,
selected_ids: *const i32,
selected: u32,
) -> 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_glm_stream_expert_cache_begin_selected_load_tensor( pub(super) fn ds4_gpu_glm_stream_expert_cache_begin_selected_load_tensor(
table: *const StreamExpertTable, table: *const StreamExpertTable,
selected: *const GpuTensor, selected: *const GpuTensor,
@@ -79,8 +100,23 @@ unsafe extern "C" {
src_offset: u64, src_offset: u64,
count: u64, count: u64,
) -> i32; ) -> i32;
pub(super) fn ds4_gpu_pack_slot_rows_f32_tensor(
out: *mut GpuTensor,
slots: *const GpuTensor,
rows: u32,
width: u32,
slot_count: u32,
slot_stride: u32,
) -> i32;
pub(super) fn ds4_gpu_begin_commands() -> i32; pub(super) fn ds4_gpu_begin_commands() -> i32;
pub(super) fn ds4_gpu_end_commands() -> i32; pub(super) fn ds4_gpu_end_commands() -> i32;
pub(super) fn ds4_gpu_signal_selected_readback_ready(event: *mut u64) -> i32;
pub(super) fn ds4_gpu_wait_selected_readback_ready(
event: u64,
label: *const std::ffi::c_char,
) -> i32;
pub(super) fn ds4_gpu_routed_moe_set_selected_override(selected: *const i32, count: u32)
-> i32;
pub(super) fn ds4_gpu_embed_tokens_hc_tensor( pub(super) fn ds4_gpu_embed_tokens_hc_tensor(
out: *mut GpuTensor, out: *mut GpuTensor,
@@ -110,6 +146,13 @@ unsafe extern "C" {
n: u32, n: u32,
eps: f32, eps: f32,
) -> i32; ) -> i32;
pub(super) fn ds4_gpu_rms_norm_plain_rows_tensor(
out: *mut GpuTensor,
x: *const GpuTensor,
n: u32,
rows: u32,
eps: f32,
) -> i32;
pub(super) fn ds4_gpu_rms_norm_weight_tensor( pub(super) fn ds4_gpu_rms_norm_weight_tensor(
out: *mut GpuTensor, out: *mut GpuTensor,
x: *const GpuTensor, x: *const GpuTensor,
@@ -119,6 +162,36 @@ unsafe extern "C" {
n: u32, n: u32,
eps: f32, eps: f32,
) -> i32; ) -> i32;
pub(super) fn ds4_gpu_rms_norm_weight_rows_tensor(
out: *mut GpuTensor,
x: *const GpuTensor,
map: *const c_void,
size: u64,
weight: u64,
n: u32,
rows: u32,
eps: f32,
) -> i32;
pub(super) fn ds4_gpu_repeat_hc_tensor(
out: *mut GpuTensor,
x: *const GpuTensor,
embd: u32,
hc: u32,
) -> i32;
pub(super) fn ds4_gpu_attention_noncausal_raw_batch_heads_tensor(
out: *mut GpuTensor,
map: *const c_void,
size: u64,
sinks: u64,
q: *const GpuTensor,
raw_cache: *const GpuTensor,
rows: u32,
visible_rows: u32,
cache_cap: u32,
raw_start: u32,
heads: u32,
head_dim: u32,
) -> i32;
pub(super) fn ds4_gpu_hc_rms_scale_project_f16_tensor( pub(super) fn ds4_gpu_hc_rms_scale_project_f16_tensor(
out: *mut GpuTensor, out: *mut GpuTensor,
scale: *mut GpuTensor, scale: *mut GpuTensor,
@@ -865,6 +938,14 @@ unsafe extern "C" {
c: *const GpuTensor, c: *const GpuTensor,
count: u32, count: u32,
) -> i32; ) -> i32;
pub(super) fn ds4_gpu_directional_steering_project_tensor(
x: *mut GpuTensor,
directions: *const GpuTensor,
layer: u32,
width: u32,
rows: u32,
scale: f32,
) -> i32;
pub(super) fn ds4_gpu_add_rms_norm_weight_tensor( pub(super) fn ds4_gpu_add_rms_norm_weight_tensor(
norm: *mut GpuTensor, norm: *mut GpuTensor,
sum: *mut GpuTensor, sum: *mut GpuTensor,
@@ -1043,6 +1124,7 @@ impl Context {
quality: bool, quality: bool,
ssd_streaming: bool, ssd_streaming: bool,
admission_bytes: u64, admission_bytes: u64,
model_spans: Option<&[(u64, u64)]>,
) -> Result<Self, String> { ) -> Result<Self, String> {
check(unsafe { ds4_gpu_init() }, "Metal initialization")?; check(unsafe { ds4_gpu_init() }, "Metal initialization")?;
unsafe { unsafe {
@@ -1059,7 +1141,19 @@ impl Context {
)); ));
} }
let data_offset = model.main.data_offset(); let data_offset = model.main.data_offset();
if let Err(error) = check( let mapped = if let Some(spans) = model_spans {
let (offsets, sizes): (Vec<_>, Vec<_>) = spans.iter().copied().unzip();
unsafe {
ds4_gpu_set_model_map_spans(
model.main.map_ptr().cast(),
model.main.len(),
offsets.as_ptr(),
sizes.as_ptr(),
spans.len() as u32,
model.main.max_tensor_bytes(),
)
}
} else {
unsafe { unsafe {
ds4_gpu_set_model_map_range( ds4_gpu_set_model_map_range(
model.main.map_ptr().cast(), model.main.map_ptr().cast(),
@@ -1068,12 +1162,27 @@ impl Context {
model.main.len() - data_offset, model.main.len() - data_offset,
model.main.max_tensor_bytes(), model.main.max_tensor_bytes(),
) )
}, }
"model mapping", };
) { if let Err(error) = check(mapped, "model mapping") {
unsafe { ds4_gpu_cleanup() }; unsafe { ds4_gpu_cleanup() };
return Err(error); return Err(error);
} }
if let Some(support) = &model.support {
let mapped = 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(),
)
};
if let Err(error) = check(mapped, "support-model mapping") {
unsafe { ds4_gpu_cleanup() };
return Err(error);
}
}
unsafe { ds4_gpu_set_quality(quality) }; unsafe { ds4_gpu_set_quality(quality) };
let model_file = File::open(model.main.path()).map_err(|error| { let model_file = File::open(model.main.path()).map_err(|error| {
unsafe { ds4_gpu_cleanup() }; unsafe { ds4_gpu_cleanup() };
@@ -1209,6 +1318,20 @@ impl Buffer {
) )
} }
pub(super) fn write_f32(&self, values: &[f32]) -> Result<(), String> {
check(
unsafe {
ds4_gpu_tensor_write(
self.raw(),
0,
values.as_ptr().cast(),
std::mem::size_of_val(values) as u64,
)
},
"uploading floats",
)
}
pub(super) fn fill(&self, value: f32, count: u64) -> Result<(), String> { pub(super) fn fill(&self, value: f32, count: u64) -> Result<(), String> {
call( call(
unsafe { ds4_gpu_tensor_fill_f32(self.raw(), value, count) }, unsafe { ds4_gpu_tensor_fill_f32(self.raw(), value, count) },

6439
src/engine/metal/hotlist.rs Normal file

File diff suppressed because it is too large Load Diff

View File

@@ -7,7 +7,7 @@ pub(crate) fn validate_model_artifact(
) -> Result<(), String> { ) -> Result<(), String> {
if support { if support {
let model = Gguf::open(path)?; let model = Gguf::open(path)?;
validate_dspark(&model, &FLASH) validate_support(&model, &FLASH).map(|_| ())
} else { } else {
let model = Model::open_main(path, expected)?; let model = Model::open_main(path, expected)?;
let summary = model.summary(); let summary = model.summary();
@@ -29,6 +29,192 @@ pub(crate) fn validate_model_artifact(
} }
} }
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(super) enum SupportKind {
LegacyMtp,
DSpark,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub(super) struct DsparkConfig {
pub(super) block_size: u32,
pub(super) markov_rank: u32,
pub(super) noise_token: u32,
pub(super) target_layers: Vec<u32>,
pub(super) stages: u32,
}
pub(super) fn dspark_config(model: &Gguf) -> Result<DsparkConfig, String> {
let block_size = first_u32(
model,
&[
"deepseek4.dspark.block_size",
"deepseek4.dspark_block_size",
"dspark.block_size",
],
)?;
let markov_rank = first_u32(
model,
&[
"deepseek4.dspark.markov_rank",
"deepseek4.dspark_markov_rank",
"dspark.markov_rank",
],
)?;
let noise_token = first_u32(
model,
&[
"deepseek4.dspark.noise_token_id",
"deepseek4.dspark_noise_token_id",
"dspark.noise_token_id",
],
)?;
let target_layers = first_u32s(
model,
&[
"deepseek4.dspark.target_layer_ids",
"deepseek4.dspark_target_layer_ids",
"dspark.target_layer_ids",
],
)?;
let stages = model
.tensors
.keys()
.filter_map(|name| {
name.strip_prefix("mtp.")?
.split('.')
.next()?
.parse::<u32>()
.ok()
})
.max()
.map_or(0, |stage| stage + 1);
Ok(DsparkConfig {
block_size,
markov_rank,
noise_token,
target_layers: target_layers.to_vec(),
stages,
})
}
pub(super) fn validate_support(model: &Gguf, shape: &Shape) -> Result<SupportKind, String> {
if model.tensors.contains_key("mtp.0.e_proj.weight")
&& model.tensors.contains_key("mtp.0.h_proj.weight")
&& model.tensors.contains_key("mtp.0.hc_head_base.weight")
{
validate_legacy_mtp(model, shape)?;
Ok(SupportKind::LegacyMtp)
} else if model.metadata.contains_key("deepseek4.dspark.block_size")
|| model.metadata.contains_key("deepseek4.dspark_block_size")
|| model.metadata.contains_key("dspark.block_size")
{
validate_dspark(model, shape)?;
Ok(SupportKind::DSpark)
} else {
Err("support GGUF is neither legacy MTP nor DSpark".into())
}
}
fn validate_legacy_mtp(model: &Gguf, shape: &Shape) -> Result<(), String> {
if shape.model != ModelChoice::DeepSeekV4Flash {
return Err("legacy MTP support is available only for DeepSeek V4 Flash".into());
}
let prefix = "mtp.0";
let hc_dim = shape.embd * shape.hc;
let hc_mix = 2 * shape.hc + shape.hc * shape.hc;
let q_dim = shape.heads * shape.head_dim;
let output_low = shape.out_groups * shape.lora_o;
for (suffix, types, dims) in [
("hc_head_base.weight", &[F32][..], vec![shape.hc]),
("hc_head_fn.weight", PLAIN, vec![hc_dim, shape.hc]),
("hc_head_scale.weight", &[F32][..], vec![1]),
("e_proj.weight", &[Q8_0][..], vec![shape.embd, shape.embd]),
("h_proj.weight", &[Q8_0][..], vec![shape.embd, shape.embd]),
("enorm.weight", &[F32][..], vec![shape.embd]),
("hnorm.weight", &[F32][..], vec![shape.embd]),
("norm.weight", &[F32][..], vec![shape.embd]),
("hc_attn_fn.weight", PLAIN, vec![hc_dim, hc_mix]),
("hc_attn_scale.weight", &[F32][..], vec![3]),
("hc_attn_base.weight", &[F32][..], vec![hc_mix]),
("attn_norm.weight", &[F32][..], vec![shape.embd]),
(
"attn_q_a.weight",
&[Q8_0][..],
vec![shape.embd, shape.lora_q],
),
("attn_q_a_norm.weight", &[F32][..], vec![shape.lora_q]),
("attn_q_b.weight", &[Q8_0][..], vec![shape.lora_q, q_dim]),
(
"attn_kv.weight",
&[Q8_0][..],
vec![shape.embd, shape.head_dim],
),
("attn_kv_a_norm.weight", &[F32][..], vec![shape.head_dim]),
("attn_sinks.weight", &[F32][..], vec![shape.heads]),
(
"attn_output_a.weight",
&[Q8_0][..],
vec![
shape.head_dim * (shape.heads / shape.out_groups),
output_low,
],
),
(
"attn_output_b.weight",
&[Q8_0][..],
vec![output_low, shape.embd],
),
("hc_ffn_fn.weight", PLAIN, vec![hc_dim, hc_mix]),
("hc_ffn_scale.weight", &[F32][..], vec![3]),
("hc_ffn_base.weight", &[F32][..], vec![hc_mix]),
("ffn_norm.weight", &[F32][..], vec![shape.embd]),
(
"ffn_gate_inp.weight",
PLAIN,
vec![shape.embd, shape.experts],
),
("exp_probs_b.bias", &[F32][..], vec![shape.experts]),
(
"ffn_gate_exps.weight",
ROUTED,
vec![shape.embd, shape.ff_expert, shape.experts],
),
(
"ffn_up_exps.weight",
ROUTED,
vec![shape.embd, shape.ff_expert, shape.experts],
),
(
"ffn_down_exps.weight",
ROUTED,
vec![shape.ff_expert, shape.embd, shape.experts],
),
(
"ffn_gate_shexp.weight",
&[Q8_0][..],
vec![shape.embd, shape.ff_expert],
),
(
"ffn_up_shexp.weight",
&[Q8_0][..],
vec![shape.embd, shape.ff_expert],
),
(
"ffn_down_shexp.weight",
&[Q8_0][..],
vec![shape.ff_expert, shape.embd],
),
] {
expect(model, &format!("{prefix}.{suffix}"), types, &dims)?;
}
same_type(
model,
"mtp.0.ffn_gate_exps.weight",
"mtp.0.ffn_up_exps.weight",
)
}
pub(super) fn validate_main(model: &Gguf, expected: ModelChoice) -> Result<Shape, String> { pub(super) fn validate_main(model: &Gguf, expected: ModelChoice) -> Result<Shape, String> {
let family = if model.bytes("general.architecture").ok() == Some(b"glm-dsa") { let family = if model.bytes("general.architecture").ok() == Some(b"glm-dsa") {
ModelFamily::Glm ModelFamily::Glm
@@ -575,38 +761,13 @@ pub(super) fn validate_dspark(model: &Gguf, shape: &Shape) -> Result<(), String>
if shape.model != ModelChoice::DeepSeekV4Flash { if shape.model != ModelChoice::DeepSeekV4Flash {
return Err("DSpark support is available only for DeepSeek V4 Flash".into()); return Err("DSpark support is available only for DeepSeek V4 Flash".into());
} }
let block_size = first_u32( let DsparkConfig {
model, block_size,
&[ markov_rank,
"deepseek4.dspark.block_size", noise_token,
"deepseek4.dspark_block_size", target_layers: targets,
"dspark.block_size", stages,
], } = dspark_config(model)?;
)?;
let markov_rank = first_u32(
model,
&[
"deepseek4.dspark.markov_rank",
"deepseek4.dspark_markov_rank",
"dspark.markov_rank",
],
)?;
let noise_token = first_u32(
model,
&[
"deepseek4.dspark.noise_token_id",
"deepseek4.dspark_noise_token_id",
"dspark.noise_token_id",
],
)?;
let targets = first_u32s(
model,
&[
"deepseek4.dspark.target_layer_ids",
"deepseek4.dspark_target_layer_ids",
"dspark.target_layer_ids",
],
)?;
if !(1..=16).contains(&block_size) || markov_rank == 0 || noise_token >= shape.vocab as u32 { if !(1..=16).contains(&block_size) || markov_rank == 0 || noise_token >= shape.vocab as u32 {
return Err("invalid DSpark block, Markov, or noise-token metadata".into()); return Err("invalid DSpark block, Markov, or noise-token metadata".into());
} }
@@ -617,18 +778,6 @@ pub(super) fn validate_dspark(model: &Gguf, shape: &Shape) -> Result<(), String>
{ {
return Err("invalid DSpark target-layer metadata".into()); return Err("invalid DSpark target-layer metadata".into());
} }
let stages = model
.tensors
.keys()
.filter_map(|name| {
name.strip_prefix("mtp.")?
.split('.')
.next()?
.parse::<u32>()
.ok()
})
.max()
.map_or(0, |stage| stage + 1);
if !(1..=8).contains(&stages) { if !(1..=8).contains(&stages) {
return Err(format!("invalid DSpark stage count: {stages}")); return Err(format!("invalid DSpark stage count: {stages}"));
} }
@@ -1079,4 +1228,15 @@ mod tests {
validate_model_artifact(path, ModelChoice::DeepSeekV4Flash, true).unwrap(); validate_model_artifact(path, ModelChoice::DeepSeekV4Flash, true).unwrap();
} }
} }
#[test]
fn installed_legacy_mtp_fixture_passes_the_target_layout() {
let path = Path::new("../ds4/gguf/DeepSeek-V4-Flash-MTP-Q4K-Q8_0-F32.gguf");
if path.exists() {
assert_eq!(
validate_support(&Gguf::open(path).unwrap(), &FLASH).unwrap(),
SupportKind::LegacyMtp
);
}
}
} }

View File

@@ -86,6 +86,20 @@ pub(crate) struct MetricsSnapshot {
pub(crate) prompt_tokens: u64, pub(crate) prompt_tokens: u64,
pub(crate) cached_tokens: u64, pub(crate) cached_tokens: u64,
pub(crate) completion_tokens: u64, pub(crate) completion_tokens: u64,
pub(crate) speculative_mode: u8,
pub(crate) speculative_cycles: u64,
pub(crate) drafted_tokens: u64,
pub(crate) accepted_draft_tokens: u64,
pub(crate) verifier_passes: u64,
pub(crate) verifier_ms: u64,
pub(crate) ssd_enabled: bool,
pub(crate) ssd_resident_bytes: u64,
pub(crate) ssd_cache_bytes: u64,
pub(crate) ssd_cache_experts: u64,
pub(crate) ssd_preloaded_experts: u64,
pub(crate) ssd_selected_requests: u64,
pub(crate) ssd_requested_bytes: u64,
pub(crate) ssd_wait_ms: u64,
pub(crate) kv_lookups: u64, pub(crate) kv_lookups: u64,
pub(crate) kv_hits: u64, pub(crate) kv_hits: u64,
pub(crate) kv_memory_hits: u64, pub(crate) kv_memory_hits: u64,
@@ -154,6 +168,20 @@ pub(crate) struct Metrics {
prompt_tokens: AtomicU64, prompt_tokens: AtomicU64,
cached_tokens: AtomicU64, cached_tokens: AtomicU64,
completion_tokens: AtomicU64, completion_tokens: AtomicU64,
speculative_mode: AtomicU8,
speculative_cycles: AtomicU64,
drafted_tokens: AtomicU64,
accepted_draft_tokens: AtomicU64,
verifier_passes: AtomicU64,
verifier_ms: AtomicU64,
ssd_enabled: AtomicBool,
ssd_resident_bytes: AtomicU64,
ssd_cache_bytes: AtomicU64,
ssd_cache_experts: AtomicU64,
ssd_preloaded_experts: AtomicU64,
ssd_selected_requests: AtomicU64,
ssd_requested_bytes: AtomicU64,
ssd_wait_ms: AtomicU64,
kv_lookups: AtomicU64, kv_lookups: AtomicU64,
kv_hits: AtomicU64, kv_hits: AtomicU64,
kv_memory_hits: AtomicU64, kv_memory_hits: AtomicU64,
@@ -227,6 +255,20 @@ impl Metrics {
prompt_tokens: AtomicU64::new(0), prompt_tokens: AtomicU64::new(0),
cached_tokens: AtomicU64::new(0), cached_tokens: AtomicU64::new(0),
completion_tokens: AtomicU64::new(0), completion_tokens: AtomicU64::new(0),
speculative_mode: AtomicU8::new(0),
speculative_cycles: AtomicU64::new(0),
drafted_tokens: AtomicU64::new(0),
accepted_draft_tokens: AtomicU64::new(0),
verifier_passes: AtomicU64::new(0),
verifier_ms: AtomicU64::new(0),
ssd_enabled: AtomicBool::new(false),
ssd_resident_bytes: AtomicU64::new(0),
ssd_cache_bytes: AtomicU64::new(0),
ssd_cache_experts: AtomicU64::new(0),
ssd_preloaded_experts: AtomicU64::new(0),
ssd_selected_requests: AtomicU64::new(0),
ssd_requested_bytes: AtomicU64::new(0),
ssd_wait_ms: AtomicU64::new(0),
kv_lookups: AtomicU64::new(0), kv_lookups: AtomicU64::new(0),
kv_hits: AtomicU64::new(0), kv_hits: AtomicU64::new(0),
kv_memory_hits: AtomicU64::new(0), kv_memory_hits: AtomicU64::new(0),
@@ -374,6 +416,52 @@ impl Metrics {
self.source.store(WorkSource::None as u8, Ordering::Relaxed); self.source.store(WorkSource::None as u8, Ordering::Relaxed);
} }
pub(crate) fn speculative_stats(
&self,
mode: u8,
cycles: u64,
drafted: u64,
accepted: u64,
verifier_passes: u64,
verifier_ms: u64,
) {
self.speculative_mode.store(mode, Ordering::Relaxed);
self.speculative_cycles.store(cycles, Ordering::Relaxed);
self.drafted_tokens.store(drafted, Ordering::Relaxed);
self.accepted_draft_tokens
.store(accepted, Ordering::Relaxed);
self.verifier_passes
.store(verifier_passes, Ordering::Relaxed);
self.verifier_ms.store(verifier_ms, Ordering::Relaxed);
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn ssd_stats(
&self,
enabled: bool,
resident_bytes: u64,
cache_bytes: u64,
cache_experts: u64,
preloaded_experts: u64,
selected_requests: u64,
requested_bytes: u64,
wait_ms: u64,
) {
self.ssd_enabled.store(enabled, Ordering::Relaxed);
self.ssd_resident_bytes
.store(resident_bytes, Ordering::Relaxed);
self.ssd_cache_bytes.store(cache_bytes, Ordering::Relaxed);
self.ssd_cache_experts
.store(cache_experts, Ordering::Relaxed);
self.ssd_preloaded_experts
.store(preloaded_experts, Ordering::Relaxed);
self.ssd_selected_requests
.store(selected_requests, Ordering::Relaxed);
self.ssd_requested_bytes
.store(requested_bytes, Ordering::Relaxed);
self.ssd_wait_ms.store(wait_ms, Ordering::Relaxed);
}
pub(crate) fn unloaded(&self) { pub(crate) fn unloaded(&self) {
self.phase self.phase
.store(RuntimePhase::Unloaded as u8, Ordering::Relaxed); .store(RuntimePhase::Unloaded as u8, Ordering::Relaxed);
@@ -385,6 +473,8 @@ impl Metrics {
self.decode_tps.store(0, Ordering::Relaxed); self.decode_tps.store(0, Ordering::Relaxed);
self.prefill_tps.store(0, Ordering::Relaxed); self.prefill_tps.store(0, Ordering::Relaxed);
self.prefill_sample.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.model_unloads.fetch_add(1, Ordering::Relaxed); self.model_unloads.fetch_add(1, Ordering::Relaxed);
} }
@@ -534,6 +624,20 @@ impl Metrics {
prompt_tokens: self.prompt_tokens.load(Ordering::Relaxed), prompt_tokens: self.prompt_tokens.load(Ordering::Relaxed),
cached_tokens: self.cached_tokens.load(Ordering::Relaxed), cached_tokens: self.cached_tokens.load(Ordering::Relaxed),
completion_tokens: self.completion_tokens.load(Ordering::Relaxed), completion_tokens: self.completion_tokens.load(Ordering::Relaxed),
speculative_mode: self.speculative_mode.load(Ordering::Relaxed),
speculative_cycles: self.speculative_cycles.load(Ordering::Relaxed),
drafted_tokens: self.drafted_tokens.load(Ordering::Relaxed),
accepted_draft_tokens: self.accepted_draft_tokens.load(Ordering::Relaxed),
verifier_passes: self.verifier_passes.load(Ordering::Relaxed),
verifier_ms: self.verifier_ms.load(Ordering::Relaxed),
ssd_enabled: self.ssd_enabled.load(Ordering::Relaxed),
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_preloaded_experts: self.ssd_preloaded_experts.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),
kv_lookups: self.kv_lookups.load(Ordering::Relaxed), kv_lookups: self.kv_lookups.load(Ordering::Relaxed),
kv_hits: self.kv_hits.load(Ordering::Relaxed), kv_hits: self.kv_hits.load(Ordering::Relaxed),
kv_memory_hits: self.kv_memory_hits.load(Ordering::Relaxed), kv_memory_hits: self.kv_memory_hits.load(Ordering::Relaxed),
@@ -866,6 +970,7 @@ mod tests {
metrics.kv_write_finished(Duration::from_millis(20), false); 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(), (2_048, 4_096));
assert_eq!(metrics.take_kv_io_sample(), (0, 0)); assert_eq!(metrics.take_kv_io_sample(), (0, 0));
metrics.ssd_stats(true, 1_024, 2_048, 16, 4, 12, 8_192, 30);
metrics.request_finished( metrics.request_finished(
WorkSource::LocalChat, WorkSource::LocalChat,
Duration::from_millis(250), Duration::from_millis(250),
@@ -890,6 +995,10 @@ mod tests {
assert_eq!(snapshot.kv_prefix_hits, 1); assert_eq!(snapshot.kv_prefix_hits, 1);
assert_eq!(snapshot.kv_read_bytes, 2_048); assert_eq!(snapshot.kv_read_bytes, 2_048);
assert_eq!(snapshot.kv_write_bytes, 4_096); assert_eq!(snapshot.kv_write_bytes, 4_096);
assert!(snapshot.ssd_enabled);
assert_eq!(snapshot.ssd_selected_requests, 12);
assert_eq!(snapshot.ssd_requested_bytes, 8_192);
assert_eq!(snapshot.ssd_wait_ms, 30);
assert_eq!(snapshot.local_kv_bytes, 4_096); assert_eq!(snapshot.local_kv_bytes, 4_096);
assert_eq!(snapshot.http_model_requests, 1); assert_eq!(snapshot.http_model_requests, 1);
} }

View File

@@ -14,8 +14,9 @@ pub(crate) const MODEL_CHOICES: [ModelChoice; 3] = [
ModelChoice::DeepSeekV4Pro, ModelChoice::DeepSeekV4Pro,
ModelChoice::Glm52, ModelChoice::Glm52,
]; ];
pub(crate) const MANAGED_ARTIFACTS: [ManagedArtifactId; 4] = [ pub(crate) const MANAGED_ARTIFACTS: [ManagedArtifactId; 5] = [
ManagedArtifactId::DeepSeekV4Flash, ManagedArtifactId::DeepSeekV4Flash,
ManagedArtifactId::DeepSeekV4FlashMtp,
ManagedArtifactId::DeepSeekV4FlashDspark, ManagedArtifactId::DeepSeekV4FlashDspark,
ManagedArtifactId::DeepSeekV4Pro, ManagedArtifactId::DeepSeekV4Pro,
ManagedArtifactId::Glm52, ManagedArtifactId::Glm52,
@@ -40,6 +41,14 @@ const FLASH_DSPARK: Artifact = Artifact {
sha256: "8b3adf5942bec22ae2ea867cd7079cf13530ba83ffcffaf00f5de48664a1a34e", sha256: "8b3adf5942bec22ae2ea867cd7079cf13530ba83ffcffaf00f5de48664a1a34e",
support: Some(true), support: Some(true),
}; };
const FLASH_MTP: Artifact = Artifact {
label: "Legacy MTP support",
file_name: "DeepSeek-V4-Flash-MTP-Q4K-Q8_0-F32.gguf",
repository: DEEPSEEK_REPOSITORY,
size: 3_807_602_400,
sha256: "afd481ee689dce9037f70f39085fcdae5a5b096d521cdad43b19fa52bf8f4083",
support: Some(true),
};
const PRO: Artifact = Artifact { const PRO: Artifact = Artifact {
label: "DeepSeek V4 Pro model", label: "DeepSeek V4 Pro model",
file_name: "DeepSeek-V4-Pro-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-Instruct-imatrix.gguf", file_name: "DeepSeek-V4-Pro-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-Instruct-imatrix.gguf",
@@ -94,9 +103,14 @@ impl ModelChoice {
} }
#[cfg(test)] #[cfg(test)]
fn artifacts(self, dspark_enabled: bool) -> impl Iterator<Item = &'static Artifact> { fn artifacts(
self,
legacy_mtp_enabled: bool,
dspark_enabled: bool,
) -> impl Iterator<Item = &'static Artifact> {
[ [
Some(self.main_artifact()), Some(self.main_artifact()),
(self.supports_dspark() && legacy_mtp_enabled).then_some(&FLASH_MTP),
(self.supports_dspark() && dspark_enabled).then_some(&FLASH_DSPARK), (self.supports_dspark() && dspark_enabled).then_some(&FLASH_DSPARK),
] ]
.into_iter() .into_iter()
@@ -112,19 +126,26 @@ pub(crate) struct EngineArtifacts {
pub(crate) fn engine_artifacts( pub(crate) fn engine_artifacts(
model: ModelChoice, model: ModelChoice,
legacy_mtp_enabled: bool,
dspark_enabled: bool, dspark_enabled: bool,
models_path: &Path, models_path: &Path,
) -> EngineArtifacts { ) -> EngineArtifacts {
EngineArtifacts { EngineArtifacts {
model: model.main_artifact().path(model, models_path), model: model.main_artifact().path(model, models_path),
mtp: (model.supports_dspark() && dspark_enabled) mtp: if model.supports_dspark() && legacy_mtp_enabled {
.then(|| FLASH_DSPARK.path(model, models_path)), Some(FLASH_MTP.path(model, models_path))
} else if model.supports_dspark() && dspark_enabled {
Some(FLASH_DSPARK.path(model, models_path))
} else {
None
},
} }
} }
#[derive(Clone, Copy, Debug, Eq, PartialEq)] #[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum ManagedArtifactId { pub(crate) enum ManagedArtifactId {
DeepSeekV4Flash, DeepSeekV4Flash,
DeepSeekV4FlashMtp,
DeepSeekV4FlashDspark, DeepSeekV4FlashDspark,
DeepSeekV4Pro, DeepSeekV4Pro,
Glm52, Glm52,
@@ -133,7 +154,9 @@ pub(crate) enum ManagedArtifactId {
impl ManagedArtifactId { impl ManagedArtifactId {
pub(crate) fn model(self) -> ModelChoice { pub(crate) fn model(self) -> ModelChoice {
match self { match self {
Self::DeepSeekV4Flash | Self::DeepSeekV4FlashDspark => ModelChoice::DeepSeekV4Flash, Self::DeepSeekV4Flash | Self::DeepSeekV4FlashMtp | Self::DeepSeekV4FlashDspark => {
ModelChoice::DeepSeekV4Flash
}
Self::DeepSeekV4Pro => ModelChoice::DeepSeekV4Pro, Self::DeepSeekV4Pro => ModelChoice::DeepSeekV4Pro,
Self::Glm52 => ModelChoice::Glm52, Self::Glm52 => ModelChoice::Glm52,
} }
@@ -142,6 +165,7 @@ impl ManagedArtifactId {
fn artifact(self) -> &'static Artifact { fn artifact(self) -> &'static Artifact {
match self { match self {
Self::DeepSeekV4Flash => &FLASH, Self::DeepSeekV4Flash => &FLASH,
Self::DeepSeekV4FlashMtp => &FLASH_MTP,
Self::DeepSeekV4FlashDspark => &FLASH_DSPARK, Self::DeepSeekV4FlashDspark => &FLASH_DSPARK,
Self::DeepSeekV4Pro => &PRO, Self::DeepSeekV4Pro => &PRO,
Self::Glm52 => &GLM, Self::Glm52 => &GLM,

View File

@@ -331,15 +331,18 @@ mod tests {
86_720_111_488 86_720_111_488
); );
assert_eq!(ModelChoice::Glm52.main_artifact().size, 211_075_856_448); assert_eq!(ModelChoice::Glm52.main_artifact().size, 211_075_856_448);
assert_eq!(ModelChoice::DeepSeekV4Flash.artifacts(true).count(), 2); assert_eq!(
assert_eq!(ModelChoice::Glm52.artifacts(true).count(), 1); ModelChoice::DeepSeekV4Flash.artifacts(true, true).count(),
3
);
assert_eq!(ModelChoice::Glm52.artifacts(true, true).count(), 1);
let id = SystemTime::now() let id = SystemTime::now()
.duration_since(UNIX_EPOCH) .duration_since(UNIX_EPOCH)
.unwrap() .unwrap()
.as_nanos(); .as_nanos();
let models_path = std::env::temp_dir().join(format!("ds4-server-models-{id}")); let models_path = std::env::temp_dir().join(format!("ds4-server-models-{id}"));
let engine = engine_artifacts(ModelChoice::DeepSeekV4Flash, true, &models_path); let engine = engine_artifacts(ModelChoice::DeepSeekV4Flash, false, true, &models_path);
assert_eq!( assert_eq!(
engine.model.file_name(), engine.model.file_name(),
Some(std::ffi::OsStr::new(FLASH.file_name)) Some(std::ffi::OsStr::new(FLASH.file_name))

View File

@@ -23,6 +23,7 @@ const DEFAULT_KV_CONTINUED_INTERVAL_TOKENS: u32 = 10_000;
pub(crate) struct SpeculativePreferences { pub(crate) struct SpeculativePreferences {
pub(crate) mtp_draft_tokens: i32, pub(crate) mtp_draft_tokens: i32,
pub(crate) mtp_margin: f32, pub(crate) mtp_margin: f32,
pub(crate) legacy_mtp_enabled: bool,
pub(crate) glm_mtp: bool, pub(crate) glm_mtp: bool,
pub(crate) glm_mtp_timing: bool, pub(crate) glm_mtp_timing: bool,
pub(crate) dspark_enabled: bool, pub(crate) dspark_enabled: bool,
@@ -35,6 +36,7 @@ impl Default for SpeculativePreferences {
Self { Self {
mtp_draft_tokens: 1, mtp_draft_tokens: 1,
mtp_margin: 3.0, mtp_margin: 3.0,
legacy_mtp_enabled: false,
glm_mtp: false, glm_mtp: false,
glm_mtp_timing: false, glm_mtp_timing: false,
dspark_enabled: false, dspark_enabled: false,
@@ -59,6 +61,12 @@ impl SpeculativePreferences {
if self.dspark_enabled && !model.supports_dspark() { if self.dspark_enabled && !model.supports_dspark() {
return Err("DSpark is not available for the selected model.".into()); return Err("DSpark is not available for the selected model.".into());
} }
if self.legacy_mtp_enabled && !model.supports_dspark() {
return Err("Legacy MTP is not available for the selected model.".into());
}
if self.legacy_mtp_enabled && self.dspark_enabled {
return Err("Legacy MTP and DSpark use different support artifacts.".into());
}
if (self.dspark_confidence_threshold.is_some() || self.dspark_strict) if (self.dspark_confidence_threshold.is_some() || self.dspark_strict)
&& !self.dspark_enabled && !self.dspark_enabled
{ {
@@ -343,6 +351,9 @@ impl DiagnosticPreferences {
if let Some(gib) = self.simulated_used_memory_gib { if let Some(gib) = self.simulated_used_memory_gib {
validate_gib("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());
}
Ok(()) Ok(())
} }
@@ -379,9 +390,6 @@ impl RuntimePreferences {
self.steering.validate(model)?; self.steering.validate(model)?;
self.diagnostics.validate()?; self.diagnostics.validate()?;
self.kv_cache.validate()?; self.kv_cache.validate()?;
if self.ssd.enabled && self.speculative.dspark_enabled {
return Err("SSD streaming is not compatible with DSpark support.".into());
}
Ok(()) Ok(())
} }
@@ -394,7 +402,12 @@ impl RuntimePreferences {
self.validate(model)?; self.validate(model)?;
Ok(EngineSettings { Ok(EngineSettings {
model, model,
artifacts: model::engine_artifacts(model, self.speculative.dspark_enabled, models_path), artifacts: model::engine_artifacts(
model,
self.speculative.legacy_mtp_enabled,
self.speculative.dspark_enabled,
models_path,
),
context_tokens, context_tokens,
execution: self.execution.engine_settings(), execution: self.execution.engine_settings(),
speculative: self.speculative.engine_settings(), speculative: self.speculative.engine_settings(),
@@ -438,6 +451,9 @@ impl ExecutionPreferences {
{ {
return Err("CPU helper threads is too large.".into()); return Err("CPU helper threads is too large.".into());
} }
if self.cpu_threads.is_some() {
return Err("CPU helper threads do not apply to the Rust Metal executor.".into());
}
if self if self
.power_percent .power_percent
.is_some_and(|power| !(1..=100).contains(&power)) .is_some_and(|power| !(1..=100).contains(&power))
@@ -680,15 +696,23 @@ mod tests {
(0, 0, 0) (0, 0, 0)
); );
let tuned = ExecutionPreferences { let unsupported_threads = ExecutionPreferences {
cpu_threads: Some(100), cpu_threads: Some(100),
..ExecutionPreferences::default()
};
assert!(
unsupported_threads
.validate(ModelChoice::DeepSeekV4Flash)
.is_err()
);
let tuned = ExecutionPreferences {
power_percent: Some(50), power_percent: Some(50),
prefill_chunk: Some(4096), prefill_chunk: Some(4096),
..ExecutionPreferences::default() ..ExecutionPreferences::default()
}; };
assert!(tuned.validate(ModelChoice::DeepSeekV4Flash).is_ok()); assert!(tuned.validate(ModelChoice::DeepSeekV4Flash).is_ok());
assert!(tuned.validate(ModelChoice::Glm52).is_err()); assert!(tuned.validate(ModelChoice::Glm52).is_err());
assert_eq!(tuned.engine_settings().cpu_threads, MAX_CPU_THREADS);
} }
#[test] #[test]
@@ -717,6 +741,21 @@ mod tests {
}; };
assert!(glm.validate(ModelChoice::Glm52).is_ok()); assert!(glm.validate(ModelChoice::Glm52).is_ok());
assert!(glm.validate(ModelChoice::DeepSeekV4Pro).is_err()); assert!(glm.validate(ModelChoice::DeepSeekV4Pro).is_err());
let legacy = SpeculativePreferences {
legacy_mtp_enabled: true,
..SpeculativePreferences::default()
};
assert!(legacy.validate(ModelChoice::DeepSeekV4Flash).is_ok());
assert!(legacy.validate(ModelChoice::DeepSeekV4Pro).is_err());
assert!(
SpeculativePreferences {
dspark_enabled: true,
..legacy
}
.validate(ModelChoice::DeepSeekV4Flash)
.is_err()
);
} }
#[test] #[test]
@@ -787,14 +826,14 @@ mod tests {
assert_eq!(engine.steering.ffn_scale, 1.0); assert_eq!(engine.steering.ffn_scale, 1.0);
assert_eq!(engine.diagnostics.simulated_used_memory_bytes, 8 * GIB); assert_eq!(engine.diagnostics.simulated_used_memory_bytes, 8 * GIB);
let incompatible = RuntimePreferences { let combined = RuntimePreferences {
speculative: SpeculativePreferences { speculative: SpeculativePreferences {
dspark_enabled: true, dspark_enabled: true,
..SpeculativePreferences::default() ..SpeculativePreferences::default()
}, },
..runtime ..runtime
}; };
assert!(incompatible.validate(ModelChoice::DeepSeekV4Flash).is_err()); assert!(combined.validate(ModelChoice::DeepSeekV4Flash).is_ok());
} }
#[test] #[test]