Merge pull request #7 from nod-ai/flash-attention

Flash attention

.github/workflows/run_bench.yml

@@ -3,7 +3,7 @@ name: Build and Test
 on:
   pull_request:
   schedule:
-    - cron: '0 8 * * *'
+    - cron: "0 8 * * *"
 
 jobs:
   build:
@@ -13,8 +13,8 @@ jobs:
       - name: Checkout repository
         uses: actions/checkout@v3
         with:
-          submodules: 'true'
-      
+          submodules: "true"
+
       - name: Update IREE
         run: |
           cd $GITHUB_WORKSPACE/third_party/iree
@@ -26,7 +26,7 @@ jobs:
       - name: Set up Ninja
         uses: ashutoshvarma/setup-ninja@v1
         env:
-          ACTIONS_ALLOW_UNSECURE_COMMANDS: 'true'
+          ACTIONS_ALLOW_UNSECURE_COMMANDS: "true"
 
       - name: Build LLVM Project
         run: |
@@ -61,7 +61,7 @@ jobs:
           cd $GITHUB_WORKSPACE
           source $GITHUB_WORKSPACE/venv/bin/activate
           for device in $(seq 3 7); do (sudo $GITHUB_WORKSPACE/build/gemm-bench --device=$device &); done
-          ./gb run --backends=rocblas --output=rocblas.hdf
+          ./gb run --backends=rocblas --repeat=1 --output=rocblas.hdf
           sudo pkill -f gemm-bench
           deactivate
 
@@ -70,10 +70,10 @@ jobs:
           cd $GITHUB_WORKSPACE
           source $GITHUB_WORKSPACE/venv/bin/activate
           for device in $(seq 3 7); do (sudo $GITHUB_WORKSPACE/build/gemm-bench --device=$device &); done
-          ./gb run --backends=iree --output=iree.hdf
+          ./gb run --backends=iree --repeat=1 --output=iree.hdf
           sudo pkill -f gemm-bench
           deactivate
-      
+
       - name: Upload IREE benchmark results
         uses: actions/upload-artifact@v4
         with:

.gitignore

@@ -1,5 +1,7 @@
 build
 rocm_gemm_venv
-*.venv
+*venv
 __pycache__
-*.hdf
+*.hdf
+attention/mlir
+attention/vmfb

.vscode/c_cpp_properties.json

@@ -0,0 +1,29 @@
+{
+  "configurations": [
+    {
+      "name": "Linux",
+      "includePath": [
+        "${workspaceFolder}/third_party/llvm-project/llvm/include",
+        "${workspaceFolder}/third_party/llvm-project/mlir/include",
+        "${workspaceFolder}/third_party/llvm-project/build/include",
+        "${workspaceFolder}/third_party/llvm-project/build/tools/mlir/include",
+        "${workspaceFolder}/src/ireekernels/include",
+        "${workspaceFolder}/third_party/iree/compiler/bindings/c",
+        "${workspaceFolder}/third_party/iree/runtime/src",
+        "${workspaceFolder}/src/benchmark",
+        "/opt/rocm/include",
+        "/usr/include",
+        "/usr/include/clang/15/include",
+        "/usr/include/c++/12",
+        "/usr/include/x86_64-linux-gnu/c++/12/",
+        "/usr/include/python3.10"
+      ],
+      "defines": [],
+      "cStandard": "c17",
+      "cppStandard": "c++14",
+      "intelliSenseMode": "linux-clang-x64",
+      "compileCommands": "${workspaceFolder}/src/ireekernelsbuild/compile_commands.json"
+    }
+  ],
+  "version": 4
+}

.vscode/settings.json

