M11: Fused MLX transformer kernels

[ausimian]

Summary

Wires MLX's handwritten mx::fast::* kernels (RMSNorm, LayerNorm, RoPE, scaled-dot-product attention) into Emily as defn-callable helpers, and ships a Bumblebee shim that swaps them in for the stock composed-defn implementations. Closes the fusion gap M10.5 noted: the two-kernel dequantize + dot cost on quantized matmuls, plus the ~5-dispatch attention chain, collapse to one kernel dispatch each.

• Native NIFs (c_src/ops/fast.cpp) over mx::fast::rms_norm, mx::fast::layer_norm, mx::fast::rope, and mx::fast::scaled_dot_product_attention. Nullable weight/bias/freqs args marshal via std::optional.
• Emily.Fast (lib/emily/fast.ex) — rms_norm/3, layer_norm/4, rope/3, rope_with_freqs/4, scaled_dot_product_attention/4 and the _with_mask variant. Each emits a Nx.Defn.Expr.optional/3 node whose op name matches a custom callback on Emily.Backend; the Evaluator dispatches to the fused kernel under Emily and falls back to a defn composition on any other backend (BinaryBackend / EXLA conformance stays green without a re-plumb).
• Emily.Backend.fast_* — six custom callbacks (outside the Nx.Backend behaviour) that unwrap refs and call Native directly.
• Emily.Bumblebee.FastKernels (test/support/) — Axon graph rewriter mirroring the M10.5 Emily.Quantization.Transform pattern. Rewrites :rms_norm and :layer_norm via Axon.map_nodes, Bumblebee.Layers.apply_rotary_embedding/5 by MFA match, and coalesces attention_weights_impl + attention_output_impl into one fused SDPA layer via Axon.rewrite_nodes. RoPE supports all four Bumblebee scaling strategies (:linear, :dynamic, :longrope, :llama3) by precomputing the inverse-frequency table at rewrite time and passing it to MLX via the freqs-override overload — no per-token BEAM-side recomputation.

Mechanism: Nx.Defn.Expr.optional/3 is Nx's documented extension point for vendor-fused kernels (the pattern EXLA uses too). Rejected the pattern-matched subgraph-fusion approach — PLAN.md M11 explicitly defers that as compiler-level work.

Scope & trade-offs

• Attention rewrite leaves the unfused attention_weights_impl node in the graph because it's referenced from Bumblebee's {output, weights} tuple. On the generation path (output_attentions=false), Nx.Defn.Evaluator.precompile walks IDs only from the returned expression — unreachable nodes never execute, so the fused SDPA is a pure win. For output_attentions=true the orphan still runs; accepted.
• head_mask fusion is approximate: applies the mask to per-head outputs post-attention. Equivalent to a 0/1 mask; fractional values diverge slightly. Bumblebee's built-in usage is 0/1 only.
• channel_index != -1 on norm layers skipped (vision-CNN heads; no transformer hits it).

Test plan

• mix precommit — 349 tests, 0 failures, credo clean
• mix test --only conformance — 17 tiny-random conformance tests green
• mix test --only fast_kernels_full test/emily/conformance/distilbert_test.exs — tiny-random DistilBERT with fused kernels matches the dense-path pinned slice
• mix test --only fast_kernels_full — full-model :fast_kernels_full variants (Qwen3 dense, Qwen3 quantized, ViT, Whisper) on a machine with the checkpoints cached
• EMILY_BENCH_FAST_KERNELS=1 mix run bench/qwen3_tokens_per_sec.exs — verify tokens/sec speedup on M3 hardware; tighten the EMILY_BENCH_PIN=<multiplier> floor after first measurement

New test surface

• test/emily/fast/{rms_norm,layer_norm,rope,sdpa}_test.exs — 16 cases across native unit vectors, defn-composability inside jitted functions, and fused-vs-composed equivalence for f32 + bf16
• test/emily/bumblebee/fast_kernels_test.exs — 3 shim unit tests on handcrafted Axon models; asserts rewrites land, init_fn succeeds, and predict output matches the unrewritten path within tolerance
• :fast_kernels_full tagged variants of every existing *_full conformance suite plus a tiny-random DistilBERT smoke. Tag excluded by default like the other *_full tags

Full details in PLAN.md §M11.