cera 0.3.1

Rust-native LLM inference engine
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141

// Q6_K GEMV — Metal kernel layout with per-byte reads.
// Wired into the wgpu GEMV dispatch (Q6K weights stay quantized in VRAM,
// ~4.9× smaller than dequantizing to f32: 210 B / 256 elems ≈ 0.82 B/elem
// vs 4 B/elem). Compute is bound by per-byte u32 load+shift+mask overhead and
// regressed vs f32 on macOS wgpu in earlier measurement, so the primary win is
// VRAM/bandwidth (matters most on mobile Adreno/Mali); byte-extraction
// throughput is a future optimization.
//
// NR=2 rows per WG, 32 threads. Dispatch: ceil(m/2).

@group(0) @binding(0) var<storage, read> a: array<u32>;
@group(0) @binding(1) var<storage, read> x: array<f32>;
@group(0) @binding(2) var<storage, read_write> y: array<f32>;
@group(0) @binding(3) var<storage, read> params: vec2<u32>;

// `get_wid` flattens the 2-D dispatch grid so m > 65535*NR rows still map to
// distinct rows (gemv_workgroups folds the row overflow into wid.y).
#include "common_decls.tmpl"

const QK_K: u32 = 256u;
const Q6K_BYTES: u32 = 210u;
const NR: u32 = 2u;
const WG_SIZE: u32 = 32u;

var<workgroup> partials: array<f32, 64>;

fn rb(off: u32) -> u32 {
    return (a[off / 4u] >> ((off % 4u) * 8u)) & 0xFFu;
}

fn ri8(off: u32) -> i32 {
    let b = rb(off);
    return i32(b) - select(0, 256, (b & 0x80u) != 0u);
}

fn rf16(off: u32) -> f32 {
    let lo = rb(off);
    let hi = rb(off + 1u);
    return unpack2x16float(lo | (hi << 8u)).x;
}

@compute @workgroup_size(32, 1, 1)
fn gemv_q6_k(
    @builtin(local_invocation_id) lid: vec3<u32>,
    @builtin(workgroup_id) wid: vec3<u32>,
) {
    let m = params.x;
    let k = params.y;
    let nb = k / QK_K;
    let row_bytes = nb * Q6K_BYTES;
    let tiisg = lid.x;
    let first_row = get_wid(wid) * NR;

    let tid_l = tiisg / 2u;
    let ix  = tiisg & 1u;
    let ip  = tid_l >> 3u;
    let il  = tid_l & 7u;
    let l0  = 4u * il;
    let is_off = 8u * ip + l0 / 16u;

    let y_offset   = 128u * ip + l0;
    let q_offset_l = 64u * ip + l0;
    let q_offset_h = 32u * ip + l0;

    var sumf0: f32 = 0.0;
    var sumf1: f32 = 0.0;

    var b = ix;
    while b < nb {
        let yb = b * QK_K + y_offset;
        var yl: array<f32, 16>;
        for (var l = 0u; l < 4u; l += 1u) {
            yl[4u * l + 0u] = x[yb + l + 0u];
            yl[4u * l + 1u] = x[yb + l + 32u];
            yl[4u * l + 2u] = x[yb + l + 64u];
            yl[4u * l + 3u] = x[yb + l + 96u];
        }

        for (var row = 0u; row < NR; row += 1u) {
            // Skip the weight reads for out-of-range rows: on an odd `m` the tail
            // workgroup's second row (`first_row + 1 == m`) would otherwise index
            // `bb = m * row_bytes + ...` past the end of the weight buffer. The
            // writes below are already guarded; sumf1 stays 0 for the skipped row.
            if first_row + row >= m {
                continue;
            }
            let bb = (first_row + row) * row_bytes + b * Q6K_BYTES;
            let ql1 = bb + q_offset_l;
            let ql2 = ql1 + 32u;
            let qh  = bb + 128u + q_offset_h;
            let sc  = bb + 192u + is_off;
            let d_off = bb + 208u;

            var sums = vec4<f32>(0.0);
            for (var l = 0u; l < 4u; l += 1u) {
                let q1 = rb(ql1 + l);
                let q2 = rb(ql2 + l);
                let qhv = rb(qh + l);

                let q6_1 = i32((q1 & 0x0Fu) | ((qhv & 0x03u) << 4u)) - 32;
                let q6_2 = i32((q2 & 0x0Fu) | ((qhv & 0x0Cu) << 2u)) - 32;
                let q6_3 = i32((q1 >> 4u)   | ( qhv & 0x30u        )) - 32;
                let q6_4 = i32((q2 >> 4u)   | ((qhv & 0xC0u) >> 2u )) - 32;

                sums[0] += yl[4u * l + 0u] * f32(q6_1);
                sums[1] += yl[4u * l + 1u] * f32(q6_2);
                sums[2] += yl[4u * l + 2u] * f32(q6_3);
                sums[3] += yl[4u * l + 3u] * f32(q6_4);
            }

            let dblk = rf16(d_off);
            let s0 = f32(ri8(sc));
            let s2 = f32(ri8(sc + 2u));
            let s4 = f32(ri8(sc + 4u));
            let s6 = f32(ri8(sc + 6u));
            let row_sum = dblk * (sums[0] * s0 + sums[1] * s2 + sums[2] * s4 + sums[3] * s6);

            if row == 0u { sumf0 += row_sum; }
            else { sumf1 += row_sum; }
        }
        b += 2u;
    }

    partials[0u * WG_SIZE + tiisg] = sumf0;
    partials[1u * WG_SIZE + tiisg] = sumf1;
    workgroupBarrier();
    for (var stride = WG_SIZE / 2u; stride > 0u; stride = stride / 2u) {
        if tiisg < stride {
            for (var r = 0u; r < NR; r += 1u) {
                let idx = r * WG_SIZE + tiisg;
                partials[idx] += partials[idx + stride];
            }
        }
        workgroupBarrier();
    }

    if tiisg == 0u {
        if first_row < m { y[first_row] = partials[0u * WG_SIZE]; }
        if first_row + 1u < m { y[first_row + 1u] = partials[1u * WG_SIZE]; }
    }
}