Match maxNumTokens to the model

What maxNumTokens means

In EngineConfig, LiteRT-LM documents maxNumTokens as "the maximum number of the sum of input and output tokens. It is equivalent to the size of the kv-cache." When it is null, the engine uses the value from the model or the engine default. So this one number decides how long a chat can get before something must give.

Three ways it goes wrong

The NPU case in more detail

In #3508, contributors found that on a Qualcomm SM8750 (Galaxy S25), a Gemma 3 270M build with cache_length 1024 lost every chunk after the first, and the same export with cache_length 896 kept all chunks. They tied it to the prefill mask size going above about 1 MiB at prefill length 128. Another contributor later repeated the 896 vs 1280 pair for Gemma 3 1B on SM8550. Two points from the thread matter for you:

• The safe value is per model. 896 works for 4-head Gemma 3 models at prefill 128. It is not a general rule.
• The thread title says Gemma 4 is fine on both NPUs. On Intel, changing to 896 did not help, so it looks like a different cause there.

Practical choice: use Gemma 4 on the NPU, or run Gemma 3 on CPU or GPU. If you must use Gemma 3 on a Qualcomm NPU, use a bundle built with cache_length 896 and set maxNumTokens to the same value.

Rules for your config

1. Find the value the bundle was built with. Ask the person who exported it, or read the export command (--cache_length in litert-torch export_hf). Do not use the model card's context length.
2. Set maxNumTokens to that value, or lower. Do not leave it at a large guess.
3. Count tokens yourself and stop the chat before the limit, with room for the answer.
4. Test a chat that gets close to the limit on each backend.

import com.google.ai.edge.litertlm.Backend
import com.google.ai.edge.litertlm.EngineConfig

/** Values come from the export, not the model card. */
data class BundleInfo(val path: String, val builtCacheLength: Int)

fun engineConfig(bundle: BundleInfo, backend: Backend): EngineConfig =
    EngineConfig(
        modelPath = bundle.path,
        backend = backend,
        maxNumTokens = bundle.builtCacheLength,
    )

/** Keep room for the reply. Start a new conversation (or summarize) before this. */
fun mustRollOver(usedTokens: Int, maxReplyTokens: Int, bundle: BundleInfo): Boolean =
    usedTokens + maxReplyTokens > (bundle.builtCacheLength * 0.9).toInt()

On iOS, Conversation.getTokenCount() exists in the Swift API on main. Use it to know how full the context is.

import LiteRTLM

func mustRollOver(_ conversation: Conversation, maxReplyTokens: Int, builtCacheLength: Int) throws -> Bool {
    let used = try conversation.getTokenCount()
    return used + maxReplyTokens > Int(Double(builtCacheLength) * 0.9)
}

The 0.9 margin is our choice. #2799 crashed close to the limit, so do not plan to fill it.

Test for silent prompt loss

Put the important fact at the end of a long prompt, after 300 tokens of filler, and ask about it. If the answer ignores it on NPU but not on CPU, you have the #3508 problem. Add this prompt to the parity set in the next lesson.

Sources

• LiteRT-LM repository (Google AI Edge)
• LiteRT-LM issue #2799