cortex_rust 0.6.0

High-performance LLM inference with 4-bit quantization and Test-Time Training (TTT)
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

//! SwiGLU - Gated MLP with SiLU activation

use candle_core::Result;
use candle_core::Tensor;
use candle_nn::VarBuilder;
use std::collections::HashMap;

use super::AdaptiveBitLinear;
use crate::model::config::QuantizationConfig;

/// SwiGLU MLP block (Gate, Down, Up projections)
pub struct SwiGLU {
    pub w1: AdaptiveBitLinear, // Gate
    pub w2: AdaptiveBitLinear, // Down
    pub w3: AdaptiveBitLinear, // Up
}

impl SwiGLU {
    /// Load directly from pre-loaded tensor HashMap (bypasses VarBuilder).
    ///
    /// # Arguments
    /// - `tensors`: Pre-loaded tensors from `candle_core::safetensors::load()`
    /// - `prefix`: Layer prefix (e.g., "model.layers.0.mlp")
    /// - `hidden_dim`: Hidden dimension
    /// - `intermediate_dim`: Intermediate dimension
    /// - `device`: Target device
    /// - `quantization`: Quantization configuration (for 4-bit support)
    pub fn load_direct(
        tensors: &HashMap<String, Tensor>,
        prefix: &str,
        hidden_dim: usize,
        intermediate_dim: usize,
        device: &candle_core::Device,
        quantization: &Option<QuantizationConfig>,
    ) -> Result<Self> {
        let w1 = AdaptiveBitLinear::load_direct(
            tensors,
            &format!("{}.gate_proj", prefix),
            hidden_dim,
            intermediate_dim,
            device,
            quantization,
        )?;
        let w2 = AdaptiveBitLinear::load_direct(
            tensors,
            &format!("{}.down_proj", prefix),
            intermediate_dim,
            hidden_dim,
            device,
            quantization,
        )?;
        let w3 = AdaptiveBitLinear::load_direct(
            tensors,
            &format!("{}.up_proj", prefix),
            hidden_dim,
            intermediate_dim,
            device,
            quantization,
        )?;
        Ok(Self { w1, w2, w3 })
    }

    pub fn load(
        hidden_dim: usize,
        intermediate_dim: usize,
        vb: VarBuilder,
        device: &candle_core::Device,
    ) -> Result<Self> {
        let w1 = AdaptiveBitLinear::load(hidden_dim, intermediate_dim, vb.pp("gate_proj"), device)?;
        let w2 = AdaptiveBitLinear::load(intermediate_dim, hidden_dim, vb.pp("down_proj"), device)?;
        let w3 = AdaptiveBitLinear::load(hidden_dim, intermediate_dim, vb.pp("up_proj"), device)?;
        Ok(Self { w1, w2, w3 })
    }

    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
        let x_gate = self.w1.forward(x)?;
        let x_up = self.w3.forward(x)?;
        let silu_gate = candle_nn::ops::silu(&x_gate)?;
        let hidden = (silu_gate * x_up)?;
        self.w2.forward(&hidden)
    }

    pub fn precompute_packed(&mut self) -> Result<()> {
        self.w1.precompute_packed()?;
        self.w2.precompute_packed()?;
        self.w3.precompute_packed()?;
        Ok(())
    }
}