qualia_core_db/lora/
webgpu_lora.rs

1//! GPU-accelerated LoRA delta application via wgpu compute shader.
2//!
3//! Computes `output += B @ (A @ x) * scaling` entirely on the GPU,
4//! avoiding a CPU round-trip for the hidden-state delta.
5//!
6//! # Bind-group layout (group 0)
7//!
8//! | binding | type            | content                           |
9//! |---------|-----------------|-----------------------------------|
10//! | 0       | storage rw      | `output` array<f32> (n_out)       |
11//! | 1       | storage r       | `input`  array<f32> (n_in)        |
12//! | 2       | storage r       | `lora_a` array<f32> (rank × n_in) |
13//! | 3       | storage r       | `lora_b` array<f32> (n_out × rank)|
14//! | 4       | uniform         | `LoraGpuParams`                   |
15
16use super::adapter_manager::{LoRAAdapter, LoRAError};
17
18// ─── GPU parameter block ─────────────────────────────────────────────────────
19
20#[repr(C)]
21#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
22pub struct LoraGpuParams {
23    pub n_in: u32,
24    pub n_out: u32,
25    pub rank: u32,
26    pub scaling: f32,
27}
28
29// ─── LoRAGpuApplicator ────────────────────────────────────────────────────────
30
31/// Holds the compiled wgpu pipeline for LoRA delta computation.
32///
33/// Constructed once and reused across inference calls.
34pub struct LoRAGpuApplicator {
35    device: wgpu::Device,
36    queue: wgpu::Queue,
37    pipeline: wgpu::ComputePipeline,
38    bgl: wgpu::BindGroupLayout,
39}
40
41impl LoRAGpuApplicator {
42    /// Initialise the wgpu device and compile the LoRA compute shader.
43    ///
44    /// Synchronous wrapper using `Box::leak` to give a `'static` tokio runtime —
45    /// the same pattern used in `QTensorEngine::new()` in `gguf_bridge.rs`.
46    pub fn new() -> Result<Self, LoRAError> {
47        let rt = Box::leak(Box::new(
48            tokio::runtime::Runtime::new()
49                .map_err(|e| LoRAError::Io(format!("tokio runtime: {e}")))?,
50        ));
51        rt.block_on(Self::new_async())
52    }
53
54    async fn new_async() -> Result<Self, LoRAError> {
55        let instance = wgpu::Instance::default();
56        let adapter = instance
57            .request_adapter(&wgpu::RequestAdapterOptions {
58                power_preference: wgpu::PowerPreference::HighPerformance,
59                ..Default::default()
60            })
61            .await
62            .map_err(|e| LoRAError::Io(format!("no wgpu adapter: {e}")))?;
63
64        // This kernel uses only baseline WebGPU storage/uniform capabilities, so it
65        // deliberately requests no optional or experimental wgpu 30 features.
66        let (device, queue) = adapter
67            .request_device(&wgpu::DeviceDescriptor::default())
68            .await
69            .map_err(|e| LoRAError::Io(format!("wgpu device: {e}")))?;
70
71        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
72            label: Some("lora_apply"),
73            source: wgpu::ShaderSource::Wgsl(include_str!("../shaders/lora_apply.wgsl").into()),
74        });
75
76        let bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
77            label: Some("lora-bgl"),
78            entries: &[
79                storage_rw_entry(0), // output
80                storage_r_entry(1),  // input
81                storage_r_entry(2),  // lora_a
82                storage_r_entry(3),  // lora_b
83                uniform_entry(4),    // params
84            ],
85        });
86
87        let pl = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
88            label: Some("lora-pl"),
89            bind_group_layouts: &[Some(&bgl)],
90            immediate_size: 0,
91        });
92
93        let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
94            label: Some("lora-pipeline"),
95            layout: Some(&pl),
96            module: &shader,
97            entry_point: Some("apply_lora"),
98            compilation_options: Default::default(),
99            cache: None,
100        });
101
102        Ok(Self {
103            device,
104            queue,
105            pipeline,
106            bgl,
107        })
108    }
109
110    /// Apply the LoRA delta to `output` on the GPU.
111    ///
112    /// `input`  — hidden state, length `n_in`
113    /// `output` — receives `+= B @ (A @ input) * scaling`, length `n_out`
114    ///
115    /// On any GPU error returns `Err`; caller should fall back to `LoRAAdapter::apply_cpu`.
116    pub fn apply(
117        &self,
118        adapter: &LoRAAdapter,
119        input: &[f32],
120        output: &mut [f32],
121    ) -> Result<(), LoRAError> {
122        let n_in = adapter.meta.n_in;
123        let n_out = adapter.meta.n_out;
124        let rank = adapter.meta.rank as usize;
125        let scaling = adapter.meta.scaling();
126
127        if input.len() != n_in || output.len() != n_out {
128            return Err(LoRAError::InferenceDimMismatch {
129                input_len: input.len(),
130                lora_n_in: n_in,
131            });
132        }
133
134        let params = LoraGpuParams {
135            n_in: n_in as u32,
136            n_out: n_out as u32,
137            rank: rank as u32,
138            scaling,
139        };
140
141        let dev = &self.device;
142
143        // ── Upload buffers ────────────────────────────────────────────────────
144        let out_buf = buf_rw(dev, bytemuck::cast_slice(output));
145        let in_buf = buf_r(dev, bytemuck::cast_slice(input));
146        let a_buf = buf_r(dev, bytemuck::cast_slice(&adapter.lora_a.data));
147        let b_buf = buf_r(dev, bytemuck::cast_slice(&adapter.lora_b.data));
148        let params_buf = buf_uniform(dev, bytemuck::bytes_of(&params));
149
150        let out_staging = dev.create_buffer(&wgpu::BufferDescriptor {
151            label: Some("lora-out-staging"),
152            size: (n_out * 4) as u64,
153            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
154            mapped_at_creation: false,
155        });
156
157        // ── Bind group ────────────────────────────────────────────────────────
158        let bg = dev.create_bind_group(&wgpu::BindGroupDescriptor {
159            label: Some("lora-bg"),
160            layout: &self.bgl,
161            entries: &[
162                wgpu::BindGroupEntry {
163                    binding: 0,
164                    resource: out_buf.as_entire_binding(),
165                },
166                wgpu::BindGroupEntry {
167                    binding: 1,
168                    resource: in_buf.as_entire_binding(),
169                },
170                wgpu::BindGroupEntry {
171                    binding: 2,
172                    resource: a_buf.as_entire_binding(),
173                },
174                wgpu::BindGroupEntry {
175                    binding: 3,
176                    resource: b_buf.as_entire_binding(),
177                },
178                wgpu::BindGroupEntry {
179                    binding: 4,
180                    resource: params_buf.as_entire_binding(),
181                },
182            ],
183        });
184
185        // ── Dispatch ─────────────────────────────────────────────────────────
186        let mut enc = dev.create_command_encoder(&wgpu::CommandEncoderDescriptor {
187            label: Some("lora-enc"),
188        });
189
190        {
191            let mut cpass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor {
192                label: Some("lora-pass"),
193                timestamp_writes: None,
194            });
195            cpass.set_pipeline(&self.pipeline);
196            cpass.set_bind_group(0, &bg, &[]);
197            cpass.dispatch_workgroups((n_out as u32 + 63) / 64, 1, 1);
198        }
199
200        enc.copy_buffer_to_buffer(&out_buf, 0, &out_staging, 0, (n_out * 4) as u64);
201        self.queue.submit(Some(enc.finish()));
202
203        // ── Readback ─────────────────────────────────────────────────────────
204        let slice = out_staging.slice(..);
205        let (tx, rx) = std::sync::mpsc::channel();
206        slice.map_async(wgpu::MapMode::Read, move |r| {
207            let _ = tx.send(r);
208        });
209        let _ = dev.poll(wgpu::PollType::wait_indefinitely());
210        rx.recv()
211            .map_err(|_| LoRAError::Io("GPU readback channel closed".into()))?
212            .map_err(|e| LoRAError::Io(format!("map_async: {e:?}")))?;
213
214        let mapped = slice
215            .get_mapped_range()
216            .expect("wgpu buffer map_range failed");
217        let result: &[f32] = bytemuck::cast_slice(&mapped);
218        output.copy_from_slice(result);
219        Ok(())
220    }
221}
222
223// ─── wgpu buffer helpers ─────────────────────────────────────────────────────
224
225fn buf_rw(dev: &wgpu::Device, data: &[u8]) -> wgpu::Buffer {
226    use wgpu::util::DeviceExt;
227    dev.create_buffer_init(&wgpu::util::BufferInitDescriptor {
228        label: Some("lora-rw"),
229        contents: data,
230        usage: wgpu::BufferUsages::STORAGE
231            | wgpu::BufferUsages::COPY_SRC
232            | wgpu::BufferUsages::COPY_DST,
233    })
234}
235
236fn buf_r(dev: &wgpu::Device, data: &[u8]) -> wgpu::Buffer {
237    use wgpu::util::DeviceExt;
238    dev.create_buffer_init(&wgpu::util::BufferInitDescriptor {
239        label: Some("lora-r"),
240        contents: data,
241        usage: wgpu::BufferUsages::STORAGE,
242    })
243}
244
245fn buf_uniform(dev: &wgpu::Device, data: &[u8]) -> wgpu::Buffer {
246    use wgpu::util::DeviceExt;
247    dev.create_buffer_init(&wgpu::util::BufferInitDescriptor {
248        label: Some("lora-uniform"),
249        contents: data,
250        usage: wgpu::BufferUsages::UNIFORM,
251    })
252}
253
254fn storage_rw_entry(binding: u32) -> wgpu::BindGroupLayoutEntry {
255    wgpu::BindGroupLayoutEntry {
256        binding,
257        visibility: wgpu::ShaderStages::COMPUTE,
258        ty: wgpu::BindingType::Buffer {
259            ty: wgpu::BufferBindingType::Storage { read_only: false },
260            has_dynamic_offset: false,
261            min_binding_size: None,
262        },
263        count: None,
264    }
265}
266
267fn storage_r_entry(binding: u32) -> wgpu::BindGroupLayoutEntry {
268    wgpu::BindGroupLayoutEntry {
269        binding,
270        visibility: wgpu::ShaderStages::COMPUTE,
271        ty: wgpu::BindingType::Buffer {
272            ty: wgpu::BufferBindingType::Storage { read_only: true },
273            has_dynamic_offset: false,
274            min_binding_size: None,
275        },
276        count: None,
277    }
278}
279
280fn uniform_entry(binding: u32) -> wgpu::BindGroupLayoutEntry {
281    wgpu::BindGroupLayoutEntry {
282        binding,
283        visibility: wgpu::ShaderStages::COMPUTE,
284        ty: wgpu::BindingType::Buffer {
285            ty: wgpu::BufferBindingType::Uniform,
286            has_dynamic_offset: false,
287            min_binding_size: None,
288        },
289        count: None,
290    }
291}
qualia_core_db/lora/webgpu_lora.rs

qualia_core_db/lora/
webgpu_lora.rs
