metrale_cache/kv_cache/paged_impl/
diagnostics.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: `PagedKvCache` BF16 fingerprint diagnostics: per-layer region
4//! checksums and per-block sums of squares, read back from the device.
5//!
6//! Owner: cache.
7//! Invariants: diagnostic only; nothing here writes the cache. Events keep
8//! the parent module's log target.
9
10use super::PagedKvCache;
11
12impl PagedKvCache {
13    /// 2026-09-25: Sum, sum of squares and sum of absolute values of a
14    /// little-endian BF16 buffer, for the two fingerprint diagnostics below.
15    /// A BF16 value is the top 16 bits of an f32.
16    fn bf16_reductions(buf: &[u8]) -> (f64, f64, f64) {
17        let (mut sum, mut ssq, mut sabs) = (0f64, 0f64, 0f64);
18        for c in buf.chunks_exact(2) {
19            let bits = u16::from_le_bytes([c[0], c[1]]);
20            let v = f32::from_bits((bits as u32) << 16) as f64;
21            sum += v;
22            ssq += v * v;
23            sabs += v.abs();
24        }
25        (sum, ssq, sabs)
26    }
27
28    /// 2026-09-25: Diagnostic. Synchronises `stream`, then logs a K and V
29    /// (sum, ssq, sabs) line per BF16 layer for each of two regions of
30    /// `blocks`: `[0, boundary_idx)` ("prefix") and the rest ("suffix").
31    /// Per-layer lines keep a local divergence from cancelling in a global
32    /// sum. Non-BF16 layers are skipped, with a warning when layer 0 is one;
33    /// a block whose copy fails is skipped. The model engine calls it only
34    /// under `METRALE_SSM_SAVE_DUMP`.
35    pub fn debug_kv_checksum_per_layer(
36        &self,
37        blocks: &[u32],
38        boundary_idx: usize,
39        gpu: &dyn metrale_gpu_runtime::gpu::GpuBackend,
40        stream: u64,
41        tag: &str,
42    ) {
43        gpu.synchronize(stream).ok();
44        let boundary = boundary_idx.min(blocks.len());
45        let regions: [(&str, &[u32]); 2] = [
46            ("prefix", &blocks[..boundary]),
47            ("suffix", &blocks[boundary..]),
48        ];
49        for (li, layer) in self.layers.iter().enumerate() {
50            if layer.dtype != super::KvCacheDtype::Bf16 {
51                if li == 0 {
52                    tracing::warn!(
53                        target: "metrale_cache::kv_cache::paged_impl",
54                        "METRALE_KV_CKSUM[{tag}] layer 0 dtype={:?} != bf16 — probe \
55                         only decodes BF16; skipping",
56                        layer.dtype
57                    );
58                }
59                continue;
60            }
61            // 2026-09-25: A BF16 layer is symmetric, so the K stride is also
62            // the V stride.
63            let nbytes = layer.k_block_stride;
64            for (rname, rblocks) in &regions {
65                let (mut k_sum, mut k_ssq, mut k_sabs) = (0f64, 0f64, 0f64);
66                let (mut v_sum, mut v_ssq, mut v_sabs) = (0f64, 0f64, 0f64);
67                for &blk in *rblocks {
68                    let mut kb = vec![0u8; nbytes];
69                    let mut vb = vec![0u8; nbytes];
70                    if gpu.copy_d2h(self.k_cache_ptr(li, blk), &mut kb).is_err()
71                        || gpu.copy_d2h(self.v_cache_ptr(li, blk), &mut vb).is_err()
72                    {
73                        continue;
74                    }
75                    let (ks, kq, ka) = Self::bf16_reductions(&kb);
76                    let (vs, vq, va) = Self::bf16_reductions(&vb);
77                    k_sum += ks;
78                    k_ssq += kq;
79                    k_sabs += ka;
80                    v_sum += vs;
81                    v_ssq += vq;
82                    v_sabs += va;
83                }
84                tracing::warn!(
85                    target: "metrale_cache::kv_cache::paged_impl",
86                    "METRALE_KV_CKSUM[{tag}] L{li} {rname} nblk={} \
87                     k_sum={k_sum:.4} k_ssq={k_ssq:.4} k_sabs={k_sabs:.4} \
88                     v_sum={v_sum:.4} v_ssq={v_ssq:.4} v_sabs={v_sabs:.4}",
89                    rblocks.len(),
90                );
91            }
92        }
93    }
94
95    /// 2026-09-25: Diagnostic. Synchronises `stream`, then logs one line per
96    /// entry of `blocks`, in order: logical index, physical block, K and V sum
97    /// of squares for layer `layer_idx`. Unlike a region sum this shows a
98    /// wrong block-to-position mapping. Does nothing for a non-BF16 layer.
99    pub fn debug_kv_per_block(
100        &self,
101        layer_idx: usize,
102        blocks: &[u32],
103        gpu: &dyn metrale_gpu_runtime::gpu::GpuBackend,
104        stream: u64,
105        tag: &str,
106    ) {
107        gpu.synchronize(stream).ok();
108        let layer = &self.layers[layer_idx];
109        if layer.dtype != super::KvCacheDtype::Bf16 {
110            return;
111        }
112        // 2026-09-25: A BF16 layer is symmetric, so the K stride is also the V
113        // stride.
114        let nbytes = layer.k_block_stride;
115        for (li, &blk) in blocks.iter().enumerate() {
116            let mut kb = vec![0u8; nbytes];
117            let mut vb = vec![0u8; nbytes];
118            if gpu
119                .copy_d2h(self.k_cache_ptr(layer_idx, blk), &mut kb)
120                .is_err()
121                || gpu
122                    .copy_d2h(self.v_cache_ptr(layer_idx, blk), &mut vb)
123                    .is_err()
124            {
125                continue;
126            }
127            let (_, k_ssq, _) = Self::bf16_reductions(&kb);
128            let (_, v_ssq, _) = Self::bf16_reductions(&vb);
129            tracing::warn!(
130                target: "metrale_cache::kv_cache::paged_impl",
131                "METRALE_KVBLK[{tag}] L{layer_idx} logical={li} phys={blk} \
132                 k_ssq={k_ssq:.4} v_ssq={v_ssq:.4}"
133            );
134        }
135    }
136}
metrale_cache/kv_cache/paged_impl/diagnostics.rs

metrale_cache/kv_cache/paged_impl/
diagnostics.rs
