Skip to content

Cache jit closures instead of re-tracing on every call - #137

Merged
polvalente merged 3 commits into
elixir-nx:mainfrom
Codcore:cache-jit-closures
Aug 29, 2026
Merged

Cache jit closures instead of re-tracing on every call#137
polvalente merged 3 commits into
elixir-nx:mainfrom
Codcore:cache-jit-closures

Conversation

@Codcore

@Codcore Codcore commented Aug 20, 2026

Copy link
Copy Markdown
Contributor

Closes #136.

__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 arrive at a program already sitting
in the program table.

This caches 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.

What is left alone

The program table and its structural key. That key is what lets structurally
identical layers share one compiled program across different call sites, and it
still does. The new table sits above it and short-circuits repeat calls of the
same call site with the same argument shapes — the shape of every decode loop and
every training step.

Calls carrying :command_queue or :hooks are not cached: those options
reference caller-owned state, and a process-wide table should not extend its
lifetime. Such calls take the previous path unchanged.

Measurements

A chain of depth cheap ops, so the timing reflects graph traversal rather than
arithmetic (script in #136):

  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

At depth 1024 that is 3.2× faster, and what remains between jit and compile
is 0.55 ms rather than 4.48 ms.

On a real workload — a 40-layer Qwen3.6-35B-A3B MoE decode step, roughly 4300
nodes — the change removes 8.35 ms per token out of 89.

Tests

mix test: 2668 passed (825 doctests, 1843 tests), 7 excluded. No failures.

A note on the trade-off

The cache is unbounded, like the program table it sits above. Each entry holds a
closure referencing a compiled program, so a caller that jits the same function
across many distinct argument shapes will accumulate entries. If that is a
concern, the same key would work with a bounded table; I kept it simple to match
the existing behaviour rather than introduce a second eviction policy.

__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.
Comment thread emlx/lib/emlx.ex Outdated
Comment thread emlx/lib/emlx.ex Outdated
@polvalente
polvalente merged commit 1d99042 into elixir-nx:main Aug 29, 2026
10 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Nx.Defn.jit re-traces the function body on every call

2 participants