Add top-k support to MLX sample

Author: goutamadwant

backends/mlx/custom_ops.py

@@ -399,9 +399,11 @@ def sample(
     temperature: Tensor,
     top_p: Tensor,
     seed: Optional[Tensor] = None,
+    top_k: Optional[Tensor] = None,
 ) -> Tensor:
     """
-    Gumbel-max sampling from softmax(logits / temperature), with top-p (nucleus).
+    Gumbel-max sampling from softmax(logits / temperature), with top-p (nucleus)
+    and optional top-k filtering.
     logits:      [B, vocab]
     temperature: scalar float tensor    (runtime input). temperature <= 0 is
                  greedy: return argmax(logits) directly (matches the device,
@@ -411,6 +413,8 @@ def sample(
     seed:        scalar int tensor or None
                  - tensor -> deterministic, keyed RNG (random::key(seed))
                  - None   -> MLX global KeySequence (non-deterministic)
+    top_k:       optional scalar int tensor. None keeps every token; otherwise
+                 sampling is restricted to the k most likely tokens.
     -> token_id: [B] int64
 
     Host/CPU reference used for export (shape/meta) and distributional checks
@@ -424,6 +428,10 @@ def sample(
     scaled = logits.float() / temperature
     probs = torch.softmax(scaled, dim=-1)
     s_probs, _ = torch.sort(probs, dim=-1, descending=True)
+    if top_k is not None:
+        k = int(top_k.item())
+        kth = s_probs[..., k - 1 : k]
+        scaled = torch.where(probs >= kth, scaled, scaled.new_tensor(float("-inf")))
     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(
@@ -440,5 +448,5 @@ def sample(
 
 
 @torch.library.register_fake("mlx::sample")
-def sample_fake(logits, temperature, top_p, seed=None):
+def sample_fake(logits, temperature, top_p, seed=None, top_k=None):
     return logits.new_empty(logits.shape[:-1], dtype=torch.long)

backends/mlx/llm/sampling.py

@@ -19,7 +19,8 @@ class SamplingHead(nn.Module):
 
       temperature: scalar float tensor, e.g. torch.tensor(0.8). Must be >= 0;
                    temperature=0 is greedy (returns argmax, no division).
-      top_k:       not implemented yet (reserved); must be None.
+      top_k:       scalar int tensor or int; keeps only the k most likely tokens.
+                   None disables top-k filtering.
       top_p:       scalar float tensor in (0, 1] for nucleus sampling. top_p=1.0
                    (the default) keeps every token, i.e. no filtering. Pass it
                    as a runtime input to tune per request.
@@ -31,10 +32,10 @@ def __init__(self, model: nn.Module):
         self.model = model
 
     def forward(self, *args, temperature, top_k=None, top_p=1.0, seed=None, **kwargs):
-        if top_k is not None:
-            raise NotImplementedError("top_k sampling is not implemented")
         logits = self.model(*args, **kwargs)  # [B, S, vocab]
         last = logits[:, -1, :]  # [B, vocab]
         if not isinstance(top_p, torch.Tensor):
             top_p = torch.tensor(float(top_p))
-        return torch.ops.mlx.sample(last, temperature, top_p, seed)
+        if top_k is not None and not isinstance(top_k, torch.Tensor):
+            top_k = torch.tensor(int(top_k), dtype=torch.int64)
+        return torch.ops.mlx.sample(last, temperature, top_p, seed, top_k)

backends/mlx/ops.py

@@ -24,6 +24,7 @@
     emit_lifted_constant,
     emit_quantized_biases,
     emit_shape,
+    emit_sub_int,
     parse_dequant_node,
     to_mlx_qparams,
     torch_dtype_to_scalar_type,
@@ -3528,10 +3529,11 @@ def _sample_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     skipping the sampling chain (so 0 is exact, not the small-epsilon approx).
     """
     args = P.args(n)
-    require_args(args, 3, 4, "mlx.sample")
+    require_args(args, 3, 5, "mlx.sample")
     require_kwargs(P.kwargs(n), set(), "mlx.sample")
     logits, temperature, top_p = args[0], args[1], args[2]
     seed = args[3] if len(args) > 3 and args[3] is not None else None
+    top_k = args[4] if len(args) > 4 and args[4] is not None else None
 
     temp_dt = n.args[1].meta["val"].dtype
     out = P.make_or_get_slot(n)
