fix(llm): chunk prefill at the bound the graph declares - #1491
Merged
Merged
Conversation
…seq_len Upstream sizes its prefill chunks from `get_max_seq_len`, which is the KV context budget rather than the widest tensor the graph accepts. gemma4_e2b_mlx_int4 advertises 2048 while its `forward` token input is DYNAMIC_BOUND to 511, so nothing ever chunks and the first prompt past 511 tokens fails in TensorImpl::internal_resize_contiguous with Error::NotSupported. The Vulkan gemma4 export bounds at 128. Read the real capacity off the method metadata, the way the legacy runner already does, and lower the prefiller's chunk size to it. The multimodal runner has no chunking at all, so split its text inputs to fit instead.
barhanc
reviewed
Sep 25, 2026
barhanc
left a comment
barhanc
left a comment
Member
There was a problem hiding this comment.
The comments should be cleaned a bit as well. Other than that it looks correct.
Comment on lines
+180
to
+181
| // Text that fits is left as text, because upstream only echoes an input it can | ||
| // still read back as a string. |
Member
There was a problem hiding this comment.
Regarding this comment, I think we should delete the echo option from LLMRunner as now it sometimes works and sometimes doesn't when the input gets chunked.
Chunked prompts reach the runner as token inputs, which upstream cannot echo, so echo only worked for short prompts.
barhanc
approved these changes
Sep 25, 2026
12 tasks
msluszniak
added a commit
that referenced
this pull request
Sep 25, 2026
…1494) (#1497) ## Description Cherry-picks five fixes from `main` onto `release/0.10` for the v0.10.3 patch, and pins `nativeLibsVersion` to `0.10.4`. - 3bf8168 `fix(libs): stop shipping backends the app opted out of` (#1466) - e76e0b9 `fix(llm): keep every terminal token out of the chat response` (#1486) - daafa64 `fix(llm): pin the model load mode instead of inheriting upstream defaults` (#1492) - 10bfbcf `fix(llm): chunk prefill at the bound the graph declares` (#1491) - 8406b98 `fix(install): make the native lib download work on Windows` (#1494) `v0.10.4-libs` carries the XNNPACK weights-cache fixes (PReLU use-after-free, and the retained weights that got Kokoro XNNPACK killed on iPhone). Docs hunks from #1466 and #1494 are left out, per the patch-release rule. Every touched file matches `main` except `package.json` and `download-libs.js`, which lacks #1479's Vulkan task map. ### Introduces a breaking change? - [ ] Yes - [x] No ### Type of change - [x] Bug fix (change which fixes an issue) - [ ] New feature (change which adds functionality) - [ ] Documentation update (improves or adds clarity to existing documentation) - [ ] Other (chores, tests, code style improvements etc.) ### Tested on - [x] iOS - [x] Android ### Testing instructions - `yarn jest __tests__`, `yarn typecheck`, `yarn lint`, `scripts/run-native-tests.sh` - `apps/speech` Release on iPhone 16 and Galaxy S26 Ultra: Kokoro EN_US XNNPACK synthesizes, and a 12-layer PReLU `.pte` matches eager output. ### Screenshots ### Related issues ### Checklist - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [x] I have updated the documentation accordingly - [x] My changes generate no new warnings ### Additional notes The version bump follows in a separate `Release v0.10.3` PR once this merges. --------- Co-authored-by: Mateusz Słuszniak <msluszniak1@gmail.com> Co-authored-by: Bartosz Hanc <bartosz.hanc02@gmail.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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
Upstream sizes prefill chunks from
get_max_seq_len, which is the KV context budget rather than the widest tensor the graph accepts.TextPrefillerthen never chunks, and any prompt past the real bound fails ininternal_resize_contiguouswithError::NotSupported. Both gemma4 builds with a sliding window disagree:get_max_seq_lenforwardboundgemma4_e2b_mlx_int4gemma_4_e2b_xnnpack_8da4wWe now read the real bound from
method_meta(...).input_tensor_meta(0)and lower the prefiller's chunk size to it, which is what the legacy runner has always done.MultimodalPrefillerhas no chunking at all, so its text inputs are split intoTOKENSpieces instead.Introduces a breaking change?
Type of change
Tested on
Testing instructions
apps/nlp, gemma-4-e2b-mlx, iPhone 16 Release. Send a prompt over 511 tokens (~6000 chars). A/B on the same binary:Error::NotSupported, KV position does not moveRelated issues
#1489
Checklist
Additional notes
Not gated on a backend. The clamp is a no-op when the two numbers agree, which is the case for the Vulkan gemma4 build (128/128) and for every other LLM checked, where the bound is just
max_seq_len - 1.The legacy runner reads the same metadata but only when
uses_backend("MLXBackend"), so it is still exposed on gemma4 XNNPACK.