Skip to content

refactor(llm)!: migrate to modular TypeScript runners and reexport models with expanded context lengths - #1508

Open
barhanc wants to merge 82 commits into
mainfrom
@bh/llm-runner-refactor
Open

barhanc wants to merge 82 commits into
mainfrom
@bh/llm-runner-refactor

Conversation

@barhanc

@barhanc barhanc commented Oct 2, 2026 •

Copy link
Copy Markdown
Member

Description

This PR undertakes an architectural overhaul of the LLM execution pipeline in react-native-executorch. The monolithic native C++ llm_runner.cpp has been replaced with specialized, modular TypeScript runners that drive the prefill and autoregressive decode loops in JS directly on top of core primitives. In tandem, sampling is now delegated to a dedicated native C++ kernel (sampler.cpp) that replicates the llama.cpp sampling pipeline with in-place logits mutations and unlocks zero-overhead constrained decoding.

Additionally, model registry definitions have been updated to point to newly re-exported models with expanded context windows (up to 16k tokens) and embedded companion reflection schemas, and the NLP example app has been upgraded with model family-tailored sampling configurations.

Key Highlights & Architectural Breakdown

1. TypeScript LLM Runners Architecture

  • Removed Native C++ Runner (cpp/extensions/llm/llm_runner.*): Removed the legacy C++ runner class and its associated test suite.
  • Modular Runners (src/extensions/llm/runners/):
    • llmTextRunner.ts: Handles text-only autoregressive models driving token_embedding and sequential/chunked prefill and single-token decode.
    • llmMultimodalRunner.ts: Handles interleaved multimodal models (e.g. LFM 2.5 VL) supporting image prefill embedding projections alongside text token streams.
    • llmGemmaRunner.ts: Dedicated runner for Gemma models with specialized position encoding and multimodal projection.
  • Polymorphic Dispatch (createLLMRunner):
    • Inspects the model methods schema via validateSpec and automatically instantiates the matching runner (text, vision, or gemma).

2. Native C++ Sampler Kernel (cpp/extensions/llm/sampler.cpp) & Constrained Decoding

Why the Sampler Had to be Implemented in C++:

  1. Large Vocabulary Overhead Across the JSI Boundary:
    Modern LLMs have vocabulary sizes ranging from $32\text{k}$ up to $151\text{k}+$ tokens (e.g., Qwen/Hammer at $151,665$ tokens). Serializing, copying, and sorting/filtering arrays of this size in JavaScript during every single decode token step introduces unacceptable Garbage Collection churn, JSI transfer overhead, and token latency regressions.
  2. In-Place Logits Mutation:
    The native C++ kernel operates directly in-place on the output logits buffer backing memory, eliminating tensor allocations entirely during generation.
  3. Execution Order Alignment with llama.cpp:
    The pipeline implements the full battle-tested llama.cpp sampling sequence with native performance.

Constrained Decoding Capabilities:

By exposing a native constraint masking phase at the very head of the sampler chain, the pipeline now natively supports dynamic constrained decoding:

  • Whitelisting & Blacklisting: Token masks (allowedTokens / bannedTokens) are applied directly in C++ by masking logits to $-\infty$ prior to penalties and softmax.
  • Grammars, JSON Schemas & Regex Engines: Higher-level libraries, CFG parsers, or user callbacks can feed token constraints per step via SamplingConstraintsCallback, enabling guaranteed structured JSON generation, regex-guided output, and strict tool-call schema compliance with zero intermediate tensor allocations.

3. Re-Exported Models with Increased Context Lengths (models.ts)

Model definitions in src/models.ts have been aligned with the re-exported models:

  • Companion Schemas & Logits Cast:
    • Re-exported .pte artifacts now embed validated get_model_schema method declarations.
    • All re-exported models now conform to single schema.

Introduces a breaking change?

  • Yes
  • No

Task Pipeline & Public API: slight changes to the runner API that was marked @experimental, higher-level chat session API and API not marked as experimental remains fully backward compatible)

Removal of Legacy Unquantized XNNPACK FP16 / BF16 Builds: legacy unquantized XNNPACK fp16 and bf16 variants have been retired and removed from active model registry definitions and Hugging Face model cards. These builds were impractical on mobile hardware due to excessive memory consumption and slow memory-bandwidth-bound decode loops.

Type of change

  • 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

  • iOS
  • Android

