feat: add fused kv_cache+sdpa (#124)

* feat: metal capture

* 200 tok/s on best case native

* format

* feat: break the 300 tok/s barrier

Author: polvalente

.gitignore

@@ -29,3 +29,6 @@ emlx-*.tar
 # Validation subproject build artifacts (not the mix.lock — that is committed).
 /**/_build/
 /**/deps/
+
+# Local working notes / investigation plans, not part of the public plan history.
+/.work/

emlx/c_src/emlx_compiler.cpp

@@ -1468,8 +1468,15 @@ static const std::unordered_map<std::string, OpFn> op_registry = {
            {}, -std::numeric_limits<float>::infinity(), mask_dtype);
        auto additive = mlx::core::where(keep, zero_val, neginf_val);
 
-       return mlx::core::fast::scaled_dot_product_attention(
+       auto attn_t = mlx::core::fast::scaled_dot_product_attention(
            q, k, v, scale, "array", additive, std::nullopt);
+
+       // Replace NaN with 0 for all-masked rows (softmax(-inf,...,-inf) = NaN
+       // in Flash-Attention when seq_len >= Metal tile size, but semantically
+       // = 0) -- happens for left-padded prefill query rows that have no
+       // causally-visible *and* key_mask-valid key at all.
+       return mlx::core::where(mlx::core::isnan(attn_t),
+                               mlx::core::zeros_like(attn_t), attn_t);
      }},
 
     // sdpa (causal + key_mask, + sinks): operands = [q, k, v, key_mask,
@@ -1504,8 +1511,12 @@ static const std::unordered_map<std::string, OpFn> op_registry = {
            {}, -std::numeric_limits<float>::infinity(), mask_dtype);
        auto additive = mlx::core::where(keep, zero_val, neginf_val);
 
-       return mlx::core::fast::scaled_dot_product_attention(
+       auto attn_t = mlx::core::fast::scaled_dot_product_attention(
            q, k, v, scale, "array", additive, sinks);
+
+       // See "fast_sdpa_causal_key_masked" above: NaN-guard all-masked rows.
+       return mlx::core::where(mlx::core::isnan(attn_t),
+                               mlx::core::zeros_like(attn_t), attn_t);
      }},
 
     // rope (scalar offset): operands = [a];
@@ -1939,6 +1950,78 @@ static const std::unordered_map<std::string, MultiOpFn> multi_op_registry = {
      [](const auto &ops, const auto &) -> std::vector<mlx::core::array> {
        return contiguous_all(mlx::core::linalg::lu(ops[0], k_linalg_cpu));
      }},
