metrale_storage/
attention_ref.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Host reference for the tiled online-softmax decode attention of
4//! `kernels/paged_decode_attn_tiled.cu` and `kernels/attention_finalize.cu`.
5//!
6//! Owner: metrale-storage.
7//! Invariants: none beyond the types.
8
9use half::bf16;
10
11#[inline]
12fn b2f(x: bf16) -> f32 {
13    x.to_f32()
14}
15
16/// 2026-09-25: Running state. `m` and `l` are `[num_seqs, num_q_heads]`, `o` is
17/// `[num_seqs, num_q_heads, head_dim]`; `new` sets m = -inf and l = o = 0.
18#[derive(Clone)]
19pub struct AttnState {
20    pub m: Vec<f32>,
21    pub l: Vec<f32>,
22    pub o: Vec<f32>,
23}
24
25impl AttnState {
26    pub fn new(num_seqs: usize, num_q_heads: usize, head_dim: usize) -> Self {
27        let n_q = num_seqs * num_q_heads;
28        Self {
29            m: vec![f32::NEG_INFINITY; n_q],
30            l: vec![0.0; n_q],
31            o: vec![0.0; n_q * head_dim],
32        }
33    }
34}
35
36/// 2026-09-25: One tile update over the paged layout `[num_blocks, block_size,
37/// num_kv_heads, head_dim]`, with no causal mask. `tile_blocks` is
38/// `[num_seqs, tile_capacity]`, `tile_block_counts` is `[num_seqs]`. It runs
39/// the kernel's per-token recurrence in fp32, but sums each dot product
40/// sequentially and uses `exp` where the kernel uses a warp-tree sum and
41/// `__expf`.
42#[allow(clippy::too_many_arguments)]
43pub fn step_tile_ref(
44    state: &mut AttnState,
45    q: &[bf16],
46    k_pool: &[bf16],
47    v_pool: &[bf16],
48    tile_blocks: &[i32],
49    tile_block_counts: &[i32],
50    num_seqs: usize,
51    num_q_heads: usize,
52    num_kv_heads: usize,
53    head_dim: usize,
54    block_size: usize,
55    tile_capacity: usize,
56    gqa_ratio: usize,
57) {
58    let inv_sqrt_d = 1.0_f32 / (head_dim as f32).sqrt();
59    let kv_token_stride = num_kv_heads * head_dim;
60    for seq in 0..num_seqs {
61        let n_blocks = tile_block_counts[seq] as usize;
62        for qh in 0..num_q_heads {
63            let kh = qh / gqa_ratio;
64            let q_off = (seq * num_q_heads + qh) * head_dim;
65            let m_idx = seq * num_q_heads + qh;
66            let o_off = m_idx * head_dim;
67            let mut m_run = state.m[m_idx];
68            let mut l_run = state.l[m_idx];
69            let mut o_run: Vec<f32> = state.o[o_off..o_off + head_dim].to_vec();
70            for b in 0..n_blocks {
71                let blk_id = tile_blocks[seq * tile_capacity + b] as usize;
72                let blk_base = blk_id * block_size * kv_token_stride;
73                for t in 0..block_size {
74                    let kv_base = blk_base + t * kv_token_stride + kh * head_dim;
75                    let mut dot = 0.0_f32;
76                    for i in 0..head_dim {
77                        dot += b2f(q[q_off + i]) * b2f(k_pool[kv_base + i]);
78                    }
79                    let logit = dot * inv_sqrt_d;
80                    let m_new = m_run.max(logit);
81                    let scale_old = (m_run - m_new).exp();
82                    let scale_new = (logit - m_new).exp();
83                    let l_new = l_run * scale_old + scale_new;
84                    for i in 0..head_dim {
85                        let v = b2f(v_pool[kv_base + i]);
86                        o_run[i] = o_run[i] * scale_old + v * scale_new;
87                    }
88                    m_run = m_new;
89                    l_run = l_new;
90                }
91            }
92            state.m[m_idx] = m_run;
93            state.l[m_idx] = l_run;
94            state.o[o_off..o_off + head_dim].copy_from_slice(&o_run);
95        }
96    }
97}
98
99pub fn finalize_ref(
100    state: &AttnState,
101    num_seqs: usize,
102    num_q_heads: usize,
103    head_dim: usize,
104) -> Vec<bf16> {
105    let mut out = vec![bf16::from_f32(0.0); num_seqs * num_q_heads * head_dim];
106    for seq in 0..num_seqs {
107        for qh in 0..num_q_heads {
108            let idx = seq * num_q_heads + qh;
109            let l = state.l[idx];
110            let inv_l = if l > 0.0 { 1.0 / l } else { 0.0 };
111            for i in 0..head_dim {
112                let val = state.o[idx * head_dim + i] * inv_l;
113                out[idx * head_dim + i] = bf16::from_f32(val);
114            }
115        }
116    }
117    out
118}
metrale_storage/attention_ref.rs

metrale_storage/
attention_ref.rs
