fix comments

Author: kuizhiqing

cmake/external/flashattn.cmake

@@ -1,4 +1,4 @@
-# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

paddle/phi/backends/dynload/flashattn.cc

@@ -1,13 +1,16 @@
-/* Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-    http://www.apache.org/licenses/LICENSE-2.0
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License. */
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
 
 #include "paddle/phi/backends/dynload/flashattn.h"
 

paddle/phi/backends/dynload/flashattn.h

@@ -1,4 +1,4 @@
-/* Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+/* Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.

paddle/phi/infermeta/backward.cc

@@ -198,6 +198,23 @@ void CropGradInferMeta(const MetaTensor& out_grad,
   }
 }
 
+void FlashAttnGradInferMeta(const MetaTensor& q,
+                            const MetaTensor& k,
+                            const MetaTensor& v,
+                            MetaTensor* dq,
+                            MetaTensor* dk,
+                            MetaTensor* dv) {
+  if (dq) {
+    dq->share_meta(q);
+  }
+  if (dk && k) {
+    dk->share_meta(k);
+  }
+  if (dv && v) {
+    dv->share_meta(v);
+  }
+}
+
 void CrossEntropyWithSoftmaxGradInferMeta(const MetaTensor& label,
                                           const MetaTensor& softmax,
                                           const MetaTensor& loss_grad,

paddle/phi/infermeta/backward.h

@@ -168,6 +168,13 @@ void FillDiagonalTensorGradInferMeta(const MetaTensor& out_grad,
                                      int dim2,
                                      MetaTensor* x_grad);
 
+void FlashAttnGradInferMeta(const MetaTensor& q,
+                            const MetaTensor& k,
+                            const MetaTensor& v,
+                            MetaTensor* dq,
+                            MetaTensor* dk,
+                            MetaTensor* dv);
+
 void GatherNdGradInferMeta(const MetaTensor& x,
                            const MetaTensor& index,
                            const MetaTensor& out_grad,

paddle/phi/infermeta/ternary.cc

@@ -267,23 +267,6 @@ void FlashAttnInferMeta(const MetaTensor& q,
   out->set_layout(q.layout());
 }
 
-void FlashAttnGradInferMeta(const MetaTensor& q,
-                            const MetaTensor& k,
-                            const MetaTensor& v,
-                            MetaTensor* dq,
-                            MetaTensor* dk,
-                            MetaTensor* dv) {
-  if (dq) {
-    dq->share_meta(q);
-  }
-  if (dk && k) {
-    dk->share_meta(k);
-  }
-  if (dv && v) {
-    dv->share_meta(v);
-  }
-}
-
 void ArangeInferMeta(const MetaTensor& start,
                      const MetaTensor& end,
                      const MetaTensor& step,

paddle/phi/infermeta/ternary.h

@@ -71,13 +71,6 @@ void FlashAttnInferMeta(const MetaTensor& q,
                         MetaTensor* softmax,
                         MetaTensor* seed_offset);
 
-void FlashAttnGradInferMeta(const MetaTensor& q,
-                            const MetaTensor& k,
-                            const MetaTensor& v,
-                            MetaTensor* dq,
-                            MetaTensor* dk,
-                            MetaTensor* dv);
-
 void InstanceNormInferMeta(const MetaTensor& x,
                            const MetaTensor& scale,
                            const MetaTensor& bias,

paddle/phi/kernels/flash_attn_grad_kernel.h

@@ -1,4 +1,4 @@
-// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.

paddle/phi/kernels/flash_attn_kernel.h

@@ -1,4 +1,4 @@
-// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.

paddle/phi/kernels/gpu/arange_kernel.cu

@@ -72,6 +72,10 @@ void ArangeNullaryKernel(const Context& dev_ctx,
   Range<T><<<grid, block, 0, stream>>>(start_value, step_value, size, out_data);
 }
 
+template decltype(ArangeNullaryKernel<int64_t, phi::GPUContext>)
+    ArangeNullaryKernel;
+template decltype(ArangeNullaryKernel<int, phi::GPUContext>)
+    ArangeNullaryKernel;
 }  // namespace phi
 
 PD_REGISTER_KERNEL(
@@ -80,6 +84,3 @@ PD_REGISTER_KERNEL(
   kernel->InputAt(1).SetBackend(phi::Backend::ALL_BACKEND);
   kernel->InputAt(2).SetBackend(phi::Backend::ALL_BACKEND);
 }
-
-PD_REGISTER_KERNEL(
-    arange_nullary, GPU, ALL_LAYOUT, phi::ArangeNullaryKernel, int64_t, int) {}