Align DeepSeek and GLM execution with DS4
This commit is contained in:
Georg Bauer
@@ -0,0 +1,31 @@
|
||||
#!/bin/bash
|
||||
# Standalone oracle only. Never add its C objects to Cargo or the app bundle.
|
||||
set -euo pipefail
|
||||
if [[ $# != 2 ]]; then
|
||||
echo "usage: bash tools/build-ds4-session-reference.sh REFERENCE_DIRECTORY OUTPUT_BINARY" >&2
|
||||
exit 2
|
||||
fi
|
||||
reference=$(cd "$1" && pwd)
|
||||
output=$2
|
||||
cd "$(dirname "$0")/.."
|
||||
[[ $(git -C "$reference" rev-parse HEAD) == ec7642cdd9ec81d01ad4b1fd8f8a3d1511533748 ]]
|
||||
[[ $(shasum -a 256 "$reference/ds4.h" | cut -d ' ' -f 1) == 1fe2491b0b709222dd41076d35b6eb85684d8c467a691ff4d45534d329f06360 ]]
|
||||
git -C "$reference" diff --exit-code HEAD -- ds4.h ds4.c ds4_metal.m ds4_image.c ds4_distributed.c ds4_tp.c ds4_ssd.c ds4_layer_pack.c metal
|
||||
# Refuse stale objects. If necessary, build the authorized reference separately.
|
||||
make -q -C "$reference" ds4
|
||||
json_library=$(cargo build --release --bin test-supervisor -j 4 --message-format=json |
|
||||
jq -rs '[.[] | select(.reason == "compiler-artifact" and .target.name == "serde_json") | .filenames[] | select(endswith(".rlib"))] | last')
|
||||
[[ -n $json_library ]]
|
||||
probe_dir=$(mktemp -d)
|
||||
trap 'rm -f "$probe_dir/canary.o"; rmdir "$probe_dir"' EXIT
|
||||
clang -c -O3 -fobjc-arc -ffast-math -mcpu=native -I native/metal \
|
||||
native/metal/ds4_canary.m -o "$probe_dir/canary.o"
|
||||
link=()
|
||||
for object in ds4 ds4_image ds4_distributed ds4_tp ds4_ssd ds4_metal ds4_layer_pack; do
|
||||
link+=(-C "link-arg=$reference/$object.o")
|
||||
done
|
||||
rustc --edition 2024 -O -D warnings tools/ds4-session-reference.rs \
|
||||
--extern "serde_json=$json_library" -L dependency=target/release/deps \
|
||||
-C "link-arg=$probe_dir/canary.o" \
|
||||
"${link[@]}" -l framework=Foundation -l framework=Metal -l objc -l pthread -l m \
|
||||
-o "$output"
|
||||
@@ -0,0 +1,775 @@
|
||||
//! Standalone reference benchmark, NEVER linked into DS4Server/Cargo targets.
|
||||
//! Uses antirez/ds4's public session API, pinned to ec7642c (ds4.h ABI, arm64).
|
||||
//! Build separately with the original reference objects; run with test-supervisor.
|
||||
//! Arguments: MODEL FAMILY(glm|deepseek) ACCELERATION(on|off) README [SUPPORT_GGUF]
|
||||
//! Power 100, Low, context 32768, temp .6/top-p .95/top-k 0/min-p 0/seed 42.
|
||||
//! Separate 32-token warmup, then one ongoing Summary/Story/Python chat to EOS.
|
||||
//! Optional DS4_REFERENCE_CANARY=/path/to/ds4-server uses the production probe
|
||||
//! in a separate process. Leave unset for clean throughput measurements.
|
||||
use serde_json::json;
|
||||
use std::ffi::{CStr, CString, c_char, c_void};
|
||||
use std::io::{self, BufRead, BufReader, Write};
|
||||
use std::mem::{offset_of, size_of};
|
||||
use std::process::{Child, Command, Stdio};
|
||||
use std::ptr::{null, null_mut};
|
||||
use std::sync::{
|
||||
Arc, Mutex,
|
||||
atomic::{AtomicBool, Ordering},
|
||||
mpsc,
|
||||
};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
#[path = "sampler-replay-benchmark.rs"]
|
||||
mod sampler_replay_benchmark;
|
||||
|
||||
// Same native probe object as DS4Server, linked only into this standalone oracle.
|
||||
#[repr(C)]
|
||||
#[derive(Default)]
|
||||
struct CanarySample {
|
||||
scheduled: f64,
|
||||
completed: f64,
|
||||
gpu_wait: f64,
|
||||
gpu_interval: f64,
|
||||
host_return: f64,
|
||||
}
|
||||
unsafe extern "C" {
|
||||
fn ds4_gpu_canary_probe(sample: *mut CanarySample) -> i32;
|
||||
fn ds4_engine_vocab_size(engine: *mut c_void) -> i32;
|
||||
fn ds4_engine_power(engine: *mut c_void) -> i32;
|
||||
fn ds4_session_prefill_cap(session: *mut c_void) -> i32;
|
||||
fn ds4_session_copy_logits(session: *mut c_void, out: *mut f32, cap: i32) -> i32;
|
||||
fn ds4_sample_logits(
|
||||
logits: *const f32,
|
||||
n: i32,
|
||||
temperature: f32,
|
||||
top_k: i32,
|
||||
top_p: f32,
|
||||
min_p: f32,
|
||||
rng: *mut u64,
|
||||
) -> i32;
|
||||
}
|
||||
|
||||
// CPU-only oracle fixture: no weights, model initialization or Metal work.
|
||||
fn sampler_fixture() {
|
||||
let mut cases = Vec::new();
|
||||
for n in [4, 64, 513, 4096] {
|
||||
let logits: Vec<f32> = if n == 4 {
|
||||
vec![2.0, 2.0, 0.0, -2.0]
|
||||
} else {
|
||||
(0..n)
|
||||
.map(|i| ((i * 37 % 101) as f32 - 50.0) / 8.0)
|
||||
.collect()
|
||||
};
|
||||
for (temperature, top_k, top_p, min_p) in [
|
||||
(0.0, 0, 0.95, 0.0),
|
||||
(0.6, 0, 0.95, 0.0),
|
||||
(0.6, 0, 1.0, 0.2),
|
||||
(0.6, 3, 0.5, 0.1),
|
||||
(0.6, 1, 0.95, 0.0),
|
||||
(0.6, 2048, 0.999, 0.0),
|
||||
(0.6, 0, 1.0, 1.1),
|
||||
(1.0, 0, 0.95, 0.0),
|
||||
] {
|
||||
for seed in [0, 42] {
|
||||
let mut rng = seed;
|
||||
let tokens: Vec<i32> = (0..32)
|
||||
.map(|_| unsafe {
|
||||
ds4_sample_logits(
|
||||
logits.as_ptr(),
|
||||
n,
|
||||
temperature,
|
||||
top_k,
|
||||
top_p,
|
||||
min_p,
|
||||
&mut rng,
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
cases.push(json!({"n":n,"temperature":temperature,"top_k":top_k,"top_p":top_p,"min_p":min_p,"seed":seed,"tokens":tokens,"rng_after":rng}));
|
||||
}
|
||||
}
|
||||
}
|
||||
emit(
|
||||
json!({"reference":"antirez/ds4", "commit":"ec7642cdd9ec81d01ad4b1fd8f8a3d1511533748", "recipe":"n=4: [2,2,0,-2]; otherwise ((i*37%101)-50)/8 in f32", "cases":cases}),
|
||||
);
|
||||
}
|
||||
struct InlineProbe {
|
||||
phase: Arc<Mutex<&'static str>>,
|
||||
stop: Arc<AtomicBool>,
|
||||
worker: Option<std::thread::JoinHandle<Result<(), String>>>,
|
||||
}
|
||||
impl InlineProbe {
|
||||
fn start() -> Result<Self, String> {
|
||||
assert_eq!(size_of::<CanarySample>(), 40);
|
||||
let phase = Arc::new(Mutex::new("startup"));
|
||||
let stop = Arc::new(AtomicBool::new(false));
|
||||
let worker_phase = Arc::clone(&phase);
|
||||
let worker_stop = Arc::clone(&stop);
|
||||
let (ready_send, ready_recv) = mpsc::sync_channel(1);
|
||||
let worker = std::thread::spawn(move || {
|
||||
let started = Instant::now();
|
||||
let mut first = true;
|
||||
let mut samples = 0;
|
||||
loop {
|
||||
let at = Instant::now();
|
||||
let phase = *worker_phase.lock().unwrap();
|
||||
let mut sample = CanarySample::default();
|
||||
let ok = unsafe { ds4_gpu_canary_probe(&mut sample) } != 0;
|
||||
let ms = |s: f64| (s >= 0.0).then_some(s * 1000.0);
|
||||
eprintln!(
|
||||
"{}",
|
||||
json!({"event":"gpu_canary_sample", "observer":"in_process", "elapsed_ms":started.elapsed().as_millis(), "started_elapsed_ms":at.duration_since(started).as_millis(), "phase":phase, "completion_phase":*worker_phase.lock().unwrap(), "ok":ok, "scheduled_ms":ms(sample.scheduled), "completed_ms":ms(sample.completed), "gpu_wait_ms":ms(sample.gpu_wait), "gpu_interval_ms":ms(sample.gpu_interval), "host_return_ms":ms(sample.host_return)})
|
||||
);
|
||||
if first {
|
||||
let _ = ready_send.send(ok);
|
||||
first = false;
|
||||
}
|
||||
if !ok {
|
||||
return Err("in-process Metal canary failed".into());
|
||||
}
|
||||
samples += 1;
|
||||
if worker_stop.load(Ordering::Relaxed) {
|
||||
emit(
|
||||
json!({"event":"gpu_canary_summary", "observer":"in_process", "samples":samples, "failures":0}),
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
if let Some(wait) = Duration::from_millis(100).checked_sub(at.elapsed()) {
|
||||
std::thread::sleep(wait);
|
||||
}
|
||||
}
|
||||
});
|
||||
let probe = Self {
|
||||
phase,
|
||||
stop,
|
||||
worker: Some(worker),
|
||||
};
|
||||
if ready_recv.recv_timeout(Duration::from_secs(2)) != Ok(true) {
|
||||
return Err("in-process canary did not become ready".into());
|
||||
}
|
||||
emit(json!({"event":"gpu_canary_ready", "observer":"in_process"}));
|
||||
Ok(probe)
|
||||
}
|
||||
fn finish(&mut self) -> Result<(), String> {
|
||||
self.stop.store(true, Ordering::Relaxed);
|
||||
let deadline = Instant::now() + Duration::from_secs(2);
|
||||
while self.worker.as_ref().is_some_and(|w| !w.is_finished()) {
|
||||
if Instant::now() >= deadline {
|
||||
return Err("in-process canary did not finish".into());
|
||||
}
|
||||
std::thread::sleep(Duration::from_millis(10));
|
||||
}
|
||||
if let Some(worker) = self.worker.take() {
|
||||
worker.join().map_err(|_| "in-process canary panicked")??;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
impl Drop for InlineProbe {
|
||||
fn drop(&mut self) {
|
||||
self.stop.store(true, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
struct Probe {
|
||||
child: Option<Child>,
|
||||
reader: Option<std::thread::JoinHandle<Result<(), String>>>,
|
||||
inline: Option<InlineProbe>,
|
||||
}
|
||||
impl Probe {
|
||||
fn start() -> Result<Self, String> {
|
||||
let child = std::env::var_os("DS4_REFERENCE_CANARY")
|
||||
.map(|path| {
|
||||
Command::new(path)
|
||||
.arg("gpu-canary")
|
||||
.stdin(Stdio::piped())
|
||||
.stdout(Stdio::piped())
|
||||
.spawn()
|
||||
})
|
||||
.transpose()
|
||||
.map_err(|e| e.to_string())?;
|
||||
let mut probe = Self {
|
||||
child,
|
||||
reader: None,
|
||||
inline: None,
|
||||
};
|
||||
match std::env::var("DS4_REFERENCE_IN_PROCESS_CANARY").as_deref() {
|
||||
Ok("1") => probe.inline = Some(InlineProbe::start()?),
|
||||
Err(std::env::VarError::NotPresent) | Ok("0") => {}
|
||||
_ => return Err("DS4_REFERENCE_IN_PROCESS_CANARY must be 0 or 1".into()),
|
||||
}
|
||||
if let Some(child) = &mut probe.child {
|
||||
let mut output = BufReader::new(child.stdout.take().ok_or("canary stdout missing")?);
|
||||
let mut ready = String::new();
|
||||
output.read_line(&mut ready).map_err(|e| e.to_string())?;
|
||||
let ready: serde_json::Value =
|
||||
serde_json::from_str(&ready).map_err(|e| e.to_string())?;
|
||||
if ready["event"] != "gpu_canary_ready" {
|
||||
return Err(format!("canary not ready: {ready}"));
|
||||
}
|
||||
emit(ready);
|
||||
probe.reader = Some(std::thread::spawn(move || {
|
||||
for line in output.lines() {
|
||||
let line = line.map_err(|e| e.to_string())?;
|
||||
emit(serde_json::from_str(&line).map_err(|e| e.to_string())?);
|
||||
}
|
||||
Ok(())
|
||||
}));
|
||||
}
|
||||
Ok(probe)
|
||||
}
|
||||
fn phase(&mut self, phase: &'static str) -> Result<(), String> {
|
||||
if let Some(probe) = &self.inline {
|
||||
if probe.worker.as_ref().is_some_and(|w| w.is_finished()) {
|
||||
return Err("in-process canary exited early".into());
|
||||
}
|
||||
*probe.phase.lock().unwrap() = phase;
|
||||
}
|
||||
if let Some(child) = &mut self.child {
|
||||
if let Some(status) = child.try_wait().map_err(|e| e.to_string())? {
|
||||
return Err(format!("canary exited early: {status}"));
|
||||
}
|
||||
writeln!(
|
||||
child.stdin.as_mut().ok_or("canary stdin closed")?,
|
||||
"{phase}"
|
||||
)
|
||||
.map_err(|e| e.to_string())?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
fn finish(&mut self) -> Result<(), String> {
|
||||
self.phase("finishing")?;
|
||||
if let Some(probe) = &mut self.inline {
|
||||
probe.finish()?;
|
||||
}
|
||||
if let Some(child) = &mut self.child {
|
||||
drop(child.stdin.take());
|
||||
let deadline = Instant::now() + Duration::from_secs(3);
|
||||
loop {
|
||||
if let Some(status) = child.try_wait().map_err(|e| e.to_string())? {
|
||||
if let Some(reader) = self.reader.take() {
|
||||
reader.join().map_err(|_| "canary reader panicked")??;
|
||||
}
|
||||
return if status.success() {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(format!("canary failed: {status}"))
|
||||
};
|
||||
}
|
||||
if Instant::now() >= deadline {
|
||||
return Err("canary did not finish".into());
|
||||
}
|
||||
std::thread::sleep(Duration::from_millis(10));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
impl Drop for Probe {
|
||||
fn drop(&mut self) {
|
||||
if let Some(child) = &mut self.child {
|
||||
let _ = child.kill();
|
||||
let _ = child.wait();
|
||||
}
|
||||
if let Some(reader) = self.reader.take() {
|
||||
let _ = reader.join();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Public reference header declarations, not a port of inference code.
|
||||
#[repr(C)]
|
||||
#[derive(Default)]
|
||||
struct Tokens {
|
||||
v: *mut i32,
|
||||
len: i32,
|
||||
cap: i32,
|
||||
}
|
||||
#[repr(C)]
|
||||
#[derive(Default)]
|
||||
struct Layers {
|
||||
start: u32,
|
||||
end: u32,
|
||||
output: bool,
|
||||
set: bool,
|
||||
}
|
||||
#[repr(C)]
|
||||
#[derive(Default)]
|
||||
struct Distributed {
|
||||
role: i32,
|
||||
layers: Layers,
|
||||
listen_host: *const c_char,
|
||||
listen_port: i32,
|
||||
coordinator_host: *const c_char,
|
||||
coordinator_port: i32,
|
||||
prefill_chunk: u32,
|
||||
prefill_window: u32,
|
||||
activation_bits: u32,
|
||||
replay_check: bool,
|
||||
debug: bool,
|
||||
}
|
||||
#[repr(C)]
|
||||
#[derive(Default)]
|
||||
struct Tp {
|
||||
role: i32,
|
||||
requested: bool,
|
||||
listen_host: *const c_char,
|
||||
listen_port: i32,
|
||||
leader_host: *const c_char,
|
||||
leader_port: i32,
|
||||
transport: i32,
|
||||
rdma_device: *const c_char,
|
||||
rdma_gid_index: i32,
|
||||
rdma_gid_index_set: bool,
|
||||
glm_token_prefill: bool,
|
||||
debug_hash: i32,
|
||||
}
|
||||
#[repr(C)]
|
||||
#[derive(Default)]
|
||||
struct Options {
|
||||
model: *const c_char,
|
||||
mtp: *const c_char,
|
||||
vision: *const c_char,
|
||||
backend: i32,
|
||||
threads: i32,
|
||||
context: i32,
|
||||
prefill_chunk: u32,
|
||||
mtp_draft_tokens: i32,
|
||||
mtp_margin: f32,
|
||||
dspark_confidence_threshold: f32,
|
||||
steering: *const c_char,
|
||||
expert_profile: *const c_char,
|
||||
steering_attn: f32,
|
||||
steering_ffn: f32,
|
||||
power: i32,
|
||||
cache_experts: u32,
|
||||
cache_bytes: u64,
|
||||
full_layers: u32,
|
||||
preload_experts: u32,
|
||||
simulated_memory: u64,
|
||||
warm_weights: bool,
|
||||
quality: bool,
|
||||
glm_mtp: bool,
|
||||
glm_mtp_timing: bool,
|
||||
dspark: bool,
|
||||
dspark_strict: bool,
|
||||
dspark_exact: bool,
|
||||
confidence_set: bool,
|
||||
cuda_tp: bool,
|
||||
ssd: bool,
|
||||
ssd_cold: bool,
|
||||
full_layers_set: bool,
|
||||
inspect: bool,
|
||||
placement_ctx: i32,
|
||||
placement_sessions: i32,
|
||||
share_prefill: bool,
|
||||
first_token_test: bool,
|
||||
graph_test: bool,
|
||||
load_slice: bool,
|
||||
load_start: u32,
|
||||
load_end: u32,
|
||||
load_output: bool,
|
||||
distributed: Distributed,
|
||||
tp: Tp,
|
||||
}
|
||||
unsafe extern "C" {
|
||||
fn ds4_engine_open(out: *mut *mut c_void, options: *const Options) -> i32;
|
||||
fn ds4_engine_close(engine: *mut c_void);
|
||||
fn ds4_engine_mtp_draft_tokens(engine: *mut c_void) -> i32;
|
||||
fn ds4_session_create(out: *mut *mut c_void, engine: *mut c_void, ctx: i32) -> i32;
|
||||
fn ds4_session_free(session: *mut c_void);
|
||||
fn ds4_session_set_progress(
|
||||
s: *mut c_void,
|
||||
f: unsafe extern "C" fn(*mut c_void, *const c_char, i32, i32),
|
||||
ud: *mut c_void,
|
||||
);
|
||||
fn ds4_session_set_display_progress(
|
||||
s: *mut c_void,
|
||||
f: unsafe extern "C" fn(*mut c_void, *const c_char, i32, i32),
|
||||
ud: *mut c_void,
|
||||
);
|
||||
fn ds4_session_sync(s: *mut c_void, tokens: *const Tokens, err: *mut c_char, len: usize)
|
||||
-> i32;
|
||||
fn ds4_session_common_prefix(s: *mut c_void, tokens: *const Tokens) -> i32;
|
||||
fn ds4_session_sample(
|
||||
s: *mut c_void,
|
||||
temp: f32,
|
||||
top_k: i32,
|
||||
top_p: f32,
|
||||
min_p: f32,
|
||||
rng: *mut u64,
|
||||
) -> i32;
|
||||
fn ds4_session_eval(s: *mut c_void, token: i32, err: *mut c_char, len: usize) -> i32;
|
||||
fn ds4_session_pos(s: *mut c_void) -> i32;
|
||||
fn ds4_session_eval_speculative(
|
||||
s: *mut c_void,
|
||||
first: i32,
|
||||
max: i32,
|
||||
eos: i32,
|
||||
temp: f32,
|
||||
top_k: i32,
|
||||
top_p: f32,
|
||||
min_p: f32,
|
||||
rng: *mut u64,
|
||||
out: *mut i32,
|
||||
cap: i32,
|
||||
err: *mut c_char,
|
||||
len: usize,
|
||||
) -> i32;
|
||||
fn ds4_session_rewind(s: *mut c_void, pos: i32);
|
||||
fn ds4_tokens_push(tokens: *mut Tokens, token: i32);
|
||||
fn ds4_tokens_free(tokens: *mut Tokens);
|
||||
fn ds4_chat_begin(e: *mut c_void, tokens: *mut Tokens);
|
||||
fn ds4_chat_append_message(
|
||||
e: *mut c_void,
|
||||
tokens: *mut Tokens,
|
||||
role: *const c_char,
|
||||
text: *const c_char,
|
||||
);
|
||||
fn ds4_chat_append_assistant_prefix(e: *mut c_void, tokens: *mut Tokens, think: i32);
|
||||
fn ds4_token_text(e: *mut c_void, token: i32, len: *mut usize) -> *mut c_char;
|
||||
fn ds4_token_eos(e: *mut c_void) -> i32;
|
||||
fn ds4_token_is_stop_for_think_mode(e: *mut c_void, token: i32, think: i32) -> bool;
|
||||
fn free(p: *mut c_void);
|
||||
}
|
||||
struct Engine(*mut c_void);
|
||||
impl Drop for Engine {
|
||||
fn drop(&mut self) {
|
||||
unsafe { ds4_engine_close(self.0) };
|
||||
}
|
||||
}
|
||||
struct Session(*mut c_void);
|
||||
impl Drop for Session {
|
||||
fn drop(&mut self) {
|
||||
unsafe { ds4_session_free(self.0) };
|
||||
}
|
||||
}
|
||||
impl Drop for Tokens {
|
||||
fn drop(&mut self) {
|
||||
unsafe { ds4_tokens_free(self) };
|
||||
}
|
||||
}
|
||||
|
||||
fn emit(value: serde_json::Value) {
|
||||
let mut out = io::stdout().lock();
|
||||
if writeln!(out, "{value}").and_then(|()| out.flush()).is_err() {
|
||||
std::process::exit(1);
|
||||
}
|
||||
}
|
||||
unsafe extern "C" fn progress(_: *mut c_void, event: *const c_char, current: i32, total: i32) {
|
||||
emit(
|
||||
json!({"event":"reference_progress", "kind":unsafe { CStr::from_ptr(event) }.to_string_lossy(), "current":current, "total":total}),
|
||||
);
|
||||
}
|
||||
fn check(code: i32, err: &[c_char]) -> Result<(), String> {
|
||||
if code == 0 {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(format!(
|
||||
"reference error {code}: {}",
|
||||
unsafe { CStr::from_ptr(err.as_ptr()) }.to_string_lossy()
|
||||
))
|
||||
}
|
||||
}
|
||||
fn session(engine: &Engine) -> Result<Session, String> {
|
||||
let mut s = null_mut();
|
||||
if unsafe { ds4_session_create(&mut s, engine.0, 32768) } != 0 || s.is_null() {
|
||||
return Err("reference session creation failed".into());
|
||||
}
|
||||
unsafe {
|
||||
ds4_session_set_progress(s, progress, null_mut());
|
||||
ds4_session_set_display_progress(s, progress, null_mut());
|
||||
}
|
||||
Ok(Session(s))
|
||||
}
|
||||
fn turn(
|
||||
engine: &Engine,
|
||||
session: &Session,
|
||||
tokens: &mut Tokens,
|
||||
prompt: &str,
|
||||
glm: bool,
|
||||
index: usize,
|
||||
max: i32,
|
||||
probe: &mut Probe,
|
||||
) -> Result<(), String> {
|
||||
let started = Instant::now();
|
||||
probe.phase("prefill")?;
|
||||
let mut err = [0; 256];
|
||||
let mut bootstrap_ms = 0.0;
|
||||
let prompt = CString::new(prompt).map_err(|e| e.to_string())?;
|
||||
if tokens.len == 0 {
|
||||
unsafe {
|
||||
ds4_chat_begin(engine.0, tokens);
|
||||
if glm {
|
||||
// Low is prompt policy; the public CLI aliases it to High.
|
||||
ds4_chat_append_message(
|
||||
engine.0,
|
||||
tokens,
|
||||
c"system".as_ptr(),
|
||||
c"Reasoning Effort: Low".as_ptr(),
|
||||
);
|
||||
}
|
||||
}
|
||||
// The UI's shared runtime prepares its system-prefix KV separately.
|
||||
// Reproduce that boundary; a cold full-prompt batch can round differently.
|
||||
let bootstrap = Instant::now();
|
||||
check(
|
||||
unsafe { ds4_session_sync(session.0, tokens, err.as_mut_ptr(), err.len()) },
|
||||
&err,
|
||||
)?;
|
||||
bootstrap_ms = bootstrap.elapsed().as_secs_f64() * 1000.0;
|
||||
} else if !glm {
|
||||
unsafe {
|
||||
ds4_tokens_push(tokens, ds4_token_eos(engine.0));
|
||||
}
|
||||
}
|
||||
unsafe {
|
||||
ds4_chat_append_message(engine.0, tokens, c"user".as_ptr(), prompt.as_ptr());
|
||||
ds4_chat_append_assistant_prefix(engine.0, tokens, 1);
|
||||
}
|
||||
let cached = unsafe { ds4_session_common_prefix(session.0, tokens) };
|
||||
let input = tokens.len;
|
||||
emit(
|
||||
json!({"event":"reference_prefill", "turn":index, "prompt_tokens":input, "cached_tokens":cached, "session_prefill_cap":unsafe { ds4_session_prefill_cap(session.0) }, "tokens":unsafe { std::slice::from_raw_parts(tokens.v, input as usize) }}),
|
||||
);
|
||||
let prefill = Instant::now();
|
||||
check(
|
||||
unsafe { ds4_session_sync(session.0, tokens, err.as_mut_ptr(), err.len()) },
|
||||
&err,
|
||||
)?;
|
||||
let prefill_ms = prefill.elapsed().as_secs_f64() * 1000.0;
|
||||
emit(json!({"event":"reference_decode", "turn":index, "prefill_ms":prefill_ms}));
|
||||
probe.phase("decode")?;
|
||||
let mut rng = 42_u64;
|
||||
let mut ids = Vec::new();
|
||||
let mut text = Vec::new();
|
||||
let decode = Instant::now();
|
||||
let mut stop = None;
|
||||
let trace_cycles = std::env::var_os("DS4_SPEC_CYCLE_TRACE").is_some();
|
||||
let mut trace = if index == 1 {
|
||||
std::env::var_os("DS4_REFERENCE_LOGITS_TRACE").map(|path| {
|
||||
let n = unsafe { ds4_engine_vocab_size(engine.0) };
|
||||
if !(1..=1_000_000).contains(&n) { return Err("invalid vocabulary size".to_owned()); }
|
||||
let file = std::fs::OpenOptions::new().create_new(true).write(true).open(&path).map_err(|e| e.to_string())?;
|
||||
emit(json!({"event":"reference_logits_trace", "path":path, "vocab":n, "rows":32, "turn":index, "format":"little-endian-f32"}));
|
||||
Ok((file, vec![0.0_f32; n as usize]))
|
||||
}).transpose()?
|
||||
} else {
|
||||
None
|
||||
};
|
||||
while ids.len() < max as usize && tokens.len < 32767 {
|
||||
if ids.len() < 32
|
||||
&& let Some((file, logits)) = &mut trace
|
||||
{
|
||||
if unsafe {
|
||||
ds4_session_copy_logits(session.0, logits.as_mut_ptr(), logits.len() as i32)
|
||||
} != logits.len() as i32
|
||||
{
|
||||
return Err("copying reference logits failed".into());
|
||||
}
|
||||
let bytes: Vec<u8> = logits.iter().flat_map(|v| v.to_le_bytes()).collect();
|
||||
file.write_all(&bytes).map_err(|e| e.to_string())?;
|
||||
}
|
||||
let first = unsafe { ds4_session_sample(session.0, 0.6, 0, 0.95, 0.0, &mut rng) };
|
||||
if unsafe { ds4_token_is_stop_for_think_mode(engine.0, first, 1) } {
|
||||
stop = Some(first);
|
||||
break;
|
||||
}
|
||||
let mut accepted = [first; 17];
|
||||
let room = (max - ids.len() as i32).min(32767 - tokens.len);
|
||||
let count = unsafe {
|
||||
if ds4_engine_mtp_draft_tokens(engine.0) > 1 {
|
||||
ds4_session_eval_speculative(
|
||||
session.0,
|
||||
first,
|
||||
room,
|
||||
ds4_token_eos(engine.0),
|
||||
0.6,
|
||||
0,
|
||||
0.95,
|
||||
0.0,
|
||||
&mut rng,
|
||||
accepted.as_mut_ptr(),
|
||||
room.min(17),
|
||||
err.as_mut_ptr(),
|
||||
err.len(),
|
||||
)
|
||||
} else {
|
||||
check(
|
||||
ds4_session_eval(session.0, first, err.as_mut_ptr(), err.len()),
|
||||
&err,
|
||||
)?;
|
||||
1
|
||||
}
|
||||
};
|
||||
if count <= 0 || count > 17 {
|
||||
return Err(format!(
|
||||
"invalid reference verify result {count}: {}",
|
||||
unsafe { CStr::from_ptr(err.as_ptr()) }.to_string_lossy()
|
||||
));
|
||||
}
|
||||
if trace_cycles {
|
||||
emit(json!({"event":"spec_cycle", "prompt_tokens":input,
|
||||
"generated":ids.len(), "first":first, "accepted":&accepted[..count as usize],
|
||||
"position":unsafe { ds4_session_pos(session.0) }}));
|
||||
}
|
||||
for &token in &accepted[..count as usize] {
|
||||
if unsafe { ds4_token_is_stop_for_think_mode(engine.0, token, 1) } {
|
||||
stop = Some(token);
|
||||
break;
|
||||
}
|
||||
unsafe {
|
||||
ds4_tokens_push(tokens, token);
|
||||
}
|
||||
ids.push(token);
|
||||
let mut len = 0;
|
||||
let piece = unsafe { ds4_token_text(engine.0, token, &mut len) };
|
||||
if piece.is_null() {
|
||||
return Err("null reference token text".into());
|
||||
}
|
||||
text.extend_from_slice(unsafe { std::slice::from_raw_parts(piece.cast::<u8>(), len) });
|
||||
unsafe {
|
||||
free(piece.cast());
|
||||
}
|
||||
}
|
||||
emit(json!({"event":"reference_tokens", "turn":index, "generated":ids.len()}));
|
||||
if stop.is_some() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
let decode_ms = decode.elapsed().as_secs_f64() * 1000.0;
|
||||
probe.phase("finishing")?;
|
||||
// Keep exactly the visible token frontier, excluding an MTP-returned EOS.
|
||||
unsafe {
|
||||
ds4_session_rewind(session.0, tokens.len);
|
||||
}
|
||||
emit(
|
||||
json!({"event":"reference_result", "turn":index, "prompt_tokens":input, "cached_tokens":cached, "completion_tokens":ids.len(), "token_ids":ids, "text":String::from_utf8(text).map_err(|e| e.to_string())?, "stop_token":stop, "finish_reason":if stop.is_some(){"stop"}else{"length"}, "bootstrap_ms":bootstrap_ms, "prefill_ms":prefill_ms, "decode_ms":decode_ms, "decode_tokens_per_second":ids.len() as f64 * 1000.0 / decode_ms, "elapsed_ms":started.elapsed().as_secs_f64()*1000.0}),
|
||||
);
|
||||
if index != 0 && stop.is_none() {
|
||||
return Err("measured turn exhausted context without EOS".into());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
fn run() -> Result<(), String> {
|
||||
// Checked against clang -fdump-record-layouts-complete for the pinned header.
|
||||
assert_eq!(
|
||||
(
|
||||
size_of::<Options>(),
|
||||
offset_of!(Options, distributed),
|
||||
offset_of!(Options, tp)
|
||||
),
|
||||
(280, 152, 216)
|
||||
);
|
||||
let args = std::env::args().skip(1).collect::<Vec<_>>();
|
||||
if args.first().is_some_and(|a| a == "--sampler-benchmark") {
|
||||
if args.len() != 3 {
|
||||
return Err("usage: reference --sampler-benchmark F32_RECORDING VOCAB".into());
|
||||
}
|
||||
let vocab = args[2].parse::<usize>().map_err(|e| e.to_string())?;
|
||||
let mut rng = 42_u64;
|
||||
emit(sampler_replay_benchmark::run(
|
||||
std::path::Path::new(&args[1]),
|
||||
vocab,
|
||||
|row| unsafe {
|
||||
ds4_sample_logits(row.as_ptr(), row.len() as i32, 0.6, 0, 0.95, 0.0, &mut rng)
|
||||
},
|
||||
)?);
|
||||
return Ok(());
|
||||
}
|
||||
if args == ["--sampler-fixture"] {
|
||||
sampler_fixture();
|
||||
return Ok(());
|
||||
}
|
||||
if args == ["--canary-self-test"] {
|
||||
let mut probe = InlineProbe::start()?;
|
||||
for phase in ["prefill", "decode"] {
|
||||
*probe.phase.lock().unwrap() = phase;
|
||||
std::thread::sleep(Duration::from_millis(300));
|
||||
}
|
||||
return probe.finish();
|
||||
}
|
||||
if !(4..=5).contains(&args.len())
|
||||
|| !["glm", "deepseek"].contains(&args[1].as_str())
|
||||
|| !["on", "off"].contains(&args[2].as_str())
|
||||
{
|
||||
return Err("usage: reference MODEL glm|deepseek on|off README [SUPPORT_GGUF]".into());
|
||||
}
|
||||
for file in [Some(&args[0]), args.get(4)].into_iter().flatten() {
|
||||
if !std::path::Path::new(file).is_file() {
|
||||
return Err(format!("installed artifact missing: {file}"));
|
||||
}
|
||||
}
|
||||
let readme = std::fs::read_to_string(&args[3]).map_err(|e| e.to_string())?;
|
||||
let model = CString::new(args[0].as_str()).map_err(|e| e.to_string())?;
|
||||
let support = args
|
||||
.get(4)
|
||||
.map(|s| CString::new(s.as_str()))
|
||||
.transpose()
|
||||
.map_err(|e| e.to_string())?;
|
||||
let glm = args[1] == "glm";
|
||||
let enabled = args[2] == "on";
|
||||
if enabled && std::env::var_os("DS4_REFERENCE_LOGITS_TRACE").is_some() {
|
||||
return Err("logits tracing requires acceleration off".into());
|
||||
}
|
||||
if !glm && enabled && support.is_none() {
|
||||
return Err("DSpark requires installed support GGUF".into());
|
||||
}
|
||||
let options = Options {
|
||||
model: model.as_ptr(),
|
||||
mtp: support.as_ref().map_or(null(), |s| s.as_ptr()),
|
||||
context: 32768,
|
||||
prefill_chunk: 0,
|
||||
mtp_draft_tokens: 1,
|
||||
mtp_margin: 3.0,
|
||||
dspark_confidence_threshold: 0.6,
|
||||
power: 100,
|
||||
glm_mtp: glm && enabled,
|
||||
dspark: !glm && enabled,
|
||||
..Options::default()
|
||||
};
|
||||
let mut probe = Probe::start()?;
|
||||
probe.phase("loading")?;
|
||||
emit(
|
||||
json!({"event":"reference_start", "family":args[1], "acceleration":enabled, "power_percent":100, "reasoning":"low", "model":args[0], "readme":args[3], "context":32768, "prefill_chunk":options.prefill_chunk, "temperature":0.6_f32, "top_p":0.95_f32, "min_p":0, "top_k":0, "seed":42, "canary":probe.child.is_some() || probe.inline.is_some(), "canary_in_process":probe.inline.is_some()}),
|
||||
);
|
||||
let mut e = null_mut();
|
||||
if unsafe { ds4_engine_open(&mut e, &options) } != 0 || e.is_null() {
|
||||
return Err("reference engine load failed".into());
|
||||
}
|
||||
let engine = Engine(e);
|
||||
let power = unsafe { ds4_engine_power(engine.0) };
|
||||
if power != 100 {
|
||||
return Err(format!("reference engine power is {power}, expected 100"));
|
||||
}
|
||||
emit(json!({"event":"reference_engine", "power_percent":power}));
|
||||
{
|
||||
let warmup = session(&engine)?;
|
||||
turn(
|
||||
&engine,
|
||||
&warmup,
|
||||
&mut Tokens::default(),
|
||||
"Reply with exactly: OK",
|
||||
glm,
|
||||
0,
|
||||
32,
|
||||
&mut probe,
|
||||
)?;
|
||||
}
|
||||
let measured = session(&engine)?;
|
||||
let mut tokens = Tokens::default();
|
||||
for (index, prompt) in [format!("Give a summary of the following text:\n\n{readme}"), "Tell me a complete short story about a lighthouse keeper. Do not ask questions.".into(), "Write a Python function is_prime(n: int) -> bool, followed by five assert examples. No tools.".into()].iter().enumerate() {
|
||||
turn(&engine, &measured, &mut tokens, prompt, glm, index + 1, i32::MAX, &mut probe)?;
|
||||
}
|
||||
drop(tokens);
|
||||
drop(measured);
|
||||
drop(engine);
|
||||
probe.finish()?;
|
||||
Ok(())
|
||||
}
|
||||
fn main() {
|
||||
if let Err(error) = run() {
|
||||
eprintln!("reference failed: {error}");
|
||||
std::process::exit(1);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,39 @@
|
||||
// CPU-only diagnostic shared by the standalone DS4 oracle and Rust tests.
|
||||
// The input is the existing 32-row, little-endian F32 logit recording.
|
||||
pub fn run(
|
||||
path: &std::path::Path,
|
||||
vocab: usize,
|
||||
mut sample: impl FnMut(&[f32]) -> i32,
|
||||
) -> Result<serde_json::Value, String> {
|
||||
if !(1..=1_000_000).contains(&vocab) {
|
||||
return Err("invalid sampler vocabulary size".into());
|
||||
}
|
||||
let expected = 32 * vocab * 4;
|
||||
if std::fs::metadata(path).map_err(|e| e.to_string())?.len() != expected as u64 {
|
||||
return Err("sampler recording must contain exactly 32 F32 rows".into());
|
||||
}
|
||||
let bytes = std::fs::read(path).map_err(|e| e.to_string())?;
|
||||
if bytes.len() != expected {
|
||||
return Err("sampler recording changed while reading".into());
|
||||
}
|
||||
let values = bytes
|
||||
.chunks_exact(4)
|
||||
.map(|b| f32::from_le_bytes(b.try_into().unwrap()))
|
||||
.collect::<Vec<_>>();
|
||||
for row in values.chunks_exact(vocab) {
|
||||
std::hint::black_box(sample(row));
|
||||
}
|
||||
let mut batch_ms = Vec::with_capacity(16);
|
||||
let mut token_ids = Vec::with_capacity(512);
|
||||
for _ in 0..16 {
|
||||
let started = std::time::Instant::now();
|
||||
for row in values.chunks_exact(vocab) {
|
||||
token_ids.push(std::hint::black_box(sample(row)));
|
||||
}
|
||||
batch_ms.push(started.elapsed().as_secs_f64() * 1000.0);
|
||||
}
|
||||
Ok(serde_json::json!({
|
||||
"event":"sampler_replay_benchmark", "vocab":vocab, "rows":32,
|
||||
"warmup_draws":32, "draws":512, "batch_ms":batch_ms, "token_ids":token_ids,
|
||||
}))
|
||||
}
|
||||
@@ -59,6 +59,8 @@ mod supervisor {
|
||||
value["event"].as_str(),
|
||||
Some(
|
||||
"gpu_canary_sample"
|
||||
| "gpu_canary_ready"
|
||||
| "gpu_canary_stall"
|
||||
| "gpu_canary_summary"
|
||||
| "reference_canary_summary"
|
||||
| "test_resource_sample"
|
||||
@@ -80,6 +82,31 @@ mod supervisor {
|
||||
progressed
|
||||
}
|
||||
|
||||
// Keep complete records together when the watchdog writes resource samples
|
||||
// to the same stream. Pipe reads can split a single worker JSON record.
|
||||
fn forward_records(
|
||||
output: &mut impl Write,
|
||||
pending: &mut Vec<u8>,
|
||||
bytes: &[u8],
|
||||
eof: bool,
|
||||
) -> io::Result<()> {
|
||||
pending.extend_from_slice(bytes);
|
||||
let end = if eof || pending.len() > 64 * 1024 {
|
||||
pending.len()
|
||||
} else {
|
||||
pending
|
||||
.iter()
|
||||
.rposition(|&b| b == b'\n')
|
||||
.map_or(0, |i| i + 1)
|
||||
};
|
||||
if end > 0 {
|
||||
output.write_all(&pending[..end])?;
|
||||
output.flush()?;
|
||||
pending.drain(..end);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn forward(
|
||||
mut input: impl Read + Send + 'static,
|
||||
mut output: impl Write + Send + 'static,
|
||||
@@ -88,15 +115,17 @@ mod supervisor {
|
||||
thread::spawn(move || {
|
||||
let mut bytes = [0; 4096];
|
||||
let mut pending = Vec::new();
|
||||
let mut records = Vec::new();
|
||||
loop {
|
||||
match input.read(&mut bytes) {
|
||||
Ok(0) => return,
|
||||
Ok(0) => {
|
||||
if forward_records(&mut output, &mut records, &[], true).is_err() {
|
||||
progress.lock().unwrap().failed = true;
|
||||
}
|
||||
return;
|
||||
}
|
||||
Ok(n) => {
|
||||
if output
|
||||
.write_all(&bytes[..n])
|
||||
.and_then(|()| output.flush())
|
||||
.is_err()
|
||||
{
|
||||
if forward_records(&mut output, &mut records, &bytes[..n], false).is_err() {
|
||||
progress.lock().unwrap().failed = true;
|
||||
return;
|
||||
}
|
||||
@@ -257,6 +286,29 @@ mod supervisor {
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn split_records_cannot_be_interleaved_with_watchdog_samples() {
|
||||
let mut output = Vec::new();
|
||||
let mut pending = Vec::new();
|
||||
forward_records(
|
||||
&mut output,
|
||||
&mut pending,
|
||||
b"{\"event\":\"gpu_canary_",
|
||||
false,
|
||||
)
|
||||
.unwrap();
|
||||
assert!(output.is_empty());
|
||||
output.extend_from_slice(b"{\"event\":\"test_resource_sample\"}\n");
|
||||
forward_records(&mut output, &mut pending, b"sample\"}\npartial", false).unwrap();
|
||||
for line in output.split(|&b| b == b'\n').filter(|l| !l.is_empty()) {
|
||||
serde_json::from_slice::<serde_json::Value>(line).unwrap();
|
||||
}
|
||||
assert_eq!(pending, b"partial");
|
||||
forward_records(&mut output, &mut pending, &[], true).unwrap();
|
||||
assert!(output.ends_with(b"partial"));
|
||||
assert!(pending.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn direct_worker_arguments_are_not_rust_test_arguments() {
|
||||
let args = |a: &[&str]| a.iter().map(std::ffi::OsString::from).collect::<Vec<_>>();
|
||||
|
||||
Reference in New Issue
Block a user