spark_model/layers/qwen3_ssm/
init_q2.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Q2_0 keep-packed qkvz init + transient-dequant prefill GEMM.
4
5use super::*;
6
7impl Qwen3SsmLayer {
8    /// Install the Tier-1c keep-packed ternary Q2_0 fused `in_proj_qkvz`
9    /// (`ATLAS_GGUF_NATIVE_Q2`). Decode dispatches `q2_0_gemv_vec`; prefill
10    /// transient-dequants via `Self::qkvz_q2_prefill_gemm`. `out_proj` is
11    /// unaffected (stays NVFP4). Requires `sequential_qkvz` (Bonsai concats
12    /// [Q|K|V|Z] at load).
13    ///
14    /// The keep-packed MMQ kernels are resolved HERE, not in the constructor:
15    /// they ship only in GGUF-serving targets, and probing them on models
16    /// that never install packed-Q2 weights fails the fail-closed boot audit
17    /// on every other GDN target.
18    pub fn set_packed_q2_qkvz(
19        &mut self,
20        qkvz: crate::weight_map::PackedQ2Weight,
21        gpu: &dyn GpuBackend,
22    ) {
23        self.qkvz_q2 = Some(qkvz);
24        self.q2_0_mmq_nc_k = super::super::try_kernel(gpu, "q2_0_mmq", "atlas_q2_0_mmq128_nc");
25        self.q2_0_mmq_wc_k = super::super::try_kernel(gpu, "q2_0_mmq", "atlas_q2_0_mmq128_wc");
26        self.q4k_quant_act_k =
27            super::super::try_kernel(gpu, "q4k_mmq", "atlas_q8_1_quantize_ds4_bf16");
28    }
29
30    /// Transient-dequant prefill GEMM for the packed qkvz: dequant the 2-bit
31    /// `[qkvz_size, h]` weight into the caller-provided PERSISTENT BF16 `scratch`
32    /// (the arena `q2_dequant_scratch`, sized to the largest packed projection),
33    /// run `dense_gemm` (`out[m, qkvz_size] = in[m, h] @ w^T`). Mirrors the
34    /// FFN/attention packed prefill — no per-matmul alloc/sync/free; the dequant
35    /// orders before the GEMM on the same `stream`. Errors if the dequant kernel
36    /// is absent.
37    #[allow(clippy::too_many_arguments)]
38    pub(crate) fn qkvz_q2_prefill_gemm(
39        &self,
40        gpu: &dyn GpuBackend,
41        input: DevicePtr,
42        out: DevicePtr,
43        scratch: DevicePtr,
44        act_q8: DevicePtr,
45        m: u32,
46        stream: u64,
47    ) -> Result<()> {
48        let w = self
49            .qkvz_q2
50            .as_ref()
51            .ok_or_else(|| anyhow::anyhow!("qkvz_q2_prefill_gemm: no packed qkvz installed"))?;
52        let (n, k) = (w.n, w.k);
53
54        // Tier-2 native MMQ (ATLAS_GGUF_NATIVE_Q2_MMQ=1): quantize `input` to q8_1
55        // then run the packed 2-bit MMQ GEMM for the fused qkvz — no BF16 weight
56        // dequant, no shared `q2_dequant_scratch` race. Group-128 only.
57        if self.q2_0_mmq_nc_k.0 != 0
58            && self.q4k_quant_act_k.0 != 0
59            && crate::layers::ops::native_q2_mmq_enabled()
60            && w.group == 128
61        {
62            crate::layers::ops::quantize_act_q8_1(
63                gpu,
64                self.q4k_quant_act_k,
65                input,
66                act_q8,
67                m,
68                k,
69                stream,
70            )?;
71            return crate::layers::ops::q2_0_mmq_gemm(
72                gpu,
73                self.q2_0_mmq_nc_k,
74                self.q2_0_mmq_wc_k,
75                act_q8,
76                w.weight,
77                out,
78                m,
79                n,
80                k,
81                stream,
82            );
83        }
84
85        if self.dequant_q2_0_gn_k.0 == 0 {
86            anyhow::bail!(
87                "dequant_q2_0_gn_to_bf16 kernel missing — packed-Q2 GDN prefill unavailable"
88            );
89        }
90        crate::layers::ops::dequant_q2_0_gn_to_bf16(
91            gpu,
92            self.dequant_q2_0_gn_k,
93            w.weight,
94            scratch,
95            n,
96            k,
97            w.group as u32,
98            stream,
99        )?;
100        let dw = DenseWeight { weight: scratch };
101        if self.dense_gemm_pipelined_k.0 != 0 {
102            crate::layers::ops::dense_gemm_bf16_pipelined(
103                gpu,
104                self.dense_gemm_pipelined_k,
105                input,
106                &dw,
107                out,
108                m,
109                n,
110                k,
111                stream,
112            )?;
113        } else {
114            crate::layers::ops::dense_gemm(
115                gpu,
116                self.dense_gemm_k,
117                input,
118                &dw,
119                out,
120                m,
121                n,
122                k,
123                stream,
124            )?;
125        }
126        Ok(())
127    }
128}
spark_model/layers/qwen3_ssm/init_q2.rs

spark_model/layers/qwen3_ssm/
init_q2.rs
