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 →
9.0 KiB249 linesraw
1#pragma OPENCL EXTENSION cl_khr_subgroups : enable
2
3#ifdef cl_intel_required_subgroup_size
4#pragma OPENCL EXTENSION cl_intel_required_subgroup_size : enable
5#define INTEL_GPU 1
6#define REQD_SUBGROUP_SIZE_16 __attribute__((intel_reqd_sub_group_size(16)))
7#define REQD_SUBGROUP_SIZE_32 __attribute__((intel_reqd_sub_group_size(32)))
8#elif defined(cl_qcom_reqd_sub_group_size)
9#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
10#define ADRENO_GPU 1
11#define REQD_SUBGROUP_SIZE_64  __attribute__((qcom_reqd_sub_group_size("half")))
12#define REQD_SUBGROUP_SIZE_128 __attribute__((qcom_reqd_sub_group_size("full")))
13#endif
14
15#ifndef S_V
16#define S_V 128
17#endif
18#ifndef KDA
19#define KDA 0
20#endif
21#ifndef SUBGROUP_SIZE
22#define SUBGROUP_SIZE 64
23#endif
24#ifndef LANES_PER_COLUMN
25#define LANES_PER_COLUMN 8
26#endif
27#ifndef COLS_PER_LANE_GROUP
28#define COLS_PER_LANE_GROUP 1
29#endif
30#ifndef SUBGROUPS_PER_WG
31#define SUBGROUPS_PER_WG 1
32#endif
33#ifndef USE_QCOM_SUBGROUP_SHUFFLE
34#define USE_QCOM_SUBGROUP_SHUFFLE 0
35#endif
36
37#define WG_SIZE             (SUBGROUP_SIZE * SUBGROUPS_PER_WG)
38#define LANE_GROUPS_PER_SG  (SUBGROUP_SIZE / LANES_PER_COLUMN)
39#define COLS_PER_SG         (LANE_GROUPS_PER_SG * COLS_PER_LANE_GROUP)
40#define COLS_PER_WG         (SUBGROUPS_PER_WG * COLS_PER_SG)
41#define ROWS_PER_LANE       (S_V / LANES_PER_COLUMN)
42
43#if USE_QCOM_SUBGROUP_SHUFFLE
44#pragma OPENCL EXTENSION cl_qcom_subgroup_shuffle : enable
45#endif
46
47// XOR-based parallel sum
48// This does a reduction across groups of LANES_PER_COLUMN
49static inline float reduce_add_shmem(float partial, __local float * temp, uint lane) {
50#if USE_QCOM_SUBGROUP_SHUFFLE
51   #pragma unroll
52    for (uint s = LANES_PER_COLUMN / 2u; s > 0u; s >>= 1u) {
53        partial += qcom_sub_group_shuffle_xor(partial, s, CLK_SUB_GROUP_SHUFFLE_WIDTH_WAVE_SIZE_QCOM, partial);
54    }
55    return partial;
56#else
57    temp[lane] = partial;
58    sub_group_barrier(CLK_LOCAL_MEM_FENCE);
59    #pragma unroll
60    for (uint s = LANES_PER_COLUMN / 2u; s > 0u; s >>= 1u) {
61        float other = temp[lane ^ s];
62        sub_group_barrier(CLK_LOCAL_MEM_FENCE);
63        temp[lane] += other;
64        sub_group_barrier(CLK_LOCAL_MEM_FENCE);
65    }
66    const float result = temp[lane];
67    sub_group_barrier(CLK_LOCAL_MEM_FENCE);
68    return result;
69#endif
70}
71
72#define REDUCE_PARTIAL(partial, temp_ptr, lid) \
73    ((LANES_PER_COLUMN == 1u) ? (partial) : reduce_add_shmem((partial), (temp_ptr), (lid)))
74
75// force compiler to optimize kernel for a specific fixed work-group size
76__attribute__((reqd_work_group_size(WG_SIZE, 1, 1)))
77#ifdef INTEL_GPU
78REQD_SUBGROUP_SIZE_32
79#elif defined (ADRENO_GPU)
80REQD_SUBGROUP_SIZE_64
81#endif
82kernel void kernel_gated_delta_net(
83        global const char * q_buf,     ulong off_q,
84        global const char * k_buf,     ulong off_k,
85        global const char * v_buf,     ulong off_v,
86        global const char * g_buf,     ulong off_g,
87        global const char * beta_buf,  ulong off_beta,
88        global const char * state_buf, ulong off_state,
89        global       char * dst_buf,   ulong off_dst,
90        uint  H_v,
91        uint  n_tokens,
92        uint  n_seqs,
93        uint  s_off,
94        uint  sq1, uint sq2, uint sq3,
95        uint  sv1, uint sv2, uint sv3,
96        uint  sb1, uint sb2, uint sb3,
97        uint  H_k,
98        uint  rq3,
99        float scale,
100        uint K) {
101
102    global const float * data_q     = (global const float *)(q_buf     + off_q);
103    global const float * data_k     = (global const float *)(k_buf     + off_k);
104    global const float * data_v     = (global const float *)(v_buf     + off_v);
105    global const float * data_g     = (global const float *)(g_buf     + off_g);
106    global const float * data_beta  = (global const float *)(beta_buf  + off_beta);
107    global const float * data_state = (global const float *)(state_buf + off_state);
108    global       float * data_dst   = (global       float *)(dst_buf   + off_dst);
109
110    const uint head_id     = get_group_id(0);
111    const uint seq_id      = get_group_id(1);
112    const uint tid         = (uint)get_local_id(0);
113
114    const uint sg_id       = get_sub_group_id(); // subgroup id
115    const uint sg_lid      = get_sub_group_local_id(); // subgroup lane id
116
117    const uint lane        = sg_lid % LANES_PER_COLUMN;
118    const uint lane_group  = sg_lid / LANES_PER_COLUMN;
119    const uint wg_col_base = get_group_id(2) * COLS_PER_WG;
120    const uint sg_col_base = wg_col_base + sg_id * COLS_PER_SG;
121
122    const uint iq1 = head_id % H_k; // head index for Q and K
123    const uint iq3 = seq_id / rq3; // seq index for Q and K
124
125    const uint state_size = S_V * S_V;
126    // input state holds s0 only [S_v, S_v, H, n_seqs]: per-seq stride is H*D.
127    const uint state_base = (seq_id * H_v + head_id) * state_size;
128    const uint q_off_base  = iq3 * sq3 + iq1 * sq1;
129    const uint v_off_base  = seq_id * sv3 + head_id * sv1;
130    const uint gb_off_base = seq_id * sb3 + head_id * sb1;
131    const uint state_out_base      = (seq_id * H_v + head_id) * state_size;
132    const uint state_size_per_snap = state_size * H_v * n_seqs;
133
134    __local float reduce_temp[WG_SIZE];
135    __local float * temp_ptr = reduce_temp + sg_id * SUBGROUP_SIZE;
136
137    float s_shard[COLS_PER_LANE_GROUP][ROWS_PER_LANE];
138    #pragma unroll
139    for (uint cg = 0; cg < COLS_PER_LANE_GROUP; cg++) {
140        const uint col = sg_col_base + cg * LANE_GROUPS_PER_SG + lane_group;
141        #pragma unroll
142        for (uint r = 0; r < ROWS_PER_LANE; r++) {
143            s_shard[cg][r] = data_state[state_base + col * S_V + r * LANES_PER_COLUMN + lane];
144        }
145    }
146
147    // snapshot slot mapping: slot 0 = most recent state, slot s = s tokens back.
148    // When n_tokens < K only slots 0..n_tokens-1 are written; older slots are caller-owned.
149    uint attn_off = (seq_id * n_tokens * H_v + head_id) * S_V;
150
151    for (uint t = 0; t < n_tokens; t++) {
152        const uint  q_off    = q_off_base + t * sq2;
153        const uint  k_off    = q_off;
154        const uint  v_off    = v_off_base + t * sv2;
155        const uint  gb_off   = gb_off_base + t * sb2;
156        const float beta_val = data_beta[gb_off];
157
158        float k_reg[ROWS_PER_LANE];
159        float q_reg[ROWS_PER_LANE];
160#if KDA
161        float g_exp[ROWS_PER_LANE];
162        #pragma unroll
163        for (uint r = 0; r < ROWS_PER_LANE; r++) {
164            const uint i = r * LANES_PER_COLUMN + lane;
165            k_reg[r] = data_k[k_off + i];
166            q_reg[r] = data_q[q_off + i];
167            g_exp[r] = exp(data_g[gb_off * S_V + i]);
168        }
169#else
170        const float g_val = exp(data_g[gb_off]);
171
172        #pragma unroll
173        for (uint r = 0; r < ROWS_PER_LANE; r++) {
174            const uint i = r * LANES_PER_COLUMN + lane;
175            k_reg[r] = data_k[k_off + i];
176            q_reg[r] = data_q[q_off + i];
177        }
178#endif
179
180        #pragma unroll
181        for (uint cg = 0; cg < COLS_PER_LANE_GROUP; cg++) {
182            const uint col = sg_col_base + cg * LANE_GROUPS_PER_SG + lane_group;
183            float v_val = data_v[v_off + col];
184
185            float kv_shard = 0.0f;
186            #pragma unroll
187            for (uint r = 0; r < ROWS_PER_LANE; r++) {
188#if KDA
189                float gs = g_exp[r] * s_shard[cg][r];
190                kv_shard += gs * k_reg[r];
191#else
192                kv_shard += s_shard[cg][r] * k_reg[r];
193#endif
194            }
195
196#if !KDA
197            kv_shard *= g_val; // Applied once instead of ROWS_PER_LANE times
198#endif
199
200            const float kv_col = REDUCE_PARTIAL(kv_shard, temp_ptr, sg_lid);
201
202            const float delta_col = (v_val - kv_col) * beta_val;
203
204            float attn_partial = 0.0f;
205            #pragma unroll
206            for (uint r = 0; r < ROWS_PER_LANE; r++) {
207#if KDA
208                float gs = g_exp[r] * s_shard[cg][r];
209#else
210                float gs = g_val * s_shard[cg][r];
211#endif
212                s_shard[cg][r] = gs + k_reg[r] * delta_col;
213                attn_partial += s_shard[cg][r] * q_reg[r];
214            }
215            const float attn_col = REDUCE_PARTIAL(attn_partial, temp_ptr, sg_lid);
216
217            if (lane == 0) {
218                data_dst[attn_off + col] = attn_col * scale;
219            }
220        }
221        attn_off += S_V * H_v;
222
223        if (K > 1u) {
224            const int target_slot = (int)n_tokens - 1 - (int)t;
225            if (target_slot >= 0 && target_slot < (int)K) {
226                #pragma unroll
227                for (uint cg = 0; cg < COLS_PER_LANE_GROUP; cg++) {
228                    const uint col = sg_col_base + cg * LANE_GROUPS_PER_SG + lane_group;
229                    const uint slot_base = s_off + (uint)target_slot * state_size_per_snap + state_out_base;
230                    #pragma unroll
231                    for (uint r = 0; r < ROWS_PER_LANE; r++) {
232                        data_dst[slot_base + col * S_V + r * LANES_PER_COLUMN + lane] = s_shard[cg][r];
233                    }
234                }
235            }
236        }
237    }
238
239    if (K == 1u) {
240        #pragma unroll
241        for (uint cg = 0; cg < COLS_PER_LANE_GROUP; cg++) {
242            const uint col = sg_col_base + cg * LANE_GROUPS_PER_SG + lane_group;
243            #pragma unroll
244            for (uint r = 0; r < ROWS_PER_LANE; r++) {
245                data_dst[s_off + state_base + col * S_V + r * LANES_PER_COLUMN + lane] = s_shard[cg][r];
246            }
247        }
248    }
249}