@@ -3674,6 +3676,46 @@ def emit_sample():
                 out=P.slot_to_tid(drop),
             )
         )
+        if top_k is not None:
+            _, top_k_val = P.make_tmp_value_slot()
+            P.emit(ItemIntNode(x=P.slot_to_tid(top_k), out=P.slot_to_vid(top_k_val)))
+            top_k_iov = P.to_int_or_vid(top_k_val)
+            top_k_index = emit_sub_int(P, top_k_iov, IntOrVid.from_literal(1))
+            _, top_k_thresh = P.make_tmp_slot()
+            P.emit(
+                TakeNode(
+                    x=P.slot_to_tid(sorted_p),
+                    index=(
+                        IntOrVidOrTid.from_vid(top_k_index.vid)
+                        if top_k_index.is_vid
+                        else IntOrVidOrTid.from_literal(top_k_index.literal)
+                    ),
+                    out=P.slot_to_tid(top_k_thresh),
+                    axis=-1,
+                )
+            )
+            P.emit(
+                ExpandDimsNode(
+                    x=P.slot_to_tid(top_k_thresh),
+                    out=P.slot_to_tid(top_k_thresh),
+                    axis=-1,
+                )
+            )
+            _, drop_k = P.make_tmp_slot()
+            P.emit(
+                LessNode(
+                    a=P.slot_to_tid(probs),
+                    b=P.slot_to_tid(top_k_thresh),
+                    out=P.slot_to_tid(drop_k),
+                )
+            )
+            P.emit(
+                LogicalOrNode(
+                    a=P.slot_to_tid(drop),
+                    b=P.slot_to_tid(drop_k),
+                    out=P.slot_to_tid(drop),
+                )
+            )
         neg_inf = emit_lifted_constant(P, float("-inf"), torch.float32)
         # masked = where(drop, -inf, scaled); then add gumbel noise in place.
         _, masked = P.make_tmp_slot()

backends/mlx/test/test_ops.py

@@ -7676,6 +7676,17 @@ def forward(self, logits, temperature, seed, top_p):
         return self.head(logits, temperature=temperature, seed=seed, top_p=top_p)
 
 
+class TopKSampleModel(nn.Module):
+    """SamplingHead with temperature, seed, and top_k as runtime inputs."""
+
+    def __init__(self):
+        super().__init__()
+        self.head = SamplingHead(_LogitsPassthrough())
+
+    def forward(self, logits, temperature, seed, top_k):
+        return self.head(logits, temperature=temperature, seed=seed, top_k=top_k)
+
+
 @register_test
 class SampleSeededTest(OpTestCase):
     """Seeded sample lowers to one MLX segment; seed threads in via ItemIntNode."""
@@ -7756,6 +7767,39 @@ def create_inputs(self) -> Tuple[torch.Tensor, ...]:
         )
 
 
