metrale_model_arch/weight_loader/
laguna.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Laguna weight loader: text layers, embedding, final norm and LM head; no MTP.
4//!
5//! Owner: model-arch weight loader (Laguna).
6//! Invariants: none beyond the types.
7
8mod load_layers;
9
10use anyhow::Result;
11use metrale_cache::kv_cache::KvCacheDtype;
12use metrale_config::ModelConfig;
13use metrale_gpu_runtime::gpu::GpuBackend;
14use metrale_model_weights::weights::WeightStore;
15
16use super::ModelWeightLoader;
17use metrale_model_layers::layer::TransformerLayer;
18use metrale_model_layers::weight_map::{DenseWeight, MtpWeights, dense};
19
20pub struct LagunaWeightLoader;
21
22impl ModelWeightLoader for LagunaWeightLoader {
23    fn supports_tp(&self) -> bool {
24        false
25    }
26
27    fn load_layers(
28        &self,
29        store: &WeightStore,
30        config: &ModelConfig,
31        gpu: &dyn GpuBackend,
32        layer_kv_dtypes: &[KvCacheDtype],
33    ) -> Result<Vec<Box<dyn TransformerLayer>>> {
34        load_layers::load_layers(store, config, gpu, layer_kv_dtypes)
35    }
36
37    fn load_embedding(
38        &self,
39        store: &WeightStore,
40        _config: &ModelConfig,
41        _gpu: &dyn GpuBackend,
42    ) -> Result<DenseWeight> {
43        dense(store, "model.embed_tokens.weight")
44    }
45
46    fn load_final_norm(
47        &self,
48        store: &WeightStore,
49        _config: &ModelConfig,
50        _gpu: &dyn GpuBackend,
51    ) -> Result<DenseWeight> {
52        dense(store, "model.norm.weight")
53    }
54
55    fn load_lm_head(
56        &self,
57        store: &WeightStore,
58        _config: &ModelConfig,
59        _gpu: &dyn GpuBackend,
60    ) -> Result<DenseWeight> {
61        dense(store, "lm_head.weight")
62    }
63
64    fn load_mtp_weights(
65        &self,
66        _store: &WeightStore,
67        _config: &ModelConfig,
68        _gpu: &dyn GpuBackend,
69    ) -> Result<Option<MtpWeights>> {
70        Ok(None)
71    }
72}
metrale_model_arch/weight_loader/laguna.rs

metrale_model_arch/weight_loader/
laguna.rs
