metrale_model_engine/model/
impl_lora.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: LoRA on `TransformerModel`: startup install (`set_lora_weights`), per-request
4//! adapter slot resolution and refs, the per-decode MoE-LoRA route, and the token-overlay build.
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: Adapter id for a request's `adapter_slot` (`< 0` resolves
18    /// to the active slot): a hash of the slot's name and generation, or `0`
19    /// when no LoRA pool is resident or the slot is empty or out of range.
20    pub fn adapter_id_for_slot(&self, slot: i32) -> u64 {
21        match self.lora.as_ref() {
22            Some(lw) => lw.adapter_id_for_slot(slot),
23            None => 0,
24        }
25    }
26
27    /// 2026-09-25: Take a ref on the pool slot a sequence uses (`< 0`
28    /// resolves to the active slot, as in [`Self::adapter_id_for_slot`]).
29    /// Returns the resolved index; the caller stores it and releases that
30    /// index at terminal free, so a rotate that changes `active` in between
31    /// cannot move the release to another counter. Returns `-1` (nothing
32    /// acquired) when no LoRA pool is resident or the slot is out of range.
33    pub fn acquire_adapter_slot(&self, slot: i32) -> i32 {
34        match self.lora.as_ref() {
35            Some(lw) => lw.acquire_slot(slot),
36            None => -1,
37        }
38    }
39
40    /// 2026-09-25: Release a ref taken by [`Self::acquire_adapter_slot`], by
41    /// the resolved index it returned. `-1` and no pool are no-ops.
42    pub fn release_adapter_slot(&self, resolved: i32) {
43        if let Some(lw) = self.lora.as_ref() {
44            lw.release_slot(resolved);
45        }
46    }
47
48    /// 2026-09-25: The MoE-LoRA route for one request's `adapter_slot`, via
49    /// `resolve_moe_lora_route`: `Skip` for a base request (`adapter_slot <
50    /// 0`), `Fold` for the active adapter, `Refuse` for any other adapter, and
51    /// `Fold` when no LoRA pool is resident.
52    pub(crate) fn moe_lora_route(
53        &self,
54        adapter_slot: i32,
55    ) -> metrale_model_layers::layer::MoeLoraRoute {
56        let (active, has) = match self.lora.as_ref() {
57            Some(lw) => (lw.active as i32, true),
58            None => (-1, false),
59        };
60        metrale_model_layers::lora::resolve_moe_lora_route(adapter_slot, active, has)
61    }
62
63    /// 2026-09-25: The MoE-LoRA route last stamped at a `Model` decode entry,
64    /// read into the decode and verify `ForwardContext`s.
65    pub(crate) fn decode_moe_route(&self) -> metrale_model_layers::layer::MoeLoraRoute {
66        match self
67            .decode_moe_route
68            .load(std::sync::atomic::Ordering::Relaxed)
69        {
70            0 => metrale_model_layers::layer::MoeLoraRoute::Skip,
71            2 => metrale_model_layers::layer::MoeLoraRoute::Refuse,
72            _ => metrale_model_layers::layer::MoeLoraRoute::Fold,
73        }
74    }
75
76    fn store_decode_moe_route(&self, route: metrale_model_layers::layer::MoeLoraRoute) {
77        let v = match route {
78            metrale_model_layers::layer::MoeLoraRoute::Skip => 0,
79            metrale_model_layers::layer::MoeLoraRoute::Fold => 1,
80            metrale_model_layers::layer::MoeLoraRoute::Refuse => 2,
81        };
82        self.decode_moe_route
83            .store(v, std::sync::atomic::Ordering::Relaxed);
84    }
85
86    /// 2026-09-25: Stamp the per-decode MoE route from a single request's `adapter_slot`.
87    pub(crate) fn stamp_decode_moe_single(&self, adapter_slot: i32) {
88        self.store_decode_moe_route(self.moe_lora_route(adapter_slot));
89    }
90
91    /// 2026-09-25: Stamp the per-decode MoE route for a decode batch: `Refuse`
92    /// if any row routes to a non-active adapter (the batched and mixed decode
93    /// entries then bail host-side through `ensure_decode_route_servable`),
94    /// else `Fold` if any row folds, else `Skip` (every row base, or an empty
95    /// batch).
96    pub(crate) fn stamp_decode_moe_batch(&self, seqs: &[&mut crate::traits::SequenceState]) {
97        use metrale_model_layers::layer::MoeLoraRoute;
98        let mut any_fold = false;
99        for s in seqs.iter() {
100            match self.moe_lora_route(s.adapter_slot) {
101                MoeLoraRoute::Refuse => {
102                    self.store_decode_moe_route(MoeLoraRoute::Refuse);
103                    return;
104                }
105                MoeLoraRoute::Fold => any_fold = true,
106                MoeLoraRoute::Skip => {}
107            }
108        }
109        self.store_decode_moe_route(if any_fold {
110            MoeLoraRoute::Fold
111        } else {
112            MoeLoraRoute::Skip
113        });
114    }
115
116    pub fn set_lora_weights(
117        &mut self,
118        mut lora: Option<metrale_model_layers::lora::LoraWeights>,
119    ) -> Result<()> {
120        if let Some(ref lw) = lora {
121            // 2026-09-25: `lora_rotatable` (the `lora_rotate` lever or
122            // METRALE_LORA_PEER) is what permits rotate and swap. A pool with
123            // several adapters does not set it by itself.
124            self.lora_rotatable =
125                self.levers.lora_rotate || metrale_model_layers::lora::lora_peer_env().is_some();
126            let kernels = ops::lora_delta::LoraKernels::new(self.gpu.as_ref())?;
127            let active = lw.active_layers().to_vec();
128            let tables = lw.tables.clone();
129            let scale_table = lw.scale_table;
130            let installed = self.install_lora_layers(&active, kernels, &tables, scale_table)?;
131            // 2026-09-25: `slots` is padded to `max_loras` with empty-named
132            // placeholders; report only the named adapters.
133            let resident: Vec<String> = lw
134                .adapter_names()
135                .into_iter()
136                .filter(|n| !n.is_empty())
137                .collect();
138            tracing::info!(
139                "LoRA: {} adapter(s) resident [{}], active '{}' installed on \
140                 {installed} layers (r={}, max_rank={}, max_loras={}, \
141                 pool={:.1} MiB, rotatable={})",
142                resident.len(),
143                resident.join(", "),
144                lw.name,
145                lw.adapter_config.r,
146                lw.max_rank,
147                lw.max_loras,
148                lw.pool_bytes as f64 / (1024.0 * 1024.0),
149                self.lora_rotatable,
150            );
151        }
152        // 2026-09-25: Build the token-overlay tables from the adapter's staged
153        // raw overlay tensors against the served embed and lm_head tables. The
154        // builder runs only when some slot shipped overlay tensors; otherwise
155        // `self.overlays` is left as it was.
156        if let Some(ref mut lw) = lora {
157            let raws = std::mem::take(&mut lw.overlay_raw);
158            if raws.iter().any(|r| r.is_some()) {
159                self.build_token_overlays(lw.max_loras, &raws)?;
160            }
161        }
162        self.lora = lora;
163        Ok(())
164    }
165
166    /// 2026-09-25: Token-overlay build: `build_overlay` row-diffs each staged
167    /// slot's raw overlay tensors against the served embed/lm_head tables and
168    /// compacts the override rows, then the `[max_loras]` device tables are
169    /// built. Sets `self.overlays` only when some slot overrides a row.
170    fn build_token_overlays(
171        &mut self,
172        max_loras: usize,
173        raws: &[Option<metrale_model_layers::lora::OverlayRawSlot>],
174    ) -> Result<()> {
175        // 2026-09-25: Treated as tied when the lm_head shares the embed buffer
176        // or an NVFP4/FP8 head is served; tied, the logit recompute reuses the
177        // embed override rows.
178        let tied = self.lm_head_weight.weight.0 == self.embed_tokens.weight.0
179            || self.lm_head_nvfp4.is_some()
180            || self.lm_head_fp8.is_some();
181        let vocab = self.config.vocab_size;
182        let h = self.config.hidden_size;
183        let served_embed = self.embed_tokens.weight;
184        let served_lmhead = self.lm_head_weight.weight;
185        let stream = self.gpu.default_stream();
186        let mut overlays: Vec<Option<metrale_model_layers::lora::EmbedOverlay>> =
187            (0..max_loras).map(|_| None).collect();
188        for (k, raw) in raws.iter().enumerate() {
189            if let Some(slot) = raw
190                && k < max_loras
191            {
192                overlays[k] = metrale_model_layers::lora::build_overlay(
193                    self.gpu.as_ref(),
194                    &self.overlay_kernels,
195                    slot,
196                    served_embed,
197                    served_lmhead,
198                    vocab,
199                    h,
200                    tied,
201                    stream,
202                )?;
203            }
204        }
205        let set = metrale_model_layers::lora::TokenOverlaySet::from_slots(
206            self.gpu.as_ref(),
207            overlays,
208            max_loras,
209            tied,
210        )?;
211        if set.any_active() {
212            tracing::info!(
213                "LoRA overlay: token-overlay tables built (max_n_override={}, tied={})",
214                set.max_n_override,
215                tied,
216            );
217            self.overlays = Some(set);
218        }
219        Ok(())
220    }
221}
metrale_model_engine/model/impl_lora.rs

metrale_model_engine/model/
impl_lora.rs
