LLM runner: select I/O contracts and accept custom buffer allocators

Choose the registered IOManager and its logits, method names and step sizes, rejecting unknown contracts and invalid metadata. Load and reset the selected methods; route prompt tails to decode with a side-effect-free capacity check. Preserve module and I/O-buffer lifetimes and expose a backend-neutral allocator hook.

Author: psiddh

extension/llm/runner/constants.h

@@ -14,6 +14,9 @@ inline constexpr auto kEnableDynamicShape = "enable_dynamic_shape";
 inline constexpr auto kBosId = "get_bos_id";
 inline constexpr auto kEosIds = "get_eos_ids";
 inline constexpr auto kMaxSeqLen = "get_max_seq_len";
+// Names the I/O contract of the program's methods,
+// which selects the IOManager that drives them.
+inline constexpr auto kIoContract = "get_io_contract";
 inline constexpr auto kMaxContextLen = "get_max_context_len";
 inline constexpr auto kVocabSize = "get_vocab_size";
 inline constexpr auto kActivationDtype = "get_activation_dtype";

extension/llm/runner/io_manager/io_manager.h

@@ -8,6 +8,13 @@
 
 #pragma once
 
+#include <functional>
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+
 #include <executorch/extension/module/module.h>
 #include <executorch/extension/tensor/tensor.h>
 
@@ -236,13 +243,162 @@ class ET_EXPERIMENTAL IOManager {
     return update_decode(model_outputs, "forward");
   }
 
+  /**
+   * @brief The logits among a method's outputs, as the sampler reads them.
+   *
+   * The default takes a method that returns only its logits. An IOManager for
+   * a model that returns more (e.g. its updated cache) picks them out, and may
+   * hand back a converted copy (e.g. dequantized).
+   *
+   * @param model_outputs The outputs of the method just executed.
+   * @param method The method that produced them.
+   * @return The logits tensor.
+   */
+  virtual runtime::Result<executorch::aten::Tensor> logits(
+      const std::vector<runtime::EValue>& model_outputs,
+      const std::string& method) {
+    (void)method;
+    if (model_outputs.size() != 1 || !model_outputs[0].isTensor()) {
+      ET_LOG(
+          Error,
+          "Expected the method to return only its logits tensor, got %zu outputs.",
+          model_outputs.size());
+      return runtime::Error::InvalidState;
+    }
+    return model_outputs[0].toTensor();
+  }
+
+  /**
+   * @brief The most tokens one call of `method` takes, or -1 for no limit of
+   * the IOManager's own (e.g. a fixed-size prompt step).
+   */
+  virtual int64_t max_tokens_per_step(const std::string& method) const {
+    (void)method;
+    return -1;
+  }
+
+  virtual int64_t max_tokens_per_step(
+      const std::string& method,
+      int64_t start_pos) const {
+    (void)start_pos;
+    return max_tokens_per_step(method);
+  }
+
+  /**
+   * @brief For an encoder-decoder contract: route `encoder_method`'s outputs
+   * into the decode step's fixed inputs. Call once after loading, then after
+   * each run of the encoder pass its outputs to `update_encoder`.
+   */
+  ET_NODISCARD virtual runtime::Error load_encoder(
+      const std::string& encoder_method) {
+    (void)encoder_method;
+    return runtime::Error::NotSupported;
+  }
+
+  ET_NODISCARD virtual runtime::Error update_encoder(
+      const std::vector<runtime::EValue>& encoder_outputs) {
+    (void)encoder_outputs;
+    return runtime::Error::NotSupported;
+  }
+
+  /**
+   * @brief The prefill and decode method names this IOManager drives, given
+   * the methods the program has. Empty strings leave the runner's defaults.
+   */
+  virtual std::pair<std::string, std::string> method_names(
+      const std::unordered_set<std::string>& available) const {
+    (void)available;
+    return {"", ""};
+  }
+
  private:
   /**
    * @brief Reference to the Module used for method metadata and execution.
    */
   ET_MODULE_NAMESPACE::Module& module_;
 };
 
