Update fused_rotary_positional_embedding.cpp

EleutherAI · Nov 27, 2023 · 5b1331d · 5b1331d
1 parent 417f55c
commit 5b1331d
Showing 1 changed file with 9 additions and 14 deletions.
diff --git a/megatron/fused_kernels/fused_rotary_positional_embedding.cpp b/megatron/fused_kernels/fused_rotary_positional_embedding.cpp
@@ -19,16 +19,14 @@
 namespace fused_rope {
 
 torch::Tensor fwd_cuda(const torch::Tensor &input, const torch::Tensor &cos,
-                       const torch::Tensor &sin);
+                       const torch::Tensor &sin, const bool transpose_output);
 
 torch::Tensor bwd_cuda(const torch::Tensor &output_grads,
-                       const torch::Tensor &cos, const torch::Tensor &sin);
+                       const torch::Tensor &cos, const torch::Tensor &sin,
+                       const bool transpose_output);
 
-torch::Tensor fwd(const at::Tensor &input_, const at::Tensor &cos_,
-                  const at::Tensor &sin_) {
-  auto input = input_.contiguous();
-  auto cos = cos_.contiguous();
-  auto sin = sin_.contiguous();
+torch::Tensor fwd(const at::Tensor &input, const at::Tensor &cos,
+                  const at::Tensor &sin, const bool transpose_output) {
   TORCH_CHECK(input.dim() == 4, "expected 4D tensor");
   TORCH_CHECK(cos.dim() == 4, "expected 4D tensor");
   TORCH_CHECK(sin.dim() == 4, "expected 4D tensor");
@@ -47,14 +45,11 @@ torch::Tensor fwd(const at::Tensor &input_, const at::Tensor &cos_,
               "expected the last dim of the input tensor is greater than the "
               "sin tensor");
 
-  return fwd_cuda(input, cos, sin);
+  return fwd_cuda(input, cos, sin, transpose_output);
 }
 
-torch::Tensor bwd(const torch::Tensor &output_grads_, const at::Tensor &cos_,
-                  const at::Tensor &sin_) {
-  auto output_grads = output_grads_.contiguous();
-  auto cos = cos_.contiguous();
-  auto sin = sin_.contiguous();
+torch::Tensor bwd(const torch::Tensor &output_grads, const at::Tensor &cos,
+                  const at::Tensor &sin, const bool transpose_output) {
   TORCH_CHECK(output_grads.dim() == 4, "expected 4D tensor");
   TORCH_CHECK(cos.dim() == 4, "expected 4D tensor");
   TORCH_CHECK(sin.dim() == 4, "expected 4D tensor");
@@ -77,7 +72,7 @@ torch::Tensor bwd(const torch::Tensor &output_grads_, const at::Tensor &cos_,
       "expected the last dim of the output_grads tensor is greater than the "
       "sin tensor");
 
-  return bwd_cuda(output_grads, cos, sin);
+  return bwd_cuda(output_grads, cos, sin, transpose_output);
 }
 
 }  // end namespace fused_rope