Update on "[ET-VK][Ops] torchao.dequantize_affine vulkan impl and shader and cleanup"

# Changes
* Implement `torchao.dequantize_affine` operator in Vulkan backend with comprehensive texture and buffer storage support
* Add block-wise dequantization mode in `dequantize_texture.glsl` and `dequantize_buffer.glsl` shaders for configurable tensor block dequantization
* Extend dequantization infrastructure in `Dequantize.cpp` to handle affine transformations with configurable block sizes and quantization parameters
* Support integer-to-floating-point conversion with precise reconstruction of original values

BE: Improved the documentation in the shader logic which is more detailed and clear


# Motivation
The existing Vulkan quantization infrastructure lacked support for the `torchao.dequantize_affine` operator, which is essential for completing the quantization-dequantization cycle in dynamic quantization workflows. The `dequantize_affine` operator provides flexible block-wise dequantization that reconstructs floating-point values from quantized integer blocks, enabling:

* **Block-wise Dequantization**: Reconstructs floating-point values from configurable tensor blocks using separate scale and zero-point parameters, enabling precise recovery of original data distributions
* **Affine Transformation**: Uses the formula `value = (qvalue - zero_point) * scale` for accurate integer-to-floating-point mapping
* **TorchAO Integration**: Seamless compatibility with TorchAO quantization workflows and completes the quantization-dequantization round-trip

# Operator Description
The `dequantize_affine` operator converts n-bit integer tensor values back to floating-point representations using pre-computed quantization parameters (scale and zero_point) applied to configurable tensor blocks. Block-wise dequantization divides tensors into blocks and applies separate dequantization parameters to each block, allowing fine-grained reconstruction of the original floating-point precision.

The dequantization formula is: `value = (qvalue - zero_point) * scale`

**Storage Requirements**: Scale and zero_point tensors must use buffer storage with width-packed layout. Input/output tensors support both buffer and texture storage with standard axis mapping. Input tensors must be integer types (kByte, kChar, kInt).

# Block-wise Dequantization Implementation
Block-wise dequantization enables fine-grained reconstruction by dividing tensors into blocks and applying separate dequantization parameters to each block. The implementation uses the same key data structures computed in `Dequantize.cpp`:

* **`block_size_vec`**: WHCN-ordered block dimensions converted from PyTorch NCHW layout (e.g., [3,3,2,1] for 3×3×2×1 blocks)
* **`tensor_size_whcn`**: Input tensor dimensions converted to WHCN layout using `utils::make_whcn_ivec4()`
* **`num_blocks_vec`**: Number of blocks per dimension calculated as `tensor_size_whcn / block_size_vec`
* **`block_stride_vec`**: Pre-computed linear strides for block grid indexing `{1, #W, #W*#H, #W*#H*#C}` to enable efficient block ID calculation

The block coordinate calculation uses: `bcoord = tidx / blockSize` where `tidx` is the tensor coordinate in WHCN layout, then the linear block ID is computed as: `block_id = bcoord.x * blockStride.x + bcoord.y * blockStride.y + bcoord.z * blockStride.z + bcoord.w * blockStride.w`

# Shader Algorithm Overview

## Texture Storage Implementation (`dequantize_texture.glsl`)

**Workgroup Configuration**:
- **Global WG Size**: Default sizing based on texture dimensions
- **Local WG Size**: Default with special handling for batch dimension dequantization (Z dimension set to 1 for proper workgroup dispatching when `global_workgroup_size[2] > 1`)

**Block-wise Mode Algorithm**:
The shader processes 3D texture positions where each position represents a texel containing 4 width-packed integer components. For each texel at position `pos`, it calculates a base tensor index `base_tidx = ivec4(pos.x * 4, pos.y, pos.z, 0)` to account for width-packing.

For each of the 4 components in the texel, it computes the actual tensor coordinate: `tidx = ivec4(base_tidx.x + i, base_tidx.y, (foldedZ % C_total), (foldedZ / C_total))` where `foldedZ = pos.z` handles batch-channel folding in 4D tensors and `C_total = numBlocks.z * blockSize.z` represents the total channel dimension.

The block coordinate is calculated using integer division: `bcoord = tidx / blockSize`, then the linear block ID uses pre-computed strides: `block_id = bcoord.x * blockStride.x + bcoord.y * blockStride.y + bcoord.z * blockStride.z + bcoord.w * blockStride.w`.

