cera 0.1.0

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

#include <metal_stdlib>
using namespace metal;

// Fused: per-head RMSnorm + RoPE for Q and K (in a GQA attention layer).
// Replaces 3 dispatches (per_head_rmsnorm Q, per_head_rmsnorm K, rope) with 1.
//
// Dispatch: max(n_heads, n_kv_heads) threadgroups × 256 threads.
//   TGs 0..n_kv_heads-1 process BOTH a Q head AND the corresponding K head.
//   TGs n_kv_heads..n_heads-1 process only their Q head.
//
// K is accessed via a byte offset into k_cache (typically pos * kv_dim * 4).

struct Params {
    uint pos;
    uint n_heads;
    uint n_kv_heads;
    uint head_dim;
    uint eps_bits;
    uint freq_base_bits;
    uint rope_type; // 0 = NeoX (pairs at [i, i+half]), 1 = interleaved (pairs at [2i, 2i+1])
};

inline void head_rmsnorm(
    device float* buf,          // buf[0..head_dim]
    const device float* w,      // w[0..head_dim]
    threadgroup float* scratch, // at least [head_dim + 8] floats
    uint tid,
    uint head_dim,
    float eps
) {
    // Sum of squares.
    float partial = 0.0f;
    for (uint i = tid; i < head_dim; i += 256u) {
        float v = buf[i];
        scratch[i] = v;  // cache x into threadgroup mem for phase 2
        partial += v * v;
    }
    // Two-stage simd reduction.
    threadgroup_barrier(mem_flags::mem_threadgroup);
    uint simd_lane = tid & 31u;
    uint simd_id = tid >> 5u;
    threadgroup float* sg_val = scratch + head_dim;
    float sg_sum = simd_sum(partial);
    if (simd_lane == 0u) sg_val[simd_id] = sg_sum;
    threadgroup_barrier(mem_flags::mem_threadgroup);
    if (simd_id == 0u) {
        float v = simd_lane < 8u ? sg_val[simd_lane] : 0.0f;
        float total = simd_sum(v);
        if (simd_lane == 0u) sg_val[0] = total;
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);
    float inv_rms = 1.0f / sqrt(sg_val[0] / float(head_dim) + eps);
    // Apply normalization back to buf.
    for (uint i = tid; i < head_dim; i += 256u) {
        buf[i] = scratch[i] * inv_rms * w[i];
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);
}

// Apply RoPE to buf[0..head_dim] in place.
// rope_type 0 = NeoX: pairs at [d, d + half_dim]
// rope_type 1 = interleaved (LLAMA_ROPE_TYPE_NORM): pairs at [2d, 2d+1]
inline void head_rope(
    device float* buf,
    uint tid,
    uint head_dim,
    uint pos,
    float freq_base,
    uint rope_type
) {
    uint half_dim = head_dim / 2u;
    // theta[d] = pos * freq_base^(-2d/head_dim). Compute with powr for O(1)
    // per thread instead of an O(d) iterative multiplication loop.
    float theta_scale = powr(freq_base, -2.0f / float(head_dim));
    for (uint d = tid; d < half_dim; d += 256u) {
        float theta = float(pos) * powr(theta_scale, float(d));
        float cos_a = cos(theta);
        float sin_a = sin(theta);
        if (rope_type == 0u) {
            // NeoX: pairs at [d, d + half_dim]
            float x0 = buf[d];
            float x1 = buf[d + half_dim];
            buf[d] = x0 * cos_a - x1 * sin_a;
            buf[d + half_dim] = x0 * sin_a + x1 * cos_a;
        } else {
            // Interleaved: pairs at [2d, 2d+1]
            float x0 = buf[2u * d];
            float x1 = buf[2u * d + 1u];
            buf[2u * d] = x0 * cos_a - x1 * sin_a;
            buf[2u * d + 1u] = x0 * sin_a + x1 * cos_a;
        }
    }
}

kernel void qk_norm_rope(
    device float* q [[buffer(0)]],        // [n_heads × head_dim]
    device float* k_cache [[buffer(1)]],  // [seq_len × kv_dim], offset to current row
    const device float* q_norm_w [[buffer(2)]],
    const device float* k_norm_w [[buffer(3)]],
    constant Params& params [[buffer(4)]],
    uint tid [[thread_position_in_threadgroup]],
    uint head [[threadgroup_position_in_grid]]
) {
    uint n_heads = params.n_heads;
    uint n_kv_heads = params.n_kv_heads;
    uint head_dim = params.head_dim;
    float eps = as_type<float>(params.eps_bits);
    float freq_base = as_type<float>(params.freq_base_bits);
    uint pos = params.pos;

    // Threadgroup scratch: head_dim floats for x-cache + 8 for simd reduction.
    // head_dim ≤ 256 for LFM2/Llama-family models.
    threadgroup float scratch[264];

    // Dispatch: (n_heads + n_kv_heads) TGs. First n_heads handle Q heads;
    // remaining n_kv_heads handle K heads. Gives every TG equal work (no
    // load imbalance from doing 2× work in the first few TGs).
    uint rope_type = params.rope_type;

    if (head < n_heads) {
        device float* q_head = q + head * head_dim;
        head_rmsnorm(q_head, q_norm_w, scratch, tid, head_dim, eps);
        head_rope(q_head, tid, head_dim, pos, freq_base, rope_type);
    } else {
        uint kh = head - n_heads;
        if (kh < n_kv_heads) {
            device float* k_head = k_cache + kh * head_dim;
            head_rmsnorm(k_head, k_norm_w, scratch, tid, head_dim, eps);
            head_rope(k_head, tid, head_dim, pos, freq_base, rope_type);
        }
    }
}