Skip to content

Commit bf6399c

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 4e0388a commit bf6399c

File tree

3 files changed

+327
-6
lines changed

3 files changed

+327
-6
lines changed

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

Lines changed: 170 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_d16;
586+
vk_pipeline pipeline_ssm_scan_f32_d128;
587+
vk_pipeline pipeline_ssm_scan_f32_d256;
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,14 @@ 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+
// Match CUDA parameter structure exactly - simplified to only needed strides
1095+
uint32_t src0_nb2, src0_nb3, src1_nb2, src1_nb3;
1096+
uint32_t src2_nb1, src2_nb2, src3_nb1;
1097+
uint32_t src4_nb2, src4_nb3, src5_nb2, src5_nb3;
1098+
uint32_t s_off;
1099+
uint32_t n_head, d_head, n_group, n_tok;
1100+
};
10901101

10911102
struct vk_op_conv2d_push_constants {
10921103
uint32_t Cout;
@@ -3588,6 +3599,10 @@ static void ggml_vk_load_shaders(vk_device& device) {
35883599

35893600
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);
35903601

3602+
ggml_vk_create_pipeline(device, device->pipeline_ssm_scan_f32_d16, "ssm_scan_f32_d16", ssm_scan_f32_d16_len, ssm_scan_f32_d16_data, "main", 8, sizeof(vk_op_ssm_scan_push_constants), {1, 1, 1}, {16, 1, 1}, 1);
3603+
ggml_vk_create_pipeline(device, device->pipeline_ssm_scan_f32_d128, "ssm_scan_f32_d128", ssm_scan_f32_d128_len, ssm_scan_f32_d128_data, "main", 8, sizeof(vk_op_ssm_scan_push_constants), {1, 1, 1}, {128, 1, 1}, 1);
3604+
ggml_vk_create_pipeline(device, device->pipeline_ssm_scan_f32_d256, "ssm_scan_f32_d256", ssm_scan_f32_d256_len, ssm_scan_f32_d256_data, "main", 8, sizeof(vk_op_ssm_scan_push_constants), {1, 1, 1}, {256, 1, 1}, 1);
3605+
35913606
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);
35923607

35933608
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 +8102,21 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
80878102
return ctx->device->pipeline_rwkv_wkv7_f32;
80888103
}
80898104
return nullptr;
8105+
case GGML_OP_SSM_SCAN:
8106+
if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
8107+
const uint32_t d_state = src0->ne[0];
8108+
if (d_state == 16) {
8109+
return ctx->device->pipeline_ssm_scan_f32_d16;
8110+
} else if (d_state == 128) {
8111+
return ctx->device->pipeline_ssm_scan_f32_d128;
8112+
} else if (d_state == 256) {
8113+
return ctx->device->pipeline_ssm_scan_f32_d256;
8114+
} else {
8115+
GGML_LOG_ERROR("Unsupported d_state=%d for SSM scan. Supported values: 16, 128, 256\n", d_state);
8116+
return nullptr;
8117+
}
8118+
}
8119+
return nullptr;
80908120
case GGML_OP_OPT_STEP_ADAMW:
80918121
if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
80928122
return ctx->device->pipeline_opt_step_adamw_f32;
@@ -9027,6 +9057,129 @@ static void ggml_vk_rwkv_wkv7(ggml_backend_vk_context * ctx, vk_context& subctx,
90279057
);
90289058
}
90299059

