spark_model/layers/qwen3_ssm/
lora.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Feature-1 MoE LoRA install on the GDN / linear-attention layer.
4//!
5//! Linear-attention layers carry NO attention LoRA (their projections are
6//! rejected at classify), but the MoE FFN exists on every layer, so a real
7//! all-layer MoE adapter installs its router + routed-expert deltas here too —
8//! the same `MoeLayer::set_lora_weights` path the full-attention layer uses.
9
10use anyhow::Result;
11use spark_runtime::gpu::GpuBackend;
12
13use super::Qwen3SsmLayer;
14use crate::layers::FfnComponent;
15use crate::layers::ops::lora_delta::{LoraKernels, LoraPair};
16use crate::lora::ExpertLoraLayer;
17
18impl Qwen3SsmLayer {
19    /// Install this GDN layer's MoE router + routed-expert LoRA onto its
20    /// `FfnComponent::Moe`. Hard-rejects (never silently drops) when the layer's
21    /// FFN is not MoE — an expert/router delta on a dense-FFN GDN layer is a
22    /// loader/adapter mismatch.
23    pub fn set_moe_lora_weights(
24        &mut self,
25        router: Option<LoraPair>,
26        experts: ExpertLoraLayer,
27        kernels: LoraKernels,
28        gpu: &dyn GpuBackend,
29    ) -> Result<()> {
30        if let FfnComponent::Moe(m) = &mut self.ffn {
31            return m.set_lora_weights(router, experts, kernels, gpu);
32        }
33        anyhow::bail!(
34            "LoRA: router/expert deltas installed on a linear-attention layer whose \
35             FFN is not MoE (loader/adapter mismatch)"
36        )
37    }
38}
39
40impl Qwen3SsmLayer {
41    /// Install this linear-attention layer's DENSE-FFN LoRA onto its
42    /// `FfnComponent::Dense`.
43    ///
44    /// The mirror of `set_moe_lora_weights` for dense-FFN hybrids. A
45    /// linear-attention layer carries no attention projections, but on
46    /// Qwen3.8-27B it does carry the SwiGLU FFN — all 64 layers do, only 16 of
47    /// which are full attention — and real adapters for that architecture ship
48    /// gate/up/down for every one of them. Rejecting those rejected three
49    /// quarters of the adapter, and the old message could only suggest
50    /// retraining with `layers_to_transform`.
51    ///
52    /// The component is the same `DenseFfnLayer` the full-attention layers
53    /// hold, so the delta path, its pinned dispatch arms and its refusals are
54    /// identical here — this only hands it the weights.
55    ///
56    /// Hard-rejects a non-dense FFN rather than dropping the pairs: a dense
57    /// delta arriving at a MoE or absent FFN is a loader/adapter mismatch, and
58    /// silently ignoring it would be an adapter that reports success and does
59    /// nothing — the exact failure this whole change removes.
60    pub fn set_ffn_lora_weights(
61        &mut self,
62        ffn: crate::layers::ops::lora_delta::LoraFfnWeights,
63    ) -> Result<()> {
64        match &mut self.ffn {
65            FfnComponent::Dense(d) => d.set_lora_weights(ffn),
66            FfnComponent::Moe(_) => anyhow::bail!(
67                "LoRA: dense-FFN delta on a linear-attention layer whose FFN is MoE — \
68                 routed-expert deltas belong on set_moe_lora_weights"
69            ),
70            FfnComponent::None => {
71                anyhow::bail!("LoRA: dense-FFN delta on a linear-attention layer that has no FFN")
72            }
73        }
74    }
75}
76
77impl Qwen3SsmLayer {
78    /// Install this layer's GDN `out_proj` delta.
79    ///
80    /// Separate from `set_ffn_lora_weights`: that one targets the block's FFN,
81    /// this one the linear-attention block's own output projection.
82    pub fn set_out_proj_lora(&mut self, pair: LoraPair, kernels: LoraKernels) {
83        self.lora_out_proj = Some((pair, kernels));
84    }
85
86    /// `out += scale * (normed_out @ A^T) @ B^T`.
87    ///
88    /// No-op without an adapter, so the base path stays byte-identical.
89    ///
90    /// The caller must invoke this AFTER any TP reduce: `out` is a partial
91    /// row-parallel product until then, and a delta added to a partial would
92    /// be summed once per rank.
93    pub(super) fn apply_lora_out_proj(
94        &self,
95        ctx: &crate::layers::ForwardContext,
96        normed_out: spark_runtime::gpu::DevicePtr,
97        out: spark_runtime::gpu::DevicePtr,
98        m: u32,
99        stream: u64,
100    ) -> Result<()> {
101        if crate::layers::ops::lora_delta::lora_no_ffn() {
102            return Ok(());
103        }
104        let Some((ref pair, ref kernels)) = self.lora_out_proj else {
105            return Ok(());
106        };
107        crate::layers::ops::lora_delta::apply_lora_delta(
108            ctx.gpu,
109            kernels,
110            pair,
111            normed_out,
112            out,
113            m,
114            ctx.buffers.lora_xa(),
115            ctx.buffers.lora_delta(),
116            stream,
117        )
118    }
119}
spark_model/layers/qwen3_ssm/lora.rs

spark_model/layers/qwen3_ssm/
lora.rs
