spark_model/layers/ple/
dump.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Highway taps for the qwen4_exp bisect.
4//!
5//! Every PIECE of this port is pinned to the reference — the n-gram ids are
6//! bit-exact, the mHC kernels and the PLE gate/conv match to cosine 0.99999,
7//! the NVMe gather is bit-exact — and the model still does not produce
8//! coherent text. That combination says the fault is in the COMPOSITION, and
9//! composition is exactly what per-kernel probes cannot see.
10//!
11//! So: dump the `hc_mult`-wide residual highway at named points and diff it
12//! against the reference layer by layer.
13//!
14//! **The sub-layer boundary is what makes this affordable.** Tapping after a
15//! block's `hc_post` but BEFORE the next `hc_pre` means the reference only has
16//! to reproduce that block — for layer 0 that is the GDN projections alone,
17//! with none of the 512-expert MoE. Only the taps that come after an MoE need
18//! experts, and even then top-10 routing over a short prompt touches a few
19//! dozen, not 512.
20//!
21//! Off unless `ATLAS_QWEN4EXP_DUMP` names a directory. Writes
22//! `<dir>/L{layer:02}_{tag}.bin` as raw little-endian FP32, `[T, hc*H]`.
23
24use spark_runtime::gpu::{DevicePtr, GpuBackend};
25
26/// Directory from `ATLAS_QWEN4EXP_DUMP`, resolved once.
27fn dump_dir() -> Option<&'static str> {
28    static DIR: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
29    DIR.get_or_init(|| {
30        let d = std::env::var("ATLAS_QWEN4EXP_DUMP")
31            .ok()
32            .filter(|s| !s.is_empty());
33        if let Some(ref path) = d {
34            let _ = std::fs::create_dir_all(path);
35            tracing::warn!(
36                "ATLAS_QWEN4EXP_DUMP={path}: taping the mHC highway to disk. \
37                 This SYNCHRONIZES and copies D2H at every tap — a debug aid, \
38                 not a serving mode."
39            );
40        }
41        d
42    })
43    .as_deref()
44}
45
46/// One-shot: refuse to overwrite a tap that already exists.
47///
48/// `SSM_LAYER_CALL_COUNTER` is a global that never resets, so the second
49/// request of a run labels its taps L36+, the third L72+, and so on. Left
50/// alone, that means a second request silently leaves the FIRST request's
51/// L00 files in place while adding mislabelled ones — and a bisect then
52/// compares a stale tap against a fresh reference and calls it a divergence.
53///
54/// So the first prefill after startup wins and everything later is ignored.
55/// The intended use is exactly that: start the server, send one request,
56/// read the taps.
57fn claim(path: &str) -> bool {
58    !std::path::Path::new(path).exists()
59}
60
61/// Tap the FP32 highway. No-op unless the dump directory is set.
62///
63/// Synchronizes before reading, so it must never run inside CUDA-graph
64/// capture — which is already true of this model's path (`ATLAS_DEBUG_NO_GRAPH`).
65pub fn tap_highway(
66    gpu: &dyn GpuBackend,
67    streams: DevicePtr,
68    layer: usize,
69    tag: &str,
70    num_tokens: usize,
71    hc_dim: usize,
72    stream: u64,
73) {
74    let Some(dir) = dump_dir() else {
75        return;
76    };
77    let path = format!("{dir}/L{layer:02}_{tag}.bin");
78    if !claim(&path) {
79        return;
80    }
81    if gpu.synchronize(stream).is_err() {
82        return;
83    }
84    let mut raw = vec![0u8; num_tokens * hc_dim * 4];
85    if gpu.copy_d2h(streams, &mut raw).is_err() {
86        return;
87    }
88    let path = format!("{dir}/L{layer:02}_{tag}.bin");
89    if let Err(e) = std::fs::write(&path, &raw) {
90        tracing::warn!("highway tap {path}: {e}");
91    }
92}
93
94/// Tap a BF16 buffer (the embedding, a block output) the same way.
95pub fn tap_bf16(
96    gpu: &dyn GpuBackend,
97    ptr: DevicePtr,
98    layer: usize,
99    tag: &str,
100    n_elements: usize,
101    stream: u64,
102) {
103    let Some(dir) = dump_dir() else {
104        return;
105    };
106    let path = format!("{dir}/L{layer:02}_{tag}.bf16.bin");
107    if !claim(&path) {
108        return;
109    }
110    if gpu.synchronize(stream).is_err() {
111        return;
112    }
113    let mut raw = vec![0u8; n_elements * 2];
114    if gpu.copy_d2h(ptr, &mut raw).is_err() {
115        return;
116    }
117    if let Err(e) = std::fs::write(&path, &raw) {
118        tracing::warn!("highway tap {path}: {e}");
119    }
120}
121
122/// Tap an FP32 buffer of `n_elements` (the injection vector, a gate).
123pub fn tap_f32(
124    gpu: &dyn GpuBackend,
125    ptr: DevicePtr,
126    layer: usize,
127    tag: &str,
128    n_elements: usize,
129    stream: u64,
130) {
131    let Some(dir) = dump_dir() else {
132        return;
133    };
134    let path = format!("{dir}/L{layer:02}_{tag}.bin");
135    if !claim(&path) {
136        return;
137    }
138    if gpu.synchronize(stream).is_err() {
139        return;
140    }
141    let mut raw = vec![0u8; n_elements * 4];
142    if gpu.copy_d2h(ptr, &mut raw).is_err() {
143        return;
144    }
145    if let Err(e) = std::fs::write(&path, &raw) {
146        tracing::warn!("highway tap {path}: {e}");
147    }
148}
spark_model/layers/ple/dump.rs

spark_model/layers/ple/
dump.rs
