Merge branch 'main' into temporary_decompose

.github/workflows/ci-llama.yaml

@@ -1,3 +1,9 @@
+# Copyright 2024 Advanced Micro Devices, Inc.
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
 name: Llama Benchmarking Tests
 
 on:

.github/workflows/ci-sdxl.yaml

@@ -56,7 +56,7 @@ jobs:
         repository: iree-org/iree
         path: ${{ env.IREE_REPO_DIR }}
         submodules: false
-        ref: 67ba1c45424d5cedc7baf7bfe8a998ee86e510af
+        ref: candidate-20241029.1062
 
     - name: Initalize IREE submodules
       working-directory: ${{ env.IREE_REPO_DIR }}

.github/workflows/ci-sharktank.yml

@@ -1,3 +1,9 @@
+# Copyright 2024 Advanced Micro Devices, Inc.
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
 name: CI - sharktank
 
 on:

.github/workflows/ci-tuner.yml

@@ -1,3 +1,9 @@
+# Copyright 2024 Advanced Micro Devices, Inc.
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
 name: CI - Tuner
 
 on:

.github/workflows/ci_eval.yaml

@@ -1,3 +1,9 @@
+# Copyright 2024 Advanced Micro Devices, Inc.
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
 name: Evaluation Tests
 
 on:

.github/workflows/ci_linux_x64-libshortfin.yml

@@ -38,6 +38,9 @@ jobs:
   build-and-test:
     name: Build and test
     runs-on: ubuntu-24.04
+    strategy:
+      matrix:
+        python-version: ["3.11", "3.12"]
 
     steps:
     - name: Install dependencies
@@ -56,7 +59,7 @@ jobs:
         repository: iree-org/iree
         path: ${{ env.IREE_REPO_DIR }}
         submodules: false
-        ref: 67ba1c45424d5cedc7baf7bfe8a998ee86e510af
+        ref: candidate-20241029.1062
 
     - name: Initalize IREE submodules
       working-directory: ${{ env.IREE_REPO_DIR }}
@@ -67,10 +70,10 @@ jobs:
         git submodule update --init --depth 1 -- third_party/googletest
         git submodule update --init --depth 1 -- third_party/hip-build-deps/
 
-    - name: Setup Python
+    - name: Setup Python ${{ matrix.python-version }}
       uses: actions/setup-python@39cd14951b08e74b54015e9e001cdefcf80e669f # v5.1.1
       with:
-        python-version: "3.12"
+        python-version: ${{ matrix.python-version }}
         cache: "pip"
     - name: Install Python packages
       # TODO: Switch to `pip install -r requirements.txt -e shortfin/`.

.github/workflows/ci_linux_x64_asan-libshortfin.yml

@@ -109,7 +109,7 @@ jobs:
         repository: iree-org/iree
         path: ${{ env.IREE_SOURCE_DIR }}
         submodules: false
-        ref: 67ba1c45424d5cedc7baf7bfe8a998ee86e510af
+        ref: candidate-20241029.1062
 
     - name: Initalize IREE submodules
       working-directory: ${{ env.IREE_SOURCE_DIR }}

.github/workflows/ci_linux_x64_nogil-libshortfin.yml

@@ -57,7 +57,7 @@ jobs:
         repository: iree-org/iree
         path: ${{ env.IREE_REPO_DIR }}
         submodules: false
-        ref: 67ba1c45424d5cedc7baf7bfe8a998ee86e510af
+        ref: candidate-20241029.1062
 
     - name: Initalize IREE submodules
       working-directory: ${{ env.IREE_REPO_DIR }}

.github/workflows/test.yaml

@@ -1,3 +1,9 @@
+# Copyright 2024 Advanced Micro Devices, Inc.
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
 name: Integration Tests
 
 on:

sharktank/sharktank/layers/kv_cache.py

