feat: compile Bumblebee Text.generation fully native (CM13)

The goal gate: Bumblebee.Text.generation — greedy and multinomial
sampling — now compiles fully native under the single-NIF replay (no
fallback), bit-identical to the Evaluator. The whole generation graph
(transformer forward, the defn while decode loop, dynamic KV-cache
writes, cumsum position ids, argmax / sampling token selection, threefry)
lowers in one piece.

Wiring this to completion needed:
- cumsum/cumprod/cummax/cummin (last-axis native fast path; interior axes
  fall through to the block's composed expansion, mirroring the backend)
- gather: single-axis via take, multi-axis via the split-indices + MLX
  multi-index gather path (ports Emily.Backend.gather/4)
- stack
- n_inputs fix: the compiled program pins n_inputs to the true flattened
  parameter count (length of the flattened vars), not IR.lower's
  referenced-parameter count — an unused input (e.g. the seed in greedy)
  no longer desyncs the eval-time input count

Gates: generation_native_test.exs (greedy + sampling, native vs
evaluator, native_fallback: :raise, default :conformance) on the cached
tiny-random Qwen3 causal LM, driven through build_generate on in-vocab
input_ids (no tokenizer, no OOB-gather pitfall). Plus compiler_indexing_test.exs
pinning cumsum/gather/stack in isolation.

Author: ausimian

RELEASE.md

@@ -50,6 +50,16 @@
   selection on the native path. Remaining gaps (`gather`/scatter,
   pooling/`window_*`, cumulative) continue to work via the graceful fallback.
 
+- **`Bumblebee.Text.generation` compiles fully native — greedy and sampling.**
+  The headline result: an end-to-end Bumblebee generation (the transformer
+  forward, the `defn while` decode loop, dynamic KV-cache writes, `cumsum`
+  position ids, `argmax`/multinomial token selection, threefry sampling) now
+  lowers to the single-NIF replay with **no fallback**, producing token ids
+  bit-identical to the Evaluator. Wiring this required `cumsum`/`cumprod`/
+  `cummax`/`cummin` (last-axis fast path), single- and multi-axis `gather`,
+  and `stack`. Greedy and multinomial sampling are gated in
+  `generation_native_test.exs`.
+
 - **`defn while` compiles native.** Data-dependent loops — including
   `Bumblebee.Text.generation`'s decode loop — now lower to the single-NIF
   replay instead of falling back. The condition and body become nested

c_src/emily/opcodes.hpp

@@ -142,9 +142,20 @@ enum class Opcode : int64_t {
   // path indexes by a loop counter). operands [a, start(s32 [naxes])];
   // iattrs [[axes...], [slice_sizes...]]
   DynSlice = 74,
+  // Inclusive cumulative reductions along an axis. operands [a];
+  // iattrs [[axis], [reverse]] (Nx cumulation is always inclusive).
+  CumSum = 75,
+  CumProd = 76,
+  CumMax = 77,
+  CumMin = 78,
+  // Multi-axis gather. operands [input, idx0, idx1, ...] (one s32 index
+  // array per gathered axis); iattrs [[axes...], [slice_sizes...]]
+  Gather = 79,
+  // Stack tensors along a new axis. operands [t0, t1, ...]; iattrs [[axis]]
+  Stack = 80,
 };
 
-inline constexpr int64_t kOpcodeCount = 75;
+inline constexpr int64_t kOpcodeCount = 81;
 
 // Quant mode code (Emily.IR @quant_modes) -> MLX mode string.
 inline std::string qmode_from_code(int64_t code) {
@@ -565,6 +576,37 @@ inline mx::array dispatch_op(Opcode op, const std::vector<mx::array> &in,
     return mx::slice(in[0], in[1],
                      emily::to_int_vec(attr_at(iattrs, 0, "dyn_slice")),
                      emily::to_mlx_shape(attr_at(iattrs, 1, "dyn_slice")), s);
+  case Opcode::CumSum:
+    return mx::cumsum(arg1(in, "cumsum"),
+                      emily::checked_int(scalar_at(iattrs, 0, "cumsum"), "axis"),
+                      scalar_at(iattrs, 1, "cumsum") != 0, /*inclusive=*/true, s);
+  case Opcode::CumProd:
+    return mx::cumprod(arg1(in, "cumprod"),
+                       emily::checked_int(scalar_at(iattrs, 0, "cumprod"), "axis"),
+                       scalar_at(iattrs, 1, "cumprod") != 0, /*inclusive=*/true, s);
+  case Opcode::CumMax:
+    return mx::cummax(arg1(in, "cummax"),
+                      emily::checked_int(scalar_at(iattrs, 0, "cummax"), "axis"),
+                      scalar_at(iattrs, 1, "cummax") != 0, /*inclusive=*/true, s);
+  case Opcode::CumMin:
+    return mx::cummin(arg1(in, "cummin"),
+                      emily::checked_int(scalar_at(iattrs, 0, "cummin"), "axis"),
+                      scalar_at(iattrs, 1, "cummin") != 0, /*inclusive=*/true, s);
+  case Opcode::Gather: {
+    if (in.size() < 2) {
+      throw std::invalid_argument("gather expects input + >=1 index operand");
+    }
+    std::vector<mx::array> indices(in.begin() + 1, in.end());
+    return mx::gather(in[0], indices,
+                      emily::to_int_vec(attr_at(iattrs, 0, "gather")),
+                      emily::to_mlx_shape(attr_at(iattrs, 1, "gather")), s);
+  }
+  case Opcode::Stack:
+    if (in.empty()) {
+      throw std::invalid_argument("stack expects >= 1 operand");
+    }
+    return mx::stack(in, emily::checked_int(scalar_at(iattrs, 0, "stack"), "axis"),
+                     s);
   }
   throw std::invalid_argument("unknown opcode " +
                               std::to_string(static_cast<int64_t>(op)));

lib/emily/compiler.ex

@@ -177,8 +177,15 @@ defmodule Emily.Compiler do
     {template, leaves_rev} =
       Composite.traverse(expr, [], fn leaf, acc -> {Nx.to_template(leaf), [leaf | acc]} end)
 
+    # The flattened parameter leaves are the true input count and slot order
+    # (the closure realises them in this order; `{:input, i}` indexes it).
+    # `IR.lower` only counts the parameters it *references*, which undercounts
+    # when an input is unused (e.g. the `seed` in greedy generation) — pin
+    # `n_inputs` to the real arity so the eval-time input count matches.
+    n_inputs = length(Composite.flatten_list(vars))
+
     case lower(Enum.reverse(leaves_rev), mode, key) do
-      {:ok, ir} -> {:ok, replay_closure(template, ir)}
+      {:ok, ir} -> {:ok, replay_closure(template, %{ir | n_inputs: n_inputs})}
       :fallback -> :fallback
     end
   end

lib/emily/ir.ex

@@ -124,7 +124,15 @@ defmodule Emily.IR do
     # RNG / dynamic indexing primitives
     bitcast: 72,
     erf_inv: 73,
-    dyn_slice: 74
+    dyn_slice: 74,
+    # inclusive cumulative reductions (iattrs [[axis],[reverse]])
+    cumsum: 75,
+    cumprod: 76,
+    cummax: 77,
+    cummin: 78,
+    # multi-axis gather: operands [input, idx0, ...]; iattrs [[axes],[slice_sizes]]
+    gather: 79,
+    stack: 80
   }
 
   # Quant mode string -> code; decoded by qmode_from_code in
@@ -409,6 +417,14 @@ defmodule Emily.IR do
     reduce_min: :min
   }
 
