spark_model/lora/
overlay_tables.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Token-overlay device pointer tables (Feature 2). The per-adapter-slot
4//! overlays are addressed by `[max_loras]` device tables of pointers/counts at
5//! LOAD-TIME-FIXED addresses, exactly like the attention BGMV `a_table`/
6//! `b_table` (`ops::lora_delta::LoraRoute`): the only per-step kernel argument
7//! is `seq_slot` (device i32, contents re-uploaded each step), so the tables are
8//! stable kernel args across CUDA-graph capture/replay and the overlay launches
9//! capture cleanly. Cell `k == 0` (null pointer) or `n_override[k] == 0` ⇒ that
10//! slot's overlay is skipped by the kernel — zero-overhead for base requests.
11
12use anyhow::Result;
13use spark_runtime::gpu::{DevicePtr, GpuBackend};
14
15use super::overlay_build::EmbedOverlay;
16
17/// The resolved overlay set for the whole adapter pool. Built once in
18/// `set_lora_weights` (Stage 2) from the per-slot [`EmbedOverlay`]s. `None` on
19/// the model ⇒ overlay feature OFF ⇒ every forward hook early-returns
20/// (byte-identical to a no-overlay build).
21pub struct TokenOverlaySet {
22    /// Per-slot compact overlays (`len == max_loras`), retained so the device
23    /// buffers they own outlive the tables that point into them.
24    pub overlays: Vec<Option<EmbedOverlay>>,
25    /// `u64[max_loras]` → `i32*[vocab]` embed slot_map (0 = slot has no overlay).
26    pub embed_slot_map_table: DevicePtr,
27    /// `u64[max_loras]` → `bf16*[n,h]` embed override rows.
28    pub embed_rows_table: DevicePtr,
29    /// `u32[max_loras]` EMBED n_override per slot — the row count of each
30    /// slot's `embed_rows_table` cell, the embed kernel's `slot < n` bound
31    /// (CWE-125 guard). Distinct from `n_override_table`, which is the
32    /// LM_HEAD count (0 for an untied slot with no lm_head overlay).
33    pub embed_n_table: DevicePtr,
34    /// `slot_map` length shared by every resident overlay (each slot records
35    /// the served vocab it was built against; `from_slots` REFUSES a mix) —
36    /// the embed kernel's `ids[r] < vocab` bound (CWE-125 guard).
37    pub vocab: u32,
38    /// `u64[max_loras]` → `bf16*[n,h]` lm_head override rows (== embed cell when tied).
39    pub lmhead_rows_table: DevicePtr,
40    /// `u64[max_loras]` → `u32*[n]` lm_head override ids (== embed cell when tied).
41    pub lmhead_ids_table: DevicePtr,
42    /// `u32[max_loras]` lm_head n_override per slot (0 ⇒ lm_head skip).
43    pub n_override_table: DevicePtr,
44    /// `max(lmhead n_override)` across slots — the lm_head kernel's `grid.y`.
45    pub max_n_override: u32,
46}
47
48/// Pack a `[max_loras]` u64 pointer array to device (le bytes → alloc → h2d),
49/// mirroring the `mk` closure in `loading.rs`.
50fn mk_u64(gpu: &dyn GpuBackend, tab: &[u64]) -> Result<DevicePtr> {
51    let bytes: Vec<u8> = tab.iter().flat_map(|p| p.to_le_bytes()).collect();
52    let d = gpu.alloc(bytes.len())?;
53    gpu.copy_h2d(&bytes, d)?;
54    Ok(d)
55}
56
57fn mk_u32(gpu: &dyn GpuBackend, tab: &[u32]) -> Result<DevicePtr> {
58    let bytes: Vec<u8> = tab.iter().flat_map(|p| p.to_le_bytes()).collect();
59    let d = gpu.alloc(bytes.len())?;
60    gpu.copy_h2d(&bytes, d)?;
61    Ok(d)
62}
63
64impl TokenOverlaySet {
65    /// Build the device tables from per-slot overlays. `tied` ⇒ the lm_head
66    /// reuses each slot's embed override rows/ids (tied vocab head, or a
67    /// quantized head derived from embed): the logit recompute `dot(hidden,
68    /// embed_row)` is exactly the tied output projection. An untied slot that
69    /// ships its own lm_head overlay uses its distinct rows/ids; an untied slot
70    /// that does NOT ⇒ `n_override[k] = 0` (embed-only correction).
71    pub fn from_slots(
72        gpu: &dyn GpuBackend,
73        overlays: Vec<Option<EmbedOverlay>>,
74        max_loras: usize,
75        tied: bool,
76    ) -> Result<Self> {
77        let mut slot_map_tab = vec![0u64; max_loras];
78        let mut embed_rows_tab = vec![0u64; max_loras];
79        let mut embed_n_tab = vec![0u32; max_loras];
80        let mut lmhead_rows_tab = vec![0u64; max_loras];
81        let mut lmhead_ids_tab = vec![0u64; max_loras];
82        let mut n_override_tab = vec![0u32; max_loras];
83        let mut max_n_override = 0u32;
84        let mut vocab = 0u32;
85
86        for (k, ov) in overlays.iter().enumerate() {
87            let Some(ov) = ov else { continue };
88            slot_map_tab[k] = ov.slot_map.0;
89            embed_rows_tab[k] = ov.rows.0;
90            embed_n_tab[k] = ov.n_override;
91            // One build pass sizes every slot_map to the same served vocab; a
92            // mix would make the kernel's single `ids[r] < vocab` bound wrong
93            // for some slot, so REFUSE it rather than guess a bound.
94            anyhow::ensure!(
95                vocab == 0 || vocab == ov.vocab,
96                "token-overlay: slot {k} was built against vocab {} but an \
97                 earlier slot against {vocab}; refusing mixed-vocab overlay tables",
98                ov.vocab
99            );
100            vocab = ov.vocab;
101            let (rows, ids, n) = match (&ov.lmhead, tied) {
102                (Some(lm), _) => (lm.rows.0, lm.ids_dev.0, lm.n_override),
103                (None, true) => (ov.rows.0, ov.ids_dev.0, ov.n_override),
104                (None, false) => (0, 0, 0),
105            };
106            lmhead_rows_tab[k] = rows;
107            lmhead_ids_tab[k] = ids;
108            n_override_tab[k] = n;
109            max_n_override = max_n_override.max(n);
110        }
111
112        Ok(Self {
113            embed_slot_map_table: mk_u64(gpu, &slot_map_tab)?,
114            embed_rows_table: mk_u64(gpu, &embed_rows_tab)?,
115            embed_n_table: mk_u32(gpu, &embed_n_tab)?,
116            vocab,
117            lmhead_rows_table: mk_u64(gpu, &lmhead_rows_tab)?,
118            lmhead_ids_table: mk_u64(gpu, &lmhead_ids_tab)?,
119            n_override_table: mk_u32(gpu, &n_override_tab)?,
120            max_n_override,
121            overlays,
122        })
123    }
124
125    /// True when at least one slot has a resident overlay (else the set is inert
126    /// and the caller can drop it to keep the hooks byte-identical to off).
127    pub fn any_active(&self) -> bool {
128        self.overlays.iter().any(|o| o.is_some())
129    }
130}
131
132#[cfg(test)]
133#[path = "overlay_tables_tests.rs"]
134mod tests;
spark_model/lora/overlay_tables.rs

spark_model/lora/
overlay_tables.rs
