[pctablet505]

Summary

torch.compile cannot trace the flash attention probe in the torch backend's
dot_product_attention, so it breaks the graph at every attention call. On a
six-layer transformer forward pass this produces 19 graph breaks and 20 separate
graphs.

MultiHeadAttention routes through ops.dot_product_attention when dropout is
0, attention scores are not returned, and inputs are 4D — the default path for
most models.

Root cause

dot_product_attention calls _can_use_flash_attention when flash_attention
is None (the default), and again with raise_error=True when it is explicitly
True. That helper constructs a pybind11 SDPAParams object to probe kernel
support:

spda_params = SDPAParams(query, key, value, mask, 0.0, is_causal, False)

Dynamo cannot trace a pybind11 constructor, so tracing stops there and resumes
afterwards — once per attention call.

Change

Report flash attention as unavailable while tracing, so the probe is never
reached and the call takes its ordinary non-flash branch. scaled_dot_product_attention
then selects the kernel through its own dispatcher, which is what it does on the
non-flash branch already. dot_product_attention itself is unchanged — the whole
change is a 12-line early return at the top of _can_use_flash_attention.

This applies to auto-detection only. An explicit flash_attention=True still
runs the probe, still breaks the graph, and still raises
Flash attention is not supported with the provided inputs — deliberately. Making
the guard raise instead does not work: dynamo discards an exception raised while
tracing and re-runs the frame eagerly, where is_compiling() is False, so the
message never reaches the user and the frame is de-optimized rather than split.

Eager is unchanged by construction: torch.compiler.is_compiling() is False
outside torch.compile, so the probe and the pinned flash-attention kernel
context still run exactly as before.

Results

torch._dynamo.explain, float32 unless noted. These counts are deterministic,
and are identical on an RTX 4050 (sm_89, torch 2.13.0+cu130) and an RTX PRO 6000
Blackwell (sm_120, torch 2.13.0+cu132):

flash_attention		graphs	breaks	ops captured
None (default)	before	20	19	178
None (default)	after	1	0	367
True	before	2	1	16
True	after	2	1	16

Op count rises because a fragmented graph only captures a fraction of the model;
with one graph, dynamo captures all of it.

The None rows are the six-layer model below. The True rows are a single
ops.dot_product_attention call in float16, since float32 with
flash_attention=True raises on both master and this branch.

harness

