ferrum_kernels/
linear.rs

1//! `Linear<B>` trait — weight-bearing projection abstraction.
2//!
3//! Lives in ferrum-kernels alongside `Backend` because:
4//!   1. `Backend::layer_forward_fused` and other "standard transformer layer"
5//!      helpers want to accept `&dyn Linear<Self>` as their projection
6//!      parameter, so the trait must be visible here.
7//!   2. Model code in `ferrum-models` depends on both ferrum-kernels and
8//!      ferrum-quantization, so keeping the trait in kernels avoids any
9//!      circular dependency between kernels and quantization.
10//!
11//! Concrete implementations (DenseLinear, GptqLinear, AwqLinear, GgufLinear)
12//! live in `ferrum-quantization`, which depends on `ferrum-kernels` for this
13//! trait and for the `Backend` it parameterises over.
14
15use crate::backend::Backend;
16
17/// Stable projection role metadata derived from model weight names.
18///
19/// This is intentionally small and backend-neutral. It lets product code
20/// choose a typed optimization path without depending on profiling labels or
21/// hidden environment variables.
22#[derive(Clone, Copy, Debug, Eq, PartialEq)]
23pub enum LinearProjectionRole {
24    Qkv,
25    Query,
26    Key,
27    Value,
28    GdnQkv,
29    GdnZ,
30    GdnQkvz,
31    GdnB,
32    GdnA,
33    GdnBa,
34    Output,
35    GateUp,
36    Gate,
37    Up,
38    Down,
39    LmHead,
40}
41
42/// Optional metadata carried by a loaded linear projection.
43#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
44pub struct LinearMetadata {
45    pub layer_index: Option<usize>,
46    pub role: Option<LinearProjectionRole>,
47}
48
49impl LinearMetadata {
50    pub const fn new(layer_index: Option<usize>, role: Option<LinearProjectionRole>) -> Self {
51        Self { layer_index, role }
52    }
53
54    pub const fn is_empty(self) -> bool {
55        self.layer_index.is_none() && self.role.is_none()
56    }
57
58    pub fn from_name(name: &str) -> Self {
59        let base = strip_tensor_suffix(name);
60        Self {
61            layer_index: parse_layer_index(base),
62            role: parse_projection_role(base),
63        }
64    }
65
66    pub fn from_fused_names<'a>(names: impl IntoIterator<Item = &'a str>) -> Self {
67        let mut layer_index = None;
68        let mut roles = Vec::new();
69
70        for name in names {
71            let metadata = Self::from_name(name);
72            if layer_index.is_none() {
73                layer_index = metadata.layer_index;
74            }
75            if let Some(role) = metadata.role {
76                roles.push(role);
77            }
78        }
79
80        let role = match roles.as_slice() {
81            [LinearProjectionRole::Query, LinearProjectionRole::Key, LinearProjectionRole::Value] => {
82                Some(LinearProjectionRole::Qkv)
83            }
84            [LinearProjectionRole::GdnQkv, LinearProjectionRole::GdnZ] => {
85                Some(LinearProjectionRole::GdnQkvz)
86            }
87            [LinearProjectionRole::GdnB, LinearProjectionRole::GdnA] => {
88                Some(LinearProjectionRole::GdnBa)
89            }
90            [LinearProjectionRole::Gate, LinearProjectionRole::Up] => {
91                Some(LinearProjectionRole::GateUp)
92            }
93            [single] => Some(*single),
94            _ => None,
95        };
96
97        Self { layer_index, role }
98    }
99}
100
101/// A weight-bearing linear projection.
102///
103/// `forward` computes `out[m, out_features] = input[m, in_features] @ W^T`.
104/// Implementations are responsible for calling the right backend kernel
105/// (`B::gemm` for dense, `B::gemm_quant` for quantized variants).
106pub trait Linear<B: Backend>: Send + Sync {
107    fn in_features(&self) -> usize;
108    fn out_features(&self) -> usize;
109
110    fn metadata(&self) -> LinearMetadata {
111        LinearMetadata::default()
112    }
113
114    #[cfg(feature = "cuda")]
115    fn cuda_marlin_touch_ref(
116        &self,
117    ) -> Option<crate::quant_linear::cuda_marlin::CudaMarlinTouchRef<'_>> {
118        None
119    }
120
121    /// Append GEMM work onto `ctx`. Caller flushes the context when results
122    /// must be materialised.
123    fn forward(&self, ctx: &mut B::Context, input: &B::Buffer, out: &mut B::Buffer, m: usize);
124}
125
126fn strip_tensor_suffix(name: &str) -> &str {
127    for suffix in [
128        ".weight", ".qweight", ".scales", ".qzeros", ".g_idx", ".bias",
129    ] {
130        if let Some(stripped) = name.strip_suffix(suffix) {
131            return stripped;
132        }
133    }
134    name
135}
136
137fn parse_layer_index(name: &str) -> Option<usize> {
138    let mut prev_was_layers = false;
139    for part in name.split('.') {
140        if prev_was_layers {
141            return part.parse::<usize>().ok();
142        }
143        prev_was_layers = part == "layers";
144    }
145    None
146}
147
148fn parse_projection_role(name: &str) -> Option<LinearProjectionRole> {
149    let tail = name.rsplit('.').next().unwrap_or(name);
150    match tail {
151        "qkv_proj" => Some(LinearProjectionRole::Qkv),
152        "q_proj" => Some(LinearProjectionRole::Query),
153        "k_proj" => Some(LinearProjectionRole::Key),
154        "v_proj" => Some(LinearProjectionRole::Value),
155        "in_proj_qkv" => Some(LinearProjectionRole::GdnQkv),
156        "in_proj_z" => Some(LinearProjectionRole::GdnZ),
157        "in_proj_qkvz" => Some(LinearProjectionRole::GdnQkvz),
158        "in_proj_b" => Some(LinearProjectionRole::GdnB),
159        "in_proj_a" => Some(LinearProjectionRole::GdnA),
160        "in_proj_ba" => Some(LinearProjectionRole::GdnBa),
161        "o_proj" => Some(LinearProjectionRole::Output),
162        "gate_up_proj" => Some(LinearProjectionRole::GateUp),
163        "gate_proj" => Some(LinearProjectionRole::Gate),
164        "up_proj" => Some(LinearProjectionRole::Up),
165        "down_proj" => Some(LinearProjectionRole::Down),
166        "lm_head" => Some(LinearProjectionRole::LmHead),
167        _ => None,
168    }
169}
170
171#[cfg(test)]
172mod tests {
173    use super::{LinearMetadata, LinearProjectionRole};
174
175    #[test]
176    fn metadata_parses_llama_layer_projection_roles() {
177        assert_eq!(
178            LinearMetadata::from_name("model.layers.17.mlp.down_proj.qweight"),
179            LinearMetadata::new(Some(17), Some(LinearProjectionRole::Down))
180        );
181        assert_eq!(
182            LinearMetadata::from_name("language_model.model.layers.3.self_attn.o_proj.weight"),
183            LinearMetadata::new(Some(3), Some(LinearProjectionRole::Output))
184        );
185        assert_eq!(
186            LinearMetadata::from_name("lm_head.weight"),
187            LinearMetadata::new(None, Some(LinearProjectionRole::LmHead))
188        );
189    }
190
191    #[test]
192    fn metadata_parses_fused_projection_roles() {
193        assert_eq!(
194            LinearMetadata::from_fused_names([
195                "model.layers.2.self_attn.q_proj",
196                "model.layers.2.self_attn.k_proj",
197                "model.layers.2.self_attn.v_proj",
198            ]),
199            LinearMetadata::new(Some(2), Some(LinearProjectionRole::Qkv))
200        );
201        assert_eq!(
202            LinearMetadata::from_fused_names([
203                "model.layers.9.mlp.gate_proj.qweight",
204                "model.layers.9.mlp.up_proj.qweight",
205            ]),
206            LinearMetadata::new(Some(9), Some(LinearProjectionRole::GateUp))
207        );
208        assert_eq!(
209            LinearMetadata::from_fused_names([
210                "model.layers.2.linear_attn.in_proj_qkv.qweight",
211                "model.layers.2.linear_attn.in_proj_z.qweight",
212            ]),
213            LinearMetadata::new(Some(2), Some(LinearProjectionRole::GdnQkvz))
214        );
215        assert_eq!(
216            LinearMetadata::from_fused_names([
217                "model.layers.2.linear_attn.in_proj_b.qweight",
218                "model.layers.2.linear_attn.in_proj_a.qweight",
219            ]),
220            LinearMetadata::new(Some(2), Some(LinearProjectionRole::GdnBa))
221        );
222    }
223}
ferrum_kernels/linear.rs

ferrum_kernels/
linear.rs
