char-fork/llama.cpp

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

fairydreamingwebgpu : reorder includes since V that appears in common_decls.tmpl may be defined as K in flash_attn_decls.tmpl if KV_OVERLAP (#27545)c1c766da5

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
9.0 KiB240 linesraw
1enable f16;
2enable subgroups;
3
4#define BYTE_HELPERS
5#include "flash_attn_decls.tmpl"
6#include "common_decls.tmpl"
7
8// Default values
9// The actual values are defined in shader-lib.
10#define HEAD_DIM_QK 64
11#define HEAD_DIM_V 64
12#define Q_TILE 4
13#define KV_TILE 64
14#define WG_SIZE 128
15
16const Q_CHUNKS: u32 = HEAD_DIM_QK / 4u;
17const V_CHUNKS: u32 = HEAD_DIM_V / 4u;
18const SCORE_REGS_PER_LANE: u32 = (KV_TILE + MIN_SUBGROUP_SIZE - 1u) / MIN_SUBGROUP_SIZE;
19const OUT_REGS_PER_LANE: u32 = (V_CHUNKS + MIN_SUBGROUP_SIZE - 1u) / MIN_SUBGROUP_SIZE;
20
21#if !defined(K_DIRECT) || !defined(V_DIRECT)
22#define STAGING_SHMEM kv_shmem
23#define STAGING_OUT_TYPE f16
24#include "flash_attn_staging.tmpl"
25const kv_shmem_size = KV_TILE * max(HEAD_DIM_QK, HEAD_DIM_V);
26var<workgroup> kv_shmem: array<f16, kv_shmem_size>;
27#endif
28
29var<workgroup> q_shmem: array<Q_TYPE, Q_TILE * HEAD_DIM_QK>;
30var<workgroup> p_shmem: array<f16, Q_TILE * KV_TILE>;
31
32@compute @workgroup_size(WG_SIZE)
33fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
34        @builtin(local_invocation_id) local_id: vec3<u32>,
35        @builtin(subgroup_id) subgroup_id: u32,
36        @builtin(subgroup_size) subgroup_size: u32,
37        @builtin(num_subgroups) num_subgroups: u32,
38        @builtin(subgroup_invocation_id) sg_inv_id: u32) {
39    if (subgroup_size == 0u || num_subgroups < Q_TILE) {
40        return;
41    }
42
43    let wg_per_head = (params.seq_len_q + Q_TILE - 1u) / Q_TILE;
44    let wg_per_batch = wg_per_head * params.n_heads;
45
46    let dst2_stride = HEAD_DIM_V * params.n_heads;
47    let dst3_stride = dst2_stride * params.seq_len_q;
48
49    let batch_idx = wg_id.x / wg_per_batch;
50    let q_batch_offset = params.offset_q + batch_idx * params.stride_q3;
51    let k_batch_offset = params.offset_k + batch_idx * params.stride_k3;
52    let v_batch_offset = params.offset_v + batch_idx * params.stride_v3;
53    let dst_batch_offset = params.offset_dst + batch_idx * dst3_stride;
54    let wg_in_batch = wg_id.x % wg_per_batch;
55
56    let head_idx = wg_in_batch / wg_per_head;
57    let q_head_offset = q_batch_offset + head_idx * params.stride_q2;
58    let k_head_idx = head_idx / params.q_per_kv;
59    let v_head_offset = v_batch_offset + k_head_idx * params.stride_v2;
60    let k_head_offset = k_batch_offset + k_head_idx * params.stride_k2;
61
62    let wg_in_head = wg_in_batch % wg_per_head;
63    let q_row_start = wg_in_head * Q_TILE;
64    let global_q_row = q_row_start + subgroup_id;
65    let row_active = subgroup_id < Q_TILE && global_q_row < params.seq_len_q;
66
67#ifdef MASK
68    let mask_global_offset = params.offset_mask + batch_idx * params.stride_mask3 + q_row_start * params.seq_len_kv;
69#endif
70
71    let dst_global_offset = dst_batch_offset + q_row_start * dst2_stride + head_idx * HEAD_DIM_V;
72
73    let head = f32(head_idx);
74    let slope = select(1.0,
75        select(pow(params.m1, 2.0 * (head - params.n_head_log2) + 1.0),
76                pow(params.m0, head + 1.0),
77                head < params.n_head_log2),
78        params.max_bias > 0.0);
79
80    for (var elem_idx = local_id.x; elem_idx < Q_TILE * HEAD_DIM_QK; elem_idx += WG_SIZE) {
81        let q_tile_row = elem_idx / HEAD_DIM_QK;
82        let q_col = elem_idx % HEAD_DIM_QK;
83        let head_q_row = q_row_start + q_tile_row;
84        let global_q_row_offset = q_head_offset + head_q_row * params.stride_q1;
85        q_shmem[elem_idx] = select(
86            0.0,
87            Q_TYPE(Q[global_q_row_offset + q_col]) * params.scale,
88            head_q_row < params.seq_len_q);
89    }
90
91    workgroupBarrier();
92
93    var row_max = FLOAT_MIN;
94    var exp_sum = 0.0;
95    var out_regs: array<vec4<f32>, OUT_REGS_PER_LANE>;
96    for (var reg_idx = 0u; reg_idx < OUT_REGS_PER_LANE; reg_idx += 1u) {
97        out_regs[reg_idx] = vec4<f32>(0.0);
98    }
99
100    let q_base = subgroup_id * HEAD_DIM_QK;
101    let subgroup_p_offset = subgroup_id * KV_TILE;
102
103    for (var kv_tile = 0u; kv_tile < params.seq_len_kv; kv_tile += KV_TILE) {
104        let kv_count = min(KV_TILE, params.seq_len_kv - kv_tile);
105        let score_slots = min(SCORE_REGS_PER_LANE, (kv_count + subgroup_size - 1u) / subgroup_size);
106        let out_slots = min(OUT_REGS_PER_LANE, (V_CHUNKS + subgroup_size - 1u) / subgroup_size);
107        var local_scores: array<f32, SCORE_REGS_PER_LANE>;
108        for (var slot = 0u; slot < SCORE_REGS_PER_LANE; slot += 1u) {
109            local_scores[slot] = FLOAT_MIN;
110        }
111
112        // The tile path stages K/V in shared memory so each tile can be reused across
113        // Q_TILE query rows. It therefore does not use the direct path.
114#ifndef K_DIRECT
115        load_k_tile_block(local_id.x, kv_count, kv_tile, k_head_offset);
116#endif
117
118        workgroupBarrier();
119
120        var local_max = FLOAT_MIN;
121        if (row_active) {
122            for (var slot = 0u; slot < score_slots; slot += 1u) {
123                let kv_local = sg_inv_id + slot * subgroup_size;
124                if (kv_local >= kv_count) {
125                    continue;
126                }
127
128                let global_k_row = kv_tile + kv_local;
129                var dot_val = 0.0;
130                for (var chunk = 0u; chunk < Q_CHUNKS; chunk += 1u) {
131                    let q_off = q_base + chunk * 4u;
132                    let qv = vec4<Q_TYPE>(
133                        q_shmem[q_off + 0u],
134                        q_shmem[q_off + 1u],
135                        q_shmem[q_off + 2u],
136                        q_shmem[q_off + 3u]);
137                    let kv_off = kv_local * HEAD_DIM_QK + chunk * 4u;
138                    let kv = vec4<f16>(
139                        kv_shmem[kv_off + 0u],
140                        kv_shmem[kv_off + 1u],
141                        kv_shmem[kv_off + 2u],
142                        kv_shmem[kv_off + 3u]);
143                    dot_val += dot(vec4<f32>(qv), vec4<f32>(kv));
144                }
145#ifdef LOGIT_SOFTCAP
146                dot_val = params.logit_softcap * tanh(dot_val);
147#endif
148#ifdef MASK
149                let mask_idx = mask_global_offset + subgroup_id * params.seq_len_kv + global_k_row;
150                dot_val += slope * f32(mask[mask_idx]);
151#endif
152                local_scores[slot] = dot_val;
153                local_max = max(local_max, dot_val);
154            }
155        }
156
157        let tile_max = subgroupMax(local_max);
158        let new_max = max(row_max, tile_max);
159        let cur_exp = exp(row_max - new_max);
160        exp_sum *= cur_exp;
161        for (var reg_idx = 0u; reg_idx < OUT_REGS_PER_LANE; reg_idx += 1u) {
162            out_regs[reg_idx] *= cur_exp;
163        }
164
165        var local_sum = 0.0;
166        for (var slot = 0u; slot < score_slots; slot += 1u) {
167            let kv_local = sg_inv_id + slot * subgroup_size;
168            if (row_active && kv_local < kv_count) {
169                let p = exp(local_scores[slot] - new_max);
170                p_shmem[subgroup_p_offset + kv_local] = f16(p);
171                local_sum += p;
172            }
173        }
174
175        workgroupBarrier();
176
177        // The tile path stages K/V in shared memory so each tile can be reused across
178        // Q_TILE query rows. It therefore does not use the direct path.
179#ifndef V_DIRECT
180        load_v_tile_block(local_id.x, kv_count, kv_tile, v_head_offset);
181#endif
182
183        workgroupBarrier();
184
185        let tile_sum = subgroupAdd(local_sum);
186        exp_sum += tile_sum;
187        row_max = new_max;
188
189        if (row_active) {
190            for (var reg_idx = 0u; reg_idx < out_slots; reg_idx += 1u) {
191                let chunk = sg_inv_id + reg_idx * subgroup_size;
192                if (chunk >= V_CHUNKS) {
193                    continue;
194                }
195
196                var acc = out_regs[reg_idx];
197                for (var kv_local = 0u; kv_local < kv_count; kv_local += 1u) {
198                    let p = f32(p_shmem[subgroup_p_offset + kv_local]);
199                    let kv_off = kv_local * HEAD_DIM_V + chunk * 4u;
200                    let v4 = vec4<f16>(
201                        kv_shmem[kv_off + 0u],
202                        kv_shmem[kv_off + 1u],
203                        kv_shmem[kv_off + 2u],
204                        kv_shmem[kv_off + 3u]);
205                    acc += p * vec4<f32>(v4);
206                }
207                out_regs[reg_idx] = acc;
208            }
209        }
210
211        workgroupBarrier();
212    }
213
214#ifdef SINKS
215    if (row_active) {
216        let sink_score = sinks[params.offset_sinks + head_idx];
217        let sink_max = max(row_max, sink_score);
218        let sink_scale = exp(row_max - sink_max);
219        for (var reg_idx = 0u; reg_idx < OUT_REGS_PER_LANE; reg_idx += 1u) {
220            out_regs[reg_idx] *= sink_scale;
221        }
222        exp_sum = exp_sum * sink_scale + exp(sink_score - sink_max);
223        row_max = sink_max;
224    }
225#endif
226
227    if (row_active) {
228        let inv_exp_sum = select(0.0, 1.0 / exp_sum, exp_sum != 0.0);
229        let row_base = dst_global_offset + subgroup_id * dst2_stride;
230        let out_slots = min(OUT_REGS_PER_LANE, (V_CHUNKS + subgroup_size - 1u) / subgroup_size);
231        for (var reg_idx = 0u; reg_idx < out_slots; reg_idx += 1u) {
232            let chunk = sg_inv_id + reg_idx * subgroup_size;
233            if (chunk >= V_CHUNKS) {
234                continue;
235            }
236            let dst_vec_index = (row_base + chunk * 4u) >> 2u;
237            dst[dst_vec_index] = vec4<DST_TYPE>(out_regs[reg_idx] * inv_exp_sum);
238        }
239    }
240}