perf(torch): eager CallSpec fast path skipping signature.bind in Layer.__call__#23198
perf(torch): eager CallSpec fast path skipping signature.bind in Layer.__call__#23198pctablet505 wants to merge 11 commits into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces an eager fast path for CallSpec to optimize the common case of a single positional tensor and no keyword arguments, bypassing the overhead of inspect.Signature.bind. Tests are added to verify that the fast path produces identical fields to the slow path and correctly excludes unsupported signatures. The review feedback suggests making the nested default argument check more robust by guarding len(v) with hasattr(v, "__len__") to prevent potential TypeError with custom PyTree classes, and optimizing the eager flag assignment in the fast path by using a fast isinstance check instead of backend.is_tensor.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| defaults_are_safe = all( | ||
| not is_backend_tensor_or_symbolic(v) | ||
| and not (tree.is_nested(v) and len(v) > 0) | ||
| for v in self._call_arg_defaults.values() | ||
| ) |
There was a problem hiding this comment.
In defaults_are_safe, len(v) is called on any default value v where tree.is_nested(v) is True. However, custom PyTree classes or other custom container types registered as nested structures might not implement __len__, which would raise a TypeError during layer initialization.\n\nTo make this check more robust and prevent potential runtime errors with custom nested default arguments, consider guarding the len(v) call with hasattr(v, "__len__").
defaults_are_safe = all(\n not is_backend_tensor_or_symbolic(v)\n and not (tree.is_nested(v) and hasattr(v, "__len__") and len(v) > 0)\n for v in self._call_arg_defaults.values()\n )| call_spec.tensor_arguments_names = [first_name] | ||
| call_spec.nested_tensor_argument_names = [] | ||
| call_spec.first_arg = x | ||
| call_spec.eager = backend.is_tensor(x) |
There was a problem hiding this comment.
Since _fast is only called when is_backend_tensor_or_symbolic(x) is True, x is guaranteed to be either a backend tensor or a KerasTensor.\n\nWe can avoid the overhead of calling backend.is_tensor(x) (which performs backend-specific checks and dispatching) by using a fast Python isinstance check against backend.KerasTensor. Since x is either a backend tensor or a KerasTensor, if it is not a KerasTensor, it must be a backend tensor (e.g., eager).
| call_spec.eager = backend.is_tensor(x) | |
| call_spec.eager = not isinstance(x, backend.KerasTensor) |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #23198 +/- ##
==========================================
+ Coverage 80.39% 84.10% +3.70%
==========================================
Files 465 465
Lines 69188 69385 +197
Branches 11379 11434 +55
==========================================
+ Hits 55622 58353 +2731
+ Misses 10742 8084 -2658
- Partials 2824 2948 +124
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
pctablet505
force-pushed
the
perf-P13-eager-callspec
branch
2 times, most recently
from
545068c to
e674f81
Compare
…ultiHeadAttention Part of keras-team#22561.
…ribution() hot paths Part of keras-team#22561.
…sor start indices
pctablet505
force-pushed
the
perf-P13-eager-callspec
branch
from
e674f81 to
23f56ce
Compare
test_quantize_with_specific_equations draws both the layer kernel (via build()) and the input from the global RNG, which TestCase.setUp does not seed. The int4 cases run on very small tensors (as few as 8 output elements), so the quantization MSE has high run-to-run variance: the typical value sits ~5x under the tolerance, but roughly 1% of unseeded draws exceed the int4 threshold, intermittently failing the numpy and openvino CI legs. Seed the RNG (set_random_seed(1337)) at the start of the test so both the kernel init and the input are deterministic. Verified across numpy, openvino, torch, tensorflow and jax: all six int8/int4 cases pass with the worst MSE still >2x under its tolerance. No production code changes. (cherry picked from commit 907e506)
3 participants
Summary
Layer.__call__builds aCallSpec— including a fullinspect.Signature.bind,apply_defaults, and a per-argument classification loop — on every layer call. For the overwhelmingly common case oflayer(x)(one positional tensor, no kwargs) this is pure Python dispatch overhead on the hot path: ~1.9 µs/call measured, onesignature.bindper layer call (≈ 70k per 32-token generation of a small transformer — see RC1 in #22561).This adds an eager fast path that reproduces the
CallSpecfields directly for that case and skipssignature.bindentirely. Anything outside the common case falls through to the unchanged slow path.Mechanism
Layer.__init__, precompute once fromself._call_signature: the argument names, the first-parameter name, the non-first defaults, and an eligibility flagself._call_spec_fast_path.Trueonly when the signature is simple: ≥ 1 parameter, the first is positional, no*args/**kwargs, every non-first parameter has a default, and no default is a tensor or nested structure. Otherwise it isFalseand the layer always takes the slow path.CallSpec._fast(...)constructs the instance via__new__and sets exactly the fields the slow__init__produces for a single positional tensor with no kwargs.Layer.__call__uses it only whenself._call_spec_fast_path and not kwargs and len(args) == 1 and (is_tensor(args[0]) or isinstance(args[0], KerasTensor)).Why it is safe
call(); it touches no tensors, ops, or the autograd/trace graph, so it is numerically identical on every backend (verified on numpy/torch/jax/tensorflow, eager and symbolic).*args/**kwargs, keyword-only-required signatures, and any call with kwargs or a non-tensor first arg all fall through to the existingCallSpec(...)._resolve_and_populate_arg): the call-context-arg resolution reads the same reproduced fields, sotrainingpropagation and the sibling-leak fix are both preserved.Tests
layer_test.py:CallSpecfield-equality matrix across representative signatures (call(inputs),call(inputs, training=None),call(inputs, training=None, mask=None), scalar-default), each for eager tensors and symbolicKerasTensors;Denseforward is bit-identical with the fast path enabled vs forced-slow.Full
layer_test.pypasses (62 passed, 2 skipped on torch).Series context
Part of a 13-PR series for #22561 (root cause RC1). Wave 1 merge order: #23182 -> #23183 -> #23184 -> #23185 -> #23186 -> #23187 -> #23188 -> #23189 -> #23190; independent follow-ups #23198, #23199, #23200, #23201 — this is #23198.
compute_output_spec's separateCallSpec(symbolic shape inference) is intentionally left unchanged. Toucheslayers/layer.py, also touched by #23200 (RC18) — the hunks are disjoint (__init__/CallSpecvs thebuilt/_calledlines in__call__).Contributor Agreement
📊 Performance & stack position
PR 10 of 13 in the #22561 torch eager-overhead series (merge order) · base
master· prev: #23190 · next: #23199Stacked on top of #23190 — this PR's own change is the top commit; the earlier commits are its parent PRs (CI here exercises the cumulative state). Numerically equivalent to
master(generate-32 token-ids identical, Δ = 0).Cumulative through this PR vs
master— Keras[torch], torch 2.12.1+cu130 / RTX 4050, GPU eager, best-of-3 of median-20:masterFramework context (generate-32 eager, same GPU): raw PyTorch 11.4 ms · raw JAX 985 ms · Keras[jax] 584 ms. Keras[torch] carries the per-call Python dispatch tax this series reduces; Keras[jax] traces it away. Full matrix + methodology: #22561 §7–§8.
Reproduction scripts: gist.