Update on "add attention_sink.py"

This PR adds `KVCacheWithAttentionSink`, which is required for `AttentionSink`. It keeps the first `sink_size` tokens as attention sinks and maintains a sliding window with `window_size` for new tokens.

Note: I am trying to implement and verify `AttentionSink` in eager mode first. So the current implementation may still have some lower errors or performance issue. For example, it does not support the case when dynamic shape is disabled. Will leave these problems to resolve when we are ready to deploy `AttentionSink` to edge.

Differential Revision: [D65235798](https://our.internmc.facebook.com/intern/diff/D65235798/)

[ghstack-poisoned]

Author: helunwencser

.ci/docker/ci_commit_pins/pytorch.txt

@@ -1 +1 @@
-bd5482c7c3e1197e10c46ff739027f917d9c1fcc
+c8a648d4dffb9f0133ff4a2ea0e660b42105d3ad

.ci/docker/common/install_clang.sh

@@ -13,7 +13,7 @@ install_ubuntu() {
   apt-get install -y --no-install-recommends clang-"$CLANG_VERSION"
   apt-get install -y --no-install-recommends llvm-"$CLANG_VERSION"
   # Also require LLD linker from llvm and libomp to build PyTorch from source
-  apt-get install -y lld "libomp-${CLANG_VERSION}-dev"
+  apt-get install -y lld "libomp-${CLANG_VERSION}-dev" "libc++-${CLANG_VERSION}-dev"
 
   # Use update-alternatives to make this version the default
   update-alternatives --install /usr/bin/clang clang /usr/bin/clang-"$CLANG_VERSION" 50

.ci/docker/requirements-ci.txt

@@ -1,5 +1,5 @@
 mpmath==1.3.0
-numpy==1.22.0; python_version == '3.10'
+numpy==1.21.3; python_version == '3.10'
 numpy==1.23.2; python_version == '3.11'
 numpy; python_version >= '3.12'
 PyYAML==6.0.1

.ci/scripts/setup-linux.sh

@@ -19,6 +19,7 @@ else
 fi
 
 # As Linux job is running inside a Docker container, all of its dependencies
-# have already been installed
-install_executorch
+# have already been installed, so we use PyTorch build from source here instead
+# of nightly. This allows CI to test against latest commits from PyTorch
+install_executorch "use-pt-pinned-commit"
 build_executorch_runner "${BUILD_TOOL}"

.ci/scripts/setup-qnn-deps.sh

@@ -31,8 +31,9 @@ install_qnn() {
 }
 
 setup_libc++() {
+  clang_version=$1
   sudo apt-get update
-  pkgs_to_check=('libc++-dev')
+  pkgs_to_check=("libc++-${clang_version}-dev")
   j=0
   while [ $j -lt ${#pkgs_to_check[*]} ]; do
     install_status=$(verify_pkg_installed ${pkgs_to_check[$j]})
@@ -47,5 +48,6 @@ setup_libc++() {
   done
 }
 
-setup_libc++
+# This needs to match with the clang version from the Docker image
+setup_libc++ 12
 install_qnn

.ci/scripts/utils.sh

@@ -20,8 +20,11 @@ install_executorch() {
   which pip
   # Install executorch, this assumes that Executorch is checked out in the
   # current directory.
-  # TODO(T199538337): clean up install scripts to use install_requirements.sh
-  ./install_requirements.sh --pybind xnnpack
+  if [[ "${1:-}" == "use-pt-pinned-commit" ]]; then
+    ./install_requirements.sh --pybind xnnpack --use-pt-pinned-commit
+  else
+    ./install_requirements.sh --pybind xnnpack
+  fi
   # Just print out the list of packages for debugging
   pip list
 }

.github/pull_request_template.md

@@ -0,0 +1,9 @@
+### Summary
+[PLEASE REMOVE] See [CONTRIBUTING.md's Pull Requests](https://github.com/pytorch/executorch/blob/main/CONTRIBUTING.md#pull-requests) for ExecuTorch PR guidelines.
+
+[PLEASE REMOVE] If this PR closes an issue, please add a `Fixes #<issue-id>` line.
+
+[PLEASE REMOVE] If this PR introduces a fix or feature that should be the upcoming release notes, please add a "Release notes: <area>" label. For a list of available release notes labels, check out [CONTRIBUTING.md's Pull Requests](https://github.com/pytorch/executorch/blob/main/CONTRIBUTING.md#pull-requests).
+
+### Test plan
+[PLEASE REMOVE] How did you test this PR? Please write down any manual commands you used and note down tests that you have written if applicable.

.github/workflows/_android.yml

@@ -66,6 +66,16 @@ jobs:
           # avoid permission issue
           sudo chown -R "${USER}" /opt/android
 
+      - name: Download Artifacts
+        shell: bash
+        run: |
+          set -eux
+          curl -O https://gha-artifacts.s3.amazonaws.com/${{ github.repository }}/${{ github.run_id }}/artifacts/llm_demo/app-debug.apk
+          curl -O https://gha-artifacts.s3.amazonaws.com/${{ github.repository }}/${{ github.run_id }}/artifacts/llm_demo/app-debug-androidTest.apk
+          curl -O https://gha-artifacts.s3.amazonaws.com/${{ github.repository }}/${{ github.run_id }}/artifacts/fp32-xnnpack-custom/model.zip
+          unzip model.zip
+          mv *.pte model.pte
+
       - name: Gradle cache
         uses: gradle/actions/setup-gradle@v3
 

.github/workflows/_unittest.yml

@@ -37,6 +37,9 @@ jobs:
         CMAKE_ARGS="-DEXECUTORCH_BUILD_XNNPACK=ON -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON" \
         .ci/scripts/setup-linux.sh cmake
 
+        # Install llama3_2_vision dependencies.
+        PYTHON_EXECUTABLE=python ./examples/models/llama3_2_vision/install_requirements.sh
+
         # Run pytest with coverage
         pytest -n auto --cov=./ --cov-report=xml
         # Run gtest
@@ -67,6 +70,10 @@ jobs:
         ${CONDA_RUN} --no-capture-output \
         .ci/scripts/setup-macos.sh cmake
 
+        # Install llama3_2_vision dependencies.
+        PYTHON_EXECUTABLE=python ${CONDA_RUN} --no-capture-output \
+        ./examples/models/llama3_2_vision/install_requirements.sh
+
         # Run pytest with coverage
         ${CONDA_RUN} pytest -n auto --cov=./ --cov-report=xml
         # Run gtest

.github/workflows/pull.yml

@@ -35,9 +35,6 @@ jobs:
     name: test-setup-linux-gcc
     uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
     strategy:
-      matrix:
-        include:
-          - build-tool: cmake
       fail-fast: false
     with:
       runner: linux.2xlarge
@@ -50,7 +47,7 @@ jobs:
         CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]")
         conda activate "${CONDA_ENV}"
 
-        BUILD_TOOL=${{ matrix.build-tool }}
+        BUILD_TOOL="cmake"
 
         PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh "${BUILD_TOOL}"
         # Build and test ExecuTorch with the add model on portable backend.
@@ -89,20 +86,11 @@ jobs:
     strategy:
       matrix:
         dtype: [fp32]
-        build-tool: [buck2, cmake]
         mode: [portable, xnnpack+custom, xnnpack+custom+qe]
         include:
           - dtype: bf16
-            build-tool: cmake
-            mode: portable
-          - dtype: bf16
-            build-tool: buck2
             mode: portable
           - dtype: bf16
-            build-tool: cmake
-            mode: custom
-          - dtype: bf16
-            build-tool: buck2
             mode: custom
       fail-fast: false
     with:
@@ -111,29 +99,30 @@ jobs:
       submodules: 'true'
       ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
       timeout: 900
+      upload-artifact: android-models
+      upload-artifact-to-s3: true
       script: |
         # The generic Linux job chooses to use base env, not the one setup by the image
         CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]")
         conda activate "${CONDA_ENV}"
 
         DTYPE=${{ matrix.dtype }}
-        BUILD_TOOL=${{ matrix.build-tool }}
+        BUILD_TOOL="cmake"
         MODE=${{ matrix.mode }}
+        ARTIFACTS_DIR_NAME="artifacts-to-be-uploaded/${DTYPE}-${MODE}"
+        ARTIFACTS_DIR_NAME="${ARTIFACTS_DIR_NAME/+/-}"
 
         # Setup executorch
-        PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh buck2
+        PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh "${BUILD_TOOL}"
         # Install requirements for export_llama
         PYTHON_EXECUTABLE=python bash examples/models/llama/install_requirements.sh
         # Test llama2
-        PYTHON_EXECUTABLE=python bash .ci/scripts/test_llama.sh stories110M "${BUILD_TOOL}" "${DTYPE}" "${MODE}"
+        PYTHON_EXECUTABLE=python bash .ci/scripts/test_llama.sh stories110M "${BUILD_TOOL}" "${DTYPE}" "${MODE}" "${ARTIFACTS_DIR_NAME}"
 
   test-llama-runner-linux-android:
     name: test-llama-runner-linux-android
     uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
     strategy:
-      matrix:
-        include:
-          - build-tool: cmake
       fail-fast: false
     with:
       runner: linux.2xlarge
@@ -146,18 +135,14 @@ jobs:
         CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]")
         conda activate "${CONDA_ENV}"
 
-        BUILD_TOOL=${{ matrix.build-tool }}
+        BUILD_TOOL="cmake"
         PYTHON_EXECUTABLE=python \
         bash .ci/scripts/build_llama_android.sh  "${BUILD_TOOL}"
 
   test-custom-ops-linux:
     name: test-custom-ops-linux
     uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
     strategy:
-      matrix:
-        include:
-          - build-tool: buck2
-          - build-tool: cmake
       fail-fast: false
     with:
       runner: linux.2xlarge
@@ -170,7 +155,7 @@ jobs:
         CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]")
         conda activate "${CONDA_ENV}"
 
-        BUILD_TOOL=${{ matrix.build-tool }}
+        BUILD_TOOL="cmake"
         PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh "${BUILD_TOOL}"
         # Test custom ops
         PYTHON_EXECUTABLE=python bash examples/portable/custom_ops/test_custom_ops.sh "${BUILD_TOOL}"
@@ -179,10 +164,6 @@ jobs:
     name: test-selective-build-linux
     uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
     strategy:
-      matrix:
-        include:
-          - build-tool: buck2
-          - build-tool: cmake
       fail-fast: false
     with:
       runner: linux.2xlarge
@@ -195,7 +176,7 @@ jobs:
         CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]")
         conda activate "${CONDA_ENV}"
 
-        BUILD_TOOL=${{ matrix.build-tool }}
+        BUILD_TOOL="cmake"
         PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh "${BUILD_TOOL}"
         # Test selective build
         PYTHON_EXECUTABLE=python bash examples/selective_build/test_selective_build.sh "${BUILD_TOOL}"
@@ -235,9 +216,6 @@ jobs:
     name: test-quantized-aot-lib-linux
     uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
     strategy:
-      matrix:
-        include:
-          - build-tool: cmake
       fail-fast: false
     with:
       runner: linux.2xlarge
@@ -250,17 +228,14 @@ jobs:
         CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]")
         conda activate "${CONDA_ENV}"
 
-        BUILD_TOOL=${{ matrix.build-tool }}
+        BUILD_TOOL="cmake"
         PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh "${BUILD_TOOL}"
         PYTHON_EXECUTABLE=python bash examples/xnnpack/quantization/test_quantize.sh "${BUILD_TOOL}" mv2
 
   test-pybind-build-linux:
     name: test-pybind-build-linux
     uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
     strategy:
-      matrix:
-        include:
-          - build-tool: cmake
       fail-fast: false
     with:
       runner: linux.2xlarge
@@ -274,7 +249,7 @@ jobs:
         conda activate "${CONDA_ENV}"
 
         # build module for executorch.extension.pybindings.portable_lib
-        BUILD_TOOL=${{ matrix.build-tool }}
+        BUILD_TOOL="cmake"
         PYTHON_EXECUTABLE=python \
         EXECUTORCH_BUILD_XNNPACK=ON \
         EXECUTORCH_BUILD_PYBIND=ON \
@@ -349,6 +324,7 @@ jobs:
 
   android:
     uses: ./.github/workflows/_android.yml
+    needs: test-llama-runner-linux
 
   unittest:
     uses: ./.github/workflows/_unittest.yml
@@ -357,10 +333,6 @@ jobs:
 
   unittest-arm:
     uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
-    strategy:
-      matrix:
-        include:
-          - build-tool: buck2
     with:
       runner: linux.2xlarge
       docker-image: executorch-ubuntu-22.04-arm-sdk
@@ -374,7 +346,7 @@ jobs:
         CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]")
         conda activate "${CONDA_ENV}"
 
-        BUILD_TOOL=${{ matrix.build-tool }}
+        BUILD_TOOL="cmake"
 
         # Setup MacOS dependencies as there is no Docker support on MacOS atm
         PYTHON_EXECUTABLE=python \
@@ -396,7 +368,6 @@ jobs:
     strategy:
       matrix:
         dtype: [fp32]
-        build-tool: [cmake]
         mode: [qnn]
       fail-fast: false
     with:
@@ -411,14 +382,14 @@ jobs:
         conda activate "${CONDA_ENV}"
 
         DTYPE=${{ matrix.dtype }}
-        BUILD_TOOL=${{ matrix.build-tool }}
+        BUILD_TOOL="cmake"
         MODE=${{ matrix.mode }}
 
         PYTHON_EXECUTABLE=python bash .ci/scripts/setup-qnn-deps.sh
         PYTHON_EXECUTABLE=python bash .ci/scripts/build-qnn-sdk.sh
 
         # Setup executorch
-        PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh buck2
+        PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh "${BUILD_TOOL}"
         # Install requirements for export_llama
         PYTHON_EXECUTABLE=python bash examples/models/llama/install_requirements.sh
         # Test llama2