mistralrs_core/xlora_models/
llama.rs

1#![allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)]
2
3use crate::{
4    amoe::AnyMoeBaseModelMixin,
5    attention::SdpaParams,
6    layers::{Llama3RotaryEmbedding, Sdpa},
7    lora::{linear_no_bias as linear, LinearLayerLike, LoraConfig, Ordering},
8    paged_attention::ModelConfigMetadata,
9    pipeline::{
10        text_models_inputs_processor::{FlashParams, PagedAttentionInputMetadata},
11        EitherCache, IsqModel,
12    },
13    utils::progress::NiceProgressBar,
14};
15use candle_core::{DType, Device, Result, Tensor};
16use candle_nn::{Embedding, Module};
17use mistralrs_quant::{QuantMethod, ShardedVarBuilder};
18use std::{collections::HashMap, sync::Arc};
19use tqdm::Iter;
20use tracing::info;
21
22use crate::{
23    device_map::DeviceMapper,
24    layers::{embedding, CausalMasker, RmsNorm},
25    models::llama::Config,
26    pipeline::{self, extract_logits, LayerCaches, NormalLoadingMetadata, NormalModel},
27};
28
29use super::{classifier::XLoraClassifier, NonGranularState, ScalingsMaker, XLoraConfig};
30
31struct CausalSelfAttention {
32    q_proj: Arc<dyn LinearLayerLike + Send + Sync>,
33    k_proj: Arc<dyn LinearLayerLike + Send + Sync>,
34    v_proj: Arc<dyn LinearLayerLike + Send + Sync>,
35    o_proj: Arc<dyn LinearLayerLike + Send + Sync>,
36    num_attention_heads: usize,
37    num_key_value_heads: usize,
38    head_dim: usize,
39    rotary_emb: Arc<Llama3RotaryEmbedding>,
40    max_seq_len: usize,
41    sdpa_params: SdpaParams,
42}
43
44impl CausalSelfAttention {
45    #[allow(clippy::too_many_arguments)]
46    fn forward(
47        &self,
48        x: &Tensor,
49        mask: &Option<Tensor>,
50        seqlen_offsets: &[usize],
51        block_idx: usize,
52        kv_cache: &mut LayerCaches,
53        scalings: Option<Tensor>,
54        global_scaling_weight: f64,
55        is_scaling_pass: Option<f64>,
56        flash_params: &FlashParams,
57    ) -> Result<Tensor> {
58        let (b_sz, seq_len, hidden_size) = x.dims3()?;
59
60        let original_dtype = x.dtype();
61        let mut x = x.clone();
62        if let Some(t) = self.q_proj.quantized_act_type() {
63            x = x.to_dtype(t)?;
64        }
65        let mut q = self.q_proj.lora_forward(
66            &x,
67            scalings.clone(),
68            global_scaling_weight,
69            is_scaling_pass,
70        )?;
71        let mut k = self.k_proj.lora_forward(
72            &x,
73            scalings.clone(),
74            global_scaling_weight,
75            is_scaling_pass,
76        )?;
77        let mut v = self.v_proj.lora_forward(
78            &x,
79            scalings.clone(),
80            global_scaling_weight,
81            is_scaling_pass,
82        )?;
83        if self.q_proj.quantized_act_type().is_some() {
84            q = q.to_dtype(original_dtype)?;
85            k = k.to_dtype(original_dtype)?;
86            v = v.to_dtype(original_dtype)?;
87        }
88
89        let (q, k, v) = if seq_len != 1 {
90            let q = q
91                .reshape((b_sz, seq_len, self.num_attention_heads, self.head_dim))?
92                .transpose(1, 2)?;
93            let k = k
94                .reshape((b_sz, seq_len, self.num_key_value_heads, self.head_dim))?
95                .transpose(1, 2)?;
96            let v = v
97                .reshape((b_sz, seq_len, self.num_key_value_heads, self.head_dim))?
98                .transpose(1, 2)?;
99            (q, k, v)
100        } else {
101            let q = q.reshape((b_sz, self.num_attention_heads, seq_len, self.head_dim))?;
102            let k = k.reshape((b_sz, self.num_key_value_heads, seq_len, self.head_dim))?;
103            let v = v.reshape((b_sz, self.num_key_value_heads, seq_len, self.head_dim))?;
104            (q, k, v)
105        };
106
107        let (q, k) = self.rotary_emb.forward(&q, &k, seqlen_offsets)?;
108
109        let (k, v) =
110            crate::pipeline::Cache::update_kv_cache(&mut kv_cache[block_idx], k, v, false)?;
111
112        let y = Sdpa.run_attention(
113            &q,
114            &k,
115            &v,
116            mask.clone().as_ref(),
117            Some(flash_params),
118            &self.sdpa_params,
119        )?;
120
121        let mut y = y.transpose(1, 2)?.reshape(&[b_sz, seq_len, hidden_size])?;
122        if let Some(t) = self.q_proj.quantized_act_type() {
123            y = y.to_dtype(t)?;
124        }
125        let mut res = self.o_proj.lora_forward(
126            &y.transpose(1, 2)?.reshape((b_sz, seq_len, ()))?,
127            scalings.clone(),
128            global_scaling_weight,
129            is_scaling_pass,
130        )?;
131        if self.q_proj.quantized_act_type().is_some() {
132            res = res.to_dtype(original_dtype)?;
133        }
134        Ok(res)
135    }
136
137    #[allow(clippy::too_many_arguments)]
138    fn load(
139        vb: ShardedVarBuilder,
140        cfg: &Config,
141        lora_config: &[((String, String), LoraConfig)],
142        count: &mut usize,
143        ord: &Ordering,
144        mapper: &dyn DeviceMapper,
145        layer_idx: usize,
146        loading_isq: bool,
147        rope: Arc<Llama3RotaryEmbedding>,
148        preload_adapters: &Option<HashMap<String, (ShardedVarBuilder, LoraConfig)>>,
149    ) -> Result<Self> {
150        let size_in = cfg.hidden_size;
151        let size_q = (cfg.hidden_size / cfg.num_attention_heads) * cfg.num_attention_heads;
152        let size_kv = (cfg.hidden_size / cfg.num_attention_heads) * cfg.num_key_value_heads;
153        let q_proj = linear(
154            size_in,
155            size_q,
156            mapper.set_device(layer_idx, vb.pp("q_proj"), loading_isq),
157            mapper.set_device(layer_idx, vb.pp("q_proj"), false),
158            lora_config,
159            count,
160            ord,
161            preload_adapters,
162        )?;
163        let k_proj = linear(
164            size_in,
165            size_kv,
166            mapper.set_device(layer_idx, vb.pp("k_proj"), loading_isq),
167            mapper.set_device(layer_idx, vb.pp("k_proj"), false),
168            lora_config,
169            count,
170            ord,
171            preload_adapters,
172        )?;
173        let v_proj = linear(
174            size_in,
175            size_kv,
176            mapper.set_device(layer_idx, vb.pp("v_proj"), loading_isq),
177            mapper.set_device(layer_idx, vb.pp("v_proj"), false),
178            lora_config,
179            count,
180            ord,
181            preload_adapters,
182        )?;
183        let o_proj = linear(
184            size_q,
185            size_in,
186            mapper.set_device(layer_idx, vb.pp("o_proj"), loading_isq),
187            mapper.set_device(layer_idx, vb.pp("o_proj"), false),
188            lora_config,
189            count,
190            ord,
191            preload_adapters,
192        )?;
193        Ok(Self {
194            q_proj,
195            k_proj,
196            v_proj,
197            o_proj,
198            num_attention_heads: cfg.num_attention_heads,
199            num_key_value_heads: cfg.num_key_value_heads,
200            head_dim: cfg.hidden_size / cfg.num_attention_heads,
201            rotary_emb: rope,
202            max_seq_len: cfg.max_position_embeddings,
203            sdpa_params: SdpaParams {
204                n_kv_groups: cfg.num_attention_heads / cfg.num_key_value_heads,
205                use_flash_attn: cfg.use_flash_attn,
206                softcap: None,
207                softmax_scale: 1.0 / ((cfg.hidden_size / cfg.num_attention_heads) as f32).sqrt(),
208                sliding_window: None,
209            },
210        })
211    }
212}
213
214#[derive(Clone)]
215struct Mlp {
216    c_fc1: Arc<dyn LinearLayerLike + Send + Sync>,
217    c_fc2: Arc<dyn LinearLayerLike + Send + Sync>,
218    c_proj: Arc<dyn LinearLayerLike + Send + Sync>,
219}
220
221impl Mlp {
222    fn forward(
223        &self,
224        x: &Tensor,
225        scalings: Option<Tensor>,
226        global_scaling_weight: f64,
227        is_scaling_pass: Option<f64>,
228    ) -> Result<Tensor> {
229        let original_dtype = x.dtype();
230        let mut x = x.clone();
231        if let Some(t) = self.c_fc1.quantized_act_type() {
232            x = x.to_dtype(t)?;
233        }
234        let x = (candle_nn::ops::silu(&self.c_fc1.lora_forward(
235            &x,
236            scalings.clone(),
237            global_scaling_weight,
238            is_scaling_pass,
239        )?)? * self.c_fc2.lora_forward(
240            &x,
241            scalings.clone(),
242            global_scaling_weight,
243            is_scaling_pass,
244        )?)?;
245        let mut res = self.c_proj.lora_forward(
246            &x,
247            scalings.clone(),
248            global_scaling_weight,
249            is_scaling_pass,
250        )?;
251        if self.c_fc1.quantized_act_type().is_some() {
252            res = res.to_dtype(original_dtype)?;
253        }
254        Ok(res)
255    }
256
257    #[allow(clippy::too_many_arguments)]
258    fn load(
259        vb: ShardedVarBuilder,
260        cfg: &Config,
261        lora_config: &[((String, String), LoraConfig)],
262        count: &mut usize,
263        ord: &Ordering,
264        mapper: &dyn DeviceMapper,
265        layer_idx: usize,
266        loading_isq: bool,
267        preload_adapters: &Option<HashMap<String, (ShardedVarBuilder, LoraConfig)>>,
268    ) -> Result<Self> {
269        let h_size = cfg.hidden_size;
270        let i_size = cfg.intermediate_size;
271        let c_fc1 = linear(
272            h_size,
273            i_size,
274            mapper.set_device(layer_idx, vb.pp("gate_proj"), loading_isq),
275            mapper.set_device(layer_idx, vb.pp("gate_proj"), false),
276            lora_config,
277            count,
278            ord,
279            preload_adapters,
280        )?;
281        let c_fc2 = linear(
282            h_size,
283            i_size,
284            mapper.set_device(layer_idx, vb.pp("up_proj"), loading_isq),
285            mapper.set_device(layer_idx, vb.pp("up_proj"), false),
286            lora_config,
287            count,
288            ord,
289            preload_adapters,
290        )?;
291        let c_proj = linear(
292            i_size,
293            h_size,
294            mapper.set_device(layer_idx, vb.pp("down_proj"), loading_isq),
295            mapper.set_device(layer_idx, vb.pp("down_proj"), false),
296            lora_config,
297            count,
298            ord,
299            preload_adapters,
300        )?;
301        Ok(Self {
302            c_fc1,
303            c_fc2,
304            c_proj,
305        })
306    }
307}
308
309struct Block {
310    rms_1: RmsNorm,
311    attn: CausalSelfAttention,
312    rms_2: RmsNorm,
313    mlp: Mlp,
314}
315
316impl Block {
317    #[allow(clippy::too_many_arguments)]
318    fn forward(
319        &self,
320        x: &Tensor,
321        mask: &Option<Tensor>,
322        seqlen_offsets: &[usize],
323        block_idx: usize,
324        kv_cache: &mut LayerCaches,
325        scalings: Option<Tensor>,
326        global_scaling_weight: f64,
327        is_scaling_pass: Option<f64>,
328        flash_params: &FlashParams,
329    ) -> Result<Tensor> {
330        let residual = x;
331        let x = self.rms_1.forward(x)?;
332        let x = (self.attn.forward(
333            &x,
334            mask,
335            seqlen_offsets,
336            block_idx,
337            kv_cache,
338            scalings.clone(),
339            global_scaling_weight,
340            is_scaling_pass,
341            flash_params,
342        )? + residual)?;
343        let residual = &x;
344        let x = (self.mlp.forward(
345            &self.rms_2.forward(&x)?,
346            scalings,
347            global_scaling_weight,
348            is_scaling_pass,
349        )? + residual)?;
350        Ok(x)
351    }
352
353    #[allow(clippy::too_many_arguments)]
354    fn load(
355        vb: ShardedVarBuilder,
356        cfg: &Config,
357        lora_config: &[((String, String), LoraConfig)],
358        count: &mut usize,
359        ord: &Ordering,
360        mapper: &dyn DeviceMapper,
361        layer_idx: usize,
362        loading_isq: bool,
363        rope: Arc<Llama3RotaryEmbedding>,
364        preload_adapters: &Option<HashMap<String, (ShardedVarBuilder, LoraConfig)>>,
365    ) -> Result<Self> {
366        let attn = CausalSelfAttention::load(
367            vb.pp("self_attn"),
368            cfg,
369            lora_config,
370            count,
371            ord,
372            mapper,
373            layer_idx,
374            loading_isq,
375            rope,
376            preload_adapters,
377        )?;
378        let mlp = Mlp::load(
379            vb.pp("mlp"),
380            cfg,
381            lora_config,
382            count,
383            ord,
384            mapper,
385            layer_idx,
386            loading_isq,
387            preload_adapters,
388        )?;
389        let rms_1 = RmsNorm::new(
390            cfg.hidden_size,
391            cfg.rms_norm_eps,
392            mapper.set_device(layer_idx, vb.pp("input_layernorm"), false),
393        )?;
394        let rms_2 = RmsNorm::new(
395            cfg.hidden_size,
396            cfg.rms_norm_eps,
397            mapper.set_device(layer_idx, vb.pp("post_attention_layernorm"), false),
398        )?;
399        Ok(Self {
400            rms_1,
401            attn,
402            rms_2,
403            mlp,
404        })
405    }
406}
407
408pub struct XLoraLlama {
409    wte: Embedding,
410    blocks: Vec<Block>,
411    ln_f: RmsNorm,
412    lm_head: Arc<dyn LinearLayerLike + Send + Sync>,
413    kv_cache: pipeline::EitherCache,
414    device: Device,
415    xlora_classifier: Option<XLoraClassifier>,
416    dtype: DType,
417    mapper: Box<dyn DeviceMapper + Send + Sync>,
418    cfg: ModelConfigMetadata,
419}
420
421impl XLoraLlama {
422    #[allow(clippy::too_many_arguments)]
423    fn inner_forward(
424        &self,
425        input_ids: &Tensor,
426        seqlen_offsets: &[usize],
427        scalings: Option<Tensor>,
428        is_full_pass: bool,
429        no_kv_cache: bool,
430        is_scaling_pass: Option<f64>,
431        flash_params: &FlashParams,
432    ) -> Result<Tensor> {
433        let mut x = self.wte.forward(input_ids)?;
434        let mut cache = if is_full_pass {
435            if no_kv_cache {
436                let mut new_cache = Vec::new();
437                for _ in 0..self.kv_cache.full().xlora_lock().len() {
438                    new_cache.push(None);
439                }
440
441                self.kv_cache.full().xlora_lock().clone_from(&new_cache);
442            }
443            self.kv_cache.full().xlora_lock()
444        } else {
445            self.kv_cache.full().lock()
446        };
447        let mask = CausalMasker.make_causal_mask_matrix(
448            input_ids,
449            &*cache,
450            x.dtype(),
451            self.cfg.num_attn_heads,
452        )?;
453        for (block_idx, block) in self.blocks.iter().enumerate() {
454            x = self.mapper.map(x, block_idx)?;
455            x = block.forward(
456                &x,
457                &mask.clone().map(|m| m.to_device(x.device()).unwrap()),
458                seqlen_offsets,
459                block_idx,
460                &mut cache,
461                scalings.clone(),
462                self.xlora_classifier
463                    .as_ref()
464                    .map(|classifier| classifier.get_global_scaling_weight())
465                    .unwrap_or(1.0),
466                is_scaling_pass,
467                flash_params,
468            )?;
469        }
470        let x = x.to_device(&self.device)?;
471        self.ln_f.forward(&x)
472    }
473
474    #[allow(clippy::too_many_arguments)]
475    pub fn forward(
476        &self,
477        input_ids: &Tensor,
478        input_ids_full: &Tensor,
479        seqlen_offsets: &[usize],
480        seqlen_offsets_full: &[usize],
481        no_kv_cache: bool,
482        non_granular_state: &Option<NonGranularState>,
483        context_lens: Vec<(usize, usize)>,
484        flash_params: &FlashParams,
485        flash_params_full: &FlashParams,
486    ) -> Result<Tensor> {
487        if self.xlora_classifier.is_some() {
488            let scalings = self.get_scalings(
489                input_ids,
490                input_ids_full,
491                seqlen_offsets,
492                seqlen_offsets_full,
493                no_kv_cache,
494                non_granular_state,
495                &vec![usize::MAX; context_lens.len()],
496                flash_params,
497                flash_params_full,
498            )?;
499
500            if no_kv_cache {
501                let mut res = self
502                    .inner_forward(
503                        input_ids_full,
504                        seqlen_offsets_full,
505                        Some(scalings),
506                        true,
507                        no_kv_cache,
508                        None,
509                        flash_params_full,
510                    )?
511                    .contiguous()?;
512                if let Some(t) = self.lm_head.quantized_act_type() {
513                    res = res.to_dtype(t)?;
514                }
515                extract_logits(
516                    &self.lm_head.lora_forward(&res, None, 1.0, None)?,
517                    context_lens,
518                )
519            } else {
520                // is_full_pass=true is ok because no_kv_cache=false
521                let mut res = self
522                    .inner_forward(
523                        input_ids,
524                        seqlen_offsets,
525                        Some(scalings),
526                        true,
527                        no_kv_cache,
528                        None,
529                        flash_params,
530                    )?
531                    .contiguous()?;
532                if let Some(t) = self.lm_head.quantized_act_type() {
533                    res = res.to_dtype(t)?;
534                }
535                extract_logits(
536                    &self.lm_head.lora_forward(&res, None, 1.0, None)?,
537                    context_lens,
538                )
539            }
540        } else {
541            let mut res = self
542                .inner_forward(
543                    input_ids,
544                    seqlen_offsets,
545                    None,
546                    false,
547                    no_kv_cache,
548                    None,
549                    flash_params,
550                )?
551                .contiguous()?;
552            if let Some(t) = self.lm_head.quantized_act_type() {
553                res = res.to_dtype(t)?;
554            }
555            extract_logits(
556                &self.lm_head.lora_forward(&res, None, 1.0, None)?,
557                context_lens,
558            )
559        }
560    }
561
562    #[allow(clippy::too_many_arguments)]
563    pub fn new(
564        cfg: &Config,
565        vb: ShardedVarBuilder,
566        lora_config: &[((String, String), LoraConfig)],
567        xlora_config: Option<XLoraConfig>,
568        xlora_ordering: Ordering,
569        is_gptx: bool,
570        normal_loading_metadata: NormalLoadingMetadata,
571        preload_adapters: &Option<HashMap<String, (ShardedVarBuilder, LoraConfig)>>,
572    ) -> Result<Self> {
573        if let Some(ref quant_cfg) = &cfg.quantization_config {
574            tracing::info!(
575                "Using {} quantization: {}.",
576                quant_cfg.name(),
577                quant_cfg.get_bits_name(&vb)
578            );
579        }
580        let mapper = normal_loading_metadata.mapper;
581        let dtype = vb.dtype();
582        let mut count = 0;
583
584        let wte = embedding(
585            cfg.vocab_size,
586            cfg.hidden_size,
587            mapper.set_nm_device(vb.pp("model.embed_tokens"), false),
588            &cfg.quantization_config,
589        )?;
590        let lm_head = linear(
591            cfg.hidden_size,
592            cfg.vocab_size,
593            mapper.set_nm_device(vb.pp("lm_head"), normal_loading_metadata.loading_isq),
594            mapper.set_nm_device(vb.pp("lm_head"), false),
595            lora_config,
596            &mut count,
597            &xlora_ordering,
598            preload_adapters,
599        )?;
600        if xlora_config.is_some() && lm_head.is_lora() {
601            // This is why we can pass dummy values (..., None, 1.0, None)?
602            candle_core::bail!("Got an adapter `lm_head` layer, this is unsupported with X-LoRA.");
603        }
604        let ln_f = RmsNorm::new(
605            cfg.hidden_size,
606            cfg.rms_norm_eps,
607            mapper.set_nm_device(vb.pp("model.norm"), false),
608        )?;
609        let mut ropes = HashMap::new();
610        for i in 0..cfg.num_hidden_layers {
611            let device = mapper
612                .device_for(i, false)
613                .unwrap_or(&normal_loading_metadata.real_device);
614            ropes.insert(
615                device.location(),
616                Arc::new(Llama3RotaryEmbedding::new_llama3(
617                    vb.dtype(),
618                    cfg,
619                    device,
620                    is_gptx,
621                )?),
622            );
623        }
624        let mut blocks: Vec<_> = NiceProgressBar::<_, 'b'>(
625            0..cfg.num_hidden_layers,
626            "Loading repeating layers",
627            &normal_loading_metadata.multi_progress,
628        )
629        .into_iter()
630        .map(|i| {
631            let device = mapper
632                .device_for(i, false)
633                .unwrap_or(&normal_loading_metadata.real_device);
634            let rotary_emb = ropes
635                .get(&device.location())
636                .expect("No RoPE for device location!")
637                .clone();
638            Block::load(
639                vb.pp(format!("model.layers.{i}")),
640                cfg,
641                lora_config,
642                &mut count,
643                &xlora_ordering,
644                &*mapper,
645                i,
646                normal_loading_metadata.loading_isq,
647                rotary_emb,
648                preload_adapters,
649            )
650            .expect("Failed to load block.")
651        })
652        .collect();
653        if xlora_config.is_none() && preload_adapters.is_none() {
654            // We are now a LoRA model so we must merge the weights
655            info!("Merging LoRA adapters.");
656            for layer in blocks.iter_mut().tqdm() {
657                Arc::get_mut(&mut layer.attn.k_proj)
658                    .unwrap()
659                    .merge_weights()?;
660                Arc::get_mut(&mut layer.attn.o_proj)
661                    .unwrap()
662                    .merge_weights()?;
663                Arc::get_mut(&mut layer.attn.q_proj)
664                    .unwrap()
665                    .merge_weights()?;
666                Arc::get_mut(&mut layer.attn.v_proj)
667                    .unwrap()
668                    .merge_weights()?;
669
670                Arc::get_mut(&mut layer.mlp.c_fc1)
671                    .unwrap()
672                    .merge_weights()?;
673                Arc::get_mut(&mut layer.mlp.c_fc2)
674                    .unwrap()
675                    .merge_weights()?;
676                Arc::get_mut(&mut layer.mlp.c_proj)
677                    .unwrap()
678                    .merge_weights()?;
679            }
680        }
681
682        Ok(Self {
683            wte,
684            blocks,
685            ln_f,
686            lm_head,
687            kv_cache: EitherCache::Full(pipeline::Cache::new(cfg.num_hidden_layers, true)),
688            device: normal_loading_metadata.real_device,
689            xlora_classifier: xlora_config.map(|xlora_config| {
690                XLoraClassifier::new(xlora_config, count, lora_config.len(), vb, false).unwrap()
691            }),
692            dtype,
693            mapper,
694            cfg: ModelConfigMetadata {
695                max_seq_len: cfg.max_position_embeddings,
696                num_layers: cfg.num_hidden_layers,
697                hidden_size: cfg.hidden_size,
698                num_kv_heads: cfg.num_key_value_heads,
699                num_attn_heads: cfg.num_attention_heads,
700                sliding_window: None,
701                k_head_dim: cfg.hidden_size / cfg.num_attention_heads,
702                v_head_dim: cfg.hidden_size / cfg.num_attention_heads,
703            },
704        })
705    }
706}
707
708impl IsqModel for XLoraLlama {
709    fn get_layers(
710        &mut self,
711    ) -> (
712        Vec<(&mut Arc<dyn QuantMethod>, Option<usize>)>,
713        &dyn DeviceMapper,
714    ) {
715        let mut tensors = Vec::new();
716        tensors.push((Arc::get_mut(&mut self.lm_head).unwrap().quant_inner(), None));
717        for (i, layer) in self.blocks.iter_mut().enumerate() {
718            tensors.push((
719                Arc::get_mut(&mut layer.attn.q_proj).unwrap().quant_inner(),
720                Some(i),
721            ));
722            tensors.push((
723                Arc::get_mut(&mut layer.attn.k_proj).unwrap().quant_inner(),
724                Some(i),
725            ));
726            tensors.push((
727                Arc::get_mut(&mut layer.attn.v_proj).unwrap().quant_inner(),
728                Some(i),
729            ));
730            tensors.push((
731                Arc::get_mut(&mut layer.attn.o_proj).unwrap().quant_inner(),
732                Some(i),
733            ));
734            tensors.push((
735                Arc::get_mut(&mut layer.mlp.c_fc1).unwrap().quant_inner(),
736                Some(i),
737            ));
738            tensors.push((
739                Arc::get_mut(&mut layer.mlp.c_fc2).unwrap().quant_inner(),
740                Some(i),
741            ));
742            tensors.push((
743                Arc::get_mut(&mut layer.mlp.c_proj).unwrap().quant_inner(),
744                Some(i),
745            ));
746        }
747        (tensors, &*self.mapper)
748    }
749
750    fn residual_tensors(&self) -> Vec<(String, Tensor)> {
751        panic!("Cannot generate UQFF for an adapter model.")
752    }
753}
754
755impl NormalModel for XLoraLlama {
756    fn forward(
757        &self,
758        _input_ids: &Tensor,
759        _seqlen_offsets: &[usize],
760        _context_lens: Vec<(usize, usize)>,
761        _position_ids: Vec<usize>,
762        _metadata: Option<(Vec<(Tensor, Tensor)>, &PagedAttentionInputMetadata)>,
763        _flash_params: &FlashParams,
764    ) -> Result<Tensor> {
765        unreachable!()
766    }
767    fn xlora_forward(
768        &self,
769        input_ids: &Tensor,
770        input_ids_full: &Tensor,
771        seqlen_offsets: &[usize],
772        seqlen_offsets_full: &[usize],
773        no_kv_cache: bool,
774        non_granular_state: &Option<crate::xlora_models::NonGranularState>,
775        context_lens: Vec<(usize, usize)>,
776        _position_ids: Vec<usize>,
777        flash_params: &FlashParams,
778        flash_params_full: &FlashParams,
779    ) -> Result<Tensor> {
780        self.forward(
781            input_ids,
782            input_ids_full,
783            seqlen_offsets,
784            seqlen_offsets_full,
785            no_kv_cache,
786            non_granular_state,
787            context_lens,
788            flash_params,
789            flash_params_full,
790        )
791    }
792    fn cache(&self) -> &super::EitherCache {
793        &self.kv_cache
794    }
795    fn cache_mut(&mut self) -> &mut super::EitherCache {
796        &mut self.kv_cache
797    }
798    fn device(&self) -> &Device {
799        &self.device
800    }
801    fn is_xlora(&self) -> bool {
802        true
803    }
804    fn max_seq_len(&self) -> usize {
805        self.blocks[0].attn.max_seq_len
806    }
807    fn config(&self) -> &ModelConfigMetadata {
808        &self.cfg
809    }
810}
811
812impl ScalingsMaker for XLoraLlama {
813    fn dtype(&self) -> DType {
814        self.dtype
815    }
816    fn get_cache(&self) -> &pipeline::EitherCache {
817        &self.kv_cache
818    }
819    fn get_classifier(&self) -> &XLoraClassifier {
820        self.xlora_classifier.as_ref().unwrap()
821    }
822    fn forward(
823        &self,
824        input_ids: &Tensor,
825        seqlen_offsets: &[usize],
826        scalings: Tensor,
827        is_full_pass: bool,
828        no_kv_cache: bool,
829        is_scaling_pass: Option<f64>,
830        _context_lens: &[usize],
831        flash_params: &FlashParams,
832    ) -> Result<Tensor> {
833        self.inner_forward(
834            input_ids,
835            seqlen_offsets,
836            Some(scalings),
837            is_full_pass,
838            no_kv_cache,
839            is_scaling_pass,
840            flash_params,
841        )
842    }
843}
844
845impl AnyMoeBaseModelMixin for XLoraLlama {}
mistralrs_core/xlora_models/llama.rs

mistralrs_core/xlora_models/
llama.rs
