Update on "[ET-VK] double, short, and uint16 dtype runtime support"

Creating support for double, short, and uint16 for quantization ops. Registering the short keyword since theres already support.

Also changing the cpu implementation to support half

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

[ghstack-poisoned]

Author: morelos

.ci/scripts/build-mediatek-sdk.sh

@@ -0,0 +1,22 @@
+#!/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.
+
+set -eux
+
+build_neuron_backend() {
+  echo "Start building neuron backend."
+  export ANDROID_NDK=/opt/ndk
+  export MEDIATEK_SDK_ROOT=/tmp/neuropilot
+  export NEURON_BUFFER_ALLOCATOR_LIB=${MEDIATEK_SDK_ROOT}/libneuron_buffer_allocator.so
+  export EXECUTORCH_ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." && pwd)"
+
+
+  cd ${EXECUTORCH_ROOT}
+  ./backends/mediatek/scripts/mtk_build.sh
+}
+
+build_neuron_backend

.ci/scripts/gather_benchmark_configs.py

@@ -32,7 +32,8 @@
 BENCHMARK_CONFIGS = {
     "xplat": [
         "xnnpack_q8",
-        "hf_xnnpack_fp32",
+        "hf_xnnpack_custom_spda_kv_cache_8da4w",
+        "et_xnnpack_custom_spda_kv_cache_8da4w",
         "llama3_fb16",
         "llama3_spinquant",
         "llama3_qlora",
@@ -129,25 +130,25 @@ def generate_compatible_configs(model_name: str, target_os=None) -> List[str]:
     """
     configs = []
     if is_valid_huggingface_model_id(model_name):
+        configs.append("hf_xnnpack_custom_spda_kv_cache_8da4w")
         if model_name.startswith("meta-llama/"):
-            # LLaMA models
+            # etLLM recipes for Llama
             repo_name = model_name.split("meta-llama/")[1]
             if "qlora" in repo_name.lower():
-                configs.append("llama3_qlora")
+                configs = ["llama3_qlora"]
             elif "spinquant" in repo_name.lower():
-                configs.append("llama3_spinquant")
+                configs = ["llama3_spinquant"]
             else:
-                configs.append("llama3_fb16")
+                configs.extend(["llama3_fb16", "et_xnnpack_custom_spda_kv_cache_8da4w"])
                 configs.extend(
                     [
                         config
                         for config in BENCHMARK_CONFIGS.get(target_os, [])
                         if config.startswith("llama")
                     ]
                 )
-        else:
-            # Non-LLaMA models
-            configs.append("hf_xnnpack_fp32")
+        if model_name.startswith("Qwen/Qwen3"):
+            configs.append("et_xnnpack_custom_spda_kv_cache_8da4w")
     elif model_name in MODEL_NAME_TO_MODEL:
         # ExecuTorch in-tree non-GenAI models
         configs.append("xnnpack_q8")

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

@@ -0,0 +1,42 @@
+#!/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.
+
+set -eux
+
+MEDIATEK_INSTALLATION_DIR=/tmp/neuropilot
+EXECUTORCH_ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." && pwd)"
+
+install_neuropilot() {
+  echo "Start installing neuropilot."
+  mkdir -p "${MEDIATEK_INSTALLATION_DIR}"
+
+  curl -Lo /tmp/neuropilot-express.tar.gz "https://s3.ap-southeast-1.amazonaws.com/mediatek.neuropilot.com/06302508-4c94-4bf2-9789-b0ee44e83e27.gz"
+  echo "Finishing downloading neuropilot sdk."
+  tar zxvf /tmp/neuropilot-express.tar.gz --strip-components=1 --directory "${MEDIATEK_INSTALLATION_DIR}"
+  echo "Finishing unzip neuropilot sdk."
+
+  # Copy NP header
+  cp ${MEDIATEK_INSTALLATION_DIR}/api/NeuronAdapter.h ${EXECUTORCH_ROOT}/backends/mediatek/runtime/include/api/
+
+  # Print the content for manual verification
+  ls -lah "${MEDIATEK_INSTALLATION_DIR}"
+}
+
+setup_neuropilot() {
+  pip3 install -r ${EXECUTORCH_ROOT}/backends/mediatek/requirements.txt
+  pip3 install ${MEDIATEK_INSTALLATION_DIR}/mtk_neuron-8.2.19-py3-none-linux_x86_64.whl
+  pip3 install ${MEDIATEK_INSTALLATION_DIR}/mtk_converter-8.13.0_public_packages/mtk_converter-8.13.0+public-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
+}
+
+setup_calibration_data() {
+  curl -Lo /tmp/imagenette2-160.tgz https://s3.amazonaws.com/fast-ai-imageclas/imagenette2-160.tgz
+  tar zxvf /tmp/imagenette2-160.tgz --strip-components=1 --directory "${MEDIATEK_INSTALLATION_DIR}"
+}
+
+install_neuropilot
+setup_neuropilot
+setup_calibration_data

.ci/scripts/test_llava.sh

@@ -147,7 +147,7 @@ run_and_verify() {
 
     # verify result.txt
     RESULT=$(cat result.txt)
-    EXPECTED_PREFIX="ASSISTANT: image captures a basketball game in progress, with several players on the court. "
+    EXPECTED_PREFIX="ASSISTANT: image captures a basketball game in progress, with"
 
     if [[ "${RESULT}" == *"${EXPECTED_PREFIX}"* ]]; then
         echo "Expected result prefix: ${EXPECTED_PREFIX}"

.ci/scripts/test_model.sh

@@ -244,6 +244,24 @@ test_model_with_mps() {
   EXPORTED_MODEL=$(find "." -type f -name "${MODEL_NAME}*.pte" -print -quit)
 }
 
+test_model_with_mediatek() {
+  if [[ "${MODEL_NAME}" == "dl3" ]]; then
+    EXPORT_SCRIPT=deeplab_v3
+  elif [[ "${MODEL_NAME}" == "mv3" ]]; then
+    EXPORT_SCRIPT=mobilenet_v3
+  elif [[ "${MODEL_NAME}" == "mv2" ]]; then
+    EXPORT_SCRIPT=mobilenet_v2
+  elif [[ "${MODEL_NAME}" == "ic4" ]]; then
+    EXPORT_SCRIPT=inception_v4
+  elif [[ "${MODEL_NAME}" == "ic3" ]]; then
+    EXPORT_SCRIPT=inception_v3
+  fi
+
+  PYTHONPATH=examples/mediatek/ "${PYTHON_EXECUTABLE}" -m examples.mediatek.model_export_scripts.${EXPORT_SCRIPT} -d /tmp/neuropilot/train -a ${EXPORT_SCRIPT}
+  EXPORTED_MODEL=$(find "./${EXPORT_SCRIPT}" -type f -name "*.pte" -print -quit)
+}
+
+
 if [[ "${BACKEND}" == "portable" ]]; then
   echo "Testing ${MODEL_NAME} with portable kernels..."
   test_model
@@ -281,6 +299,12 @@ elif [[ "${BACKEND}" == *"xnnpack"* ]]; then
   if [[ $? -eq 0 ]]; then
     prepare_artifacts_upload
   fi
+elif [[ "${BACKEND}" == "mediatek" ]]; then
+  echo "Testing ${MODEL_NAME} with mediatek..."
+  test_model_with_mediatek
+  if [[ $? -eq 0 ]]; then
+    prepare_artifacts_upload
+  fi
 else
   set +e
   if [[ "${BACKEND}" == *"quantization"* ]]; then

.ci/scripts/tests/test_gather_benchmark_configs.py

@@ -112,15 +112,24 @@ def test_generate_compatible_configs_llama_model(self):
         result = self.gather_benchmark_configs.generate_compatible_configs(
             model_name, target_os
         )
-        expected = ["llama3_fb16", "llama3_coreml_ane"]
-        self.assertEqual(result, expected)
+        expected = [
+            "llama3_fb16",
+            "llama3_coreml_ane",
+            "et_xnnpack_custom_spda_kv_cache_8da4w",
+            "hf_xnnpack_custom_spda_kv_cache_8da4w",
+        ]
+        self.assertCountEqual(result, expected)
 
         target_os = "android"
         result = self.gather_benchmark_configs.generate_compatible_configs(
             model_name, target_os
         )
-        expected = ["llama3_fb16"]
-        self.assertEqual(result, expected)
+        expected = [
+            "llama3_fb16",
+            "et_xnnpack_custom_spda_kv_cache_8da4w",
+            "hf_xnnpack_custom_spda_kv_cache_8da4w",
+        ]
+        self.assertCountEqual(result, expected)
 
     def test_generate_compatible_configs_quantized_llama_model(self):
         model_name = "meta-llama/Llama-3.2-1B-Instruct-SpinQuant_INT4_EO8"

.ci/scripts/unittest-buck2.sh

@@ -15,8 +15,10 @@ buck2 query "//backends/apple/... + //backends/example/... + \
 //kernels/optimized/... + //kernels/portable/... + //kernels/quantized/... + \
 //kernels/test/... + //runtime/... + //schema/... + //test/... + //util/..."
 
+# TODO: optimized ops are unbuildable because they now use ATen; put
+# them back after we can use PyTorch in OSS buck.
 UNBUILDABLE_OPTIMIZED_OPS_REGEX="_elu|gelu|fft|log_softmax"
-BUILDABLE_OPTIMIZED_OPS=$(buck2 query //kernels/optimized/cpu/... | grep -E -v $UNBUILDABLE_OPTIMIZED_OPS_REGEX)
+BUILDABLE_OPTIMIZED_OPS= #$(buck2 query //kernels/optimized/cpu/... | grep -E -v $UNBUILDABLE_OPTIMIZED_OPS_REGEX)
 
 # TODO: build prim_ops_test_cpp again once supported_features works in
 # OSS buck.

.github/workflows/android-perf-private-device-experiment.yml

@@ -18,7 +18,7 @@ on:
         description: Models to be benchmarked
         required: false
         type: string
-        default: mv3,meta-llama/Llama-3.2-1B-Instruct-SpinQuant_INT4_EO8,meta-llama/Llama-3.2-1B-Instruct-QLORA_INT4_EO8
+        default: Qwen/Qwen3-0.6B
       devices:
         description: Target devices to run benchmark
         required: false
@@ -34,7 +34,7 @@ on:
         description: Models to be benchmarked
         required: false
         type: string
-        default: mv3,meta-llama/Llama-3.2-1B-Instruct-SpinQuant_INT4_EO8,meta-llama/Llama-3.2-1B-Instruct-QLORA_INT4_EO8
+        default: Qwen/Qwen3-0.6B
       devices:
         description: Target devices to run benchmark
         required: false
@@ -57,6 +57,6 @@ jobs:
       id-token: write
       contents: read
     with:
-      models: ${{ inputs.models || 'mv3,meta-llama/Llama-3.2-1B-Instruct-SpinQuant_INT4_EO8,meta-llama/Llama-3.2-1B-Instruct-QLORA_INT4_EO8' }}
+      models: ${{ inputs.models || github.event_name == 'schedule' && 'Qwen/Qwen3-0.6B,HuggingFaceTB/SmolLM2-135M,meta-llama/Llama-3.2-1B,allenai/OLMo-1B-hf' || 'Qwen/Qwen3-0.6B' }}
       devices: samsung_galaxy_s22_private
       benchmark_configs: ${{ inputs.benchmark_configs }}

.github/workflows/android-perf.yml

@@ -6,12 +6,14 @@ on:
   pull_request:
     paths:
       - .github/workflows/android-perf.yml
+      - .ci/scripts/gather_benchmark_configs.py
       - extension/benchmark/android/benchmark/android-llm-device-farm-test-spec.yml.j2
   push:
     branches:
       - main
     paths:
       - .github/workflows/android-perf.yml
+      - .ci/scripts/gather_benchmark_configs.py
       - extension/benchmark/android/benchmark/android-llm-device-farm-test-spec.yml.j2
   # Note: GitHub has an upper limit of 10 inputs
   workflow_dispatch:
@@ -20,7 +22,7 @@ on:
         description: Models to be benchmarked
         required: false
         type: string
-        default: llama
+        default: Qwen/Qwen3-0.6B
       devices:
         description: Target devices to run benchmark
         required: false
@@ -36,7 +38,7 @@ on:
         description: Models to be benchmarked
         required: false
         type: string
-        default: llama
+        default: Qwen/Qwen3-0.6B
       devices:
         description: Target devices to run benchmark
         required: false
@@ -70,7 +72,7 @@ jobs:
           # Separate default values from the workflow dispatch. To ensure defaults are accessible
           # during scheduled runs and to provide flexibility for different defaults between
           # on-demand and periodic benchmarking.
-          CRON_DEFAULT_MODELS: ${{ github.event_name == 'schedule' && 'llama,mv3,mv2,ic4,ic3,resnet50,edsr,mobilebert,w2l,meta-llama/Llama-3.2-1B,meta-llama/Llama-3.2-1B-Instruct-SpinQuant_INT4_EO8,meta-llama/Llama-3.2-1B-Instruct-QLORA_INT4_EO8' || 'llama' }}
+          CRON_DEFAULT_MODELS: ${{ github.event_name == 'schedule' && 'mv3,mv2,ic4,ic3,resnet50,edsr,mobilebert,w2l,meta-llama/Llama-3.2-1B,meta-llama/Llama-3.2-1B-Instruct-SpinQuant_INT4_EO8,meta-llama/Llama-3.2-1B-Instruct-QLORA_INT4_EO8,Qwen/Qwen3-0.6B,HuggingFaceTB/SmolLM2-135M,allenai/OLMo-1B-hf' || 'Qwen/Qwen3-0.6B' }}
           CRON_DEFAULT_DEVICES: samsung_galaxy_s22
         run: |
           set -eux
@@ -201,8 +203,8 @@ jobs:
             HF_MODEL_REPO=${{ matrix.model }}
             OUT_ET_MODEL_NAME="$(echo "$HF_MODEL_REPO" | awk -F'/' '{print $2}' | sed 's/_/-/g' | tr '[:upper:]' '[:lower:]')_${{ matrix.config }}"
 
+            # Convert HF checkpoint to ET via etLLM path
             if [[ "$HF_MODEL_REPO" == meta-llama/* ]]; then
-                # Llama models on Hugging Face
                 if [[ ${{ matrix.config }} == "llama3_spinquant" ]]; then
                     # SpinQuant
                     # Download prequantized chceckpoint from Hugging Face
@@ -272,6 +274,21 @@ jobs:
                       --metadata '{"get_bos_id":128000, "get_eos_ids":[128009, 128001]}' \
                       --output_name="${OUT_ET_MODEL_NAME}.pte"
                     ls -lh "${OUT_ET_MODEL_NAME}.pte"
+                elif [[ ${{ matrix.config }} == "et_xnnpack_custom_spda_kv_cache_8da4w" ]]; then
+                    DOWNLOADED_PATH=$(bash .ci/scripts/download_hf_hub.sh --model_id "${HF_MODEL_REPO}" --subdir "original" --files "tokenizer.model" "params.json" "consolidated.00.pth")
+                    python -m examples.models.llama.export_llama \
+                      --model llama3_2 \
+                      --checkpoint "${DOWNLOADED_PATH}/consolidated.00.pth" \
+                      --params "${DOWNLOADED_PATH}/params.json" \
+                      -kv \
+                      --use_sdpa_with_kv_cache \
+                      -d fp32 \
+                      -X \
+                      --xnnpack-extended-ops \
+                      -qmode 8da4w -G 32 -E 8,0 \
+                      --metadata '{"get_bos_id":128000, "get_eos_ids":[128009, 128001]}' \
+                      --output_name="${OUT_ET_MODEL_NAME}.pte"
+                    ls -lh "${OUT_ET_MODEL_NAME}.pte"
                 elif [[ ${{ matrix.config }} == "llama3_qnn_htp" ]]; then
                     export QNN_SDK_ROOT=/tmp/qnn/2.28.0.241029
                     export LD_LIBRARY_PATH=$QNN_SDK_ROOT/lib/x86_64-linux-clang/
@@ -292,21 +309,75 @@ jobs:
                     OUT_ET_MODEL_NAME="llama3_2_qnn" # Qualcomm hard-coded it in their script
                     find . -name "${OUT_ET_MODEL_NAME}.pte" -not -path "./${OUT_ET_MODEL_NAME}.pte" -exec mv {} ./ \;
                     ls -lh "${OUT_ET_MODEL_NAME}.pte"
-                else
-                    # By default, test with the Hugging Face model and the xnnpack recipe
-                    DOWNLOADED_PATH=$(bash .ci/scripts/download_hf_hub.sh --model_id "${HF_MODEL_REPO}" --subdir "original" --files "tokenizer.model")
-                    python -m extension.export_util.export_hf_model -hfm="$HF_MODEL_REPO" -o "$OUT_ET_MODEL_NAME"
-                    ls -lh "${OUT_ET_MODEL_NAME}.pte"
                 fi
-            else
-                echo "Unsupported model ${{ matrix.model }}"
-                exit 1
+            elif [[ "$HF_MODEL_REPO" == "Qwen/Qwen3-0.6B" ]]; then
+              if [[ ${{ matrix.config }} == "et_xnnpack_custom_spda_kv_cache_8da4w" ]]; then
+                DOWNLOADED_PATH=$(bash .ci/scripts/download_hf_hub.sh --model_id "${HF_MODEL_REPO}" --subdir "." --files "tokenizer.json")
+                python -m examples.models.llama.export_llama \
+                  --model qwen3-0_6b \
+                  --params examples/models/qwen3/0_6b_config.json \
+                  -kv \
+                  --use_sdpa_with_kv_cache \
+                  -d fp32 \
+                  -X \
+                  --xnnpack-extended-ops \
+                  -qmode 8da4w \
+                  -G 32 \
+                  -E 8,0 \
+                  --metadata '{"get_bos_id": 151644, "get_eos_ids":[151645]}' \
+                  --output_name="${OUT_ET_MODEL_NAME}.pte"
+                ls -lh "${OUT_ET_MODEL_NAME}.pte"
+              fi
+            fi
+
+            if [[ ${{ matrix.config }} == "hf_xnnpack_custom_spda_kv_cache_8da4w" ]]; then
+              DOWNLOADED_PATH=$(
+                bash .ci/scripts/download_hf_hub.sh \
+                  --model_id "${HF_MODEL_REPO}" \
+                  --files "tokenizer.json"
+              )
+              echo "tokenizer.json is downloaded to $DOWNLOADED_PATH"
+
+              # Install optimum-executorch
+              git clone https://github.com/huggingface/optimum-executorch
+              pushd optimum-executorch
+              # There is no release yet, for CI stability, always test from the same commit on main
+              git checkout 1c653dc49812fc431a22312c7295d97005d22e12
+              python install_dev.py
+              pip list
+
+              ARGS=(
+                "--model" "${HF_MODEL_REPO}"
+                "--task" "text-generation"
+                "--recipe" "xnnpack"
+                "--use_custom_sdpa"
+                "--qlinear"
+                "--qembedding"
+                "--output_dir" ".."
+              )
+
+              # Add conditional arguments based on model
+              case "${HF_MODEL_REPO}" in
+                *"google/gemma-3-1b-it"*)
+                  echo "--use_custom_kv_cache can not be used for HybridCache"
+                  ;;
+                *)
+                  ARGS+=("--use_custom_kv_cache")
+                  ;;
+              esac
+
+              optimum-cli export executorch "${ARGS[@]}"
+              popd
+
+              mv model.pte ${OUT_ET_MODEL_NAME}.pte
+              ls -lh "${OUT_ET_MODEL_NAME}.pte"
             fi
 
-            zip -j model.zip "${OUT_ET_MODEL_NAME}.pte" "${DOWNLOADED_PATH}/tokenizer.model"
+            zip -j model.zip ${OUT_ET_MODEL_NAME}.pte ${DOWNLOADED_PATH}/tokenizer.*
             ls -lh model.zip
-            mkdir -p "${ARTIFACTS_DIR_NAME}"
-            mv model.zip "${ARTIFACTS_DIR_NAME}"
+            mkdir -p ${ARTIFACTS_DIR_NAME}
+            mv model.zip ${ARTIFACTS_DIR_NAME}
+            ls -lh ${ARTIFACTS_DIR_NAME}
         elif [[ ${{ matrix.model }} == "llama" ]]; then
             # Install requirements for export_llama
             PYTHON_EXECUTABLE=python bash examples/models/llama/install_requirements.sh