ferrotensor_graph/
lfm.rs

1use crate::{Attention, KVCacheLayer, Linear, RMSNorm, Tensor, MLP};
2use ferrotensor_backend::{ComputeEngine, UnifiedTensor};
3use std::cell::OnceCell;
4use std::collections::HashMap;
5use std::sync::Arc;
6
7/// LFM2 short-convolution operator: projection, depthwise conv, projection.
8pub struct Lfm2Conv {
9    /// Input projection.
10    pub in_proj: Linear,
11    /// Depthwise kernel `[dim, 3]`.
12    pub conv_weight: Tensor, // [dim, 3]
13    /// Bias `[dim]`.
14    pub conv_bias: Tensor, // [dim]
15    /// Output projection.
16    pub out_proj: Linear,
17}
18
19impl Lfm2Conv {
20    /// Applies the operator to `x`.
21    pub fn forward(&self, engine: &ComputeEngine, x: &Tensor) -> Tensor {
22        let projected = self.in_proj.forward(engine, x);
23
24        let out_data = UnifiedTensor::empty(&x.ctx, &projected.data.shape);
25        engine.conv1d(
26            &x.ctx,
27            &projected.data,
28            &self.conv_weight.data,
29            &self.conv_bias.data,
30            &out_data,
31        );
32
33        let conv_out = Tensor::new(x.ctx.clone(), out_data);
34        self.out_proj.forward(engine, &conv_out)
35    }
36
37    /// Trainable tensors (weights, then any LoRA adapters).
38    pub fn params(&self) -> Vec<Tensor> {
39        let mut p = self.in_proj.params();
40        p.push(self.conv_weight.clone());
41        p.push(self.conv_bias.clone());
42        p.extend(self.out_proj.params());
43        p
44    }
45
46    /// Only the LoRA adapter tensors (what LoRA fine-tuning updates).
47    pub fn lora_params(&self) -> Vec<Tensor> {
48        let mut p = self.in_proj.lora_params();
49        p.extend(self.out_proj.lora_params());
50        p
51    }
52
53    /// Parameters keyed by Hugging Face-style names under `prefix`.
54    pub fn named_params(&self, prefix: &str) -> HashMap<String, Tensor> {
55        let mut map = HashMap::new();
56        map.extend(self.in_proj.named_params(&format!("{}.in_proj", prefix)));
57        map.insert(format!("{}.conv.weight", prefix), self.conv_weight.clone());
58        map.insert(format!("{}.conv.bias", prefix), self.conv_bias.clone());
59        map.extend(self.out_proj.named_params(&format!("{}.out_proj", prefix)));
60        map
61    }
62
63    /// LoRA tensors keyed by name under `prefix`.
64    pub fn named_lora_params(&self, prefix: &str) -> HashMap<String, Tensor> {
65        let mut map = HashMap::new();
66        map.extend(
67            self.in_proj
68                .named_params(&format!("{}.in_proj", prefix))
69                .into_iter()
70                .filter(|(name, _)| name.contains("lora_")),
71        );
72        map.extend(
73            self.out_proj
74                .named_params(&format!("{}.out_proj", prefix))
75                .into_iter()
76                .filter(|(name, _)| name.contains("lora_")),
77        );
78        map
79    }
80}
81
82/// The token-mixing operator of an LFM2 block.
83#[allow(clippy::large_enum_variant)]
84pub enum Lfm2LayerOp {
85    /// Short convolution.
86    Conv(Lfm2Conv),
87    /// Grouped-query attention.
88    Attention(Attention),
89}
90
91/// LFM2 block: `x + op(norm(x))` then `h + mlp(norm(h))`.
92pub struct Lfm2Block {
93    /// Norm before the operator.
94    pub operator_norm: RMSNorm,
95    /// Convolution or attention.
96    pub op: Lfm2LayerOp,
97    /// Norm before the MLP.
98    pub ffn_norm: RMSNorm,
99    /// Feed-forward.
100    pub mlp: MLP,
101}
102
103impl Lfm2Block {
104    /// Runs the block; `pos`, `cache` and `mask` are used by attention layers.
105    pub fn forward(
106        &self,
107        engine: &ComputeEngine,
108        x: &Tensor,
109        pos: usize,
110        cache: Option<&KVCacheLayer>,
111        mask: Option<&UnifiedTensor>,
112    ) -> Tensor {
113        // Residual 1: x + Op(Norm(x))
114        let h = self.operator_norm.forward(engine, x);
115        let h = match &self.op {
116            Lfm2LayerOp::Conv(conv) => conv.forward(engine, &h),
117            Lfm2LayerOp::Attention(attn) => attn.forward(engine, &h, pos, cache, mask),
118        };
119        let h = Tensor::add(x, &h, engine);
120
121        // Residual 2: h + MLP(Norm(h))
122        let m = self.ffn_norm.forward(engine, &h);
123        let m = self.mlp.forward(engine, &m);
124        Tensor::add(&h, &m, engine)
125    }
126
127    /// Trainable tensors (weights, then any LoRA adapters).
128    pub fn params(&self) -> Vec<Tensor> {
129        let mut p = self.operator_norm.params();
130        match &self.op {
131            Lfm2LayerOp::Conv(c) => p.extend(c.params()),
132            Lfm2LayerOp::Attention(a) => p.extend(a.params()),
133        }
134        p.extend(self.ffn_norm.params());
135        p.extend(self.mlp.params());
136        p
137    }
138
139    /// Only the LoRA adapter tensors (what LoRA fine-tuning updates).
140    pub fn lora_params(&self) -> Vec<Tensor> {
141        let mut p = match &self.op {
142            Lfm2LayerOp::Conv(c) => c.lora_params(),
143            Lfm2LayerOp::Attention(a) => a.lora_params(),
144        };
145        p.extend(self.mlp.lora_params());
146        p
147    }
148
149    /// Parameters keyed by Hugging Face-style names under `prefix`.
150    pub fn named_params(&self, prefix: &str) -> HashMap<String, Tensor> {
151        let mut map = HashMap::new();
152        map.extend(
153            self.operator_norm
154                .named_params(&format!("{}.operator_norm", prefix)),
155        );
156        match &self.op {
157            Lfm2LayerOp::Conv(c) => map.extend(c.named_params(&format!("{}.conv", prefix))),
158            Lfm2LayerOp::Attention(a) => {
159                map.extend(a.named_params(&format!("{}.self_attn", prefix)))
160            }
161        }
162        map.extend(self.ffn_norm.named_params(&format!("{}.ffn_norm", prefix)));
163        map.extend(self.mlp.named_params(&format!("{}.feed_forward", prefix)));
164        map
165    }
166
167    /// LoRA tensors keyed by name under `prefix`.
168    pub fn named_lora_params(&self, prefix: &str) -> HashMap<String, Tensor> {
169        let mut map = match &self.op {
170            Lfm2LayerOp::Conv(c) => c.named_lora_params(&format!("{}.conv", prefix)),
171            Lfm2LayerOp::Attention(a) => a
172                .named_params(&format!("{}.self_attn", prefix))
173                .into_iter()
174                .filter(|(name, _)| name.contains("lora_"))
175                .collect(),
176        };
177        map.extend(
178            self.mlp
179                .named_params(&format!("{}.feed_forward", prefix))
180                .into_iter()
181                .filter(|(name, _)| name.contains("lora_")),
182        );
183        map
184    }
185}
186
187/// A decoder-only language model: LFM2 (mixed conv/attention blocks) or
188/// Llama-style (every block is attention), as built by
189/// [`Lfm2ModelBuilder`](crate::model_builder::Lfm2ModelBuilder).
190pub struct Lfm2Model {
191    /// Embedding table `[vocab, hidden]`.
192    pub embed_tokens: Tensor, // Embedding weight
193    /// Optional norm applied right after the embedding (neither Llama nor
194    /// Hugging Face LFM2 use one; LFM2's `embedding_norm` is its final norm).
195    pub embedding_norm: Option<RMSNorm>,
196    /// Blocks.
197    pub layers: Vec<Lfm2Block>,
198    /// Final norm.
199    pub norm_f: RMSNorm,
200    /// Projection to vocabulary logits.
201    pub lm_head: Linear,
202    /// Host copy of `embed_tokens`, read once on first use (the embedding
203    /// lookup runs on the CPU).
204    embed_host: OnceCell<Vec<f32>>,
205}
206
207impl Lfm2Model {
208    /// Assemble a model from its parts.
209    pub fn new(
210        embed_tokens: Tensor,
211        embedding_norm: Option<RMSNorm>,
212        layers: Vec<Lfm2Block>,
213        norm_f: RMSNorm,
214        lm_head: Linear,
215    ) -> Self {
216        Self {
217            embed_tokens,
218            embedding_norm,
219            layers,
220            norm_f,
221            lm_head,
222            embed_host: OnceCell::new(),
223        }
224    }
225
226    /// Whether every block is attention, so incremental decoding with a
227    /// [`KVCache`](crate::KVCache) is exact (conv blocks keep no state).
228    pub fn supports_kv_cache(&self) -> bool {
229        self.layers
230            .iter()
231            .all(|l| matches!(l.op, Lfm2LayerOp::Attention(_)))
232    }
233
234    /// An empty KV cache for sequences of up to `max_seq` tokens, or `None`
235    /// when [`supports_kv_cache`](Self::supports_kv_cache) is false.
236    pub fn new_kv_cache(&self, max_seq: usize) -> Option<crate::KVCache> {
237        if !self.supports_kv_cache() {
238            return None;
239        }
240        let ctx = &self.embed_tokens.ctx;
241        let layers = self
242            .layers
243            .iter()
244            .filter_map(|l| match &l.op {
245                Lfm2LayerOp::Attention(a) => {
246                    let shape = [1, max_seq, a.num_kv_heads * a.head_dim];
247                    Some(crate::KVCacheLayer {
248                        key: UnifiedTensor::zeros(ctx, &shape),
249                        value: UnifiedTensor::zeros(ctx, &shape),
250                    })
251                }
252                Lfm2LayerOp::Conv(_) => None,
253            })
254            .collect();
255        Some(crate::KVCache { layers })
256    }
257
258    /// Logits `[1, seq, vocab]` for `input_ids`, the first of which is at `pos`.
259    pub fn forward(
260        &self,
261        engine: &ComputeEngine,
262        input_ids: &[u32],
263        pos: usize,
264        cache: Option<&mut crate::KVCache>,
265        mask: Option<&UnifiedTensor>,
266    ) -> Tensor {
267        let ctx = self.embed_tokens.ctx.clone();
268
269        // 1. Embedding
270        // Note: Generic Embedding.forward usually handles this, manual here for LFM specific
271        let mut x = self.lookup_embeddings(&ctx, engine, input_ids);
272        if let Some(norm) = &self.embedding_norm {
273            x = norm.forward(engine, &x);
274        }
275
276        // 2. Layers
277        for (i, layer) in self.layers.iter().enumerate() {
278            let layer_cache = cache.as_ref().map(|c| &c.layers[i]);
279            x = layer.forward(engine, &x, pos, layer_cache, mask);
280        }
281
282        // 3. Final Norm & Head
283        x = self.norm_f.forward(engine, &x);
284        self.lm_head.forward(engine, &x)
285    }
286
287    fn lookup_embeddings(
288        &self,
289        ctx: &Arc<ferrotensor_backend::WgpuContext>,
290        _engine: &ComputeEngine,
291        ids: &[u32],
292    ) -> Tensor {
293        let hidden_dim = self.embed_tokens.data.shape[1];
294        let seq_len = ids.len();
295        let mut data = Vec::with_capacity(seq_len * hidden_dim);
296
297        // Host-side lookup against a cached copy of the table (there is no
298        // GPU gather kernel yet). Out-of-range ids map to zeros.
299        let w = self
300            .embed_host
301            .get_or_init(|| pollster::block_on(self.embed_tokens.data.to_vec(ctx)));
302        for &id in ids {
303            let start = (id as usize) * hidden_dim;
304            match w.get(start..start + hidden_dim) {
305                Some(row) => data.extend_from_slice(row),
306                None => data.extend(std::iter::repeat_n(0.0, hidden_dim)),
307            }
308        }
309        // The lookup is not differentiable, so its output is a constant.
310        Tensor::frozen(
311            ctx.clone(),
312            UnifiedTensor::new(ctx, &data, &[1, seq_len, hidden_dim]),
313        )
314    }
315
316    /// Only the LoRA adapter tensors (what LoRA fine-tuning updates).
317    pub fn lora_params(&self) -> Vec<Tensor> {
318        let mut p = Vec::new();
319        for layer in &self.layers {
320            p.extend(layer.lora_params());
321        }
322        p.extend(self.lm_head.lora_params());
323        p
324    }
325
326    /// All trainable tensors.
327    pub fn params(&self) -> Vec<Tensor> {
328        let mut p = vec![self.embed_tokens.clone()];
329        if let Some(norm) = &self.embedding_norm {
330            p.extend(norm.params());
331        }
332        for layer in &self.layers {
333            p.extend(layer.params());
334        }
335        p.extend(self.norm_f.params());
336        p.extend(self.lm_head.params());
337        p
338    }
339
340    /// All parameters keyed by Hugging Face-style names.
341    pub fn named_params(&self) -> HashMap<String, Tensor> {
342        let mut map = HashMap::new();
343        map.insert(
344            "model.embed_tokens.weight".to_string(),
345            self.embed_tokens.clone(),
346        );
347        if let Some(norm) = &self.embedding_norm {
348            map.extend(norm.named_params("model.embedding_norm"));
349        }
350        for (idx, layer) in self.layers.iter().enumerate() {
351            map.extend(layer.named_params(&format!("model.layers.{}", idx)));
352        }
353        map.extend(self.norm_f.named_params("model.norm"));
354        map.extend(self.lm_head.named_params("lm_head"));
355        map
356    }
357
358    /// LoRA tensors keyed by name.
359    pub fn named_lora_params(&self) -> HashMap<String, Tensor> {
360        let mut map = HashMap::new();
361        for (idx, layer) in self.layers.iter().enumerate() {
362            map.extend(layer.named_lora_params(&format!("model.layers.{}", idx)));
363        }
364        map.extend(
365            self.lm_head
366                .named_params("lm_head")
367                .into_iter()
368                .filter(|(name, _)| name.contains("lora_")),
369        );
370        map
371    }
372}
ferrotensor_graph/lfm.rs

ferrotensor_graph/
lfm.rs
