Debugging Litert Gpu Accuracy

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).

Copy code
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:

   bashCopy code
   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:

   bashCopy code
   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:

bashCopy code
/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:

   bashCopy code
   # 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:

bashCopy code
./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 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:

Use two targeted probes (full code in 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:

textCopy code
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):

   bashCopy code
   /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

Contributions

To contribute or modify this skill or its references, follow the contribution guidelines before making changes.