Update

Author: larryliu0820

backends/cuda/cuda_partitioner.py

@@ -16,6 +16,7 @@
     PartitionResult,
 )
 from executorch.exir.backend.utils import tag_constant_data, tag_mutated_buffer
+from torch._export.utils import is_buffer, is_lifted_tensor_constant, is_param
 from torch.export.exported_program import ExportedProgram
 
 
@@ -56,6 +57,18 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
         tag_constant_data(exported_program)
         tag_mutated_buffer(exported_program)
 
+        # Tag constant placeholders that have no users
+        # tag_constant_data only tags constants that have users with delegation_tag
+        # but we need to tag all constants for this partition
+        for node in exported_program.graph.nodes:
+            if node.op == "placeholder" and (
+                is_param(exported_program, node)
+                or is_buffer(exported_program, node)
+                or is_lifted_tensor_constant(exported_program, node)
+            ):
+                if "delegation_tag" not in node.meta:
+                    node.meta["delegation_tag"] = tag
+
         return PartitionResult(
             tagged_exported_program=exported_program, partition_tags=partition_tags
         )

examples/models/whisper/CMakeLists.txt

@@ -0,0 +1,101 @@
+# 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.
+
+cmake_minimum_required(VERSION 3.24)
+project(whisper_runner)
+
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+
+set(EXECUTORCH_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/../../..")
+include(${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake)
+
+# Let files say "include <executorch/path/to/header.h>"
+set(_common_include_directories ${EXECUTORCH_ROOT}/..)
+
+# Need this for gflags for some reason
+set(gflags_DIR ${CMAKE_CURRENT_BINARY_DIR}/../../../third-party/gflags)
+find_package(gflags REQUIRED)
+
+list(APPEND CMAKE_FIND_ROOT_PATH ${CMAKE_CURRENT_BINARY_DIR}/../../..)
+find_package(executorch CONFIG REQUIRED FIND_ROOT_PATH_BOTH)
+executorch_target_link_options_shared_lib(executorch)
+
+set(link_libraries executorch gflags)
+set(_srcs multimodal.cpp)
+
+list(
+  APPEND
+  link_libraries
+  optimized_native_cpu_ops_lib
+  quantized_ops_lib
+  custom_ops
+  cpublas
+  eigen_blas
+)
+executorch_target_link_options_shared_lib(optimized_native_cpu_ops_lib)
+executorch_target_link_options_shared_lib(quantized_ops_lib)
+executorch_target_link_options_shared_lib(custom_ops)
+
+# XNNPACK
+if(TARGET xnnpack_backend)
+  set(xnnpack_backend_libs xnnpack_backend XNNPACK xnnpack-microkernels-prod)
+  if(TARGET kleidiai)
+    list(APPEND xnnpack_backend_libs kleidiai)
+  endif()
+  list(APPEND link_libraries ${xnnpack_backend_libs})
+  executorch_target_link_options_shared_lib(xnnpack_backend)
+endif()
+
+# Add LLM runner and extension module
+if(NOT TARGET extension_llm_runner)
+  message(
+    FATAL_ERROR
+      "ExecuTorch must be installed with EXECUTORCH_BUILD_EXTENSION_LLM_RUNNER enabled."
+  )
+endif()
+
+# Needed for cpuinfo where it uses android specific log lib
+if(ANDROID)
+  list(APPEND link_libraries log)
+endif()
+
+# Add the required ExecuTorch extensions for multimodal LLM runner
+list(
+  APPEND
+  link_libraries
+  extension_llm_runner
+  extension_module
+  extension_data_loader
+  extension_tensor
+  extension_flat_tensor
+)
+
+# Link CUDA backend
+if(EXECUTORCH_BUILD_CUDA)
+  find_package(CUDAToolkit REQUIRED)
+  list(APPEND link_libraries aoti_cuda)
+  executorch_target_link_options_shared_lib(aoti_cuda)
+endif()
+
+if(EXECUTORCH_BUILD_METAL)
+  list(APPEND link_libraries metal_backend)
+  executorch_target_link_options_shared_lib(metal_backend)
+endif()
+
+# Add tokenizers
+list(APPEND link_libraries tokenizers::tokenizers)
+
+add_executable(whisper_runner runner.cpp main.cpp)
+
+target_include_directories(whisper_runner PUBLIC ${_common_include_directories})
+
+target_link_libraries(
+  whisper_runner
+  PUBLIC
+  ${link_libraries}
+)
+target_compile_options(whisper_runner PUBLIC ${_common_compile_options})

examples/models/whisper/README.md

@@ -0,0 +1,69 @@
+# Whisper Runner
+
+This directory hosts a lightweight C++ helper that drives Whisper models
+exported to ExecuTorch. The `WhisperRunner` owns the `Module` instance that
+wraps a bundled `.pte` program and optional `.ptd` weight file, loads the
+`encoder` and `text_decoder` methods, and exposes a `transcribe()` loop that
+streams decoded text pieces through a callback.
+
+The runner assumes:
+- `model.pte` contains both Whisper encoder and decoder entry points named
+  `encoder` and `text_decoder`.
+- External parameters (for example KV cache blocks) are stored in a companion
+  `model.ptd`.
+- A tokenizer JSON compatible with the ExecuTorch tokenizers shim is available.
+
+Audio preprocessing is not part of the runner itself. To transform raw audio
+into the mel features expected by the encoder, reuse the pattern in
+`examples/models/voxtral/multimodal.cpp`, which loads a `preprocessor.pte`
+module to generate the spectrogram tensor.
+
+## Build
+
+```bash
+cmake -G Ninja \
+  -B cmake-out/examples/models/whisper \
+  -S examples/models/whisper
+cmake --build cmake-out/examples/models/whisper -j
+```
+
+The build produces a static library named `whisper_runner`. Link it into your
+application together with the standard ExecuTorch runtime libraries and the
+tokenizer target (`tokenizers::tokenizers`).
+
+## Usage
+
+```cpp
+#include <executorch/examples/models/whisper/runner.h>
+#include <executorch/extension/tensor/tensor_ptr.h>
+
+using example::WhisperRunner;
+using example::WhisperTranscribeConfig;
+
+WhisperRunner runner("model.pte", "model.ptd", "tokenizer.json");
+ET_CHECK_OK(runner.load());
+
+// `features` is the mel spectrogram tensor produced by the preprocessor.
+executorch::aten::Tensor features = load_features_somehow();
+
+WhisperTranscribeConfig config;
+config.max_new_tokens = 128; // stop after 128 generated tokens
+config.temperature = 0.7f;  // optional: enable stochastic sampling
+
+auto tokens_result = runner.transcribe(
+    features,
+    config,
+    [](const std::string& piece) {
+      std::cout << piece;
+    });
+
+if (!tokens_result.ok()) {
+  ET_LOG(Error, "Transcription failed: %d", static_cast<int>(tokens_result.error()));
+}
+```
+
+`transcribe()` returns the full token history (prompt + generated tokens) and
+invokes the callback every time a new token is emitted. Provide a non-empty
+`decoder_input_ids` vector if you want to seed the decoder with a custom prompt,
+and override `WhisperTranscribeConfig::eos_token_ids` when the model exposes
+custom termination ids.

examples/models/whisper/main.cpp

@@ -0,0 +1,194 @@
+/*
+ * 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 <fstream>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include <gflags/gflags.h>
+
+#include <executorch/examples/models/whisper/runner.h>
+#include <executorch/extension/llm/runner/util.h>
+#include <executorch/extension/llm/runner/wav_loader.h>
+#include <executorch/extension/module/module.h>
+#include <executorch/extension/tensor/tensor_ptr_maker.h>
+#include <executorch/runtime/core/evalue.h>
+#include <executorch/runtime/platform/log.h>
+
+DEFINE_string(model_path, "model.pte", "Path to Whisper model (.pte).");
+DEFINE_string(data_path, "", "Optional path to Whisper weights (.ptd).");
+DEFINE_string(
+    tokenizer_path,
+    ".",
+    "Path to tokenizer directory containing tokenizer.json, tokenizer_config.json, and special_tokens_map.json.");
+DEFINE_string(
+    preprocessor_path,
+    "",
+    "Path to preprocessor .pte for converting raw audio.");
+DEFINE_string(
+    audio_path,
+    "",
+    "Path to input audio file. Accepts .wav or raw float .bin.");
+DEFINE_string(
+    preprocessed_audio_path,
+    "",
+    "Path to preprocessed audio features file (.bin). If provided, skips preprocessing.");
+DEFINE_double(
+    temperature,
+    0.0,
+    "Sampling temperature. 0.0 performs greedy decoding.");
+DEFINE_int32(max_new_tokens, 128, "Maximum number of tokens to generate.");
+
+namespace {
+
+using ::executorch::extension::from_blob;
+using ::executorch::extension::Module;
+
+bool ends_with(const std::string& value, const std::string& suffix) {
+  return value.size() >= suffix.size() &&
+      value.compare(value.size() - suffix.size(), suffix.size(), suffix) == 0;
+}
+
+std::vector<float> load_preprocessed_audio(
+    const std::string& preprocessed_audio_path) {
+  ET_LOG(
+      Info,
+      "Loading preprocessed audio from: %s",
+      preprocessed_audio_path.c_str());
+
+  std::ifstream stream(
+      preprocessed_audio_path, std::ios::binary | std::ios::ate);
+  if (!stream.is_open()) {
+    ET_LOG(
+        Error,
+        "Failed to open preprocessed audio file: %s",
+        preprocessed_audio_path.c_str());
+    throw std::runtime_error("Failed to open preprocessed audio file");
+  }
+
+  std::size_t byte_size = static_cast<std::size_t>(stream.tellg());
+  stream.seekg(0, std::ios::beg);
+
+  const int64_t batch_size = 1;
+  const int64_t feature_dim = 128;
+  const int64_t time_steps = 3000;
+  const int64_t expected_elements = batch_size * feature_dim * time_steps;
+  const std::size_t expected_bytes = expected_elements * sizeof(float);
+
+  if (byte_size != expected_bytes) {
+    ET_LOG(
+        Error,
+        "Preprocessed audio file size mismatch. Expected %zu bytes, got %zu bytes",
+        expected_bytes,
+        byte_size);
+    throw std::runtime_error("Preprocessed audio file size mismatch");
+  }
+
+  std::vector<float> feature_data(expected_elements);
+  stream.read(reinterpret_cast<char*>(feature_data.data()), byte_size);
+  stream.close();
+
+  return feature_data;
+}
+
+} // namespace
+
+int main(int argc, char** argv) {
+  gflags::ParseCommandLineFlags(&argc, &argv, true);
+
+  ::executorch::extension::TensorPtr features;
+  std::vector<float> audio_data;
+  std::unique_ptr<Module> processor;
+
+  if (!FLAGS_preprocessed_audio_path.empty()) {
+    audio_data = load_preprocessed_audio(FLAGS_preprocessed_audio_path);
+
+    const int64_t batch_size = 1;
+    const int64_t feature_dim = 128;
+    const int64_t time_steps = 3000;
+    features = from_blob(
+        audio_data.data(),
+        /*sizes=*/{batch_size, feature_dim, time_steps},
+        /*strides=*/{feature_dim * time_steps, feature_dim, 1},
+        ::executorch::aten::ScalarType::Float);
+  } else {
+    // Original preprocessing path
+    if (FLAGS_audio_path.empty()) {
+      ET_LOG(
+          Error,
+          "Either audio_path or preprocessed_audio_path flag must be provided.");
+      return 1;
+    }
+
+    audio_data =
+        executorch::extension::llm::load_wav_audio_data(FLAGS_audio_path);
+    ET_LOG(
+        Info,
+        "First 2 values of audio data: %f, %f",
+        audio_data[0],
+        audio_data[1]);
+    // Preprocess audio
+    processor = std::make_unique<Module>(
+        FLAGS_preprocessor_path, Module::LoadMode::Mmap);
+    auto load_error = processor->load();
+    if (load_error != ::executorch::runtime::Error::Ok) {
+      ET_LOG(Error, "Failed to load preprocessor module.");
+      return 1;
+    }
+
+    auto audio_tensor = from_blob(
+        audio_data.data(),
+        {static_cast<::executorch::aten::SizesType>(audio_data.size())},
+        ::executorch::aten::ScalarType::Float);
+
+    auto processed_result = processor->execute("forward", audio_tensor);
+    if (processed_result.error() != ::executorch::runtime::Error::Ok) {
+      ET_LOG(Error, "Audio preprocessing failed.");
+      return 1;
+    }
+    auto outputs = std::move(processed_result.get());
+    if (outputs.empty() || !outputs[0].isTensor()) {
+      ET_LOG(Error, "Preprocessor returned unexpected outputs.");
+      return 1;
+    }
+    auto tensor = outputs[0].toTensor();
+    ET_LOG(
+        Info,
+        "Result scalar_type: %s, first value %f",
+        ::executorch::runtime::toString(tensor.scalar_type()),
+        tensor.mutable_data_ptr<float>()[0]);
+    features = std::make_shared<::executorch::aten::Tensor>(std::move(tensor));
+  }
+
+  example::WhisperRunner runner(
+      FLAGS_model_path, FLAGS_data_path, FLAGS_tokenizer_path);
+  auto load_err = runner.load();
+  if (load_err != ::executorch::runtime::Error::Ok) {
+    ET_LOG(Error, "Failed to load Whisper model.");
+    return 1;
+  }
+
+  example::WhisperTranscribeConfig config;
+  config.max_new_tokens = FLAGS_max_new_tokens;
+  config.temperature = static_cast<float>(FLAGS_temperature);
+
+  std::string transcript;
+  auto result =
+      runner.transcribe(features, config, [&](const std::string& piece) {
+        ::executorch::extension::llm::safe_printf(piece.c_str());
+        fflush(stdout);
+      });
+
+  if (!result.ok()) {
+    ET_LOG(Error, "Transcription failed.");
+    return 1;
+  }
+
+  return 0;
+}