metrale_model_layers/layers/ple/
dump.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Debug taps: dump the `hc_mult`-wide residual highway and other
4//! buffers at named points, for a layer-by-layer diff against the reference.
5//!
6//! Off unless `METRALE_QWEN4EXP_DUMP` names a directory. `tap_highway` and
7//! `tap_f32` write `<dir>/L{layer:02}_{tag}.bin` as raw little-endian FP32;
8//! `tap_bf16` writes `<dir>/L{layer:02}_{tag}.bf16.bin` as raw BF16.
9//!
10//! Owner: model-layers (PLE).
11//! Invariants:
12//! - A tap never overwrites an existing file.
13
14use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
15
16/// 2026-09-25: Directory from `METRALE_QWEN4EXP_DUMP` (empty means unset),
17/// resolved once.
18fn dump_dir() -> Option<&'static str> {
19    static DIR: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
20    DIR.get_or_init(|| {
21        let d = std::env::var("METRALE_QWEN4EXP_DUMP")
22            .ok()
23            .filter(|s| !s.is_empty());
24        if let Some(ref path) = d {
25            let _ = std::fs::create_dir_all(path);
26            tracing::warn!(
27                "METRALE_QWEN4EXP_DUMP={path}: taping the mHC highway to disk. \
28                 This SYNCHRONIZES and copies D2H at every tap — a debug aid, \
29                 not a serving mode."
30            );
31        }
32        d
33    })
34    .as_deref()
35}
36
37/// 2026-09-25: Whether `path` is free to write; an existing tap is never
38/// overwritten. The SSM taps are labelled by `SSM_LAYER_CALL_COUNTER`, which
39/// counts every SSM layer call and never resets, so only the first prefill
40/// call after startup carries true layer numbers; later calls write under
41/// higher labels.
42fn claim(path: &str) -> bool {
43    !std::path::Path::new(path).exists()
44}
45
46/// 2026-09-25: Tap the FP32 highway, `num_tokens * hc_dim` values. No-op
47/// unless the dump directory is set. It synchronizes the stream first, so it
48/// must not run inside CUDA-graph capture. Errors are logged or dropped,
49/// never returned.
50pub fn tap_highway(
51    gpu: &dyn GpuBackend,
52    streams: DevicePtr,
53    layer: usize,
54    tag: &str,
55    num_tokens: usize,
56    hc_dim: usize,
57    stream: u64,
58) {
59    let Some(dir) = dump_dir() else {
60        return;
61    };
62    let path = format!("{dir}/L{layer:02}_{tag}.bin");
63    if !claim(&path) {
64        return;
65    }
66    if gpu.synchronize(stream).is_err() {
67        return;
68    }
69    let mut raw = vec![0u8; num_tokens * hc_dim * 4];
70    if gpu.copy_d2h(streams, &mut raw).is_err() {
71        return;
72    }
73    let path = format!("{dir}/L{layer:02}_{tag}.bin");
74    if let Err(e) = std::fs::write(&path, &raw) {
75        tracing::warn!("highway tap {path}: {e}");
76    }
77}
78
79/// 2026-09-25: Tap a BF16 buffer of `n_elements` the same way.
80pub fn tap_bf16(
81    gpu: &dyn GpuBackend,
82    ptr: DevicePtr,
83    layer: usize,
84    tag: &str,
85    n_elements: usize,
86    stream: u64,
87) {
88    let Some(dir) = dump_dir() else {
89        return;
90    };
91    let path = format!("{dir}/L{layer:02}_{tag}.bf16.bin");
92    if !claim(&path) {
93        return;
94    }
95    if gpu.synchronize(stream).is_err() {
96        return;
97    }
98    let mut raw = vec![0u8; n_elements * 2];
99    if gpu.copy_d2h(ptr, &mut raw).is_err() {
100        return;
101    }
102    if let Err(e) = std::fs::write(&path, &raw) {
103        tracing::warn!("highway tap {path}: {e}");
104    }
105}
106
107/// 2026-09-25: Tap an FP32 buffer of `n_elements` the same way.
108pub fn tap_f32(
109    gpu: &dyn GpuBackend,
110    ptr: DevicePtr,
111    layer: usize,
112    tag: &str,
113    n_elements: usize,
114    stream: u64,
115) {
116    let Some(dir) = dump_dir() else {
117        return;
118    };
119    let path = format!("{dir}/L{layer:02}_{tag}.bin");
120    if !claim(&path) {
121        return;
122    }
123    if gpu.synchronize(stream).is_err() {
124        return;
125    }
126    let mut raw = vec![0u8; n_elements * 4];
127    if gpu.copy_d2h(ptr, &mut raw).is_err() {
128        return;
129    }
130    if let Err(e) = std::fs::write(&path, &raw) {
131        tracing::warn!("highway tap {path}: {e}");
132    }
133}
metrale_model_layers/layers/ple/dump.rs

metrale_model_layers/layers/ple/
dump.rs
