c10::optional -> std::optional (#58)

Signed-off-by: Ye (Charlotte) Qi <yeq@meta.com>

Author: yeqcharlotte

csrc/common/pytorch_shim.h

@@ -15,13 +15,13 @@
  * that the following types are not support by PyTorch libary bindings:
  *  - `int`
  *  - `float`
- *  - `c10::optional<T> &`
- *  - `c10::optional<const at::Tensor> &`
+ *  - `std::optional<T> &`
+ *  - `std::optional<const at::Tensor> &`
  * So we convert them to (respectively):
  *  - `int64_t`
  *  - `double`
- *  - `const c10::optional<T>&`
- *  - `const c10::optional<at::Tensor>&`
+ *  - `const std::optional<T>&`
+ *  - `const std::optional<at::Tensor>&`
  */
 
 template<typename T>
@@ -38,36 +38,36 @@ template<typename T>
 T convert_from_pytorch_compatible_type(pytorch_library_compatible_type_t<T> arg) 
     { return pytorch_library_compatible_type<T>::convert_from_type(arg); }
 
-// Map `c10::optional<T> &` -> `const c10::optional<T>&`
+// Map `std::optional<T> &` -> `const std::optional<T>&`
 //  (NOTE: this is bit unsafe but non of the ops in flash_attn mutate 
 //   the optional container)
 template<typename T>
-struct pytorch_library_compatible_type<c10::optional<T> &> { 
-    using type = const c10::optional<T>&;
-    static c10::optional<T>& convert_from_type(const c10::optional<T> &arg) { 
-        return const_cast<c10::optional<T>&>(arg); 
+struct pytorch_library_compatible_type<std::optional<T> &> { 
+    using type = const std::optional<T>&;
+    static std::optional<T>& convert_from_type(const std::optional<T> &arg) { 
+        return const_cast<std::optional<T>&>(arg); 
     }
 };
 
-// Map `c10::optional<T>` -> 
-//          `c10::optional<pytorch_library_compatible_type_t<T>>`
-//  (NOTE: tested for `c10::optional<int>` -> `c10::optional<int64_t>`)
+// Map `std::optional<T>` -> 
+//          `std::optional<pytorch_library_compatible_type_t<T>>`
+//  (NOTE: tested for `std::optional<int>` -> `std::optional<int64_t>`)
 template<typename T>
-struct pytorch_library_compatible_type<c10::optional<T>> { 
-    using type = c10::optional<pytorch_library_compatible_type_t<T>>;
-    static c10::optional<pytorch_library_compatible_type_t<T>> convert_from_type(c10::optional<T> arg) { 
+struct pytorch_library_compatible_type<std::optional<T>> { 
+    using type = std::optional<pytorch_library_compatible_type_t<T>>;
+    static std::optional<pytorch_library_compatible_type_t<T>> convert_from_type(std::optional<T> arg) { 
         return arg; 
     }
 };
 
-// Map `c10::optional<const at::Tensor>&` -> `const c10::optional<at::Tensor>&`
+// Map `std::optional<const at::Tensor>&` -> `const std::optional<at::Tensor>&`
 template<>
-struct pytorch_library_compatible_type<c10::optional<const at::Tensor> &> { 
-    using type = const c10::optional<at::Tensor>&;
-    static c10::optional<const at::Tensor>& convert_from_type(
-        const c10::optional<at::Tensor> &arg) {
-        return const_cast<c10::optional<const at::Tensor>&>(
-            reinterpret_cast<const c10::optional<const at::Tensor>&>(arg)); 
+struct pytorch_library_compatible_type<std::optional<const at::Tensor> &> { 
+    using type = const std::optional<at::Tensor>&;
+    static std::optional<const at::Tensor>& convert_from_type(
+        const std::optional<at::Tensor> &arg) {
+        return const_cast<std::optional<const at::Tensor>&>(
+            reinterpret_cast<const std::optional<const at::Tensor>&>(arg)); 
     }
 };
 

csrc/flash_attn/flash_api_sparse.cpp

@@ -318,11 +318,11 @@ mha_varlen_fwd_sparse(at::Tensor &q,  // total_q x num_heads x head_size, total_
                       const at::Tensor &block_offset,
                       const at::Tensor &column_count,
                       const at::Tensor &column_index,
-                      const c10::optional<at::Tensor> &out_, // total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
+                      const std::optional<at::Tensor> &out_, // total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
                       const at::Tensor &cu_seqlens_q,  // b+1
                       const at::Tensor &cu_seqlens_k,  // b+1
-                      const c10::optional<at::Tensor> &seqused_k, // b. If given, only this many elements of each batch element's keys are used.
-                      const c10::optional<at::Tensor> &alibi_slopes_, // num_heads or b x num_heads
+                      const std::optional<at::Tensor> &seqused_k, // b. If given, only this many elements of each batch element's keys are used.
+                      const std::optional<at::Tensor> &alibi_slopes_, // num_heads or b x num_heads
                       int64_t max_seqlen_q,
                       const int64_t max_seqlen_k,
                       const double p_dropout,
@@ -331,7 +331,7 @@ mha_varlen_fwd_sparse(at::Tensor &q,  // total_q x num_heads x head_size, total_
                       bool is_causal,
                       const double softcap,
                       const bool return_softmax,
-                      c10::optional<at::Generator> gen_) {
+                      std::optional<at::Generator> gen_) {
 
     auto [cc_major, cc_minor] = get_compute_capability(get_current_device());
     bool is_sm8x = cc_major == 8 && cc_minor >= 0;

csrc/flash_attn/flash_api_torch_lib.cpp

@@ -86,11 +86,11 @@ mha_varlen_fwd_sparse(at::Tensor &q,  // total_q x num_heads x head_size, total_
                       const at::Tensor &block_offset,
                       const at::Tensor &column_count,
                       const at::Tensor &column_index,
-                      const c10::optional<at::Tensor> &out_, // total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
+                      const std::optional<at::Tensor> &out_, // total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
                       const at::Tensor &cu_seqlens_q,  // b+1
                       const at::Tensor &cu_seqlens_k,  // b+1
-                      const c10::optional<at::Tensor> &seqused_k, // b. If given, only this many elements of each batch element's keys are used.
-                      const c10::optional<at::Tensor> &alibi_slopes_, // num_heads or b x num_heads
+                      const std::optional<at::Tensor> &seqused_k, // b. If given, only this many elements of each batch element's keys are used.
+                      const std::optional<at::Tensor> &alibi_slopes_, // num_heads or b x num_heads
                       int64_t max_seqlen_q,
                       const int64_t max_seqlen_k,
                       const double p_dropout,
@@ -99,7 +99,7 @@ mha_varlen_fwd_sparse(at::Tensor &q,  // total_q x num_heads x head_size, total_
                       bool is_causal,
                       const double softcap,
                       const bool return_softmax,
-                      c10::optional<at::Generator> gen_);
+                      std::optional<at::Generator> gen_);
 
 /**
  *  Torch Library Registration