spark_model/layers/
qsa_select.rs

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

spark_model/layers/
qsa_select.rs
