proof_stream/viz/
attention.rs

1//! Helpers for building `ProofEvent::AttentionHeatmap` from raw score matrices.
2
3use crate::events::ProofEvent;
4
5/// Maximum attention grid dimension to stream. Larger matrices are downsampled.
6pub const MAX_ATTN_DIM: usize = 64;
7
8/// Build an `AttentionHeatmap` event by sampling the attention score matrix.
9///
10/// `scores_row_major` must have length `seq_len * seq_len`.
11/// If `seq_len > MAX_ATTN_DIM` the matrix is subsampled (stride sampling).
12pub fn attention_heatmap_event(
13    layer_idx: usize,
14    head_idx: usize,
15    num_heads: usize,
16    seq_len: usize,
17    scores_row_major: &[f32],
18) -> ProofEvent {
19    let out_dim = seq_len.min(MAX_ATTN_DIM);
20    let stride = if seq_len > MAX_ATTN_DIM {
21        (seq_len + MAX_ATTN_DIM - 1) / MAX_ATTN_DIM
22    } else {
23        1
24    };
25
26    let mut sampled = Vec::with_capacity(out_dim * out_dim);
27    for r in (0..seq_len).step_by(stride).take(out_dim) {
28        for c in (0..seq_len).step_by(stride).take(out_dim) {
29            let idx = r * seq_len + c;
30            sampled.push(if idx < scores_row_major.len() {
31                scores_row_major[idx]
32            } else {
33                0.0
34            });
35        }
36    }
37
38    ProofEvent::AttentionHeatmap {
39        layer_idx,
40        head_idx,
41        num_heads,
42        seq_len: out_dim,
43        scores: sampled,
44    }
45}
46
47/// Build an `AttentionHeatmap` from raw M31 values (u32) converted to f32.
48pub fn attention_heatmap_from_u32(
49    layer_idx: usize,
50    head_idx: usize,
51    num_heads: usize,
52    seq_len: usize,
53    scores_u32: &[u32],
54) -> ProofEvent {
55    let scores_f32: Vec<f32> = scores_u32
56        .iter()
57        .map(|&v| v as f32 / 0x7fff_ffff_u32 as f32)
58        .collect();
59    attention_heatmap_event(layer_idx, head_idx, num_heads, seq_len, &scores_f32)
60}
proof_stream/viz/attention.rs

proof_stream/viz/
attention.rs
