Add top-k support to MLX sample

[goutamadwant]

Fixes#20548

Summary

• Thread optional top_k through SamplingHead and mlx::sample.
• Apply top-k filtering as an additional threshold mask that composes with the existing top-p nucleus mask.
• Add eager, export, end-to-end, and MLX lowering coverage for top-k sampling.

Test plan

• PYTHONPATH=src:. python3 -m unittest executorch.backends.mlx.test.test_sample.TestSampleOp executorch.backends.mlx.test.test_sample.TestSampleExport
• python3 -m compileall -q backends/mlx/llm/sampling.py backends/mlx/custom_ops.py backends/mlx/ops.py backends/mlx/test/test_sample.py backends/mlx/test/test_ops.py

cc @metascroy

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20564

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 12 Pending

As of commit 7e9ae35 with merge base 55a71e6 ():

NEW FAILURE - The following job has failed:

• pull / test-eval_llama-wikitext-linux / linux-job (gh)
  RuntimeError: Command docker exec -t f07b8ecee9be9cc829d92ad63b572fb3fb94908a1e2f16b39fadf2553b965213 /exec failed with exit code 1

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: goutamadwant / name: goutamadwant (c944a5a)

[goutamadwant]

@pytorchbot label "release notes: llm"

[metascroy]

Nit: Can we reorder the args here to be temp, top_k, top_p, seed?

We also need to modify this custom ops behavior to work with top_k correctly now.

Do huggingface style, where topk happens before topp, which requires renormalization, roughly something like this.

Something similar is required on the emit path.

def sample(logits, temperature, top_k=None, top_p=1.0, seed=None):
if float(temperature) <= 0:
return torch.argmax(logits, dim=-1)
scaled = logits.float() / temperature
# ── Top-k FIRST (on logits; monotonicity ⇒ top-k logits = top-k probs) ──
if top_k is not None:
k = int(top_k.item())
s_scaled, _ = torch.sort(scaled, dim=-1, descending=True)
kth = s_scaled[..., k - 1 : k]
scaled = torch.where(scaled >= kth, scaled, scaled.new_tensor(float("-inf")))
# ── Top-p on the *renormalized* distribution ──
probs = torch.softmax(scaled, dim=-1) # exp(-inf)=0 → renormalized over top-k
s_probs, _ = torch.sort(probs, dim=-1, descending=True)
cum = torch.cumsum(s_probs, dim=-1)
keep = (cum - s_probs) <= top_p
thresh = torch.where(keep, s_probs, s_probs.new_tensor(float("inf"))).amin(
dim=-1, keepdim=True
)
scaled = torch.where(probs >= thresh, scaled, scaled.new_tensor(float("-inf")))
# ── Gumbel-max ──
if seed is None:
u = torch.rand(scaled.shape)
else:
gen = torch.Generator().manual_seed(int(seed.item()))
u = torch.rand(scaled.shape, generator=gen)
gumbel = -torch.log(-torch.log(u))
return torch.argmax(scaled + gumbel, dim=-1)

[metascroy]

Can we have test case in non-sorted order, e.g., 0.15, 0.5, .05, 0.3?

[metascroy]

Let's keep top_k specified at runtime, with default value int max (select everything)

[metascroy]

Let's do Int clipping to vocab size if topk > vocab size.

[metascroy]

top_k_val is already intorvid?

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 3m 58s —— View job

---

Code review: Add top-k support to MLX sample

I read through all five changed files plus the emit path in ops.py and the helpers it relies on. Scott's inline comments already cover the main direction; below I confirm them with a concrete numeric case and add a few additional findings.

1. Top-k / top-p composition diverges from HuggingFace semantics (correctness) — confirms Scott's sampling.py:41 comment

Both the reference op (custom_ops.py) and the emit path (ops.py) compute the top-k mask and the top-p mask independently on the same un-renormalized probs and then OR the drop masks together. HF does top-k first, renormalizes, then applies top-p to the renormalized distribution. These are not equivalent. Worked example, probs = [0.5, 0.3, 0.15, 0.05], top_k=3, top_p=0.83:

• This PR: top-k keeps {0,1,2}, top-p (prefix mass [0, 0.5, 0.8, 0.95] ≤ 0.83) keeps {0,1,2} → final keep {0,1,2}.
• HF: top-k → renormalize over 0.95 → [0.526, 0.316, 0.158]; top-p on that (prefix [0, 0.526, 0.842] ≤ 0.83) keeps {0,1} → final keep {0,1}.

So the result genuinely differs once both filters are active. Scott's suggested sample() rewrite (top-k on logits → softmax/renormalize → top-p on the renormalized probs) is the fix, and the same restructuring is needed on the emit path — the current LogicalOrNode combine mirrors the reference's OR-compose, so both must move to sequential top-k-then-renormalize-then-top-p. The new test_top_k_and_top_p_compose test passes only because the chosen values happen to agree under both semantics; a case like the one above would catch the divergence.

2. top_k > vocab_size is an out-of-bounds bug (correctness) — relates to Scott's ops.py:3681 clipping comment

custom_ops.py:431:

k=int(top_k.item())
kth=s_probs[..., k-1 : k]

