Merge branch 'main' into arm-backend-vgf-kv-cache-export-aliasing

Author: zingo

.ci/docker/build.sh

@@ -116,6 +116,10 @@ fi
 # Copy requirements-lintrunner.txt from root to here
 cp ../../requirements-lintrunner.txt ./
 
+# Copy the Arm setup scripts that pin the FVP and toolchain downloads
+mkdir -p arm-setup
+cp ../../backends/arm/scripts/{utils,fvp_utils,toolchain_utils}.sh arm-setup/
+
 # OSDC runners have no docker daemon, so the build runs on the in-cluster
 # BuildKit pool via a remote buildx builder. That builder has nowhere to load
 # an image into, so the result has to go straight to the registry.

.ci/docker/common/install_arm_downloads.sh

@@ -0,0 +1,34 @@
+#!/bin/bash
+# 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.
+
+# Cache the developer.arm.com downloads of examples/arm/setup.sh, so CI jobs
+# copy them from ARM_SETUP_DOWNLOAD_CACHE instead of fetching them each run.
+# Jobs download whatever is missing, so a failed download only costs the cache.
+
+set -exu
+
+ARCH="$(uname -m)"
+OS="$(uname -s)"
+ARM_SETUP_CURL_PROGRESS_ARGS=(--no-progress-meter)
+# Fill the cache rather than read from it
+cache_dir="${ARM_SETUP_DOWNLOAD_CACHE}"
+unset ARM_SETUP_DOWNLOAD_CACHE
+source ./arm-setup/fvp_utils.sh
+source ./arm-setup/toolchain_utils.sh
+
+cache_download() {
+    local url=$1
+    local md5=$2
+    local output="${cache_dir}/$(basename "${url}")"
+    download_with_retry "cache" "${url}" "${output}" "${md5}" || rm -f "${output}"
+}
+
+mkdir -p "${cache_dir}"
+cache_download "${corstone300_url}" "${corstone300_md5_checksum}"
+cache_download "${corstone320_url}" "${corstone320_md5_checksum}"
+gcc_select_toolchain
+cache_download "${toolchain_url}" "${toolchain_md5_checksum}"

.ci/docker/requirements-ci.txt

@@ -9,7 +9,7 @@ torchsr==1.0.4
 transformers==5.3.0
 zstd==1.5.5.1
 pandas>=2.2.2; python_version >= '3.10'
-pytest==7.2.0
+pytest==7.3.0
 pytest-cov==4.1.0
 expecttest==0.1.6
 hypothesis==6.84.2

.ci/docker/ubuntu/Dockerfile

@@ -103,6 +103,16 @@ COPY ./common/utils.sh utils.sh
 RUN if [ -n "${ZEPHYR_SDK}" ]; then bash ./install_zephyr.sh; fi
 RUN rm install_zephyr.sh zephyr_sdk_release_proxy.py utils.sh
 
+# Cache the developer.arm.com downloads of examples/arm/setup.sh, which the
+# Arm and Zephyr jobs run.
+ENV ARM_SETUP_DOWNLOAD_CACHE=/opt/arm-download-cache
+COPY ./common/install_arm_downloads.sh install_arm_downloads.sh
+COPY ./arm-setup arm-setup
+RUN if [ -n "${ARM_SDK}${ZEPHYR_SDK}" ]; then \
+    bash ./install_arm_downloads.sh || echo "Arm download cache not populated"; \
+    fi
+RUN rm -r install_arm_downloads.sh arm-setup
+
 ARG QNN_SDK
 
 ARG MEDIATEK_SDK

.ci/scripts/download_hf_hub.sh

@@ -7,13 +7,15 @@ export HF_HUB_DISABLE_XET=1
 # Arguments:
 # 1. model_id: The Hugging Face repository ID (e.g., "organization/model_name")
 # 2. subdir: The optional subdirectory in the repo to look for files (pass "" if not used)
