Emily.Fast.* fused kernels fail when composed inside a caller's own defn (fallback functions aren't defnp)

[lostbean]

Summary

Every Emily.Fast fused kernel (rms_norm, layer_norm, rope,
rope_with_freqs, and all four scaled_dot_product_attention* variants)
raises RuntimeError: cannot invoke Emily.Fast.<kernel> inside defn because it was not defined with defn when called from inside a defn-defined
function nested one level below the top-level jit_apply/defn entry
point — even though the module's own moduledoc documents this exact usage
as supported ("Call these from inside a defn or Nx.Defn.jit-traced
function, alongside regular Nx ops").

Environment

• emily 1.0.0 (Hex), MLX 0.32.0
• macOS, Apple Silicon
• elixir 1.18.4 / OTP 28
• nx 0.12.1
• Both Emily.Compiler with native: true and without (plain
  Nx.Defn.Evaluator) — same failure either way.

Repro

The moduledoc's own example, run two ways:

defmodule DebugFastDefn do
  import Nx.Defn

  defn block(x, w) do
    Emily.Fast.rms_norm(x, w, eps: 1.0e-5)
  end
end

x = Nx.tensor([[1.0, 2.0, 3.0, 4.0]], backend: Emily.Backend)
w = Nx.tensor([1.0, 1.0, 1.0, 1.0], backend: Emily.Backend)

Works — calling the moduledoc's Nx.Defn.jit_apply form directly, with
the anonymous function as the top-level jit'd unit:

Nx.Defn.jit_apply(fn x, w -> Emily.Fast.rms_norm(x, w, eps: 1.0e-5) end, [x, w])
#=> #Nx.Tensor<...> -- succeeds

Fails — the exact same call, one level deeper, wrapped in a
module-level defn (i.e. the moduledoc's own defn block(x, w) do ... end
example, called via DebugFastDefn.block(x, w) or via
Nx.Defn.jit_apply(&DebugFastDefn.block/2, [x, w], compiler: Emily.Compiler, native: true)):

** (RuntimeError) cannot invoke Emily.Fast.rms_norm/3 inside defn because it was not defined with defn
    (nx 0.12.1) lib/nx/defn/compiler.ex:190: Nx.Defn.Compiler.runtime_fun/3
    ...

Confirmed reproducible under both compiler: Emily.Compiler (with and
without native: true) and plain Nx.Defn.Evaluator (no Emily.Compiler
configured at all) — identical error either way.

Root cause (traced through lib/emily/fast.ex)

Every fused kernel wraps its Nx.block/4 fallback in a plain def,
not defnp:

# lib/emily/fast.ex
def rms_norm(x, weight, opts \\ []) do
  opts = Keyword.validate!(opts, eps: 1.0e-6)
  block = struct!(FB.RMSNorm, opts)

  Nx.block(block, [x, weight], output_like(x), fn ^block, x, weight ->
    rms_norm_fallback(x, weight, opts)   # <-- plain def, not defnp
  end)
end

defp rms_norm_fallback(x, weight, opts) do
  ...
end

Nx.block/4's fallback function argument must itself be defn-traceable.
When the outer call is the literal top-level unit handed to
Nx.Defn.jit_apply/defn, tracing apparently short-circuits before this
matters. Once Emily.Fast.rms_norm/3 is called from inside an
already-defn-traced function (i.e. genuine composition — the moduledoc's
own stated use case), the trace has to actually enter Nx.block's fallback
path, and hits the plain-def rms_norm_fallback/2 — hence "was not
defined with defn."

This isn't rms_norm-specific — the identical pattern (Nx.block(..., fn ... -> some_plain_def_fallback(...) end)) appears in every fused kernel in
this module: layer_norm/layer_norm_fallback,
rope/rope_fallback, rope_with_freqs/rope_freqs_fallback,
scaled_dot_product_attention/sdpa_fallback and sdpa_sinks_fallback,
scaled_dot_product_attention_with_mask/sdpa_masked_fallback and
sdpa_masked_sinks_fallback.

Impact

This blocks a real use case: composing multiple fused kernels (e.g. several
transformer layers' worth of rms_norm + rope + SDPA) into one larger
defn-traced graph, to get Emily.Compiler's native single-NIF
whole-graph replay speedup (per #163's own numbers, ~5x on a decode loop)
across the composed graph rather than dispatching each fused kernel as its
own separate top-level call. Right now, the only way to call any
Emily.Fast kernel successfully is as the literal top-level jit'd unit —
never nested inside a caller's own defn.

Suggested fix

Change each *_fallback function from defp to defnp (and correspondingly
any helper functions those fallbacks call that aren't already defn-safe).
Given the fallback bodies already appear to be pure Nx composition (matmuls,
elementwise ops, reductions — nothing outside defn's supported op set, based
on reading rms_norm_fallback/layer_norm_fallback), this looks like it
should be a mechanical, low-risk change, but I haven't attempted the fix
myself or checked whether it's contained to this file.

Acceptance criteria (suggested)

• The moduledoc's own defn block(x, w) do Emily.Fast.layer_norm(w, b, ...) end
  example succeeds when called through a caller's defn, not just as the
  top-level jit'd unit.
• Same for rope, rope_with_freqs, and both SDPA variants.
• A regression test nesting at least one Emily.Fast call two levels deep
  (a defn calling a defnp calling Emily.Fast.*) so this doesn't
  regress silently.

[lostbean]

Update: tried the suggested defp → defnp fix, it's not sufficient

Attempted this in a fork. Converting each Nx.block fallback (and its
shape-manipulation helpers) from defp/def to defnp/deftransformp
compiles cleanly, but the original error persists at runtime — because
defn's call-dispatch checks the callee's own definition kind at the
call site, and the public API functions (Emily.Fast.rms_norm/3 etc.) are
still plain def.

Making those defn too hits a second wall: Nx.block's own callback
signature — fn ^block, x, weight -> ... end, a pin-matched anonymous
function — isn't itself expressible inside a defn body (defn rejects
pin-patterns as anonymous function parameters: "only variables, tuples,
maps, and structs are allowed as patterns in defn").

So this isn't a mechanical rename. It looks like Nx.block's current
construction pattern doesn't support being called from a caller's own
defn at all — which is a real gap against the moduledoc's own claim
("Call these from inside a defn or Nx.Defn.jit-traced function").
Whether the right fix is a different Nx.block construction, or a
statement that composition genuinely isn't supported yet, isn't something
I can resolve without knowing more about Nx.block's intended contract.

Why this matters (real use case)

I'm porting a real ~450M-parameter transformer (SmolVLA, a
vision-language-action model) to Elixir/Nx against emily. Each layer
already calls Emily.Fast.rms_norm/rope/scaled_dot_product_attention_with_mask
as separate top-level dispatches (the only way that currently works), and
that's the dominant cost in a ~1.6s forward pass — 16 layers × 10
integration steps, each layer several separate NIF-crossing calls. Composing
a whole layer (or the whole model) into one defn-traced graph is the
natural next step to get Emily.Compiler's native whole-graph replay
speedup (the ~5x you measured in #163) — but it's blocked as long as fused
kernels can only be called as the literal top-level jit'd unit, never nested
inside a caller's own defn.

[lostbean]

I dug into this and have a fix up (PR to follow from lostbean/emily#fix/205-fast-defn-composition). Two findings worth recording here:

The suggested fix (defp → defnp on the fallbacks) doesn't work — and can't

The fallback callbacks pin-match the block struct (fn ^block, x, weight -> ...), and Nx.block/4's callback is applied as plain Elixir at trace time by Nx.Defn.Expr.block/4 — the fallbacks are supposed to be ordinary functions. Converting them to defnp hits the pin-match (illegal in defn) and misdiagnoses the failure.

Actual root cause

defn rewrites every remote call to Nx.Defn.Compiler.__remote__(mod, fun, :"__defn:fun__", args), which dispatches through a generated __defn:name__/arity delegator. Only defn/deftransform definitions export those delegators. The public kernels (rms_norm, layer_norm, rope, rope_with_freqs, the four SDPA variants) are plain def, so any call from inside a caller's own defn raises "was not defined with defn" before the body ever runs. That's why the top-level jit_apply(fn ...) form works (no remote-call rewrite happens) but the moduledoc's own defn block(x, w) example fails.

The fix is one level up from the suggestion: convert the eight public kernels to deftransform. That exports the delegator (making them defn-composable) while keeping the bodies plain Elixir — so the Nx.block construction, the pin-matched callbacks, and the entire fast path are untouched.

Verified against this issue's exact repro

On the fix branch, all three previously-failing forms now pass, with identical values across all lanes (max diff 0.0 vs Nx.Defn.Evaluator/Nx.BinaryBackend):

• DebugFastDefn.block(x, w) called directly ✅
• Nx.Defn.jit_apply(&DebugFastDefn.block/2, [x, w], compiler: Emily.Compiler, native: true, native_fallback: :raise) ✅ — with :raise set, so it provably ran the fused native lane
• plain Nx.Defn.Evaluator, no Emily.Compiler ✅

All acceptance criteria are covered by a new test/emily/fast/composed_defn_test.exs (14 tests): the two-levels-deep regression tests (defn → defnp → kernel) for rms_norm, rope, and scaled_dot_product_attention_with_mask, plus composed fused-vs-BinaryBackend conformance for all 8 kernel variants, with the fused path asserted via IR opcode inspection, native_fallback: :raise, and a negative [:emily, :block, :fallback] telemetry assertion.

Bonus: the new conformance tests exposed two pre-existing fused-vs-fallback bugs

Fixed in separate commits on the same branch, since they sit on the same code paths:

1. SDPA fallback's causal mask was inverted (attended future keys, masked past ones). Now matches mx::fast::scaled_dot_product_attention's bottom-right-aligned convention: query row i attends keys j <= i + (k_len - q_len).
2. rope_with_freqs passed the HF inverse-frequency table straight to mx::fast::rope, which expects its reciprocal (MLX computes θ = position / freqs — see mlx/fast.cpp and the Metal rope kernels — while emily documents the HF θ = position × inv_freq convention). Fused rotations were wrong at every position > 0. Both native lanes (Emily.Backend.block/4 and the Emily.Compiler IR lowering) now invert the table before the NIF call.

These went unnoticed because the existing per-kernel conformance tests set Nx.default_backend(Emily.Backend) in setup, which routes their oracle lane through the fused kernel too — they compare fused against fused (diff exactly 0.0), so the fallback/fused divergence was invisible. Worth a follow-up to rebase those oracles onto Nx.BinaryBackend; I didn't rewrite them in this PR to keep it scoped.

Real-world validation

Converted a SmolVLA (16-layer joint-attention transformer, real lerobot/smolvla_base checkpoint) consumer's per-layer def pipeline to a single composed defn using the fixed kernels: numerical parity held (0.646% mean relative error, unchanged from the eager baseline), and warm end-to-end inference went from ~1.26s to ~611ms median — the whole stack now replays as one native graph per step, which is precisely the use case this issue says is blocked.