Testing instructions

  • All 38 test suites (4044 tests) pass cleanly across TypeScript API tests and mock JSI environments.
  • Native sampler test suites passed via LlmOpsTest.cpp.
  • TypeScript type checks.
  • Run LLM example app on-device (both Android and iOS) and empirically test at least one model from each family. Pay special attention to Gemma models, and MLX / Vulkan exports as these are the most fragile.
  • Inspect all public HF repos in the LLM collection: https://huggingface.co/collections/software-mansion/llm-multimodal after update.

Screenshots

Related issues

Closes #1458 #1483 #1377 #1351

Checklist

  • I have performed a self-review of my code
  • I have commented my code, particularly in hard-to-understand areas
  • I have updated the documentation accordingly
  • My changes generate no new warnings

Additional notes

@barhanc barhanc self-assigned this Oct 2, 2026
@barhanc
barhanc force-pushed the @bh/llm-runner-refactor branch 2 times, most recently from cf97f59 to 14a6fff Compare October 6, 2026 10:44
@barhanc barhanc changed the title @bh/llm runner refactor refactor(llm): migrate to modular TypeScript runners and reexport models with expanded context lengths Oct 12, 2026
@barhanc
barhanc marked this pull request as ready for review October 12, 2026 15:51
@barhanc
barhanc requested a review from msluszniak October 12, 2026 15:51
@barhanc

barhanc commented Oct 12, 2026

Copy link
Copy Markdown
Member Author

@msluszniak Everything except the llmGemmaRunner.ts is ready for review. When I'm done with Gemma Omni export I will let you know.

@barhanc
barhanc requested a review from kfaracik October 12, 2026 15:52
@barhanc barhanc changed the title refactor(llm): migrate to modular TypeScript runners and reexport models with expanded context lengths refactor(llm)!: migrate to modular TypeScript runners and reexport models with expanded context lengths Oct 12, 2026
@msluszniak msluszniak linked an issue Oct 13, 2026 that may be closed by this pull request
@msluszniak

Copy link
Copy Markdown
Member

Ok, regarding testing:

  • The multimodal runner has almost no TS tests - we should probably handle this one.
  • The fake sampler ignores every option except allowedTokens/bannedTokens. So the tests never check that temperature/topK/dry/xtc reach native. The merge in sampler.ts:164-169 isn't asserted anywhere. One test that records the options passed to llm.sample and checks the merged defaults would close that gap cheaply.
  • Variant dispatch is never tested (llmRunner.ts:197-205).
  • constructionFailure.test.ts dropped the runners bucket but didn't add an LLM factory.
  • Just a reminder to test llmGemmaRunner.ts.

@msluszniak

Copy link
Copy Markdown
Member

Native Sampler tests:

  • NaN options get through validation and give wrong output silently.
    • temperature: NaN always returns the least likely surviving token. sampleMultinomial never satisfies cum >= r and falls through to candidates.back().
    • topP, minP and repetitionPenalty set to NaN are ignored silently. temperature: Infinity is accepted.
    • Each range check is written so NaN passes it. Adding std::isfinite (plus one test) to every float option would fix it.
  • DRY and XTC options are checked after the user's constraints JS callback has run and after the repetition penalty has written to the buffer. Bad XTC options are only caught after top-k/top-p have run. Parsing and validating every option up front, would fix this.
  • Repetition penalty has no window. In llama.cpp there is a window, now we apply it to every token. Probably we should handle this as well.
  • Every decode step: runs Int32Array.from(generatedTokens) which is O(n^2) copying over a generation, copies that array into a std::vector, and builds a fresh std::unordered_set. Preallocating Int32Array (and passed as view) would remove the JS side of this.
  • There's no seed option. llama.cpp has it so same as with windom repetition penatly, maybe worth adding.
  • Missing tests: the full chain together aka constraints + penalty + DRY + top-k/p + XTC in one call to lock the order in, and the NaN cases from the first point.

@msluszniak

Copy link
Copy Markdown
Member

Typecheck:

  • I know that this one is preexisting but I think we should fix it. yarn typecheck fails with Cannot find module 'react-native-executorch/legacy' / ...-resource-fetcher. Adding a top-level typecheck step that builds the packages first, or paths mappings to src, would stop it looking like a failure.
  • Runners have {...} as unknown as LLMRunner. I'm thinking if there is an elegant way to deal with this and do not switching off the type check on them.
  • To omit return totStats! in multimodal (and gemma) we can start with { numTokens: 0, durationMs: 0, tokensPerSecond: 0 }. This would avoid dependency on guard we have rn. WDYT?

