DEV Community

Cover image for Keras #22561 speedups: compiled Keras on PyTorch now generates text 11.6x faster
Rahul Kumar
Rahul Kumar

Posted on

Keras #22561 speedups: compiled Keras on PyTorch now generates text 11.6x faster

#python #pytorch #performance #machinelearning

Between master on Jul 18 and master on Oct 15, 2026, 15 of our PRs landed in Keras. Same GPU, same torch and jax versions, same models; only the Keras source changes. On the issue's test LM, compiled Keras[torch] now generates text 11.6× faster, torch.compile went from slowing Keras down to making it 6.56× faster than eager Keras, and an LLM forward makes 12–23% fewer Python calls on all four backends.

Here is compiled generation at every merge. Pick a config or a workload, and hover a point to see its runs:

Generate-32 throughput, July master vs today

Tokens per second from the 2.1M-parameter test LM, 128 tokens per call (4 sequences × 32 steps); longer bars are faster. In July, compiling a Keras[torch] model made it slower: 1,458 tok/s compiled against 2,139 eager. Today the compiled model runs at 16,975 tok/s: 6.56× faster than eager Keras, within 1.53× of plain compiled PyTorch, down from 17.8×.

Generate-32 throughput in tokens per second, July versus October master. Keras on PyTorch with torch.compile: 1,458 to 16,975 (11.6 times); eager: 2,139 to 2,588; plain PyTorch 26,003 compiled and 21,870 eager. Keras on JAX under jit: 55,734 and 52,200; eager 437 to 479; plain JAX 57,642 jit and 1,438 eager.

Three merges that made torch.compile work with Keras

Three merges each turned one compiled workload from slow to fast in a single step.

Three cards: #23327, merged Aug 30, compiled LLM forward 2.50 to 0.326 ms, 7.68 times faster; #23201, merged Sep 11, compiled CNN forward 0.945 to 0.156 ms, 6.05 times faster; #23647, merged Sep 30, compiled generate-32 57.2 to 7.82 ms, 7.32 times faster.

The three PRs: #23327 (attention), #23201 (dict ordering in the tree code) and #23647 (identity checks in the tree code). Graph breaks per forward pass, at every merge:

Step chart of torch.compile graph breaks per forward pass at every merge: the LLM drops from 11 to 0 at #23327 and the CNN from 9 to 0 at #23201; both stay at 0 through October 15.

Less Python in every Keras call, on all four backends

Most of the series changed shared code (Layer.call, op tracebacks, dtype handling, input checks), so the saving reaches every backend, not just torch. Python function calls per forward pass, July to today: PyTorch −23% on the LM and −33% on the CNN; JAX −12% on the LM and −17% on the CNN; TensorFlow −13% on the LM and −19% on the CNN; NumPy −16% on the LM and −18% on the CNN. Adding #23336 takes NumPy to −36% on the LM. These counts are exact and repeat call for call on any machine.

Keras[jax]

Under jax.jit the Python runs once at trace time, so Keras[jax] was already close to plain JAX and stays there: 2.30 → 2.45 ms per generate-32 against 2.22 ms for plain JAX. Eager Keras[jax] is where the saving shows in the exact counts: −12% Python calls per forward. The timing agrees: the sweep gives 293 → 267 ms (1.10×) and the back-to-back A/B rounds give 1.08×.

Every config, every workload

Show all 12 configs and workloads (ms per call)
Keras config Workload Plain framework (ms) Jul 18 (ms) Oct 15 (ms) Change + #23336 (ms) A/B check
Keras[torch] · eager CNN forward 0.064 0.668 0.477 1.40× 0.466 1.39×
Keras[torch] · eager LLM forward 0.189 1.88 1.55 1.22× 1.48 1.20×
Keras[torch] · eager LLM generate-32 5.85 59.8 49.5 1.21× 47.2 1.20×
Keras[torch] · torch.compile CNN forward 0.076 1.28 0.160 8.03× 0.161 8.41×
Keras[torch] · torch.compile LLM forward 0.160 2.69 0.302 8.93× 0.299 9.60×
Keras[torch] · torch.compile LLM generate-32 4.92 87.8 7.54 11.6× 7.66 12.4×
Keras[jax] · eager CNN forward 0.508 1.34 1.23 1.08× 1.18 1.15×
Keras[jax] · eager LLM forward 2.79 8.46 8.38 1.01× 7.90 1.12×
Keras[jax] · eager LLM generate-32 89.0 293 267 1.10× 269 1.08×
Keras[jax] · jax.jit CNN forward 0.077 0.079 0.081 0.97× 0.080 1.02×
Keras[jax] · jax.jit LLM forward 0.120 0.131 0.135 0.97× 0.130 1.05×
Keras[jax] · jax.jit LLM generate-32 2.22 2.30 2.45 0.94× 2.27 1.05×