import keras, torch
from keras import layers
B, S, DM, H, LAYERS = 8, 128, 512, 8, 6
inp = keras.Input(shape=(S, DM)); x = inp
for _ in range(LAYERS):
    a = layers.MultiHeadAttention(num_heads=H, key_dim=DM // H)(x, x)
    x = layers.LayerNormalization()(x + a)
    f = layers.Dense(DM * 2, activation="relu")(x)
    x = layers.LayerNormalization()(x + layers.Dense(DM)(f))
model = keras.Model(inp, x).to("cuda")
t = torch.randn(B, S, DM, device="cuda")
print(torch._dynamo.explain(lambda z: model(z, training=False))(t))

Wall clock for the same six-layer forward, float32, on the RTX PRO 6000
Blackwell. Median of 8 A/B pairs with the order alternated per pair, each pair
30 repeats of 20 iterations timed with CUDA events. Running the base against an
identical checkout of itself measures run-to-run noise at 1.7% eager and 7.6%
compiled, putting the significance threshold at 2.5% and 11.4% respectively:

	eager	compiled
before	5.98 ms	8.05 ms
after	5.98 ms	0.98 ms

Compiled is 8.2x faster, all 8 pairs agreeing in direction, 1.03x spread across
them. Eager lands at 1.005x, inside its threshold. Before this change,
torch.compile makes this model slower than its own eager forward.

Behavior

Kernel dispatch is unchanged. Measured across fp16 / bf16 / fp32, with and without
a mask, and at head_dim 64, 128 and 320, the kernel selected by
scaled_dot_product_attention is identical to master in every case, in both eager
and compiled mode. On sm_120 that spans 72 configurations, with no difference in
any of them.

Under compile the auto-detect path no longer pins the kernel via
sdpa_kernel([FLASH_ATTENTION]), leaving the selection to the dispatcher. Where
the dispatcher lands on the kernel the pin would have chosen, output is
bit-identical: on sm_120, max|compiled − eager| is exactly 0.0 across 18
configurations spanning fp32 / fp16 / bf16, head_dim 64 and 128, plain, masked
and causal. Where it lands elsewhere the two can differ slightly; on sm_89 fp16
shows about 1.2e-04 and fp32 about 6e-08.

Tests

keras/src/backend/torch/nn_test.py adds one parameterized test that compiles
ops.dot_product_attention with fullgraph=True, so a graph break becomes an
error, and compares against eager across six argument shapes: plain, causal,
masked, masked-and-causal, bias, and custom scale.

All six fail against master and pass with this change, both on a CPU-only torch
build as CI runs it and on CUDA (sm_120).

Scope, stated honestly: the cases are float32 and flash_attention=None, so they
cover the auto-detect path only. Grouped-query inputs, explicit flash_attention
values, and the hardware-dependent flash-vs-dispatcher divergence above are not
covered, and CI runs CPU-only so the kernel-dispatch matrix is not covered there
either.

• I am a human, and not a bot.
• I will be responsible for responding to review comments in a timely manner.
• I will work with the maintainers to push this PR forward until submission.

[gemini-code-assist]

Code Review

This pull request prevents graph breaks during torch.compile in the PyTorch backend's dot_product_attention by bypassing the flash-attention support probe when compiling is active, deferring the kernel choice to PyTorch's native scaled_dot_product_attention. It also adds comprehensive compilation tests. The feedback suggests making the compilation check more robust by verifying that torch has the compiler attribute before accessing it, which avoids potential AttributeError in older PyTorch environments.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 0% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 71.16%. Comparing base (6944627) to head (b150544).
⚠️ Report is 22 commits behind head on master.

Files with missing lines	Patch %	Lines
keras/src/backend/torch/nn.py	0.00%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #23327      +/-   ##
==========================================
+ Coverage   70.91%   71.16%   +0.24%     
==========================================
  Files         496      496              
  Lines       84908    85536     +628     
  Branches    13893    14086     +193     
==========================================
+ Hits        60215    60870     +655     
+ Misses      21664    21638      -26     
+ Partials     3029     3028       -1     

Flag	Coverage Δ
keras	71.00% <0.00%> (+0.24%)	⬆️
keras-cpu	70.23% <0.00%> (+5.13%)	⬆️
keras-gpu	69.95% <0.00%> (+11.56%)	⬆️
keras-jax	49.24% <0.00%> (+0.31%)	⬆️
keras-numpy	45.07% <0.00%> (-0.01%)	⬇️
keras-openvino	49.84% <0.00%> (-0.01%)	⬇️
keras-tensorflow	50.56% <0.00%> (+1.30%)	⬆️
keras-torch	50.15% <0.00%> (+0.31%)	⬆️
keras-tpu	57.46% <0.00%> (+9.51%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[hertschuh]

Thanks for looking into this!

[hertschuh]

Instead of replicating the code to scaled_dot_product_attention, I would simply move this test at the beginning of _can_use_flash_attention.

This would be more true to the spirit of what the flash_attention argument means. If you compile, you cannot use flash attention, and the user should be aware of that.

[pctablet505]

Moved it, dot_product_attention is now identical to master, and the change is just a guard at the top of _can_use_flash_attention.

Two corrections from measuring it on an RTX 4050 (sm_89), though:

"If you compile, you cannot use flash attention" isn't the case. With the probe skipped and fp16 inputs, scaled_dot_product_attention still dispatches to pytorch_flash::flash_fwd_kernel (torch.profiler, probe executions 0). Only the pybind11 SDPAParams construction is untraceable, not the kernel, so returning False here costs no flash attention.

Making the user aware via raise_error doesn't work either; I tried it first. Dynamo discards an exception raised while tracing and re-runs the frame eagerly, where is_compiling() is False. So the probe runs anyway, the message never surfaces, and op_count drops to 1 from master's 16. Worse than the break it was removing.

So the guard only covers auto-detection (not raise_error). An explicit flash_attention=True still behaves exactly as on master and reports "Flash attention is not supported with the provided inputs".

[hertschuh]

"If you compile, you cannot use flash attention" isn't the case. With the probe skipped and fp16 inputs, scaled_dot_product_attention still dispatches to pytorch_flash::flash_fwd_kernel (torch.profiler, probe executions 0). Only the pybind11 SDPAParams construction is untraceable, not the kernel, so returning False here costs no flash attention.

So what you're saying is that scaled_dot_product_attention is always and automatically "upgraded" to flash attention when you compile? So we cannot really apply flash_attention=True/False, it's always going to be True?

[pctablet505]

Short answer: no. flash_attention=True still works exactly as before, and
flash_attention=False never turned flash attention off — that is already true
on master, in eager. This PR changes neither of it.

What the flag does

• True — pins the flash kernel and raises if the input cannot use it. This PR
  leaves it alone (the guard is not raise_error), so it still raises under
  torch.compile.
• False — does not forbid flash. It skips the pin and calls
  scaled_dot_product_attention with no sdpa_kernel(...) context, so PyTorch's
  dispatcher chooses — and for fp16/bf16 it chooses flash.
• None — probes, then does one of the above. Under compile this PR skips the
  probe, landing on exactly the same code path as False.

Measured dispatch, classified by the CUDA kernel that actually runs:

input	None	True	False
fp16/bf16, no mask	flash	flash	flash
fp16/bf16, masked	mem-efficient	raises	mem-efficient
fp32	mem-efficient	raises	mem-efficient

Same in eager and compiled. I ran all 72 combinations — 3 dtypes × mask/no-mask ×
head_dim 64/128 × 3 flag values × eager/compiled — on master and on this branch:
both pick the same kernel in every one, 0 differences.

Two things follow from the bold cell:

• It is master, in eager: flash_attention=False runs
  pytorch_flash::flash_fwd_kernel. So "False does not disable flash" is
  pre-existing behavior, not something this PR introduces.
• fp32 never uses flash at all (flash requires fp16/bf16), so on the default
  float32 path this change alters the selected kernel in zero cases.

If False is meant to genuinely forbid flash, that is a real gap — but a
separate one on master, and I am happy to open a PR for it.

Wall clock. The PR body omits this because my previous machine could not hold
a stable measurement. This one can. Six-layer transformer forward, float32,
median of 8 alternating A/B pairs against the merge-base:

	eager	compiled
without this change	5.98 ms	8.05 ms
with this change	5.98 ms	0.98 ms

• compiled: 8.2x faster, 8/8 pairs agreeing, per-pair spread 1.03x
• eager: unchanged (1.005x, inside a 2.5% noise floor); the guard costs 0.1 µs
  per call
• without the change, torch.compile makes this model slower than eager

[hertschuh]

See #23327 (comment)

[hertschuh]

Thanks for the fix!