Merge pull request #169 from ausimian/feat/expr-compiler-cm14-docs-drift

docs: scope the fused-while drift caveats honestly

Author: ausimian

RELEASE.md

@@ -80,10 +80,17 @@
   On Qwen3-0.6B this lifts greedy decode to **~5.4× the evaluator (~1.1× over
   the plain native lane**, ~68 vs ~62 tok/s on an M-series Mac). The trade-off:
   `mx::compile` reassociates f32, so logits drift by a few ULP and the output
-  is **not** bit-identical to the evaluator — greedy argmax is stable under
-  that, so the generated token ids still match exactly (the completions are
-  byte-identical). The `native-fused` lane in `bench/qwen3_tokens_per_sec.exs`
-  measures it; `generation_native_test.exs` gates the greedy token match.
+  is **not** bit-identical to the evaluator. Greedy argmax is robust to that
+  drift, so in our Qwen3-0.6B run the generated token ids matched the
+  evaluator's exactly (byte-identical completions) — but that is an *empirical*
+  result, not a guarantee: a near-tie top-2 logit can flip a token, and any
+  decision the drift can tip over — argmax, or a loop trip count whose
+  condition reads a reassociated reduction — diverges by more than a few ULP
+  once it flips. **Sampling strategies (e.g. multinomial) will diverge from the
+  evaluator under fusion** even with a fixed seed; the gate and bench cover the
+  greedy lane only. The `native-fused` lane in `bench/qwen3_tokens_per_sec.exs`
+  measures throughput; `generation_native_test.exs` gates the greedy token
+  match.
 
 - **`defn while` compiles native.** Data-dependent loops — including
   `Bumblebee.Text.generation`'s decode loop — now lower to the single-NIF

lib/emily/compiler.ex

@@ -88,9 +88,14 @@ defmodule Emily.Compiler do
       the decode loop host-controlled but fuses each loop **body** under
       `mx::compile`, replaying the cached fused callable every token. Defaults
       to `false`; a non-boolean raises `ArgumentError`. Opt-in because the
-      fusion reassociates f32 to within a few ULP — greedy token ids still
-      match the evaluator, but logits are not bit-identical. Only the native
-      path consults it, so it is ignored unless `native: true`.
+      fusion reassociates f32 to within a few ULP — logits are not
+      bit-identical to the evaluator. Greedy argmax is robust to that drift
+      (greedy token ids matched the evaluator in our tests), but the match is
+      empirical, not guaranteed: any discrete decision the drift can tip —
+      argmax on a near-tie, or a `while` trip count whose condition reads a
+      reassociated reduction — diverges once it flips. **Sampling strategies
+      diverge from the evaluator under fusion** even with a fixed seed. Only
+      the native path consults it, so it is ignored unless `native: true`.
 
   Any other option is silently dropped. This matches how
   `Nx.Defn.Evaluator` and EXLA handle their own option lists, and is
@@ -228,7 +233,7 @@ defmodule Emily.Compiler do
   # guarded step: it raises `ArgumentError` on an op or construct it can't
   # lower yet, which we turn into a graceful `:fallback` (or re-raise in
   # `:raise` mode). `Program.compile/1` is deliberately kept outside the
-  # rescue (in `replay_closure/2`) — it raises only on malformed IR, i.e. a
+  # rescue (in `replay_closure/3`) — it raises only on malformed IR, i.e. a
   # compiler bug, which must surface loudly rather than be masked as an
   # "unsupported op" fallback.
   defp lower(leaves, mode, key) do

lib/emily/program.ex

@@ -75,9 +75,14 @@ defmodule Emily.Program do
           can't trace) is instead replayed host-controlled with each loop
           *body* fused under `mx::compile` and cached per stream — the body
           is shape-stable, so the fused callable cache-hits across
-          iterations rather than recompiling per step. Either way the
-          fusion reassociates f32 to within a few ULP, so the result is
-          not bit-identical to `:sync`.
+          iterations rather than recompiling per step. The loop condition
+          is left as a raw (un-fused) replay. Either way the fusion
+          reassociates f32 to within a few ULP, so the result is not
+          bit-identical to `:sync` — and for a `while` whose condition
+          reads a reassociated reduction of the fused body, the drift can
+          change the *trip count*, diverging by more than a few ULP. Safe
+          for the generation loop (its condition is an integer offset, not
+          a function of the drifted logits).
   """
   @spec eval(Native.worker(), t(), [Native.tensor()], keyword()) :: [Native.tensor()]
   def eval(worker, program, inputs, opts \\ []) do

test/emily/conformance/generation_native_test.exs

@@ -82,9 +82,12 @@ defmodule Emily.Conformance.GenerationNativeTest do
     # The opt-in fused-while lane (CM14): the `defn while` decode loop runs
     # host-controlled, but each per-token forward (the loop body) replays
     # through a per-stream-cached `mx::compile`'d callable. `mx::compile`
-    # reassociates f32 so the logits drift by a few ULP — but greedy argmax
-    # is stable under that, so the generated token ids still match the
-    # evaluator exactly. (Not a binary-identical gate, by construction.)
+    # reassociates f32 so the logits drift by a few ULP — greedy argmax is
+    # robust to that drift, so the token ids match here. This is a
+    # token-id gate, not a binary-identical one (the logits aren't), and
+    # the match is empirical: a near-tie top-2 logit could in principle
+    # flip a token on another model/prompt. Sampling strategies would
+    # diverge under fusion, so only greedy is gated.
     gc = configure(ctx.gen_config, %{type: :greedy_search})
     fused = generate_ids(ctx.model_info, gc, @native_compiled)
     eval = generate_ids(ctx.model_info, gc, @eval)