Wire text-only prefillPrompt() to C++ runner->prefill()#17724
Wire text-only prefillPrompt() to C++ runner->prefill()#17724kirklandsign wants to merge 4 commits intomainfrom
Conversation
Fix prompt double-prefill in the image generate() overload: the method called prefillPrompt(prompt) which appends to the JNI buffer, then passed the same prompt to native generate() which appends it again on the multimodal path, corrupting KV cache state. Also fix all 5 null-input checks in JNI append methods that incorrectly returned Error::EndOfMethod (means "method execution finished") instead of Error::InvalidArgument. Replace deep copy of prefill_inputs_ with std::move to avoid unnecessary allocation for large image/audio inputs.
Previously, prefillPrompt() only appended to a JNI-side buffer
that was never consumed on the text-only code path. This was a
silent no-op: the buffered data sat there forever while generate()
passed its own prompt directly to the runner.
Now for text-only models (MODEL_TYPE_CATEGORY_LLM), prefillPrompt()
calls runner_->prefill() directly, which runs real model computation
and populates the KV cache. This enables chat history reload:
module.prefillPrompt("system: ..."); // fills KV cache
module.prefillPrompt("user: hello"); // fills KV cache
module.generate("user: new q", cb); // generates from pos_
Add prefill() to IRunner with a default NotSupported return so
vendor runners (QNN, MediaTek) are unaffected. Mark the existing
TextLLMRunner::prefill() as override.
Also clear prefill_inputs_ in resetContext() so stale multimodal
buffer data doesn't persist across context resets.
The method now invokes real C++ prefill for text-only models, so "append" no longer describes what it does.
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/17724
Note: Links to docs will display an error until the docs builds have been completed. ❌ 3 New FailuresAs of commit bf12659 with merge base a29539d ( NEW FAILURES - The following jobs have failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
bot
added
the
CLA Signed
This PR needs a
|
kirklandsign
requested review from
larryliu0820 and
mergennachin
as code owners
There was a problem hiding this comment.
Pull request overview
This PR enables the prefillPrompt() method to actually populate the KV cache for text-only (LLM) models by calling the C++ runner's prefill() method directly, rather than just buffering the input. This fixes a silent no-op where prefilled data was never consumed.
Changes:
- Added a virtual
prefill()method to theIRunnerinterface with a default implementation returningError::NotSupportedfor backward compatibility with vendor runners - Updated
TextLLMRunner::prefill()to be marked asoverride - Modified JNI layer to call
runner_->prefill()for LLM models instead of just buffering - Cleared
prefill_inputs_buffer inresetContext()to prevent stale multimodal data from persisting
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| extension/llm/runner/irunner.h | Added virtual prefill() method with default NotSupported implementation |
| extension/llm/runner/text_llm_runner.h | Marked existing prefill() as override |
| extension/android/jni/jni_layer_llama.cpp | Renamed method and added logic to call runner->prefill() for LLM models; cleared prefill_inputs_ in reset |
| extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/llm/LlmModule.java | Renamed JNI method from appendTextInput to prefillTextInput |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| /** | ||
| * Prefill a multimodal Module with the given text input. | ||
| * | ||
| * @param prompt The text prompt to prefill. | ||
| * @return 0, as the updated starting position in KV cache of the input in the LLM is no longer | ||
| * exposed to user. | ||
| * @throws RuntimeException if the prefill failed | ||
| */ |
There was a problem hiding this comment.
The javadoc comment states "Prefill a multimodal Module with the given text input" but this method now also works for text-only (LLM) models. The documentation should be updated to reflect that this method works for both multimodal and text-only models, or just say "Prefill the Module with the given text input" without specifically mentioning multimodal.
Summary
Previously, prefillPrompt() only appended to a JNI-side buffer
that was never consumed on the text-only code path. This was a
silent no-op: the buffered data sat there forever while generate()
passed its own prompt directly to the runner.
Now for text-only models (MODEL_TYPE_CATEGORY_LLM), prefillPrompt()
calls runner_->prefill() directly, which runs real model computation
and populates the KV cache. This enables chat history reload:
module.prefillPrompt("system: ..."); // fills KV cache
module.prefillPrompt("user: hello"); // fills KV cache
module.generate("user: new q", cb); // generates from pos_
Add prefill() to IRunner with a default NotSupported return so
vendor runners (QNN, MediaTek) are unaffected. Mark the existing
TextLLMRunner::prefill() as override.
Also clear prefill_inputs_ in resetContext() so stale multimodal
buffer data doesn't persist across context resets.
Test plan
CI