Skip to content

Commit 3d4e86b

Browse files
giuseppegiuseppe
authored
vulkan: Add State Space Model (SSM) Operations Support (ggml-org#16463)
* vulkan: implement SSM scan operation Add State Space Model scan operation to the Vulkan backend. Signed-off-by: Giuseppe Scrivano <[email protected]> * vulkan: implement SSM conv operation Add State Space Model conv operation to the Vulkan backend. Signed-off-by: Giuseppe Scrivano <[email protected]> --------- Signed-off-by: Giuseppe Scrivano <[email protected]>
1 parent 342c728 commit 3d4e86b

File tree

5 files changed

+396
-8
lines changed

5 files changed

+396
-8
lines changed

‎docs/ops.md‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -100,8 +100,8 @@ Legend:
100100
| SOFT_MAX_BACK | ❌ | ❌ | 🟡 | 🟡 | ❌ | ❌ | 🟡 | ✅ | ❌ |
101101
| SQR | ❌ | ✅ | ✅ | ✅ | 🟡 | ❌ | ✅ | 🟡 | ❌ |
102102
| SQRT | ❌ | ✅ | ✅ | ✅ | 🟡 | ❌ | ✅ | ❌ | ❌ |
103-
| SSM_CONV | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
104-
| SSM_SCAN | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
103+
| SSM_CONV | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
104+
| SSM_SCAN | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
105105
| STEP | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | 🟡 | ❌ | ❌ |
106106
| SUB | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ❌ |
107107
| SUM | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |

‎ggml/src/ggml-vulkan/ggml-vulkan.cpp‎

Lines changed: 221 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -582,6 +582,9 @@ struct vk_device_struct {
582582
vk_pipeline pipeline_pool2d_f32;
583583
vk_pipeline pipeline_rwkv_wkv6_f32;
584584
vk_pipeline pipeline_rwkv_wkv7_f32;
585+
vk_pipeline pipeline_ssm_scan_f32_d128;
586+
vk_pipeline pipeline_ssm_scan_f32_d256;
587+
vk_pipeline pipeline_ssm_conv_f32;
585588
vk_pipeline pipeline_opt_step_adamw_f32;
586589
vk_pipeline pipeline_opt_step_sgd_f32;
587590
vk_pipeline pipeline_conv2d_f32[CONV_SHAPE_COUNT];
@@ -1087,6 +1090,19 @@ struct vk_op_rwkv_wkv7_push_constants {
10871090
uint32_t C;
10881091
uint32_t H;
10891092
};
1093+
struct vk_op_ssm_scan_push_constants {
1094+
uint32_t nb02, nb03, nb12, nb13;
1095+
uint32_t nb21, nb22, nb31;
1096+
uint32_t nb42, nb43, nb52, nb53;
1097+
uint32_t s_off;
1098+
uint32_t n_head, d_head, n_group, n_tok;
1099+
};
1100+
struct vk_op_ssm_conv_push_constants {
1101+
uint32_t nb01, nb02;
1102+
uint32_t nb11;
1103+
uint32_t dst_nb0, dst_nb1, dst_nb2;
1104+
uint32_t nc, ncs, nr, n_t, n_s;
1105+
};
10901106

10911107
struct vk_op_conv2d_push_constants {
10921108
uint32_t Cout;
@@ -3591,6 +3607,11 @@ static void ggml_vk_load_shaders(vk_device& device) {
35913607

35923608
ggml_vk_create_pipeline(device, device->pipeline_rwkv_wkv7_f32, "rwkv_wkv7_f32", rwkv_wkv7_f32_len, rwkv_wkv7_f32_data, "main", 8, sizeof(vk_op_rwkv_wkv7_push_constants), {1, 1, 1}, {device->subgroup_size}, 1);
35933609

3610+
ggml_vk_create_pipeline(device, device->pipeline_ssm_scan_f32_d128, "ssm_scan_f32", ssm_scan_f32_len, ssm_scan_f32_data, "main", 8, sizeof(vk_op_ssm_scan_push_constants), {1, 1, 1}, {128, device->subgroup_size, 16}, 1);
3611+
ggml_vk_create_pipeline(device, device->pipeline_ssm_scan_f32_d256, "ssm_scan_f32", ssm_scan_f32_len, ssm_scan_f32_data, "main", 8, sizeof(vk_op_ssm_scan_push_constants), {1, 1, 1}, {256, device->subgroup_size, 16}, 1);
3612+
3613+
ggml_vk_create_pipeline(device, device->pipeline_ssm_conv_f32, "ssm_conv_f32", ssm_conv_f32_len, ssm_conv_f32_data, "main", 3, sizeof(vk_op_ssm_conv_push_constants), {32, 1, 1}, {32}, 1);
3614+
35943615
ggml_vk_create_pipeline(device, device->pipeline_opt_step_adamw_f32, "opt_step_adamw_f32", opt_step_adamw_f32_len, opt_step_adamw_f32_data, "main", 5, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1);
35953616

35963617
ggml_vk_create_pipeline(device, device->pipeline_opt_step_sgd_f32, "opt_step_sgd_f32", opt_step_sgd_f32_len, opt_step_sgd_f32_data, "main", 3, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1);
@@ -8098,6 +8119,21 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
80988119
return ctx->device->pipeline_rwkv_wkv7_f32;
80998120
}
81008121
return nullptr;
8122+
case GGML_OP_SSM_SCAN:
8123+
if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
8124+
const uint32_t d_state = src0->ne[0];
8125+
if (d_state == 128) {
8126+
return ctx->device->pipeline_ssm_scan_f32_d128;
8127+
} else if (d_state == 256) {
8128+
return ctx->device->pipeline_ssm_scan_f32_d256;
8129+
}
8130+
}
8131+
return nullptr;
8132+
case GGML_OP_SSM_CONV:
8133+
if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
8134+
return ctx->device->pipeline_ssm_conv_f32;
8135+
}
8136+
return nullptr;
81018137
case GGML_OP_OPT_STEP_ADAMW:
81028138
if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
81038139
return ctx->device->pipeline_opt_step_adamw_f32;
@@ -8592,6 +8628,14 @@ static void ggml_vk_op_f32(ggml_backend_vk_context * ctx, vk_context& subctx, co
85928628
}
85938629
}
85948630
break;
8631+
case GGML_OP_SSM_CONV:
8632+
{
8633+
const uint32_t nr = src0->ne[1];
8634+
const uint32_t n_t = dst->ne[1];
8635+
const uint32_t n_s = dst->ne[2];
8636+
elements = { nr, n_t, n_s };
8637+
}
8638+
break;
85958639
default:
85968640
elements = { (uint32_t)ggml_nelements(src0), 1, 1 };
85978641
break;
@@ -9038,6 +9082,117 @@ static void ggml_vk_rwkv_wkv7(ggml_backend_vk_context * ctx, vk_context& subctx,
90389082
);
90399083
}
90409084

