[ROCm] Bump AOTriton to 0.10b (pytorch#156290)

Notable new features/optimizations for SDPA operators on AMD systems from AOTriton 0.10b:

* Official support of gfx950/gfx1201
* Experimental support of gfx1101/gfx1151/gfx1150/gfx1200
* Reduce libaotriton.so binary size by over 80%.
  + Without this optimization the binary size of `libaotriton.so` could be
    over 100MiB due to 2x more supported architectures compared with 0.9b.
    Now it is only about 11MiB.
* Support sliding window attention (SWA) in
  `_flash_attention_forward/backward`. Should fix pytorch#154582

See https://github.com/ROCm/aotriton/releases/tag/0.10b for full details,
including Known Problems.

Notable changes to SDPA backend:

* `std::optional<int64_t>` `window_size_left/right` are directly passed to
  ROCM's SDPA backend, because the default value `-1` is meaningful to
  AOTriton's backend and bottom-right aligned causal mask is implemented with
  negative `window_size_left/right`
* Some code clean up around `USE_CK_FLASH_ATTENTION`

Pull Request resolved: pytorch#156290
Approved by: https://github.com/jithunnair-amd, https://github.com/jeffdaily

Author: xinyazhang

aten/src/ATen/native/transformers/cuda/attention.cu

@@ -1113,8 +1113,10 @@ _flash_attention_forward(
   std::optional<Tensor> alibi_slopes = _alibi_slopes;
   const float softcap = 0.0;
 
-  const int non_null_window_left = window_size_left.has_value() ? window_size_left.value() : -1;
-  const int non_null_window_right = window_size_right.has_value() ? window_size_right.value() : -1;
+#ifndef USE_ROCM  // ROCM backend accepts std::optional for window_size_left/right directly.
+  const int non_null_window_left = window_size_left.value_or(-1);
+  const int non_null_window_right = window_size_right.value_or(-1);
+#endif
 
   // We are going to have two paths:
   // 1. The standard MHA path for dense tensors
@@ -1151,8 +1153,13 @@ _flash_attention_forward(
             softmax_scale,
             false /*zero_tensors*/,
             is_causal,
+#ifdef USE_ROCM
+            window_size_left,
+            window_size_right,
+#else
             non_null_window_left,
             non_null_window_right,
+#endif
             softcap,
             return_debug_mask,
             std::nullopt /*gen_*/);
@@ -1175,8 +1182,13 @@ _flash_attention_forward(
             dropout_p,
             softmax_scale,
             is_causal,
+#ifdef USE_ROCM
+            window_size_left,
+            window_size_right,
+#else
             non_null_window_left,
             non_null_window_right,
+#endif
             softcap,
             return_debug_mask, /*return_softmax (this is used for testing)*/
             std::nullopt);

aten/src/ATen/native/transformers/cuda/attention_backward.cu

@@ -87,8 +87,10 @@ std::tuple<Tensor, Tensor, Tensor> _flash_attention_backward(
   auto contiguous_grad_out = grad_out.contiguous();
   auto contiguous_out = out.contiguous();
 
+#ifndef USE_ROCM  // ROCM backend accepts std::optional for window_size_left/right directly.
   const int non_null_window_left = window_size_left.has_value() ? window_size_left.value() : -1;
   const int non_null_window_right = window_size_right.has_value() ? window_size_right.value() : -1;
+#endif
 
   std::optional<at::Tensor> dq{std::nullopt};
   std::optional<at::Tensor> dk{std::nullopt};
@@ -136,8 +138,13 @@ std::tuple<Tensor, Tensor, Tensor> _flash_attention_backward(
         softmax_scale,
         false /*zero_tensors*/,
         is_causal,
+#ifdef USE_ROCM
+        window_size_left,
+        window_size_right,
+#else
         non_null_window_left,
         non_null_window_right,
+#endif
         softcap,
         determinisitic,
         philox_seed,
@@ -159,8 +166,13 @@ std::tuple<Tensor, Tensor, Tensor> _flash_attention_backward(
         dropout_p,
         softmax_scale,
         is_causal,
+#ifdef USE_ROCM
+        window_size_left,
+        window_size_right,
+#else
         non_null_window_left,
         non_null_window_right,
+#endif
         softcap,
         determinisitic,
         philox_seed,