metrale_model_layers/lora/
loading.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: LoRA adapter loading: the model-family check, the pack of
4//! audited adapters into the rank-padded slot pool and the expert pool, the
5//! pointer and scale tables, and the in-place disk swap of one slot.
6//!
7//! Owner: model-layers (lora).
8//! Invariants:
9//! - Every pool byte an adapter does not write is zero: both pools are zeroed
10//!   after allocation, and `pack_store_into_slot` zeroes the slot before it
11//!   packs.
12
13use std::collections::BTreeMap;
14use std::sync::atomic::{AtomicU64, AtomicUsize};
15
16use anyhow::{Result, bail};
17use metrale_config::{ModelConfig, PeftAdapterConfig};
18use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
19use metrale_model_weights::weights::WeightStore;
20
21use super::*;
22use crate::layers::ops::lora_delta::LoraPair;
23use crate::weight_map::DenseWeight;
24
25/// 2026-09-25: Refuses every model family except dense `qwen3_5`
26/// (`is_qwen35_dense`), `holo3_1_moe` and `qwen3_6_moe`. It checks the
27/// `model_type` after config dispatch, which can rewrite it.
28fn check_family(cfg: &ModelConfig) -> Result<()> {
29    if !(cfg.is_qwen35_dense()
30        || cfg.model_type == "holo3_1_moe"
31        || cfg.model_type == "qwen3_6_moe")
32    {
33        bail!(
34            "REJECT[unvalidated-family]: LoRA v0 is validated on qwen3_5 dense \
35             (holo-3.1-0.8b), holo3_1_moe (holo-3.1-35b-a3b), and qwen3_6_moe \
36             (Qwen3.6-35B-A3B) only; model_type='{}', num_experts={}",
37            cfg.model_type,
38            cfg.num_experts
39        );
40    }
41    Ok(())
42}
43
44/// 2026-09-25: Pack one audited adapter into pool `slot`, the byte region at
45/// `slot * pool_slot_bytes`. The walk within a slot is the same for every
46/// slot: layers ascending, then [`LoraModule::ALL`] filtered by
47/// `applies_to_layer`, A then B. A is copied as `[r, in]` into the head of its
48/// `[max_rank, in]` region; B's rows are re-strided from r to max_rank.
49/// Returns the slot's pairs indexed by global layer and, per applicable
50/// (layer, module), the (a_ptr, b_ptr) as raw u64, (0, 0) where the adapter
51/// has no pair, for the pointer tables.
52#[allow(clippy::type_complexity)]
53fn pack_slot(
54    slot: usize,
55    name: &str,
56    adapter_store: &WeightStore,
57    peft: &PeftAdapterConfig,
58    found: &BTreeMap<(usize, LoraModule), [Option<String>; 2]>,
59    cfg: &ModelConfig,
60    gpu: &dyn GpuBackend,
61    pool: DevicePtr,
62    max_lora_rank: usize,
63) -> Result<(
64    Vec<Option<LoraLayerWeights>>,
65    BTreeMap<(usize, LoraModule), (u64, u64)>,
66)> {
67    let scale = peft.scaling();
68    let slot_bytes = pool_slot_bytes(cfg, max_lora_rank);
69    let mut layers: Vec<Option<LoraLayerWeights>> =
70        (0..cfg.num_hidden_layers).map(|_| None).collect();
71    let mut slot_ptrs: BTreeMap<(usize, LoraModule), (u64, u64)> = BTreeMap::new();
72    let mut off = slot * slot_bytes;
73    // 2026-09-25: The same order and predicate as `pool_slot_bytes` and
74    // `module_slot_offsets`, so the offsets stay inside the slot and match the
75    // RDMA landing offsets.
76    for layer_idx in 0..cfg.num_hidden_layers {
77        let mut lw = LoraLayerWeights::empty(layer_idx);
78        let mut any = false;
79        for module in LoraModule::ALL {
80            if !module.applies_to_layer(cfg, layer_idx) {
81                continue;
82            }
83            let (out_dim, in_dim) = module.dims(cfg);
84            let a_off = off;
85            let b_off = off + max_lora_rank * in_dim * BF16_BYTES;
86            off = b_off + out_dim * max_lora_rank * BF16_BYTES;
87            let a_ptr = DevicePtr(pool.0 + a_off as u64);
88            let b_ptr = DevicePtr(pool.0 + b_off as u64);
89
90            let mut this = (0u64, 0u64);
91            if let Some([Some(a_key), Some(b_key)]) = found.get(&(layer_idx, module)) {
92                let a_t = adapter_store.get(a_key)?;
93                let mut a_host = vec![0u8; peft.r * in_dim * BF16_BYTES];
94                gpu.copy_d2h(a_t.ptr, &mut a_host)?;
95                gpu.copy_h2d(&a_host, a_ptr)?;
96                let b_t = adapter_store.get(b_key)?;
97                let mut b_src = vec![0u8; out_dim * peft.r * BF16_BYTES];
98                gpu.copy_d2h(b_t.ptr, &mut b_src)?;
99                let mut b_host = vec![0u8; out_dim * max_lora_rank * BF16_BYTES];
100                for row in 0..out_dim {
101                    let d = row * max_lora_rank * BF16_BYTES;
102                    let s = row * peft.r * BF16_BYTES;
103                    b_host[d..d + peft.r * BF16_BYTES]
104                        .copy_from_slice(&b_src[s..s + peft.r * BF16_BYTES]);
105                }
106                gpu.copy_h2d(&b_host, b_ptr)?;
107
108                let pair = LoraPair {
109                    a: DenseWeight { weight: a_ptr },
110                    b: DenseWeight { weight: b_ptr },
111                    rank: peft.r as u32,
112                    k_in: in_dim as u32,
113                    n_out: out_dim as u32,
114                    scale,
115                    max_rank: max_lora_rank as u32,
116                };
117                tracing::info!(
118                    "LoRA: slot {slot} '{name}' layer {layer_idx} {module:?} r={} \
119                     scale={:.6} A=[{},{}] B=[{},{}] (padded to max_rank={})",
120                    peft.r,
121                    scale,
122                    peft.r,
123                    in_dim,
124                    out_dim,
125                    peft.r,
126                    max_lora_rank
127                );
128                match module {
129                    LoraModule::QProj => lw.q_proj = Some(pair),
130                    LoraModule::KProj => lw.k_proj = Some(pair),
131                    LoraModule::VProj => lw.v_proj = Some(pair),
132                    LoraModule::OProj => lw.o_proj = Some(pair),
133                    LoraModule::GateProj => lw.gate_proj = Some(pair),
134                    LoraModule::UpProj => lw.up_proj = Some(pair),
135                    LoraModule::DownProj => lw.down_proj = Some(pair),
136                    LoraModule::OutProj => lw.out_proj = Some(pair),
137                }
138                this = (a_ptr.0, b_ptr.0);
139                any = true;
140            }
141            slot_ptrs.insert((layer_idx, module), this);
142        }
143        if any {
144            layers[layer_idx] = Some(lw);
145        }
146    }
147    debug_assert_eq!(off, (slot + 1) * slot_bytes);
148    Ok((layers, slot_ptrs))
149}
150
151/// 2026-09-25: Load startup adapters: audit each one, check free memory, pack
152/// adapter k into slot k, then build the per-module `[max_loras]` pointer
153/// tables (0 where a slot has no pair) and the `[max_loras]` scale table.
154/// Slots past the adapter count are empty placeholders.
155///
156/// Called through `ModelWeightLoader::load_lora_adapters` from model-engine
157/// `factory::build_model` before `BufferArena::new` and before the free-memory
158/// read that sizes the KV cache, so the pools come out of the KV budget.
159pub fn load_lora_adapters_multi(
160    adapters: &[LoraAdapterInput<'_>],
161    cfg: &ModelConfig,
162    gpu: &dyn GpuBackend,
163    max_loras: usize,
164    max_lora_rank: usize,
165) -> Result<LoraWeights> {
166    check_family(cfg)?;
167    if adapters.is_empty() {
168        bail!("REJECT[no-adapters]: load_lora_adapters_multi called with an empty set");
169    }
170    if adapters.len() > max_loras {
171        bail!(
172            "REJECT[too-many-adapters]: {} --lora-adapter given but --max-loras={} \
173             (pool has {} slots); raise --max-loras or stage the extras on an \
174             $METRALE_LORA_PEER for on-demand RDMA swap",
175            adapters.len(),
176            max_loras,
177            max_loras
178        );
179    }
180
181    // 2026-09-25: Audit every adapter before any allocation.
182    let mut audited: Vec<AuditedAdapter> = Vec::with_capacity(adapters.len());
183    for a in adapters {
184        audited.push(audit_adapter(a.store, &a.peft, cfg, max_lora_rank)?);
185    }
186
187    // 2026-09-25: The expert/router pool is sized from each adapter's audited
188    // keys at `max_lora_expert_rank`, summed over adapters.
189    let expert_rank = max_lora_expert_rank();
190    let expert_total: usize = adapters
191        .iter()
192        .zip(&audited)
193        .map(|(_, au)| {
194            let (ek, rl) = expert_pack::key_lists(&au.router, &au.experts);
195            expert_router_bytes(cfg, &ek, &rl, expert_rank)
196        })
197        .sum();
198
199    // 2026-09-25: Refuse unless free memory is at least twice both pools. Then
200    // one allocation for all slots, zeroed, so padding rows and columns and
201    // unused slots contribute nothing to the padded-rank contraction.
202    let pool_bytes = pool_slot_bytes(cfg, max_lora_rank) * max_loras;
203    let free = gpu.free_memory()?;
204    if (pool_bytes + expert_total) * 2 > free {
205        bail!(
206            "OOM pre-flight (LoRA pool): {:.1} MiB attn pool ({} slots) + {:.1} MiB \
207             expert/router pool would leave < 1× headroom of {:.1} MiB free; every \
208             pool byte comes directly out of the KV-cache budget on GB10 unified memory",
209            pool_bytes as f64 / (1024.0 * 1024.0),
210            max_loras,
211            expert_total as f64 / (1024.0 * 1024.0),
212            free as f64 / (1024.0 * 1024.0),
213        );
214    }
215    let pool = gpu.alloc(pool_bytes)?;
216    gpu.memset(pool, 0, pool_bytes)?;
217    // 2026-09-25: One zeroed expert/router pool for all adapters, allocated
218    // only when some adapter has router or expert pairs.
219    let expert_pool = if expert_total > 0 {
220        let ep = gpu.alloc(expert_total)?;
221        gpu.memset(ep, 0, expert_total)?;
222        Some(ep)
223    } else {
224        None
225    };
226    let mut expert_off = 0usize;
227
228    let mut slots: Vec<AdapterSlot> = Vec::with_capacity(adapters.len());
229    let mut a_tabs: BTreeMap<(usize, LoraModule), Vec<u64>> = BTreeMap::new();
230    let mut b_tabs: BTreeMap<(usize, LoraModule), Vec<u64>> = BTreeMap::new();
231    // 2026-09-25: Upload each adapter's raw token-overlay tensors, `None` when
232    // it ships none. The model builds the overlay tables from them later, in
233    // `set_lora_weights`.
234    let mut overlay_raw: Vec<Option<OverlayRawSlot>> = Vec::with_capacity(adapters.len());
235    for (k, a) in adapters.iter().enumerate() {
236        overlay_raw.push(stage_overlay_raw(
237            a.store,
238            &audited[k].overlay,
239            &a.peft,
240            cfg.hidden_size,
241            gpu,
242        )?);
243        let (mut layers, slot_ptrs) = pack_slot(
244            k,
245            &a.name,
246            a.store,
247            &a.peft,
248            &audited[k].attn,
249            cfg,
250            gpu,
251            pool,
252            max_lora_rank,
253        )?;
254        if let Some(ep) = expert_pool {
255            let packed = expert_pack::pack_into(
256                &mut layers,
257                a.store,
258                &a.peft,
259                &audited[k].router,
260                &audited[k].experts,
261                cfg,
262                gpu,
263                ep,
264                expert_rank,
265                &mut expert_off,
266            )?;
267            if packed > 0 {
268                tracing::info!(
269                    "LoRA: slot {k} '{}' packed {packed} router/expert pair(s) \
270                     (expert_rank={expert_rank})",
271                    a.name
272                );
273            }
274        }
275        for ((layer, module), (a_ptr, b_ptr)) in slot_ptrs {
276            a_tabs
277                .entry((layer, module))
278                .or_insert_with(|| vec![0u64; max_loras])[k] = a_ptr;
279            b_tabs
280                .entry((layer, module))
281                .or_insert_with(|| vec![0u64; max_loras])[k] = b_ptr;
282        }
283        slots.push(AdapterSlot {
284            name: a.name.clone(),
285            adapter_config: a.peft.clone(),
286            layers,
287            generation: 0,
288        });
289    }
290
291    // 2026-09-25: Slots `[pinned, max_loras)` are the promotion cache. They get
292    // empty-named placeholders so a swap can write any slot index
293    // (`swap_lora_slot_from_peer` refuses a slot that is not in `slots`). A
294    // placeholder's pool region is zero, its pointer-table and scale entries
295    // are 0, and its empty name maps to the base id. `pinned >= 1`: an empty
296    // adapter set was refused above.
297    let pinned = slots.len();
298    let num_layers = cfg.num_hidden_layers;
299    while slots.len() < max_loras {
300        slots.push(AdapterSlot {
301            name: String::new(),
302            adapter_config: PeftAdapterConfig {
303                r: 1,
304                lora_alpha: 0.0,
305                target_modules: Vec::new(),
306                target_modules_pattern: None,
307                use_rslora: false,
308                layers_to_transform: None,
309                trainable_token_indices: Vec::new(),
310                modules_to_save: Vec::new(),
311                lora_embedding: false,
312            },
313            layers: vec![None; num_layers],
314            generation: 0,
315        });
316    }
317
318    // 2026-09-25: The per-module `[max_loras]` u64 pointer tables. The model
319    // builds each layer's bgmv route from them (model-engine
320    // `impl_lora_rotate.rs`, `install_lora_layers`).
321    let mk = |tab: &[u64]| -> Result<DevicePtr> {
322        let bytes: Vec<u8> = tab.iter().flat_map(|p| p.to_le_bytes()).collect();
323        let d = gpu.alloc(bytes.len())?;
324        gpu.copy_h2d(&bytes, d)?;
325        Ok(d)
326    };
327    let mut tables = BTreeMap::new();
328    for (key, a_tab) in &a_tabs {
329        let b_tab = &b_tabs[key];
330        tables.insert(*key, (mk(a_tab)?, mk(b_tab)?));
331    }
332
333    // 2026-09-25: The `[max_loras]` f32 scale table, 0.0 for unused slots;
334    // `lora_bgmv.cu` reads `scale_table[s]` for the row's slot.
335    debug_assert_eq!(expert_off, expert_total, "expert pool filled exactly");
336    let scale_vals = scale_table_values(adapters, max_loras);
337    let scale_bytes: Vec<u8> = scale_vals.iter().flat_map(|s| s.to_le_bytes()).collect();
338    let scale_table = gpu.alloc(scale_bytes.len())?;
339    gpu.copy_h2d(&scale_bytes, scale_table)?;
340
341    Ok(LoraWeights {
342        name: slots[0].name.clone(),
343        adapter_config: slots[0].adapter_config.clone(),
344        max_rank: max_lora_rank,
345        max_loras,
346        pool,
347        pool_bytes,
348        expert_pool,
349        expert_pool_bytes: expert_total,
350        slots,
351        active: 0,
352        tables,
353        scale_table,
354        // 2026-09-25: One counter per pool index, `max_loras` of them.
355        ref_counts: (0..max_loras).map(|_| AtomicUsize::new(0)).collect(),
356        pinned,
357        last_used: (0..max_loras).map(|_| AtomicU64::new(0)).collect(),
358        lru_tick: AtomicU64::new(0),
359        overlay_raw,
360    })
361}
362
363/// 2026-09-25: Disk swap: audit `store` and pack it into existing pool `slot`
364/// of `lw` in place, with the same `audit_adapter` and `pack_slot` as the
365/// startup load. The slot region is zeroed first, since it still holds the
366/// previous adapter. Then the slot's name, config and pairs are replaced, its
367/// pointer and scale table entries refreshed, and its generation bumped.
368/// Returns the new pairs so the caller can re-install them when the slot is
369/// active.
370///
371/// Refused, before anything is written: `slot >= max_loras`, a slot with
372/// in-flight sequences, and adapters with router, expert or token-overlay
373/// tensors.
374pub fn pack_store_into_slot(
375    lw: &mut LoraWeights,
376    slot: usize,
377    name: &str,
378    store: &WeightStore,
379    peft: &PeftAdapterConfig,
380    cfg: &ModelConfig,
381    gpu: &dyn GpuBackend,
382) -> Result<Vec<Option<LoraLayerWeights>>> {
383    if slot >= lw.max_loras {
384        bail!(
385            "LoRA disk swap: slot {slot} >= max_loras {} (pool has {} slots)",
386            lw.max_loras,
387            lw.max_loras
388        );
389    }
390    let busy = lw.slot_ref_count(slot);
391    if busy > 0 {
392        bail!(
393            "LoRA disk swap REFUSED: slot {slot} has {busy} in-flight sequence(s) \
394             (ref_count>0); cannot replace an adapter mid-decode"
395        );
396    }
397    validate_peft_config(peft, lw.max_rank)?;
398    let audited = audit_adapter(store, peft, cfg, lw.max_rank)?;
399    if expert_pack::present(&audited.router, &audited.experts) {
400        bail!(
401            "LoRA disk swap REFUSED: adapter '{name}' carries router/expert deltas \
402             (Feature-1); runtime slot-swap of the expert pool is a phase-2 followup"
403        );
404    }
405    if !audited.overlay.is_empty() {
406        bail!(
407            "LoRA disk swap REFUSED: adapter '{name}' ships token-overlay tensors \
408             (Feature-2); runtime slot-swap of the overlay tables is a phase-2 \
409             followup (would silently drop the overlay otherwise)"
410        );
411    }
412    let found = audited.attn;
413    let slot_bytes = pool_slot_bytes(cfg, lw.max_rank);
414    gpu.memset(
415        DevicePtr(lw.pool.0 + (slot * slot_bytes) as u64),
416        0,
417        slot_bytes,
418    )?;
419    let (layers, _slot_ptrs) = pack_slot(
420        slot,
421        name,
422        store,
423        peft,
424        &found,
425        cfg,
426        gpu,
427        lw.pool,
428        lw.max_rank,
429    )?;
430    lw.slots[slot].name = name.to_string();
431    lw.slots[slot].adapter_config = peft.clone();
432    lw.slots[slot].layers = layers.clone();
433    lw.refresh_slot_tables(slot, &layers, peft.scaling(), gpu)?;
434    // 2026-09-25: A new generation gives the slot a new adapter id, so a
435    // same-name request cannot hit prefix-cache entries of the old contents.
436    lw.slots[slot].generation = lw.slots[slot].generation.wrapping_add(1);
437    Ok(layers)
438}
439
440/// 2026-09-25: `load_lora_adapters_multi` with one adapter, whose slot name is
441/// empty.
442pub fn load_lora_adapters_generic(
443    adapter_store: &WeightStore,
444    peft: &PeftAdapterConfig,
445    cfg: &ModelConfig,
446    gpu: &dyn GpuBackend,
447    max_loras: usize,
448    max_lora_rank: usize,
449) -> Result<LoraWeights> {
450    let inputs = [LoraAdapterInput {
451        name: String::new(),
452        store: adapter_store,
453        peft: peft.clone(),
454    }];
455    load_lora_adapters_multi(&inputs, cfg, gpu, max_loras, max_lora_rank)
456}
metrale_model_layers/lora/loading.rs

metrale_model_layers/lora/
loading.rs
