pctablet505
Work / Performance and inference

Keras on PyTorch: compiled generation 11.6x faster

Traced why Keras on PyTorch ran about 9x slower than plain PyTorch, and merged the fixes upstream, one root cause each.

  • Performance and inference
  • Google
  • Jul 2026 – Oct 2026
  • Merged
  • 2 min read
11.6x
faster 32-token LLM generation under torch.compile, Keras master on July 18 against October 15 (87.8 → 7.54 ms per call)
● checked 28 Oct 2026 · source
20 → 0
torch.compile graph breaks per forward pass (CNN 9 → 0, LLM 11 → 0)
● checked 28 Oct 2026 · source
−12 to −23%
Python function calls per forward pass, on all four Keras backends
● checked 28 Oct 2026 · source
1.21x
faster 32-token generation in eager mode over the same months (59.8 → 49.5 ms per call)
● checked 28 Oct 2026 · source
8.45x
the time eager Keras on PyTorch still takes against plain PyTorch on the same benchmark (10.2x in July)
● checked 28 Oct 2026 · source
9x
slower than plain PyTorch on small models in eager mode, about, when the issue was opened
● checked 28 Oct 2026 · source
Headline numbers, Keras master on July 18 against October 15: compiled and eager generation, Python calls per forward, graph breaksHeadline numbers, Keras master on July 18 against October 15: compiled and eager generation, Python calls per forward, graph breaks
Headline numbers, Keras master on July 18 against October 15: compiled and eager generation, Python calls per forward, graph breaks. From the progress update on keras#22561, posted 18 Oct 2026.

Keras runs on PyTorch, but on small models it ran about 9x slower than the same model in plain PyTorch. I traced the cost to Python work Keras repeated on every call, opened keras#22561 with the root causes, and landed the fixes as small pull requests, one cause each. Each was timed on the commit just before and just after its merge.

Problem

Most of each call went to Keras's own Python: building traceback details that were only needed when an op failed, re-setting layer state, checking dtypes and inputs. torch.compile could not help either. Graph breaks split every forward pass and compiled steps kept recompiling, so compiling a Keras model made generation slower than running it eagerly.

Approach

  • Counted every Python call with cProfile and every torch.compile graph break with torch._dynamo.explain, on the issue's two models: a small CNN and a two-layer transformer language model.
  • Split one large fix into small pull requests, one root cause each, each with tests and a before-and-after benchmark.
  • Three merges made torch.compile work with Keras: one removed the graph break in attention, one made the PyTorch tree code traceable, and one stopped compiled steps from recompiling.
  • Most fixes changed shared code (Layer.__call__, op tracebacks, dtype handling, input checks), so the JAX, TensorFlow and NumPy backends gained too.

Results

By the progress report of October 18, 15 pull requests were merged. Between Keras master on July 18 and on October 15, compiled 32-token generation got 11.6x faster and eager generation 1.21x faster. Graph breaks per forward pass fell from 20 to 0, and Python calls per forward pass fell 12 to 23% on all four backends. The progress report times every pull request on its own; its interactive version on DEV is under Links.

What did not work

The eager gap did not close: eager Keras on PyTorch still takes 8.45x the time of plain PyTorch on this benchmark, so the issue stays open. When the report was posted, two more fixes were in review. Training, CPU timing and larger models were not measured.

Charts from the report

From the progress update on keras#22561, posted 18 Oct 2026.

Three merges that each made one compiled Keras-on-PyTorch workload several times fasterThree merges that each made one compiled Keras-on-PyTorch workload several times faster
Three merges that each made one compiled Keras-on-PyTorch workload several times faster.
Compiled 32-token generation from July to October, one bar per pull request, the rest of master as one grey barCompiled 32-token generation from July to October, one bar per pull request, the rest of master as one grey bar
Compiled 32-token generation from July to October, one bar per pull request, the rest of master as one grey bar.
Tokens per second for 32-token generation, July against October, Keras on PyTorch and JAX, with the plain frameworks for referenceTokens per second for 32-token generation, July against October, Keras on PyTorch and JAX, with the plain frameworks for reference
Tokens per second for 32-token generation, July against October, Keras on PyTorch and JAX, with the plain frameworks for reference.
Python function calls per LLM forward on PyTorch, JAX, TensorFlow and NumPy, July against October and at every commitPython function calls per LLM forward on PyTorch, JAX, TensorFlow and NumPy, July against October and at every commit
Python function calls per LLM forward on PyTorch, JAX, TensorFlow and NumPy, July against October and at every commit.
Compiled 32-token generation at every commit from July to October, and torch.compile graph breaks per forwardCompiled 32-token generation at every commit from July to October, and torch.compile graph breaks per forward
Compiled 32-token generation at every commit from July to October, and torch.compile graph breaks per forward.
Keras on JAX, eager and under jax.jit, 32-token generation at every commitKeras on JAX, eager and under jax.jit, 32-token generation at every commit
Keras on JAX, eager and under jax.jit, 32-token generation at every commit.

More in performance and inference

  • GLM-5.3-Flash on one GPU2 Sep 2026

    Served a 181 GiB mixture-of-experts LLM from one 96 GiB GPU and took decoding from 10.8 to 15.0 tokens/s shipped, 16.5 at peak. · Independent · Published

  • One call from a Keras model to a phone: LiteRT exportDec 2025 – Jul 2026

    Built the LiteRT export path in Keras and KerasHub: a trained model to a phone-ready file in one call, on TensorFlow and PyTorch. · Google · Merged

  • End-to-end UI tests for XProfAug 2026 – now

    Wrote the end-to-end UI test suite for XProf, the open-source ML profiler, with 107 browser tests and a side-by-side check of two builds. · Google · Shipped

  • Llama 3.1 in KerasHub, and Gemma 3 fixesApr 2025 – Jul 2026

    Implemented Llama 3.1 in KerasHub and fixed Gemma 3 bugs, including a preprocessor crash across Gemma 3, 3n and 4. · Google · Merged