spark_model/lora/
target.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Feature-1 (MoE expert + router LoRA) target surface.
4//!
5//! [`classify_key`](super::classify_key) historically decoded a PEFT tensor
6//! name to `(layer, LoraModule, A|B)` where [`LoraModule`] is the flat 7-variant
7//! dense-projection enum (q/k/v/o + gate/up/down). MoE routed-expert and router
8//! deltas need a third dimension the dense enum cannot carry: the EXPERT INDEX
9//! (`mlp.experts.{N}.{gate,up,down}_proj`) and the router (`mlp.gate`, distinct
10//! from the dense `mlp.gate_proj`). [`LoraTarget`] wraps the dense enum and adds
11//! those two variants so the classifier return keeps `LoraModule` intact for the
12//! attention/dense path (zero blast-radius on the S-LoRA BGMV route tables) while
13//! routing expert/router deltas into their own sparse per-layer storage.
14//!
15//! CORRECTNESS-FIRST (Feature-1 phase 1): the internal representation is always
16//! PER-EXPERT ([`ExpertLoraLayer`], a sparse `BTreeMap<(expert, proj), LoraPair>`),
17//! applied via the existing `apply_lora_delta` fold (no new CUDA kernel). Fused
18//! on-disk import (`target_parameters` / fused `experts.gate_up_proj`) and the
19//! grouped-BGMV / fused-epilogue kernels are explicit follow-ups (phases 2/3).
20
21use std::collections::BTreeMap;
22
23use atlas_core::config::ModelConfig;
24
25use crate::layers::ops::lora_delta::LoraPair;
26
27use super::LoraModule;
28
29/// Which routed-expert projection a delta targets. Ordered so
30/// `BTreeMap<(u16, ExpertProj), _>` has a stable, testable key order.
31#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
32pub enum ExpertProj {
33    Gate,
34    Up,
35    Down,
36}
37
38impl ExpertProj {
39    /// PEFT suffix (the leaf `.`-segment of the on-disk tensor name).
40    pub fn peft_name(&self) -> &'static str {
41        match self {
42            Self::Gate => "gate_proj",
43            Self::Up => "up_proj",
44            Self::Down => "down_proj",
45        }
46    }
47
48    /// (out_dim, in_dim) of the base routed-expert projection on `layer`.
49    ///
50    /// Uses the PER-LAYER routed intermediate ([`ModelConfig::moe_intermediate_size_for`],
51    /// e.g. 512 on Holo-3.1-35B-A3B) — NEVER the dense `intermediate_size`
52    /// (5120), which belongs to the standalone SwiGLU FFN. gate/up map hidden→
53    /// inter, down maps inter→hidden.
54    pub fn dims(&self, cfg: &ModelConfig, layer: usize) -> (usize, usize) {
55        let h = cfg.hidden_size;
56        let inter = cfg.moe_intermediate_size_for(layer);
57        match self {
58            Self::Gate | Self::Up => (inter, h),
59            Self::Down => (h, inter),
60        }
61    }
62}
63
64/// (out_dim, in_dim) of the base router (`mlp.gate`) projection: `[num_experts,
65/// hidden]`. A router LoRA perturbs the pre-selection routing logits.
66pub fn router_dims(cfg: &ModelConfig) -> (usize, usize) {
67    (cfg.num_experts, cfg.hidden_size)
68}
69
70/// The decoded target of one PEFT LoRA tensor. `Attn` keeps the existing dense
71/// [`LoraModule`] path byte-identical; `Router` / `Expert` are the Feature-1
72/// additions routed into per-layer [`ExpertLoraLayer`] / router storage.
73#[derive(Debug, Clone, Copy, PartialEq, Eq)]
74pub enum LoraTarget {
75    /// One of the 7 dense projections (q/k/v/o/gate/up/down_proj).
76    Attn(LoraModule),
77    /// The MoE router `mlp.gate` (`[num_experts, hidden]`).
78    Router,
79    /// One routed expert's projection (`mlp.experts.{n}.{proj}`).
80    Expert { n: u16, proj: ExpertProj },
81}
82
83/// One MoE layer's routed-expert LoRA coverage: a SPARSE map keyed by
84/// `(expert_index, projection)`. Real adapters adapt a subset of the (up to 512)
85/// experts, so a dense `Vec<[LoraPair; 3]>` sized `num_experts` would waste
86/// storage and force a 512-wide static walk; the map only holds the packed pairs.
87/// `LoraPair` is `Copy`, so the whole struct is cheap to `Clone` on install.
88#[derive(Clone, Default)]
89pub struct ExpertLoraLayer {
90    pub pairs: BTreeMap<(u16, ExpertProj), LoraPair>,
91}
92
93impl ExpertLoraLayer {
94    /// This layer's `(expert, proj)` pair, if adapted.
95    pub fn pair(&self, expert: u16, proj: ExpertProj) -> Option<&LoraPair> {
96        self.pairs.get(&(expert, proj))
97    }
98
99    /// Sorted, deduped list of the expert indices this layer adapts — the
100    /// SSOT for "which experts get a delta applied" in the forward side-path.
101    pub fn adapted_experts(&self) -> Vec<u16> {
102        let mut v: Vec<u16> = self.pairs.keys().map(|(e, _)| *e).collect();
103        v.dedup();
104        v
105    }
106
107    pub fn is_empty(&self) -> bool {
108        self.pairs.is_empty()
109    }
110}
111
112/// Pure padded-byte estimator for the (separate) expert/router pool, used by the
113/// VRAM preflight and pinned by a golden unit test — mirrors
114/// `super::pool_slot_bytes` but over the audited routed-expert + router key
115/// set (real adapters target a SUBSET, so this is sized from the audit, never
116/// from `num_experts × num_layers` maxima). Per (layer, expert, proj) and per
117/// router layer: `(stride·in + out·stride)·2` BF16 bytes, where `stride` is
118/// the DERIVED uint4-aligned `expert_pack::packed_stride` of
119/// `max_rank` — the same derivation the pack loop uses (SSOT), so sizing and
120/// packing agree byte-for-byte even at a non-multiple-of-8 rank cap.
121pub fn expert_router_bytes(
122    cfg: &ModelConfig,
123    expert_keys: &[(usize, ExpertProj)],
124    router_layers: &[usize],
125    max_rank: usize,
126) -> usize {
127    let stride = super::expert_pack::packed_stride(max_rank);
128    let per = |out: usize, inp: usize| (stride * inp + out * stride) * 2;
129    let experts: usize = expert_keys
130        .iter()
131        .map(|(layer, proj)| {
132            let (out, inp) = proj.dims(cfg, *layer);
133            per(out, inp)
134        })
135        .sum();
136    let routers: usize = router_layers
137        .iter()
138        .map(|_| {
139            let (out, inp) = router_dims(cfg);
140            per(out, inp)
141        })
142        .sum();
143    experts + routers
144}
145
146#[cfg(test)]
147#[path = "target_tests.rs"]
148mod tests;
spark_model/lora/target.rs

spark_model/lora/
target.rs