+  # Cumulative reductions arrive as `Nx.block/4` nodes. Block struct -> opcode.
+  @cumulative_blocks %{
+    Nx.Block.CumulativeSum => :cumsum,
+    Nx.Block.CumulativeProduct => :cumprod,
+    Nx.Block.CumulativeMax => :cummax,
+    Nx.Block.CumulativeMin => :cummin
+  }
+
   defp lower_op(%T{data: %Nx.Defn.Expr{op: op, args: [a, opts]}} = t, state)
        when is_map_key(@reductions, op) do
     {ra, state} = lower_node(a, state)
@@ -568,6 +584,41 @@ defmodule Emily.IR do
     end
   end
 
+  # gather(input, indices, opts). Mirrors Emily.Backend.gather/4: single-axis
+  # gathers `take` along the one axis (indices cast to s32); multi-axis
+  # gathers split the `{..., R}` index tensor into R per-axis index arrays
+  # and use MLX's multi-index gather. Both reshape to the output shape (token
+  # selection in sampling is this shape). A layout MLX gather can't express
+  # raises, so the graceful fallback handles it.
+  defp lower_op(%T{data: %Nx.Defn.Expr{op: :gather, args: [input, indices, opts]}} = t, state) do
+    axes = opts[:axes]
+    indices_shape = Tuple.to_list(indices.shape)
+
+    {ri, state} = lower_node(input, state)
+    {rx, state} = lower_node(indices, state)
+
+    {r, state} =
+      case axes do
+        [axis] ->
+          {ix, state} = emit(state, :astype, [rx], [[dtype_code({:s, 32})]])
+          emit(state, :take, [ri, ix], [[axis]])
+
+        _ when is_list(axes) ->
+          unless scatter_gather_compatible?(indices_shape, axes) do
+            raise ArgumentError,
+                  "Emily Expr compiler: gather index layout #{inspect(indices_shape)} " <>
+                    "for axes #{inspect(axes)} is not MLX-gather-compatible."
+          end
+
+          {idx_refs, state} = split_indices_for_gather(rx, indices_shape, length(axes), state)
+          slice_sizes = slice_sizes_for_gather(input.shape, axes)
+          emit(state, :gather, [ri | idx_refs], [axes, slice_sizes])
+      end
+
+    {r, state} = emit(state, :reshape, [r], [Tuple.to_list(t.shape)])
+    coerce(r, t.type, state)
+  end
+
   # put_slice(src, start_indices, slice): write `slice` into `src` at
   # `start_indices`. Mirrors Emily.Backend.put_slice/4 (cast src + update
   # to out.type), but supports RUNTIME (tensor) start indices — the decode
