776 lines
27 KiB
Rust
776 lines
27 KiB
Rust
//! 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);
|
|
}
|
|
}
|