@msluszniak

Copy link
Copy Markdown
Member

Regarding HF and PR description:

  • I don't understand this:
    image
    and
    image
    So what is the context 2k or 8k?

  • PR claims that: "Qwen 3 0.6B has a Vulkan 8da4w build" but no Vulkan file on HF and none in the registry.

  • PR claims that "Llama 3.2 1B/3B MLX max seq is 4096" but from what I checked it is 2048, please double check this one.

  • NIT: PR says that: "LFM 2.5 1.2B max seq is 2048" but MLX is only 65.

  • API not marked as experimental remains fully backward compatible and FP16/BF16 baselines kept

    But what about XNNPACK_BF16/FP16 variants that are removed? I would change the PR wording here.

  • The per-variant config.json on HF files are stale in 38 of 45 variants. Most still say get_max_context_len: 2048. Please update them. Also stale f16/bf16 are still listed.

  • Qwen3 stop tokens differ by backend. XNNPACK has [151645, 151643], MLX has [151645], so MLX won't stop on <|endoftext|>.

  • The cards don't mention context length, the main feature of this PR. The variants table would be the natural place for it.

  • Five card sizes are stale: the three LFM VL MLX files (for example 1.4 GB listed vs 2.2 GiB real), and Gemma MLX and Vulkan. The other rows only disagree because of GB vs GiB, and the cards mix the two units. We should unify GBs vs GibS.

  • MLX int4 is labelled 4w in some repos and int4 in others. We should unify it.

  • Gemma's MLX file is gemma4_… while its other files are gemma_4_….

  • The collection still lists qwen-3.5 but is not in registery.

@barhanc

barhanc commented Oct 13, 2026

Copy link
Copy Markdown
Member Author

Native Sampler tests:

  • NaN options get through validation and give wrong output silently.

    • temperature: NaN always returns the least likely surviving token. sampleMultinomial never satisfies cum >= r and falls through to candidates.back().
    • topP, minP and repetitionPenalty set to NaN are ignored silently. temperature: Infinity is accepted.
    • Each range check is written so NaN passes it. Adding std::isfinite (plus one test) to every float option would fix it.

NaNs can't pass from TS as the asType conversions guard against that.

  • DRY and XTC options are checked after the user's constraints JS callback has run and after the repetition penalty has written to the buffer. Bad XTC options are only caught after top-k/top-p have run. Parsing and validating every option up front, would fix this.

The problem is that there is no clean way to do this upfront, since whole dry/xtc objects are optional, so we would have to define quite a lot of optional C++ values at the top. Imo not parsing upfront isn't that big of a problem, wrong config isn't something that the user can recover from so prior logits modification is fine imo.

  • Repetition penalty has no window. In llama.cpp there is a window, now we apply it to every token. Probably we should handle this as well.

Will fix this.

  • Every decode step: runs Int32Array.from(generatedTokens) which is O(n^2) copying over a generation, copies that array into a std::vector, and builds a fresh std::unordered_set. Preallocating Int32Array (and passed as view) would remove the JS side of this.

Not worth it imo. The transformer inference itself is O(n^2) and I extensively tested the sampling pipeline---in release build it amounts to < 1% of total inference time.

  • There's no seed option. llama.cpp has it so same as with windom repetition penatly, maybe worth adding.
  • Missing tests: the full chain together aka constraints + penalty + DRY + top-k/p + XTC in one call to lock the order in, and the NaN cases from the first point.

Will add those.

@barhanc

barhanc commented Oct 13, 2026

Copy link
Copy Markdown
Member Author

Typecheck:

  • I know that this one is preexisting but I think we should fix it. yarn typecheck fails with Cannot find module 'react-native-executorch/legacy' / ...-resource-fetcher. Adding a top-level typecheck step that builds the packages first, or paths mappings to src, would stop it looking like a failure.

yarn typecheck works fine for me ???

  • Runners have {...} as unknown as LLMRunner. I'm thinking if there is an elegant way to deal with this and do not switching off the type check on them.

This is deliberate. LLM runner has readonly [llmRunnerBrand]: never; (same as tensor, model) to prevent accidental as LLMRunner so in places where it is truly correct (like in the constructor) one needs as unknown as LLMRunner. I can remove the brand and this won't be needed, so whichever you think is better.

  • To omit return totStats! in multimodal (and gemma) we can start with { numTokens: 0, durationMs: 0, tokensPerSecond: 0 }. This would avoid dependency on guard we have rn. WDYT?

