ferrum_quantization/
gptq.rs

1//! GPTQ linear projection — thin factory wrapper.
2//!
3//! Phase 3e/2: the actual kernel dispatch lives inside the boxed
4//! `Linear<B>` returned by `B::load_gptq` (`CudaMarlinLinear` on
5//! CUDA, `CpuGptqLinear` on CPU). This module just re-exposes the
6//! historical constructor names so callers don't have to switch.
7
8use ferrum_kernels::backend::{Backend, BackendQuantMarlin};
9use ferrum_kernels::Linear;
10use ferrum_kernels::LinearMetadata;
11#[cfg(feature = "cuda")]
12use ferrum_kernels::LinearProjectionRole;
13use ferrum_types::Result;
14use std::sync::Arc;
15
16#[cfg(feature = "cuda")]
17fn cuda_marlin_profile_label(metadata: LinearMetadata) -> Option<&'static str> {
18    match metadata.role? {
19        LinearProjectionRole::Qkv
20        | LinearProjectionRole::Query
21        | LinearProjectionRole::Key
22        | LinearProjectionRole::Value
23        | LinearProjectionRole::GdnQkv
24        | LinearProjectionRole::GdnZ
25        | LinearProjectionRole::GdnQkvz => Some("gptq.linear.qkv_proj"),
26        LinearProjectionRole::Output => Some("gptq.linear.o_proj"),
27        LinearProjectionRole::GateUp | LinearProjectionRole::Gate | LinearProjectionRole::Up => {
28            Some("gptq.linear.gate_up_proj")
29        }
30        LinearProjectionRole::Down => Some("gptq.linear.down_proj"),
31        LinearProjectionRole::LmHead => Some("gptq.linear.lm_head"),
32        LinearProjectionRole::GdnB | LinearProjectionRole::GdnA | LinearProjectionRole::GdnBa => {
33            Some("gptq.linear.other_proj")
34        }
35    }
36}
37
38/// GPTQ-format Linear projection, polymorphic over backend.
39///
40/// Holds a boxed backend-specific `Linear<B>` produced by `B::load_gptq`.
41/// `forward()` delegates straight through.
42pub struct GptqLinear<B: Backend + BackendQuantMarlin> {
43    inner: Box<dyn Linear<B> + Send + Sync>,
44    metadata: LinearMetadata,
45}
46
47impl<B: Backend + BackendQuantMarlin> GptqLinear<B> {
48    /// Build from raw host-side GPTQ tensors. The Backend repacks into
49    /// its preferred format once (Marlin tiles on CUDA, dequant on CPU)
50    /// and returns a boxed Linear; inference uses the boxed forward.
51    ///
52    /// `qweight`: `[k/8, n]` i32 (packed int4)
53    /// `scales`:  `[k/group_size, n]` f32 (converted from f16 by caller)
54    /// `qzeros`:  `[k/group_size, n/8]` i32
55    /// `g_idx`:   `[k]` i32 — optional, only used for desc_act=true
56    /// `bias`:    `[n]` f32 — optional fused bias (Qwen2.5 attention)
57    #[allow(clippy::too_many_arguments)]
58    pub fn from_raw(
59        qweight: &[i32],
60        scales: &[f32],
61        qzeros: &[i32],
62        g_idx: Option<&[i32]>,
63        bias: Option<&[f32]>,
64        bits: u32,
65        group_size: usize,
66        in_features: usize,
67        out_features: usize,
68    ) -> Result<Self> {
69        Self::from_raw_with_metadata(
70            qweight,
71            scales,
72            qzeros,
73            g_idx,
74            bias,
75            bits,
76            group_size,
77            in_features,
78            out_features,
79            LinearMetadata::default(),
80        )
81    }
82
83    #[allow(clippy::too_many_arguments)]
84    pub fn from_raw_with_metadata(
85        qweight: &[i32],
86        scales: &[f32],
87        qzeros: &[i32],
88        g_idx: Option<&[i32]>,
89        bias: Option<&[f32]>,
90        bits: u32,
91        group_size: usize,
92        in_features: usize,
93        out_features: usize,
94        metadata: LinearMetadata,
95    ) -> Result<Self> {
96        let inner = B::load_gptq(
97            qweight,
98            scales,
99            qzeros,
100            g_idx,
101            bias,
102            bits,
103            group_size,
104            in_features,
105            out_features,
106        )?;
107        Ok(Self { inner, metadata })
108    }
109}
110
111impl<B: Backend + BackendQuantMarlin> Linear<B> for GptqLinear<B> {
112    fn in_features(&self) -> usize {
113        self.inner.in_features()
114    }
115
116    fn out_features(&self) -> usize {
117        self.inner.out_features()
118    }
119
120    fn metadata(&self) -> LinearMetadata {
121        if self.metadata.is_empty() {
122            self.inner.metadata()
123        } else {
124            self.metadata
125        }
126    }
127
128    #[cfg(feature = "cuda")]
129    fn cuda_marlin_touch_ref(
130        &self,
131    ) -> Option<ferrum_kernels::quant_linear::cuda_marlin::CudaMarlinTouchRef<'_>> {
132        self.inner.cuda_marlin_touch_ref()
133    }
134
135    fn forward(&self, ctx: &mut B::Context, input: &B::Buffer, out: &mut B::Buffer, m: usize) {
136        #[cfg(feature = "cuda")]
137        let _cuda_alloc_label = if ferrum_kernels::backend::cuda::marlin::profile_marlin() {
138            cuda_marlin_profile_label(self.metadata())
139                .map(ferrum_kernels::backend::cuda::push_alloc_label)
140        } else {
141            None
142        };
143        self.inner.forward(ctx, input, out, m);
144    }
145}
146
147/// View into a single column-slice of a shared stacked GPTQ store.
148///
149/// Phase 3e/2: backed by a `Box<dyn Linear<B>>` produced by
150/// `B::make_stacked_expert_linear` (CUDA: `CudaMarlinStackedExpertLinear`;
151/// CPU: `CpuGptqLinear` over a sliced row range). The store itself is
152/// `Arc<B::GptqStore>` so cloning a view is cheap; dropping all views
153/// drops the underlying store.
154pub struct StackedExpertLinear<B: Backend + BackendQuantMarlin> {
155    inner: Box<dyn Linear<B> + Send + Sync>,
156    /// Kept for in_features() reporting.
157    k: usize,
158    /// Kept for out_features() reporting.
159    expert_n: usize,
160}
161
162impl<B: Backend + BackendQuantMarlin> StackedExpertLinear<B> {
163    /// Phase C step 4b: takes the trait-object MarlinExpertStack
164    /// directly (was `Arc<B::GptqStore>` + `B::make_stacked_expert_linear`).
165    pub fn new(
166        stack: Arc<dyn ferrum_kernels::MarlinExpertStack<B>>,
167        expert_offset: usize,
168        expert_n: usize,
169    ) -> Result<Self> {
170        let k = stack.k();
171        let inner = stack.make_expert_linear(expert_offset, expert_n, None)?;
172        Ok(Self { inner, k, expert_n })
173    }
174
175    pub fn new_with_bias(
176        stack: Arc<dyn ferrum_kernels::MarlinExpertStack<B>>,
177        expert_offset: usize,
178        expert_n: usize,
179        bias: &[f32],
180    ) -> Result<Self> {
181        let k = stack.k();
182        let inner = stack.make_expert_linear(expert_offset, expert_n, Some(bias))?;
183        Ok(Self { inner, k, expert_n })
184    }
185}
186
187impl<B: Backend + BackendQuantMarlin> Linear<B> for StackedExpertLinear<B> {
188    fn in_features(&self) -> usize {
189        self.k
190    }
191
192    fn out_features(&self) -> usize {
193        self.expert_n
194    }
195
196    fn forward(&self, ctx: &mut B::Context, input: &B::Buffer, out: &mut B::Buffer, m: usize) {
197        self.inner.forward(ctx, input, out, m);
198    }
199}
ferrum_quantization/gptq.rs

ferrum_quantization/
gptq.rs
