metrale_model_engine/model/
impl_lora_rotate.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: LoRA per-layer install walk, graph-cache drain and runtime rotation
4//! (`install_lora_layers`, `destroy_lora_decode_graphs`, `rotate_lora_to`).
5//!
6//! Owner: model-engine.
7//! Invariants: none beyond the types.
8
9use anyhow::{Context, Result};
10
11use metrale_gpu_runtime::gpu::DevicePtr;
12
13use super::types::TransformerModel;
14use metrale_model_layers::layers::ops;
15
16impl TransformerModel {
17    /// 2026-09-25: Install one slot's per-layer pairs onto the layer structs;
18    /// the startup install, rotation and every swap path use it. `layers` is
19    /// indexed by global layer index. Returns the number of layers installed.
20    pub(super) fn install_lora_layers(
21        &mut self,
22        layers: &[Option<metrale_model_layers::lora::LoraLayerWeights>],
23        kernels: ops::lora_delta::LoraKernels,
24        tables: &std::collections::BTreeMap<
25            (usize, metrale_model_layers::lora::LoraModule),
26            (
27                metrale_gpu_runtime::gpu::DevicePtr,
28                metrale_gpu_runtime::gpu::DevicePtr,
29            ),
30        >,
31        scale_table: metrale_gpu_runtime::gpu::DevicePtr,
32    ) -> Result<usize> {
33        use metrale_model_layers::lora::LoraModule;
34        // 2026-09-25: Per-module route: the pool's a/b tables for (layer,
35        // module) plus this slot's pair dimensions. `None` when the slot has no
36        // pair for the module or the pool has no table for it.
37        let mk_route = |layer_idx: usize,
38                        module: LoraModule,
39                        pair: &Option<ops::lora_delta::LoraPair>|
40         -> Option<ops::lora_delta::LoraRoute> {
41            let p = pair.as_ref()?;
42            let (a_table, b_table) = *tables.get(&(layer_idx, module))?;
43            Some(ops::lora_delta::LoraRoute {
44                a_table,
45                b_table,
46                scale_table,
47                k_in: p.k_in,
48                n_out: p.n_out,
49                max_rank: p.max_rank,
50            })
51        };
52        let mut installed = 0usize;
53        for (idx, layer) in self.layers.iter_mut().enumerate() {
54            let Some(layer_weights) = layers.get(idx).and_then(|o| o.as_ref()) else {
55                continue;
56            };
57            let has_moe = layer_weights.router.is_some()
58                || layer_weights
59                    .experts
60                    .as_ref()
61                    .is_some_and(|e| !e.is_empty());
62            // 2026-09-25: Built before the downcast because both layer kinds
63            // install the dense FFN delta; building it once keeps them from
64            // disagreeing.
65            let ffn_weights = if layer_weights.gate_proj.is_some()
66                || layer_weights.up_proj.is_some()
67                || layer_weights.down_proj.is_some()
68            {
69                Some(ops::lora_delta::LoraFfnWeights {
70                    gate: layer_weights.gate_proj,
71                    up: layer_weights.up_proj,
72                    down: layer_weights.down_proj,
73                    kernels,
74                })
75            } else {
76                None
77            };
78            let any = layer
79                .as_any_mut()
80                .ok_or_else(|| anyhow::anyhow!("LoRA: adapted layer {idx} is not downcastable"))?;
81            // 2026-09-25: A full-attention layer takes attention, dense-FFN and
82            // MoE deltas. A linear-attention layer takes dense-FFN, GDN
83            // `out_proj` and MoE deltas; `classify_key` rejects attention
84            // projections on it.
85            if let Some(attn) =
86                any.downcast_mut::<metrale_model_layers::layers::Qwen3AttentionLayer>()
87            {
88                let attn_weights = ops::lora_delta::LoraAttnWeights {
89                    layer_idx: idx,
90                    q: layer_weights.q_proj,
91                    k: layer_weights.k_proj,
92                    v: layer_weights.v_proj,
93                    o: layer_weights.o_proj,
94                    kernels,
95                    q_route: mk_route(idx, LoraModule::QProj, &layer_weights.q_proj),
96                    k_route: mk_route(idx, LoraModule::KProj, &layer_weights.k_proj),
97                    v_route: mk_route(idx, LoraModule::VProj, &layer_weights.v_proj),
98                    o_route: mk_route(idx, LoraModule::OProj, &layer_weights.o_proj),
99                };
100                attn.set_lora_weights(attn_weights, ffn_weights)?;
101                if has_moe {
102                    attn.set_moe_lora_weights(
103                        layer_weights.router,
104                        layer_weights.experts.clone().unwrap_or_default(),
105                        kernels,
106                        self.gpu.as_ref(),
107                    )?;
108                }
109            } else if let Some(ssm) =
110                any.downcast_mut::<metrale_model_layers::layers::Qwen3SsmLayer>()
111            {
112                // 2026-09-25: `classify_key` rejects q/k/v/o on this layer
113                // kind, so one reaching here is a loader bug.
114                let has_attn_proj = layer_weights.q_proj.is_some()
115                    || layer_weights.k_proj.is_some()
116                    || layer_weights.v_proj.is_some()
117                    || layer_weights.o_proj.is_some();
118                if has_attn_proj {
119                    anyhow::bail!(
120                        "LoRA: attention-projection delta on linear-attention layer {idx} — \
121                         classify should have rejected this (that layer has no q/k/v/o)"
122                    );
123                }
124                if let Some(ffn) = ffn_weights {
125                    ssm.set_ffn_lora_weights(ffn).with_context(|| {
126                        format!("LoRA: installing dense-FFN delta on linear-attention layer {idx}")
127                    })?;
128                }
129                // 2026-09-25: GDN out_proj: only linear-attention layers have one.
130                if let Some(pair) = layer_weights.out_proj {
131                    ssm.set_out_proj_lora(pair, kernels);
132                }
133                if has_moe {
134                    ssm.set_moe_lora_weights(
135                        layer_weights.router,
136                        layer_weights.experts.clone().unwrap_or_default(),
137                        kernels,
138                        self.gpu.as_ref(),
139                    )?;
140                }
141            } else {
142                anyhow::bail!(
143                    "LoRA: adapted layer {idx} is neither a Qwen3AttentionLayer nor a \
144                     Qwen3SsmLayer (loader/adapter layer-type mismatch)"
145                );
146            }
147            installed += 1;
148        }
149        Ok(installed)
150    }
151
152    /// 2026-09-25: After a rotate or swap re-points the installed LoRA pairs,
153    /// drain the decode, batched-decode, K=2/3/4 verify, K=γ verify and fused
154    /// decode+verify graph caches and destroy each graph, so none replays the
155    /// old pair pointers. `GraphHandle` has no `Drop`, so clearing a map alone
156    /// would leak the graphs. A failed destroy is logged. The batched-verify
157    /// cache (`verify_batched_graphs`) is not drained here.
158    pub(super) fn destroy_lora_decode_graphs(&self) {
159        let drain = |name: &str, graphs: Vec<metrale_gpu_runtime::gpu::GraphHandle>| {
160            for g in graphs {
161                if g.0 != 0
162                    && let Err(e) = self.gpu.destroy_graph(g)
163                {
164                    tracing::warn!("LoRA graph clear: destroy {name}: {e:#}");
165                }
166            }
167        };
168        drain(
169            "decode_graph",
170            self.decode_graph.lock().drain().map(|(_, g)| g).collect(),
171        );
172        drain(
173            "batch_decode_graph",
174            self.batch_decode_graphs
175                .lock()
176                .0
177                .drain()
178                .map(|(_, (g, _))| g)
179                .collect(),
180        );
181        drain(
182            "verify2_graph",
183            self.verify2_graph.lock().drain().map(|(_, g)| g).collect(),
184        );
185        drain(
186            "verify3_graph",
187            self.verify3_graph.lock().drain().map(|(_, g)| g).collect(),
188        );
189        drain(
190            "verify4_graph",
191            self.verify4_graph.lock().drain().map(|(_, g)| g).collect(),
192        );
193        drain(
194            "verify_kgamma_graph",
195            self.verify_kgamma_graph
196                .lock()
197                .drain()
198                .map(|(_, g)| g)
199                .collect(),
200        );
201        drain(
202            "fused_graph",
203            self.fused_graph.lock().drain().map(|(_, g)| g).collect(),
204        );
205    }
206
207    /// 2026-09-25: Make the resident adapter `name` the active one: re-install
208    /// its per-layer pairs with `install_lora_layers`, then
209    /// `destroy_lora_decode_graphs`. Call it only at a scheduler-quiescent
210    /// point.
211    ///
212    /// Refuses, changing nothing, when no adapter is loaded, `name` is not
213    /// resident, rotation is not armed (`lora_rotatable`), or the current
214    /// active slot has in-flight sequences. If the install walk fails, the
215    /// active slot has already switched.
216    pub fn rotate_lora_to(&mut self, name: &str) -> Result<()> {
217        let slot = {
218            let lw = self
219                .lora
220                .as_ref()
221                .ok_or_else(|| anyhow::anyhow!("LoRA rotation: no adapter loaded"))?;
222            lw.slot_of(name).ok_or_else(|| {
223                anyhow::anyhow!(
224                    "LoRA rotation: adapter '{name}' is not resident (have [{}])",
225                    lw.adapter_names().join(", ")
226                )
227            })?
228        };
229        if !self.lora_rotatable {
230            anyhow::bail!(
231                "LoRA rotation not armed (single adapter, METRALE_LORA_ROTATE unset); \
232                 set METRALE_LORA_ROTATE=1 (forces eager decode) to rotate at runtime"
233            );
234        }
235        // 2026-09-25: Rotation re-installs the new slot's pairs onto the layer
236        // structs, so a sequence still decoding on the old active adapter
237        // would continue with the new delta. Refuse while the active slot has
238        // in-flight sequences.
239        {
240            let lw = self.lora.as_ref().unwrap();
241            let cur = lw.active;
242            if lw.slot_ref_count(cur) > 0 {
243                anyhow::bail!(
244                    "LoRA rotation refused: active slot {cur} has in-flight \
245                     sequences (ref_count>0); rotate at a quiescent point"
246                );
247            }
248        }
249        let (layers, active_name, r, tables, scale_table) = {
250            let lw = self.lora.as_mut().unwrap();
251            lw.active = slot;
252            lw.name = lw.slots[slot].name.clone();
253            lw.adapter_config = lw.slots[slot].adapter_config.clone();
254            (
255                lw.slots[slot].layers.clone(),
256                lw.name.clone(),
257                lw.adapter_config.r,
258                lw.tables.clone(),
259                lw.scale_table,
260            )
261        };
262        let kernels = ops::lora_delta::LoraKernels::new(self.gpu.as_ref())?;
263        let installed = self.install_lora_layers(&layers, kernels, &tables, scale_table)?;
264        self.destroy_lora_decode_graphs();
265        tracing::info!(
266            "LoRA rotation → slot {slot} '{active_name}' (r={r}) re-installed on \
267             {installed} layers"
268        );
269        Ok(())
270    }
271}
metrale_model_engine/model/impl_lora_rotate.rs

metrale_model_engine/model/
impl_lora_rotate.rs
