On-device LiteRT / TFLite conversion of the prompt-conditioned mask decoder of SAM 2.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)

This is the lightweight, per-click half of the SAM 2 image path. Pair it with the SAM 2.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:

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).
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).
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).
LayerNorm (×9): scale-before-square SafeLayerNorm (fp16-overflow-safe, mathematically identical).
Constants baked: image_positional_embeddings and the no-mask dense prompt are baked as buffers.
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 and paper for full dataset details.

Performance

Measured on a Pixel 8a (Tensor G3, Android 16) with the standard TFLite benchmark_model 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.

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.

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 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. This is a format conversion; all credit to the original authors (Meta AI).

litert-community/SAM2.1-Hiera-Tiny-Mask-Decoder

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