qualia_core_db/inference/
topk_gpu.rs

1//! STELLAR §A A1a — **native GPU dispatch of the top-K reduction** + on-device parity.
2//!
3//! Runs `shaders/topk_reduction.wgsl` on a real wgpu device: each workgroup reduces a block of
4//! the logits to its top-K candidates, the host merges to the global top-K. The `#[test]` below
5//! verifies it **on silicon** against the CPU oracle `topk::topk_cpu` (skips cleanly with no
6//! adapter). This is the reusable core that the decode-loop splice (behind `QUALIA_LLM_GPU_TOPK`)
7//! will call instead of the full-logit-readback `dispatch_output_argmax_chunked`.
8//!
9//! Native only — mirrors `ternary_gpu.rs`.
10
11use crate::topk::{
12    merge_block_candidates, topk_params_bytes, TopKItem, TOPK_BLOCK_SIZE, TOPK_REDUCTION_WGSL,
13};
14
15/// Reduce `logits` to its global top-K on the GPU. `block_size` is elements per workgroup
16/// (clamped to `TOPK_BLOCK_SIZE`, the WGSL `var<workgroup>` cap). Blocking (native readback).
17pub fn topk_gpu(
18    device: &wgpu::Device,
19    queue: &wgpu::Queue,
20    logits: &[f32],
21    k: usize,
22    block_size: usize,
23) -> Vec<TopKItem> {
24    let n = logits.len();
25    let k = k.max(1);
26    let block_size = block_size.clamp(1, TOPK_BLOCK_SIZE);
27    if n == 0 {
28        return Vec::new();
29    }
30    let num_blocks = n.div_ceil(block_size);
31    let cand_count = num_blocks * k;
32
33    let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
34        label: Some("topk_reduction"),
35        source: wgpu::ShaderSource::Wgsl(TOPK_REDUCTION_WGSL.into()),
36    });
37    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
38        label: Some("topk_pipeline"),
39        layout: None,
40        module: &shader,
41        entry_point: Some("topk_block"),
42        compilation_options: Default::default(),
43        cache: None,
44    });
45
46    let logits_buf = device.create_buffer(&wgpu::BufferDescriptor {
47        label: Some("topk_logits"),
48        size: (n * 4) as u64,
49        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
50        mapped_at_creation: false,
51    });
52    queue.write_buffer(&logits_buf, 0, bytemuck::cast_slice(logits));
53
54    let params = topk_params_bytes(n as u32, k as u32, block_size as u32);
55    let params_buf = device.create_buffer(&wgpu::BufferDescriptor {
56        label: Some("topk_params"),
57        size: params.len() as u64,
58        usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
59        mapped_at_creation: false,
60    });
61    queue.write_buffer(&params_buf, 0, &params);
62
63    let mk_io = |label: &str| {
64        device.create_buffer(&wgpu::BufferDescriptor {
65            label: Some(label),
66            size: (cand_count * 4) as u64,
67            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
68            mapped_at_creation: false,
69        })
70    };
71    let cand_val = mk_io("topk_cand_val");
72    let cand_idx = mk_io("topk_cand_idx");
73    let mk_stg = |label: &str| {
74        device.create_buffer(&wgpu::BufferDescriptor {
75            label: Some(label),
76            size: (cand_count * 4) as u64,
77            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
78            mapped_at_creation: false,
79        })
80    };
81    let stg_val = mk_stg("topk_stg_val");
82    let stg_idx = mk_stg("topk_stg_idx");
83
84    let bind = device.create_bind_group(&wgpu::BindGroupDescriptor {
85        label: Some("topk_bind"),
86        layout: &pipeline.get_bind_group_layout(0),
87        entries: &[
88            wgpu::BindGroupEntry {
89                binding: 0,
90                resource: logits_buf.as_entire_binding(),
91            },
92            wgpu::BindGroupEntry {
93                binding: 1,
94                resource: params_buf.as_entire_binding(),
95            },
96            wgpu::BindGroupEntry {
97                binding: 2,
98                resource: cand_val.as_entire_binding(),
99            },
100            wgpu::BindGroupEntry {
101                binding: 3,
102                resource: cand_idx.as_entire_binding(),
103            },
104        ],
105    });
106
107    let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
108        label: Some("topk_enc"),
109    });
110    {
111        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
112            label: Some("topk_pass"),
113            timestamp_writes: None,
114        });
115        pass.set_pipeline(&pipeline);
116        pass.set_bind_group(0, &bind, &[]);
117        pass.dispatch_workgroups(num_blocks as u32, 1, 1);
118    }
119    let bytes = (cand_count * 4) as u64;
120    encoder.copy_buffer_to_buffer(&cand_val, 0, &stg_val, 0, bytes);
121    encoder.copy_buffer_to_buffer(&cand_idx, 0, &stg_idx, 0, bytes);
122    queue.submit(Some(encoder.finish()));
123
124    let sv = stg_val.slice(..);
125    let si = stg_idx.slice(..);
126    let (tx_v, rx_v) = std::sync::mpsc::channel();
127    let (tx_i, rx_i) = std::sync::mpsc::channel();
128    sv.map_async(wgpu::MapMode::Read, move |r| {
129        let _ = tx_v.send(r);
130    });
131    si.map_async(wgpu::MapMode::Read, move |r| {
132        let _ = tx_i.send(r);
133    });
134    let _ = device.poll(wgpu::PollType::wait_indefinitely());
135    rx_v.recv().expect("map val").expect("map topk val");
136    rx_i.recv().expect("map idx").expect("map topk idx");
137
138    let dv = sv.get_mapped_range().expect("wgpu buffer map_range failed");
139    let di = si.get_mapped_range().expect("wgpu buffer map_range failed");
140    let vals: Vec<f32> = bytemuck::cast_slice(&dv)[..cand_count].to_vec();
141    let idxs: Vec<u32> = bytemuck::cast_slice(&di)[..cand_count].to_vec();
142    drop(dv);
143    drop(di);
144    stg_val.unmap();
145    stg_idx.unmap();
146
147    merge_block_candidates(&vals, &idxs, k, None)
148}
149
150#[cfg(test)]
151mod tests {
152    use super::*;
153    use crate::topk::topk_cpu;
154
155    fn try_gpu() -> Option<(wgpu::Device, wgpu::Queue)> {
156        let instance = wgpu::Instance::default();
157        let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
158            power_preference: wgpu::PowerPreference::HighPerformance,
159            ..Default::default()
160        }))
161        .ok()?;
162        pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor::default())).ok()
163    }
164
165    fn assert_parity(gpu: &[TopKItem], cpu: &[TopKItem]) {
166        assert_eq!(gpu.len(), cpu.len(), "top-k length");
167        for (g, c) in gpu.iter().zip(cpu.iter()) {
168            assert_eq!(
169                g.token_id, c.token_id,
170                "token id mismatch: gpu {g:?} cpu {c:?}"
171            );
172            assert!(
173                (g.logit - c.logit).abs() < 1e-5,
174                "logit mismatch: gpu {g:?} cpu {c:?}"
175            );
176        }
177    }
178
179    /// ON-DEVICE PARITY: small n spanning multiple blocks; top-K == CPU oracle (incl. tie-break).
180    #[test]
181    fn topk_gpu_matches_cpu_small_multiblock() {
182        let Some((device, queue)) = try_gpu() else {
183            eprintln!("topk_gpu: no wgpu adapter — skipping");
184            return;
185        };
186        // 50 logits, block_size 16 → 4 workgroups; deterministic with an embedded tie.
187        let mut logits: Vec<f32> = (0..50).map(|i| ((i * 13 + 7) % 31) as f32 - 15.0).collect();
188        logits[8] = 99.0;
189        logits[40] = 99.0; // tie at the top → lower id (8) must win
190        logits[3] = f32::NAN; // must never be selected
191
192        for k in [1usize, 5, 12] {
193            let gpu = topk_gpu(&device, &queue, &logits, k, 16);
194            let cpu = topk_cpu(&logits, k);
195            assert_parity(&gpu, &cpu);
196        }
197        eprintln!("topk_gpu: small multi-block parity OK");
198    }
199
200    /// ON-DEVICE PARITY at vocab scale (49 152, block 1024) — the real decode shape.
201    #[test]
202    fn topk_gpu_matches_cpu_vocab_scale() {
203        let Some((device, queue)) = try_gpu() else {
204            eprintln!("topk_gpu vocab: no wgpu adapter — skipping");
205            return;
206        };
207        let n = 49_152usize;
208        let logits: Vec<f32> = (0..n)
209            .map(|i| (((i * 1103515245 + 12345) >> 7) % 1000) as f32 * 0.01 - 5.0)
210            .collect();
211        for k in [1usize, 32, 64] {
212            let gpu = topk_gpu(&device, &queue, &logits, k, TOPK_BLOCK_SIZE);
213            let cpu = topk_cpu(&logits, k);
214            assert_parity(&gpu, &cpu);
215        }
216        eprintln!("topk_gpu: vocab-scale (49152) parity OK for k∈{{1,32,64}}");
217    }
218}
qualia_core_db/inference/topk_gpu.rs

qualia_core_db/inference/
topk_gpu.rs
