Module lora 

Startup-static PEFT LoRA adapter: remap/validate/pack into the fixed-address rank-padded pool. v0 = one adapter, slot 0, always on.

NAMING: everything here is Peft*/adapter_*/Lora* (adapter sense) — kv_lora_rank/q_lora_rank (atlas-core/src/config.rs:182-207) are MLA vocabulary, not this.

NOTE on leaks: the intermediate WeightStore device copies of the unpadded A/B tensors become garbage after pool packing and are never freed (no dealloc on weight structs anywhere in Atlas). Accepted at holo adapter scale (~tens of MiB).

SDD facade: the surface is split by functional seam into types (the module/AB enums + weight/slot structs + LoraWeights impl), slot_math (pure slot/offset placement + routing), key (classify + adapter identity), env (the $ATLAS_LORA_* hatches + validate_peft_config), and loading (audit/pack + the load entry points). Every public name re-exports at its own visibility so crate::lora::X / spark_model::lora::X paths are stable.

Modules

rdma_stage

RDMA LoRA staging (spark-model half): turn a peer-staged adapter’s manifest into a set of pool-slot LANDING TARGETS (the only place classify_key + the per-slot offset math live), then drive spark_storage::RdmaLoraLoader to RDMA-load the adapter’s A/B straight into a resident slot for fast rotation. Landing is byte-identical to the disk pack (the loader does the same F16/F32→BF16 convert + B row-repack).

Structs

AdapterSlot

One packed pool slot: a resident adapter’s own name/config + its per-layer pairs (a/b DevicePtrs into that slot’s byte sub-region of the shared pool). layers is GLOBAL-layer-indexed (len = num_hidden_layers), the same index the install walk uses.

EmbedOverlay

Stage-2 compact result for one adapter slot’s embed overlay. rows holds the n_override replacement embedding rows in ascending-ids order; slot_map is the [vocab] i32 lookup (-1 default, [id] = compact-row-index) the embed kernel indexes by token id. lmhead is Some only for an UNTIED head that ships its own overlay tensors.

ExpertLoraLayer

One MoE layer’s routed-expert LoRA coverage: a SPARSE map keyed by (expert_index, projection). Real adapters adapt a subset of the (up to 512) experts, so a dense Vec<[LoraPair; 3]> sized num_experts would waste storage and force a 512-wide static walk; the map only holds the packed pairs. LoraPair is Copy, so the whole struct is cheap to Clone on install.

ExpertWork

One expert’s contiguous sorted-row range (the grouped-GEMM row block for that expert). rows == 0 experts are dropped by the planner.

LmHeadOverlay

Distinct output-projection overlay for an untied lm_head. Rows are recomputed into logits by a dot(hidden, row) per overridden id (no slot_map).

LoraAdapterInput

One adapter to pack, for the multi-adapter entry point. store is the adapter’s on-device BF16 WeightStore (host F16/F32→BF16 already done by spark_runtime::weights::adapter::load_adapter_safetensors).

LoraLayerWeights

One full-attention layer’s adapted modules. None = module not adapted. Pairs are the CANONICAL LoraPair from layers::ops::lora_delta (Copy — installed by copy into the layer structs at model build).

LoraWeights

The loaded adapter set: one fixed-address rank-padded pool holding up to max_loras equal-size slots, one AdapterSlot per resident adapter, and per-module [max_loras] device u64 pointer tables (the frozen M2 BGMV contract — filled index k for each packed slot, NULL for the rest).

OverlayRaw

Stage-1 raw upload: the adapter’s own overlay tensors copied into owned device scratch (the source WeightStore is freed after the loader). Dims are recorded so Stage 2 can shape the row-diff and compaction without the store. embed_base/lmhead_base also carry a modules_to_save full-weight replacement (row-diff finds the changed rows either way).

OverlayRawSlot

Stage-1 raw upload + the (host) clamped trainable ids the delta rows align to.

OverlayTensor

A classified overlay tensor: its target embedding and its role.

OverlayTensors

Tensor names the loader has collected for one adapter, partitioned by (module, role). Option<String> = the safetensors key, filled at most once.

SlotView

Task #27: a per-slot snapshot for the pure victim-selection policy. Taken on the model thread at a scheduler-quiescent point (the only place ref_count is authoritative), then handed to select_victim_slot.

