Update on "[ET Device Support] DeviceAllocator interface and DeviceAllocatorRegistry"

This diff introduces the `DeviceAllocator` abstract interface and `DeviceAllocatorRegistry` for device-specific memory allocation. This is a foundational abstraction that enables the runtime to dispatch memory operations to the appropriate device backend other than CPU (CUDA, etc.).

**DeviceAllocator interface provides:**
- `init_buffer()` - Initialize memory buffer pools for memory-planned tensors
- `get_offset_address()` - Get pointer to offset within pre-allocated buffer
- `allocate()` / `deallocate()` - Dynamic device memory allocation
- `copy_host_to_device()` / `copy_device_to_host()` - Data transfer between host and device
- `device_type()` - Returns the device type this allocator handles

**DeviceAllocatorRegistry provides:**
- Singleton registry mapping DeviceType → DeviceAllocator
- `register_allocator()` / `get_allocator()` methods
- Fixed-size array indexed by device type (no dynamic allocation, embedded-friendly)

**Design notes:**
- Registry stores raw pointers (non-owning) - allocators are expected to be singletons with static lifetime
- Follows ExecuTorch's embedded-first philosophy (no std::unique_ptr, no heap allocation in registry)
- Convenience free functions `register_device_allocator()` and `get_device_allocator()` for ease of use

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

[ghstack-poisoned]

Author: Gasoonjia

.ci/docker/build.sh

@@ -40,6 +40,15 @@ case "${IMAGE_NAME}" in
     LINTRUNNER=""
     GCC_VERSION=11
     ;;
+  executorch-ubuntu-22.04-gcc11-aarch64-android)
+    LINTRUNNER=""
+    GCC_VERSION=11
+    ANDROID_NDK_VERSION=r28c
+    ;;
+  executorch-ubuntu-22.04-gcc11-aarch64-arm-sdk)
+    ARM_SDK=yes
+    GCC_VERSION=11
+    ;;
   executorch-ubuntu-22.04-linter)
     LINTRUNNER=yes
     CLANG_VERSION=12

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

@@ -1 +1 @@
-5bf1aeb587e9b1f3572b0bd60265c5dafd007b73
+a9592258daacad7423fd5f39aaa59c6e36471520

.ci/docker/ci_commit_pins/pytorch.txt

@@ -1 +1 @@
-7a79b41e29a790ebb4b530eb98a89381e2d7de29
+659af3c353e49b35c191cdd2dba3b3c79d0e6822

.ci/docker/common/install_android.sh

