metrale_model_layers/lora/
overlay_build.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The device build of the token overlay, in two steps because
4//! the adapter's tensors and the served embed/lm_head tables are not alive at
5//! the same time:
6//! - [`stage_overlay_raw`], at adapter load: copy the overlay tensors into
7//!   device buffers of their own ([`OverlayRaw`]) while the adapter's
8//!   [`WeightStore`] exists;
9//! - [`build_overlay`], from model-engine's `set_lora_weights`: diff against
10//!   the served tables, compact the overridden rows, and build the `slot_map`
11//!   ([`EmbedOverlay`]).
12//!
13//! Owner: model-layers (lora).
14//! Invariants: none beyond the types.
15
16use anyhow::{Result, bail};
17use metrale_config::PeftAdapterConfig;
18use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
19use metrale_model_weights::weights::WeightStore;
20
21use super::overlay::{
22    OverlayTensors, ROWDIFF_THRESH, build_override_set, clamp_trainable_to_vocab, override_source,
23};
24use crate::layers::ops::token_overlay::{self, OverlayKernels};
25
26const BF16_BYTES: usize = 2;
27const F32_BYTES: usize = 4;
28
29/// 2026-09-25: The adapter's overlay tensors copied to device buffers of their
30/// own, with their row counts: `*_base` is `[*_r, h]` and `*_delta` is
31/// `[*_t, h]`. `*_base` holds either the `token_adapter` base table or a
32/// `modules_to_save` full table.
33#[derive(Debug, Clone, Copy, Default)]
34pub struct OverlayRaw {
35    pub embed_base: Option<DevicePtr>,
36    pub embed_delta: Option<DevicePtr>,
37    pub embed_r: u32,
38    pub embed_t: u32,
39    pub lmhead_base: Option<DevicePtr>,
40    pub lmhead_delta: Option<DevicePtr>,
41    pub lmhead_r: u32,
42    pub lmhead_t: u32,
43}
44
45/// 2026-09-25: One adapter's [`OverlayRaw`] and its trainable ids.
46pub struct OverlayRawSlot {
47    pub raw: OverlayRaw,
48    /// 2026-09-25: PEFT `trainable_token_indices` as given; `build_overlay`
49    /// clamps them to the served vocab.
50    pub trainable: Vec<u32>,
51}
52
53/// 2026-09-25: One adapter slot's embed overlay: `rows` (`[n_override, h]`
54/// BF16) in the order of `ids_dev` (ascending u32 vocab ids), and `slot_map`
55/// (i32 `[vocab]`, the row index for an overridden id, else -1), which the
56/// embed kernel indexes by token id. `lmhead` is `Some` only for an untied head
57/// with its own overlay tensors.
58#[derive(Debug)]
59pub struct EmbedOverlay {
60    pub rows: DevicePtr,
61    pub ids_dev: DevicePtr,
62    pub slot_map: DevicePtr,
63    pub n_override: u32,
64    /// 2026-09-25: The length of `slot_map`, the served vocab; the embed
65    /// kernel skips token ids `>= vocab` (`token_overlay.cu`).
66    pub vocab: u32,
67    pub lmhead: Option<LmHeadOverlay>,
68}
69
70/// 2026-09-25: The overlay of an untied lm_head: `rows` (`[n_override, h]`
71/// BF16) and their vocab ids. The lm_head kernel writes `dot(hidden, row)` to
72/// each overridden id's logit.
73#[derive(Debug)]
74pub struct LmHeadOverlay {
75    pub rows: DevicePtr,
76    pub ids_dev: DevicePtr,
77    pub n_override: u32,
78}
79
80/// 2026-09-25: f32 → BF16 bits, round to nearest even.
81fn f32_to_bf16(x: f32) -> u16 {
82    let bits = x.to_bits();
83    let round = ((bits >> 16) & 1) + 0x7fff;
84    (bits.wrapping_add(round) >> 16) as u16
85}
86
87/// 2026-09-25: Copy one store tensor to a new device buffer; returns it and
88/// `shape[0]`. An error unless the shape is `[rows, h]`.
89fn stage_tensor(
90    store: &WeightStore,
91    name: &str,
92    h: usize,
93    gpu: &dyn GpuBackend,
94) -> Result<(DevicePtr, u32)> {
95    let t = store.get(name)?;
96    if t.shape.len() != 2 || t.shape[1] != h {
97        bail!(
98            "REJECT[overlay-shape]: '{name}' is {:?}, expected [rows, {h}] (hidden)",
99            t.shape
100        );
101    }
102    let bytes: usize = t.shape.iter().product::<usize>() * t.dtype.byte_size();
103    let dst = gpu.alloc(bytes)?;
104    gpu.copy_d2d(t.ptr, dst, bytes)?;
105    Ok((dst, t.shape[0] as u32))
106}
107
108/// 2026-09-25: Copy one adapter's overlay tensors to device buffers of their
109/// own. `Ok(None)` when it has none, or has a `lora_embedding` tensor (the
110/// audit refuses those). A `modules_to_save` table goes into `*_base` when
111/// there is no `token_adapter` base. An error when neither table has a base.
112pub fn stage_overlay_raw(
113    store: &WeightStore,
114    overlay: &OverlayTensors,
115    peft: &PeftAdapterConfig,
116    h: usize,
117    gpu: &dyn GpuBackend,
118) -> Result<Option<OverlayRawSlot>> {
119    if overlay.is_empty() || overlay.lora_embedding_seen {
120        return Ok(None);
121    }
122    let mut raw = OverlayRaw::default();
123    let embed_base_name = overlay.embed_base.as_ref().or(overlay.embed_full.as_ref());
124    if let Some(name) = embed_base_name {
125        let (p, r) = stage_tensor(store, name, h, gpu)?;
126        raw.embed_base = Some(p);
127        raw.embed_r = r;
128    }
129    if let Some(name) = &overlay.embed_delta {
130        let (p, t) = stage_tensor(store, name, h, gpu)?;
131        raw.embed_delta = Some(p);
132        raw.embed_t = t;
133    }
134    let lmhead_base_name = overlay
135        .lmhead_base
136        .as_ref()
137        .or(overlay.lmhead_full.as_ref());
138    if let Some(name) = lmhead_base_name {
139        let (p, r) = stage_tensor(store, name, h, gpu)?;
140        raw.lmhead_base = Some(p);
141        raw.lmhead_r = r;
142    }
143    if let Some(name) = &overlay.lmhead_delta {
144        let (p, t) = stage_tensor(store, name, h, gpu)?;
145        raw.lmhead_delta = Some(p);
146        raw.lmhead_t = t;
147    }
148    if raw.embed_base.is_none() && raw.lmhead_base.is_none() {
149        bail!(
150            "REJECT[overlay-no-base]: overlay tensors present but no embed/lm_head base row table"
151        );
152    }
153    Ok(Some(OverlayRawSlot {
154        raw,
155        trainable: peft.trainable_token_indices.clone(),
156    }))
157}
158
159/// 2026-09-25: The overridden ids (ascending) and their `[n, h]` BF16 rows on
160/// the device, for the embed and lm_head builds.
161struct Compact {
162    ids: Vec<u32>,
163    rows_dev: DevicePtr,
164    ids_dev: DevicePtr,
165    n: u32,
166}
167
168/// 2026-09-25: Diff the first `min(r, vocab)` rows of `base` against `served`,
169/// add the `kept` trainable ids, and build the BF16 replacement rows, each from
170/// the delta or the base row per `override_source`. `None` when no row is
171/// overridden. Delta rows are read as f32.
172#[allow(clippy::too_many_arguments)]
173fn compact_override(
174    gpu: &dyn GpuBackend,
175    kernels: &OverlayKernels,
176    base: DevicePtr,
177    delta: Option<DevicePtr>,
178    r: u32,
179    served: DevicePtr,
180    vocab: usize,
181    h: usize,
182    kept: &[u32],
183    stream: u64,
184) -> Result<Option<Compact>> {
185    let r_eff = (r as usize).min(vocab);
186    let flags_dev = gpu.alloc(r_eff.max(1))?;
187    if r_eff > 0 {
188        token_overlay::embed_rowdiff(
189            gpu,
190            kernels.rowdiff,
191            base,
192            served,
193            flags_dev,
194            r_eff as u32,
195            h as u32,
196            ROWDIFF_THRESH,
197            stream,
198        )?;
199        gpu.synchronize(stream)?;
200    }
201    let mut flags = vec![0u8; r_eff];
202    if r_eff > 0 {
203        gpu.copy_d2h(flags_dev, &mut flags)?;
204    }
205    let _ = gpu.free(flags_dev);
206    let row_diff: Vec<bool> = flags.iter().map(|&b| b != 0).collect();
207    let ids = build_override_set(&row_diff, kept);
208    if ids.is_empty() {
209        return Ok(None);
210    }
211    let n = ids.len();
212    let mut compact = vec![0u8; n * h * BF16_BYTES];
213    let mut frow = vec![0u8; h * F32_BYTES];
214    for (ci, &id) in ids.iter().enumerate() {
215        let dst = &mut compact[ci * h * BF16_BYTES..(ci + 1) * h * BF16_BYTES];
216        match override_source(id, kept) {
217            Some(k) => {
218                let d = delta.ok_or_else(|| {
219                    anyhow::anyhow!(
220                        "REJECT[overlay-delta-missing]: trainable id {id} but no delta tensor"
221                    )
222                })?;
223                gpu.copy_d2h(d.offset(k * h * F32_BYTES), &mut frow)?;
224                for i in 0..h {
225                    let x = f32::from_le_bytes([
226                        frow[i * 4],
227                        frow[i * 4 + 1],
228                        frow[i * 4 + 2],
229                        frow[i * 4 + 3],
230                    ]);
231                    dst[i * 2..i * 2 + 2].copy_from_slice(&f32_to_bf16(x).to_le_bytes());
232                }
233            }
234            None => {
235                gpu.copy_d2h(base.offset(id as usize * h * BF16_BYTES), dst)?;
236            }
237        }
238    }
239    let rows_dev = gpu.alloc(compact.len())?;
240    gpu.copy_h2d(&compact, rows_dev)?;
241    let ids_bytes: Vec<u8> = ids.iter().flat_map(|i| i.to_le_bytes()).collect();
242    let ids_dev = gpu.alloc(ids_bytes.len())?;
243    gpu.copy_h2d(&ids_bytes, ids_dev)?;
244    Ok(Some(Compact {
245        ids,
246        rows_dev,
247        ids_dev,
248        n: n as u32,
249    }))
250}
251
252/// 2026-09-25: Build one slot's [`EmbedOverlay`]: clamp the trainable ids,
253/// diff and compact the embed rows, build the `slot_map`, and, for an untied
254/// head (`!tied`) with its own base, the lm_head rows. The staged buffers are
255/// freed when an overlay is returned. `Ok(None)` when the slot has no embed
256/// base or overrides no row. An error when the overlay kernels are not loaded.
257/// Model-engine passes `tied` when the lm_head shares the embed buffer or is
258/// served quantized; `TokenOverlaySet::from_slots` then points the lm_head at
259/// the embed rows.
260#[allow(clippy::too_many_arguments)]
261pub fn build_overlay(
262    gpu: &dyn GpuBackend,
263    kernels: &OverlayKernels,
264    slot: &OverlayRawSlot,
265    served_embed: DevicePtr,
266    served_lmhead: DevicePtr,
267    vocab: usize,
268    h: usize,
269    tied: bool,
270    stream: u64,
271) -> Result<Option<EmbedOverlay>> {
272    if kernels.rowdiff.0 == 0 || kernels.embed_overlay.0 == 0 {
273        bail!(
274            "REJECT[overlay-kernels-missing]: adapter ships token-overlay tensors but the \
275             token_overlay CUDA kernels are not loaded (rebuild with the kernel image)"
276        );
277    }
278    let raw = &slot.raw;
279    let Some(embed_base) = raw.embed_base else {
280        return Ok(None);
281    };
282    let (kept, skipped) = clamp_trainable_to_vocab(&slot.trainable, raw.embed_r as usize, vocab)?;
283    if skipped > 0 {
284        tracing::warn!(
285            "LoRA overlay: dropped {skipped} vocab-extension trainable id(s) beyond served vocab {vocab}"
286        );
287    }
288    let Some(embed) = compact_override(
289        gpu,
290        kernels,
291        embed_base,
292        raw.embed_delta,
293        raw.embed_r,
294        served_embed,
295        vocab,
296        h,
297        &kept,
298        stream,
299    )?
300    else {
301        return Ok(None);
302    };
303    let mut slot_map = vec![-1i32; vocab];
304    for (ci, &id) in embed.ids.iter().enumerate() {
305        slot_map[id as usize] = ci as i32;
306    }
307    let sm_bytes: Vec<u8> = slot_map.iter().flat_map(|i| i.to_le_bytes()).collect();
308    let slot_map_dev = gpu.alloc(sm_bytes.len())?;
309    gpu.copy_h2d(&sm_bytes, slot_map_dev)?;
310
311    let lmhead = if let Some(base) = raw.lmhead_base.filter(|_| !tied) {
312        compact_override(
313            gpu,
314            kernels,
315            base,
316            raw.lmhead_delta,
317            raw.lmhead_r,
318            served_lmhead,
319            vocab,
320            h,
321            &kept,
322            stream,
323        )?
324        .map(|c| LmHeadOverlay {
325            rows: c.rows_dev,
326            ids_dev: c.ids_dev,
327            n_override: c.n,
328        })
329    } else {
330        None
331    };
332
333    for p in [
334        raw.embed_base,
335        raw.embed_delta,
336        raw.lmhead_base,
337        raw.lmhead_delta,
338    ]
339    .into_iter()
340    .flatten()
341    {
342        let _ = gpu.free(p);
343    }
344
345    Ok(Some(EmbedOverlay {
346        rows: embed.rows_dev,
347        ids_dev: embed.ids_dev,
348        slot_map: slot_map_dev,
349        n_override: embed.n,
350        vocab: vocab as u32,
351        lmhead,
352    }))
353}
354
355#[cfg(test)]
356#[path = "overlay_build_tests.rs"]
357mod tests;
metrale_model_layers/lora/overlay_build.rs

metrale_model_layers/lora/
overlay_build.rs
