MLX: on-device token sampling with Gumbel-max

backends/mlx/custom_ops.py

@@ -391,3 +391,35 @@ def gather_qmm_fake(
     else:
         batch = w.shape[:-2]
     return x.new_empty((*batch, M, N))
+
+
+@torch.library.custom_op("mlx::sample", mutates_args=())
+def sample(
+    logits: Tensor, temperature: Tensor, seed: Optional[Tensor] = None
+) -> Tensor:
+    """
+    Gumbel-max sampling from softmax(logits / temperature).
+    logits:      [B, vocab]
+    temperature: scalar float tensor    (runtime input)
+    seed:        scalar int tensor or None
+                 - tensor -> deterministic, keyed RNG (random::key(seed))
+                 - None   -> MLX global KeySequence (non-deterministic)
+    -> token_id: [B] int64
+
+    Host/CPU reference used for export (shape/meta) and distributional checks
+    only. It is NOT bit-identical to the lowered on-device graph: this uses torch
+    RNG (plain torch.rand, no uint32/nextafter uniform) while the delegate uses
+    MLX RNG, so a given seed does not reproduce the same tokens host vs. device.
+    """
+    if seed is None:
+        u = torch.rand(logits.shape)  # global RNG
+    else:
+        gen = torch.Generator().manual_seed(int(seed.item()))
+        u = torch.rand(logits.shape, generator=gen)
+    gumbel = -torch.log(-torch.log(u))
+    return torch.argmax(logits / temperature + gumbel, dim=-1)
+
+
+@torch.library.register_fake("mlx::sample")
+def sample_fake(logits, temperature, seed=None):
+    return logits.new_empty(logits.shape[:-1], dtype=torch.long)

backends/mlx/llm/sampling.py

@@ -0,0 +1,32 @@
+#
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+#
+
+import torch
+import torch.nn as nn
+
+
+class SamplingHead(nn.Module):
+    """
+    Wraps a model that returns logits and samples a token id on-device.
+
+        forward(*model_args, temperature, seed=None, **model_kwargs) -> token_id
+
+      temperature: scalar float tensor, e.g. torch.tensor(0.8). Must be > 0;
+                   logits are divided by it, so 0.0 yields inf/nan. For greedy,
+                   pass a small epsilon (e.g. 1e-4), not 0.
+      seed:        scalar int tensor (seeded) or None (unseeded export)
+    """
+
+    def __init__(self, model: nn.Module):
+        super().__init__()
+        self.model = model
+
+    def forward(self, *args, temperature, seed=None, **kwargs):
+        logits = self.model(*args, **kwargs)  # [B, S, vocab]
+        last = logits[:, -1, :]  # [B, vocab]
+        return torch.ops.mlx.sample(last, temperature, seed)

backends/mlx/ops.py

@@ -22,6 +22,7 @@
 from executorch.backends.mlx.builder.op_helpers import (
     emit_lifted_constant,
     emit_quantized_biases,
+    emit_shape,
     parse_dequant_node,
     to_mlx_qparams,
     torch_dtype_to_scalar_type,
@@ -115,6 +116,7 @@
     PartitionNode,
     PowerNode,
     ProdNode,
+    RandomBitsNode,
     ReciprocalNode,
     RemainderNode,
     RepeatNode,
@@ -3454,6 +3456,105 @@ def _argmax_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     return out
 
 
+@REGISTRY.register(target=[torch.ops.mlx.sample.default])
+def _sample_handler(P: MLXProgramBuilder, n: Node) -> Slot:
+    """Gumbel-max sampling: argmax(logits / temperature + gumbel_noise).
+
+    Reproduces MLX's uniform -> gumbel -> argmax layering in the IR using the
+    new RandomBitsNode plus existing elementwise nodes, so a sampled token id is
+    produced on-device instead of returning the full logits tensor.
+    """
+    args = P.args(n)
+    require_args(args, 2, 3, "mlx.sample")
+    require_kwargs(P.kwargs(n), set(), "mlx.sample")
+    logits, temperature = args[0], args[1]
+    seed = args[2] if len(args) > 2 and args[2] is not None else None
+
+    dt = n.args[0].meta["val"].dtype
+    shape = emit_shape(P, n.args[0], logits)
+
+    # Optional runtime seed: tensor -> SymInt (Vid) via ItemIntNode. Absent ->
+    # leave RandomBitsNode.seed unset (MLX global RNG).
+    seed_field = None
+    if seed is not None:
+        _, seed_val = P.make_tmp_value_slot()
+        P.emit(ItemIntNode(x=P.slot_to_tid(seed), out=P.slot_to_vid(seed_val)))
+        seed_field = P.to_int_or_vid(seed_val)
+
+    # uniform u in [0, 1): bits/uint32_max, clamped just below 1 (random.cpp:95)
+    _, bits = P.make_tmp_slot()
+    P.emit(
+        RandomBitsNode(out=P.slot_to_tid(bits), shape=shape, width=4, seed=seed_field)
+    )
+    _, bits_f = P.make_tmp_slot()
+    P.emit(
+        AsTypeNode(
+            x=P.slot_to_tid(bits),
+            out=P.slot_to_tid(bits_f),
+            scalar_type=torch_dtype_to_scalar_type(torch.float32),
+        )
+    )
+    umax = emit_lifted_constant(P, 4294967295.0, torch.float32)
+    _, div0 = P.make_tmp_slot()
+    P.emit(
+        DivideNode(
+            a=P.slot_to_tid(bits_f), b=P.slot_to_tid(umax), out=P.slot_to_tid(div0)
+        )
+    )
+    prev1 = emit_lifted_constant(
+        P, float(torch.nextafter(torch.tensor(1.0), torch.tensor(0.0))), torch.float32
+    )
+    _, clamp = P.make_tmp_slot()
+    P.emit(
+        MinimumNode(
+            a=P.slot_to_tid(div0), b=P.slot_to_tid(prev1), out=P.slot_to_tid(clamp)
+        )
+    )
+    # gumbel noise g = -log(-log(u)) (random.cpp:367), computed in float32.
+    # Keep the uniform in float32 through the log chain: casting it down to a
+    # low-precision dtype (e.g. bf16) can round the clamp (~0.99999994) back up
+    # to 1.0 -> log(1.0)=0 -> -log(0)=+inf, which then dominates argmax (a fixed
+    # seed makes that the same token every step). Only the finite gumbel is cast
+    # to the logits dtype.
+    _, l1 = P.make_tmp_slot()
+    P.emit(LogNode(x=P.slot_to_tid(clamp), out=P.slot_to_tid(l1)))
+    _, g1 = P.make_tmp_slot()
+    P.emit(NegNode(x=P.slot_to_tid(l1), out=P.slot_to_tid(g1)))
+    _, l2 = P.make_tmp_slot()
+    P.emit(LogNode(x=P.slot_to_tid(g1), out=P.slot_to_tid(l2)))
+    _, g_f32 = P.make_tmp_slot()
+    P.emit(NegNode(x=P.slot_to_tid(l2), out=P.slot_to_tid(g_f32)))
+    _, g = P.make_tmp_slot()
+    P.emit(
+        AsTypeNode(
+            x=P.slot_to_tid(g_f32),
+            out=P.slot_to_tid(g),
+            scalar_type=torch_dtype_to_scalar_type(dt),
+        )
+    )
+
+    # sample: argmax(logits / temperature + g) over the vocab axis
+    _, scaled = P.make_tmp_slot()
+    P.emit(
+        DivideNode(
+            a=P.slot_to_tid(logits),
+            b=P.slot_to_tid(temperature),
+            out=P.slot_to_tid(scaled),
+        )
+    )
+    _, noisy = P.make_tmp_slot()
+    P.emit(
+        AddNode(a=P.slot_to_tid(scaled), b=P.slot_to_tid(g), out=P.slot_to_tid(noisy))
+    )
+    out = P.make_or_get_slot(n)
+    P.emit(
+        ArgmaxNode(
+            x=P.slot_to_tid(noisy), out=P.slot_to_tid(out), axis=-1, keepdims=False
+        )
+    )
+    return out
+
+
 @REGISTRY.register(target=[torch.ops.aten.argmin.default])
 def _argmin_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     """Handle aten.argmin - index of min element along axis."""

backends/mlx/runtime/MLXInterpreter.h

@@ -1695,6 +1695,23 @@ exec_argmax(const ArgmaxNode& n, ExecutionState& st, StreamOrDevice s) {
   st.set_tensor(n.out, argmax(x, n.axis, n.keepdims, s));
 }
 
+inline void exec_random_bits(
+    const RandomBitsNode& n,
+    ExecutionState& st,
+    StreamOrDevice s) {
+  // Only width=4 (uint32) is supported; reject other widths.
+  if (n.width != 4) {
+    throw std::runtime_error("random_bits: only width=4 (uint32) is supported");
+  }
+  auto shape = to_shape(n.shape, st);
+  check_allocation_bounded(shape, uint32, "random_bits");
+  std::optional<array> key = std::nullopt;
+  if (n.seed.has_value()) {
+    key = random::key(static_cast<uint64_t>(resolve_int(n.seed.value(), st)));
+  }
+  st.set_tensor(n.out, random::bits(shape, n.width, key, s));
+}
+
 inline void
 exec_argmin(const ArgminNode& n, ExecutionState& st, StreamOrDevice s) {
   const auto& x = st.const_tensor_ref(n.x);
@@ -2057,6 +2074,9 @@ class Interpreter {
       case OpCode::ARGMAX:
         ops::exec_argmax(std::get<ArgmaxNode>(instr.node), st, s);
         break;
+      case OpCode::RANDOM_BITS:
+        ops::exec_random_bits(std::get<RandomBitsNode>(instr.node), st, s);
+        break;
       case OpCode::SLICE_UPDATE:
         ops::exec_slice_update(std::get<SliceUpdateNode>(instr.node), st, s);
         break;

backends/mlx/serialization/generate.py

@@ -831,7 +831,12 @@ def _emit_py_prebuild(kind: str, fld: FBSField) -> List[str]:
     if kind in _PY_PREBUILD_OFFSET:
         suffix = "_off"
         expr = _PY_PREBUILD_OFFSET[kind].format(name=n)
-        return [f"        {n}{suffix} = {expr}"]
+        # optional_str carries its own None handling; other compound offset
+        # fields (int_or_vid, etc.) must be guarded when optional so a None
+        # value is serialized as an absent field rather than crashing.
+        if fld.required or kind == "optional_str":
+            return [f"        {n}{suffix} = {expr}"]
+        return [f"        {n}{suffix} = {expr} if op.{n} is not None else None"]
     return []
 
 
@@ -855,7 +860,12 @@ def _emit_py_add(
         return [f"        {add}(builder, op.{n})"]
     # Pre-built offsets (string, compound types)
     if kind in ("str", "int_or_vid", "float_or_vid", "vid_or_tid", "int_or_vid_or_tid"):
-        return [f"        {add}(builder, {n}_off)"]
+        if fld.required:
+            return [f"        {add}(builder, {n}_off)"]
+        return [
+            f"        if {n}_off is not None:",
+            f"            {add}(builder, {n}_off)",
+        ]
     # Pre-built vectors (required vs optional)
     if kind in (
         "list_int",
@@ -1056,6 +1066,8 @@ def _fbs_type_to_cpp(
             return "std::optional<Tid>"
         if fbs_type == "Vid":
             return "std::optional<Vid>"
+        if fbs_type in ("IntOrVid", "FloatOrVid", "VidOrTid", "IntOrVidOrTid"):
+            return f"std::optional<{cpp_type}>"
         if fld is not None and fld.default == "null" and fbs_type in FBS_TO_CPP:
             return f"std::optional<{cpp_type}>"
 
@@ -1140,7 +1152,7 @@ def _generate_loader_case(table: FBSTable) -> List[str]:
 
         fb_field_name = fld.name
         kind = _get_field_kind(fld, table)
-        load_lines = _emit_cpp_load(kind, fld.name, fb_field_name, table)
+        load_lines = _emit_cpp_load(kind, fld.name, fb_field_name, table, fld)
         if load_lines is None:
             raise ValueError(
                 f"Unhandled field kind '{kind}' for field '{fld.name}' in table '{table.name}'. "
@@ -1172,16 +1184,24 @@ def _generate_loader_case(table: FBSTable) -> List[str]:
 }
 
 
-def _emit_cpp_load(
-    kind: str, name: str, fb_name: str, table=None
+def _emit_cpp_load(  # noqa: C901
+    kind: str, name: str, fb_name: str, table=None, fld=None
 ) -> "List[str] | None":
     """Emit C++ load lines for a field kind, or None if kind is unrecognized."""
     # Interned string fields share one std::string via the load-time pool.
     if _is_interned_str(table, name) and kind in ("str", "optional_str"):
         return [f"      node.{name} = strpool.intern(fb->{fb_name}());"]
-    # Required struct / compound via converter
+    # Struct / compound via converter
     if kind in _CPP_CONVERTER:
         conv = _CPP_CONVERTER[kind]
+        # Optional compound fields (e.g. an optional IntOrVid) must be
+        # presence-guarded; convert_* throws on a null FlatBuffer pointer.
+        if fld is not None and not fld.required:
+            return [
+                f"      if (fb->{fb_name}()) {{",
+                f"        node.{name} = {conv}(fb->{fb_name}());",
+                "      }",
+            ]
         return [f"      node.{name} = {conv}(fb->{fb_name}());"]
     # Scalars (direct value)
     if kind in ("int", "float", "bool"):

backends/mlx/serialization/schema.fbs

@@ -985,6 +985,14 @@ table IfNode {
     else_chain_idx: uint32;            // index into MLXGraph.instruction_chains
 }
 
+table RandomBitsNode {
+    out: Tid (required);
+    shape: [IntOrVid] (required);
+    seed: IntOrVid;              // OPTIONAL: present -> random::key(seed);
+                                 //           absent  -> MLX global KeySequence
+    width: int32 = 4;            // bytes per element (4 -> uint32)
+}
+
 // Custom Metal kernel execution via mlx::core::fast::metal_kernel().
 // Two-phase API:
 //   1. Factory: metal_kernel(name, input_names, output_names, source, header,
@@ -1161,7 +1169,8 @@ union OpNode {
     BitwiseAndNode,
     BitwiseOrNode,
     BitwiseXorNode,
-    IfNode
+    IfNode,
+    RandomBitsNode
     // BC: Add new op nodes here (append only)
 }