+    // The head-transpose needed for `fast::scaled_dot_product_attention`
+    // (which requires {B, N, T, D}) is done internally on q and on the
+    // *already-written* k_upd/v_upd -- this is the same single full-cache
+    // transpose the unfused path already pays for (its own
+    // `build_sdpa_layer` transposes the post-update K/V cache before
+    // calling SDPA); fusing does not add a second one, since k_upd/v_upd
+    // are returned pre-transpose.
+    {"kv_cache_sdpa_update",
+     [](const auto &ops, const auto &attrs) -> std::vector<mlx::core::array> {
+       const auto &q = ops[0];
+       const auto &new_k = ops[1];
+       const auto &new_v = ops[2];
+       const auto &k_cache = ops[3];
+       const auto &v_cache = ops[4];
+       const auto &offset = ops[5];
+       const auto &key_mask = ops[6];
+       float scale = attr_to_float(attrs[0]);
+       int kv_offset = static_cast<int>(attrs[1]);
+
+       int T_max = static_cast<int>(k_cache.shape(1));
+       int T_q_write = static_cast<int>(new_k.shape(1));
+
+       // Dynamic write-start {0, clamp(offset, 0, T_max-T_q_write), 0, 0} --
+       // same clamp-and-concatenate construction as op_registry["put_slice"]
+       // (T_max-T_q_write, not T_max-1, so the write itself never runs past
+       // the cache's last valid row when T_q_write > 1, e.g. prefill).
+       auto s = mlx::core::reshape(mlx::core::astype(offset, mlx::core::int32), {1});
+       auto max_s = mlx::core::full({1}, T_max - T_q_write, mlx::core::int32);
+       auto zero1 = mlx::core::zeros({1}, mlx::core::int32);
+       auto offset_clamped = mlx::core::minimum(mlx::core::maximum(s, zero1), max_s);
+       auto start = mlx::core::concatenate({zero1, offset_clamped, zero1, zero1}, 0);
+
+       std::vector<int> all_axes = {0, 1, 2, 3};
+       auto k_upd = mlx::core::slice_update(k_cache, new_k, start, all_axes);
+       auto v_upd = mlx::core::slice_update(v_cache, new_v, start, all_axes);
+
+       // Head-transpose for SDPA: {B, T, N, D} -> {B, N, T, D}. Mirrors
+       // build_sdpa_layer's own Nx.transpose(axes: [0, 2, 1, 3]) exactly.
+       std::vector<int> head_transpose = {0, 2, 1, 3};
+       auto q_t = mlx::core::transpose(q, head_transpose);
+       auto k_t = mlx::core::transpose(k_upd, head_transpose);
+       auto v_t = mlx::core::transpose(v_upd, head_transpose);
+
+       // Combined causal + key_mask additive mask -- verbatim copy of
+       // "fast_sdpa_causal_key_masked"'s formula above (T_kv = T_max here).
+       auto km = mlx::core::reshape(key_mask, {key_mask.shape(0), 1, 1, key_mask.shape(1)});
+       int T_q = q_t.shape(2);
+       auto row = mlx::core::reshape(mlx::core::arange(T_q, mlx::core::int32), {1, 1, T_q, 1});
+       auto col = mlx::core::reshape(mlx::core::arange(T_max, mlx::core::int32), {1, 1, 1, T_max});
+       auto causal_bool = mlx::core::less_equal(
+           col, mlx::core::add(row, mlx::core::array(kv_offset, mlx::core::int32)));
+       auto keep = mlx::core::logical_and(km, causal_bool);
+       auto mask_dtype = q_t.dtype();
+       auto zero_val = mlx::core::zeros({}, mask_dtype);
+       auto neginf_val =
+           mlx::core::full({}, -std::numeric_limits<float>::infinity(), mask_dtype);
+       auto additive = mlx::core::where(keep, zero_val, neginf_val);
+
+       auto attn_out = mlx::core::fast::scaled_dot_product_attention(
+           q_t, k_t, v_t, scale, "array", additive, std::nullopt);
+
+       // Replace NaN with 0 for all-masked rows (softmax(-inf,...,-inf) = NaN
+       // in Flash-Attention when seq_len >= Metal tile size, but semantically
+       // = 0) -- happens for left-padded prefill query rows that have no
+       // causally-visible *and* key_mask-valid key at all. See
+       // "fast_sdpa_causal_key_masked" above for the unfused version of this
+       // same guard.
+       attn_out = mlx::core::where(mlx::core::isnan(attn_out),
+                                   mlx::core::zeros_like(attn_out), attn_out);
+
+       return {attn_out, k_upd, v_upd};
+     }},
 };
 
 // ── interpret_instructions ────────────────────────────────────────────────

emlx/c_src/emlx_fast.cpp

@@ -339,8 +339,14 @@ mlx::core::array fast_sdpa_causal_key_masked_impl(
     auto neginf_val = full({}, -std::numeric_limits<float>::infinity(), mask_dtype);
     auto additive = where(keep, zero_val, neginf_val);
 
-    return fast::scaled_dot_product_attention(
+    auto attn_t = fast::scaled_dot_product_attention(
         *q, *k, *v, (float)scale, "array", additive, sinks_opt, device);
+
+    // Replace NaN with 0 for all-masked rows (softmax(-inf,...,-inf) = NaN in
+    // Flash-Attention when seq_len >= Metal tile size, but semantically = 0).
+    // This happens for left-padded prefill rows, where a padding query
+    // position has no causally-visible *and* key_mask-valid keys at all.
+    return where(isnan(attn_t), zeros_like(attn_t, device), attn_t, device);
   }
 }
 FINE_ASYNC_NIF(fast_sdpa_causal_key_masked)

emlx/c_src/emlx_nif.cpp

@@ -1335,6 +1335,33 @@ NIF(set_cache_limit) {
   return nx::nif::ok(env, enif_make_uint64(env, prev));
 }
 
