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

// Common WGSL declarations for cera kernels

#ifndef COMMON_DECLS_TMPL
#define COMMON_DECLS_TMPL

// Max workgroups per dimension (standard wgpu limit)
const MAX_WG: u32 = 65535u;

// Flatten a 2D workgroup ID into a linear index.
// Used when dispatches exceed MAX_WG in the X dimension.
fn get_wid(wid: vec3<u32>) -> u32 {
    return wid.x + wid.y * MAX_WG;
}

// Tree-reduce `shared_sum[0..256]` in-place; result lands in `shared_sum[0]`.
// Caller must have already populated `shared_sum[tid]` and issued a
// `workgroupBarrier()` so all writes are visible.
// NOTE: Requires `var<workgroup> shared_sum: array<f32, 256>;` to be defined
// in the calling shader.
#ifdef WG_SUM_REDUCE
fn workgroup_sum_reduce(tid: u32) {
    if tid < 128u { shared_sum[tid] += shared_sum[tid + 128u]; }
    workgroupBarrier();
    if tid < 64u { shared_sum[tid] += shared_sum[tid + 64u]; }
    workgroupBarrier();
    if tid < 32u { shared_sum[tid] += shared_sum[tid + 32u]; }
    workgroupBarrier();
    if tid < 16u { shared_sum[tid] += shared_sum[tid + 16u]; }
    workgroupBarrier();
    if tid < 8u { shared_sum[tid] += shared_sum[tid + 8u]; }
    workgroupBarrier();
    if tid < 4u { shared_sum[tid] += shared_sum[tid + 4u]; }
    workgroupBarrier();
    if tid < 2u { shared_sum[tid] += shared_sum[tid + 2u]; }
    workgroupBarrier();
    if tid < 1u { shared_sum[tid] += shared_sum[tid + 1u]; }
    workgroupBarrier();
}
#endif

// RoPE: Compute rotary angle for a given position and dimension.
fn rope_angle(pos: u32, d: u32, head_dim: u32, freq_base: f32) -> f32 {
    return f32(pos) * pow(freq_base, -2.0 * f32(d) / f32(head_dim));
}

// RoPE: Rotate a pair of values by a given angle.
fn rotate_rope(x0: f32, x1: f32, angle: f32) -> vec2<f32> {
    let cos_a = cos(angle);
    let sin_a = sin(angle);
    return vec2<f32>(
        x0 * cos_a - x1 * sin_a,
        x0 * sin_a + x1 * cos_a
    );
}

#ifdef BYTE_HELPERS
fn get_byte(value: u32, index: u32) -> u32 {
    return (value >> (index * 8)) & 0xFF;
}
#endif

// Q8_0 dequant helpers. Shared by gemv_q8_0 / gemm_q8_0.
//
// Contract: the including shader MUST declare the Q8_0 weight buffer as
//   @group(0) @binding(0) var<storage, read> a: array<u32>;
// (these helpers index `a` directly, relying on WGSL module-scope order
// independence).
//
// Q8_0 block layout (34 bytes per 32 elements):
//   bytes 0-1:  f16 scale
//   bytes 2-33: 32 signed i8 quants
#ifdef Q8_0_HELPERS
// Read a (possibly unaligned) u32 from the byte-addressed `a` buffer.
//
// When `byte_offset` is not 4-aligned this also touches `a[word_idx + 1]`.
// That index is always in bounds: a Q8_0 row is `nb * 34` bytes and the
// u32 buffer length is `ceil(row_bytes * m / 4)`. Because 34 is not a
// multiple of 4, rounding up always leaves at least one extra u32 of
// slack past the final quant word, so the +1 read never exceeds the
// buffer. Do not call this on a buffer whose length is not derived from
// whole Q8_0 rows.
fn get_u32_at(byte_offset: u32) -> u32 {
    let word_idx = byte_offset / 4u;
    let shift = (byte_offset & 3u) * 8u;
    let lo = a[word_idx];
    if shift == 0u {
        return lo;
    }
    let hi = a[word_idx + 1u];
    return (lo >> shift) | (hi << (32u - shift));
}

// Dot one Q8_0 block (32 weights) of row `row`, block `bi`, against the
// 32 staged activations in `xl`. `row_bytes` is the per-row byte stride
// (`nb * 34`). Returns the scaled partial sum for this block.
fn process_block_q8_0(
    row: u32,
    bi: u32,
    row_bytes: u32,
    xl: ptr<function, array<f32, 32>>,
) -> f32 {
    let block_byte = row * row_bytes + bi * 34u;
    let scale_bits = get_u32_at(block_byte) & 0xFFFFu;
    let scale = unpack2x16float(scale_bits).x;

    var sum = 0.0;
    for (var i = 0u; i < 32u; i += 4u) {
        let packed = get_u32_at(block_byte + 2u + i);
        sum += f32(bitcast<i32>((packed & 0x000000FFu) << 24u) >> 24u) * (*xl)[i + 0u];
        sum += f32(bitcast<i32>((packed & 0x0000FF00u) << 16u) >> 24u) * (*xl)[i + 1u];
        sum += f32(bitcast<i32>((packed & 0x00FF0000u) << 8u) >> 24u) * (*xl)[i + 2u];
        sum += f32(bitcast<i32>(packed & 0xFF000000u) >> 24u) * (*xl)[i + 3u];
    }

    return sum * scale;
}

#endif

#endif // COMMON_DECLS_TMPL