Skip to content

Commit 2a6f4cc

Browse files
giuseppegiuseppe
committed
vulkan: implement SSM scan operation
Add State Space Model scan operation to the Vulkan backend. Signed-off-by: Giuseppe Scrivano <[email protected]>
1 parent 9c7185d commit 2a6f4cc

File tree

3 files changed

+320
-6
lines changed

3 files changed

+320
-6
lines changed

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

Lines changed: 182 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -582,6 +582,8 @@ 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;
585587
vk_pipeline pipeline_opt_step_adamw_f32;
586588
vk_pipeline pipeline_opt_step_sgd_f32;
587589
vk_pipeline pipeline_conv2d_f32[CONV_SHAPE_COUNT];
@@ -1087,6 +1089,13 @@ struct vk_op_rwkv_wkv7_push_constants {
10871089
uint32_t C;
10881090
uint32_t H;
10891091
};
1092+
struct vk_op_ssm_scan_push_constants {
1093+
uint32_t src0_nb2, src0_nb3, src1_nb2, src1_nb3;
1094+
uint32_t src2_nb1, src2_nb2, src3_nb1;
1095+
uint32_t src4_nb2, src4_nb3, src5_nb2, src5_nb3;
1096+
uint32_t s_off;
1097+
uint32_t n_head, d_head, n_group, n_tok;
1098+
};
10901099

10911100
struct vk_op_conv2d_push_constants {
10921101
uint32_t Cout;
@@ -3588,6 +3597,9 @@ static void ggml_vk_load_shaders(vk_device& device) {
35883597

35893598
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);
35903599

3600+
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);
3601+
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);
3602+
35913603
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);
35923604

35933605
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);
@@ -8087,6 +8099,16 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
80878099
return ctx->device->pipeline_rwkv_wkv7_f32;
80888100
}
80898101
return nullptr;
8102+
case GGML_OP_SSM_SCAN:
8103+
if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
8104+
const uint32_t d_state = src0->ne[0];
8105+
if (d_state == 128) {
8106+
return ctx->device->pipeline_ssm_scan_f32_d128;
8107+
} else if (d_state == 256) {
8108+
return ctx->device->pipeline_ssm_scan_f32_d256;
8109+
}
8110+
}
8111+
return nullptr;
80908112
case GGML_OP_OPT_STEP_ADAMW:
80918113
if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
80928114
return ctx->device->pipeline_opt_step_adamw_f32;
@@ -9027,6 +9049,101 @@ static void ggml_vk_rwkv_wkv7(ggml_backend_vk_context * ctx, vk_context& subctx,
90279049
);
90289050
}
90299051

