Daniel HiltgenandGitHub 4d22e65d49 mlx: sort long top-k rows in chunks (#19392) ...

MLX's argpartition sorts the whole row, and a row longer than one block
of the GPU sort takes a multi-pass merge sort. TopKIndices sorts each
2048-value chunk in a single block, then only the chunks' top k, and
returns the same indices in the same order.

The sampler's top_k and top_logprobs and the DFlash2 candidate selector
use it in place of a full-vocabulary argpartition.

2026-10-29 18:53:49 -07:00

..

logprob_test.go

…

sample_test.go

…

sample.go

mlx: sort long top-k rows in chunks (#19392)

2026-10-29 18:53:49 -07:00