metrale_model_layers/lora/
expert_apply.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Host-driven MoE LoRA folds built on `apply_lora_delta`: the
4//! router delta onto the routing logits (`apply_router_lora`), and per-expert
5//! deltas onto the expert-sorted grouped-GEMM output
6//! (`apply_expert_lora_sorted`), whose row blocks come from a host copy of
7//! `expert_offsets`. Their only caller in the tree is model-engine's
8//! `tests/moe_lora_delta_parity.rs`; the MoE layer folds through
9//! `layers/moe/lora.rs` and `ops::moe_lora_grouped`.
10//!
11//! Owner: model-layers (lora).
12//! Invariants:
13//! - Each `apply_lora_delta` call gets at most `max(max_rows, 1)` rows.
14
15use anyhow::Result;
16use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
17
18use super::{ExpertLoraLayer, ExpertProj};
19use crate::layers::ops::lora_delta::{LoraKernels, LoraPair, apply_lora_delta};
20
21const BF16_BYTES: u64 = 2;
22
23/// 2026-09-25: One expert's contiguous block of sorted rows. The planner,
24/// `expert_delta_workitems`, emits none with `rows == 0`.
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
26pub struct ExpertWork {
27    pub expert: u16,
28    pub row_off: u32,
29    pub rows: u32,
30}
31
32/// 2026-09-25: Work items for the adapted experts, from the `expert_offsets`
33/// prefix sum (`[num_experts + 1]`; expert `e`'s sorted rows are
34/// `expert_offsets[e]..expert_offsets[e + 1]`). An expert index past the
35/// table, or with `end <= start`, is skipped.
36pub fn expert_delta_workitems(expert_offsets: &[u32], adapted: &[u16]) -> Vec<ExpertWork> {
37    let n_experts = expert_offsets.len().saturating_sub(1);
38    let mut work = Vec::with_capacity(adapted.len());
39    for &e in adapted {
40        let e_us = e as usize;
41        if e_us >= n_experts {
42            continue;
43        }
44        let start = expert_offsets[e_us];
45        let end = expert_offsets[e_us + 1];
46        if end <= start {
47            continue;
48        }
49        work.push(ExpertWork {
50            expert: e,
51            row_off: start,
52            rows: end - start,
53        });
54    }
55    work
56}
57
58/// 2026-09-25: Fold one projection's delta over `rows` contiguous rows, one
59/// `apply_lora_delta` call per block of at most `max_rows` rows, so scratch of
60/// `max_rows * max_rank` (`lora_xa`) and `max_rows * n_out` (`lora_delta`)
61/// BF16 suffices.
62#[allow(clippy::too_many_arguments)]
63fn fold_chunked(
64    gpu: &dyn GpuBackend,
65    kernels: &LoraKernels,
66    pair: &LoraPair,
67    x: DevicePtr,
68    base_out: DevicePtr,
69    rows: u32,
70    max_rows: u32,
71    lora_xa: DevicePtr,
72    lora_delta: DevicePtr,
73    stream: u64,
74) -> Result<()> {
75    let step = max_rows.max(1);
76    let mut done = 0u32;
77    while done < rows {
78        let m = (rows - done).min(step);
79        let x_row = x.offset((done as u64 * pair.k_in as u64 * BF16_BYTES) as usize);
80        let out_row = base_out.offset((done as u64 * pair.n_out as u64 * BF16_BYTES) as usize);
81        apply_lora_delta(
82            gpu, kernels, pair, x_row, out_row, m, lora_xa, lora_delta, stream,
83        )?;
84        done += m;
85    }
86    Ok(())
87}
88
89/// 2026-09-25: Fold the router (`mlp.gate`) delta onto `gate_logits`
90/// (`[n, num_experts]`, in place) from `router_in` (`[n, hidden]`), in blocks
91/// of at most `max_rows` rows.
92#[allow(clippy::too_many_arguments)]
93pub fn apply_router_lora(
94    gpu: &dyn GpuBackend,
95    kernels: &LoraKernels,
96    pair: &LoraPair,
97    router_in: DevicePtr,
98    gate_logits: DevicePtr,
99    n: u32,
100    max_rows: u32,
101    lora_xa: DevicePtr,
102    lora_delta: DevicePtr,
103    stream: u64,
104) -> Result<()> {
105    fold_chunked(
106        gpu,
107        kernels,
108        pair,
109        router_in,
110        gate_logits,
111        n,
112        max_rows,
113        lora_xa,
114        lora_delta,
115        stream,
116    )
117}
118
119/// 2026-09-25: Fold `proj`'s per-expert deltas onto the expert-sorted
120/// grouped-GEMM output. `x` is the sorted input (`[total_expanded, k_in]`),
121/// `base_out` the sorted output (`[total_expanded, n_out]`, in place), and
122/// `expert_offsets_host` a host copy of `expert_offsets`. Only experts with a
123/// pair for `proj` and at least one row are folded, each over its own row
124/// block.
125#[allow(clippy::too_many_arguments)]
126pub fn apply_expert_lora_sorted(
127    gpu: &dyn GpuBackend,
128    kernels: &LoraKernels,
129    layer: &ExpertLoraLayer,
130    proj: ExpertProj,
131    expert_offsets_host: &[u32],
132    x: DevicePtr,
133    base_out: DevicePtr,
134    max_rows: u32,
135    lora_xa: DevicePtr,
136    lora_delta: DevicePtr,
137    stream: u64,
138) -> Result<()> {
139    let work = expert_delta_workitems(expert_offsets_host, &layer.adapted_experts());
140    for w in work {
141        let Some(pair) = layer.pair(w.expert, proj) else {
142            continue;
143        };
144        let x_row = x.offset((w.row_off as u64 * pair.k_in as u64 * BF16_BYTES) as usize);
145        let out_row = base_out.offset((w.row_off as u64 * pair.n_out as u64 * BF16_BYTES) as usize);
146        fold_chunked(
147            gpu, kernels, pair, x_row, out_row, w.rows, max_rows, lora_xa, lora_delta, stream,
148        )?;
149    }
150    Ok(())
151}
152
153#[cfg(test)]
154#[path = "expert_apply_tests.rs"]
155mod tests;
metrale_model_layers/lora/expert_apply.rs

metrale_model_layers/lora/
expert_apply.rs
