From 5a68c300c8a95dca52fbf9ebb99bed033a728404 Mon Sep 17 00:00:00 2001 From: Codcore Date: Thu, 20 Aug 2026 16:27:20 +0200 Subject: [PATCH 1/2] Cache jit closures instead of re-tracing on every call MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit __jit__/5 forwarded straight to __compile__/4, which discards the key Nx supplies and calls native_compile/3. That runs fun.(vars) and walks the resulting expression on every call, only to look up a program already in the program table. Cache the closure __compile__/4 returns, keyed by that key plus the shapes and types of vars and the device. On a hit the trace never runs. The program table is untouched. Its structural key is what lets structurally identical layers share one compiled program across different call sites; this table sits above it and short-circuits repeat calls of the same call site — the shape of every decode loop and training step. Calls carrying a :command_queue or :hooks are not cached: those options reference caller-owned state whose lifetime should not be extended by a process-wide table. Measured on a chain of cheap ops, where traversal rather than arithmetic dominates: depth before after Nx.Defn.compile 32 0.59 ms 0.43 ms 0.37 ms 128 1.01 ms 0.53 ms 0.38 ms 512 3.12 ms 1.18 ms 0.67 ms 1024 5.89 ms 1.85 ms 1.30 ms On a 40-layer MoE decode step the same change removes 8.35 ms per token out of 89. --- emlx/lib/emlx.ex | 68 +++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 67 insertions(+), 1 deletion(-) diff --git a/emlx/lib/emlx.ex b/emlx/lib/emlx.ex index e1a847c..38dfcb1 100644 --- a/emlx/lib/emlx.ex +++ b/emlx/lib/emlx.ex @@ -1764,9 +1764,75 @@ defmodule EMLX do # `dispatch_key/3` still has to fall back to the structural walk on a miss.) @dispatch_key_by_id_table :emlx_dispatch_key_by_id + # Closures returned by `__compile__/4`, keyed by the `key` Nx supplies plus + # the shapes/types of `vars` and the options that affect lowering. + # + # Without this, every `Nx.Defn.jit/2` call re-runs `fun.(vars)` and walks the + # resulting expression, only to arrive at a program that is already in + # `@program_table`. The program cache stays as it is — its structural key is + # what lets structurally identical layers share one program across *different* + # call sites. This table is the layer above: it short-circuits *repeat* calls + # of the same call site with the same argument shapes, which is the shape of + # every decode loop and every training step. + @jit_closure_table :emlx_jit_closures + @impl Nx.Defn.Compiler def __jit__(key, vars, fun, args_list, opts) do - __compile__(key, vars, fun, opts).(args_list) + case jit_cache_key(key, vars, opts) do + nil -> + __compile__(key, vars, fun, opts).(args_list) + + cache_key -> + closure = + case :ets.lookup(jit_closure_table(), cache_key) do + [{^cache_key, cached}] -> + cached + + [] -> + built = __compile__(key, vars, fun, opts) + :ets.insert(jit_closure_table(), {cache_key, built}) + built + end + + closure.(args_list) + end + end + + # `nil` means "do not cache": an unhashable key, or options carrying a + # caller-owned resource whose lifetime we must not extend. + defp jit_cache_key(key, vars, opts) do + if Keyword.has_key?(opts, :command_queue) or Keyword.has_key?(opts, :hooks) do + nil + else + templates = + vars + |> Nx.Defn.Composite.flatten_list() + |> Enum.map(fn %Nx.Tensor{shape: shape, type: type} -> {shape, type} end) + + {key, templates, Keyword.get(opts, :device, default_device())} + end + rescue + _ -> nil + end + + defp jit_closure_table do + case :ets.whereis(@jit_closure_table) do + :undefined -> + try do + :ets.new(@jit_closure_table, [ + :named_table, + :public, + :set, + read_concurrency: true, + write_concurrency: true + ]) + rescue + ArgumentError -> @jit_closure_table + end + + _ -> + @jit_closure_table + end end @impl Nx.Defn.Compiler From 32b601fd9e5f8d2178274bb15a241eaa658c57d9 Mon Sep 17 00:00:00 2001 From: Paulo Valente <16843419+polvalente@users.noreply.github.com> Date: Sat, 29 Aug 2026 00:00:40 -0300 Subject: [PATCH 2/2] refactor: cache compile instead of jit --- emlx/lib/emlx.ex | 128 ++++++++++++----------------------- emlx/lib/emlx/application.ex | 1 + 2 files changed, 44 insertions(+), 85 deletions(-) diff --git a/emlx/lib/emlx.ex b/emlx/lib/emlx.ex index 38dfcb1..7cd453d 100644 --- a/emlx/lib/emlx.ex +++ b/emlx/lib/emlx.ex @@ -1738,112 +1738,57 @@ defmodule EMLX do @behaviour Nx.Defn.Compiler - # Known EMLX-specific compiler opts. `:command_queue` is injected by - # `__partitions_options__/1` but may also be passed directly by callers - # that manage their own queues (equivalent to a manual `with_queue`). @valid_compiler_keys [:device, :max_concurrency, :command_queue, :hooks] - # Process-lifetime dispatch cache backing `dispatch_key/3` + - # `get_or_compile_program/6` (see their docs) — a compiled program is keyed - # by a *structural* signature of its `Expr` (not object identity), so - # it survives across `Nx.Defn.Graph.run/3`'s per-call re-tracing and is - # shared across structurally-identical call sites (e.g. every one of - # Qwen3's 28 attention layers), not just within one closure's lifetime. + # Structural program cache for `get_or_compile_program/6`. @native_dispatch_cache_table :emlx_native_dispatch_cache - # A second, process-lifetime cache in front of `dispatch_key/3`'s own - # (expensive — O(nodes), plus per-opaque-scope SHA256 hashing) structural - # walk, keyed by `output_expr`'s own node identity rather than its - # structural signature. This matters for `run_while_loop/3`'s host-driven - # `cond_fn`/`body_fn`: `Nx.Defn.jit/2` retraces `fn _ -> body_expr end` - # once and caches *that* trace by argument template, so every subsequent - # call re-enters `__jit__`/`build_eval_fn` with the *exact same* `Expr` - # (identical ids, not just structurally identical) — walking it again on - # every decode step is pure waste. (This is unlike `Nx.Defn.Graph.run/3`'s - # per-stage re-tracing, which *does* mint fresh ids each call — hence - # `dispatch_key/3` still has to fall back to the structural walk on a miss.) + # Expr-id cache in front of `dispatch_key/3`. @dispatch_key_by_id_table :emlx_dispatch_key_by_id - # Closures returned by `__compile__/4`, keyed by the `key` Nx supplies plus - # the shapes/types of `vars` and the options that affect lowering. - # - # Without this, every `Nx.Defn.jit/2` call re-runs `fun.(vars)` and walks the - # resulting expression, only to arrive at a program that is already in - # `@program_table`. The program cache stays as it is — its structural key is - # what lets structurally identical layers share one program across *different* - # call sites. This table is the layer above: it short-circuits *repeat* calls - # of the same call site with the same argument shapes, which is the shape of - # every decode loop and every training step. - @jit_closure_table :emlx_jit_closures + # `native_compile/3` closures. Created by `init/0`. + @compile_closure_table :emlx_compile_closures - @impl Nx.Defn.Compiler - def __jit__(key, vars, fun, args_list, opts) do - case jit_cache_key(key, vars, opts) do - nil -> - __compile__(key, vars, fun, opts).(args_list) - - cache_key -> - closure = - case :ets.lookup(jit_closure_table(), cache_key) do - [{^cache_key, cached}] -> - cached - - [] -> - built = __compile__(key, vars, fun, opts) - :ets.insert(jit_closure_table(), {cache_key, built}) - built - end + @doc false + def init do + case :ets.whereis(@compile_closure_table) do + :undefined -> + :ets.new(@compile_closure_table, [ + :named_table, + :public, + :set, + read_concurrency: true, + write_concurrency: true + ]) - closure.(args_list) + _ -> + :ok end end - # `nil` means "do not cache": an unhashable key, or options carrying a - # caller-owned resource whose lifetime we must not extend. - defp jit_cache_key(key, vars, opts) do - if Keyword.has_key?(opts, :command_queue) or Keyword.has_key?(opts, :hooks) do - nil - else - templates = - vars - |> Nx.Defn.Composite.flatten_list() - |> Enum.map(fn %Nx.Tensor{shape: shape, type: type} -> {shape, type} end) - - {key, templates, Keyword.get(opts, :device, default_device())} - end - rescue - _ -> nil + @impl Nx.Defn.Compiler + def __jit__(key, vars, fun, args_list, opts) do + __compile__(key, vars, fun, opts).(args_list) end - defp jit_closure_table do - case :ets.whereis(@jit_closure_table) do - :undefined -> - try do - :ets.new(@jit_closure_table, [ - :named_table, - :public, - :set, - read_concurrency: true, - write_concurrency: true - ]) - rescue - ArgumentError -> @jit_closure_table - end + defp compile_cache_key(key, vars, hooks, device) do + templates = + Enum.map(vars, fn var -> Nx.Defn.Composite.traverse(var, &Nx.to_template/1) end) - _ -> - @jit_closure_table - end + {key, templates, hooks, device} end @impl Nx.Defn.Compiler - def __compile__(_key, vars, fun, opts) do + def __compile__(key, vars, fun, opts) do Keyword.validate!(opts, @valid_compiler_keys) - case Keyword.get(opts, :hooks, %{}) do + hooks = Keyword.get(opts, :hooks, %{}) + + case hooks do empty when empty == %{} -> :ok - hooks -> + _ -> raise ArgumentError, "EMLX does not support the :hooks named-override map (got callbacks for " <> "#{inspect(Map.keys(hooks))}) — :hook/:io_call expr nodes lower natively " <> @@ -1854,7 +1799,20 @@ defmodule EMLX do queue = Keyword.get(opts, :command_queue) device = Keyword.get(opts, :device, default_device()) - wrap_with_queue(queue, native_compile(vars, fun, device)) + cache_key = compile_cache_key(key, vars, hooks, device) + + eval_fn = + case :ets.lookup(@compile_closure_table, cache_key) do + [{^cache_key, cached}] -> + cached + + [] -> + built = native_compile(vars, fun, device) + :ets.insert(@compile_closure_table, {cache_key, built}) + built + end + + wrap_with_queue(queue, eval_fn) end # Attempts to lower `fun.(vars)` to an `EMLX.Native.Expr` program and build a diff --git a/emlx/lib/emlx/application.ex b/emlx/lib/emlx/application.ex index 706fd0b..2d874e5 100644 --- a/emlx/lib/emlx/application.ex +++ b/emlx/lib/emlx/application.ex @@ -47,6 +47,7 @@ defmodule EMLX.Application do @doc false def start(_type, _args) do EMLX.Profiling.init() + EMLX.init() ensure_default_worker!(:cpu, _gpu_optional? = false) ensure_default_worker!(:gpu, _gpu_optional? = true) ensure_worker!(:runtime_call_worker, :cpu, _gpu_optional? = false)