SDPA decomposition does not preserve intermediate dtype for fp16/bf16 inputs

When exporting an fp16/bf16 model, aten.scaled_dot_product_attention is decomposed during to_edge before backend partitioning. For fp16/bf16 inputs, the current decomposition promotes SDPA intermediates to fp32, then casts the final result back to the original output dtype. In other words, the public output dtype is preserved, but the intermediate dtype is not.

Is this fp32 intermediate behavior intended to be unconditional, or should users/backends have a way to request dtype-preserving SDPA decomposition?

More generally, what should users/backends expect from Core ATen decompositions with respect to dtype? Should decompositions generally preserve the dtype of intermediate operations when the original op is exported with fp16/bf16 inputs, or is it expected that decompositions may introduce higher-precision opmath intermediates as long as the public operator output dtype is preserved?

Minimal example:

import torch
from executorch.exir import to_edge

class M(torch.nn.Module):
    def forward(self, q, k, v):
        return torch.nn.functional.scaled_dot_product_attention(
            q, k, v, dropout_p=0.0
        )

inputs = tuple(torch.randn(1, 2, 4, 8, dtype=torch.float16) for _ in range(3))
ep = torch.export.export(M(), inputs)

edge = to_edge(ep)
print(edge.exported_program().graph_module)

After to_edge, the graph contains fp32 intermediates from the SDPA decomposition even though the model inputs and final SDPA output are fp16. The same applies to bf16.

This is understandable from a numerical-stability perspective, especially around the score matmul and softmax. However, since SDPA is decomposed before backend partitioning, backend-specific passes do not see the original SDPA op and cannot easily choose a lower-precision SDPA implementation.

cc @digantdesai @freddan80 @per @zingo @mansnils @Sebastian-Larsson @robell @rascani