@@ -191,8 +191,8 @@ def write_timestep(
         update_count = len(cache_partitions)
 
         for b in range(bs):
-            row_index = torch.tensor(b, dtype=torch.int64)
-            row_start_pos = seq_positions[row_index]
+            row_index = torch.tensor([b], dtype=torch.int64)
+            row_start_pos = seq_positions[row_index].unsqueeze(0)
 
             for i, update in enumerate(cache_partitions):
                 cache = state[transformer_block_index * update_count + i]

shortfin/CMakeLists.txt

@@ -183,7 +183,7 @@ elseif (SHORTFIN_BUNDLE_DEPS)
   FetchContent_Declare(
     shortfin_iree
     GIT_REPOSITORY https://github.com/iree-org/iree.git
-    GIT_TAG 67ba1c45424d5cedc7baf7bfe8a998ee86e510af
+    GIT_TAG candidate-20241029.1062
     GIT_SUBMODULES ${IREE_SUBMODULES}
     GIT_SHALLOW TRUE
     SYSTEM

shortfin/python/shortfin_apps/llm/components/generate.py

@@ -54,10 +54,12 @@ async def run(self):
             token = sfnp.argmax(exec.result_logits)
             token_int = token.items[0]
 
+            self.append_token(token_int)
             # Decode loop.
-            # TODO: Use correct eot token from config.
             exec.start_position = len(self.input_token_ids) - 1
-            while token_int != 128001:
+            # TODO: Use correct eot token from config.
+            # while token_int != 128001:
+            for i in range(15):
                 exec.reset(InferencePhase.DECODE)
                 exec.input_token_ids = [token_int]
                 exec.start_position += 1

shortfin/python/shortfin_apps/llm/components/service.py

@@ -337,21 +337,30 @@ async def run(self):
                         m.items = self.exec_requests[i].input_token_ids
             tokens_host.copy_to(tokens)
 
-            # Populate seq_lens.
-            seq_lens_host = seq_lens.for_transfer()
-            with seq_lens_host.map(discard=True) as m:
-                m.fill(0)
-                m.items = [len(req.input_token_ids) for req in self.exec_requests]
-            seq_lens_host.copy_to(seq_lens)
-
-            # For decode, populate start_positions.
+            # For prefill, populate seq_lens
+            if self.phase == InferencePhase.PREFILL:
+                seq_lens_host = seq_lens.for_transfer()
+                with seq_lens_host.map(discard=True) as m:
+                    m.fill(0)
+                    m.items = [len(req.input_token_ids) for req in self.exec_requests]
+                seq_lens_host.copy_to(seq_lens)
+
+            # For decode, populate start_positions and seq_lens.
+            # paged_llm_v1 and export_paged_llm_v1 do some funky things with start_positions and seq_lens
+            # TODO: make them not so funky
             if self.phase == InferencePhase.DECODE:
                 start_positions_host = start_positions.for_transfer()
                 with start_positions_host.map(discard=True) as m:
                     m.fill(0)
                     m.items = [req.start_position for req in self.exec_requests]
                 start_positions_host.copy_to(start_positions)
 
+                seq_lens_host = seq_lens.for_transfer()
+                with seq_lens_host.map(discard=True) as m:
+                    m.fill(0)
+                    m.items = [req.start_position + 1 for req in self.exec_requests]
+                seq_lens_host.copy_to(seq_lens)
+
             # Populate cache pages.
             seq_block_ids_host = seq_block_ids.for_transfer()
             for i in range(bs):

shortfin/requirements-iree-compiler.txt

@@ -1,3 +1,4 @@
 # Keep in sync with IREE_REF in CI and GIT_TAG in CMakeLists.txt
 -f https://iree.dev/pip-release-links.html
-iree-compiler==20240904.1006
+iree-compiler==20241029.1062
+iree-runtime==20241029.1062

shortfin/setup.py

@@ -111,12 +111,12 @@ def is_cpp_prebuilt():
 
 # Due to a quirk of setuptools, that package_dir map must only contain
 # paths relative to the directory containing setup.py. Why? No one knows.
-REL_SOURCE_DIR = SOURCE_DIR.relative_to(SETUPPY_DIR, walk_up=True)
-REL_BINARY_DIR = BINARY_DIR.relative_to(SETUPPY_DIR, walk_up=True)
-REL_CMAKE_DEFAULT_BUILD_DIR = CMAKE_DEFAULT_BUILD_DIR.relative_to(
-    SETUPPY_DIR, walk_up=True
+REL_SOURCE_DIR = Path(os.path.relpath(SOURCE_DIR, SETUPPY_DIR))
+REL_BINARY_DIR = Path(os.path.relpath(BINARY_DIR, SETUPPY_DIR))
+REL_CMAKE_DEFAULT_BUILD_DIR = Path(
+    os.path.relpath(CMAKE_DEFAULT_BUILD_DIR, SETUPPY_DIR)
 )
-REL_CMAKE_TRACY_BUILD_DIR = CMAKE_TRACY_BUILD_DIR.relative_to(SETUPPY_DIR, walk_up=True)
+REL_CMAKE_TRACY_BUILD_DIR = Path(os.path.relpath(CMAKE_TRACY_BUILD_DIR, SETUPPY_DIR))
 
 
 class CMakeExtension(Extension):

shortfin/tests/apps/sd/e2e_test.py

@@ -1,3 +1,9 @@
+# Copyright 2024 Advanced Micro Devices, Inc.
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
 import json
 import requests
 import time