metrale_model_layers/layers/ops/
qsa.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the Qwen3.8-Flash-Next QSA indexer kernels
4//! (`kernels/gb10/qwen3.8-flash-next/nvfp4/qsa_indexer.cu`): block-key pooling,
5//! query prep, block scoring, the selected-token K/V gather, and the per-row
6//! prefill selection passes. The gather writes a contiguous scratch that the
7//! paged decode attention reads through an identity block table.
8//!
9//! Owner: model-layers ops.
10//! Invariants: none beyond the types.
11
12use anyhow::Result;
13use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
14use metrale_gpu_runtime::kernel_args::KernelLaunch;
15
16/// 2026-09-25: Pool `n_new` newly complete blocks starting at `first_block`: the
17/// mean of `ratio` raw keys, then RMSNorm scaled by `1 + w`, then rope at the
18/// block's first token position. Returns without launching when `n_new == 0`.
19#[allow(clippy::too_many_arguments)]
20pub fn qsa_block_pool(
21    gpu: &dyn GpuBackend,
22    kernel: KernelHandle,
23    raw_keys: DevicePtr,
24    k_norm_w: DevicePtr,
25    block_keys: DevicePtr,
26    first_block: u32,
27    n_new: u32,
28    ratio: u32,
29    hd: u32,
30    rot: u32,
31    theta: f32,
32    eps: f32,
33    stream: u64,
34) -> Result<()> {
35    if n_new == 0 {
36        return Ok(());
37    }
38    KernelLaunch::new(gpu, kernel)
39        .grid([n_new, 1, 1])
40        .block([hd, 1, 1])
41        .shared_mem((hd + 32) * 4)
42        .arg_ptr(raw_keys)
43        .arg_ptr(k_norm_w)
44        .arg_ptr(block_keys)
45        .arg_u32(first_block)
46        .arg_u32(ratio)
47        .arg_u32(hd)
48        .arg_u32(rot)
49        .arg_f32(theta)
50        .arg_f32(eps)
51        .launch(stream)
52}
53
54/// 2026-09-25: One decode query: per head, RMSNorm scaled by `1 + w`, then
55/// partial rope at `pos`; FP32 output.
56#[allow(clippy::too_many_arguments)]
57pub fn qsa_qprep(
58    gpu: &dyn GpuBackend,
59    kernel: KernelHandle,
60    q_in: DevicePtr,
61    q_norm_w: DevicePtr,
62    q_out: DevicePtr,
63    n_heads: u32,
64    hd: u32,
65    rot: u32,
66    pos: u32,
67    theta: f32,
68    eps: f32,
69    stream: u64,
70) -> Result<()> {
71    KernelLaunch::new(gpu, kernel)
72        .grid([n_heads, 1, 1])
73        .block([hd, 1, 1])
74        .shared_mem((hd + 32) * 4)
75        .arg_ptr(q_in)
76        .arg_ptr(q_norm_w)
77        .arg_ptr(q_out)
78        .arg_u32(hd)
79        .arg_u32(rot)
80        .arg_u32(pos)
81        .arg_f32(theta)
82        .arg_f32(eps)
83        .launch(stream)
84}
85
86/// 2026-09-25: `scores[b] = sum_h relu(q_h . k_b) / sqrt(hd)` for `n_blocks`
87/// blocks.
88#[allow(clippy::too_many_arguments)]
89pub fn qsa_score(
90    gpu: &dyn GpuBackend,
91    kernel: KernelHandle,
92    q: DevicePtr,
93    block_keys: DevicePtr,
94    scores: DevicePtr,
95    n_blocks: u32,
96    n_heads: u32,
97    hd: u32,
98    stream: u64,
99) -> Result<()> {
100    KernelLaunch::new(gpu, kernel)
101        .grid([n_blocks, 1, 1])
102        .block([hd, 1, 1])
103        .shared_mem(32 * 4)
104        .arg_ptr(q)
105        .arg_ptr(block_keys)
106        .arg_ptr(scores)
107        .arg_u32(n_heads)
108        .arg_u32(hd)
109        .launch(stream)
110}
111
112/// 2026-09-25: Copy the selected tokens' K/V rows into contiguous NHD scratch:
113/// slot `i` holds position `sel[i]`.
114#[allow(clippy::too_many_arguments)]
115pub fn qsa_gather(
116    gpu: &dyn GpuBackend,
117    kernel: KernelHandle,
118    k_cache: DevicePtr,
119    v_cache: DevicePtr,
120    block_table: DevicePtr,
121    sel: DevicePtr,
122    k_out: DevicePtr,
123    v_out: DevicePtr,
124    n_sel: u32,
125    block_size: u32,
126    nkv: u32,
127    hd: u32,
128    stream: u64,
129) -> Result<()> {
130    KernelLaunch::new(gpu, kernel)
131        .grid([n_sel, 1, 1])
132        .block([256, 1, 1])
133        .arg_ptr(k_cache)
134        .arg_ptr(v_cache)
135        .arg_ptr(block_table)
136        .arg_ptr(sel)
137        .arg_ptr(k_out)
138        .arg_ptr(v_out)
139        .arg_u32(block_size)
140        .arg_u32(nkv)
141        .arg_u32(hd)
142        .launch(stream)
143}
144
145/// 2026-09-25: [`qsa_qprep`] for `rows` prefill rows at positions
146/// `first_pos + row`, reading q from rows of the indexer projection `qk`.
147#[allow(clippy::too_many_arguments)]
148pub fn qsa_qprep_rows(
149    gpu: &dyn GpuBackend,
150    kernel: KernelHandle,
151    qk: DevicePtr,
152    q_norm_w: DevicePtr,
153    q_out: DevicePtr,
154    rows: u32,
155    first_pos: u32,
156    qkw: u32,
157    n_heads: u32,
158    hd: u32,
159    rot: u32,
160    theta: f32,
161    eps: f32,
162    stream: u64,
163) -> Result<()> {
164    KernelLaunch::new(gpu, kernel)
165        .grid([rows, n_heads, 1])
166        .block([hd, 1, 1])
167        .shared_mem((hd + 32) * 4)
168        .arg_ptr(qk)
169        .arg_ptr(q_norm_w)
170        .arg_ptr(q_out)
171        .arg_u32(first_pos)
172        .arg_u32(qkw)
173        .arg_u32(n_heads)
174        .arg_u32(hd)
175        .arg_u32(rot)
176        .arg_f32(theta)
177        .arg_f32(eps)
178        .launch(stream)
179}
180
181/// 2026-09-25: Tensor-core [`qsa_score_rows`]: one block of 8 warps covers 16
182/// rows by 64 blocks instead of one block per (row, block). q is split into
183/// BF16 high and low halves for the MMA, so the scores are not bit-identical to
184/// [`qsa_score_rows`]; they feed a top-k.
185#[allow(clippy::too_many_arguments)]
186pub fn qsa_score_rows_tc(
187    gpu: &dyn GpuBackend,
188    kernel: KernelHandle,
189    q: DevicePtr,
190    block_keys: DevicePtr,
191    scores: DevicePtr,
192    rows: u32,
193    n_blocks_max: u32,
194    first_pos: u32,
195    score_stride: u32,
196    ratio: u32,
197    stream: u64,
198) -> Result<()> {
199    KernelLaunch::new(gpu, kernel)
200        .grid([rows.div_ceil(16), n_blocks_max.div_ceil(64), 1])
201        .block([256, 1, 1])
202        .arg_ptr(q)
203        .arg_ptr(block_keys)
204        .arg_ptr(scores)
205        .arg_u32(first_pos)
206        .arg_u32(score_stride)
207        .arg_u32(ratio)
208        .arg_u32(n_blocks_max)
209        .launch(stream)
210}
211
212/// 2026-09-25: Per-row block scores for prefill rows; blocks at or past a row's
213/// complete-block count score `-inf`.
214#[allow(clippy::too_many_arguments)]
215pub fn qsa_score_rows(
216    gpu: &dyn GpuBackend,
217    kernel: KernelHandle,
218    q: DevicePtr,
219    block_keys: DevicePtr,
220    scores: DevicePtr,
221    rows: u32,
222    n_blocks_max: u32,
223    first_pos: u32,
224    score_stride: u32,
225    ratio: u32,
226    n_heads: u32,
227    hd: u32,
228    stream: u64,
229) -> Result<()> {
230    KernelLaunch::new(gpu, kernel)
231        .grid([rows, n_blocks_max, 1])
232        .block([hd, 1, 1])
233        .shared_mem(32 * 4)
234        .arg_ptr(q)
235        .arg_ptr(block_keys)
236        .arg_ptr(scores)
237        .arg_u32(first_pos)
238        .arg_u32(score_stride)
239        .arg_u32(ratio)
240        .arg_u32(n_heads)
241        .arg_u32(hd)
242        .launch(stream)
243}
244
245/// 2026-09-25: Per (row, q head), attention over exactly the selected blocks in
246/// `lists` plus the incomplete tail, read from the paged cache; the result
247/// overwrites that row's context in `attn_out`.
248#[allow(clippy::too_many_arguments)]
249pub fn qsa_prefill_attn(
250    gpu: &dyn GpuBackend,
251    kernel: KernelHandle,
252    q: DevicePtr,
253    k_cache: DevicePtr,
254    v_cache: DevicePtr,
255    block_table: DevicePtr,
256    lists: DevicePtr,
257    attn_out: DevicePtr,
258    rows: u32,
259    first_pos: u32,
260    topk: u32,
261    ratio: u32,
262    block_size: u32,
263    nq: u32,
264    nkv: u32,
265    hd: u32,
266    inv_sqrt_d: f32,
267    stream: u64,
268) -> Result<()> {
269    // 2026-09-25: Per-warp `[hd]` partials for 8 warps, then 8 `m` and 8 `l`.
270    let smem = (8 * hd + 16) * 4;
271    KernelLaunch::new(gpu, kernel)
272        .grid([rows, nq, 1])
273        .block([256, 1, 1])
274        .shared_mem(smem)
275        .arg_ptr(q)
276        .arg_ptr(k_cache)
277        .arg_ptr(v_cache)
278        .arg_ptr(block_table)
279        .arg_ptr(lists)
280        .arg_ptr(attn_out)
281        .arg_u32(first_pos)
282        .arg_u32(topk)
283        .arg_u32(ratio)
284        .arg_u32(block_size)
285        .arg_u32(nq)
286        .arg_u32(nkv)
287        .arg_u32(hd)
288        .arg_f32(inv_sqrt_d)
289        .launch(stream)
290}
metrale_model_layers/layers/ops/qsa.rs

metrale_model_layers/layers/ops/
qsa.rs