+@register_test
+class SampleTopKTest(OpTestCase):
+    """Top-k sample emits the extra threshold and combined mask nodes."""
+
+    name = "sample_top_k"
+    skip_comparison = True  # sampling RNG is not host/device bit-identical
+    expected_node_counts = {
+        "IfNode": 1,
+        "RandomBitsNode": 1,
+        "ArgmaxNode": 2,
+        "ItemIntNode": 3,  # seed + top_k + temperature>0 condition
+        "SoftmaxNode": 1,
+        "SortNode": 1,
+        "CumsumNode": 1,
+        "MinNode": 1,
+        "TakeNode": 1,
+        "ExpandDimsNode": 1,
+        "LogicalOrNode": 1,
+        "WhereNode": 2,
+    }
+
+    def create_model(self) -> nn.Module:
+        return TopKSampleModel()
+
+    def create_inputs(self) -> Tuple[torch.Tensor, ...]:
+        return (
+            torch.randn(1, 4, 256),
+            torch.tensor(0.8),
+            torch.tensor(0, dtype=torch.int64),
+            torch.tensor(2, dtype=torch.int64),
+        )
+
+
 @register_test
 class SampleGreedyTest(OpTestCase):
     """Greedy argmax(logits) is bit-exact host/device, so verify the token with the

backends/mlx/test/test_sample.py

@@ -63,6 +63,17 @@ def forward(self, logits, temperature, seed, top_p):
         return self.head(logits, temperature=temperature, seed=seed, top_p=top_p)
 
 
+class TopKSampleModel(nn.Module):
+    """SamplingHead with temperature, seed, and top_k as runtime inputs."""
+
+    def __init__(self):
+        super().__init__()
+        self.head = SamplingHead(_LogitsPassthrough())
+
+    def forward(self, logits, temperature, seed, top_k):
+        return self.head(logits, temperature=temperature, seed=seed, top_k=top_k)
+
+
 def _ref_gumbel_max(logits: torch.Tensor, temperature: float, seed: int):
     """Independent Gumbel-max reference using the same torch RNG as the op."""
     gen = torch.Generator().manual_seed(seed)
@@ -76,11 +87,18 @@ def _tv_distance(p: torch.Tensor, q: torch.Tensor) -> float:
     return 0.5 * torch.abs(p - q).sum().item()
 
 
-def _sample(logits, temperature, seed: Optional[int], top_p: float = 1.0):
+def _sample(
+    logits,
+    temperature,
+    seed: Optional[int],
+    top_p: float = 1.0,
+    top_k: Optional[int] = None,
+):
     t = torch.tensor(float(temperature))
     s = None if seed is None else torch.tensor(int(seed), dtype=torch.int64)
     p = torch.tensor(float(top_p))  # 1.0 = off
-    return torch.ops.mlx.sample(logits, t, p, s)
+    k = None if top_k is None else torch.tensor(int(top_k), dtype=torch.int64)
+    return torch.ops.mlx.sample(logits, t, p, s, k)
 
 
 class TestSampleOp(unittest.TestCase):
@@ -142,6 +160,25 @@ def test_top_p_one_keeps_all(self):
         tokens = _sample(base.expand(20000, 4), 1.0, seed=0, top_p=1.0)
         self.assertTrue((tokens == 3).any())
 
+    def test_top_k_restricts_to_top_k(self):
+        # probs [0.5, 0.3, 0.15, 0.05]; top_k=2 keeps {0,1}.
+        base = torch.log(torch.tensor([0.5, 0.3, 0.15, 0.05]))
+        tokens = _sample(base.expand(5000, 4), 1.0, seed=0, top_k=2)
+        self.assertTrue(torch.isin(tokens, torch.tensor([0, 1])).all())
+        self.assertEqual(set(tokens.tolist()), {0, 1})
+
+    def test_top_k_none_keeps_all(self):
+        # top_k=None -> no filtering; the tail token (index 3) is reachable.
+        base = torch.log(torch.tensor([0.5, 0.3, 0.15, 0.05]))
+        tokens = _sample(base.expand(20000, 4), 1.0, seed=0, top_k=None)
+        self.assertTrue((tokens == 3).any())
+
+    def test_top_k_and_top_p_compose(self):
+        # top_p=0.7 keeps {0,1}; top_k=1 intersects that to {0}.
+        base = torch.log(torch.tensor([0.5, 0.3, 0.15, 0.05]))
+        tokens = _sample(base.expand(5000, 4), 1.0, seed=0, top_p=0.7, top_k=1)
+        self.assertEqual(set(tokens.tolist()), {0})
+
 
 class TestSampleExport(unittest.TestCase):
     """Runtime-input semantics that survive export: temperature and seed stay
@@ -218,6 +255,24 @@ def test_top_p_end_to_end(self):
         (token,) = load_tensors_from_bin(out_bin)
         self.assertIn(int(token), {0, 1, 2})  # tail token (index 3) excluded
 
+    def test_top_k_end_to_end(self):
+        # On-device top-k: probs [0.5,0.3,0.15,0.05], top_k=2 -> token in {0,1}.
+        logits = torch.log(torch.tensor([0.5, 0.3, 0.15, 0.05])).view(1, 1, 4)
+        inputs = (
+            logits,
+            torch.tensor(1.0),
+            torch.tensor(0, dtype=torch.int64),
+            torch.tensor(2, dtype=torch.int64),
+        )
+        tmp = Path(self._tmp)
+        pte, in_bin, out_bin = tmp / "topk.pte", tmp / "in.bin", tmp / "out.bin"
+        export_model_to_pte(TopKSampleModel(), inputs, pte)
+        save_tensors_to_bin(list(inputs), in_bin)
+
+        self.assertTrue(run_cpp_test_runner(pte, in_bin, out_bin))
+        (token,) = load_tensors_from_bin(out_bin)
+        self.assertIn(int(token), {0, 1})  # tail tokens excluded
+
 
 if __name__ == "__main__":
     unittest.main()