cera 0.2.2

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

#include <metal_stdlib>
using namespace metal;

// F32 GEMV: y[row] = dot(A[row, :], x).
// One threadgroup per row (2D dispatch: row = tg.x + tg.y * 65535).
// 32 threads per threadgroup (one simdgroup).

struct Params { uint m; uint k; };

kernel void gemv_f32(
    const device float* a [[buffer(0)]],
    const device float* x [[buffer(1)]],
    device float* y [[buffer(2)]],
    constant Params& params [[buffer(3)]],
    uint3 tid_v [[thread_position_in_threadgroup]],
    uint3 tg_id [[threadgroup_position_in_grid]]
) {
    uint tid = tid_v.x;
    uint row = tg_id.x + tg_id.y * 65535u;
    uint m = params.m;
    uint k = params.k;
    if (row >= m) return;

    uint row_offset = row * k;
    float partial = 0.0f;
    for (uint col = tid; col < k; col += 32u) {
        partial += a[row_offset + col] * x[col];
    }
    float total = simd_sum(partial);
    if (tid == 0u) y[row] = total;
}

// F32 GEMV with accumulate: y[row] += dot(A[row, :], x). Same layout/dispatch as
// `gemv_f32`; used for the LoRA up-projection epilogue (out += B_scaled ยท tmp),
// where `scale` is pre-folded into the uploaded B so no separate scale pass is
// needed.
kernel void gemv_f32_accum(
    const device float* a [[buffer(0)]],
    const device float* x [[buffer(1)]],
    device float* y [[buffer(2)]],
    constant Params& params [[buffer(3)]],
    uint3 tid_v [[thread_position_in_threadgroup]],
    uint3 tg_id [[threadgroup_position_in_grid]]
) {
    uint tid = tid_v.x;
    uint row = tg_id.x + tg_id.y * 65535u;
    uint m = params.m;
    uint k = params.k;
    if (row >= m) return;

    uint row_offset = row * k;
    float partial = 0.0f;
    for (uint col = tid; col < k; col += 32u) {
        partial += a[row_offset + col] * x[col];
    }
    float total = simd_sum(partial);
    if (tid == 0u) y[row] += total;
}