metrale_model_layers/lora/
slot_math.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: GPU-free LoRA slot math: victim selection for the promotion
4//! cache, the per-row `seq_slot` buffer, scale-table values, the routed-prefill
5//! slot and pair, and the per-slot byte layout of the adapter pool.
6//!
7//! Owner: model-layers (lora).
8//! Invariants:
9//! - `pool_slot_bytes`, `module_slot_offsets` and `pack_slot` (`loading.rs`)
10//!   walk one layout: layers ascending, then `LoraModule::ALL` filtered by
11//!   `applies_to_layer`, A (`max_rank * in` elements) then B
12//!   (`out * max_rank`), BF16.
13
14use metrale_config::ModelConfig;
15
16use super::*;
17use crate::layers::ops::lora_delta::LoraPair;
18
19pub(crate) const BF16_BYTES: usize = 2;
20
21/// 2026-09-25: Pick the slot a promotion overwrites, from `cache`, the
22/// `(slot_index, view)` pairs the caller passes (model-engine passes
23/// `cache_slot_views()`, slots `[pinned, max_loras)`):
24///   1. the first unfilled slot;
25///   2. else the `ref_count == 0` slot with the smallest `last_used`;
26///   3. else `Err(PoolFull)`.
27pub fn select_victim_slot(cache: &[(usize, SlotView)]) -> Result<usize, VictimError> {
28    if let Some((idx, _)) = cache.iter().find(|(_, v)| !v.filled) {
29        return Ok(*idx);
30    }
31    cache
32        .iter()
33        .filter(|(_, v)| v.ref_count == 0)
34        .min_by_key(|(_, v)| v.last_used)
35        .map(|(idx, _)| *idx)
36        .ok_or(VictimError::PoolFull)
37}
38
39/// 2026-09-25: The host `seq_slot[padded_n]` buffer the bgmv reads, one entry
40/// per row:
41///   row i < n: `adapter_slots[i]` if `>= 0`, else `active` (a negative slot
42///     means the active adapter);
43///   row i >= n (padding): `-1`, which `lora_bgmv.cu` treats as no delta.
44pub fn build_seq_slot_host(adapter_slots: &[i32], padded_n: usize, active: i32) -> Vec<i32> {
45    let n = adapter_slots.len();
46    (0..padded_n)
47        .map(|i| {
48            if i < n {
49                let s = adapter_slots[i];
50                if s >= 0 { s } else { active }
51            } else {
52                -1
53            }
54        })
55        .collect()
56}
57
58/// 2026-10-03: Which attention LoRA sites a step of one sequence runs.
59#[derive(Debug, Clone, Copy, PartialEq, Eq)]
60pub enum LoraSites {
61    /// 2026-10-03: Sites that fold the installed pair when no slot buffer is given: one-row decode
62    /// (`apply_q_lora`, `apply_kv_lora`, the o_proj fold) and prefill (`prefill_one_proj`,
63    /// `paged_oproj`).
64    PairFallback,
65    /// 2026-10-03: The multi-sequence sites an MTP verify's K rows run (`ms_qkv_apply_lora`,
66    /// `ms_o_proj_lora`): they have no pair fallback and fold nothing without a slot buffer.
67    MultiSeq,
68}
69
70/// 2026-10-03: The `count`-row slot buffer a step of one sequence uploads, or `None` for the
71/// installed-pair path: only a request that resolves to the active adapter, on sites with a pair
72/// fallback, takes the pair. Multi-sequence sites always get their rows' slots, or the attention
73/// would run without the adapter (the active-adapter MTP verify did until 2026-10-03).
74pub fn uniform_seq_slots(
75    adapter_slot: i32,
76    active: i32,
77    count: usize,
78    sites: LoraSites,
79) -> Option<Vec<i32>> {
80    let resolved = if adapter_slot >= 0 {
81        adapter_slot
82    } else {
83        active
84    };
85    if sites == LoraSites::PairFallback && resolved == active {
86        return None;
87    }
88    Some(build_seq_slot_host(
89        &vec![adapter_slot; count],
90        count,
91        active,
92    ))
93}
94
95/// 2026-09-25: Values of the `[max_loras]` f32 scale table: entry `k` is
96/// adapter `k`'s `scaling()` (alpha/r, or alpha/√r under rsLoRA), 0.0 for
97/// `k >= adapters.len()`.
98pub(crate) fn scale_table_values(adapters: &[LoraAdapterInput<'_>], max_loras: usize) -> Vec<f32> {
99    let mut v = vec![0.0f32; max_loras];
100    for (k, a) in adapters.iter().enumerate() {
101        v[k] = a.peft.scaling();
102    }
103    v
104}
105
106/// 2026-09-25: The predicate behind [`LoraWeights::routed_prefill_slot`].
107/// Resolves `adapter_slot` (`>= 0` is that slot, negative is `active`) and
108/// returns it only when it is not `active` and is `< num_slots`; otherwise
109/// `None`.
110pub fn routed_prefill_slot_of(adapter_slot: i32, active: usize, num_slots: usize) -> Option<usize> {
111    let resolved = if adapter_slot >= 0 {
112        adapter_slot as usize
113    } else {
114        active
115    };
116    (resolved != active && resolved < num_slots).then_some(resolved)
117}
118
119/// 2026-09-25: The [`LoraPair`] for (global layer, module) in a slot's
120/// global-layer-indexed `layers`. `None` when the index is out of range, the
121/// layer has no pairs, or the adapter has no pair for that module.
122pub fn select_routed_pair(
123    layers: &[Option<LoraLayerWeights>],
124    global_layer_idx: usize,
125    module: LoraModule,
126) -> Option<&LoraPair> {
127    layers
128        .get(global_layer_idx)
129        .and_then(|o| o.as_ref())
130        .and_then(|l| l.module_pair(module))
131}
132
133/// 2026-09-25: Bytes of one pool slot: over every layer and every module that
134/// `applies_to_layer`, `(max_rank * in + out * max_rank) * 2`.
135pub(crate) fn pool_slot_bytes(cfg: &ModelConfig, max_rank: usize) -> usize {
136    (0..cfg.num_hidden_layers)
137        .map(|layer| {
138            LoraModule::ALL
139                .iter()
140                .filter(|m| m.applies_to_layer(cfg, layer))
141                .map(|m| {
142                    let (out, inp) = m.dims(cfg);
143                    (max_rank * inp + out * max_rank) * BF16_BYTES
144                })
145                .sum::<usize>()
146        })
147        .sum()
148}
149
150/// 2026-09-25: Byte offset of slot `slot` in the pool:
151/// `slot * pool_slot_bytes`.
152// 2026-09-25: Its callers are `rdma_stage` and tests; `rdma_stage` compiles only
153// on unix, with `cuda` or under test (`mod.rs`).
154#[cfg_attr(not(all(feature = "cuda", unix)), allow(dead_code))]
155pub(crate) fn slot_base_offset(slot: usize, cfg: &ModelConfig, max_rank: usize) -> usize {
156    slot * pool_slot_bytes(cfg, max_rank)
157}
158
159/// 2026-09-25: The (a_off, b_off) of (layer, module) within a slot: the
160/// offsets `pack_slot` reaches for that pair. `None` when `applies_to_layer`
161/// is false for it or the layer is out of range. Called by `rdma_stage` and
162/// tests.
163#[cfg_attr(not(all(feature = "cuda", unix)), allow(dead_code))]
164pub(crate) fn module_slot_offsets(
165    cfg: &ModelConfig,
166    max_rank: usize,
167    target_layer: usize,
168    target_module: LoraModule,
169) -> Option<(usize, usize)> {
170    // 2026-09-25: The same walk as `pack_slot` and `pool_slot_bytes`; RDMA
171    // landing writes at these offsets.
172    let mut off = 0usize;
173    for layer_idx in 0..cfg.num_hidden_layers {
174        for module in LoraModule::ALL {
175            if !module.applies_to_layer(cfg, layer_idx) {
176                continue;
177            }
178            let (out_dim, in_dim) = module.dims(cfg);
179            let a_off = off;
180            let b_off = off + max_rank * in_dim * BF16_BYTES;
181            off = b_off + out_dim * max_rank * BF16_BYTES;
182            if layer_idx == target_layer && module == target_module {
183                return Some((a_off, b_off));
184            }
185        }
186    }
187    None
188}
189
190#[cfg(test)]
191#[path = "slot_math_tests.rs"]
192mod tests;
metrale_model_layers/lora/slot_math.rs

metrale_model_layers/lora/
slot_math.rs
