metal: prune top-k argsort merge rounds to top_k

.gitignore

@@ -28,6 +28,7 @@
 /tests/test_mxfp4_cuda
 /tests/test_mxfp4_dot
 /tests/test_mxfp4_metal
+/tests/test_argsort_metal
 /tests/test_q4k_dot
 /tests/test_sampling
 /tests/test_glm53_kda

Makefile

@@ -83,6 +83,7 @@ help:
 	@echo "  make metal-prefill-variant-bench  Build the balanced Metal prefill variant benchmark"
 	@echo "  make check-mxfp4-half-lut  Verify the checked-in MXFP4 half LUT matches the generator"
 	@echo "  make test-mxfp4-metal  Check the MXFP4 half LUT, then run Metal MXFP4 exactness tests"
+	@echo "  make test-argsort-metal  Run the Metal top-k argsort correctness test"
 	@echo "  make dspark-verify-depth  Run DSpark speculative verification smoke if support GGUF is present"
 	@echo "  make mtp-verify-depth  Run legacy MTP speculative verification smoke if MTP GGUF is present"
 	@echo "  make clean        Remove build outputs"
@@ -142,6 +143,15 @@ check-mxfp4-half-lut:
 test-mxfp4-metal: check-mxfp4-half-lut tests/test_mxfp4_metal
 	./tests/test_mxfp4_metal
 
+tests/test_argsort_metal.o: tests/test_argsort_metal.c ds4_gpu.h
+	$(CC) $(CFLAGS) -I. -c -o $@ $<
+
+tests/test_argsort_metal: tests/test_argsort_metal.o ds4_metal.o ds4_image.o
+	$(CC) $(CFLAGS) -o $@ $^ $(METAL_LDLIBS)
+
+test-argsort-metal: tests/test_argsort_metal
+	./tests/test_argsort_metal
+
 cpu: ds4_cli_cpu.o ds4_server_cpu.o ds4_bench_cpu.o ds4_eval_cpu.o ds4_agent_cpu.o ds4_help.o ds4_prompt_prefix.o ds4_web.o ds4_kvstore.o linenoise.o rax.o ds4_gpu_args_cpu.o $(CPU_CORE_OBJS)
 	$(CC) $(CFLAGS) -o ds4 ds4_cli_cpu.o ds4_help.o ds4_prompt_prefix.o linenoise.o ds4_gpu_args_cpu.o $(CPU_CORE_OBJS) $(LDLIBS)
 	$(CC) $(CFLAGS) -o ds4-server ds4_server_cpu.o ds4_help.o ds4_kvstore.o rax.o ds4_gpu_args_cpu.o $(CPU_CORE_OBJS) $(LDLIBS)
@@ -611,4 +621,4 @@ mxfp4-dot-test: tests/test_mxfp4_dot.c
 	./tests/test_mxfp4_dot
 
 clean:
-	rm -f ds4 ds4-server ds4-bench ds4-eval ds4-agent ds4_cpu ds4_native ds4_server_test ds4_test ds4_agent_test gguf-tools/quality-testing/score_official gguf-tools/quality-testing/score_official.o speed-bench/metal_decode_schedule_bench speed-bench/metal_prefill_variant_bench speed-bench/*.o tests/test_q4k_dot tests/test_mxfp4_dot tests/test_mxfp4_metal tests/test_mxfp4_rocm tests/test_mxfp4_cuda tests/test_metal_session_batch tests/test_glm53_kda tests/test_glm53_kda_rocm tests/test_glm53_vision_engine tests/test_glm53_vision_prompt tests/test_deepseek4_vision_image tests/test_prompt_prefix tests/test_gpu_xdev tests/test_gpu_model_cache tests/test_gpu_lookup_cache_strict tests/test_engine_mgpu_refusal tests/test_engine_mgpu_runtime tests/test_engine_correctness tests/test_sampling tests/test_cuda_session_batch tests/test_cuda_mixed_batch tests/*.o *.o tests/cuda_long_context_smoke tests/cuda_long_context_smoke.o
+	rm -f ds4 ds4-server ds4-bench ds4-eval ds4-agent ds4_cpu ds4_native ds4_server_test ds4_test ds4_agent_test gguf-tools/quality-testing/score_official gguf-tools/quality-testing/score_official.o speed-bench/metal_decode_schedule_bench speed-bench/metal_prefill_variant_bench speed-bench/*.o tests/test_q4k_dot tests/test_mxfp4_dot tests/test_mxfp4_metal tests/test_argsort_metal tests/test_mxfp4_rocm tests/test_mxfp4_cuda tests/test_metal_session_batch tests/test_glm53_kda tests/test_glm53_kda_rocm tests/test_glm53_vision_engine tests/test_glm53_vision_prompt tests/test_deepseek4_vision_image tests/test_prompt_prefix tests/test_gpu_xdev tests/test_gpu_model_cache tests/test_gpu_lookup_cache_strict tests/test_engine_mgpu_refusal tests/test_engine_mgpu_runtime tests/test_engine_correctness tests/test_sampling tests/test_cuda_session_batch tests/test_cuda_mixed_batch tests/*.o *.o tests/cuda_long_context_smoke tests/cuda_long_context_smoke.o

ds4_metal.m

@@ -6333,6 +6333,8 @@ static int ds4_gpu_encode_rope_tail_inplace(
     int32_t  ne3;
     int32_t  top_k;
     int32_t  len;
+    int32_t  total;
+    int32_t  keep_k;
 } ds4_gpu_kargs_argsort_merge;
 
 typedef struct {
@@ -19087,15 +19089,29 @@ int ds4_gpu_indexer_topk_tensor(
              threadsPerThreadgroup:MTLSizeMake((NSUInteger)nth, 1, 1)];
         ds4_gpu_end_compute_encoder(cb, enc);
 
+        /* Runs are packed at stride len, all full but the last; total tracks the
+         * valid elements per row.  A round merges pairs into runs of new_len,
+         * so total shrinks toward top_k instead of staying at work_width. */
         int32_t len = block_top_k;
