create nested spcecific aten function

drisspg · drisspg · commit 2b2247ea8aa6 · 2023-10-04T13:18:38.000-07:00
diff --git a/aten/src/ATen/native/cpu/FlashAttentionKernel.cpp b/aten/src/ATen/native/cpu/FlashAttentionKernel.cpp
@@ -55,8 +55,6 @@ void cpu_flash_attention(
     const Tensor& logsumexp,
     const Tensor& cum_seq_q,
     const Tensor& cum_seq_k,
-    int64_t& max_q,
-    int64_t& max_k,
     const Tensor& philox_seed,
     const Tensor& philox_offset,
     const Tensor& debug_attn_mask,
@@ -279,8 +277,6 @@ void cpu_flash_attention_backward(
     const at::Tensor& logsumexp,
     const Tensor& cumulative_sequence_length_q,
     const Tensor& cumulative_sequence_length_k,
-    const int64_t max_seqlen_batch_q,
-    const int64_t max_seqlen_batch_k,
     double dropout_p,
     bool is_causal,
     const at::Tensor& philox_seed,
@@ -540,8 +536,6 @@ void flash_attention_kernel_impl(
     const Tensor& logsumexp,
     const Tensor& cum_seq_q,
     const Tensor& cum_seq_k,
-    int64_t& max_q,
-    int64_t& max_k,
     const Tensor& philox_seed,
     const Tensor& philox_offset,
     const Tensor& debug_attn_mask,
@@ -558,17 +552,17 @@ void flash_attention_kernel_impl(
     if (q_seq_len >= 768) {
       cpu_flash_attention<scalar_t, 256, 512>(
         output, logsumexp, cum_seq_q, cum_seq_k,
-        max_q, max_k, philox_seed, philox_offset, debug_attn_mask,
+        philox_seed, philox_offset, debug_attn_mask,
         query, key, value, dropout_p, is_causal, return_debug_mask, scale);
     } else if (q_seq_len >= 192) {
       cpu_flash_attention<scalar_t, 64, 512>(
         output, logsumexp, cum_seq_q, cum_seq_k,
-        max_q, max_k, philox_seed, philox_offset, debug_attn_mask,
+        philox_seed, philox_offset, debug_attn_mask,
         query, key, value, dropout_p, is_causal, return_debug_mask, scale);
     } else {
       cpu_flash_attention<scalar_t, 32, 512>(
         output, logsumexp, cum_seq_q, cum_seq_k,
-        max_q, max_k, philox_seed, philox_offset, debug_attn_mask,
+        philox_seed, philox_offset, debug_attn_mask,
         query, key, value, dropout_p, is_causal, return_debug_mask, scale);
     }
   });
@@ -586,8 +580,6 @@ void flash_attention_backward_kernel_impl(
     const at::Tensor& logsumexp,
     const Tensor& cum_seq_q,
     const Tensor& cum_seq_k,
-    const int64_t max_q,
-    const int64_t max_k,
     double dropout_p,
     bool is_causal,
     const at::Tensor& philox_seed,
@@ -604,19 +596,19 @@ void flash_attention_backward_kernel_impl(
       cpu_flash_attention_backward<scalar_t, 256, 512>(
         grad_q, grad_k, grad_v, grad_out_contig,
         query, key, value, out, logsumexp,
-        cum_seq_q, cum_seq_k, max_q, max_k, dropout_p,
+        cum_seq_q, cum_seq_k, dropout_p,
         is_causal, philox_seed, philox_offset, scale);
     } else if (q_seq_len >= 192) {
       cpu_flash_attention_backward<scalar_t, 64, 512>(
         grad_q, grad_k, grad_v, grad_out_contig,
         query, key, value, out, logsumexp,
-        cum_seq_q, cum_seq_k, max_q, max_k, dropout_p,
+        cum_seq_q, cum_seq_k, dropout_p,
         is_causal, philox_seed, philox_offset, scale);
     } else {
       cpu_flash_attention_backward<scalar_t, 32, 512>(
         grad_q, grad_k, grad_v, grad_out_contig,
         query, key, value, out, logsumexp,
-        cum_seq_q, cum_seq_k, max_q, max_k, dropout_p,
+        cum_seq_q, cum_seq_k, dropout_p,
         is_causal, philox_seed, philox_offset, scale);
     }
   });
diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
@@ -14349,14 +14349,18 @@
   variants: function
   tags: nondeterministic_seeded
 
-- func: _scaled_dot_product_flash_attention(Tensor query, Tensor key, Tensor value, float dropout_p=0.0, bool is_causal=False, bool return_debug_mask=False, *, float? scale=None) -> (Tensor output, Tensor logsumexp, Tensor cum_seq_q, Tensor cum_seq_k, int max_q, int max_k, Tensor philox_seed, Tensor philox_offset, Tensor debug_attn_mask)
+- func: _scaled_dot_product_flash_attention(Tensor query, Tensor key, Tensor value, float dropout_p=0.0, bool is_causal=False, bool return_debug_mask=False, *, float? scale=None) -> (Tensor output, Tensor logsumexp, Tensor cum_seq_q, Tensor cum_seq_k, Tensor philox_seed, Tensor philox_offset, Tensor debug_attn_mask)
   dispatch:
     CPU: _scaled_dot_product_flash_attention_cpu
     CUDA: _scaled_dot_product_flash_attention_cuda
+  tags: nondeterministic_seeded
+
+- func: _scaled_dot_product_flash_attention_nested(Tensor query, Tensor key, Tensor value, float dropout_p=0.0, bool is_causal=False, bool return_debug_mask=False, *, float? scale=None) -> (Tensor output, Tensor logsumexp, Tensor cum_seq_q, Tensor cum_seq_k, int max_q, int max_k, Tensor philox_seed, Tensor philox_offset, Tensor debug_attn_mask)
+  dispatch:
     NestedTensorCUDA: _scaled_dot_product_flash_attention_nestedtensor_cuda
   tags: nondeterministic_seeded
 
-- func: _scaled_dot_product_flash_attention_backward(Tensor grad_out, Tensor query, Tensor key, Tensor value, Tensor out, Tensor logsumexp, Tensor cum_seq_q, Tensor cum_seq_k, int max_q, int max_k, float dropout_p, bool is_causal, Tensor philox_seed, Tensor philox_offset, *, float? scale=None) -> (Tensor grad_query, Tensor grad_key, Tensor grad_value)
+- func: _scaled_dot_product_flash_attention_backward(Tensor grad_out, Tensor query, Tensor key, Tensor value, Tensor out, Tensor logsumexp, Tensor cum_seq_q, Tensor cum_seq_k, SymInt max_q, SymInt max_k, float dropout_p, bool is_causal, Tensor philox_seed, Tensor philox_offset, *, float? scale=None) -> (Tensor grad_query, Tensor grad_key, Tensor grad_value)
   device_check: NoCheck
   variants: function
   dispatch:
diff --git a/aten/src/ATen/native/transformers/attention.cpp b/aten/src/ATen/native/transformers/attention.cpp
@@ -561,6 +561,10 @@ at::Tensor post_process_flash_output(
   return out;
 }
 
+bool has_nested_inputs(const Tensor& query, const Tensor& key, const Tensor& value){
+  return query.is_nested() || key.is_nested() || value.is_nested();
+}
+
 } // namespace
 
 // Computes scaled dot product attention on query, key and value tensors, using
