File size: 2,697 Bytes
6c46b3f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
use std::path::Path;
use anyhow::{Result, anyhow};
use super::backend::{ComputeBackend, BackendType};
use llama_cpp_2::model::LlamaModel;
use llama_cpp_2::llama_backend::LlamaBackend;
use llama_cpp_2::model::params::LlamaModelParams;
use std::sync::Arc;
use tracing::info;

pub struct LlamaCppBackend {
    backend: Arc<LlamaBackend>,
    model: Option<LlamaModel>,
    start_layer: usize,
    end_layer: usize,
}

impl LlamaCppBackend {
    pub fn new() -> Result<Self> {
        let backend = LlamaBackend::init().map_err(|e| anyhow!("Failed to init llama backend: {}", e))?;
        Ok(Self {
            backend: Arc::new(backend),
            model: None,
            start_layer: 0,
            end_layer: 0,
        })
    }
}

impl ComputeBackend for LlamaCppBackend {
    fn load_layer_slice(&mut self, model_path: &Path, start_layer: usize, end_layer: usize) -> Result<()> {
        let model_params = LlamaModelParams::default();
        // NOTE: Dans une vraie implémentation de slicing custom, on utiliserait 
        // les paramètres de llama.cpp pour ne charger que certaines couches en VRAM.
        let model = LlamaModel::load_from_file(&self.backend, model_path, &model_params)
            .map_err(|e| anyhow!("Failed to load model: {}", e))?;
        
        self.model = Some(model);
        self.start_layer = start_layer;
        self.end_layer = end_layer;
        Ok(())
    }

    fn forward_slice(&mut self, hidden_states: &[f32], seq_len: usize) -> Result<Vec<f32>> {
        // Dans une architecture Swarm, cette fonction :
        // 1. Reçoit l'hidden_state du noeud précédent
        // 2. L'injecte dans le contexte local (start_layer)
        // 3. Calcule jusqu'à end_layer
        // 4. Extrait le nouvel hidden_state pour le suivant
        info!("Calcul Swarm sur couches {} à {} (seq_len: {})", self.start_layer, self.end_layer, seq_len);
        
        // Simulation pour le protocole : on retourne un vecteur modifié
        Ok(hidden_states.iter().map(|x| x + 0.01).collect())
    }

    fn get_backend_type(&self) -> BackendType {
        // Logique de priorité : Native > Vulkan > CPU
        if cfg!(feature = "metal") {
            BackendType::NativeMetal
        } else if cfg!(feature = "cuda") {
            BackendType::NativeCuda
        } else if cfg!(feature = "rocm") {
            BackendType::NativeRocm
        } else if cfg!(feature = "vulkan") {
            // Vulkan est notre pont universel pour AMD Windows et Android (Snapdragon/Adreno)
            BackendType::VulkanUniversal
        } else if cfg!(feature = "opencl") {
            BackendType::OpenCL
        } else {
            BackendType::Cpu
        }
    }
}