From 58df04c454bc91327af62195ac1bc6474e40883b Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Mon, 31 Aug 2026 09:53:28 +0300 Subject: [PATCH] metal : add top-k radix implementation Assisted-by: DeepSeek-v4-Flash-0731 --- ggml/src/ggml-metal/ggml-metal-device.cpp | 19 +++- ggml/src/ggml-metal/ggml-metal-device.h | 1 + ggml/src/ggml-metal/ggml-metal-impl.h | 11 +++ ggml/src/ggml-metal/ggml-metal-ops.cpp | 72 ++++++++++++++- ggml/src/ggml-metal/kernels/argsort.metal | 105 ++++++++++++++++++++++ 5 files changed, 206 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index 5a2f01f1dba..b8d2ef9ce27 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -1336,7 +1336,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_fwht(ggml_metal_ return res; } -// note: reuse the argsort kernel for top_k +// note: reuse the argsort kernel for the bitonic top_k fallback ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k(ggml_metal_library_t lib, const ggml_tensor * op) { assert(op->op == GGML_OP_TOP_K); @@ -1364,6 +1364,23 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k(ggml_metal return res; } +ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_radix(ggml_metal_library_t lib, const ggml_tensor * op) { + assert(op->op == GGML_OP_TOP_K); + + char base[256]; + char name[256]; + + snprintf(base, 256, "kernel_top_k_%s_%s", ggml_type_name(op->src[0]->type), ggml_type_name(op->type)); + snprintf(name, 256, "%s", base); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { + res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); + } + + return res; +} + ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_merge(ggml_metal_library_t lib, const ggml_tensor * op) { assert(op->op == GGML_OP_TOP_K); diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index 003b688dbac..7f6520103d1 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h @@ -145,6 +145,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort_merge (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_fwht (ggml_metal_library_t lib, int n); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k (ggml_metal_library_t lib, const struct ggml_tensor * op); +struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_radix (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_merge (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_bin (ggml_metal_library_t lib, const struct ggml_tensor * op, int32_t n_fuse ); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_bin_one (ggml_metal_library_t lib, enum ggml_op op); diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index 49102afe9c0..bdcd9c9e3d8 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -1189,6 +1189,17 @@ typedef struct { int32_t len; } ggml_metal_kargs_argsort_merge; +typedef struct { + int32_t ne00; // number of columns (elements per row) + int32_t ne01; // rows + int32_t ne02; + int32_t ne03; + uint64_t nb01; // row stride in src0 + uint64_t nb02; + uint64_t nb03; + int32_t top_k; // k +} ggml_metal_kargs_top_k; + typedef struct { int32_t nrows; } ggml_metal_kargs_fwht; diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 7671d1d0156..30ea4ec2745 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -5091,7 +5091,9 @@ int ggml_metal_op_argsort(ggml_metal_op_t ctx, int idx) { return 1; } -int ggml_metal_op_top_k(ggml_metal_op_t ctx, int idx) { +// bitonic-sort + merge fallback: efficient when k is small and there are few rows, +// where the single-workgroup-per-row radix-select cannot reach enough parallelism +static void ggml_metal_op_top_k_bitonic(ggml_metal_op_t ctx, int idx) { ggml_tensor * op = ctx->node(idx); ggml_metal_library_t lib = ctx->lib; @@ -5199,6 +5201,74 @@ int ggml_metal_op_top_k(ggml_metal_op_t ctx, int idx) { len <<= 1; } +} + +// radix-select: one workgroup per row. Maps each float to an order-preserving unsigned +// key, finds the k-th largest via 4 radix-8 histogram passes, then compacts the top-k +// indices. Fast for large k and/or many rows. +static void ggml_metal_op_top_k_radix(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; + ggml_metal_encoder_t enc = ctx->enc; + + GGML_ASSERT(ggml_is_contiguous_rows(op->src[0])); + + GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); + + auto pipeline = ggml_metal_library_get_pipeline_top_k_radix(lib, op); + + // one workgroup per row; radix-select the k-th largest value + const int nth = std::min(1024, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); + + ggml_metal_kargs_top_k args = { + /*.ne00 =*/ ne00, + /*.ne01 =*/ ne01, + /*.ne02 =*/ ne02, + /*.ne03 =*/ ne03, + /*.nb01 =*/ nb01, + /*.nb02 =*/ nb02, + /*.nb03 =*/ nb03, + /*.top_k =*/ (int32_t) op->ne[0], + }; + + // shared memory: 256-entry histogram + bucket/above scalars + output counter + const size_t smem_histo = GGML_PAD(256*sizeof(uint32_t), 16); + const size_t smem_bucket = GGML_PAD( sizeof(uint32_t), 16); + const size_t smem_above = GGML_PAD( sizeof(uint32_t), 16); + const size_t smem_out = GGML_PAD( sizeof(uint32_t), 16); + + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); + + ggml_metal_encoder_set_threadgroup_memory_size(enc, smem_histo, 0); + ggml_metal_encoder_set_threadgroup_memory_size(enc, smem_bucket, 1); + ggml_metal_encoder_set_threadgroup_memory_size(enc, smem_above, 2); + ggml_metal_encoder_set_threadgroup_memory_size(enc, smem_out, 3); + + ggml_metal_encoder_dispatch_threadgroups(enc, ne01, ne02, ne03, nth, 1, 1); +} + +int ggml_metal_op_top_k(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + // radix-select has a fixed single-workgroup-per-row cost (~50-60us) that is only + // amortized for long rows, many rows, or a large k; otherwise the bitonic path wins + const int ncols = op->src[0]->ne[0]; + const int k = op->ne[0]; + const int nrows = ggml_nrows(op->src[0]); + + const bool use_radix = + ncols > 2048 && (k > 64 || (nrows > 4 && ncols >= 8192)); + + if (use_radix) { + ggml_metal_op_top_k_radix(ctx, idx); + } else { + ggml_metal_op_top_k_bitonic(ctx, idx); + } return 1; } diff --git a/ggml/src/ggml-metal/kernels/argsort.metal b/ggml/src/ggml-metal/kernels/argsort.metal index 7d144fbd755..e81d194c339 100644 --- a/ggml/src/ggml-metal/kernels/argsort.metal +++ b/ggml/src/ggml-metal/kernels/argsort.metal @@ -230,3 +230,108 @@ kernel void kernel_argsort_merge_f32_i32( template [[host_name("kernel_argsort_merge_f32_i32_asc")]] kernel argsort_merge_t kernel_argsort_merge_f32_i32; template [[host_name("kernel_argsort_merge_f32_i32_desc")]] kernel argsort_merge_t kernel_argsort_merge_f32_i32; + +static inline uint ggml_top_k_f2ui(float x) { + uint y = as_type(x); + if ((y & 0x80000000u) != 0u) { + y ^= 0xFFFFFFFFu; // negative floats: flip all bits + } else { + y |= 0x80000000u; // positive floats: set the sign bit + } + return y; +} + +kernel void kernel_top_k_f32_i32( + constant ggml_metal_kargs_top_k & args, + device const char * src0, + device int32_t * dst, + threadgroup atomic_uint * histo [[threadgroup(0)]], + threadgroup uint * sh_bucket [[threadgroup(1)]], + threadgroup uint * sh_above [[threadgroup(2)]], + threadgroup atomic_uint * out_count [[threadgroup(3)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort3 tpitg[[thread_position_in_threadgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { + + const uint ncols = args.ne00; + const uint top_k = args.top_k; + const uint i01 = tgpig[0]; + const uint i02 = tgpig[1]; + const uint i03 = tgpig[2]; + + device const float * src0_row = (device const float *) (src0 + args.nb01*i01 + args.nb02*i02 + args.nb03*i03); + + device int32_t * dst_row = dst + top_k*(i01 + args.ne01*i02 + args.ne01*args.ne02*i03); + + const uint tid = tpitg.x; + const uint ntg_x = ntg.x; + + uint prefix = 0; // fixed high bits of the threshold key + uint desired = top_k; // count still needed from the candidate range + + for (int shift = 24; shift >= 0; shift -= 8) { + for (uint i = tid; i < 256; i += ntg_x) { + atomic_store_explicit(&histo[i], 0u, memory_order_relaxed); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + const uint hi_mask = (shift + 8 >= 32) ? 0u : (0xFFFFFFFFu << uint(shift + 8)); + const uint prefix_hi = prefix & hi_mask; + + for (uint i = tid; i < ncols; i += ntg_x) { + const uint key = ggml_top_k_f2ui(src0_row[i]); + if ((key & hi_mask) == prefix_hi) { + atomic_fetch_add_explicit(&histo[(key >> uint(shift)) & 0xFFu], 1u, memory_order_relaxed); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // top-down scan for the bucket holding the k-th value + if (tid == 0) { + uint acc = 0; + uint b = 0; + for (int bb = 255; bb >= 0; --bb) { + const uint c = atomic_load_explicit(&histo[bb], memory_order_relaxed); + if (acc + c >= desired) { + b = uint(bb); + break; + } + acc += c; + } + *sh_bucket = b; + *sh_above = acc; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + prefix |= *sh_bucket << uint(shift); + desired -= *sh_above; + + // ensure every thread has consumed sh_bucket/sh_above before the next pass + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + if (tid == 0) { + atomic_store_explicit(out_count, 0u, memory_order_relaxed); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // emit everything above the threshold, then fill the rest from ties + const uint threshold = prefix; + + for (uint i = tid; i < ncols; i += ntg_x) { + if (ggml_top_k_f2ui(src0_row[i]) > threshold) { + const uint pos = atomic_fetch_add_explicit(out_count, 1u, memory_order_relaxed); + dst_row[pos] = (int32_t) i; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint i = tid; i < ncols; i += ntg_x) { + if (ggml_top_k_f2ui(src0_row[i]) == threshold) { + const uint pos = atomic_fetch_add_explicit(out_count, 1u, memory_order_relaxed); + if (pos < top_k) { + dst_row[pos] = (int32_t) i; + } + } + } +}