TokenOverlaySet

The resolved overlay set for the whole adapter pool. Built once in set_lora_weights (Stage 2) from the per-slot EmbedOverlays. None on the model ⇒ overlay feature OFF ⇒ every forward hook early-returns (byte-identical to a no-overlay build).

Enums

AdapterAb

ExpertProj

Which routed-expert projection a delta targets. Ordered so BTreeMap<(u16, ExpertProj), _> has a stable, testable key order.

LoraModule

LoraTarget

The decoded target of one PEFT LoRA tensor. Attn keeps the existing dense LoraModule path byte-identical; Router / Expert are the Feature-1 additions routed into per-layer ExpertLoraLayer / router storage.

OverlayModule

Which tied embedding an overlay tensor belongs to.

OverlayTensorKind

The role an overlay tensor plays in build_overlay.

VictimError

Why a promotion cannot find a victim slot.

Constants

ROWDIFF_THRESH

Rows whose max abs difference from the served embed table exceeds this are treated as “overridden” by a modules_to_save/baked base_layer. Clears bf16 rounding noise (matching rows land ≤0.05; a real differing row ≥1.3).

Functions

adapter_id_hash

Stable u64 identity for an adapter, derived from its human NAME (never the runtime pool slot index, which is reused across swap/rotation). Task #24: this is the cache-identity key that keeps the KV/prefix cache adapter-correct so a request reuses ONLY blocks computed under the same adapter.

allow_partial_targets

ATLAS_LORA_ALLOW_PARTIAL=1 — load an adapter that names target modules Atlas cannot apply, skipping those and applying the rest.

apply_expert_lora_sorted

Fold proj’s per-expert LoRA deltas onto the SORTED grouped-GEMM output.

apply_router_lora

Fold the router (mlp.gate) LoRA delta onto the routing logits in place, BEFORE top-k. router_in = [n, hidden], gate_logits = [n, num_experts] (modified in place). Chunked by max_rows (the scratch capacity in rows).

build_moe_row_adapter_decode

SOLID Incr-4 (batched decode fold): build the per-row [padded_n] i32 adapter map the device-side MoE gather-BGMV fold reads. One token per sequence at decode, so a per-row map IS a per-seq map (no top_k expansion — the kernel indexes row_adapter[row / top_k]).

build_moe_row_adapter_host

Build the [total_tokens] per-packed-row adapter map for the device-side grouped fold. cu_seqlens_host is the [batch + 1] prefix-sum of per-stream token counts (the packing SSOT); adapter_slots[b] is stream b’s adapter_slot. Each stream’s slot is broadcast across its [cu_seqlens_host[b], cu_seqlens_host[b + 1]) row span. A base stream (adapter_slot < 0) writes -1 (device kernel skips those rows); a stream deferring to the active adapter would resolve < 0 → active at the call site before this, so a genuine base row is a real -1 here (distinct base sentinel — see lora-solid.md §6).

build_overlay

Stage 2 (set_lora_weights): row-diff against the served tables, compact the override rows, build the slot_map, and free the Stage-1 raw scratch. Ok(None) when the slot overrides nothing (silently inert = correct).

build_override_set

Union of (rows that differ from the served base) and (trainable ids), sorted ascending and deduped. This is the final set of vocab rows the overlay replaces. row_diff[i] = row i of the adapter base differs from served.

build_seq_slot_host

Build the per-step seq_slot[padded_n] host buffer the batched bgmv reads, from each real sequence’s adapter_slot. Resolution rules (graph-safe: contents vary per step, buffer address is fixed): real row i (< n): adapter_slots[i] if >= 0, else active — a request with no adapter field carries -1 and DEFERS to the installed active adapter, so a single global adapter (or a rotate re-point) applies to every default row exactly like the n==1 path. pad row i (n..padded_n): -1 — base / no delta (bgmv early-returns). A row that explicitly names the base model (some future -1-means-base convention) is out of scope here; -1 uniformly means “defer to active”.

clamp_trainable_to_vocab

Clamp trainable ids to the served vocab, preserving list order (the delta tensor’s rows align positionally to it).

classify_key