@@ -617,6 +621,11 @@ Tensor scaled_dot_product_attention(
         Tensor value_padded = pad_last_dim<8, false>(value);
         // We need to calculate the scale based off the OG head dim size
         auto og_scale = sdp::calculate_scale(query_, scale);
+        if (has_nested_inputs(query_padded, key_padded, value_padded)) {
+            auto out_lse_softmax = at::_scaled_dot_product_flash_attention_nested(
+            query_padded, key_padded, value_padded, dropout_p, is_causal, false /*return_debug_mask*/, og_scale.as_float_unchecked());
+            return post_process_flash_output(std::get<0>(out_lse_softmax), og_size);
+        }
         auto out_lse_softmax = at::_scaled_dot_product_flash_attention(
             query_padded, key_padded, value_padded, dropout_p, is_causal, false /*return_debug_mask*/, og_scale.as_float_unchecked());
         return post_process_flash_output(std::get<0>(out_lse_softmax), og_size);
@@ -715,8 +724,6 @@ std::tuple<
     at::Tensor,
     at::Tensor,
     at::Tensor,
-    int64_t,
-    int64_t,
     at::Tensor,
     at::Tensor,
     at::Tensor>
@@ -751,21 +758,19 @@ _scaled_dot_product_flash_attention_cpu(
       query.options().dtype(accumulate_dtype));
   at::Tensor cum_seq_q = at::empty({}, at::kLong);
   at::Tensor cum_seq_k = at::empty({}, at::kLong);
-  int64_t max_q = 0;
-  int64_t max_k = 0;
   at::Tensor philox_seed = at::empty({}, at::kLong);
   at::Tensor philox_offset = at::empty({}, at::kLong);
   at::Tensor debug_attn_mask = at::empty({}, query.options());
 
   flash_attention_kernel(kCPU, output, logsumexp, cum_seq_q, cum_seq_k,
-      max_q, max_k, philox_seed, philox_offset, debug_attn_mask,
+      philox_seed, philox_offset, debug_attn_mask,
       query, key, value, dropout_p, is_causal, return_debug_mask, scale);
 
   output = output.transpose(1, 2);
   logsumexp = logsumexp.transpose(1, 2);
 
   return std::make_tuple(std::move(output), std::move(logsumexp),
-      std::move(cum_seq_q), std::move(cum_seq_k), max_q, max_k,
+      std::move(cum_seq_q), std::move(cum_seq_k),
       std::move(philox_seed), std::move(philox_offset), std::move(debug_attn_mask));
 }
 
