metrale_model_arch/glm5next_dsa/select/
launch.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `select_tokens`, the launcher for one DSA selection pass.
4//!
5//! Owner: model-arch (GLM-5.3 DSA).
6//! Invariants: none beyond the types.
7
8use super::*;
9
10/// 2026-09-25: Run the selection kernels, leaving `[q_rows, out_width]` token ids in
11/// [`DsaSelectScratch::tokens`].
12///
13/// Fails before any launch when `scratch` is too small for `geom`, or when a ceiling launch
14/// has no `geom_dev` or has `q_rows > 1`. The kernels are enqueued on `stream`; nothing here
15/// synchronises.
16#[allow(clippy::too_many_arguments)]
17pub fn select_tokens(
18    gpu: &dyn GpuBackend,
19    kernels: &Glm5NextDsaKernels,
20    cfg: &Glm5NextDsaConfig,
21    geom: &DsaSelectGeometry,
22    inputs: &DsaSelectInputs,
23    scratch: &DsaSelectScratch,
24    launch: DsaSelectLaunch,
25    stream: u64,
26) -> Result<()> {
27    scratch.fits(cfg, geom)?;
28
29    let d = geom.index_head_dim;
30    let kp = geom.index_kpool;
31
32    let ceiling = match launch {
33        DsaSelectLaunch::Exact => None,
34        DsaSelectLaunch::Ceiling { max_pools } => {
35            if inputs.geom_dev.0 == 0 {
36                bail!(
37                    "DSA select: a ceiling launch has no host geometry to fall back on and \
38                     needs `geom_dev`; passing NULL would run the frozen scalars."
39                );
40            }
41            if geom.q_rows != 1 {
42                bail!(
43                    "DSA select: ceiling launch is decode-only (q_rows must be 1, got {}). \
44                     At q_rows > 1 the row strides vary with the context and a replayed \
45                     graph would index the wrong rows.",
46                    geom.q_rows
47                );
48            }
49            Some(max_pools)
50        }
51    };
52    let gd = inputs.geom_dev;
53
54    // 2026-09-25: Under a ceiling launch the kernels read S, the pool counts, the tile and
55    // select_k from `geom_dev`, so the scalar arguments are unused. They are set from the
56    // ceiling, so every capture at the same ceiling passes the same values.
57    let (seq_a, npools_a, np2_a, selk_a) = match ceiling {
58        Some(m) => (
59            m * kp,
60            m,
61            m.next_power_of_two().max(2).min(topk_tile()),
62            cfg.select_k(m),
63        ),
64        None => (geom.seq, geom.n_pools, geom.topk_np2, geom.select_k),
65    };
66
67    // 2026-09-25: Below `index_kpool` tokens there are no pools to score or sort, and a
68    // zero-extent grid is not a valid launch, so `dsa_index_scores` and `dsa_topk_pools` are
69    // skipped. `dsa_kpool_compress` still runs (`n_pools_full >= 1`), and
70    // `dsa_expand_selection` writes the row to -1 and appends the visible tail, which is the
71    // whole selection here. Under a ceiling launch both always run; with no live pools
72    // `dsa_index_scores` blocks return at once and `dsa_topk_pools` selects nothing.
73    let has_pools = geom.n_pools > 0 || ceiling.is_some();
74
75    // 2026-09-25: Pool compression over the full pool count; the trailing partial pool is
76    // written and marked invalid, and no later kernel reads it.
77    KernelLaunch::new(gpu, kernels.kpool_compress)
78        .grid([ceiling.map_or(geom.n_pools_full, |m| m + 1) as u32, 1, 1])
79        .block([d.min(1024) as u32, 1, 1])
80        .arg_ptr(inputs.k_normed)
81        .arg_ptr(inputs.gate)
82        .arg_ptr(inputs.valid)
83        .arg_ptr(inputs.ape)
84        .arg_ptr(scratch.pool_keys)
85        .arg_ptr(scratch.pool_indices)
86        .arg_ptr(scratch.pool_valid)
87        .arg_u32(seq_a as u32)
88        .arg_u32(d as u32)
89        .arg_u32(kp as u32)
90        .arg_i32(inputs.first_key)
91        .arg_ptr(gd)
92        .launch(stream)?;
93
94    if has_pools {
95        KernelLaunch::new(gpu, kernels.index_scores)
96            .grid([
97                ceiling.unwrap_or(geom.n_pools) as u32,
98                geom.q_rows as u32,
99                1,
100            ])
101            .block([SCORES_BLOCK, 1, 1])
102            // 2026-09-25: `dsa_index_scores` keeps one f32 per index head in shared memory and
103            // sums them in head order.
104            .shared_mem(SCORES_BLOCK.max((geom.index_heads * 4) as u32))
105            .arg_ptr(inputs.q)
106            .arg_ptr(scratch.pool_keys)
107            .arg_ptr(inputs.weights)
108            .arg_ptr(scratch.pool_indices)
109            .arg_ptr(scratch.pool_valid)
110            .arg_ptr(inputs.valid)
111            .arg_ptr(inputs.q_pos)
112            .arg_ptr(scratch.scores)
113            .arg_ptr(scratch.valid_cand)
114            .arg_u32(geom.q_rows as u32)
115            .arg_u32(npools_a as u32)
116            .arg_u32(geom.index_heads as u32)
117            .arg_u32(d as u32)
118            .arg_u32(kp as u32)
119            .arg_u32(seq_a as u32)
120            .arg_f32((d as f32).powf(-0.5))
121            .arg_ptr(gd)
122            .launch(stream)?;
123
124        // 2026-09-25: `np2_a` is at most `topk_tile()`, so the request is at most
125        // `topk_smem_for_tile(topk_tile())`, within `TOPK_SMEM_CEILING`.
126        KernelLaunch::new(gpu, kernels.topk_pools)
127            .grid([geom.q_rows as u32, 1, 1])
128            .block([ROW_BLOCK, 1, 1])
129            .shared_mem(topk_smem_for_tile(np2_a) as u32)
130            .arg_ptr(scratch.scores)
131            .arg_ptr(scratch.selected)
132            .arg_u32(geom.q_rows as u32)
133            .arg_u32(npools_a as u32)
134            .arg_u32(np2_a as u32)
135            .arg_u32(selk_a as u32)
136            .arg_ptr(gd)
137            .launch(stream)?;
138    }
139
140    KernelLaunch::new(gpu, kernels.expand_selection)
141        .grid([geom.q_rows as u32, 1, 1])
142        .block([ROW_BLOCK, 1, 1])
143        .arg_ptr(scratch.selected)
144        .arg_ptr(scratch.pool_indices)
145        .arg_ptr(scratch.valid_cand)
146        .arg_ptr(inputs.valid)
147        .arg_ptr(inputs.q_pos)
148        .arg_ptr(inputs.q_mask)
149        .arg_ptr(scratch.tokens)
150        .arg_u32(geom.q_rows as u32)
151        .arg_u32(npools_a as u32)
152        .arg_u32(kp as u32)
153        .arg_u32(seq_a as u32)
154        .arg_u32(selk_a as u32)
155        .arg_u32(geom.out_width as u32)
156        .arg_i32(inputs.first_key)
157        .arg_i32(cfg.always_select_tail as i32)
158        .arg_ptr(gd)
159        .launch(stream)?;
160
161    Ok(())
162}
metrale_model_arch/glm5next_dsa/select/launch.rs

metrale_model_arch/glm5next_dsa/select/
launch.rs