9060+
static void ggml_vk_ssm_scan(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst, bool dryrun = false) {
9061+
const ggml_tensor * src0 = dst->src[0];
9062+
const ggml_tensor * src1 = dst->src[1];
9063+
const ggml_tensor * src2 = dst->src[2];
9064+
const ggml_tensor * src3 = dst->src[3];
9065+
const ggml_tensor * src4 = dst->src[4];
9066+
const ggml_tensor * src5 = dst->src[5];
9067+
const ggml_tensor * src6 = dst->src[6];
9068+
9069+
GGML_ASSERT(!ggml_is_quantized(src0->type));
9070+
GGML_ASSERT(!ggml_is_quantized(src1->type));
9071+
GGML_ASSERT(!ggml_is_quantized(src2->type));
9072+
GGML_ASSERT(!ggml_is_quantized(src3->type));
9073+
GGML_ASSERT(!ggml_is_quantized(src4->type));
9074+
GGML_ASSERT(!ggml_is_quantized(src5->type));
9075+
GGML_ASSERT(src6->type == GGML_TYPE_I32);
9076+
GGML_ASSERT(dst->buffer != nullptr);
9077+
9078+
const uint32_t d_state = src0->ne[0];
9079+
const uint32_t head_dim = src0->ne[1];
9080+
const uint32_t n_head = src1->ne[1];
9081+
const uint32_t n_group = src4->ne[1];
9082+
const uint32_t n_tok = src1->ne[2];
9083+
const uint32_t n_seq = src1->ne[3];
9084+
9085+
bool is_mamba2 = (src3->nb[1] == sizeof(float));
9086+
if (is_mamba2) {
9087+
if (d_state == 128 || d_state == 256) {
9088+
GGML_ASSERT(head_dim % 16 == 0);
9089+
} else {
9090+
GGML_ABORT("doesn't support d_state!=(128 or 256).");
9091+
}
9092+
} else {
9093+
GGML_ASSERT(n_head % 128 == 0);
9094+
GGML_ASSERT(head_dim == 1);
9095+
GGML_ASSERT(n_group == 1);
9096+
if (d_state != 16) {
9097+
GGML_ABORT("doesn't support d_state!=16.");
9098+
}
9099+
}
9100+
9101+
vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, src0, src1, src2, dst, dst->op);
9102+
GGML_ASSERT(pipeline != nullptr);
9103+
9104+
if (dryrun) {
9105+
ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
9106+
return;
9107+
}
9108+
9109+
const int64_t s_off = ggml_nelements(src1) * sizeof(float);
9110+
9111+
const vk_op_ssm_scan_push_constants pc = {
9112+
(uint32_t)src0->nb[2], (uint32_t)src0->nb[3],
9113+
(uint32_t)src1->nb[2], (uint32_t)src1->nb[3],
9114+
(uint32_t)src2->nb[1], (uint32_t)src2->nb[2],
9115+
(uint32_t)src3->nb[1],
9116+
(uint32_t)src4->nb[2], (uint32_t)src4->nb[3],
9117+
(uint32_t)src5->nb[2], (uint32_t)src5->nb[3],
9118+
(uint32_t)s_off,
9119+
n_head, head_dim, n_group, n_tok
9120+
};
9121+
9122+
ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context;
9123+
ggml_backend_vk_buffer_context * src_buf_ctxs[7];
9124+
for (int i = 0; i < 7; i++) {
9125+
src_buf_ctxs[i] = (ggml_backend_vk_buffer_context *)dst->src[i]->buffer->context;
9126+
}
9127+
9128+
vk_buffer d_D = nullptr, d_srcs[7] = { nullptr };
9129+
size_t dst_offset = 0, src_offsets[7] = { 0 };
9130+
bool dst_uma = false, srcs_uma[7] = { false };
9131+
9132+
if (ctx->device->uma) {
9133+
for (int i = 0; i < 7; i++) {
9134+
ggml_vk_host_get(ctx->device, dst->src[i]->data, d_srcs[i], src_offsets[i]);
9135+
srcs_uma[i] = d_srcs[i] != nullptr;
9136+
}
9137+
ggml_vk_host_get(ctx->device, dst->data, d_D, dst_offset);
9138+
dst_uma = d_D != nullptr;
9139+
}
9140+
9141+
if (!dst_uma) {
9142+
d_D = dst_buf_ctx->dev_buffer;
9143+
dst_offset = vk_tensor_offset(dst) + dst->view_offs;
9144+
}
9145+
for (int i = 0; i < 7; i++) {
9146+
if (!srcs_uma[i]) {
9147+
d_srcs[i] = src_buf_ctxs[i]->dev_buffer;
9148+
src_offsets[i] = vk_tensor_offset(dst->src[i]) + dst->src[i]->view_offs;
9149+
}
9150+
}
9151+
9152+
size_t dst_size = ggml_nbytes(dst);
9153+
size_t src_sizes[7];
9154+
for (int i = 0; i < 7; i++) {
9155+
src_sizes[i] = ggml_nbytes(dst->src[i]);
9156+
}
9157+
9158+
std::array<uint32_t, 3> elements;
9159+
9160+
if (is_mamba2) {
9161+
const int splitH = 16;
9162+
const uint32_t num_workgroups_x = (n_head * head_dim + (splitH - 1)) / splitH;
9163+
const uint32_t num_workgroups_y = n_seq;
9164+
elements = { num_workgroups_x, num_workgroups_y, 1 };
9165+
} else {
9166+
const uint32_t num_workgroups_x = n_seq;
9167+
const uint32_t num_workgroups_y = (n_head + 128 - 1) / 128;
9168+
elements = { num_workgroups_x, num_workgroups_y, 1 };
9169+
}
9170+
9171+
ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, {
9172+
vk_subbuffer{ d_srcs[0], src_offsets[0], src_sizes[0] },
9173+
vk_subbuffer{ d_srcs[1], src_offsets[1], src_sizes[1] },
9174+
vk_subbuffer{ d_srcs[2], src_offsets[2], src_sizes[2] },
9175+
vk_subbuffer{ d_srcs[3], src_offsets[3], src_sizes[3] },
9176+
vk_subbuffer{ d_srcs[4], src_offsets[4], src_sizes[4] },
9177+
vk_subbuffer{ d_srcs[5], src_offsets[5], src_sizes[5] },
9178+
vk_subbuffer{ d_srcs[6], src_offsets[6], src_sizes[6] },
9179+
vk_subbuffer{ d_D, dst_offset, dst_size }
9180+
}, pc, elements);
9181+
}
9182+
90309183
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) {
90319184
const ggml_tensor * x = dst->src[0];
90329185
const ggml_tensor * g = dst->src[1];
@@ -10859,6 +11012,7 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr
1085911012
case GGML_OP_CONV_2D_DW:
1086011013
case GGML_OP_RWKV_WKV6:
1086111014
case GGML_OP_RWKV_WKV7:
11015+
case GGML_OP_SSM_SCAN:
1086211016
case GGML_OP_LEAKY_RELU:
1086311017
case GGML_OP_FLASH_ATTN_EXT:
1086411018
case GGML_OP_OPT_STEP_ADAMW:
@@ -11276,6 +11430,11 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr
1127611430

1127711431
break;
1127811432

11433+
case GGML_OP_SSM_SCAN:
11434+
ggml_vk_ssm_scan(ctx, compute_ctx, node, dryrun);
11435+
11436+
break;
11437+
1127911438
case GGML_OP_OPT_STEP_ADAMW:
1128011439
ggml_vk_opt_step_adamw(ctx, compute_ctx, node, dryrun);
1128111440

@@ -11387,6 +11546,7 @@ static bool ggml_vk_compute_forward(ggml_backend_vk_context * ctx, ggml_cgraph *
1138711546
case GGML_OP_CONV_2D_DW:
1138811547
case GGML_OP_RWKV_WKV6:
1138911548
case GGML_OP_RWKV_WKV7:
11549+
case GGML_OP_SSM_SCAN:
1139011550
case GGML_OP_LEAKY_RELU:
1139111551
case GGML_OP_REPEAT:
1139211552
case GGML_OP_REPEAT_BACK:
@@ -12866,6 +13026,7 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
1286613026
case GGML_OP_POOL_2D:
1286713027
case GGML_OP_RWKV_WKV6:
1286813028
case GGML_OP_RWKV_WKV7:
13029+
case GGML_OP_SSM_SCAN:
1286913030
return true;
1287013031
case GGML_OP_CONV_TRANSPOSE_1D:
1287113032
return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32;
@@ -13211,14 +13372,14 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph *
1321113372

1321213373
struct ggml_context * ggml_ctx = ggml_init(iparams);
1321313374

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"};
13375+
std::array<struct ggml_tensor *, 7> src_clone = {nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr};
13376+
std::array<size_t, 7> src_size = {0, 0, 0, 0, 0, 0, 0};
13377+
std::array<void *, 7> src_buffer = {nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr};
13378+
const char * srci_name[7] = {"src0", "src1", "src2", "src3", "src4", "src5", "src6"};
1321813379

1321913380
struct ggml_tensor * tensor_clone = nullptr;
1322013381

13221-
for (int i = 0; i < 6; i++) {
13382+
for (int i = 0; i < 7; i++) {
1322213383
ggml_tensor * srci = tensor->src[i];
1322313384
if (fused_rms_norm_mul) {
1322413385
rms_norm_idx = tensor->src[0]->op == GGML_OP_RMS_NORM ? 0 : 1;
@@ -13525,6 +13686,9 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph *
1352513686
src_clone[2]);
1352613687
} else if (tensor->op == GGML_OP_ADD_ID) {
1352713688
tensor_clone = ggml_add_id(ggml_ctx, src_clone[0], src_clone[1], src_clone[2]);
13689+
} else if (tensor->op == GGML_OP_SSM_SCAN) {
13690+
tensor_clone = ggml_ssm_scan(ggml_ctx, src_clone[0], src_clone[1], src_clone[2],
13691+
src_clone[3], src_clone[4], src_clone[5], src_clone[6]);
1352813692
}
1352913693
else {
1353013694
std::cerr << "Missing vk_check_results OP: " << ggml_op_name(tensor->op) << std::endl;
@@ -13546,7 +13710,7 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph *
1354613710
memcpy(comp_result, tensor_clone->data, comp_size);
1354713711
memcpy(comp_nb, tensor_clone->nb, sizeof(size_t) * GGML_MAX_DIMS);
1354813712

13549-
for (int i = 0; i < 6; i++) {
13713+
for (int i = 0; i < 7; i++) {
1355013714
if (src_buffer[i] != nullptr) {
1355113715
free(src_buffer[i]);
1355213716
}

0 commit comments

Comments
 (0)