ruDNN 0.1.49

Neural network operators for Ruda.
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

use ruda_kernel::dsl as kernel_dsl;
use ruda_kernel::dsl::prelude::*;

#[ruda(launch)]
pub(crate) fn swiglu<F: Float>(
    gate: &mut Array<F>,
    up: &Array<F>,
    #[define(F)] _dtype: StorageType,
) {
    let position = ABSOLUTE_POS;
    if position >= gate.len() {
        terminate!();
    }
    let value = f32::cast_from(gate[position]);
    let silu = F::cast_from(value / (1.0f32 + f32::exp(-value)));
    gate[position] = F::cast_from(f32::cast_from(silu) * f32::cast_from(up[position]));
}

/// Training forward must retain gate/up for backward, but need not copy gate.
#[ruda(launch)]
pub(crate) fn swiglu_out<F: Float>(
    gate: &Array<F>, up: &Array<F>, out: &mut Array<F>,
    #[define(F)] _dtype: StorageType,
) {
    let i = ABSOLUTE_POS;
    if i < out.len() {
        let g = f32::cast_from(gate[i]);
        let silu = F::cast_from(g / (1.0f32 + f32::exp(-g)));
        out[i] = F::cast_from(f32::cast_from(silu) * f32::cast_from(up[i]));
    }
}

#[ruda(launch)]
pub(crate) fn swiglu_backward<F: Float>(
    gate: &Array<F>, up: &Array<F>, grad: &Array<F>,
    dgate: &mut Array<F>, dup: &mut Array<F>,
    #[define(F)] _dtype: StorageType,
) {
    let i = ABSOLUTE_POS;
    if i >= grad.len() { terminate!(); }
    let g = f32::cast_from(gate[i]);
    let u = f32::cast_from(up[i]);
    let dy = f32::cast_from(grad[i]);
    let sigmoid = 1.0f32 / (1.0f32 + f32::exp(-g));
    // Match forward's stored SiLU and mul-backward storage rounding. Keep
    // this contract aligned with ruda-torch's native silu_mul training kernel.
    let silu = F::cast_from(g / (1.0f32 + f32::exp(-g)));
    let intermediate = F::cast_from(dy * u);
    dgate[i] = F::cast_from(f32::cast_from(intermediate) * sigmoid * (1.0f32 + g * (1.0f32 - sigmoid)));
    dup[i] = F::cast_from(dy * f32::cast_from(silu));
}

#[ruda(launch)]
pub(crate) fn add<F: Float>(a: &Array<F>, b: &Array<F>, out: &mut Array<F>, #[define(F)] _dtype: StorageType) {
    let i = ABSOLUTE_POS;
    if i < out.len() { out[i] = F::cast_from(f32::cast_from(a[i]) + f32::cast_from(b[i])); }
}

/// Original gate VJP only, preserving the same intermediate storage rounding.
#[ruda(launch)]
pub(crate) fn swiglu_backward_gate<F:Float>(gate:&Array<F>,up:&Array<F>,grad:&Array<F>,dgate:&mut Array<F>,#[define(F)] _dtype:StorageType) {
    let i=ABSOLUTE_POS;if i>=grad.len() {terminate!();}
    let g=f32::cast_from(gate[i]);let u=f32::cast_from(up[i]);let dy=f32::cast_from(grad[i]);
    let sigmoid=1.0f32/(1.0f32+f32::exp(-g));let intermediate=F::cast_from(dy*u);
    dgate[i]=F::cast_from(f32::cast_from(intermediate)*sigmoid*(1.0f32+g*(1.0f32-sigmoid)));
}
/// Original up VJP only, retaining forward's actual stored SiLU.
#[ruda(launch)]
pub(crate) fn swiglu_backward_up<F:Float>(gate:&Array<F>,grad:&Array<F>,dup:&mut Array<F>,#[define(F)] _dtype:StorageType) {
    let i=ABSOLUTE_POS;if i>=grad.len() {terminate!();}
    let g=f32::cast_from(gate[i]);let dy=f32::cast_from(grad[i]);let silu=F::cast_from(g/(1.0f32+f32::exp(-g)));
    dup[i]=F::cast_from(dy*f32::cast_from(silu));
}