char-fork/llama.cpp

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

Hongqiang Wangopencl: port fused ssm_scan kernel (Mamba-2, d_state in {128, 256}) to GPU (#26439)b062ba735

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
8.7 KiB216 linesraw
1// Mamba2 fused SSM scan kernel. One workgroup per (head, dim, seq); WG size =
2// 64 threads. Each thread owns c_factor = d_state/64 state elements in
3// private registers; the state stays resident across the n_tokens t-loop
4//
5// References:
6//   ggml/src/ggml-cuda/ssm-scan.cu:117 ssm_scan_f32_group
7//   ggml/src/ggml-cpu/ops.cpp:9368 ggml_compute_forward_ssm_scan_f32
8
9#pragma OPENCL EXTENSION cl_khr_fp16 : enable
10
11#ifdef cl_khr_subgroups
12#pragma OPENCL EXTENSION cl_khr_subgroups : enable
13#endif
14
15#if defined(cl_qcom_reqd_sub_group_size)
16#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
17#define REQD_SUBGROUP_SIZE_64 __attribute__((qcom_reqd_sub_group_size("half")))
18#else
19#define REQD_SUBGROUP_SIZE_64
20#endif
21
22inline float softplus_f32(float x) {
23    return (x <= 20.0f) ? log(1.0f + exp(x)) : x;
24}
25
26// d_state = 128 (most Mamba-2 models, e.g. mamba2-2.7B, Codestral-Mamba).
27// WG = 64 threads, each holds 2 state elements (tid and tid+64).
28REQD_SUBGROUP_SIZE_64
29kernel void kernel_ssm_scan_f32_mamba2_d128(
30    global const char * src0_base, ulong src0_off,
31    global const char * src1_base, ulong src1_off,
32    global const char * src2_base, ulong src2_off,
33    global const char * src3_base, ulong src3_off,
34    global const char * src4_base, ulong src4_off,
35    global const char * src5_base, ulong src5_off,
36    global const char * src6_base, ulong src6_off,
37    global       char * dst_base,  ulong dst_off,
38    ulong s0_nb2, ulong s0_nb3,
39    ulong x_nb2,  ulong x_nb3,
40    ulong dt_nb1, ulong dt_nb2,
41    ulong A_nb1,
42    ulong B_nb2,  ulong B_nb3,
43    ulong C_nb2,  ulong C_nb3,
44    ulong s_off_bytes,
45    int   head_dim, int n_head, int n_group, int n_tokens
46) {
47    const int d_state = 128;
48
49    const int tid     = (int) get_local_id(0);
50    const int wg_x    = (int) get_group_id(0);
51    const int seq_id  = (int) get_group_id(1);
52
53    const int head_id = wg_x / head_dim;
54    const int dim_id  = wg_x - head_id * head_dim;
55    const int g       = head_id / (n_head / n_group);
56
57    src0_base += src0_off;
58    src1_base += src1_off;
59    src2_base += src2_off;
60    src3_base += src3_off;
61    src4_base += src4_off;
62    src5_base += src5_off;
63    src6_base += src6_off;
64    dst_base  += dst_off;
65
66    const int seq_slot = ((global const int *) src6_base)[seq_id];
67
68    const ulong state_base_off = (ulong)seq_slot * s0_nb3 + (ulong)head_id * s0_nb2
69                                + (ulong)dim_id * d_state * sizeof(float);
70    global const float * s0_warp = (global const float *)(src0_base + state_base_off);
71    const ulong state_out_off = (ulong)seq_id * s0_nb3 + (ulong)head_id * s0_nb2
72                              + (ulong)dim_id * d_state * sizeof(float);
73    global float * s_warp = (global float *)(dst_base + s_off_bytes + state_out_off);
74
75    global const char * x_seq  = src1_base + (ulong)seq_id * x_nb3;
76    global const char * dt_seq = src2_base + (ulong)seq_id * dt_nb2;
77    global const char * B_seq  = src4_base + (ulong)seq_id * B_nb3 + (ulong)g * d_state * sizeof(float);
78    global const char * C_seq  = src5_base + (ulong)seq_id * C_nb3 + (ulong)g * d_state * sizeof(float);
79
80    const ulong y_dim_total = (ulong)n_head * head_dim;
81    global float * y_seq = (global float *)dst_base
82                           + (ulong)seq_id * (ulong)n_tokens * y_dim_total;
83
84    const float A_val = ((global const float *)src3_base)[(ulong)head_id * A_nb1 / sizeof(float)];
85
86    // c_factor = 2: each thread owns 2 state elements (tid and tid+64).
87    float state0 = s0_warp[tid];
88    float state1 = s0_warp[tid + 64];
89
90    for (int t = 0; t < n_tokens; ++t) {
91        const float dt_h        = ((global const float *)(dt_seq + (ulong)t * dt_nb1))[head_id];
92        const float dt_softplus = softplus_f32(dt_h);
93        const float dA          = exp(dt_softplus * A_val);
94        const float x_val       = ((global const float *)(x_seq + (ulong)t * x_nb2))[(ulong)head_id * head_dim + dim_id];
95        const float x_dt        = x_val * dt_softplus;
96
97        const float B0 = ((global const float *)(B_seq + (ulong)t * B_nb2))[tid];
98        const float B1 = ((global const float *)(B_seq + (ulong)t * B_nb2))[tid + 64];
99        const float C0 = ((global const float *)(C_seq + (ulong)t * C_nb2))[tid];
100        const float C1 = ((global const float *)(C_seq + (ulong)t * C_nb2))[tid + 64];
101
102        state0 = state0 * dA + B0 * x_dt;
103        state1 = state1 * dA + B1 * x_dt;
104        const float partial = state0 * C0 + state1 * C1;
105
106        const float sum = sub_group_reduce_add(partial);
107        if (tid == 0) {
108            y_seq[(ulong)t * y_dim_total + (ulong)head_id * head_dim + dim_id] = sum;
109        }
110    }
111
112    s_warp[tid]      = state0;
113    s_warp[tid + 64] = state1;
114}
115
116// d_state = 256 (Falcon-H1). WG = 64 threads, each holds 4 state elements.
117REQD_SUBGROUP_SIZE_64
118kernel void kernel_ssm_scan_f32_mamba2_d256(
119    global const char * src0_base, ulong src0_off,
120    global const char * src1_base, ulong src1_off,
121    global const char * src2_base, ulong src2_off,
122    global const char * src3_base, ulong src3_off,
123    global const char * src4_base, ulong src4_off,
124    global const char * src5_base, ulong src5_off,
125    global const char * src6_base, ulong src6_off,
126    global       char * dst_base,  ulong dst_off,
127    ulong s0_nb2, ulong s0_nb3,
128    ulong x_nb2,  ulong x_nb3,
129    ulong dt_nb1, ulong dt_nb2,
130    ulong A_nb1,
131    ulong B_nb2,  ulong B_nb3,
132    ulong C_nb2,  ulong C_nb3,
133    ulong s_off_bytes,
134    int   head_dim, int n_head, int n_group, int n_tokens
135) {
136    const int d_state = 256;
137
138    const int tid     = (int) get_local_id(0);
139    const int wg_x    = (int) get_group_id(0);
140    const int seq_id  = (int) get_group_id(1);
141
142    const int head_id = wg_x / head_dim;
143    const int dim_id  = wg_x - head_id * head_dim;
144    const int g       = head_id / (n_head / n_group);
145
146    src0_base += src0_off;
147    src1_base += src1_off;
148    src2_base += src2_off;
149    src3_base += src3_off;
150    src4_base += src4_off;
151    src5_base += src5_off;
152    src6_base += src6_off;
153    dst_base  += dst_off;
154
155    const int seq_slot = ((global const int *) src6_base)[seq_id];
156
157    const ulong state_base_off = (ulong)seq_slot * s0_nb3 + (ulong)head_id * s0_nb2
158                                + (ulong)dim_id * d_state * sizeof(float);
159    global const float * s0_warp = (global const float *)(src0_base + state_base_off);
160    const ulong state_out_off = (ulong)seq_id * s0_nb3 + (ulong)head_id * s0_nb2
161                              + (ulong)dim_id * d_state * sizeof(float);
162    global float * s_warp = (global float *)(dst_base + s_off_bytes + state_out_off);
163
164    global const char * x_seq  = src1_base + (ulong)seq_id * x_nb3;
165    global const char * dt_seq = src2_base + (ulong)seq_id * dt_nb2;
166    global const char * B_seq  = src4_base + (ulong)seq_id * B_nb3 + (ulong)g * d_state * sizeof(float);
167    global const char * C_seq  = src5_base + (ulong)seq_id * C_nb3 + (ulong)g * d_state * sizeof(float);
168
169    const ulong y_dim_total = (ulong)n_head * head_dim;
170    global float * y_seq = (global float *)dst_base
171                           + (ulong)seq_id * (ulong)n_tokens * y_dim_total;
172
173    const float A_val = ((global const float *)src3_base)[(ulong)head_id * A_nb1 / sizeof(float)];
174
175    // c_factor = 4: each thread owns 4 state elements.
176    float state0 = s0_warp[tid];
177    float state1 = s0_warp[tid + 64];
178    float state2 = s0_warp[tid + 128];
179    float state3 = s0_warp[tid + 192];
180
181    for (int t = 0; t < n_tokens; ++t) {
182        const float dt_h        = ((global const float *)(dt_seq + (ulong)t * dt_nb1))[head_id];
183        const float dt_softplus = softplus_f32(dt_h);
184        const float dA          = exp(dt_softplus * A_val);
185        const float x_val       = ((global const float *)(x_seq + (ulong)t * x_nb2))[(ulong)head_id * head_dim + dim_id];
186        const float x_dt        = x_val * dt_softplus;
187
188        global const float * B_t = (global const float *)(B_seq + (ulong)t * B_nb2);
189        global const float * C_t = (global const float *)(C_seq + (ulong)t * C_nb2);
190
191        const float B0 = B_t[tid];
192        const float B1 = B_t[tid + 64];
193        const float B2 = B_t[tid + 128];
194        const float B3 = B_t[tid + 192];
195        const float C0 = C_t[tid];
196        const float C1 = C_t[tid + 64];
197        const float C2 = C_t[tid + 128];
198        const float C3 = C_t[tid + 192];
199
200        state0 = state0 * dA + B0 * x_dt;
201        state1 = state1 * dA + B1 * x_dt;
202        state2 = state2 * dA + B2 * x_dt;
203        state3 = state3 * dA + B3 * x_dt;
204        const float partial = state0 * C0 + state1 * C1 + state2 * C2 + state3 * C3;
205
206        const float sum = sub_group_reduce_add(partial);
207        if (tid == 0) {
208            y_seq[(ulong)t * y_dim_total + (ulong)head_id * head_dim + dim_id] = sum;
209        }
210    }
211
212    s_warp[tid]       = state0;
213    s_warp[tid + 64]  = state1;
214    s_warp[tid + 128] = state2;
215    s_warp[tid + 192] = state3;
216}