Will fix.

@msluszniak

Copy link
Copy Markdown
Member

NaNs can't pass from TS as the asType conversions guard against that.

I think you are only partially right here and it indeed works for integers, but nor for floats. asType<float> is just static_cast<float>(asType<double>(...)), and asType<double> only checks isNumber(). NaN is a number.

This is deliberate. LLM runner has readonly [llmRunnerBrand]: never; (same as tensor, model) to prevent accidental as LLMRunner so in places where it is truly correct (like in the constructor) one needs as unknown as LLMRunner. I can remove the brand and this won't be needed, so whichever you think is better.

I though more about:
return { ... } satisfies Omit<LLMRunner, typeof llmRunnerBrand> as unknown as LLMRunner;

yarn typecheck works fine for me ???

It only fails on a fresh clone. I meant a situation when lib hasn't been build yet.

@barhanc

barhanc commented Oct 13, 2026 •

Copy link
Copy Markdown
Member Author

I've addressed the review findings regarding code implementation. I'll fix the HF now.

  • Multimodal, Dispatch & Sampler Test Coverage (b2dfa05)

    • Added comprehensive TypeScript unit tests for the multimodal runner (llmMultimodalRunner).
    • Added tests verifying variant dispatch across LLM runners.
    • Added tests asserting that sampling options (temperature, topK, dry, xtc, etc.) correctly merge and reach native llm.sample.
    • Added LLM factory failure assertions to constructionFailure.test.ts.
  • Multimodal & Gemma Prefill Stats Cleanup (e31b197)

    • Initialized prefill totStats to { numTokens: 0, durationMs: 0, tokensPerSecond: 0 } in llmMultimodalRunner.ts and llmGemmaRunner.ts, eliminating the non-null assertion (totStats!).
  • Double & Float Non-Finite Validation (3a1f0cf)

    • Added std::isfinite checks to asType<double> (and float) conversions in cpp/core/conversions.cpp to reject NaN and Infinity, along with unit tests in ConversionsTest.cpp.
  • Repetition Penalty Window & Sampler Seed (b343f8a)

    • Added a penaltyLastN sliding-window option for repetition penalty in sampler.cpp and sampler.ts (matching llama.cpp behavior instead of penalizing all history).
    • Added a deterministic seed option to the sampler and native LLM sampling ops.
    • Added full sampling pipeline integration tests in LlmOpsTest.cpp asserting the full execution chain (constraints + penalty window + DRY + top-k/p + XTC + seed) and non-finite number handling.
  • Root yarn typecheck Clean-Checkout Failure

    • Fixed the missing module / type declarations failure (react-native-executorch/legacy, resource fetchers) on fresh clones by chaining yarn prepare before parallel workspace typechecks in the root package.json.

@msluszniak

Copy link
Copy Markdown
Member

I'm in the middle of the model sweep but one finding: <|python_tag|> leaks into the chat reply after a Llama 3.2 1B tool call.

@msluszniak

msluszniak commented Oct 13, 2026 •

Copy link
Copy Markdown
Member

More from the on-device sweep:

  • EOS tokens reach onToken before the stop check (llmTextRunner.ts:194-197, same loop in the other runners). Phi-4 Mini shows <|end|> in the reply, since stopTokens comes from tokenizer_config.json while get_eos_ids has more ids. Checking eosIds before emitting would fix it.
  • Hammer 2.1 answers non-tool questions with [], and after a tool result it keeps calling until max tool turns. The app should treat [] as "no call".
  • Qwen 3 <think> blocks show raw in the chat. Worth collapsing or stripping them in the app.

@barhanc

barhanc commented Oct 13, 2026

Copy link
Copy Markdown
Member Author

I'm in the middle of the model sweep but one finding: <|python_tag|> leaks into the chat reply after a Llama 3.2 1B tool call.

Fixed in Llama's tool parser in example app.

More from the on-device sweep:

  • EOS tokens reach onToken before the stop check (llmTextRunner.ts:194-197, same loop in the other runners). Phi-4 Mini shows <|end|> in the reply, since stopTokens comes from tokenizer_config.json while get_eos_ids has more ids. Checking eosIds before emitting would fix it.

Fixed.

  • Hammer 2.1 answers non-tool questions with [], and after a tool result it keeps calling until max tool turns. The app should treat [] as "no call".

Not sure this is in the scope of this PR, the example app is a crude one for general testing. Handling this would require special logic just for these 3 models.

  • Qwen 3 <think> blocks show raw in the chat. Worth collapsing or stripping them in the app.