+// Starts a Metal GPU frame capture, writing the trace to `path` (a
+// `.gputrace` bundle). Requires `MTL_CAPTURE_ENABLED=1` in the process
+// environment *before* the BEAM (and therefore the Metal device) started —
+// Apple's capture gate is read once at process startup and cannot be
+// enabled lazily. On MLX 0.31.2, missing this precondition throws (caught
+// by CATCH() below and surfaced as a normal EMLX.NIFError), not a silent
+// no-op. See EMLX.metal_start_capture/1's moduledoc for details.
+NIF(metal_start_capture) {
+  std::string path;
+  if (!nx::nif::get(env, argv[0], path)) {
+    return nx::nif::error(env, "Unable to get path param.");
+  }
+  try {
+    mlx::core::metal::start_capture(path);
+    return nx::nif::ok(env);
+  }
+  CATCH()
+}
+
+NIF(metal_stop_capture) {
+  try {
+    mlx::core::metal::stop_capture();
+    return nx::nif::ok(env);
+  }
+  CATCH()
+}
+
 NIF(strides) {
   TENSOR_PARAM(0, t);
 
@@ -1828,8 +1855,12 @@ NIF(eval_program) { return emlx::native::eval_program(env, argc, argv); }
 ASYNC_NIF(eval_program)
 
 static ErlNifFunc nif_funcs[] = {
-    {"eval", 2, eval_async, ERL_NIF_DIRTY_JOB_IO_BOUND},
-    {"eval_many", 2, eval_many_async, ERL_NIF_DIRTY_JOB_IO_BOUND},
+    // No dirty-scheduler flag: eval/eval_many post to a dedicated Worker
+    // OS thread via async_dispatch (emlx_async.hpp) and return immediately —
+    // the calling BEAM scheduler never blocks on the actual MLX work, same
+    // as the ~150 sibling ASYNC_NIF ops below (item 3.8 fold-in).
+    {"eval", 2, eval_async},
+    {"eval_many", 2, eval_many_async},
     {"to_device", 3, to_device_async},
     {"to_blob", 2, to_blob_async},
     {"to_blob", 3, to_blob_async},
@@ -1979,6 +2010,8 @@ static ErlNifFunc nif_funcs[] = {
     {"reset_peak_memory", 0, reset_peak_memory},
     {"set_memory_limit", 1, set_memory_limit},
     {"set_cache_limit", 1, set_cache_limit},
+    {"metal_start_capture", 1, metal_start_capture},
+    {"metal_stop_capture", 0, metal_stop_capture},
 
     // ── Worker control NIFs.
     {"command_queue_new", 1, command_queue_new},
@@ -2006,9 +2039,11 @@ static ErlNifFunc nif_funcs[] = {
     {"kv_cache_sdpa_update", 9, kv_cache_sdpa_update_async},
 
     // ── Native compiler NIFs.
-    {"compile_program", 2, emlx::native::compile_program_async,
-     ERL_NIF_DIRTY_JOB_CPU_BOUND},
-    {"eval_program", 3, eval_program_async, ERL_NIF_DIRTY_JOB_IO_BOUND},
+    // No dirty-scheduler flag on either: both are FINE_ASYNC_NIF/ASYNC_NIF-
+    // wrapped (post to a Worker OS thread, return immediately) — same
+    // reasoning as eval/eval_many above (item 3.8 fold-in).
+    {"compile_program", 2, emlx::native::compile_program_async},
+    {"eval_program", 3, eval_program_async},
     // resolve_runtime_call is NOT worker-routed: it only decodes a reply and
     // memcpy's it into pre-registered buffers/notifies a condvar — no MLX
     // graph work, so it can run directly on the calling BEAM scheduler.

emlx/lib/emlx.ex

@@ -2406,4 +2406,38 @@ defmodule EMLX do
   def set_cache_limit(limit) when is_integer(limit) and limit >= 0 do
     EMLX.NIF.set_cache_limit(limit) |> unwrap!()
   end
+
+  @doc """
+  Starts a Metal GPU frame capture, writing the trace to `path` (a
+  `.gputrace` bundle openable in Xcode's GPU Debugger via File → Open).
+
+  Requires `MTL_CAPTURE_ENABLED=1` to be set in the process environment
+  *before* the BEAM starts — Apple's Metal framework reads this gate once at
+  process startup (macOS 14+), and setting it later (e.g. from within
+  Elixir) has no effect. On MLX 0.31.2, forgetting to set it raises a clear
+  `EMLX.NIFError` (`"[metal::start_capture] Failed to start: Capture layer
+  is not inserted."`) rather than silently no-op'ing — confirmed by direct
+  test, contradicting this NIF's original design assumption that
+  `start_capture`/`stop_capture` never throw on ordinary misconfiguration.
+  Still worth confirming `path` exists and is non-empty after
+  `metal_stop_capture/0` as a belt-and-suspenders check, since a future MLX
+  version could change this behavior back to a silent no-op.
+
+  ## Examples
+
+      # export MTL_CAPTURE_ENABLED=1   (before starting the BEAM)
+      EMLX.metal_start_capture("/tmp/trace.gputrace")
+      # ... run some MLX work ...
+      EMLX.metal_stop_capture()
+  """
+  def metal_start_capture(path) when is_binary(path) do
+    EMLX.NIF.metal_start_capture(path) |> unwrap!()
+  end
+
+  @doc """
+  Stops a Metal GPU frame capture started with `metal_start_capture/1`.
+
+  See `metal_start_capture/1` for the `MTL_CAPTURE_ENABLED` precondition.
+  """
+  def metal_stop_capture, do: EMLX.NIF.metal_stop_capture() |> unwrap!()
 end

emlx/lib/emlx/fast.ex

@@ -56,6 +56,9 @@ defmodule EMLX.Fast do
   - `scaled_dot_product_attention_causal_key_masked/5` — causal SDPA; checks key_mask at C++ level, fast-paths to pure causal when all-ones
   - `scaled_dot_product_attention_causal_key_masked/6` — same, plus `opts` (`:sinks`)
   - `swiglu/2` — fused SwiGLU: `silu(gate) * up`
+  - `kv_cache_sdpa_update/8` — fused decode-time KV-cache write + causal SDPA
+    read, 3 outputs (`{attn_out, k_cache_updated, v_cache_updated}`) —
+    inference/decode-loop only, no grad support
   - `einsum/2` — variadic-operand Einstein summation (`mlx::core::einsum`);
     **eager-only, not defn-safe** (see its docs)
 
@@ -456,6 +459,133 @@ defmodule EMLX.Fast do
     if Nx.type(out) != dtype, do: Nx.as_type(out, dtype), else: out
   end
 
+  # ── Fused decode-time KV-cache update + SDPA ────────────────────────────────
+  #
+  # SPIKE for `.work/performance_gap/01-attention-block-fusion-via-compiler.md`
+  # (Procedure step 3), built after that stage's Results already concluded the
+  # write-donation fix has no plausible mechanism to help (donation cannot
+  # survive `:while`-carry storage regardless of which primitive performs the
+  # write — see that doc's "Fix surface"/"Recommendation" sections) — kept as
+  # a real, working artifact per explicit instruction to build it anyway, not
+  # because a throughput win is expected. Unlike every other function in this
+  # module, this one does NOT use the single-output `:__EMLX__` metadata
+  # mechanism (`emlx_metadata/4` above) — that mechanism produces exactly one
+  # physical `ref` (see `EMLX.Native.Expr`'s `:metadata` expand_node clause),
+  # and this op has 3 outputs (attn_out, k_upd, v_upd). Instead it uses
+  # `Nx.runtime_call/4`'s `:__emlx_native_multi_op__` escape hatch (see that
+  # same module's `:runtime_call` expand_node clause), which reuses
+  # `Nx.runtime_call`'s existing multi-output container/`:elem` machinery
+  # while still lowering to one `multi_op_registry` C++ instruction — no BEAM
+  # round-trip, same as every other op here, just without the extra metadata
+  # wrapper layer (which the multi-output case can't use).
+  #
+  # No `Nx.Defn.grad` support — this op only targets the (non-differentiable,
+  # forward-only) decode/generation loop.
+
+  @doc """
+  Fused decode-time KV-cache write (2× `put_slice`) + causal-masked SDPA read,
+  as one native op — the traced-graph analog of the hand-written
+  `EMLX.kv_cache_sdpa_update/7` eager NIF, but operating on the *full*
+  fixed-size preallocated cache with a *dynamic* (array-valued) `offset`
+  instead of a host int, so it can sit inside a compiled `:while` decode
+  loop without forcing a host readback per iteration.
+
+  All tensors use **Bumblebee-native, heads-NOT-transposed** layout (matching
+  `Bumblebee.Layers.Decoder.attention_cache/4`'s `{batch, seq, heads, head_dim}`
+  cache shape and `EMLXAxon`'s `build_sdpa_layer/5` pre-transpose inputs) — the
+  head-transpose SDPA itself needs is done internally, once, on `q` and on the
+  already-written `k_upd`/`v_upd`; the returned caches stay untransposed so
+  they can feed a `:while` carry directly with no extra transpose.
+
+  - `q`        — `{B, T_q,  N_q,  D}` (T_q = 1 for decode; not enforced)
+  - `new_k`    — `{B, T_q,  N_kv, D}`
+  - `new_v`    — `{B, T_q,  N_kv, D}`
+  - `k_cache`  — `{B, T_max, N_kv, D}`
+  - `v_cache`  — `{B, T_max, N_kv, D}`
+  - `offset`   — scalar/`{1}` int tensor — write position along the `T_max` axis
+  - `key_mask` — `{B, T_max}` — 1 = valid (readable), 0 = masked
+  - `scale`    — pre-computed scalar
+
+  Returns `{attn_out, k_cache_updated, v_cache_updated}` — `attn_out` is
+  **head-transposed**, `{B, N_q, T_q, D}` (matching what the unfused
+  `scaled_dot_product_attention_causal_key_masked/6` metadata node itself
+  already returns; callers transpose it back same as today), while the
+  updated caches keep `k_cache`/`v_cache`'s own (untransposed) shape.
+
+  Eager: composes plain `Nx.put_slice` ×2 (clamped host offset) +
+  `scaled_dot_product_attention_causal_key_masked/6` — deliberately
+  independent of the traced path's C++ formula below, so it doubles as an
+  correctness oracle for the native `multi_op_registry["kv_cache_sdpa_update"]`
+  kernel (same combined causal+key_mask formula, reimplemented once in
+  Elixir/Nx and once in C++).
+
+  Traced: lowers to one `multi_op_registry["kv_cache_sdpa_update"]`
+  (`emlx_compiler.cpp`) native IR instruction.
+  """
+  deftransform kv_cache_sdpa_update(q, new_k, new_v, k_cache, v_cache, offset, key_mask, scale) do
+    t_q = elem(Nx.shape(q), 1)
+    t_max = elem(Nx.shape(k_cache), 1)
+    kv_offset = if t_q == 1, do: t_max - 1, else: 0
+
+    offset = offset |> Nx.reshape({1}) |> Nx.as_type({:s, 32})
+    args = {q, new_k, new_v, k_cache, v_cache, offset, key_mask}
+
+    if traced?([q, new_k, new_v, k_cache, v_cache, offset, key_mask]) do
+      {qb, qt, qn, qd} = Nx.shape(q)
+      attn_template = Nx.template({qb, qn, qt, qd}, Nx.type(q))
+
+      out_template = {attn_template, Nx.to_template(k_cache), Nx.to_template(v_cache)}
+
+      Nx.runtime_call(
+        out_template,
+        args,
+        [
+          scale: scale,
+          kv_offset: kv_offset,
+          __emlx_native_multi_op__:
+            {:kv_cache_sdpa_update, [NativeExpr.f64_bits(scale), kv_offset]}
+        ],
+        &kv_cache_sdpa_update_callback/2
+      )
+    else
+      kv_cache_sdpa_update_callback(args, scale: scale, kv_offset: kv_offset)
+    end
+  end
+
+  @doc false
+  def kv_cache_sdpa_update_callback(
+        {%Nx.Tensor{} = q, %Nx.Tensor{} = new_k, %Nx.Tensor{} = new_v, %Nx.Tensor{} = k_cache,
+         %Nx.Tensor{} = v_cache, %Nx.Tensor{} = offset, %Nx.Tensor{} = key_mask},
+        opts
+      ) do
+    scale = opts[:scale]
+    kv_offset = opts[:kv_offset]
+    t_max = elem(Nx.shape(k_cache), 1)
+
+    start =
+      offset
+      |> Nx.reshape({})
+      |> Nx.as_type({:s, 64})
+      |> Nx.to_number()
+      |> max(0)
+      |> min(t_max - 1)
+
+    k_upd = Nx.put_slice(k_cache, [0, start, 0, 0], new_k)
+    v_upd = Nx.put_slice(v_cache, [0, start, 0, 0], new_v)
+
+    q_t = Nx.transpose(q, axes: [0, 2, 1, 3])
+    k_t = Nx.transpose(k_upd, axes: [0, 2, 1, 3])
+    v_t = Nx.transpose(v_upd, axes: [0, 2, 1, 3])
+
+    attn_out =
+      sdpa_causal_key_masked_callback({q_t, k_t, v_t, key_mask},
+        scale: scale,
+        kv_offset: kv_offset
+      )
+
+    {attn_out, k_upd, v_upd}
+  end
+
   # ── RoPE ────────────────────────────────────────────────────────────────────
 
   @doc """