Document the rank-3 attention GPU miscompute and the v2 fix

README.md

@@ -34,16 +34,25 @@ turns a **point prompt** into segmentation masks **per tap** (a few ms each) —
 | Inputs | `image_embeddings [1,256,64,64]`, `sparse_prompt [1,2,256]`, `feat_s1 [1,64,128,128]`, `feat_s0 [1,32,256,256]` |
 | Outputs | `pred_masks [1,3,256,256]` (logits, 3 multimask candidates), `iou_scores [1,3]` |
 | Precision / size | FP16, **17 MB** |
-| Device | Pixel 8a — GPU-resident (358/358 LITERT_CL) but run on **CPU** for correct masks (see below) |
+| File | **`sam2_tiny_mask_decoder_v2_fp16.tflite`** (recommended). `sam2_tiny_mask_decoder_fp16.tflite` is the earlier build, kept for reference — see the note below. |
+| Device | Pixel 8a — **fully GPU** (LITERT_CL), correct masks, ~7 ms/tap |
 | Op set | `banned ops = NONE`, `>4-D tensors = 0` (BATCH_MATMUL ×15, SOFTMAX ×7, GELU ×2, CONV_2D ×2) |
 
-> **⚠ Residency ≠ correctness.** This decoder fully delegates to the LiteRT GPU (358/358 LITERT_CL
-> nodes), but on the Pixel 8a its **GPU fp16 output is numerically wrong** — a face tap that the CPU
-> decoder segments at IoU ≈ 0.62 collapses to ≈ 0.10 with the mask on the background under the GPU
-> delegate. The companion encoder's GPU output is fine (encoder-GPU + decoder-CPU matches all-CPU
-> exactly). It is **not** LayerNorm (plain vs. overflow-safe give the same wrong GPU result); the
-> offending op is still being localized with a per-op GPU dump. **Run this decoder on CPU** (it is tiny
-> and fast); the heavy image encoder is the part that benefits from the GPU.
+> **⚠ Residency ≠ correctness — and why v2 exists.** The first build (`sam2_tiny_mask_decoder_fp16.tflite`)
+> fully delegated to the GPU (358/358 LITERT_CL nodes, `banned ops = NONE`, `>4-D = 0`, desktop parity
+> corr 1.0) yet returned **silently wrong masks on the Pixel 8a GPU** (corr 0.265 vs CPU; a face tap at
+> IoU ≈ 0.62 on CPU collapsed to ≈ 0.10 with the mask on the background).
+>
+> The cause was found by device A/B bisection: its attention was written with the **batch dim collapsed**
+> (q/k/v shaped `[heads, N, d]`, rank 3). The GPU delegate mis-computes that form. It is **not** an fp16
+> problem (forcing fp32 GPU compute still gives corr 0.473) and **not** LayerNorm (plain and
+> overflow-safe LN give the same wrong result). The mask head's rank-2 matmul is innocent.
+>
+> **`v2` keeps the leading batch dim (rank-4 SDPA, `[1, heads, N, d]`).** Host numerics are identical
+> (eager cos 0.999999); on the Pixel 8a GPU it restores **corr 0.9998 / binary-IoU 0.999** vs CPU and is
+> **~20 % faster** (6.8 ms vs 8.5 ms). Inputs and outputs are unchanged, so v2 is a drop-in replacement.
+> Note the companion encoder's rank-3 SDPA *is* GPU-correct — a healthy sibling graph proves nothing;
+> only a numeric GPU-vs-CPU check on device catches this.
 
 ## Pipeline (how the inputs are produced)
 
@@ -112,9 +121,9 @@ card). Mask boundaries are carried by the near-exact high-resolution features, s
 // once per image - encoder on GPU (decoder-ready v2 variant from the companion repo)
 val enc = CompiledModel.create(context.assets, "sam2_tiny_image_encoder_v2_fp16.tflite",
     CompiledModel.Options(Accelerator.GPU), null)
-// per tap - decoder on CPU (GPU-resident but fp16-incorrect on device; see the note above)
-val dec = CompiledModel.create(context.assets, "sam2_tiny_mask_decoder_fp16.tflite",
-    CompiledModel.Options(Accelerator.CPU), null)
+// per tap - decoder on GPU (v2: rank-4 attention, GPU-correct; see the note above)
+val dec = CompiledModel.create(context.assets, "sam2_tiny_mask_decoder_v2_fp16.tflite",
+    CompiledModel.Options(Accelerator.GPU), null)
 // dec inputs: 0 image_embeddings[1,256,64,64], 1 sparse[1,2,256],
 //             2 feat_s1[1,64,128,128], 3 feat_s0[1,32,256,256]
 // dec outputs: pred_masks[1,3,256,256] logits, iou_scores[1,3] -> argmax(iou), threshold 0
@@ -143,8 +152,8 @@ coord = 2 * np.pi * ((2 * (np.array([px, py], np.float32) + 0.5) / 1024 - 1) @ p
 tok0 = np.concatenate([np.sin(coord), np.cos(coord)]) + pe1
 sparse = np.stack([tok0, nap])[None].astype(np.float32)               # [1,2,256]
 
-# 3) decode masks (CPU)
-dec = Interpreter(model_path="sam2_tiny_mask_decoder_fp16.tflite"); dec.allocate_tensors()
+# 3) decode masks
+dec = Interpreter(model_path="sam2_tiny_mask_decoder_v2_fp16.tflite"); dec.allocate_tensors()
 feed = {(1,2,256): sparse}; feed.update(eo)                           # match inputs by shape
 for d in dec.get_input_details(): dec.set_tensor(d["index"], feed[tuple(d["shape"])])
 dec.invoke()