char-fork/llama.cpp

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

Gaurav GargRemove padding and multiple D2D copies for MTP (#24086)e95dae18d

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
3.8 KiB149 linesraw
1@group(0) @binding(0)
2var<storage, read_write> src_q: array<f32>;
3
4@group(0) @binding(1)
5var<storage, read_write> src_k: array<f32>;
6
7@group(0) @binding(2)
8var<storage, read_write> src_v: array<f32>;
9
10@group(0) @binding(3)
11var<storage, read_write> src_g: array<f32>;
12
13@group(0) @binding(4)
14var<storage, read_write> src_beta: array<f32>;
15
16@group(0) @binding(5)
17var<storage, read_write> src_state: array<f32>;
18
19@group(0) @binding(6)
20var<storage, read_write> dst: array<f32>;
21
22struct Params {
23    h: u32,
24    n_tokens: u32,
25    n_seqs: u32,
26    s_off: u32,
27
28    sq1: u32,
29    sq2: u32,
30    sq3: u32,
31
32    sv1: u32,
33    sv2: u32,
34    sv3: u32,
35
36    sb1: u32,
37    sb2: u32,
38    sb3: u32,
39
40    neq1: u32,
41    rq3: u32,
42    K: u32,
43    scale: f32,
44};
45
46@group(0) @binding(7)
47var<uniform> params: Params;
48
49var<workgroup> sh_k: array<f32, S_V>;
50var<workgroup> sh_q: array<f32, S_V>;
51var<workgroup> sh_g: array<f32, S_V>;
52
53@compute @workgroup_size(WG_SIZE)
54fn main(
55    @builtin(workgroup_id) workgroup_id: vec3<u32>,
56    @builtin(local_invocation_id) local_id: vec3<u32>
57) {
58    let head_id = workgroup_id.x;
59    let seq_id = workgroup_id.y;
60    let col = local_id.x;
61
62    let iq1 = head_id % params.neq1;
63    let iq3 = seq_id / params.rq3;
64
65    let state_size = S_V * S_V;
66    // input state holds s0 only [S_v, S_v, H, n_seqs]: per-seq stride is H*D.
67    let state_in_base = (seq_id * params.h + head_id) * state_size;
68    let state_out_base = (seq_id * params.h + head_id) * state_size;
69    let state_size_per_snap = state_size * params.h * params.n_seqs;
70
71    var state: array<f32, S_V>;
72    for (var i = 0u; i < S_V; i++) {
73        state[i] = src_state[state_in_base + col * S_V + i];
74    }
75
76    var attn_off = (seq_id * params.n_tokens * params.h + head_id) * S_V;
77
78    for (var t = 0u; t < params.n_tokens; t++) {
79        let q_off = iq3 * params.sq3 + t * params.sq2 + iq1 * params.sq1;
80        let k_off = q_off;
81        let v_off = seq_id * params.sv3 + t * params.sv2 + head_id * params.sv1;
82        let gb_off = seq_id * params.sb3 + t * params.sb2 + head_id * params.sb1;
83
84        sh_q[col] = src_q[q_off + col];
85        sh_k[col] = src_k[k_off + col];
86
87#ifdef KDA
88        let g_base = gb_off * S_V;
89        sh_g[col] = exp(src_g[g_base + col]);
90#endif
91
92        workgroupBarrier();
93
94        let v_val = src_v[v_off + col];
95        let beta_val = src_beta[gb_off];
96
97        var kv_col = 0.0;
98        var delta_col = 0.0;
99        var attn_col = 0.0;
100
101#ifdef KDA
102        for (var i = 0u; i < S_V; i++) {
103            kv_col += (sh_g[i] * state[i]) * sh_k[i];
104        }
105
106        delta_col = (v_val - kv_col) * beta_val;
107
108        for (var i = 0u; i < S_V; i++) {
109            state[i] = sh_g[i] * state[i] + sh_k[i] * delta_col;
110            attn_col += state[i] * sh_q[i];
111        }
112#else
113        let g_val = exp(src_g[gb_off]);
114
115        for (var i = 0u; i < S_V; i++) {
116            kv_col += state[i] * sh_k[i];
117        }
118
119        delta_col = (v_val - g_val * kv_col) * beta_val;
120
121        for (var i = 0u; i < S_V; i++) {
122            state[i] = g_val * state[i] + sh_k[i] * delta_col;
123            attn_col += state[i] * sh_q[i];
124        }
125#endif
126
127        dst[attn_off + col] = attn_col * params.scale;
128        attn_off += S_V * params.h;
129
130        if (params.K > 1u) {
131            // snapshot slot mapping: slot 0 = most recent state, slot s = s tokens back.
132            let target_slot = i32(params.n_tokens) - 1 - i32(t);
133            if (target_slot >= 0 && target_slot < i32(params.K)) {
134                let slot_base = params.s_off + u32(target_slot) * state_size_per_snap + state_out_base;
135                for (var i = 0u; i < S_V; i++) {
136                    dst[slot_base + col * S_V + i] = state[i];
137                }
138            }
139        }
140
141        workgroupBarrier();
142    }
143
144    if (params.K == 1u) {
145        for (var i = 0u; i < S_V; i++) {
146            dst[params.s_off + state_out_base + col * S_V + i] = state[i];
147        }
148    }
149}