Optimize on-device AI: downscale images, smaller KV cache, capped history, NPU memory, CPU threads

[animeshahilya]

Detailed description

This PR stacks on #162 (Snapdragon NPU support) and optimizes the full on-device inference path: what bytes go in, how big the KV cache is, how history grows, how memory is gated, and how CPU threads are used. Cloud behavior is untouched. Every change is confined to the on-device branch with fallback to the original bytes and settings on any failure.

1. Downscale screenshots for the on-device model (largest latency win)

Problem in detail: every Describe image and Describe screen request encodes the screenshot at full phone resolution (typically 1080x2400, about 2.5 megapixels) via DataFieldUtils.encodeImageToByteArray at JPEG quality 80. That exact byte array travels through GeminiActor and GeminiRestEndpoint into LocalGemmaRunner, which base64-decodes it and hands it to LiteRT-LM as Content.ImageBytes. The vision encoder bills every pixel as prefill work: vision tokens scale with image area, prefill time scales with total input tokens, and peak memory scales with both. On a mid-range phone the model card shows 1.8 s time-to-first-token on CPU vs 0.3 s on GPU; most of that gap on image tasks is vision prefill, not the text prompt. Sending a 2.5 MP image to a 2-4B phone model is pure waste: Gemma 4 E2B vision was trained to read UI screenshots at far lower resolution, and text in screenshots stays legible at 1024 px on the long side.

Change: LocalGemmaRunner now runs downscaleForOnDevice on the JPEG just before model.generate, only on the local path. It reads bounds with inJustDecodeBounds, returns the original when already at or under 1024 px on the long side, otherwise decodes with a power-of-two inSampleSize for cheap close scaling, then createScaledBitmap to exactly 1024 max side and re-encodes at JPEG 72. Any exception or null decode returns the original bytes, so unit tests without native bitmap support and corrupt images never break. Cloud requests never see this function: the downscale sits after the cloud/local fork in LocalGemmaRequestPerformer, inside the worker that already runs off the main thread.

Effect: roughly a quarter of the pixels (1024 vs 2400 long side at same aspect), proportionally fewer vision tokens, faster prefill, smaller binder-exempt temp file in RemoteLocalLlm, and less heat per describe. Text legibility is preserved because 1024 px still resolves UI type, and the text prompt already carries the accessibility node list with exact labels.

2. Smaller KV cache for E2B: maxTokens 6144 to 4096

Problem in detail: LocalModel sets maxTokens 6144 for E2B, which becomes EngineConfig.maxNumTokens, which sizes the KV cache. The cache is allocated at load time, so it inflates load latency and resident memory on every launch, and on 6 GB phones it is the difference between surviving and being killed by Android (MSG_LOW_MEMORY and MSG_STOPPED paths). Actual need: instructions are about 600 characters, compact node lists are capped at 120 lines of about 60 characters (ScreenTree COMPACT_MAX_LINES and COMPACT_MAX_LABEL_LENGTH), history is now capped (see 5), and answers are 2-4 sentences plus up to 3 image blurbs, a few hundred tokens. Even with vision tokens, the working set fits in about 3000 tokens. 6144 is roughly twice what describe tasks use.

Change: E2B portable plus both E2B NPU builds go to 4096, with a comment showing the budget math. E4B and 12B stay at 6144 (larger models, same prompt shape, keep conservative). NPU comment notes the saving matters more there because NPU KV cache sits in limited on-device SRAM. No API change: RemoteLocalLlm already forwards model.maxTokens, and null still means model default for experimental entries.

3. Cap Screen Q&A history at 2000 chars

Problem in detail: ScreenOverviewRequester.requestQuery joins the entire chatHistory with no bound, then resends it with the full screenshot on every turn. Each turn therefore grows the prompt: turn N pays prefill for all N-1 previous turns plus the image. Long sessions blow past the KV cache set in (2), slow every answer linearly, and eventually OOM the :ondeviceai process, which the user hears as the model stopped message.