@@ -40,8 +40,14 @@ install_ndk() {
   rm -rf "${NDK_INSTALLATION_DIR}" && mkdir -p "${NDK_INSTALLATION_DIR}"
 
   pushd /tmp
-  # The NDK installation is cached on ossci-android S3 bucket
-  curl -Os --retry 3 "https://ossci-android.s3.amazonaws.com/android-ndk-${ANDROID_NDK_VERSION}-linux.zip"
+  ARCH=$(uname -m)
+  if [ "${ARCH}" = "aarch64" ]; then
+    # aarch64 NDK is not cached on S3, download from Google directly
+    curl -Os --retry 3 "https://dl.google.com/android/repository/android-ndk-${ANDROID_NDK_VERSION}-linux.zip"
+  else
+    # The NDK installation is cached on ossci-android S3 bucket
+    curl -Os --retry 3 "https://ossci-android.s3.amazonaws.com/android-ndk-${ANDROID_NDK_VERSION}-linux.zip"
+  fi
   unzip -qo "android-ndk-${ANDROID_NDK_VERSION}-linux.zip"
 
   # Print the content for manual verification
@@ -73,7 +79,10 @@ install_sdk() {
   yes | /opt/cmdline-tools/bin/sdkmanager --sdk_root="${SDK_INSTALLATION_DIR}" --install "build-tools;35.0.0"
   # And some more tools for future emulator tests
   yes | /opt/cmdline-tools/bin/sdkmanager --sdk_root="${SDK_INSTALLATION_DIR}" --install "platform-tools"
-  yes | /opt/cmdline-tools/bin/sdkmanager --sdk_root="${SDK_INSTALLATION_DIR}" --install "tools"
+  # The 'tools' package (emulator) is not available on aarch64
+  if [ "$(uname -m)" != "aarch64" ]; then
+    yes | /opt/cmdline-tools/bin/sdkmanager --sdk_root="${SDK_INSTALLATION_DIR}" --install "tools"
+  fi
 }
 
 install_prerequiresites

.ci/docker/requirements-ci.txt

@@ -30,6 +30,6 @@ sphinx-reredirects==0.1.4
 matplotlib>=3.9.4
 sphinx-copybutton==0.5.2
 # PyTorch Theme
-pytorch_sphinx_theme2==0.2.0
+pytorch_sphinx_theme2==0.4.4
 # script unit test requirements
 yaspin==3.1.0

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

@@ -18,7 +18,7 @@ build_qnn_backend() {
   export EXECUTORCH_ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." && pwd)"
 
   parallelism=$(( $(nproc) - 1 ))
-  bash backends/qualcomm/scripts/build.sh --skip_linux_android --skip_linux_embedded --job_number ${parallelism} --release
+  bash backends/qualcomm/scripts/build.sh --skip_linux_android --job_number ${parallelism} --release
 }
 
 set_up_aot() {

.ci/scripts/export_model_artifact.sh

@@ -9,7 +9,7 @@
 
 show_help() {
   cat << EOF
-Usage: export_model_artifact.sh <device> <hf_model> [quant_name] [output_dir]
+Usage: export_model_artifact.sh <device> <hf_model> [quant_name] [output_dir] [mode]
 
 Export a HuggingFace model to CUDA/Metal/XNNPACK format with optional quantization.
 
@@ -22,6 +22,7 @@ Arguments:
                  - mistralai/Voxtral-Mini-4B-Realtime-2602
                  - openai/whisper series (whisper-{small, medium, large, large-v2, large-v3, large-v3-turbo})
                  - google/gemma-3-4b-it
+                 - nvidia/diar_streaming_sortformer_4spk-v2
                  - nvidia/parakeet-tdt
 
   quant_name   Quantization type (optional, default: non-quantized)
@@ -34,13 +35,23 @@ Arguments:
 
   output_dir   Output directory for artifacts (optional, default: current directory)
 
+  mode         Export mode (optional, default: vr-streaming)
+               Supported modes:
+                 - vr-streaming: Voxtral Realtime streaming mode
+                 - vr-offline: Voxtral Realtime offline mode
+
 Examples:
   export_model_artifact.sh metal "openai/whisper-small"
   export_model_artifact.sh metal "nvidia/parakeet-tdt" "quantized-int4-metal"
+  export_model_artifact.sh metal "mistralai/Voxtral-Mini-4B-Realtime-2602" "quantized-int4-metal"
+  export_model_artifact.sh metal "mistralai/Voxtral-Mini-4B-Realtime-2602" "non-quantized" "." "vr-streaming"
   export_model_artifact.sh cuda "mistralai/Voxtral-Mini-3B-2507" "quantized-int4-tile-packed"
+  export_model_artifact.sh cuda-windows "nvidia/diar_streaming_sortformer_4spk-v2" "non-quantized" "./output"
   export_model_artifact.sh cuda "google/gemma-3-4b-it" "non-quantized" "./output"
   export_model_artifact.sh cuda "nvidia/parakeet-tdt" "non-quantized" "./output"
   export_model_artifact.sh xnnpack "nvidia/parakeet-tdt" "quantized-8da4w" "./output"
+  export_model_artifact.sh xnnpack "mistralai/Voxtral-Mini-4B-Realtime-2602" "quantized-8da4w" "./output"
+  export_model_artifact.sh xnnpack "mistralai/Voxtral-Mini-4B-Realtime-2602" "non-quantized" "./output" "vr-offline"
 EOF
 }
 
@@ -61,6 +72,26 @@ DEVICE="$1"
 HF_MODEL="$2"
 QUANT_NAME="${3:-non-quantized}"
 OUTPUT_DIR="${4:-.}"
+MODE="${5:-}"
+
+# Validate mode if specified
+if [ -n "$MODE" ]; then
+  case "$MODE" in
+    vr-streaming|vr-offline)
+      # Voxtral Realtime modes require Voxtral Realtime model
+      if [ "$HF_MODEL" != "mistralai/Voxtral-Mini-4B-Realtime-2602" ]; then
+        echo "Error: Mode '$MODE' can only be used with Voxtral Realtime model"
+        echo "Provided model: $HF_MODEL"
+        exit 1
+      fi
+      ;;
+    *)
+      echo "Error: Unsupported mode '$MODE'"
+      echo "Supported modes: vr-streaming, vr-offline"
+      exit 1
+      ;;
+  esac
+fi
 
 case "$DEVICE" in
   cuda)
