diff --git a/.gitignore b/.gitignore index e706a1484..0d9e3f4b6 100644 --- a/.gitignore +++ b/.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 diff --git a/Makefile b/Makefile index c58acb939..d4c1f0390 100644 --- a/Makefile +++ b/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 diff --git a/ds4_metal.m b/ds4_metal.m index 47bea0821..54000b90a 100644 --- a/ds4_metal.m +++ b/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; diff --git a/metal/argsort.metal b/metal/argsort.metal index 77c473f96..ae977375d 100644 --- a/metal/argsort.metal +++ b/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; } diff --git a/tests/test_argsort_metal.c b/tests/test_argsort_metal.c new file mode 100644 index 000000000..9b40c2bd3 --- /dev/null +++ b/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 +#include +#include +#include + +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; +}