candle_transformers/models/voxtral/
voxtral_llama.rs

1use crate::models::with_tracing::{linear_no_bias as linear, Linear, RmsNorm};
2use candle::{DType, Device, IndexOp, Result, Tensor, D};
3use candle_nn::{embedding, Embedding, Module, VarBuilder};
4use serde::Deserialize;
5use std::collections::HashMap;
6
7pub const DEFAULT_MAX_SEQ_LEN: usize = 4096;
8
9#[derive(Debug, Clone, PartialEq, Deserialize)]
10pub struct VoxtralLlamaConfig {
11    pub hidden_size: usize,
12    pub intermediate_size: usize,
13    pub vocab_size: usize,
14    pub num_hidden_layers: usize,
15    pub num_attention_heads: usize,
16    pub num_key_value_heads: usize,
17    pub head_dim: Option<usize>, // explicit head_dim from config
18    pub use_flash_attn: bool,
19    pub rms_norm_eps: f64,
20    pub rope_theta: f32,
21    pub max_position_embeddings: usize,
22    pub tie_word_embeddings: bool,
23}
24
25impl VoxtralLlamaConfig {
26    /// Voxtral 3B text model configuration
27    pub fn voxtral_3b() -> Self {
28        Self {
29            hidden_size: 3072,
30            intermediate_size: 8192,
31            vocab_size: 131072,
32            num_hidden_layers: 30,
33            num_attention_heads: 32,
34            num_key_value_heads: 8,
35            head_dim: Some(128), // Voxtral uses explicit head_dim=128
36            use_flash_attn: true,
37            rms_norm_eps: 1e-5,
38            rope_theta: 100_000_000.0,
39            max_position_embeddings: 131072,
40            tie_word_embeddings: false,
41        }
42    }
43
44    /// Voxtral 24B text model configuration
45    pub fn voxtral_24b() -> Self {
46        Self {
47            hidden_size: 5120,
48            intermediate_size: 32768,
49            vocab_size: 131072,
50            num_hidden_layers: 40,
51            num_attention_heads: 32,
52            num_key_value_heads: 8,
53            head_dim: Some(128), // Voxtral uses explicit head_dim=128
54            use_flash_attn: true,
55            rms_norm_eps: 1e-5,
56            rope_theta: 100_000_000.0,
57            max_position_embeddings: 131072,
58            tie_word_embeddings: false,
59        }
60    }
61}
62
63#[derive(Debug, Clone)]
64pub struct VoxtralLlamaCache {
65    masks: HashMap<(usize, usize), Tensor>,
66    pub use_kv_cache: bool,
67    kvs: Vec<Option<(Tensor, Tensor)>>,
68    cos: Tensor,
69    sin: Tensor,
70    device: Device,
71}
72
73fn calculate_default_inv_freq(cfg: &VoxtralLlamaConfig) -> Vec<f32> {
74    let head_dim = cfg
75        .head_dim
76        .unwrap_or(cfg.hidden_size / cfg.num_attention_heads);
77    (0..head_dim)
78        .step_by(2)
79        .map(|i| 1f32 / cfg.rope_theta.powf(i as f32 / head_dim as f32))
80        .collect()
81}
82
83impl VoxtralLlamaCache {
84    pub fn new(
85        use_kv_cache: bool,
86        dtype: DType,
87        config: &VoxtralLlamaConfig,
88        device: &Device,
89    ) -> Result<Self> {
90        // precompute freqs_cis
91        let theta = calculate_default_inv_freq(config);
92
93        let theta = Tensor::new(theta, device)?;
94
95        let idx_theta = Tensor::arange(0, config.max_position_embeddings as u32, device)?
96            .to_dtype(DType::F32)?
97            .reshape((config.max_position_embeddings, 1))?
98            .matmul(&theta.reshape((1, theta.elem_count()))?)?;
99        // This is different from the paper, see:
100        // https://github.com/huggingface/transformers/blob/6112b1c6442aaf7affd2b0676a1cd4eee30c45cf/src/transformers/models/llama/modeling_llama.py#L112 # trufflehog:ignore
101        let cos = idx_theta.cos()?.to_dtype(dtype)?;
102        let sin = idx_theta.sin()?.to_dtype(dtype)?;
103        Ok(Self {
104            masks: HashMap::new(),
105            use_kv_cache,
106            kvs: vec![None; config.num_hidden_layers],
107            device: device.clone(),
108            cos,
109            sin,
110        })
111    }
112
113    fn mask(&mut self, seq_len: usize, index_pos: usize) -> Result<Tensor> {
114        let kv_len = index_pos + seq_len;
115        if let Some(mask) = self.masks.get(&(seq_len, kv_len)) {
116            Ok(mask.clone())
117        } else {
118            let mask = crate::utils::build_causal_mask(seq_len, index_pos, &self.device)?;
119            self.masks.insert((seq_len, kv_len), mask.clone());
120            Ok(mask)
121        }
122    }
123}
124
125#[derive(Debug, Clone)]
126struct CausalSelfAttention {
127    q_proj: Linear,
128    k_proj: Linear,
129    v_proj: Linear,
130    o_proj: Linear,
131    num_attention_heads: usize,
132    num_key_value_heads: usize,
133    head_dim: usize,
134    use_flash_attn: bool,
135    span: tracing::Span,
136    span_rot: tracing::Span,
137    max_position_embeddings: usize,
138}
139
140#[cfg(feature = "flash-attn")]
141fn flash_attn(
142    q: &Tensor,
143    k: &Tensor,
144    v: &Tensor,
145    softmax_scale: f32,
146    causal: bool,
147) -> Result<Tensor> {
148    candle_flash_attn::flash_attn(q, k, v, softmax_scale, causal)
149}
150
151#[cfg(not(feature = "flash-attn"))]
152fn flash_attn(_: &Tensor, _: &Tensor, _: &Tensor, _: f32, _: bool) -> Result<Tensor> {
153    unimplemented!("compile with '--features flash-attn'")
154}
155
156impl CausalSelfAttention {
157    fn apply_rotary_emb(
158        &self,
159        x: &Tensor,
160        index_pos: usize,
161        cache: &VoxtralLlamaCache,
162    ) -> Result<Tensor> {
163        let _enter = self.span_rot.enter();
164        let (_b_sz, _, seq_len, _hidden_size) = x.dims4()?;
165        let cos = cache.cos.narrow(0, index_pos, seq_len)?;
166        let sin = cache.sin.narrow(0, index_pos, seq_len)?;
167
168        // Ensure dtype consistency between input tensor and position embeddings
169        let x_dtype = x.dtype();
170        let cos = if cos.dtype() != x_dtype {
171            cos.to_dtype(x_dtype)?
172        } else {
173            cos
174        };
175        let sin = if sin.dtype() != x_dtype {
176            sin.to_dtype(x_dtype)?
177        } else {
178            sin
179        };
180
181        candle_nn::rotary_emb::rope(x, &cos, &sin)
182    }
183
184    fn forward(
185        &self,
186        x: &Tensor,
187        index_pos: usize,
188        block_idx: usize,
189        cache: &mut VoxtralLlamaCache,
190    ) -> Result<Tensor> {
191        let _enter = self.span.enter();
192        let (b_sz, seq_len, _hidden_size) = x.dims3()?;
193        let q = self.q_proj.forward(x)?;
194        let k = self.k_proj.forward(x)?;
195        let v = self.v_proj.forward(x)?;
196
197        let q = q
198            .reshape((b_sz, seq_len, self.num_attention_heads, self.head_dim))?
199            .transpose(1, 2)?
200            .contiguous()?;
201        let k = k
202            .reshape((b_sz, seq_len, self.num_key_value_heads, self.head_dim))?
203            .transpose(1, 2)?
204            .contiguous()?;
205        let mut v = v
206            .reshape((b_sz, seq_len, self.num_key_value_heads, self.head_dim))?
207            .transpose(1, 2)?;
208
209        let q = self.apply_rotary_emb(&q, index_pos, cache)?;
210        let mut k = self.apply_rotary_emb(&k, index_pos, cache)?;
211
212        if cache.use_kv_cache {
213            if let Some((cache_k, cache_v)) = &cache.kvs[block_idx] {
214                k = Tensor::cat(&[cache_k, &k], 2)?.contiguous()?;
215                v = Tensor::cat(&[cache_v, &v], 2)?.contiguous()?;
216                let k_seq_len = k.dims()[1];
217                if k_seq_len > self.max_position_embeddings {
218                    k = k
219                        .narrow(
220                            D::Minus1,
221                            k_seq_len - self.max_position_embeddings,
222                            self.max_position_embeddings,
223                        )?
224                        .contiguous()?
225                }
226                let v_seq_len = v.dims()[1];
227                if v_seq_len > 2 * self.max_position_embeddings {
228                    v = v
229                        .narrow(
230                            D::Minus1,
231                            v_seq_len - self.max_position_embeddings,
232                            self.max_position_embeddings,
233                        )?
234                        .contiguous()?
235                }
236            }
237            cache.kvs[block_idx] = Some((k.clone(), v.clone()))
238        }
239
240        let k = self.repeat_kv(k)?;
241        let v = self.repeat_kv(v)?;
242
243        let y = if self.use_flash_attn {
244            // flash-attn expects (b_sz, seq_len, nheads, head_dim)
245            let q = q.transpose(1, 2)?;
246            let k = k.transpose(1, 2)?;
247            let v = v.transpose(1, 2)?;
248            let softmax_scale = 1f32 / (self.head_dim as f32).sqrt();
249            flash_attn(&q, &k, &v, softmax_scale, seq_len > 1)?.transpose(1, 2)?
250        } else {
251            let in_dtype = q.dtype();
252            let q = q.to_dtype(DType::F32)?;
253            let k = k.to_dtype(DType::F32)?;
254            let v = v.to_dtype(DType::F32)?;
255            let att = (q.matmul(&k.t()?)? / (self.head_dim as f64).sqrt())?;
256            let att = if seq_len == 1 {
257                att
258            } else {
259                let mask = cache.mask(seq_len, index_pos)?.broadcast_as(att.shape())?;
260                masked_fill(&att, &mask, f32::NEG_INFINITY)?
261            };
262
263            let att = candle_nn::ops::softmax_last_dim(&att)?;
264            // Convert to contiguous as matmul doesn't support strided vs for now.
265            att.matmul(&v.contiguous()?)?.to_dtype(in_dtype)?
266        };
267        // Use the actual tensor dimensions from attention computation
268        let actual_hidden_size = self.num_attention_heads * self.head_dim;
269        let y = y
270            .transpose(1, 2)?
271            .reshape(&[b_sz, seq_len, actual_hidden_size])?;
272        let y = self.o_proj.forward(&y)?;
273        Ok(y)
274    }
275
276    fn repeat_kv(&self, x: Tensor) -> Result<Tensor> {
277        crate::utils::repeat_kv(x, self.num_attention_heads / self.num_key_value_heads)
278    }
279
280    fn load(vb: VarBuilder, cfg: &VoxtralLlamaConfig) -> Result<Self> {
281        let span = tracing::span!(tracing::Level::TRACE, "attn");
282        let span_rot = tracing::span!(tracing::Level::TRACE, "attn-rot");
283        let size_in = cfg.hidden_size;
284
285        // Use explicit head_dim if provided, otherwise calculate from hidden_size
286        let head_dim = cfg
287            .head_dim
288            .unwrap_or(cfg.hidden_size / cfg.num_attention_heads);
289        let size_q = head_dim * cfg.num_attention_heads;
290        let size_kv = head_dim * cfg.num_key_value_heads;
291
292        let q_proj = linear(size_in, size_q, vb.pp("q_proj"))?;
293        let k_proj = linear(size_in, size_kv, vb.pp("k_proj"))?;
294        let v_proj = linear(size_in, size_kv, vb.pp("v_proj"))?;
295        let o_proj = linear(size_q, size_in, vb.pp("o_proj"))?;
296        Ok(Self {
297            q_proj,
298            k_proj,
299            v_proj,
300            o_proj,
301            num_attention_heads: cfg.num_attention_heads,
302            num_key_value_heads: cfg.num_key_value_heads,
303            head_dim, // use the calculated head_dim from above
304            use_flash_attn: cfg.use_flash_attn,
305            span,
306            span_rot,
307            max_position_embeddings: cfg.max_position_embeddings,
308        })
309    }
310}
311
312fn masked_fill(on_false: &Tensor, mask: &Tensor, on_true: f32) -> Result<Tensor> {
313    let shape = mask.shape();
314    let on_true = Tensor::new(on_true, on_false.device())?.broadcast_as(shape.dims())?;
315    let m = mask.where_cond(&on_true, on_false)?;
316    Ok(m)
317}
318
319#[derive(Debug, Clone)]
320struct Mlp {
321    c_fc1: Linear,
322    c_fc2: Linear,
323    c_proj: Linear,
324    span: tracing::Span,
325}
326
327impl Mlp {
328    fn forward(&self, x: &Tensor) -> Result<Tensor> {
329        let _enter = self.span.enter();
330        let x = (candle_nn::ops::silu(&self.c_fc1.forward(x)?)? * self.c_fc2.forward(x)?)?;
331        self.c_proj.forward(&x)
332    }
333
334    fn load(vb: VarBuilder, cfg: &VoxtralLlamaConfig) -> Result<Self> {
335        let span = tracing::span!(tracing::Level::TRACE, "mlp");
336        let h_size = cfg.hidden_size;
337        let i_size = cfg.intermediate_size;
338        let c_fc1 = linear(h_size, i_size, vb.pp("gate_proj"))?;
339        let c_fc2 = linear(h_size, i_size, vb.pp("up_proj"))?;
340        let c_proj = linear(i_size, h_size, vb.pp("down_proj"))?;
341        Ok(Self {
342            c_fc1,
343            c_fc2,
344            c_proj,
345            span,
346        })
347    }
348}
349
350#[derive(Debug, Clone)]
351struct Block {
352    rms_1: RmsNorm,
353    attn: CausalSelfAttention,
354    rms_2: RmsNorm,
355    mlp: Mlp,
356    span: tracing::Span,
357}
358
359impl Block {
360    fn forward(
361        &self,
362        x: &Tensor,
363        index_pos: usize,
364        block_idx: usize,
365        cache: &mut VoxtralLlamaCache,
366    ) -> Result<Tensor> {
367        let _enter = self.span.enter();
368        let residual = x;
369        let x = self.rms_1.forward(x)?;
370        let x = (self.attn.forward(&x, index_pos, block_idx, cache)? + residual)?;
371        let residual = &x;
372        let x = (self.mlp.forward(&self.rms_2.forward(&x)?)? + residual)?;
373        Ok(x)
374    }
375
376    fn load(vb: VarBuilder, cfg: &VoxtralLlamaConfig) -> Result<Self> {
377        let span = tracing::span!(tracing::Level::TRACE, "block");
378        let attn = CausalSelfAttention::load(vb.pp("self_attn"), cfg)?;
379        let mlp = Mlp::load(vb.pp("mlp"), cfg)?;
380        let rms_1 = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?;
381        let rms_2 = RmsNorm::new(
382            cfg.hidden_size,
383            cfg.rms_norm_eps,
384            vb.pp("post_attention_layernorm"),
385        )?;
386        Ok(Self {
387            rms_1,
388            attn,
389            rms_2,
390            mlp,
391            span,
392        })
393    }
394}
395
396#[derive(Debug, Clone)]
397pub struct VoxtralLlama {
398    wte: Embedding,
399    blocks: Vec<Block>,
400    ln_f: RmsNorm,
401    lm_head: Linear,
402}
403
404impl VoxtralLlama {
405    // required by LLaVA
406    pub fn embed(&self, x: &Tensor) -> Result<Tensor> {
407        self.wte.forward(x)
408    }
409    // required by LLaVA
410    pub fn forward_input_embed(
411        &self,
412        input_embed: &Tensor,
413        index_pos: usize,
414        cache: &mut VoxtralLlamaCache,
415    ) -> Result<Tensor> {
416        let (_, seq_len, _) = input_embed.dims3()?;
417        let mut x = input_embed.clone();
418        for (block_idx, block) in self.blocks.iter().enumerate() {
419            x = block.forward(&x, index_pos, block_idx, cache)?;
420        }
421        let x = self.ln_f.forward(&x)?;
422        // Handle both single token and multi-token sequences properly
423        let x = if seq_len == 1 {
424            x.i((.., 0, ..))?
425        } else {
426            x.i((.., seq_len - 1, ..))?
427        }
428        .contiguous()?;
429        let logits = self.lm_head.forward(&x)?;
430        logits.to_dtype(DType::F32)
431    }
432
433    pub fn forward(
434        &self,
435        x: &Tensor,
436        index_pos: usize,
437        cache: &mut VoxtralLlamaCache,
438    ) -> Result<Tensor> {
439        let (_b_sz, seq_len) = x.dims2()?;
440        let mut x = self.wte.forward(x)?;
441        for (block_idx, block) in self.blocks.iter().enumerate() {
442            x = block.forward(&x, index_pos, block_idx, cache)?;
443        }
444        let x = self.ln_f.forward(&x)?;
445        let x = x.i((.., seq_len - 1, ..))?.contiguous()?;
446        let logits = self.lm_head.forward(&x)?;
447        logits.to_dtype(DType::F32)
448    }
449
450    pub fn load(vb: VarBuilder, cfg: &VoxtralLlamaConfig) -> Result<Self> {
451        let wte = embedding(cfg.vocab_size, cfg.hidden_size, vb.pp("model.embed_tokens"))?;
452        let lm_head = if cfg.tie_word_embeddings {
453            Linear::from_weights(wte.embeddings().clone(), None)
454        } else {
455            linear(cfg.hidden_size, cfg.vocab_size, vb.pp("lm_head"))?
456        };
457        let ln_f = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("model.norm"))?;
458        let blocks: Vec<_> = (0..cfg.num_hidden_layers)
459            .map(|i| Block::load(vb.pp(format!("model.layers.{i}")), cfg).unwrap())
460            .collect();
461
462        Ok(Self {
463            wte,
464            blocks,
465            ln_f,
466            lm_head,
467        })
468    }
469}
candle_transformers/models/voxtral/voxtral_llama.rs

candle_transformers/models/voxtral/
voxtral_llama.rs