@@ -112,6 +143,14 @@ case "$HF_MODEL" in
     PREPROCESSOR_FEATURE_SIZE=""
     PREPROCESSOR_OUTPUT=""
     ;;
+  Qwen/Qwen3-0.6B)
+    MODEL_NAME="qwen3"
+    TASK="text-generation"
+    MAX_SEQ_LEN="64"
+    EXTRA_PIP=""
+    PREPROCESSOR_FEATURE_SIZE=""
+    PREPROCESSOR_OUTPUT=""
+    ;;
   nvidia/parakeet-tdt)
     MODEL_NAME="parakeet"
     TASK=""
@@ -120,6 +159,14 @@ case "$HF_MODEL" in
     PREPROCESSOR_FEATURE_SIZE=""
     PREPROCESSOR_OUTPUT=""
     ;;
+  nvidia/diar_streaming_sortformer_4spk-v2)
+    MODEL_NAME="sortformer"
+    TASK=""
+    MAX_SEQ_LEN=""
+    EXTRA_PIP=""
+    PREPROCESSOR_FEATURE_SIZE=""
+    PREPROCESSOR_OUTPUT=""
+    ;;
   mistralai/Voxtral-Mini-4B-Realtime-2602)
     MODEL_NAME="voxtral_realtime"
     TASK=""
@@ -130,7 +177,7 @@ case "$HF_MODEL" in
     ;;
   *)
     echo "Error: Unsupported model '$HF_MODEL'"
-    echo "Supported models: mistralai/Voxtral-Mini-3B-2507, mistralai/Voxtral-Mini-4B-Realtime-2602, openai/whisper-{small, medium, large, large-v2, large-v3, large-v3-turbo}, google/gemma-3-4b-it, nvidia/parakeet-tdt"
+    echo "Supported models: mistralai/Voxtral-Mini-3B-2507, mistralai/Voxtral-Mini-4B-Realtime-2602, openai/whisper-{small, medium, large, large-v2, large-v3, large-v3-turbo}, google/gemma-3-4b-it, Qwen/Qwen3-0.6B, nvidia/diar_streaming_sortformer_4spk-v2, nvidia/parakeet-tdt"
     exit 1
     ;;
 esac
@@ -210,7 +257,43 @@ if [ "$MODEL_NAME" = "parakeet" ]; then
   exit 0
 fi
 
-# Voxtral Realtime uses a custom export script (streaming mode)
+# Sortformer uses a custom export script
+if [ "$MODEL_NAME" = "sortformer" ]; then
+  if [ "$QUANT_NAME" != "non-quantized" ]; then
+    echo "Error: Sortformer currently supports only non-quantized export"
+    exit 1
+  fi
+
+  pip install -r examples/models/sortformer/install_requirements.txt
+
+  SORTFORMER_BACKEND="$DEVICE"
+  if [ "$DEVICE" = "cuda-windows" ]; then
+    SORTFORMER_BACKEND="cuda-windows"
+  elif [ "$DEVICE" = "cuda" ]; then
+    SORTFORMER_BACKEND="cuda"
+  elif [ "$DEVICE" = "xnnpack" ]; then
+    SORTFORMER_BACKEND="xnnpack"
+  else
+    SORTFORMER_BACKEND="portable"
+  fi
+
+  python -m executorch.examples.models.sortformer.export_sortformer \
+      --hf-model "${HF_MODEL}" \
+      --backend "${SORTFORMER_BACKEND}" \
+      --output-dir "${OUTPUT_DIR}"
+
+  test -f "${OUTPUT_DIR}/sortformer.pte"
+  mv "${OUTPUT_DIR}/sortformer.pte" "${OUTPUT_DIR}/model.pte"
+  # CUDA saves named data to separate .ptd file, XNNPACK/portable do not.
+  if [ "$DEVICE" = "cuda" ] || [ "$DEVICE" = "cuda-windows" ]; then
+    test -f "${OUTPUT_DIR}/aoti_cuda_blob.ptd"
+  fi
+  ls -al "${OUTPUT_DIR}"
+  echo "::endgroup::"
+  exit 0
+fi
+
+# Voxtral Realtime uses a custom export script
 if [ "$MODEL_NAME" = "voxtral_realtime" ]; then
   pip install safetensors huggingface_hub
 
