cera 0.2.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

// tanh-approximation GELU, in-place over an f32 buffer.
//   gelu(x) = 0.5 * x * (1 + tanh(sqrt(2/π) * (x + 0.044715 * x^3)))
//
// Mirrors `cpu::gelu_inplace` (ggml's default GELU), which is what CLIP-family
// ViTs trained with `clip.use_gelu = true` expect. NOT the erf form.
//
// Dispatch: (ceil(n / 256), 1, 1) workgroups.
//
// Bind group 0:
//   @binding(0) x: array<f32>      (read-write, activated in-place)
//   @binding(1) params: vec2<u32>  (n, unused)

@group(0) @binding(0) var<storage, read_write> x: array<f32>;
@group(0) @binding(1) var<storage, read> params: vec2<u32>;

@compute @workgroup_size(256, 1, 1)
fn gelu_inplace(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.x;
    let n = params.x;
    if i >= n { return; }
    let xv = x[i];
    // Clamp the tanh argument: tanh saturates to ±1 by |arg|≈15, but a GPU
    // tanh computed as (exp(2a)-1)/(exp(2a)+1) overflows to inf/inf = NaN for
    // large `a` (the cubic term makes `a` ~180 for |x|~17). Clamping is
    // numerically identical to the CPU f32::tanh on the saturated tail.
    let inner = clamp(0.7978845608 * (xv + 0.044715 * xv * xv * xv), -15.0, 15.0);
    x[i] = 0.5 * xv * (1.0 + tanh(inner));
}