public_git
/
whisper.cpp
mirror of https://github.com/ggml-org/whisper.cpp.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

opencl: port fused ssm_scan kernel (Mamba-2, d_state in {128, 256}) to GPU (llama/26439) 

* opencl: port fused ssm_scan kernel (Mamba-2, d_state in {128, 256})

Fold the fused per-token SSM_SCAN recurrent step from opencl/gdn-qwen36-35b
onto the unified base. Previously SSM_SCAN fell back to CPU here; now scalar-A
Mamba-2 with d_state in {128,256}, all-f32, runs on GPU. Other shapes (incl.
Mamba-1 element-wise A) still fall back. test-backend-ops -o SSM_SCAN passes on
Adreno X2-90. opt-out via GGML_OPENCL_DISABLE_SSM_SCAN=1.

* opencl: cleanup

* opencl: require K == 1

---------

Co-authored-by: Li He <lih@qti.qualcomm.com>
This commit is contained in:
Hongqiang Wang 2026-08-19 13:35:17 -07:00 committed by Georgi Gerganov
parent 3d277427a7
commit cd340ea832
3 changed files with 360 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
ggml/src/ggml-opencl/CMakeLists.txt
View File

@ -202,6 +202,7 @@ set(GGML_OPENCL_KERNELS
sqr
sqrt
ssm_conv
ssm_scan
gated_delta_net
sub
sum_rows

143
ggml/src/ggml-opencl/ggml-opencl.cpp
View File

@ -866,6 +866,9 @@ struct ggml_backend_opencl_context {
// [size_idx][kda][tgpp] where size_idx: 0=S_V=16, 1=32, 2=64, 3=128; kda: 0 or 1.
// tgpp 0 = TG variant (COLS_PER_LANE_GROUP=1), tgpp 1 = prefill variant (COLS_PER_LANE_GROUP=4).
cl_kernel kernel_gated_delta_net_f32[4][2][2] = {};
cl_kernel kernel_ssm_scan_f32_mamba2_d128 = nullptr;
cl_kernel kernel_ssm_scan_f32_mamba2_d256 = nullptr;
cl_kernel kernel_timestep_embedding;
cl_kernel kernel_gemv_moe_q4_0_f32_ns, kernel_gemm_moe_q4_0_f32_ns, kernel_gemm_moe_q4_0_f32_ns_bin;
cl_kernel kernel_gemm_moe_q8_0_f32_ns;
@ -3154,6 +3157,24 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
GGML_LOG_CONT(".");
}
// ssm_scan (Mamba-2 fused per-token recurrent step; d_state in {128, 256})
{
#ifdef GGML_OPENCL_EMBED_KERNELS
const std::string kernel_src {
#include "ssm_scan.cl.h"
};
#else
const std::string kernel_src = read_file("ssm_scan.cl");
#endif
cl_program prog =
build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_ssm_scan_f32_mamba2_d128 = clCreateKernel(prog, "kernel_ssm_scan_f32_mamba2_d128", &err), err));
CL_CHECK((backend_ctx->kernel_ssm_scan_f32_mamba2_d256 = clCreateKernel(prog, "kernel_ssm_scan_f32_mamba2_d256", &err), err));
CL_CHECK(clReleaseProgram(prog));
GGML_LOG_CONT(".");
}
// gated_delta_net: one kernel per (S_V, KDA, tgpp) triple.
{
#ifdef GGML_OPENCL_EMBED_KERNELS
@ -7301,6 +7322,23 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te
(op->src[0]->type == GGML_TYPE_F16 && op->src[1]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32);
case GGML_OP_SSM_CONV:
return (op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32);
case GGML_OP_SSM_SCAN: {
// Mamba-2 fused per-token scan. Requires src3->ne[0] == 1 (scalar
// A per head); d_state in {128, 256}; all sources f32. Falls back
// to CPU otherwise (incl. Mamba-1 element-wise A).
for (int i = 0; i < 6; ++i) {
if (op->src[i]->type != GGML_TYPE_F32) {
return false;
}
}
if (op->type != GGML_TYPE_F32) {
return false;
}
const int K = ggml_get_op_params_i32(op, 0);
const int d_state = (int) op->src[0]->ne[0];
const bool is_mamba2 = (op->src[3]->ne[0] == 1);
return is_mamba2 && (d_state == 128 || d_state == 256) && (K == 1);
}
case GGML_OP_GATED_DELTA_NET:
{
// Match the Vulkan backend: only F32 -> F32, S_v in {16, 32, 64, 128}.
@ -12260,6 +12298,103 @@ static void ggml_cl_mean(ggml_backend_t backend, const ggml_tensor * src0, const
backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
}
static void ggml_cl_ssm_scan(ggml_backend_t backend, ggml_tensor * dst) {
const ggml_tensor * src0 = dst->src[0]; // s
const ggml_tensor * src1 = dst->src[1]; // x
const ggml_tensor * src2 = dst->src[2]; // dt
const ggml_tensor * src3 = dst->src[3]; // A
const ggml_tensor * src4 = dst->src[4]; // B
const ggml_tensor * src5 = dst->src[5]; // C
const ggml_tensor * src6 = dst->src[6]; // ids
GGML_ASSERT(src0 && src1 && src2 && src3 && src4 && src5 && src6 && dst);
ggml_backend_opencl_context * backend_ctx = (ggml_backend_opencl_context *) backend->context;
ggml_tensor_extra_cl * e0 = (ggml_tensor_extra_cl *) src0->extra;
ggml_tensor_extra_cl * e1 = (ggml_tensor_extra_cl *) src1->extra;
ggml_tensor_extra_cl * e2 = (ggml_tensor_extra_cl *) src2->extra;
ggml_tensor_extra_cl * e3 = (ggml_tensor_extra_cl *) src3->extra;
ggml_tensor_extra_cl * e4 = (ggml_tensor_extra_cl *) src4->extra;
ggml_tensor_extra_cl * e5 = (ggml_tensor_extra_cl *) src5->extra;
ggml_tensor_extra_cl * e6 = (ggml_tensor_extra_cl *) src6->extra;
ggml_tensor_extra_cl * ed = (ggml_tensor_extra_cl *) dst->extra;
cl_ulong o0 = e0->offset + src0->view_offs;
cl_ulong o1 = e1->offset + src1->view_offs;
cl_ulong o2 = e2->offset + src2->view_offs;
cl_ulong o3 = e3->offset + src3->view_offs;
cl_ulong o4 = e4->offset + src4->view_offs;
cl_ulong o5 = e5->offset + src5->view_offs;
cl_ulong o6 = e6->offset + src6->view_offs;
cl_ulong od = ed->offset + dst->view_offs;
const int d_state = (int) src0->ne[0];
const int head_dim = (int) src0->ne[1];
const int n_head = (int) src1->ne[1];
const int n_group = (int) src4->ne[1];
const int n_tokens = (int) src1->ne[2];
const int n_seqs = (int) src1->ne[3];
// Mirror CPU ref: s_off = ggml_nelements(src1) * sizeof(float)
const cl_ulong s_off_bytes = (cl_ulong) ggml_nelements(src1) * sizeof(float);
cl_kernel kernel = (d_state == 128)
? backend_ctx->kernel_ssm_scan_f32_mamba2_d128
: backend_ctx->kernel_ssm_scan_f32_mamba2_d256;
GGML_ASSERT(kernel != nullptr);
cl_ulong s0_nb2 = src0->nb[2];
cl_ulong s0_nb3 = src0->nb[3];
cl_ulong x_nb2 = src1->nb[2];
cl_ulong x_nb3 = src1->nb[3];
cl_ulong dt_nb1 = src2->nb[1];
cl_ulong dt_nb2 = src2->nb[2];
cl_ulong A_nb1 = src3->nb[1];
cl_ulong B_nb2 = src4->nb[2];
cl_ulong B_nb3 = src4->nb[3];
cl_ulong C_nb2 = src5->nb[2];
cl_ulong C_nb3 = src5->nb[3];
CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &e0->data_device));
CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_ulong), &o0));
CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &e1->data_device));
CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_ulong), &o1));
CL_CHECK(clSetKernelArg(kernel, 4, sizeof(cl_mem), &e2->data_device));
CL_CHECK(clSetKernelArg(kernel, 5, sizeof(cl_ulong), &o2));
CL_CHECK(clSetKernelArg(kernel, 6, sizeof(cl_mem), &e3->data_device));
CL_CHECK(clSetKernelArg(kernel, 7, sizeof(cl_ulong), &o3));
CL_CHECK(clSetKernelArg(kernel, 8, sizeof(cl_mem), &e4->data_device));
CL_CHECK(clSetKernelArg(kernel, 9, sizeof(cl_ulong), &o4));
CL_CHECK(clSetKernelArg(kernel, 10, sizeof(cl_mem), &e5->data_device));
CL_CHECK(clSetKernelArg(kernel, 11, sizeof(cl_ulong), &o5));
CL_CHECK(clSetKernelArg(kernel, 12, sizeof(cl_mem), &e6->data_device));
CL_CHECK(clSetKernelArg(kernel, 13, sizeof(cl_ulong), &o6));
CL_CHECK(clSetKernelArg(kernel, 14, sizeof(cl_mem), &ed->data_device));
CL_CHECK(clSetKernelArg(kernel, 15, sizeof(cl_ulong), &od));
CL_CHECK(clSetKernelArg(kernel, 16, sizeof(cl_ulong), &s0_nb2));
CL_CHECK(clSetKernelArg(kernel, 17, sizeof(cl_ulong), &s0_nb3));
CL_CHECK(clSetKernelArg(kernel, 18, sizeof(cl_ulong), &x_nb2));
CL_CHECK(clSetKernelArg(kernel, 19, sizeof(cl_ulong), &x_nb3));
CL_CHECK(clSetKernelArg(kernel, 20, sizeof(cl_ulong), &dt_nb1));
CL_CHECK(clSetKernelArg(kernel, 21, sizeof(cl_ulong), &dt_nb2));
CL_CHECK(clSetKernelArg(kernel, 22, sizeof(cl_ulong), &A_nb1));
CL_CHECK(clSetKernelArg(kernel, 23, sizeof(cl_ulong), &B_nb2));
CL_CHECK(clSetKernelArg(kernel, 24, sizeof(cl_ulong), &B_nb3));
CL_CHECK(clSetKernelArg(kernel, 25, sizeof(cl_ulong), &C_nb2));
CL_CHECK(clSetKernelArg(kernel, 26, sizeof(cl_ulong), &C_nb3));
CL_CHECK(clSetKernelArg(kernel, 27, sizeof(cl_ulong), &s_off_bytes));
CL_CHECK(clSetKernelArg(kernel, 28, sizeof(int), &head_dim));
CL_CHECK(clSetKernelArg(kernel, 29, sizeof(int), &n_head));
CL_CHECK(clSetKernelArg(kernel, 30, sizeof(int), &n_group));
CL_CHECK(clSetKernelArg(kernel, 31, sizeof(int), &n_tokens));
size_t global_work_size[] = { (size_t)n_head * head_dim * 64, (size_t)n_seqs, 1 };
size_t local_work_size[] = { 64, 1, 1 };
backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
}
static void ggml_cl_ssm_conv(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
GGML_ASSERT(src0);
GGML_ASSERT(src0->extra);
@ -24746,6 +24881,14 @@ bool ggml_cl_compute_forward(ggml_backend_t backend, struct ggml_tensor * tensor
}
func = ggml_cl_ssm_conv;
break;
case GGML_OP_SSM_SCAN:
if (!any_on_device) {
return false;
}
// SSM_SCAN has 7 source tensors, so it cannot use the standard
// (src0, src1, dst) func signature. Dispatch directly and return.
ggml_cl_ssm_scan(backend, tensor);
return true;
case GGML_OP_GATED_DELTA_NET:
if (!any_on_device) {
return false;

216
ggml/src/ggml-opencl/kernels/ssm_scan.cl Normal file
View File

@ -0,0 +1,216 @@
// Mamba2 fused SSM scan kernel. One workgroup per (head, dim, seq); WG size =
// 64 threads. Each thread owns c_factor = d_state/64 state elements in
// private registers; the state stays resident across the n_tokens t-loop
//
// References:
// ggml/src/ggml-cuda/ssm-scan.cu:117 ssm_scan_f32_group
// ggml/src/ggml-cpu/ops.cpp:9368 ggml_compute_forward_ssm_scan_f32
#pragma OPENCL EXTENSION cl_khr_fp16 : enable
#ifdef cl_khr_subgroups
#pragma OPENCL EXTENSION cl_khr_subgroups : enable
#endif
#if defined(cl_qcom_reqd_sub_group_size)
#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
#define REQD_SUBGROUP_SIZE_64 __attribute__((qcom_reqd_sub_group_size("half")))
#else
#define REQD_SUBGROUP_SIZE_64
#endif
inline float softplus_f32(float x) {
return (x <= 20.0f) ? log(1.0f + exp(x)) : x;
}
// d_state = 128 (most Mamba-2 models, e.g. mamba2-2.7B, Codestral-Mamba).
// WG = 64 threads, each holds 2 state elements (tid and tid+64).
REQD_SUBGROUP_SIZE_64
kernel void kernel_ssm_scan_f32_mamba2_d128(
global const char * src0_base, ulong src0_off,
global const char * src1_base, ulong src1_off,
global const char * src2_base, ulong src2_off,
global const char * src3_base, ulong src3_off,
global const char * src4_base, ulong src4_off,
global const char * src5_base, ulong src5_off,
global const char * src6_base, ulong src6_off,
global char * dst_base, ulong dst_off,
ulong s0_nb2, ulong s0_nb3,
ulong x_nb2, ulong x_nb3,
ulong dt_nb1, ulong dt_nb2,
ulong A_nb1,
ulong B_nb2, ulong B_nb3,
ulong C_nb2, ulong C_nb3,
ulong s_off_bytes,
int head_dim, int n_head, int n_group, int n_tokens
) {
const int d_state = 128;
const int tid = (int) get_local_id(0);
const int wg_x = (int) get_group_id(0);
const int seq_id = (int) get_group_id(1);
const int head_id = wg_x / head_dim;
const int dim_id = wg_x - head_id * head_dim;
const int g = head_id / (n_head / n_group);
src0_base += src0_off;
src1_base += src1_off;
src2_base += src2_off;
src3_base += src3_off;
src4_base += src4_off;
src5_base += src5_off;
src6_base += src6_off;
dst_base += dst_off;
const int seq_slot = ((global const int *) src6_base)[seq_id];
const ulong state_base_off = (ulong)seq_slot * s0_nb3 + (ulong)head_id * s0_nb2
+ (ulong)dim_id * d_state * sizeof(float);
global const float * s0_warp = (global const float *)(src0_base + state_base_off);
const ulong state_out_off = (ulong)seq_id * s0_nb3 + (ulong)head_id * s0_nb2
+ (ulong)dim_id * d_state * sizeof(float);
global float * s_warp = (global float *)(dst_base + s_off_bytes + state_out_off);
global const char * x_seq = src1_base + (ulong)seq_id * x_nb3;
global const char * dt_seq = src2_base + (ulong)seq_id * dt_nb2;
global const char * B_seq = src4_base + (ulong)seq_id * B_nb3 + (ulong)g * d_state * sizeof(float);
global const char * C_seq = src5_base + (ulong)seq_id * C_nb3 + (ulong)g * d_state * sizeof(float);
const ulong y_dim_total = (ulong)n_head * head_dim;
global float * y_seq = (global float *)dst_base
+ (ulong)seq_id * (ulong)n_tokens * y_dim_total;
const float A_val = ((global const float *)src3_base)[(ulong)head_id * A_nb1 / sizeof(float)];
// c_factor = 2: each thread owns 2 state elements (tid and tid+64).
float state0 = s0_warp[tid];
float state1 = s0_warp[tid + 64];
for (int t = 0; t < n_tokens; ++t) {
const float dt_h = ((global const float *)(dt_seq + (ulong)t * dt_nb1))[head_id];
const float dt_softplus = softplus_f32(dt_h);
const float dA = exp(dt_softplus * A_val);
const float x_val = ((global const float *)(x_seq + (ulong)t * x_nb2))[(ulong)head_id * head_dim + dim_id];
const float x_dt = x_val * dt_softplus;
const float B0 = ((global const float *)(B_seq + (ulong)t * B_nb2))[tid];
const float B1 = ((global const float *)(B_seq + (ulong)t * B_nb2))[tid + 64];
const float C0 = ((global const float *)(C_seq + (ulong)t * C_nb2))[tid];
const float C1 = ((global const float *)(C_seq + (ulong)t * C_nb2))[tid + 64];
state0 = state0 * dA + B0 * x_dt;
state1 = state1 * dA + B1 * x_dt;
const float partial = state0 * C0 + state1 * C1;
const float sum = sub_group_reduce_add(partial);
if (tid == 0) {
y_seq[(ulong)t * y_dim_total + (ulong)head_id * head_dim + dim_id] = sum;
}
}
s_warp[tid] = state0;
s_warp[tid + 64] = state1;
}
// d_state = 256 (Falcon-H1). WG = 64 threads, each holds 4 state elements.
REQD_SUBGROUP_SIZE_64
kernel void kernel_ssm_scan_f32_mamba2_d256(
global const char * src0_base, ulong src0_off,
global const char * src1_base, ulong src1_off,
global const char * src2_base, ulong src2_off,
global const char * src3_base, ulong src3_off,
global const char * src4_base, ulong src4_off,
global const char * src5_base, ulong src5_off,
global const char * src6_base, ulong src6_off,
global char * dst_base, ulong dst_off,
ulong s0_nb2, ulong s0_nb3,
ulong x_nb2, ulong x_nb3,
ulong dt_nb1, ulong dt_nb2,
ulong A_nb1,
ulong B_nb2, ulong B_nb3,
ulong C_nb2, ulong C_nb3,
ulong s_off_bytes,
int head_dim, int n_head, int n_group, int n_tokens
) {
const int d_state = 256;
const int tid = (int) get_local_id(0);
const int wg_x = (int) get_group_id(0);
const int seq_id = (int) get_group_id(1);
const int head_id = wg_x / head_dim;
const int dim_id = wg_x - head_id * head_dim;
const int g = head_id / (n_head / n_group);
src0_base += src0_off;
src1_base += src1_off;
src2_base += src2_off;
src3_base += src3_off;
src4_base += src4_off;
src5_base += src5_off;
src6_base += src6_off;
dst_base += dst_off;
const int seq_slot = ((global const int *) src6_base)[seq_id];
const ulong state_base_off = (ulong)seq_slot * s0_nb3 + (ulong)head_id * s0_nb2
+ (ulong)dim_id * d_state * sizeof(float);
global const float * s0_warp = (global const float *)(src0_base + state_base_off);
const ulong state_out_off = (ulong)seq_id * s0_nb3 + (ulong)head_id * s0_nb2
+ (ulong)dim_id * d_state * sizeof(float);
global float * s_warp = (global float *)(dst_base + s_off_bytes + state_out_off);
global const char * x_seq = src1_base + (ulong)seq_id * x_nb3;
global const char * dt_seq = src2_base + (ulong)seq_id * dt_nb2;
global const char * B_seq = src4_base + (ulong)seq_id * B_nb3 + (ulong)g * d_state * sizeof(float);
global const char * C_seq = src5_base + (ulong)seq_id * C_nb3 + (ulong)g * d_state * sizeof(float);
const ulong y_dim_total = (ulong)n_head * head_dim;
global float * y_seq = (global float *)dst_base
+ (ulong)seq_id * (ulong)n_tokens * y_dim_total;
const float A_val = ((global const float *)src3_base)[(ulong)head_id * A_nb1 / sizeof(float)];
// c_factor = 4: each thread owns 4 state elements.
float state0 = s0_warp[tid];
float state1 = s0_warp[tid + 64];
float state2 = s0_warp[tid + 128];
float state3 = s0_warp[tid + 192];
for (int t = 0; t < n_tokens; ++t) {
const float dt_h = ((global const float *)(dt_seq + (ulong)t * dt_nb1))[head_id];
const float dt_softplus = softplus_f32(dt_h);
const float dA = exp(dt_softplus * A_val);
const float x_val = ((global const float *)(x_seq + (ulong)t * x_nb2))[(ulong)head_id * head_dim + dim_id];
const float x_dt = x_val * dt_softplus;
global const float * B_t = (global const float *)(B_seq + (ulong)t * B_nb2);
global const float * C_t = (global const float *)(C_seq + (ulong)t * C_nb2);
const float B0 = B_t[tid];
const float B1 = B_t[tid + 64];
const float B2 = B_t[tid + 128];
const float B3 = B_t[tid + 192];
const float C0 = C_t[tid];
const float C1 = C_t[tid + 64];
const float C2 = C_t[tid + 128];
const float C3 = C_t[tid + 192];
state0 = state0 * dA + B0 * x_dt;
state1 = state1 * dA + B1 * x_dt;
state2 = state2 * dA + B2 * x_dt;
state3 = state3 * dA + B3 * x_dt;
const float partial = state0 * C0 + state1 * C1 + state2 * C2 + state3 * C3;
const float sum = sub_group_reduce_add(partial);
if (tid == 0) {
y_seq[(ulong)t * y_dim_total + (ulong)head_id * head_dim + dim_id] = sum;
}
}
s_warp[tid] = state0;
s_warp[tid + 64] = state1;
s_warp[tid + 128] = state2;
s_warp[tid + 192] = state3;
}