You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
[LLM:Bugfix] Refresh PLE for successive Omni text requests - #4894
Gemma 4 Omni instances can retain the one-token PLE input produced by the final decode step. A later multi-token text prefill then sees a non-null mPleInput, so Llm::embedding() keeps the stale tensor instead of rebuilding PLE for the new request.
This change:
clears mPleInput before a multi-token prefill only when Omni is on the text-only path and PLE is enabled; and
leaves the multimodal path unchanged so its full-input precomputed PLE is preserved.
Module
LLM / Omni (Gemma 4 PLE)
Type
Feature
Bugfix
Perf
Refact
Style
Doc
Test
Chore
Testing
Validated with a temporary test_omni_ple_lifecycle executable built using MNN_LLM_BUILD_TEST=ON; the test source and CMake target were removed per maintainer review.
Negative control: after reverting only the fix, the temporary test fails with Text prefill reused the stale one-token PLE input.
Fixed implementation: the same test passes and observes the PLE sequence length changing from 1 to 3.
Thanks for the focused PLE lifecycle fix and the A/B validation!
As with #4888, please remove transformers/llm/engine/test/test_omni_ple_lifecycle.cpp and its corresponding executable, include-directory, and link entries in transformers/llm/engine/CMakeLists.txt. This test covers a very narrow scenario, and we would prefer not to retain a separate test executable for it. The validation already documented in the PR is sufficient for this change; there is no need to move or expand the test.
Please keep the core PLE refresh fix in Omni::embedding() unchanged. Thanks again!
I’ve removed test_omni_ple_lifecycle.cpp and its corresponding executable, include-directory, and link entries from transformers/llm/engine/CMakeLists.txt. No replacement test was added.
The core PLE refresh fix in Omni::embedding() remains unchanged. I updated the PR description to clarify that the temporary test was used only for A/B validation and is no longer part of the change. The PR now contains only the core omni.cpp fix. Thanks!
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Fixes #4890.
Gemma 4 Omni instances can retain the one-token PLE input produced by the final decode step. A later multi-token text prefill then sees a non-null
mPleInput, soLlm::embedding()keeps the stale tensor instead of rebuilding PLE for the new request.This change:
mPleInputbefore a multi-token prefill only when Omni is on the text-only path and PLE is enabled; andModule
LLM / Omni (Gemma 4 PLE)
Type
Testing
test_omni_ple_lifecycleexecutable built usingMNN_LLM_BUILD_TEST=ON; the test source and CMake target were removed per maintainer review.Text prefill reused the stale one-token PLE input.Checklist
[Module:Type] Descriptionformat