cera 0.2.7

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

#include <metal_stdlib>
using namespace metal;

// Element-wise operations on f32 buffers.
// 256 threads per threadgroup. Dispatch ceil(n/256) threadgroups.

struct Params { uint n; uint _pad; };

kernel void memcpy_f32(
    const device float* src [[buffer(0)]],
    device float* dst [[buffer(1)]],
    constant Params& params [[buffer(2)]],
    uint gid [[thread_position_in_grid]]
) {
    if (gid >= params.n) return;
    dst[gid] = src[gid];
}

kernel void add_inplace(
    device float* a [[buffer(0)]],
    const device float* b [[buffer(1)]],
    constant Params& params [[buffer(2)]],
    uint gid [[thread_position_in_grid]]
) {
    if (gid >= params.n) return;
    a[gid] = a[gid] + b[gid];
}

kernel void mul_inplace(
    device float* a [[buffer(0)]],
    const device float* b [[buffer(1)]],
    constant Params& params [[buffer(2)]],
    uint gid [[thread_position_in_grid]]
) {
    if (gid >= params.n) return;
    a[gid] = a[gid] * b[gid];
}

// Out-of-place mul: dst[i] = a[i] * b[i]. Lets callers pass a/b with byte
// offsets into a bigger buffer, avoiding two prior memcpy dispatches.
kernel void mul_out(
    const device float* a [[buffer(0)]],
    const device float* b [[buffer(1)]],
    device float* dst [[buffer(2)]],
    constant Params& params [[buffer(3)]],
    uint gid [[thread_position_in_grid]]
) {
    if (gid >= params.n) return;
    dst[gid] = a[gid] * b[gid];
}

// Cast f32 → f16 (for writing to f16 KV cache).
kernel void cast_f32_to_f16(
    const device float* src [[buffer(0)]],
    device half* dst [[buffer(1)]],
    constant Params& params [[buffer(2)]],
    uint gid [[thread_position_in_grid]]
) {
    if (gid >= params.n) return;
    dst[gid] = half(src[gid]);
}

kernel void silu_mul_inplace(
    device float* a [[buffer(0)]],
    const device float* b [[buffer(1)]],
    constant Params& params [[buffer(2)]],
    uint gid [[thread_position_in_grid]]
) {
    if (gid >= params.n) return;
    float g = a[gid];
    a[gid] = (g / (1.0f + exp(-g))) * b[gid];
}

// Params carrying a scalar multiplier (raw f32 bits) for the Granite-style
// scaled residual add and logit scaling. Identity (scale = 1.0) for every
// other arch.
struct ScaleParams { uint n; uint scale_bits; };

// In-place scaled add: a[i] += scale * b[i]. Used for the decode-path residual
// adds when residual_multiplier != 1.0 (Granite).
kernel void scaled_add_inplace(
    device float* a [[buffer(0)]],
    const device float* b [[buffer(1)]],
    constant ScaleParams& params [[buffer(2)]],
    uint gid [[thread_position_in_grid]]
) {
    if (gid >= params.n) return;
    a[gid] = a[gid] + as_type<float>(params.scale_bits) * b[gid];
}

// In-place scale: a[i] *= scale. Used to divide final logits by Granite's
// logit scaling factor (caller passes 1/logit_scale).
kernel void scale_f32(
    device float* a [[buffer(0)]],
    constant ScaleParams& params [[buffer(1)]],
    uint gid [[thread_position_in_grid]]
) {
    if (gid >= params.n) return;
    a[gid] = a[gid] * as_type<float>(params.scale_bits);
}