cera 0.5.5

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

#include <metal_stdlib>
using namespace metal;

// Batched qk_norm_rope: processes N tokens in one dispatch.
// Dispatch: (n_tokens * (n_heads + n_kv_heads)) TGs × 256 threads.

struct BatchParams {
    uint start_pos;
    uint n_tokens;
    uint n_heads;
    uint n_kv_heads;
    uint head_dim;
    uint eps_bits;
    uint freq_base_bits;
    uint rope_type;
    uint q_stride;  // floats between Q vectors (typically hs)
    uint k_stride;  // floats between K vectors (typically kv_dim)
    uint has_freq_factors; // 1 ⇒ divide each pair's angle by freq_factors[d] (Llama-3)
    uint has_qk_norm;      // 1 ⇒ per-head RMSnorm (Qwen3/LFM2); 0 ⇒ rope-only (LLaMA/Qwen2)
};

inline void head_rmsnorm(
    device float* buf,
    const device float* w,
    threadgroup float* scratch,
    uint tid,
    uint head_dim,
    float eps
) {
    float partial = 0.0f;
    for (uint i = tid; i < head_dim; i += 256u) {
        float v = buf[i];
        scratch[i] = v;
        partial += v * v;
    }
    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);
    for (uint i = tid; i < head_dim; i += 256u) {
        buf[i] = scratch[i] * inv_rms * w[i];
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);
}

inline void head_rope(
    device float* buf,
    uint tid,
    uint head_dim,
    uint pos,
    float freq_base,
    uint rope_type,
    const device float* freq_factors, // [head_dim/2], valid only if has_freq_factors
    uint has_freq_factors
) {
    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));
        // Llama-3 long-context RoPE scaling (mirrors the decode kernel).
        if (has_freq_factors != 0u) {
            theta = theta / freq_factors[d];
        }
        float cos_a = cos(theta);
        float sin_a = sin(theta);
        if (rope_type == 0u) {
            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 {
            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_batch(
    device float* q_batch [[buffer(0)]],
    device float* k_batch [[buffer(1)]],
    const device float* q_norm_w [[buffer(2)]],
    const device float* k_norm_w [[buffer(3)]],
    constant BatchParams& params [[buffer(4)]],
    const device float* freq_factors [[buffer(5)]],
    uint tid [[thread_position_in_threadgroup]],
    uint tg_idx [[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 rope_type = params.rope_type;
    uint has_ff = params.has_freq_factors;
    uint has_qk_norm = params.has_qk_norm;

    uint heads_per_token = n_heads + n_kv_heads;
    uint token_idx = tg_idx / heads_per_token;
    uint head = tg_idx % heads_per_token;
    uint pos = params.start_pos + token_idx;

    threadgroup float scratch[264];

    if (head < n_heads) {
        device float* q_head = q_batch + token_idx * params.q_stride + head * head_dim;
        if (has_qk_norm != 0u) {
            head_rmsnorm(q_head, q_norm_w, scratch, tid, head_dim, eps);
        }
        head_rope(q_head, tid, head_dim, pos, freq_base, rope_type, freq_factors, has_ff);
    } else {
        uint kh = head - n_heads;
        if (kh < n_kv_heads) {
            device float* k_head = k_batch + token_idx * params.k_stride + kh * head_dim;
            if (has_qk_norm != 0u) {
                head_rmsnorm(k_head, k_norm_w, scratch, tid, head_dim, eps);
            }
            head_rope(k_head, tid, head_dim, pos, freq_base, rope_type, freq_factors, has_ff);
        }
    }
}