perf(torch): SDPA is_causal dispatch + bounded causal-mask cache in MultiHeadAttention#23186
perf(torch): SDPA is_causal dispatch + bounded causal-mask cache in MultiHeadAttention#23186pctablet505 wants to merge 5 commits into
Conversation
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
![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 optimizes the MultiHeadAttention layer by introducing a bounded cache for computed causal masks and adding an ultra-fast path that directly calls PyTorch's native scaled dot-product attention when using causal masks. The review feedback suggests restricting the causal mask cache strictly to the PyTorch backend to prevent tracer leaks and graph mismatch errors in JAX and TensorFlow. Additionally, it recommends removing unnecessary .contiguous() calls on transposed tensors in the PyTorch fast-path to avoid redundant memory copies.
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.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #23186 +/- ##
==========================================
- Coverage 84.70% 84.09% -0.61%
==========================================
Files 465 465
Lines 69335 69735 +400
Branches 11407 11508 +101
==========================================
- Hits 58729 58645 -84
- Misses 7672 8145 +473
- Partials 2934 2945 +11
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:
|
2b936aa to
2dcbf92
Compare
2dcbf92 to
0838c6c
Compare
27923d3 to
ff2f9f0
Compare
…ultiHeadAttention Part of keras-team#22561.
_compute_causal_mask evicted the oldest entry via next(iter(cache)) followed by cache.pop(key) with no default. That check-then-act sequence could KeyError if a shared layer instance is accessed concurrently and another thread empties the cache between the two calls. Guard against an empty iterator and pass a default to pop so eviction is safe even in that unlikely case.
ff2f9f0 to
b360a76
Compare
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request optimizes MultiHeadAttention for the PyTorch backend by introducing a fast path that routes causal attention directly to PyTorch's native scaled_dot_product_attention (SDPA) with is_causal=True, bypassing mask materialization. It also adds a bounded FIFO cache (up to 8 entries) for causal masks on PyTorch with static shapes to avoid redundant mask creation. Unit tests are added to verify correctness, fallback paths, and cache eviction. The feedback notes that the transposed tensors are not made contiguous before being passed to SDPA, which contradicts the PR description and could lead to performance issues or runtime errors on certain CUDA kernels.
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.
Rewrite the two comments around the torch SDPA causal dispatch as flowing sentences instead of a colon-led topic heading, and drop a stray series-tag comment above the new tests.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request optimizes the MultiHeadAttention layer for the PyTorch backend by dispatching directly to the native causal scaled_dot_product_attention kernel and introducing a bounded cache for causal masks to avoid redundant computations. The review feedback highlights two important issues in the caching implementation: first, using isinstance checks for sequence lengths can mistakenly allow symbolic integers (torch.SymInt) to be cached under torch.compile, which should be replaced with strict type checks; second, the FIFO eviction mechanism is not thread-safe and could raise a RuntimeError under concurrent execution, which can be resolved by atomically clearing the cache.
Require sequence lengths to be exact Python ints (type(x) is int) before caching, so a symbolic torch.SymInt under torch.compile is excluded and falls through to recomputation instead of being used as an unhashable cache key. Back the cache with an OrderedDict and evict the oldest entry via popitem(last=False), a single atomic pop that avoids iterating a dict another thread could resize concurrently. Add a regression test that a symbolic SymInt length stays out of the cache while concrete int lengths are cached under exact-int keys.
3 participants
Closes #23277
Torch backend: when
MultiHeadAttentiondoes eager causal self-attention with no explicit mask, dispatch straight totorch.nn.functional.scaled_dot_product_attention(..., is_causal=True)— skipping the[T, S]mask materialization and theops.dot_product_attentiondispatch layer. Off that fast path, the fallback causal mask is cached (bounded at 8 entries) instead of rebuilt on every call.Why it's safe
use_causal_mask=True, no explicitattention_mask, zero dropout, noreturn_attention_scores, flash attention off. Anything else falls through to the existing path untouched, so other backends and configurations cannot regress.(B, T, N, H)to SDPA's(B, N, T, H)layout and back. The head dimension stays contiguous, so no.contiguous()copy is needed, and results match the explicit lower-triangular-mask path (verified, including a T ≠ S case).(q_len, kv_len, device, "bool"); dtype is deliberately omitted since the bool mask depends only on lengths and device. Only exact Pythonintlengths are cached — undertorch.compilelengths become symbolicSymInts (unhashable), which simply fall through to recomputation, so compiled models keep working. JAX and TF are excluded entirely: caching there would leak a tracer acrossjit/tf.functiontraces, and those backends constant-fold the mask anyway.call()overrides and never reach this code.Benchmarks (GPU, 3-way: master / this PR alone / this PR + parents #23182→#23185)
Interleaved A/B process pairs against an A/A null calibration; "~noise" means below the measured noise floor, not claimed as a win. Forward outputs are bit-identical to master with this PR in isolation, and generate-32 token ids match in all states (the small stacked-state float32 delta, 1.67e-06, comes from a parent PR's einsum reordering already accepted there). torch 2.13.0+cu130, RTX 4050 Laptop GPU, master a2f8eac.
Tests
New unit tests cover fast-path activation, numeric equivalence with the explicit-mask path, every fallback gate (non-torch backend, flash attention, attention scores, explicit mask), the cache size bound, and
SymIntexclusion.Contributor Agreement