Skills DirectorySearch all skills ↗AI agent skill · by @google-ai-edge
Debugging Litert Gpu Accuracy
Diagnoses and fixes on-device GPU numerical corruption, garbage generation, zeroed KV caches, and ML Drift delegate lowering bugs for LiteRT and LiteRT-LM models on Android (OpenCL and WebGPU). Use when a model produces correct output on CPU but generates garbage, NaNs, or wrong outputs on GPU; when comparing CPU vs GPU execution op-by-op or state-buffer-by-state-buffer; or when debugging composite ops (cache_update, runtime_bmm, sdpa), ring-buffer sliding-window KV caches, or TensorStorageType (BUFFER vs TEXTURE_2D) mismatches in ml_drift. Don't use for NPU/QNN dispatch issues, CPU-only model conversion failures, or adding standard E2E device test targets (use adding-litert-lm-e2e-tests).
Archived 2026-10-23 · Source updated 2026-10-12 · License: Apache-2.0. Check the original repository for the latest version.
SKILL.md instructions
# Debugging LiteRT / LiteRT-LM On-Device GPU Accuracy & ML Drift Bugs
Use this 6-phase funnel to isolate why a LiteRT-LM model produces correct
output on CPU (`--backend=cpu`) but generates garbage, NaNs, or empty/corrupted
state on Android GPU (`--backend=gpu`, OpenCL or WebGPU).
```
Phase 1: Unpack & diff model metadata (.litertlm) against a GPU-good model
│
Phase 2: On-device repro matrix (CPU vs OpenCL vs WebGPU + prompt sweep)
│
Phase 3: Op-by-op accuracy_debugger on extracted TFLite subgraph
│ ├─ Op diverges in isolation ──► Fix op kernel / precision
│ └─ All ops pass (maxdiff≈0) ──► Composite / graph-lowering / state bug
▼
Phase 4: Host-side CPU vs GPU state-buffer dump after prefill_128 & decode
│ (Pinpoints exact corrupted KV cache layers: e.g. local ring vs global)
▼
Phase 5: ML Drift descriptor dump & in-shader sentinel readback probes
│ (Bisects runtime param chain & TensorStorageType mismatches)
▼
Phase 6: Fix in ml_drift / delegate, unit-test GpuModel invariants, & validate
```
--------------------------------------------------------------------------------
## Critical Environment Gotchas (Read First)
1. **Delete the on-device ML Drift program cache before every GPU run after
editing any kernel, parser, or selector code.** Otherwise the device
silently reuses the cached compiled OpenCL/WebGPU program and your C++ edits
never execute:
```bash
adb shell "rm -f /data/local/tmp/{model_dir}/*mldrift_program_cache.bin"
```
Keep `*_mldrift_weight_cache.bin` and `*.xnnpack_cache` so weight conversion
stays fast.
2. **Injecting device environment variables:** `run_litert_lm.sh` does not
forward custom host env vars or provide a build-only flag. Run
`run_litert_lm.sh --skip_model_push` once to build and push
`libLiteRtOpenClAccelerator.so` and `litert_lm_advanced_main`, then invoke
`adb shell` directly with your probe env vars:
```bash
OUT=/data/local/tmp/{model_dir}
adb shell "rm -f ${OUT}/*mldrift_program_cache.bin"
adb shell "ADSP_LIBRARY_PATH=${OUT} LD_LIBRARY_PATH=${OUT} \
LITERT_LM_DEBUG_DUMP_STATE=1 MLDRIFT_CACHE_PROBE=1 \
${OUT}/litert_lm_advanced_main \
--model_path=${OUT}/model.litertlm --backend=gpu \
--async=false --clear_kv_cache_before_prefill=true \
--convert_weights_on_gpu=true --minloglevel=0 \
--input_prompt=\"hi\""
```
3. **Layering check in `third_party/ml_drift` and `delegate/composite`:**
Adding `#include "third_party/absl/log/absl_log.h"` to kernel files fails
Bazel/Blaze `layering_check`. Use `#include <cstdio>` + `fprintf(stderr,
...)` and `#include <cstdlib>` + `std::getenv(...)` for temporary probes.
--------------------------------------------------------------------------------
## Phase 1: Unpack & Diff Model Metadata Against a Known-Good Model
Unpack the failing `.litertlm` bundle and a known-good GPU model to compare
their `ExecutorMetadataProto`, `LlmMetadataProto`, and embedded `.tflite`
subgraphs:
```bash
/google/bin/releases/arca9-local-blaze-cli/blaze-for-agents run \
//third_party/odml/litert_lm/runtime/util:litertlm_builder_cli -- \
unpack --input_file=/tmp/dbg/failing.litertlm --output_dir=/tmp/dbg/unpacked
```
Check `ExecutorMetadataProto.pbtext` and `LlmMetadataProto.pbtext` for:
- **Hybrid KV cache topology**: Does the model mix `TYPE_LOCAL_*_CACHE`
(sliding-window / ring buffer, e.g. `minimum_sequence_length: 1024`,
`maximum_sequence_length: 1280`) with global caches
(`maximum_sequence_length: 32771`)?
- **`attention_mask_settings`**: Presence of `sliding_window_size` triggers
the `gpu_optimized_single_buffer_cache_` ring-buffer path in LiteRT-LM and
`ring_buffer_sdpa` in the exported graph.
--------------------------------------------------------------------------------
## Phase 2: Reproduce & Narrow the Failure Surface on Device
Run the repro prompt across backends and configurations to eliminate whole
subsystems up front:
1. **CPU vs OpenCL GPU vs WebGPU:**
```bash
# OpenCL GPU
./third_party/odml/litert_lm/runtime/engine/run_litert_lm.sh \
--models={model_name} --target_os=android --backend=gpu \
--skip_model_push --input_prompt="{prompt}"
# WebGPU (--use_webgpu)
./third_party/odml/litert_lm/runtime/engine/run_litert_lm.sh \
--models={model_name} --target_os=android --backend=gpu --use_webgpu \
--skip_model_push --input_prompt="{prompt}"
```
- If **both OpenCL and WebGPU fail identically**, the bug is in
backend-agnostic code: the exported graph, LiteRT-LM host tensor
filling, or shared `ml_drift/common` / `delegate/composite` graph
lowering (such as `LiteRtOpSelector`), **not** an OpenCL-specific driver
or shader bug.
2. **Prompt-length sweep:** Test a minimal prompt (`"hi"` — note that system
prompt wrapping still yields ~14 tokens) vs >128 tokens (multi-chunk
prefill) vs >`sliding_window_size`. If even a 14-token prompt fails on the
very first generated token, wrap-around and multi-chunk boundary math are
ruled out.
--------------------------------------------------------------------------------
## Phase 3: Run Op-by-Op `accuracy_debugger` (and Know Its Blind Spots)
Run `accuracy_debugger` on the unpacked
`Section3_TFLiteModel_tf_lite_prefill_decode.tflite` (`prefill_128` subgraph) to
compare GPU vs CPU reference op-by-op:
```bash
./litert/tools/accuracy_debugger/google/debug_accuracy_qc.sh \
--model_path=/tmp/dbg/unpacked/Section3_TFLiteModel_tf_lite_prefill_decode.tflite \
--signature_name=prefill_128 \
--num_ops=400
```
Inspect the generated CSV (`Index, Op Code, Tensor Name, Max Diff, MSE, Cos Sim,
SNR, PSNR, NaN, Min Ref, Max Ref, Min Accel, Max Accel, Status`):
- **If a standard op shows low `Cos Sim` / high `Max Diff` / `NaN`**: You have
an isolated op numerical/precision bug.
- **Blind spot to remember**: `accuracy_debugger` verifies each op in
isolation and reports `ACCEL_COMPILE_FAILED` on
`shlo.composite(odml.runtime_bmm)` and `shlo.composite(odml.cache_update)`.
Consequently, an entire parameter computation chain (e.g. `SLICE` →
`FLOOR_MOD` → `CONCAT`) can pass with `maxdiff = 0` in `accuracy_debugger`
while still producing garbage in the full compiled graph due to **cross-op
storage-type / producer-rewiring bugs** triggered only when composite ops
are lowered alongside standard ops.
--------------------------------------------------------------------------------
## Phase 4: Diff CPU vs GPU State Buffers After `prefill_128` & `decode`
When `accuracy_debugger` shows standard ops are healthy, instrument
`google3/third_party/odml/litert_lm/runtime/executor/llm_litert_compiled_model_executor.cc`
(see [references/mldrift_shader_probes.md](references/mldrift_shader_probes.md)
for the drop-in snippet) to dump `n`, `sum`, `sumabs`, `min`, `max`, `nan`, and
the first 6 elements (`head=[...]`) of every input and output buffer in
`BindTensorsAndRunPrefill` and `BindTensorsAndRunDecode` under `--async=false`.
Run once with `--backend=cpu` and once with `--backend=gpu` on `"hi"` and diff:
- **Check `param_tensor` (`prefill_in`)**: Verify the 7-channel `int32`
runtime parameter tensor `[start, end, end, ...]` matches on host.
- **Check `kv_cache_k_{layer}` / `kv_cache_v_{layer}` (`prefill_out`)**:
- If **global layers** (e.g. layers 3, 7, 11, ...) match CPU
(`sumabs > 0`) while **local/sliding-window layers** (e.g. layers 0, 1,
2, ...) have `sumabs = 0` (`min=0, max=0`), the ring-buffer branch of
`add_values_to_cache_kernel.cc` is early-exiting without writing.
--------------------------------------------------------------------------------
## Phase 5: Bisect with ML Drift Descriptor & Shader Readback Probes
Understand the parameter contract before probing:
| Slot | Meaning | Consumer |
| :--- | :--- | :--- |
| `0` | `token_index_offset` (`start` or `start % S`) | `add_values_to_cache` (both branches) |
| `1` | `active_tokens` (`end` or `min(end, S)`) | `add_values_to_cache` non-ring bound |
| `2` | `kActiveTokensAlignedIndex` | `runtime_batched_matmul` (`src_end_ch_index` / `dst_end_ch_index`) & `sdpa_transposed` |
| `3` | Ring `update_length` (`end - start`) | `add_values_to_cache` **ring branch only** (`if (X >= update_length) return;`) |
Use two targeted probes (full code in
[references/mldrift_shader_probes.md](references/mldrift_shader_probes.md)):
1. **In-shader param readback (`MLDRIFT_CACHE_PROBE`)**: In
`ml_drift/delegate/composite/add_values_to_cache_kernel.cc`,
bypass the `X >= update_length` gate and overwrite `final_value_k` with
`float4(params.Read(0..3))` and `final_value_v` with
`float4(params.Read(4..7))`. The host `[DBG]` dump of
`kv_cache_k_0 head=[...]` and `kv_cache_v_0 head=[...]` then directly prints
the exact `int32` values the GPU shader read from `args.params`!
2. **Host-side `OperationDef` / `TensorStorageType` dump
(`MLDRIFT_SLICE_PROBE`)**: In
`google3/third_party/ml_drift/common/kernels/strided_slice.cc` (inside
`GetStridedSliceCode`), print `starts.c`, `ends.c`, `src.Channels()`, and
`src.GetStorageType()` (`1 = BUFFER`, `3 = TEXTURE_2D`) for small int32
tensors (`c <= 16`).
### Classic Root-Cause Signature (`b/565413009`)
If `MLDRIFT_CACHE_PROBE` shows `params[0..2]` are valid (`0, 14, 14`) while
`params[3..6]` are garbage (`±inf` in fp16), and `[SLICEDBG]` shows:
```text
starts.c=0 ends.c=3 | src c=7 storage=1 (BUFFER) -> valid
starts.c=3 ends.c=7 | src c=7 storage=3 (TEXTURE_2D) -> GARBAGE (unproduced tensor!)
```
Look at `LiteRtOpSelector::ParamTensorToBuffer` in
`ml_drift/delegate/composite/litert_op_selector.cc`:
- Composite ops (`add_values_to_cache`, `runtime_batched_matmul`,
`sdpa_transposed`) declare `params` via `AddSrcBuffer` (`BUFFER` storage),
whereas standard ops like `StridedSlice` (`rest = param'[3:7]`) consume the
original tensor via `AddSrcTensor` (`TEXTURE_2D`) without consulting
`replaced_tensors_`.
- Calling
`model_builder->UpdateOutputTensor(param_tensor, new_param_tensor.id)`
**re-points the producer node's output** from `param_id` to
`new_param_tensor.id`, leaving `param_id` with **no producer in the
`GpuModel`** so any non-buffer consumer reads uninitialized GPU memory.
- **Correct fix**: Emit `model_builder->Copy(param_tensor, new_param_tensor)`
instead of `UpdateOutputTensor`. When `param_tensor` has no other consumers,
`LinkNodes()` in `google3/third_party/ml_drift/common/merge_nodes.cc`
automatically fuses the single-consumer copy back into the producer.
--------------------------------------------------------------------------------
## Phase 6: Unit-Test the `GpuModel` Invariant & Validate End-to-End
1. **Revert all temporary probes** (`hg revert` on kernel/executor files)
before running final validation.
2. **Write a hermetic `cc_test`** (see
`ml_drift/delegate/composite/litert_op_selector_test.cc`)
that builds a `GpuModel` via `GpuModelBuilder` + `LiteRtOpSelector` and
asserts that **every `in_id` of every `GpuNode` in `gpu_model.nodes` is
present in `defined_tensors`** (graph inputs ∪ const tensors ∪ node
outputs):
```bash
/google/bin/releases/arca9-local-blaze-cli/blaze-for-agents test \
//ml_drift/delegate/composite:litert_op_selector_test
```
3. **Validate on device** across:
- Failing model on OpenCL (`--backend=gpu`)
- Failing model on WebGPU (`--backend=gpu --use_webgpu`)
- A baseline non-sliding-window model
(`--models=kanana_i4_107_kv_1024 --backend=gpu`)
--------------------------------------------------------------------------------
## References
- **Drop-in C++ probes (Executor state dumper, Shader param readback, Slice
storage dumper)**:
[mldrift_shader_probes.md](references/mldrift_shader_probes.md)
## Contributions
To contribute or modify this skill or its references, follow the
[contribution guidelines](references/contributing.md) before making changes.