Change prefill() to return Result<uint64_t>

prefill() now returns the sampled next token directly, making the
API more natural. Callers get the token if they want it, and
internally it's also stored in prefill_next_token_ for the
generate("") workflow.

This eliminates the need for the private prefill_and_sample() helper
since prefill() itself returns the token — generate(vector) can
call prefill() directly.

This PR was authored with the assistance of Claude.

Author: kirklandsign

extension/llm/runner/irunner.h

@@ -18,6 +18,7 @@
 
 #include <executorch/extension/llm/runner/stats.h>
 #include <executorch/runtime/core/error.h>
+#include <executorch/runtime/core/result.h>
 
 namespace executorch {
 namespace extension {
@@ -140,7 +141,7 @@ class ET_EXPERIMENTAL IRunner {
    * @param num_eos Number of EOS tokens to append during encoding
    * @return Error::Ok if successful, an error otherwise
    */
-  virtual runtime::Error prefill(
+  virtual runtime::Result<uint64_t> prefill(
       const std::vector<MultimodalInput>& inputs,
       int32_t num_bos = 0,
       int32_t num_eos = 0) {

extension/llm/runner/multimodal_runner.cpp

@@ -85,10 +85,13 @@ Error MultimodalRunner::load() {
     ET_LOG(Info, format, __VA_ARGS__);     \
   }
 
-Result<uint64_t> MultimodalRunner::prefill_and_sample(
+Result<uint64_t> MultimodalRunner::prefill(
     const std::vector<MultimodalInput>& inputs,
     int32_t num_bos,
     int32_t num_eos) {
+  if (!is_loaded()) {
+    ET_CHECK_OK_OR_RETURN_ERROR(load());
+  }
   uint64_t last_token = 0;
   for (size_t i = 0; i < inputs.size(); ++i) {
     const auto& input = inputs[i];
@@ -119,24 +122,10 @@ Result<uint64_t> MultimodalRunner::prefill_and_sample(
     }
     last_token = prefill_result.get();
   }
+  prefill_next_token_ = last_token;
   return last_token;
 }
 
-Error MultimodalRunner::prefill(
-    const std::vector<MultimodalInput>& inputs,
-    int32_t num_bos,
-    int32_t num_eos) {
-  if (!is_loaded()) {
-    ET_CHECK_OK_OR_RETURN_ERROR(load());
-  }
-  auto result = prefill_and_sample(inputs, num_bos, num_eos);
-  if (!result.ok()) {
-    return result.error();
-  }
-  prefill_next_token_ = result.get();
-  return Error::Ok;
-}
-
 Error MultimodalRunner::generate(
     const std::string& prompt,
     const GenerationConfig& config,
@@ -282,8 +271,7 @@ Error MultimodalRunner::generate(
   }
 
   // Prefill all inputs and get the first decode token
-  auto prefill_result =
-      prefill_and_sample(inputs, config.num_bos, config.num_eos);
+  auto prefill_result = prefill(inputs, config.num_bos, config.num_eos);
   ET_CHECK_OK_OR_RETURN_ERROR(prefill_result.error());
   uint64_t cur_token = prefill_result.get();
 

extension/llm/runner/multimodal_runner.h

@@ -147,17 +147,18 @@ class ET_EXPERIMENTAL MultimodalRunner : public IRunner {
    * text.
    * @param num_bos Number of BOS tokens to prepend during encoding.
    * @param num_eos Number of EOS tokens to append during encoding.
-   * @return The error code. KV cache position is tracked internally in pos_.
+   * @return The next token predicted after prefill, or an error.
+   *         KV cache position is tracked internally in pos_.
    */
-  ::executorch::runtime::Error prefill(
+  ::executorch::runtime::Result<uint64_t> prefill(
       const std::vector<MultimodalInput>& inputs,
       int32_t num_bos = 0,
       int32_t num_eos = 0) override;
 
   /**
    * Convenience overload: prefill a single text prompt.
    */
-  ::executorch::runtime::Error
+  ::executorch::runtime::Result<uint64_t>
   prefill(const std::string& prompt, int32_t num_bos = 0, int32_t num_eos = 0) {
     std::vector<MultimodalInput> inputs;
     inputs.emplace_back(MultimodalInput(prompt));
@@ -195,12 +196,6 @@ class ET_EXPERIMENTAL MultimodalRunner : public IRunner {
   // Internal state
   std::optional<uint64_t> prefill_next_token_;
   int64_t pos_;
-
- private:
-  ::executorch::runtime::Result<uint64_t> prefill_and_sample(
-      const std::vector<MultimodalInput>& inputs,
-      int32_t num_bos,
-      int32_t num_eos);
 };
 
 } // namespace llm

extension/llm/runner/pybindings.cpp

@@ -121,8 +121,8 @@ class PyTextLLMRunner {
     }
     {
       py::gil_scoped_release release;
-      Error error = runner_->prefill(prompt, config.num_bos, config.num_eos);
-      THROW_IF_ERROR(error, "Prefill failed");
+      auto result = runner_->prefill(prompt, config.num_bos, config.num_eos);
+      THROW_IF_ERROR(result.error(), "Prefill failed");
     }
   }
 
@@ -234,8 +234,8 @@ class PyMultimodalRunner {
     }
     {
       py::gil_scoped_release release;
-      Error error = runner_->prefill(inputs, 0, 0);
-      THROW_IF_ERROR(error, "Prefill failed");
+      auto result = runner_->prefill(inputs, 0, 0);
+      THROW_IF_ERROR(result.error(), "Prefill failed");
     }
   }
 

extension/llm/runner/text_llm_runner.cpp

@@ -257,7 +257,7 @@ Error TextLLMRunner::generate(
   return Error::Ok;
 }
 
-Error TextLLMRunner::prefill(
+Result<uint64_t> TextLLMRunner::prefill(
     const std::vector<MultimodalInput>& inputs,
     int32_t num_bos,
     int32_t num_eos) {
@@ -283,7 +283,10 @@ Error TextLLMRunner::prefill(
     // Skip non-text inputs — text-only runner
   }
 
-  return Error::Ok;
+  if (!prefill_next_token_.has_value()) {
+    return Error::InvalidArgument;
+  }
+  return prefill_next_token_.value();
 }
 
 Error TextLLMRunner::warmup(const std::string& prompt, int32_t max_new_tokens) {

extension/llm/runner/text_llm_runner.h

@@ -109,17 +109,18 @@ class ET_EXPERIMENTAL TextLLMRunner : public IRunner {
    * @param inputs A vector of MultimodalInput objects.
    * @param num_bos Number of BOS tokens to prepend during text encoding.
    * @param num_eos Number of EOS tokens to append during text encoding.
-   * @return The error code. KV cache position is tracked internally in pos_.
+   * @return The next token predicted after prefill, or an error.
+   *         KV cache position is tracked internally in pos_.
    */
-  ::executorch::runtime::Error prefill(
+  ::executorch::runtime::Result<uint64_t> prefill(
       const std::vector<MultimodalInput>& inputs,
       int32_t num_bos = 0,
       int32_t num_eos = 0) override;
 
   /**
    * Convenience overload: prefill a single text prompt.
    */
-  ::executorch::runtime::Error
+  ::executorch::runtime::Result<uint64_t>
   prefill(const std::string& prompt, int32_t num_bos = 0, int32_t num_eos = 0) {
     std::vector<MultimodalInput> inputs;
     inputs.emplace_back(MultimodalInput(prompt));
@@ -130,7 +131,7 @@ class ET_EXPERIMENTAL TextLLMRunner : public IRunner {
    * Prefill a text prompt using GenerationConfig.
    * Deprecated: prefer prefill(prompt, num_bos, num_eos).
    */
-  ::executorch::runtime::Error prefill(
+  ::executorch::runtime::Result<uint64_t> prefill(
       const std::string& prompt,
       const GenerationConfig& config) {
     return prefill(prompt, config.num_bos, config.num_eos);