ferrotensor_backend/engine/
norm.rs

1//! Normalisation and softmax kernels.
2
3use super::{grid, uniform, ComputeEngine};
4use crate::{UnifiedTensor, WgpuContext};
5
6impl ComputeEngine {
7    /// RMSNorm over the last dimension with eps = 1e-5:
8    /// `out = x / sqrt(mean(x^2) + eps) * weight`.
9    pub fn rms_norm(
10        &self,
11        ctx: &WgpuContext,
12        input: &UnifiedTensor,
13        weight: &UnifiedTensor,
14        out: &UnifiedTensor,
15    ) {
16        self.rms_norm_eps(ctx, input, weight, out, 1e-5);
17    }
18
19    /// RMSNorm over the last dimension with a caller-chosen `eps`.
20    pub fn rms_norm_eps(
21        &self,
22        ctx: &WgpuContext,
23        input: &UnifiedTensor,
24        weight: &UnifiedTensor,
25        out: &UnifiedTensor,
26        eps: f32,
27    ) {
28        let size = *input.shape.last().unwrap_or(&1);
29        let rows = input.shape.iter().product::<usize>() / size;
30        let params = uniform(
31            ctx,
32            bytemuck::cast_slice(&[size as u32, rows as u32, eps.to_bits(), 0]),
33        );
34        self.run(
35            ctx,
36            &self.rms_norm,
37            &[
38                (0, input.buffer()),
39                (1, weight.buffer()),
40                (2, out.buffer()),
41                (3, &params),
42            ],
43            grid(rows as u32),
44        );
45    }
46
47    /// Fused RMSNorm + Residual: out = RMSNorm(x + residual) * weight
48    pub fn rms_norm_residual(
49        &self,
50        ctx: &WgpuContext,
51        x: &UnifiedTensor,
52        residual: &UnifiedTensor,
53        weight: &UnifiedTensor,
54        out: &UnifiedTensor,
55    ) {
56        let size = x.shape.iter().product::<usize>() as u32;
57        let params = uniform(ctx, bytemuck::cast_slice(&[size]));
58        self.run(
59            ctx,
60            &self.rms_residual,
61            &[
62                (0, x.buffer()),
63                (1, residual.buffer()),
64                (2, weight.buffer()),
65                (3, out.buffer()),
66                (4, &params),
67            ],
68            (1, 1, 1),
69        );
70    }
71
72    /// Softmax over the last dimension of `input`.
73    pub fn softmax(&self, ctx: &WgpuContext, input: &UnifiedTensor, out: &UnifiedTensor) {
74        let size = *input.shape.last().unwrap_or(&1);
75        let rows = input.shape.iter().product::<usize>() / size;
76        let params = uniform(ctx, bytemuck::cast_slice(&[size as u32, rows as u32]));
77        self.run(
78            ctx,
79            &self.softmax,
80            &[(0, input.buffer()), (1, out.buffer()), (2, &params)],
81            grid(rows as u32),
82        );
83    }
84
85    /// Gradient of softmax over the last dimension, given the forward
86    /// output `final_out`: `grad_in = s * (grad_out - sum(grad_out * s))`.
87    pub fn softmax_backward(
88        &self,
89        ctx: &WgpuContext,
90        grad_out: &UnifiedTensor,
91        final_out: &UnifiedTensor,
92        grad_in: &UnifiedTensor,
93    ) {
94        let size = *grad_out.shape.last().unwrap_or(&1);
95        let rows = grad_out.shape.iter().product::<usize>() / size;
96        let params = uniform(ctx, bytemuck::cast_slice(&[size as u32, rows as u32]));
97        self.run(
98            ctx,
99            &self.softmax_bw,
100            &[
101                (2, &params),
102                (3, grad_out.buffer()),
103                (4, final_out.buffer()),
104                (5, grad_in.buffer()),
105            ],
106            grid(rows as u32),
107        );
108    }
109}
ferrotensor_backend/engine/norm.rs

ferrotensor_backend/engine/
norm.rs
