rlx-wgpu 0.2.12

Cross-platform GPU backend for RLX via wgpu (Metal/Vulkan/DX12/WebGPU)
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
//
// GPL-3.0-only. See LICENSE.

//! IQ-family grid LUTs staged into a wgpu storage buffer.
//!
//! Same byte layout as `rlx_metal::kernels::iq_grid_buffer` /
//! `rlx_cuda::iq_grid::cuda_iq_grid_buffer`.

#[cfg(not(target_arch = "wasm32"))]
use std::sync::{Mutex, OnceLock};

fn build_bytes() -> Vec<u8> {
    use rlx_gguf::iq_grids::{
        IQ1S_GRID, IQ2S_GRID, IQ2XS_GRID, IQ2XXS_GRID, IQ3S_GRID, IQ3XXS_GRID, KMASK_IQ2XS,
        KSIGNS_IQ2XS,
    };
    let mut bytes = Vec::with_capacity(33_944);
    bytes.extend_from_slice(&KMASK_IQ2XS);
    bytes.extend_from_slice(&KSIGNS_IQ2XS);
    for v in IQ2XXS_GRID.iter() {
        bytes.extend_from_slice(&v.to_le_bytes());
    }
    for v in IQ2XS_GRID.iter() {
        bytes.extend_from_slice(&v.to_le_bytes());
    }
    for v in IQ2S_GRID.iter() {
        bytes.extend_from_slice(&v.to_le_bytes());
    }
    for v in IQ3XXS_GRID.iter() {
        bytes.extend_from_slice(&v.to_le_bytes());
    }
    for v in IQ3S_GRID.iter() {
        bytes.extend_from_slice(&v.to_le_bytes());
    }
    for v in IQ1S_GRID.iter() {
        bytes.extend_from_slice(&v.to_le_bytes());
    }
    bytes
}

/// Build the IQ grid LUT storage buffer (uncached). The byte layout matches
/// the Metal / CUDA equivalents.
fn create_grid_buffer(device: &wgpu::Device, queue: &wgpu::Queue) -> wgpu::Buffer {
    let bytes = build_bytes();
    let padded = bytes.len().div_ceil(4) * 4;
    let mut upload = bytes;
    upload.resize(padded, 0);
    let buf = device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("rlx-wgpu iq_grid_lut"),
        size: padded as u64,
        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    });
    queue.write_buffer(&buf, 0, &upload);
    buf
}

/// Cached IQ grid LUT (read-only storage). Built once per process.
///
/// Native uses a `Sync` global; on wasm `wgpu::Buffer` is `!Send + !Sync`
/// (the browser is single-threaded), so the cache is a `thread_local` —
/// equivalent here since there is only ever one thread.
#[cfg(not(target_arch = "wasm32"))]
pub fn wgpu_iq_grid_buffer(device: &wgpu::Device, queue: &wgpu::Queue) -> wgpu::Buffer {
    static CACHE: OnceLock<Mutex<Option<wgpu::Buffer>>> = OnceLock::new();
    let cache = CACHE.get_or_init(|| Mutex::new(None));
    let mut guard = cache.lock().expect("iq_grid cache poisoned");
    if let Some(buf) = guard.as_ref() {
        return buf.clone();
    }
    let buf = create_grid_buffer(device, queue);
    guard.replace(buf.clone());
    buf
}

/// Cached IQ grid LUT (read-only storage). Built once per thread on wasm.
#[cfg(target_arch = "wasm32")]
pub fn wgpu_iq_grid_buffer(device: &wgpu::Device, queue: &wgpu::Queue) -> wgpu::Buffer {
    thread_local! {
        static CACHE: std::cell::RefCell<Option<wgpu::Buffer>> =
            const { std::cell::RefCell::new(None) };
    }
    CACHE.with(|cache| {
        let mut guard = cache.borrow_mut();
        if let Some(buf) = guard.as_ref() {
            return buf.clone();
        }
        let buf = create_grid_buffer(device, queue);
        guard.replace(buf.clone());
        buf
    })
}