Repository navigation
Conversation
… kernel The server defaults (temperature 1, top_p 0.95, top_k 20) ran mlx_lm's argsort for top_p and argpartition (a full stable sort in MLX) for top_k after every resident step: about 280 us of a 4.9 ms token. A chunked radix-select kernel now finds the 48 largest values with MLX's tie order, walks the nucleus mass over them and draws with the same gumbel noise. A step whose top_p threshold lies within MARGIN (1e-5) of the kernel's mass estimate, or whose tie group is too large, is flagged; the packet stops there like an EOS and the caller draws that step with the sort-based sampler (one shared sort for both masks). The grammar never sees a flagged token. A rerun of the whole packet is not possible: its writes evict sliding-window entries the rerun needs. Exactness is empirical, not proven. mlx_lm compares a float32 cumsum of 150K ascending probabilities with 1 - top_p; the kernel sums the same terms in another order, so the worst-case bound is about n*eps = 1e-2. Measured over 31.6K decode steps (code, prose, Chinese, tool JSON, a 16K-token context; three runs), the largest difference between the kernel's value and MLX's cumsum at the candidate positions was 8.3e-7 (p99.99 7.2e-7), 8% of MARGIN. Tokens drawn with the kernel and with mlx_lm's sampler under the same key differed in 0 of those steps; 1 in ~10000 steps is flagged and redone. A token could still differ if the two sums disagreed by more than MARGIN at a step where the threshold falls that close to a token's mass. Server decode at 1K, T=1/top_p 0.95/top_k 20 (bench-server, 3 runs, alternating with the unchanged tree): 188.2, 187.7 -> 196.9, 197.3 tok/s. Greedy and unfiltered T=1 are unchanged.
…ts own change top_p_top_k reproduces apply_top_k(apply_top_p(x)) exactly: the k-th value comes from the top_p sort, and ties keep the lowest indices as MLX's stable argpartition does. Tokens equal make_sampler's for the same key by construction. The radix-select kernel and the flagged-step redo agree with make_sampler only empirically, so they move to a separate change that builds on this one.
tc3oliver
marked this pull request as ready for review
October 24, 2026 09:11
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
With top_p and top_k both set (the server default is 0.95 / 20),
make_samplersorts the 150K vocabulary twice per step:argsortinapply_top_pandargpartitioninapply_top_k. MLX implements the latter as a full stable sort.top_p_top_kreuses the top_p sort. The k-th largest value islogprobs[order[-top_k]], and ties at that value keep the lowest indices, which is what the stableargpartitionkeeps. The mask equalsapply_top_k(apply_top_p(x))exactly, so tokens equalmake_sampler's for the same key. Other settings keepmake_sampler.The sampler goes from 324 µs to 277 µs per step.
Results
M4 Max (40-core GPU), mlx 0.32.0,
RecursiveSI/flint-int4, on top of #17 (mainrejects the M4 Max).rec-mlx serve, 1K prompt, T=1, top_p 0.95, top_k 20, 3 runs alternated with the unchanged tree: decode goes from 189.7 to 192.0 tok/s (+1.2%).Tests
test_resident_sampler.pycompares the mask withapply_top_k(apply_top_p(x))and the token withmake_samplerunder the same seed:It fails if ties keep the higher index or if k is off by one.
validate_resident.py: 61/61.make testpasses.The earlier version of this PR also had a candidate-selection kernel. Its agreement with
make_sampleris only measured, not proven, so it will follow as a separate draft PR on top of this one.