From 5c39c8d58754fc2966f26a7ad0aef08d4adf32d0 Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Tue, 3 Nov 2026 08:16:52 -0800 Subject: [PATCH 1/2] Support Whisper positional indexing (#23903) Summary: Lower matrix row gathers from `aten.index.Tensor` through the existing embedding path and lower positive bounded `aten.repeat` shapes with `ggml_repeat`. Preserve snapshots when repeat aliases its input so later cache writes cannot change previously computed values. Allow row-index validation to follow repeats and validate integer `aten.arange.start_step` bounds when building the graph. This supports Whisper's encoder position embeddings and decoder cache-position gathers while retaining input bounds checks on cached graph replay. Review row-index validation and gather lowering first, then repeat lowering and the operator regressions. Authored with assistance from Codex. Differential Revision: D125479769 --- .../runtime/engine/ggml/GgmlGraphBuilder.cpp | 57 ++++++++++--------- .../runtime/engine/ggml/GgmlGraphBuilder.h | 3 +- .../runtime/engine/ggml/GgmlLowerLinear.cpp | 26 +++++++-- .../runtime/engine/ggml/GgmlLowerShape.cpp | 27 +++++++++ .../native/test/runtime/ggml/test_ggml_ops.py | 50 ++++++++++++++++ 5 files changed, 132 insertions(+), 31 deletions(-) diff --git a/backends/native/runtime/engine/ggml/GgmlGraphBuilder.cpp b/backends/native/runtime/engine/ggml/GgmlGraphBuilder.cpp index 372176f8d2d..ee6718ebacc 100644 --- a/backends/native/runtime/engine/ggml/GgmlGraphBuilder.cpp +++ b/backends/native/runtime/engine/ggml/GgmlGraphBuilder.cpp @@ -649,8 +649,8 @@ void GraphBuilder::lower(const Node& node) { const std::string key = op_key(node.target); if (key == "aten.linear.default") { lower_linear(node); - } else if (key == "aten.embedding.default") { - lower_embedding(node); + } else if (key == "aten.embedding.default" || key == "aten.index.Tensor") { + lower_row_gather(node); } else if (key == "aten.rms_norm.default") { lower_rms_norm(node); } else if ( @@ -713,6 +713,8 @@ void GraphBuilder::lower(const Node& node) { } else if ( key == "aten.expand_copy.default" || key == "aten.expand.default") { lower_expand(node); + } else if (key == "aten.repeat.default") { + lower_repeat(node); } else if (key == "aten.cat.default") { lower_cat(node); } else if ( @@ -753,31 +755,21 @@ void GraphBuilder::lower(const Node& node) { } // Row indices are checked on the host, which needs them to come from a -// graph input, possibly through integer casts and views; any other source -// is rejected. +// graph input through integer casts, views or repeats, or an integer arange +// whose bounds are resolved when building the graph. void GraphBuilder::check_row_indices(ValueId id, int64_t rows) { static const std::unordered_set kSameValues = { - "aten._to_copy.default", - "aten.alias.default", - "aten.clone.default", - "aten.contiguous.default", - "aten.expand.default", - "aten.expand_copy.default", - "aten.permute.default", - "aten.permute_copy.default", - "aten.reshape.default", - "aten._unsafe_view.default", - "aten.select.int", - "aten.select_copy.int", - "aten.slice.Tensor", - "aten.slice_copy.Tensor", - "aten.squeeze.dims", - "aten.squeeze_copy.dims", - "aten.transpose_copy.int", - "aten.unsqueeze.default", - "aten.unsqueeze_copy.default", - "aten.view.default", - "aten.view_copy.default", + "aten._to_copy.default", "aten.alias.default", + "aten.clone.default", "aten.contiguous.default", + "aten.expand.default", "aten.expand_copy.default", + "aten.permute.default", "aten.permute_copy.default", + "aten.repeat.default", "aten.reshape.default", + "aten._unsafe_view.default", "aten.select.int", + "aten.select_copy.int", "aten.slice.Tensor", + "aten.slice_copy.Tensor", "aten.squeeze.dims", + "aten.squeeze_copy.dims", "aten.transpose_copy.int", + "aten.unsqueeze.default", "aten.unsqueeze_copy.default", + "aten.view.default", "aten.view_copy.default", }; const auto is_index = [this](ValueId value) { const ScalarType dtype = method_.graph.value(value).tensor_meta().dtype; @@ -796,6 +788,19 @@ void GraphBuilder::check_row_indices(ValueId id, int64_t rows) { break; } const Node& node = method_.graph.node(producer); + if (node.is_call() && op_key(node.target) == "aten.arange.start_step") { + const int64_t start = int_arg(node, "start", 0, 0); + const int64_t step = int_arg(node, "step", 2, 1, true); + const int64_t count = ggml_nelements(tensor(id).tensor); + int64_t last = 0; + if (count > 0 && + (__builtin_mul_overflow(count - 1, step, &last) || + __builtin_add_overflow(start, last, &last) || + std::min(start, last) < 0 || std::max(start, last) >= rows)) { + throw std::out_of_range("ggml: arange row index is outside the table"); + } + return; + } const NamedArgument* source = find_arg(node, "self", 0); if (!node.is_call() || !kSameValues.contains(op_key(node.target)) || source == nullptr || source->arg.kind() != ArgKind::Tensor) { @@ -805,7 +810,7 @@ void GraphBuilder::check_row_indices(ValueId id, int64_t rows) { } throw std::runtime_error( "ggml: row indices must come from a graph input, through integer " - "casts and views only"); + "casts, views or repeats, or an integer arange"); } ggml_tensor* GraphBuilder::as_f32(ggml_tensor* tensor) { diff --git a/backends/native/runtime/engine/ggml/GgmlGraphBuilder.h b/backends/native/runtime/engine/ggml/GgmlGraphBuilder.h index ee8db65dfc7..2f69e368fe0 100644 --- a/backends/native/runtime/engine/ggml/GgmlGraphBuilder.h +++ b/backends/native/runtime/engine/ggml/GgmlGraphBuilder.h @@ -186,7 +186,7 @@ class GraphBuilder { void set_single_output(const Node& node, TensorRef result); void lower(const Node& node); void lower_linear(const Node& node); - void lower_embedding(const Node& node); + void lower_row_gather(const Node& node); void check_row_indices(ValueId id, int64_t rows); void lower_rms_norm(const Node& node); void lower_layer_norm(const Node& node); @@ -244,6 +244,7 @@ class GraphBuilder { void lower_unsqueeze(const Node& node); void lower_squeeze(const Node& node); void lower_expand(const Node& node); + void lower_repeat(const Node& node); void lower_cat(const Node& node); void lower_matmul(const Node& node); void lower_update_cache(const Node& node); diff --git a/backends/native/runtime/engine/ggml/GgmlLowerLinear.cpp b/backends/native/runtime/engine/ggml/GgmlLowerLinear.cpp index 93281445b80..01625b42a09 100644 --- a/backends/native/runtime/engine/ggml/GgmlLowerLinear.cpp +++ b/backends/native/runtime/engine/ggml/GgmlLowerLinear.cpp @@ -47,11 +47,29 @@ void GraphBuilder::lower_linear(const Node& node) { set_single_output(node, TensorRef{result, std::move(sizes)}); } -void GraphBuilder::lower_embedding(const Node& node) { +void GraphBuilder::lower_row_gather(const Node& node) { const TensorRef weight = tensor_arg(node, "weight", 0); - TensorRef indices = tensor_arg(node, "indices", 1); - check_row_indices( - require_arg(node, "indices", 1).arg.as_tensor().id, weight.sizes.at(0)); + const Argument& arg = require_arg(node, "indices", 1).arg; + ValueId indices_id = kInvalid; + if (op_key(node.target) == "aten.embedding.default") { + indices_id = arg.as_tensor().id; + } else { + const std::vector ids = arg.kind() == ArgKind::TensorList + ? arg.as_tensor_list().ids + : arg.as_optional_tensor_list().ids; + if (ids.size() != 1 || !valid(ids.front())) { + throw std::runtime_error("ggml: index supports one row-index tensor"); + } + indices_id = ids.front(); + } + const ScalarType dtype = method_.graph.value(indices_id).tensor_meta().dtype; + if (weight.sizes.size() != 2 || + (dtype != ScalarType::Int && dtype != ScalarType::Long)) { + throw std::runtime_error( + "ggml: row gather needs a matrix and integer indices"); + } + TensorRef indices = tensor(indices_id); + check_row_indices(indices_id, weight.sizes[0]); if (indices.tensor->type != GGML_TYPE_I32) { indices.tensor = ggml_cast(instance_.context.get(), indices.tensor, GGML_TYPE_I32); diff --git a/backends/native/runtime/engine/ggml/GgmlLowerShape.cpp b/backends/native/runtime/engine/ggml/GgmlLowerShape.cpp index 2927f5b5dba..96382fd6e94 100644 --- a/backends/native/runtime/engine/ggml/GgmlLowerShape.cpp +++ b/backends/native/runtime/engine/ggml/GgmlLowerShape.cpp @@ -8,6 +8,7 @@ #include #include +#include #include #include #include @@ -294,6 +295,32 @@ void GraphBuilder::lower_expand(const Node& node) { set_single_output(node, TensorRef{result, std::move(sizes)}); } +void GraphBuilder::lower_repeat(const Node& node) { + const TensorRef input = tensor_arg(node, "self", 0); + if (input.tensor == nullptr) { + throw std::runtime_error("ggml: repeat input has no tensor"); + } + std::vector sizes = int_list(require_arg(node, "repeats", 1).arg); + if (sizes.size() < input.sizes.size() || sizes.size() > 4) { + throw std::runtime_error( + "ggml: repeat needs input rank <= repeats rank <= 4"); + } + const size_t offset = sizes.size() - input.sizes.size(); + for (size_t i = 0; i < sizes.size(); ++i) { + const int64_t extent = i < offset ? 1 : input.sizes[i - offset]; + if (extent <= 0 || sizes[i] <= 0 || + sizes[i] > std::numeric_limits::max() / extent) { + throw std::runtime_error("ggml: repeat requires positive, bounded sizes"); + } + sizes[i] *= extent; + } + ggml_context* ctx = instance_.context.get(); + ggml_tensor* shape = new_tensor(ctx, input.tensor->type, sizes); + ggml_tensor* result = ggml_repeat(ctx, input.tensor, shape); + // A no-op repeat must still snapshot storage that can be written later. + set_single_output(node, TensorRef{snapshot(result), std::move(sizes)}); +} + void GraphBuilder::lower_cat(const Node& node) { const Argument& tensors_arg = require_arg(node, "tensors", 0).arg; if (tensors_arg.kind() != ArgKind::TensorList) { diff --git a/backends/native/test/runtime/ggml/test_ggml_ops.py b/backends/native/test/runtime/ggml/test_ggml_ops.py index 54d31653a77..2a7169ffecd 100644 --- a/backends/native/test/runtime/ggml/test_ggml_ops.py +++ b/backends/native/test/runtime/ggml/test_ggml_ops.py @@ -884,6 +884,56 @@ def forward(self, indices): self._check(SqueezedEmbedding(), (torch.tensor([[1, 3, 9]]),)) + def test_embedding_arange_indices(self): + class PositionEmbedding(Embedding): + def forward(self, x): + return self.embedding(torch.arange(1, 10, 2)) + x + + self._check(PositionEmbedding(), (torch.randn(5, 8),)) + + def test_embedding_arange_indices_rejects_out_of_bounds(self): + class InvalidPositionEmbedding(Embedding): + def forward(self, x): + return self.embedding(torch.arange(-1, 4)) + x + + with self.assertRaisesRegex(AssertionError, "arange row index"): + # Export does not execute the invalid gather; the runner must reject it. + with tempfile.TemporaryDirectory() as tmp: + outdir = Path(tmp) + x = torch.randn(5, 8) + package = outdir / "model.ptn" + ep = torch.export.export(InvalidPositionEmbedding().eval(), (x,)) + to_native(ep).save(str(package)) + _run_ggml(package, [(x,)], outdir, 1) + + def test_repeated_row_gather_rechecks_indices(self): + class RepeatedRowGather(Embedding): + def forward(self, indices): + return self.embedding.weight[indices.unsqueeze(0).repeat(2, 1)] + + model = RepeatedRowGather().eval() + inputs = (torch.tensor([1, 3, 9]),) + self._check(model, inputs) + with tempfile.TemporaryDirectory() as tmp: + outdir = Path(tmp) + package = outdir / "model.ptn" + to_native(torch.export.export(model, inputs)).save(str(package)) + for invalid in (-1, 10): + with self.subTest(invalid=invalid): + with self.assertRaisesRegex(AssertionError, "row index is outside"): + _run_ggml( + package, + [inputs, (torch.tensor([1, invalid, 3]),)], + outdir, + 1, + ) + + def test_repeat_tiles_non_singleton_dimensions(self): + self._check( + Lambda(lambda x: x.repeat(2, 3, 2)), + (torch.arange(6, dtype=torch.float32).reshape(2, 3),), + ) + def test_rms_norm(self): self._check(RMSNorm(), (torch.randn(2, 3, 16),)) From 789b93c32ff530929f511e8d92dcd779ceeebad5 Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Tue, 3 Nov 2026 08:16:52 -0800 Subject: [PATCH 2/2] Add Whisper runner under examples/audio (#23904) Summary: Add `run_whisper` to encode a raw fp32 log-mel window, project cross K/V once, prefill the forced prefix into the static KV cache, and greedily decode while advancing `cache_position`. Cross K/V input staging persists across executions. The default token budget is 32, matching the exporter, and is capped by remaining decoder context. A full-capacity prefix stays unchanged, and decoding stops without computing an unused token at the budget limit. `--expected` compares the complete ID sequence with eager vectors and returns exit 3 on mismatch. `EXECUTORCH_BUILD_NATIVE_AUDIO_RUNNER` builds and installs the runner with Native Vulkan. Audio and vision runners share gflags resolution and `runner_utils`; either option enables the examples directory. The README documents matching export/run token budgets, and the Vulkan TODO records future zero-copy chaining and attention-weight storage work. Token files and EOS lists reject partially numeric or invalid IDs with exit 1; comma-only EOS lists are rejected. Reference IDs are validated before executing the model. Integer parsing accepts surrounding whitespace and CRLF input. Authored with Claude Code; updated with Codex. Differential Revision: D125296073 --- backends/native/CMakeLists.txt | 5 +- backends/native/README.md | 31 +- backends/native/examples/CMakeLists.txt | 3 + backends/native/examples/audio/BUCK | 27 ++ backends/native/examples/audio/CMakeLists.txt | 26 ++ .../native/examples/audio/WhisperRunner.cpp | 337 ++++++++++++++++++ .../native/examples/audio/WhisperRunner.h | 20 ++ .../examples/audio/run_whisper_ggml_cpu.cpp | 35 ++ backends/native/test/BUCK | 3 + backends/native/test/test_hf_whisper.py | 52 +++ tools/cmake/preset/default.cmake | 6 + 11 files changed, 543 insertions(+), 2 deletions(-) create mode 100644 backends/native/examples/audio/CMakeLists.txt create mode 100644 backends/native/examples/audio/WhisperRunner.cpp create mode 100644 backends/native/examples/audio/WhisperRunner.h create mode 100644 backends/native/examples/audio/run_whisper_ggml_cpu.cpp diff --git a/backends/native/CMakeLists.txt b/backends/native/CMakeLists.txt index 1a2911aab89..f285b3e6b9a 100644 --- a/backends/native/CMakeLists.txt +++ b/backends/native/CMakeLists.txt @@ -46,7 +46,10 @@ install( PATTERN "ggml" EXCLUDE ) -if(EXECUTORCH_BUILD_NATIVE_VISION_RUNNER OR EXECUTORCH_BUILD_NATIVE_LLM_RUNNER) +if(EXECUTORCH_BUILD_NATIVE_VISION_RUNNER + OR EXECUTORCH_BUILD_NATIVE_AUDIO_RUNNER + OR EXECUTORCH_BUILD_NATIVE_LLM_RUNNER +) add_subdirectory(examples) endif() diff --git a/backends/native/README.md b/backends/native/README.md index 9ce1b69d277..d894530affa 100644 --- a/backends/native/README.md +++ b/backends/native/README.md @@ -204,7 +204,7 @@ 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_cpu_engine_test` covers the CPU engine when `EXECUTORCH_BUILD_NATIVE_CPU` is enabled. `native_vulkan_passes_test` covers the -Vulkan graph rewrites without requiring a GPU. When either runner below is +Vulkan graph rewrites without requiring a GPU. When any runner below is enabled, `native_runner_utils_test` covers their shared runner utilities. Run these with `ctest --test-dir cmake-out -R native_`. @@ -310,6 +310,35 @@ and prints the selected engine and device and the top-5 logits. When invalid arguments, 2 for a loading or execution error, and 3 for a reference mismatch. +## Whisper speech recognition runner + +`EXECUTORCH_BUILD_NATIVE_AUDIO_RUNNER=ON` builds `run_whisper_ggml_cpu` from +`examples/audio` with Native GGML on CPU. It requires gflags. Initialize the +GGML submodule as described above, then: + +```sh +cmake -S . -B cmake-out -DCMAKE_BUILD_TYPE=Release \ + -DEXECUTORCH_BUILD_NATIVE=ON -DEXECUTORCH_BUILD_NATIVE_GGML=ON \ + -DEXECUTORCH_BUILD_NATIVE_AUDIO_RUNNER=ON +cmake --build cmake-out --target run_whisper_ggml_cpu +cmake --install cmake-out --prefix install +install/bin/run_whisper_ggml_cpu --cpu_threads 4 \ + --package whisper_tiny.ptn --input input_features.bin \ + --prompt prompt_ids.txt --expected expected_ids.txt --eos_ids 50257 +``` + +Export the `.ptn` with +[`export_hf_whisper.py`](examples/audio/export_hf_whisper.py); `--dump_dir` also +writes the log-mel input, forced prompt, and eager greedy ids, and the exporter +prints the model's EOS ids. The runner projects cross-attention K/V once, then +preserves those inputs and its self-attention cache across prefill and decode, +including GGML graph rebuilds. It greedily decodes until EOS or +`--max_new_tokens` (default 32, matching the exporter), capped by the remaining +decoder context. `--max_cached_graphs` defaults to 1; 0 rebuilds each call. +When `--expected` is given, it exits nonzero if the generated ids differ from +the reference. Use the same `--max_new_tokens` value when exporting reference +vectors and running the comparison. + ## Native LLM runner `EXECUTORCH_BUILD_NATIVE_LLM_RUNNER=ON` builds `run_hf_llm_vulkan` with Native diff --git a/backends/native/examples/CMakeLists.txt b/backends/native/examples/CMakeLists.txt index 2f6235040eb..77113aca47a 100644 --- a/backends/native/examples/CMakeLists.txt +++ b/backends/native/examples/CMakeLists.txt @@ -17,6 +17,9 @@ add_subdirectory(runner_utils) if(EXECUTORCH_BUILD_NATIVE_VISION_RUNNER) add_subdirectory(vision) endif() +if(EXECUTORCH_BUILD_NATIVE_AUDIO_RUNNER) + add_subdirectory(audio) +endif() if(EXECUTORCH_BUILD_NATIVE_LLM_RUNNER) add_subdirectory(hf_llm/runtime) endif() diff --git a/backends/native/examples/audio/BUCK b/backends/native/examples/audio/BUCK index f6e7ab8a11f..cc881522ad4 100644 --- a/backends/native/examples/audio/BUCK +++ b/backends/native/examples/audio/BUCK @@ -26,3 +26,30 @@ fbcode_target( "fbsource//third-party/pypi/safetensors:safetensors", ], ) + +runtime.cxx_library( + name = "whisper_runner", + srcs = ["WhisperRunner.cpp"], + exported_headers = ["WhisperRunner.h"], + deps = [ + "//executorch/backends/native/examples/runner_utils:runner_utils", + "//executorch/backends/native/runtime:runtime", + "//executorch/backends/native/runtime/deserialize:package", + "//executorch/backends/native/runtime/graph:scalar_type", + ], + exported_deps = ["//executorch/backends/native/runtime/engine:engine"], + external_deps = ["gflags"], + visibility = ["PUBLIC"], +) + +fbcode_target( + _kind = runtime.cxx_binary, + name = "run_whisper_ggml_cpu", + srcs = ["run_whisper_ggml_cpu.cpp"], + deps = [ + ":whisper_runner", + "//executorch/backends/native/runtime/engine:ggml_engine_cpu", + ], + external_deps = ["gflags"], + visibility = ["PUBLIC"], +) diff --git a/backends/native/examples/audio/CMakeLists.txt b/backends/native/examples/audio/CMakeLists.txt new file mode 100644 index 00000000000..5aabdeb5501 --- /dev/null +++ b/backends/native/examples/audio/CMakeLists.txt @@ -0,0 +1,26 @@ +# 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. + +add_library(native_whisper_runner STATIC WhisperRunner.cpp) +target_link_libraries( + native_whisper_runner + PUBLIC native_engine + PRIVATE ${_native_runner_gflags} native_package native_runner_utils + native_runtime +) + +add_executable(run_whisper_ggml_cpu run_whisper_ggml_cpu.cpp) +target_link_libraries( + run_whisper_ggml_cpu PRIVATE native_whisper_runner native_ggml_engine + ${_native_runner_gflags} +) +foreach(target native_whisper_runner run_whisper_ggml_cpu) + target_compile_options( + ${target} PRIVATE $<$:/EHsc> + $<$>:-fexceptions> + ) +endforeach() +install(TARGETS run_whisper_ggml_cpu RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR}) diff --git a/backends/native/examples/audio/WhisperRunner.cpp b/backends/native/examples/audio/WhisperRunner.cpp new file mode 100644 index 00000000000..39443b57895 --- /dev/null +++ b/backends/native/examples/audio/WhisperRunner.cpp @@ -0,0 +1,337 @@ +// 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. + +// Run a Whisper PTN (encoder + cross_kv + decoder methods, see +// export_hf_whisper) on a native engine: encode a raw fp32 log-mel +// window once, project cross-attention K/V once, prefill the forced +// decoder prefix into the static KV cache, then greedily decode one token +// per call, advancing cache_position. Cross K/V are set as decoder inputs +// once before prefill; input staging persists across execute() calls, so +// decode steps only set the token and its position. Optionally compare +// the id sequence with an eager reference. The decoder executable's cache +// buffers are process-lifetime state: one run decodes one utterance. +// Exit codes: 0 success, 1 invalid arguments, 2 execution error, 3 mismatch. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +DEFINE_string(package, "", "Path to the *.ptn to run. Required."); +DEFINE_string(input, "", "Log-mel input tensor, raw float32. Required."); +DEFINE_string(prompt, "", "Forced decoder prefix ids, one per line. Required."); +DEFINE_string( + expected, + "", + "Reference id sequence to compare against, one per line."); +DEFINE_string(eos_ids, "", "Comma-separated end-of-sequence ids. Required."); +DEFINE_string(encoder_method, "encoder", "Encoder method to run."); +DEFINE_string( + cross_kv_method, + "cross_kv", + "Cross-KV projection method to run."); +DEFINE_string(decoder_method, "decoder", "Decoder method to run."); +DEFINE_int64(max_new_tokens, 32, "Stop after generating this many tokens."); + +namespace { + +int64_t parse_id(const std::string& token) { + size_t consumed = 0; + const int64_t id = std::stoll(token, &consumed); + if (token.find_first_not_of(" \t\r\n", consumed) != std::string::npos) { + throw std::invalid_argument("not an integer id: " + token); + } + return id; +} + +std::unordered_set parse_id_list(const std::string& flag) { + std::unordered_set ids; + size_t start = 0; + while (start <= flag.size()) { + const size_t comma = flag.find(',', start); + const std::string token = + flag.substr(start, comma == std::string::npos ? comma : comma - start); + if (!token.empty()) { + ids.insert(parse_id(token)); + } + if (comma == std::string::npos) { + break; + } + start = comma + 1; + } + return ids; +} + +std::vector read_ids(const std::string& path) { + std::vector ids; + for (const std::string& line : ptn::runner_utils::read_lines(path)) { + if (!line.empty()) { + ids.push_back(parse_id(line)); + } + } + return ids; +} + +void print_ids(const char* label, const std::vector& ids) { + std::printf("%s (%zu):", label, ids.size()); + for (const int64_t id : ids) { + std::printf(" %lld", static_cast(id)); + } + std::printf("\n"); +} + +} // namespace + +namespace ptn { + +int run_whisper(const std::function()>& make_host) { + if (FLAGS_package.empty() || FLAGS_input.empty() || FLAGS_prompt.empty() || + FLAGS_eos_ids.empty() || FLAGS_max_new_tokens < 1) { + std::fprintf( + stderr, + "error: --package, --input, --prompt, --eos_ids and a positive " + "--max_new_tokens are required\n"); + return 1; + } + + std::unordered_set eos_ids; + std::vector ids; + std::vector expected; + try { + eos_ids = parse_id_list(FLAGS_eos_ids); + if (eos_ids.empty()) { + std::fprintf(stderr, "error: --eos_ids holds no ids\n"); + return 1; + } + ids = read_ids(FLAGS_prompt); + if (ids.empty()) { + std::fprintf(stderr, "error: --prompt holds no ids\n"); + return 1; + } + + if (!FLAGS_expected.empty()) { + expected = read_ids(FLAGS_expected); + } + } catch (const std::exception& e) { + std::fprintf(stderr, "error: invalid token ids: %s\n", e.what()); + return 1; + } + + try { + const auto package = + std::make_shared(ptn::Package::load(FLAGS_package)); + const ptn::ByteSpan program_bytes = package->program_bytes(); + const auto program = std::make_shared( + ptn::Program::load(program_bytes.data(), program_bytes.size())); + std::printf( + "package: %s (program %zu bytes, %zu constants / %zu bytes)\n", + FLAGS_package.c_str(), + program_bytes.size(), + package->keys().size(), + package->constant_bytes()); + + const std::unique_ptr host = make_host(); + if (!host) { + throw std::runtime_error("failed to create engine host"); + } + const std::unique_ptr context = + host->create_context(program, package); + if (!context) { + throw std::runtime_error("failed to create engine context"); + } + std::printf( + "engine: %s on %s\n", + host->name().c_str(), + host->device_name().c_str()); + + const std::unique_ptr encoder = + context->compile(FLAGS_encoder_method); + const std::unique_ptr cross_kv = + context->compile(FLAGS_cross_kv_method); + const std::unique_ptr decoder = + context->compile(FLAGS_decoder_method); + if (!encoder || !cross_kv || !decoder) { + throw std::runtime_error("failed to compile Whisper methods"); + } + std::printf( + "methods: %s (%zu in, %zu out), %s (%zu in, %zu out), " + "%s (%zu in, %zu out)\n", + FLAGS_encoder_method.c_str(), + encoder->num_inputs(), + encoder->num_outputs(), + FLAGS_cross_kv_method.c_str(), + cross_kv->num_inputs(), + cross_kv->num_outputs(), + FLAGS_decoder_method.c_str(), + decoder->num_inputs(), + decoder->num_outputs()); + std::printf("constants: %zu package bytes\n", package->constant_bytes()); + + const size_t num_cross = cross_kv->num_outputs(); + if (encoder->num_inputs() != 1 || encoder->num_outputs() != 1 || + cross_kv->num_inputs() != 1 || num_cross == 0 || num_cross % 2 != 0 || + decoder->num_inputs() != 2 + num_cross || decoder->num_outputs() != 1) { + std::fprintf( + stderr, + "error: expected encoder 1->1, cross_kv 1->%zu and decoder " + "%zu->1\n", + num_cross, + 2 + num_cross); + return 2; + } + bool dtypes_ok = encoder->input_dtype(0) == ptn::kFloat && + encoder->output_dtype(0) == ptn::kFloat && + cross_kv->input_dtype(0) == ptn::kFloat && + decoder->input_dtype(0) == ptn::kLong && + decoder->input_dtype(1) == ptn::kLong && + decoder->output_dtype(0) == ptn::kFloat; + for (size_t i = 0; i < num_cross; ++i) { + dtypes_ok = dtypes_ok && cross_kv->output_dtype(i) == ptn::kFloat && + decoder->input_dtype(2 + i) == ptn::kFloat && + cross_kv->output_sizes(i) == decoder->input_sizes(2 + i); + } + if (!dtypes_ok) { + std::fprintf( + stderr, + "error: expected fp32 mel/hidden/cross/logits and int64 " + "ids/cache_position, with cross shapes matching decoder inputs\n"); + return 2; + } + const int64_t cache_capacity = decoder->input_sizes(1).front(); + if (static_cast(ids.size()) > cache_capacity) { + std::fprintf( + stderr, + "error: prompt has %zu ids but the KV cache holds %lld\n", + ids.size(), + static_cast(cache_capacity)); + return 1; + } + + const std::vector mel_sizes = encoder->input_sizes(0); + const std::vector mel = ptn::runner_utils::read_floats(FLAGS_input); + if (static_cast(mel.size()) != + ptn::runner_utils::numel_of(mel_sizes)) { + std::fprintf( + stderr, + "error: input has %zu floats, the encoder wants %lld %s\n", + mel.size(), + static_cast(ptn::runner_utils::numel_of(mel_sizes)), + ptn::runner_utils::shape_str(mel_sizes).c_str()); + return 2; + } + encoder->set_input(0, mel.data(), mel.size(), ptn::kFloat); + encoder->execute(); + + const std::vector hidden_sizes = encoder->output_sizes(0); + std::vector hidden( + static_cast(ptn::runner_utils::numel_of(hidden_sizes))); + encoder->get_output(0, hidden.data(), hidden.size(), ptn::kFloat); + std::printf( + "encoder: %s -> %s\n", + ptn::runner_utils::shape_str(mel_sizes).c_str(), + ptn::runner_utils::shape_str(hidden_sizes).c_str()); + if (hidden_sizes != cross_kv->input_sizes(0)) { + std::fprintf( + stderr, + "error: encoder output %s does not match cross_kv input %s\n", + ptn::runner_utils::shape_str(hidden_sizes).c_str(), + ptn::runner_utils::shape_str(cross_kv->input_sizes(0)).c_str()); + return 2; + } + + // Project cross-attention K/V once and stage them as decoder inputs. + // Staging persists across execute() calls, so decode steps below only + // set the token id and cache_position. + cross_kv->set_input(0, hidden.data(), hidden.size(), ptn::kFloat); + cross_kv->execute(); + for (size_t i = 0; i < num_cross; ++i) { + std::vector cross(static_cast( + ptn::runner_utils::numel_of(cross_kv->output_sizes(i)))); + cross_kv->get_output(i, cross.data(), cross.size(), ptn::kFloat); + decoder->set_input(2 + i, cross.data(), cross.size(), ptn::kFloat); + } + std::printf( + "cross_kv: %zu tensors of %s\n", + num_cross, + ptn::runner_utils::shape_str(cross_kv->output_sizes(0)).c_str()); + + print_ids("prompt", ids); + const auto decode_step = [&decoder_executable = *decoder]( + const std::vector& window_ids, + int64_t first_position) { + std::vector positions(window_ids.size()); + for (size_t i = 0; i < positions.size(); ++i) { + positions[i] = first_position + static_cast(i); + } + decoder_executable.resize_input( + 0, {1, static_cast(window_ids.size())}); + decoder_executable.resize_input( + 1, {static_cast(positions.size())}); + decoder_executable.set_input( + 0, window_ids.data(), window_ids.size(), ptn::kLong); + decoder_executable.set_input( + 1, positions.data(), positions.size(), ptn::kLong); + decoder_executable.execute(); + + const std::vector logit_sizes = + decoder_executable.output_sizes(0); + const int64_t vocab = logit_sizes.back(); + std::vector logits( + static_cast(ptn::runner_utils::numel_of(logit_sizes))); + decoder_executable.get_output( + 0, logits.data(), logits.size(), ptn::kFloat); + const size_t last_row = logits.size() - static_cast(vocab); + const std::vector row( + logits.begin() + static_cast(last_row), logits.end()); + return static_cast(ptn::runner_utils::top_k(row, 1).front()); + }; + + const int64_t max_new_tokens = std::min( + FLAGS_max_new_tokens, + cache_capacity - static_cast(ids.size())); + if (max_new_tokens > 0) { + int64_t next = decode_step(ids, 0); + for (int64_t step = 0; step < max_new_tokens; ++step) { + ids.push_back(next); + if (eos_ids.count(next) != 0 || step + 1 == max_new_tokens) { + break; + } + next = decode_step({next}, static_cast(ids.size()) - 1); + } + } + print_ids("transcript", ids); + + if (FLAGS_expected.empty()) { + std::printf("result: ran (no reference given)\n"); + return 0; + } + const bool pass = ids == expected; + if (!pass) { + print_ids("expected", expected); + } + std::printf("result: %s\n", pass ? "PASS" : "FAIL"); + return pass ? 0 : 3; + } catch (const std::exception& e) { + std::fprintf(stderr, "error: %s\n", e.what()); + return 2; + } +} + +} // namespace ptn diff --git a/backends/native/examples/audio/WhisperRunner.h b/backends/native/examples/audio/WhisperRunner.h new file mode 100644 index 00000000000..3f7bc31377c --- /dev/null +++ b/backends/native/examples/audio/WhisperRunner.h @@ -0,0 +1,20 @@ +// 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. + +#pragma once + +#include +#include + +#include + +namespace ptn { + +// Runs encoder, cross_kv and cached decoder methods using parsed command-line +// flags and the engine supplied by the caller. Returns the process exit status. +int run_whisper(const std::function()>& make_host); + +} // namespace ptn diff --git a/backends/native/examples/audio/run_whisper_ggml_cpu.cpp b/backends/native/examples/audio/run_whisper_ggml_cpu.cpp new file mode 100644 index 00000000000..46cf9f536b8 --- /dev/null +++ b/backends/native/examples/audio/run_whisper_ggml_cpu.cpp @@ -0,0 +1,35 @@ +// 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. + +#include + +#include +#include +#include + +DEFINE_int32( + cpu_threads, + 0, + "ggml CPU threads; zero keeps the backend default."); +DEFINE_int32( + max_cached_graphs, + 1, + "Graphs kept per executable by input shapes; 0 rebuilds every call."); + +int main(int argc, char** argv) { + gflags::SetUsageMessage( + "Transcribe a Whisper .ptn with the native ggml CPU engine."); + gflags::ParseCommandLineFlags(&argc, &argv, true); + return ptn::run_whisper([] { + ptn::GgmlEngineOptions options; + options.cpu_threads = FLAGS_cpu_threads; + if (FLAGS_max_cached_graphs < 0) { + throw std::invalid_argument("--max_cached_graphs must be non-negative"); + } + options.max_cached_graphs = static_cast(FLAGS_max_cached_graphs); + return ptn::GgmlEngineHost::create(options); + }); +} diff --git a/backends/native/test/BUCK b/backends/native/test/BUCK index a1e520843ae..4b49f9400dd 100644 --- a/backends/native/test/BUCK +++ b/backends/native/test/BUCK @@ -92,6 +92,9 @@ fbcode_target( _kind = runtime.python_test, name = "test_hf_whisper", srcs = ["test_hf_whisper.py"], + resources = { + "//executorch/backends/native/examples/audio:run_whisper_ggml_cpu": "run_whisper_ggml_cpu", + }, deps = [ "//caffe2:torch", "//executorch/backends/native:lib", diff --git a/backends/native/test/test_hf_whisper.py b/backends/native/test/test_hf_whisper.py index 7e53594c9f7..d3810ae40eb 100644 --- a/backends/native/test/test_hf_whisper.py +++ b/backends/native/test/test_hf_whisper.py @@ -11,10 +11,15 @@ incremental decoding. Runtime operator support is covered by engine tests. """ +from __future__ import annotations + import operator import os +import subprocess import tempfile import unittest +from importlib.resources import as_file, files +from pathlib import Path import torch from executorch.backends.native import get_default_compile_config @@ -226,6 +231,53 @@ def test_exported_programs_match_eager(self): actual = actual[0] torch.testing.assert_close(actual, expected) + def test_ggml_runner_matches_eager(self): + with tempfile.TemporaryDirectory() as tmp: + outdir = Path(tmp) + package = outdir / "whisper.ptn" + to_native(self.programs).save(str(package)) + features = outdir / "input.bin" + features.write_bytes(self.features.contiguous().numpy().tobytes()) + prompt = outdir / "prompt.txt" + prefix = prompt_ids(self.model) + prompt.write_text("".join(f"{token}\n" for token in prefix)) + expected_path = outdir / "expected.txt" + for budget, cached_graphs in ((1, 1), (32, 0), (32, 1), (32, 2)): + with self.subTest(budget=budget, cached_graphs=cached_graphs): + expected = _greedy_transcript( + self.encoder, + self.cross_kv, + WhisperDecoder(self.model).eval(), + self.features, + prefix, + {2}, + budget, + ) + expected_path.write_text( + "".join(f"{token}\n" for token in expected) + ) + with as_file(files(__package__) / "run_whisper_ggml_cpu") as runner: + result = subprocess.run( + [ + str(runner), + f"--package={package}", + f"--input={features}", + f"--prompt={prompt}", + f"--expected={expected_path}", + "--eos_ids=2", + "--cpu_threads=2", + f"--max_new_tokens={budget}", + f"--max_cached_graphs={cached_graphs}", + ], + capture_output=True, + text=True, + timeout=120, + ) + self.assertEqual( + result.returncode, 0, result.stdout + result.stderr + ) + self.assertIn("result: PASS", result.stdout) + def test_three_method_package_saves(self): # Regression: each method lifts its own scalar constants under the # same auto-generated names; the package must still save. diff --git a/tools/cmake/preset/default.cmake b/tools/cmake/preset/default.cmake index 2e5436ff8b4..a0d52fb5bce 100644 --- a/tools/cmake/preset/default.cmake +++ b/tools/cmake/preset/default.cmake @@ -205,6 +205,11 @@ cmake_dependent_option( "EXECUTORCH_BUILD_NATIVE_CPU OR EXECUTORCH_BUILD_NATIVE_VULKAN OR EXECUTORCH_BUILD_NATIVE_GGML" OFF ) +cmake_dependent_option( + EXECUTORCH_BUILD_NATIVE_AUDIO_RUNNER + "Build the native GGML CPU Whisper speech recognition runner" OFF + "EXECUTORCH_BUILD_NATIVE_GGML" OFF +) cmake_dependent_option( EXECUTORCH_BUILD_NATIVE_LLM_RUNNER "Build the native LLM runners" OFF "EXECUTORCH_BUILD_NATIVE_VULKAN OR EXECUTORCH_BUILD_NATIVE_GGML" OFF @@ -213,6 +218,7 @@ announce_configured_options(EXECUTORCH_BUILD_NATIVE_CPU) announce_configured_options(EXECUTORCH_BUILD_NATIVE_VULKAN) announce_configured_options(EXECUTORCH_BUILD_NATIVE_GGML) announce_configured_options(EXECUTORCH_BUILD_NATIVE_VISION_RUNNER) +announce_configured_options(EXECUTORCH_BUILD_NATIVE_AUDIO_RUNNER) announce_configured_options(EXECUTORCH_BUILD_NATIVE_LLM_RUNNER) define_overridable_option( EXECUTORCH_BUILD_WEBGPU "Build the WebGPU backend" BOOL OFF