From 4951447aea1d52620701f48168b6ae5efdfba163 Mon Sep 17 00:00:00 2001 From: Jeff Bolz Date: Sun, 30 Aug 2026 10:49:49 -0500 Subject: [PATCH 1/2] vulkan: topk --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 76 ++++- .../vulkan-shaders/topk_global.comp | 293 ++++++++++++++++++ .../vulkan-shaders/vulkan-shaders-gen.cpp | 2 + tests/test-backend-ops.cpp | 11 +- 4 files changed, 377 insertions(+), 5 deletions(-) create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/topk_global.comp diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 8fbb1359f40..0791cdbf187 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -626,6 +626,11 @@ enum shader_reduction_mode { static constexpr uint32_t num_argsort_pipelines = 11; static constexpr uint32_t num_topk_moe_pipelines = 10; static constexpr uint32_t num_topk_pipelines = 11; +// Keep these values in sync with topk_global.comp. +static constexpr uint32_t topk_global_radix_bits = 5; +static constexpr uint32_t topk_global_num_buckets = 1u << topk_global_radix_bits; +static constexpr uint32_t topk_global_num_passes = (32 + topk_global_radix_bits - 1) / topk_global_radix_bits; +static constexpr uint32_t topk_global_chunk_size = 256; static constexpr std::initializer_list topk_moe_early_softmax_norm{ GGML_OP_SOFT_MAX, GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS, GGML_OP_RESHAPE, @@ -1056,6 +1061,8 @@ struct vk_device_struct { vk_pipeline pipeline_argsort_f32[num_argsort_pipelines]; vk_pipeline pipeline_argsort_large_f32[num_argsort_pipelines]; vk_pipeline pipeline_topk_f32[num_topk_pipelines]; + vk_pipeline pipeline_topk_global_f32; + vk_pipeline pipeline_topk_global_subgroup_f32; vk_pipeline pipeline_sum_rows_f32; vk_pipeline pipeline_cross_entropy_loss_f32, pipeline_cross_entropy_loss_f32_wg512; vk_pipeline pipeline_cross_entropy_loss_back_f32, pipeline_cross_entropy_loss_back_f32_wg512; @@ -1748,6 +1755,13 @@ struct vk_op_topk_push_constants { uint32_t last_pass; }; +struct vk_op_topk_global_push_constants { + uint32_t ncols; + uint32_t k; + uint32_t nrows; + uint32_t pass; +}; + struct vk_op_im2col_push_constants { uint64_t dst_addr; uint32_t batch_offset; uint32_t offset_delta; @@ -5812,6 +5826,15 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { } } + const uint32_t topk_global_block_size = 1u << std::min(device->max_workgroup_size_log2, 7u); + const uint32_t topk_global_items_per_thread = topk_global_chunk_size / topk_global_block_size; + ggml_vk_create_pipeline2(device, device->pipeline_topk_global_f32, "topk_global_f32", topk_global_f32_len, topk_global_f32_data, "main", 3, sizeof(vk_op_topk_global_push_constants), {topk_global_chunk_size, 1, 1}, + {topk_global_block_size, topk_global_items_per_thread}, 1, true); + if (device->subgroup_basic && device->subgroup_arithmetic && device->subgroup_ballot && device->subgroup_require_full_support && device->subgroup_size >= topk_global_num_buckets) { + ggml_vk_create_pipeline2(device, device->pipeline_topk_global_subgroup_f32, "topk_global_subgroup_f32", topk_global_subgroup_f32_len, topk_global_subgroup_f32_data, "main", 3, + sizeof(vk_op_topk_global_push_constants), {topk_global_chunk_size, 1, 1}, {topk_global_block_size, topk_global_items_per_thread}, 1, true, true, device->subgroup_size); + } + ggml_vk_create_pipeline(device, device->pipeline_argmax_f32, "argmax_f32", argmax_f32_len, argmax_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, { device->subgroup_size }, 1); ggml_vk_create_pipeline(device, device->pipeline_sum_rows_f32, "sum_rows_f32", sum_rows_f32_len, sum_rows_f32_data, "main", 2, sizeof(vk_op_sum_rows_push_constants), {1, 1, 1}, { device->subgroup_size }, 1); @@ -13931,11 +13954,60 @@ static void ggml_vk_argsort(ggml_backend_vk_context * ctx, vk_context& subctx, c } } +static void ggml_vk_topk_global(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst) { + const uint32_t ncols = src0->ne[0]; + const uint32_t nrows = ggml_nrows(src0); + const uint32_t k = dst->ne[0]; + const size_t scratch_size = size_t{nrows} * (topk_global_num_passes * topk_global_num_buckets + 8) * sizeof(uint32_t); + + if (ctx->prealloc_size_x < scratch_size) { + ctx->prealloc_size_x = scratch_size; + ggml_vk_preallocate_buffers(ctx, subctx); + } + if (ctx->prealloc_x_need_sync) { + ggml_vk_sync_buffers(ctx, subctx); + } + + vk_subbuffer src_buf = ggml_vk_tensor_subbuffer(ctx, src0); + vk_subbuffer scratch_buf = { ctx->prealloc_x, 0, scratch_size }; + vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst); + vk_pipeline pipeline = ctx->device->pipeline_topk_global_subgroup_f32 ? ctx->device->pipeline_topk_global_subgroup_f32 : ctx->device->pipeline_topk_global_f32; + GGML_ASSERT(pipeline); + + subctx->s->buffer->buf.fillBuffer(scratch_buf.buffer->buffer, scratch_buf.offset, scratch_buf.size, 0); + ggml_vk_sync_buffers(ctx, subctx); + + std::array elements = { + ncols, + std::min(nrows, ctx->device->properties.limits.maxComputeWorkGroupCount[1]), + 1, + }; + vk_op_topk_global_push_constants pc { ncols, k, nrows, 0 }; + + for (uint32_t pass = 0; pass < topk_global_num_passes; ++pass) { + pc.pass = pass; + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { src_buf, scratch_buf, dst_buf }, pc, elements); + ggml_vk_sync_buffers(ctx, subctx); + } + + pc.pass = topk_global_num_passes; + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { src_buf, scratch_buf, dst_buf }, pc, elements); + ctx->prealloc_x_need_sync = true; +} + static void ggml_vk_topk(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst) { uint32_t ncols = src0->ne[0]; uint32_t nrows = ggml_nrows(src0); uint32_t k = dst->ne[0]; + uint32_t min_pipeline = (uint32_t)log2f(float(k)) + 1; + if (min_pipeline >= num_topk_pipelines || !ctx->device->pipeline_topk_f32[min_pipeline]) { + ggml_vk_topk_global(ctx, subctx, src0, dst); + return; + } + vk_op_topk_push_constants pc { ncols, ncols, ncols, k, nrows, 0, 0 }; if (ctx->prealloc_x_need_sync) { @@ -18717,12 +18789,10 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm if (!ggml_is_contiguous(op) || !ggml_is_contiguous(op->src[0])) { return false; } - // We could potentially support larger, using argsort to sort the - // whole thing. Not clear if this is needed. uint32_t min_pipeline = (uint32_t)log2f(float(op->ne[0])) + 1; if (min_pipeline >= num_topk_pipelines || !device->pipeline_topk_f32[min_pipeline]) { - return false; + return device->pipeline_topk_global_subgroup_f32 != nullptr || device->pipeline_topk_global_f32 != nullptr; } } return true; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/topk_global.comp b/ggml/src/ggml-vulkan/vulkan-shaders/topk_global.comp new file mode 100644 index 00000000000..3f4284d764f --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/topk_global.comp @@ -0,0 +1,293 @@ +#version 450 + +#if USE_SUBGROUP +#extension GL_KHR_shader_subgroup_arithmetic : require +#extension GL_KHR_shader_subgroup_ballot : require +#extension GL_KHR_shader_subgroup_basic : require +#endif + +#include "types.glsl" + +layout(constant_id = 0) const int BLOCK_SIZE = 128; +layout(constant_id = 1) const int ITEMS_PER_THREAD = 2; + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) buffer S {uint data_s[];}; +layout (binding = 2) writeonly buffer D {int data_d[];}; + +layout (push_constant) uniform parameter { + uint ncols; + uint k; + uint nrows; + uint pass; +} p; + +const uint RADIX_BITS = 5; +const uint NUM_BUCKETS = 1 << RADIX_BITS; +const uint NUM_PASSES = (32 + RADIX_BITS - 1) / RADIX_BITS; +const uint LAST_BITS = 32 - (NUM_PASSES - 1) * RADIX_BITS; +const uint CHUNK_SIZE = BLOCK_SIZE * ITEMS_PER_THREAD; + +shared uint histogram[NUM_BUCKETS]; +#if !USE_SUBGROUP +shared uint scan_values[BLOCK_SIZE]; +shared uint scan_total; +#endif +shared uint selected_prefix; +shared uint selected_rank; +shared uint selected_greater; +shared uint output_greater_base; +shared uint output_equal_base; +#if USE_SUBGROUP +shared uint subgroup_greater_counts[BLOCK_SIZE]; +shared uint subgroup_equal_counts[BLOCK_SIZE]; +#endif + +uint float_key(float x) { + uint y = floatBitsToUint(x); + return (y & 0x80000000) != 0 ? ~y : y | 0x80000000; +} + +uint histogram_index(uint row, uint pass, uint bucket) { + return (row * NUM_PASSES + pass) * NUM_BUCKETS + bucket; +} + +uint state_index(uint row, uint slot) { + return p.nrows * NUM_PASSES * NUM_BUCKETS + row * 6 + slot * 3; +} + +uint counter_index(uint row) { + return p.nrows * (NUM_PASSES * NUM_BUCKETS + 6) + row * 2; +} + +void find_next_prefix(uint row, uint num_passes) { + uint tid = gl_LocalInvocationID.x; + if (tid == 0) { + selected_prefix = 0; + selected_rank = p.k; + selected_greater = 0; + + if (num_passes > 1) { + uint state = state_index(row, (num_passes - 2) & 1); + selected_prefix = data_s[state]; + selected_rank = data_s[state + 1]; + selected_greater = data_s[state + 2]; + } + } + barrier(); + + if (num_passes == 0) { + return; + } + + uint pass = num_passes - 1; + uint shift = pass == NUM_PASSES - 1 ? 0 : 32 - (pass + 1) * RADIX_BITS; +#if USE_SUBGROUP + if (gl_SubgroupID == 0) { + uint lane = gl_SubgroupInvocationID; + uint prefix = 0; + uint rank = 0; + uint greater = 0; + if (lane == 0) { + prefix = selected_prefix; + rank = selected_rank; + greater = selected_greater; + } + prefix = subgroupBroadcastFirst(prefix); + rank = subgroupBroadcastFirst(rank); + greater = subgroupBroadcastFirst(greater); + for (uint base = 0; base < NUM_BUCKETS; base += gl_SubgroupSize) { + uint bucket_index = base + lane; + uint count = bucket_index < NUM_BUCKETS ? data_s[histogram_index(row, pass, NUM_BUCKETS - 1 - bucket_index)] : 0; + uint sum = subgroupInclusiveAdd(count); + uint total = subgroupAdd(count); + uvec4 candidates = subgroupBallot(bucket_index < NUM_BUCKETS && sum >= rank); + if (subgroupBallotBitCount(candidates) != 0) { + uint selected_lane = subgroupBallotFindLSB(candidates); + uint above = subgroupBroadcast(sum - count, selected_lane); + prefix |= (NUM_BUCKETS - 1 - (base + selected_lane)) << shift; + rank -= above; + greater += above; + break; + } + rank -= total; + greater += total; + } + if (lane == 0) { + selected_prefix = prefix; + selected_rank = rank; + selected_greater = greater; + } + } +#else + if (tid == 0) { + for (int bucket = int(NUM_BUCKETS) - 1; bucket >= 0; --bucket) { + uint count = data_s[histogram_index(row, pass, uint(bucket))]; + if (selected_rank > count) { + selected_rank -= count; + selected_greater += count; + } else { + selected_prefix |= uint(bucket) << shift; + break; + } + } + } +#endif + barrier(); +} + +#if !USE_SUBGROUP +uint workgroup_exclusive_scan(uint value) { + uint tid = gl_LocalInvocationID.x; + scan_values[tid] = value; + barrier(); + + for (uint offset = 1; offset < BLOCK_SIZE; offset <<= 1) { + uint add = tid >= offset ? scan_values[tid - offset] : 0; + barrier(); + scan_values[tid] += add; + barrier(); + } + + if (tid == BLOCK_SIZE - 1) { + scan_total = scan_values[tid]; + } + uint result = scan_values[tid] - value; + barrier(); + return result; +} +#endif + +void count_pass(uint row) { + uint tid = gl_LocalInvocationID.x; + for (uint bucket = tid; bucket < NUM_BUCKETS; bucket += BLOCK_SIZE) { + histogram[bucket] = 0; + } + + find_next_prefix(row, p.pass); + if (p.pass > 0 && tid == 0 && gl_WorkGroupID.x == 0) { + uint state = state_index(row, (p.pass - 1) & 1); + data_s[state] = selected_prefix; + data_s[state + 1] = selected_rank; + data_s[state + 2] = selected_greater; + } + + uint prefix_mask = 0; + if (p.pass > 0) { + prefix_mask = ~0u << (32 - p.pass * RADIX_BITS); + } + uint shift = p.pass == NUM_PASSES - 1 ? 0 : 32 - (p.pass + 1) * RADIX_BITS; + uint bucket_mask = p.pass == NUM_PASSES - 1 ? (1u << LAST_BITS) - 1 : NUM_BUCKETS - 1; + uint chunk_base = gl_WorkGroupID.x * CHUNK_SIZE; + + for (uint item = 0; item < ITEMS_PER_THREAD; ++item) { + uint col = chunk_base + item * BLOCK_SIZE + tid; + if (col < p.ncols) { + uint key = float_key(data_a[row * p.ncols + col]); + if (p.pass == 0 || (key & prefix_mask) == selected_prefix) { + atomicAdd(histogram[(key >> shift) & bucket_mask], 1); + } + } + } + barrier(); + + for (uint bucket = tid; bucket < NUM_BUCKETS; bucket += BLOCK_SIZE) { + if (histogram[bucket] != 0) { + atomicAdd(data_s[histogram_index(row, p.pass, bucket)], histogram[bucket]); + } + } +} + +void compact_pass(uint row) { + uint tid = gl_LocalInvocationID.x; + find_next_prefix(row, NUM_PASSES); + + uint chunk_base = gl_WorkGroupID.x * CHUNK_SIZE; + uint keys[ITEMS_PER_THREAD]; + uint greater_count = 0; + uint equal_count = 0; + for (uint item = 0; item < ITEMS_PER_THREAD; ++item) { + uint col = chunk_base + item * BLOCK_SIZE + tid; + keys[item] = col < p.ncols ? float_key(data_a[row * p.ncols + col]) : 0; + greater_count += col < p.ncols && keys[item] > selected_prefix ? 1 : 0; + equal_count += col < p.ncols && keys[item] == selected_prefix ? 1 : 0; + } + +#if USE_SUBGROUP + uint greater_offset = subgroupExclusiveAdd(greater_count); + uint equal_offset = subgroupExclusiveAdd(equal_count); + uint subgroup_greater = subgroupAdd(greater_count); + uint subgroup_equal = subgroupAdd(equal_count); + if (subgroupElect()) { + subgroup_greater_counts[gl_SubgroupID] = subgroup_greater; + subgroup_equal_counts[gl_SubgroupID] = subgroup_equal; + } + barrier(); + + if (tid == 0) { + uint greater_total = 0; + uint equal_total = 0; + for (uint subgroup = 0; subgroup < gl_NumSubgroups; ++subgroup) { + uint next_greater = greater_total + subgroup_greater_counts[subgroup]; + uint next_equal = equal_total + subgroup_equal_counts[subgroup]; + subgroup_greater_counts[subgroup] = greater_total; + subgroup_equal_counts[subgroup] = equal_total; + greater_total = next_greater; + equal_total = next_equal; + } + uint counters = counter_index(row); + output_greater_base = atomicAdd(data_s[counters], greater_total); + output_equal_base = atomicAdd(data_s[counters + 1], equal_total); + } + barrier(); + greater_offset += subgroup_greater_counts[gl_SubgroupID]; + equal_offset += subgroup_equal_counts[gl_SubgroupID]; +#else + uint greater_offset = workgroup_exclusive_scan(greater_count); + uint workgroup_greater = scan_total; + uint equal_offset = workgroup_exclusive_scan(equal_count); + uint workgroup_equal = scan_total; + + if (tid == 0) { + uint counters = counter_index(row); + output_greater_base = atomicAdd(data_s[counters], workgroup_greater); + output_equal_base = atomicAdd(data_s[counters + 1], workgroup_equal); + } + barrier(); +#endif + + uint local_greater = 0; + uint local_equal = 0; + for (uint item = 0; item < ITEMS_PER_THREAD; ++item) { + uint col = chunk_base + item * BLOCK_SIZE + tid; + if (col >= p.ncols) { + continue; + } + if (keys[item] > selected_prefix) { + uint output_idx = output_greater_base + greater_offset + local_greater++; + if (output_idx < selected_greater) { + data_d[row * p.k + output_idx] = int(col); + } + } else if (keys[item] == selected_prefix) { + uint output_idx = output_equal_base + equal_offset + local_equal++; + if (output_idx < selected_rank) { + data_d[row * p.k + selected_greater + output_idx] = int(col); + } + } + } +} + +void main() { + uint row = gl_WorkGroupID.y; + while (row < p.nrows) { + if (p.pass < NUM_PASSES) { + count_pass(row); + } else { + compact_pass(row); + } + barrier(); + row += gl_NumWorkGroups.y; + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp index d375c2d1277..072412ab746 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -1025,6 +1025,8 @@ void process_shaders() { string_to_spv("argsort_large_f32", "argsort_large.comp", {{"A_TYPE", "float"}}); string_to_spv("topk_argsort_f32", "topk_argsort.comp", {{"A_TYPE", "float"}}); + string_to_spv("topk_global_f32", "topk_global.comp", {{"A_TYPE", "float"}, {"USE_SUBGROUP", "0"}}); + string_to_spv("topk_global_subgroup_f32", "topk_global.comp", {{"A_TYPE", "float"}, {"USE_SUBGROUP", "1"}}); string_to_spv("topk_nary_search_f32", "topk_nary_search.comp", {{"A_TYPE", "float"}}); string_to_spv("argmax_f32", "argmax.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "int"}})); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 4a7a0623174..bc7c6310cc0 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9723,6 +9723,11 @@ static std::vector> make_test_cases_eval() { } } } + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {1024, 1, 1, 1}, 1024)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {2048, 2, 1, 1}, 1024)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {2048, 2, 1, 1}, 2048, true)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {4096, 1, 1, 1}, 2048)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {8192, 2, 1, 1}, 2051, true)); for (int k : {4, 8, 16, 32}) { for (int nrows : {1, 8, 16}) { test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {202048, nrows, 1, 1}, k)); @@ -10477,10 +10482,12 @@ static std::vector> make_test_cases_perf() { test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {cols, nrows, 1, 1}, 16)); } } - for (auto k : {1, 4, 8, 10, 16, 32, 40, 400}) { + for (auto k : {1, 4, 8, 10, 16, 32, 40, 400, 2048, 9999}) { for (auto nrows : {1, 16}) { for (auto cols : {k, 1000, 65000, 200000}) { - test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {cols, nrows, 1, 1}, k)); + if (cols >= k) { + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {cols, nrows, 1, 1}, k)); + } } } } From 7a3ae9b68da0dcdcbd917cb8525d7c591fd66b21 Mon Sep 17 00:00:00 2001 From: Jeff Bolz Date: Sun, 30 Aug 2026 12:47:43 -0500 Subject: [PATCH 2/2] tuning based on rows and num SMs --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 39 ++++++++++++++++++++-------- tests/test-backend-ops.cpp | 2 ++ 2 files changed, 30 insertions(+), 11 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 0791cdbf187..5146519e19e 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -630,7 +630,10 @@ static constexpr uint32_t num_topk_pipelines = 11; static constexpr uint32_t topk_global_radix_bits = 5; static constexpr uint32_t topk_global_num_buckets = 1u << topk_global_radix_bits; static constexpr uint32_t topk_global_num_passes = (32 + topk_global_radix_bits - 1) / topk_global_radix_bits; -static constexpr uint32_t topk_global_chunk_size = 256; +static constexpr uint32_t topk_global_items_per_thread[] = { 2, 4, 8 }; +static constexpr uint32_t num_topk_global_pipelines = sizeof(topk_global_items_per_thread) / sizeof(topk_global_items_per_thread[0]); +static constexpr uint32_t topk_global_default_pipeline = 0; +static constexpr uint32_t topk_global_max_auto_pipeline = 2; static constexpr std::initializer_list topk_moe_early_softmax_norm{ GGML_OP_SOFT_MAX, GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS, GGML_OP_RESHAPE, @@ -1061,8 +1064,8 @@ struct vk_device_struct { vk_pipeline pipeline_argsort_f32[num_argsort_pipelines]; vk_pipeline pipeline_argsort_large_f32[num_argsort_pipelines]; vk_pipeline pipeline_topk_f32[num_topk_pipelines]; - vk_pipeline pipeline_topk_global_f32; - vk_pipeline pipeline_topk_global_subgroup_f32; + vk_pipeline pipeline_topk_global_f32[num_topk_global_pipelines]; + vk_pipeline pipeline_topk_global_subgroup_f32[num_topk_global_pipelines]; vk_pipeline pipeline_sum_rows_f32; vk_pipeline pipeline_cross_entropy_loss_f32, pipeline_cross_entropy_loss_f32_wg512; vk_pipeline pipeline_cross_entropy_loss_back_f32, pipeline_cross_entropy_loss_back_f32_wg512; @@ -5827,12 +5830,16 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { } const uint32_t topk_global_block_size = 1u << std::min(device->max_workgroup_size_log2, 7u); - const uint32_t topk_global_items_per_thread = topk_global_chunk_size / topk_global_block_size; - ggml_vk_create_pipeline2(device, device->pipeline_topk_global_f32, "topk_global_f32", topk_global_f32_len, topk_global_f32_data, "main", 3, sizeof(vk_op_topk_global_push_constants), {topk_global_chunk_size, 1, 1}, - {topk_global_block_size, topk_global_items_per_thread}, 1, true); - if (device->subgroup_basic && device->subgroup_arithmetic && device->subgroup_ballot && device->subgroup_require_full_support && device->subgroup_size >= topk_global_num_buckets) { - ggml_vk_create_pipeline2(device, device->pipeline_topk_global_subgroup_f32, "topk_global_subgroup_f32", topk_global_subgroup_f32_len, topk_global_subgroup_f32_data, "main", 3, - sizeof(vk_op_topk_global_push_constants), {topk_global_chunk_size, 1, 1}, {topk_global_block_size, topk_global_items_per_thread}, 1, true, true, device->subgroup_size); + for (uint32_t i = 0; i < num_topk_global_pipelines; ++i) { + const uint32_t items_per_thread = topk_global_items_per_thread[i]; + const uint32_t chunk_size = topk_global_block_size * items_per_thread; + const std::string suffix = "_i" + std::to_string(items_per_thread); + ggml_vk_create_pipeline2(device, device->pipeline_topk_global_f32[i], "topk_global_f32" + suffix, topk_global_f32_len, topk_global_f32_data, "main", 3, + sizeof(vk_op_topk_global_push_constants), {chunk_size, 1, 1}, {topk_global_block_size, items_per_thread}, 1, true); + if (device->subgroup_basic && device->subgroup_arithmetic && device->subgroup_ballot && device->subgroup_require_full_support && device->subgroup_size >= topk_global_num_buckets) { + ggml_vk_create_pipeline2(device, device->pipeline_topk_global_subgroup_f32[i], "topk_global_subgroup_f32" + suffix, topk_global_subgroup_f32_len, topk_global_subgroup_f32_data, "main", 3, + sizeof(vk_op_topk_global_push_constants), {chunk_size, 1, 1}, {topk_global_block_size, items_per_thread}, 1, true, true, device->subgroup_size); + } } ggml_vk_create_pipeline(device, device->pipeline_argmax_f32, "argmax_f32", argmax_f32_len, argmax_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, { device->subgroup_size }, 1); @@ -13971,7 +13978,17 @@ static void ggml_vk_topk_global(ggml_backend_vk_context * ctx, vk_context& subct vk_subbuffer src_buf = ggml_vk_tensor_subbuffer(ctx, src0); vk_subbuffer scratch_buf = { ctx->prealloc_x, 0, scratch_size }; vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst); - vk_pipeline pipeline = ctx->device->pipeline_topk_global_subgroup_f32 ? ctx->device->pipeline_topk_global_subgroup_f32 : ctx->device->pipeline_topk_global_f32; + uint32_t pipeline_index = topk_global_default_pipeline; + const uint64_t min_workgroups = 4ull * std::max(1u, ctx->device->shader_core_count); + for (uint32_t i = topk_global_max_auto_pipeline; i > topk_global_default_pipeline; --i) { + vk_pipeline candidate = ctx->device->pipeline_topk_global_subgroup_f32[i] ? ctx->device->pipeline_topk_global_subgroup_f32[i] : ctx->device->pipeline_topk_global_f32[i]; + const uint64_t workgroups = uint64_t(nrows) * ((ncols + candidate->wg_denoms[0] - 1) / candidate->wg_denoms[0]); + if (workgroups >= min_workgroups) { + pipeline_index = i; + break; + } + } + vk_pipeline pipeline = ctx->device->pipeline_topk_global_subgroup_f32[pipeline_index] ? ctx->device->pipeline_topk_global_subgroup_f32[pipeline_index] : ctx->device->pipeline_topk_global_f32[pipeline_index]; GGML_ASSERT(pipeline); subctx->s->buffer->buf.fillBuffer(scratch_buf.buffer->buffer, scratch_buf.offset, scratch_buf.size, 0); @@ -18792,7 +18809,7 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm uint32_t min_pipeline = (uint32_t)log2f(float(op->ne[0])) + 1; if (min_pipeline >= num_topk_pipelines || !device->pipeline_topk_f32[min_pipeline]) { - return device->pipeline_topk_global_subgroup_f32 != nullptr || device->pipeline_topk_global_f32 != nullptr; + return device->pipeline_topk_global_subgroup_f32[topk_global_default_pipeline] != nullptr || device->pipeline_topk_global_f32[topk_global_default_pipeline] != nullptr; } } return true; diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index bc7c6310cc0..887668ca5ef 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10476,6 +10476,8 @@ static std::vector> make_test_cases_perf() { test_cases.emplace_back(new test_argsort(GGML_TYPE_F32, {200000, 16, 1, 1})); test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {2, 1, 1, 1}, 1)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {16384, 1, 1, 1}, 2051)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {16384, 512, 1, 1}, 2051)); // widths around the tiling threshold for (auto cols : {4096, 8192, 12288, 16384, 24576, 32768, 65536, 131072}) { for (auto nrows : {1, 16}) {