@@ -623,6 +674,12 @@ defmodule Emily.IR do
     emit_coerced(state, :concatenate, refs, [[axis]], t.type)
   end
 
+  # stack(tensors, axis): join along a NEW axis. Mirrors Emily.Backend.stack/3.
+  defp lower_op(%T{data: %Nx.Defn.Expr{op: :stack, args: [tensors, axis]}} = t, state) do
+    {refs, state} = Enum.map_reduce(tensors, state, &lower_node/2)
+    emit_coerced(state, :stack, refs, [[axis]], t.type)
+  end
+
   # conv: ports Emily.Backend.conv/4 — permute input -> NHWC and kernel ->
   # OHWI (casting both to out.type), mx::conv_general, then permute the
   # result NHWC -> NCHW -> the user's output layout. batch_group_size > 1
@@ -858,11 +915,26 @@ defmodule Emily.IR do
     emit_coerced(state, :take, [ri, rx], [[axis]], t.type)
   end
 
+  # Cumulative families. Like Emily.Backend.block/4, the last-axis case uses
+  # the native MLX `cumsum`/`cumprod`/`cummax`/`cummin` kernel; interior axes
+  # (which MLX can't always factor) fall back to the block's composed
+  # expansion. Nx cumulation is always inclusive.
+  defp lower_block(%mod{axis: axis, reverse: reverse}, [t], expr, out, state)
+       when is_map_key(@cumulative_blocks, mod) do
+    if axis == tuple_size(out.shape) - 1 do
+      {rt, state} = lower_node(t, state)
+      op = Map.fetch!(@cumulative_blocks, mod)
+      emit_coerced(state, op, [rt], [[axis], [bool_int(reverse)]], out.type)
+    else
+      lower_node(expr, state)
+    end
+  end
+
   # Any other block struct raises. Lowering the block's composed
   # expansion would silently diverge from the Evaluator whenever
   # Emily.Backend.block/4 dispatches that struct through a fused / native
