metrale_model_layers/lora/
overlay_tables.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The token overlay's `[max_loras]` device tables of pointers
4//! and counts, one cell per adapter slot, allocated once when the set is
5//! built. The kernels pick a row's cell by its slot (`seq_slot`, or `active`)
6//! and skip a slot with a null `slot_map` or a zero count
7//! (`token_overlay.cu`).
8//!
9//! Owner: model-layers (lora).
10//! Invariants:
11//! - `from_slots` returns a set only when every overlay has the same `vocab`.
12
13use anyhow::Result;
14use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
15
16use super::overlay_build::EmbedOverlay;
17
18/// 2026-09-25: The overlay tables for the whole adapter pool, built from the
19/// per-slot [`EmbedOverlay`]s. Model-engine keeps it only when some slot has an
20/// overlay (`any_active`).
21pub struct TokenOverlaySet {
22    /// 2026-09-25: The per-slot overlays, kept so the buffers the tables point
23    /// into stay allocated.
24    pub overlays: Vec<Option<EmbedOverlay>>,
25    /// 2026-09-25: `u64[max_loras]` of `i32[vocab]` `slot_map` pointers, 0 for
26    /// a slot without an overlay.
27    pub embed_slot_map_table: DevicePtr,
28    /// 2026-09-25: `u64[max_loras]` of embed row pointers (`[n, h]` BF16).
29    pub embed_rows_table: DevicePtr,
30    /// 2026-09-25: `u32[max_loras]` embed row counts; the embed kernel skips a
31    /// `slot_map` entry not below it. `n_override_table` is the lm_head count.
32    pub embed_n_table: DevicePtr,
33    /// 2026-09-25: The `slot_map` length of every overlay (0 when there is
34    /// none); the kernels skip token ids `>= vocab`.
35    pub vocab: u32,
36    /// 2026-09-25: `u64[max_loras]` of lm_head row pointers.
37    pub lmhead_rows_table: DevicePtr,
38    /// 2026-09-25: `u64[max_loras]` of lm_head id pointers (`u32[n]`).
39    pub lmhead_ids_table: DevicePtr,
40    /// 2026-09-25: `u32[max_loras]` lm_head row counts; 0 skips the slot.
41    pub n_override_table: DevicePtr,
42    /// 2026-09-25: The largest lm_head row count, the lm_head kernel's
43    /// `grid.y` (`token_overlay.cu`).
44    pub max_n_override: u32,
45}
46
47fn mk_u64(gpu: &dyn GpuBackend, tab: &[u64]) -> Result<DevicePtr> {
48    let bytes: Vec<u8> = tab.iter().flat_map(|p| p.to_le_bytes()).collect();
49    let d = gpu.alloc(bytes.len())?;
50    gpu.copy_h2d(&bytes, d)?;
51    Ok(d)
52}
53
54fn mk_u32(gpu: &dyn GpuBackend, tab: &[u32]) -> Result<DevicePtr> {
55    let bytes: Vec<u8> = tab.iter().flat_map(|p| p.to_le_bytes()).collect();
56    let d = gpu.alloc(bytes.len())?;
57    gpu.copy_h2d(&bytes, d)?;
58    Ok(d)
59}
60
61impl TokenOverlaySet {
62    /// 2026-09-25: Build the tables. A slot's lm_head cells are its own lm_head
63    /// overlay when it has one; else, when `tied`, its embed rows and ids;
64    /// else empty (count 0). An error when two overlays were built against
65    /// different vocab sizes.
66    pub fn from_slots(
67        gpu: &dyn GpuBackend,
68        overlays: Vec<Option<EmbedOverlay>>,
69        max_loras: usize,
70        tied: bool,
71    ) -> Result<Self> {
72        let mut slot_map_tab = vec![0u64; max_loras];
73        let mut embed_rows_tab = vec![0u64; max_loras];
74        let mut embed_n_tab = vec![0u32; max_loras];
75        let mut lmhead_rows_tab = vec![0u64; max_loras];
76        let mut lmhead_ids_tab = vec![0u64; max_loras];
77        let mut n_override_tab = vec![0u32; max_loras];
78        let mut max_n_override = 0u32;
79        let mut vocab = 0u32;
80
81        for (k, ov) in overlays.iter().enumerate() {
82            let Some(ov) = ov else { continue };
83            slot_map_tab[k] = ov.slot_map.0;
84            embed_rows_tab[k] = ov.rows.0;
85            embed_n_tab[k] = ov.n_override;
86            // 2026-09-25: The kernels take one `vocab` bound for all slots.
87            anyhow::ensure!(
88                vocab == 0 || vocab == ov.vocab,
89                "token-overlay: slot {k} was built against vocab {} but an \
90                 earlier slot against {vocab}; refusing mixed-vocab overlay tables",
91                ov.vocab
92            );
93            vocab = ov.vocab;
94            let (rows, ids, n) = match (&ov.lmhead, tied) {
95                (Some(lm), _) => (lm.rows.0, lm.ids_dev.0, lm.n_override),
96                (None, true) => (ov.rows.0, ov.ids_dev.0, ov.n_override),
97                (None, false) => (0, 0, 0),
98            };
99            lmhead_rows_tab[k] = rows;
100            lmhead_ids_tab[k] = ids;
101            n_override_tab[k] = n;
102            max_n_override = max_n_override.max(n);
103        }
104
105        Ok(Self {
106            embed_slot_map_table: mk_u64(gpu, &slot_map_tab)?,
107            embed_rows_table: mk_u64(gpu, &embed_rows_tab)?,
108            embed_n_table: mk_u32(gpu, &embed_n_tab)?,
109            vocab,
110            lmhead_rows_table: mk_u64(gpu, &lmhead_rows_tab)?,
111            lmhead_ids_table: mk_u64(gpu, &lmhead_ids_tab)?,
112            n_override_table: mk_u32(gpu, &n_override_tab)?,
113            max_n_override,
114            overlays,
115        })
116    }
117
118    /// 2026-09-25: True when at least one slot has an overlay.
119    pub fn any_active(&self) -> bool {
120        self.overlays.iter().any(|o| o.is_some())
121    }
122}
123
124#[cfg(test)]
125#[path = "overlay_tables_tests.rs"]
126mod tests;
metrale_model_layers/lora/overlay_tables.rs

metrale_model_layers/lora/
overlay_tables.rs
