metrale_model_layers/layers/qwen3_ssm/
init_q2.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The packed Q2_0 fused `in_proj_qkvz`: its install and its
4//! prefill GEMM.
5//!
6//! Owner: model-layers (qwen3_ssm).
7//! Invariants: none beyond the types.
8
9use super::*;
10
11impl Qwen3SsmLayer {
12    /// 2026-09-25: Install the packed Q2_0 fused `in_proj_qkvz`. Decode runs
13    /// `q2_0_gemv_vec` on it and prefill runs `qkvz_q2_prefill_gemm`. The layer
14    /// must be `sequential_qkvz`: both paths write the projection straight into
15    /// the deinterleaved buffer, and decode reads `qkvz_q2` only on the
16    /// sequential branch.
17    ///
18    /// The MMQ kernels are looked up here rather than in `new`, so a layer that
19    /// never installs a packed weight issues no lookup and leaves no failed
20    /// row in the boot kernel audit.
21    pub fn set_packed_q2_qkvz(
22        &mut self,
23        qkvz: crate::weight_map::PackedQ2Weight,
24        gpu: &dyn GpuBackend,
25    ) {
26        self.qkvz_q2 = Some(qkvz);
27        self.q2_0_mmq_nc_k = super::super::try_kernel(gpu, "q2_0_mmq", "metrale_q2_0_mmq128_nc");
28        self.q2_0_mmq_wc_k = super::super::try_kernel(gpu, "q2_0_mmq", "metrale_q2_0_mmq128_wc");
29        self.q4k_quant_act_k =
30            super::super::try_kernel(gpu, "q4k_mmq", "metrale_q8_1_quantize_ds4_bf16");
31    }
32
33    /// 2026-09-25: `out[m, n] = input[m, k] @ W^T` for the packed qkvz.
34    ///
35    /// With both MMQ handles resolved, `METRALE_GGUF_NATIVE_Q2_MMQ=1` and group
36    /// size 128, it quantizes `input` to q8_1 into `act_q8` and runs
37    /// `q2_0_mmq_gemm`. Otherwise it dequantizes the weight to BF16 into
38    /// `scratch` and runs `dense_gemm_bf16_pipelined` (or `dense_gemm` when that
39    /// handle is 0) on the same `stream`. Allocates nothing. Errors when no
40    /// packed qkvz is installed or the dequant kernel is missing.
41    #[allow(clippy::too_many_arguments)]
42    pub(crate) fn qkvz_q2_prefill_gemm(
43        &self,
44        gpu: &dyn GpuBackend,
45        input: DevicePtr,
46        out: DevicePtr,
47        scratch: DevicePtr,
48        act_q8: DevicePtr,
49        m: u32,
50        stream: u64,
51    ) -> Result<()> {
52        let w = self
53            .qkvz_q2
54            .as_ref()
55            .ok_or_else(|| anyhow::anyhow!("qkvz_q2_prefill_gemm: no packed qkvz installed"))?;
56        let (n, k) = (w.n, w.k);
57
58        if self.q2_0_mmq_nc_k.0 != 0
59            && self.q4k_quant_act_k.0 != 0
60            && crate::layers::ops::native_q2_mmq_enabled()
61            && w.group == 128
62        {
63            crate::layers::ops::quantize_act_q8_1(
64                gpu,
65                self.q4k_quant_act_k,
66                input,
67                act_q8,
68                m,
69                k,
70                stream,
71            )?;
72            return crate::layers::ops::q2_0_mmq_gemm(
73                gpu,
74                self.q2_0_mmq_nc_k,
75                self.q2_0_mmq_wc_k,
76                act_q8,
77                w.weight,
78                out,
79                m,
80                n,
81                k,
82                stream,
83            );
84        }
85
86        if self.dequant_q2_0_gn_k.0 == 0 {
87            anyhow::bail!(
88                "dequant_q2_0_gn_to_bf16 kernel missing — packed-Q2 GDN prefill unavailable"
89            );
90        }
91        crate::layers::ops::dequant_q2_0_gn_to_bf16(
92            gpu,
93            self.dequant_q2_0_gn_k,
94            w.weight,
95            scratch,
96            n,
97            k,
98            w.group as u32,
99            stream,
100        )?;
101        let dw = DenseWeight { weight: scratch };
102        if self.dense_gemm_pipelined_k.0 != 0 {
103            crate::layers::ops::dense_gemm_bf16_pipelined(
104                gpu,
105                self.dense_gemm_pipelined_k,
106                input,
107                &dw,
108                out,
109                m,
110                n,
111                k,
112                stream,
113            )?;
114        } else {
115            crate::layers::ops::dense_gemm(
116                gpu,
117                self.dense_gemm_k,
118                input,
119                &dw,
120                out,
121                m,
122                n,
123                k,
124                stream,
125            )?;
126        }
127        Ok(())
128    }
129}
metrale_model_layers/layers/qwen3_ssm/init_q2.rs

metrale_model_layers/layers/qwen3_ssm/
init_q2.rs