-        while (len < work_width) {
-            const int32_t nm = (work_width + 2 * len - 1) / (2 * len);
+        int32_t total = work_width;
+        int32_t nruns = npr;
+        while (nruns > 1) {
+            const int32_t nm = (nruns + 1) / 2;   /* odd leftover passes through */
             const bool final_merge = nm == 1;
+            /* Final round: nruns == 2 so total <= 2*len, and total >= top_k, hence new_len == top_k. */
+            const int32_t new_len = (2 * len < (int32_t)top_k) ? 2 * len : (int32_t)top_k;
             NSUInteger merge_threads = g_argsort_merge_f32_i32_desc_pipeline.maxTotalThreadsPerThreadgroup;
             if (merge_threads == 0 || merge_threads > 512u) merge_threads = 512u;
-            if (merge_threads > (NSUInteger)len) merge_threads = (NSUInteger)len;
+            if (merge_threads > (NSUInteger)new_len) merge_threads = (NSUInteger)new_len;
             if (merge_threads == 0) merge_threads = 1;
 
+            /* Full pairs write new_len each, the last partial pair writes
+             * min(r, new_len); q is clamped to the pairs actually dispatched. */
+            const int32_t q = total / (2 * len);
+            const int32_t r = total % (2 * len);
+            const int32_t total_next =
+                MIN(q, nm) * new_len + (q < nm ? MIN(r, new_len) : 0);
+
             ds4_gpu_kargs_argsort_merge merge_args = {
                 .ne00 = (int64_t)n_comp,
                 .ne01 = (int64_t)n_tokens,
@@ -19109,8 +19125,10 @@ int ds4_gpu_indexer_topk_tensor(
                 .ne1 = (int32_t)n_tokens,
                 .ne2 = 1,
                 .ne3 = 1,
-                .top_k = nm == 1 ? (int32_t)top_k : work_width,
+                .top_k = final_merge ? (int32_t)top_k : work_width,
                 .len = len,
+                .total = total,
+                .keep_k = (int32_t)top_k,
             };
 
             enc = ds4_gpu_compute_encoder(cb);
@@ -19128,7 +19146,9 @@ int ds4_gpu_indexer_topk_tensor(
             const NSUInteger tmp = cur_off;
             cur_off = next_off;
             next_off = tmp;
-            len <<= 1;
+            len = new_len;
+            total = total_next;
+            nruns = nm;
         }
 
         if (!ds4_gpu_finish_command_buffer(cb, owned, "indexer top-k")) return 0;

metal/argsort.metal

@@ -29,6 +29,8 @@ struct ds4_metal_args_argsort_merge {
     int32_t  ne3;
     int32_t  top_k;
     int32_t  len;
+    int32_t  total;
+    int32_t  keep_k;
 };
 
 typedef void (argsort_t)(
@@ -142,21 +144,25 @@ kernel void kernel_argsort_merge_f32_i32(
     const int i02 = tgpig[1];
     const int i03 = tgpig[2];
 
-    const int start = im * (2 * args.len);
+    // Runs are read at stride 2*len within the args.total valid elements of the
+    // row; output is packed at stride new_len so a round keeps at most keep_k.
+    const int start_read  = im * (2 * args.len);
+    const int new_len     = MIN(2 * args.len, args.keep_k);
+    const int start_write = im * new_len;
 
-    const int len0 = MIN(args.len, MAX(0, args.ne0 - (int)(start)));
-    const int len1 = MIN(args.len, MAX(0, args.ne0 - (int)(start + args.len)));
+    const int len0 = MIN(args.len, MAX(0, args.total - (int)(start_read)));
+    const int len1 = MIN(args.len, MAX(0, args.total - (int)(start_read + args.len)));
 
     const int total = len0 + len1;
 
-    device const int32_t * tmp0 = tmp + start
+    device const int32_t * tmp0 = tmp + start_read
         + i01*args.ne0
         + i02*args.ne0*args.ne01
         + i03*args.ne0*args.ne01*args.ne02;
 
     device const int32_t * tmp1 = tmp0 + args.len;
 
-    dst += start
+    dst += start_write
         + i01*args.top_k
         + i02*args.top_k*args.ne01
         + i03*args.top_k*args.ne01*args.ne02;
@@ -170,16 +176,13 @@ kernel void kernel_argsort_merge_f32_i32(
         return;
     }
 
-    const int chunk = (total + ntg.x - 1) / ntg.x;
+    const int out   = MIN(total, new_len);
+    const int chunk = (out + ntg.x - 1) / ntg.x;
 
     const int k0 = tpitg.x * chunk;
-    const int k1 = MIN(MIN(k0 + chunk, total), args.top_k);
+    const int k1 = MIN(k0 + chunk, out);
 
-    if (k0 >= args.top_k) {
-        return;
-    }
-
-    if (k0 >= total) {
+    if (k0 >= out) {
         return;
     }
 

tests/test_argsort_metal.c

@@ -0,0 +1,186 @@
+/*
+ * Metal-only top-k argsort test (ds4_gpu_indexer_topk_tensor).
+ *
+ * Exercises the Metal merge kernel's pruning path: intermediate rounds keep
+ * top_k instead of work_width, including odd run counts (leftover pass-through)
+ * and cases where top_k is small relative to n_comp (prune from the first
+ * merge). The existing run_topk2048 in test_gpu_xdev.c is CUDA-only and cannot
+ * cover this. Each case's comment names the path it actually exercises; the
+ * per-case print (see print_path) mirrors the host shape math so the path is
+ * confirmed at run time rather than inferred.
+ */
+
+#include "ds4_gpu.h"
+
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+bool ds4_log_is_tty(FILE *fp) {
+    (void)fp;
+    return false;
+}
+
+typedef struct {
+    float score;
+    uint32_t index;
+} topk_ref_entry;
+
+static int topk_ref_cmp(const void *ap, const void *bp) {
+    const topk_ref_entry *a = (const topk_ref_entry *)ap;
+    const topk_ref_entry *b = (const topk_ref_entry *)bp;
+    if (a->score > b->score) return -1;
+    if (a->score < b->score) return 1;
+    return a->index < b->index ? -1 : (a->index > b->index ? 1 : 0);
+}
+
+/* Mirror ds4_gpu_indexer_topk_tensor's shape math and merge-round loop
+ * (ds4_metal.m) to print which path each case exercises. The argsort kernel is
+ * a plain Metal kernel, so maxTotalThreadsPerThreadgroup is the 1024 default. */
+static void print_path(uint32_t n_comp, uint32_t top_k) {
+    const uint32_t max_threads = 1024u;
+    uint32_t nth = 1u;
+    while (nth < n_comp && 2u * nth <= max_threads) nth *= 2u;
+    const uint32_t npr = (n_comp + nth - 1u) / nth;
+    const uint32_t block_top_k = top_k < nth ? top_k : nth;
+    uint32_t work_width = top_k;
+    if (npr > 1) {
+        const uint32_t last_block = n_comp - (npr - 1u) * nth;
+        work_width = (npr - 1u) * block_top_k +
+                     (last_block < block_top_k ? last_block : block_top_k);
+    }
+    uint32_t rounds = 0, len = block_top_k, total = work_width, nruns = npr;
+    while (nruns > 1) {
+        const uint32_t nm = (nruns + 1u) / 2u;
+        const uint32_t new_len = (2u * len < top_k) ? 2u * len : top_k;
+        /* full pairs write new_len; last partial pair min(remainder, new_len);
+         * q clamped to nm keeps the accounting self-limiting like the host. */
+        const uint32_t q = total / (2u * len);
+        const uint32_t r = total % (2u * len);
+        total = (q < nm ? q : nm) * new_len + (q < nm ? (r < new_len ? r : new_len) : 0u);
+        rounds++;
+        len = new_len;
+        nruns = nm;
+    }
+    if (npr <= 1) {
+        fprintf(stderr, "  -> nth=%u npr=%u: one-pass (merge kernel never runs)\n",
+                nth, npr);
+    } else {
+        fprintf(stderr, "  -> nth=%u npr=%u work_width=%u: %u merge round(s)\n",
+                nth, npr, work_width, rounds);
+    }
+}
+
+int main(void) {
+    /* {n_comp, top_k, n_tokens} */
+    const uint32_t cases[][3] = {
+        { 3355u, 2048u, 2u },   /* odd npr=4, pruning from an intermediate merge */
+        { 5003u, 2048u, 2u },   /* odd npr=5, larger odd count */
+        { 300u,  8u,    2u },   /* one-pass (n_comp < nth => npr=1), multi-token */
+        { 4096u, 2048u, 2u },   /* power-of-two n_comp, even rounds */
+        { 8192u, 8u,    1u },   /* npr=8 maximal-pruning tree (total 64->8), 3 rounds */
+        { 8192u, 300u,  1u },   /* non-power-of-two top_k: len stays 300, no transition */
+        { 8192u, 1u,    1u },   /* degenerate top_k=1 max-reduction (new_len==1) */
+    };
+    const uint32_t n_cases = (uint32_t)(sizeof(cases) / sizeof(cases[0]));
+    uint32_t ci;
+
+    if (!ds4_gpu_init()) {
+        fprintf(stderr, "argsort Metal: ds4_gpu_init failed\n");
+        return 1;
+    }
+
+    for (ci = 0; ci < n_cases; ci++) {
+        const uint32_t n_comp = cases[ci][0];
+        const uint32_t top_k = cases[ci][1];
+        const uint32_t n_tokens = cases[ci][2];
+        print_path(n_comp, top_k);
+        const uint64_t n_scores = (uint64_t)n_tokens * n_comp;
+        float *host_scores = (float *)malloc((size_t)n_scores * sizeof(float));
+        uint32_t *host_selected =
+            (uint32_t *)malloc((size_t)n_tokens * top_k * sizeof(uint32_t));
+        topk_ref_entry *ref =
+            (topk_ref_entry *)malloc((size_t)n_comp * sizeof(topk_ref_entry));
+        if (!host_scores || !host_selected || !ref) {
+            fprintf(stderr, "argsort Metal: host alloc failed\n");
+            return 1;
+        }
+        /* Deterministic permutation of 0..n_comp-1 per token, so scores are
+         * unique (no ties) and float-exact. This validates the pruning/merge
+         * logic without depending on tie-break order, which the Metal merge
+         * resolves differently from the reference (a pre-existing discrepancy
+         * unrelated to pruning). */
+        uint32_t *perm = (uint32_t *)malloc((size_t)n_comp * sizeof(uint32_t));
+        if (!perm) {
+            fprintf(stderr, "argsort Metal: perm alloc failed\n");
+            return 1;
+        }
+        for (uint32_t t = 0; t < n_tokens; t++) {
+            for (uint32_t i = 0; i < n_comp; i++) perm[i] = i;
+            uint32_t seed = 0x1234567u + 0x9e3779b9u * t;
+            for (uint32_t i = n_comp; i > 1; i--) {
+                seed = seed * 1664525u + 1013904223u;
+                uint32_t j = seed % i;
+                uint32_t tmp = perm[i - 1];
+                perm[i - 1] = perm[j];
+                perm[j] = tmp;
+            }
+            for (uint32_t i = 0; i < n_comp; i++) {
+                host_scores[(uint64_t)t * n_comp + i] = (float)perm[i];
+            }
+        }
+        free(perm);
+
+        ds4_gpu_tensor *scores = ds4_gpu_tensor_alloc(n_scores * sizeof(float));
+        ds4_gpu_tensor *selected =
+            ds4_gpu_tensor_alloc((uint64_t)n_tokens * top_k * sizeof(uint32_t));
+        if (scores == NULL || selected == NULL) {
+            fprintf(stderr, "argsort Metal: tensor alloc failed\n");
+            return 1;
+        }
+
+        int ok = ds4_gpu_tensor_write(
+                     scores, 0, host_scores, n_scores * sizeof(float)) &&
+                 ds4_gpu_indexer_topk_tensor(
+                     selected, scores, n_comp, n_tokens, top_k) &&
+                 ds4_gpu_tensor_read(
+                     selected, 0, host_selected,
+                     (uint64_t)n_tokens * top_k * sizeof(uint32_t));
+        if (!ok) {
+            fprintf(stderr,
+                    "argsort Metal: FAIL compute n=%u top_k=%u\n", n_comp, top_k);
+            return 1;
+        }
+
+        for (uint32_t t = 0; t < n_tokens; t++) {
+            for (uint32_t i = 0; i < n_comp; i++) {
+                ref[i].score = host_scores[(uint64_t)t * n_comp + i];
+                ref[i].index = i;
+            }
+            qsort(ref, n_comp, sizeof(ref[0]), topk_ref_cmp);
+            for (uint32_t i = 0; i < top_k; i++) {
+                if (host_selected[(uint64_t)t * top_k + i] != ref[i].index) {
+                    fprintf(stderr,
+                            "argsort Metal: FAIL n=%u token=%u rank=%u "
+                            "got=%u want=%u\n",
+                            n_comp, t, i,
+                            host_selected[(uint64_t)t * top_k + i],
+                            ref[i].index);
+                    return 1;
+                }
+            }
+        }
+
+        ds4_gpu_tensor_free(scores);
+        ds4_gpu_tensor_free(selected);
+        free(host_scores);
+        free(host_selected);
+        free(ref);
+        fprintf(stderr, "argsort Metal: n=%u top_k=%u OK\n", n_comp, top_k);
+    }
+
+    ds4_gpu_cleanup();
+    fprintf(stderr, "argsort Metal: all cases PASS\n");
+    return 0;
+}