Add Whisper HF export (encoder, cross_kv, static-cache decoder) (#23901)

.ci/scripts/test-native.sh

@@ -20,17 +20,22 @@ case "${component}" in
     test_regex='^native_runtime_test$'
     ;;
   ggml)
-    cmake_args+=(-DEXECUTORCH_BUILD_NATIVE_GGML=ON -DGGML_NATIVE=OFF)
-    targets+=(native_ggml_engine_test)
-    test_regex='^native_ggml_engine_test$'
+    cmake_args+=(
+      -DEXECUTORCH_BUILD_NATIVE_GGML=ON
+      -DEXECUTORCH_BUILD_NATIVE_LLM_RUNNER=ON
+      -DGGML_NATIVE=OFF
+    )
+    targets+=(run_hf_llm_ggml native_ggml_engine_test native_llm_runner_utils_test)
+    test_regex='^native_(ggml_engine|llm_runner_utils)_test$'
     ;;
   vulkan)
     cmake_args+=(
       -DEXECUTORCH_BUILD_NATIVE_VULKAN=ON
       -DEXECUTORCH_BUILD_NATIVE_VISION_RUNNER=ON
+      -DEXECUTORCH_BUILD_NATIVE_LLM_RUNNER=ON
       -DEXECUTORCH_VULKAN_SHADER_COMPILE_NTHREADS="${build_jobs}"
     )
-    targets+=(native_vulkan_engine run_classification native_vulkan_passes_test native_runner_utils_test)
+    targets+=(native_vulkan_engine run_classification run_hf_llm_vulkan native_vulkan_passes_test native_runner_utils_test)
     test_regex='^native_(vulkan_passes|runner_utils)_test$'
     ;;
   *)
@@ -53,3 +58,10 @@ cmake --build "${build_dir}" \
   --target "${targets[@]}"
 ctest --test-dir "${build_dir}/backends/native/test" \
   -R "${test_regex}" --output-on-failure --no-tests=error --timeout 120
+
+if [[ "${component}" == "ggml" ]]; then
+  PYTHONPATH="${PWD}/src:${PYTHONPATH:-}" \
+    FLATC_EXECUTABLE="${FLATC_EXECUTABLE:-${PWD}/${build_dir}/third-party/flatc_ep/bin/flatc}" \
+    python .ci/scripts/test_native_llm_runner.py \
+      --runner "${build_dir}/backends/native/examples/hf_llm/runtime/run_hf_llm_ggml"
+fi

.ci/scripts/test_native_llm_runner.py

