metrale_model_layers/layers/qwen3_ssm/
carry.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: Carried-state GDN verify (`gdn_carry_wy{2,3,4}` and `gdn_carry_flush`,
4//! kernels/gb10/common/gated_delta_rule_carry.cu): the layer-side handles and binding, the
5//! engage decision as a pure function, the verify launch, and the fold a declining layer
6//! runs before its parent kernels.
7//!
8//! Owner: model-layers, GDN/SSM layer (`qwen3_ssm`).
9//! Invariants:
10//! - The carry kernels run only when `carry_decision` holds, which requires the caller's
11//!   request (`ForwardContext::gdn_write_on_accept` with a carry binding present).
12//! - A layer asked to carry that declines folds the batch's pending rows into H before
13//!   any parent kernel reads H (`carry_flush_run`).
14//!
15//! ## What the model owns
16//!
17//! The model allocates the stash (`[layers][slots][seq_floats]` f32), the pending counts
18//! (`[layers][slots]` u32), the slot table (`u32` per batch position) and one engaged
19//! word per layer, binds them once before any capture, stages the slot table per verify
20//! and sets the pending counts after each verdict (`model-engine gdn_carry.rs`).
21
22use anyhow::Result;
23use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
24
25use super::Qwen3SsmLayer;
26use crate::layer::GdnCarryBinding;
27use crate::layers::ops;
28
29/// 2026-09-26: The kernel handles, each 0 when the module is absent: the verify for
30/// K = 2, 3, 4 and the fold.
31pub(super) struct CarryKernels {
32    pub wy: [KernelHandle; 3],
33    pub wy_lazy: [KernelHandle; 3],
34    pub flush: KernelHandle,
35    pub conv: KernelHandle,
36    pub conv_flush: KernelHandle,
37    /// 2026-10-01: The exact verify's twins (`gdn_exact_carry{2,3,4}` and `_lazy` in the model
38    /// directory, `gdn_carry_conv_f32`), each 0 where the kernel set has none. A kernel set
39    /// without the `gdn_exact_carry` module issues no lookup for them (`try_target_kernel`).
40    pub exact: [KernelHandle; 3],
41    pub exact_lazy: [KernelHandle; 3],
42    pub conv_f32: KernelHandle,
43    /// 2026-10-01: The single-sequence exact verify (`gdn_exact_chain{2,3,4}`,
44    /// `gdn_conv_chain_f32`), each 0 where the kernel set has none.
45    pub chain: [KernelHandle; 3],
46    pub conv_chain: KernelHandle,
47    /// 2026-10-01: The FP16 h-state exact verify (`gdn_exact_chain_f16_{2,3,4}`), looked up
48    /// only when the exact verify runs on an FP16 h-state, so a kernel set without it refuses
49    /// that serve at the boot audit; and the batched FP32 conv chain beside it
50    /// (`gdn_conv_chain_f32_batched`).
51    pub chain_f16: [KernelHandle; 3],
52    pub conv_chain_batched: KernelHandle,
53}
54
55pub(super) fn carry_kernels(gpu: &dyn GpuBackend) -> CarryKernels {
56    let m = "gated_delta_rule_carry";
57    CarryKernels {
58        wy: [
59            crate::layers::try_kernel(gpu, m, "gdn_carry_wy2"),
60            crate::layers::try_kernel(gpu, m, "gdn_carry_wy3"),
61            crate::layers::try_kernel(gpu, m, "gdn_carry_wy4"),
62        ],
63        wy_lazy: [
64            crate::layers::try_kernel(gpu, m, "gdn_carry_wy2_lazy"),
65            crate::layers::try_kernel(gpu, m, "gdn_carry_wy3_lazy"),
66            crate::layers::try_kernel(gpu, m, "gdn_carry_wy4_lazy"),
67        ],
68        flush: carry_flush_kernel(gpu),
69        conv: crate::layers::try_kernel(gpu, m, "gdn_carry_conv"),
70        conv_flush: crate::layers::try_kernel(gpu, m, "gdn_carry_conv_flush"),
71        exact: [
72            crate::layers::try_target_kernel(gpu, "gdn_exact_carry", "gdn_exact_carry2"),
73            crate::layers::try_target_kernel(gpu, "gdn_exact_carry", "gdn_exact_carry3"),
74            crate::layers::try_target_kernel(gpu, "gdn_exact_carry", "gdn_exact_carry4"),
75        ],
76        exact_lazy: [
77            crate::layers::try_target_kernel(gpu, "gdn_exact_carry", "gdn_exact_carry2_lazy"),
78            crate::layers::try_target_kernel(gpu, "gdn_exact_carry", "gdn_exact_carry3_lazy"),
79            crate::layers::try_target_kernel(gpu, "gdn_exact_carry", "gdn_exact_carry4_lazy"),
80        ],
81        conv_f32: crate::layers::try_kernel(gpu, m, "gdn_carry_conv_f32"),
82        chain: [
83            crate::layers::try_target_kernel(gpu, "gdn_exact_carry", "gdn_exact_chain2"),
84            crate::layers::try_target_kernel(gpu, "gdn_exact_carry", "gdn_exact_chain3"),
85            crate::layers::try_target_kernel(gpu, "gdn_exact_carry", "gdn_exact_chain4"),
86        ],
87        conv_chain: crate::layers::try_kernel(gpu, m, "gdn_conv_chain_f32"),
88        chain_f16: if super::ssm_h_fp16_enabled() && super::verify_exact_enabled() {
89            [2, 3, 4].map(|k| {
90                crate::layers::try_kernel(
91                    gpu,
92                    "gdn_exact_carry",
93                    &format!("gdn_exact_chain_f16_{k}"),
94                )
95            })
96        } else {
97            [KernelHandle(0); 3]
98        },
99        conv_chain_batched: crate::layers::try_kernel(gpu, m, "gdn_conv_chain_f32_batched"),
100    }
101}
102
103/// 2026-10-01: The fold of pending rows (`gdn_carry_flush`), or under the exact verify the model
104/// directory's twin (`gdn_exact_carry_flush`), which every directory with the exact twins
105/// defines: one whose strided decode clamps the state after each update folds with the same
106/// clamp. Shared by the layer's fold and the model's (model-engine `gdn_carry.rs`).
107pub fn carry_flush_kernel(gpu: &dyn GpuBackend) -> KernelHandle {
108    if super::verify_exact_enabled() {
109        let exact =
110            crate::layers::try_target_kernel(gpu, "gdn_exact_carry", "gdn_exact_carry_flush");
111        if exact.0 != 0 {
112            return exact;
113        }
114    }
115    crate::layers::try_kernel(gpu, "gated_delta_rule_carry", "gdn_carry_flush")
116}
117
118/// 2026-09-26: Whether a run of `n` sequences takes the lazy kernels
119/// (`ops::GDN_CARRY_LAZY_MIN_SEQS`).
120pub(super) fn carry_lazy(n: usize) -> bool {
121    n >= ops::GDN_CARRY_LAZY_MIN_SEQS
122}
123
124/// 2026-09-26: The per-layer carry state: handles, head dims and the binding.
125pub(super) struct CarryState {
126    pub kernels: CarryKernels,
127    /// 2026-09-26: `[nk, nv, kd, vd]`.
128    pub dims: [usize; 4],
129    pub binding: std::sync::OnceLock<GdnCarryBinding>,
130}
131
132/// 2026-09-26: Everything the engage decision depends on, so tests can pin it without a
133/// GPU.
134#[derive(Clone, Copy, Debug)]
135pub(super) struct CarryRequest {
136    /// 2026-09-26: `ForwardContext::gdn_write_on_accept` in a serve that bound carry.
137    pub requested: bool,
138    pub kk: usize,
139    pub h_f16: bool,
140    pub kd: usize,
141    pub vd: usize,
142    pub kernel_linked: bool,
143    pub bound: bool,
144}
145
146/// 2026-09-26: The one place that decides whether a carry kernel runs.
147pub(super) const fn carry_decision(r: CarryRequest) -> bool {
148    r.requested
149        && r.kk >= 2
150        && r.kk <= 4
151        && !r.h_f16
152        && r.kd == 128
153        && r.vd == 128
154        && r.kernel_linked
155        && r.bound
156}
157
158impl CarryState {
159    pub(super) fn new(gpu: &dyn GpuBackend, dims: [usize; 4]) -> Self {
160        Self {
161            kernels: carry_kernels(gpu),
162            dims,
163            binding: std::sync::OnceLock::new(),
164        }
165    }
166
167    /// 2026-09-26: Stash width when this layer can carry: every kernel linked and
168    /// kd == vd == 128. `None` otherwise.
169    pub(super) fn seq_floats(&self) -> Option<usize> {
170        let [_, nv, kd, vd] = self.dims;
171        let k = &self.kernels;
172        let linked = k.wy.iter().chain(&k.wy_lazy).all(|h| h.0 != 0)
173            && [k.flush, k.conv, k.conv_flush].iter().all(|h| h.0 != 0);
174        (linked && kd == 128 && vd == 128).then(|| ops::gdn_carry_seq_floats(nv, kd, vd))
175    }
176
177    fn kernel_for(&self, kk: usize, lazy: bool) -> KernelHandle {
178        match (kk, lazy) {
179            (2..=4, false) => self.kernels.wy[kk - 2],
180            (2..=4, true) => self.kernels.wy_lazy[kk - 2],
181            _ => KernelHandle(0),
182        }
183    }
184
185    /// 2026-10-01: The exact verify's chain twin for width `kk`, 0 when absent.
186    pub(super) fn exact_kernel_for(&self, kk: usize, lazy: bool) -> KernelHandle {
187        match (kk, lazy) {
188            (2..=4, false) => self.kernels.exact[kk - 2],
189            (2..=4, true) => self.kernels.exact_lazy[kk - 2],
190            _ => KernelHandle(0),
191        }
192    }
193}
194
195impl Qwen3SsmLayer {
196    /// 2026-09-26: The binding when this call asked for carry and a binding exists.
197    pub(super) fn carry_requested(&self, write_on_accept: bool) -> Option<GdnCarryBinding> {
198        if !write_on_accept {
199            return None;
200        }
201        self.carry.binding.get().copied()
202    }
203
204    /// 2026-09-26: Whether the carry kernel for verify width `kk` runs on this call.
205    pub(super) fn carry_now(&self, write_on_accept: bool, kk: usize) -> bool {
206        let [_, _, kd, vd] = self.carry.dims;
207        carry_decision(CarryRequest {
208            requested: write_on_accept,
209            kk,
210            h_f16: super::ssm_h_fp16_enabled(),
211            kd,
212            vd,
213            kernel_linked: self.carry.kernel_for(kk, false).0 != 0,
214            bound: self.carry.binding.get().is_some(),
215        })
216    }
217
218    /// 2026-09-26: Launch the carry verify for one run of `n` sequences at width `kk`.
219    /// `slot_tab` and `flag` are the run's slices of the shared slot table and of this
220    /// layer's engaged words.
221    #[allow(clippy::too_many_arguments)]
222    pub(super) fn carry_launch(
223        &self,
224        gpu: &dyn GpuBackend,
225        kk: usize,
226        h_table: DevicePtr,
227        slot_tab: DevicePtr,
228        flag: DevicePtr,
229        q_ptr: DevicePtr,
230        k_ptr: DevicePtr,
231        v_ptr: DevicePtr,
232        gate_ptr: DevicePtr,
233        beta_ptr: DevicePtr,
234        gdn_out_buf: DevicePtr,
235        n: usize,
236        conv_dim: usize,
237        stream: u64,
238    ) -> Result<()> {
239        let b = self
240            .carry
241            .binding
242            .get()
243            .ok_or_else(|| anyhow::anyhow!("carry_launch: no carry binding"))?;
244        let [nk, nv, kd, _] = self.carry.dims;
245        ops::gdn_carry_wy(
246            gpu,
247            self.carry.kernel_for(kk, carry_lazy(n)),
248            h_table,
249            q_ptr,
250            k_ptr,
251            v_ptr,
252            gate_ptr,
253            beta_ptr,
254            gdn_out_buf,
255            b.stash,
256            slot_tab,
257            b.pend,
258            b.seq_floats as u32,
259            n as u32,
260            nk as u32,
261            nv as u32,
262            conv_dim as u32,
263            conv_dim as u32,
264            (nv * 2) as u32,
265            kd as u32,
266            flag,
267            stream,
268        )?;
269        static LOGGED: std::sync::atomic::AtomicU32 = std::sync::atomic::AtomicU32::new(0);
270        let bit = 1u32 << kk;
271        if LOGGED.fetch_or(bit, std::sync::atomic::Ordering::Relaxed) & bit == 0 {
272            tracing::info!(
273                "batched-verify GDN CARRY ENGAGED (n={n}, k={kk}): pending rows folded in the \
274                 verify, no h intermediates written (kill switch METRALE_NO_GDN_CARRY)"
275            );
276        }
277        Ok(())
278    }
279
280    /// 2026-09-26: The carried conv launch for one run (`gdn_carry_conv`), with the batched
281    /// parent's arguments: `conv_base` is the run's first conv state, the rest follow at
282    /// `conv_bytes` strides; `slot_tab` is the run's slice of the slot table.
283    pub(super) fn carry_conv_launch(
284        &self,
285        gpu: &dyn GpuBackend,
286        conv_base: DevicePtr,
287        slot_tab: DevicePtr,
288        n: usize,
289        a: &super::trait_decode_batched_conv_gdn::ConvGdnArgs,
290    ) -> Result<()> {
291        let b = self
292            .carry
293            .binding
294            .get()
295            .ok_or_else(|| anyhow::anyhow!("carry_conv_launch: no carry binding"))?;
296        ops::gdn_carry_conv(
297            gpu,
298            self.carry.kernels.conv,
299            conv_base,
300            a.deinterleaved,
301            &self.ssm.conv1d,
302            a.conv_out_buf,
303            b.conv_stash,
304            slot_tab,
305            b.pend,
306            b.conv_seq_elems as u32,
307            a.num_tokens as u32,
308            a.conv_dim as u32,
309            a.d_conv as u32,
310            a.qk_ch,
311            a.kd as u32,
312            a.qkvz_size as u32,
313            a.conv_dim as u32,
314            1e-6,
315            n as u32,
316            (self.conv_state_bytes / 4) as u32,
317            (a.num_tokens * a.qkvz_size) as u32,
318            (a.num_tokens * a.conv_dim) as u32,
319            carry_lazy(n),
320            a.stream,
321        )
322    }
323
324    /// 2026-09-26: Fold the pending rows of `n` sequences into this layer's H and conv
325    /// window before a parent kernel reads them. `h_table`, `slot_tab` and `run_first`
326    /// locate the run in the WY table slab 0, the slot table and the conv-state table. A
327    /// no-op without a binding.
328    pub(super) fn carry_flush_run(
329        &self,
330        gpu: &dyn GpuBackend,
331        h_table: DevicePtr,
332        run_first: usize,
333        n: usize,
334        stream: u64,
335    ) -> Result<()> {
336        let Some(b) = self.carry.binding.get().copied() else {
337            return Ok(());
338        };
339        anyhow::ensure!(
340            !h_table.is_null(),
341            "carry_flush_run: a layer asked to carry declined without WY tables"
342        );
343        let [nk, nv, kd, vd] = self.carry.dims;
344        let slot_tab = b.slot_tab.offset(run_first * 4);
345        ops::gdn_carry_flush(
346            gpu,
347            self.carry.kernels.flush,
348            h_table,
349            0,
350            b.stash,
351            0,
352            slot_tab,
353            b.pend,
354            0,
355            b.seq_floats as u32,
356            n as u32,
357            nv as u32,
358            1,
359            stream,
360        )?;
361        ops::gdn_carry_conv_flush(
362            gpu,
363            self.carry.kernels.conv_flush,
364            b.conv_tab.offset(run_first * 8),
365            0,
366            b.conv_stash,
367            0,
368            slot_tab,
369            b.pend,
370            0,
371            b.conv_seq_elems as u32,
372            n as u32,
373            (nk * kd * 2 + nv * vd) as u32,
374            (self.conv_state_bytes / 4 / (nk * kd * 2 + nv * vd)) as u32,
375            1,
376            stream,
377        )
378    }
379}
380
381#[cfg(test)]
382mod tests {
383    use super::{CarryRequest, carry_decision};
384
385    const GO: CarryRequest = CarryRequest {
386        requested: true,
387        kk: 2,
388        h_f16: false,
389        kd: 128,
390        vd: 128,
391        kernel_linked: true,
392        bound: true,
393    };
394
395    #[test]
396    fn widths_two_to_four_engage() {
397        for kk in 2..=4 {
398            assert!(carry_decision(CarryRequest { kk, ..GO }));
399        }
400    }
401
402    #[test]
403    fn each_gate_declines_alone() {
404        assert!(!carry_decision(CarryRequest {
405            requested: false,
406            ..GO
407        }));
408        assert!(!carry_decision(CarryRequest { kk: 1, ..GO }));
409        assert!(!carry_decision(CarryRequest { kk: 5, ..GO }));
410        assert!(!carry_decision(CarryRequest { h_f16: true, ..GO }));
411        assert!(!carry_decision(CarryRequest { kd: 64, ..GO }));
412        assert!(!carry_decision(CarryRequest { vd: 256, ..GO }));
413        assert!(!carry_decision(CarryRequest {
414            kernel_linked: false,
415            ..GO
416        }));
417        assert!(!carry_decision(CarryRequest { bound: false, ..GO }));
418    }
419}
metrale_model_layers/layers/qwen3_ssm/carry.rs

metrale_model_layers/layers/qwen3_ssm/
carry.rs
