metrale_model_arch/glm5next_dsa/layer/
rows.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: Three `Glm5NextDsaLayer` methods: `store_indexer_row`, the last step of
4//! `indexer_forward`; `select_rows_batched`, the one-pass selection for all rows of a
5//! prefill sub-chunk; and `write_kv_row`, one MTP drafter context row.
6//!
7//! Owner: model-arch (GLM-5.3 DSA).
8//! Invariants:
9//! - `write_kv_row` writes no KV latent when the indexer cache is behind `seq_len`.
10
11use anyhow::{Result, bail};
12use metrale_cache::kv_cache::PagedKvCache;
13use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
14use metrale_gpu_runtime::kernel_args::KernelLaunch;
15use metrale_model_layers::layer::{ForwardContext, LayerState};
16
17use super::super::select::{DsaSelectInputs, select_tokens};
18use super::super::state::Glm5NextDsaState;
19use super::{Glm5NextDsaLayer, gemm};
20
21impl Glm5NextDsaLayer {
22    /// 2026-09-26: `indexer_forward`'s last step: places the projected `k_normed` and `gate`
23    /// at row `pos` (through `dsa_indexer_store` at `pos_dev` when given), marks the row
24    /// valid, and advances `state` by one row.
25    pub(super) fn store_indexer_row(
26        &self,
27        gpu: &dyn GpuBackend,
28        state: &mut Glm5NextDsaState,
29        pos_dev: Option<DevicePtr>,
30        pos: usize,
31        d: usize,
32        stream: u64,
33    ) -> Result<()> {
34        let w = &self.workspace;
35        match pos_dev {
36            // 2026-09-25: Placement and the validity mark both use the device-side position;
37            // a memset at `valid.offset(pos)` would fix a host address in a captured graph.
38            Some(pd) => {
39                KernelLaunch::new(gpu, self.select_kernels.indexer_store)
40                    .grid([1, 1, 1])
41                    .block([d.min(1024) as u32, 1, 1])
42                    .arg_ptr(w.stage_k)
43                    .arg_ptr(w.stage_gate)
44                    .arg_ptr(pd)
45                    .arg_ptr(state.k_normed)
46                    .arg_ptr(state.gate)
47                    .arg_ptr(state.valid)
48                    .arg_u32(d as u32)
49                    .launch(stream)?;
50            }
51            None => gpu.memset_async(state.valid.offset(pos), 1, 1, stream)?,
52        }
53        state.advance(1)
54    }
55
56    /// 2026-09-25: The selection for all `k` rows in one pass (`q_rows = k`), used when
57    /// [`batch_select_enabled`](super::batch_select_enabled) holds; see there for why it
58    /// selects what per-row passes do.
59    ///
60    /// The `wq_b` projection is one M = 1 GEMV per row, or one cuBLASLt GEMM for all rows
61    /// when `METRALE_GLM_DSA_BATCH_QIDX=1` (`glm5next_layer::dsa_batch_qidx`, off by
62    /// default).
63    #[allow(clippy::too_many_arguments)]
64    pub(super) fn select_rows_batched(
65        &self,
66        gpu: &dyn GpuBackend,
67        k: usize,
68        state: &Glm5NextDsaState,
69        q_pos_host: &[i32],
70        stream: u64,
71    ) -> Result<()> {
72        let w = &self.workspace;
73        // 2026-09-25: `decode_k` calls this after all k indexer writes, so `P` is the group's
74        // final pool count; each row limits itself to its own position through `q_pos[r]`.
75        let geom = state.geometry(&self.cfg, k)?;
76        let idx_row = self.cfg.index_heads * self.cfg.index_head_dim;
77        // 2026-09-25: `q_resid` is a contiguous `[k, q_lora_rank]` BF16 block and `q_idx_rows`
78        // a contiguous `[k, idx_row]` FP32 one, so one GEMM covers all rows. It sums in a
79        // different order from the M = 1 GEMV, so near-tied index scores can select other
80        // pools; that is why it is opt-in. Decode and the speculative verify never reach this
81        // function (`batch_select_enabled`).
82        if crate::glm5next_layer::dsa_batch_qidx() {
83            metrale_model_layers::layers::ops::cublas_bf16_proj_dense_f32_out(
84                w.q_resid,
85                self.weights.wq_b,
86                w.q_idx_rows,
87                k as u32,
88                idx_row as u32,
89                self.cfg.q_lora_rank as u32,
90                stream,
91            )?;
92        } else {
93            for row in 0..k {
94                gemm(
95                    gpu,
96                    self.kernels.gemm_f32,
97                    self.kernels.gemv_f32,
98                    // 2026-09-25: No FP32-out batched GEMV kernel exists.
99                    KernelHandle(0),
100                    w.q_resid.offset(row * self.cfg.q_lora_rank * 2),
101                    self.weights.wq_b,
102                    w.q_idx_rows.offset(row * idx_row * 4),
103                    1,
104                    idx_row,
105                    self.cfg.q_lora_rank,
106                    stream,
107                )?;
108            }
109        }
110        let q_pos_bytes: Vec<u8> = q_pos_host.iter().flat_map(|p| p.to_le_bytes()).collect();
111        gpu.copy_h2d(&q_pos_bytes, w.q_pos_rows)?;
112        let inputs = DsaSelectInputs {
113            k_normed: state.k_normed,
114            gate: state.gate,
115            valid: state.valid,
116            ape: self.weights.ape,
117            q: w.q_idx_rows,
118            weights: w.head_weights_rows,
119            q_pos: w.q_pos_rows,
120            q_mask: w.q_mask_rows,
121            first_key: 0,
122            // 2026-09-25: Host geometry only: `select_tokens` refuses the device-geometry
123            // (ceiling) launch at `q_rows > 1`.
124            geom_dev: DevicePtr::NULL,
125        };
126        let t = crate::glm5next_layer::profile::start();
127        // 2026-09-25: The base of the `[max_rows, out_width]` output, not a row slice: the
128        // kernels index rows themselves, so rows 0..k land in their own slots.
129        select_tokens(
130            gpu,
131            &self.select_kernels,
132            &self.cfg,
133            &geom,
134            &inputs,
135            &w.select,
136            super::super::select::DsaSelectLaunch::Exact,
137            stream,
138        )?;
139        crate::glm5next_layer::profile::end(
140            crate::glm5next_layer::profile::DSA_SELECT,
141            t,
142            gpu,
143            stream,
144        );
145        Ok(())
146    }
147
148    /// 2026-09-25: One MTP drafter context row: writes the KV latent and the indexer row, with
149    /// no query, selection or attention.
150    ///
151    /// The row goes to KV slot `seq_len`, and the indexer writes at `state.len()`. An indexer
152    /// cache ahead of `seq_len` is rewound; one behind is an error, so rows must be written
153    /// densely from 0.
154    #[allow(clippy::too_many_arguments)]
155    pub fn write_kv_row(
156        &self,
157        hidden: DevicePtr,
158        state: &mut dyn LayerState,
159        kv_cache: &mut PagedKvCache,
160        seq_len: usize,
161        block_table: &mut Vec<u32>,
162        ctx: &ForwardContext,
163        stream: u64,
164    ) -> Result<()> {
165        let gpu = ctx.gpu;
166        let st = state
167            .as_any_mut()
168            .downcast_mut::<Glm5NextDsaState>()
169            .ok_or_else(|| {
170                anyhow::anyhow!("Glm5NextDsaLayer got a state that is not Glm5NextDsaState")
171            })?;
172        match st.len().cmp(&seq_len) {
173            std::cmp::Ordering::Greater => st.rewind_to(seq_len)?,
174            std::cmp::Ordering::Less => bail!(
175                "DSA layer {}: indexer cache holds {} rows but the drafter is at {seq_len} — \
176                 rows are MISSING, not merely stale.",
177                self.layer_idx,
178                st.len()
179            ),
180            std::cmp::Ordering::Equal => {}
181        }
182        let w = &self.workspace;
183        gemm(
184            gpu,
185            self.kernels.gemm,
186            self.kernels.gemv,
187            self.kernels.gemv_batchm,
188            hidden,
189            self.weights.kv_a_proj,
190            w.kv_a,
191            1,
192            self.cfg.kv_lora_rank,
193            self.cfg.hidden,
194            stream,
195        )?;
196        let block_size = kv_cache.config().block_size;
197        let logical = seq_len / block_size;
198        let physical = *block_table.get(logical).ok_or_else(|| {
199            anyhow::anyhow!(
200                "DSA layer {}: block table has {} entries, needs logical block {logical} for \
201                 drafter row {seq_len}",
202                self.layer_idx,
203                block_table.len()
204            )
205        })? as usize;
206        let slot = (physical * block_size + seq_len % block_size) as i64;
207        gpu.copy_h2d(&slot.to_le_bytes(), w.slot)?;
208        KernelLaunch::new(gpu, self.kernels.latent_write)
209            .grid([1, 1, 1])
210            .block([self.cfg.kv_lora_rank as u32, 1, 1])
211            .arg_ptr(w.kv_a)
212            .arg_ptr(self.weights.kv_a_layernorm)
213            .arg_ptr(kv_cache.k_pool_ptr(self.attn_layer_idx))
214            .arg_ptr(w.slot)
215            .arg_u32(self.cfg.kv_lora_rank as u32)
216            .arg_f32(self.rms_eps)
217            .arg_f32(1.0 / self.kv_scale)
218            .launch(stream)?;
219        self.indexer_forward(gpu, hidden, st, None, stream)
220    }
221}
metrale_model_arch/glm5next_dsa/layer/rows.rs

metrale_model_arch/glm5next_dsa/layer/
rows.rs