@@ -0,0 +1,109 @@
+# 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.
+
+"""Compare GGML LLM runner prefill and decoding with a tiny eager HF Llama."""
+
+import argparse
+import base64
+import re
+import subprocess
+import tempfile
+from pathlib import Path
+
+import torch
+from executorch.backends.native.examples.hf_llm.export_hf_llm import LastTokenLM
+from executorch.exir.native import to_native
+from transformers import LlamaConfig, LlamaForCausalLM
+
+
+@torch.no_grad()
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--runner", required=True, type=Path)
+    args = parser.parse_args()
+
+    torch.manual_seed(0)
+    torch.set_num_threads(4)
+    max_seq_length = 48
+    max_new_tokens = 8
+    config = LlamaConfig(
+        vocab_size=512,
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        max_position_embeddings=max_seq_length,
+        bos_token_id=256,
+        eos_token_id=257,
+        attn_implementation="sdpa",
+    )
+    model = LlamaForCausalLM(config).eval()
+    model.generation_config.use_cache = True
+    model.generation_config.cache_implementation = "static"
+    wrapper = LastTokenLM(model, max_seq_length).eval()
+    seq = torch.export.Dim("seq", min=1, max=max_seq_length - 1)
+    exported = torch.export.export(
+        wrapper,
+        (torch.zeros((1, 6), dtype=torch.long), torch.arange(6)),
+        dynamic_shapes={"input_ids": {1: seq}, "cache_position": {0: seq}},
+        strict=True,
+    )
+
+    with tempfile.TemporaryDirectory() as directory:
+        root = Path(directory)
+        package = root / "llama.ptn"
+        to_native(exported).save(str(package))
+        tokenizer = root / "tokenizer.model"
+        # The byte vocabulary puts Llama's BOS at 256, followed by 255 specials.
+        tokenizer.write_text(
+            "".join(
+                f"{base64.b64encode(bytes([i])).decode()} {i}\n" for i in range(256)
+            )
+        )
+        prompt = "abc"
+        unreachable_eos_id = config.vocab_size
+        tokens = [config.bos_token_id, *prompt.encode()]
+        expected = []
+        # Recompute the full prefix without a cache for an independent reference.
+        for _ in range(max_new_tokens):
+            logits = model(torch.tensor([tokens]), use_cache=False).logits
+            token = int(logits[0, -1].argmax())
+            expected.append(token)
+            tokens.append(token)
+
+        result = subprocess.run(
+            [
+                str(args.runner.resolve()),
+                "--device=cpu",
+                "--cpu_threads=4",
+                f"--package={package}",
+                f"--tokenizer={tokenizer}",
+                f"--prompt={prompt}",
+                f"--eos_ids={unreachable_eos_id}",
+                f"--max_new_tokens={max_new_tokens}",
+                f"--max_seq_length={max_seq_length}",
+                "--print_token_ids",
+            ],
+            capture_output=True,
+            text=True,
+            errors="replace",
+            timeout=120,
+        )
+        print(result.stdout, end="")
+        print(result.stderr, end="")
+        result.check_returncode()
+        match = re.search(r"^token_ids: ([0-9,]+)$", result.stdout, re.MULTILINE)
+        if match is None:
+            raise AssertionError("Runner did not report generated token IDs")
+        actual = [int(token) for token in match[1].split(",")]
+        if actual != expected:
+            raise AssertionError(f"Generated tokens {actual} != eager {expected}")
+        print(f"GGML HF LLM runner matched all {max_new_tokens} eager token IDs")
+
+
+if __name__ == "__main__":
+    main()

.github/workflows/test-native.yml