Change: cappedHistory keeps the tail within MAX_HISTORY_CHARS 2000, cutting at a newline so the model never sees a half turn. The current screen nodes, screenshot, and latest question are untouched; only stale turns are dropped, which are the cheapest tokens to lose because the image is resent anyway. Constant is in the companion object next to TAG for easy tuning.

4. NPU-aware free-memory gate

Problem in detail: OnDeviceAiSettings.hasFreeMemoryFor requires availableBytes >= modelBytes, and LocalModelService.checkFreeMemory enforces it before every load. That is correct for CPU, where weights live in CPU RAM (model card: 1733 MB for E2B). NPU weights live on Hexagon SRAM and the CPU footprint tracks the GPU one (676 MB), about 40 percent of file size. With the 1:1 check, Snapdragon NPU phones fail the gate and speak not enough free memory even though the load would succeed, negating PR #162.

Change: new hasFreeMemoryForModel scales required bytes to 60 percent for requiresNpu models, documented with the card numbers. checkFreeMemory takes an isNpu flag plumbed from the existing KEY_USE_NPU bundle value, logs it, and applies the 0.6 factor. Old hasFreeMemoryFor is kept for existing callers and tests. Factor 0.6 is conservative (true ratio is about 0.39) to leave headroom for the vision encoder and binder temp file.

5. Explicit CPU thread count

Problem in detail: Backend.CPU() with no thread count leaves XNNPack thread selection to the runtime default, which can oversubscribe on 8-core flagships and contend with TalkBack's own main-thread event handling on the same cores. Model-card CPU numbers use 4 threads.

Change: Backend.CPU(threadCount = minOf(4, availableProcessors)) in LiteRtLmLocalLlm.loadEngine, with a comment. GPU and NPU paths are untouched. One line, no behavior change on 4-core phones, less jitter on 8-core ones.

What was deliberately left out

Speculative decoding, conversation reuse across turns, constrained JSON decoding, APK flavor slimming, and parallel downloads were audited and excluded to keep this PR reviewable and risk-free. They are good follow-ups and are noted here so reviewers can see the scoping: speculative decoding needs post-May-2026 model files plus memory A/B; conversation reuse needs service lifecycle work; constrained decoding needs LiteRT-LM ResponseFormat wiring; APK slimming needs a wear-build check because the fragment is in the shared source set.

Verification

• Existing OnDeviceAiSettingsTest unaffected: old hasFreeMemoryFor signature kept, null SoC still filters NPU models, defaults still fall back to E2B.
• LocalGemmaRunnerTest unaffected: downscale catches all exceptions and returns input, so fake JPEG bytes pass through.
• Manual checks needed: on-device Describe image latency before and after (logcat Answered in X ms), prompt length log (Prompt is X characters), NPU load with tight memory, and a 10-turn Screen Q&A session to confirm history stops growing.
• Base: this branch stacks on snapdragon-npu-gemma (Add Snapdragon NPU acceleration with SoC-specific Gemma 4 E2B models #162). Merge Add Snapdragon NPU acceleration with SoC-specific Gemma 4 E2B models #162 first, then retarget this PR to master.

[trypsynth]

Thanks! Several of these are good wins on their own, and none of them depend on the NPU:

• Downscaling screenshots before they reach the on-device model
• Capping Screen Q&A history
• Setting an explicit CPU thread count
• Probably the smaller E2B token budget too, if it holds up on real describe requests

Could you split those into their own PR against master? That way they can land without waiting on #162. Keep the NPU-specific memory check here, or fold it into #162.

A couple of things for the split PR:

• Please confirm the 4096-token budget with a real busy screen. A long node list, plus an image and a couple of Q&A turns, would show whether it fits. If a request overflows, it should fail clearly, not just get cut off.
• Please add a small test for cappedHistory. It should check that the history keeps whole turns and always keeps the latest one.