unillm-runtime 0.1.0

Core inference runtime for UniLLM with 47 model architectures
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

//! SwiGLU Activation Operations
//!
//! SwiGLU is used in LLaMA/Mistral/Qwen MLP layers.
//! Formula: out = silu(gate) * up = (gate * sigmoid(gate)) * up

use crate::simd::get_simd_backend;

/// Compute fused SwiGLU: out = silu(gate) * up
///
/// This fuses SiLU activation with element-wise multiplication,
/// reducing memory bandwidth by avoiding intermediate tensor.
#[inline]
pub fn fused_swiglu(gate: &[f32], up: &[f32], out: &mut [f32]) {
    get_simd_backend().fused_swiglu(gate, up, out);
}

/// Compute SiLU (Swish) activation: out = x * sigmoid(x)
#[inline]
pub fn silu(x: &[f32], out: &mut [f32]) {
    get_simd_backend().silu(x, out);
}

/// In-place SiLU
pub fn silu_inplace(x: &mut [f32]) {
    for v in x.iter_mut() {
        *v = *v / (1.0 + (-*v).exp());
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_fused_swiglu() {
        let gate = vec![0.0, 1.0, 2.0, -1.0];
        let up = vec![1.0, 1.0, 1.0, 1.0];
        let mut out = vec![0.0; 4];

        fused_swiglu(&gate, &up, &mut out);

        // silu(0) = 0
        assert!((out[0] - 0.0).abs() < 1e-5);

        // silu(1) ≈ 0.731
        assert!((out[1] - 0.731).abs() < 0.01);

        // silu(2) ≈ 1.762
        assert!((out[2] - 1.762).abs() < 0.01);

        // silu(-1) ≈ -0.269
        assert!((out[3] - (-0.269)).abs() < 0.01);
    }

    #[test]
    fn test_silu() {
        let x = vec![0.0, 1.0, -1.0];
        let mut out = vec![0.0; 3];

        silu(&x, &mut out);

        assert!((out[0] - 0.0).abs() < 1e-5);
        assert!((out[1] - 0.731).abs() < 0.01);
        assert!((out[2] - (-0.269)).abs() < 0.01);
    }
}