From 4a38f8f9f3fc58abaa1d7b017e5814faa1909d9f Mon Sep 17 00:00:00 2001 From: Mergen Nachin Date: Fri, 12 Jun 2026 12:57:22 -0700 Subject: [PATCH] llm_runner: plumb prefill temperature Session-based serving drives generation as prefill plus token steps instead of one monolithic generate call. For that path to be correct, the first sampled token produced during prefill must honor the same sampling inputs as the rest of the decode loop; otherwise requests using temperature can silently start greedily and then switch behavior on later tokens. This threads optional temperature through TextPrefiller and exposes the existing TextTokenGenerator logit-processor application so token-step callers can reuse the same sampling preparation as generate(). The goal is to remove a divergence point before session-backed serving starts depending on these primitives. Default behavior remains greedy, so existing callers that do not pass temperature keep the same semantics. The added tests focus on the new non-default path and on sharing the logit-processor logic rather than duplicating it. --- extension/llm/runner/README.md | 4 +- .../llm/runner/test/test_text_llm_runner.cpp | 20 +++ .../llm/runner/test/test_text_prefiller.cpp | 156 +++++++++++++++++- extension/llm/runner/text_prefiller.cpp | 42 ++++- extension/llm/runner/text_prefiller.h | 20 +++ extension/llm/runner/text_token_generator.h | 18 +- 6 files changed, 242 insertions(+), 18 deletions(-) diff --git a/extension/llm/runner/README.md b/extension/llm/runner/README.md index 4fa3b079039..9e1b5c307b9 100644 --- a/extension/llm/runner/README.md +++ b/extension/llm/runner/README.md @@ -731,7 +731,7 @@ std::unordered_map get_llm_metadata( |-----------|------|---------|-------------| | `max_new_tokens` | `int32_t` | `-1` | Maximum new tokens to generate (-1 = use available context) | | `seq_len` | `int32_t` | `1024` | Total sequence length including prompt | -| `temperature` | `float` | `0.8f` | Sampling temperature (0.0 = deterministic, 1.0+ = creative) | +| `temperature` | `float` | `0.8f` | Sampling temperature in [0.0, 1.0] (0.0 = deterministic) | | `echo` | `bool` | `true` | Whether to echo the input prompt | | `num_bos` | `int8_t` | `1` | Number of beginning-of-sequence tokens | | `num_eos` | `int8_t` | `1` | Number of end-of-sequence tokens | @@ -824,7 +824,7 @@ GenerationConfig config; config.temperature = 0.1f; // Very deterministic runner->generate(factual_prompt, config, callback); -config.temperature = 1.2f; // Very creative +config.temperature = 1.0f; // Highest supported temperature runner->generate(creative_prompt, config, callback); ``` diff --git a/extension/llm/runner/test/test_text_llm_runner.cpp b/extension/llm/runner/test/test_text_llm_runner.cpp index 5a98c55bb2f..9a967dd82c1 100644 --- a/extension/llm/runner/test/test_text_llm_runner.cpp +++ b/extension/llm/runner/test/test_text_llm_runner.cpp @@ -709,6 +709,26 @@ TEST_F(RunnerTest, TextTokenGeneratorProcessorChainMasksMultipleTokens) { EXPECT_EQ(generated_tokens, expected); } +TEST_F(RunnerTest, TextTokenGeneratorRejectsTemperatureOutOfRange) { + auto tokenizer = createMockTokenizer(); + auto text_decoder_runner = createMockTextDecoderRunner(); + Stats stats; + auto generator = createTextTokenGenerator( + tokenizer.get(), text_decoder_runner.get(), &stats); + + std::vector tokens = {1, 2, 3}; + EXPECT_CALL(*text_decoder_runner, step(_, _)).Times(0); + + EXPECT_EQ( + generator->generate(tokens, 3, 3, -0.1f, [](const std::string&) {}) + .error(), + Error::InvalidArgument); + EXPECT_EQ( + generator->generate(tokens, 3, 3, 1.1f, [](const std::string&) {}) + .error(), + Error::InvalidArgument); +} + // Without any processors, greedy argmax picks token 3 (zero-overhead path). TEST_F(RunnerTest, TextTokenGeneratorWithoutProcessorPicksArgmax) { auto tokenizer = createMockTokenizer(); diff --git a/extension/llm/runner/test/test_text_prefiller.cpp b/extension/llm/runner/test/test_text_prefiller.cpp index 5ed5031dace..e859db36346 100644 --- a/extension/llm/runner/test/test_text_prefiller.cpp +++ b/extension/llm/runner/test/test_text_prefiller.cpp @@ -80,7 +80,12 @@ class TextPrefillerTest : public Test { ::executorch::runtime::Result, prefill_chunk, (std::vector&, int64_t&), - ()); + (override)); + MOCK_METHOD( + ::executorch::runtime::Result, + prefill_chunk, + (std::vector&, int64_t&, float), + (override)); }; // Create a mock TextPrefiller @@ -112,9 +117,10 @@ TEST_F(TextPrefillerTest, PrefillCallsPrefillChunkOnceWhenPromptFits) { int64_t start_pos = 0; // Expect prefill_chunk to be called exactly once with the entire prompt - EXPECT_CALL(*prefiller, prefill_chunk(_, _)) + constexpr float temperature = 0.7f; + EXPECT_CALL(*prefiller, prefill_chunk(_, _, FloatEq(temperature))) .Times(1) - .WillOnce([&](std::vector& tokens, int64_t& pos) { + .WillOnce([&](std::vector& tokens, int64_t& pos, float temp) { // Verify the tokens passed to prefill_chunk EXPECT_EQ(tokens.size(), prompt_tokens.size()); for (size_t i = 0; i < tokens.size(); i++) { @@ -122,17 +128,134 @@ TEST_F(TextPrefillerTest, PrefillCallsPrefillChunkOnceWhenPromptFits) { } // Verify the position EXPECT_EQ(pos, start_pos); + EXPECT_EQ(temp, temperature); return Result(42); }); // Call prefill - auto result = prefiller->prefill(prompt_tokens, start_pos); + auto result = prefiller->prefill(prompt_tokens, start_pos, temperature); // Verify the result EXPECT_EQ(result.error(), Error::Ok); EXPECT_EQ(result.get(), 42); } +TEST_F(TextPrefillerTest, TwoArgumentPrefillUsesGreedyTemperature) { + auto prefiller = createMockTextPrefiller(10); + + std::vector prompt_tokens = {1, 2, 3}; + int64_t start_pos = 0; + + EXPECT_CALL(*prefiller, prefill_chunk(_, _, FloatEq(0.0f))) + .Times(1) + .WillOnce([](std::vector&, int64_t&, float) { + return Result(42); + }); + + auto result = prefiller->prefill(prompt_tokens, start_pos); + + EXPECT_EQ(result.error(), Error::Ok); + EXPECT_EQ(result.get(), 42); +} + +TEST_F(TextPrefillerTest, PrefillAcceptsTemperatureBounds) { + auto prefiller = createMockTextPrefiller(10); + + std::vector prompt_tokens = {1, 2, 3}; + int64_t start_pos = 0; + + { + InSequence seq; + EXPECT_CALL(*prefiller, prefill_chunk(_, _, FloatEq(0.0f))) + .WillOnce([](std::vector&, int64_t&, float) { + return Result(41); + }); + EXPECT_CALL(*prefiller, prefill_chunk(_, _, FloatEq(1.0f))) + .WillOnce([](std::vector&, int64_t&, float) { + return Result(42); + }); + } + + auto greedy = prefiller->prefill(prompt_tokens, start_pos, 0.0f); + auto max_temp = prefiller->prefill(prompt_tokens, start_pos, 1.0f); + + EXPECT_EQ(greedy.error(), Error::Ok); + EXPECT_EQ(greedy.get(), 41); + EXPECT_EQ(max_temp.error(), Error::Ok); + EXPECT_EQ(max_temp.get(), 42); +} + +TEST_F(TextPrefillerTest, PrefillRejectsTemperatureOutOfRange) { + auto prefiller = createMockTextPrefiller(10); + + std::vector prompt_tokens = {1, 2, 3}; + int64_t start_pos = 0; + + EXPECT_CALL(*prefiller, prefill_chunk(_, _, _)).Times(0); + + EXPECT_EQ( + prefiller->prefill(prompt_tokens, start_pos, -0.1f).error(), + Error::InvalidArgument); + EXPECT_EQ( + prefiller->prefill(prompt_tokens, start_pos, 1.1f).error(), + Error::InvalidArgument); +} + +TEST_F(TextPrefillerTest, TwoArgumentPrefillChunkOverrideStillDispatches) { + class LegacyPrefiller final : public TextPrefiller { + public: + explicit LegacyPrefiller(TextDecoderRunner* text_decoder_runner) + : TextPrefiller(text_decoder_runner, true, true, 10) {} + + Result prefill_chunk(std::vector&, int64_t&) override { + called = true; + return Result(42); + } + + bool called = false; + }; + + LegacyPrefiller prefiller(&text_decoder_runner_); + TextPrefiller* base = &prefiller; + std::vector prompt_tokens = {1, 2, 3}; + int64_t start_pos = 0; + + auto result = base->prefill_chunk(prompt_tokens, start_pos); + + EXPECT_EQ(result.error(), Error::Ok); + EXPECT_EQ(result.get(), 42); + EXPECT_TRUE(prefiller.called); +} + +TEST_F(TextPrefillerTest, ChunkedPrefillSamplesOnlyLastChunkWithTemperature) { + auto prefiller = createMockTextPrefiller(3); + + std::vector prompt_tokens = {1, 2, 3, 4, 5, 6, 7, 8}; + int64_t start_pos = 0; + constexpr float temperature = 0.9f; + + { + InSequence seq; + EXPECT_CALL(*prefiller, prefill_chunk(_, _, FloatEq(0.0f))) + .WillOnce([](std::vector&, int64_t&, float) { + return Result(10); + }); + EXPECT_CALL(*prefiller, prefill_chunk(_, _, FloatEq(0.0f))) + .WillOnce([](std::vector&, int64_t&, float) { + return Result(11); + }); + EXPECT_CALL(*prefiller, prefill_chunk(_, _, FloatEq(temperature))) + .WillOnce([](std::vector&, int64_t&, float) { + return Result(12); + }); + } + + auto result = prefiller->prefill(prompt_tokens, start_pos, temperature); + + EXPECT_EQ(result.error(), Error::Ok); + EXPECT_EQ(result.get(), 12); +} + // Test that prefill() calls prefill_chunk() multiple times when prompt tokens > // max_seq_len TEST_F( @@ -217,14 +340,14 @@ TEST_F(TextPrefillerTest, PrefillHandlesPrefillChunkErrorsCorrectly) { InSequence seq; // First chunk: tokens [1, 2, 3] - succeeds - EXPECT_CALL(*prefiller, prefill_chunk(_, _)) - .WillOnce([&](std::vector& tokens, int64_t& pos) { + EXPECT_CALL(*prefiller, prefill_chunk(_, _, _)) + .WillOnce([&](std::vector& tokens, int64_t& pos, float) { return Result(10); }); // Second chunk: tokens [4, 5] - fails - EXPECT_CALL(*prefiller, prefill_chunk(_, _)) - .WillOnce([&](std::vector& tokens, int64_t& pos) { + EXPECT_CALL(*prefiller, prefill_chunk(_, _, _)) + .WillOnce([&](std::vector& tokens, int64_t& pos, float) { return Result(Error::InvalidArgument); }); } @@ -236,6 +359,23 @@ TEST_F(TextPrefillerTest, PrefillHandlesPrefillChunkErrorsCorrectly) { EXPECT_EQ(result.error(), Error::InvalidArgument); } +TEST_F(TextPrefillerTest, PrefillChunkRejectsTemperatureOutOfRange) { + auto prefiller = createTextPrefiller(10, true, true); + + std::vector prompt_tokens = {1, 2, 3}; + int64_t start_pos = 0; + + EXPECT_CALL(text_decoder_runner_, step(_, _)).Times(0); + + EXPECT_EQ( + prefiller->prefill_chunk(prompt_tokens, start_pos, -0.1f).error(), + Error::InvalidArgument); + EXPECT_EQ( + prefiller->prefill_chunk(prompt_tokens, start_pos, 1.1f).error(), + Error::InvalidArgument); + EXPECT_EQ(start_pos, 0); +} + // Test that prefill_chunk() works correctly with parallel prefill enabled TEST_F(TextPrefillerTest, PrefillChunkWorksWithParallelPrefill) { // Create a TextPrefiller with parallel prefill enabled diff --git a/extension/llm/runner/text_prefiller.cpp b/extension/llm/runner/text_prefiller.cpp index a391cef01de..a3bfcced4c7 100644 --- a/extension/llm/runner/text_prefiller.cpp +++ b/extension/llm/runner/text_prefiller.cpp @@ -29,7 +29,19 @@ TextPrefiller::TextPrefiller( ::executorch::runtime::Result TextPrefiller::prefill( std::vector& prompt_tokens, int64_t& start_pos) { + return prefill(prompt_tokens, start_pos, 0.0f); +} + +::executorch::runtime::Result TextPrefiller::prefill( + std::vector& prompt_tokens, + int64_t& start_pos, + float temperature) { ET_CHECK_MSG(!prompt_tokens.empty(), "Prompt cannot be null"); + ET_CHECK_OR_RETURN_ERROR( + temperature >= 0.0f && temperature <= 1.0f, + InvalidArgument, + "Temperature must be in [0, 1], got %f", + static_cast(temperature)); if (!text_decoder_runner_->is_method_loaded()) { ET_CHECK_OK_OR_RETURN_ERROR(text_decoder_runner_->load()); } @@ -54,8 +66,14 @@ ::executorch::runtime::Result TextPrefiller::prefill( num_tokens_to_prefill_with, prompt_tokens_to_process.begin()); - // Process this chunk - auto chunk_result = prefill_chunk(prompt_tokens_to_process, start_pos); + // Only the final chunk samples the first generated token. + const bool is_last_chunk = + num_tokens_to_process + num_tokens_to_prefill_with >= + num_prompt_tokens; + auto chunk_result = prefill_chunk( + prompt_tokens_to_process, + start_pos, + is_last_chunk ? temperature : 0.0f); ET_CHECK_OK_OR_RETURN_ERROR(chunk_result.error()); cur_token = chunk_result.get(); @@ -65,13 +83,25 @@ ::executorch::runtime::Result TextPrefiller::prefill( return cur_token; } else { // If prompt tokens don't exceed max_seq_len_, process them directly - return prefill_chunk(prompt_tokens, start_pos); + return prefill_chunk(prompt_tokens, start_pos, temperature); } } ::executorch::runtime::Result TextPrefiller::prefill_chunk( std::vector& prompt_tokens, int64_t& start_pos) { + return prefill_chunk(prompt_tokens, start_pos, 0.0f); +} + +::executorch::runtime::Result TextPrefiller::prefill_chunk( + std::vector& prompt_tokens, + int64_t& start_pos, + float temperature) { + ET_CHECK_OR_RETURN_ERROR( + temperature >= 0.0f && temperature <= 1.0f, + InvalidArgument, + "Temperature must be in [0, 1], got %f", + static_cast(temperature)); // enable_parallel_prefill_ maybe set even when not using kv cache // When kv cache is not used, start pos is ignored int32_t num_prompt_tokens = prompt_tokens.size(); @@ -92,7 +122,8 @@ ::executorch::runtime::Result TextPrefiller::prefill_chunk( Info, "Prefill token result numel(): %zu", outputs_res.get().numel()); start_pos += num_prompt_tokens; - cur_token = text_decoder_runner_->logits_to_token(outputs_res.get()); + cur_token = + text_decoder_runner_->logits_to_token(outputs_res.get(), temperature); } else { // sequential prefill int64_t pos = 0; // position in the sequence // NOLINTNEXTLINE(facebook-hte-ParameterUncheckedArrayBounds) @@ -128,7 +159,8 @@ ::executorch::runtime::Result TextPrefiller::prefill_chunk( start_pos++; } - cur_token = text_decoder_runner_->logits_to_token(logits_tensor); + cur_token = + text_decoder_runner_->logits_to_token(logits_tensor, temperature); } return cur_token; } diff --git a/extension/llm/runner/text_prefiller.h b/extension/llm/runner/text_prefiller.h index a02cd3d1bf4..9f032122349 100644 --- a/extension/llm/runner/text_prefiller.h +++ b/extension/llm/runner/text_prefiller.h @@ -32,23 +32,43 @@ class ET_EXPERIMENTAL TextPrefiller { * tokenizer. * @param start_pos The starting position in KV cache of the input in the LLM * Module. + * Equivalent to `prefill(prompt_tokens, start_pos, 0.0f)`. * @return The next token of the LLM Module after prefill. */ virtual ::executorch::runtime::Result prefill( std::vector& prompt_tokens, int64_t& start_pos); + /** + * Like `prefill(prompt_tokens, start_pos)`, but samples the first generated + * token with `temperature` in [0.0, 1.0]. + */ + virtual ::executorch::runtime::Result prefill( + std::vector& prompt_tokens, + int64_t& start_pos, + float temperature); + /** * Helper method to prefill a chunk of tokens. * @param prompt_tokens The chunk of text prompt tokens to process. * @param start_pos The starting position in KV cache of the input in the LLM * Module. + * Equivalent to `prefill_chunk(prompt_tokens, start_pos, 0.0f)`. * @return The next token of the LLM Module after prefilling this chunk. */ virtual ::executorch::runtime::Result prefill_chunk( std::vector& prompt_tokens, int64_t& start_pos); + /** + * Like `prefill_chunk(prompt_tokens, start_pos)`, but samples the produced + * token with `temperature` in [0.0, 1.0]. + */ + virtual ::executorch::runtime::Result prefill_chunk( + std::vector& prompt_tokens, + int64_t& start_pos, + float temperature); + /** * Load the necessary resources for the TextPrefiller. * This method should be called before using the prefill methods. diff --git a/extension/llm/runner/text_token_generator.h b/extension/llm/runner/text_token_generator.h index 3627cacf3c3..d64927f6559 100644 --- a/extension/llm/runner/text_token_generator.h +++ b/extension/llm/runner/text_token_generator.h @@ -55,6 +55,15 @@ class ET_EXPERIMENTAL TextTokenGenerator { return logit_processors_.size(); } + /// Apply the registered logit processors before sampling. + inline ::executorch::runtime::Error apply_logit_processors( + executorch::aten::Tensor& logits) { + for (auto& processor : logit_processors_) { + ET_CHECK_OK_OR_RETURN_ERROR(processor->process(logits)); + } + return ::executorch::runtime::Error::Ok; + } + virtual ~TextTokenGenerator() = default; /** @@ -79,6 +88,11 @@ class ET_EXPERIMENTAL TextTokenGenerator { const std::function& token_callback = {}) { ET_CHECK_MSG( !tokens.empty(), "Token generation loop shouldn't take empty tokens"); + ET_CHECK_OR_RETURN_ERROR( + temperature >= 0.0f && temperature <= 1.0f, + InvalidArgument, + "Temperature must be in [0, 1], got %f", + static_cast(temperature)); int64_t pos = start_pos; // position in the sequence std::vector token_data; // allocate space for the tokens @@ -126,9 +140,7 @@ class ET_EXPERIMENTAL TextTokenGenerator { prev_token = cur_token; - for (auto& processor : logit_processors_) { - ET_CHECK_OK_OR_RETURN_ERROR(processor->process(logits_tensor)); - } + ET_CHECK_OK_OR_RETURN_ERROR(apply_logit_processors(logits_tensor)); stats_->on_sampling_begin(); cur_token =