Same as above.

@msluszniak

msluszniak commented Oct 14, 2026 •

Copy link
Copy Markdown
Member

iOS (iPhone 16, MLX) follow-up, checked on host with the published .pte files and greedy decoding:

  • Qwen 3 0.6B MLX int4 looks degraded: it answers "The capital of Poland is Krakow" (device and host), while bf16 Qwen3-0.6B answers Warsaw on the same prompt.

@msluszniak

Copy link
Copy Markdown
Member

Checked with react-native-rag + @react-native-rag/executorch@0.10.0 (Galaxy S26 Ultra, iPhone 16 MLX/XNNPACK).

One issue in the PR: runner reset and getKVCacheState aren't 'worklet' (all three runners), while prefill/generate/stop are. The docs say the runner runs on a worklet runtime and suggest reset(targetPos) for branching, which throws inside a worklet. Adding 'worklet' to both fixes it.

Compatibility only (RAG side): createLLMRunner is now async, so the adapter's wrapAsync(llm.createLLMRunner) fails, and maxSeqLen -> maxContextLen makes it silently drop chat history. With those patched, retrieval, interrupt and reload pass on both platforms. We'll need an adapter release for 0.11.

@barhanc

barhanc commented Oct 14, 2026

Copy link
Copy Markdown
Member Author

@msluszniak I've updated PR description, fixed HF model cards and published missing Gemma model. Everything is ready for next pass.

One issue in the PR: runner reset and getKVCacheState aren't 'worklet' (all three runners), while prefill/generate/stop are. The docs say the runner runs on a worklet runtime and suggest reset(targetPos) for branching, which throws inside a worklet. Adding 'worklet' to both fixes it.

Fixed.

  • The per-variant config.json on HF files are stale in 38 of 45 variants. Most still say get_max_context_len: 2048. Please update them. Also stale f16/bf16 are still listed.

Fixed.

  • Qwen3 stop tokens differ by backend. XNNPACK has [151645, 151643], MLX has [151645], so MLX won't stop on <|endoftext|>.

Fixed.

  • The cards don't mention context length, the main feature of this PR. The variants table would be the natural place for it.

Added context length columns.

  • Five card sizes are stale: the three LFM VL MLX files (for example 1.4 GB listed vs 2.2 GiB real), and Gemma MLX and Vulkan. The other rows only disagree because of GB vs GiB, and the cards mix the two units. We should unify GBs vs GibS.

Fixed.

  • MLX int4 is labelled 4w in some repos and int4 in others. We should unify it.

Fixed. It's int4 / int8 consistently now.

  • Gemma's MLX file is gemma4_… while its other files are gemma_4_….

Fixed. Now all files are named gemma_4_....

@msluszniak

Copy link
Copy Markdown
Member

@barhanc okay, I will check that tommorow with a focus on gemma runner.

@msluszniak

Copy link
Copy Markdown
Member

Second pass, Gemma runner:

  • Media positions get zero PLE (llmGemmaRunner.ts:257). HF Gemma 4 uses the pad token's PLE there (embed_tokens_per_layer[0], nonzero in the checkpoint). Running token_embedding([0]) once at load and reusing its PLE output would match.
  • The audio_encoder mask output is unused, so the padded tail of the last block (up to 11 tokens per chunk) gets prefilled. HF drops the masked tokens.
  • Omni XNNPACK reports get_max_context_len 4048 (MAX_SEQ_LEN = 4048 in both Gemma XNNPACK export scripts), probably meant 4096.

@msluszniak

Copy link
Copy Markdown
Member

Android (S26 Ultra), Gemma:

  • Both Gemma Vulkan exports give garbage on Adreno. A raw forward of the same prompt outside the runner: E2B Vulkan returns all-zero logits, Omni Vulkan picks <|turn> first and then every logit sits at -30 (softcap). XNNPACK is fine on the same prompt.
  • XNNPACK (both E2B and Omni): after a tool result the answer is empty, the first sampled token is <eos>. The rendered prompt matches the HF template, so this looks like the model or export, not the runner. Audio transcription is unreliable too ("Transcribe this audio exactly." gives <eos>).

…ol parsing

- Check eosIds before emitting to onToken and pushing to generatedTokens across text, gemma, and multimodal runners
- Support empty array [] tool call responses in parseHammerToolCalls
- Strip <|python_tag|> in parseLlamaToolCalls
@barhanc
barhanc force-pushed the @bh/llm-runner-refactor branch from 848bbfe to a715834 Compare October 21, 2026 11:32
@barhanc barhanc linked an issue Oct 21, 2026 that may be closed by this pull request
@barhanc

