mlx-native 0.8.1

Pure-Rust Metal GPU compute library for MLX-compatible inference on Apple Silicon
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

#include <metal_stdlib>
using namespace metal;

/// Elementwise SiLU (swish) forward.
///
///   silu(x) = x · sigmoid(x) = x / (1 + exp(-x))
///
/// Buffer layout:
///   buffer(0): input  — float[n]
///   buffer(1): output — float[n]
///   buffer(2): params — uint[1]: n
///
/// Grid: 1D threads across n.
kernel void silu_f32(
    device const float *input  [[buffer(0)]],
    device float       *output [[buffer(1)]],
    device const uint  *params [[buffer(2)]],
    uint tid [[thread_position_in_grid]]
) {
    const uint n = params[0];
    if (tid >= n) return;
    const float x = input[tid];
    const float s = 1.0f / (1.0f + metal::exp(-x));
    output[tid] = x * s;
}

/// Elementwise SiLU backward.
///
///   silu'(x) = sigmoid(x) + x · sigmoid(x) · (1 − sigmoid(x))
///            = sigmoid(x) · (1 + x · (1 − sigmoid(x)))
///   dx[i]    = dy[i] · silu'(x[i])
///
/// `x` is the FORWARD INPUT (not the forward output).
///
/// Buffer layout:
///   buffer(0): x      — float[n]    (forward input)
///   buffer(1): dy     — float[n]    (upstream gradient)
///   buffer(2): dx     — float[n]    (output gradient)
///   buffer(3): params — uint[1]: n
///
/// Grid: 1D threads across n.
kernel void silu_backward_f32(
    device const float *x      [[buffer(0)]],
    device const float *dy     [[buffer(1)]],
    device float       *dx     [[buffer(2)]],
    device const uint  *params [[buffer(3)]],
    uint tid [[thread_position_in_grid]]
) {
    const uint n = params[0];
    if (tid >= n) return;
    const float xv = x[tid];
    const float s = 1.0f / (1.0f + metal::exp(-xv));
    // silu'(x) = s · (1 + x · (1 − s))
    const float deriv = s * (1.0f + xv * (1.0f - s));
    dx[tid] = dy[tid] * deriv;
}