metrale_model_layers/layers/
qsa_select.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Per-query prefill selection for the QSA indexer.
4//!
5//! Owner: model-layers (QSA).
6//! Invariants:
7//! - Only rows at global positions `>= inert_bound()` are overwritten; rows
8//!   below keep the dense attention output.
9//!
10//! A child module of `qsa` (via `#[path]`), so the indexer's private fields
11//! are reachable without widening their visibility.
12
13use anyhow::{Context, Result};
14use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
15
16use super::{QsaIndexer, QsaSeqState};
17use crate::layers::ops;
18
19impl QsaIndexer {
20    /// 2026-09-25: Per-query prefill selection for a prefill chunk. Chunk rows whose
21    /// global position (`seq_start + row`) is at or past the inert bound get
22    /// their attention context rows (pre-gate, pre-o_proj) overwritten with
23    /// attention over their selected set, read from the paged KV cache, which
24    /// must already hold every prior chunk plus this one. Rows below the bound
25    /// keep the dense output, which is identical there. `prefill_ingest` must
26    /// have run for this chunk (`prefill_inner.rs` calls it before attention).
27    /// `METRALE_QSA_NO_PREFILL_SELECT=1` makes this a no-op.
28    #[allow(clippy::too_many_arguments)]
29    pub fn prefill_select(
30        &self,
31        st: &mut QsaSeqState,
32        normed: DevicePtr,
33        q_roped: DevicePtr,
34        attn_ctx: DevicePtr,
35        k_pool: DevicePtr,
36        v_pool: DevicePtr,
37        seq_block_table: &[u32],
38        seq_start: usize,
39        num_tokens: usize,
40        nq: u32,
41        block_size: u32,
42        inv_sqrt_d: f32,
43        scratch: DevicePtr,
44        gpu: &dyn GpuBackend,
45        stream: u64,
46    ) -> Result<()> {
47        let bound = self.inert_bound(); // 2026-09-25: first selective global position
48        let total = seq_start + num_tokens;
49        if total <= bound {
50            return Ok(());
51        }
52        static S2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
53        if *S2_OFF
54            .get_or_init(|| std::env::var("METRALE_QSA_NO_PREFILL_SELECT").as_deref() == Ok("1"))
55        {
56            return Ok(());
57        }
58        let diag = {
59            static D: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
60            *D.get_or_init(|| std::env::var("METRALE_QSA_S2_DIAG").as_deref() == Ok("1"))
61        };
62        // 2026-09-25: METRALE_QSA_S2_DIAG=1: keep the dense context of the last
63        // row before the overwrite, and log cosine(dense, selected) after.
64        let q_row = nq as usize * self.hd_attn as usize;
65        let mut dense_last = Vec::new();
66        if diag {
67            dense_last = vec![0u8; q_row * 2];
68            gpu.copy_d2h_on_stream(
69                attn_ctx.offset((num_tokens - 1) * q_row * 2),
70                &mut dense_last,
71                stream,
72            )?;
73            // 2026-09-25: Norm probes: row 100, the first selective row and the
74            // last row.
75            let probe = |row: usize| -> Result<f64> {
76                let mut b = vec![0u8; q_row * 2];
77                gpu.copy_d2h_on_stream(attn_ctx.offset(row * q_row * 2), &mut b, stream)?;
78                Ok(b.chunks_exact(2)
79                    .map(|c| {
80                        let v =
81                            f32::from_bits((u16::from_le_bytes([c[0], c[1]]) as u32) << 16) as f64;
82                        v * v
83                    })
84                    .sum::<f64>()
85                    .sqrt())
86            };
87            tracing::warn!(
88                "QSA S2 DIAG norms: row100={:.3} first_sel(row {bound})={:.3} last={:.3} q_row={q_row}",
89                probe(100)?,
90                probe(bound)?,
91                probe(num_tokens - 1)?
92            );
93            let probe_at = |base: DevicePtr, row: usize| -> Result<f64> {
94                let mut b = vec![0u8; q_row * 2];
95                gpu.copy_d2h_on_stream(base.offset(row * q_row * 2), &mut b, stream)?;
96                Ok(b.chunks_exact(2)
97                    .map(|c| {
98                        let v =
99                            f32::from_bits((u16::from_le_bytes([c[0], c[1]]) as u32) << 16) as f64;
100                        v * v
101                    })
102                    .sum::<f64>()
103                    .sqrt())
104            };
105            let mut ctx_line = String::new();
106            let mut q_line = String::new();
107            for row in [128usize, 256, 512, 768, 1024, 1280, 1536, 1792, 1900, 2000] {
108                ctx_line += &format!(" {row}:{:.2}", probe_at(attn_ctx, row)?);
109            }
110            tracing::warn!("QSA S2 DIAG wide:{ctx_line}");
111            ctx_line = String::new();
112            for row in (2040..2056).step_by(2) {
113                ctx_line += &format!(" {row}:{:.2}", probe_at(attn_ctx, row)?);
114                q_line += &format!(" {row}:{:.2}", probe_at(q_roped, row)?);
115            }
116            tracing::warn!("QSA S2 DIAG ctx rows:{ctx_line}");
117            tracing::warn!("QSA S2 DIAG   q rows:{q_line}");
118        }
119        // 2026-09-25: Upload the physical block table for the full context (a
120        // selective query attends blocks from every prior chunk).
121        let pages_needed = total.div_ceil(block_size as usize);
122        anyhow::ensure!(
123            seq_block_table.len() >= pages_needed,
124            "QSA: block table has {} pages for {} tokens",
125            seq_block_table.len(),
126            pages_needed
127        );
128        let tbytes: Vec<u8> = seq_block_table[..pages_needed]
129            .iter()
130            .flat_map(|b| (*b as i32).to_le_bytes())
131            .collect();
132        gpu.copy_h2d_async(&tbytes, self.prefill_table_dev, stream)?;
133        let block_table_dev = self.prefill_table_dev;
134        // 2026-09-25: Must equal ROWS in sizes.rs `qsa_select_scratch`.
135        const ROWS: usize = 2048;
136        let ratio = self.ratio as usize;
137        let topk = self.block_topk as usize;
138        let heads = self.n_heads as usize;
139        let hd = self.hd as usize;
140        let hd_attn = self.hd_attn as usize;
141        let qkw = self.qk_width();
142        let q_row = nq as usize * hd_attn;
143
144        // 2026-09-25: Scratch layout. The per-call score stride is
145        // `total.div_ceil(ratio)`; sizes.rs reserves `max_seq_len.div_ceil(ratio)`.
146        let stride = total.div_ceil(ratio);
147        let qk_buf = scratch;
148        let qpost = scratch.offset(ROWS * qkw * 2);
149        let scores = qpost.offset(ROWS * heads * hd * 4);
150        let lists = scores.offset(ROWS * stride * 4);
151
152        let first_sel_pos = bound.max(seq_start);
153        let n_sel_total = total - first_sel_pos;
154        let mut slab = 0usize;
155        while slab < n_sel_total {
156            let rows = ROWS.min(n_sel_total - slab);
157            let first_pos = first_sel_pos + slab;
158            let first_row = first_pos - seq_start;
159
160            ops::cublas_bf16_proj_dense(
161                normed.offset(first_row * self.hidden as usize * 2),
162                self.qk_proj_w,
163                qk_buf,
164                rows as u32,
165                qkw as u32,
166                self.hidden,
167                stream,
168            )
169            .context("QSA qk projection (prefill select)")?;
170            ops::qsa_qprep_rows(
171                gpu,
172                self.k_qprep_rows_k,
173                qk_buf,
174                self.q_norm_w,
175                qpost,
176                rows as u32,
177                first_pos as u32,
178                qkw as u32,
179                self.n_heads,
180                self.hd,
181                self.rot,
182                self.theta,
183                self.eps,
184                stream,
185            )?;
186            // 2026-09-25: `n_blocks_max` is the last row's complete-block count.
187            let n_blocks_max = (first_pos + rows) / ratio;
188            // 2026-09-25: The tensor-core scorer is used when the target ships it and
189            // the shape is 4 heads x 128; METRALE_QSA_SCORE_SCALAR=1 forces the
190            // scalar kernel.
191            let tc = self.k_score_rows_tc_k.0 != 0
192                && self.n_heads == 4
193                && self.hd == 128
194                && std::env::var("METRALE_QSA_SCORE_SCALAR").as_deref() != Ok("1");
195            if tc {
196                ops::qsa_score_rows_tc(
197                    gpu,
198                    self.k_score_rows_tc_k,
199                    qpost,
200                    st.block_keys,
201                    scores,
202                    rows as u32,
203                    n_blocks_max as u32,
204                    first_pos as u32,
205                    stride as u32,
206                    self.ratio,
207                    stream,
208                )?;
209            } else {
210                ops::qsa_score_rows(
211                    gpu,
212                    self.k_score_rows_k,
213                    qpost,
214                    st.block_keys,
215                    scores,
216                    rows as u32,
217                    n_blocks_max as u32,
218                    first_pos as u32,
219                    stride as u32,
220                    self.ratio,
221                    self.n_heads,
222                    self.hd,
223                    stream,
224                )?;
225            }
226
227            // 2026-09-25: Host top-k per row; the D2H is ordered after the scoring
228            // kernel on `stream`. Larger score first, lower index on ties.
229            let mut raw = vec![0u8; rows * stride * 4];
230            gpu.copy_d2h_on_stream(scores, &mut raw, stream)?;
231            let sc: Vec<f32> = raw
232                .chunks_exact(4)
233                .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
234                .collect();
235            let mut host_lists = vec![0u8; rows * topk * 4];
236            for r in 0..rows {
237                let complete = (first_pos + r + 1) / ratio;
238                let row_sc = &sc[r * stride..r * stride + complete];
239                let mut order: Vec<u32> = (0..complete as u32).collect();
240                order.sort_by(|&a, &b| {
241                    row_sc[b as usize]
242                        .partial_cmp(&row_sc[a as usize])
243                        .unwrap_or(std::cmp::Ordering::Equal)
244                        .then(a.cmp(&b))
245                });
246                for (i, b) in order[..topk].iter().enumerate() {
247                    host_lists[(r * topk + i) * 4..(r * topk + i) * 4 + 4]
248                        .copy_from_slice(&(*b as i32).to_le_bytes());
249                }
250            }
251            gpu.copy_h2d_async(&host_lists, lists, stream)?;
252
253            ops::qsa_prefill_attn(
254                gpu,
255                self.k_prefill_attn_k,
256                q_roped.offset(first_row * q_row * 2),
257                k_pool,
258                v_pool,
259                block_table_dev,
260                lists,
261                attn_ctx.offset(first_row * q_row * 2),
262                rows as u32,
263                first_pos as u32,
264                topk as u32,
265                self.ratio,
266                block_size,
267                nq,
268                self.nkv_attn,
269                self.hd_attn,
270                inv_sqrt_d,
271                stream,
272            )?;
273            slab += rows;
274        }
275        if diag {
276            let mut sel_last = vec![0u8; q_row * 2];
277            gpu.copy_d2h_on_stream(
278                attn_ctx.offset((num_tokens - 1) * q_row * 2),
279                &mut sel_last,
280                stream,
281            )?;
282            let f = |b: &[u8]| -> Vec<f32> {
283                b.chunks_exact(2)
284                    .map(|c| f32::from_bits((u16::from_le_bytes([c[0], c[1]]) as u32) << 16))
285                    .collect()
286            };
287            let (a, b) = (f(&dense_last), f(&sel_last));
288            let dot: f64 = a.iter().zip(&b).map(|(x, y)| *x as f64 * *y as f64).sum();
289            let na: f64 = a.iter().map(|x| (*x as f64).powi(2)).sum::<f64>().sqrt();
290            let nb: f64 = b.iter().map(|x| (*x as f64).powi(2)).sum::<f64>().sqrt();
291            tracing::warn!(
292                "QSA S2 DIAG: last-row ctx dense-vs-selected cos={:.6} |dense|={:.3} |sel|={:.3}",
293                dot / (na * nb).max(1e-30),
294                na,
295                nb
296            );
297        }
298        tracing::debug!(
299            "QSA prefill select: {} selective rows over {} tokens",
300            n_sel_total,
301            num_tokens
302        );
303        Ok(())
304    }
305}
metrale_model_layers/layers/qsa_select.rs

metrale_model_layers/layers/
qsa_select.rs
