spark_model/layers/ops/
qsa.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Launchers for the Qwen3.8-Flash-Next QSA indexer kernels
4//! (`qsa_indexer.cu`): block-key pooling, decode-query prep, block scoring
5//! and the selected-token K/V gather. See the .cu header for the semantics
6//! and the scratch-as-paged-cache trick that lets the EXISTING paged decode
7//! attention consume the selection.
8
9use anyhow::Result;
10use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
11use spark_runtime::kernel_args::KernelLaunch;
12
13/// Pool `n_new` freshly complete blocks starting at `first_block`:
14/// mean over `ratio` raw keys -> RMSNorm*(1+w) -> rope at block-start pos.
15#[allow(clippy::too_many_arguments)]
16pub fn qsa_block_pool(
17    gpu: &dyn GpuBackend,
18    kernel: KernelHandle,
19    raw_keys: DevicePtr,
20    k_norm_w: DevicePtr,
21    block_keys: DevicePtr,
22    first_block: u32,
23    n_new: u32,
24    ratio: u32,
25    hd: u32,
26    rot: u32,
27    theta: f32,
28    eps: f32,
29    stream: u64,
30) -> Result<()> {
31    if n_new == 0 {
32        return Ok(());
33    }
34    KernelLaunch::new(gpu, kernel)
35        .grid([n_new, 1, 1])
36        .block([hd, 1, 1])
37        .shared_mem((hd + 32) * 4)
38        .arg_ptr(raw_keys)
39        .arg_ptr(k_norm_w)
40        .arg_ptr(block_keys)
41        .arg_u32(first_block)
42        .arg_u32(ratio)
43        .arg_u32(hd)
44        .arg_u32(rot)
45        .arg_f32(theta)
46        .arg_f32(eps)
47        .launch(stream)
48}
49
50/// One decode query: per head, RMSNorm*(1+w) + partial rope at `pos` -> FP32.
51#[allow(clippy::too_many_arguments)]
52pub fn qsa_qprep(
53    gpu: &dyn GpuBackend,
54    kernel: KernelHandle,
55    q_in: DevicePtr,
56    q_norm_w: DevicePtr,
57    q_out: DevicePtr,
58    n_heads: u32,
59    hd: u32,
60    rot: u32,
61    pos: u32,
62    theta: f32,
63    eps: f32,
64    stream: u64,
65) -> Result<()> {
66    KernelLaunch::new(gpu, kernel)
67        .grid([n_heads, 1, 1])
68        .block([hd, 1, 1])
69        .shared_mem((hd + 32) * 4)
70        .arg_ptr(q_in)
71        .arg_ptr(q_norm_w)
72        .arg_ptr(q_out)
73        .arg_u32(hd)
74        .arg_u32(rot)
75        .arg_u32(pos)
76        .arg_f32(theta)
77        .arg_f32(eps)
78        .launch(stream)
79}
80
81/// `scores[b]` = sum_h relu(q_h . k_b) / sqrt(hd) over `n_blocks` blocks.
82#[allow(clippy::too_many_arguments)]
83pub fn qsa_score(
84    gpu: &dyn GpuBackend,
85    kernel: KernelHandle,
86    q: DevicePtr,
87    block_keys: DevicePtr,
88    scores: DevicePtr,
89    n_blocks: u32,
90    n_heads: u32,
91    hd: u32,
92    stream: u64,
93) -> Result<()> {
94    KernelLaunch::new(gpu, kernel)
95        .grid([n_blocks, 1, 1])
96        .block([hd, 1, 1])
97        .shared_mem(32 * 4)
98        .arg_ptr(q)
99        .arg_ptr(block_keys)
100        .arg_ptr(scores)
101        .arg_u32(n_heads)
102        .arg_u32(hd)
103        .launch(stream)
104}
105
106/// Pack the selected tokens' K/V rows into contiguous NHD scratch.
107#[allow(clippy::too_many_arguments)]
108pub fn qsa_gather(
109    gpu: &dyn GpuBackend,
110    kernel: KernelHandle,
111    k_cache: DevicePtr,
112    v_cache: DevicePtr,
113    block_table: DevicePtr,
114    sel: DevicePtr,
115    k_out: DevicePtr,
116    v_out: DevicePtr,
117    n_sel: u32,
118    block_size: u32,
119    nkv: u32,
120    hd: u32,
121    stream: u64,
122) -> Result<()> {
123    KernelLaunch::new(gpu, kernel)
124        .grid([n_sel, 1, 1])
125        .block([256, 1, 1])
126        .arg_ptr(k_cache)
127        .arg_ptr(v_cache)
128        .arg_ptr(block_table)
129        .arg_ptr(sel)
130        .arg_ptr(k_out)
131        .arg_ptr(v_out)
132        .arg_u32(block_size)
133        .arg_u32(nkv)
134        .arg_u32(hd)
135        .launch(stream)
136}
137
138/// Stage 2: per-row q prep for a contiguous selective row range.
139#[allow(clippy::too_many_arguments)]
140pub fn qsa_qprep_rows(
141    gpu: &dyn GpuBackend,
142    kernel: KernelHandle,
143    qk: DevicePtr,
144    q_norm_w: DevicePtr,
145    q_out: DevicePtr,
146    rows: u32,
147    first_pos: u32,
148    qkw: u32,
149    n_heads: u32,
150    hd: u32,
151    rot: u32,
152    theta: f32,
153    eps: f32,
154    stream: u64,
155) -> Result<()> {
156    KernelLaunch::new(gpu, kernel)
157        .grid([rows, n_heads, 1])
158        .block([hd, 1, 1])
159        .shared_mem((hd + 32) * 4)
160        .arg_ptr(qk)
161        .arg_ptr(q_norm_w)
162        .arg_ptr(q_out)
163        .arg_u32(first_pos)
164        .arg_u32(qkw)
165        .arg_u32(n_heads)
166        .arg_u32(hd)
167        .arg_u32(rot)
168        .arg_f32(theta)
169        .arg_f32(eps)
170        .launch(stream)
171}
172
173/// Tensor-core `qsa_score_rows` (split-q). Geometry differs from the
174/// per-(row,block) kernel: one CTA covers 16 rows x 64 blocks with 8 warps,
175/// so the launch count drops from rows*blocks to ~1/1000th of that.
176///
177/// Measured 39.010 → 1.831 ms average per call under nsys on a 28K prefill
178/// (qwen4_exp, GB10): ~21x. Selection-equivalent, not bit-exact — split-q
179/// carries ~17 mantissa bits, and the consumer is a top-k.
180#[allow(clippy::too_many_arguments)]
181pub fn qsa_score_rows_tc(
182    gpu: &dyn GpuBackend,
183    kernel: KernelHandle,
184    q: DevicePtr,
185    block_keys: DevicePtr,
186    scores: DevicePtr,
187    rows: u32,
188    n_blocks_max: u32,
189    first_pos: u32,
190    score_stride: u32,
191    ratio: u32,
192    stream: u64,
193) -> Result<()> {
194    KernelLaunch::new(gpu, kernel)
195        .grid([rows.div_ceil(16), n_blocks_max.div_ceil(64), 1])
196        .block([256, 1, 1])
197        .arg_ptr(q)
198        .arg_ptr(block_keys)
199        .arg_ptr(scores)
200        .arg_u32(first_pos)
201        .arg_u32(score_stride)
202        .arg_u32(ratio)
203        .arg_u32(n_blocks_max)
204        .launch(stream)
205}
206
207/// Stage 2: per-row block scores, -inf beyond each row's complete count.
208#[allow(clippy::too_many_arguments)]
209pub fn qsa_score_rows(
210    gpu: &dyn GpuBackend,
211    kernel: KernelHandle,
212    q: DevicePtr,
213    block_keys: DevicePtr,
214    scores: DevicePtr,
215    rows: u32,
216    n_blocks_max: u32,
217    first_pos: u32,
218    score_stride: u32,
219    ratio: u32,
220    n_heads: u32,
221    hd: u32,
222    stream: u64,
223) -> Result<()> {
224    KernelLaunch::new(gpu, kernel)
225        .grid([rows, n_blocks_max, 1])
226        .block([hd, 1, 1])
227        .shared_mem(32 * 4)
228        .arg_ptr(q)
229        .arg_ptr(block_keys)
230        .arg_ptr(scores)
231        .arg_u32(first_pos)
232        .arg_u32(score_stride)
233        .arg_u32(ratio)
234        .arg_u32(n_heads)
235        .arg_u32(hd)
236        .launch(stream)
237}
238
239/// Stage 2: per-row selected-set attention, overwriting the context rows.
240#[allow(clippy::too_many_arguments)]
241pub fn qsa_prefill_attn(
242    gpu: &dyn GpuBackend,
243    kernel: KernelHandle,
244    q: DevicePtr,
245    k_cache: DevicePtr,
246    v_cache: DevicePtr,
247    block_table: DevicePtr,
248    lists: DevicePtr,
249    attn_out: DevicePtr,
250    rows: u32,
251    first_pos: u32,
252    topk: u32,
253    ratio: u32,
254    block_size: u32,
255    nq: u32,
256    nkv: u32,
257    hd: u32,
258    inv_sqrt_d: f32,
259    stream: u64,
260) -> Result<()> {
261    // 8 warps x [hd] acc partials + m/l per warp.
262    let smem = (8 * hd + 16) * 4;
263    KernelLaunch::new(gpu, kernel)
264        .grid([rows, nq, 1])
265        .block([256, 1, 1])
266        .shared_mem(smem)
267        .arg_ptr(q)
268        .arg_ptr(k_cache)
269        .arg_ptr(v_cache)
270        .arg_ptr(block_table)
271        .arg_ptr(lists)
272        .arg_ptr(attn_out)
273        .arg_u32(first_pos)
274        .arg_u32(topk)
275        .arg_u32(ratio)
276        .arg_u32(block_size)
277        .arg_u32(nq)
278        .arg_u32(nkv)
279        .arg_u32(hd)
280        .arg_f32(inv_sqrt_d)
281        .launch(stream)
282}
spark_model/layers/ops/qsa.rs

spark_model/layers/ops/
qsa.rs
