Store prefill output token to enable prefill() → generate("") workflow

Both TextLLMRunner and MultimodalRunner now store the sampled next
token from prefill() in prefill_next_token_. When generate() is
called with an empty prompt, it consumes this token and starts
decoding directly without re-prefilling.

This enables the workflow:
  runner->prefill("system prompt", 1, 0);
  runner->prefill("user turn", 0, 0);
  runner->generate("", config, callback);  // decode from KV cache

This PR was authored with the assistance of Claude.

Author: kirklandsign

extension/llm/runner/multimodal_runner.cpp

@@ -133,6 +133,7 @@ Error MultimodalRunner::prefill(
   if (!result.ok()) {
     return result.error();
   }
+  prefill_next_token_ = result.get();
   return Error::Ok;
 }
 
@@ -141,11 +142,100 @@ Error MultimodalRunner::generate(
     const GenerationConfig& config,
     std::function<void(const std::string&)> token_callback,
     std::function<void(const Stats&)> stats_callback) {
+  if (!prompt.empty()) {
+    std::vector<MultimodalInput> inputs;
+    inputs.emplace_back(MultimodalInput(prompt));
+    return generate(inputs, config, token_callback, stats_callback);
+  }
+
+  // Empty prompt: consume prefill_next_token_ and go straight to decode
   ET_CHECK_OR_RETURN_ERROR(
-      !prompt.empty(), InvalidArgument, "Prompt cannot be empty");
-  std::vector<MultimodalInput> inputs;
-  inputs.emplace_back(MultimodalInput(prompt));
-  return generate(inputs, config, token_callback, stats_callback);
+      prefill_next_token_.has_value(),
+      InvalidState,
+      "Empty prompt requires a prior prefill() call");
+
+  if (!is_loaded()) {
+    ET_CHECK_OK_OR_RETURN_ERROR(load());
+  }
+
+  // Wrap the token_callback with print function
+  std::function<void(const std::string&)> wrapped_callback =
+      [token_callback, config](const std::string& piece) {
+        if (!config.warming) {
+          safe_printf(piece.c_str());
+          fflush(stdout);
+        }
+        if (token_callback) {
+          token_callback(piece);
+        }
+      };
+
+  stats_->inference_start_ms = time_in_ms();
+
+  uint64_t cur_token = prefill_next_token_.value();
+  prefill_next_token_.reset();
+
+  stats_->first_token_ms = time_in_ms();
+  stats_->prompt_eval_end_ms = time_in_ms();
+  stats_->num_prompt_tokens = pos_;
+
+  auto decode_result = tokenizer_->decode(cur_token, cur_token);
+  if (!decode_result.ok()) {
+    ET_LOG(
+        Error,
+        "Tokenizers error code %d",
+        static_cast<uint32_t>(decode_result.error()));
+    return Error::InvalidArgument;
+  }
+  wrapped_callback(std::move(*decode_result));
+
+  // Resolve max_new_tokens based on config
+  int64_t max_context_len = metadata_.at(kMaxContextLen);
+  int32_t max_new_tokens = config.resolve_max_new_tokens(max_context_len, pos_);
+
+  ET_CHECK_OR_RETURN_ERROR(
+      max_new_tokens > 0,
+      InvalidArgument,
+      "Max new tokens %d is less than or equal to 0",
+      max_new_tokens);
+
+  // Set ignore_eos based on config
+  text_token_generator_->set_ignore_eos(config.ignore_eos);
+
+  // Generate tokens using the text token generator
+  std::vector<uint64_t> prompt_tokens = {cur_token};
+  auto generate_result = text_token_generator_->generate(
+      prompt_tokens,
+      pos_,
+      max_new_tokens -
+          1, // Subtract 1 because prefill already generated 1 token
+      config.temperature,
+      wrapped_callback);
+  if (!generate_result.ok()) {
+    return generate_result.error();
+  }
+  int64_t num_generated_tokens = generate_result.get();
+
+  pos_ += num_generated_tokens;
+  // Update stats
+  stats_->num_generated_tokens = num_generated_tokens;
+  // Finalize stats and call callback
+  stats_->inference_end_ms = time_in_ms();
+
+  if (!config.warming) {
+    printf("\n");
+  }
+  if (config.warming) {
+    ET_LOG(Info, "Warmup run finished!");
+  } else {
+    // Do not print report during warmup
+    print_report(*stats_);
+  }
+  if (stats_callback) {
+    stats_callback(*stats_);
+  }
+
+  return Error::Ok;
 }
 
 Error MultimodalRunner::generate(

extension/llm/runner/multimodal_runner.h

@@ -15,6 +15,7 @@
 #include <cstdint>
 #include <functional>
 #include <memory>
+#include <optional>
 #include <string>
 #include <unordered_map>
 
@@ -170,6 +171,7 @@ class ET_EXPERIMENTAL MultimodalRunner : public IRunner {
   void reset() override {
     pos_ = 0;
     stats_->reset();
+    prefill_next_token_.reset();
   }
 
   ~MultimodalRunner() override = default;
@@ -191,6 +193,7 @@ class ET_EXPERIMENTAL MultimodalRunner : public IRunner {
 #endif
 
   // Internal state
+  std::optional<uint64_t> prefill_next_token_;
   int64_t pos_;
 
  private:

extension/llm/runner/text_llm_runner.cpp

@@ -76,9 +76,6 @@ Error TextLLMRunner::generate(
     const GenerationConfig& config,
     std::function<void(const std::string&)> token_callback,
     std::function<void(const Stats&)> stats_callback) {
-  // Prepare the inputs.
-  // Use ones-initialized inputs.
-  ET_CHECK_MSG(!prompt.empty(), "Prompt cannot be null");
   if (!is_loaded()) {
     stats_->model_load_start_ms = time_in_ms();
     ET_CHECK_OK_OR_RETURN_ERROR(load());
@@ -112,37 +109,66 @@ Error TextLLMRunner::generate(
   stats_->inference_start_ms = time_in_ms();
   shouldStop_ = false;
 
-  ::tokenizers::Result<std::vector<uint64_t>> encode_res = tokenizer_->encode(
-      prompt,
-      /*bos=*/config.num_bos,
-      /*eos=*/config.num_eos);
+  // Capture remaining KV cache capacity before prefill (pos_ will change)
+  int64_t max_context_len = metadata_.at(kMaxContextLen) - pos_;
 
-  if (!encode_res.ok()) {
-    ET_LOG(
-        Error,
-        "Failed to encode prompt %s. Tokenizers error code %d",
-        prompt.c_str(),
-        static_cast<uint32_t>(encode_res.error()));
-    return Error::InvalidArgument;
-  }
+  uint64_t cur_token = 0;
+  int num_prompt_tokens = 0;
+  std::vector<uint64_t> prompt_tokens;
+
+  if (!prompt.empty()) {
+    ::tokenizers::Result<std::vector<uint64_t>> encode_res =
+        tokenizer_->encode(
+            prompt,
+            /*bos=*/config.num_bos,
+            /*eos=*/config.num_eos);
+
+    if (!encode_res.ok()) {
+      ET_LOG(
+          Error,
+          "Failed to encode prompt %s. Tokenizers error code %d",
+          prompt.c_str(),
+          static_cast<uint32_t>(encode_res.error()));
+      return Error::InvalidArgument;
+    }
 
-  // encode the (string) prompt into tokens sequence
-  std::vector<uint64_t> prompt_tokens = encode_res.get();
-  int num_prompt_tokens = prompt_tokens.size();
+    // encode the (string) prompt into tokens sequence
+    prompt_tokens = encode_res.get();
+    num_prompt_tokens = prompt_tokens.size();
+
+    ET_CHECK_OR_RETURN_ERROR(
+        num_prompt_tokens >= 1,
+        InvalidArgument,
+        "Expected at least 1 prompt token");
+    ET_CHECK_OR_RETURN_ERROR(
+        num_prompt_tokens < max_context_len,
+        InvalidArgument,
+        "num_prompt_tokens %d >= max_context_len %" PRId64
+        ", Max seq length exceeded - please increase max seq len value in your export script",
+        num_prompt_tokens,
+        max_context_len);
+
+    // print prompts
+    if (config.echo) {
+      wrapped_callback(prompt);
+    }
 
-  // Reduce max_context_len by pos_
-  int64_t max_context_len = metadata_.at(kMaxContextLen) - pos_;
-  ET_CHECK_OR_RETURN_ERROR(
-      num_prompt_tokens >= 1,
-      InvalidArgument,
-      "Expected at least 1 prompt token");
-  ET_CHECK_OR_RETURN_ERROR(
-      num_prompt_tokens < max_context_len,
-      InvalidArgument,
-      "num_prompt_tokens %d >= max_context_len %" PRId64
-      ", Max seq length exceeded - please increase max seq len value in your export script",
-      num_prompt_tokens,
-      max_context_len);
+    // Prefill first
+    // Here feed all tokens to the model and get the next predicted token
+    // after the prompt. After that we will enter generate loop.
+    auto prefill_res = text_prefiller_->prefill(prompt_tokens, pos_);
+    ET_CHECK_OK_OR_RETURN_ERROR(prefill_res.error());
+    cur_token = prefill_res.get();
+    prefill_next_token_.reset();
+  } else {
+    // Empty prompt: consume token from a prior prefill() call
+    ET_CHECK_OR_RETURN_ERROR(
+        prefill_next_token_.has_value(),
+        InvalidState,
+        "Empty prompt requires a prior prefill() call");
+    cur_token = prefill_next_token_.value();
+    prefill_next_token_.reset();
+  }
 
   // Determine max_new_tokens using the GenerationConfig's resolve method,
   // then subtract pos_ for max_new_tokens.
@@ -152,28 +178,17 @@ Error TextLLMRunner::generate(
   ET_LOG(
       Info,
       "Max new tokens resolved: %d, given pos_ %" PRId64
-      ", num_prompt_tokens %zu, max_context_len %" PRId64,
+      ", num_prompt_tokens %d, max_context_len %" PRId64,
       max_new_tokens,
       pos_,
-      prompt_tokens.size(),
+      num_prompt_tokens,
       max_context_len);
   ET_CHECK_OR_RETURN_ERROR(
       max_new_tokens > 0,
       InvalidArgument,
       "Max new tokens %d is less than or equal to 0",
       max_new_tokens);
 
-  // Prefill first
-  // Here feed all tokens to the model and get the next predicted token
-  // after the prompt. After that we will enter generate loop.
-
-  // print prompts
-  if (config.echo) {
-    wrapped_callback(prompt);
-  }
-  auto prefill_res = text_prefiller_->prefill(prompt_tokens, pos_);
-  ET_CHECK_OK_OR_RETURN_ERROR(prefill_res.error());
-  uint64_t cur_token = prefill_res.get();
   stats_->first_token_ms = time_in_ms();
   stats_->prompt_eval_end_ms = time_in_ms();
 
@@ -225,7 +240,8 @@ Error TextLLMRunner::generate(
     RUNNER_ET_LOG(config.warming, "Max new tokens %i reached!", max_new_tokens);
   }
 
-  stats_->num_prompt_tokens = num_prompt_tokens;
+  stats_->num_prompt_tokens =
+      prompt.empty() ? static_cast<int64_t>(pos_) : num_prompt_tokens;
   stats_->num_generated_tokens = num_generated_tokens;
 
   if (config.warming) {
@@ -260,6 +276,7 @@ Error TextLLMRunner::prefill(
       std::vector<uint64_t> tokens = encode_res.get();
       auto prefill_res = text_prefiller_->prefill(tokens, pos_);
       ET_CHECK_OK_OR_RETURN_ERROR(prefill_res.error());
+      prefill_next_token_ = prefill_res.get();
       num_bos = 0;
       num_eos = 0;
     }
@@ -295,6 +312,7 @@ void TextLLMRunner::stop() {
 void TextLLMRunner::reset() {
   stats_->reset();
   pos_ = 0;
+  prefill_next_token_.reset();
 }
 
 } // namespace executorch::extension::llm

extension/llm/runner/text_llm_runner.h

@@ -14,6 +14,7 @@
 #include <cstdint>
 #include <functional>
 #include <memory>
+#include <optional>
 #include <string>
 #include <unordered_map>
 
@@ -189,6 +190,9 @@ class ET_EXPERIMENTAL TextLLMRunner : public IRunner {
   // Deprecated, we should rely on the temperature in GenerationConfig instead.
   float temperature_ = -1.0f;
 
+  // Token predicted by the last prefill() call, consumed by generate("").
+  std::optional<uint64_t> prefill_next_token_;
+
   // The position in KV cache of the input, starting from 0.
   int64_t pos_ = 0;
 };