-# 3. file_names: A space-separated list of filenames to be downloaded
+# 3. revision: The commit hash to download
+# 4. file_names: A space-separated list of filenames to be downloaded
 # Returns:
 # The directory containing the downloaded files
 function download_hf_files() {
   local model_id="$1"
   local subdir="$2"
-  shift 2
+  local revision="$3"
+  shift 3
   local file_names=("$@")  # Capture all remaining arguments as an array
 
   local download_dir
@@ -24,7 +26,8 @@ from huggingface_hub import hf_hub_download
 # Download the first file and get its directory
 path = hf_hub_download(
     repo_id='${model_id}',
-    filename='${subdir:+${subdir}/}${file_names[0]}'
+    filename='${subdir:+${subdir}/}${file_names[0]}',
+    revision='${revision}'
 )
 import os
 print(os.path.dirname(path))")
@@ -41,7 +44,8 @@ from huggingface_hub import hf_hub_download
 # Download the file
 hf_hub_download(
     repo_id='${model_id}',
-    filename='${subdir:+${subdir}/}${file_name}'
+    filename='${subdir:+${subdir}/}${file_name}',
+    revision='${revision}'
 )"
 
     if [ $? -ne 0 ]; then
@@ -67,6 +71,10 @@ if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
         SUBDIR="$2"
         shift 2
         ;;
+      --revision)
+        REVISION="$2"
+        shift 2
+        ;;
       --files)
         shift
         FILES_TO_DOWNLOAD=()
@@ -83,13 +91,13 @@ if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
   done
 
   # Validate required arguments
-  if [ -z "$MODEL_ID" ] || [ ${#FILES_TO_DOWNLOAD[@]} -eq 0 ]; then
-    echo "Usage: $0 --model_id <model_id> --subdir <subdir> --files <file1> [<file2> ...]" >&2
+  if [ -z "$MODEL_ID" ] || [ -z "$REVISION" ] || [ ${#FILES_TO_DOWNLOAD[@]} -eq 0 ]; then
+    echo "Usage: $0 --model_id <model_id> --revision <commit> --subdir <subdir> --files <file1> [<file2> ...]" >&2
     exit 1
   fi
 
   # Call the function
-  DOWNLOAD_DIR=$(download_hf_files "$MODEL_ID" "$SUBDIR" "${FILES_TO_DOWNLOAD[@]}")
+  DOWNLOAD_DIR=$(download_hf_files "$MODEL_ID" "$SUBDIR" "$REVISION" "${FILES_TO_DOWNLOAD[@]}")
   if [ $? -eq 0 ]; then
     echo "$DOWNLOAD_DIR"
   else

.ci/scripts/export_model_artifact.sh

@@ -50,7 +50,7 @@ Arguments:
                  - vr-streaming: Voxtral Realtime streaming mode
                  - vr-offline: Voxtral Realtime offline mode
                  - solo-text: Muse Glimmer solo text mode
-                 - solo-text-batching: Muse Glimmer text on the CUDA batching runtime
+                 - dflash-image-batching: Muse Glimmer DFlash text and image on the CUDA batching runtime
                  - dflash-image: Muse Glimmer DFlash vision mode
 
 Examples:
@@ -105,7 +105,7 @@ if [ -n "$MODE" ]; then
         exit 1
       fi
       ;;
-    solo-text|solo-text-batching|dflash-image)
+    solo-text|dflash-image-batching|dflash-image)
       if [ "$HF_MODEL" != "meta-models/Muse-Glimmer-30B-GGUF" ]; then
         echo "Error: Mode '$MODE' can only be used with Muse Glimmer model"
         echo "Provided model: $HF_MODEL"
@@ -114,7 +114,7 @@ if [ -n "$MODE" ]; then
       ;;
     *)
       echo "Error: Unsupported mode '$MODE'"
-      echo "Supported modes: vr-streaming, vr-offline, solo-text, solo-text-batching, dflash-image"
+      echo "Supported modes: vr-streaming, vr-offline, solo-text, dflash-image-batching, dflash-image"
       exit 1
       ;;
   esac
@@ -354,8 +354,8 @@ if [ "$MODEL_NAME" = "muse_glimmer" ]; then
     echo "Error: Muse Glimmer requires quantization 'kquant-17gb' or 'kquant-dynamic'"
     exit 1
   fi
-  if [ "$MODE" != "solo-text" ] && [ "$MODE" != "solo-text-batching" ] && [ "$MODE" != "dflash-image" ]; then
-    echo "Error: Muse Glimmer requires mode 'solo-text', 'solo-text-batching' or 'dflash-image'"
+  if [ "$MODE" != "solo-text" ] && [ "$MODE" != "dflash-image-batching" ] && [ "$MODE" != "dflash-image" ]; then
+    echo "Error: Muse Glimmer requires mode 'solo-text', 'dflash-image-batching' or 'dflash-image'"
     exit 1
   fi
 fi
@@ -707,23 +707,17 @@ if [ "$MODEL_NAME" = "muse_glimmer" ]; then
           --backend cuda \
           --output-dir "${OUTPUT_DIR}"
       ;;
-    solo-text-batching)
-      # The off-graph cache exports the batching methods CudaExecutor drives:
-      # Static forward_{N}, dynamic forward_others and the device samplers.
-      EXPORT_START_SECONDS=$SECONDS
-      TMPDIR="$INDUCTOR_TMPDIR" \
-      TORCHINDUCTOR_CACHE_DIR="$INDUCTOR_CACHE" \
-      python -m executorch.examples.models.muse_glimmer.export.export_solo \
-          --gguf "$TARGET_GGUF_PATH" \
-          --backend cuda \
-          --use-offgraph-kv-cache \
-          --output-dir "${OUTPUT_DIR}"
-      ;;
-    dflash-image)
+    dflash-image|dflash-image-batching)
       DRAFT_GGUF_FILE="dflash-Muse-Glimmer-30B-Q4_K_M.gguf"
       MMPROJ_GGUF_FILE="mmproj-Muse-Glimmer-30B-Q4_K_M.gguf"
       python -c "from huggingface_hub import hf_hub_download; hf_hub_download('${HF_MODEL}', '${DRAFT_GGUF_FILE}', local_dir='${LOCAL_MODEL_DIR}')"
       python -c "from huggingface_hub import hf_hub_download; hf_hub_download('${HF_MODEL}', '${MMPROJ_GGUF_FILE}', local_dir='${LOCAL_MODEL_DIR}')"
