spark_model/layers/ops/
gemv_q2.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Launch wrappers for the native keep-packed ternary Q2_0 decode GEMV.
4//!
5//! Mirrors the `w8a16_gemv` / `w4a16_gemv` wrapper shape (`KernelLaunch`,
6//! grid `(ceil(N/4),1,1)`, block `(256,1,1)`), but the Q2_0 weight carries its
7//! fp16 scale INLINE in each `block_q2_0`, so there is no separate scale-pointer
8//! argument — the kernel reads `d` from every block. Kernel source:
9//! `kernels/gb10/common/q2_0_gemv.cu` (module stem `q2_0_gemv`).
10
11use anyhow::Result;
12use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
13use spark_runtime::kernel_args::{KernelLaunch, div_ceil};
14
15use crate::weight_map::PackedQ2Weight;
16
17/// Q2_0 GEMV (M=1 decode): `C[1,N] = A[1,K] @ dequant(B)`, weights kept packed.
18///
19/// `A` is BF16 `[1, K]`; `B` is the raw `block_q2_0` buffer; `C` is BF16
20/// `[1, N]`. Dequant `(code-1)*d` happens inside the dot-product.
21///
22/// Kernel: `q2_0_gemv(A, B, C, N, K, group)`  Grid: (ceil(N/4),1,1) Block: (256,1,1)
23pub fn q2_0_gemv(
24    gpu: &dyn GpuBackend,
25    kernel: KernelHandle,
26    input: DevicePtr,
27    weight: &PackedQ2Weight,
28    output: DevicePtr,
29    stream: u64,
30) -> Result<()> {
31    KernelLaunch::new(gpu, kernel)
32        .grid([div_ceil(weight.n, 4), 1, 1])
33        .block([256, 1, 1])
34        .arg_ptr(input)
35        .arg_ptr(weight.weight)
36        .arg_ptr(output)
37        .arg_u32(weight.n)
38        .arg_u32(weight.k)
39        .arg_u32(weight.group as u32)
40        .launch(stream)
41}
42
43/// Q2_0 batched GEMV (M=1..8 decode): `C[M,N] = A[M,K] @ dequant(B)`.
44///
45/// Reads each weight block once and accumulates across all `m` activation rows,
46/// amortizing the (2-bit) weight-byte read across the batch. `A` is BF16
47/// `[M, K]` row-major, `C` is BF16 `[M, N]` row-major.
48///
49/// Kernel: `q2_0_gemv_batchm(A, B, C, N, K, group, M)`.
50#[allow(clippy::too_many_arguments)]
51pub fn q2_0_gemv_batchm(
52    gpu: &dyn GpuBackend,
53    kernel: KernelHandle,
54    input: DevicePtr,
55    weight: &PackedQ2Weight,
56    output: DevicePtr,
57    m: u32,
58    stream: u64,
59) -> Result<()> {
60    KernelLaunch::new(gpu, kernel)
61        .grid([div_ceil(weight.n, 4), 1, 1])
62        .block([256, 1, 1])
63        .arg_ptr(input)
64        .arg_ptr(weight.weight)
65        .arg_ptr(output)
66        .arg_u32(weight.n)
67        .arg_u32(weight.k)
68        .arg_u32(weight.group as u32)
69        .arg_u32(m)
70        .launch(stream)
71}
72
73/// Dequant a packed Q2_0 weight `[N, K]` (contiguous `block_q2_0` blocks) into a
74/// pre-allocated BF16 scratch buffer `[N, K]` on `stream`, IN PLACE (no alloc,
75/// no host sync). Reuses the load-time `dequant_q2_0_gn_to_bf16` kernel
76/// (`dequant_gguf_bf16` module). Used by packed-Q2 PREFILL: dequant → transient
77/// BF16 → normal BF16 GEMM → free scratch (the resident weight stays 2-bit).
78///
79/// `n_blocks = n * (k / group)`; each block is `2 + group/4` bytes and expands
80/// to `group` BF16 elements. Kernel: grid `(n_blocks,1,1)` block `(256,1,1)`.
81#[allow(clippy::too_many_arguments)]
82pub fn dequant_q2_0_gn_to_bf16(
83    gpu: &dyn GpuBackend,
84    kernel: KernelHandle,
85    blocks: DevicePtr,
86    out: DevicePtr,
87    n: u32,
88    k: u32,
89    group: u32,
90    stream: u64,
91) -> Result<()> {
92    let n_blocks = n * (k / group);
93    let block_bytes = 2 + group / 4;
94    KernelLaunch::new(gpu, kernel)
95        .grid([n_blocks, 1, 1])
96        .block([256, 1, 1])
97        .arg_ptr(blocks)
98        .arg_ptr(out)
99        .arg_u32(n_blocks)
100        .arg_u32(group)
101        .arg_u32(block_bytes)
102        .launch(stream)
103}
104
105#[cfg(test)]
106mod tests {
107    use half::f16;
108
109    /// Build one `block_q2_0` (group @ 34/18 bytes) from `group` codes in
110    /// {0,1,2,3} plus an fp16 scale — the exact on-disk PrismML layout the
111    /// `q2_0_gemv` kernel reads: `[fp16 d @ front][group/4 bytes, 4 codes/byte,
112    /// low-bits-first]`.
113    fn pack_block(d: f32, codes: &[u8]) -> Vec<u8> {
114        let group = codes.len();
115        let mut b = Vec::with_capacity(2 + group / 4);
116        b.extend_from_slice(&f16::from_f32(d).to_le_bytes());
117        for chunk in codes.chunks(4) {
118            let mut byte = 0u8;
119            for (t, &c) in chunk.iter().enumerate() {
120                debug_assert!(c < 4);
121                byte |= (c & 3) << (2 * t as u8);
122            }
123            b.push(byte);
124        }
125        b
126    }
127
128    /// Pure Rust mirror of the kernel's dequant-in-dot-product:
129    /// `out = sum_k a[k] * (code(k)-1) * d(k/group)`, reading blocks exactly as
130    /// `q2_0_gemv.cu` does. One weight row = `k/group` contiguous blocks.
131    fn packed_dot(row_bytes: &[u8], a: &[f32], group: usize) -> f32 {
132        let block_bytes = 2 + group / 4;
133        let blocks = a.len() / group;
134        let mut acc = 0.0f32;
135        for b in 0..blocks {
136            let blk = &row_bytes[b * block_bytes..(b + 1) * block_bytes];
137            let d = f16::from_le_bytes([blk[0], blk[1]]).to_f32();
138            let qs = &blk[2..];
139            for j in 0..group {
140                let code = (qs[j >> 2] >> (2 * (j & 3))) & 3;
141                acc += a[b * group + j] * ((code as i32 - 1) as f32) * d;
142            }
143        }
144        acc
145    }
146
147    /// Golden: dequant each code to `(code-1)*d` first, then a plain dense dot.
148    /// This is the numeric oracle the on-device kernel must match; here it locks
149    /// the CPU-side bit layout + symbol mapping the kernel spec depends on.
150    fn dense_dot(codes: &[u8], d: &[f32], a: &[f32], group: usize) -> f32 {
151        (0..a.len())
152            .map(|k| a[k] * ((codes[k] as i32 - 1) as f32) * d[k / group])
153            .sum()
154    }
155
156    #[test]
157    fn q2_block_layout_reference_matches_dense_for_supported_groups() {
158        for group in [64usize, 128] {
159            let k = group * 2;
160            // Deterministic pseudo-random codes {0,1,2,3} and activations.
161            let codes: Vec<u8> = (0..k).map(|i| ((i * 7 + 3) % 4) as u8).collect();
162            let a: Vec<f32> = (0..k).map(|i| ((i % 11) as f32 - 5.0) * 0.25).collect();
163            let d = [0.0123f32, -0.0456f32];
164
165            let mut row = Vec::new();
166            row.extend(pack_block(d[0], &codes[0..group]));
167            row.extend(pack_block(d[1], &codes[group..]));
168            assert_eq!(row.len(), 2 * (2 + group / 4), "group {group}");
169
170            let got = packed_dot(&row, &a, group);
171            let want = dense_dot(&codes, &d, &a, group);
172            assert!(
173                (got - want).abs() < 1e-3,
174                "group {group}: packed {got} vs dense {want}"
175            );
176        }
177    }
178
179    #[test]
180    fn ternary_symbols_are_code_minus_one() {
181        // code {0,1,2,3} → {-1, 0, +1, +2}: asymmetric "ternary+" per the spec.
182        let group = 4usize;
183        let codes = [0u8, 1, 2, 3];
184        let a = [1.0f32, 1.0, 1.0, 1.0];
185        let d = [2.0f32];
186        let row = pack_block(d[0], &codes);
187        // sum a*(code-1)*d = (-1 + 0 + 1 + 2) * 2 = 4.
188        assert!((packed_dot(&row, &a, group) - 4.0).abs() < 1e-4);
189        // low-bits-first: byte 0 packs codes[0..4] = 0|1<<2|2<<4|3<<6 = 0xE4.
190        assert_eq!(row[2], 0xE4);
191    }
192
193    #[test]
194    fn shipped_kernels_use_the_q2_layout_and_symbol_mapping() {
195        let baseline = include_str!("../../../../../kernels/gb10/common/q2_0_gemv.cu");
196        let vector = include_str!("../../../../../kernels/gb10/common/q2_0_gemv_vec.cu");
197        let strip_comments = |source: &str| {
198            source
199                .lines()
200                .map(|line| line.split_once("//").map_or(line, |(code, _)| code))
201                .collect::<Vec<_>>()
202                .join("\n")
203        };
204        let baseline = strip_comments(baseline);
205        let vector = strip_comments(vector);
206
207        assert_eq!(
208            baseline
209                .matches("const unsigned int block_bytes = 2u + group / 4u;")
210                .count(),
211            2,
212            "baseline single-row and batch kernels must use the inline-scale layout"
213        );
214        assert_eq!(
215            baseline.matches("const float d = q2_rd_f16(blk);").count(),
216            2
217        );
218        assert_eq!(
219            baseline
220                .matches("const unsigned char* qs = blk + 2;")
221                .count(),
222            2
223        );
224        assert!(baseline.contains("acc += a * (float)(code - 1) * d;"));
225        assert!(baseline.contains("const float wv = (float)(code - 1) * d;"));
226
227        assert_eq!(
228            vector
229                .matches("const unsigned int block_bytes = 2u + group / 4u;")
230                .count(),
231            2,
232            "vector single-row and batch kernels must use the inline-scale layout"
233        );
234        assert_eq!(
235            vector.matches("const float d = q2v_rd_f16(blk);").count(),
236            2
237        );
238        assert_eq!(vector.matches("q2v_rd_u32(blk + 2 + jg * 4u)").count(), 2);
239        assert!(vector.contains("acc += a[j] * (float)(code - 1) * d;"));
240        assert!(vector.contains("wv[j] = (float)((int)((codes >> (2 * j)) & 3u) - 1) * d;"));
241    }
242}
spark_model/layers/ops/gemv_q2.rs

spark_model/layers/ops/
gemv_q2.rs