@@ -7,8 +7,13 @@ on:
       - .github/workflows/_docker-image.yml
       - .ci/docker/**
       - .ci/scripts/test-native.sh
+      - .ci/scripts/test_native_llm_runner.py
       - .gitmodules
       - CMakeLists.txt
+      - install_requirements.py
+      - install_requirements.sh
+      - requirements-dev.txt
+      - requirements-examples.txt
       - backends/native/**
       - backends/vulkan/CMakeLists.txt
       - backends/vulkan/cmake/**
@@ -25,6 +30,10 @@ on:
       - backends/xnnpack/third-party/FXdiv
       - backends/xnnpack/third-party/cpuinfo
       - backends/xnnpack/third-party/pthreadpool
+      - examples/models/llama/tokenizer/**
+      - extension/llm/tokenizers
+      - extension/llm/export/**
+      - exir/**
       - runtime/**
       - third-party/**
       - tools/cmake/**
@@ -93,4 +102,12 @@ jobs:
           export CMAKE_PREFIX_PATH="/tmp/native-deps:${CMAKE_PREFIX_PATH:-}"
           export PATH="/tmp/native-deps/bin:${PATH}"
 
+          if [[ "${NATIVE_COMPONENT}" == "ggml" ]]; then
+            ./install_requirements.sh --use-pt-pinned-commit
+            # The LLM runner test imports executorch from src/ without installing the
+            # wheel, so install its export-time dependencies directly.
+            pip install flatbuffers ruamel.yaml tabulate \
+              "$(grep -E '^transformers\b' requirements-examples.txt)"
+          fi
+
           CMAKE_BUILD_PARALLEL_LEVEL=8 bash .ci/scripts/test-native.sh "${NATIVE_COMPONENT}"

CMakeLists.txt

@@ -980,7 +980,10 @@ if(EXECUTORCH_BUILD_EXTENSION_NAMED_DATA_MAP)
 endif()
 
 if(EXECUTORCH_BUILD_EXTENSION_LLM)
-  if(EXECUTORCH_BUILD_EXTENSION_LLM_RUNNER)
+  if(EXECUTORCH_BUILD_EXTENSION_LLM_RUNNER
+     OR EXECUTORCH_BUILD_NATIVE_LLM_RUNNER
+  )
+    # Both LLM runners need tokenizers with regex lookahead support.
     set(SUPPORT_REGEX_LOOKAHEAD ON)
     # llama/runner/CMakeLists.txt builds a shared library libllama_runner.so
     # that transitively depends on tokenizers. Need to build tokenizers with
@@ -991,7 +994,9 @@ if(EXECUTORCH_BUILD_EXTENSION_LLM)
     set(CMAKE_POSITION_INDEPENDENT_CODE ON)
   endif()
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/llm/tokenizers)
-  if(EXECUTORCH_BUILD_EXTENSION_LLM_RUNNER)
+  if(EXECUTORCH_BUILD_EXTENSION_LLM_RUNNER
+     OR EXECUTORCH_BUILD_NATIVE_LLM_RUNNER
+  )
     set(CMAKE_POSITION_INDEPENDENT_CODE
         ${ORIGINAL_CMAKE_POSITION_INDEPENDENT_CODE_FLAG}
     )

backends/native/CMakeLists.txt

@@ -36,7 +36,7 @@ install(
   PATTERN "ggml" EXCLUDE
 )
 
-if(EXECUTORCH_BUILD_NATIVE_VISION_RUNNER)
+if(EXECUTORCH_BUILD_NATIVE_VISION_RUNNER OR EXECUTORCH_BUILD_NATIVE_LLM_RUNNER)
   add_subdirectory(examples)
 endif()
 

backends/native/README.md

@@ -33,6 +33,11 @@ quantize/dequantize ops. TorchAO int4 group-wise weights of `linear` and
 stored packed with their scheme in tensor metadata. Each method must lower wholly
 to one delegate; `to_native` raises otherwise.
 
+Immutable constant keys are namespaced as `<method>/<key>`, even for a single
+method. A key mutated by any method becomes `state/<key>` in every method, so
+methods share its runtime state. Method names may not be `state` or start with
+`state/`. `to_native` raises if namespaced keys collide.
+
 A `.ptn` is an uncompressed zip with fixed entry names:
 
 | Entry | Contents |
@@ -110,16 +115,21 @@ target_link_libraries(my_app PRIVATE native_api native_vulkan_engine)
 
 With `EXECUTORCH_BUILD_TESTS=ON`, `native_runtime_test` covers the shared runtime
 and API. `native_vulkan_passes_test` covers the Vulkan graph rewrites without
-requiring a GPU. When the vision runner below is enabled,
-`native_runner_utils_test` covers its shared runner utilities. Run these with
+requiring a GPU. When either runner below is enabled,
+`native_runner_utils_test` covers their shared runner utilities. Run these with
 `ctest --test-dir cmake-out -R native_`.
 
 The `Test Native Engines` CI workflow has separate Linux jobs for the shared
 runtime, GGML, and Vulkan, with each test suite run once. To reproduce with
 dependencies installed, run `bash .ci/scripts/test-native.sh runtime`,
 `bash .ci/scripts/test-native.sh ggml`, or
-`bash .ci/scripts/test-native.sh vulkan` from the OSS checkout. The Vulkan job
-also builds `run_classification`; its tests do not require a GPU.
+`bash .ci/scripts/test-native.sh vulkan` from the OSS checkout. The GGML job
+also runs `native_llm_runner_utils_test` and executes `run_hf_llm_ggml` on a tiny
+seeded Hugging Face Llama with a generated byte tokenizer. Its eight generated
+token IDs must match eager PyTorch, covering prefill and KV-cache decoding without
+downloading a model. The GGML job requires the Python export dependencies
+installed by the `Test Native Engines` workflow. The Vulkan job builds
+`run_classification` and `run_hf_llm_vulkan`; its tests do not require a GPU.
 
 ### GGML engine
 
@@ -177,3 +187,31 @@ Export the `.ptn`, input, reference logits, and labels with
 [`export.py`](../../examples/native/vision/export.py). The runner prints the
 top-5 logits and, when `--expected` is given, exits nonzero if the logits
 differ from the reference beyond `--rtol`/`--atol` (both default to `1e-3`).
+
+## Native LLM runner
+
+`EXECUTORCH_BUILD_NATIVE_LLM_RUNNER=ON` builds `run_hf_llm_vulkan` with Native
+Vulkan and `run_hf_llm_ggml` with Native GGML. It requires gflags and
+`tokenizers::tokenizers` (the vendored tokenizer is built automatically in an
+OSS checkout; otherwise provide an installed tokenizers package with regex
+lookahead support). It accepts Llama 3 TikToken tokenizer files.
+
+To build the text runner for GGML CPU, initialize the llama.cpp submodule as above:
+
+```sh
+cmake -S . -B cmake-out -DCMAKE_BUILD_TYPE=Release \
+  -DEXECUTORCH_BUILD_NATIVE=ON -DEXECUTORCH_BUILD_NATIVE_GGML=ON \
+  -DEXECUTORCH_BUILD_NATIVE_LLM_RUNNER=ON
+cmake --build cmake-out --target run_hf_llm_ggml
+cmake --install cmake-out --prefix install
+install/bin/run_hf_llm_ggml --device cpu --cpu_threads 4 \
+  --package llama.ptn --tokenizer tokenizer.model \
+  --prompt "The capital of France is" --max_new_tokens 32
+```
+
+The GGML runner accepts `--device cpu|gpu`, `--device_index`, `--cpu_threads`,
+and `--max_cached_graphs` (default 1, 0 rebuilds each call). For GPU execution,
+enable the desired `GGML_*` backend when configuring ExecuTorch, as above.
+
+With `EXECUTORCH_BUILD_TESTS=ON`, `native_llm_runner_utils_test` exercises the
+LLM generation utilities.

backends/native/examples/CMakeLists.txt

@@ -4,5 +4,19 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+if(NOT TARGET gflags AND NOT TARGET gflags::gflags)
+  find_package(gflags REQUIRED)
+endif()
+if(TARGET gflags::gflags)
+  set(_native_runner_gflags gflags::gflags)
+else()
+  set(_native_runner_gflags gflags)
+endif()
+
 add_subdirectory(runner_utils)
-add_subdirectory(vision)
+if(EXECUTORCH_BUILD_NATIVE_VISION_RUNNER)
+  add_subdirectory(vision)
+endif()
+if(EXECUTORCH_BUILD_NATIVE_LLM_RUNNER)
+  add_subdirectory(hf_llm/runtime)
+endif()

backends/native/examples/audio/BUCK

@@ -0,0 +1,28 @@
+load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target")
+load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
+
+oncall("executorch")
+
+fbcode_target(
+    _kind = runtime.python_library,
+    name = "hf_whisper",
+    srcs = ["export_hf_whisper.py"],
+    visibility = ["//executorch/backends/native/test/..."],
+    deps = [
+        "//caffe2:torch",
+        "//executorch/backends/native/examples/hf_llm:hf_utils",
+        "//executorch/exir:lib",
+        "//executorch/exir/native:native",
+        "fbsource//third-party/pypi/transformers:transformers",
+    ],
+)
+
+fbcode_target(
+    _kind = runtime.python_binary,
+    name = "export_hf_whisper",
+    main_function = "executorch.backends.native.examples.audio.export_hf_whisper.main",
+    deps = [
+        ":hf_whisper",
+        "fbsource//third-party/pypi/safetensors:safetensors",
+    ],
+)