Merge pull request #25 from ausimian/m14-stream-per-process

M14: stream-per-process concurrent inference

Author: ausimian

PLAN.md

@@ -26,7 +26,9 @@ checklist so future-us understands the trade-offs.
   scope from M9; mixed-precision master weights: in scope from M16.)
 - Drop-in replacement for EMLX. We borrow where it's clearly right, but
   we're not constrained by its API.
-- `Emily.Stream` as a public API — MLX streams stay internal in v1.
+- `Emily.Stream` beyond the narrow `with_stream/2` + `new/1` +
+  `synchronize/1` surface M14 introduces for the documented
+  "big model, multi-process serving" pattern.
 
 ## Architecture
 
@@ -635,7 +637,7 @@ checked in alongside the goldens.
   macOS (CPU) or Linux+CUDA. Emits a complete `ExlaGoldenData` module:
   `elixir bench/exla_golden_gen.exs`.
 
-### M14 — Serving concurrency cookbook + pooled-serving helper
+### M14 — Serving concurrency: stream-per-process + cookbook
 
 `Emily.Compiler.__partitions_options__/1` raises on
 `max_concurrency > 1` — correct (Metal isn't safe for concurrent
@@ -644,36 +646,33 @@ single Emily-backed `Nx.Serving` cannot scale past one concurrent
 request. Production users will hit this in week one. M14 stops being
 silent about it and ships a tested pattern.
 
-- **`Emily.Serving.start_pool/2`**: helper that starts N
-  `Nx.Serving` instances behind a pool (poolboy or a hand-rolled
-  Registry round-robin; pick the lighter dep). Each Serving runs in
-  its own worker, dispatching requests in parallel across pool
-  members rather than within one member.
-- **Documented pattern**: "for K concurrent inference requests, start
-  K servings". Trade-off: each pool member loads its own weights —
-  fine for small models, painful for Qwen3-7B+.
-- **Alternative for large models**: stream-per-process via MLX's
-  `mx::scheduler::new_stream`. Expose `Native.set_default_stream/1`
-  and `Emily.Stream.with_stream/2`; multi-process serving with one
-  shared model + per-process MLX streams becomes the documented
-  "big model" path. (Promotes streams from internal-only — see
+- **Stream-per-process**: the primary deliverable. Expose
+  `Native.set_default_stream/1` and `Emily.Stream.with_stream/2`
+  via MLX's `mx::scheduler::new_stream`. Each process gets its own
+  Metal command queue; one shared model, per-process streams, no
+  weight duplication. (Promotes streams from internal-only — see
   Project Decisions — to a narrowly-scoped public surface.)
+- **Cookbook: pooled servings**: documented pattern — "for K
+  concurrent inference requests, start K `Nx.Serving` instances
+  behind your own pool (poolboy, Registry round-robin, etc.)".
+  No library code; clients bring their own pool since Emily already
+  behaves correctly under that model. Trade-off: each pool member
+  loads its own weights — fine for small models, painful for
+  Qwen3-7B+.
 - **README + moduledoc updates**: surface the limitation and both
   patterns prominently. Today neither is mentioned outside a buried
   comment in `Emily.Compiler`.
 
 **Testing**:
-- Pool helper test: K servings, K parallel requests, assert wall-clock
-  is ~K× faster than serial.
 - Stream test: two processes, two streams, same model loaded once;
   no SIGSEGV under sustained parallel load
   (`test/soak/backend_concurrency_test.exs` documents the SIGSEGV
   story for the unstreamed case — this is the negative control).
-- `:serving_full` opt-in: end-to-end Qwen3 pool plus per-stream
-  large-model pattern.
+- `:serving_full` opt-in: end-to-end per-stream large-model
+  pattern.
 
-**Exit:** both patterns documented; pool helper shipped; concurrency
-soak demonstrates the streamed path is stable.
+**Exit:** both patterns documented; concurrency soak demonstrates
+the streamed path is stable.
 
 ### M15 — Native linalg
 

README.md

@@ -2,11 +2,10 @@
 
 Elixir bindings and Nx backend for Apple's [MLX](https://github.com/ml-explore/mlx).
 
-**Status: M7 — Bumblebee conformance breadth.** Backend (M2), Defn
-compiler (M5), and four Bumblebee models (DistilBERT, Qwen3, ViT,
-Whisper) run end-to-end. M8 (native `conv`) and M9 (1.0 release) are
-next. See [`PLAN.md`](PLAN.md) for the full roadmap and
-[`RELEASE.md`](RELEASE.md) for unreleased-version notes.
+**Status: M14 — Serving concurrency.** Stream-per-process
+concurrent inference via `Emily.Stream`. See [`PLAN.md`](PLAN.md) for
+the full roadmap and [`RELEASE.md`](RELEASE.md) for unreleased-version
+notes.
 
 ## Why
 
@@ -63,6 +62,36 @@ The low-level tensor API (`Emily.from_binary/3`, `to_binary/1`,
 `shape/1`, `dtype/1`, `eval/1`) remains available for diagnostics and
 direct MLX round-trips, but most users should go through Nx.
 
+## Concurrency
+
+MLX dispatches GPU work through Metal command queues. By default all
+ops share one queue (the default stream), which is not safe for
+concurrent dispatch from multiple OS threads.
+
+**Stream-per-process** — for concurrent inference on a shared model:
+
+```elixir
+stream = Emily.Stream.new(:gpu)
+
+Emily.Stream.with_stream(stream, fn ->
+  # All Emily ops here dispatch on this stream's command queue.
+  model.(input)
+end)
+```
+
+Each stream maps to its own Metal command queue. Multiple processes
+can run inference concurrently — one shared model, no weight
+duplication. Create streams at init time (one per serving process),
+not per-request.
+
+**Pooled servings** — for simpler setups with small models, start K
+`Nx.Serving` instances behind a pool (poolboy, Registry, etc.). Each
+instance loads its own weights and runs on the default stream. No
+`Emily.Stream` needed. Trade-off: each pool member holds its own
+weight copy.
+
+See `Emily.Stream` moduledoc for details.
+
 ## Milestones shipped
 
 - **M0** — NIF scaffold, MLX prebuilt fetch, tensor round-trip.

RELEASE.md

@@ -2,6 +2,24 @@
 
 ## Added
 
+- M14 — Serving concurrency: stream-per-process. `Emily.Stream` lets
+  each BEAM process use its own Metal command queue for concurrent
+  inference. `Emily.Stream.new/1` creates a stream,
+  `Emily.Stream.with_stream/2` scopes all ops in a block to that
+  stream, and `Emily.Stream.synchronize/1` waits for completion.
+  The stream index is passed explicitly to every op NIF (no
+  thread-local race) via a `-1` sentinel for "use default stream"
+  (backwards-compatible). `Emily.Compiler.__partitions_options__/1`
+  error message now points to `Emily.Stream`.
+  - **New files**: `c_src/stream.cpp` (4 stream management NIFs),
+    `lib/emily/stream.ex` (`Emily.Stream` struct + API),
+    `test/emily/stream_test.exs`, `test/soak/stream_concurrency_test.exs`.
+  - **Modified**: every op NIF gained a trailing `int64_t s` stream
+    parameter; `Emily.Native` stubs, `Emily.Backend`, and all test
+    files updated accordingly.
+  - README now documents both concurrency patterns (stream-per-process
+    and pooled servings).
+
 - M13 — EXLA gradient conformance. Adds a third gradient oracle —
   EXLA (XLA CPU backend) — to catch bugs where Emily and BinaryBackend
   agree on the wrong gradient (they share the same `Nx.Defn.grad`

c_src/emily/tensor.hpp

@@ -59,4 +59,13 @@ unwrap_all(const std::vector<fine::ResourcePtr<Tensor>> &tensors) {
   return out;
 }
 
+// Resolve a stream index from Elixir into an mx::Stream.
+// -1 (the sentinel for "no explicit stream") falls through to the
+// thread-local default — backwards-compatible with pre-M14 code paths.
+inline mx::Stream resolve_stream(int64_t stream_index) {
+  if (stream_index < 0)
+    return mx::default_stream(mx::default_device());
+  return mx::get_stream(static_cast<int>(stream_index));
+}
+
 } // namespace emily

c_src/emily_nif.cpp

@@ -76,8 +76,9 @@ FINE_NIF(from_binary, 0);
 // cumulative reductions on interior axes of some 4-D shapes) raise
 // "Unable to safely factor shape" here; the Backend layer routes the
 // known cases around us.
-fine::Term to_binary(ErlNifEnv *env, fine::ResourcePtr<Tensor> tensor) {
-  auto materialized = mx::contiguous(tensor->array);
+fine::Term to_binary(ErlNifEnv *env, fine::ResourcePtr<Tensor> tensor, int64_t s) {
+  auto stream = emily::resolve_stream(s);
+  auto materialized = mx::contiguous(tensor->array, false, stream);
   mx::eval(materialized);
 
   // Defensive: mx::contiguous is supposed to give a row-contiguous

c_src/ops/binary.cpp

@@ -15,8 +15,9 @@ namespace {
   fine::ResourcePtr<Tensor> nif_name(                                          \
       ErlNifEnv *,                                                             \
       fine::ResourcePtr<Tensor> a,                                             \
-      fine::ResourcePtr<Tensor> b) {                                           \
-    return wrap(mlx_fn(a->array, b->array));                                   \
+      fine::ResourcePtr<Tensor> b,                                             \
+      int64_t s) {                                                             \
+    return wrap(mlx_fn(a->array, b->array, emily::resolve_stream(s)));         \
   }                                                                            \
   FINE_NIF(nif_name, 0);
 

c_src/ops/cast.cpp

@@ -18,8 +18,10 @@ namespace {
 fine::ResourcePtr<Tensor> astype(
     ErlNifEnv *,
     fine::ResourcePtr<Tensor> a,
-    std::tuple<fine::Atom, int64_t> dtype) {
-  return wrap(mx::astype(a->array, to_mlx_dtype(dtype)));
+    std::tuple<fine::Atom, int64_t> dtype,
+    int64_t s) {
+  return wrap(mx::astype(a->array, to_mlx_dtype(dtype),
+                         emily::resolve_stream(s)));
 }
 FINE_NIF(astype, 0);
 
@@ -28,8 +30,10 @@ FINE_NIF(astype, 0);
 fine::ResourcePtr<Tensor> bitcast(
     ErlNifEnv *,
     fine::ResourcePtr<Tensor> a,
-    std::tuple<fine::Atom, int64_t> dtype) {
-  return wrap(mx::view(a->array, to_mlx_dtype(dtype)));
+    std::tuple<fine::Atom, int64_t> dtype,
+    int64_t s) {
+  return wrap(mx::view(a->array, to_mlx_dtype(dtype),
+                       emily::resolve_stream(s)));
 }
 FINE_NIF(bitcast, 0);
 

c_src/ops/conv.cpp

@@ -19,28 +19,29 @@ using emily::wrap;
 
 namespace {
 
-// `padding` bundles low/high padding into a single tuple so the NIF
-// arity stays manageable; both have length == spatial rank.
+// `padding` bundles low/high into a tuple; `dilation` bundles
+// kernel/input dilation — both keep the NIF arity manageable.
 fine::ResourcePtr<Tensor> conv_general(
     ErlNifEnv *,
     fine::ResourcePtr<Tensor> input,
     fine::ResourcePtr<Tensor> weight,
     std::vector<int64_t> stride,
     std::tuple<std::vector<int64_t>, std::vector<int64_t>> padding,
-    std::vector<int64_t> kernel_dilation,
-    std::vector<int64_t> input_dilation,
+    std::tuple<std::vector<int64_t>, std::vector<int64_t>> dilation,
     int64_t groups,
-    bool flip) {
+    bool flip,
+    int64_t s) {
   return wrap(mx::conv_general(
       input->array,
       weight->array,
       to_int_vec(stride),
       to_int_vec(std::get<0>(padding)),
       to_int_vec(std::get<1>(padding)),
-      to_int_vec(kernel_dilation),
-      to_int_vec(input_dilation),
+      to_int_vec(std::get<0>(dilation)),
+      to_int_vec(std::get<1>(dilation)),
       static_cast<int>(groups),
-      flip));
+      flip,
+      emily::resolve_stream(s)));
 }
 FINE_NIF(conv_general, 0);
 

c_src/ops/creation.cpp

@@ -20,16 +20,20 @@ namespace {
 fine::ResourcePtr<Tensor> zeros(
     ErlNifEnv *,
     std::vector<int64_t> shape,
-    std::tuple<fine::Atom, int64_t> dtype) {
-  return wrap(mx::zeros(to_mlx_shape(shape), to_mlx_dtype(dtype)));
+    std::tuple<fine::Atom, int64_t> dtype,
+    int64_t s) {
+  return wrap(mx::zeros(to_mlx_shape(shape), to_mlx_dtype(dtype),
+                        emily::resolve_stream(s)));
 }
 FINE_NIF(zeros, 0);
 
 fine::ResourcePtr<Tensor> ones(
     ErlNifEnv *,
     std::vector<int64_t> shape,
-    std::tuple<fine::Atom, int64_t> dtype) {
-  return wrap(mx::ones(to_mlx_shape(shape), to_mlx_dtype(dtype)));
+    std::tuple<fine::Atom, int64_t> dtype,
+    int64_t s) {
+  return wrap(mx::ones(to_mlx_shape(shape), to_mlx_dtype(dtype),
+                       emily::resolve_stream(s)));
 }
 FINE_NIF(ones, 0);
 
@@ -39,8 +43,10 @@ fine::ResourcePtr<Tensor> full(
     ErlNifEnv *,
     std::vector<int64_t> shape,
     fine::ResourcePtr<Tensor> value,
-    std::tuple<fine::Atom, int64_t> dtype) {
-  return wrap(mx::full(to_mlx_shape(shape), value->array, to_mlx_dtype(dtype)));
+    std::tuple<fine::Atom, int64_t> dtype,
+    int64_t s) {
+  return wrap(mx::full(to_mlx_shape(shape), value->array, to_mlx_dtype(dtype),
+                       emily::resolve_stream(s)));
 }
 FINE_NIF(full, 0);
 
@@ -51,8 +57,10 @@ fine::ResourcePtr<Tensor> arange(
     double start,
     double stop,
     double step,
-    std::tuple<fine::Atom, int64_t> dtype) {
-  return wrap(mx::arange(start, stop, step, to_mlx_dtype(dtype)));
+    std::tuple<fine::Atom, int64_t> dtype,
+    int64_t s) {
+  return wrap(mx::arange(start, stop, step, to_mlx_dtype(dtype),
+                         emily::resolve_stream(s)));
 }
 FINE_NIF(arange, 0);
 
@@ -62,11 +70,13 @@ fine::ResourcePtr<Tensor> eye(
     int64_t n,
     int64_t m,
     int64_t k,
-    std::tuple<fine::Atom, int64_t> dtype) {
+    std::tuple<fine::Atom, int64_t> dtype,
+    int64_t s) {
   return wrap(mx::eye(static_cast<int>(n),
                       static_cast<int>(m),
                       static_cast<int>(k),
-                      to_mlx_dtype(dtype)));
+                      to_mlx_dtype(dtype),
+                      emily::resolve_stream(s)));
 }
 FINE_NIF(eye, 0);
 

c_src/ops/fast.cpp

@@ -52,9 +52,11 @@ fine::ResourcePtr<Tensor> fast_rms_norm(
     ErlNifEnv *,
     fine::ResourcePtr<Tensor> x,
     std::optional<fine::ResourcePtr<Tensor>> weight,
-    double eps) {
+    double eps,
+    int64_t s) {
   return wrap(mx::fast::rms_norm(
-      x->array, opt_array(weight), static_cast<float>(eps)));
+      x->array, opt_array(weight), static_cast<float>(eps),
+      emily::resolve_stream(s)));
 }
 FINE_NIF(fast_rms_norm, 0);
 
@@ -71,12 +73,14 @@ fine::ResourcePtr<Tensor> fast_layer_norm(
     fine::ResourcePtr<Tensor> x,
     std::optional<fine::ResourcePtr<Tensor>> weight,
     std::optional<fine::ResourcePtr<Tensor>> bias,
-    double eps) {
+    double eps,
+    int64_t s) {
   return wrap(mx::fast::layer_norm(
       x->array,
       opt_array(weight),
       opt_array(bias),
-      static_cast<float>(eps)));
+      static_cast<float>(eps),
+      emily::resolve_stream(s)));
 }
 FINE_NIF(fast_layer_norm, 0);
 