-  # kernel (e.g. SDPAWithSinks, the Nx.Block.LinAlg.* / Take / FFT /
-  # cumulative families) — a worse failure than a clear "unsupported".
+  # kernel (e.g. SDPAWithSinks, the Nx.Block.LinAlg.* / FFT families) — a
+  # worse failure than a clear "unsupported".
   # Additional fused blocks are added alongside their opcode.
   defp lower_block(struct, _in_args, _expr, _t, _state) do
     raise ArgumentError,
@@ -1007,4 +1079,39 @@ defmodule Emily.IR do
     {r, state} = emit(state, :maximum, [ref, lo_c])
     emit(state, :minimum, [r, hi_c])
   end
+
+  # MLX's multi-index gather needs the index tensor's leading dims to be the
+  # batch and the last axis to select across `axes` (mirrors
+  # Emily.Backend.scatter_gather_compatible?/2).
+  defp scatter_gather_compatible?(indices_shape, axes) do
+    is_list(axes) and axes != [] and length(indices_shape) >= 2 and
+      List.last(indices_shape) == length(axes)
+  end
+
+  # Split an `{..., R}` index tensor into R per-axis s32 index arrays (each
+  # the leading batch with the last axis dropped) — ports
+  # Emily.Backend.split_indices_per_axis/4 with static slices.
+  defp split_indices_for_gather(indices_ref, indices_shape, n_axes, state) do
+    rank = length(indices_shape)
+    last_axis = rank - 1
+    batch_shape = Enum.take(indices_shape, last_axis)
+    strides = List.duplicate(1, rank)
+    batch_zeros = List.duplicate(0, last_axis)
+
+    Enum.map_reduce(0..(n_axes - 1)//1, state, fn i, state ->
+      {r, state} =
+        emit(state, :slice, [indices_ref], [batch_zeros ++ [i], batch_shape ++ [i + 1], strides])
+
+      {r, state} = emit(state, :squeeze, [r], [[last_axis]])
+      emit(state, :astype, [r], [[dtype_code({:s, 32})]])
+    end)
+  end
+
+  # Per-axis slice size for gather: 1 on a gathered axis, the full extent
+  # otherwise (mirrors Emily.Backend.slice_sizes_for_gather/2).
+  defp slice_sizes_for_gather(input_shape, axes) do
+    axes_set = MapSet.new(axes)
+    rank = tuple_size(input_shape)
+    for i <- 0..(rank - 1)//1, do: if(i in axes_set, do: 1, else: elem(input_shape, i))
+  end
 end

test/emily/compiler_indexing_test.exs

