metal: prune top-k argsort merge rounds to top_k

[emilianbold]

Summary

kernel_argsort_merge_f32_i32 carried work_width elements per row through every intermediate merge round, so the full row was re-merged until the final merge clamped to top_k. It also divided thread work by the number of elements read rather than written, so once a round truncates, the high threads exit immediately and whole SIMD groups idle.

This PR prunes each round to top_k and splits the work over the outputs that actually exist. Same result, less work, better occupancy.

The work was started with DeepSeek Flash on ds4 then refined with Claude.

The change

metal/argsort.metal:

• Reads walk the input runs at stride 2*len, bounded by a new args.total (the valid elements per row) instead of args.ne0 (the buffer row stride).
• Writes pack at start_write = im * new_len where new_len = MIN(2*len, keep_k), so a round keeps at most top_k per pair.
• The work split keys off out = MIN(total, new_len), and one guard (k0 >= out) replaces the two old early returns.
• Two new kernel args: total, keep_k.

ds4_metal.m (ds4_gpu_indexer_topk_tensor):

• Tracks len / total / nruns per round and drives the loop off the run count rather than len < work_width.
• total_next is O(1): full pairs write new_len, the last partial pair writes MIN(r, new_len), with q clamped to the dispatched pair count.

Odd run counts need no special case: the last run has no partner, which falls out of len1 == 0, and the merge copies it through.

Why it is correct

Runs stay packed at stride len with every run full but the last, so:

• total >= top_k is preserved every round;
• the final round has nruns == 2, hence total <= 2*len, hence 2*len >= top_k, hence new_len == top_k and exactly top_k elements are written to the output tensor;
• writes are bounded by total_next <= total <= work_width, so they stay in the row.

Truncating a pair to new_len is exact rather than approximate: a dropped element had >= top_k elements above it inside the pair, hence >= top_k in the whole row, so it was not in the top-k.

Performance

Merge-element writes per row, across all rounds, at nth = 1024:

shape (n_comp / top_k)	npr	before	after	ratio
8192 / 512	8	8704	3584	2.43x
16384 / 1024	16	50176	15360	3.27x
32768 / 2048	32	133120	63488	2.10x
8192 / 2048	8	18432	14336	1.29x

Measured merge-GPU time per call, median over 7 runs on an idle GPU:

shape	n_tokens	before	after	speedup
8192 / 512	1	13.5 us	11.0 us	1.22x
8192 / 512	256	76.1 us	59.4 us	1.28x
16384 / 1024	1	28.3 us	16.3 us	1.74x
16384 / 1024	256	402.2 us	175.3 us	2.29x
32768 / 2048	1	49.1 us	23.5 us	2.09x
32768 / 2048	256	1265.4 us	560.9 us	2.26x

Two caveats on the measurement. It was taken with the merge rounds isolated in their own command buffer, which is not how they are encoded in production, and the harness that does that is not part of this PR.

End to end this is modest in an ordinary context: the saving is about 0.10 ms/token at Flash 32k summed over 43 layers, roughly 0.3% of a ~31 ms token. It grows with context though.

Testing

make test-argsort-metal runs tests/test_argsort_metal.c