PEFT key → (layer, module, A|B). Every unsupported shape is a NAMED hard rejection — never a skip. Prefix-agnostic on purpose: the Holo base checkpoint keys are model.language_model.layers.{i}.* (weight_prefix auto-detected server-side), but a PEFT trainer wrapping the text trunk emits model.layers.{i}.*; both carry the layer index right after “.layers.”.

classify_overlay_key

Classify a PEFT adapter tensor as a token-overlay tensor, or None if it is an ordinary lora_A/lora_B weight (which super::classify_key handles).

ensure_decode_route_servable

SOLID Incr-4 host-side pre-lookup guard for the BATCHED decode entries (decode_batch_compute_main, mixed_forward): a batch containing a row routed to a NON-active adapter (Refuse) cannot be served by the single-active fold — build_moe_row_adapter_decode defensively maps such rows to base, so proceeding would SILENTLY serve base weights for an adapter-routed request. Call BEFORE any graph lookup/capture so captured padded_n graphs stay route-agnostic. Pure (no self, no GPU) so the decision is unit-testable without hardware.

expert_delta_workitems

PURE: from the base MoE expert_offsets prefix-sum ([num_experts + 1], expert_offsets[e]..expert_offsets[e+1] = expert e’s sorted rows) and the adapter’s adapted-expert set, produce the (expert, row_off, rows) work-items for the delta side-path. Experts with zero routed rows or a malformed offset pair are skipped (never a panic). This is the correctness-critical mapping and is unit-tested without a GPU.

expert_router_bytes

Pure padded-byte estimator for the (separate) expert/router pool, used by the VRAM preflight and pinned by a golden unit test — mirrors super::pool_slot_bytes but over the audited routed-expert + router key set (real adapters target a SUBSET, so this is sized from the audit, never from num_experts × num_layers maxima). Per (layer, expert, proj) and per router layer: (stride·in + out·stride)·2 BF16 bytes, where stride is the DERIVED uint4-aligned expert_pack::packed_stride of max_rank — the same derivation the pack loop uses (SSOT), so sizing and packing agree byte-for-byte even at a non-multiple-of-8 rank cap.

full_attention_layers

is_gdn_key

Whether key names a GDN / linear-attention tensor — the family classify_key rejects outright.

load_lora_adapters_generic

Single-adapter convenience wrapper (packs slot 0 only) — byte-identical to the pre-multi-adapter path. Kept for the unit tests and any single-adapter caller. The name is stamped onto the sole slot.

load_lora_adapters_multi

Model-agnostic MULTI-adapter PEFT load: audit every adapter, VRAM-preflight the N-slot pool, pack each adapter into its slot (0..N-1), and build the per-module [max_loras] pointer tables (index k filled per packed slot, rest NULL). One resident adapter is byte-identical to the single-adapter path (slot 0, off starts at 0).

lora_eager_env

Permanent LoRA debugging hatch: ATLAS_LORA_EAGER=1 (or true) forces eager decode (no CUDA-graph capture) when an adapter is active, so graph-vs-eager output parity can be compared in the field. Read ONCE — the decode graph gate runs per token. Resolved at the point of use rather than cached in a static: the model carries this as ModelLevers::lora_eager for the per-token decode gate, and the remaining callers are one-shot startup checks where a getenv is free.

lora_experts_env

Feature-1 (MoE expert + router LoRA) master switch. ATLAS_LORA_EXPERTS=1 (or true) opts INTO loading + applying routed-expert / router deltas. DEFAULT OFF: an adapter that targets mlp.experts.* / mlp.gate is a NAMED reject at load unless this is set, so the base path stays byte-identical and the (correctness-first, host-synced, non-graphable) expert side-path is never silently on. Read once.

lora_peer_env

$ATLAS_LORA_PEER (host:port of an atlas-weight-peer staging a rotation set) — when set, arms rotation (eager decode) even for a single resident slot, because an RDMA swap re-points that slot in place. Unset = disk path only, byte-identical to today.

lora_rotate_env

ATLAS_LORA_ROTATE=1 (or true) ARMS runtime adapter rotation: it forces eager decode (no CUDA-graph capture) so a set_active_lora re-point is immediately live (eager-on-rotate — the graph would otherwise replay the previously-captured slot pointers). A pool with >1 resident adapter arms this automatically (see TransformerModel::lora_rotatable), so this env is only needed to arm rotation on a SINGLE resident adapter (e.g. RDMA slot-swap-in-place). Unset + a single startup adapter = today’s behaviour exactly (graphs ON, slot-0 pointers baked). See lora_eager_env on why this is not cached.