+      # The off-graph cache exports the batching methods the CUDA DFlash
+      # executor drives, plus the vision encoder and forward_from_embeddings.
+      OFFGRAPH_FLAG=""
+      if [ "$MODE" = "dflash-image-batching" ]; then
+        OFFGRAPH_FLAG="--use-offgraph-kv-cache"
+      fi
       EXPORT_START_SECONDS=$SECONDS
       TMPDIR="$INDUCTOR_TMPDIR" \
       TORCHINDUCTOR_CACHE_DIR="$INDUCTOR_CACHE" \
@@ -732,10 +726,11 @@ if [ "$MODEL_NAME" = "muse_glimmer" ]; then
           --draft-gguf "${LOCAL_MODEL_DIR}/${DRAFT_GGUF_FILE}" \
           --mmproj "${LOCAL_MODEL_DIR}/${MMPROJ_GGUF_FILE}" \
           --backend cuda \
+          $OFFGRAPH_FLAG \
           --output-dir "${OUTPUT_DIR}"
       ;;
     *)
-      echo "Error: Muse Glimmer requires mode 'solo-text', 'solo-text-batching' or 'dflash-image'"
+      echo "Error: Muse Glimmer requires mode 'solo-text', 'dflash-image-batching' or 'dflash-image'"
       exit 1
       ;;
   esac
@@ -746,7 +741,7 @@ if [ "$MODEL_NAME" = "muse_glimmer" ]; then
 
   test -f "${OUTPUT_DIR}/model.pte"
   test -f "${OUTPUT_DIR}/aoti_cuda_blob.ptd"
-  if [ "$MODE" = "dflash-image" ]; then
+  if [ "$MODE" = "dflash-image" ] || [ "$MODE" = "dflash-image-batching" ]; then
     test -f "${OUTPUT_DIR}/pos_embed.bin"
   fi
   ls -al "${OUTPUT_DIR}"

.ci/scripts/muse_glimmer_batching_smoke.cpp

@@ -20,6 +20,7 @@ DECLARE_string(pte);
 DECLARE_string(tokenizer);
 DECLARE_string(data_path);
 DECLARE_string(backend);
+DECLARE_bool(dflash);
 DECLARE_int32(max_sessions);
 DECLARE_int32(max_session_tokens);
 DECLARE_uint64(bos_id);
