ferrotensor-backend 0.1.0

wgpu compute backend with hand-written WGSL kernels for the Ferrotensor ML framework
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

//! Normalisation and softmax kernels.

use super::{grid, uniform, ComputeEngine};
use crate::{UnifiedTensor, WgpuContext};

impl ComputeEngine {
    /// RMSNorm over the last dimension with eps = 1e-5:
    /// `out = x / sqrt(mean(x^2) + eps) * weight`.
    pub fn rms_norm(
        &self,
        ctx: &WgpuContext,
        input: &UnifiedTensor,
        weight: &UnifiedTensor,
        out: &UnifiedTensor,
    ) {
        self.rms_norm_eps(ctx, input, weight, out, 1e-5);
    }

    /// RMSNorm over the last dimension with a caller-chosen `eps`.
    pub fn rms_norm_eps(
        &self,
        ctx: &WgpuContext,
        input: &UnifiedTensor,
        weight: &UnifiedTensor,
        out: &UnifiedTensor,
        eps: f32,
    ) {
        let size = *input.shape.last().unwrap_or(&1);
        let rows = input.shape.iter().product::<usize>() / size;
        let params = uniform(
            ctx,
            bytemuck::cast_slice(&[size as u32, rows as u32, eps.to_bits(), 0]),
        );
        self.run(
            ctx,
            &self.rms_norm,
            &[
                (0, input.buffer()),
                (1, weight.buffer()),
                (2, out.buffer()),
                (3, &params),
            ],
            grid(rows as u32),
        );
    }

    /// Fused RMSNorm + Residual: out = RMSNorm(x + residual) * weight
    pub fn rms_norm_residual(
        &self,
        ctx: &WgpuContext,
        x: &UnifiedTensor,
        residual: &UnifiedTensor,
        weight: &UnifiedTensor,
        out: &UnifiedTensor,
    ) {
        let size = x.shape.iter().product::<usize>() as u32;
        let params = uniform(ctx, bytemuck::cast_slice(&[size]));
        self.run(
            ctx,
            &self.rms_residual,
            &[
                (0, x.buffer()),
                (1, residual.buffer()),
                (2, weight.buffer()),
                (3, out.buffer()),
                (4, &params),
            ],
            (1, 1, 1),
        );
    }

    /// Softmax over the last dimension of `input`.
    pub fn softmax(&self, ctx: &WgpuContext, input: &UnifiedTensor, out: &UnifiedTensor) {
        let size = *input.shape.last().unwrap_or(&1);
        let rows = input.shape.iter().product::<usize>() / size;
        let params = uniform(ctx, bytemuck::cast_slice(&[size as u32, rows as u32]));
        self.run(
            ctx,
            &self.softmax,
            &[(0, input.buffer()), (1, out.buffer()), (2, &params)],
            grid(rows as u32),
        );
    }

    /// Gradient of softmax over the last dimension, given the forward
    /// output `final_out`: `grad_in = s * (grad_out - sum(grad_out * s))`.
    pub fn softmax_backward(
        &self,
        ctx: &WgpuContext,
        grad_out: &UnifiedTensor,
        final_out: &UnifiedTensor,
        grad_in: &UnifiedTensor,
    ) {
        let size = *grad_out.shape.last().unwrap_or(&1);
        let rows = grad_out.shape.iter().product::<usize>() / size;
        let params = uniform(ctx, bytemuck::cast_slice(&[size as u32, rows as u32]));
        self.run(
            ctx,
            &self.softmax_bw,
            &[
                (2, &params),
                (3, grad_out.buffer()),
                (4, final_out.buffer()),
                (5, grad_in.buffer()),
            ],
            grid(rows as u32),
        );
    }
}