spark_model/lora/
rdma_stage.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! RDMA LoRA staging (spark-model half): turn a peer-staged adapter's manifest
4//! into a set of pool-slot LANDING TARGETS (the only place `classify_key` +
5//! the per-slot offset math live), then drive `spark_storage::RdmaLoraLoader`
6//! to RDMA-load the adapter's A/B straight into a resident slot for fast
7//! rotation. Landing is byte-identical to the disk pack (the loader does the
8//! same F16/F32→BF16 convert + B row-repack).
9//!
10//! Gated behind `$ATLAS_LORA_PEER` at the call site; when unset the disk
11//! rotation path is unchanged.
12
13use std::collections::BTreeMap;
14
15use anyhow::{Result, anyhow, bail};
16use atlas_core::config::{ModelConfig, PeftAdapterConfig};
17use spark_runtime::gpu::DevicePtr;
18use spark_storage::weight_peer::WeightManifest;
19use spark_storage::{LoraAbKind, LoraLandTarget};
20
21use super::{
22    AdapterAb, LoraLayerWeights, LoraModule, LoraTarget, classify_key, module_slot_offsets,
23    pool_slot_bytes, slot_base_offset,
24};
25use crate::layers::ops::lora_delta::LoraPair;
26use crate::weight_map::DenseWeight;
27
28/// Build the landing targets for one adapter's manifest into pool `slot`. Each
29/// `lora_A/lora_B` tensor is classified to (layer, module, A|B) and mapped to
30/// its byte sub-region `pool + slot*slot_bytes + a_off|b_off`. The adapter's
31/// real rank r is read from the tensor shape (A=`[r,in]`, B=`[out,r]`). Rejections
32/// from `classify_key` (GDN / wrong-layer / non-PEFT key) fire here
33/// too — never a silent skip.
34pub fn build_land_targets(
35    manifest: &WeightManifest,
36    cfg: &ModelConfig,
37    pool: DevicePtr,
38    slot: usize,
39    max_rank: usize,
40) -> Result<Vec<LoraLandTarget>> {
41    let base = pool.0 + slot_base_offset(slot, cfg, max_rank) as u64;
42    let mut targets = Vec::with_capacity(manifest.tensors.len());
43    let mut pairs: BTreeMap<(usize, LoraModule), [Option<usize>; 2]> = BTreeMap::new();
44    for rec in &manifest.tensors {
45        let (layer, target, ab) = classify_key(&rec.name, cfg)?;
46        // RDMA slot-swap stages ONLY the equal-size attention/dense pool. Router
47        // and routed-expert LoRA (Feature-1) live in a separate expert pool with
48        // its own offset math and are not RDMA-swappable in P1 — reject by name.
49        let module = match target {
50            LoraTarget::Attn(m) => m,
51            LoraTarget::Router | LoraTarget::Expert { .. } => bail!(
52                "lora rdma: '{}' is a router/expert delta (Feature-1); RDMA \
53                 slot-swap stages the attention pool only",
54                rec.name
55            ),
56        };
57        let (a_off, b_off) = module_slot_offsets(cfg, max_rank, layer, module)
58            .ok_or_else(|| anyhow!("lora rdma: layer {layer} not a full-attention slot layer"))?;
59        let (out_dim, in_dim) = module.dims(cfg);
60        // Audit the complete on-wire geometry before deriving r. The landing
61        // transforms trust these dimensions when copying into a fixed-size
62        // pool sub-region, so accepting an extra/missing/wrong dimension here
63        // can otherwise become a panic or a mispacked adapter later.
64        let rank = match ab {
65            AdapterAb::A if rec.shape.len() == 2 && rec.shape[1] == in_dim as u64 => {
66                rec.shape[0] as usize
67            }
68            AdapterAb::B if rec.shape.len() == 2 && rec.shape[0] == out_dim as u64 => {
69                rec.shape[1] as usize
70            }
71            AdapterAb::A => bail!(
72                "REJECT[shape-mismatch]: '{}' is {:?}, expected [r, {}]",
73                rec.name,
74                rec.shape,
75                in_dim
76            ),
77            AdapterAb::B => bail!(
78                "REJECT[shape-mismatch]: '{}' is {:?}, expected [{}, r]",
79                rec.name,
80                rec.shape,
81                out_dim
82            ),
83        };
84        if rank == 0 {
85            bail!("REJECT[shape-mismatch]: '{}' has zero rank", rec.name);
86        }
87        if rank > max_rank {
88            bail!(
89                "lora rdma: adapter rank {rank} for {} exceeds pool max_rank {max_rank}",
90                rec.name
91            );
92        }
93        let (kind, off) = match ab {
94            AdapterAb::A => (LoraAbKind::A, a_off),
95            AdapterAb::B => (LoraAbKind::B, b_off),
96        };
97        let pair = pairs.entry((layer, module)).or_default();
98        let cell = &mut pair[ab as usize];
99        if cell.is_some() {
100            bail!("REJECT[duplicate-tensor]: two tensors map to layer {layer} {module:?} {ab:?}");
101        }
102        *cell = Some(rank);
103        targets.push(LoraLandTarget {
104            tensor_name: rec.name.clone(),
105            kind,
106            dst: base + off as u64,
107            out_dim,
108            in_dim,
109            rank,
110            max_rank,
111        });
112    }
113    if targets.is_empty() {
114        bail!("lora rdma: adapter manifest has no lora_A/lora_B tensors");
115    }
116    for ((layer, module), pair) in pairs {
117        let [Some(a_rank), Some(b_rank)] = pair else {
118            bail!(
119                "REJECT[unpaired-tensor]: layer {layer} {module:?} has only one of lora_A/lora_B"
120            );
121        };
122        if a_rank != b_rank {
123            bail!(
124                "REJECT[rank-mismatch]: layer {layer} {module:?} has A rank {a_rank}, B rank {b_rank}"
125            );
126        }
127    }
128    Ok(targets)
129}
130
131/// Rebuild a slot's per-layer [`LoraLayerWeights`] after an in-place RDMA
132/// reload — the A/B bytes changed AND the adapter's r/scale may differ, so the
133/// `LoraPair`s (which bake rank + scale) must be rebuilt, not just re-pointed.
134/// Pointers are deterministic (`pool + slot*slot_bytes + off`); this does NOT
135/// touch the GPU. Modules present are those with a target of the matching kind.
136pub fn rebuild_slot_layers(
137    targets: &[LoraLandTarget],
138    cfg: &ModelConfig,
139    peft: &PeftAdapterConfig,
140    pool: DevicePtr,
141    slot: usize,
142    max_rank: usize,
143) -> Result<Vec<Option<LoraLayerWeights>>> {
144    let scale = peft.scaling();
145    let base = pool.0 + slot_base_offset(slot, cfg, max_rank) as u64;
146    let mut layers: Vec<Option<LoraLayerWeights>> =
147        (0..cfg.num_hidden_layers).map(|_| None).collect();
148    // Group targets by (layer, module): we need both A and B present to build a
149    // pair. Re-derive from classify (targets carry only geometry, not keys' layer).
150    // Simpler: walk the pool layout and, for each (layer, module), find whether a
151    // target lands there (by matching dst).
152    // Same walk as `pool_slot_bytes` / `pack_slot` / `module_slot_offsets`:
153    // every layer, applicable modules only. Walking full-attention layers x
154    // ALL would ask for the offset of a module that layer cannot carry (e.g.
155    // the GDN `out_proj` on an attention layer), which now correctly has none.
156    for rec_layer in 0..cfg.num_hidden_layers {
157        let mut lw = LoraLayerWeights::empty(rec_layer);
158        let mut any = false;
159        for module in LoraModule::ALL {
160            if !module.applies_to_layer(cfg, rec_layer) {
161                continue;
162            }
163            let (a_off, b_off) = module_slot_offsets(cfg, max_rank, rec_layer, module)
164                .expect("applicable module has a slot offset");
165            let a_dst = base + a_off as u64;
166            let b_dst = base + b_off as u64;
167            let a_t = targets
168                .iter()
169                .find(|t| t.kind == LoraAbKind::A && t.dst == a_dst);
170            let b_t = targets
171                .iter()
172                .find(|t| t.kind == LoraAbKind::B && t.dst == b_dst);
173            if let (Some(a), Some(b)) = (a_t, b_t) {
174                let (out_dim, in_dim) = module.dims(cfg);
175                if a.rank != b.rank || a.rank != peft.r {
176                    bail!(
177                        "REJECT[rank-mismatch]: layer {rec_layer} {module:?} has target ranks A={} B={}, config r={}",
178                        a.rank,
179                        b.rank,
180                        peft.r
181                    );
182                }
183                if a.max_rank != max_rank
184                    || b.max_rank != max_rank
185                    || a.out_dim != out_dim
186                    || b.out_dim != out_dim
187                    || a.in_dim != in_dim
188                    || b.in_dim != in_dim
189                {
190                    bail!(
191                        "REJECT[landing-geometry]: layer {rec_layer} {module:?} target geometry does not match the pool layout"
192                    );
193                }
194                let pair = LoraPair {
195                    a: DenseWeight {
196                        weight: DevicePtr(a_dst),
197                    },
198                    b: DenseWeight {
199                        weight: DevicePtr(b_dst),
200                    },
201                    rank: a.rank as u32,
202                    k_in: in_dim as u32,
203                    n_out: out_dim as u32,
204                    scale,
205                    max_rank: max_rank as u32,
206                };
207                match module {
208                    LoraModule::QProj => lw.q_proj = Some(pair),
209                    LoraModule::KProj => lw.k_proj = Some(pair),
210                    LoraModule::VProj => lw.v_proj = Some(pair),
211                    LoraModule::OProj => lw.o_proj = Some(pair),
212                    LoraModule::GateProj => lw.gate_proj = Some(pair),
213                    LoraModule::UpProj => lw.up_proj = Some(pair),
214                    LoraModule::DownProj => lw.down_proj = Some(pair),
215                    LoraModule::OutProj => lw.out_proj = Some(pair),
216                }
217                any = true;
218            }
219        }
220        if any {
221            layers[rec_layer] = Some(lw);
222        }
223    }
224    Ok(layers)
225}
226
227/// The per-slot byte length (re-exported so the swap path can re-zero exactly
228/// one slot's sub-region before an in-place reload).
229pub fn slot_bytes(cfg: &ModelConfig, max_rank: usize) -> usize {
230    pool_slot_bytes(cfg, max_rank)
231}
232
233/// Fetch a peer-staged adapter's manifest over the `weight_peer` control
234/// channel (connect → request → read manifest, then drop the connection).
235/// Needed to build landing targets before the loader's own verbs handshake.
236#[cfg(feature = "cuda")]
237pub fn fetch_adapter_manifest(peer_addr: &str, adapter_id: &str) -> Result<WeightManifest> {
238    use std::net::TcpStream;
239
240    use anyhow::Context;
241    use spark_storage::weight_peer::{read_weight_manifest, write_model_request};
242
243    let mut stream =
244        TcpStream::connect(peer_addr).with_context(|| format!("connect lora peer {peer_addr}"))?;
245    stream.set_nodelay(true).ok();
246    write_model_request(&mut stream, adapter_id).context("send adapter request")?;
247    let manifest = read_weight_manifest(&mut stream).context("read adapter manifest")?;
248    // Drop the connection without a transport handshake; the loader reconnects
249    // for the actual one-sided read.
250    let _ = std::io::Write::write_all(&mut stream, &[]);
251    Ok(manifest)
252}
253
254#[cfg(test)]
255#[path = "rdma_stage_tests.rs"]
256mod tests;
spark_model/lora/rdma_stage.rs

spark_model/lora/
rdma_stage.rs