@@ -220,25 +303,48 @@ if [ "$MODEL_NAME" = "voxtral_realtime" ]; then
 
   # Per-component quantization flags
   VR_QUANT_ARGS=""
+  VR_DTYPE_ARGS=""
   if [ "$QUANT_NAME" = "quantized-8da4w" ]; then
     VR_QUANT_ARGS="--qlinear-encoder 8da4w --qlinear 8da4w --qlinear-group-size 32 --qembedding 8w"
+  elif [ "$QUANT_NAME" = "quantized-int4-metal" ]; then
+    VR_QUANT_ARGS="--qlinear-encoder fpa4w --qlinear fpa4w"
+  elif [ "$QUANT_NAME" = "quantized-int4-tile-packed" ]; then
+    VR_QUANT_ARGS="--qlinear-encoder 4w --qlinear-encoder-packing-format tile_packed_to_4d --qlinear 4w --qlinear-packing-format tile_packed_to_4d --qembedding 8w"
+    VR_DTYPE_ARGS="--dtype bf16"
+  fi
+
+  # Determine streaming mode based on MODE parameter
+  USE_STREAMING="true"
+  if [ "$MODE" = "vr-offline" ]; then
+    USE_STREAMING="false"
+  fi
+
+  # Configure export and preprocessor based on streaming mode
+  STREAMING_ARG=""
+  PREPROCESSOR_ARGS="--feature_size 128 --output_file ${OUTPUT_DIR}/preprocessor.pte"
+  if [ "$USE_STREAMING" = "true" ]; then
+    STREAMING_ARG="--streaming"
+    PREPROCESSOR_ARGS="$PREPROCESSOR_ARGS --streaming"
+  else
+    PREPROCESSOR_ARGS="$PREPROCESSOR_ARGS --stack_output --max_audio_len 300"
   fi
 
   python -m executorch.examples.models.voxtral_realtime.export_voxtral_rt \
       --model-path "$LOCAL_MODEL_DIR" \
-      --backend xnnpack \
-      --streaming \
+      --backend "$DEVICE" \
+      ${STREAMING_ARG} \
       --output-dir "${OUTPUT_DIR}" \
-      ${VR_QUANT_ARGS}
+      ${VR_QUANT_ARGS} \
+      ${VR_DTYPE_ARGS}
 
-  # Export streaming preprocessor (no chunk padding)
-  python -m executorch.extension.audio.mel_spectrogram \
-      --feature_size 128 \
-      --streaming \
-      --output_file "${OUTPUT_DIR}/preprocessor.pte"
+  # Export preprocessor
+  python -m executorch.extension.audio.mel_spectrogram ${PREPROCESSOR_ARGS}
 
   test -f "${OUTPUT_DIR}/model.pte"
   test -f "${OUTPUT_DIR}/preprocessor.pte"
+  if [ "$DEVICE" = "cuda" ] || [ "$DEVICE" = "cuda-windows" ]; then
+    test -f "${OUTPUT_DIR}/aoti_cuda_blob.ptd"
+  fi
   # Copy tokenizer from downloaded model weights
   cp "$LOCAL_MODEL_DIR/tekken.json" "${OUTPUT_DIR}/tekken.json"
   ls -al "${OUTPUT_DIR}"

.ci/scripts/test_backend.sh

@@ -85,7 +85,10 @@ else
 fi
 CMAKE_ARGS="$EXTRA_BUILD_ARGS" ${CONDA_RUN_CMD} $SETUP_SCRIPT --build-tool cmake --build-mode Release --editable true
 
+GOLDEN_DIR="${ARTIFACT_DIR}/golden-artifacts"
+export GOLDEN_ARTIFACTS_DIR="${GOLDEN_DIR}"
+
 EXIT_CODE=0