9052+
static void ggml_vk_ssm_scan(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst, bool dryrun = false) {
9053+
const ggml_tensor * src0 = dst->src[0];
9054+
const ggml_tensor * src1 = dst->src[1];
9055+
const ggml_tensor * src2 = dst->src[2];
9056+
const ggml_tensor * src3 = dst->src[3];
9057+
const ggml_tensor * src4 = dst->src[4];
9058+
const ggml_tensor * src5 = dst->src[5];
9059+
9060+
GGML_ASSERT(dst->buffer != nullptr);
9061+
9062+
const uint32_t head_dim = src0->ne[1];
9063+
const uint32_t n_head = src1->ne[1];
9064+
const uint32_t n_group = src4->ne[1];
9065+
const uint32_t n_tok = src1->ne[2];
9066+
const uint32_t n_seq = src1->ne[3];
9067+
9068+
bool is_mamba2 = (src3->nb[1] == sizeof(float));
9069+
GGML_ASSERT(is_mamba2);
9070+
9071+
vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, src0, src1, src2, dst, dst->op);
9072+
GGML_ASSERT(pipeline != nullptr);
9073+
9074+
if (dryrun) {
9075+
ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
9076+
return;
9077+
}
9078+
9079+
const int64_t s_off = ggml_nelements(src1) * sizeof(float);
9080+
9081+
const vk_op_ssm_scan_push_constants pc = {
9082+
(uint32_t)src0->nb[2], (uint32_t)src0->nb[3],
9083+
(uint32_t)src1->nb[2], (uint32_t)src1->nb[3],
9084+
(uint32_t)src2->nb[1], (uint32_t)src2->nb[2],
9085+
(uint32_t)src3->nb[1],
9086+
(uint32_t)src4->nb[2], (uint32_t)src4->nb[3],
9087+
(uint32_t)src5->nb[2], (uint32_t)src5->nb[3],
9088+
(uint32_t)s_off,
9089+
n_head, head_dim, n_group, n_tok
9090+
};
9091+
9092+
ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context;
9093+
ggml_backend_vk_buffer_context * src_buf_ctxs[GGML_MAX_SRC];
9094+
for (int i = 0; i < GGML_MAX_SRC && dst->src[i] != nullptr; i++) {
9095+
src_buf_ctxs[i] = (ggml_backend_vk_buffer_context *)dst->src[i]->buffer->context;
9096+
}
9097+
9098+
vk_buffer d_D = nullptr, d_srcs[GGML_MAX_SRC] = { nullptr };
9099+
size_t dst_offset = 0, src_offsets[GGML_MAX_SRC] = { 0 };
9100+
bool dst_uma = false, srcs_uma[GGML_MAX_SRC] = { false };
9101+
9102+
if (ctx->device->uma) {
9103+
for (int i = 0; i < GGML_MAX_SRC && dst->src[i] != nullptr; i++) {
9104+
ggml_vk_host_get(ctx->device, dst->src[i]->data, d_srcs[i], src_offsets[i]);
9105+
srcs_uma[i] = d_srcs[i] != nullptr;
9106+
}
9107+
ggml_vk_host_get(ctx->device, dst->data, d_D, dst_offset);
9108+
dst_uma = d_D != nullptr;
9109+
}
9110+
9111+
if (!dst_uma) {
9112+
d_D = dst_buf_ctx->dev_buffer;
9113+
dst_offset = vk_tensor_offset(dst) + dst->view_offs;
9114+
}
9115+
for (int i = 0; i < GGML_MAX_SRC && dst->src[i] != nullptr; i++) {
9116+
if (!srcs_uma[i]) {
9117+
d_srcs[i] = src_buf_ctxs[i]->dev_buffer;
9118+
src_offsets[i] = vk_tensor_offset(dst->src[i]) + dst->src[i]->view_offs;
9119+
}
9120+
}
9121+
9122+
size_t dst_size = ggml_nbytes(dst);
9123+
size_t src_sizes[GGML_MAX_SRC];
9124+
for (int i = 0; i < GGML_MAX_SRC && dst->src[i] != nullptr; i++) {
9125+
src_sizes[i] = ggml_nbytes(dst->src[i]);
9126+
}
9127+
9128+
std::array<uint32_t, 3> elements;
9129+
9130+
const int splitH = 16;
9131+
const uint32_t num_workgroups_x = CEIL_DIV(n_head * head_dim, splitH);
9132+
const uint32_t num_workgroups_y = n_seq;
9133+
elements = { num_workgroups_x, num_workgroups_y, 1 };
9134+
9135+
ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, {
9136+
vk_subbuffer{ d_srcs[0], src_offsets[0], src_sizes[0] },
9137+
vk_subbuffer{ d_srcs[1], src_offsets[1], src_sizes[1] },
9138+
vk_subbuffer{ d_srcs[2], src_offsets[2], src_sizes[2] },
9139+
vk_subbuffer{ d_srcs[3], src_offsets[3], src_sizes[3] },
9140+
vk_subbuffer{ d_srcs[4], src_offsets[4], src_sizes[4] },
9141+
vk_subbuffer{ d_srcs[5], src_offsets[5], src_sizes[5] },
9142+
vk_subbuffer{ d_srcs[6], src_offsets[6], src_sizes[6] },
9143+
vk_subbuffer{ d_D, dst_offset, dst_size }
9144+
}, pc, elements);
9145+
}
9146+
90309147
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) {
90319148
const ggml_tensor * x = dst->src[0];
90329149
const ggml_tensor * g = dst->src[1];
@@ -10859,6 +10976,7 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr
1085910976
case GGML_OP_CONV_2D_DW:
1086010977
case GGML_OP_RWKV_WKV6:
1086110978
case GGML_OP_RWKV_WKV7:
10979+
case GGML_OP_SSM_SCAN:
1086210980
case GGML_OP_LEAKY_RELU:
1086310981
case GGML_OP_FLASH_ATTN_EXT:
1086410982
case GGML_OP_OPT_STEP_ADAMW:
@@ -11276,6 +11394,11 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr
1127611394

1127711395
break;
1127811396

11397+
case GGML_OP_SSM_SCAN:
11398+
ggml_vk_ssm_scan(ctx, compute_ctx, node, dryrun);
11399+
11400+
break;
11401+
1127911402
case GGML_OP_OPT_STEP_ADAMW:
1128011403
ggml_vk_opt_step_adamw(ctx, compute_ctx, node, dryrun);
1128111404

@@ -11387,6 +11510,7 @@ static bool ggml_vk_compute_forward(ggml_backend_vk_context * ctx, ggml_cgraph *
1138711510
case GGML_OP_CONV_2D_DW:
1138811511
case GGML_OP_RWKV_WKV6:
1138911512
case GGML_OP_RWKV_WKV7:
11513+
case GGML_OP_SSM_SCAN:
1139011514
case GGML_OP_LEAKY_RELU:
1139111515
case GGML_OP_REPEAT:
1139211516
case GGML_OP_REPEAT_BACK:
@@ -12867,6 +12991,55 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
1286712991
case GGML_OP_RWKV_WKV6:
1286812992
case GGML_OP_RWKV_WKV7:
1286912993
return true;
12994+
case GGML_OP_SSM_SCAN:
12995+
{
12996+
if (ggml_is_quantized(op->src[0]->type) || ggml_is_quantized(op->src[1]->type) || ggml_is_quantized(op->src[2]->type)) {
12997+
return false;
12998+
}
12999+
if (op->src[3] && ggml_is_quantized(op->src[3]->type)) {
13000+
return false;
13001+
}
13002+
if (op->src[4] && ggml_is_quantized(op->src[4]->type)) {
13003+
return false;
13004+
}
13005+
if (op->src[5] && ggml_is_quantized(op->src[5]->type)) {
13006+
return false;
13007+
}
13008+
if (op->src[6] && op->src[6]->type != GGML_TYPE_I32) {
13009+
return false;
13010+
}
13011+
if (op->src[0]->type != GGML_TYPE_F32 || op->type != GGML_TYPE_F32) {
13012+
return false;
13013+
}
13014+
13015+
const uint32_t d_state = op->src[0]->ne[0];
13016+
const uint32_t head_dim = op->src[0]->ne[1];
13017+
13018+
bool is_mamba2 = (op->src[3] && op->src[3]->nb[1] == sizeof(float));
13019+
if (!is_mamba2) {
13020+
return false;
13021+
}
13022+
13023+
if ((d_state != 128 && d_state != 256) || head_dim % 16 != 0) {
13024+
return false;
13025+
}
13026+
13027+
ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context;
13028+
const vk_device& device = ggml_vk_get_device(ctx->device);
13029+
13030+
const uint32_t SPLIT_H = 16;
13031+
const uint32_t MAX_D_STATE = 256;
13032+
13033+
size_t stateC_size = SPLIT_H * MAX_D_STATE * sizeof(float);
13034+
size_t warp_sdata_size = MAX_D_STATE * sizeof(float);
13035+
size_t total_shared_memory = stateC_size + warp_sdata_size;
13036+
13037+
if (total_shared_memory > device->properties.limits.maxComputeSharedMemorySize) {
13038+
return false;
13039+
}
13040+
13041+
return true;
13042+
}
1287013043
case GGML_OP_CONV_TRANSPOSE_1D:
1287113044
return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32;
1287213045
case GGML_OP_CONV_2D:
@@ -13211,14 +13384,14 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph *
1321113384

1321213385
struct ggml_context * ggml_ctx = ggml_init(iparams);
1321313386

13214-
std::array<struct ggml_tensor *, 6> src_clone = {nullptr, nullptr, nullptr, nullptr, nullptr, nullptr};
13215-
std::array<size_t, 6> src_size = {0, 0, 0, 0, 0, 0};
13216-
std::array<void *, 6> src_buffer = {nullptr, nullptr, nullptr, nullptr, nullptr, nullptr};
13217-
const char * srci_name[6] = {"src0", "src1", "src2", "src3", "src4", "src5"};
13387+
std::array<struct ggml_tensor *, GGML_MAX_SRC> src_clone = {nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr};
13388+
std::array<size_t, GGML_MAX_SRC> src_size = {};
13389+
std::array<void *, GGML_MAX_SRC> src_buffer = {};
13390+
const char * srci_name[GGML_MAX_SRC] = {"src0", "src1", "src2", "src3", "src4", "src5", "src6", "src7", "src8", "src9"};
1321813391

1321913392
struct ggml_tensor * tensor_clone = nullptr;
1322013393

13221-
for (int i = 0; i < 6; i++) {
13394+
for (int i = 0; i < GGML_MAX_SRC; i++) {
1322213395
ggml_tensor * srci = tensor->src[i];
1322313396
if (fused_rms_norm_mul) {
1322413397
rms_norm_idx = tensor->src[0]->op == GGML_OP_RMS_NORM ? 0 : 1;
@@ -13525,6 +13698,9 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph *
1352513698
src_clone[2]);
1352613699
} else if (tensor->op == GGML_OP_ADD_ID) {
1352713700
tensor_clone = ggml_add_id(ggml_ctx, src_clone[0], src_clone[1], src_clone[2]);
13701+
} else if (tensor->op == GGML_OP_SSM_SCAN) {
13702+
tensor_clone = ggml_ssm_scan(ggml_ctx, src_clone[0], src_clone[1], src_clone[2],
13703+
src_clone[3], src_clone[4], src_clone[5], src_clone[6]);
1352813704
}
1352913705
else {
1353013706
std::cerr << "Missing vk_check_results OP: " << ggml_op_name(tensor->op) << std::endl;
@@ -13546,7 +13722,7 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph *
1354613722
memcpy(comp_result, tensor_clone->data, comp_size);
1354713723
memcpy(comp_nb, tensor_clone->nb, sizeof(size_t) * GGML_MAX_DIMS);
1354813724

13549-
for (int i = 0; i < 6; i++) {
13725+
for (int i = 0; i < GGML_MAX_SRC; i++) {
1355013726
if (src_buffer[i] != nullptr) {
1355113727
free(src_buffer[i]);
1355213728
}

0 commit comments

Comments
 (0)