A/B check: July ÷ today from 3 separate back-to-back rounds (higher is a bigger speedup).

Each of our PRs, measured where it landed

Each PR was timed on the commit just before and just after it merged, so every row is that one PR's own effect. Δ is just after the merge minus just before it; negative is better.

Show the per-PR table
PR Merged What it changed Δ calls/fwd torch JAX TF NumPy Δ eager gen-32 (ms) Δ compiled gen-32 (ms) Graph breaks
#23183 07-28 Traceback details are built only when an op fails, not on every call −598 −586 −622 −546 −2.9 −5.2 ·
#23327 08-30 Removed the torch.compile graph break in the torch backend's dot_product_attention +6 0 0 0 +0.40 −12.3 LLM 11 → 0
#23189 09-03 Fixed an always-true device check that sent every tensor through .to() 0 +38 0 0 +0.28 −1.1 ·
#23200 09-05 Layer.call stops re-setting built/_called on every call −1,848 −1,098 −1,188 −1,100 −4.9 −10.5 ·
#23182 09-10 O(1) dtype membership and a fast path for canonical dtype strings −300 −82 −497 −116 −0.89 −1.5 ·
#23199 09-10 Dropped a redundant channels_last copy on torch conv inputs 0 −50 0 0 +0.40 −0.29 ·
#23201 09-11 Torch tree code made traceable: canonical dict key order −142 +8 0 0 +0.55 −5.2 CNN 9 → 0
#23296 09-11 Keeps call-context kwargs on the recorded Node (fix) −109 −78 −14 −46 −0.24 −0.09 ·
#23187 09-12 TorchTrainer skips a full flatten walk to find the batch size 0 −26 0 0 +0.14 −0.21 ·
#23301 09-12 Layer name scope opened once per call, not twice −330 −568 −638 −330 −0.89 −0.21 ·
#23190 09-15 torch slice() keeps dynamic dims under torch.export 0 −16 0 0 −0.23 −0.16 ·
#23647 09-30 Torch tree code stops comparing id()s 0 −8 0 0 −0.58 −49.4 ·
#23645 10-07 torch pooling ops run in channels_last 0 −16 0 0 +0.10 −0.05 ·
#23184 10-13 Fast path for single-InputSpec input checks −102 −59 −63 −63 +0.06 +0.05 ·
#23920 10-15 Functional models reject a mismatched last axis (fix) 0 −54 0 0 +0.10 +0.02 ·
#23336 not merged yet (approved) Memoized dtype promotion in result_type −713 −378 −1,471 −3,399 −2.2 +0.12 ·

Hot functions in one Keras[torch] LLM forward

Show the hot-functions table
Function Jul 18 Oct 15 Change + #23336
isinstance 2,600 1,740 −33% 1,680
traceback error wrappers 74 48 −35% 48
thread-local config lookups 624 513 −18% 513
standardize_dtype 177 244 +38% 184
tree.flatten 93 67 −28% 67
tree.map_structure 102 62 −39% 62
InputSpec checks 14 14 +0% 14
CallSpec objects built 23 23 +0% 23
inspect.Signature.bind 22 22 +0% 22
convert_to_tensor 86 86 +0% 86
get_device 138 138 +0% 138

How this was measured

  • Machine: one NVIDIA RTX PRO 6000 Blackwell, kept idle during the runs.
  • Software: the same for every run: torch 2.13.0, jax 0.11.2, Python 3.14.7. Only the Keras source changes.
  • Commits: 31 snapshots of master: Jul 18 (a2f8eac8), just before and just after each of the 15 merges, Oct 15 (cc9de2b31), and Oct 15 plus #23336.
  • Models: the issue's two test models: a small CNN (114k parameters) and a 2-layer transformer LM (2.1M parameters, 32-token greedy generate).
  • Timing: the issue's benchmark script, one process at a time, warm-up first, a GPU sync after every call. Each number is a median, and the whole sweep ran three times for PyTorch and twice for JAX.
  • Cross-check: a separate back-to-back run of July against October gives 12.4× compiled and 1.20× eager (the sweep: 11.6× and 1.21×).
  • Counts: Python calls from cProfile, graph breaks from torch._dynamo.explain. These are exact and repeat on any machine.
  • Not covered: training, CPU, larger models, and TensorFlow or NumPy timing (call counts only).

The work is tracked in issue #22561. Thank you to the Keras maintainers who reviewed and merged these PRs.

Top comments (0)