Merge pull request #145 from ausimian/feat/async-eval

feat: async_eval and integer-offset RoPE for non-blocking decode

Author: ausimian

RELEASE.md

@@ -0,0 +1,16 @@
+### Added
+
+- `Emily.async_eval/1` (and `Emily.Native.async_eval/2`) schedule evaluation of
+  one or more lazy graphs **without blocking on the GPU**, wrapping
+  `mlx::core::async_eval`. The work is handed to the device's command queue and
+  the call returns as soon as it is enqueued — not when it finishes. This lets a
+  caller keep dispatching the next step's ops while the device computes the
+  current one (e.g. an autoregressive decode loop), blocking only when a value
+  is actually read back on the host via `to_binary/1` / `eval/1`. Pass every
+  output of a step (logits plus all KV-cache buffers) in one call.
+- `Emily.Native.fast_rope_int/8` — RoPE with an **integer** absolute-position
+  `offset` (routing to MLX's int-offset `rope` overload), for incremental decode
+  where the caller tracks position host-side. Complements the existing
+  tensor-offset `fast_rope/8`. Note: feed the kernel the 4-D
+  `{batch, heads, seq, head_dim}` layout — in 3-D, MLX 0.31 mis-rotates
+  single-token (`seq == 1`) inputs.

bench/native/compile_microbench.cpp

@@ -39,6 +39,7 @@ namespace mx = mlx::core;
 // 1000 iterations is fast to run yet long enough that fusion matters.
 constexpr int kBatch = 1;
 static int kSeq = 128;                    // overridable via --seq
+static int kLayers = 1;                   // overridable via --layers (stack the block)
 constexpr int kHidden = 1024;
 constexpr int kHeads = 16;
 constexpr int kHeadDim = 64;             // kHeads * kHeadDim = 1024 = kHidden
@@ -129,6 +130,18 @@ static std::vector<mx::array> block(const std::vector<mx::array>& in) {
   return {out};
 }
 
+// Stack the block kLayers times (reusing weights — fine for timing). This
+// reaches forward-scale op counts (~15 ops/block × 48 ≈ 720, matching the
+// Gemma decode forward's ~750) so dispatch (build+encode) is large enough to
+// dominate, isolating whether mx::compile amortizes it.
+static std::vector<mx::array> multi_block(const std::vector<mx::array>& in) {
+  std::vector<mx::array> cur = in;
+  for (int l = 0; l < kLayers; ++l) {
+    cur[0] = block(cur)[0];
+  }
+  return {cur[0]};
+}
+
 // ---------------------------------------------------------------------
 
 struct Stats {
@@ -200,20 +213,48 @@ static Stats time_runs(Fn&& fn, int warmup, int iters) {
   return summarise(samples);
 }
 
+// Dispatch-only timing: time [build graph + async_eval] WITHOUT waiting for
+// the GPU (synchronize happens outside the timed region). This isolates the
+// host-side dispatch cost (graph build + MLX encode/schedule) — the ~68 ms/tok
+// that dominates Gemma decode — from GPU compute. The question: does compile
+// shrink THIS, not the GPU.
+template <typename Fn>
+static Stats time_dispatch(Fn&& fn, int warmup, int iters) {
+  for (int i = 0; i < warmup; ++i) {
+    auto out = fn();
+    mx::async_eval(out);
+    mx::synchronize();
+  }
+
+  std::vector<double> samples;
+  samples.reserve(iters);
+
+  for (int i = 0; i < iters; ++i) {
+    auto start = std::chrono::high_resolution_clock::now();
+    auto out = fn();
+    mx::async_eval(out);  // schedule; returns after encode, before GPU completes
+    auto end = std::chrono::high_resolution_clock::now();
+    mx::synchronize();    // drain GPU OUTSIDE the timed region
+    samples.push_back(std::chrono::duration<double, std::milli>(end - start).count());
+  }
+
+  return summarise(samples);
+}
+
 static void run_on_device(mx::Device::DeviceType dev_type, const char* label,
                           int warmup, int iters) {
   mx::set_default_device(mx::Device(dev_type));
 
   auto inputs = make_inputs();
 
-  // Uncompiled baseline: call block() directly on each iteration.
+  // Uncompiled baseline: rebuild the (kLayers-stacked) graph each iteration.
   auto uncompiled_fn = [&inputs]() -> std::vector<mx::array> {
-    return block(inputs);
+    return multi_block(inputs);
   };
 
-  // Compiled: wrap block in mx::compile. First call of the returned
-  // closure does the trace; subsequent calls hit the fused tape.
-  auto compiled_closure = mx::compile(block, /*shapeless=*/false);
+  // Compiled: wrap in mx::compile. First call traces; subsequent calls hit the
+  // cached tape (skipping the graph build/simplify).
+  auto compiled_closure = mx::compile(multi_block, /*shapeless=*/false);
   auto compiled_fn = [&compiled_closure, &inputs]() -> std::vector<mx::array> {
     return compiled_closure(inputs);
   };
@@ -236,8 +277,17 @@ static void run_on_device(mx::Device::DeviceType dev_type, const char* label,
   std::printf("  speedup (median) = %.3fx    (min) = %.3fx\n",
               median_speedup, min_speedup);
 
-  bool passes = median_speedup >= 1.20;
-  std::printf("  gate (>=1.20x on median): %s\n", passes ? "PASS" : "FAIL");
+  // Dispatch-only: the host-side build+encode cost (no GPU wait) — the thing
+  // that dominates Gemma decode. Does compile amortize it?
+  auto unc_disp = time_dispatch(uncompiled_fn, warmup, iters);
+  auto cmp_disp = time_dispatch(compiled_fn, warmup, iters);
+  std::printf("  [dispatch-only, no GPU wait]\n");
+  std::printf("    uncompiled  min=%.3f ms  median=%.3f ms\n",
+              unc_disp.min_ms, unc_disp.median_ms);
+  std::printf("    compiled    min=%.3f ms  median=%.3f ms\n",
+              cmp_disp.min_ms, cmp_disp.median_ms);
+  std::printf("    dispatch speedup (median) = %.3fx\n",
+              unc_disp.median_ms / cmp_disp.median_ms);
 }
 
 // ---------------------------------------------------------------------
@@ -310,6 +360,8 @@ int main(int argc, char** argv) {
       iters = std::atoi(argv[++i]);
     } else if (std::strcmp(argv[i], "--seq") == 0 && i + 1 < argc) {
       kSeq = std::atoi(argv[++i]);
+    } else if (std::strcmp(argv[i], "--layers") == 0 && i + 1 < argc) {
+      kLayers = std::atoi(argv[++i]);
     }
   }
 

c_src/emily_nif.cpp

@@ -144,4 +144,26 @@ fine::Term eval_nif(ErlNifEnv *env,
 }
 FINE_NIF(eval_nif, 0);
 
+// async_eval_nif/2 — schedule evaluation of the lazy graphs rooted at
+// these tensors *without* blocking on the GPU. Unlike eval_nif (which
+// calls mx::eval and waits for the GPU before replying), this calls
+// mx::async_eval: the worker hands the graphs to the Metal command queue
+// and the NIF replies {ref, {:ok, :ok}} as soon as the work is *queued*,
+// not finished. Callers drive an autoregressive loop by scheduling the
+// next step's outputs here and only blocking (via to_binary/eval) on the
+// value they actually need on the host — so dispatch of step N+1 overlaps
+// the GPU compute of step N. Takes a list so one call schedules a step's
+// logits + every KV-cache buffer together.
+fine::Term async_eval_nif(ErlNifEnv *env,
+                          fine::ResourcePtr<WorkerThread> w,
+                          std::vector<fine::ResourcePtr<Tensor>> tensors) {
+  return emily::async_reply(
+      env, w,
+      [tensors = std::move(tensors)](mx::Stream &, ErlNifEnv *msg_env) {
+        mx::async_eval(emily::unwrap_all(tensors));
+        return fine::encode(msg_env, emily::atoms::ok);
+      });
+}
+FINE_NIF(async_eval_nif, 0);
+
 FINE_INIT("Elixir.Emily.Native");

c_src/ops/fast.cpp

@@ -84,6 +84,33 @@ fine::Term fast_rope_nif(
 }
 FINE_NIF(fast_rope_nif, 0);
 
+// Int-offset variant. `offset` is a plain integer absolute position (the
+// caller tracks it host-side). Uses MLX's int-offset rope overload, which is
+// correct for single-token (seq == 1) inputs — the array-offset overload
+// (fast_rope_nif) mis-rotates seq == 1, breaking incremental decode. `base`
+// is nullopt when `freqs` is supplied.
+fine::Term fast_rope_int_nif(
+    ErlNifEnv *env,
+    fine::ResourcePtr<WorkerThread> w,
+    fine::ResourcePtr<Tensor> x,
+    int64_t dims,
+    bool traditional,
+    std::optional<double> base,
+    double scale,
+    int64_t offset,
+    std::optional<fine::ResourcePtr<Tensor>> freqs) {
+  return async_encoded(env, w,
+      [x = std::move(x), dims, traditional, base, scale, offset,
+       freqs = std::move(freqs)](mx::Stream &s) {
+        std::optional<float> base_f;
+        if (base) base_f = static_cast<float>(*base);
+        return wrap(mx::fast::rope(x->array, emily::checked_int(dims, "dims"), traditional,
+                                   base_f, static_cast<float>(scale),
+                                   emily::checked_int(offset, "offset"), opt_array(freqs), s));
+      });
+}
+FINE_NIF(fast_rope_int_nif, 0);
+
 fine::Term fast_scaled_dot_product_attention_nif(
     ErlNifEnv *env,
     fine::ResourcePtr<WorkerThread> w,

lib/emily.ex

@@ -187,4 +187,29 @@ defmodule Emily do
       {Native.eval(Emily.MlxStream.default_worker(), tensor), %{}}
     end)
   end
+
+  @doc """
+  Schedule evaluation of `tensors` *without* blocking on the GPU.
+
+  Hands the lazy graphs to the device's command queue and returns as soon
+  as the work is enqueued — not when it finishes (contrast `eval/1`, which
+  waits). This lets a caller keep dispatching the next step's ops while the
+  GPU computes the current one; block only when you actually read a value
+  on the host (`to_binary/1` / `eval/1`). Pass every output of a step
+  (e.g. logits plus all KV-cache buffers) in one call.
+
+  ## Examples
+
+      iex> a = Nx.tensor([1.0, 2.0, 3.0], backend: Emily.Backend)
+      iex> b = a |> Nx.multiply(2.0) |> Nx.add(1.0)
+      iex> Emily.async_eval([b.data.ref])
+      :ok
+
+  """
+  @spec async_eval([t()]) :: :ok
+  def async_eval(tensors) when is_list(tensors) do
+    :telemetry.span([:emily, :async_eval], %{}, fn ->
+      {Native.async_eval(Emily.MlxStream.default_worker(), tensors), %{}}
+    end)
+  end
 end

lib/emily/native.ex

@@ -83,6 +83,16 @@ defmodule Emily.Native do
   @spec eval(worker(), tensor()) :: :ok
   def eval(w, tensor), do: await(eval_nif(w, tensor))
 
+  @doc false
+  @spec async_eval_nif(worker(), [tensor()]) :: reference()
+  def async_eval_nif(_w, _tensors), do: nif()
+
+  # Schedule (non-blocking) evaluation of several lazy graphs at once: the
+  # GPU work is queued and this returns as soon as it's enqueued, not when
+  # it finishes. See `async_eval_nif` in c_src/emily_nif.cpp.
+  @spec async_eval(worker(), [tensor()]) :: :ok
+  def async_eval(w, tensors), do: await(async_eval_nif(w, tensors))
+
   # --- Worker ------------------------------------------------------
 
   # Default per-worker queue depth. Each op is awaited synchronously, so
@@ -771,6 +781,45 @@ defmodule Emily.Native do
         )
       )
 