9085+
static void ggml_vk_ssm_scan(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst, bool dryrun = false) {
9086+
const ggml_tensor * src0 = dst->src[0];
9087+
const ggml_tensor * src1 = dst->src[1];
9088+
const ggml_tensor * src2 = dst->src[2];
9089+
const ggml_tensor * src3 = dst->src[3];
9090+
const ggml_tensor * src4 = dst->src[4];
9091+
const ggml_tensor * src5 = dst->src[5];
9092+
9093+
GGML_ASSERT(dst->buffer != nullptr);
9094+
9095+
const uint32_t head_dim = src0->ne[1];
9096+
const uint32_t n_head = src1->ne[1];
9097+
const uint32_t n_group = src4->ne[1];
9098+
const uint32_t n_tok = src1->ne[2];
9099+
const uint32_t n_seq = src1->ne[3];
9100+
9101+
bool is_mamba2 = (src3->nb[1] == sizeof(float));
9102+
GGML_ASSERT(is_mamba2);
9103+
9104+
vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, src0, src1, src2, dst, dst->op);
9105+
GGML_ASSERT(pipeline != nullptr);
9106+
9107+
if (dryrun) {
9108+
ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
9109+
return;
9110+
}
9111+
9112+
const int64_t s_off = ggml_nelements(src1) * sizeof(float);
9113+
9114+
const vk_op_ssm_scan_push_constants pc = {
9115+
(uint32_t)src0->nb[2], (uint32_t)src0->nb[3],
9116+
(uint32_t)src1->nb[2], (uint32_t)src1->nb[3],
9117+
(uint32_t)src2->nb[1], (uint32_t)src2->nb[2],
9118+
(uint32_t)src3->nb[1],
9119+
(uint32_t)src4->nb[2], (uint32_t)src4->nb[3],
9120+
(uint32_t)src5->nb[2], (uint32_t)src5->nb[3],
9121+
(uint32_t)s_off,
9122+
n_head, head_dim, n_group, n_tok
9123+
};
9124+
9125+
ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context;
9126+
ggml_backend_vk_buffer_context * src_buf_ctxs[GGML_MAX_SRC];
9127+
for (int i = 0; i < GGML_MAX_SRC && dst->src[i] != nullptr; i++) {
9128+
src_buf_ctxs[i] = (ggml_backend_vk_buffer_context *)dst->src[i]->buffer->context;
9129+
}
9130+
9131+
vk_buffer d_D = nullptr, d_srcs[GGML_MAX_SRC] = { nullptr };
9132+
size_t dst_offset = 0, src_offsets[GGML_MAX_SRC] = { 0 };
9133+
bool dst_uma = false, srcs_uma[GGML_MAX_SRC] = { false };
9134+
9135+
if (ctx->device->uma) {
9136+
for (int i = 0; i < GGML_MAX_SRC && dst->src[i] != nullptr; i++) {
9137+
ggml_vk_host_get(ctx->device, dst->src[i]->data, d_srcs[i], src_offsets[i]);
9138+
srcs_uma[i] = d_srcs[i] != nullptr;
9139+
}
9140+
ggml_vk_host_get(ctx->device, dst->data, d_D, dst_offset);
9141+
dst_uma = d_D != nullptr;
9142+
}
9143+
9144+
if (!dst_uma) {
9145+
d_D = dst_buf_ctx->dev_buffer;
9146+
dst_offset = vk_tensor_offset(dst) + dst->view_offs;
9147+
}
9148+
for (int i = 0; i < GGML_MAX_SRC && dst->src[i] != nullptr; i++) {
9149+
if (!srcs_uma[i]) {
9150+
d_srcs[i] = src_buf_ctxs[i]->dev_buffer;
9151+
src_offsets[i] = vk_tensor_offset(dst->src[i]) + dst->src[i]->view_offs;
9152+
}
9153+
}
9154+
9155+
size_t dst_size = ggml_nbytes(dst);
9156+
size_t src_sizes[GGML_MAX_SRC];
9157+
for (int i = 0; i < GGML_MAX_SRC && dst->src[i] != nullptr; i++) {
9158+
src_sizes[i] = ggml_nbytes(dst->src[i]);
9159+
}
9160+
9161+
std::array<uint32_t, 3> elements;
9162+
9163+
const int splitH = 16;
9164+
const uint32_t num_workgroups_x = CEIL_DIV(n_head * head_dim, splitH);
9165+
const uint32_t num_workgroups_y = n_seq;
9166+
elements = { num_workgroups_x, num_workgroups_y, 1 };
9167+
9168+
ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, {
9169+
vk_subbuffer{ d_srcs[0], src_offsets[0], src_sizes[0] },
9170+
vk_subbuffer{ d_srcs[1], src_offsets[1], src_sizes[1] },
9171+
vk_subbuffer{ d_srcs[2], src_offsets[2], src_sizes[2] },
9172+
vk_subbuffer{ d_srcs[3], src_offsets[3], src_sizes[3] },
9173+
vk_subbuffer{ d_srcs[4], src_offsets[4], src_sizes[4] },
9174+
vk_subbuffer{ d_srcs[5], src_offsets[5], src_sizes[5] },
9175+
vk_subbuffer{ d_srcs[6], src_offsets[6], src_sizes[6] },
9176+
vk_subbuffer{ d_D, dst_offset, dst_size }
9177+
}, pc, elements);
9178+
}
9179+
9180+
static void ggml_vk_ssm_conv(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst, bool dryrun = false) {
9181+
const ggml_tensor * src0 = dst->src[0];
9182+
const ggml_tensor * src1 = dst->src[1];
9183+
9184+
ggml_vk_op_f32<vk_op_ssm_conv_push_constants>(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_SSM_CONV, {
9185+
(uint32_t)src0->nb[1], (uint32_t)src0->nb[2],
9186+
(uint32_t)src1->nb[1],
9187+
(uint32_t)dst->nb[0], (uint32_t)dst->nb[1], (uint32_t)dst->nb[2],
9188+
(uint32_t)src1->ne[0],
9189+
(uint32_t)src0->ne[0],
9190+
(uint32_t)src0->ne[1],
9191+
(uint32_t)dst->ne[1],
9192+
(uint32_t)dst->ne[2],
9193+
}, dryrun);
9194+
}
9195+
90419196
static void ggml_vk_op_f32_opt_step_adamw(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst, const vk_op_push_constants&& pc, bool dryrun = false) {
90429197
const ggml_tensor * x = dst->src[0];
90439198
const ggml_tensor * g = dst->src[1];
@@ -10870,6 +11025,8 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr
1087011025
case GGML_OP_CONV_2D_DW:
1087111026
case GGML_OP_RWKV_WKV6:
1087211027
case GGML_OP_RWKV_WKV7:
11028+
case GGML_OP_SSM_SCAN:
11029+
case GGML_OP_SSM_CONV:
1087311030
case GGML_OP_LEAKY_RELU:
1087411031
case GGML_OP_FLASH_ATTN_EXT:
1087511032
case GGML_OP_OPT_STEP_ADAMW:
@@ -11287,6 +11444,16 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr
1128711444

1128811445
break;
1128911446

11447+
case GGML_OP_SSM_SCAN:
11448+
ggml_vk_ssm_scan(ctx, compute_ctx, node, dryrun);
11449+
11450+
break;
11451+
11452+
case GGML_OP_SSM_CONV:
11453+
ggml_vk_ssm_conv(ctx, compute_ctx, node, dryrun);
11454+
11455+
break;
11456+
1129011457
case GGML_OP_OPT_STEP_ADAMW:
1129111458
ggml_vk_opt_step_adamw(ctx, compute_ctx, node, dryrun);
1129211459

@@ -11398,6 +11565,8 @@ static bool ggml_vk_compute_forward(ggml_backend_vk_context * ctx, ggml_cgraph *
1139811565
case GGML_OP_CONV_2D_DW:
1139911566
case GGML_OP_RWKV_WKV6:
1140011567
case GGML_OP_RWKV_WKV7:
11568+
case GGML_OP_SSM_SCAN:
11569+
case GGML_OP_SSM_CONV:
1140111570
case GGML_OP_LEAKY_RELU:
1140211571
case GGML_OP_REPEAT:
1140311572
case GGML_OP_REPEAT_BACK:
@@ -12879,6 +13048,47 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
1287913048
case GGML_OP_RWKV_WKV6:
1288013049
case GGML_OP_RWKV_WKV7:
1288113050
return true;
13051+
case GGML_OP_SSM_SCAN:
13052+
{
13053+
for (int i = 0; i < 6; i++) {
13054+
if (op->src[i] && ggml_is_quantized(op->src[i]->type)) {
13055+
return false;
13056+
}
13057+
}
13058+
if (op->src[6] && op->src[6]->type != GGML_TYPE_I32) {
13059+
return false;
13060+
}
13061+
if (op->src[0]->type != GGML_TYPE_F32 || op->type != GGML_TYPE_F32) {
13062+
return false;
13063+
}
13064+
13065+
const uint32_t d_state = op->src[0]->ne[0];
13066+
const uint32_t head_dim = op->src[0]->ne[1];
13067+
13068+
bool is_mamba2 = (op->src[3] && op->src[3]->nb[1] == sizeof(float));
13069+
if (!is_mamba2) {
13070+
return false;
13071+
}
13072+
13073+
if ((d_state != 128 && d_state != 256) || head_dim % 16 != 0) {
13074+
return false;
13075+
}
13076+
13077+
ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context;
13078+
const vk_device& device = ggml_vk_get_device(ctx->device);
13079+
13080+
const uint32_t SPLIT_H = 16;
13081+
13082+
size_t stateC_size = SPLIT_H * d_state * sizeof(float);
13083+
13084+
if (stateC_size > device->properties.limits.maxComputeSharedMemorySize) {
13085+
return false;
13086+
}
13087+
13088+
return true;
13089+
}
13090+
case GGML_OP_SSM_CONV:
13091+
return true;
1288213092
case GGML_OP_CONV_TRANSPOSE_1D:
1288313093
return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32;
1288413094
case GGML_OP_CONV_2D:
@@ -13223,14 +13433,14 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph *
1322313433

1322413434
struct ggml_context * ggml_ctx = ggml_init(iparams);
1322513435

13226-
std::array<struct ggml_tensor *, 6> src_clone = {nullptr, nullptr, nullptr, nullptr, nullptr, nullptr};
13227-
std::array<size_t, 6> src_size = {0, 0, 0, 0, 0, 0};
13228-
std::array<void *, 6> src_buffer = {nullptr, nullptr, nullptr, nullptr, nullptr, nullptr};
13229-
const char * srci_name[6] = {"src0", "src1", "src2", "src3", "src4", "src5"};
13436+
std::array<struct ggml_tensor *, GGML_MAX_SRC> src_clone = {nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr};
13437+
std::array<size_t, GGML_MAX_SRC> src_size = {};
13438+
std::array<void *, GGML_MAX_SRC> src_buffer = {};
13439+
const char * srci_name[GGML_MAX_SRC] = {"src0", "src1", "src2", "src3", "src4", "src5", "src6", "src7", "src8", "src9"};
1323013440

1323113441
struct ggml_tensor * tensor_clone = nullptr;
1323213442

13233-
for (int i = 0; i < 6; i++) {
13443+
for (int i = 0; i < GGML_MAX_SRC; i++) {
1323413444
ggml_tensor * srci = tensor->src[i];
1323513445
if (fused_rms_norm_mul) {
1323613446
rms_norm_idx = tensor->src[0]->op == GGML_OP_RMS_NORM ? 0 : 1;
@@ -13537,6 +13747,11 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph *
1353713747
src_clone[2]);
1353813748
} else if (tensor->op == GGML_OP_ADD_ID) {
1353913749
tensor_clone = ggml_add_id(ggml_ctx, src_clone[0], src_clone[1], src_clone[2]);
13750+
} else if (tensor->op == GGML_OP_SSM_SCAN) {
13751+
tensor_clone = ggml_ssm_scan(ggml_ctx, src_clone[0], src_clone[1], src_clone[2],
13752+
src_clone[3], src_clone[4], src_clone[5], src_clone[6]);
13753+
} else if (tensor->op == GGML_OP_SSM_CONV) {
13754+
tensor_clone = ggml_ssm_conv(ggml_ctx, src_clone[0], src_clone[1]);
1354013755
}
1354113756
else {
1354213757
std::cerr << "Missing vk_check_results OP: " << ggml_op_name(tensor->op) << std::endl;
@@ -13558,7 +13773,7 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph *
1355813773
memcpy(comp_result, tensor_clone->data, comp_size);
1355913774
memcpy(comp_nb, tensor_clone->nb, sizeof(size_t) * GGML_MAX_DIMS);
1356013775

13561-
for (int i = 0; i < 6; i++) {
13776+
for (int i = 0; i < GGML_MAX_SRC; i++) {
1356213777
if (src_buffer[i] != nullptr) {
1356313778
free(src_buffer[i]);
1356413779
}

‎ggml/src/ggml-vulkan/vulkan-shaders/ssm_conv.comp‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
#version 450
2+
3+
#extension GL_EXT_control_flow_attributes : require
4+
5+
#include "types.glsl"
6+
7+
layout(constant_id = 0) const uint BLOCK_SIZE = 32;
8+
9+
layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
10+
11+
layout(binding = 0) readonly buffer Src0 { float src0[]; };
12+
layout(binding = 1) readonly buffer Src1 { float src1[]; };
13+
layout(binding = 2) buffer Dst { float dst[]; };
14+
15+
layout(push_constant) uniform PushConstants {
16+
uint nb01; uint nb02;
17+
uint nb11;
18+
uint dst_nb0; uint dst_nb1; uint dst_nb2;
19+
uint nc; uint ncs; uint nr; uint n_t; uint n_s;
20+
};
21+
22+
void main() {
23+
const uint global_thread_id = gl_GlobalInvocationID.x;
24+
const uint i2 = gl_WorkGroupID.y;
25+
const uint i3 = gl_WorkGroupID.z;
26+
27+
if (global_thread_id >= nr || i2 >= n_t || i3 >= n_s) {
28+
return;
29+
}
30+
31+
const uint i1 = global_thread_id;
32+
const uint src0_base = i3 * (nb02 / 4) + i2 + i1 * (nb01 / 4);
33+
const uint src1_base = i1 * (nb11 / 4);
34+
const uint dst_idx = i3 * (dst_nb2 / 4) + i2 * (dst_nb1 / 4) + i1;
35+
36+
float sum = 0.0;
37+
[[unroll]] for (uint i0 = 0; i0 < nc; i0++) {
38+
const uint src0_idx = src0_base + i0;
39+
const uint src1_idx = src1_base + i0;
40+
sum += src0[src0_idx] * src1[src1_idx];
41+
}
42+
43+
dst[dst_idx] = sum;
44+
}

0 commit comments

Comments
 (0)