metrale_model_layers/layers/ops/
q2_0_mmq.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launcher for the packed Q2_0 (ternary) MMQ prefill GEMM,
4//! `metrale_q2_0_mmq128_nc/_wc` in `kernels/gb10/qwen3.6-27b/nvfp4/q2_0_mmq.cu`.
5//! The 2-bit weight stays packed: the vendored `load_tiles_q2_0` unpacks
6//! `code - 1` into the int8 tile, and `vec_dot_q8_0_q8_1_mma` multiplies it by a
7//! q8_1 activation and applies both scales. The activation quantizer and scratch
8//! size are Q4_K's ([`super::quantize_act_q8_1`], [`super::q8_1_scratch_bytes`]),
9//! in the DS4 layout.
10//!
11//! Owner: model-layers ops.
12//! Invariants: none beyond the types.
13
14use anyhow::Result;
15use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
16use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
17
18use crate::weight_map::PackedQ2Weight;
19
20/// 2026-09-25: Weights per `block_q2_0`.
21pub const QK2_0: u32 = 128;
22/// 2026-09-25: `sizeof(block_q2_0)`: an fp16 scale `d` and 128 codes at 4 per byte.
23pub const Q2_0_BLOCK_BYTES: usize = 34;
24
25/// 2026-09-25: True only when `METRALE_GGUF_NATIVE_Q2_MMQ` is exactly `1`. With
26/// it, and with the kernels resolved and group 128 on gate, up and down, the
27/// dense FFN prefill runs packed Q2 weights through [`q2_0_mmq_gemm`] instead of
28/// dequantizing them to BF16. Whether weights stay packed at all is
29/// `METRALE_GGUF_NATIVE_Q2`'s decision.
30pub fn native_q2_mmq_enabled() -> bool {
31    std::env::var("METRALE_GGUF_NATIVE_Q2_MMQ").ok().as_deref() == Some("1")
32}
33
34/// 2026-09-25: Bytes of the packed `block_q2_0` form of an `[n, k]` weight; `k`
35/// must be a multiple of 128.
36pub fn q2_0_weight_bytes(n: u32, k: u32) -> usize {
37    (n as usize) * (k as usize / QK2_0 as usize) * Q2_0_BLOCK_BYTES
38}
39
40/// 2026-09-25: `C[m, n] = A[m, k] x W[n, k]` in BF16. `a_q8` is the DS4 q8_1
41/// activation from [`super::quantize_act_q8_1`]; `w_q2_0` is the packed weight as
42/// loaded. The tile is 128 x 128 like the Q4_K GEMM, and Q2_0's x tile row
43/// (`MMQ_MMA_TILE_X_K_Q8_0`) is the same 76 ints as Q4_K's, so it reuses
44/// [`Q4K_MMQ_SMEM`](super::Q4K_MMQ_SMEM). `kernel_wc` is used when `n` is not a
45/// multiple of 128.
46#[allow(clippy::too_many_arguments)]
47pub fn q2_0_mmq_gemm(
48    gpu: &dyn GpuBackend,
49    kernel_nc: KernelHandle,
50    kernel_wc: KernelHandle,
51    a_q8: DevicePtr,
52    w_q2_0: DevicePtr,
53    out_bf16: DevicePtr,
54    m: u32,
55    n: u32,
56    k: u32,
57    stream: u64,
58) -> Result<()> {
59    let kernel = if !n.is_multiple_of(128) {
60        kernel_wc
61    } else {
62        kernel_nc
63    };
64    KernelLaunch::new(gpu, kernel)
65        .grid([div_ceil(n, 128), div_ceil(m, 128), 1])
66        .block([32, 8, 1])
67        .shared_mem(super::q4k_mmq::Q4K_MMQ_SMEM)
68        .arg_ptr(w_q2_0)
69        .arg_ptr(a_q8)
70        .arg_ptr(out_bf16)
71        .arg_u32(n)
72        .arg_u32(m)
73        .arg_u32(k)
74        .arg_u32(k / QK2_0)
75        .arg_u32(m)
76        .arg_u32(n)
77        .launch(stream)
78}
79
80/// 2026-09-25: [`q2_0_mmq_gemm`] on a [`PackedQ2Weight`]. Returns an error unless
81/// `w.group == 128`, the group size of `block_q2_0`.
82pub fn q2_0_mmq_gemm_packed(
83    gpu: &dyn GpuBackend,
84    kernel_nc: KernelHandle,
85    kernel_wc: KernelHandle,
86    a_q8: DevicePtr,
87    w: &PackedQ2Weight,
88    out_bf16: DevicePtr,
89    m: u32,
90    stream: u64,
91) -> Result<()> {
92    anyhow::ensure!(
93        w.group == 128,
94        "Q2_0 MMQ requires group 128 (got {}); use the transient-dequant path for group 64",
95        w.group
96    );
97    q2_0_mmq_gemm(
98        gpu, kernel_nc, kernel_wc, a_q8, w.weight, out_bf16, m, w.n, w.k, stream,
99    )
100}
101
102#[cfg(test)]
103mod tests {
104    use super::*;
105
106    use half::f16;
107
108    // 2026-09-25: The block stores the weight scale `d` as fp16, so both models
109    // round it through fp16.
110    fn f32_to_f16_bits(x: f32) -> u16 {
111        f16::from_f32(x).to_bits()
112    }
113    fn f16_bits_to_f32(bits: u16) -> f32 {
114        f16::from_bits(bits).to_f32()
115    }
116
117    /// 2026-09-25: Pack an `[n, k]` code matrix (values 0..=2, weight `code - 1`)
118    /// and one scale per row and 128 codes into `block_q2_0` bytes: fp16 `d`, then
119    /// code `j` at bits `2 * (j % 4)` of byte `2 + j / 4`.
120    fn pack_q2_0(codes: &[u8], scales: &[f32], n: usize, k: usize) -> Vec<u8> {
121        assert_eq!(k % 128, 0);
122        let blocks_per_row = k / 128;
123        let mut out = vec![0u8; n * blocks_per_row * Q2_0_BLOCK_BYTES];
124        for row in 0..n {
125            for b in 0..blocks_per_row {
126                let blk = (row * blocks_per_row + b) * Q2_0_BLOCK_BYTES;
127                let dbits = f32_to_f16_bits(scales[row * blocks_per_row + b]);
128                out[blk] = (dbits & 0xff) as u8;
129                out[blk + 1] = (dbits >> 8) as u8;
130                for j in 0..128 {
131                    let c = codes[row * k + b * 128 + j] & 0x3;
132                    let byte = blk + 2 + j / 4;
133                    out[byte] |= c << (2 * (j % 4));
134                }
135            }
136        }
137        out
138    }
139
140    /// 2026-09-25: CPU model of the Q2_0 MMQ arithmetic: quantize the activation
141    /// per 32 values (`d_a = absmax / 127`, round to nearest int8), sum int8
142    /// products per 32, then scale by `d_w * d_a`. Output is row-major `[m, n]`.
143    fn mmq_cpu(
144        act: &[f32],
145        codes: &[u8],
146        scales: &[f32],
147        m: usize,
148        n: usize,
149        k: usize,
150    ) -> Vec<f32> {
151        let bpr = k / 128;
152        let ng = k / 32;
153        let mut a_q = vec![0i8; m * k];
154        let mut a_d = vec![0f32; m * ng];
155        for r in 0..m {
156            for g in 0..ng {
157                let mut amax = 0f32;
158                for t in 0..32 {
159                    amax = amax.max(act[r * k + g * 32 + t].abs());
160                }
161                let d = amax / 127.0;
162                a_d[r * ng + g] = d;
163                for t in 0..32 {
164                    let q = if d > 0.0 {
165                        (act[r * k + g * 32 + t] / d).round().clamp(-127.0, 127.0)
166                    } else {
167                        0.0
168                    };
169                    a_q[r * k + g * 32 + t] = q as i8;
170                }
171            }
172        }
173        let mut out = vec![0f32; m * n];
174        for r in 0..m {
175            for col in 0..n {
176                let mut acc = 0f32;
177                for g in 0..ng {
178                    let b = g / 4;
179                    let dw = f16_bits_to_f32(f32_to_f16_bits(scales[col * bpr + b]));
180                    let da = a_d[r * ng + g];
181                    let mut isum = 0i32;
182                    for t in 0..32 {
183                        let ki = g * 32 + t;
184                        let w = (codes[col * k + ki] & 0x3) as i32 - 1;
185                        isum += w * a_q[r * k + ki] as i32;
186                    }
187                    acc += isum as f32 * dw * da;
188                }
189                out[r * n + col] = acc;
190            }
191        }
192        out
193    }
194
195    /// 2026-09-25: FP oracle: dequantized `(code - 1) * d` times the unquantized
196    /// activation.
197    fn oracle(act: &[f32], codes: &[u8], scales: &[f32], m: usize, n: usize, k: usize) -> Vec<f32> {
198        let bpr = k / 128;
199        let mut out = vec![0f32; m * n];
200        for r in 0..m {
201            for col in 0..n {
202                let mut acc = 0f32;
203                for ki in 0..k {
204                    let b = ki / 128;
205                    let dw = f16_bits_to_f32(f32_to_f16_bits(scales[col * bpr + b]));
206                    let w = ((codes[col * k + ki] & 0x3) as i32 - 1) as f32 * dw;
207                    acc += w * act[r * k + ki];
208                }
209                out[r * n + col] = acc;
210            }
211        }
212        out
213    }
214
215    #[test]
216    fn q2_0_block_layout_matches_spec() {
217        assert_eq!(Q2_0_BLOCK_BYTES, 2 + 128 / 4);
218        assert_eq!(QK2_0, 128);
219        assert_eq!(q2_0_weight_bytes(256, 512), 256 * (512 / 128) * 34);
220        let codes: Vec<u8> = (0..128).map(|j| (j % 3) as u8).collect();
221        let packed = pack_q2_0(&codes, &[0.5], 1, 128);
222        assert_eq!(packed.len(), 34);
223        for j in 0..128usize {
224            let got = (packed[2 + j / 4] >> (2 * (j % 4))) & 0x3;
225            assert_eq!(got, (j % 3) as u8, "code {j} mispacked");
226        }
227
228        let launcher = include_str!("../../../../../kernels/gb10/qwen3.6-27b/nvfp4/q2_0_mmq.cu");
229        let vendor =
230            include_str!("../../../../../kernels/gb10/qwen3.6-27b/nvfp4/q4k_vendor/mmq.cuh");
231        assert!(launcher.contains("constexpr ggml_type type = GGML_TYPE_Q2_0;"));
232        assert!(launcher.contains("metrale_q2_0_tile<128, false>"));
233        assert!(launcher.contains("metrale_q2_0_tile<128, true>"));
234        assert!(vendor.contains("load_tiles   = load_tiles_q2_0<mmq_y, need_check>;"));
235        assert!(vendor.contains(
236            "vec_dot_mma  = vec_dot_q8_0_q8_1_mma<mmq_x, mmq_y, MMQ_Q8_1_DS_LAYOUT_DS4>;"
237        ));
238        assert_eq!(
239            vendor.matches("& 0x3) - 1;").count(),
240            4,
241            "the Q2 unpack must subtract one from all four packed codes"
242        );
243    }
244
245    #[test]
246    fn q2_0_mmq_reference_arithmetic_matches_fp_oracle() {
247        // 2026-09-25: K spans two 128-blocks, so the block boundary is crossed.
248        let (m, n, k) = (3usize, 5usize, 256usize);
249        let bpr = k / 128;
250        let mut codes = vec![0u8; n * k];
251        for (i, c) in codes.iter_mut().enumerate() {
252            *c = (((i * 2654435761usize) >> 5) % 3) as u8;
253        }
254        let mut scales = vec![0f32; n * bpr];
255        for (i, s) in scales.iter_mut().enumerate() {
256            *s = 0.015 + 0.01 * ((i % 7) as f32);
257        }
258        let mut act = vec![0f32; m * k];
259        for (i, a) in act.iter_mut().enumerate() {
260            let x = (i as f32 * 0.12345).sin();
261            *a = x * 1.7;
262        }
263
264        let mmq = mmq_cpu(&act, &codes, &scales, m, n, k);
265        let orc = oracle(&act, &codes, &scales, m, n, k);
266
267        // 2026-09-25: The two differ only by the int8 activation quantization.
268        let mut max_rel = 0f32;
269        let mut denom = 0f32;
270        let mut num = 0f32;
271        for i in 0..m * n {
272            let e = (mmq[i] - orc[i]).abs();
273            assert!(
274                e <= 0.02 * orc[i].abs().max(0.5),
275                "idx {i}: mmq {} vs oracle {}",
276                mmq[i],
277                orc[i]
278            );
279            num += e * e;
280            denom += orc[i] * orc[i];
281            let r = e / orc[i].abs().max(1e-3);
282            max_rel = max_rel.max(r);
283        }
284        let l2_rel = (num / denom.max(1e-12)).sqrt();
285        assert!(
286            l2_rel < 1e-2,
287            "Q2_0 MMQ L2 rel_err {l2_rel:.4e} exceeds 1e-2 (max pointwise {max_rel:.4e})"
288        );
289    }
290}
metrale_model_layers/layers/ops/q2_0_mmq.rs

metrale_model_layers/layers/ops/
q2_0_mmq.rs
