Update base for Update on "[Executorch][SDPA] Fix bug in sdpa"

This diff fixes two bugs
1. When doing flash attention, the partical q @ k block may contain some entries that needs to be masked out. This logic had a bug. Maybe this bug also exist in PT core. I will look into that to add test and see if I can prove it.
2. Due to special handling via start_pos in SDPA it also exposed the bug in 1 when doing really long sequence prefill in chunked manner.

It is probably better to just use mask though.

Code has detail comments on the issue and fix.

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

cc larryliu0820 mergennachin cccclai helunwencser jackzhxng

[ghstack-poisoned]

Author: kimishpatel

.ci/scripts/gather_benchmark_configs.py

@@ -263,7 +263,8 @@ def is_valid_huggingface_model_id(model_name: str) -> bool:
 def get_benchmark_configs() -> Dict[str, Dict]:  # noqa: C901
     """
     Gather benchmark configurations for a given set of models on the target operating system and devices.
-
+    CHANGE IF this function's return changed:
+        extract_model_info() in executorch/.github/scripts/extract_benchmark_results.py IF YOU CHANGE THE RESULT OF THIS FUNCTION.
     Args:
         None
 

.ci/scripts/wheel/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# This file is sourced into the environment before building a pip wheel. It
+# should typically only contain shell variable assignments. Be sure to export
+# any variables so that subprocesses will see them.
+
+source "${GITHUB_WORKSPACE}/${REPOSITORY}/.ci/scripts/wheel/envvar_base.sh"

build/packaging/env_var_script_linux.sh renamed to .ci/scripts/wheel/envvar_base.sh

@@ -8,16 +8,7 @@
 # should typically only contain shell variable assignments. Be sure to export
 # any variables so that subprocesses will see them.
 
-# Enable pybindings so that users can execute ExecuTorch programs from python.
-export EXECUTORCH_BUILD_PYBIND=1
-
-# Ensure that CMAKE_ARGS is defined before referencing it. Defaults to empty
-# if not defined.
-export CMAKE_ARGS="${CMAKE_ARGS:-}"
-
-# Link the XNNPACK backend into the pybindings runtime so that users can execute
-# ExecuTorch programs that delegate to it.
-CMAKE_ARGS="${CMAKE_ARGS} -DEXECUTORCH_BUILD_XNNPACK=ON"
+source "${GITHUB_WORKSPACE}/${REPOSITORY}/.ci/scripts/wheel/envvar_base.sh"
 
 # When building for macOS, link additional backends into the pybindings runtime.
 CMAKE_ARGS="${CMAKE_ARGS} -DEXECUTORCH_BUILD_COREML=ON"

.ci/scripts/wheel/envvar_linux.sh

@@ -14,4 +14,4 @@ set -euxo pipefail
 # which does install them. Though we'd need to disable build isolation to be
 # able to see the installed torch package.
 
-pip install --progress-bar off -r requirements-dev.txt
+"${GITHUB_WORKSPACE}/${REPOSITORY}/install_requirements.sh"

build/packaging/env_var_script_macos.sh renamed to .ci/scripts/wheel/envvar_macos.sh

@@ -0,0 +1,66 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+import subprocess
+import sys
+from dataclasses import dataclass
+from functools import cache
+from typing import List
+
+
+@cache
+def _unsafe_get_env(key: str) -> str:
+    value = os.getenv(key)
+    if value is None:
+        raise RuntimeError(f"environment variable '{key}' is not set")
+    return value
+
+
+@cache
+def _repository_root_dir() -> str:
+    return os.path.join(
+        _unsafe_get_env("GITHUB_WORKSPACE"),
+        _unsafe_get_env("REPOSITORY"),
+    )
+
+
+# For some reason, we are unable to see the entire repo in the python path.
+# So manually add it.
+sys.path.append(_repository_root_dir())
+from examples.models import Backend, Model
+
+
+@dataclass
+class ModelTest:
+    model: Model
+    backend: Backend
+
+
+def run_tests(model_tests: List[ModelTest]) -> None:
+    # Why are we doing this envvar shenanigans? Since we build the testers, which
+    # uses buck, we cannot run as root. This is a sneaky of getting around that
+    # test.
+    #
+    # This can be reverted if either:
+    #   - We remove usage of buck in our builds
+    #   - We stop running the Docker image as root: https://github.com/pytorch/test-infra/issues/5091
+    envvars = os.environ.copy()
+    envvars.pop("HOME")
+
+    for model_test in model_tests:
+        subprocess.run(
+            [
+                os.path.join(_repository_root_dir(), ".ci/scripts/test_model.sh"),
+                str(model_test.model),
+                # What to build `executor_runner` with for testing.
+                "cmake",
+                str(model_test.backend),
+            ],
+            env=envvars,
+            check=True,
+            cwd=_repository_root_dir(),
+        )

build/packaging/post_build_script.sh renamed to .ci/scripts/wheel/post_build_script.sh

@@ -0,0 +1,19 @@
+#!/usr/bin/env python
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import test_base
+from examples.models import Backend, Model
+
+if __name__ == "__main__":
+    test_base.run_tests(
+        model_tests=[
+            test_base.ModelTest(
+                model=Model.Mv3,
+                backend=Backend.XnnpackQuantizationDelegation,
+            )
+        ]
+    )

build/packaging/pre_build_script.sh renamed to .ci/scripts/wheel/pre_build_script.sh

@@ -0,0 +1,19 @@
+#!/usr/bin/env python
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import test_base
+from examples.models import Backend, Model
+
+if __name__ == "__main__":
+    test_base.run_tests(
+        model_tests=[
+            test_base.ModelTest(
+                model=Model.Mv3,
+                backend=Backend.XnnpackQuantizationDelegation,
+            )
+        ]
+    )