@@ -108,7 +112,8 @@ fine::ResourcePtr<Tensor> fast_rope(
     std::optional<double> base,
     double scale,
     fine::ResourcePtr<Tensor> offset,
-    std::optional<fine::ResourcePtr<Tensor>> freqs) {
+    std::optional<fine::ResourcePtr<Tensor>> freqs,
+    int64_t s) {
   std::optional<float> base_f;
   if (base) base_f = static_cast<float>(*base);
 
@@ -119,7 +124,8 @@ fine::ResourcePtr<Tensor> fast_rope(
       base_f,
       static_cast<float>(scale),
       offset->array,
-      opt_array(freqs)));
+      opt_array(freqs),
+      emily::resolve_stream(s)));
 }
 FINE_NIF(fast_rope, 0);
 
@@ -147,14 +153,16 @@ fine::ResourcePtr<Tensor> fast_scaled_dot_product_attention(
     fine::ResourcePtr<Tensor> v,
     double scale,
     std::string mask_mode,
-    std::vector<fine::ResourcePtr<Tensor>> mask_arrs) {
+    std::vector<fine::ResourcePtr<Tensor>> mask_arrs,
+    int64_t s) {
   return wrap(mx::fast::scaled_dot_product_attention(
       q->array,
       k->array,
       v->array,
       static_cast<float>(scale),
       mask_mode,
-      unwrap_all(mask_arrs)));
+      unwrap_all(mask_arrs),
+      emily::resolve_stream(s)));
 }
 FINE_NIF(fast_scaled_dot_product_attention, 0);