memra_reference/
hidden_trace.rs

1//! Env-gated hidden-state trace, shared by the reference executor and the CUDA trunk.
2//!
3//! `MEMRA_HYPER_TRACE=<path>` makes BOTH sides append the same stage names, in the same
4//! format, from the same token ids — so a bisect compares like with like instead of two
5//! hand-rolled dumps that agree on nothing but the layer index. It lives in the reference
6//! crate because `memra-engine` already depends on it, so one emitter serves both.
7//!
8//! Format (`memra-hidden-trace-v1`), one line per stage, the LAST token row only:
9//!
10//! ```text
11//! stage\t<name>\t<layer|-1>\t<width>\t<f32 bits hex>,<f32 bits hex>,...
12//! ```
13//!
14//! Only the last row is emitted: it is the row the logits come from, and it is the row every
15//! banked oracle in this lane already pins. Width is `streams * hidden` at the residual-stream
16//! stages and `hidden` at the branch stages.
17//!
18//! Off by default and, when off, costs one `OnceLock` read per call — no allocation, no
19//! device work, and no arm selection changes anywhere (unlike the MoE traces, which move
20//! `observation_mode`). Turning it on cannot change what the model computes.
21
22use std::fmt::Write as _;
23use std::fs::File;
24use std::io::Write as _;
25use std::sync::{Mutex, OnceLock};
26
27static SINK: OnceLock<Option<Mutex<File>>> = OnceLock::new();
28
29fn sink() -> Option<&'static Mutex<File>> {
30    SINK.get_or_init(|| {
31        let path = std::env::var_os("MEMRA_HYPER_TRACE")?;
32        let mut file = File::create(&path)
33            .unwrap_or_else(|error| panic!("MEMRA_HYPER_TRACE={path:?}: {error}"));
34        writeln!(file, "format\tmemra-hidden-trace-v1").ok();
35        Some(Mutex::new(file))
36    })
37    .as_ref()
38}
39
40/// True when `MEMRA_HYPER_TRACE` named a path. Callers use it to skip a device readback.
41pub fn enabled() -> bool {
42    sink().is_some()
43}
44
45/// Record the run's token ids once, so a trace file can be checked against its oracle TSV.
46pub fn emit_tokens(token_ids: &[u32]) {
47    let Some(sink) = sink() else { return };
48    let ids = token_ids
49        .iter()
50        .map(u32::to_string)
51        .collect::<Vec<_>>()
52        .join(",");
53    let mut file = sink.lock().expect("hidden-trace sink poisoned");
54    writeln!(file, "tokens\t{ids}").ok();
55}
56
57/// Emit the last token row of a `[rows, width]` row-major activation.
58///
59/// `layer` is the plan layer index, or `-1` for the trunk-level stages (`expand`, `collapse`).
60pub fn emit_last_row(stage: &str, layer: i64, rows: usize, width: usize, data: &[f32]) {
61    let Some(sink) = sink() else { return };
62    if rows == 0 || width == 0 || data.len() != rows * width {
63        panic!(
64            "hidden-trace {stage}[{layer}]: {} values is not rows {rows} x width {width}",
65            data.len()
66        );
67    }
68    let row = &data[(rows - 1) * width..];
69    let mut line = String::with_capacity(width * 9 + 64);
70    let _ = write!(line, "stage\t{stage}\t{layer}\t{width}\t");
71    for (index, value) in row.iter().enumerate() {
72        if index != 0 {
73            line.push(',');
74        }
75        let _ = write!(line, "{:08x}", value.to_bits());
76    }
77    line.push('\n');
78    let mut file = sink.lock().expect("hidden-trace sink poisoned");
79    file.write_all(line.as_bytes()).ok();
80}
memra_reference/hidden_trace.rs

memra_reference/
hidden_trace.rs
