[fusilli] Emit custom scale SDPA through flex attention

Remove the regular SDPA scale fallback now that torch.hop_flex_attention supports an explicit scale operand.

Update the custom-scale emitter test to expect flex_attention and remove the stale custom-scale xfail.

Co-authored-by: GPT-5 Codex <noreply@openai.com>
Signed-off-by: Keshav Vinayak Jha <keshavvinayakjha@gmail.com>

Author: keshavvinayak01

include/fusilli/support/asm_emitter.h

@@ -2478,12 +2478,7 @@ inline bool SdpaNode::useLegacySdpaAsm() const {
   if (sdpaAttr.getMASK() || sdpaAttr.getDropout() != 0.0f)
     return true;
 
-  // The legacy SDPA op does not produce logsumexp, so keep generate_stats on
-  // flex_attention even when extra attributes are present.
-  if (sdpaAttr.getGenerateStats())
-    return false;
-
-  return sdpaAttr.getScale().has_value();
+  return false;
 }
 
 inline std::string SdpaNode::emitModuleScopeAsm() const {

samples/CMakeLists.txt

@@ -198,8 +198,3 @@ add_fusilli_samples(
     libsdpautils
     Catch2::Catch2WithMain
 )
-
-# XFAIL: TODO(#404): Remove once regular SDPA supports non-default scales.
-set_tests_properties(
-  fusilli_sdpa_samples_sdpa_fprop_custom_scale
-  PROPERTIES WILL_FAIL TRUE)

samples/sdpa_utils.h

@@ -94,7 +94,7 @@ inline constexpr std::string_view kSdpaWithMask = R"mlir(
 )mlir";
 // clang-format on
 
-/// Builds the MLIR template for torch.aten.scaled_dot_product_attention.
+/// Builds the MLIR template for SDPA custom-op samples.
 ///
 /// Selects the appropriate R-string template (with/without attn_mask) and
 /// resolves scalar placeholders. Standard CustomOp dtype/name placeholders

tests/lit/test_sdpa_asm_emitter_custom_scale.cpp

@@ -4,9 +4,6 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-// XFAIL: *
-// TODO(iree-org/fusilli#404): Remove XFAIL once IREE supports non-default SDPA
-// scale values.
 // RUN: %{TEST_EXE} | iree-opt --verify-roundtrip
 // RUN: %{TEST_EXE} | FileCheck %s --check-prefix=TORCH-CHECK
 // RUN: %{TEST_EXE} stats | FileCheck %s --check-prefix=%{BACKEND}-STATS-CHECK
@@ -33,12 +30,10 @@
 // TORCH-CHECK:       %permute_V_val_3_sdpa = torch.constant.int 3
 // TORCH-CHECK:       %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<int>
 // TORCH-CHECK:       %v_sdpa_perm = torch.aten.permute %v, %permute_V_sdpa : !torch.vtensor<[1,8,64,64],f16>, !torch.list<int> -> !torch.vtensor<[1,8,64,64],f16>
-// TORCH-CHECK:       %none_mask_sdpa = torch.constant.none
-// TORCH-CHECK:       %dropout_sdpa = torch.constant.float 0.000000e+00
-// TORCH-CHECK:       %is_causal_sdpa = torch.constant.bool false
 // TORCH-CHECK:       %scale_sdpa = torch.constant.float 5.000000e-02
-// TORCH-CHECK:       %enable_gqa_sdpa = torch.constant.bool false
-// TORCH-CHECK:       %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>
+// TORCH-CHECK:       %return_lse_sdpa = torch.constant.bool false
+// TORCH-CHECK:       %return_max_scores_sdpa = torch.constant.bool false
+// TORCH-CHECK:       %sdpa_O_sdpa_perm, %sdpa_logsumexp_sdpa, %sdpa_max_scores_sdpa = torch.hop_flex_attention %q_sdpa_perm, %k_sdpa_perm, %v_sdpa_perm, %scale_sdpa, %return_lse_sdpa, %return_max_scores_sdpa : !torch.vtensor<[1,8,64,64],f16>, !torch.vtensor<[1,8,64,64],f16>, !torch.vtensor<[1,8,64,64],f16>, !torch.float, !torch.bool, !torch.bool -> !torch.vtensor<[1,8,64,64],f16>, !torch.none, !torch.none
 // TORCH-CHECK:       %permute_O_val_0_sdpa = torch.constant.int 0
 // TORCH-CHECK:       %permute_O_val_1_sdpa = torch.constant.int 1
 // TORCH-CHECK:       %permute_O_val_2_sdpa = torch.constant.int 2
@@ -51,9 +46,9 @@
 // TORCH-CHECK:   }
 //
 // AMDGPU-STATS-CHECK: "transient-memory-size": 0
-// AMDGPU-STATS-CHECK: "dispatch-count": 1
+// AMDGPU-STATS-CHECK: "dispatch-count": {{[1-9][0-9]*}}
 // CPU-STATS-CHECK: "transient-memory-size": 0
-// CPU-STATS-CHECK: "dispatch-count": 1
+// CPU-STATS-CHECK: "dispatch-count": {{[1-9][0-9]*}}
 //
 // clang-format on