---
license: apache-2.0
library_name: litert
pipeline_tag: mask-generation
tags:
  - litert
  - tflite
  - sam2
  - segment-anything
  - mask-decoder
  - interactive-segmentation
  - on-device
  - gpu
base_model: facebook/sam2.1-hiera-tiny
---

# SAM 2.1 (Hiera-Tiny) mask decoder — LiteRT GPU

On-device **LiteRT / TFLite** conversion of the **prompt-conditioned mask decoder** of
[**SAM 2.1 Hiera-Tiny**](https://huggingface.co/facebook/sam2.1-hiera-tiny) (Meta, Apache-2.0),
running **fully on the mobile GPU** via the LiteRT `CompiledModel` API (ML Drift / `LITERT_CL` delegate).

![SAM 2.1 tap-to-segment running on-device (LiteRT)](demo.gif)

This is the lightweight, per-click half of the SAM 2 image path. Pair it with the
[**SAM 2.1 Hiera-Tiny image encoder**](https://huggingface.co/litert-community/SAM2.1-Hiera-Tiny-Image-Encoder)
(run **once** per image, and it is the heavy half — 208.0 ms on a Galaxy S26 GPU against this decoder's 10.08 ms): the encoder produces the multi-scale feature pyramid, and this decoder
turns a **point prompt** into segmentation masks **per tap** (a few ms each) — interactive "tap to segment".

| | |
|---|---|
| Task | Mask decoder for promptable segmentation (SAM 2 image path) |
| Architecture | 2-layer two-way transformer (token↔image cross-attention) + mask up-sampler |
| 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** |
| 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, 425 / 425 nodes), correct masks, **30.9 ms/tap** (CPU 222.8 ms on the same phone); Galaxy S26 GPU 10.55 ms — see below |
| Op set | `banned ops = NONE`, `>4-D tensors = 0` (BATCH_MATMUL ×15, SOFTMAX ×7, GELU ×2, CONV_2D ×2) |

> **⚠ 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)

```
RGB image ──> image encoder (run once) ──> image_embeddings[1,256,64,64], feat_s1[1,64,128,128], feat_s0[1,32,256,256]
tap (x,y in 1024-space) ──> prompt encode (host-side, see below) ──> sparse_prompt[1,2,256]
                                              │
       image_embeddings + feat_s0/s1 + sparse_prompt ──> THIS decoder ──> 3 masks + 3 IoU
       pick argmax(IoU) ──> upsample 256×256 logits to image size ──> threshold > 0 ──> overlay
```

The decoder uses the **encoder variant that already folds `conv_s0` / `conv_s1` + `no_memory`** so its
outputs are directly decoder-ready (no host reshaping between the two models).

### Host-side prompt encoding (single positive point)

The tiny point→token step (a sin/cos positional encoding) is done on the host to keep the GPU graph
sin/cos-free. For a positive click `(x, y)` in 1024×1024 model space, with the bundled constants
`posmat [2,128]`, `point_embed[1] [256]`, `not_a_point [256]`:

```
c      = (([x, y]) + 0.5) / 1024          # normalize, half-pixel shift
c      = 2*c - 1
coord  = 2*pi * (c @ posmat)              # [128]
token0 = concat(sin(coord), cos(coord)) + point_embed[1]   # the positive point
token1 = not_a_point                      # the padding point
sparse_prompt = [[token0, token1]]        # [1, 2, 256]
```

This matches the upstream `Sam2PromptEncoder` to ~3.7e-7.

## GPU-clean conversion (what was re-authored)

Converted with `litert-torch`, **model-side rewrites only — no converter patch**, each weights-faithful:

1. **Two-way attention (×7)**: re-expressed as **3-D batched SDPA** `[heads, N, d]` (a 4-D SDPA makes the
   delegate emit a `BROADCAST_TO`).
2. **Mask up-sampler `ConvTranspose2d` (×2)**: replaced with the exact **zero-stuff + `Conv2d`** identity
   (`TRANSPOSE_CONV` is rejected on Pixel 8a; this is numerically identical, not a bilinear approximation).
3. **Mask head**: the `hyper_in @ upscaled` mask projection is kept **≤4-D** (the upstream
   `[1,1,4,256,256]` 5-D tensor is collapsed; batch/point-batch are 1).
4. **LayerNorm (×9)**: scale-before-square **SafeLayerNorm** (fp16-overflow-safe, mathematically identical).
5. **Constants baked**: `image_positional_embeddings` and the no-mask dense prompt are baked as buffers.
6. **Multimask path**: static slice `[1:]` of the 3 candidate masks — no dynamic-stability
   `argmax` / `gather` / `where`.

## Fidelity (honest)

Eager re-authoring is **numerically exact** (`cos = 1.000`). End-to-end through the two FP16 tflite models
(encoder → host prompt-encode → decoder) vs the PyTorch reference, for a center click:

| Metric | value |
|---|---|
| mask logits cosine | **0.999999** |
| binary mask IoU (threshold 0) | **0.99964** |
| IoU-score head | ref `[0.936, 0.022, 0.399]` vs got `[0.936, 0.022, 0.399]` |

The deepest `64×64` image embedding drifts slightly on the GPU (true-fp16 deep attention; see the encoder
card). Mask boundaries are carried by the near-exact high-resolution features, so mask quality holds.

## Minimal usage

**Android (Kotlin, CompiledModel GPU)**

```kotlin
// 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 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
```

**Python (desktop verification)**

```python
MEAN = np.array([0.485, 0.456, 0.406], np.float32)
STD  = np.array([0.229, 0.224, 0.225], np.float32)
import numpy as np
from PIL import Image
from ai_edge_litert.interpreter import Interpreter

# 1) encode once (decoder-ready v2 encoder from the companion encoder repo)
img = Image.open("photo.jpg").convert("RGB").resize((1024, 1024))
x = ((np.asarray(img, np.float32) / 255 - MEAN) / STD).transpose(2, 0, 1)[None]
enc = Interpreter(model_path="sam2_tiny_image_encoder_v2_fp16.tflite"); enc.allocate_tensors()
enc.set_tensor(enc.get_input_details()[0]["index"], x); enc.invoke()
eo = {tuple(d["shape"]): enc.get_tensor(d["index"]) for d in enc.get_output_details()}

# 2) host prompt-encode one positive tap (px, py) in 1024-space (constants: this repo)
px, py = 512, 384
posmat, pe1, nap = np.split(np.fromfile("prompt_encode_const.bin", np.float32), [256, 512])
coord = 2 * np.pi * ((2 * (np.array([px, py], np.float32) + 0.5) / 1024 - 1) @ posmat.reshape(2, 128))
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
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()
o = {len(d["shape"]): dec.get_tensor(d["index"]) for d in dec.get_output_details()}
masks, iou = o[4], o[2]                                               # [1,3,256,256], [1,3]
best = masks[0, iou[0].argmax()] > 0                                  # [256,256] binary mask
Image.fromarray(best.astype(np.uint8) * 255).resize(Image.open("photo.jpg").size).save("mask.png")
```

## Training data & PII

SAM 2 was trained by Meta on **SA-1B** (licensed photos) and **SA-V** (licensed videos) with
model-in-the-loop mask annotation. No new training was performed for this conversion — it is a
weights-faithful format change of the public `facebook/sam2.1-hiera-tiny` checkpoint. Because the source
data is real-world imagery it may incidentally contain people, faces, vehicles, signage and other PII; no
PII was deliberately collected and this conversion adds none. Apply your own content/PII filtering as
appropriate. See the [SAM 2 release](https://github.com/facebookresearch/sam2) and
[paper](https://arxiv.org/abs/2408.00714) for full dataset details.

## Performance

Measured on a **Pixel 8a** (Tensor G3, Android 16) with the standard TFLite [`benchmark_model`](https://ai.google.dev/edge/litert/models/measurement) tool — 10 warm-up runs then 50 timed runs, reported as the tool's mean.

| Runtime | Backend | Graph on GPU | Latency |
|---|---|---|---|
| TFLite `benchmark_model` (`TfLiteGpuDelegateV2`) — `sam2_tiny_mask_decoder_v2_fp16.tflite` | GPU (OpenCL) | 71 / 425 | 247.6 ms |
| TFLite `benchmark_model` (`TfLiteGpuDelegateV2`) — `sam2_tiny_mask_decoder_fp16.tflite` | GPU (OpenCL) | 294 / 378 | 123.7 ms |
| TFLite `benchmark_model` — `sam2_tiny_mask_decoder_v2_fp16.tflite` | CPU (XNNPACK, 4 threads) | — | 88.1 ms |
| TFLite `benchmark_model` — `sam2_tiny_mask_decoder_fp16.tflite` | CPU (XNNPACK, 4 threads) | — | 86.8 ms |

**Any on-device figure recorded when this model shipped came from a different runtime.** It was taken through LiteRT's own `CompiledModel` accelerator (logcat reports it as `LITERT_CL`), which is the path the Kotlin sample app and the LiteRT API use, and it appears elsewhere on this card. The rows above are the classic TFLite OpenCL delegate, measured with a tool anyone can download and re-run. The two are not comparable, so read the rows above as a reproducible floor rather than as this model's speed on LiteRT.

On this delegate the CPU is the faster choice for `sam2_tiny_mask_decoder_v2_fp16.tflite` (88.1 ms on CPU against 247.6 ms on GPU), `sam2_tiny_mask_decoder_fp16.tflite` (86.8 ms on CPU against 123.7 ms on GPU) — worth knowing before you reach for the GPU on a mid-range phone.

Note that the GPU does not take the whole graph here (71 / 425 in `sam2_tiny_mask_decoder_v2_fp16.tflite`, 294 / 378 in `sam2_tiny_mask_decoder_fp16.tflite`); the remainder runs on the CPU and the split costs a per-partition round trip.

## Pixel 8a — LiteRT `CompiledModel`

Both files take the **whole graph** on the Pixel 8a GPU through LiteRT's own accelerator:
`Replacing 425 out of 425 node(s) with delegate (LITERT_CL)` for the v2 build, 378 / 378 for the
earlier one. That is the path the Kotlin snippet above uses, and it is not the classic delegate
measured in the previous section — same phone, same file, 425 / 425 here against 71 / 425 there.

| file | backend | graph on GPU | inference (median / min) | load |
|---|---|---|---:|---:|
| `sam2_tiny_mask_decoder_v2_fp16.tflite` | GPU (`LITERT_CL`) | 425 / 425 | 30.88 ms / 29.51 ms | 1679 ms |
| `sam2_tiny_mask_decoder_v2_fp16.tflite` | CPU (XNNPACK) | — | 222.8 ms / 216.1 ms | 48 ms |
| `sam2_tiny_mask_decoder_fp16.tflite` | GPU (`LITERT_CL`) | 378 / 378 | 30.06 ms / 29.34 ms | 2558 ms |
| `sam2_tiny_mask_decoder_fp16.tflite` | CPU (XNNPACK) | — | 218.7 ms / 214.0 ms | 47 ms |

Measured on a **Pixel 8a** (Tensor G3, Android 16) with LiteRT `CompiledModel` 2.2.0, one
accelerator per process, 5 warm-up runs then N=50 timed runs, median reported. Every run held
thermal status `NONE` throughout. Headroom 0.58–0.62, where 1.0 is the throttling threshold.

On this phone the GPU is **7.2x faster than the CPU** for the v2 build (30.88 ms against
222.8 ms) and pays for it at load — 1.7 s against 48 ms. A Galaxy S26 runs the same file at
10.55 ms on its GPU (see below), about a 3x device gap on the same code path.

Two notes on figures elsewhere on this card. The correctness note below records 358 / 358 nodes
for `sam2_tiny_mask_decoder_fp16.tflite`; this run sees 378 / 378 on LiteRT 2.2.0, and node
counts move with the runtime version. And **correctness was not re-checked here** — this harness
feeds synthetic input buffers and times them, so the mask-quality note still decides which file
to ship.

## Snapdragon NPU (Hexagon)

- `sam2_tiny_mask_decoder_fp16.tflite` — the NPU is **1.25x faster** than the GPU (8.08 ms against 10.08 ms) and loads 11.35x faster (111 ms against 1258 ms).
- `sam2_tiny_mask_decoder_v2_fp16.tflite` — the NPU is **1.31x faster** than the GPU (8.04 ms against 10.55 ms) and loads 10.69x faster (110 ms against 1177 ms).

| file | backend | compiled | inference (median / min) | load |
|---|---|---|---:|---:|
| `sam2_tiny_mask_decoder_fp16.tflite` | NPU (Hexagon v81) | on-device JIT | 8.08 ms / 7.96 ms | 111 ms |
| `sam2_tiny_mask_decoder_fp16.tflite` | GPU (Adreno) | — | 10.08 ms / 9.23 ms | 1258 ms |
| `sam2_tiny_mask_decoder_v2_fp16.tflite` | NPU (Hexagon v81) | on-device JIT | 8.04 ms / 7.96 ms | 110 ms |
| `sam2_tiny_mask_decoder_v2_fp16.tflite` | GPU (Adreno) | — | 10.55 ms / 10.08 ms | 1177 ms |

Measured on a **Samsung Galaxy S26** (Snapdragon 8 Elite Gen 5 / SM8850, Hexagon v81, Android 16) with LiteRT `CompiledModel` 2.2.0, one accelerator per process, 5 warm-up runs then N=50 timed runs, median reported. Every run held thermal status `NONE` throughout. Headroom 0.79–0.80, where 1.0 is the throttling threshold.

**The NPU rows ran the published file unchanged.** LiteRT compiled it for the Hexagon on the device at first load. Each first compile took 4.4 s here. The `load` column above is the cached load every later run pays. Recipe and the runtime libraries it needs: [NPU guide](https://github.com/john-rocky/hf-to-litertlm/blob/main/docs/android-npu.md).

The sweep timed both files; it did not re-check mask correctness on Adreno. The correctness note above still decides which file to ship.

GPU wiring: [GPU guide](https://github.com/john-rocky/hf-to-litertlm/blob/main/docs/android-gpu.md).

## Raspberry Pi 5 (CPU)

Measured on a Raspberry Pi 5 Model B Rev 1.1 (8 GB, Raspberry Pi OS 64-bit) with the LiteRT [`benchmark_model`](https://ai.google.dev/edge/litert/models/measurement) tool from litert-cli-nightly 0.2.0.dev20260805: CPU inference (XNNPACK, 4 threads), 3 invocations per file of 10 warm-up plus 50 timed runs (the tool caps a phase at 150 s, so very slow graphs run fewer — the Runs column is the actual timed total). The latency is the median across invocations; the spread is the min–max over all timed runs. No thermal throttling occurred during these runs (`vcgencmd get_throttled` stayed `0x0`).

| File | Inference (median) | Spread (min–max) | Runs | Peak memory |
|---|---:|---:|---:|---:|
| `sam2_tiny_mask_decoder_fp16.tflite` | 156.2 ms | 154.3–158.7 ms | 150 | 135 MB |
| `sam2_tiny_mask_decoder_v2_fp16.tflite` | 158.3 ms | 157.1–161.4 ms | 150 | 135 MB |

## License

Apache-2.0, inherited from the upstream [SAM 2.1](https://huggingface.co/facebook/sam2.1-hiera-tiny).
This is a format conversion; all credit to the original authors (Meta AI).