+using IOManagerFactory =
+    std::function<std::unique_ptr<IOManager>(ET_MODULE_NAMESPACE::Module&)>;
+
+/**
+ * @brief IOManagers by the I/O contract they speak.
+ *
+ * A program names its contract in the `get_io_contract` constant method; the
+ * runner builds the IOManager registered for it, or the default IOManager when
+ * the program names none.
+ */
+inline std::unordered_map<std::string, IOManagerFactory>&
+io_manager_registry() {
+  static std::unordered_map<std::string, IOManagerFactory> registry;
+  return registry;
+}
+
+inline bool register_io_manager(
+    const std::string& contract,
+    IOManagerFactory factory) {
+  return io_manager_registry().emplace(contract, std::move(factory)).second;
+}
+
+inline runtime::Result<std::unique_ptr<IOManager>> make_io_manager(
+    ET_MODULE_NAMESPACE::Module& module,
+    const std::string& contract) {
+  if (!contract.empty()) {
+    auto it = io_manager_registry().find(contract);
+    if (it != io_manager_registry().end()) {
+      auto manager = it->second(module);
+      if (!manager) {
+        return runtime::Error::InvalidState;
+      }
+      return manager;
+    }
+    ET_LOG(
+        Error,
+        "No IOManager registered for I/O contract '%s'.",
+        contract.c_str());
+    return runtime::Error::NotSupported;
+  }
+  return std::make_unique<IOManager>(module);
+}
+
+/// Registers the IOManagers this library implements. Idempotent.
+void register_builtin_io_managers();
+
+/**
+ * @brief Memory for the I/O buffers an IOManager owns.
+ *
+ * A backend whose device reads graph I/O from memory it can map (for example
+ * QNN's rpcmem) registers an allocator, so the device uses these buffers in
+ * place instead of copying them on every call. Without one, buffers come from
+ * the heap.
+ */
+class ET_EXPERIMENTAL IOBufferAllocator {
+ public:
+  virtual ~IOBufferAllocator() = default;
+  /// A block of at least `nbytes`, or nullptr to fall back to the heap.
+  virtual void* allocate(size_t nbytes, size_t alignment) = 0;
+  virtual void deallocate(void* block) = 0;
+  /// Declares a tensor whose data starts at `tensor_data` inside `block`.
+  virtual void bind(void* block, void* tensor_data) {
+    (void)block;
+    (void)tensor_data;
+  }
+};
+
+using IOBufferAllocatorFactory =
+    std::function<std::unique_ptr<IOBufferAllocator>()>;
+
+inline IOBufferAllocatorFactory& io_buffer_allocator_factory() {
+  static IOBufferAllocatorFactory factory;
+  return factory;
+}
+
+/// Sets the allocator IOManagers use for their I/O buffers. The last call wins.
+inline bool set_io_buffer_allocator(IOBufferAllocatorFactory factory) {
+  io_buffer_allocator_factory() = std::move(factory);
+  return true;
+}
+
 } // namespace llm
 } // namespace extension
 } // namespace executorch

extension/llm/runner/llm_runner_helper.cpp

@@ -10,6 +10,7 @@
 
 #include <executorch/extension/llm/runner/image_prefiller.h>
 #include <executorch/extension/llm/runner/llm_runner_helper.h>
+
 #include <executorch/extension/llm/runner/multimodal_decoder_runner.h>
 #include <executorch/extension/llm/runner/multimodal_prefiller.h>
 #include <executorch/extension/llm/runner/multimodal_runner.h>
@@ -26,11 +27,15 @@
 #include <pytorch/tokenizers/tekken.h>
 #include <pytorch/tokenizers/tiktoken.h>
 
