Update setup.py

megatron/fused_kernels/setup.py

@@ -42,27 +42,35 @@ def _get_cuda_bare_metal_version(cuda_dir):
 }
 setup(
     name="fused_kernels",
-    version="0.0.1",
-    author="Sid Black & Alejandro Molina et al.",
-    author_email="[email protected]",
+    version="0.0.2",
+    author="EleutherAI",
+    author_email="[email protected]",
     include_package_data=False,
     ext_modules=[
         CUDAExtension(
-            "scaled_upper_triang_masked_softmax_cuda",
-            [
+            name="scaled_upper_triang_masked_softmax_cuda",
+            sources=[
                 str(srcpath / "scaled_upper_triang_masked_softmax.cpp"),
                 str(srcpath / "scaled_upper_triang_masked_softmax_cuda.cu"),
             ],
             extra_compile_args=cuda_ext_args,
         ),
         CUDAExtension(
-            "scaled_masked_softmax_cuda",
-            [
+            name="scaled_masked_softmax_cuda",
+            sources=[
                 str(srcpath / "scaled_masked_softmax.cpp"),
                 str(srcpath / "scaled_masked_softmax_cuda.cu"),
             ],
             extra_compile_args=cuda_ext_args,
         ),
+        CUDAExtension(
+            name="fused_rotary_positional_embedding",
+            sources=[
+                str(srcpath / "fused_rotary_positional_embedding.cpp"),
+                str(srcpath / "csrc/megatron/fused_rotary_positional_embedding_cuda.cu"),
+            ],
+            extra_compile_args=cuda_ext_args,
+        ),
     ],
     cmdclass={"build_ext": BuildExtension},
 )