max_lora_expert_rank

Feature-1 padded expert/router LoRA rank cap (ATLAS_LORA_EXPERT_RANK, default 16). Separate from --max-lora-rank (the attention pool) because the per-(layer,expert,proj) pool grows ~num_experts × num_layers faster, so a low cap bounds the expert-pool VRAM blow-up. An adapter with r above this is a named reject.

no_batch_verify

ATLAS_LORA_NO_BATCH_VERIFY=1 — restore the old refusal of cross-sequence batched speculative verify while a LoRA adapter is resident.

override_source

Where an overridden id’s replacement row comes from: Some(k) = trainable delta row k (delta WINS when an id is both trainable and baked-different); None = the adapter’s baked base_layer[id].

pack_store_into_slot

Runtime disk swap: audit + pack an already-loaded adapter store into an EXISTING pool slot of lw, in place, and stamp that slot’s name/config/layers. Byte-identical to a startup pack of the same adapter into that slot — same audit, A-contiguous copy, and B row-repack via pack_slot. The slot sub-region is re-zeroed first (a reused slot still holds the prior adapter’s bytes, and pad rows/cols must stay 0 for padded-K correctness). Returns the rebuilt per-layer pairs so the caller can re-install them if the slot is currently active. Like the startup pack, the intermediate store’s device copies leak (small, one-off per swap). Used for the pool-size-1 dynamic-load demo (load a different adapter into the single slot at runtime).

prefill_bgmv_forced

ATLAS_LORA_PREFILL_BGMV=1 — force prefill LoRA through the per-row BGMV instead of the tensor-core GEMM.

reject_pending_overlay

Feature 2 load gate, called from the loader once overlay tensors are collected. The device-side overlay apply is now WIRED (Stage-1 super::overlay_build::stage_overlay_raw upload → Stage-2 super::overlay_build::build_overlay row-diff/compact → the embed_tokens / lm_head forward hooks in crate::model::token_overlay), so trainable_tokens / modules_to_save {embed_tokens, lm_head} tensors are LOADED rather than rejected.

resolve_moe_lora_route

Resolve the Feature-1 MoE-LoRA fold decision for a single-request pass.

routed_prefill_slot_of

#30 (routed-prefill precision): the pure predicate behind LoraWeights::routed_prefill_slot, split out for unit testing. Resolves a request’s adapter_slot (>= 0 → that slot, -1 → active) and returns Some(resolved) ONLY when it routes to a NON-active, in-range slot. Returns None for an active/base request (byte-identical installed-pair path) and for out-of-range slots. Kept in exact lockstep with upload_seq_slot_uniform (resolved == active → DevicePtr(0)).

router_dims

(out_dim, in_dim) of the base router (mlp.gate) projection: [num_experts, hidden]. A router LoRA perturbs the pre-selection routing logits.

select_routed_pair

#30 (routed-prefill precision): pure selector for a routed prefill’s (global_layer, module) LoraPair out of a request slot’s GLOBAL-layer-indexed layers. None when the index is out of range, the layer is unadapted, or the routed adapter does not adapt that module (the caller then falls back to the bgmv/installed path — no delta if the slot’s a_table cell is base). GPU-free + unit-tested so the (layer, module) indexing is verifiable without hardware.

select_victim_slot

Task #27 pure victim-selection policy over the CACHE region only (the caller passes (slot_index, view) for slots [pinned, max_loras) — pinned startup adapters are never candidates, so the resident set and its position-based resolver can never desync). Tiers:

stage_overlay_raw

Stage 1 (loader): upload the classified overlay tensors of one adapter into owned device scratch. Ok(None) when the adapter ships no overlay tensors. embed_full/lmhead_full (modules_to_save) map onto the *_base slot.

validate_peft_config

Adapter-config gates that need build-time context (--max-lora-rank). Parse-time gates (peft_type/DoRA/bias/regex target_modules/…) already ran in atlas_core::config::parse_peft_adapter_config.