+  @doc false
+  @spec fast_rope_int_nif(
+          worker(),
+          tensor(),
+          integer(),
+          boolean(),
+          float() | nil,
+          float(),
+          integer(),
+          tensor() | nil
+        ) :: reference()
+  def fast_rope_int_nif(_w, _x, _dims, _traditional, _base, _scale, _offset, _freqs), do: nif()
+
+  # Like `fast_rope/8` but `offset` is a plain integer (absolute position),
+  # routing to MLX's int-offset rope overload. Correct for single-token decode
+  # (seq == 1), where the tensor-offset `fast_rope/8` mis-rotates.
+  @spec fast_rope_int(
+          worker(),
+          tensor(),
+          integer(),
+          boolean(),
+          float() | nil,
+          float(),
+          integer(),
+          tensor() | nil
+        ) :: tensor()
+  def fast_rope_int(w, x, dims, traditional, base, scale, offset, freqs),
+    do:
+      await(
+        fast_rope_int_nif(w, x, dims, traditional, base, scale, offset, freqs),
+        native_context(:fast_rope_int, w, [x: x, freqs: freqs],
+          dims: dims,
+          traditional: traditional,
+          base: base,
+          scale: scale,
+          offset: offset
+        )
+      )
+
   @doc false
   @spec fast_scaled_dot_product_attention_nif(
           worker(),

test/emily/fast/rope_test.exs

@@ -127,4 +127,48 @@ defmodule Emily.Fast.RoPETest do
       assert_close(fused, expected, tol: @f32_tol)
     end
   end
+
+  describe "fast_rope_int/8 (integer offset, incremental decode)" do
+    # Incremental decode ropes one token at a time at an integer absolute
+    # position. It must be correct for seq == 1 — fed the 4D
+    # {1, heads, seq, head_dim} layout (in 3D, MLX 0.31 mis-rotates seq == 1).
+    # This pins the decode-vs-prefill consistency a generation loop relies on:
+    # a single token roped at offset k must equal position k of the full
+    # sequence roped at offset 0.
+    test "single-token rope at offset k equals position k of a full-sequence rope" do
+      dims = 32
+      heads = 2
+      seq = 8
+      k = 5
+      w = Emily.MlxStream.default_worker()
+
+      freqs =
+        Nx.iota({div(dims, 2)}, type: :f32)
+        |> Nx.multiply(2.0)
+        |> Nx.divide(dims)
+        |> then(&Nx.pow(10_000.0, &1))
+
+      x = Nx.iota({1, heads, seq, dims}, type: :f32) |> Nx.divide(100)
+
+      wrap = fn r, shape ->
+        %Nx.Tensor{
+          data: %Emily.Backend{ref: r},
+          shape: shape,
+          type: Emily.Native.dtype(r),
+          names: List.duplicate(nil, tuple_size(shape))
+        }
+      end
+
+      rope_int = fn t, offset, shape ->
+        Emily.Native.fast_rope_int(w, t.data.ref, dims, false, nil, 1.0, offset, freqs.data.ref)
+        |> wrap.(shape)
+      end
+
+      full = rope_int.(x, 0, {1, heads, seq, dims})
+      tok = Nx.slice(x, [0, 0, k, 0], [1, heads, 1, dims])
+      single = rope_int.(tok, k, {1, heads, 1, dims})
+
+      assert_close(single, Nx.slice(full, [0, 0, k, 0], [1, heads, 1, dims]), tol: @f32_tol)
+    end
+  end
 end

test/emily_test.exs

@@ -87,6 +87,24 @@ defmodule EmilyTest do
       assert Emily.to_binary(t) == bin
     end
 
+    test "async_eval/1 schedules without blocking and yields the same values as eager" do
+      a = Nx.tensor([1.0, 2.0, 3.0, 4.0], backend: Emily.Backend)
+      b = a |> Nx.multiply(2.0) |> Nx.add(1.0)
+      c = Nx.subtract(b, a)
+
+      # Non-blocking schedule of several outputs at once returns :ok.
+      assert :ok = Emily.async_eval([b.data.ref, c.data.ref])
+
+      # The later host read (to_binary, via Nx) sees the awaited values —
+      # identical to the eager path, only the eval timing differs.
+      assert Nx.to_flat_list(b) == [3.0, 5.0, 7.0, 9.0]
+      assert Nx.to_flat_list(c) == [2.0, 3.0, 4.0, 5.0]
+    end
+
+    test "async_eval/1 of an empty list is a no-op" do
+      assert :ok = Emily.async_eval([])
+    end
+
     test "tensor survives GC as long as a ref is held" do
       bin = <<42.0::float-32-native>>
       t = Emily.from_binary(bin, [1], {:f, 32})