candding 0.1.0

Pure candle embeddings for Rust with no ONNX: dense, sparse, multi-vector and rerank models today, image models on the roadmap
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

use candle_core::Module;
use candle_core::Tensor;
use candle_nn::Linear;
use candle_nn::VarBuilder;
use candle_nn::linear_no_bias;

use super::config::Qwen3Config;
use crate::error::Result;
use crate::layers::attention::scaled_dot_product;
use crate::layers::mask::AttentionMask;
use crate::layers::rms_norm::RmsNorm;
use crate::layers::rms_norm::RmsNormStyle;
use crate::layers::rope::RotaryEmbedding;

pub struct Qwen3Attention {
  q_proj: Linear,
  k_proj: Linear,
  v_proj: Linear,
  o_proj: Linear,
  q_norm: RmsNorm,
  k_norm: RmsNorm,
  num_heads: usize,
  num_kv_heads: usize,
  head_dim: usize,
}

impl Qwen3Attention {
  pub fn load(config: &Qwen3Config, vb: VarBuilder) -> Result<Self> {
    let hidden = config.hidden_size;
    let heads = config.num_attention_heads;
    let kv_heads = config.num_key_value_heads;
    let head_dim = config.head_dim;
    Ok(Self {
      q_proj: linear_no_bias(hidden, heads * head_dim, vb.pp("q_proj"))?,
      k_proj: linear_no_bias(hidden, kv_heads * head_dim, vb.pp("k_proj"))?,
      v_proj: linear_no_bias(hidden, kv_heads * head_dim, vb.pp("v_proj"))?,
      o_proj: linear_no_bias(heads * head_dim, hidden, vb.pp("o_proj"))?,
      q_norm: RmsNorm::load(
        vb.pp("q_norm"),
        head_dim,
        config.rms_norm_eps,
        RmsNormStyle::Scale,
      )?,
      k_norm: RmsNorm::load(
        vb.pp("k_norm"),
        head_dim,
        config.rms_norm_eps,
        RmsNormStyle::Scale,
      )?,
      num_heads: heads,
      num_kv_heads: kv_heads,
      head_dim,
    })
  }

  /// Queries and keys are normalized per head before the rotation, as in transformers' `Qwen3Attention`.
  pub fn forward(
    &self,
    hidden: &Tensor,
    rope: &RotaryEmbedding,
    mask: &AttentionMask,
  ) -> Result<Tensor> {
    let (batch, seq, _) = hidden.dims3()?;
    let q = self
      .q_norm
      .forward(&self.q_proj.forward(hidden)?.reshape((
        batch,
        seq,
        self.num_heads,
        self.head_dim,
      ))?)?
      .transpose(1, 2)?;
    let k = self
      .k_norm
      .forward(&self.k_proj.forward(hidden)?.reshape((
        batch,
        seq,
        self.num_kv_heads,
        self.head_dim,
      ))?)?
      .transpose(1, 2)?;
    let v = self
      .v_proj
      .forward(hidden)?
      .reshape((batch, seq, self.num_kv_heads, self.head_dim))?
      .transpose(1, 2)?
      .contiguous()?;
    let (q, k) = rope.apply(&q, &k)?;
    let scale = 1.0 / (self.head_dim as f64).sqrt();
    let context = scaled_dot_product(&q, &k, &v, mask, scale)?
      .transpose(1, 2)?
      .reshape((batch, seq, self.num_heads * self.head_dim))?;
    Ok(self.o_proj.forward(&context)?)
  }
}