@@ -722,6 +722,11 @@ struct vk_device_struct {
722722 vk_pipeline pipeline_soft_max_f32, pipeline_soft_max_f32_f16;
723723 vk_pipeline pipeline_soft_max_f32_wg512, pipeline_soft_max_f32_f16_wg512;
724724 vk_pipeline pipeline_soft_max_back_f32;
725+
726+ vk_pipeline pipeline_soft_max_large1_f32, pipeline_soft_max_large1_f32_f16;
727+ vk_pipeline pipeline_soft_max_large2_f32, pipeline_soft_max_large2_f32_f16;
728+ vk_pipeline pipeline_soft_max_large3_f32, pipeline_soft_max_large3_f32_f16;
729+
725730 vk_pipeline pipeline_rope_norm_f32, pipeline_rope_norm_f16, pipeline_rope_norm_f32_f16;
726731 vk_pipeline pipeline_rope_neox_f32, pipeline_rope_neox_f16, pipeline_rope_neox_f32_f16;
727732 vk_pipeline pipeline_rope_multi_f32, pipeline_rope_multi_f16;
@@ -3998,6 +4003,13 @@ static void ggml_vk_load_shaders(vk_device& device) {
39984003 ggml_vk_create_pipeline(device, device->pipeline_soft_max_f32_f16_wg512, "soft_max_f32_f16_wg512", soft_max_f32_f16_len, soft_max_f32_f16_data, "main", 4, sizeof(vk_op_soft_max_push_constants), {1, 1, 1}, { 512 }, 1);
39994004 ggml_vk_create_pipeline(device, device->pipeline_soft_max_back_f32, "soft_max_back_f32", soft_max_back_f32_len, soft_max_back_f32_data, "main", 3, sizeof(vk_op_push_constants), {1, 1, 1}, { device->subgroup_size }, 1, true);
40004005
4006+ ggml_vk_create_pipeline(device, device->pipeline_soft_max_large1_f32, "soft_max_large1_f32", soft_max_large1_f32_len, soft_max_large1_f32_data, "main", 6, sizeof(vk_op_soft_max_push_constants), {1, 1, 1}, { 128, 4 }, 1, true);
4007+ ggml_vk_create_pipeline(device, device->pipeline_soft_max_large2_f32, "soft_max_large2_f32", soft_max_large2_f32_len, soft_max_large2_f32_data, "main", 6, sizeof(vk_op_soft_max_push_constants), {1, 1, 1}, { 128, 4 }, 1, true);
4008+ ggml_vk_create_pipeline(device, device->pipeline_soft_max_large3_f32, "soft_max_large3_f32", soft_max_large3_f32_len, soft_max_large3_f32_data, "main", 6, sizeof(vk_op_soft_max_push_constants), {1, 1, 1}, { 128, 4 }, 1, true);
4009+ ggml_vk_create_pipeline(device, device->pipeline_soft_max_large1_f32_f16, "soft_max_large1_f32_f16", soft_max_large1_f32_f16_len, soft_max_large1_f32_f16_data, "main", 6, sizeof(vk_op_soft_max_push_constants), {1, 1, 1}, { 128, 4 }, 1, true);
4010+ ggml_vk_create_pipeline(device, device->pipeline_soft_max_large2_f32_f16, "soft_max_large2_f32_f16", soft_max_large2_f32_f16_len, soft_max_large2_f32_f16_data, "main", 6, sizeof(vk_op_soft_max_push_constants), {1, 1, 1}, { 128, 4 }, 1, true);
4011+ ggml_vk_create_pipeline(device, device->pipeline_soft_max_large3_f32_f16, "soft_max_large3_f32_f16", soft_max_large3_f32_f16_len, soft_max_large3_f32_f16_data, "main", 6, sizeof(vk_op_soft_max_push_constants), {1, 1, 1}, { 128, 4 }, 1, true);
4012+
40014013 ggml_vk_create_pipeline(device, device->pipeline_rope_norm_f32, "rope_norm_f32", rope_norm_f32_len, rope_norm_f32_data, "main", 5, sizeof(vk_op_rope_push_constants), {1, 512, 1}, {}, 1);
40024014 ggml_vk_create_pipeline(device, device->pipeline_rope_neox_f32, "rope_neox_f32", rope_neox_f32_len, rope_neox_f32_data, "main", 5, sizeof(vk_op_rope_push_constants), {1, 512, 1}, {}, 1);
40034015 ggml_vk_create_pipeline(device, device->pipeline_rope_multi_f32, "rope_multi_f32", rope_multi_f32_len, rope_multi_f32_data, "main", 5, sizeof(vk_op_rope_push_constants), {1, 512, 1}, {}, 1);
@@ -10117,7 +10129,7 @@ static void ggml_vk_soft_max(ggml_backend_vk_context * ctx, vk_context& subctx,
1011710129 const float m0 = powf(2.0f, -(max_bias ) / n_head_log2);
1011810130 const float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2);
1011910131
10120- ggml_vk_op_f32< vk_op_soft_max_push_constants>(ctx, subctx, src0, src1, src2, nullptr, dst, GGML_OP_SOFT_MAX, {
10132+ vk_op_soft_max_push_constants pc {
1012110133 ncols,
1012210134 src1 != nullptr ? nrows_y : (uint32_t)0,
1012310135 (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],
@@ -10128,7 +10140,55 @@ static void ggml_vk_soft_max(ggml_backend_vk_context * ctx, vk_context& subctx,
1012810140 n_head_log2,
1012910141 nrows_x,
1013010142 src2 != nullptr
10131- });
10143+ };
10144+
10145+ if (ncols <= 16384) {
10146+ ggml_vk_op_f32<vk_op_soft_max_push_constants>(ctx, subctx, src0, src1, src2, nullptr, dst, GGML_OP_SOFT_MAX, std::move(pc));
10147+ } else {
10148+
10149+ vk_subbuffer buf_a = ggml_vk_tensor_subbuffer(ctx, src0);
10150+ vk_subbuffer buf_b = src1 ? ggml_vk_tensor_subbuffer(ctx, src1) : buf_a;
10151+ vk_subbuffer buf_c = src2 ? ggml_vk_tensor_subbuffer(ctx, src2) : buf_a;
10152+ vk_subbuffer buf_d = ggml_vk_tensor_subbuffer(ctx, dst);
10153+
10154+ uint32_t elems_per_wg = 128 * 4;
10155+ uint32_t num_wgs = CEIL_DIV(ncols, elems_per_wg);
10156+ size_t tmp_size = num_wgs * nrows_x * sizeof(float);
10157+
10158+ if (ctx->prealloc_size_x < tmp_size) {
10159+ ctx->prealloc_size_x = tmp_size;
10160+ ggml_vk_preallocate_buffers(ctx, subctx);
10161+ }
10162+ if (ctx->prealloc_size_y < tmp_size) {
10163+ ctx->prealloc_size_y = tmp_size;
10164+ ggml_vk_preallocate_buffers(ctx, subctx);
10165+ }
10166+ if (ctx->prealloc_x_need_sync || ctx->prealloc_y_need_sync) {
10167+ ggml_vk_sync_buffers(ctx, subctx);
10168+ }
10169+
10170+ vk_subbuffer buf_x = { ctx->prealloc_x, 0, tmp_size };
10171+ vk_subbuffer buf_y = { ctx->prealloc_y, 0, tmp_size };
10172+
10173+ std::array<uint32_t, 3> elements = { num_wgs, nrows_x, 1 };
10174+
10175+ vk_pipeline pipeline1 = src1 && src1->type == GGML_TYPE_F16 ? ctx->device->pipeline_soft_max_large1_f32_f16 : ctx->device->pipeline_soft_max_large1_f32;
10176+ vk_pipeline pipeline2 = src1 && src1->type == GGML_TYPE_F16 ? ctx->device->pipeline_soft_max_large2_f32_f16 : ctx->device->pipeline_soft_max_large2_f32;
10177+ vk_pipeline pipeline3 = src1 && src1->type == GGML_TYPE_F16 ? ctx->device->pipeline_soft_max_large3_f32_f16 : ctx->device->pipeline_soft_max_large3_f32;
10178+
10179+ ggml_pipeline_request_descriptor_sets(ctx, pipeline1, 1);
10180+ ggml_pipeline_request_descriptor_sets(ctx, pipeline2, 1);
10181+ ggml_pipeline_request_descriptor_sets(ctx, pipeline3, 1);
10182+
10183+ ggml_vk_dispatch_pipeline(ctx, subctx, pipeline1, { buf_a, buf_b, buf_c, buf_d, buf_x, buf_y }, pc, elements);
10184+ ggml_vk_sync_buffers(ctx, subctx);
10185+ ggml_vk_dispatch_pipeline(ctx, subctx, pipeline2, { buf_a, buf_b, buf_c, buf_d, buf_x, buf_y }, pc, elements);
10186+ ggml_vk_sync_buffers(ctx, subctx);
10187+ ggml_vk_dispatch_pipeline(ctx, subctx, pipeline3, { buf_a, buf_b, buf_c, buf_d, buf_x, buf_y }, pc, elements);
10188+
10189+ ctx->prealloc_x_need_sync = true;
10190+ ctx->prealloc_y_need_sync = true;
10191+ }
1013210192}
1013310193
1013410194static void ggml_vk_soft_max_back(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
0 commit comments