[SDPA] Forward fails to compile with non-default scale
- Dominant language
- C++
- Stars
- 25
- Forks
- 16
- PR merge metrics
- No merged PRs in 30d
Description
SDPA forward fails to compile on AMDGPU when `set_attn_scale_value` is set to a value other than `1/sqrt(head_dim)`.
The existing `sdpa_fprop_custom_scale` sample passes today only because it sets `scale=0.125f` with `head_dim=64`, which happens to equal the default `1/sqrt(64) = 0.125`. Replacing the scale with any other value reproduces.
## Repro (sample)
```diff
--- a/samples/sdpa/sdpa_fprop_custom_scale.cpp
+++ b/samples/sdpa/sdpa_fprop_custom_scale.cpp
@@ -16,5 +16,5 @@ TEST_CASE("SDPA forward: custom scale f16", "[sdpa][graph]") {
executeSdpa(handle, DataType::Half,
/*batch=*/1, /*headsQ=*/8, /*headsK=*/8, /*headsV=*/8,
/*seqQ=*/64, /*seqKV=*/64, /*headDim=*/64,
- /*isCausal=*/false, /*scale=*/0.125f);
+ /*isCausal=*/false, /*scale=*/0.05f);
}
```
## Repro (standalone iree-compile)
`fusilli_sdpa_scale005.mlir` (verbatim from the failing sample's `iree-compile-input.mlir`):
```mlir
module @module {
func.func @main(%sdpa_O_: !torch.tensor<[1,8,64,64],f16>, %k: !torch.vtensor<[1,8,64,64],f16>, %q: !torch.vtensor<[1,8,64,64],f16>, %v: !torch.vtensor<[1,8,64,64],f16>) attributes {torch.assume_strict_symbolic_shapes} {
%permute_Q_val_0_sdpa = torch.constant.int 0
%permute_Q_val_1_sdpa = torch.constant.int 1
%permute_Q_val_2_sdpa = torch.constant.int 2
%permute_Q_val_3_sdpa = torch.constant.int 3
%permute_Q_sdpa = torch.prim.ListConstruct %permute_Q_val_0_sdpa, %permute_Q_val_1_sdpa, %permute_Q_val_2_sdpa, %permute_Q_val_3_sdpa : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list
%q_sdpa_perm = torch.aten.permute %q, %permute_Q_sdpa : !torch.vtensor<[1,8,64,64],f16>, !torch.list -> !torch.vtensor<[1,8,64,64],f16>
%permute_K_val_0_sdpa = torch.constant.int 0
%permute_K_val_1_sdpa = torch.constant.int 1
%permute_K_val_2_sdpa = torch.constant.int 2
%permute_K_val_3_sdpa = torch.constant.int 3
%permute_K_sdpa = torch.prim.ListConstruct %permute_K_val_0_sdpa, %permute_K_val_1_sdpa, %permute_K_val_2_sdpa, %permute_K_val_3_sdpa : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list
%k_sdpa_perm = torch.aten.permute %k, %permute_K_sdpa : !torch.vtensor<[1,8,64,64],f16>, !torch.list -> !torch.vtensor<[1,8,64,64],f16>
%permute_V_val_0_sdpa = torch.constant.int 0
%permute_V_val_1_sdpa = torch.constant.int 1
%permute_V_val_2_sdpa = torch.constant.int 2
%permute_V_val_3_sdpa = torch.constant.int 3
%permute_V_sdpa = torch.prim.ListConstruct %permute_V_val_0_sdpa, %permute_V_val_1_sdpa, %permute_V_val_2_sdpa, %permute_V_val_3_sdpa : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list
%v_sdpa_perm = torch.aten.permute %v, %permute_V_sdpa : !torch.vtensor<[1,8,64,64],f16>, !torch.list -> !torch.vtensor<[1,8,64,64],f16>
%none_mask_sdpa = torch.constant.none
%dropout_sdpa = torch.constant.float 0.000000e+00
%is_causal_sdpa = torch.constant.bool false
%scale_sdpa = torch.constant.float 5.000000e-02
%enable_gqa_sdpa = torch.constant.bool false
%sdpa_O_sdpa_perm = torch.aten.scaled_dot_product_attention %q_sdpa_perm, %k_sdpa_perm, %v_sdpa_perm, %none_mask_sdpa, %dropout_sdpa, %is_causal_sdpa, %scale_sdpa, %enable_gqa_sdpa : !torch.vtensor<[1,8,64,64],f16>, !torch.vtensor<[1,8,64,64],f16>, !torch.vtensor<[1,8,64,64],f16>, !torch.none, !torch.float, !torch.bool, !torch.float, !torch.bool -> !torch.vtensor<[1,8,64,64],f16>
%permute_O_val_0_sdpa = torch.constant.int 0
%permute_O_val_1_sdpa = torch.constant.int 1
%permute_O_val_2_sdpa = torch.constant.int 2
%permute_O_val_3_sdpa = torch.constant.int 3
%permute_O_sdpa = torch.prim.ListConstruct %permute_O_val_0_sdpa, %permute_O_val_1_sdpa, %permute_O_val_2_sdpa, %permute_O_val_3_sdpa : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list
%sdpa_O = torch.aten.permute %sdpa_O_sdpa_perm, %permute_O_sdpa : !torch.vtensor<[1,8,64,64],f16>, !torch.list -> !torch.vtensor<[1,8,64,64],f16>
torch.overwrite.tensor.contents %sdpa_O overwrites %sdpa_O_ : !torch.vtensor<[1,8,64,64],f16>, !torch.tensor<[1,8,64,64],f16>
return
}
}
```
Compile command (verbatim from fusilli's `iree-compile-command.txt`, with workspace-local paths replaced):
```
iree-compile fusilli_sdpa_scale005.mlir \
--iree-hal-target-backends=rocm --iree-rocm-target=mi300x \
--iree-opt-level=O3 \
'--iree-preprocessing-pass-pipeline=builtin.module(iree-preprocessing-backward-data-conv-pipeline)' \
--iree-flow-enable-pad-handling \
--iree-global-opt-propagate-transposes-through-conv \
--iree-global-opt-enable-sink-transpose-through-pad \
--iree-dispatch-creation-enable-fuse-padding-into-linalg-consumer-ops \
--iree-dispatch-creation-enable-aggressive-reshape-movement \
--iree-dispatch-creation-enable-split-reduction \
--iree-torch-externalize-transients \
-o out.vmfb
```
Replacing the `5.000000e-02` constant with `1.250000e-01` (= `1/sqrt(64)`) makes the same command succeed.
## Error output
```
fusilli_sdpa_scale005.mlir:35:25: error: failed to legalize operation 'torch.aten.scaled_dot_product_attention' that was explicitly marked illegal: %14 = "torch.aten.scaled_dot_product_attention"(%9, %11, %13, %3, %2, %1, %0, %1) : (!torch.vtensor<[1,8,64,64],f16>, !torch.vtensor<[1,8,64,64],f16>, !torch.vtensor<[1,8,64,64],f16>, !torch.none, !torch.float, !torch.bool, !torch.float, !torch.bool) -> !torch.vtensor<[1,8,64,64],f16>
```
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.