metrale_model_layers/layers/qwen3_ssm/
lora.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: LoRA on the GDN layer: install of the MoE router/expert, dense
4//! FFN and GDN `out_proj` deltas, and the `out_proj` delta's apply step.
5//!
6//! Owner: model-layers (qwen3_ssm).
7//! Invariants: none beyond the types.
8
9use anyhow::Result;
10use metrale_gpu_runtime::gpu::GpuBackend;
11
12use super::Qwen3SsmLayer;
13use crate::layers::FfnComponent;
14use crate::layers::ops::lora_delta::{LoraKernels, LoraPair};
15use crate::lora::ExpertLoraLayer;
16
17impl Qwen3SsmLayer {
18    /// 2026-09-25: Install router and routed-expert LoRA on this layer's MoE FFN
19    /// (`MoeLayer::set_lora_weights`). Errors when the FFN is not MoE.
20    pub fn set_moe_lora_weights(
21        &mut self,
22        router: Option<LoraPair>,
23        experts: ExpertLoraLayer,
24        kernels: LoraKernels,
25        gpu: &dyn GpuBackend,
26    ) -> Result<()> {
27        if let FfnComponent::Moe(m) = &mut self.ffn {
28            return m.set_lora_weights(router, experts, kernels, gpu);
29        }
30        anyhow::bail!(
31            "LoRA: router/expert deltas installed on a linear-attention layer whose \
32             FFN is not MoE (loader/adapter mismatch)"
33        )
34    }
35}
36
37impl Qwen3SsmLayer {
38    /// 2026-09-25: Install dense-FFN LoRA on this layer's `FfnComponent::Dense`
39    /// (`DenseFfnLayer::set_lora_weights`). Errors when the FFN is MoE or
40    /// absent, rather than dropping the deltas.
41    pub fn set_ffn_lora_weights(
42        &mut self,
43        ffn: crate::layers::ops::lora_delta::LoraFfnWeights,
44    ) -> Result<()> {
45        match &mut self.ffn {
46            FfnComponent::Dense(d) => d.set_lora_weights(ffn),
47            FfnComponent::Moe(_) => anyhow::bail!(
48                "LoRA: dense-FFN delta on a linear-attention layer whose FFN is MoE — \
49                 routed-expert deltas belong on set_moe_lora_weights"
50            ),
51            FfnComponent::None => {
52                anyhow::bail!("LoRA: dense-FFN delta on a linear-attention layer that has no FFN")
53            }
54        }
55    }
56}
57
58impl Qwen3SsmLayer {
59    /// 2026-09-25: Install this layer's GDN `out_proj` delta, which
60    /// `apply_lora_out_proj` adds.
61    pub fn set_out_proj_lora(&mut self, pair: LoraPair, kernels: LoraKernels) {
62        self.lora_out_proj = Some((pair, kernels));
63    }
64
65    /// 2026-09-25: `out += pair.scale * (normed_out @ A^T) @ B^T`, in place.
66    ///
67    /// Returns without launching when no delta is installed or
68    /// `METRALE_LORA_NO_FFN=1` (`lora_no_ffn`). Its caller,
69    /// `ssm_tp_all_reduce`, runs it after the TP all-reduce: before that `out`
70    /// is a per-rank partial, and a delta added there would be summed once per
71    /// rank.
72    pub(super) fn apply_lora_out_proj(
73        &self,
74        ctx: &crate::layers::ForwardContext,
75        normed_out: metrale_gpu_runtime::gpu::DevicePtr,
76        out: metrale_gpu_runtime::gpu::DevicePtr,
77        m: u32,
78        stream: u64,
79    ) -> Result<()> {
80        if crate::layers::ops::lora_delta::lora_no_ffn() {
81            return Ok(());
82        }
83        let Some((ref pair, ref kernels)) = self.lora_out_proj else {
84            return Ok(());
85        };
86        crate::layers::ops::lora_delta::apply_lora_delta(
87            ctx.gpu,
88            kernels,
89            pair,
90            normed_out,
91            out,
92            m,
93            ctx.buffers.lora_xa(),
94            ctx.buffers.lora_delta(),
95            stream,
96        )
97    }
98}
metrale_model_layers/layers/qwen3_ssm/lora.rs

metrale_model_layers/layers/qwen3_ssm/
lora.rs