If k > vocab, the slice [k-1:k] is empty, so probs >= kth broadcasts against a size-0 tensor and the mask is wrong/raises. On the emit path the corresponding TakeNode(index = top_k - 1) indexes out of bounds. Clipping k = min(k, vocab) (Scott's point) fixes both. Worth a test with top_k larger than the vocab to lock this in.

3. top_k as runtime default instead of Optional — Scott's ops.py:3679 comment

Making top_k always-present with a default of INT_MAX (clipped to vocab per #2) removes the if top_k is not None branch in both custom_ops.py and the _sample_handler emit block, and the is not None checks in sampling.py. That collapses two code paths into one and naturally yields "keep everything" when unset.

4. Argument ordering — Scott's sampling.py:41 comment

Reorder to temperature, top_k, top_p, seed consistently across mlx::sample, sample_fake, _sample_handler's args[...] unpacking, and the _sample test helper. Note this is a signature change to the registered custom op, so every call site and the require_args(args, 3, 5, ...) indices need to move together.

5. to_int_or_vid already handles the literal/vid split — Scott's ops.py:3682 comment

ops.py:3688-3692:

index=(
IntOrVidOrTid.from_vid(top_k_index.vid)
iftop_k_index.is_videlseIntOrVidOrTid.from_literal(top_k_index.literal)
)

top_k_val comes from ItemIntNode so it's always a Vid here — top_k_iov/top_k_index will never be a literal in this path, making the else branch dead. P.to_int_or_vid_or_tid(top_k_val) (program_builder.py:318) already does this conversion in one call; prefer it over the manual ternary.

6. Test coverage — Scott's test_sample.py:165 comment

• Add a non-sorted-logits case (e.g. probs 0.15, 0.5, 0.05, 0.3) so the test exercises the sort path rather than already-descending input.
• Add a top_k > vocab case (Re-sync with internal repository #2) and a combined case where HF vs OR-compose semantics actually differ (Add support for quantized LeakyReLU #1), so the renormalization fix is verified rather than coincidentally passing.

Minor

• custom_ops.py:431int(top_k.item()) is fine for the eager/meta reference, but confirm the .item() data-dependent read doesn't trip export's symbolic tracing on this op (the seed path already uses ItemIntNode on device, so the device side is fine).
• Docstring in sampling.py still reads cleanly, but once Rename _pt2e to pt2e #3 lands ("None disables top-k") the wording should reflect the always-present default instead.

Overall the lowering scaffolding (Take → ExpandDims → Less → LogicalOr) is well-built and the node-count test is a nice touch; the substantive work is switching both the reference and emit paths to HF-style sequential filtering and adding the vocab-size clip.
· branch fix-mlx-sample-top-k

[goutamadwant]

thanks for your quick review @metascroy will take a look and address them.

[goutamadwant]

@metascroy Addressed below changes in the latest commit.

• Reordered the custom op args to temperature, top_k, top_p, seed.
• Switched sampling to top-k first.. and then top-p on the renormalized probabilities.
• Made the default top_k path use a max-int sentinel clipped to vocab size.
• Removed the independent OR-composed top-k/top-p masks in the emit path.
• Added coverage for non-sorted logits, top_k > vocab, and the combined top-k/top-p case where renormalization changes the result.
  let me know if you have any more suggestions. thanks!

[goutamadwant]

@claude review this code

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 2m 11s —— View job

---

Code review: top-k support for MLX sample (updated)

I re-read all five changed files plus the _sample_handler emit path. The latest commit cleanly addresses every point from the prior round. Summary below.

Prior review items — all resolved ✅

#	Item	Status
1	HF-style sequential filtering (top-k → renormalize → top-p)	✅ custom_ops.py:430-444 does top-k on logits, then softmax (renormalizes over the kept set), then top-p. Emit path mirrors this: top-k WhereNode feeds the SoftmaxNode at ops.py:3695, and the old LogicalOrNode OR-compose is gone (test asserts "LogicalOrNode": 0).
2	top_k > vocab clipping	✅ min(int(top_k.item()), scaled.shape[-1]) (custom_ops.py:430) and MinimumNode against vocab (ops.py:3622-3628).
3	top_k runtime arg with int-max default	✅ sampling.py:40-43 defaults to torch.iinfo(torch.int64).max; the op signature now always takes top_k.
4	Arg ordering temperature, top_k, top_p, seed	✅ Consistent across sample, sample_fake, _sample_handler unpacking, and the _sample test helper. require_args(args, 4, 5, ...) updated.
5	Use to_int_or_vid_or_tid	✅ ops.py:3660 now calls P.to_int_or_vid_or_tid(top_k_index) directly — dead literal branch removed.
6	Test coverage	✅ Added non-sorted input (test_top_k_restricts_to_top_k), top_k > vocab (test_top_k_clips_to_vocab_size), and a true HF-vs-OR divergence case (test_top_k_and_top_p_compose, top_k=3, top_p=0.83 → {0,1}). Plus an end-to-end on-device top-k test and a SampleTopKTest node-count case.

I worked the test_top_k_and_top_p_compose numbers by hand: top_k=3 keeps {0.5,0.3,0.15}, renormalizes to [0.526,0.316,0.158], and top-p=0.83 (prefix [0,0.526,0.842] ≤ 0.83) keeps {0,1}. This is the case that would have failed under the old OR-compose, so the renormalization fix is genuinely exercised now, not coincidentally passing.

Minor observations (non-blocking)

1. custom_ops.py:432 empty-slice edge when k=0.k is clipped above by vocab but not below. A top_k=0 would make s_scaled[..., -1:0] empty and broadcast oddly. This isn't reachable through SamplingHead (default is int-max, callers pass ≥1), so it's only a concern if the op is called directly with 0. A max(k, 1) clamp or a one-line docstring note that top_k ≥ 1 is expected would close it. Not worth blocking on.

2. Emit-path top-k threshold vs. reference — confirmed equivalent. The reference keeps scaled >= kth (>=), while the emit path drops via LessNode(scaled, top_k_thresh) i.e. keeps scaled >= thresh too — consistent. Good. The negate→sort→negate dance to get a descending sort (since SortNode is ascending-only) and the TakeNode(index = k-1) correctly pick the k-th largest. Matches the top-p chain's existing pattern.

3. int(top_k.item()) data-dependent read (custom_ops.py:430): fine for the eager/meta reference; the device path uses ItemIntNode so symbolic tracing is unaffected, same as the existing seed handling.

4. Buffer-reuse comments in ops.py are accurate (verified scaled/probs/sorted_p live-range notes against their actual reads).

Verdict

Looks good to me — the substantive correctness fix (sequential HF-style filtering + vocab clip) is in on both the reference and the lowered graph, and the tests lock in the behaviors that distinguish it from the old implementation. I couldn't execute the test suite in this sandbox (command approval), so please confirm the documented unittest + compileall runs are green in CI before merging. Only the k=0 clamp (#1) is worth a quick thought; everything else is optional polish.
· branch fix-mlx-sample-top-k

[metascroy]

Looks great! Thanks for the contribution @goutamadwant!

[goutamadwant]

Fixes#20548

Summary

• Thread optional top_k through SamplingHead and mlx::sample.
• Apply top-k filtering as an additional threshold mask that composes with the existing top-p nucleus mask.
• Add eager, export, end-to-end, and MLX lowering coverage for top-k sampling.

Test plan

• PYTHONPATH=src:. python3 -m unittest executorch.backends.mlx.test.test_sample.TestSampleOp executorch.backends.mlx.test.test_sample.TestSampleExport
• python3 -m compileall -q backends/mlx/llm/sampling.py backends/mlx/custom_ops.py backends/mlx/ops.py backends/mlx/test/test_sample.py backends/mlx/test/test_ops.py

cc @metascroy

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20564

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 12 Pending

As of commit 7e9ae35 with merge base 55a71e6 ():

NEW FAILURE - The following job has failed:

• pull / test-eval_llama-wikitext-linux / linux-job (gh)
  RuntimeError: Command docker exec -t f07b8ecee9be9cc829d92ad63b572fb3fb94908a1e2f16b39fadf2553b965213 /exec failed with exit code 1

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: goutamadwant / name: goutamadwant (c944a5a)

[goutamadwant]

@pytorchbot label "release notes: llm"

[metascroy]

Nit: Can we reorder the args here to be temp, top_k, top_p, seed?

We also need to modify this custom ops behavior to work with top_k correctly now.

Do huggingface style, where topk happens before topp, which requires renormalization, roughly something like this.

Something similar is required on the emit path.

def sample(logits, temperature, top_k=None, top_p=1.0, seed=None):
if float(temperature) <= 0:
return torch.argmax(logits, dim=-1)
scaled = logits.float() / temperature
# ── Top-k FIRST (on logits; monotonicity ⇒ top-k logits = top-k probs) ──
if top_k is not None:
k = int(top_k.item())
s_scaled, _ = torch.sort(scaled, dim=-1, descending=True)
kth = s_scaled[..., k - 1 : k]
scaled = torch.where(scaled >= kth, scaled, scaled.new_tensor(float("-inf")))
# ── Top-p on the *renormalized* distribution ──
probs = torch.softmax(scaled, dim=-1) # exp(-inf)=0 → renormalized over top-k
s_probs, _ = torch.sort(probs, dim=-1, descending=True)
cum = torch.cumsum(s_probs, dim=-1)
keep = (cum - s_probs) <= top_p
thresh = torch.where(keep, s_probs, s_probs.new_tensor(float("inf"))).amin(
dim=-1, keepdim=True
)
scaled = torch.where(probs >= thresh, scaled, scaled.new_tensor(float("-inf")))
# ── Gumbel-max ──
if seed is None:
u = torch.rand(scaled.shape)
else:
gen = torch.Generator().manual_seed(int(seed.item()))
u = torch.rand(scaled.shape, generator=gen)
gumbel = -torch.log(-torch.log(u))
return torch.argmax(scaled + gumbel, dim=-1)

[metascroy]

Can we have test case in non-sorted order, e.g., 0.15, 0.5, .05, 0.3?

[metascroy]

Let's keep top_k specified at runtime, with default value int max (select everything)

[metascroy]

Let's do Int clipping to vocab size if topk > vocab size.

[metascroy]

top_k_val is already intorvid?

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 3m 58s —— View job

---

Code review: Add top-k support to MLX sample

I read through all five changed files plus the emit path in ops.py and the helpers it relies on. Scott's inline comments already cover the main direction; below I confirm them with a concrete numeric case and add a few additional findings.

1. Top-k / top-p composition diverges from HuggingFace semantics (correctness) — confirms Scott's sampling.py:41 comment

Both the reference op (custom_ops.py) and the emit path (ops.py) compute the top-k mask and the top-p mask independently on the same un-renormalized probs and then OR the drop masks together. HF does top-k first, renormalizes, then applies top-p to the renormalized distribution. These are not equivalent. Worked example, probs = [0.5, 0.3, 0.15, 0.05], top_k=3, top_p=0.83:

• This PR: top-k keeps {0,1,2}, top-p (prefix mass [0, 0.5, 0.8, 0.95] ≤ 0.83) keeps {0,1,2} → final keep {0,1,2}.
• HF: top-k → renormalize over 0.95 → [0.526, 0.316, 0.158]; top-p on that (prefix [0, 0.526, 0.842] ≤ 0.83) keeps {0,1} → final keep {0,1}.

So the result genuinely differs once both filters are active. Scott's suggested sample() rewrite (top-k on logits → softmax/renormalize → top-p on the renormalized probs) is the fix, and the same restructuring is needed on the emit path — the current LogicalOrNode combine mirrors the reference's OR-compose, so both must move to sequential top-k-then-renormalize-then-top-p. The new test_top_k_and_top_p_compose test passes only because the chosen values happen to agree under both semantics; a case like the one above would catch the divergence.

2. top_k > vocab_size is an out-of-bounds bug (correctness) — relates to Scott's ops.py:3681 clipping comment

custom_ops.py:431:

k=int(top_k.item())
kth=s_probs[..., k-1 : k]

If k > vocab, the slice [k-1:k] is empty, so probs >= kth broadcasts against a size-0 tensor and the mask is wrong/raises. On the emit path the corresponding TakeNode(index = top_k - 1) indexes out of bounds. Clipping k = min(k, vocab) (Scott's point) fixes both. Worth a test with top_k larger than the vocab to lock this in.

3. top_k as runtime default instead of Optional — Scott's ops.py:3679 comment

Making top_k always-present with a default of INT_MAX (clipped to vocab per #2) removes the if top_k is not None branch in both custom_ops.py and the _sample_handler emit block, and the is not None checks in sampling.py. That collapses two code paths into one and naturally yields "keep everything" when unset.

4. Argument ordering — Scott's sampling.py:41 comment

Reorder to temperature, top_k, top_p, seed consistently across mlx::sample, sample_fake, _sample_handler's args[...] unpacking, and the _sample test helper. Note this is a signature change to the registered custom op, so every call site and the require_args(args, 3, 5, ...) indices need to move together.

5. to_int_or_vid already handles the literal/vid split — Scott's ops.py:3682 comment

ops.py:3688-3692:

index=(
IntOrVidOrTid.from_vid(top_k_index.vid)
iftop_k_index.is_videlseIntOrVidOrTid.from_literal(top_k_index.literal)
)

top_k_val comes from ItemIntNode so it's always a Vid here — top_k_iov/top_k_index will never be a literal in this path, making the else branch dead. P.to_int_or_vid_or_tid(top_k_val) (program_builder.py:318) already does this conversion in one call; prefer it over the manual ternary.

6. Test coverage — Scott's test_sample.py:165 comment

• Add a non-sorted-logits case (e.g. probs 0.15, 0.5, 0.05, 0.3) so the test exercises the sort path rather than already-descending input.
• Add a top_k > vocab case (Re-sync with internal repository #2) and a combined case where HF vs OR-compose semantics actually differ (Add support for quantized LeakyReLU #1), so the renormalization fix is verified rather than coincidentally passing.

Minor

• custom_ops.py:431int(top_k.item()) is fine for the eager/meta reference, but confirm the .item() data-dependent read doesn't trip export's symbolic tracing on this op (the seed path already uses ItemIntNode on device, so the device side is fine).
• Docstring in sampling.py still reads cleanly, but once Rename _pt2e to pt2e #3 lands ("None disables top-k") the wording should reflect the always-present default instead.

Overall the lowering scaffolding (Take → ExpandDims → Less → LogicalOr) is well-built and the node-count test is a nice touch; the substantive work is switching both the reference and emit paths to HF-style sequential filtering and adding the vocab-size clip.
· branch fix-mlx-sample-top-k

[goutamadwant]

thanks for your quick review @metascroy will take a look and address them.

[goutamadwant]

@metascroy Addressed below changes in the latest commit.

• Reordered the custom op args to temperature, top_k, top_p, seed.
• Switched sampling to top-k first.. and then top-p on the renormalized probabilities.
• Made the default top_k path use a max-int sentinel clipped to vocab size.
• Removed the independent OR-composed top-k/top-p masks in the emit path.
• Added coverage for non-sorted logits, top_k > vocab, and the combined top-k/top-p case where renormalization changes the result.
  let me know if you have any more suggestions. thanks!

[goutamadwant]

@claude review this code

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 2m 11s —— View job

---

Code review: top-k support for MLX sample (updated)

I re-read all five changed files plus the _sample_handler emit path. The latest commit cleanly addresses every point from the prior round. Summary below.

Prior review items — all resolved ✅

#	Item	Status
1	HF-style sequential filtering (top-k → renormalize → top-p)	✅ custom_ops.py:430-444 does top-k on logits, then softmax (renormalizes over the kept set), then top-p. Emit path mirrors this: top-k WhereNode feeds the SoftmaxNode at ops.py:3695, and the old LogicalOrNode OR-compose is gone (test asserts "LogicalOrNode": 0).
2	top_k > vocab clipping	✅ min(int(top_k.item()), scaled.shape[-1]) (custom_ops.py:430) and MinimumNode against vocab (ops.py:3622-3628).
3	top_k runtime arg with int-max default	✅ sampling.py:40-43 defaults to torch.iinfo(torch.int64).max; the op signature now always takes top_k.
4	Arg ordering temperature, top_k, top_p, seed	✅ Consistent across sample, sample_fake, _sample_handler unpacking, and the _sample test helper. require_args(args, 4, 5, ...) updated.
5	Use to_int_or_vid_or_tid	✅ ops.py:3660 now calls P.to_int_or_vid_or_tid(top_k_index) directly — dead literal branch removed.
6	Test coverage	✅ Added non-sorted input (test_top_k_restricts_to_top_k), top_k > vocab (test_top_k_clips_to_vocab_size), and a true HF-vs-OR divergence case (test_top_k_and_top_p_compose, top_k=3, top_p=0.83 → {0,1}). Plus an end-to-end on-device top-k test and a SampleTopKTest node-count case.

I worked the test_top_k_and_top_p_compose numbers by hand: top_k=3 keeps {0.5,0.3,0.15}, renormalizes to [0.526,0.316,0.158], and top-p=0.83 (prefix [0,0.526,0.842] ≤ 0.83) keeps {0,1}. This is the case that would have failed under the old OR-compose, so the renormalization fix is genuinely exercised now, not coincidentally passing.

Minor observations (non-blocking)

1. custom_ops.py:432 empty-slice edge when k=0.k is clipped above by vocab but not below. A top_k=0 would make s_scaled[..., -1:0] empty and broadcast oddly. This isn't reachable through SamplingHead (default is int-max, callers pass ≥1), so it's only a concern if the op is called directly with 0. A max(k, 1) clamp or a one-line docstring note that top_k ≥ 1 is expected would close it. Not worth blocking on.

2. Emit-path top-k threshold vs. reference — confirmed equivalent. The reference keeps scaled >= kth (>=), while the emit path drops via LessNode(scaled, top_k_thresh) i.e. keeps scaled >= thresh too — consistent. Good. The negate→sort→negate dance to get a descending sort (since SortNode is ascending-only) and the TakeNode(index = k-1) correctly pick the k-th largest. Matches the top-p chain's existing pattern.

3. int(top_k.item()) data-dependent read (custom_ops.py:430): fine for the eager/meta reference; the device path uses ItemIntNode so symbolic tracing is unaffected, same as the existing seed handling.

4. Buffer-reuse comments in ops.py are accurate (verified scaled/probs/sorted_p live-range notes against their actual reads).

Verdict

Looks good to me — the substantive correctness fix (sequential HF-style filtering + vocab clip) is in on both the reference and the lowered graph, and the tests lock in the behaviors that distinguish it from the old implementation. I couldn't execute the test suite in this sandbox (command approval), so please confirm the documented unittest + compileall runs are green in CI before merging. Only the k=0 clamp (#1) is worth a quick thought; everything else is optional polish.
· branch fix-mlx-sample-top-k

[metascroy]

Looks great! Thanks for the contribution @goutamadwant!

[goutamadwant]

Fixes#20548

Summary

• Thread optional top_k through SamplingHead and mlx::sample.
• Apply top-k filtering as an additional threshold mask that composes with the existing top-p nucleus mask.
• Add eager, export, end-to-end, and MLX lowering coverage for top-k sampling.

Test plan

• PYTHONPATH=src:. python3 -m unittest executorch.backends.mlx.test.test_sample.TestSampleOp executorch.backends.mlx.test.test_sample.TestSampleExport
• python3 -m compileall -q backends/mlx/llm/sampling.py backends/mlx/custom_ops.py backends/mlx/ops.py backends/mlx/test/test_sample.py backends/mlx/test/test_ops.py

cc @metascroy

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20564

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 12 Pending

As of commit 7e9ae35 with merge base 55a71e6 ():

NEW FAILURE - The following job has failed:

• pull / test-eval_llama-wikitext-linux / linux-job (gh)
  RuntimeError: Command docker exec -t f07b8ecee9be9cc829d92ad63b572fb3fb94908a1e2f16b39fadf2553b965213 /exec failed with exit code 1

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: goutamadwant / name: goutamadwant (c944a5a)

[goutamadwant]

@pytorchbot label "release notes: llm"

[metascroy]

Nit: Can we reorder the args here to be temp, top_k, top_p, seed?

We also need to modify this custom ops behavior to work with top_k correctly now.

Do huggingface style, where topk happens before topp, which requires renormalization, roughly something like this.

Something similar is required on the emit path.

def sample(logits, temperature, top_k=None, top_p=1.0, seed=None):
if float(temperature) <= 0:
return torch.argmax(logits, dim=-1)
scaled = logits.float() / temperature
# ── Top-k FIRST (on logits; monotonicity ⇒ top-k logits = top-k probs) ──
if top_k is not None:
k = int(top_k.item())
s_scaled, _ = torch.sort(scaled, dim=-1, descending=True)
kth = s_scaled[..., k - 1 : k]
scaled = torch.where(scaled >= kth, scaled, scaled.new_tensor(float("-inf")))
# ── Top-p on the *renormalized* distribution ──
probs = torch.softmax(scaled, dim=-1) # exp(-inf)=0 → renormalized over top-k
s_probs, _ = torch.sort(probs, dim=-1, descending=True)
cum = torch.cumsum(s_probs, dim=-1)
keep = (cum - s_probs) <= top_p
thresh = torch.where(keep, s_probs, s_probs.new_tensor(float("inf"))).amin(
dim=-1, keepdim=True
)
scaled = torch.where(probs >= thresh, scaled, scaled.new_tensor(float("-inf")))
# ── Gumbel-max ──
if seed is None:
u = torch.rand(scaled.shape)
else:
gen = torch.Generator().manual_seed(int(seed.item()))
u = torch.rand(scaled.shape, generator=gen)
gumbel = -torch.log(-torch.log(u))
return torch.argmax(scaled + gumbel, dim=-1)

[metascroy]

Can we have test case in non-sorted order, e.g., 0.15, 0.5, .05, 0.3?

[metascroy]

Let's keep top_k specified at runtime, with default value int max (select everything)

[metascroy]

Let's do Int clipping to vocab size if topk > vocab size.

[metascroy]

top_k_val is already intorvid?

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 3m 58s —— View job

---

Code review: Add top-k support to MLX sample

I read through all five changed files plus the emit path in ops.py and the helpers it relies on. Scott's inline comments already cover the main direction; below I confirm them with a concrete numeric case and add a few additional findings.

1. Top-k / top-p composition diverges from HuggingFace semantics (correctness) — confirms Scott's sampling.py:41 comment

Both the reference op (custom_ops.py) and the emit path (ops.py) compute the top-k mask and the top-p mask independently on the same un-renormalized probs and then OR the drop masks together. HF does top-k first, renormalizes, then applies top-p to the renormalized distribution. These are not equivalent. Worked example, probs = [0.5, 0.3, 0.15, 0.05], top_k=3, top_p=0.83:

• This PR: top-k keeps {0,1,2}, top-p (prefix mass [0, 0.5, 0.8, 0.95] ≤ 0.83) keeps {0,1,2} → final keep {0,1,2}.
• HF: top-k → renormalize over 0.95 → [0.526, 0.316, 0.158]; top-p on that (prefix [0, 0.526, 0.842] ≤ 0.83) keeps {0,1} → final keep {0,1}.

So the result genuinely differs once both filters are active. Scott's suggested sample() rewrite (top-k on logits → softmax/renormalize → top-p on the renormalized probs) is the fix, and the same restructuring is needed on the emit path — the current LogicalOrNode combine mirrors the reference's OR-compose, so both must move to sequential top-k-then-renormalize-then-top-p. The new test_top_k_and_top_p_compose test passes only because the chosen values happen to agree under both semantics; a case like the one above would catch the divergence.

2. top_k > vocab_size is an out-of-bounds bug (correctness) — relates to Scott's ops.py:3681 clipping comment

custom_ops.py:431:

k=int(top_k.item())
kth=s_probs[..., k-1 : k]

If k > vocab, the slice [k-1:k] is empty, so probs >= kth broadcasts against a size-0 tensor and the mask is wrong/raises. On the emit path the corresponding TakeNode(index = top_k - 1) indexes out of bounds. Clipping k = min(k, vocab) (Scott's point) fixes both. Worth a test with top_k larger than the vocab to lock this in.

3. top_k as runtime default instead of Optional — Scott's ops.py:3679 comment

Making top_k always-present with a default of INT_MAX (clipped to vocab per #2) removes the if top_k is not None branch in both custom_ops.py and the _sample_handler emit block, and the is not None checks in sampling.py. That collapses two code paths into one and naturally yields "keep everything" when unset.

4. Argument ordering — Scott's sampling.py:41 comment

Reorder to temperature, top_k, top_p, seed consistently across mlx::sample, sample_fake, _sample_handler's args[...] unpacking, and the _sample test helper. Note this is a signature change to the registered custom op, so every call site and the require_args(args, 3, 5, ...) indices need to move together.

5. to_int_or_vid already handles the literal/vid split — Scott's ops.py:3682 comment

ops.py:3688-3692:

index=(
IntOrVidOrTid.from_vid(top_k_index.vid)
iftop_k_index.is_videlseIntOrVidOrTid.from_literal(top_k_index.literal)
)

top_k_val comes from ItemIntNode so it's always a Vid here — top_k_iov/top_k_index will never be a literal in this path, making the else branch dead. P.to_int_or_vid_or_tid(top_k_val) (program_builder.py:318) already does this conversion in one call; prefer it over the manual ternary.

6. Test coverage — Scott's test_sample.py:165 comment

• Add a non-sorted-logits case (e.g. probs 0.15, 0.5, 0.05, 0.3) so the test exercises the sort path rather than already-descending input.
• Add a top_k > vocab case (Re-sync with internal repository #2) and a combined case where HF vs OR-compose semantics actually differ (Add support for quantized LeakyReLU #1), so the renormalization fix is verified rather than coincidentally passing.

Minor

• custom_ops.py:431int(top_k.item()) is fine for the eager/meta reference, but confirm the .item() data-dependent read doesn't trip export's symbolic tracing on this op (the seed path already uses ItemIntNode on device, so the device side is fine).
• Docstring in sampling.py still reads cleanly, but once Rename _pt2e to pt2e #3 lands ("None disables top-k") the wording should reflect the always-present default instead.

Overall the lowering scaffolding (Take → ExpandDims → Less → LogicalOr) is well-built and the node-count test is a nice touch; the substantive work is switching both the reference and emit paths to HF-style sequential filtering and adding the vocab-size clip.
· branch fix-mlx-sample-top-k

[goutamadwant]

thanks for your quick review @metascroy will take a look and address them.

[goutamadwant]

@metascroy Addressed below changes in the latest commit.

• Reordered the custom op args to temperature, top_k, top_p, seed.
• Switched sampling to top-k first.. and then top-p on the renormalized probabilities.
• Made the default top_k path use a max-int sentinel clipped to vocab size.
• Removed the independent OR-composed top-k/top-p masks in the emit path.
• Added coverage for non-sorted logits, top_k > vocab, and the combined top-k/top-p case where renormalization changes the result.
  let me know if you have any more suggestions. thanks!

[goutamadwant]

@claude review this code

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 2m 11s —— View job

---

Code review: top-k support for MLX sample (updated)

I re-read all five changed files plus the _sample_handler emit path. The latest commit cleanly addresses every point from the prior round. Summary below.

Prior review items — all resolved ✅

#	Item	Status
1	HF-style sequential filtering (top-k → renormalize → top-p)	✅ custom_ops.py:430-444 does top-k on logits, then softmax (renormalizes over the kept set), then top-p. Emit path mirrors this: top-k WhereNode feeds the SoftmaxNode at ops.py:3695, and the old LogicalOrNode OR-compose is gone (test asserts "LogicalOrNode": 0).
2	top_k > vocab clipping	✅ min(int(top_k.item()), scaled.shape[-1]) (custom_ops.py:430) and MinimumNode against vocab (ops.py:3622-3628).
3	top_k runtime arg with int-max default	✅ sampling.py:40-43 defaults to torch.iinfo(torch.int64).max; the op signature now always takes top_k.
4	Arg ordering temperature, top_k, top_p, seed	✅ Consistent across sample, sample_fake, _sample_handler unpacking, and the _sample test helper. require_args(args, 4, 5, ...) updated.
5	Use to_int_or_vid_or_tid	✅ ops.py:3660 now calls P.to_int_or_vid_or_tid(top_k_index) directly — dead literal branch removed.
6	Test coverage	✅ Added non-sorted input (test_top_k_restricts_to_top_k), top_k > vocab (test_top_k_clips_to_vocab_size), and a true HF-vs-OR divergence case (test_top_k_and_top_p_compose, top_k=3, top_p=0.83 → {0,1}). Plus an end-to-end on-device top-k test and a SampleTopKTest node-count case.

I worked the test_top_k_and_top_p_compose numbers by hand: top_k=3 keeps {0.5,0.3,0.15}, renormalizes to [0.526,0.316,0.158], and top-p=0.83 (prefix [0,0.526,0.842] ≤ 0.83) keeps {0,1}. This is the case that would have failed under the old OR-compose, so the renormalization fix is genuinely exercised now, not coincidentally passing.

Minor observations (non-blocking)

1. custom_ops.py:432 empty-slice edge when k=0.k is clipped above by vocab but not below. A top_k=0 would make s_scaled[..., -1:0] empty and broadcast oddly. This isn't reachable through SamplingHead (default is int-max, callers pass ≥1), so it's only a concern if the op is called directly with 0. A max(k, 1) clamp or a one-line docstring note that top_k ≥ 1 is expected would close it. Not worth blocking on.

2. Emit-path top-k threshold vs. reference — confirmed equivalent. The reference keeps scaled >= kth (>=), while the emit path drops via LessNode(scaled, top_k_thresh) i.e. keeps scaled >= thresh too — consistent. Good. The negate→sort→negate dance to get a descending sort (since SortNode is ascending-only) and the TakeNode(index = k-1) correctly pick the k-th largest. Matches the top-p chain's existing pattern.

3. int(top_k.item()) data-dependent read (custom_ops.py:430): fine for the eager/meta reference; the device path uses ItemIntNode so symbolic tracing is unaffected, same as the existing seed handling.

4. Buffer-reuse comments in ops.py are accurate (verified scaled/probs/sorted_p live-range notes against their actual reads).

Verdict

Looks good to me — the substantive correctness fix (sequential HF-style filtering + vocab clip) is in on both the reference and the lowered graph, and the tests lock in the behaviors that distinguish it from the old implementation. I couldn't execute the test suite in this sandbox (command approval), so please confirm the documented unittest + compileall runs are green in CI before merging. Only the k=0 clamp (#1) is worth a quick thought; everything else is optional polish.
· branch fix-mlx-sample-top-k

[metascroy]

Looks great! Thanks for the contribution @goutamadwant!

[goutamadwant]

Fixes#20548

Summary

• Thread optional top_k through SamplingHead and mlx::sample.
• Apply top-k filtering as an additional threshold mask that composes with the existing top-p nucleus mask.
• Add eager, export, end-to-end, and MLX lowering coverage for top-k sampling.

Test plan

• PYTHONPATH=src:. python3 -m unittest executorch.backends.mlx.test.test_sample.TestSampleOp executorch.backends.mlx.test.test_sample.TestSampleExport
• python3 -m compileall -q backends/mlx/llm/sampling.py backends/mlx/custom_ops.py backends/mlx/ops.py backends/mlx/test/test_sample.py backends/mlx/test/test_ops.py

cc @metascroy

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20564

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 12 Pending

As of commit 7e9ae35 with merge base 55a71e6 ():

NEW FAILURE - The following job has failed:

• pull / test-eval_llama-wikitext-linux / linux-job (gh)
  RuntimeError: Command docker exec -t f07b8ecee9be9cc829d92ad63b572fb3fb94908a1e2f16b39fadf2553b965213 /exec failed with exit code 1

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: goutamadwant / name: goutamadwant (c944a5a)

[goutamadwant]

@pytorchbot label "release notes: llm"

[metascroy]

Nit: Can we reorder the args here to be temp, top_k, top_p, seed?

We also need to modify this custom ops behavior to work with top_k correctly now.

Do huggingface style, where topk happens before topp, which requires renormalization, roughly something like this.

Something similar is required on the emit path.

def sample(logits, temperature, top_k=None, top_p=1.0, seed=None):
if float(temperature) <= 0:
return torch.argmax(logits, dim=-1)
scaled = logits.float() / temperature
# ── Top-k FIRST (on logits; monotonicity ⇒ top-k logits = top-k probs) ──
if top_k is not None:
k = int(top_k.item())
s_scaled, _ = torch.sort(scaled, dim=-1, descending=True)
kth = s_scaled[..., k - 1 : k]
scaled = torch.where(scaled >= kth, scaled, scaled.new_tensor(float("-inf")))
# ── Top-p on the *renormalized* distribution ──
probs = torch.softmax(scaled, dim=-1) # exp(-inf)=0 → renormalized over top-k
s_probs, _ = torch.sort(probs, dim=-1, descending=True)
cum = torch.cumsum(s_probs, dim=-1)
keep = (cum - s_probs) <= top_p
thresh = torch.where(keep, s_probs, s_probs.new_tensor(float("inf"))).amin(
dim=-1, keepdim=True
)
scaled = torch.where(probs >= thresh, scaled, scaled.new_tensor(float("-inf")))
# ── Gumbel-max ──
if seed is None:
u = torch.rand(scaled.shape)
else:
gen = torch.Generator().manual_seed(int(seed.item()))
u = torch.rand(scaled.shape, generator=gen)
gumbel = -torch.log(-torch.log(u))
return torch.argmax(scaled + gumbel, dim=-1)

[metascroy]

Can we have test case in non-sorted order, e.g., 0.15, 0.5, .05, 0.3?

[metascroy]

Let's keep top_k specified at runtime, with default value int max (select everything)

[metascroy]

Let's do Int clipping to vocab size if topk > vocab size.

[metascroy]

top_k_val is already intorvid?

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 3m 58s —— View job

---

Code review: Add top-k support to MLX sample

I read through all five changed files plus the emit path in ops.py and the helpers it relies on. Scott's inline comments already cover the main direction; below I confirm them with a concrete numeric case and add a few additional findings.

1. Top-k / top-p composition diverges from HuggingFace semantics (correctness) — confirms Scott's sampling.py:41 comment

Both the reference op (custom_ops.py) and the emit path (ops.py) compute the top-k mask and the top-p mask independently on the same un-renormalized probs and then OR the drop masks together. HF does top-k first, renormalizes, then applies top-p to the renormalized distribution. These are not equivalent. Worked example, probs = [0.5, 0.3, 0.15, 0.05], top_k=3, top_p=0.83:

• This PR: top-k keeps {0,1,2}, top-p (prefix mass [0, 0.5, 0.8, 0.95] ≤ 0.83) keeps {0,1,2} → final keep {0,1,2}.
• HF: top-k → renormalize over 0.95 → [0.526, 0.316, 0.158]; top-p on that (prefix [0, 0.526, 0.842] ≤ 0.83) keeps {0,1} → final keep {0,1}.

So the result genuinely differs once both filters are active. Scott's suggested sample() rewrite (top-k on logits → softmax/renormalize → top-p on the renormalized probs) is the fix, and the same restructuring is needed on the emit path — the current LogicalOrNode combine mirrors the reference's OR-compose, so both must move to sequential top-k-then-renormalize-then-top-p. The new test_top_k_and_top_p_compose test passes only because the chosen values happen to agree under both semantics; a case like the one above would catch the divergence.

2. top_k > vocab_size is an out-of-bounds bug (correctness) — relates to Scott's ops.py:3681 clipping comment

custom_ops.py:431:

k=int(top_k.item())
kth=s_probs[..., k-1 : k]

If k > vocab, the slice [k-1:k] is empty, so probs >= kth broadcasts against a size-0 tensor and the mask is wrong/raises. On the emit path the corresponding TakeNode(index = top_k - 1) indexes out of bounds. Clipping k = min(k, vocab) (Scott's point) fixes both. Worth a test with top_k larger than the vocab to lock this in.

3. top_k as runtime default instead of Optional — Scott's ops.py:3679 comment

Making top_k always-present with a default of INT_MAX (clipped to vocab per #2) removes the if top_k is not None branch in both custom_ops.py and the _sample_handler emit block, and the is not None checks in sampling.py. That collapses two code paths into one and naturally yields "keep everything" when unset.

4. Argument ordering — Scott's sampling.py:41 comment

Reorder to temperature, top_k, top_p, seed consistently across mlx::sample, sample_fake, _sample_handler's args[...] unpacking, and the _sample test helper. Note this is a signature change to the registered custom op, so every call site and the require_args(args, 3, 5, ...) indices need to move together.

5. to_int_or_vid already handles the literal/vid split — Scott's ops.py:3682 comment

ops.py:3688-3692:

index=(
IntOrVidOrTid.from_vid(top_k_index.vid)
iftop_k_index.is_videlseIntOrVidOrTid.from_literal(top_k_index.literal)
)

top_k_val comes from ItemIntNode so it's always a Vid here — top_k_iov/top_k_index will never be a literal in this path, making the else branch dead. P.to_int_or_vid_or_tid(top_k_val) (program_builder.py:318) already does this conversion in one call; prefer it over the manual ternary.

6. Test coverage — Scott's test_sample.py:165 comment

• Add a non-sorted-logits case (e.g. probs 0.15, 0.5, 0.05, 0.3) so the test exercises the sort path rather than already-descending input.
• Add a top_k > vocab case (Re-sync with internal repository #2) and a combined case where HF vs OR-compose semantics actually differ (Add support for quantized LeakyReLU #1), so the renormalization fix is verified rather than coincidentally passing.

Minor

• custom_ops.py:431int(top_k.item()) is fine for the eager/meta reference, but confirm the .item() data-dependent read doesn't trip export's symbolic tracing on this op (the seed path already uses ItemIntNode on device, so the device side is fine).
• Docstring in sampling.py still reads cleanly, but once Rename _pt2e to pt2e #3 lands ("None disables top-k") the wording should reflect the always-present default instead.

Overall the lowering scaffolding (Take → ExpandDims → Less → LogicalOr) is well-built and the node-count test is a nice touch; the substantive work is switching both the reference and emit paths to HF-style sequential filtering and adding the vocab-size clip.
· branch fix-mlx-sample-top-k

[goutamadwant]

thanks for your quick review @metascroy will take a look and address them.

[goutamadwant]

@metascroy Addressed below changes in the latest commit.

• Reordered the custom op args to temperature, top_k, top_p, seed.
• Switched sampling to top-k first.. and then top-p on the renormalized probabilities.
• Made the default top_k path use a max-int sentinel clipped to vocab size.
• Removed the independent OR-composed top-k/top-p masks in the emit path.
• Added coverage for non-sorted logits, top_k > vocab, and the combined top-k/top-p case where renormalization changes the result.
  let me know if you have any more suggestions. thanks!

[goutamadwant]

@claude review this code

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 2m 11s —— View job

---

Code review: top-k support for MLX sample (updated)

I re-read all five changed files plus the _sample_handler emit path. The latest commit cleanly addresses every point from the prior round. Summary below.

Prior review items — all resolved ✅

#	Item	Status
1	HF-style sequential filtering (top-k → renormalize → top-p)	✅ custom_ops.py:430-444 does top-k on logits, then softmax (renormalizes over the kept set), then top-p. Emit path mirrors this: top-k WhereNode feeds the SoftmaxNode at ops.py:3695, and the old LogicalOrNode OR-compose is gone (test asserts "LogicalOrNode": 0).
2	top_k > vocab clipping	✅ min(int(top_k.item()), scaled.shape[-1]) (custom_ops.py:430) and MinimumNode against vocab (ops.py:3622-3628).
3	top_k runtime arg with int-max default	✅ sampling.py:40-43 defaults to torch.iinfo(torch.int64).max; the op signature now always takes top_k.
4	Arg ordering temperature, top_k, top_p, seed	✅ Consistent across sample, sample_fake, _sample_handler unpacking, and the _sample test helper. require_args(args, 4, 5, ...) updated.
5	Use to_int_or_vid_or_tid	✅ ops.py:3660 now calls P.to_int_or_vid_or_tid(top_k_index) directly — dead literal branch removed.
6	Test coverage	✅ Added non-sorted input (test_top_k_restricts_to_top_k), top_k > vocab (test_top_k_clips_to_vocab_size), and a true HF-vs-OR divergence case (test_top_k_and_top_p_compose, top_k=3, top_p=0.83 → {0,1}). Plus an end-to-end on-device top-k test and a SampleTopKTest node-count case.

I worked the test_top_k_and_top_p_compose numbers by hand: top_k=3 keeps {0.5,0.3,0.15}, renormalizes to [0.526,0.316,0.158], and top-p=0.83 (prefix [0,0.526,0.842] ≤ 0.83) keeps {0,1}. This is the case that would have failed under the old OR-compose, so the renormalization fix is genuinely exercised now, not coincidentally passing.

Minor observations (non-blocking)

1. custom_ops.py:432 empty-slice edge when k=0.k is clipped above by vocab but not below. A top_k=0 would make s_scaled[..., -1:0] empty and broadcast oddly. This isn't reachable through SamplingHead (default is int-max, callers pass ≥1), so it's only a concern if the op is called directly with 0. A max(k, 1) clamp or a one-line docstring note that top_k ≥ 1 is expected would close it. Not worth blocking on.

2. Emit-path top-k threshold vs. reference — confirmed equivalent. The reference keeps scaled >= kth (>=), while the emit path drops via LessNode(scaled, top_k_thresh) i.e. keeps scaled >= thresh too — consistent. Good. The negate→sort→negate dance to get a descending sort (since SortNode is ascending-only) and the TakeNode(index = k-1) correctly pick the k-th largest. Matches the top-p chain's existing pattern.

3. int(top_k.item()) data-dependent read (custom_ops.py:430): fine for the eager/meta reference; the device path uses ItemIntNode so symbolic tracing is unaffected, same as the existing seed handling.

4. Buffer-reuse comments in ops.py are accurate (verified scaled/probs/sorted_p live-range notes against their actual reads).

Verdict

Looks good to me — the substantive correctness fix (sequential HF-style filtering + vocab clip) is in on both the reference and the lowered graph, and the tests lock in the behaviors that distinguish it from the old implementation. I couldn't execute the test suite in this sandbox (command approval), so please confirm the documented unittest + compileall runs are green in CI before merging. Only the k=0 clamp (#1) is worth a quick thought; everything else is optional polish.
· branch fix-mlx-sample-top-k

[metascroy]

Looks great! Thanks for the contribution @goutamadwant!

[goutamadwant]

Fixes#20548

Summary

• Thread optional top_k through SamplingHead and mlx::sample.
• Apply top-k filtering as an additional threshold mask that composes with the existing top-p nucleus mask.
• Add eager, export, end-to-end, and MLX lowering coverage for top-k sampling.

Test plan

• PYTHONPATH=src:. python3 -m unittest executorch.backends.mlx.test.test_sample.TestSampleOp executorch.backends.mlx.test.test_sample.TestSampleExport
• python3 -m compileall -q backends/mlx/llm/sampling.py backends/mlx/custom_ops.py backends/mlx/ops.py backends/mlx/test/test_sample.py backends/mlx/test/test_ops.py

cc @metascroy

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20564

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 12 Pending

As of commit 7e9ae35 with merge base 55a71e6 ():

NEW FAILURE - The following job has failed:

• pull / test-eval_llama-wikitext-linux / linux-job (gh)
  RuntimeError: Command docker exec -t f07b8ecee9be9cc829d92ad63b572fb3fb94908a1e2f16b39fadf2553b965213 /exec failed with exit code 1

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: goutamadwant / name: goutamadwant (c944a5a)

[goutamadwant]

@pytorchbot label "release notes: llm"

[metascroy]

Nit: Can we reorder the args here to be temp, top_k, top_p, seed?

We also need to modify this custom ops behavior to work with top_k correctly now.

Do huggingface style, where topk happens before topp, which requires renormalization, roughly something like this.

Something similar is required on the emit path.

def sample(logits, temperature, top_k=None, top_p=1.0, seed=None):
if float(temperature) <= 0:
return torch.argmax(logits, dim=-1)
scaled = logits.float() / temperature
# ── Top-k FIRST (on logits; monotonicity ⇒ top-k logits = top-k probs) ──
if top_k is not None:
k = int(top_k.item())
s_scaled, _ = torch.sort(scaled, dim=-1, descending=True)
kth = s_scaled[..., k - 1 : k]
scaled = torch.where(scaled >= kth, scaled, scaled.new_tensor(float("-inf")))
# ── Top-p on the *renormalized* distribution ──
probs = torch.softmax(scaled, dim=-1) # exp(-inf)=0 → renormalized over top-k
s_probs, _ = torch.sort(probs, dim=-1, descending=True)
cum = torch.cumsum(s_probs, dim=-1)
keep = (cum - s_probs) <= top_p
thresh = torch.where(keep, s_probs, s_probs.new_tensor(float("inf"))).amin(
dim=-1, keepdim=True
)
scaled = torch.where(probs >= thresh, scaled, scaled.new_tensor(float("-inf")))
# ── Gumbel-max ──
if seed is None:
u = torch.rand(scaled.shape)
else:
gen = torch.Generator().manual_seed(int(seed.item()))
u = torch.rand(scaled.shape, generator=gen)
gumbel = -torch.log(-torch.log(u))
return torch.argmax(scaled + gumbel, dim=-1)

[metascroy]

Can we have test case in non-sorted order, e.g., 0.15, 0.5, .05, 0.3?

[metascroy]

Let's keep top_k specified at runtime, with default value int max (select everything)

[metascroy]

Let's do Int clipping to vocab size if topk > vocab size.

[metascroy]

top_k_val is already intorvid?

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 3m 58s —— View job

---

Code review: Add top-k support to MLX sample

I read through all five changed files plus the emit path in ops.py and the helpers it relies on. Scott's inline comments already cover the main direction; below I confirm them with a concrete numeric case and add a few additional findings.

1. Top-k / top-p composition diverges from HuggingFace semantics (correctness) — confirms Scott's sampling.py:41 comment

Both the reference op (custom_ops.py) and the emit path (ops.py) compute the top-k mask and the top-p mask independently on the same un-renormalized probs and then OR the drop masks together. HF does top-k first, renormalizes, then applies top-p to the renormalized distribution. These are not equivalent. Worked example, probs = [0.5, 0.3, 0.15, 0.05], top_k=3, top_p=0.83:

• This PR: top-k keeps {0,1,2}, top-p (prefix mass [0, 0.5, 0.8, 0.95] ≤ 0.83) keeps {0,1,2} → final keep {0,1,2}.
• HF: top-k → renormalize over 0.95 → [0.526, 0.316, 0.158]; top-p on that (prefix [0, 0.526, 0.842] ≤ 0.83) keeps {0,1} → final keep {0,1}.

So the result genuinely differs once both filters are active. Scott's suggested sample() rewrite (top-k on logits → softmax/renormalize → top-p on the renormalized probs) is the fix, and the same restructuring is needed on the emit path — the current LogicalOrNode combine mirrors the reference's OR-compose, so both must move to sequential top-k-then-renormalize-then-top-p. The new test_top_k_and_top_p_compose test passes only because the chosen values happen to agree under both semantics; a case like the one above would catch the divergence.

2. top_k > vocab_size is an out-of-bounds bug (correctness) — relates to Scott's ops.py:3681 clipping comment

custom_ops.py:431:

k=int(top_k.item())
kth=s_probs[..., k-1 : k]

If k > vocab, the slice [k-1:k] is empty, so probs >= kth broadcasts against a size-0 tensor and the mask is wrong/raises. On the emit path the corresponding TakeNode(index = top_k - 1) indexes out of bounds. Clipping k = min(k, vocab) (Scott's point) fixes both. Worth a test with top_k larger than the vocab to lock this in.

3. top_k as runtime default instead of Optional — Scott's ops.py:3679 comment

Making top_k always-present with a default of INT_MAX (clipped to vocab per #2) removes the if top_k is not None branch in both custom_ops.py and the _sample_handler emit block, and the is not None checks in sampling.py. That collapses two code paths into one and naturally yields "keep everything" when unset.

4. Argument ordering — Scott's sampling.py:41 comment

Reorder to temperature, top_k, top_p, seed consistently across mlx::sample, sample_fake, _sample_handler's args[...] unpacking, and the _sample test helper. Note this is a signature change to the registered custom op, so every call site and the require_args(args, 3, 5, ...) indices need to move together.

5. to_int_or_vid already handles the literal/vid split — Scott's ops.py:3682 comment

ops.py:3688-3692:

index=(
IntOrVidOrTid.from_vid(top_k_index.vid)
iftop_k_index.is_videlseIntOrVidOrTid.from_literal(top_k_index.literal)
)

top_k_val comes from ItemIntNode so it's always a Vid here — top_k_iov/top_k_index will never be a literal in this path, making the else branch dead. P.to_int_or_vid_or_tid(top_k_val) (program_builder.py:318) already does this conversion in one call; prefer it over the manual ternary.

6. Test coverage — Scott's test_sample.py:165 comment

• Add a non-sorted-logits case (e.g. probs 0.15, 0.5, 0.05, 0.3) so the test exercises the sort path rather than already-descending input.
• Add a top_k > vocab case (Re-sync with internal repository #2) and a combined case where HF vs OR-compose semantics actually differ (Add support for quantized LeakyReLU #1), so the renormalization fix is verified rather than coincidentally passing.

Minor

• custom_ops.py:431int(top_k.item()) is fine for the eager/meta reference, but confirm the .item() data-dependent read doesn't trip export's symbolic tracing on this op (the seed path already uses ItemIntNode on device, so the device side is fine).
• Docstring in sampling.py still reads cleanly, but once Rename _pt2e to pt2e #3 lands ("None disables top-k") the wording should reflect the always-present default instead.

Overall the lowering scaffolding (Take → ExpandDims → Less → LogicalOr) is well-built and the node-count test is a nice touch; the substantive work is switching both the reference and emit paths to HF-style sequential filtering and adding the vocab-size clip.
· branch fix-mlx-sample-top-k

[goutamadwant]

thanks for your quick review @metascroy will take a look and address them.

[goutamadwant]

@metascroy Addressed below changes in the latest commit.

• Reordered the custom op args to temperature, top_k, top_p, seed.
• Switched sampling to top-k first.. and then top-p on the renormalized probabilities.
• Made the default top_k path use a max-int sentinel clipped to vocab size.
• Removed the independent OR-composed top-k/top-p masks in the emit path.
• Added coverage for non-sorted logits, top_k > vocab, and the combined top-k/top-p case where renormalization changes the result.
  let me know if you have any more suggestions. thanks!

[goutamadwant]

@claude review this code

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 2m 11s —— View job

---

Code review: top-k support for MLX sample (updated)

I re-read all five changed files plus the _sample_handler emit path. The latest commit cleanly addresses every point from the prior round. Summary below.

Prior review items — all resolved ✅

#	Item	Status
1	HF-style sequential filtering (top-k → renormalize → top-p)	✅ custom_ops.py:430-444 does top-k on logits, then softmax (renormalizes over the kept set), then top-p. Emit path mirrors this: top-k WhereNode feeds the SoftmaxNode at ops.py:3695, and the old LogicalOrNode OR-compose is gone (test asserts "LogicalOrNode": 0).
2	top_k > vocab clipping	✅ min(int(top_k.item()), scaled.shape[-1]) (custom_ops.py:430) and MinimumNode against vocab (ops.py:3622-3628).
3	top_k runtime arg with int-max default	✅ sampling.py:40-43 defaults to torch.iinfo(torch.int64).max; the op signature now always takes top_k.
4	Arg ordering temperature, top_k, top_p, seed	✅ Consistent across sample, sample_fake, _sample_handler unpacking, and the _sample test helper. require_args(args, 4, 5, ...) updated.
5	Use to_int_or_vid_or_tid	✅ ops.py:3660 now calls P.to_int_or_vid_or_tid(top_k_index) directly — dead literal branch removed.
6	Test coverage	✅ Added non-sorted input (test_top_k_restricts_to_top_k), top_k > vocab (test_top_k_clips_to_vocab_size), and a true HF-vs-OR divergence case (test_top_k_and_top_p_compose, top_k=3, top_p=0.83 → {0,1}). Plus an end-to-end on-device top-k test and a SampleTopKTest node-count case.

I worked the test_top_k_and_top_p_compose numbers by hand: top_k=3 keeps {0.5,0.3,0.15}, renormalizes to [0.526,0.316,0.158], and top-p=0.83 (prefix [0,0.526,0.842] ≤ 0.83) keeps {0,1}. This is the case that would have failed under the old OR-compose, so the renormalization fix is genuinely exercised now, not coincidentally passing.

Minor observations (non-blocking)

1. custom_ops.py:432 empty-slice edge when k=0.k is clipped above by vocab but not below. A top_k=0 would make s_scaled[..., -1:0] empty and broadcast oddly. This isn't reachable through SamplingHead (default is int-max, callers pass ≥1), so it's only a concern if the op is called directly with 0. A max(k, 1) clamp or a one-line docstring note that top_k ≥ 1 is expected would close it. Not worth blocking on.

2. Emit-path top-k threshold vs. reference — confirmed equivalent. The reference keeps scaled >= kth (>=), while the emit path drops via LessNode(scaled, top_k_thresh) i.e. keeps scaled >= thresh too — consistent. Good. The negate→sort→negate dance to get a descending sort (since SortNode is ascending-only) and the TakeNode(index = k-1) correctly pick the k-th largest. Matches the top-p chain's existing pattern.

3. int(top_k.item()) data-dependent read (custom_ops.py:430): fine for the eager/meta reference; the device path uses ItemIntNode so symbolic tracing is unaffected, same as the existing seed handling.

4. Buffer-reuse comments in ops.py are accurate (verified scaled/probs/sorted_p live-range notes against their actual reads).

Verdict

Looks good to me — the substantive correctness fix (sequential HF-style filtering + vocab clip) is in on both the reference and the lowered graph, and the tests lock in the behaviors that distinguish it from the old implementation. I couldn't execute the test suite in this sandbox (command approval), so please confirm the documented unittest + compileall runs are green in CI before merging. Only the k=0 clamp (#1) is worth a quick thought; everything else is optional polish.
· branch fix-mlx-sample-top-k

[metascroy]

Looks great! Thanks for the contribution @goutamadwant!

[goutamadwant]

Fixes#20548

Summary

• Thread optional top_k through SamplingHead and mlx::sample.
• Apply top-k filtering as an additional threshold mask that composes with the existing top-p nucleus mask.
• Add eager, export, end-to-end, and MLX lowering coverage for top-k sampling.

Test plan

• PYTHONPATH=src:. python3 -m unittest executorch.backends.mlx.test.test_sample.TestSampleOp executorch.backends.mlx.test.test_sample.TestSampleExport
• python3 -m compileall -q backends/mlx/llm/sampling.py backends/mlx/custom_ops.py backends/mlx/ops.py backends/mlx/test/test_sample.py backends/mlx/test/test_ops.py

cc @metascroy

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20564

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 12 Pending

As of commit 7e9ae35 with merge base 55a71e6 ():

NEW FAILURE - The following job has failed:

• pull / test-eval_llama-wikitext-linux / linux-job (gh)
  RuntimeError: Command docker exec -t f07b8ecee9be9cc829d92ad63b572fb3fb94908a1e2f16b39fadf2553b965213 /exec failed with exit code 1

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: goutamadwant / name: goutamadwant (c944a5a)

[goutamadwant]

@pytorchbot label "release notes: llm"

[metascroy]

Nit: Can we reorder the args here to be temp, top_k, top_p, seed?

We also need to modify this custom ops behavior to work with top_k correctly now.

Do huggingface style, where topk happens before topp, which requires renormalization, roughly something like this.

Something similar is required on the emit path.

def sample(logits, temperature, top_k=None, top_p=1.0, seed=None):
if float(temperature) <= 0:
return torch.argmax(logits, dim=-1)
scaled = logits.float() / temperature
# ── Top-k FIRST (on logits; monotonicity ⇒ top-k logits = top-k probs) ──
if top_k is not None:
k = int(top_k.item())
s_scaled, _ = torch.sort(scaled, dim=-1, descending=True)
kth = s_scaled[..., k - 1 : k]
scaled = torch.where(scaled >= kth, scaled, scaled.new_tensor(float("-inf")))
# ── Top-p on the *renormalized* distribution ──
probs = torch.softmax(scaled, dim=-1) # exp(-inf)=0 → renormalized over top-k
s_probs, _ = torch.sort(probs, dim=-1, descending=True)
cum = torch.cumsum(s_probs, dim=-1)
keep = (cum - s_probs) <= top_p
thresh = torch.where(keep, s_probs, s_probs.new_tensor(float("inf"))).amin(
dim=-1, keepdim=True
)
scaled = torch.where(probs >= thresh, scaled, scaled.new_tensor(float("-inf")))
# ── Gumbel-max ──
if seed is None:
u = torch.rand(scaled.shape)
else:
gen = torch.Generator().manual_seed(int(seed.item()))
u = torch.rand(scaled.shape, generator=gen)
gumbel = -torch.log(-torch.log(u))
return torch.argmax(scaled + gumbel, dim=-1)

[metascroy]

Can we have test case in non-sorted order, e.g., 0.15, 0.5, .05, 0.3?

[metascroy]

Let's keep top_k specified at runtime, with default value int max (select everything)

[metascroy]

Let's do Int clipping to vocab size if topk > vocab size.

[metascroy]

top_k_val is already intorvid?

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 3m 58s —— View job

---

Code review: Add top-k support to MLX sample

I read through all five changed files plus the emit path in ops.py and the helpers it relies on. Scott's inline comments already cover the main direction; below I confirm them with a concrete numeric case and add a few additional findings.

1. Top-k / top-p composition diverges from HuggingFace semantics (correctness) — confirms Scott's sampling.py:41 comment

Both the reference op (custom_ops.py) and the emit path (ops.py) compute the top-k mask and the top-p mask independently on the same un-renormalized probs and then OR the drop masks together. HF does top-k first, renormalizes, then applies top-p to the renormalized distribution. These are not equivalent. Worked example, probs = [0.5, 0.3, 0.15, 0.05], top_k=3, top_p=0.83:

• This PR: top-k keeps {0,1,2}, top-p (prefix mass [0, 0.5, 0.8, 0.95] ≤ 0.83) keeps {0,1,2} → final keep {0,1,2}.
• HF: top-k → renormalize over 0.95 → [0.526, 0.316, 0.158]; top-p on that (prefix [0, 0.526, 0.842] ≤ 0.83) keeps {0,1} → final keep {0,1}.

So the result genuinely differs once both filters are active. Scott's suggested sample() rewrite (top-k on logits → softmax/renormalize → top-p on the renormalized probs) is the fix, and the same restructuring is needed on the emit path — the current LogicalOrNode combine mirrors the reference's OR-compose, so both must move to sequential top-k-then-renormalize-then-top-p. The new test_top_k_and_top_p_compose test passes only because the chosen values happen to agree under both semantics; a case like the one above would catch the divergence.

2. top_k > vocab_size is an out-of-bounds bug (correctness) — relates to Scott's ops.py:3681 clipping comment

custom_ops.py:431:

k=int(top_k.item())
kth=s_probs[..., k-1 : k]

If k > vocab, the slice [k-1:k] is empty, so probs >= kth broadcasts against a size-0 tensor and the mask is wrong/raises. On the emit path the corresponding TakeNode(index = top_k - 1) indexes out of bounds. Clipping k = min(k, vocab) (Scott's point) fixes both. Worth a test with top_k larger than the vocab to lock this in.

3. top_k as runtime default instead of Optional — Scott's ops.py:3679 comment

Making top_k always-present with a default of INT_MAX (clipped to vocab per #2) removes the if top_k is not None branch in both custom_ops.py and the _sample_handler emit block, and the is not None checks in sampling.py. That collapses two code paths into one and naturally yields "keep everything" when unset.

4. Argument ordering — Scott's sampling.py:41 comment

Reorder to temperature, top_k, top_p, seed consistently across mlx::sample, sample_fake, _sample_handler's args[...] unpacking, and the _sample test helper. Note this is a signature change to the registered custom op, so every call site and the require_args(args, 3, 5, ...) indices need to move together.

5. to_int_or_vid already handles the literal/vid split — Scott's ops.py:3682 comment

ops.py:3688-3692:

index=(
IntOrVidOrTid.from_vid(top_k_index.vid)
iftop_k_index.is_videlseIntOrVidOrTid.from_literal(top_k_index.literal)
)

top_k_val comes from ItemIntNode so it's always a Vid here — top_k_iov/top_k_index will never be a literal in this path, making the else branch dead. P.to_int_or_vid_or_tid(top_k_val) (program_builder.py:318) already does this conversion in one call; prefer it over the manual ternary.

6. Test coverage — Scott's test_sample.py:165 comment

• Add a non-sorted-logits case (e.g. probs 0.15, 0.5, 0.05, 0.3) so the test exercises the sort path rather than already-descending input.
• Add a top_k > vocab case (Re-sync with internal repository #2) and a combined case where HF vs OR-compose semantics actually differ (Add support for quantized LeakyReLU #1), so the renormalization fix is verified rather than coincidentally passing.

Minor

• custom_ops.py:431int(top_k.item()) is fine for the eager/meta reference, but confirm the .item() data-dependent read doesn't trip export's symbolic tracing on this op (the seed path already uses ItemIntNode on device, so the device side is fine).
• Docstring in sampling.py still reads cleanly, but once Rename _pt2e to pt2e #3 lands ("None disables top-k") the wording should reflect the always-present default instead.

Overall the lowering scaffolding (Take → ExpandDims → Less → LogicalOr) is well-built and the node-count test is a nice touch; the substantive work is switching both the reference and emit paths to HF-style sequential filtering and adding the vocab-size clip.
· branch fix-mlx-sample-top-k

[goutamadwant]

thanks for your quick review @metascroy will take a look and address them.

[goutamadwant]

@metascroy Addressed below changes in the latest commit.

• Reordered the custom op args to temperature, top_k, top_p, seed.
• Switched sampling to top-k first.. and then top-p on the renormalized probabilities.
• Made the default top_k path use a max-int sentinel clipped to vocab size.
• Removed the independent OR-composed top-k/top-p masks in the emit path.
• Added coverage for non-sorted logits, top_k > vocab, and the combined top-k/top-p case where renormalization changes the result.
  let me know if you have any more suggestions. thanks!

[goutamadwant]

@claude review this code

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 2m 11s —— View job

---

Code review: top-k support for MLX sample (updated)

I re-read all five changed files plus the _sample_handler emit path. The latest commit cleanly addresses every point from the prior round. Summary below.

Prior review items — all resolved ✅

#	Item	Status
1	HF-style sequential filtering (top-k → renormalize → top-p)	✅ custom_ops.py:430-444 does top-k on logits, then softmax (renormalizes over the kept set), then top-p. Emit path mirrors this: top-k WhereNode feeds the SoftmaxNode at ops.py:3695, and the old LogicalOrNode OR-compose is gone (test asserts "LogicalOrNode": 0).
2	top_k > vocab clipping	✅ min(int(top_k.item()), scaled.shape[-1]) (custom_ops.py:430) and MinimumNode against vocab (ops.py:3622-3628).
3	top_k runtime arg with int-max default	✅ sampling.py:40-43 defaults to torch.iinfo(torch.int64).max; the op signature now always takes top_k.
4	Arg ordering temperature, top_k, top_p, seed	✅ Consistent across sample, sample_fake, _sample_handler unpacking, and the _sample test helper. require_args(args, 4, 5, ...) updated.
5	Use to_int_or_vid_or_tid	✅ ops.py:3660 now calls P.to_int_or_vid_or_tid(top_k_index) directly — dead literal branch removed.
6	Test coverage	✅ Added non-sorted input (test_top_k_restricts_to_top_k), top_k > vocab (test_top_k_clips_to_vocab_size), and a true HF-vs-OR divergence case (test_top_k_and_top_p_compose, top_k=3, top_p=0.83 → {0,1}). Plus an end-to-end on-device top-k test and a SampleTopKTest node-count case.

I worked the test_top_k_and_top_p_compose numbers by hand: top_k=3 keeps {0.5,0.3,0.15}, renormalizes to [0.526,0.316,0.158], and top-p=0.83 (prefix [0,0.526,0.842] ≤ 0.83) keeps {0,1}. This is the case that would have failed under the old OR-compose, so the renormalization fix is genuinely exercised now, not coincidentally passing.

Minor observations (non-blocking)

1. custom_ops.py:432 empty-slice edge when k=0.k is clipped above by vocab but not below. A top_k=0 would make s_scaled[..., -1:0] empty and broadcast oddly. This isn't reachable through SamplingHead (default is int-max, callers pass ≥1), so it's only a concern if the op is called directly with 0. A max(k, 1) clamp or a one-line docstring note that top_k ≥ 1 is expected would close it. Not worth blocking on.

2. Emit-path top-k threshold vs. reference — confirmed equivalent. The reference keeps scaled >= kth (>=), while the emit path drops via LessNode(scaled, top_k_thresh) i.e. keeps scaled >= thresh too — consistent. Good. The negate→sort→negate dance to get a descending sort (since SortNode is ascending-only) and the TakeNode(index = k-1) correctly pick the k-th largest. Matches the top-p chain's existing pattern.

3. int(top_k.item()) data-dependent read (custom_ops.py:430): fine for the eager/meta reference; the device path uses ItemIntNode so symbolic tracing is unaffected, same as the existing seed handling.

4. Buffer-reuse comments in ops.py are accurate (verified scaled/probs/sorted_p live-range notes against their actual reads).

Verdict

Looks good to me — the substantive correctness fix (sequential HF-style filtering + vocab clip) is in on both the reference and the lowered graph, and the tests lock in the behaviors that distinguish it from the old implementation. I couldn't execute the test suite in this sandbox (command approval), so please confirm the documented unittest + compileall runs are green in CI before merging. Only the k=0 clamp (#1) is worth a quick thought; everything else is optional polish.
· branch fix-mlx-sample-top-k

[metascroy]

Looks great! Thanks for the contribution @goutamadwant!

[goutamadwant]

Fixes#20548

Summary

• Thread optional top_k through SamplingHead and mlx::sample.
• Apply top-k filtering as an additional threshold mask that composes with the existing top-p nucleus mask.
• Add eager, export, end-to-end, and MLX lowering coverage for top-k sampling.

Test plan

• PYTHONPATH=src:. python3 -m unittest executorch.backends.mlx.test.test_sample.TestSampleOp executorch.backends.mlx.test.test_sample.TestSampleExport
• python3 -m compileall -q backends/mlx/llm/sampling.py backends/mlx/custom_ops.py backends/mlx/ops.py backends/mlx/test/test_sample.py backends/mlx/test/test_ops.py

cc @metascroy

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20564

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 12 Pending

As of commit 7e9ae35 with merge base 55a71e6 ():

NEW FAILURE - The following job has failed:

• pull / test-eval_llama-wikitext-linux / linux-job (gh)
  RuntimeError: Command docker exec -t f07b8ecee9be9cc829d92ad63b572fb3fb94908a1e2f16b39fadf2553b965213 /exec failed with exit code 1

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: goutamadwant / name: goutamadwant (c944a5a)

[goutamadwant]

@pytorchbot label "release notes: llm"

[metascroy]

Nit: Can we reorder the args here to be temp, top_k, top_p, seed?

We also need to modify this custom ops behavior to work with top_k correctly now.

Do huggingface style, where topk happens before topp, which requires renormalization, roughly something like this.

Something similar is required on the emit path.

def sample(logits, temperature, top_k=None, top_p=1.0, seed=None):
if float(temperature) <= 0:
return torch.argmax(logits, dim=-1)
scaled = logits.float() / temperature
# ── Top-k FIRST (on logits; monotonicity ⇒ top-k logits = top-k probs) ──
if top_k is not None:
k = int(top_k.item())
s_scaled, _ = torch.sort(scaled, dim=-1, descending=True)
kth = s_scaled[..., k - 1 : k]
scaled = torch.where(scaled >= kth, scaled, scaled.new_tensor(float("-inf")))
# ── Top-p on the *renormalized* distribution ──
probs = torch.softmax(scaled, dim=-1) # exp(-inf)=0 → renormalized over top-k
s_probs, _ = torch.sort(probs, dim=-1, descending=True)
cum = torch.cumsum(s_probs, dim=-1)
keep = (cum - s_probs) <= top_p
thresh = torch.where(keep, s_probs, s_probs.new_tensor(float("inf"))).amin(
dim=-1, keepdim=True
)
scaled = torch.where(probs >= thresh, scaled, scaled.new_tensor(float("-inf")))
# ── Gumbel-max ──
if seed is None:
u = torch.rand(scaled.shape)
else:
gen = torch.Generator().manual_seed(int(seed.item()))
u = torch.rand(scaled.shape, generator=gen)
gumbel = -torch.log(-torch.log(u))
return torch.argmax(scaled + gumbel, dim=-1)

[metascroy]

Can we have test case in non-sorted order, e.g., 0.15, 0.5, .05, 0.3?

[metascroy]

Let's keep top_k specified at runtime, with default value int max (select everything)

[metascroy]

Let's do Int clipping to vocab size if topk > vocab size.

[metascroy]

top_k_val is already intorvid?

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 3m 58s —— View job

---

Code review: Add top-k support to MLX sample

I read through all five changed files plus the emit path in ops.py and the helpers it relies on. Scott's inline comments already cover the main direction; below I confirm them with a concrete numeric case and add a few additional findings.

1. Top-k / top-p composition diverges from HuggingFace semantics (correctness) — confirms Scott's sampling.py:41 comment

Both the reference op (custom_ops.py) and the emit path (ops.py) compute the top-k mask and the top-p mask independently on the same un-renormalized probs and then OR the drop masks together. HF does top-k first, renormalizes, then applies top-p to the renormalized distribution. These are not equivalent. Worked example, probs = [0.5, 0.3, 0.15, 0.05], top_k=3, top_p=0.83:

• This PR: top-k keeps {0,1,2}, top-p (prefix mass [0, 0.5, 0.8, 0.95] ≤ 0.83) keeps {0,1,2} → final keep {0,1,2}.
• HF: top-k → renormalize over 0.95 → [0.526, 0.316, 0.158]; top-p on that (prefix [0, 0.526, 0.842] ≤ 0.83) keeps {0,1} → final keep {0,1}.

So the result genuinely differs once both filters are active. Scott's suggested sample() rewrite (top-k on logits → softmax/renormalize → top-p on the renormalized probs) is the fix, and the same restructuring is needed on the emit path — the current LogicalOrNode combine mirrors the reference's OR-compose, so both must move to sequential top-k-then-renormalize-then-top-p. The new test_top_k_and_top_p_compose test passes only because the chosen values happen to agree under both semantics; a case like the one above would catch the divergence.

2. top_k > vocab_size is an out-of-bounds bug (correctness) — relates to Scott's ops.py:3681 clipping comment

custom_ops.py:431:

k=int(top_k.item())
kth=s_probs[..., k-1 : k]

If k > vocab, the slice [k-1:k] is empty, so probs >= kth broadcasts against a size-0 tensor and the mask is wrong/raises. On the emit path the corresponding TakeNode(index = top_k - 1) indexes out of bounds. Clipping k = min(k, vocab) (Scott's point) fixes both. Worth a test with top_k larger than the vocab to lock this in.

3. top_k as runtime default instead of Optional — Scott's ops.py:3679 comment

Making top_k always-present with a default of INT_MAX (clipped to vocab per #2) removes the if top_k is not None branch in both custom_ops.py and the _sample_handler emit block, and the is not None checks in sampling.py. That collapses two code paths into one and naturally yields "keep everything" when unset.

4. Argument ordering — Scott's sampling.py:41 comment

Reorder to temperature, top_k, top_p, seed consistently across mlx::sample, sample_fake, _sample_handler's args[...] unpacking, and the _sample test helper. Note this is a signature change to the registered custom op, so every call site and the require_args(args, 3, 5, ...) indices need to move together.

5. to_int_or_vid already handles the literal/vid split — Scott's ops.py:3682 comment

ops.py:3688-3692:

index=(
IntOrVidOrTid.from_vid(top_k_index.vid)
iftop_k_index.is_videlseIntOrVidOrTid.from_literal(top_k_index.literal)
)

top_k_val comes from ItemIntNode so it's always a Vid here — top_k_iov/top_k_index will never be a literal in this path, making the else branch dead. P.to_int_or_vid_or_tid(top_k_val) (program_builder.py:318) already does this conversion in one call; prefer it over the manual ternary.

6. Test coverage — Scott's test_sample.py:165 comment

• Add a non-sorted-logits case (e.g. probs 0.15, 0.5, 0.05, 0.3) so the test exercises the sort path rather than already-descending input.
• Add a top_k > vocab case (Re-sync with internal repository #2) and a combined case where HF vs OR-compose semantics actually differ (Add support for quantized LeakyReLU #1), so the renormalization fix is verified rather than coincidentally passing.

Minor

• custom_ops.py:431int(top_k.item()) is fine for the eager/meta reference, but confirm the .item() data-dependent read doesn't trip export's symbolic tracing on this op (the seed path already uses ItemIntNode on device, so the device side is fine).
• Docstring in sampling.py still reads cleanly, but once Rename _pt2e to pt2e #3 lands ("None disables top-k") the wording should reflect the always-present default instead.

Overall the lowering scaffolding (Take → ExpandDims → Less → LogicalOr) is well-built and the node-count test is a nice touch; the substantive work is switching both the reference and emit paths to HF-style sequential filtering and adding the vocab-size clip.
· branch fix-mlx-sample-top-k

[goutamadwant]

thanks for your quick review @metascroy will take a look and address them.

[goutamadwant]

@metascroy Addressed below changes in the latest commit.

• Reordered the custom op args to temperature, top_k, top_p, seed.
• Switched sampling to top-k first.. and then top-p on the renormalized probabilities.
• Made the default top_k path use a max-int sentinel clipped to vocab size.
• Removed the independent OR-composed top-k/top-p masks in the emit path.
• Added coverage for non-sorted logits, top_k > vocab, and the combined top-k/top-p case where renormalization changes the result.
  let me know if you have any more suggestions. thanks!

[goutamadwant]

@claude review this code

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 2m 11s —— View job

---

Code review: top-k support for MLX sample (updated)

I re-read all five changed files plus the _sample_handler emit path. The latest commit cleanly addresses every point from the prior round. Summary below.

Prior review items — all resolved ✅

#	Item	Status
1	HF-style sequential filtering (top-k → renormalize → top-p)	✅ custom_ops.py:430-444 does top-k on logits, then softmax (renormalizes over the kept set), then top-p. Emit path mirrors this: top-k WhereNode feeds the SoftmaxNode at ops.py:3695, and the old LogicalOrNode OR-compose is gone (test asserts "LogicalOrNode": 0).
2	top_k > vocab clipping	✅ min(int(top_k.item()), scaled.shape[-1]) (custom_ops.py:430) and MinimumNode against vocab (ops.py:3622-3628).
3	top_k runtime arg with int-max default	✅ sampling.py:40-43 defaults to torch.iinfo(torch.int64).max; the op signature now always takes top_k.
4	Arg ordering temperature, top_k, top_p, seed	✅ Consistent across sample, sample_fake, _sample_handler unpacking, and the _sample test helper. require_args(args, 4, 5, ...) updated.
5	Use to_int_or_vid_or_tid	✅ ops.py:3660 now calls P.to_int_or_vid_or_tid(top_k_index) directly — dead literal branch removed.
6	Test coverage	✅ Added non-sorted input (test_top_k_restricts_to_top_k), top_k > vocab (test_top_k_clips_to_vocab_size), and a true HF-vs-OR divergence case (test_top_k_and_top_p_compose, top_k=3, top_p=0.83 → {0,1}). Plus an end-to-end on-device top-k test and a SampleTopKTest node-count case.

I worked the test_top_k_and_top_p_compose numbers by hand: top_k=3 keeps {0.5,0.3,0.15}, renormalizes to [0.526,0.316,0.158], and top-p=0.83 (prefix [0,0.526,0.842] ≤ 0.83) keeps {0,1}. This is the case that would have failed under the old OR-compose, so the renormalization fix is genuinely exercised now, not coincidentally passing.

Minor observations (non-blocking)

1. custom_ops.py:432 empty-slice edge when k=0.k is clipped above by vocab but not below. A top_k=0 would make s_scaled[..., -1:0] empty and broadcast oddly. This isn't reachable through SamplingHead (default is int-max, callers pass ≥1), so it's only a concern if the op is called directly with 0. A max(k, 1) clamp or a one-line docstring note that top_k ≥ 1 is expected would close it. Not worth blocking on.

2. Emit-path top-k threshold vs. reference — confirmed equivalent. The reference keeps scaled >= kth (>=), while the emit path drops via LessNode(scaled, top_k_thresh) i.e. keeps scaled >= thresh too — consistent. Good. The negate→sort→negate dance to get a descending sort (since SortNode is ascending-only) and the TakeNode(index = k-1) correctly pick the k-th largest. Matches the top-p chain's existing pattern.

3. int(top_k.item()) data-dependent read (custom_ops.py:430): fine for the eager/meta reference; the device path uses ItemIntNode so symbolic tracing is unaffected, same as the existing seed handling.

4. Buffer-reuse comments in ops.py are accurate (verified scaled/probs/sorted_p live-range notes against their actual reads).

Verdict

Looks good to me — the substantive correctness fix (sequential HF-style filtering + vocab clip) is in on both the reference and the lowered graph, and the tests lock in the behaviors that distinguish it from the old implementation. I couldn't execute the test suite in this sandbox (command approval), so please confirm the documented unittest + compileall runs are green in CI before merging. Only the k=0 clamp (#1) is worth a quick thought; everything else is optional polish.
· branch fix-mlx-sample-top-k

[metascroy]

Looks great! Thanks for the contribution @goutamadwant!

[goutamadwant]

Fixes#20548

Summary

• Thread optional top_k through SamplingHead and mlx::sample.
• Apply top-k filtering as an additional threshold mask that composes with the existing top-p nucleus mask.
• Add eager, export, end-to-end, and MLX lowering coverage for top-k sampling.

Test plan

• PYTHONPATH=src:. python3 -m unittest executorch.backends.mlx.test.test_sample.TestSampleOp executorch.backends.mlx.test.test_sample.TestSampleExport
• python3 -m compileall -q backends/mlx/llm/sampling.py backends/mlx/custom_ops.py backends/mlx/ops.py backends/mlx/test/test_sample.py backends/mlx/test/test_ops.py

cc @metascroy

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20564

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 12 Pending

As of commit 7e9ae35 with merge base 55a71e6 ():

NEW FAILURE - The following job has failed:

• pull / test-eval_llama-wikitext-linux / linux-job (gh)
  RuntimeError: Command docker exec -t f07b8ecee9be9cc829d92ad63b572fb3fb94908a1e2f16b39fadf2553b965213 /exec failed with exit code 1

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: goutamadwant / name: goutamadwant (c944a5a)

[goutamadwant]

@pytorchbot label "release notes: llm"

[metascroy]

Nit: Can we reorder the args here to be temp, top_k, top_p, seed?

We also need to modify this custom ops behavior to work with top_k correctly now.

Do huggingface style, where topk happens before topp, which requires renormalization, roughly something like this.

Something similar is required on the emit path.

def sample(logits, temperature, top_k=None, top_p=1.0, seed=None):
if float(temperature) <= 0:
return torch.argmax(logits, dim=-1)
scaled = logits.float() / temperature
# ── Top-k FIRST (on logits; monotonicity ⇒ top-k logits = top-k probs) ──
if top_k is not None:
k = int(top_k.item())
s_scaled, _ = torch.sort(scaled, dim=-1, descending=True)
kth = s_scaled[..., k - 1 : k]
scaled = torch.where(scaled >= kth, scaled, scaled.new_tensor(float("-inf")))
# ── Top-p on the *renormalized* distribution ──
probs = torch.softmax(scaled, dim=-1) # exp(-inf)=0 → renormalized over top-k
s_probs, _ = torch.sort(probs, dim=-1, descending=True)
cum = torch.cumsum(s_probs, dim=-1)
keep = (cum - s_probs) <= top_p
thresh = torch.where(keep, s_probs, s_probs.new_tensor(float("inf"))).amin(
dim=-1, keepdim=True
)
scaled = torch.where(probs >= thresh, scaled, scaled.new_tensor(float("-inf")))
# ── Gumbel-max ──
if seed is None:
u = torch.rand(scaled.shape)
else:
gen = torch.Generator().manual_seed(int(seed.item()))
u = torch.rand(scaled.shape, generator=gen)
gumbel = -torch.log(-torch.log(u))
return torch.argmax(scaled + gumbel, dim=-1)

[metascroy]

Can we have test case in non-sorted order, e.g., 0.15, 0.5, .05, 0.3?

[metascroy]

Let's keep top_k specified at runtime, with default value int max (select everything)

[metascroy]

Let's do Int clipping to vocab size if topk > vocab size.

[metascroy]

top_k_val is already intorvid?

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 3m 58s —— View job

---

Code review: Add top-k support to MLX sample

I read through all five changed files plus the emit path in ops.py and the helpers it relies on. Scott's inline comments already cover the main direction; below I confirm them with a concrete numeric case and add a few additional findings.

1. Top-k / top-p composition diverges from HuggingFace semantics (correctness) — confirms Scott's sampling.py:41 comment

Both the reference op (custom_ops.py) and the emit path (ops.py) compute the top-k mask and the top-p mask independently on the same un-renormalized probs and then OR the drop masks together. HF does top-k first, renormalizes, then applies top-p to the renormalized distribution. These are not equivalent. Worked example, probs = [0.5, 0.3, 0.15, 0.05], top_k=3, top_p=0.83:

• This PR: top-k keeps {0,1,2}, top-p (prefix mass [0, 0.5, 0.8, 0.95] ≤ 0.83) keeps {0,1,2} → final keep {0,1,2}.
• HF: top-k → renormalize over 0.95 → [0.526, 0.316, 0.158]; top-p on that (prefix [0, 0.526, 0.842] ≤ 0.83) keeps {0,1} → final keep {0,1}.

So the result genuinely differs once both filters are active. Scott's suggested sample() rewrite (top-k on logits → softmax/renormalize → top-p on the renormalized probs) is the fix, and the same restructuring is needed on the emit path — the current LogicalOrNode combine mirrors the reference's OR-compose, so both must move to sequential top-k-then-renormalize-then-top-p. The new test_top_k_and_top_p_compose test passes only because the chosen values happen to agree under both semantics; a case like the one above would catch the divergence.

2. top_k > vocab_size is an out-of-bounds bug (correctness) — relates to Scott's ops.py:3681 clipping comment

custom_ops.py:431:

k=int(top_k.item())
kth=s_probs[..., k-1 : k]

If k > vocab, the slice [k-1:k] is empty, so probs >= kth broadcasts against a size-0 tensor and the mask is wrong/raises. On the emit path the corresponding TakeNode(index = top_k - 1) indexes out of bounds. Clipping k = min(k, vocab) (Scott's point) fixes both. Worth a test with top_k larger than the vocab to lock this in.

3. top_k as runtime default instead of Optional — Scott's ops.py:3679 comment

Making top_k always-present with a default of INT_MAX (clipped to vocab per #2) removes the if top_k is not None branch in both custom_ops.py and the _sample_handler emit block, and the is not None checks in sampling.py. That collapses two code paths into one and naturally yields "keep everything" when unset.

4. Argument ordering — Scott's sampling.py:41 comment

Reorder to temperature, top_k, top_p, seed consistently across mlx::sample, sample_fake, _sample_handler's args[...] unpacking, and the _sample test helper. Note this is a signature change to the registered custom op, so every call site and the require_args(args, 3, 5, ...) indices need to move together.

5. to_int_or_vid already handles the literal/vid split — Scott's ops.py:3682 comment

ops.py:3688-3692:

index=(
IntOrVidOrTid.from_vid(top_k_index.vid)
iftop_k_index.is_videlseIntOrVidOrTid.from_literal(top_k_index.literal)
)

top_k_val comes from ItemIntNode so it's always a Vid here — top_k_iov/top_k_index will never be a literal in this path, making the else branch dead. P.to_int_or_vid_or_tid(top_k_val) (program_builder.py:318) already does this conversion in one call; prefer it over the manual ternary.

6. Test coverage — Scott's test_sample.py:165 comment

• Add a non-sorted-logits case (e.g. probs 0.15, 0.5, 0.05, 0.3) so the test exercises the sort path rather than already-descending input.
• Add a top_k > vocab case (Re-sync with internal repository #2) and a combined case where HF vs OR-compose semantics actually differ (Add support for quantized LeakyReLU #1), so the renormalization fix is verified rather than coincidentally passing.

Minor

• custom_ops.py:431int(top_k.item()) is fine for the eager/meta reference, but confirm the .item() data-dependent read doesn't trip export's symbolic tracing on this op (the seed path already uses ItemIntNode on device, so the device side is fine).
• Docstring in sampling.py still reads cleanly, but once Rename _pt2e to pt2e #3 lands ("None disables top-k") the wording should reflect the always-present default instead.

Overall the lowering scaffolding (Take → ExpandDims → Less → LogicalOr) is well-built and the node-count test is a nice touch; the substantive work is switching both the reference and emit paths to HF-style sequential filtering and adding the vocab-size clip.
· branch fix-mlx-sample-top-k

[goutamadwant]

thanks for your quick review @metascroy will take a look and address them.

[goutamadwant]

@metascroy Addressed below changes in the latest commit.

• Reordered the custom op args to temperature, top_k, top_p, seed.
• Switched sampling to top-k first.. and then top-p on the renormalized probabilities.
• Made the default top_k path use a max-int sentinel clipped to vocab size.
• Removed the independent OR-composed top-k/top-p masks in the emit path.
• Added coverage for non-sorted logits, top_k > vocab, and the combined top-k/top-p case where renormalization changes the result.
  let me know if you have any more suggestions. thanks!

[goutamadwant]

@claude review this code

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 2m 11s —— View job

---

Code review: top-k support for MLX sample (updated)

I re-read all five changed files plus the _sample_handler emit path. The latest commit cleanly addresses every point from the prior round. Summary below.

Prior review items — all resolved ✅

#	Item	Status
1	HF-style sequential filtering (top-k → renormalize → top-p)	✅ custom_ops.py:430-444 does top-k on logits, then softmax (renormalizes over the kept set), then top-p. Emit path mirrors this: top-k WhereNode feeds the SoftmaxNode at ops.py:3695, and the old LogicalOrNode OR-compose is gone (test asserts "LogicalOrNode": 0).
2	top_k > vocab clipping	✅ min(int(top_k.item()), scaled.shape[-1]) (custom_ops.py:430) and MinimumNode against vocab (ops.py:3622-3628).
3	top_k runtime arg with int-max default	✅ sampling.py:40-43 defaults to torch.iinfo(torch.int64).max; the op signature now always takes top_k.
4	Arg ordering temperature, top_k, top_p, seed	✅ Consistent across sample, sample_fake, _sample_handler unpacking, and the _sample test helper. require_args(args, 4, 5, ...) updated.
5	Use to_int_or_vid_or_tid	✅ ops.py:3660 now calls P.to_int_or_vid_or_tid(top_k_index) directly — dead literal branch removed.
6	Test coverage	✅ Added non-sorted input (test_top_k_restricts_to_top_k), top_k > vocab (test_top_k_clips_to_vocab_size), and a true HF-vs-OR divergence case (test_top_k_and_top_p_compose, top_k=3, top_p=0.83 → {0,1}). Plus an end-to-end on-device top-k test and a SampleTopKTest node-count case.

I worked the test_top_k_and_top_p_compose numbers by hand: top_k=3 keeps {0.5,0.3,0.15}, renormalizes to [0.526,0.316,0.158], and top-p=0.83 (prefix [0,0.526,0.842] ≤ 0.83) keeps {0,1}. This is the case that would have failed under the old OR-compose, so the renormalization fix is genuinely exercised now, not coincidentally passing.

Minor observations (non-blocking)

1. custom_ops.py:432 empty-slice edge when k=0.k is clipped above by vocab but not below. A top_k=0 would make s_scaled[..., -1:0] empty and broadcast oddly. This isn't reachable through SamplingHead (default is int-max, callers pass ≥1), so it's only a concern if the op is called directly with 0. A max(k, 1) clamp or a one-line docstring note that top_k ≥ 1 is expected would close it. Not worth blocking on.

2. Emit-path top-k threshold vs. reference — confirmed equivalent. The reference keeps scaled >= kth (>=), while the emit path drops via LessNode(scaled, top_k_thresh) i.e. keeps scaled >= thresh too — consistent. Good. The negate→sort→negate dance to get a descending sort (since SortNode is ascending-only) and the TakeNode(index = k-1) correctly pick the k-th largest. Matches the top-p chain's existing pattern.

3. int(top_k.item()) data-dependent read (custom_ops.py:430): fine for the eager/meta reference; the device path uses ItemIntNode so symbolic tracing is unaffected, same as the existing seed handling.

4. Buffer-reuse comments in ops.py are accurate (verified scaled/probs/sorted_p live-range notes against their actual reads).

Verdict

Looks good to me — the substantive correctness fix (sequential HF-style filtering + vocab clip) is in on both the reference and the lowered graph, and the tests lock in the behaviors that distinguish it from the old implementation. I couldn't execute the test suite in this sandbox (command approval), so please confirm the documented unittest + compileall runs are green in CI before merging. Only the k=0 clamp (#1) is worth a quick thought; everything else is optional polish.
· branch fix-mlx-sample-top-k

[metascroy]

Looks great! Thanks for the contribution @goutamadwant!