Add tilelang install and flash_attention impl (#180)

Summary:
To install tilelang:

```
python install.py --tile
```

To test tilelang:

```
python run.py --op flash_attention --only tile,triton_tutorial_flash_v2 --metrics tflops

  (Batch, Heads, SeqLen, Dhead)    tile-tflops    triton_tutorial_flash_v2-tflops
-------------------------------  -------------  ---------------------------------
               (4, 48, 128, 64)        40.7874                            57.0654
               (4, 48, 256, 64)        97.2592                           132.104
               (4, 48, 512, 64)       159.593                            256.794
              (4, 48, 1024, 64)       223.976                            331.811
              (4, 48, 2048, 64)       263.711                            345.755
              (4, 48, 4096, 64)       277.59                             354.224
              (4, 48, 8192, 64)       288.024                            350.316
             (4, 48, 16384, 64)       292.826                            351.27
                        average       205.471                            272.418
```

Pull Request resolved: #180

Reviewed By: FindHao

Differential Revision: D72250350

Pulled By: xuzhao9

fbshipit-source-id: 42947fd43d3e3fadce7adb70141c061787ee9e48

Author: xuzhao9

.ci/tritonbench/test-install.sh

@@ -19,3 +19,4 @@ torch_compile_debug/
 build/
 /*.csv
 *.hatchet
+autotuner.log

.gitignore

@@ -53,11 +53,12 @@ We depend on the following projects as a source of customized Triton or CUTLASS
 * (CUDA, HIP) [kernels](https://github.com/triton-lang/kernels)
 * (CUDA, HIP) [generative-recommenders](https://github.com/facebookresearch/generative-recommenders)
 * (CUDA, HIP) [Liger-Kernel](https://github.com/linkedin/Liger-Kernel)
+* (CUDA, HIP) [tilelang](https://github.com/tile-ai/tilelang)
 * (CUDA) [xformers](https://github.com/facebookresearch/xformers)
 * (CUDA) [flash-attention](https://github.com/Dao-AILab/flash-attention)
 * (CUDA) [FBGEMM](https://github.com/pytorch/FBGEMM)
 * (CUDA) [ThunderKittens](https://github.com/HazyResearch/ThunderKittens)
-* (CUDA) [cutlass-kernels](https://github.com/ColfaxResearch/cutlass-kernels)
+
 
 
 ## License

README.md

@@ -60,10 +60,6 @@ RUN cd /workspace/tritonbench && \
 RUN cd /workspace/tritonbench && \
     bash .ci/tritonbench/install.sh
 
-# Test Tritonbench
-RUN cd /workspace/tritonbench && \
-    bash .ci/tritonbench/test-install.sh
-
 # Remove NVIDIA driver library - they are supposed to be mapped at runtime
 RUN sudo apt-get purge -y libnvidia-compute-550
 

docker/tritonbench-nightly.dockerfile

@@ -125,10 +125,10 @@ def setup_hip(args: argparse.Namespace):
     parser.add_argument("--tk", action="store_true", help="Install ThunderKittens")
     parser.add_argument("--liger", action="store_true", help="Install Liger-kernel")
     parser.add_argument("--xformers", action="store_true", help="Install xformers")
+    parser.add_argument("--tile", action="store_true", help="install tile lang")
     parser.add_argument(
         "--all", action="store_true", help="Install all custom kernel repos"
     )
-    parser.add_argument("--test", action="store_true", help="Run tests")
     args = parser.parse_args()
 
     if args.all and is_hip():
@@ -166,6 +166,7 @@ def setup_hip(args: argparse.Namespace):
     if args.fbgemm or args.fbgemm_all or args.all:
         logger.info("[tritonbench] installing FBGEMM...")
         install_fbgemm(genai=(not args.fbgemm_all))
+        test_fbgemm()
     if args.fa2 or args.all:
         logger.info("[tritonbench] installing fa2 from source...")
         install_fa2(compile=True)
@@ -182,6 +183,11 @@ def setup_hip(args: argparse.Namespace):
         from tools.tk.install import install_tk
 
         install_tk()
+    if args.tile:
+        logger.info("[tritonbench] installing tilelang...")
+        from tools.tilelang.install import install_tile
+
+        install_tile()
     if args.liger or args.all:
         logger.info("[tritonbench] installing liger-kernels...")
         install_liger()
@@ -191,6 +197,3 @@ def setup_hip(args: argparse.Namespace):
 
         install_xformers()
     logger.info("[tritonbench] installation complete!")
-    # run tests to check installation
-    if args.test:
-        test_fbgemm()

install.py

@@ -0,0 +1,23 @@
+import os
+import subprocess
+import sys
+
+REQUIREMENTS_FILE = os.path.join(
+    os.path.dirname(os.path.abspath(__file__)), "requirements.txt"
+)
+
+
+def install_requirements(requirements_txt: str):
+    # ignore dependencies to bypass reinstalling pytorch stable version
+    cmd = ["pip", "install", "-r", requirements_txt, "--no-deps"]
+    subprocess.check_call(cmd)
+
+
+def check_install():
+    cmd = [sys.executable, "-c", "import tilelang"]
+    subprocess.check_call(cmd)
+
+
+def install_tile():
+    install_requirements(REQUIREMENTS_FILE)
+    check_install()

tools/tilelang/install.py

@@ -0,0 +1,3 @@
+tilelang==0.1.3
+Cython
+decorator

tools/tilelang/requirements.txt

@@ -103,6 +103,15 @@
 except (ImportError, IOError, AttributeError, TypeError):
     HAS_XFORMERS = False
 
+try:
+    import tilelang
+
+    from .tilelang_mha import tilelang_mha
+
+    HAS_TILELANG = True
+except (ImportError, IOError, AttributeError, TypeError):
+    HAS_TILELANG = False
+
 # [Optional] colfax cutlass backend
 try:
     if not hasattr(torch.version, "git_version"):
@@ -467,6 +476,37 @@ def _inner():
 
         return _inner
 
+    @register_benchmark(enabled=HAS_TILELANG)
+    def tile(self, q, k, v):
+        # [B, H, S, D] -> [B, S, H, D]
+        q = q.transpose(1, 2).contiguous()
+        k = k.transpose(1, 2).contiguous()
+        v = v.transpose(1, 2).contiguous()
+        best_config = tilelang_mha(
+            self.BATCH,
+            self.H,
+            self.N_CTX,
+            self.D_HEAD,
+            self.causal,
+            self.dtype,
+            tune=True,
+        )[1]
+        func = tilelang_mha(
+            self.BATCH,
+            self.H,
+            self.N_CTX,
+            self.D_HEAD,
+            self.causal,
+            self.dtype,
+        )(*best_config)
+        jit_kernel = tilelang.compile(func, out_idx=[3])
+
+        def _inner():
+            o = jit_kernel(q, k, v)
+            return o
+
+        return _inner
+
     @register_benchmark(enabled=False, label=f"cudnn")
     def cudnn(self, q, k, v):
         os.environ["TORCH_CUDNN_SDPA_ENABLED"] = "1"