40 lines
1.5 KiB
Rust
40 lines
1.5 KiB
Rust
// 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,
|
|
}))
|
|
}
|