diff --git a/extension/llm/runner/multimodal_runner.cpp b/extension/llm/runner/multimodal_runner.cpp index 9afaab0b97e..694e2021993 100644 --- a/extension/llm/runner/multimodal_runner.cpp +++ b/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(); } diff --git a/extension/llm/runner/test/test_text_llm_runner.cpp b/extension/llm/runner/test/test_text_llm_runner.cpp index f6f5137a90f..55eb64375c5 100644 --- a/extension/llm/runner/test/test_text_llm_runner.cpp +++ b/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&, int64_t&) { + runner_ptr->stop(); // e.g. the user cancels while the prompt runs + return Result(4); + }); + + auto stats = std::make_unique(); + auto text_token_generator = createTextTokenGenerator( + tokenizer.get(), text_decoder_runner.get(), stats.get()); + auto module = std::make_unique(); + auto io_manager = + std::make_unique(*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(); + auto text_token_generator = createTextTokenGenerator( + tokenizer.get(), text_decoder_runner.get(), stats.get()); + auto module = std::make_unique(); + auto io_manager = + std::make_unique(*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(); diff --git a/extension/llm/runner/text_llm_runner.cpp b/extension/llm/runner/text_llm_runner.cpp index ce4f6a08ea2..e097871d084 100644 --- a/extension/llm/runner/text_llm_runner.cpp +++ b/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(); diff --git a/extension/llm/runner/text_token_generator.h b/extension/llm/runner/text_token_generator.h index d64927f6559..a33cbd6ef41 100644 --- a/extension/llm/runner/text_token_generator.h +++ b/extension/llm/runner/text_token_generator.h @@ -78,6 +78,9 @@ 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 generate( @@ -85,7 +88,8 @@ class ET_EXPERIMENTAL TextTokenGenerator { int64_t start_pos, int32_t max_new_tokens, float temperature = 0.0f, - const std::function& token_callback = {}) { + const std::function& 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.