@@ -0,0 +1,124 @@
+{
+  "cmake.configureOnOpen": false,
+  "C_Cpp.default.compilerPath": "/usr/bin/clang++",
+  "files.associations": {
+    "*.inc": "cpp",
+    "string": "cpp",
+    "filesystem": "cpp",
+    "cstring": "cpp",
+    "print": "cpp",
+    "__nullptr": "cpp",
+    "exception": "cpp",
+    "initializer_list": "cpp",
+    "new": "cpp",
+    "stdexcept": "cpp",
+    "typeinfo": "cpp",
+    "algorithm": "cpp",
+    "cmath": "cpp",
+    "cstddef": "cpp",
+    "cstdint": "cpp",
+    "cstdio": "cpp",
+    "cstdlib": "cpp",
+    "limits": "cpp",
+    "type_traits": "cpp",
+    "__bit_reference": "cpp",
+    "__locale": "cpp",
+    "__verbose_abort": "cpp",
+    "array": "cpp",
+    "bitset": "cpp",
+    "cctype": "cpp",
+    "clocale": "cpp",
+    "condition_variable": "cpp",
+    "cstdarg": "cpp",
+    "ctime": "cpp",
+    "cwchar": "cpp",
+    "execution": "cpp",
+    "memory": "cpp",
+    "ios": "cpp",
+    "iosfwd": "cpp",
+    "iostream": "cpp",
+    "istream": "cpp",
+    "locale": "cpp",
+    "mutex": "cpp",
+    "ratio": "cpp",
+    "semaphore": "cpp",
+    "sstream": "cpp",
+    "streambuf": "cpp",
+    "string_view": "cpp",
+    "atomic": "cpp",
+    "bit": "cpp",
+    "*.tcc": "cpp",
+    "chrono": "cpp",
+    "compare": "cpp",
+    "concepts": "cpp",
+    "vector": "cpp",
+    "functional": "cpp",
+    "iterator": "cpp",
+    "random": "cpp",
+    "system_error": "cpp",
+    "tuple": "cpp",
+    "utility": "cpp",
+    "ostream": "cpp",
+    "stop_token": "cpp",
+    "thread": "cpp",
+    "compile.h": "c",
+    "*.ipp": "c",
+    "expected": "c",
+    "utils.h": "c",
+    "embedding_api.h": "c",
+    "string_builder.h": "c",
+    "__hash_table": "cpp",
+    "__split_buffer": "cpp",
+    "__tree": "cpp",
+    "deque": "cpp",
+    "forward_list": "cpp",
+    "list": "cpp",
+    "map": "cpp",
+    "regex": "cpp",
+    "set": "cpp",
+    "span": "cpp",
+    "unordered_map": "cpp",
+    "unordered_set": "cpp",
+    "valarray": "cpp",
+    "ranges": "cpp",
+    "charconv": "cpp",
+    "api.h": "c",
+    "loader.h": "c",
+    "*.def": "cpp",
+    "future": "cpp",
+    "runtime.h": "c",
+    "buffer_view.h": "c",
+    "complex": "cpp",
+    "__node_handle": "cpp",
+    "any": "cpp",
+    "cinttypes": "cpp",
+    "codecvt": "cpp",
+    "csetjmp": "cpp",
+    "csignal": "cpp",
+    "cwctype": "cpp",
+    "fstream": "cpp",
+    "iomanip": "cpp",
+    "numbers": "cpp",
+    "optional": "cpp",
+    "queue": "cpp",
+    "shared_mutex": "cpp",
+    "stack": "cpp",
+    "variant": "cpp",
+    "memory_resource": "cpp",
+    "numeric": "cpp",
+    "cfenv": "cpp",
+    "__functional_base": "c",
+    "__memory": "c",
+    "libc-header-start.h": "c",
+    "typeindex": "cpp",
+    "syncstream": "c",
+    "runtime.hpp": "c",
+    "strstream": "cpp",
+    "__functional_03": "cpp"
+  },
+  "editor.formatOnSave": true,
+  "editor.rulers": [80],
+  "clang-format.executable": "usr/bin/clang-format",
+  "C_Cpp.errorSquiggles": "enabled",
+  "dotnet.defaultSolution": "disable"
+}