metrale_gpu_runtime/buffers/
debug_checksum.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `BufferArena::debug_buffer_checksum`, a debug fingerprint of the arena's buffers.
4//!
5//! Owner: gpu-runtime.
6//! Invariants: none beyond the types.
7
8use super::BufferArena;
9use crate::gpu::{DevicePtr, GpuBackend};
10
11impl BufferArena {
12    /// 2026-09-25: Synchronize `stream`, then log at warn level the sum, sum of
13    /// squares and sum of absolute values of each listed buffer read as f32
14    /// lanes over its whole allocation (non-finite lanes skipped), tagged with
15    /// `tag`. A failed synchronize or copy is ignored; a failed copy skips that
16    /// buffer. The model engine calls it at the first decode step when
17    /// `ssm_save_dump` is set.
18    pub fn debug_buffer_checksum(&self, gpu: &dyn GpuBackend, stream: u64, tag: &str) {
19        gpu.synchronize(stream).ok();
20        let probe = |name: &str, ptr: DevicePtr, bytes: usize| {
21            let mut hb = vec![0u8; bytes];
22            if gpu.copy_d2h(ptr, &mut hb).is_err() {
23                return;
24            }
25            let (mut sum, mut ssq, mut sabs) = (0f64, 0f64, 0f64);
26            for c in hb.chunks_exact(4) {
27                let v = f32::from_le_bytes([c[0], c[1], c[2], c[3]]) as f64;
28                if v.is_finite() {
29                    sum += v;
30                    ssq += v * v;
31                    sabs += v.abs();
32                }
33            }
34            tracing::warn!(
35                "METRALE_BUF_CKSUM[{tag}] {name} bytes={bytes} sum={sum:.6} ssq={ssq:.6} sabs={sabs:.6}"
36            );
37        };
38        probe(
39            "hidden_states",
40            self.hidden_states,
41            self.sizes.hidden_states,
42        );
43        probe("residual", self.residual, self.sizes.residual);
44        probe("norm_output", self.norm_output, self.sizes.norm_output);
45        probe("qkv_output", self.qkv_output, self.sizes.qkv_output);
46        probe("attn_output", self.attn_output, self.sizes.attn_output);
47        probe("gate_logits", self.gate_logits, self.sizes.gate_logits);
48        probe("moe_output", self.moe_output, self.sizes.moe_output);
49        probe("ssm_qkvz", self.ssm_qkvz, self.sizes.ssm_qkvz);
50        probe("ssm_ba", self.ssm_ba, self.sizes.ssm_ba);
51        probe(
52            "ssm_deinterleaved",
53            self.ssm_deinterleaved,
54            self.sizes.ssm_deinterleaved,
55        );
56        probe("ssm_gates", self.ssm_gates, self.sizes.ssm_gates);
57        probe(
58            "ssm_conv_out_f32",
59            self.ssm_conv_out_f32,
60            self.sizes.ssm_conv_out_f32,
61        );
62        probe(
63            "expert_gate_out",
64            self.expert_gate_out,
65            self.sizes.expert_gate_out,
66        );
67        probe(
68            "expert_up_out",
69            self.expert_up_out,
70            self.sizes.expert_up_out,
71        );
72        probe(
73            "expert_down_out",
74            self.expert_down_out,
75            self.sizes.expert_down_out,
76        );
77        probe(
78            "splitk_workspace",
79            self.splitk_workspace,
80            self.sizes.splitk_workspace,
81        );
82    }
83}
metrale_gpu_runtime/buffers/debug_checksum.rs

metrale_gpu_runtime/buffers/
debug_checksum.rs