Each integer component is dequantized using its corresponding block's parameters: `value = dequantize_val(qvalue, t_scale[block_id], t_zero_point[block_id])` where `dequantize_val()` applies the formula `(qvalue - zero_point) * scale`. The reconstructed floating-point values are written to the output texel with proper type handling for double precision outputs.

## Buffer Storage Implementation (`dequantize_buffer.glsl`)

**Workgroup Configuration**:
- **Global WG Size**: Default sizing based on buffer element count
- **Local WG Size**: Default sizing without special constraints

**Block-wise Mode Algorithm**:
The shader processes linear buffer indices using `gl_GlobalInvocationID.x` as the output buffer index. It converts this to tensor coordinates using `bufi_to_tidx(out_bufi, t_out_strides, out_dim_order)` which handles the buffer-to-tensor index mapping with proper stride calculations.

For each element, it computes the block coordinate directly: `bcoord = out_tidx / blockSize` where `out_tidx` is the 4D tensor coordinate in WHCN layout. The linear block ID calculation uses the same pre-computed stride approach: `block_id = bcoord.x * blockStride.x + bcoord.y * blockStride.y + bcoord.z * blockStride.z + bcoord.w * blockStride.w`.

The quantized integer value is loaded using the corresponding input buffer index: `qvalue = t_in[in_bufi]` where `in_bufi = tidx_to_bufi(out_tidx, t_in_strides)`. Dequantization applies the block-specific parameters: `value = dequantize_val(qvalue, t_scale[block_id], t_zero_point[block_id])` to reconstruct the original floating-point value.

**Future Improvements**: Dynamic workgroup sizing based on block dimensions

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

cc SS-JIA manuelcandales cbilgin

[ghstack-poisoned]

Author: morelos

.Package.swift/executorch_llm/dummy.swift

@@ -0,0 +1 @@
+

.Package.swift/executorch_llm_debug/dummy.swift

@@ -0,0 +1 @@
+

.ci/docker/build.sh

@@ -7,9 +7,11 @@
 
 set -exu
 
-IMAGE_NAME="$1"
+FULL_IMAGE_NAME="$1"
 shift
 
+IMAGE_NAME=$(echo "${FULL_IMAGE_NAME}" | sed 's/ci-image://')
+
 echo "Building ${IMAGE_NAME} Docker image"
 
 OS=ubuntu
@@ -41,6 +43,10 @@ case "${IMAGE_NAME}" in
     ARM_SDK=yes
     CLANG_VERSION=12
     ;;
+  executorch-ubuntu-22.04-zephyr-sdk)
+    ZEPHYR_SDK=yes
+    GCC_VERSION=11
+    ;;
   executorch-ubuntu-22.04-qnn-sdk)
     QNN_SDK=yes
     CLANG_VERSION=12
@@ -85,6 +91,7 @@ docker build \
   --build-arg "LINTRUNNER=${LINTRUNNER:-}" \
   --build-arg "BUILD_DOCS=${BUILD_DOCS}" \
   --build-arg "ARM_SDK=${ARM_SDK:-}" \
+  --build-arg "ZEPHYR_SDK=${ZEPHYR_SDK:-}" \
   --build-arg "QNN_SDK=${QNN_SDK:-}" \
   --build-arg "MEDIATEK_SDK=${MEDIATEK_SDK:-}" \
   --build-arg "ANDROID_NDK_VERSION=${ANDROID_NDK_VERSION:-}" \

.ci/docker/ci_commit_pins/optimum-executorch.txt

@@ -1 +1 @@
-a3942627f5ac048e06b4b1d703b0a6a53bf6da5b
+eea657ddbdeb1118943a92fb73c289985c3ee1ba

.ci/docker/ci_commit_pins/pytorch.txt

@@ -1 +1 @@
-90f1e7bed15ca5e48c61c5b6dc5ad4810524f82f
+6fc0ad22f0a07b6f38d138861c56a765d5a9bb02

.ci/docker/common/install_zephyr.sh

