Merge pull request #112 from HazyResearch/ring-attn

Add Ring Attention Kernel

Author: StuartSul

kernels/ring_attention/Makefile

@@ -0,0 +1,13 @@
+torch:
+	python3 setup.py build_ext --inplace --verbose
+	rm -rf build/
+
+# If there is an error with libc10 not found, try:
+#   import torch; print(torch.__path__[0])
+# Say whatever was printed above was PATH, run:
+#  export LD_LIBRARY_PATH=PATH/lib:$LD_LIBRARY_PATH
+# For me it was:
+#  export LD_LIBRARY_PATH=/usr/local/lib/python3.10/dist-packages/torch/lib:$LD_LIBRARY_PATH
+
+clean:
+	rm -rf build/ *.so __pycache__/ $(TARGET)

kernels/ring_attention/setup.py

@@ -0,0 +1,52 @@
+import os
+import subprocess
+from setuptools import setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+# Set environment variables
+thunderkittens_root = os.getenv('THUNDERKITTENS_ROOT', os.path.abspath(os.path.join(os.getcwd(), '../../')))
+python_include = subprocess.check_output(['python3', '-c', "import sysconfig; print(sysconfig.get_path('include'))"]).decode().strip()
+torch_include = subprocess.check_output(['python3', '-c', "import torch; from torch.utils.cpp_extension import include_paths; print(' '.join(['-I' + p for p in include_paths()]))"]).decode().strip()
+
+# CUDA flags
+cuda_flags = [
+    '-DNDEBUG',
+    '-Xcompiler=-Wno-psabi',
+    '-Xcompiler=-fno-strict-aliasing',
+    '--expt-extended-lambda',
+    '--expt-relaxed-constexpr',
+    '-forward-unknown-to-host-compiler',
+    '--use_fast_math',
+    '-std=c++20',
+    '-O3',
+    '-Xnvlink=--verbose',
+    '-Xptxas=--verbose',
+    '-Xptxas=--warn-on-spills',
+    f'-I{thunderkittens_root}/include',
+    f'-I{thunderkittens_root}/prototype',
+    f'-I{python_include}',
+    '-DTORCH_COMPILE',
+    '-DKITTENS_HOPPER', # assume H100 for ring attn
+    '-arch=sm_90a',     # assume H100 for ring attn
+] + torch_include.split()
+cpp_flags = [
+    '-std=c++20',
+    '-O3'
+]
+source_files = ['tk_ring_attention.cu']
+
+setup(
+    name='tk_ring_attention',
+    ext_modules=[
+        CUDAExtension(
+            'tk_ring_attention',
+            sources=source_files, 
+            extra_compile_args={'cxx' : cpp_flags,
+                                'nvcc' : cuda_flags}, 
+            libraries=['cuda']
+        )
+    ],
+    cmdclass={
+        'build_ext': BuildExtension
+    }
+)