char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Reese Levineggml-webgpu: add q4_0/q8_0 SET_ROWS (#23760)b22da2588

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
6.5 KiB224 linesraw
1#ifdef DST_Q8_0
2#define BLOCK_SIZE 32u
3#define BLOCK_BYTES 34u
4#define QS_WORDS 8u
5#elif defined(DST_Q4_0)
6#define BLOCK_SIZE 32u
7#define BLOCK_BYTES 18u
8#define QS_WORDS 4u
9#endif
10
11@group(0) @binding(0)
12var<storage, read_write> src: array<f32>;
13
14@group(0) @binding(1)
15var<storage, read_write> idx: array<u32>;
16
17@group(0) @binding(2)
18#ifdef PAIR_BLOCKS
19var<storage, read_write> dst: array<u32>;
20#else
21var<storage, read_write> dst: array<atomic<u32>>;
22#endif
23
24#ifdef I64_IDX
25@group(0) @binding(3)
26var<storage, read_write> error: atomic<u32>;
27#define PARAMS_BINDING 4
28#else
29#define PARAMS_BINDING 3
30#endif
31
32struct Params {
33    offset_src: u32, // in elements
34    offset_idx: u32, // in elements
35    offset_dst: u32, // in blocks
36
37    // Strides (in elements / blocks)
38    stride_src1: u32,
39    stride_src2: u32,
40    stride_src3: u32,
41
42    stride_idx0: u32,
43    stride_idx1: u32,
44    stride_idx2: u32,
45
46    stride_dst1: u32,
47    stride_dst2: u32,
48    stride_dst3: u32,
49
50    // Shape of src
51    ne0: u32,
52    n_rows: u32,
53    ne2: u32,
54    ne3: u32,
55
56    // Shape of idx
57    idx1: u32,
58    idx2: u32,
59};
60
61@group(0) @binding(PARAMS_BINDING)
62var<uniform> params: Params;
63
64// if the quantization type is unaligned and there are an odd number of blocks per row, we need to store atomically
65#ifndef PAIR_BLOCKS
66fn merge_store_dst_word(word_idx: u32, mask: u32, bits: u32) {
67    loop {
68        let old = atomicLoad(&dst[word_idx]);
69        let merged = (old & ~mask) | (bits & mask);
70        let result = atomicCompareExchangeWeak(&dst[word_idx], old, merged);
71        if (result.exchanged) {
72            return;
73        }
74    }
75}
76#else
77fn merge_store_dst_word(word_idx: u32, mask: u32, bits: u32) {
78    let old = dst[word_idx];
79    dst[word_idx] = (old & ~mask) | (bits & mask);
80}
81#endif
82
83fn store_u16(dst_word_idx: u32, block_byte_offset: u32, byte_offset: u32, value: u32) {
84    let total_byte_offset = block_byte_offset + byte_offset;
85    let word_idx = dst_word_idx + total_byte_offset / 4u;
86    let shift = (total_byte_offset & 2u) * 8u;
87    let mask = 0xFFFFu << shift;
88    merge_store_dst_word(word_idx, mask, (value & 0xFFFFu) << shift);
89}
90
91fn store_u32(dst_word_idx: u32, block_byte_offset: u32, byte_offset: u32, value: u32) {
92    let total_byte_offset = block_byte_offset + byte_offset;
93    let word_idx = dst_word_idx + total_byte_offset / 4u;
94    let shift = (total_byte_offset & 3u) * 8u;
95
96    if (shift == 0u) {
97#ifdef PAIR_BLOCKS
98        dst[word_idx] = value;
99#else
100        atomicStore(&dst[word_idx], value);
101#endif
102        return;
103    }
104
105    let lo_mask = 0xFFFFFFFFu << shift;
106    let hi_mask = (1u << shift) - 1u;
107    merge_store_dst_word(word_idx, lo_mask, value << shift);
108    merge_store_dst_word(word_idx + 1u, hi_mask, value >> (32u - shift));
109}
110
111fn quantize_block_params(src_block: u32) -> vec2<f32> {
112#ifdef DST_Q8_0
113    var amax = 0.0;
114    for (var j: u32 = 0u; j < BLOCK_SIZE; j++) {
115        amax = max(amax, abs(src[src_block + j]));
116    }
117
118    let d = amax / 127.0;
119    let id = select(0.0, 1.0 / d, d > 0.0);
120    return vec2(d, id);
121#elif defined(DST_Q4_0)
122    var amax = 0.0;
123    var max_val = 0.0;
124    for (var j: u32 = 0u; j < BLOCK_SIZE; j++) {
125        let v = src[src_block + j];
126        let av = abs(v);
127        if (amax < av) {
128            amax = av;
129            max_val = v;
130        }
131    }
132
133    let d = max_val / -8.0;
134    let id = select(0.0, 1.0 / d, d != 0.0);
135    return vec2(d, id);
136#endif
137}
138
139fn quantize_block_word(src_block: u32, j: u32, id: f32) -> u32 {
140#ifdef DST_Q8_0
141    let base = src_block + j * 4u;
142    return (u32(i32(round(src[base + 0u] * id)) & 0xFF) << 0u) |
143           (u32(i32(round(src[base + 1u] * id)) & 0xFF) << 8u) |
144           (u32(i32(round(src[base + 2u] * id)) & 0xFF) << 16u) |
145           (u32(i32(round(src[base + 3u] * id)) & 0xFF) << 24u);
146#elif defined(DST_Q4_0)
147    var packed_q = 0u;
148    for (var k: u32 = 0u; k < 4u; k++) {
149        let x0 = src[src_block + j * 4u + k] * id;
150        let x1 = src[src_block + 16u + j * 4u + k] * id;
151        let q0 = u32(clamp(i32(x0 + 8.5), 0, 15));
152        let q1 = u32(clamp(i32(x1 + 8.5), 0, 15));
153        packed_q |= (q0 & 0xFu) << (8u * k);
154        packed_q |= (q1 & 0xFu) << (8u * k + 4u);
155    }
156    return packed_q;
157#endif
158}
159
160fn quantize_block(src_block: u32, dst_word_idx: u32, block_byte_offset: u32) {
161    let params = quantize_block_params(src_block);
162    let d = params.x;
163    let id = params.y;
164    let packed_d = pack2x16float(vec2(d, 0.0)) & 0xFFFFu;
165    store_u16(dst_word_idx, block_byte_offset, 0u, packed_d);
166
167    for (var j: u32 = 0u; j < QS_WORDS; j++) {
168        store_u32(dst_word_idx, block_byte_offset, 2u + j * 4u, quantize_block_word(src_block, j, id));
169    }
170}
171
172@compute @workgroup_size(WG_SIZE)
173fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
174    let blocks_per_row = params.ne0 / BLOCK_SIZE;
175#ifdef PAIR_BLOCKS
176    let blocks_per_invocation = 2u;
177#else
178    let blocks_per_invocation = 1u;
179#endif
180    let invocations_per_row = blocks_per_row / blocks_per_invocation;
181    let total_invocations = params.ne3 * params.ne2 * params.n_rows * invocations_per_row;
182    if (gid.x >= total_invocations) {
183        return;
184    }
185
186    var i = gid.x / invocations_per_row;
187    let block_in_row = (gid.x % invocations_per_row) * blocks_per_invocation;
188
189    let i_src3 = i / (params.ne2 * params.n_rows);
190    i = i % (params.ne2 * params.n_rows);
191    let i_src2 = i / params.n_rows;
192    let i_src1 = i % params.n_rows;
193
194    let i_idx2 = i_src3 % params.idx2;
195    let i_idx1 = i_src2 % params.idx1;
196    let i_idx0 = i_src1;
197
198#ifdef I64_IDX
199    let idx_high = (params.offset_idx + i_idx0 * params.stride_idx0 + i_idx1 * params.stride_idx1 + i_idx2 * params.stride_idx2) * 2u;
200    let idx_val = idx[idx_high];
201    let idx_low_val = idx[idx_high + 1u];
202
203    if (idx_low_val != 0u) {
204        atomicStore(&error, 1u);
205        return;
206    }
207#else
208    let idx_i = params.offset_idx + i_idx0 * params.stride_idx0 + i_idx1 * params.stride_idx1 + i_idx2 * params.stride_idx2;
209    let idx_val = idx[idx_i];
210#endif
211
212    let dst_row_blocks = params.offset_dst + idx_val * params.stride_dst1 + i_src2 * params.stride_dst2 + i_src3 * params.stride_dst3;
213    let src_row = params.offset_src + i_src1 * params.stride_src1 + i_src2 * params.stride_src2 + i_src3 * params.stride_src3;
214    let src_block = src_row + block_in_row * BLOCK_SIZE;
215    let dst_block_byte = (dst_row_blocks + block_in_row) * BLOCK_BYTES;
216
217    let dst_word_idx = dst_block_byte / 4u;
218#ifdef PAIR_BLOCKS
219    quantize_block(src_block, dst_word_idx, 0u);
220    quantize_block(src_block + BLOCK_SIZE, dst_word_idx, BLOCK_BYTES);
221#else
222    quantize_block(src_block, dst_word_idx, dst_block_byte & 3u);
223#endif
224}