@@ -0,0 +1,87 @@
+
+#!/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 -ex
+
+# shellcheck source=/dev/null
+source "$(dirname "${BASH_SOURCE[0]}")/utils.sh"
+
+# Double check if the NDK version is set
+[ -n "${ZEPHYR_SDK}" ]
+
+install_prerequiresites() {
+    rm /var/lib/dpkg/info/libc-bin.*
+    apt-get clean
+    apt-get -y update
+    apt-get install -y libc-bin
+    apt-get -y update
+    apt-get clean
+    apt-get install --no-install-recommends -y dos2unix
+    apt-get install --no-install-recommends -y ca-certificates
+    apt-get install -y --reinstall libc-bin
+    apt-get install --no-install-recommends -y file
+    apt-get install --no-install-recommends -y locales
+    apt-get install --no-install-recommends -y git
+    apt-get install --no-install-recommends -y build-essential
+    apt-get install --no-install-recommends -y cmake
+    apt-get install --no-install-recommends -y ninja-build gperf
+    apt-get install --no-install-recommends -y device-tree-compiler
+    apt-get install --no-install-recommends -y wget
+    apt-get install --no-install-recommends -y curl
+    apt-get install --no-install-recommends -y xz-utils
+    apt-get install --no-install-recommends -y dos2unix
+    apt-get install --no-install-recommends -y vim
+    apt-get install --no-install-recommends -y nano
+    apt-get install --no-install-recommends -y mc
+    apt-get install --no-install-recommends -y openssh-server
+    apt-get install -y gdb
+
+    # Zephyr SDK relies on python 3.12
+    apt install software-properties-common -y
+    add-apt-repository ppa:deadsnakes/ppa -y
+    apt update
+    apt install -y python3.12 python3.12-dev python3.12-venv python3-pip
+    update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.12 1
+
+    # Upgrade cmake ot 3.24
+    apt update
+    apt install cmake
+    apt install software-properties-common lsb-release
+    apt update
+    test -f /usr/share/doc/kitware-archive-keyring/copyright || \
+        wget -O - https://apt.kitware.com/keys/kitware-archive-latest.asc 2>/dev/null | gpg --dearmor - | tee /usr/share/keyrings/kitware-archive-keyring.gpg >/dev/null
+    "deb [signed-by=/usr/share/keyrings/kitware-archive-keyring.gpg] https://apt.kitware.com/ubuntu/ $(lsb_release -cs) main" | tee /etc/apt/sources.list.d/kitware.list > /dev/null
+    apt update
+    apt install cmake
+
+    # Install additional required software for Zephyr
+    apt install --no-install-recommends -y ccache \
+        dfu-util \
+        python3-setuptools \
+        python3-tk \
+        python3-wheel \
+        make \
+        gcc \
+        libsdl2-dev \
+        libmagic1 \
+        xterm \
+        telnet \
+        net-tools
+    apt install --no-install-recommends -y gcc-multilib g++-multilib
+    apt-get clean -y
+    apt-get autoremove --purge -y
+    rm -rf /var/lib/apt/lists/*
+    wget https://apt.kitware.com/kitware-archive.sh && \
+        chmod +x kitware-archive.sh && \
+        ./kitware-archive.sh && \
+        rm -f kitware-archive.sh
+    pip_install --no-cache-dir west
+    pip_install pyelftools
+}
+
+install_prerequiresites

.ci/docker/ubuntu/Dockerfile

@@ -84,6 +84,12 @@ RUN rm install_android.sh
 
 ARG ARM_SDK
 
+ARG ZEPHYR_SDK
+COPY ./common/install_zephyr.sh install_zephyr.sh
+COPY ./common/utils.sh utils.sh
+RUN if [ -n "${ZEPHYR_SDK}" ]; then bash ./install_zephyr.sh; fi
+RUN rm install_zephyr.sh utils.sh
+
 ARG QNN_SDK
 
 ARG MEDIATEK_SDK

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

@@ -33,6 +33,8 @@ set_up_aot() {
       -DEXECUTORCH_BUILD_DEVTOOLS=ON \
       -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \
       -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \
+      -DEXECUTORCH_BUILD_EXTENSION_EXTENSION_LLM=ON \
+      -DEXECUTORCH_BUILD_EXTENSION_EXTENSION_LLM_RUNNER=ON \
       -DEXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR=ON \
       -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \
       -DEXECUTORCH_ENABLE_EVENT_TRACER=ON \

.ci/scripts/build_llama_android.sh

@@ -19,7 +19,7 @@ install_executorch_and_backend_lib() {
   echo "Installing executorch and xnnpack backend"
   clean_executorch_install_folders
   mkdir cmake-android-out
-  ANDROID_NDK=/opt/ndk
+  ANDROID_NDK=${ANDROID_NDK:-/opt/ndk}
   BUCK2=buck2
   ANDROID_ABI=arm64-v8a
   cmake --preset llm \

.ci/scripts/gather_benchmark_configs.py

@@ -16,16 +16,65 @@
 from examples.models import MODEL_NAME_TO_MODEL
 
 
-# Device pools for AWS Device Farm
+DEVICE_POOLS_REGEX = re.compile(r"(?P<device_name>[^\+]+)\+(?P<variant>[^\+]+)")
+# Device pools for AWS Device Farm. Initially, I choose to distribute models to these pool
+# round-robin for simplicity. For public pool, only one per device type is needed because
+# AWS will scale the number of devices there for us. However, for private pool, we need to
+# manually maintain multiple pools of the same device to evenly distribute models there.
+# The pool ARNs are extracted from the output of the following command:
+#   aws devicefarm list-device-pools \
+#    --arn arn:aws:devicefarm:us-west-2:308535385114:project:02a2cf0f-6d9b-45ee-ba1a-a086587469e6 \
+#    --region us-west-2
 DEVICE_POOLS = {
-    "apple_iphone_15": "arn:aws:devicefarm:us-west-2:308535385114:devicepool:02a2cf0f-6d9b-45ee-ba1a-a086587469e6/3b5acd2e-92e2-4778-b651-7726bafe129d",
-    "apple_iphone_15+ios_18": "arn:aws:devicefarm:us-west-2:308535385114:devicepool:02a2cf0f-6d9b-45ee-ba1a-a086587469e6/12c8b15c-8d03-4e07-950d-0a627e7595b4",
-    "samsung_galaxy_s22": "arn:aws:devicefarm:us-west-2:308535385114:devicepool:02a2cf0f-6d9b-45ee-ba1a-a086587469e6/e59f866a-30aa-4aa1-87b7-4510e5820dfa",
-    "samsung_galaxy_s22_private": "arn:aws:devicefarm:us-west-2:308535385114:devicepool:02a2cf0f-6d9b-45ee-ba1a-a086587469e6/ea6b049d-1508-4233-9a56-5d9eacbe1078",
-    "samsung_galaxy_s24": "arn:aws:devicefarm:us-west-2:308535385114:devicepool:02a2cf0f-6d9b-45ee-ba1a-a086587469e6/98f8788c-2e25-4a3c-8bb2-0d1e8897c0db",
-    "google_pixel_8_pro": "arn:aws:devicefarm:us-west-2:308535385114:devicepool:02a2cf0f-6d9b-45ee-ba1a-a086587469e6/d65096ab-900b-4521-be8b-a3619b69236a",
-    "google_pixel_3_private_rooted": "arn:aws:devicefarm:us-west-2:308535385114:devicepool:02a2cf0f-6d9b-45ee-ba1a-a086587469e6/98d23ca8-ea9e-4fb7-b725-d402017b198d",
-    "apple_iphone_15_private": "arn:aws:devicefarm:us-west-2:308535385114:devicepool:02a2cf0f-6d9b-45ee-ba1a-a086587469e6/55929353-2f28-4ee5-bdff-d1a95f58cb28",
+    "apple_iphone_15": {
+        "public": [
+            "arn:aws:devicefarm:us-west-2:308535385114:devicepool:02a2cf0f-6d9b-45ee-ba1a-a086587469e6/3b5acd2e-92e2-4778-b651-7726bafe129d",
+        ],
+        "ios_18_public": [
+            "arn:aws:devicefarm:us-west-2:308535385114:devicepool:02a2cf0f-6d9b-45ee-ba1a-a086587469e6/12c8b15c-8d03-4e07-950d-0a627e7595b4",
+        ],
+        "private": [
+            "arn:aws:devicefarm:us-west-2:308535385114:devicepool:02a2cf0f-6d9b-45ee-ba1a-a086587469e6/55929353-2f28-4ee5-bdff-d1a95f58cb28",
+        ],
+        "plus_private": [
+            "arn:aws:devicefarm:us-west-2:308535385114:devicepool:02a2cf0f-6d9b-45ee-ba1a-a086587469e6/767bfb3e-a00e-4d92-998b-4eafdcf7213b",
+        ],
+        "pro_private": [
+            "arn:aws:devicefarm:us-west-2:308535385114:devicepool:02a2cf0f-6d9b-45ee-ba1a-a086587469e6/1394f34c-2981-4c55-aaa2-246871ac713b",
+            "arn:aws:devicefarm:us-west-2:308535385114:devicepool:02a2cf0f-6d9b-45ee-ba1a-a086587469e6/099e8def-4609-4383-8787-76b88e500c1d",
+            "arn:aws:devicefarm:us-west-2:308535385114:devicepool:02a2cf0f-6d9b-45ee-ba1a-a086587469e6/d6707270-b009-479e-a83a-7bdb255f9de5",
+        ],
+    },
+    "samsung_galaxy_s22": {
+        "public": [
+            "arn:aws:devicefarm:us-west-2:308535385114:devicepool:02a2cf0f-6d9b-45ee-ba1a-a086587469e6/e59f866a-30aa-4aa1-87b7-4510e5820dfa",
+        ],
+        "private": [
+            "arn:aws:devicefarm:us-west-2:308535385114:devicepool:02a2cf0f-6d9b-45ee-ba1a-a086587469e6/ea6b049d-1508-4233-9a56-5d9eacbe1078",
+            "arn:aws:devicefarm:us-west-2:308535385114:devicepool:02a2cf0f-6d9b-45ee-ba1a-a086587469e6/1fa924a1-5aff-475b-8f4d-f7c6d8de4fe9",
+        ],
+        "ultra_private": [
+            "arn:aws:devicefarm:us-west-2:308535385114:devicepool:02a2cf0f-6d9b-45ee-ba1a-a086587469e6/5f79d72e-e229-4f9c-962f-5d37196fcfe7",
+        ],
+    },
+    "samsung_galaxy_s24": {
+        "public": [
+            "arn:aws:devicefarm:us-west-2:308535385114:devicepool:02a2cf0f-6d9b-45ee-ba1a-a086587469e6/98f8788c-2e25-4a3c-8bb2-0d1e8897c0db",
+        ],
+        "ultra_private": [
+            "arn:aws:devicefarm:us-west-2:308535385114:devicepool:02a2cf0f-6d9b-45ee-ba1a-a086587469e6/5f79d72e-e229-4f9c-962f-5d37196fcfe7",
+        ],
+    },
+    "google_pixel_8": {
+        "pro_public": [
+            "arn:aws:devicefarm:us-west-2:308535385114:devicepool:02a2cf0f-6d9b-45ee-ba1a-a086587469e6/d65096ab-900b-4521-be8b-a3619b69236a",
+        ],
+    },
+    "google_pixel_3": {
+        "rooted_private": [
+            "arn:aws:devicefarm:us-west-2:308535385114:devicepool:02a2cf0f-6d9b-45ee-ba1a-a086587469e6/98d23ca8-ea9e-4fb7-b725-d402017b198d",
+        ],
+    },
 }
 
 # Predefined benchmark configurations
@@ -318,25 +367,56 @@ def get_benchmark_configs() -> Dict[str, Dict]:  # noqa: C901
 
         # Add configurations for each valid device
         for device in devices:
+            # Parse the device name
+            m = re.match(DEVICE_POOLS_REGEX, device)
+            if not m:
+                logging.warning(
+                    f"Invalid device name: {device} is not in DEVICE_NAME+VARIANT format. Skipping."
+                )
+                continue
+
+            device_name = m.group("device_name")
+            variant = m.group("variant")
+
+            if device_name not in DEVICE_POOLS:
+                logging.warning(f"Unsupported device '{device}'. Skipping.")
+                continue
+
+            if variant not in DEVICE_POOLS[device_name]:
+                logging.warning(
+                    f"Unsupported {device}'s variant '{variant}'. Skipping."
+                )
+                continue
+
+            device_pool_count = len(DEVICE_POOLS[device_name][variant])
+            if not device_pool_count:
+                logging.warning(
+                    f"No device pool defined for {device}'s variant '{variant}'. Skipping."
+                )
+                continue
+
+            device_pool_index = 0
             for config in configs:
-                if config == "llama3_coreml_ane" and not device.endswith("+ios_18"):
-                    device = f"{device}+ios_18"
+                if config == "llama3_coreml_ane" and "ios_18" not in variant:
+                    variant = "ios_18_public"
                     logging.info(
-                        f"Benchmark config '{config}' only works on iOS 18+, auto-upgraded device pool to '{device}'"
+                        f"Benchmark config '{config}' only works on iOS 18+, auto-upgraded device variant to '{variant}'"
                     )
 
-                if device not in DEVICE_POOLS:
-                    logging.warning(f"Unsupported device '{device}'. Skipping.")
-                    continue
-
                 record = {
                     "model": model_name,
                     "config": config,
-                    "device_name": device,
-                    "device_arn": DEVICE_POOLS[device],
+                    "device_name": device_name,
+                    "variant": variant,
+                    "device_arn": DEVICE_POOLS[device_name][variant][
+                        device_pool_index % device_pool_count
+                    ],
                 }
                 benchmark_configs["include"].append(record)
 
+                # Distribute configs to pools of the same device round-robin
+                device_pool_index += 1
+
     set_output("benchmark_configs", json.dumps(benchmark_configs))