Make Emily.Fast fused kernels composable inside a caller's defn (fixes #205)

[lostbean]

Fixes #205.

What

Three commits, smallest-possible diffs:

1. 36fbb04 — make Emily.Fast fused kernels callable inside a caller's own defn. Converts the eight public kernels (rms_norm/3, layer_norm/4, rope/3, rope_with_freqs/4, and the four scaled_dot_product_attention* variants) from plain def to deftransform, plus a new test/emily/fast/composed_defn_test.exs (14 tests). einsum stays def (eager-only); every *_fallback stays plain defp.
2. e32e3c4 — correct the causal-mask direction in the SDPA composed fallback. The mask was inverted (attended future keys). Now bottom-right-aligned to match mx::fast::sdpa: query row i attends keys j <= i + (k_len - q_len).
3. 5ade402 — invert the freqs table before mx::fast::rope in both native lanes. emily documents the HF convention (θ = position × inv_freq) but MLX's rope computes θ = position / freqs, so the fused rope_with_freqs was wrong at every position > 0. Both Emily.Backend.block/4 and the Emily.Compiler IR lowering now pass the reciprocal.

Why deftransform, not the issue's suggested defp → defnp

The failure is not in the fallbacks. defn rewrites remote calls to dispatch through a generated __defn:name__/arity delegator that only defn/deftransform definitions export — a plain-def kernel raises "was not defined with defn" before its body runs, which is why only the literal top-level jit_apply(fn ...) form ever worked. deftransform exports the delegator while keeping the body plain Elixir, which is required: Nx.block/4's callback is applied as ordinary Elixir at trace time (Nx.Defn.Expr.block/4) and pin-matches the block struct (fn ^block, ... ->), which is illegal inside defn. The Nx.block construction and the entire fast path are byte-for-byte untouched.

Bugs 2 and 3 were exposed by the new conformance tests and are pre-existing (present in 1.0.0), not composition regressions — the composed and top-level lanes agreed with each other on the wrong values. They went unnoticed because the existing per-kernel tests set Nx.default_backend(Emily.Backend) in setup, so their oracle lane also runs fused and the comparison is fused-vs-fused (diff exactly 0.0). I left those tests as-is to keep this PR scoped, but rebasing their oracles onto Nx.BinaryBackend is worth a follow-up — happy to split any of the three commits into separate PRs if preferred.

Test coverage (composed_defn_test.exs)

• The issue's acceptance criteria: regression tests nesting a kernel two levels deep (defn → defnp → Emily.Fast.*) for rms_norm, rope, and scaled_dot_product_attention_with_mask, including under global Emily.Backend + Emily.Compiler defaults.
• Composed fused-vs-fallback conformance for all 8 kernel variants: compiler: Emily.Compiler, native: true, native_fallback: :raise vs Nx.Defn.Evaluator on Nx.BinaryBackend, within f32 tolerance.
• The fused path is asserted three ways: composed graphs lower to the :fast_* IR opcodes, the fused lane runs with native_fallback: :raise, and a telemetry handler flunks on any [:emily, :block, :fallback] event for an Emily.Fast.Block.* struct.
• A mixed-kernel test composing several kernels in one defn.

Validation

• Full suite: 817 tests, 0 failures (40 doctests, 79 properties).
• The issue's exact repro (all three previously-failing forms) passes; all lanes agree to 0.0.
• Real-world check: converting a SmolVLA consumer (16-layer joint-attention stack, real lerobot/smolvla_base checkpoint) from per-layer eager calls to one composed defn kept numerical parity (0.646% MRE vs the Python reference, unchanged) and cut warm inference from ~1.26s to ~611ms median — the composed-graph speedup this issue is after.

Note: all changes are Elixir-only; the native suite was run against the released, checksum-verified 1.0.0 NIF (my dev machine lacks the Metal toolchain for an MLX source build). Nothing under c_src/ is touched.

🤖 Generated with Claude Code