[LLM runner] Keep stop() requests made during prefill

extension/llm/runner/multimodal_runner.cpp

@@ -239,6 +239,9 @@ Error MultimodalRunner::generate(
     ET_CHECK_OK_OR_RETURN_ERROR(load());
   }
 
+  // A stop() from here on, including during prefill, ends this generation.
+  text_token_generator_->reset_stop();
+
   if (config.warming) {
     ET_LOG(Info, "Doing a warmup run...");
   }
@@ -333,7 +336,8 @@ Error MultimodalRunner::generate(
       /*max_new_tokens=*/max_new_tokens -
           1, // Subtract 1 because prefill already generated 1 token
       /*temperature=*/config.temperature,
-      /*token_callback=*/wrapped_callback);
+      /*token_callback=*/wrapped_callback,
+      /*reset_stop=*/false);
   if (!generate_result.ok()) {
     return generate_result.error();
   }

extension/llm/runner/test/test_text_llm_runner.cpp

@@ -301,6 +301,90 @@ TEST_F(RunnerTest, GenerateCallsCallbackExactlyMaxNewTokensTimes) {
 }
 
 // Test that warmup() calls generate with the warming flag set
+// stop() called while the prompt is being prefilled must end the generation:
+// only the token prefill produced is emitted, and the decoder never runs.
+TEST_F(RunnerTest, StopDuringPrefillEndsGeneration) {
+  auto tokenizer = createMockTokenizer();
+  auto text_decoder_runner = createMockTextDecoderRunner();
+  EXPECT_CALL(*text_decoder_runner, step(_, _)).Times(0);
+  auto text_prefiller = createMockTextPrefiller(text_decoder_runner.get());
+
+  TextLLMRunner* runner_ptr = nullptr;
+  ON_CALL(*text_prefiller, prefill(_, _))
+      .WillByDefault([&](std::vector<uint64_t>&, int64_t&) {
+        runner_ptr->stop(); // e.g. the user cancels while the prompt runs
+        return Result<uint64_t>(4);
+      });
+
+  auto stats = std::make_unique<executorch::llm::Stats>();
+  auto text_token_generator = createTextTokenGenerator(
+      tokenizer.get(), text_decoder_runner.get(), stats.get());
+  auto module = std::make_unique<MockModule>();
+  auto io_manager =
+      std::make_unique<executorch::extension::llm::IOManager>(*module);
+  TextLLMRunner runner(
+      createDefaultMetadata(),
+      std::unique_ptr<::tokenizers::Tokenizer>(tokenizer.release()),
+      std::move(module),
+      std::move(text_decoder_runner),
+      std::unique_ptr<::executorch::extension::llm::TextPrefiller>(
+          text_prefiller.release()),
+      std::move(io_manager),
+      std::move(text_token_generator),
+      std::move(stats));
+  runner_ptr = &runner;
+  runner.load();
+
+  GenerationConfig config;
+  config.max_new_tokens = 10;
+  config.echo = false;
+  CallbackCounter counter;
+  Error err = runner.generate(
+      "test prompt", config, [&counter](const std::string& token) {
+        counter.callback(token);
+      });
+
+  EXPECT_EQ(err, Error::Ok);
+  EXPECT_EQ(counter.getCount(), 1);
+}
+
+// A stop() from before generate() must not carry into it.
+TEST_F(RunnerTest, StopBeforeGenerateDoesNotCarryOver) {
+  auto tokenizer = createMockTokenizer();
+  auto text_decoder_runner = createMockTextDecoderRunner();
+  auto text_prefiller = createMockTextPrefiller(text_decoder_runner.get());
+  auto stats = std::make_unique<executorch::llm::Stats>();
+  auto text_token_generator = createTextTokenGenerator(
+      tokenizer.get(), text_decoder_runner.get(), stats.get());
+  auto module = std::make_unique<MockModule>();
+  auto io_manager =
+      std::make_unique<executorch::extension::llm::IOManager>(*module);
+  TextLLMRunner runner(
+      createDefaultMetadata(),
+      std::unique_ptr<::tokenizers::Tokenizer>(tokenizer.release()),
+      std::move(module),
+      std::move(text_decoder_runner),
+      std::unique_ptr<::executorch::extension::llm::TextPrefiller>(
+          text_prefiller.release()),
+      std::move(io_manager),
+      std::move(text_token_generator),
+      std::move(stats));
+  runner.load();
+  runner.stop();
+
+  GenerationConfig config;
+  config.max_new_tokens = 10;
+  config.echo = false;
+  CallbackCounter counter;
+  Error err = runner.generate(
+      "test prompt", config, [&counter](const std::string& token) {
+        counter.callback(token);
+      });
+
+  EXPECT_EQ(err, Error::Ok);
+  EXPECT_EQ(counter.getCount(), config.max_new_tokens);
+}
+
 TEST_F(RunnerTest, WarmupCallsGenerateWithWarmingFlag) {
   // Create mock instances using helper functions
   auto tokenizer = createMockTokenizer();

extension/llm/runner/text_llm_runner.cpp

@@ -82,6 +82,9 @@ Error TextLLMRunner::generate(
     stats_->model_load_end_ms = time_in_ms();
   }
 
+  // A stop() from here on, including during prefill, ends this generation.
+  text_token_generator_->reset_stop();
+
   if (config.warming) {
     ET_LOG(Info, "Doing a warmup run...");
   }
@@ -242,7 +245,12 @@ Error TextLLMRunner::generate(
 
   // Generate max_new_tokens - 1 because prefill already generated 1 token.
   auto generate_result = text_token_generator_->generate(
-      prompt_tokens, pos_, max_new_tokens - 1, resolved_temp, wrapped_callback);
+      prompt_tokens,
+      pos_,
+      max_new_tokens - 1,
+      resolved_temp,
+      wrapped_callback,
+      /*reset_stop=*/false);
 
   if (!generate_result.ok()) {
     return generate_result.error();

extension/llm/runner/text_token_generator.h

@@ -78,14 +78,18 @@ class ET_EXPERIMENTAL TextTokenGenerator {
    * random predictions, while a lower temperature results in more deterministic
    * predictions.
    * @param token_callback what to do after a token is generated.
+   * @param reset_stop clear an earlier stop() before generating. Runners that
+   * call reset_stop() themselves before prefill pass false, so a stop() that
+   * arrives during prefill still ends the generation.
    * @return how many tokens are generated.
    */
   inline ::executorch::runtime::Result<int64_t> generate(
       std::vector<uint64_t> tokens,
       int64_t start_pos,
       int32_t max_new_tokens,
       float temperature = 0.0f,
-      const std::function<void(const std::string&)>& token_callback = {}) {
+      const std::function<void(const std::string&)>& token_callback = {},
+      bool reset_stop = true) {
     ET_CHECK_MSG(
         !tokens.empty(), "Token generation loop shouldn't take empty tokens");
     ET_CHECK_OR_RETURN_ERROR(
@@ -128,10 +132,15 @@ class ET_EXPERIMENTAL TextTokenGenerator {
           resize_tensor_ptr(tokens_managed, token_shape));
     }
 
-    should_stop_.store(false, std::memory_order_relaxed);
+    if (reset_stop) {
+      should_stop_.store(false, std::memory_order_relaxed);
+    }
 
     // Generate our tokens
     while (pos < start_pos + max_new_tokens) {
+      if (should_stop_.load(std::memory_order_relaxed)) {
+        break;
+      }
       // Run the model
       auto logits_res = text_decoder_runner_->step(tokens_managed, pos);
 
@@ -192,6 +201,14 @@ class ET_EXPERIMENTAL TextTokenGenerator {
     should_stop_.store(true, std::memory_order_relaxed);
   }
 
+  /**
+   * Clear an earlier stop(). Runners call this when a generation starts,
+   * before prefill, and then pass reset_stop = false to generate().
+   */
+  inline void reset_stop() {
+    should_stop_.store(false, std::memory_order_relaxed);
+  }
+
   /**
    * Load the necessary resources for TextTokenGenerator.
    * This method should be called before using the generate() method.