@@ -0,0 +1,54 @@
+defmodule Emily.CompilerIndexingTest do
+  @moduledoc """
+  CM13 — the gather / stack / cumulative ops Bumblebee's generation loop
+  needs, lowered single-NIF and bit-identical to the Evaluator. (The
+  generation conformance test exercises them in anger; these pin each op in
+  isolation.) Run under `native_fallback: :raise` so a regression raises
+  rather than silently falling back.
+  """
+  use ExUnit.Case, async: true
+
+  @native [compiler: Emily.Compiler, native: true, native_fallback: :raise]
+  @eval [compiler: Emily.Compiler]
+
+  defp t(data, opts \\ []), do: Nx.tensor(data, [backend: Emily.Backend] ++ opts)
+
+  defp equiv(fun, args) do
+    native = apply(Nx.Defn.jit(fun, @native), args)
+    eval = apply(Nx.Defn.jit(fun, @eval), args)
+    assert %Emily.Backend{} = native.data
+    assert native.shape == eval.shape and native.type == eval.type
+    assert Nx.to_binary(native) == Nx.to_binary(eval)
+    native
+  end
+
+  test "cumulative sum/product/max/min along the last axis (forward and reverse)" do
+    x = t([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+
+    for op <- [:cumulative_sum, :cumulative_product, :cumulative_max, :cumulative_min] do
+      equiv(fn x -> apply(Nx, op, [x, [axis: 1]]) end, [x])
+      equiv(fn x -> apply(Nx, op, [x, [axis: 1, reverse: true]]) end, [x])
+    end
+  end
+
+  test "single-axis gather" do
+    x = t([10.0, 20.0, 30.0, 40.0])
+    idx = t([[3], [1], [0]], type: :s64)
+    out = equiv(fn x, i -> Nx.gather(x, i) end, [x, idx])
+    assert Nx.to_flat_list(out) == [40.0, 20.0, 10.0]
+  end
+
+  test "multi-axis gather (the sampling path)" do
+    x = t([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
+    idx = t([[0, 1], [2, 0]], type: :s64)
+    out = equiv(fn x, i -> Nx.gather(x, i) end, [x, idx])
+    assert Nx.to_flat_list(out) == [2.0, 5.0]
+  end
+
+  test "stack along a new axis" do
+    a = t([1.0, 2.0, 3.0])
+    b = t([4.0, 5.0, 6.0])
+    equiv(fn a, b -> Nx.stack([a, b]) end, [a, b])
+    equiv(fn a, b -> Nx.stack([a, b], axis: 1) end, [a, b])
+  end
+end

test/emily/conformance/generation_native_test.exs

@@ -0,0 +1,70 @@
+defmodule Emily.Conformance.GenerationNativeTest do
+  @moduledoc """
+  CM13 — the goal gate: `Bumblebee.Text.generation`'s own decode loop (a
+  `defn while`) compiles **fully native** under `compiler: Emily.Compiler,
+  native: true`, producing token ids bit-identical to the Evaluator path.
+
+  `native_fallback: :raise` makes this a no-fallback gate: every op in the
+  whole generation graph — the transformer forward, the `while` loop, the
+  dynamic KV-cache writes, `cumsum` for position ids, `argmax`/sampling for
+  token selection, threefry for the sampling seed — must lower, or the build
+  raises rather than silently degrading to the evaluator.
+
+  Driven through `build_generate` on small in-vocabulary `input_ids` (not
+  text) so the test needs no tokenizer and stays clear of the OOB-gather
+  pitfall of pairing a tiny-random model with a full-vocab tokenizer.
+  """
+  use ExUnit.Case, async: false
+  @moduletag :conformance
+
+  alias Bumblebee.Text.Generation
+
+  @repo {:hf, "bumblebee-testing/tiny-random-Qwen3ForCausalLM"}
+
+  setup_all do
+    Nx.global_default_backend(Emily.Backend)
+    {:ok, model_info} = Bumblebee.load_model(@repo)
+    {:ok, gen_config} = Bumblebee.load_generation_config(@repo)
+    %{model_info: model_info, gen_config: gen_config}
+  end
+
+  @native [compiler: Emily.Compiler, native: true, native_fallback: :raise]
+  @eval [compiler: Emily.Compiler]
+
+  # Run `model`'s generation through `build_generate` on a fixed in-vocab
+  # prompt, returning the generated token ids.
+  defp generate_ids(%{model: model, params: params, spec: spec}, gen_config, defn_options) do
+    generate = Generation.build_generate(model, spec, gen_config)
+
+    inputs = %{
+      "input_ids" => Nx.tensor([[1, 2, 3, 4, 5, 6]], type: :s64, backend: Emily.Backend),
+      "seed" => Nx.tensor([0], type: :s64, backend: Emily.Backend)
+    }
+
+    Nx.Defn.jit(generate, defn_options).(params, inputs).token_ids
+    |> Nx.to_flat_list()
+  end
+
+  defp configure(gen_config, strategy) do
+    Bumblebee.configure(gen_config,
+      max_new_tokens: 6,
+      strategy: strategy,
+      pad_token_id: 0,
+      eos_token_id: 0
+    )
+  end
+
+  test "greedy generation compiles fully native, bit-identical to the evaluator", ctx do
+    gc = configure(ctx.gen_config, %{type: :greedy_search})
+    native = generate_ids(ctx.model_info, gc, @native)
+    eval = generate_ids(ctx.model_info, gc, @eval)
+    assert native == eval
+  end
+
+  test "multinomial sampling compiles fully native, bit-identical to the evaluator", ctx do
+    gc = configure(ctx.gen_config, %{type: :multinomial_sampling})
+    native = generate_ids(ctx.model_info, gc, @native)
+    eval = generate_ids(ctx.model_info, gc, @eval)
+    assert native == eval
+  end
+end