Catch silent GPU errors with a parity test

What the issues show

For #3913, a contributor traced the cause to the rotary position being cast to a float tensor in the model graph. float16 cannot represent every integer above 2048, so errors start after position 2048. The reporter said float32 fixed it but roughly halved decode speed (8.0 s to 16.1 s for a 250 word answer) and added about 0.45 GB of memory, so they use float32 only for long contexts. Google said they reproduced it.

Lesson: GPU float16 can be right on short prompts and wrong on long ones. You only see it if you compare.

The harness

Use greedy decoding so runs repeat: SamplerConfig(topK = 1, topP = 1.0, temperature = 0.0). Each case has a checker that does not need a judge model: the exact number, the exact word, or "no <pad>".

{"id": "copy-48-digits-2500", "prompt_file": "prompts/copy_digits_2500.txt", "must_contain": ["743486394834120"], "must_not_contain": ["<pad>"]}
{"id": "fact-at-end-300", "prompt_file": "prompts/fact_at_end_300.txt", "must_contain": ["4.2 million"], "must_not_contain": []}
{"id": "short-math", "prompt": "What is 17 plus 26? Answer with the number only.", "must_contain": ["43"], "must_not_contain": ["<pad>"]}

Include prompts below 1,000 tokens, around 2,000, and above 3,000. #3913 and the follow-up suite there showed failures grow with prompt length.

Run it on the phone (Kotlin)

import com.google.ai.edge.litertlm.ActivationDataType
import com.google.ai.edge.litertlm.Backend
import com.google.ai.edge.litertlm.ConversationConfig
import com.google.ai.edge.litertlm.Engine
import com.google.ai.edge.litertlm.EngineConfig
import com.google.ai.edge.litertlm.SamplerConfig
import org.json.JSONObject

data class Variant(val name: String, val backend: Backend, val activation: ActivationDataType?)

val variants = listOf(
    Variant("cpu", Backend.CPU(), null),
    Variant("gpu_fp16", Backend.GPU(), ActivationDataType.FLOAT16),
    Variant("gpu_fp32", Backend.GPU(), ActivationDataType.FLOAT32),
)

fun runParity(modelPath: String, maxTokens: Int, cases: List<Pair<String, String>>): List<JSONObject> {
    val rows = mutableListOf<JSONObject>()
    for (v in variants) {
        val config = EngineConfig(modelPath = modelPath, backend = v.backend,
            maxNumTokens = maxTokens, activationDataType = v.activation)
        Engine(config).use { engine ->
            engine.initialize()
            for ((id, prompt) in cases) {
                val greedy = ConversationConfig(samplerConfig = SamplerConfig(topK = 1, topP = 1.0, temperature = 0.0))
                engine.createConversation(greedy).use { c ->   // fresh conversation per case
                    val t0 = System.nanoTime()
                    val reply = c.sendMessage(prompt).toString()
                    rows += JSONObject()
                        .put("soc", android.os.Build.SOC_MODEL)
                        .put("variant", v.name).put("case", id)
                        .put("ms", (System.nanoTime() - t0) / 1_000_000)
                        .put("reply", reply)
                }
            }
        }
    }
    return rows
}

Build.SOC_MODEL needs API level 31. EngineConfig.activationDataType exists in Config.kt on main; its docs say the default is usually float16 for GPU and float32 for CPU, and that an unsupported type makes initialize() fail. Write the rows to a JSONL file and pull it with adb.

Run it on iOS (Swift)

import LiteRTLM

struct ParityRow: Codable { let variant: String; let caseID: String; let ms: Int; let reply: String }

func runParity(modelPath: String, maxTokens: Int, cases: [(String, String)]) async throws -> [ParityRow] {
    let variants: [(String, Backend, ActivationDataType?)] = [
        ("cpu", .cpu(), nil), ("gpu_fp16", .gpu, .float16), ("gpu_fp32", .gpu, .float32),
    ]
    var rows: [ParityRow] = []
    for (name, backend, activation) in variants {
        let config = try EngineConfig(modelPath: modelPath, backend: backend,
                                      maxNumTokens: maxTokens, activationDataType: activation)
        let engine = Engine(engineConfig: config)
        try await engine.initialize()
        for (id, prompt) in cases {
            let conversation = try await engine.createConversation()
            let start = Date()
            let reply = try await conversation.sendMessage(Message(prompt))
            rows.append(ParityRow(variant: name, caseID: id,
                                  ms: Int(Date().timeIntervalSince(start) * 1000),
                                  reply: reply.toString))
        }
    }
    return rows
}

Pass a greedy SamplerConfig through ConversationConfig in your version so runs repeat. Check the initializer labels in your installed package.

Score and build the allowlist (Python on the Mac)

import json
import sys
from collections import defaultdict
from pathlib import Path

cases = {}
for line in Path("golden.jsonl").read_text().splitlines():
    case = json.loads(line)
    cases[case["id"]] = case

def passed(case: dict, reply: str) -> bool:
    return all(s in reply for s in case["must_contain"]) and not any(s in reply for s in case["must_not_contain"])

score = defaultdict(lambda: [0, 0])  # (soc, variant) -> [passed, total]
for path in sys.argv[1:]:
    for line in Path(path).read_text().splitlines():
        row = json.loads(line)
        key = (row["soc"], row["variant"])
        score[key][1] += 1
        score[key][0] += passed(cases[row["case"]], row["reply"])

allowlist = defaultdict(list)
for (soc, variant), (ok, total) in sorted(score.items()):
    print(f"{soc:12} {variant:9} {ok}/{total}")
    cpu_ok, cpu_total = score[(soc, "cpu")]
    if variant != "cpu" and ok == total and cpu_ok == cpu_total:
        allowlist[soc].append(variant)

Path("gpu_allowlist.json").write_text(json.dumps(allowlist, indent=2))

adb pull /sdcard/Android/data/your.app/files/parity.jsonl parity-s24.jsonl
uv run python score.py parity-*.jsonl

The rule is strict on purpose: a GPU variant goes on the list for a SoC only if it passes every case that the CPU passes.

Pick the backend at runtime

fun chooseBackend(allowlist: Map<String, List<String>>, promptTokens: Int): Variant {
    val allowed = allowlist[android.os.Build.SOC_MODEL].orEmpty()
    return when {
        "gpu_fp16" in allowed && promptTokens < 2048 -> variants[1]
        "gpu_fp32" in allowed -> variants[2]
        else -> variants[0]   // CPU is the safe default
    }
}

The 2048 cut follows the #3913 analysis. Even on an allowlisted SoC, long prompts go to float32 or CPU. Ship the allowlist as a remote config, so you can remove a SoC without an app update. Another GPU problem, #2421, ends multi-turn GPU chats on Mali-G715 with CL_INVALID_COMMAND_QUEUE. If sendMessage throws on GPU, close safely and retry once on CPU.

For a broader test setup with CI gates, see On-Device AI Eval Harness for CI.

Sources

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