barhanc commented Oct 21, 2026

Copy link
Copy Markdown
Member Author

stop() is lost if it arrives before the reply or during the user prefill. The full reply is generated anyway (10/10 runs per platform). With a long prompt, the prefill stops after the first chunk and the model answers the truncated message, while getHistory() still holds the full one. prefill and generate both clear the cancel flag on entry (llmTextRunner.ts:115,168, multimodal :131,183,250,295, Gemma :199,309,336,398,442). Suggestion: clear it once per sendMessage, then roll back the half-prefilled message and return finishReason: 'stopped'.

Fixed.

I've also added the LFM 2.5 2.6B model (with 8k context), so you can test it as well.

@kfaracik

Copy link
Copy Markdown
Collaborator

Thanks @barhanc, the docs for local paths are fine for us, we'll just strip file:// from the picker URIs on our side.

BOS: I think we got our wires crossed. The id is fine, the problem is that the variable never gets to the template. Llama 3.2's template starts with {{- bos_token }}, but chatPreprocessor.process only passes tools, messages and add_generation_prompt to render(). So bos_token renders as an empty string, and the tokenizer doesn't add it later either (kNumAddedBosTokens = 0). On the device the Llama prompt starts with <|start_header_id|>system… and the ids go 128006, 9125, … with no 128000 in front. Gemma 4, LFM 2.5 and Bielik have the same issue. Passing bos_token / eos_token from tokenizer_config.json into render() should fix all four, which is what HF's apply_chat_template does.

Stop: nice, the session rollback works. One thing is left in the runner: prefill() and generate() still reset isCancelled on entry. Since we're going with the raw runner, a stop that lands right before one of those calls gets lost on our side. In the session, a stop that comes before the prefill starts still waits for the whole prompt to be read before it rolls back, and for Gemma under memory pressure that can be minutes. If the flag were cleared once per turn instead of on entry, both would be fixed. And a small one: a cancelled turn returns finishReason: 'stop', so a 'cancelled' value would make it easier to tell the two apart.

Seed: seed++ per call doesn't help within one reply, because generate() passes the same config to sample() at every step (llmTextRunner.ts:195), so every token is drawn at the same quantile. Doing seed + step inside generate() would fix it, and sample stays pure.

From my side that's it for this PR. Maybe one more thing: a 'length' finish reason at the context limit, if it's not much work.

@barhanc

barhanc commented Oct 21, 2026 •

Copy link
Copy Markdown
Member Author

BOS: I think we got our wires crossed. The id is fine, the problem is that the variable never gets to the template. Llama 3.2's template starts with {{- bos_token }}, but chatPreprocessor.process only passes tools, messages and add_generation_prompt to render(). So bos_token renders as an empty string, and the tokenizer doesn't add it later either (kNumAddedBosTokens = 0). On the device the Llama prompt starts with <|start_header_id|>system… and the ids go 128006, 9125, … with no 128000 in front. Gemma 4, LFM 2.5 and Bielik have the same issue. Passing bos_token / eos_token from tokenizer_config.json into render() should fix all four, which is what HF's apply_chat_template does.

Added.

Stop: nice, the session rollback works. One thing is left in the runner: prefill() and generate() still reset isCancelled on entry. Since we're going with the raw runner, a stop that lands right before one of those calls gets lost on our side. In the session, a stop that comes before the prefill starts still waits for the whole prompt to be read before it rolls back, and for Gemma under memory pressure that can be minutes. If the flag were cleared once per turn instead of on entry, both would be fixed. And a small one: a cancelled turn returns finishReason: 'stop', so a 'cancelled' value would make it easier to tell the two apart.

As per docs, the runner's stop is there to stop active inference, not to schedule a cancellation. In chat session stop schedules an eventual cancellation and a stop that comes before first prefill is a very unlikely timing.

Added the 'cancelled' reason for turn end.

Seed: seed++ per call doesn't help within one reply, because generate() passes the same config to sample() at every step (llmTextRunner.ts:195), so every token is drawn at the same quantile. Doing seed + step inside generate() would fix it, and sample stays pure.

Fixed.

@msluszniak msluszniak left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Tested new lfm2.5 2.6B model on iOS with XNNPACK and MLX and worked correctly on my tests :))

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

3 participants