qualia_core_db/inference/
ternary_gpu.rs

1//! Task #12 / STELLAR §A — **native GPU dispatch of the ternary GEMM** + on-device parity.
2//!
3//! Runs `shaders/ternary_gemm.wgsl` (`ternary::TERNARY_GEMM_WGSL`) on a real wgpu device. This is
4//! the piece the FFN inference loop calls to execute a ternary-packed weight on the GPU; the
5//! `#[test]` below verifies it **on silicon** against the byte-exact CPU oracle
6//! `ternary::ternary_gemm_cpu` (it skips cleanly when no adapter is present, e.g. headless CI).
7//!
8//! Native only — the wasm WebGPU path reuses the same WGSL through `gguf_bridge`'s pipeline set
9//! when the kernel is spliced into the layer loop (the remaining integration step).
10
11use crate::ternary::{rebake_ternary_blob_to_2bit, TERNARY_GEMM_2BIT_WGSL, TERNARY_GEMM_WGSL};
12use std::collections::HashMap;
13
14/// 32-byte `TernaryParams` uniform matching `ternary_gemm.wgsl` (n_in, n_out, n_batch,
15/// in_row_stride, out_row_stride, scale, + 2 pad words).
16fn ternary_params_bytes(n_in: u32, n_out: u32, n_batch: u32, scale: f32) -> [u8; 32] {
17    let mut b = [0u8; 32];
18    b[0..4].copy_from_slice(&n_in.to_le_bytes());
19    b[4..8].copy_from_slice(&n_out.to_le_bytes());
20    b[8..12].copy_from_slice(&n_batch.to_le_bytes());
21    // in_row_stride / out_row_stride = 0 → dense (the shader falls back to n_in / n_out)
22    b[20..24].copy_from_slice(&scale.to_le_bytes());
23    b
24}
25
26/// Execute the **base-3** ternary GEMM on the GPU (`ternary_gemm.wgsl`).
27pub fn ternary_gemm_gpu(
28    device: &wgpu::Device,
29    queue: &wgpu::Queue,
30    activations: &[f32],
31    packed: &[u8],
32    scale: f32,
33    n_in: usize,
34    n_out: usize,
35    n_batch: usize,
36) -> Vec<f32> {
37    run_gemm(
38        device,
39        queue,
40        TERNARY_GEMM_WGSL,
41        activations,
42        packed,
43        scale,
44        n_in,
45        n_out,
46        n_batch,
47    )
48}
49
50/// Execute the **2-bit branchless** ternary GEMM on the GPU (`ternary_gemm_2bit.wgsl`). `packed`
51/// must be 2-bit packed (`ternary::pack_trits_2bit`).
52pub fn ternary_gemm_gpu_2bit(
53    device: &wgpu::Device,
54    queue: &wgpu::Queue,
55    activations: &[f32],
56    packed: &[u8],
57    scale: f32,
58    n_in: usize,
59    n_out: usize,
60    n_batch: usize,
61) -> Vec<f32> {
62    run_gemm(
63        device,
64        queue,
65        TERNARY_GEMM_2BIT_WGSL,
66        activations,
67        packed,
68        scale,
69        n_in,
70        n_out,
71        n_batch,
72    )
73}
74
75/// Shared dispatch for both ternary GEMM kernels (identical bindings/params; only the WGSL differs).
76/// Returns the `n_batch × n_out` output. Strides default to dense. Blocking (native readback).
77#[allow(clippy::too_many_arguments)]
78fn run_gemm(
79    device: &wgpu::Device,
80    queue: &wgpu::Queue,
81    wgsl: &str,
82    activations: &[f32],
83    packed: &[u8],
84    scale: f32,
85    n_in: usize,
86    n_out: usize,
87    n_batch: usize,
88) -> Vec<f32> {
89    let n_batch = n_batch.max(1);
90    let out_elems = n_batch * n_out;
91
92    let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
93        label: Some("ternary_gemm"),
94        source: wgpu::ShaderSource::Wgsl(wgsl.into()),
95    });
96    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
97        label: Some("ternary_gemm_pipeline"),
98        layout: None, // auto bind-group layout (native)
99        module: &shader,
100        entry_point: Some("ternary_gemm"),
101        compilation_options: Default::default(),
102        cache: None,
103    });
104
105    let mk_buf = |label: &str, contents: &[u8], usage: wgpu::BufferUsages| {
106        let buf = device.create_buffer(&wgpu::BufferDescriptor {
107            label: Some(label),
108            size: contents.len().max(4) as u64,
109            usage,
110            mapped_at_creation: false,
111        });
112        if !contents.is_empty() {
113            queue.write_buffer(&buf, 0, contents);
114        }
115        buf
116    };
117
118    let act_buf = mk_buf(
119        "ternary_act",
120        bytemuck::cast_slice(activations),
121        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
122    );
123    // trit_words is array<u32> → pad the packed bytes to a 4-byte multiple.
124    let mut trits = packed.to_vec();
125    while trits.len() % 4 != 0 || trits.is_empty() {
126        trits.push(0);
127    }
128    let trit_buf = mk_buf(
129        "ternary_trits",
130        &trits,
131        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
132    );
133    let params = ternary_params_bytes(n_in as u32, n_out as u32, n_batch as u32, scale);
134    let param_buf = mk_buf(
135        "ternary_params",
136        &params,
137        wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
138    );
139    let out_buf = device.create_buffer(&wgpu::BufferDescriptor {
140        label: Some("ternary_out"),
141        size: (out_elems * 4).max(4) as u64,
142        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
143        mapped_at_creation: false,
144    });
145    let staging = device.create_buffer(&wgpu::BufferDescriptor {
146        label: Some("ternary_staging"),
147        size: (out_elems * 4).max(4) as u64,
148        usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
149        mapped_at_creation: false,
150    });
151
152    let bind = device.create_bind_group(&wgpu::BindGroupDescriptor {
153        label: Some("ternary_bind"),
154        layout: &pipeline.get_bind_group_layout(0),
155        entries: &[
156            wgpu::BindGroupEntry {
157                binding: 0,
158                resource: act_buf.as_entire_binding(),
159            },
160            wgpu::BindGroupEntry {
161                binding: 1,
162                resource: trit_buf.as_entire_binding(),
163            },
164            wgpu::BindGroupEntry {
165                binding: 2,
166                resource: param_buf.as_entire_binding(),
167            },
168            wgpu::BindGroupEntry {
169                binding: 3,
170                resource: out_buf.as_entire_binding(),
171            },
172        ],
173    });
174
175    let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
176        label: Some("ternary_enc"),
177    });
178    {
179        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
180            label: Some("ternary_pass"),
181            timestamp_writes: None,
182        });
183        pass.set_pipeline(&pipeline);
184        pass.set_bind_group(0, &bind, &[]);
185        let wg_x = (n_out as u32).div_ceil(64).max(1);
186        pass.dispatch_workgroups(wg_x, n_batch as u32, 1);
187    }
188    encoder.copy_buffer_to_buffer(&out_buf, 0, &staging, 0, (out_elems * 4).max(4) as u64);
189    queue.submit(Some(encoder.finish()));
190
191    let slice = staging.slice(..);
192    let (tx, rx) = std::sync::mpsc::channel();
193    slice.map_async(wgpu::MapMode::Read, move |r| {
194        let _ = tx.send(r);
195    });
196    let _ = device.poll(wgpu::PollType::wait_indefinitely());
197    rx.recv()
198        .expect("map channel")
199        .expect("map ternary staging");
200    let data = slice
201        .get_mapped_range()
202        .expect("wgpu buffer map_range failed");
203    let out: Vec<f32> = bytemuck::cast_slice(&data)[..out_elems].to_vec();
204    drop(data);
205    staging.unmap();
206    out
207}
208
209// ── A1b inc 2b: resident 2-bit ternary-FFN GEMM (persistent pipeline + resident weights) ─────────
210
211/// Storage-buffer binding offset alignment. wgpu's `min_storage_buffer_offset_alignment` is 256 on
212/// the engine's baselines (A2000/Vulkan, DX12, WebGPU downlevel), so every resident tensor starts on
213/// this boundary — a sub-range binding at a non-aligned offset is a validation error.
214const TERNARY_RESIDENT_ALIGN: u64 = 256;
215
216#[inline]
217fn align_up_u64(x: u64, align: u64) -> u64 {
218    (x + align - 1) / align * align
219}
220
221/// One resident ternary tensor's location + metadata inside [`TernaryFfnResident::weights`].
222#[derive(Clone, Copy, Debug)]
223struct ResidentTernaryTensor {
224    /// 256-aligned byte offset of this tensor's 2-bit packed weights in the resident buffer.
225    gpu_offset: u64,
226    /// 2-bit packed byte length (padded to a 4-byte multiple = the bound `size`).
227    packed_len: u32,
228    n_in: u32,
229    n_out: u32,
230    /// Per-tensor BitNet absmean scale (applied once per output element).
231    scale: f32,
232}
233
234/// A1b (STELLAR §A): the resident **2-bit branchless** ternary-FFN GEMM dispatcher — the perf core.
235///
236/// **The D7 fix.** The 2-bit kernel pipeline + bind-group layout are built **once**; every FFN
237/// ternary weight is rebaked base-3 → 2-bit (lossless; see [`rebake_ternary_blob_to_2bit`]) and
238/// uploaded **once** into a single resident VRAM buffer. (The prior 1.02× failure rebuilt the
239/// pipeline per call.) At decode each FFN GEMV binds its own weight sub-range — no re-upload —
240/// writes only the activation row + params, dispatches, and reads back `n_out` floats. Tensors are
241/// keyed by their P64 blob offset (`GgufTensorInfo::byte_offset`), the unique per-tensor handle.
242///
243/// On-disk the FFN weights stay base-3 (densest, 1.6 bit); 2-bit is the GPU-resident layout only
244/// (2.0 bit, shift/mask, divergence-free → 1.77× vs F16 on the A2000). Native-only; the wasm WebGPU
245/// ternary path reuses the same WGSL through the MC8 resident arena (a later step).
246pub struct TernaryFfnResident {
247    pipeline: wgpu::ComputePipeline,
248    weights: wgpu::Buffer,
249    act_buf: wgpu::Buffer,
250    out_buf: wgpu::Buffer,
251    params_buf: wgpu::Buffer,
252    staging: wgpu::Buffer,
253    map: HashMap<u64, ResidentTernaryTensor>,
254    max_in: u32,
255    max_out: u32,
256    resident_bytes: u64,
257}
258
259impl TernaryFfnResident {
260    /// Build the resident set from `(key, n_in, n_out, base3_blob)` tuples. Each base-3 blob
261    /// (`[scale f32][5-trits/byte]`) is rebaked to the 2-bit runtime layout here — load-time heap is
262    /// the sanctioned path; the hot loop stays zero-heap. `key` is the tensor's P64 blob offset.
263    /// Returns `None` if `tensors` is empty or any blob is malformed.
264    pub fn build(
265        device: &wgpu::Device,
266        queue: &wgpu::Queue,
267        tensors: &[(u64, usize, usize, &[u8])],
268    ) -> Option<Self> {
269        if tensors.is_empty() {
270            return None;
271        }
272        // 1. Rebake each base-3 blob → 2-bit; lay out 256-aligned within one resident buffer.
273        let mut map: HashMap<u64, ResidentTernaryTensor> = HashMap::with_capacity(tensors.len());
274        let mut uploads: Vec<(u64, Vec<u8>)> = Vec::with_capacity(tensors.len());
275        let mut cursor: u64 = 0;
276        let (mut max_in, mut max_out) = (0u32, 0u32);
277        for &(key, n_in, n_out, blob) in tensors {
278            let count = n_in.checked_mul(n_out)?;
279            if count == 0 {
280                return None;
281            }
282            let (scale, mut packed) = rebake_ternary_blob_to_2bit(blob, count);
283            if packed.is_empty() {
284                return None;
285            }
286            // trit_words is array<u32> → pad the packed bytes to a 4-byte multiple.
287            while packed.len() % 4 != 0 {
288                packed.push(0);
289            }
290            let gpu_offset = align_up_u64(cursor, TERNARY_RESIDENT_ALIGN);
291            let packed_len = packed.len() as u32;
292            map.insert(
293                key,
294                ResidentTernaryTensor {
295                    gpu_offset,
296                    packed_len,
297                    n_in: n_in as u32,
298                    n_out: n_out as u32,
299                    scale,
300                },
301            );
302            uploads.push((gpu_offset, packed));
303            cursor = gpu_offset + packed_len as u64;
304            max_in = max_in.max(n_in as u32);
305            max_out = max_out.max(n_out as u32);
306        }
307        let resident_bytes = cursor.max(4);
308
309        // 2. Persistent pipeline (auto layout) — built ONCE.
310        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
311            label: Some("ternary_ffn_2bit"),
312            source: wgpu::ShaderSource::Wgsl(TERNARY_GEMM_2BIT_WGSL.into()),
313        });
314        let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
315            label: Some("ternary_ffn_resident_pipeline"),
316            layout: None,
317            module: &shader,
318            entry_point: Some("ternary_gemm"),
319            compilation_options: Default::default(),
320            cache: None,
321        });
322
323        // 3. Resident weight buffer + reusable IO buffers.
324        let weights = device.create_buffer(&wgpu::BufferDescriptor {
325            label: Some("TernaryFfnResidentWeights"),
326            size: resident_bytes,
327            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
328            mapped_at_creation: false,
329        });
330        for (off, bytes) in &uploads {
331            queue.write_buffer(&weights, *off, bytes);
332        }
333        let act_buf = device.create_buffer(&wgpu::BufferDescriptor {
334            label: Some("TernaryFfnAct"),
335            size: (max_in as u64 * 4).max(4),
336            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
337            mapped_at_creation: false,
338        });
339        let out_buf = device.create_buffer(&wgpu::BufferDescriptor {
340            label: Some("TernaryFfnOut"),
341            size: (max_out as u64 * 4).max(4),
342            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
343            mapped_at_creation: false,
344        });
345        let params_buf = device.create_buffer(&wgpu::BufferDescriptor {
346            label: Some("TernaryFfnParams"),
347            size: 32,
348            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
349            mapped_at_creation: false,
350        });
351        let staging = device.create_buffer(&wgpu::BufferDescriptor {
352            label: Some("TernaryFfnStaging"),
353            size: (max_out as u64 * 4).max(4),
354            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
355            mapped_at_creation: false,
356        });
357
358        Some(Self {
359            pipeline,
360            weights,
361            act_buf,
362            out_buf,
363            params_buf,
364            staging,
365            map,
366            max_in,
367            max_out,
368            resident_bytes,
369        })
370    }
371
372    /// Whether a tensor with `key` is resident (and its shape, for the dispatch guard).
373    pub fn contains(&self, key: u64) -> bool {
374        self.map.contains_key(&key)
375    }
376
377    /// Total resident weight bytes (the VRAM footprint of the ternary FFN).
378    pub fn resident_bytes(&self) -> u64 {
379        self.resident_bytes
380    }
381
382    /// Number of resident ternary tensors.
383    pub fn len(&self) -> usize {
384        self.map.len()
385    }
386
387    pub fn is_empty(&self) -> bool {
388        self.map.is_empty()
389    }
390
391    /// Run `out[..n_out] = scale · W·act` for the resident tensor `key` (batch 1 = decode GEMV).
392    /// Returns `false` (caller falls back) if the key is absent, the shape mismatches the resident
393    /// metadata, the IO bounds are exceeded, or the GPU readback fails — fail-closed, never garbage.
394    pub fn gemv(
395        &self,
396        device: &wgpu::Device,
397        queue: &wgpu::Queue,
398        key: u64,
399        input: &[f32],
400        out: &mut [f32],
401        n_in: usize,
402        n_out: usize,
403    ) -> bool {
404        let t = match self.map.get(&key) {
405            Some(t) => *t,
406            None => return false,
407        };
408        if t.n_in as usize != n_in
409            || t.n_out as usize != n_out
410            || n_in > input.len()
411            || n_out > out.len()
412            || n_in as u32 > self.max_in
413            || n_out as u32 > self.max_out
414        {
415            return false;
416        }
417
418        queue.write_buffer(&self.act_buf, 0, bytemuck::cast_slice(&input[..n_in]));
419        let params = ternary_params_bytes(n_in as u32, n_out as u32, 1, t.scale);
420        queue.write_buffer(&self.params_buf, 0, &params);
421
422        let weight_binding = wgpu::BindingResource::Buffer(wgpu::BufferBinding {
423            buffer: &self.weights,
424            offset: t.gpu_offset,
425            size: std::num::NonZeroU64::new(t.packed_len as u64),
426        });
427        let bind = device.create_bind_group(&wgpu::BindGroupDescriptor {
428            label: Some("ternary_ffn_resident_bind"),
429            layout: &self.pipeline.get_bind_group_layout(0),
430            entries: &[
431                wgpu::BindGroupEntry {
432                    binding: 0,
433                    resource: self.act_buf.as_entire_binding(),
434                },
435                wgpu::BindGroupEntry {
436                    binding: 1,
437                    resource: weight_binding,
438                },
439                wgpu::BindGroupEntry {
440                    binding: 2,
441                    resource: self.params_buf.as_entire_binding(),
442                },
443                wgpu::BindGroupEntry {
444                    binding: 3,
445                    resource: self.out_buf.as_entire_binding(),
446                },
447            ],
448        });
449
450        let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
451            label: Some("ternary_ffn_enc"),
452        });
453        {
454            let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
455                label: Some("ternary_ffn_pass"),
456                timestamp_writes: None,
457            });
458            pass.set_pipeline(&self.pipeline);
459            pass.set_bind_group(0, &bind, &[]);
460            pass.dispatch_workgroups((n_out as u32).div_ceil(64).max(1), 1, 1);
461        }
462        let out_bytes = (n_out * 4) as wgpu::BufferAddress;
463        encoder.copy_buffer_to_buffer(&self.out_buf, 0, &self.staging, 0, out_bytes);
464        queue.submit(Some(encoder.finish()));
465
466        let slice = self.staging.slice(..out_bytes);
467        let (tx, rx) = std::sync::mpsc::channel();
468        slice.map_async(wgpu::MapMode::Read, move |r| {
469            let _ = tx.send(r);
470        });
471        let _ = device.poll(wgpu::PollType::wait_indefinitely());
472        match rx.recv() {
473            Ok(Ok(())) => {}
474            _ => return false,
475        }
476        let data = slice
477            .get_mapped_range()
478            .expect("wgpu buffer map_range failed");
479        out[..n_out].copy_from_slice(&bytemuck::cast_slice(&data)[..n_out]);
480        drop(data);
481        self.staging.unmap();
482        true
483    }
484}
485
486#[cfg(test)]
487mod tests {
488    use super::*;
489    use crate::ternary::{pack_trits, ternary_gemm_cpu};
490
491    fn try_gpu() -> Option<(wgpu::Device, wgpu::Queue)> {
492        let instance = wgpu::Instance::default();
493        let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
494            power_preference: wgpu::PowerPreference::HighPerformance,
495            ..Default::default()
496        }))
497        .ok()?;
498        pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor::default())).ok()
499    }
500
501    /// ON-DEVICE PARITY: `ternary_gemm.wgsl` on a real GPU == `ternary_gemm_cpu`. Skips cleanly with
502    /// no adapter (headless CI); runs for real on a machine with a GPU.
503    #[test]
504    fn ternary_gemm_gpu_matches_cpu_oracle() {
505        let Some((device, queue)) = try_gpu() else {
506            eprintln!("ternary_gemm_gpu: no wgpu adapter — skipping on-device parity");
507            return;
508        };
509
510        // (n_out=5 × n_in=7), 3 batch rows; deterministic trits + activations.
511        let (n_in, n_out, n_batch) = (7usize, 5usize, 3usize);
512        let scale = 0.37_f32;
513        let trits: Vec<i8> = (0..n_in * n_out).map(|k| (k % 3) as i8 - 1).collect();
514        let packed = pack_trits(&trits);
515        let act: Vec<f32> = (0..n_in * n_batch)
516            .map(|j| (j as f32) * 0.25 - 1.5)
517            .collect();
518
519        let gpu = ternary_gemm_gpu(&device, &queue, &act, &packed, scale, n_in, n_out, n_batch);
520
521        let mut cpu = vec![0.0f32; n_batch * n_out];
522        ternary_gemm_cpu(&act, &packed, scale, n_in, n_out, n_batch, 0, 0, &mut cpu);
523
524        assert_eq!(gpu.len(), cpu.len());
525        for i in 0..cpu.len() {
526            assert!(
527                (gpu[i] - cpu[i]).abs() < 1e-4,
528                "elem {i}: gpu {} vs cpu {}",
529                gpu[i],
530                cpu[i]
531            );
532        }
533        eprintln!(
534            "ternary_gemm_gpu: on-device parity OK ({} elems)",
535            cpu.len()
536        );
537    }
538
539    /// ON-DEVICE PARITY (2-bit branchless): `ternary_gemm_2bit.wgsl` == `ternary_gemm_cpu_2bit`.
540    #[test]
541    fn ternary_gemm_gpu_2bit_matches_cpu_oracle() {
542        use crate::ternary::{pack_trits_2bit, ternary_gemm_cpu_2bit};
543        let Some((device, queue)) = try_gpu() else {
544            eprintln!("ternary_gemm_gpu_2bit: no wgpu adapter — skipping");
545            return;
546        };
547        let (n_in, n_out, n_batch) = (7usize, 5usize, 3usize);
548        let scale = 0.37_f32;
549        let trits: Vec<i8> = (0..n_in * n_out).map(|k| (k % 3) as i8 - 1).collect();
550        let packed = pack_trits_2bit(&trits);
551        let act: Vec<f32> = (0..n_in * n_batch)
552            .map(|j| (j as f32) * 0.25 - 1.5)
553            .collect();
554
555        let gpu =
556            ternary_gemm_gpu_2bit(&device, &queue, &act, &packed, scale, n_in, n_out, n_batch);
557        let mut cpu = vec![0.0f32; n_batch * n_out];
558        ternary_gemm_cpu_2bit(&act, &packed, scale, n_in, n_out, n_batch, 0, 0, &mut cpu);
559        for i in 0..cpu.len() {
560            assert!(
561                (gpu[i] - cpu[i]).abs() < 1e-4,
562                "elem {i}: gpu {} vs cpu {}",
563                gpu[i],
564                cpu[i]
565            );
566        }
567        eprintln!(
568            "ternary_gemm_gpu_2bit: on-device parity OK ({} elems)",
569            cpu.len()
570        );
571    }
572
573    /// Build a base-3 ternary blob (`[scale f32][5-trits/byte]`) from explicit trits + scale.
574    fn base3_blob(scale: f32, trits: &[i8]) -> Vec<u8> {
575        let mut b = scale.to_le_bytes().to_vec();
576        b.extend_from_slice(&pack_trits(trits));
577        b
578    }
579
580    /// A1b inc 2b ON-DEVICE GATE: the resident 2-bit dispatcher (persistent pipeline + ONE resident
581    /// weight buffer + 256-aligned sub-range bindings, keyed by P64 blob offset) reproduces the
582    /// base-3 CPU oracle byte-for-byte — for TWO differently-shaped tensors at distinct keys, proving
583    /// the per-key lookup, sub-range binding, rebake, and GEMV are all correct. Skips with no GPU.
584    #[test]
585    fn ternary_ffn_resident_matches_cpu_oracle() {
586        let Some((device, queue)) = try_gpu() else {
587            eprintln!("ternary_ffn_resident: no wgpu adapter — skipping");
588            return;
589        };
590        // Two FFN-shaped tensors with DIFFERENT shapes + scales + distinct keys (blob offsets).
591        let a = (0x1000u64, 64usize, 40usize, 0.30f32);
592        let b = (0x2000u64, 96usize, 24usize, 0.11f32);
593        let trits_a: Vec<i8> = (0..a.1 * a.2).map(|k| (k % 3) as i8 - 1).collect();
594        let trits_b: Vec<i8> = (0..b.1 * b.2)
595            .map(|k| ((k * 7 + 2) % 3) as i8 - 1)
596            .collect();
597        let blob_a = base3_blob(a.3, &trits_a);
598        let blob_b = base3_blob(b.3, &trits_b);
599
600        let resident = TernaryFfnResident::build(
601            &device,
602            &queue,
603            &[(a.0, a.1, a.2, &blob_a), (b.0, b.1, b.2, &blob_b)],
604        )
605        .expect("build resident ternary set");
606        assert_eq!(resident.len(), 2);
607        assert!(resident.contains(a.0) && resident.contains(b.0));
608
609        for (key, n_in, n_out, scale, blob) in
610            [(a.0, a.1, a.2, a.3, &blob_a), (b.0, b.1, b.2, b.3, &blob_b)]
611        {
612            let act: Vec<f32> = (0..n_in).map(|j| (j as f32) * 0.25 - 1.0).collect();
613            let mut gpu = vec![0f32; n_out];
614            assert!(
615                resident.gemv(&device, &queue, key, &act, &mut gpu, n_in, n_out),
616                "resident gemv must succeed for key {key:#x}"
617            );
618            let mut cpu = vec![0f32; n_out];
619            ternary_gemm_cpu(&act, &blob[4..], scale, n_in, n_out, 1, 0, 0, &mut cpu);
620            for i in 0..n_out {
621                assert!(
622                    (gpu[i] - cpu[i]).abs() < 1e-4,
623                    "key {key:#x} row {i}: gpu {} vs cpu {}",
624                    gpu[i],
625                    cpu[i]
626                );
627            }
628        }
629        // fail-closed (false, never garbage): an absent key, and a present key with the WRONG shape.
630        let mut tmp = vec![0f32; 64];
631        assert!(
632            !resident.gemv(&device, &queue, 0xDEAD, &[0.0; 64], &mut tmp, 64, 40),
633            "absent key must fail-closed"
634        );
635        assert!(
636            !resident.gemv(&device, &queue, a.0, &[0.0; 64], &mut tmp, 64, 64),
637            "present key with mismatched n_out must fail-closed (a is 64x40, not 64x64)"
638        );
639        eprintln!(
640            "ternary_ffn_resident: on-device parity OK (2 tensors, {} resident bytes)",
641            resident.resident_bytes()
642        );
643    }
644
645    /// INDICATIVE A/B timing: base-3 branchy vs 2-bit branchless on a large GEMV. Wall-clock incl.
646    /// per-dispatch submit/readback overhead — a *relative* signal, not a rigorous TPS number
647    /// (real measurement needs timestamp queries inside the fused FFN loop). Skips with no GPU.
648    #[test]
649    fn ternary_gemm_2bit_vs_base3_indicative_timing() {
650        use crate::ternary::{pack_trits, pack_trits_2bit};
651        use std::time::Instant;
652        let Some((device, queue)) = try_gpu() else {
653            eprintln!("ternary timing: no wgpu adapter — skipping");
654            return;
655        };
656        let (n_in, n_out, n_batch) = (4096usize, 4096usize, 1usize); // decode-shape GEMV
657        let scale = 0.05_f32;
658        let trits: Vec<i8> = (0..n_in * n_out)
659            .map(|k| ((k * 7 + 1) % 3) as i8 - 1)
660            .collect();
661        let p3 = pack_trits(&trits);
662        let p2 = pack_trits_2bit(&trits);
663        let act: Vec<f32> = (0..n_in).map(|j| (j as f32 % 13.0) * 0.1 - 0.6).collect();
664        let iters = 60;
665
666        // warmup
667        let _ = ternary_gemm_gpu(&device, &queue, &act, &p3, scale, n_in, n_out, n_batch);
668        let _ = ternary_gemm_gpu_2bit(&device, &queue, &act, &p2, scale, n_in, n_out, n_batch);
669
670        let t0 = Instant::now();
671        for _ in 0..iters {
672            let _ = ternary_gemm_gpu(&device, &queue, &act, &p3, scale, n_in, n_out, n_batch);
673        }
674        let base3 = t0.elapsed().as_secs_f64() / iters as f64 * 1e3;
675
676        let t1 = Instant::now();
677        for _ in 0..iters {
678            let _ = ternary_gemm_gpu_2bit(&device, &queue, &act, &p2, scale, n_in, n_out, n_batch);
679        }
680        let bit2 = t1.elapsed().as_secs_f64() / iters as f64 * 1e3;
681
682        eprintln!(
683            "ternary GEMV {}x{} (indicative, incl. overhead): base-3 branchy {:.3} ms/iter | 2-bit branchless {:.3} ms/iter | speedup {:.2}x",
684            n_out, n_in, base3, bit2, base3 / bit2.max(1e-9)
685        );
686    }
687
688    /// F16 GEMV baseline shader (same bindings/params as the ternary kernels).
689    const F16_GEMV_WGSL: &str = include_str!("../shaders/f16_gemv.wgsl");
690
691    /// Persistent-pipeline, batched-dispatch GEMV timing — fixes the per-call rebuild flaw of the
692    /// indicative test. Pipeline + buffers are created ONCE; `K` dispatches are encoded per submit
693    /// and `S` submits are timed, so per-dispatch time is GPU-execution-dominated (submit/alloc
694    /// overhead amortized away). Returns ms per dispatch.
695    fn bench_kernel(
696        device: &wgpu::Device,
697        queue: &wgpu::Queue,
698        wgsl: &str,
699        entry: &str,
700        weight_bytes: &[u8],
701        n_in: usize,
702        n_out: usize,
703    ) -> f64 {
704        use std::time::Instant;
705        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
706            label: Some("bench"),
707            source: wgpu::ShaderSource::Wgsl(wgsl.into()),
708        });
709        let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
710            label: Some("bench_pipeline"),
711            layout: None,
712            module: &shader,
713            entry_point: Some(entry),
714            compilation_options: Default::default(),
715            cache: None,
716        });
717        let mk = |contents: &[u8], usage: wgpu::BufferUsages| {
718            let b = device.create_buffer(&wgpu::BufferDescriptor {
719                label: None,
720                size: contents.len().max(4) as u64,
721                usage,
722                mapped_at_creation: false,
723            });
724            if !contents.is_empty() {
725                queue.write_buffer(&b, 0, contents);
726            }
727            b
728        };
729        let act = mk(
730            bytemuck::cast_slice(&vec![0.1f32; n_in]),
731            wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
732        );
733        let mut w = weight_bytes.to_vec();
734        while w.len() % 4 != 0 || w.is_empty() {
735            w.push(0);
736        }
737        let wbuf = mk(
738            &w,
739            wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
740        );
741        let params = super::ternary_params_bytes(n_in as u32, n_out as u32, 1, 1.0);
742        let pbuf = mk(
743            &params,
744            wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
745        );
746        let obuf = device.create_buffer(&wgpu::BufferDescriptor {
747            label: None,
748            size: (n_out * 4).max(4) as u64,
749            usage: wgpu::BufferUsages::STORAGE,
750            mapped_at_creation: false,
751        });
752        let bind = device.create_bind_group(&wgpu::BindGroupDescriptor {
753            label: None,
754            layout: &pipeline.get_bind_group_layout(0),
755            entries: &[
756                wgpu::BindGroupEntry {
757                    binding: 0,
758                    resource: act.as_entire_binding(),
759                },
760                wgpu::BindGroupEntry {
761                    binding: 1,
762                    resource: wbuf.as_entire_binding(),
763                },
764                wgpu::BindGroupEntry {
765                    binding: 2,
766                    resource: pbuf.as_entire_binding(),
767                },
768                wgpu::BindGroupEntry {
769                    binding: 3,
770                    resource: obuf.as_entire_binding(),
771                },
772            ],
773        });
774        let wg_x = (n_out as u32).div_ceil(64).max(1);
775        let (k, s) = (32u32, 8u32);
776        let submit_batch = || {
777            let mut enc =
778                device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
779            {
780                let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor {
781                    label: None,
782                    timestamp_writes: None,
783                });
784                pass.set_pipeline(&pipeline);
785                pass.set_bind_group(0, &bind, &[]);
786                for _ in 0..k {
787                    pass.dispatch_workgroups(wg_x, 1, 1);
788                }
789            }
790            queue.submit(Some(enc.finish()));
791            let _ = device.poll(wgpu::PollType::wait_indefinitely());
792        };
793        submit_batch(); // warmup
794        let t0 = Instant::now();
795        for _ in 0..s {
796            submit_batch();
797        }
798        t0.elapsed().as_secs_f64() * 1e3 / (k * s) as f64
799    }
800
801    /// A2000 KERNEL BENCHMARK (real numbers): F16 vs base-3 ternary vs 2-bit branchless ternary, at a
802    /// decode-shape GEMV, with persistent pipeline + buffer reuse. Skips with no GPU.
803    #[test]
804    fn ternary_kernel_benchmark() {
805        use crate::ternary::{pack_trits, pack_trits_2bit};
806        let Some((device, queue)) = try_gpu() else {
807            eprintln!("ternary_kernel_benchmark: no wgpu adapter — skipping");
808            return;
809        };
810        let (n_in, n_out) = (4096usize, 4096usize); // decode GEMV (batch 1)
811        let trits: Vec<i8> = (0..n_in * n_out)
812            .map(|k| ((k * 7 + 1) % 3) as i8 - 1)
813            .collect();
814        let p3 = pack_trits(&trits);
815        let p2 = pack_trits_2bit(&trits);
816        let f16 = vec![0u8; n_in * n_out * 2]; // f16 weights (size = bandwidth; values irrelevant to timing)
817
818        let f16ms = bench_kernel(
819            &device,
820            &queue,
821            F16_GEMV_WGSL,
822            "f16_gemv",
823            &f16,
824            n_in,
825            n_out,
826        );
827        let base3 = bench_kernel(
828            &device,
829            &queue,
830            TERNARY_GEMM_WGSL,
831            "ternary_gemm",
832            &p3,
833            n_in,
834            n_out,
835        );
836        let bit2 = bench_kernel(
837            &device,
838            &queue,
839            TERNARY_GEMM_2BIT_WGSL,
840            "ternary_gemm",
841            &p2,
842            n_in,
843            n_out,
844        );
845
846        eprintln!(
847            "── A2000 GEMV {}x{} batch=1 (persistent pipeline, {}MB/{}KB/{}KB weights) ──",
848            n_out,
849            n_in,
850            f16.len() / (1 << 20),
851            p3.len() >> 10,
852            p2.len() >> 10
853        );
854        eprintln!("  F16 baseline        : {:.4} ms/dispatch", f16ms);
855        eprintln!(
856            "  ternary base-3      : {:.4} ms/dispatch  ({:.2}x vs F16)",
857            base3,
858            f16ms / base3.max(1e-12)
859        );
860        eprintln!(
861            "  ternary 2-bit branchless: {:.4} ms/dispatch  ({:.2}x vs F16, {:.2}x vs base-3)",
862            bit2,
863            f16ms / bit2.max(1e-12),
864            base3 / bit2.max(1e-12)
865        );
866    }
867}
qualia_core_db/inference/ternary_gpu.rs

qualia_core_db/inference/
ternary_gpu.rs
