char-fork/llama.cpp

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

Masashi Yoshimuraggml-webgpu: improve MTP inference by using mat-vec path for small batches (#24811)7c908502e

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
5.2 KiB178 linesraw
1#ifdef USE_SUBGROUP_REDUCTION
2enable subgroups;
3#endif
4enable f16;
5
6requires packed_4x8_integer_dot_product;
7
8#include "common_decls.tmpl"
9
10struct Params {
11    offset_src1: u32,
12    stride_11: u32,
13    stride_12: u32,
14    stride_13: u32,
15    ne0: u32,
16    ne1: u32,
17    ne2: u32,
18    ne3: u32,
19};
20
21#define SRC1_TYPE vec4<SRC1_INNER_TYPE>
22
23@group(0) @binding(0) var<storage, read_write> src1: array<SRC1_TYPE>;
24@group(0) @binding(1) var<storage, read_write> src1q: array<q8_1>;
25
26@group(0) @binding(2) var<uniform> params: Params;
27
28#ifdef USE_SUBGROUP_REDUCTION
29fn cluster_max_8(v: f32) -> f32 {
30    var r = v;
31    r = max(r, subgroupShuffleXor(r, 1u));
32    r = max(r, subgroupShuffleXor(r, 2u));
33    r = max(r, subgroupShuffleXor(r, 4u));
34    return r;
35}
36
37#if defined(MUL_ACC_Q4_0) || defined(MUL_ACC_Q4_1) || defined(MUL_ACC_Q4_K)
38fn cluster_add_i4x8(v: i32) -> i32 {
39    var r= v;
40    r += subgroupShuffleXor(r, 1u);
41    r += subgroupShuffleXor(r, 2u);
42    r += subgroupShuffleXor(r, 4u);
43    return r;
44}
45#endif
46#endif
47
48#ifdef USE_WORKGROUP_REDUCTION
49#define CLUSTER_SIZE 8
50
51var<workgroup> partial_amaxs: array<array<f32, CLUSTER_SIZE>, WG_SIZE / CLUSTER_SIZE>;
52var<workgroup> partial_sums:  array<array<i32, CLUSTER_SIZE>, WG_SIZE / CLUSTER_SIZE>;
53#endif
54
55@compute @workgroup_size(WG_SIZE)
56fn main(
57    @builtin(local_invocation_id) local_id: vec3<u32>,
58    @builtin(workgroup_id) wg_id: vec3<u32>,
59    @builtin(num_workgroups) num_wg: vec3<u32>
60) {
61    let thread_id = local_id.x;
62    let ne0_vec4 = params.ne0 / 4u;
63
64    let wg_per_vec = (ne0_vec4 + (WG_SIZE - 1u)) / WG_SIZE;
65    let total_batches = wg_per_vec * params.ne1 * params.ne2 * params.ne3;
66
67    let wg_linear = wg_id.y * num_wg.x + wg_id.x;
68    if (wg_linear >= total_batches) {
69        return;
70    }
71
72    let vec_idx = wg_linear / wg_per_vec;
73    let src13_idx = vec_idx / (params.ne2 * params.ne1);
74    let vec_ne12_num       = vec_idx % (params.ne2 * params.ne1);
75    let src12_idx = vec_ne12_num / params.ne1;
76    let src11_idx = vec_ne12_num % params.ne1;
77    let src1_idx_base = params.offset_src1 + src13_idx * params.stride_13 + src12_idx * params.stride_12 + src11_idx * params.stride_11;
78    let src1_idx_vec4_base = src1_idx_base / 4u;
79
80    let blocks_per_row = params.ne0 / 32u;
81    let blocks_per_wg = (WG_SIZE * 4u) / 32u;
82    let src1q_idx_base = ((src13_idx * params.ne2 + src12_idx) * params.ne1 + src11_idx) * blocks_per_row;
83    let src11_wg_idx = wg_linear % wg_per_vec;
84    let src1q_idx = src1q_idx_base + src11_wg_idx * blocks_per_wg + thread_id / 8u;
85    let qs_idx = thread_id % 8u;
86
87    // reduction
88    var q4 = vec4<f32>(0.0);
89    var q4_quants = 0u;
90    var thread_amax = 0.0;
91
92    let src11_vec4_idx = src11_wg_idx * WG_SIZE + thread_id;
93    let is_valid = src11_vec4_idx < ne0_vec4;
94
95#ifdef USE_SUBGROUP_REDUCTION
96
97    var d = 0.0;
98
99    if (is_valid) {
100        q4 = src1[src1_idx_vec4_base + src11_vec4_idx];
101        let abs_q4 = abs(q4);
102        thread_amax = max(max(abs_q4[0u], abs_q4[1u]), max(abs_q4[2], abs_q4[3]));
103    }
104
105    d = cluster_max_8(thread_amax) / 127.0;
106
107    if (is_valid) {
108        let id = select(0.0, 1.0 / d, d > 0.0);
109        q4_quants = pack4xI8(vec4<i32>(round(q4 * id)));
110        if (qs_idx == 0u) {
111            src1q[src1q_idx].d = f16(d);
112        }
113        src1q[src1q_idx].qs[qs_idx] = q4_quants;
114    }
115
116#if defined(MUL_ACC_Q4_0) || defined(MUL_ACC_Q4_1) || defined(MUL_ACC_Q4_K)
117    let q4_quants_sum = dot4I8Packed(q4_quants, 0x01010101u);
118    let s = f16(d * f32(cluster_add_i4x8(q4_quants_sum)));
119
120    if (is_valid) {
121        if (qs_idx == 0u) {
122            src1q[src1q_idx].s = s;
123        }
124    }
125#endif
126#endif
127
128#ifdef USE_WORKGROUP_REDUCTION
129
130    var d = 0.0;
131    let cluster_id = thread_id / 8u;
132
133    if (is_valid) {
134        q4 = src1[src1_idx_vec4_base + src11_vec4_idx];
135        let abs_q4 = abs(q4);
136        thread_amax = max(max(abs_q4[0], abs_q4[1]), max(abs_q4[2], abs_q4[3]));
137        partial_amaxs[cluster_id][qs_idx] = thread_amax;
138    }
139
140    workgroupBarrier();
141
142    if (is_valid) {
143        let amax = max(
144                    max(
145                        max(partial_amaxs[cluster_id][0], partial_amaxs[cluster_id][1]), max(partial_amaxs[cluster_id][2], partial_amaxs[cluster_id][3])),
146                    max(
147                        max(partial_amaxs[cluster_id][4], partial_amaxs[cluster_id][5]), max(partial_amaxs[cluster_id][6], partial_amaxs[cluster_id][7]))
148                );
149
150        d = amax / 127.0;
151        let id = select(0.0f, 1.0f / d, d > 0.0f);
152
153        q4_quants = pack4xI8(vec4<i32>(round(q4 * id)));
154        src1q[src1q_idx].qs[qs_idx] = q4_quants;
155
156        if (qs_idx == 0u) {
157            src1q[src1q_idx].d = f16(d);
158        }
159    }
160
161#if defined(MUL_ACC_Q4_0) || defined(MUL_ACC_Q4_1) || defined(MUL_ACC_Q4_K)
162
163    partial_sums[cluster_id][qs_idx] = dot4I8Packed(q4_quants, 0x01010101u);
164
165    workgroupBarrier();
166
167    if (is_valid) {
168        if (qs_idx == 0u) {
169            let s = d * f32(partial_sums[cluster_id][0] + partial_sums[cluster_id][1] + partial_sums[cluster_id][2] + partial_sums[cluster_id][3]
170                                    + partial_sums[cluster_id][4] + partial_sums[cluster_id][5] + partial_sums[cluster_id][6] + partial_sums[cluster_id][7]);
171            src1q[src1q_idx].s = f16(s);
172        }
173    }
174
175#endif
176#endif
177
178}