+#include <limits>
+
 namespace executorch::extension::llm {
 
 using ::executorch::extension::Module;
 using ::executorch::runtime::Error;
 
+void register_builtin_io_managers() {}
+
 std::unique_ptr<tokenizers::Tokenizer> load_tokenizer(
     const std::string& tokenizer_path,
     std::unique_ptr<std::vector<std::string>> special_tokens,
@@ -126,6 +131,13 @@ get_llm_metadata(tokenizers::Tokenizer* tokenizer, Module* module) {
       if (!get_result.ok()) {
         return get_result.error();
       }
+      if (!get_result->isInt() && !get_result->isBool()) {
+        ET_LOG(
+            Error,
+            "Metadata method %s must return an integer or boolean.",
+            method_name.c_str());
+        return Error::InvalidArgument;
+      }
       value = get_result->toScalar().to<decltype(metadata)::mapped_type>();
     } else {
       ET_LOG(
@@ -148,6 +160,15 @@ get_llm_metadata(tokenizers::Tokenizer* tokenizer, Module* module) {
         metadata[llm::kMaxContextLen]);
   }
 
+  for (const auto* name : {kMaxSeqLen, kMaxContextLen}) {
+    const auto value = metadata.at(name);
+    ET_CHECK_OR_RETURN_ERROR(
+        value > 0 && value <= std::numeric_limits<int32_t>::max(),
+        InvalidArgument,
+        "Metadata %s must be in [1, INT32_MAX].",
+        name);
+  }
+
   // Set tokenizer-related metadata
   metadata[llm::kBosId] = tokenizer->bos_tok();
   metadata[llm::kVocabSize] = tokenizer->vocab_size();
@@ -266,8 +287,32 @@ std::unique_ptr<TextLLMRunner> create_text_llm_runner(
   auto eos_ids = std::make_unique<std::unordered_set<uint64_t>>(
       llm::get_eos_ids(tokenizer.get(), module.get()));
 
-  // Create IOManager
-  std::unique_ptr<IOManager> io_manager = std::make_unique<IOManager>(*module);
+  // Create the IOManager for the I/O contract the program declares, if any.
+  register_builtin_io_managers();
+  std::string io_contract;
+  if (auto names = module->method_names();
+      names.ok() && names->count(kIoContract)) {
+    auto contract = module->get(kIoContract);
+    if (!contract.ok() || !contract->isString()) {
+      ET_LOG(Error, "Failed to read the program's I/O contract label.");
+      return nullptr;
+    }
+    io_contract = std::string(contract->toString());
+    if (io_contract.empty()) {
+      ET_LOG(Error, "The program's I/O contract label is empty.");
+      return nullptr;
+    }
+    ET_LOG(Info, "I/O contract: %s", io_contract.c_str());
+  }
+  auto io_manager_result = make_io_manager(*module, io_contract);
+  if (!io_manager_result.ok()) {
+    ET_LOG(
+        Error,
+        "Failed to create the IOManager: 0x%x.",
+        static_cast<unsigned>(io_manager_result.error()));
+    return nullptr;
+  }
+  auto io_manager = std::move(*io_manager_result);
 
   // Read vocab_size for Sampler
   int32_t vocab_size = static_cast<int32_t>(metadata.at(kVocabSize));
@@ -289,6 +334,17 @@ std::unique_ptr<TextLLMRunner> create_text_llm_runner(
       method_names.count(kPrefillMethod) > 0) {
     prefill_method_name = kPrefillMethod;
   }
+  // An IOManager may drive methods of its own contract's names.
+  auto [io_prefill, io_decode] = io_manager->method_names(method_names);
+  if (!io_prefill.empty()) {
+    prefill_method_name = io_prefill;
+  }
+  if (io_manager->load(
+          prefill_method_name,
+          io_decode.empty() ? prefill_method_name : io_decode) != Error::Ok) {
+    ET_LOG(Error, "Failed to load the IOManager");
+    return nullptr;
+  }
 
   // Create text_decoder_runner
   ET_LOG(Info, "Using method: %s", prefill_method_name.c_str());
@@ -302,7 +358,15 @@ std::unique_ptr<TextLLMRunner> create_text_llm_runner(
   // The decode stage gets its own decoder runner when the PTE exports a
   // separate decode method; otherwise both stages share text_decoder_runner.
   std::unique_ptr<TextDecoderRunner> decode_text_decoder_runner;
-  if (method_names.count(kDecodeMethod) > 0 &&
+  if (!io_decode.empty() && io_decode != prefill_method_name) {
+    decode_text_decoder_runner = std::make_unique<TextDecoderRunner>(
+        module.get(),
+        io_manager.get(),
+        io_decode,
+        std::make_unique<Sampler>(vocab_size, init_temp),
+        stats.get());
+  } else if (
+      io_decode.empty() && method_names.count(kDecodeMethod) > 0 &&
       prefill_method_name != kDecodeMethod) {
     ET_LOG(Info, "Using decode method: %s", kDecodeMethod);
     decode_text_decoder_runner = std::make_unique<TextDecoderRunner>(
@@ -316,12 +380,17 @@ std::unique_ptr<TextLLMRunner> create_text_llm_runner(
       ? decode_text_decoder_runner.get()
       : text_decoder_runner.get();
 
-  // Create text_prefiller
+  // Create text_prefiller. A prompt method with a fixed step size takes the
+  // prompt in chunks of that size.
+  const int64_t step_tokens =
+      io_manager->max_tokens_per_step(prefill_method_name);
   auto text_prefiller = std::make_unique<TextPrefiller>(
       text_decoder_runner.get(),
       metadata.at(kUseKVCache),
-      metadata.at(kEnableDynamicShape),
-      metadata.at(kMaxSeqLen));
+      step_tokens > 1 ? true
+                      : static_cast<bool>(metadata.at(kEnableDynamicShape)),
+      step_tokens > 1 ? step_tokens : metadata.at(kMaxSeqLen),
+      step_tokens > 1 ? decode_runner : nullptr);
 
   // Create text_token_generator with stats
   auto text_token_generator = std::make_unique<TextTokenGenerator>(

extension/llm/runner/test/CMakeLists.txt

@@ -19,6 +19,7 @@ include(${EXECUTORCH_ROOT}/tools/cmake/Test.cmake)
 
 set(_test_srcs
     test_generation_config.cpp
+    test_io_manager.cpp
     test_model_metadata.cpp
     test_text_llm_runner.cpp
     test_text_prefiller.cpp
@@ -49,14 +50,21 @@ set(_metadata_pte_files
     "${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_geometry_invalid_windows.pte"
     "${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_geometry_empty.pte"
 )
+# Local builds may share a host with other model exports.
+set(_export_prefix)
+if(EXECUTORCH_TEST_EXPORT_LOCK)
+  find_program(_flock flock REQUIRED)
+  set(_export_prefix ${_flock} ${EXECUTORCH_TEST_EXPORT_LOCK})
+endif()
+
 add_custom_command(
   OUTPUT ${_metadata_pte_files}
   # Run against the installed executorch (which stages exir/_serialize/*.fbs at
   # install time). Do NOT prepend the source tree to PYTHONPATH: it shadows the
   # installed package with a source copy that lacks the staged .fbs, breaking
   # to_executorch() under a non-editable install.
   COMMAND
-    ${PYTHON_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/export_model_metadata.py
+    ${_export_prefix} ${PYTHON_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/export_model_metadata.py
     --outdir ${CMAKE_CURRENT_BINARY_DIR}
   WORKING_DIRECTORY ${EXECUTORCH_ROOT}
   DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/export_model_metadata.py

extension/llm/runner/test/targets.bzl

@@ -27,6 +27,14 @@ def define_common_targets():
         ],
     )
 
+    runtime.cxx_test(
+        name = "test_io_manager",
+        srcs = ["test_io_manager.cpp"],
+        deps = [
+            "//executorch/extension/llm/runner/io_manager:io_manager",
+        ],
+    )
+
     runtime.cxx_test(
         name = "test_text_prefiller",
         srcs = ["test_text_prefiller.cpp"],