-${CONDA_RUN_CMD} pytest -c /dev/nul -n auto backends/test/suite/$SUITE/ -m flow_$FLOW --json-report --json-report-file="$REPORT_FILE" || EXIT_CODE=$?
+${CONDA_RUN_CMD} pytest -c /dev/null -n auto backends/test/suite/$SUITE/ -m flow_$FLOW --json-report --json-report-file="$REPORT_FILE" || EXIT_CODE=$?
 # Generate markdown summary.
 ${CONDA_RUN_CMD} python -m executorch.backends.test.suite.generate_markdown_summary_json "$REPORT_FILE" > ${GITHUB_STEP_SUMMARY:-"step_summary.md"} --exit-code $EXIT_CODE

.ci/scripts/test_huggingface_optimum_model.py

@@ -11,7 +11,6 @@
 from datasets import load_dataset
 
 from optimum.executorch import (
-    ExecuTorchModelForCausalLM,
     ExecuTorchModelForImageClassification,
     ExecuTorchModelForMaskedLM,
     ExecuTorchModelForSeq2SeqLM,
@@ -143,27 +142,62 @@ def test_text_generation(model_id, model_dir, recipe, *, quantize=True, run_only
                 "--qembedding",
                 "8w",
             ]
+    elif recipe == "cuda":
+        command += [
+            "--dtype",
+            "bfloat16",
+            "--device",
+            "cuda",
+        ]
+        if quantize:
+            command += [
+                "--qlinear",
+                "4w",
+                "--qlinear_packing_format",
+                "tile_packed_to_4d",
+                "--qembedding",
+                "8w",
+            ]
     else:
         assert (
             not quantize
-        ), "Quantization is only supported for XnnPack and CoreML recipes at the moment."
+        ), "Quantization is only supported for XnnPack, CoreML, and CUDA recipes at the moment."
 
     if not run_only:
         cli_export(command, model_dir)
 
+    if recipe == "cuda":
+        model_path = Path(model_dir) / "model.pte"
+        cuda_blob_path = Path(model_dir) / "aoti_cuda_blob.ptd"
+        assert model_path.exists(), f"Main model file not found: {model_path}"
+        assert cuda_blob_path.exists(), f"CUDA blob not found: {cuda_blob_path}"
+
     tokenizer = AutoTokenizer.from_pretrained(model_id)
-    tokenizer.save_pretrained(model_dir)
-    model = ExecuTorchModelForCausalLM.from_pretrained(model_dir)
-    generated_text = model.text_generation(
-        tokenizer=tokenizer,
-        prompt="Simply put, the theory of relativity states that",
-        max_seq_len=64,
+    saved_files = tokenizer.save_pretrained(model_dir)
+    tokenizer_path = get_tokenizer_path(model_dir, saved_files)
+
+    from executorch.extension.llm.runner import GenerationConfig, TextLLMRunner
+
+    if recipe == "cuda":
+        runner = TextLLMRunner(
+            f"{model_dir}/model.pte",
+            tokenizer_path,
+            f"{model_dir}/aoti_cuda_blob.ptd",
+        )
+    else:
+        runner = TextLLMRunner(f"{model_dir}/model.pte", tokenizer_path)
+    tokens = []
+    runner.generate(
+        "Simply put, the theory of relativity states that",
+        GenerationConfig(seq_len=64, temperature=0, echo=True),
+        token_callback=lambda t: tokens.append(t),
     )
+    generated_text = "".join(tokens)
     print(f"\nGenerated text:\n\t{generated_text}")
     generated_tokens = tokenizer(generated_text, return_tensors="pt").input_ids
 
     # Free memory before loading eager for quality check
-    del model
+    del runner
     del tokenizer
     gc.collect()
 

.ci/scripts/test_lora.sh

@@ -12,7 +12,8 @@ source "$(dirname "${BASH_SOURCE[0]}")/utils.sh"
 cmake_install_executorch_libraries() {
     echo "Installing libexecutorch.a, libextension_module.so, libportable_ops_lib.a"
     rm -rf cmake-out
-    cmake --workflow llm-release
+    cmake --preset llm-release -DEXECUTORCH_ENABLE_LOGGING=ON
+    cmake --build --preset llm-release-install
 }
 
 cmake_build_llama_runner() {