metrale_model_layers/layers/ops/
gemv_q2.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the packed Q2_0 GEMVs, which keep the 2-bit weight packed, and for its dequant to BF16.
4//!
5//! Owner: model-layers ops.
6//! Invariants:
7//! - A Q2_0 weight carries its fp16 scale inside each `block_q2_0`, so no launcher here takes
8//!   a separate scale pointer (kernels/gb10/common/q2_0_gemv.cu).
9//! - The GEMV grids are `ceil(N/4)` blocks of 256 threads: 4 outputs per block, 64 threads per
10//!   output (`N_PER_BLOCK`, `BLOCK_SIZE` in that file).
11
12use anyhow::Result;
13use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
14use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
15
16use crate::weight_map::PackedQ2Weight;
17
18/// 2026-09-25: Q2_0 GEMV for M=1 decode: `C[1,N] = A[1,K] @ dequant(B)^T`. `A` and `C` are
19/// BF16; `B` is the raw `block_q2_0` buffer, and each weight is dequantized as `(code-1)*d`
20/// inside the dot product.
21pub fn q2_0_gemv(
22    gpu: &dyn GpuBackend,
23    kernel: KernelHandle,
24    input: DevicePtr,
25    weight: &PackedQ2Weight,
26    output: DevicePtr,
27    stream: u64,
28) -> Result<()> {
29    KernelLaunch::new(gpu, kernel)
30        .grid([div_ceil(weight.n, 4), 1, 1])
31        .block([256, 1, 1])
32        .arg_ptr(input)
33        .arg_ptr(weight.weight)
34        .arg_ptr(output)
35        .arg_u32(weight.n)
36        .arg_u32(weight.k)
37        .arg_u32(weight.group as u32)
38        .launch(stream)
39}
40
41/// 2026-09-25: Batched Q2_0 GEMV: `C[M,N] = A[M,K] @ dequant(B)^T`, `A` and `C` row-major
42/// BF16. Each weight block is read once and applied to all `m` rows.
43///
44/// `m` must be in `1..=8`: the kernel keeps `acc[MAX_M]` with `MAX_M = 8` and indexes it up to
45/// `M`. This launcher does not check it.
46#[allow(clippy::too_many_arguments)]
47pub fn q2_0_gemv_batchm(
48    gpu: &dyn GpuBackend,
49    kernel: KernelHandle,
50    input: DevicePtr,
51    weight: &PackedQ2Weight,
52    output: DevicePtr,
53    m: u32,
54    stream: u64,
55) -> Result<()> {
56    KernelLaunch::new(gpu, kernel)
57        .grid([div_ceil(weight.n, 4), 1, 1])
58        .block([256, 1, 1])
59        .arg_ptr(input)
60        .arg_ptr(weight.weight)
61        .arg_ptr(output)
62        .arg_u32(weight.n)
63        .arg_u32(weight.k)
64        .arg_u32(weight.group as u32)
65        .arg_u32(m)
66        .launch(stream)
67}
68
69/// 2026-09-25: Dequantize a packed Q2_0 weight `[N, K]` into the caller's BF16 buffer `[N, K]`
70/// on `stream`, with no allocation and no host sync (kernel `dequant_q2_0_gn_to_bf16` in
71/// kernels/gb10/common/dequant_gguf_bf16.cu). Packed-Q2 prefill uses it to feed a BF16 GEMM
72/// while the resident weight stays packed (`dense_ffn.rs`, `qwen3_attention/prefill_weights.rs`).
73///
74/// `k` must be a multiple of `group`. One CTA per `block_q2_0`: `n * (k / group)` blocks of
75/// `2 + group/4` bytes, each expanding to `group` BF16 values.
76#[allow(clippy::too_many_arguments)]
77pub fn dequant_q2_0_gn_to_bf16(
78    gpu: &dyn GpuBackend,
79    kernel: KernelHandle,
80    blocks: DevicePtr,
81    out: DevicePtr,
82    n: u32,
83    k: u32,
84    group: u32,
85    stream: u64,
86) -> Result<()> {
87    let n_blocks = n * (k / group);
88    let block_bytes = 2 + group / 4;
89    KernelLaunch::new(gpu, kernel)
90        .grid([n_blocks, 1, 1])
91        .block([256, 1, 1])
92        .arg_ptr(blocks)
93        .arg_ptr(out)
94        .arg_u32(n_blocks)
95        .arg_u32(group)
96        .arg_u32(block_bytes)
97        .launch(stream)
98}
99
100#[cfg(test)]
101mod tests {
102    use half::f16;
103
104    /// 2026-09-25: Build one `block_q2_0` from `group` codes in {0,1,2,3} and an fp16 scale, in
105    /// the layout the `q2_0_gemv` kernel reads: `[fp16 d][group/4 bytes, 4 codes per byte,
106    /// low bits first]` (34 bytes at group 128, 18 at group 64).
107    fn pack_block(d: f32, codes: &[u8]) -> Vec<u8> {
108        let group = codes.len();
109        let mut b = Vec::with_capacity(2 + group / 4);
110        b.extend_from_slice(&f16::from_f32(d).to_le_bytes());
111        for chunk in codes.chunks(4) {
112            let mut byte = 0u8;
113            for (t, &c) in chunk.iter().enumerate() {
114                debug_assert!(c < 4);
115                byte |= (c & 3) << (2 * t as u8);
116            }
117            b.push(byte);
118        }
119        b
120    }
121
122    /// 2026-09-25: Host model of the kernel's dot product,
123    /// `out = sum_k a[k] * (code(k)-1) * d(k/group)`, walking one weight row as `k/group`
124    /// contiguous blocks as `q2_0_gemv.cu` does.
125    fn packed_dot(row_bytes: &[u8], a: &[f32], group: usize) -> f32 {
126        let block_bytes = 2 + group / 4;
127        let blocks = a.len() / group;
128        let mut acc = 0.0f32;
129        for b in 0..blocks {
130            let blk = &row_bytes[b * block_bytes..(b + 1) * block_bytes];
131            let d = f16::from_le_bytes([blk[0], blk[1]]).to_f32();
132            let qs = &blk[2..];
133            for j in 0..group {
134                let code = (qs[j >> 2] >> (2 * (j & 3))) & 3;
135                acc += a[b * group + j] * ((code as i32 - 1) as f32) * d;
136            }
137        }
138        acc
139    }
140
141    /// 2026-09-25: Oracle: dequantize each code to `(code-1)*d` first, then a plain dense dot.
142    /// Comparing it with `packed_dot` checks the block layout and symbol mapping.
143    fn dense_dot(codes: &[u8], d: &[f32], a: &[f32], group: usize) -> f32 {
144        (0..a.len())
145            .map(|k| a[k] * ((codes[k] as i32 - 1) as f32) * d[k / group])
146            .sum()
147    }
148
149    #[test]
150    fn q2_block_layout_reference_matches_dense_for_supported_groups() {
151        for group in [64usize, 128] {
152            let k = group * 2;
153            let codes: Vec<u8> = (0..k).map(|i| ((i * 7 + 3) % 4) as u8).collect();
154            let a: Vec<f32> = (0..k).map(|i| ((i % 11) as f32 - 5.0) * 0.25).collect();
155            let d = [0.0123f32, -0.0456f32];
156
157            let mut row = Vec::new();
158            row.extend(pack_block(d[0], &codes[0..group]));
159            row.extend(pack_block(d[1], &codes[group..]));
160            assert_eq!(row.len(), 2 * (2 + group / 4), "group {group}");
161
162            let got = packed_dot(&row, &a, group);
163            let want = dense_dot(&codes, &d, &a, group);
164            assert!(
165                (got - want).abs() < 1e-3,
166                "group {group}: packed {got} vs dense {want}"
167            );
168        }
169    }
170
171    #[test]
172    fn ternary_symbols_are_code_minus_one() {
173        // 2026-09-25: code {0,1,2,3} maps to {-1, 0, +1, +2}.
174        let group = 4usize;
175        let codes = [0u8, 1, 2, 3];
176        let a = [1.0f32, 1.0, 1.0, 1.0];
177        let d = [2.0f32];
178        let row = pack_block(d[0], &codes);
179        // 2026-09-25: sum a*(code-1)*d = (-1 + 0 + 1 + 2) * 2 = 4.
180        assert!((packed_dot(&row, &a, group) - 4.0).abs() < 1e-4);
181        // 2026-09-25: Low bits first: byte 0 packs codes[0..4] as 0|1<<2|2<<4|3<<6 = 0xE4.
182        assert_eq!(row[2], 0xE4);
183    }
184
185    #[test]
186    fn shipped_kernels_use_the_q2_layout_and_symbol_mapping() {
187        let baseline = include_str!("../../../../../kernels/gb10/common/q2_0_gemv.cu");
188        let vector = include_str!("../../../../../kernels/gb10/common/q2_0_gemv_vec.cu");
189        let strip_comments = |source: &str| {
190            source
191                .lines()
192                .map(|line| line.split_once("//").map_or(line, |(code, _)| code))
193                .collect::<Vec<_>>()
194                .join("\n")
195        };
196        let baseline = strip_comments(baseline);
197        let vector = strip_comments(vector);
198
199        assert_eq!(
200            baseline
201                .matches("const unsigned int block_bytes = 2u + group / 4u;")
202                .count(),
203            2,
204            "baseline single-row and batch kernels must use the inline-scale layout"
205        );
206        assert_eq!(
207            baseline.matches("const float d = q2_rd_f16(blk);").count(),
208            2
209        );
210        assert_eq!(
211            baseline
212                .matches("const unsigned char* qs = blk + 2;")
213                .count(),
214            2
215        );
216        assert!(baseline.contains("acc += a * (float)(code - 1) * d;"));
217        assert!(baseline.contains("const float wv = (float)(code - 1) * d;"));
218
219        assert_eq!(
220            vector
221                .matches("const unsigned int block_bytes = 2u + group / 4u;")
222                .count(),
223            2,
224            "vector single-row and batch kernels must use the inline-scale layout"
225        );
226        assert_eq!(
227            vector.matches("const float d = q2v_rd_f16(blk);").count(),
228            2
229        );
230        assert_eq!(vector.matches("q2v_rd_u32(blk + 2 + jg * 4u)").count(), 2);
231        assert!(vector.contains("acc += a[j] * (float)(code - 1) * d;"));
232        assert!(vector.contains("wv[j] = (float)((int)((codes >> (2 * j)) & 3u) - 1) * d;"));
233    }
234}
metrale_model_layers/layers/ops/gemv_q2.rs

metrale_model_layers/layers/ops/
gemv_q2.rs
