[pytorchbot]

Merge after #24448. This pick is stacked on it, because #24426 builds on CUDA backend changes that the #24298, #24320, #24324 and #24448 picks bring to release/1.6.

What is wrong today

Decomposing scaled_dot_product_attention on a CUDA device turns it into one of its fused variants: _scaled_dot_product_flash_attention without a mask, or _scaled_dot_product_efficient_attention with one. torch.export's default decompositions do this. So does Torch-TensorRT, for an attention it leaves to the CUDA backend, for example one whose keys are a constant. The CUDA backend has no kernel for these ops:

RuntimeError: Method forward missing fallback kernels (1 total):
  - at::_ops::_scaled_dot_product_flash_attention::call

What this change does

A new pass, ReplaceFusedSdpaWithSdpaPass, turns the flash, efficient and cuDNN variants back into scaled_dot_product_attention, with the same query, key, value, bias, dropout, causal flag and scale. The Triton SDPA kernel or Inductor's math decomposition then lowers it, as it lowers any other attention. The pass also looks inside control flow branches.

• It sets enable_gqa unless the query and key head counts are known to be equal, because the fused variants group fewer key heads the same way.
• It leaves out a scale of None, because the Triton SDPA kernels take a float scale.

The pass rewrites a node only when scaled_dot_product_attention computes the same result:

• Only the attention output may be used, which is the case for inference.
• A causal call must have no bias, because scaled_dot_product_attention takes no mask together with is_causal.
• A causal call must have query and key lengths that are known to be equal. Flash attention aligns a causal mask to the bottom right, the other variants and scaled_dot_product_attention to the top left, and the Triton decode kernel for one query ignores it.

Otherwise the node stays, the pass logs why, and the error above still names it.

What was tested

With torch 2.14 on an A100:

• test_fused_sdpa decomposes attention on a CUDA device into a causal flash attention with fewer key heads and a custom scale, an efficient attention with a mask, and a flash attention of one query over 256 keys. It then lowers the program, checks that the first and last calls get the Triton kernels, runs it, and compares each output with eager.
• test_replace_fused_sdpa.py checks on the CPU which calls the pass rewrites and with which keyword arguments, which calls it keeps, that the program stays valid, and that the pass runs before the Triton kernel replacement.

They fail without this change. The other CUDA tests in the same CI job pass.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24455

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 Pending, 1 Unrelated Failure, 1 Unclassified Failure

As of commit 13abd1d with merge base 4db778e ():

UNCLASSIFIED FAILURE - DrCI could not classify the following job because the workflow did not run on the merge base. The failure may be pre-existing on trunk or introduced by this PR:

• Test CUDA Builds / unittest-cuda-quantized-gemm (int4) / linux-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / nxp-mcuxpresso-test / linux-job (gh) (trunk failure)
  RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1

This comment was automatically generated by Dr. CI and updates every 15 minutes.