@@ -802,7 +807,7 @@ _scaled_dot_product_flash_attention_backward_cpu(
 
   flash_attention_backward_kernel(kCPU, grad_q, grad_k, grad_v,
       grad_out_t, q_t, k_t, v_t, o_t, lse_t, cum_seq_q, cum_seq_k,
-      max_q, max_k, dropout_p, is_causal, philox_seed, philox_offset, scale);
+      dropout_p, is_causal, philox_seed, philox_offset, scale);
 
   grad_q = grad_q.transpose(1, 2);
   grad_k = grad_k.transpose(1, 2);
diff --git a/aten/src/ATen/native/transformers/attention.h b/aten/src/ATen/native/transformers/attention.h
@@ -52,7 +52,7 @@ TORCH_API Tensor qkv_projection(
 using flash_attention_fn = void (*)(
     const Tensor& output, const Tensor& logsumexp,
     const Tensor& cum_seq_q, const Tensor& cum_seq_k,
-    int64_t& max_q, int64_t& max_k, const Tensor& philox_seed,
+    const Tensor& philox_seed,
     const Tensor& philox_offset, const Tensor& debug_attn_mask,
     const Tensor& query, const Tensor& key, const Tensor& value,
     double dropout_p, bool is_causal, bool return_debug_mask,
@@ -64,7 +64,6 @@ using flash_attention_backward_fn = void (*)(
     const Tensor& query, const Tensor& key,
     const Tensor& value, const Tensor& out, const Tensor& logsumexp,
     const Tensor& cum_seq_q, const Tensor& cum_seq_k,
-    const int64_t max_q, const int64_t max_k,
     double dropout_p, bool is_causal,
     const Tensor& philox_seed, const Tensor& philox_offset,
     c10::optional<double> scale);
diff --git a/aten/src/ATen/native/transformers/cuda/attention.cu b/aten/src/ATen/native/transformers/cuda/attention.cu
@@ -642,7 +642,7 @@ std::tuple<Tensor, Tensor> native_multi_head_attention_cuda(
   }
   return std::make_tuple(std::move(proj), std::move(qkt));
 }
-std::tuple<Tensor, Tensor, Tensor, Tensor, int64_t, int64_t, Tensor, Tensor, Tensor> _scaled_dot_product_flash_attention_cuda(
+std::tuple<Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor> _scaled_dot_product_flash_attention_cuda(
     const Tensor& query,
     const Tensor& key,
     const Tensor& value,
@@ -691,7 +691,7 @@ std::tuple<Tensor, Tensor, Tensor, Tensor, int64_t, int64_t, Tensor, Tensor, Ten
   // Reshape output to convert nnz to batch_size and seq_len
   Tensor attention = output.transpose(1,2);
 
-  return std::make_tuple(attention, logsumexp, Tensor(), Tensor(), max_seqlen_batch_q, max_seqlen_batch_k, philox_seed, philox_offset, debug_attn_mask);
+  return std::make_tuple(attention, logsumexp, Tensor(), Tensor(), philox_seed, philox_offset, debug_attn_mask);
 }
 
 std::tuple<Tensor, Tensor, Tensor, Tensor> _scaled_dot_product_efficient_attention_cuda(
@@ -828,11 +828,11 @@ _flash_attention_forward(
   debug_attn_mask =
       return_debug_mask ? debug_attn_mask : at::empty({0}, query.options());
   return std::make_tuple(
-      output,
-      logsumexp,
-      philox_seed,
-      philox_offset,
-      debug_attn_mask);
+      std::move(output),
+      std::move(logsumexp),
+      std::move(philox_seed),
+      std::move(philox_offset),
+      std::move(debug_attn_mask));
 
 #endif
   TORCH_CHECK(false, "USE_FLASH_ATTENTION was not enabled for build.")
diff --git a/test/test_transformers.py b/test/test_transformers.py
@@ -1189,6 +1189,7 @@ def ones_tensor(*shape):
                         _ = mha_f(qkv_f, qkv_f, qkv_f, attn_mask=mask, need_weights=False, is_causal=True)
                         torch.cuda.synchronize()
 
+    @slowTest
     @unittest.skipIf(
         not PLATFORM_SUPPORTS_FLASH_ATTENTION, "Platform does not supposrt fused SDPA or pre-SM80 hardware"
     )
@@ -2496,6 +2497,7 @@ def test_flash_attention_vs_math_ref_grads(self, device, batch_size: int, seq_le
         self.assertEqual(value.grad, value_ref.grad.to(value.grad.dtype),
                          atol=grad_v_ref_atol, rtol=grad_v_ref_rtol)
 
+    @slowTest
     @unittest.skipIf(not PLATFORM_SUPPORTS_FLASH_ATTENTION, "Does not support SDPA or pre-SM80 hardware")
     @parametrize("batch_size", [1, 8])
     @parametrize("seq_len_q", [256, 512, 1024])
diff --git a/tools/autograd/derivatives.yaml b/tools/autograd/derivatives.yaml
@@ -2764,9 +2764,9 @@
   output_differentiability: [True, False, False, False]
   query, key, value, attn_bias: _scaled_dot_product_efficient_attention_backward(grad, query, key, value, attn_bias, output, log_sumexp, philox_seed, philox_offset, dropout_p, grad_input_mask, is_causal, scale)
 
-- name: _scaled_dot_product_flash_attention(Tensor query, Tensor key, Tensor value, float dropout_p=0.0, bool is_causal=False, bool return_debug_mask=False, *, float? scale=None) -> (Tensor output, Tensor logsumexp, Tensor cum_seq_q, Tensor cum_seq_k, int max_q, int max_k, Tensor philox_seed, Tensor philox_offset, Tensor debug_attn_mask)
-  output_differentiability: [True, False, False, False, False, False, False, False, False]
-  query, key, value: _scaled_dot_product_flash_attention_backward(grad, query, key, value, output, logsumexp, cum_seq_q, cum_seq_k, max_q, max_k, dropout_p, is_causal, philox_seed, philox_offset, scale)
+- name: _scaled_dot_product_flash_attention(Tensor query, Tensor key, Tensor value, float dropout_p=0.0, bool is_causal=False, bool return_debug_mask=False, *, float? scale=None) -> (Tensor output, Tensor logsumexp, Tensor cum_seq_q, Tensor cum_seq_k, Tensor philox_seed, Tensor philox_offset, Tensor debug_attn_mask)
+  output_differentiability: [True, False, False, False, False, False, False]
+  query, key, value: _scaled_dot_product_flash_attention_backward_symint(grad, query, key, value, output, logsumexp, cum_seq_q, cum_seq_k, query.sym_size(2), key.sym_size(2), dropout_p, is_causal, philox_seed, philox_offset, scale)
 
 # - name: _flash_attention_forward(Tensor query, Tensor key, Tensor value, Tensor? cum_seq_q, Tensor? cum_seq_k, int? max_q, int? max_k, float dropout_p, bool is_causal, bool return_debug_mask, *, float? scale=None) -> (Tensor output, Tensor query_padded, Tensor key_padded, Tensor value_padded, Tensor softmax_logsumexp, Tensor philox_seed, Tensor philox_offset, Tensor debug_attn_mask)
 #   output_differentiability: [True, False, False, False, False, False, False, False]
diff --git a/torch/_decomp/decompositions.py b/torch/_decomp/decompositions.py
@@ -4093,7 +4093,7 @@ def scaled_dot_product_flash_attention(
     return_debug_mask: bool = False,
     *,
     scale: Optional[float] = None,
-) -> Tuple[Tensor, Tensor, Tensor, Tensor, int, int, Tensor, Tensor, Tensor]:
+) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
     dtype = query.dtype
     batchSize, num_head, qSize, headSize = (
         query.shape[0],
@@ -4123,7 +4123,6 @@ def scaled_dot_product_flash_attention(
     cum_seq_q, cum_seq_k = torch.empty([], dtype=torch.long), torch.empty(
         [], dtype=torch.long
     )
-    max_q, max_k = 0, 0
     philox_seed, philox_offset = torch.empty([], dtype=torch.long), torch.empty(
         [], dtype=torch.long
     )
@@ -4175,8 +4174,6 @@ def scaled_dot_product_flash_attention(
         logsumexp,
         cum_seq_q,
         cum_seq_k,
-        max_q,
-        max_k,
         philox_seed,
         philox_offset,
         debug_attn_mask,
diff --git a/torch/_meta_registrations.py b/torch/_meta_registrations.py
@@ -4900,8 +4900,6 @@ def meta__scaled_dot_product_flash(
             logsumexp,
             torch.empty((), dtype=torch.int32, device="meta"),
             torch.empty((), dtype=torch.int32, device="meta"),
-            0,
-            0,
             torch.empty((), dtype=torch.long, device="meta"),
             torch.empty((), dtype=torch.long, device="meta"),
             torch.empty((), dtype=query.dtype, device=query.device),
@@ -4941,8 +4939,6 @@ def meta__scaled_dot_product_flash(
         logsumexp,
         None,
         None,
-        max_seqlen_batch_q,
-        max_seqlen_batch_k,
         torch.empty((), dtype=torch.long, device="meta"),
         torch.empty((), dtype=torch.long, device="meta"),
         debug_mask,
diff --git a/torch/csrc/inductor/aoti_torch/c/shim.h b/torch/csrc/inductor/aoti_torch/c/shim.h
@@ -164,11 +164,9 @@ AOTI_TORCH_EXPORT AOTITorchError aoti_torch__scaled_dot_product_flash_attention(
     AtenTensorHandle* ret1, // returns new reference
     AtenTensorHandle* ret2, // returns new reference
     AtenTensorHandle* ret3, // returns new reference
-    int64_t* ret4,
-    int64_t* ret5,
-    AtenTensorHandle* ret6, // returns new reference
-    AtenTensorHandle* ret7, // returns new reference
-    AtenTensorHandle* ret8 // returns new reference
+    AtenTensorHandle* ret4, // returns new reference
+    AtenTensorHandle* ret5, // returns new reference
+    AtenTensorHandle* ret6 // returns new reference
 );
 
 // This function will create a new uninitialized tensor object
diff --git a/torch/csrc/inductor/aoti_torch/shim_common.cpp b/torch/csrc/inductor/aoti_torch/shim_common.cpp
@@ -195,25 +195,22 @@ AOTITorchError aoti_torch__scaled_dot_product_flash_attention(
     AtenTensorHandle* ret1, // returns new reference
     AtenTensorHandle* ret2, // returns new reference
     AtenTensorHandle* ret3, // returns new reference
-    int64_t* ret4,
-    int64_t* ret5,
-    AtenTensorHandle* ret6, // returns new reference
-    AtenTensorHandle* ret7, // returns new reference
-    AtenTensorHandle* ret8 // returns new reference
+    AtenTensorHandle* ret4, // returns new reference
+    AtenTensorHandle* ret5, // returns new reference
+    AtenTensorHandle* ret6 // returns new reference
 ) {
   AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({
     at::Tensor* query_tensor = tensor_handle_to_tensor_pointer(query);
     at::Tensor* key_tensor = tensor_handle_to_tensor_pointer(key);
     at::Tensor* value_tensor = tensor_handle_to_tensor_pointer(value);
-    auto [r0, r1, r2, r3, r4, r5, r6, r7, r8] =
-        at::_scaled_dot_product_flash_attention(
-            *query_tensor,
-            *key_tensor,
-            *value_tensor,
-            dropout_p,
-            is_causal,
-            return_debug_mask,
-            scale);
+    auto [r0, r1, r2, r3, r4, r5, r6] = at::_scaled_dot_product_flash_attention(
+        *query_tensor,
+        *key_tensor,
+        *value_tensor,
+        dropout_p,
+        is_causal,
+        return_debug_mask,
+        scale);
 
     at::Tensor* ret0_tensor = new at::Tensor(std::move(r0));
     *ret0 = tensor_pointer_to_tensor_handle(ret0_tensor);
@@ -228,14 +225,12 @@ AOTITorchError aoti_torch__scaled_dot_product_flash_attention(
       at::Tensor* ret3_tensor = new at::Tensor(std::move(r3));
       *ret3 = tensor_pointer_to_tensor_handle(ret3_tensor);
     }
-    *ret4 = r4;
-    *ret5 = r5;
+    at::Tensor* ret4_tensor = new at::Tensor(std::move(r4));
+    *ret4 = tensor_pointer_to_tensor_handle(ret4_tensor);
+    at::Tensor* ret5_tensor = new at::Tensor(std::move(r5));
+    *ret5 = tensor_pointer_to_tensor_handle(ret5_tensor);
     at::Tensor* ret6_tensor = new at::Tensor(std::move(r6));
     *ret6 = tensor_pointer_to_tensor_handle(ret6_tensor);
-    at::Tensor* ret7_tensor = new at::Tensor(std::move(r7));
-    *ret7 = tensor_pointer_to_tensor_handle(ret7_tensor);
-    at::Tensor* ret8_tensor = new at::Tensor(std::move(r8));
-    *ret8 = tensor_pointer_to_tensor_handle(ret8_tensor);
   });
 }
 
diff --git a/torchgen/api/python.py b/torchgen/api/python.py
@@ -1130,6 +1130,7 @@ def dispatch_lambda_arg(cpp_arg: Binding) -> DispatchLambdaArgument:
     "::std::vector<at::Tensor>",
     # Needed for flash attention forw/backward
     "::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor,int64_t,int64_t,at::Tensor,at::Tensor,at::Tensor>",
+    "::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor>",
     "at::Scalar",
     "bool",
     "int64_t",