@@ -64,6 +65,7 @@ int main(int argc, char** argv) {
     config.max_sessions = FLAGS_max_sessions;
     config.max_session_tokens = FLAGS_max_session_tokens;
     config.bos_id = FLAGS_bos_id;
+    config.use_dflash = FLAGS_dflash;
     auto backend = llm::create_muse_glimmer_backend(config);
     if (!backend.ok()) {
       throw std::runtime_error(

.ci/scripts/test-native.sh

@@ -20,14 +20,18 @@ case "${component}" in
     test_regex='^native_runtime_test$'
     ;;
   cpu)
+    # The CPU plan test covers ETDump tracing, including operator shapes.
     cmake_args+=(
       -DEXECUTORCH_BUILD_NATIVE_CPU=ON
       -DEXECUTORCH_BUILD_XNNPACK=ON
       -DEXECUTORCH_BUILD_KERNELS_OPTIMIZED=ON
       -DEXECUTORCH_BUILD_NATIVE_VISION_RUNNER=ON
+      -DEXECUTORCH_BUILD_DEVTOOLS=ON
+      -DEXECUTORCH_ENABLE_EVENT_TRACER=ON
+      -DEXECUTORCH_CPU_XNNPACK_OPERATOR_SHAPES=ON
     )
-    targets+=(native_cpu_engine_test run_classification_cpu)
-    test_regex='^native_cpu_engine_test$'
+    targets+=(native_cpu_engine_test cpu_plan_test run_classification_cpu)
+    test_regex='^(native_cpu_engine|cpu_plan)_test$'
     ;;
   ggml)
     cmake_args+=(
@@ -36,8 +40,9 @@ case "${component}" in
       -DEXECUTORCH_BUILD_NATIVE_VISION_RUNNER=ON
       -DGGML_NATIVE=OFF
     )
-    targets+=(run_hf_llm_ggml native_ggml_engine_test native_llm_runner_utils_test run_classification_ggml_cpu)
-    test_regex='^native_(ggml_engine|llm_runner_utils)_test$'
+    # TODO: run the other sample tests once they pass.
+    targets+=(run_hf_llm_ggml native_ggml_engine_test native_llm_runner_utils_test run_classification_ggml_cpu native_output_check_test native_ggml_sample_test)
+    test_regex='^native_(ggml_engine|llm_runner_utils|output_check)_test$|^native_ggml_sample_test_act_quant$'
     ;;
   vulkan)
     cmake_args+=(
@@ -67,7 +72,7 @@ cmake -S . -B "${build_dir}" -G Ninja \
 cmake --build "${build_dir}" \
   --parallel "${build_jobs}" \
   --target "${targets[@]}"
-ctest --test-dir "${build_dir}/backends/native/test" \
+ctest --test-dir "${build_dir}" \
   -R "${test_regex}" --output-on-failure --no-tests=error --timeout 120
 
 if [[ "${component}" == "ggml" ]]; then

.ci/scripts/test_backend.sh

@@ -76,11 +76,13 @@ if [[ "$FLOW" == *arm* ]]; then
         PYTEST_RETRY_ARGS=(--reruns 2 --reruns-delay 1)
     fi
 
-    # Setup ARM deps.
+    # Setup ARM deps. Only the Ethos-U flows run on an FVP.
     if [[ "$FLOW" == *vgf* ]]; then
-        .ci/scripts/setup-arm-baremetal-tools.sh --enable-mlsdk-deps --install-mlsdk-deps-with-pip
-    else
+        .ci/scripts/setup-arm-baremetal-tools.sh --disable-ethos-u-deps --enable-mlsdk-deps --install-mlsdk-deps-with-pip
+    elif [[ "$FLOW" == *ethos_u* ]]; then
         .ci/scripts/setup-arm-baremetal-tools.sh
+    else
+        .ci/scripts/setup-arm-baremetal-tools.sh --disable-ethos-u-deps
     fi
     source examples/arm/arm-scratch/setup_path.sh
 
@@ -98,7 +100,7 @@ fi
 if [[ "$FLOW" == *cortex_m* ]]; then
     # Cortex-M runs on the Corstone-300 FVP, using the same Arm toolchain as the
     # Ethos-U flows but its own semihosting runner.
-    .ci/scripts/setup-arm-baremetal-tools.sh
+    .ci/scripts/setup-arm-baremetal-tools.sh --fvp corstone300
     source examples/arm/arm-scratch/setup_path.sh
 
     backends/cortex_m/test/build_test_runner.sh

.ci/scripts/test_cuda_runtime_windows.ps1

@@ -34,7 +34,9 @@ $tests = @(
     "test_cuda_mutable_state",
     "test_cuda_weight_cache",
     "test_cuda_output_copy",
+    "test_cuda_backend_execute",
     "test_cuda_sort_rand",
+    "test_cuda_current_stream",
     "test_cuda_guard",
     "test_cuda_stream_guard"
 )