char-fork/llama.cpp

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

lnigamggml : recurrent state rollback for ggml_ssm_scan (#26623)1692f9e50

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
7.9 KiB247 linesraw
1#ifdef USE_SUBGROUP_REDUCTION
2enable subgroups;
3#endif
4
5struct Params {
6    offset_s: u32,
7    offset_x: u32,
8    offset_dt: u32,
9    offset_A: u32,
10    offset_B: u32,
11    offset_C: u32,
12    offset_ids: u32,
13    offset_dst: u32,
14
15    stride_s1: u32,
16    stride_s2: u32,
17    stride_s3: u32,
18
19    stride_x1: u32,
20    stride_x2: u32,
21    stride_x3: u32,
22
23    stride_dt1: u32,
24    stride_dt2: u32,
25
26    a_ne0: u32,
27    stride_A1: u32,
28
29    stride_B1: u32,
30    stride_B2: u32,
31    stride_B3: u32,
32
33    stride_C1: u32,
34    stride_C2: u32,
35    stride_C3: u32,
36
37    d_state: u32,
38    d_inner: u32,
39    n_head: u32,
40    n_group: u32,
41    n_seq_tokens: u32,
42
43    y_elems: u32,
44    K: u32,
45};
46
47@group(0) @binding(0) var<storage, read_write> s_in: array<f32>;
48#ifdef XBC_OVERLAP
49#ifdef IDS_OVERLAP
50@group(0) @binding(1) var<storage, read_write> x_dt_B_C_ids_merged: array<u32>;
51#ifdef A_OVERLAP
52@group(0) @binding(2) var<storage, read_write> dst: array<f32>;
53@group(0) @binding(3) var<uniform> params: Params;
54#else
55@group(0) @binding(2) var<storage, read_write> A: array<f32>;
56@group(0) @binding(3) var<storage, read_write> dst: array<f32>;
57@group(0) @binding(4) var<uniform> params: Params;
58#endif
59#else
60@group(0) @binding(1) var<storage, read_write> x_dt_B_C_merged: array<f32>;
61#ifdef A_OVERLAP
62@group(0) @binding(2) var<storage, read_write> ids: array<i32>;
63@group(0) @binding(3) var<storage, read_write> dst: array<f32>;
64@group(0) @binding(4) var<uniform> params: Params;
65#else
66@group(0) @binding(2) var<storage, read_write> A: array<f32>;
67@group(0) @binding(3) var<storage, read_write> ids: array<i32>;
68@group(0) @binding(4) var<storage, read_write> dst: array<f32>;
69@group(0) @binding(5) var<uniform> params: Params;
70#endif
71#endif
72#else
73@group(0) @binding(1) var<storage, read_write> x: array<f32>;
74@group(0) @binding(2) var<storage, read_write> dt: array<f32>;
75@group(0) @binding(3) var<storage, read_write> A: array<f32>;
76@group(0) @binding(4) var<storage, read_write> B: array<f32>;
77@group(0) @binding(5) var<storage, read_write> C: array<f32>;
78@group(0) @binding(6) var<storage, read_write> ids: array<i32>;
79@group(0) @binding(7) var<storage, read_write> dst: array<f32>;
80@group(0) @binding(8) var<uniform> params: Params;
81#endif
82
83var<workgroup> shared_x_dt: array<f32, TOKENS_PER_TILE>;
84var<workgroup> shared_dtsp: array<f32, TOKENS_PER_TILE>;
85var<workgroup> shared_reduce: array<f32, TOKENS_PER_TILE * WG_SIZE>;
86
87fn reduce_base(token_in_tile: u32) -> u32 {
88    return token_in_tile * WG_SIZE;
89}
90
91#ifdef XBC_OVERLAP
92fn read_merged_f32(idx: u32) -> f32 {
93#ifdef IDS_OVERLAP
94    return bitcast<f32>(x_dt_B_C_ids_merged[idx]);
95#else
96    return x_dt_B_C_merged[idx];
97#endif
98}
99#endif
100
101fn read_state_slot(i3: u32) -> u32 {
102#ifdef IDS_OVERLAP
103    return x_dt_B_C_ids_merged[params.offset_ids + i3];
104#else
105    return u32(ids[params.offset_ids + i3]);
106#endif
107}
108
109@compute @workgroup_size(WG_SIZE)
110fn main(
111    @builtin(local_invocation_id) local_id: vec3<u32>,
112    @builtin(workgroup_id) wg_id: vec3<u32>,
113    @builtin(num_workgroups) num_wg: vec3<u32>
114#ifdef USE_SUBGROUP_REDUCTION
115  , @builtin(subgroup_id) subgroup_id: u32,
116    @builtin(subgroup_invocation_id) subgroup_invocation_id: u32,
117    @builtin(num_subgroups) num_subgroups: u32
118#endif
119) {
120    let tid = local_id.x;
121    let wg_linear = wg_id.y * num_wg.x + wg_id.x;
122
123    let i1 = wg_linear % params.d_inner;
124    let head_seq = wg_linear / params.d_inner;
125    let ir = head_seq % params.n_head;
126    let i3 = head_seq / params.n_head;
127    let n_seqs = params.y_elems / (params.n_seq_tokens * params.n_head * params.d_inner);
128
129    let state_slot = read_state_slot(i3);
130    let g = ir / (params.n_head / params.n_group);
131
132    let s_idx = params.offset_s + tid + i1 * params.stride_s1 + ir * params.stride_s2 + state_slot * params.stride_s3;
133    var s_prev = s_in[s_idx];
134
135    let a_idx = params.offset_A + (tid % params.a_ne0) + ir * params.stride_A1;
136#ifdef A_OVERLAP
137    let A0 = read_merged_f32(a_idx);
138#else
139    let A0 = A[a_idx];
140#endif
141
142    for (var token_base = 0u; token_base < params.n_seq_tokens; token_base += TOKENS_PER_TILE) {
143        if (tid < TOKENS_PER_TILE) {
144            let token = token_base + tid;
145            if (token < params.n_seq_tokens) {
146                let x_idx = params.offset_x + i1 + ir * params.stride_x1 + token * params.stride_x2 + i3 * params.stride_x3;
147                let dt_idx = params.offset_dt + ir + token * params.stride_dt1 + i3 * params.stride_dt2;
148#ifdef XBC_OVERLAP
149                let dt0 = read_merged_f32(dt_idx);
150#else
151                let dt0 = dt[dt_idx];
152#endif
153                let dtsp = select(log(1.0 + exp(dt0)), dt0, dt0 > 20.0);
154                shared_dtsp[tid] = dtsp;
155#ifdef XBC_OVERLAP
156                shared_x_dt[tid] = read_merged_f32(x_idx) * dtsp;
157#else
158                shared_x_dt[tid] = x[x_idx] * dtsp;
159#endif
160            }
161        }
162
163        workgroupBarrier();
164
165        for (var token_in_tile = 0u; token_in_tile < TOKENS_PER_TILE; token_in_tile++) {
166            let token = token_base + token_in_tile;
167            if (token >= params.n_seq_tokens) {
168                break;
169            }
170
171            let x_dt = shared_x_dt[token_in_tile];
172            let dA = exp(shared_dtsp[token_in_tile] * A0);
173            let reduce_idx = reduce_base(token_in_tile) + tid;
174
175            let b_idx = params.offset_B + tid + g * params.stride_B1 + token * params.stride_B2 + i3 * params.stride_B3;
176            let c_idx = params.offset_C + tid + g * params.stride_C1 + token * params.stride_C2 + i3 * params.stride_C3;
177#ifdef XBC_OVERLAP
178            let s = s_prev * dA + read_merged_f32(b_idx) * x_dt;
179#else
180            let s = s_prev * dA + B[b_idx] * x_dt;
181#endif
182            s_prev = s;
183
184            let slot = params.n_seq_tokens - 1u - token;
185            if (slot > 0u && slot < params.K) {
186                let snapshot_idx =
187                    params.offset_dst + params.y_elems + tid + i1 * params.d_state +
188                    ir * (params.d_state * params.d_inner) +
189                    (slot * n_seqs + i3) * (params.d_state * params.d_inner * params.n_head);
190                dst[snapshot_idx] = s;
191            }
192
193#ifdef USE_SUBGROUP_REDUCTION
194#ifdef XBC_OVERLAP
195            let subgroup_partial = subgroupAdd(s * read_merged_f32(c_idx));
196#else
197            let subgroup_partial = subgroupAdd(s * C[c_idx]);
198#endif
199            if (subgroup_invocation_id == 0u) {
200                shared_reduce[reduce_idx - tid + subgroup_id] = subgroup_partial;
201            }
202#else
203#ifdef XBC_OVERLAP
204            shared_reduce[reduce_idx] = s * read_merged_f32(c_idx);
205#else
206            shared_reduce[reduce_idx] = s * C[c_idx];
207#endif
208#endif
209
210            workgroupBarrier();
211
212#ifdef USE_SUBGROUP_REDUCTION
213            if (tid == 0u) {
214                var sum = 0.0;
215                for (var sg = 0u; sg < num_subgroups; sg++) {
216                    sum += shared_reduce[reduce_base(token_in_tile) + sg];
217                }
218                let y_idx =
219                    params.offset_dst + i1 + ir * params.d_inner + token * (params.n_head * params.d_inner) +
220                    i3 * (params.n_seq_tokens * params.n_head * params.d_inner);
221                dst[y_idx] = sum;
222            }
223#else
224            for (var stride = WG_SIZE / 2u; stride > 0u; stride >>= 1u) {
225                if (tid < stride) {
226                    shared_reduce[reduce_idx] += shared_reduce[reduce_idx + stride];
227                }
228                workgroupBarrier();
229            }
230
231            if (tid == 0u) {
232                let y_idx =
233                    params.offset_dst + i1 + ir * params.d_inner + token * (params.n_head * params.d_inner) +
234                    i3 * (params.n_seq_tokens * params.n_head * params.d_inner);
235                dst[y_idx] = shared_reduce[reduce_base(token_in_tile)];
236            }
237#endif
238
239            workgroupBarrier();
240        }
241    }
242
243    let state_idx =
244        params.offset_dst + params.y_elems + tid + i1 * params.d_state + ir * (params.d_state * params.d_inner) +
245        i3 * (params.d_state * params.d_inner * params.n_head);
246    dst[state_idx] = s_prev;
247}