Re-export with public ExecuTorch 1.5.1 (recipe in export/)

NOTICE

@@ -1,4 +1,4 @@
-This repository contains a modified version of facebook/sam2-hiera-tiny.
+This repository contains a modified version of facebook/sam2.1-hiera-tiny.
 
 Modifications by Meta Platforms, Inc.: conversion to the ExecuTorch program format (.pte) for the
 ExecuTorch WebGPU backend.

README.md

@@ -1,7 +1,7 @@
 ---
 license: apache-2.0
 base_model:
-- facebook/sam2-hiera-tiny
+- facebook/sam2.1-hiera-tiny
 library_name: executorch
 pipeline_tag: mask-generation
 inference: false
@@ -11,27 +11,27 @@ tags:
 - on-device
 - browser
 ---
-# SAM 2 Hiera Tiny for ExecuTorch WebGPU
+# SAM 2.1 Hiera Tiny for ExecuTorch WebGPU
 
 This repository contains an [ExecuTorch](https://github.com/pytorch/executorch) conversion of
-[facebook/sam2-hiera-tiny](https://huggingface.co/facebook/sam2-hiera-tiny) for the ExecuTorch **WebGPU** backend, so the model runs entirely in a web browser on
+[facebook/sam2.1-hiera-tiny](https://huggingface.co/facebook/sam2.1-hiera-tiny) for the ExecuTorch **WebGPU** backend, so the model runs entirely in a web browser on
 the visitor's GPU. It powers the corresponding demo in the ExecuTorch WebGPU Playground (launching soon on [executorch.ai](https://executorch.ai)).
 
-- **Base model:** [facebook/sam2-hiera-tiny](https://huggingface.co/facebook/sam2-hiera-tiny)
+- **Base model:** [facebook/sam2.1-hiera-tiny](https://huggingface.co/facebook/sam2.1-hiera-tiny)
 - **Format:** ExecuTorch `.pte` programs delegated to the WebGPU backend
-- **Quantization:** fp32 (as exported)
-- **Total size:** 149.3 MB
+- **Quantization:** fp32 (no quantization)
+- **Total size:** 128.1 MB
 
 ## About the original model
 
-The tiny Hiera variant of SAM 2, a foundation model for promptable visual segmentation in images and videos.
+The tiny Hiera variant of SAM 2.1, the improved checkpoint release of SAM 2, a foundation model for promptable visual segmentation in images and videos.
 
 - **Developed by:** Meta FAIR
 - **Paper:** [SAM 2: Segment Anything in Images and Videos](https://arxiv.org/abs/2408.00714)
-- 38.9M parameters
+- 39.0M parameters
 - Segments any object from point or box prompts
 
-For training data, evaluation results and full details, see the original model card: [facebook/sam2-hiera-tiny](https://huggingface.co/facebook/sam2-hiera-tiny).
+For training data, evaluation results and full details, see the original model card: [facebook/sam2.1-hiera-tiny](https://huggingface.co/facebook/sam2.1-hiera-tiny).
 
 ## Intended use and limitations
 
@@ -41,14 +41,18 @@ This conversion covers single-image segmentation only (no video memory). Masks c
 
 ## What changed in this conversion
 
-- Image encoder and mask decoder exported as separate programs
-- Fixed 512x512 input (upstream: 1024x1024)
-- Prompt embeddings are precomputed and shipped as `prompt_*.bin` files
+- Image encoder and mask decoder exported as two separate programs; the prompt encoder runs on the host from `prompt_*.bin` tables generated from the upstream weights
+- Fixed 512x512 input (upstream: 1024x1024), so masks are 128x128 (upstream: 256x256)
+- The encoder adds `no_memory_embedding` to the image embedding, as upstream does for single images
+- Three prompt tokens: a click is [point, padding, padding] and a box is [top-left, bottom-right, padding]. The decoder ignores a padding token that follows another padding token, so every prompt type matches upstream
+- No mask-prompt input. All weights fp32. Attention uses ExecuTorch's fused `et_vk.sdpa` op (within 5e-6 of PyTorch attention), and Hiera's windowed max-pooling runs as one batched pool (exact)
 
 ## Inputs and outputs
 
-- `sam2_encoder.pte`: one 512x512 RGB image, ImageNet-normalized (mean 0.485/0.456/0.406, std 0.229/0.224/0.225), to image embeddings
-- `sam2_decoder.pte`: image embeddings plus prompt tokens built from `prompt_*.bin`, to mask logits
+- `sam2_encoder.pte`: `forward(pixels float32 [1,3,512,512])`, an RGB image resized to 512x512 (aspect ratio not kept), scaled to [0, 1] and ImageNet-normalized (mean 0.485/0.456/0.406, std 0.229/0.224/0.225) -> `image_embeddings [1,256,32,32]`, `hr0 [1,256,128,128]`, `hr1 [1,256,64,64]` (high-resolution features)
+- `sam2_decoder.pte`: `forward(image_embeddings, image_pe [1,256,32,32], sparse [1,1,3,256], dense [1,256,32,32], hr0, hr1)` -> `masks [1,3,128,128]` (logits of 3 candidate masks), `iou_scores [1,1,3]` (predicted IoU, after sigmoid), `object_score_logits [1,1,1]`
+- `image_pe` is `prompt_img_pos.bin` and `dense` is `prompt_dense.bin` (no mask prompt). Each sparse token is `[sin(2πcG), cos(2πcG)] + e`, where `G` [2,128] is `prompt_gaussian.bin`, `c` is the point mapped to the 512 grid (+0.5 pixel) and normalized to [-1, 1], and `e` is `prompt_point_embed_fg.bin` / `prompt_point_embed_bg.bin` / `prompt_box_tl.bin` / `prompt_box_br.bin`; a padding token is `prompt_not_a_point.bin` alone; fill unused slots with padding tokens
+- All tensors are float32, contiguous NCHW; `.bin` files are raw little-endian float32
 
 ## Files
 
@@ -62,16 +66,17 @@ This conversion covers single-image segmentation only (no video memory). Masks c
 | `prompt_not_a_point.bin` | 1.0 KB | `19771c067b922d9d…` |
 | `prompt_point_embed_bg.bin` | 1.0 KB | `91c432d10fcef8e3…` |
 | `prompt_point_embed_fg.bin` | 1.0 KB | `2e037dab5bee002d…` |
-| `sam2_decoder.pte` | 16.2 MB | `4d6f790ccf3c7ae8…` |
-| `sam2_encoder.pte` | 131.2 MB | `46626cd828cca2f2…` |
+| `sam2_decoder.pte` | 16.2 MB | `57b91034acc71be2…` |
+| `sam2_encoder.pte` | 109.9 MB | `30f9fb469b779558…` |
 
 Full checksums are in `manifest.json`. The playground verifies them after download.
 
 ## Running it
 
 These programs were exported with the ExecuTorch WebGPU backend
-([backends/webgpu](https://github.com/pytorch/executorch/tree/main/backends/webgpu)) and validated with the
-WebAssembly runtime build used by the ExecuTorch WebGPU Playground (launching soon on [executorch.ai](https://executorch.ai)). That runtime build is not published yet; until it is,
+([backends/webgpu](https://github.com/pytorch/executorch/tree/main/backends/webgpu)) and checked against the upstream transformers model at 512x512 (click and box prompts give masks with IoU 1.000),
+and run in a browser on WebGPU with the WebAssembly runtime build used by the ExecuTorch WebGPU Playground (launching soon on [executorch.ai](https://executorch.ai)), where the masks match the CPU run with IoU 1.000.
+They have not yet been tested on a hardware GPU. That runtime build is not published yet; until it is,
 the playground is the supported way to run them. It downloads these files once, caches them in the browser,
 and runs everything locally.
 
@@ -79,8 +84,11 @@ Requirements: a browser with WebGPU enabled.
 
 ## Provenance
 
-- Converted from the upstream checkpoint by the ExecuTorch team with the WebGPU export path.
-- A reproducible export recipe will be published alongside the open-source ExecuTorch WebGPU runtime.
+- **Upstream:** `facebook/sam2.1-hiera-tiny` at revision `de431c4043854a71d8101e17995dfe596bf101a5` (official)
+- **Checked against the official checkpoint:** `model.safetensors` (sha256 `48c14467e5cf9e51870511feb72c89688e82dd74523142c0538b663e193ac2a7`) has 469/469 tensors bit-identical to the authors' original checkpoint `sam2.1_hiera_tiny.pt` (sha256 `7402e0d864fa82708a20fbd15bc84245c2f26dff0eb43a4b5b93452deb34be69`) in the same repo at the same revision; `python export/_sam2.py verify-official` repeats the check
+- **Exported with:** executorch 1.5.1+cpu, torch 2.14.1+cpu, torchao 0.18.0+cpu, transformers 5.19.0 (executorch git `3b6068392324`)
+- **Reproduce:** `python recipes/sam2-hiera-tiny.py --out out/sam2.1-hiera-tiny` (recipe in `export/`)
+- Re-running the recipe produces byte-identical files.
 
 ## Citation
 

export/_common.py

@@ -0,0 +1,330 @@
+"""Shared helpers for the ExecuTorch WebGPU export recipes in this directory.
+
+Each recipe builds an eager model from a pinned upstream revision and calls
+`export_vision` (or its own LLM path), which writes:
+
+  <out>/<name>_webgpu.pte    program lowered with WebGPUPartitioner (delegate id "VulkanBackend")
+  <out>/<name>_portable.pte  same exported graph without a delegate, for CPU validation only
+  <out>/recipe.json          upstream id + revision, versions, shapes, dtype, command line
+  <out>/validation.json      eager vs exported-program vs portable .pte (and goldens, if given)
+"""
+
+import hashlib
+import io
+import json
+import platform
+import shlex
+import sys
+import urllib.request
+import zipfile
+from pathlib import Path
+
+import numpy as np
+import torch
+
+PROXY_HINT = (
+    "set https_proxy if the machine needs a proxy for huggingface.co / pypi.org"
+)
+
+
+def versions():
+    import executorch
+    import torchao
+    from executorch.version import git_version
+
+    out = {
+        "python": platform.python_version(),
+        "torch": torch.__version__,
+        "torchao": torchao.__version__,
+        "executorch": getattr(executorch, "__version__", None)
+        or __import__("executorch.version").version.__version__,
+        "executorch_git": git_version,
+    }
+    try:
+        import transformers
+
+        out["transformers"] = transformers.__version__
+    except ImportError:
+        pass
+    try:
+        import huggingface_hub
+
+        out["huggingface_hub"] = huggingface_hub.__version__
+    except ImportError:
+        pass
+    return out
+
+
+def sha256_file(path):
+    h = hashlib.sha256()
+    with open(path, "rb") as fh:
+        for chunk in iter(lambda: fh.read(1 << 20), b""):
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def hf_file(repo_id, filename, revision, expected_sha256=None, cache_dir=None):
+    from huggingface_hub import hf_hub_download
+
+    path = hf_hub_download(
+        repo_id=repo_id, filename=filename, revision=revision, cache_dir=cache_dir
+    )
+    digest = sha256_file(path)
+    if expected_sha256 and digest != expected_sha256:
+        sys.exit(
+            f"{repo_id}@{revision}:{filename} sha256 {digest} != {expected_sha256}"
+        )
+    return Path(path), digest
+
+
+def pypi_wheel_member(url, expected_sha256, member):
+    """Read one file out of a PyPI wheel, verifying the wheel's sha256."""
+    with urllib.request.urlopen(url, timeout=120) as resp:
+        blob = resp.read()
+    digest = hashlib.sha256(blob).hexdigest()
+    if digest != expected_sha256:
+        sys.exit(f"{url}: sha256 {digest} != {expected_sha256} ({PROXY_HINT})")
+    with zipfile.ZipFile(io.BytesIO(blob)) as zf:
+        return zf.read(member).decode()
+
+
+def scrub(text, out_dir):
+    """Replace machine-specific paths so recipe.json can ship with the model."""
+    text = text.replace(str(Path(out_dir).resolve()), "<out>")
+    text = text.replace(str(Path(__file__).resolve().parent), "recipes")
+    return text.replace(str(Path.home()), "~")
+
+
+def command_line(out_dir):
+    argv = [f"recipes/{Path(sys.argv[0]).name}", *sys.argv[1:]]
+    args = []
+    for a in argv:
+        if a.startswith("/") or a.startswith("."):
+            a = scrub(str(Path(a).resolve()), out_dir)
+        args.append(a)
+    return scrub(" ".join(["python", *(shlex.quote(a) for a in args)]), out_dir)
+
+
+def lower_webgpu(ep):
+    from executorch.backends.webgpu.partitioner import WebGPUPartitioner
+    from executorch.exir import to_edge_transform_and_lower
+
+    return to_edge_transform_and_lower(
+        ep, partitioner=[WebGPUPartitioner()]
+    ).to_executorch()
+
+
+def lower_portable(ep):
+    from executorch.exir import to_edge_transform_and_lower
+
+    return to_edge_transform_and_lower(ep).to_executorch()
+
+
+def delegate_summary(pte_bytes):
+    """Non-delegated kernels and per-delegate op histograms of a serialized program."""
+    import collections
+
+    from executorch.backends.vulkan.serialization.vulkan_graph_serialize import (
+        extract_vk_flatbuffer,
+        flatbuffer_to_vk_graph,
+    )
+    from executorch.exir._serialize._program import deserialize_pte_binary
+
+    pte = deserialize_pte_binary(pte_bytes)
+    prog = pte.program if hasattr(pte, "program") else pte
+    seg_base = int.from_bytes(pte_bytes[24:32], "little")
+    plan = prog.execution_plan[0]
+    ops = [f"{o.name}.{o.overload}" for o in plan.operators]
+    kernels = collections.Counter()
+    for chain in plan.chains:
+        for ins in chain.instructions:
+            if type(ins.instr_args).__name__ == "KernelCall":
+                kernels[ops[ins.instr_args.op_index]] += 1
+    delegates = []
+    for d in plan.delegates:
+        proc = d.processed
+        if proc.location.name == "INLINE":
+            blob = prog.backend_delegate_data[proc.index].data
+        else:
+            seg = prog.segments[proc.index]
+            blob = pte_bytes[seg_base + seg.offset : seg_base + seg.offset + seg.size]
+        g = flatbuffer_to_vk_graph(extract_vk_flatbuffer(blob))
+        delegates.append(
+            {
+                "id": d.id,
+                "ops": dict(collections.Counter(c.name for c in g.chain)),
+                "inputs": len(g.input_ids),
+                "outputs": len(g.output_ids),
+            }
+        )
+    return {"non_delegated_kernels": dict(kernels), "delegates": delegates}
+
+
+def compare(a, b):
+    a = np.asarray(a, dtype=np.float64).ravel()
+    b = np.asarray(b, dtype=np.float64).ravel()
+    n = min(a.size, b.size)
+    a, b = a[:n], b[:n]
+    cos = float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-30))
+    return {
+        "max_abs": float(np.max(np.abs(a - b))),
+        "mean_abs": float(np.mean(np.abs(a - b))),
+        "cosine": cos,
+    }
+
+
+def run_pte(pte_bytes, inputs):
+    from executorch.extension.pybindings.portable_lib import (
+        _load_for_executorch_from_buffer,
+    )
+
+    m = _load_for_executorch_from_buffer(pte_bytes)
+    return [
+        o.detach().numpy() if isinstance(o, torch.Tensor) else o
+        for o in m.forward(inputs)
+    ]
+
+
+def try_run_webgpu_pte_on_cpu(pte_bytes, inputs):
+    try:
+        return run_pte(pte_bytes, inputs), None
+    except Exception as e:  # noqa: BLE001 - the error text is the result
+        return None, f"{type(e).__name__}: {str(e).splitlines()[0][:300]}"
+
+
+def export_vision(
+    *,
+    name,
+    model,
+    example_input,
+    out_dir,
+    upstream,
+    extra_recipe,
+    golden=None,
+    output_index=None,
+):
+    """Export, lower, write both .pte files, recipe.json, and validation.json.
+
+    `golden` is an optional (input ndarray, output ndarray) pair; when its input
+    shape matches `example_input`, all paths are also compared against it.
+    `output_index` picks one tensor when the eager model returns a tuple.
+    """
+    out_dir = Path(out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    torch.manual_seed(0)
+    model = model.eval()
+
+    class Wrapped(torch.nn.Module):
+        def __init__(self, m):
+            super().__init__()
+            self.m = m
+
+        def forward(self, x):
+            y = self.m(x)
+            return y[output_index] if output_index is not None else y
+
+    wrapped = Wrapped(model).eval()
+    with torch.no_grad():
+        ep = torch.export.export(wrapped, (example_input,), strict=False)
+        webgpu = lower_webgpu(ep)
+        portable = lower_portable(ep)
+
+    webgpu_path = out_dir / f"{name}_webgpu.pte"
+    portable_path = out_dir / f"{name}_portable.pte"
+    webgpu_path.write_bytes(webgpu.buffer)
+    portable_path.write_bytes(portable.buffer)
+
+    summary = delegate_summary(webgpu.buffer)
+    out_shape = list(ep.module()(example_input).shape)
+    recipe = {
+        "model": name,
+        "upstream": upstream,
+        "versions": versions(),
+        "export": {
+            "api": "torch.export.export(strict=False) -> to_edge_transform_and_lower(partitioner=[WebGPUPartitioner()]) -> to_executorch()",
+            "partitioner": "executorch.backends.webgpu.partitioner.WebGPUPartitioner() (default compile options; wraps VulkanPartitioner)",
+            "inputs": [
+                {
+                    "shape": list(example_input.shape),
+                    "dtype": "float32",
+                    "layout": "NCHW",
+                }
+            ],
+            "outputs": [{"shape": out_shape, "dtype": "float32"}],
+            "dtype": "float32",
+            "quantization": "none (fp32)",
+            **extra_recipe,
+        },
+        "files": {
+            webgpu_path.name: {
+                "bytes": webgpu_path.stat().st_size,
+                "sha256": sha256_file(webgpu_path),
+            },
+            portable_path.name: {
+                "bytes": portable_path.stat().st_size,
+                "sha256": sha256_file(portable_path),
+                "note": "CPU-only validation artifact (portable kernels, no delegate); not for release",
+            },
+        },
+        "graph": summary,
+        "command": command_line(out_dir),
+    }
+    (out_dir / "recipe.json").write_text(json.dumps(recipe, indent=2) + "\n")
+
+    # Validation on CPU.
+    gen = torch.Generator().manual_seed(1234)
+    rand_input = torch.rand(example_input.shape, generator=gen)
+    cases = {"random": rand_input}
+    if golden is not None and tuple(golden[0].shape) == tuple(example_input.shape):
+        cases["golden"] = torch.from_numpy(golden[0].copy())
+    report = {
+        "cpu_runtime": "executorch pybindings (portable_lib) from the pip wheel",
+        "cases": {},
+    }
+    webgpu_cpu_error = None
+    for label, x in cases.items():
+        with torch.no_grad():
+            eager = wrapped(x).numpy()
+            exported = ep.module()(x).numpy()
+        port = run_pte(portable.buffer, (x,))[0]
+        res = {
+            "exported_vs_eager": compare(exported, eager),
+            "portable_pte_vs_eager": compare(port, eager),
+        }
+        wg, err = try_run_webgpu_pte_on_cpu(webgpu.buffer, (x,))
+        if wg is not None:
+            res["webgpu_pte_vs_eager"] = compare(wg[0], eager)
+        else:
+            webgpu_cpu_error = err
+        if label == "golden":
+            ref = golden[1]
+            res["internal_golden_vs_eager"] = compare(eager, ref)
+            res["internal_golden_vs_portable_pte"] = compare(port, ref)
+            res["golden_note"] = (
+                "reference output from the internal goldens.manifest.json"
+            )
+        report["cases"][label] = res
+    report["webgpu_pte_on_cpu"] = (
+        "ran"
+        if webgpu_cpu_error is None
+        else f"not runnable on CPU pybindings: {webgpu_cpu_error}"
+    )
+    (out_dir / "validation.json").write_text(json.dumps(report, indent=2) + "\n")
+    print(
+        json.dumps(
+            {"recipe": recipe["files"], "graph": summary, "validation": report},
+            indent=2,
+        )
+    )
+    return ep, webgpu, portable
+
+
+def load_golden(manifest_path, key):
+    man = json.loads(Path(manifest_path).read_text())[key]["tensors"]
+    base = Path(manifest_path).parent
+
+    def ld(t):
+        return np.fromfile(base / t["file"], dtype=np.float32).reshape(t["shape"])
+
+    inp = man.get("input")
+    return ld(inp), ld(man["output"])

export/_sam2.py

@@ -0,0 +1,288 @@
+#!/usr/bin/env python3
+"""Inspection helpers for the SAM 2 recipe (sam2-hiera-tiny.py). CPU only, read-only on inputs.
+
+_sam2.py ops A.pte [B.pte]          every delegated op call with its argument types/shapes; with B,
+                                    the op-signature set difference (ops/arity/ranks in A not in B)
+_sam2.py consts A.pte B.pte         how many delegate constants of A are byte-identical in B
+_sam2.py io A.pte                   method names, program inputs/outputs, non-delegated kernels
+"""
+
+import collections
+import hashlib
+import json
+import sys
+
+
+def _program(data):
+    from executorch.exir._serialize._program import deserialize_pte_binary
+
+    pte = deserialize_pte_binary(data)
+    return pte, (pte.program if hasattr(pte, "program") else pte)
+
+
+def _blobs(data):
+    pte, prog = _program(data)
+    seg_base = int.from_bytes(data[24:32], "little")
+    out = []
+    for plan in prog.execution_plan:
+        for d in plan.delegates:
+            proc = d.processed
+            if proc.location.name == "INLINE":
+                out.append(
+                    (
+                        plan.name,
+                        d.id,
+                        bytes(prog.backend_delegate_data[proc.index].data),
+                    )
+                )
+            else:
+                seg = prog.segments[proc.index]
+                out.append(
+                    (
+                        plan.name,
+                        d.id,
+                        data[seg_base + seg.offset : seg_base + seg.offset + seg.size],
+                    )
+                )
+    nd = getattr(pte, "named_data", None)
+    resolve = (
+        (
+            lambda k: bytes(nd.buffers[nd.pte_data[k].buffer_index])
+            if k in nd.pte_data
+            else None
+        )
+        if nd
+        else (lambda k: None)
+    )
+    return out, resolve
+
+
+def vk_graph(blob):
+    from executorch.backends.vulkan.serialization.vulkan_graph_serialize import (
+        extract_vk_flatbuffer,
+        flatbuffer_to_vk_graph,
+    )
+
+    return flatbuffer_to_vk_graph(extract_vk_flatbuffer(blob))
+
+
+def describe(g, vid):
+    v = g.values[vid].value
+    n = type(v).__name__
+    if n == "VkTensor":
+        kind = "const" if v.constant_id >= 0 else "tensor"
+        return f"{kind}{list(v.dims)}:{v.datatype.name if hasattr(v.datatype, 'name') else v.datatype}"
+    if n == "Int":
+        return f"Int={v.int_val}"
+    if n == "Double":
+        return f"Double={v.double_val:g}"
+    if n == "Bool":
+        return f"Bool={v.bool_val}"
+    if n in ("IntList", "DoubleList", "BoolList"):
+        return f"{n}={list(v.items)}"
+    if n == "ValueList":
+        return "ValueList[" + ", ".join(describe(g, i) for i in v.items) + "]"
+    if n == "String":
+        return f"String={v.string_val!r}"
+    return n
+
+
+def ops(path):
+    data = open(path, "rb").read()
+    blobs, _ = _blobs(data)
+    calls = []
+    for plan, did, blob in blobs:
+        g = vk_graph(blob)
+        for c in g.chain:
+            calls.append(
+                {"plan": plan, "op": c.name, "args": [describe(g, a) for a in c.args]}
+            )
+    return calls
+
+
+def signature(call):
+    """Op name + arity + per-argument kind/rank (shape-agnostic) for set comparisons."""
+    parts = []
+    for a in call["args"]:
+        if a.startswith(("tensor[", "const[")):
+            dims = json.loads(a[a.index("[") : a.index("]") + 1])
+            parts.append(f"{a.split('[')[0]}{len(dims)}d:{a.rsplit(':', 1)[1]}")
+        else:
+            parts.append(a.split("=")[0])
+    return f"{call['op']}({', '.join(parts)})"
+
+
+def consts(path):
+    from executorch.backends.vulkan.serialization.vulkan_graph_serialize import (
+        VulkanDelegateHeader,
+    )
+
+    data = open(path, "rb").read()
+    blobs, resolve = _blobs(data)
+    out = []
+    for plan, did, blob in blobs:
+        h = VulkanDelegateHeader.from_bytes(
+            blob[: VulkanDelegateHeader.EXPECTED_LENGTH]
+        )
+        g = vk_graph(blob)
+        raw = blob[h.bytes_offset : h.bytes_offset + h.bytes_size]
+        for v in g.values:
+            t = v.value
+            if type(t).__name__ != "VkTensor" or t.constant_id < 0:
+                continue
+            c = g.constants[t.constant_id]
+            b = (
+                resolve(c.named_key)
+                if c.named_key
+                else raw[c.offset : c.offset + c.length]
+            )
+            out.append(
+                {
+                    "dims": list(t.dims),
+                    "sha": hashlib.sha256(b).hexdigest(),
+                    "bytes": len(b),
+                }
+            )
+    return out
+
+
+def io(path):
+    data = open(path, "rb").read()
+    _, prog = _program(data)
+    res = {"file": path, "bytes": len(data), "methods": {}}
+    for plan in prog.execution_plan:
+
+        def d(i):
+            v = plan.values[i].val
+            return (
+                f"{list(v.sizes)}:{v.scalar_type.name}"
+                if hasattr(v, "sizes")
+                else type(v).__name__
+            )
+
+        names = [f"{o.name}.{o.overload}" for o in plan.operators]
+        kernels = collections.Counter(
+            names[ins.instr_args.op_index]
+            for ch in plan.chains
+            for ins in ch.instructions
+            if type(ins.instr_args).__name__ == "KernelCall"
+        )
+        res["methods"][plan.name] = {
+            "inputs": [d(i) for i in plan.inputs],
+            "outputs": [d(i) for i in plan.outputs],
+            "delegates": [x.id for x in plan.delegates],
+            "non_delegated_kernels": dict(kernels),
+        }
+    return res
+
+
+OFFICIAL = {
+    "sam2.1": ("facebook/sam2.1-hiera-tiny", "de431c4043854a71d8101e17995dfe596bf101a5", "sam2.1_hiera_tiny.pt",
+               "7402e0d864fa82708a20fbd15bc84245c2f26dff0eb43a4b5b93452deb34be69"),
+}
+
+
+def verify_official(weights="sam2.1"):
+    import importlib.util
+    import pathlib
+    import tempfile
+
+    import torch
+    from huggingface_hub import hf_hub_download
+    from safetensors.torch import load_file
+
+    here = pathlib.Path(__file__).resolve().parent
+    spec = importlib.util.spec_from_file_location("sam2_recipe", here / "sam2-hiera-tiny.py")
+    recipe = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(recipe)
+    snap, digest = recipe.snapshot(weights)
+    hf = load_file(str(snap / "model.safetensors"))
+    repo, rev, name, want = OFFICIAL[weights]
+    with tempfile.TemporaryDirectory() as tmp:
+        p = hf_hub_download(repo, name, revision=rev, local_dir=tmp)
+        got = hashlib.sha256(open(p, "rb").read()).hexdigest()
+        if got != want:
+            sys.exit(f"{name}: sha256 {got} != {want}")
+        orig = torch.load(p, map_location="cpu", weights_only=True)["model"]
+    h = lambda t: hashlib.sha256(t.detach().contiguous().numpy().tobytes()).hexdigest()  # noqa: E731
+    pool = {h(v): k for k, v in orig.items()}
+    matched, unmatched = 0, []
+    for k, v in hf.items():
+        if h(v) in pool:
+            matched += 1
+        else:
+            unmatched.append(k)
+    notes = {}
+    if "prompt_encoder.point_embed.weight" in unmatched:
+        pe = hf["prompt_encoder.point_embed.weight"]
+        if all(torch.equal(pe[i], orig[f"sam_prompt_encoder.point_embeddings.{i}.weight"][0]) for i in range(pe.shape[0])):
+            unmatched.remove("prompt_encoder.point_embed.weight")
+            matched += 1
+            notes["prompt_encoder.point_embed.weight"] = "rows == sam_prompt_encoder.point_embeddings.{0..3}.weight"
+    return {"hf": f"{repo}@{rev}/model.safetensors", "hf_sha256": digest, "official": f"{repo}@{rev}/{name}",
+            "official_sha256": want, "hf_tensors": len(hf), "official_tensors": len(orig),
+            "hf_tensors_bit_identical_to_official": matched, "unmatched": unmatched, "notes": notes}
+
+
+def main():
+    cmd, args = sys.argv[1], sys.argv[2:]
+    if cmd == "io":
+        print(json.dumps(io(args[0]), indent=1))
+    elif cmd == "ops":
+        a = ops(args[0])
+        if len(args) == 1:
+            for c in a:
+                print(c["op"], ", ".join(c["args"]))
+            return
+        b = ops(args[1])
+        sa, sb = (
+            collections.Counter(map(signature, a)),
+            collections.Counter(map(signature, b)),
+        )
+        print(
+            json.dumps(
+                {
+                    "ops_count": [sum(sa.values()), sum(sb.values())],
+                    "op_types_only_in_a": sorted(
+                        {c["op"] for c in a} - {c["op"] for c in b}
+                    ),
+                    "op_types_only_in_b": sorted(
+                        {c["op"] for c in b} - {c["op"] for c in a}
+                    ),
+                    "signatures_only_in_a": {
+                        k: v for k, v in sorted(sa.items()) if k not in sb
+                    },
+                    "signatures_only_in_b": {
+                        k: v for k, v in sorted(sb.items()) if k not in sa
+                    },
+                },
+                indent=1,
+            )
+        )
+    elif cmd == "verify-official":
+        print(json.dumps(verify_official(*(args or ["sam2.1"])), indent=1))
+    elif cmd == "consts":
+        a, b = consts(args[0]), consts(args[1])
+        hb, ha = (
+            collections.Counter(c["sha"] for c in b),
+            collections.Counter(c["sha"] for c in a),
+        )
+        print(
+            json.dumps(
+                {
+                    "a_constants": len(a),
+                    "b_constants": len(b),
+                    "a_found_in_b": sum(c["sha"] in hb for c in a),
+                    "b_found_in_a": sum(c["sha"] in ha for c in b),
+                    "a_missing": [c["dims"] for c in a if c["sha"] not in hb],
+                    "b_missing": [c["dims"] for c in b if c["sha"] not in ha],
+                },
+                indent=1,
+            )
+        )
+    else:
+        sys.exit(__doc__)
+
+
+if __name__ == "__main__":
+    main()

export/recipe.json

@@ -0,0 +1,589 @@
+{
+  "model": "sam2.1-hiera-tiny (512 encoder/decoder split)",
+  "upstream": {
+    "model": "facebook/sam2.1-hiera-tiny",
+    "revision": "de431c4043854a71d8101e17995dfe596bf101a5",
+    "official": true,
+    "license": "Apache-2.0",
+    "files": {
+      "model.safetensors": {
+        "sha256": "48c14467e5cf9e51870511feb72c89688e82dd74523142c0538b663e193ac2a7"
+      }
+    },
+    "note": "transformers Sam2Model port published by the SAM 2 authors on the Hub"
+  },
+  "versions": {
+    "python": "3.12.15+meta",
+    "torch": "2.14.1+cpu",
+    "torchao": "0.18.0+cpu",
+    "executorch": "1.5.1+cpu",
+    "executorch_git": "3b60683923245cf472b7323426920e15623ba361",
+    "transformers": "5.19.0",
+    "huggingface_hub": "2.2.0"
+  },
+  "export": {
+    "api": "torch.export.export(strict=False) -> to_edge_transform_and_lower(partitioner=[WebGPUPartitioner()]) -> to_executorch(), one program per module",
+    "partitioner": "executorch.backends.webgpu.partitioner.WebGPUPartitioner() (default compile options; wraps VulkanPartitioner)",
+    "model_code": "transformers Sam2Model; attention through torch.ops.et_vk.sdpa (transformers attention backend, recipes/sam2-hiera-tiny.py:et_vk_attention); config image_size overridden 1024 -> 512",
+    "programs": {
+      "sam2_encoder.pte": {
+        "method": "forward",
+        "inputs": [
+          {
+            "name": "pixels",
+            "shape": [
+              1,
+              3,
+              512,
+              512
+            ],
+            "dtype": "float32"
+          }
+        ],
+        "outputs": [
+          {
+            "name": "image_embeddings",
+            "shape": [
+              1,
+              256,
+              32,
+              32
+            ],
+            "dtype": "float32",
+            "note": "no_memory_embedding added"
+          },
+          {
+            "name": "hr0",
+            "shape": [
+              1,
+              256,
+              128,
+              128
+            ],
+            "dtype": "float32",
+            "note": "before mask_decoder.conv_s0"
+          },
+          {
+            "name": "hr1",
+            "shape": [
+              1,
+              256,
+              64,
+              64
+            ],
+            "dtype": "float32",
+            "note": "before mask_decoder.conv_s1"
+          }
+        ]
+      },
+      "sam2_decoder.pte": {
+        "method": "forward",
+        "prompt_tokens": 3,
+        "inputs": [
+          {
+            "name": "image_embeddings",
+            "shape": [
+              1,
+              256,
+              32,
+              32
+            ],
+            "dtype": "float32"
+          },
+          {
+            "name": "image_pe",
+            "shape": [
+              1,
+              256,
+              32,
+              32
+            ],
+            "dtype": "float32"
+          },
+          {
+            "name": "sparse",
+            "shape": [
+              1,
+              1,
+              3,
+              256
+            ],
+            "dtype": "float32"
+          },
+          {
+            "name": "dense",
+            "shape": [
+              1,
+              256,
+              32,
+              32
+            ],
+            "dtype": "float32"
+          },
+          {
+            "name": "hr0",
+            "shape": [
+              1,
+              256,
+              128,
+              128
+            ],
+            "dtype": "float32"
+          },
+          {
+            "name": "hr1",
+            "shape": [
+              1,
+              256,
+              64,
+              64
+            ],
+            "dtype": "float32"
+          }
+        ],
+        "outputs": [
+          {
+            "name": "masks",
+            "shape": [
+              1,
+              3,
+              128,
+              128
+            ],
+            "dtype": "float32",
+            "note": "logits, mask tokens 1..3"
+          },
+          {
+            "name": "iou_scores",
+            "shape": [
+              1,
+              1,
+              3
+            ],
+            "dtype": "float32"
+          },
+          {
+            "name": "object_score_logits",
+            "shape": [
+              1,
+              1,
+              1
+            ],
+            "dtype": "float32"
+          }
+        ]
+      }
+    },
+    "prompt_tables": {
+      "prompt_gaussian.bin": {
+        "shape": [
+          2,
+          128
+        ],
+        "dtype": "float32 little-endian"
+      },
+      "prompt_point_embed_bg.bin": {
+        "shape": [
+          256
+        ],
+        "dtype": "float32 little-endian"
+      },
+      "prompt_point_embed_fg.bin": {
+        "shape": [
+          256
+        ],
+        "dtype": "float32 little-endian"
+      },
+      "prompt_box_tl.bin": {
+        "shape": [
+          256
+        ],
+        "dtype": "float32 little-endian"
+      },
+      "prompt_box_br.bin": {
+        "shape": [
+          256
+        ],
+        "dtype": "float32 little-endian"
+      },
+      "prompt_not_a_point.bin": {
+        "shape": [
+          256
+        ],
+        "dtype": "float32 little-endian"
+      },
+      "prompt_img_pos.bin": {
+        "shape": [
+          1,
+          256,
+          32,
+          32
+        ],
+        "dtype": "float32 little-endian"
+      },
+      "prompt_dense.bin": {
+        "shape": [
+          1,
+          256,
+          32,
+          32
+        ],
+        "dtype": "float32 little-endian"
+      }
+    },
+    "prompt_table_sources": {
+      "prompt_gaussian.bin": "shared_image_embedding.positional_embedding",
+      "prompt_point_embed_bg.bin": "prompt_encoder.point_embed.weight[0]",
+      "prompt_point_embed_fg.bin": "prompt_encoder.point_embed.weight[1]",
+      "prompt_box_tl.bin": "prompt_encoder.point_embed.weight[2]",
+      "prompt_box_br.bin": "prompt_encoder.point_embed.weight[3]",
+      "prompt_not_a_point.bin": "prompt_encoder.not_a_point_embed.weight[0]",
+      "prompt_img_pos.bin": "Sam2Model.get_image_wide_positional_embeddings() with image_size=512",
+      "prompt_dense.bin": "prompt_encoder.no_mask_embed.weight broadcast to [1,256,32,32]"
+    },
+    "graph_rewrites": [
+      "Hiera q-pooling of windowed tensors [nW,h,w,C] done as one batch-1 max_pool2d over the windows stacked along H (exact: h is a multiple of the stride); keeps all pools in the delegate",
+      "attention via torch.ops.et_vk.sdpa (same math as torch SDPA; checked in validation.json)",
+      "decoder: a not-a-point token directly after another one gets an additive -1e9 key bias in every attention (computed in-graph from the sparse input), so padding up to the fixed token count does not change the result",
+      "Hiera position embedding (bicubic pos_embed resize + tiled window embedding) for the fixed input size computed in eager and stored as a constant"
+    ],
+    "differences_from_upstream_pipeline": [
+      "input 512x512 instead of 1024x1024 (masks 128x128 instead of 256x256)",
+      "fixed 3 prompt tokens per decode: points, box corners, one not-a-point pad, then padding pads that the decoder masks out (exactly upstream's prompt for any prompt that fits; larger prompts are composed by the site)",
+      "no mask prompt input (dense = no_mask_embed)"
+    ],
+    "preprocessing_outside_graph": "RGB, resize to 512x512 without keeping aspect, /255, ImageNet mean/std",
+    "dtype": "float32",
+    "quantization": "none (fp32)"
+  },
+  "files": {
+    "prompt_box_br.bin": {
+      "bytes": 1024,
+      "sha256": "a0df40a29af4a9b65d390ae46913d4170762a306fc8061ea87f1d8f4a4a1b0d5"
+    },
+    "prompt_box_tl.bin": {
+      "bytes": 1024,
+      "sha256": "96c208c44e95f7044bc480951be257c415cfa9c91db907e4575ae3c92d09e6b7"
+    },
+    "prompt_dense.bin": {
+      "bytes": 1048576,
+      "sha256": "52cece57b549dc81439c1156a07a25ff02db8cb81329e8eefb68c90457a884f3"
+    },
+    "prompt_gaussian.bin": {
+      "bytes": 1024,
+      "sha256": "c37a6a9361886286fc3e868dd279dfa4703919e3fee65c07a07e5a75e980d84f"
+    },
+    "prompt_img_pos.bin": {
+      "bytes": 1048576,
+      "sha256": "89dc4182ee4ac6a366dbfd0e9b825e4d7b6e88ded91b2ab06604eb5139e1cc2c"
+    },
+    "prompt_not_a_point.bin": {
+      "bytes": 1024,
+      "sha256": "19771c067b922d9d652e6f3dd1f6e2976c68abb8e6fd9db571255ae37687ea50"
+    },
+    "prompt_point_embed_bg.bin": {
+      "bytes": 1024,
+      "sha256": "91c432d10fcef8e32ce561eb4820beb6a772a61a3a937d7c2af8597061759cfa"
+    },
+    "prompt_point_embed_fg.bin": {
+      "bytes": 1024,
+      "sha256": "2e037dab5bee002d94bdb12d3141ab952cfe2315f0c9feb53a7c78a7e656ed33"
+    },
+    "sam2_decoder.pte": {
+      "bytes": 16949252,
+      "sha256": "57b91034acc71be2ce1835ea546a174295daee1046e40db0743e0238ef11e278"
+    },
+    "sam2_decoder_portable.pte": {
+      "bytes": 16968388,
+      "sha256": "cac19779cd445105fdfed5599bf4f6ed041fbc4d006c931be8bc2c3cdf3c4655",
+      "note": "CPU-only validation artifact (portable kernels, no delegate; torch SDPA instead of et_vk.sdpa); not for release"
+    },
+    "sam2_encoder.pte": {
+      "bytes": 115248512,
+      "sha256": "30f9fb469b779558b6b5c7c3eb995c50e1213ffbe082e31430c929ba2b3ffe20"
+    },
+    "sam2_encoder_portable.pte": {
+      "bytes": 115320192,
+      "sha256": "b18fc39b847470555f8b69401a2acc67a39f635124ce5ac9c2d6e25917fe44f7",
+      "note": "CPU-only validation artifact (portable kernels, no delegate; torch SDPA instead of et_vk.sdpa); not for release"
+    }
+  },
+  "graph": {
+    "sam2_encoder": {
+      "non_delegated_kernels": {},
+      "delegates": [
+        {
+          "id": "VulkanBackend",
+          "ops": {
+            "et_vk.prepack.default": 2,
+            "aten.convolution.default": 5,
+            "aten.clone.default": 32,
+            "aten.permute_copy.default": 83,
+            "aten.add.Tensor": 27,
+            "aten.native_layer_norm.default": 24,
+            "aten.constant_pad_nd.default": 9,
+            "aten.view_copy.default": 114,
+            "aten.linear.default": 51,
+            "aten.slice_copy.Tensor": 46,
+            "aten.squeeze_copy.dims": 36,
+            "et_vk.sdpa.default": 12,
+            "aten.gelu.default": 12,
+            "aten.max_pool2d_with_indices.default": 6,
+            "aten.upsample_nearest2d.vec": 1
+          },
+          "inputs": 1,
+          "outputs": 3
+        }
+      ],
+      "io": {
+        "methods": [
+          "forward"
+        ],
+        "inputs": [
+          {
+            "shape": [
+              1,
+              3,
+              512,
+              512
+            ],
+            "dtype": "FLOAT",
+            "dim_order": [
+              0,
+              1,
+              2,
+              3
+            ]
+          }
+        ],
+        "outputs": [
+          {
+            "shape": [
+              1,
+              256,
+              32,
+              32
+            ],
+            "dtype": "FLOAT",
+            "dim_order": [
+              0,
+              1,
+              2,
+              3
+            ]
+          },
+          {
+            "shape": [
+              1,
+              256,
+              128,
+              128
+            ],
+            "dtype": "FLOAT",
+            "dim_order": [
+              0,
+              1,
+              2,
+              3
+            ]
+          },
+          {
+            "shape": [
+              1,
+              256,
+              64,
+              64
+            ],
+            "dtype": "FLOAT",
+            "dim_order": [
+              0,
+              1,
+              2,
+              3
+            ]
+          }
+        ]
+      },
+      "io_contract": "matches the previous export (shapes, fp32, contiguous dim order, single forward method)"
+    },
+    "sam2_decoder": {
+      "non_delegated_kernels": {},
+      "delegates": [
+        {
+          "id": "VulkanBackend",
+          "ops": {
+            "et_vk.prepack.default": 13,
+            "aten.sub.Tensor": 2,
+            "aten.abs.default": 1,
+            "aten.sum.dim_IntList": 1,
+            "aten.mul.Tensor": 3,
+            "aten.relu.default": 15,
+            "aten.slice_copy.Tensor": 5,
+            "aten.full.default": 1,
+            "aten.clone.default": 23,
+            "aten.cat.default": 4,
+            "aten.unsqueeze_copy.default": 16,
+            "aten.convolution.default": 4,
+            "aten.repeat.default": 1,
+            "aten.add.Tensor": 25,
+            "aten.view_copy.default": 57,
+            "aten.permute_copy.default": 36,
+            "aten.squeeze_copy.dims": 34,
+            "aten.linear.default": 47,
+            "aten.expand_copy.default": 7,
+            "et_vk.sdpa.default": 7,
+            "aten.native_layer_norm.default": 10,
+            "aten.bmm.default": 4,
+            "aten.select_copy.int": 7,
+            "aten.gelu.default": 2,
+            "aten.sigmoid.default": 1
+          },
+          "inputs": 6,
+          "outputs": 3
+        }
+      ],
+      "io": {
+        "methods": [
+          "forward"
+        ],
+        "inputs": [
+          {
+            "shape": [
+              1,
+              256,
+              32,
+              32
+            ],
+            "dtype": "FLOAT",
+            "dim_order": [
+              0,
+              1,
+              2,
+              3
+            ]
+          },
+          {
+            "shape": [
+              1,
+              256,
+              32,
+              32
+            ],
+            "dtype": "FLOAT",
+            "dim_order": [
+              0,
+              1,
+              2,
+              3
+            ]
+          },
+          {
+            "shape": [
+              1,
+              1,
+              3,
+              256
+            ],
+            "dtype": "FLOAT",
+            "dim_order": [
+              0,
+              1,
+              2,
+              3
+            ]
+          },
+          {
+            "shape": [
+              1,
+              256,
+              32,
+              32
+            ],
+            "dtype": "FLOAT",
+            "dim_order": [
+              0,
+              1,
+              2,
+              3
+            ]
+          },
+          {
+            "shape": [
+              1,
+              256,
+              128,
+              128
+            ],
+            "dtype": "FLOAT",
+            "dim_order": [
+              0,
+              1,
+              2,
+              3
+            ]
+          },
+          {
+            "shape": [
+              1,
+              256,
+              64,
+              64
+            ],
+            "dtype": "FLOAT",
+            "dim_order": [
+              0,
+              1,
+              2,
+              3
+            ]
+          }
+        ],
+        "outputs": [
+          {
+            "shape": [
+              1,
+              3,
+              128,
+              128
+            ],
+            "dtype": "FLOAT",
+            "dim_order": [
+              0,
+              1,
+              2,
+              3
+            ]
+          },
+          {
+            "shape": [
+              1,
+              1,
+              3
+            ],
+            "dtype": "FLOAT",
+            "dim_order": [
+              0,
+              1,
+              2
+            ]
+          },
+          {
+            "shape": [
+              1,
+              1,
+              1
+            ],
+            "dtype": "FLOAT",
+            "dim_order": [
+              0,
+              1,
+              2
+            ]
+          }
+        ]
+      },
+      "io_contract": "matches the previous export (shapes, fp32, contiguous dim order, single forward method)"
+    }
+  },
+  "command": "python recipes/sam2-hiera-tiny.py --out out/sam2.1-hiera-tiny"
+}

export/requirements.txt

@@ -0,0 +1,15 @@
+# Environment the re-exports in REEXPORT.md were produced with (Python 3.12, Linux x86_64, CPU only).
+#   python -m venv venv
+#   venv/bin/pip install -r recipes/requirements.txt --extra-index-url https://download.pytorch.org/whl/cpu
+# executorch 1.5.1 = pytorch/executorch git 3b60683923245cf472b7323426920e15623ba361; its wheel ships
+# backends/webgpu (WebGPUPartitioner) and examples/models/llama (export_llm).
+executorch==1.5.1
+torch==2.14.1+cpu
+torchvision==0.29.1+cpu
+torchao==0.18.0+cpu
+transformers==5.19.0
+huggingface_hub==2.2.0
+safetensors==0.8.0
+tokenizers==0.23.3
+numpy==2.5.4
+pillow==12.3.0

export/sam2-hiera-tiny.py

@@ -0,0 +1,939 @@
+#!/usr/bin/env python3
+"""SAM 2.1 Hiera Tiny -> two ExecuTorch WebGPU programs + prompt tables, 512x512 input.
+
+Same split and I/O contract as the playground's previous export (sam2_encoder.pte, sam2_decoder.pte,
+prompt_*.bin); the prompt encoder runs in JS (src/runtime/sam.js) from the tables.
+
+  sam2_encoder.pte  forward(pixels[1,3,512,512] f32: RGB, square resize, /255, ImageNet mean/std)
+                    -> (image_embeddings[1,256,32,32], hr0[1,256,128,128], hr1[1,256,64,64])
+                    hr0/hr1 are the FPN levels *before* conv_s0/conv_s1 (the decoder applies them).
+  sam2_decoder.pte  forward(image_embeddings[1,256,32,32], image_pe[1,256,32,32], sparse[1,1,3,256],
+                            dense[1,256,32,32], hr0[1,256,128,128], hr1[1,256,64,64])
+                    -> (masks[1,3,128,128] logits, iou_scores[1,1,3], object_score_logits[1,1,1])
+                    multimask output (mask tokens 1..3), three prompt tokens (--prompt-tokens): points,
+                    box corners, one not-a-point pad, then pads; pads after the first are masked, so
+                    clicks, click+exclude and boxes decode exactly as upstream.
+  prompt_gaussian.bin          [2,128] f32  random-Fourier matrix (shared_image_embedding)
+  prompt_point_embed_bg.bin    [256]   f32  point_embed[0] (negative click)
+  prompt_point_embed_fg.bin    [256]   f32  point_embed[1] (positive click)
+  prompt_box_tl.bin / _br.bin  [256]   f32  point_embed[2] / [3] (box corners)
+  prompt_not_a_point.bin       [256]   f32  not_a_point_embed (padding token)
+  prompt_img_pos.bin   [1,256,32,32]   f32  get_image_wide_positional_embeddings() at 512
+  prompt_dense.bin     [1,256,32,32]   f32  no_mask_embed broadcast (no mask prompt)
+
+Upstream: the HF transformers port (Apache-2.0) of facebook/sam2.1-hiera-tiny, which is what the
+previous export was built from (its decoder weights and all prompt tables are bit-identical to it).
+`--weights sam2` uses facebook/sam2-hiera-tiny (SAM 2.0) instead.
+
+Graph changes, all output-preserving: attention goes through the public fused op
+torch.ops.et_vk.sdpa (registered as a transformers attention backend); the Hiera windowed absolute
+position embedding (bicubic resize of pos_embed + tiled window embedding) for the fixed input size is
+computed once in eager and stored as a constant; Hiera's q-pooling of windowed tensors runs as one
+batch-1 max_pool2d over the windows stacked along H (exact; keeps every pool in the delegate). Unlike the previous export, the encoder adds
+`no_memory_embedding` to the 32x32 embedding, as upstream does for single images
+(transformers Sam2Model.forward; `directly_add_no_mem_embed: true` in the official yaml).
+`--no-memory-embedding omit` reproduces the previous behaviour.
+"""
+
+import argparse
+import json
+import sys
+import types
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+import _common  # noqa: E402
+
+WEIGHTS = {
+    "sam2.1": {
+        "repo": "facebook/sam2.1-hiera-tiny",
+        "revision": "de431c4043854a71d8101e17995dfe596bf101a5",
+        "files": {
+            "model.safetensors": "48c14467e5cf9e51870511feb72c89688e82dd74523142c0538b663e193ac2a7",
+        },
+    },
+    "sam2": {
+        "repo": "facebook/sam2-hiera-tiny",
+        "revision": "7c218beaf0bb87874785f32b582f640134fc1c09",
+        "files": {
+            "model.safetensors": "1449035ba6ba2ed524c646da989e5cb3a9b93b97f1c1b2d474e598c3b748d99f",
+        },
+    },
+}
+CONFIG_FILES = ("config.json", "preprocessor_config.json", "processor_config.json")
+
+# Validation images. Prompts are in original-image pixels, without padding; site_tokens() turns
+# them into the decoder's tokens the way planPrompt() in sam.js does (points, box corners, one
+# not-a-point pad, then pads up to the decoder's token count).
+IMAGE = {
+    "name": "sam-car.png",
+    "repo": "huggingface/documentation-images",
+    "repo_type": "dataset",
+    "revision": "79244b5169b24b4e1444393d9b87349c9182cfa4",
+    "file": "transformers/model_doc/sam-car.png",
+    "sha256": "c6ae6440e3862dd9768cd9f7098b783ac980629dd7bc1fe69c10b39fca5a84cd",
+    "prompts": {
+        "door_click": [("pos", 1680, 840)],
+        "wheel_click": [("pos", 832, 1088)],
+        "door_plus_exclude": [("pos", 1680, 840), ("neg", 1000, 200)],
+        "car_box": [("tl", 120, 400), ("br", 2528, 1248)],
+        "door_and_wheel_clicks": [("pos", 1680, 840), ("pos", 832, 1088)],
+    },
+}
+# Optional: extra sample images for validation (`--assets-dir`).
+# Not public inputs; pinned by sha256 so the numbers are reproducible.
+ASSET_IMAGES = {
+    "real_pets.jpg": {
+        "sha256": "85e4f2f27368a203c30eac68a71e4dde40edac6417b4a8c9df80af5d30b191b2",
+        "prompts": {
+            "pug_face_click": [("pos", 505, 150)],
+            "blanket_click": [("pos", 500, 500)],
+            "blanket_minus_face": [("pos", 500, 500), ("neg", 505, 150)],
+            "pug_box": [("tl", 165, 65), ("br", 850, 720)],
+        },
+    },
+    "yolo_people.jpg": {
+        "sha256": "16d73869e3267a7d4ed00de8e860833bd1657c1b252e94c0c348277adc7b6edb",
+        "prompts": {
+            "left_man_head_click": [("pos", 570, 260)],
+            "right_man_face_click": [("pos", 980, 180)],
+            "left_jacket_minus_right": [("pos", 300, 520), ("neg", 880, 520)],
+            "right_man_box": [("tl", 750, 40), ("br", 1150, 720)],
+        },
+    },
+}
+SIZE = 512
+MEAN = (0.485, 0.456, 0.406)
+STD = (0.229, 0.224, 0.225)
+INPUT_NAMES = ["image_embeddings", "image_pe", "sparse", "dense", "hr0", "hr1"]
+
+
+# Additive key bias [1, 1, 1, tokens] set by Decoder.forward while the mask decoder runs: it hides
+# extra padding tokens (see Decoder.pad_bias). None everywhere else (encoder, references).
+_TOKEN_KEY_BIAS = None
+
+
+def _key_bias(query, key, attention_mask):
+    bias = _TOKEN_KEY_BIAS
+    if bias is None or key.shape[-2] != bias.shape[-1]:
+        return attention_mask
+    full = bias.expand(query.shape[0], query.shape[1], query.shape[2], key.shape[-2])
+    return full if attention_mask is None else attention_mask + full
+
+
+def et_vk_attention(
+    module, query, key, value, attention_mask, scaling=None, dropout=0.0, **kwargs
+):
+    """transformers attention backend emitting the fused et_vk.sdpa op (public ExecuTorch Vulkan custom op,
+    registered by the WebGPU runtime). q/k/v are [B, H, S, D]; returns [B, S, H, D] like the built-in backends."""
+    import executorch.backends.vulkan.custom_ops_lib  # noqa: F401  (defines torch.ops.et_vk.*)
+
+    mask = _key_bias(query, key, attention_mask)
+    out = torch.ops.et_vk.sdpa(query, key, value, mask, scaling)
+    return out.transpose(1, 2).contiguous(), None
+
+
+def torch_attention(
+    module, query, key, value, attention_mask, scaling=None, dropout=0.0, **kwargs
+):
+    """Same as et_vk_attention with torch SDPA (for the CPU-only portable decoder)."""
+    mask = _key_bias(query, key, attention_mask)
+    out = F.scaled_dot_product_attention(query, key, value, attn_mask=mask, scale=scaling)
+    return out.transpose(1, 2).contiguous(), None
+
+
+def register_attention():
+    from transformers import AttentionInterface
+
+    AttentionInterface.register("et_vk_sdpa", et_vk_attention)
+    AttentionInterface.register("torch_sdpa_padmask", torch_attention)
+
+
+def site_tokens(prompt, n_tokens):
+    """planPrompt() in sam.js for a prompt that fits one decode: points (positive first), box corners,
+    one not-a-point pad (SAM 2's own padding), then pads up to the decoder's token count."""
+    pos = [t for t in prompt if t[0] == "pos"]
+    neg = [t for t in prompt if t[0] == "neg"]
+    corners = [t for t in prompt if t[0] in ("tl", "br")]
+    native = pos + neg + corners + [("pad",)]
+    if len(native) > n_tokens:
+        sys.exit(f"prompt {prompt} needs {len(native)} tokens; the decoder takes {n_tokens}")
+    return native + [("pad",)] * (n_tokens - len(native))
+
+
+def snapshot(weights):
+    """Pinned upstream snapshot: weights (sha256-checked) + configs at the same revision."""
+    w = WEIGHTS[weights]
+    path, digest = _common.hf_file(
+        w["repo"], "model.safetensors", w["revision"], w["files"]["model.safetensors"]
+    )
+    for f in CONFIG_FILES:
+        _common.hf_file(w["repo"], f, w["revision"])
+    return path.parent, digest
+
+
+def load_model(snap, size, attn="sdpa"):
+    from transformers import Sam2Config, Sam2Model
+
+    cfg = Sam2Config.from_pretrained(snap)
+    if size != cfg.image_size:
+        f = size // cfg.prompt_encoder_config.patch_size
+        cfg.image_size = size
+        cfg.prompt_encoder_config.image_size = size
+        cfg.vision_config.backbone_config.image_size = [size, size]
+        cfg.vision_config.backbone_feature_sizes = [
+            [4 * f, 4 * f],
+            [2 * f, 2 * f],
+            [f, f],
+        ]
+    model, info = Sam2Model.from_pretrained(
+        snap, config=cfg, attn_implementation=attn, output_loading_info=True
+    )
+    if info["missing_keys"] or info["mismatched_keys"] or info["unexpected_keys"]:
+        sys.exit(f"weights did not load cleanly: {info}")
+    return model.eval()
+
+
+def pool_windows_as_one_image():
+    """Hiera's q-pooling (2x2 max pool, stride 2) on windowed tensors [nW, h, w, C] is rewritten to pool
+    the windows stacked along H as a single [1, C, nW*h, w] image. h is a multiple of the stride, so no
+    pooling cell straddles two windows and the result is identical; batch-1 pooling is what the Vulkan
+    partitioner can represent (a [256, 384, 4, 4] pool is left on CPU otherwise)."""
+    from transformers.models.sam2 import modeling_sam2
+
+    original = getattr(modeling_sam2, "_upstream_do_pool", modeling_sam2.do_pool)
+    modeling_sam2._upstream_do_pool = original
+
+    def do_pool(x, query_stride=None):
+        if query_stride is None or x.shape[0] == 1:
+            return original(x, query_stride)
+        b, h, w, c = x.shape
+        sh, sw = (query_stride, query_stride) if isinstance(query_stride, int) else tuple(query_stride)
+        if h % sh or w % sw:
+            return original(x, query_stride)
+        y = x.reshape(1, b * h, w, c).permute(0, 3, 1, 2)
+        y = F.max_pool2d(y, kernel_size=query_stride, stride=query_stride, ceil_mode=False)
+        return y.permute(0, 2, 3, 1).reshape(b, h // sh, w // sw, c)
+
+    modeling_sam2.do_pool = do_pool
+
+
+def upstream_pool():
+    """Undo pool_windows_as_one_image() (upstream code path, for references)."""
+    from transformers.models.sam2 import modeling_sam2
+
+    modeling_sam2.do_pool = getattr(modeling_sam2, "_upstream_do_pool", modeling_sam2.do_pool)
+
+
+def freeze_hiera_pos_embed(model, size):
+    bb = model.vision_encoder.backbone
+    with torch.no_grad():
+        frozen = bb._get_pos_embed((size // 4, size // 4)).detach().contiguous().clone()
+    bb.register_buffer("frozen_pos_embed", frozen, persistent=False)
+    bb._get_pos_embed = types.MethodType(lambda self, _hw: self.frozen_pos_embed, bb)
+    return model
+
+
+class Encoder(torch.nn.Module):
+    def __init__(self, model, add_no_memory=True):
+        super().__init__()
+        self.vision = model.vision_encoder
+        self.add_no_memory = add_no_memory
+        self.register_buffer(
+            "no_mem",
+            model.no_memory_embedding.detach().reshape(1, -1, 1, 1).clone(),
+            persistent=False,
+        )
+
+    def forward(self, pixels):
+        hr0, hr1, emb = self.vision(pixels).fpn_hidden_states
+        if self.add_no_memory:
+            emb = emb + self.no_mem
+        return emb, hr0, hr1
+
+
+class Decoder(torch.nn.Module):
+    """Mask decoder on precomputed prompt embeddings.
+
+    A pad (not-a-point) token that directly follows another pad is hidden from every attention as a
+    key (additive -1e9, which underflows to an exact 0 weight). Nothing reads such a token's own
+    output, so a fixed-size decode of [pos, pad, pad] equals upstream's [pos, pad], and [tl, br, pad]
+    equals upstream's box prompt. The first pad is kept: it is SAM 2's own padding token."""
+
+    def __init__(self, model):
+        super().__init__()
+        self.md = model.mask_decoder
+        self.n_out = self.md.obj_score_token.weight.shape[0] + 1 + self.md.num_mask_tokens
+        self.register_buffer(
+            "not_a_point",
+            model.prompt_encoder.not_a_point_embed.weight[0].detach().clone(),
+            persistent=False,
+        )
+
+    def pad_bias(self, sparse):
+        d = (sparse - self.not_a_point).abs().sum(-1)  # [1, 1, N]; 0 only for an exact pad token
+        is_pad = torch.relu(1.0 - d * 1e6)  # 1.0 for a pad, 0.0 otherwise
+        repeated = is_pad[..., 1:] * is_pad[..., :-1]
+        keep = torch.zeros(1, 1, self.n_out + 1, dtype=sparse.dtype)
+        return (torch.cat([keep, repeated], dim=-1) * -1e9).unsqueeze(2)  # [1, 1, 1, n_out + N]
+
+    def forward(self, emb, image_pe, sparse, dense, hr0, hr1):
+        global _TOKEN_KEY_BIAS
+        _TOKEN_KEY_BIAS = self.pad_bias(sparse) if sparse.shape[2] > 1 else None
+        try:
+            return self._decode(emb, image_pe, sparse, dense, hr0, hr1)
+        finally:
+            _TOKEN_KEY_BIAS = None
+
+    def _decode(self, emb, image_pe, sparse, dense, hr0, hr1):
+        masks, iou, _, obj = self.md(
+            image_embeddings=emb,
+            image_positional_embeddings=image_pe,
+            sparse_prompt_embeddings=sparse,
+            dense_prompt_embeddings=dense,
+            multimask_output=True,
+            high_resolution_features=[self.md.conv_s0(hr0), self.md.conv_s1(hr1)],
+        )
+        # [B, point_batch=1, 3, H, W] -> [B, 3, H, W], as the previous export.
+        return masks[:, 0], iou, obj
+
+
+def prompt_tables(model):
+    pe = model.prompt_encoder
+    with torch.no_grad():
+        dense = (
+            pe.no_mask_embed.weight.reshape(1, -1, 1, 1)
+            .expand(1, -1, *pe.image_embedding_size)
+            .contiguous()
+        )
+        tabs = {
+            "gaussian": model.shared_image_embedding.positional_embedding,
+            "point_embed_bg": pe.point_embed.weight[0],
+            "point_embed_fg": pe.point_embed.weight[1],
+            "box_tl": pe.point_embed.weight[2],
+            "box_br": pe.point_embed.weight[3],
+            "not_a_point": pe.not_a_point_embed.weight[0],
+            "img_pos": model.get_image_wide_positional_embeddings(),
+            "dense": dense,
+        }
+    return {
+        k: v.detach().to(torch.float32).contiguous().numpy() for k, v in tabs.items()
+    }
+
+
+def encode_tokens(tokens, width, height, tabs, size=SIZE):
+    """Python port of the playground's browser-side prompt encoding (float64 math, f32 out)."""
+    g = tabs["gaussian"].reshape(-1).astype(np.float64)
+    hidden = tabs["not_a_point"].size
+    D = hidden // 2
+    kind = {
+        "neg": tabs["point_embed_bg"],
+        "pos": tabs["point_embed_fg"],
+        "tl": tabs["box_tl"],
+        "br": tabs["box_br"],
+    }
+    out = np.zeros((len(tokens), hidden), np.float64)
+    for r, t in enumerate(tokens):
+        if t[0] == "pad":
+            out[r] = tabs["not_a_point"]
+            continue
+        x, y = round(t[1] * 10) / 10, round(t[2] * 10) / 10
+        cx = (2 * (x * (size / width) + 0.5)) / size - 1
+        cy = (2 * (y * (size / height) + 0.5)) / size - 1
+        a = 2 * np.pi * (cx * g[:D] + cy * g[D:])
+        out[r, :D] = np.sin(a) + kind[t[0]][:D]
+        out[r, D:] = np.cos(a) + kind[t[0]][D:]
+    return torch.from_numpy(out.astype(np.float32)).reshape(1, 1, len(tokens), hidden)
+
+
+def load_image(path, size):
+    from PIL import Image
+
+    img = Image.open(path).convert("RGB")
+    a = np.asarray(img.resize((size, size), Image.BILINEAR), dtype=np.float32) / 255.0
+    a = (a - np.array(MEAN, np.float32)) / np.array(STD, np.float32)
+    return torch.from_numpy(a.transpose(2, 0, 1).copy()).unsqueeze(0), img
+
+
+def iou(a, b):
+    a, b = np.asarray(a) > 0, np.asarray(b) > 0
+    u = np.logical_or(a, b).sum()
+    return round(float(np.logical_and(a, b).sum() / u), 4) if u else 1.0
+
+
+def upsample(masks, w, h):
+    """Bilinear, align_corners=False: the same mapping as upsampleMask() in sam.js."""
+    t = torch.from_numpy(np.ascontiguousarray(masks, np.float32)).reshape(
+        1, -1, *np.asarray(masks).shape[-2:]
+    )
+    return F.interpolate(t, size=(h, w), mode="bilinear", align_corners=False)[
+        0
+    ].numpy()
+
+
+def upstream_inputs(tokens):
+    """The same prompt in transformers' native form (Sam2PromptEncoder pads points with a not-a-point
+    token and always adds one after box corners, like the original SAM 2 predictor)."""
+    pts = [t for t in tokens if t[0] in ("pos", "neg")]
+    corners = [t for t in tokens if t[0] in ("tl", "br")]
+    kw = {}
+    if corners:
+        (_, x0, y0), (_, x1, y1) = corners
+        kw["input_boxes"] = [[[x0, y0, x1, y1]]]
+    if pts:
+        kw["input_points"] = [[[[t[1], t[2]] for t in pts]]]
+        kw["input_labels"] = [[[1 if t[0] == "pos" else 0 for t in pts]]]
+    return kw
+
+
+def upstream_reference(model, processor, img, tokens):
+    """Upstream transformers pipeline end to end (Sam2Processor -> Sam2Model -> post_process_masks):
+    mask logits at the original image size, and the IoU head's scores."""
+    inputs = processor(images=img, return_tensors="pt", **upstream_inputs(tokens))
+    with torch.no_grad():
+        o = model(
+            pixel_values=inputs["pixel_values"],
+            input_points=inputs.get("input_points"),
+            input_labels=inputs.get("input_labels"),
+            input_boxes=inputs.get("input_boxes"),
+            multimask_output=True,
+        )
+    m = processor.post_process_masks(
+        o.pred_masks, inputs["original_sizes"], binarize=False
+    )[0]
+    return m[0].numpy(), o.iou_scores[0, 0].numpy()
+
+
+def upstream_512_same_pixels(model, pixels, tokens, w, h):
+    """Upstream Sam2Model.forward (512 config) on exactly the pixels the export sees."""
+    s = model.prompt_encoder.input_image_size
+    kw = upstream_inputs(tokens)
+    sx, sy = s / w, s / h
+    args = {}
+    if "input_boxes" in kw:
+        x0, y0, x1, y1 = kw["input_boxes"][0][0]
+        args["input_boxes"] = torch.tensor(
+            [[[x0 * sx, y0 * sy, x1 * sx, y1 * sy]]], dtype=torch.float32
+        )
+    if "input_points" in kw:
+        args["input_points"] = torch.tensor(
+            [[[[x * sx, y * sy] for x, y in kw["input_points"][0][0]]]],
+            dtype=torch.float32,
+        )
+        args["input_labels"] = torch.tensor(kw["input_labels"], dtype=torch.int64)
+    with torch.no_grad():
+        o = model(pixel_values=pixels, multimask_output=True, **args)
+    return o.pred_masks[0, 0].numpy(), o.iou_scores[0, 0].numpy()
+
+
+CONTRACT = {
+    "sam2_encoder": (
+        [[1, 3, SIZE, SIZE]],
+        [[1, 256, SIZE // 16, SIZE // 16], [1, 256, SIZE // 4, SIZE // 4], [1, 256, SIZE // 8, SIZE // 8]],
+    ),
+    "sam2_decoder": (
+        [[1, 256, SIZE // 16, SIZE // 16], [1, 256, SIZE // 16, SIZE // 16], [1, 1, "N", 256],
+         [1, 256, SIZE // 16, SIZE // 16], [1, 256, SIZE // 4, SIZE // 4], [1, 256, SIZE // 8, SIZE // 8]],
+        [[1, 3, SIZE // 4, SIZE // 4], [1, 1, 3], [1, 1, 1]],
+    ),
+}
+
+
+def program_io(pte_bytes):
+    """Method names and forward() I/O (shape, dtype, dim order) as serialized in the program."""
+    from executorch.exir._serialize._program import deserialize_pte_binary
+
+    pte = deserialize_pte_binary(pte_bytes)
+    prog = pte.program if hasattr(pte, "program") else pte
+
+    def desc(plan, i):
+        v = plan.values[i].val
+        return {"shape": list(v.sizes), "dtype": v.scalar_type.name, "dim_order": list(v.dim_order)}
+
+    plan = {p.name: p for p in prog.execution_plan}
+    f = plan["forward"]
+    return {
+        "methods": sorted(plan),
+        "inputs": [desc(f, i) for i in f.inputs],
+        "outputs": [desc(f, i) for i in f.outputs],
+    }
+
+
+def check_contract(name, io, size, n_tokens):
+    """The playground's programs: one `forward`, fp32, contiguous NCHW, these shapes."""
+    if size != SIZE:
+        return "skipped (non-default size)"
+    want_in, want_out = CONTRACT[name]
+    want_in = [[n_tokens if d == "N" else d for d in s] for s in want_in]
+    ok = (
+        io["methods"] == ["forward"]
+        and [t["shape"] for t in io["inputs"]] == want_in
+        and [t["shape"] for t in io["outputs"]] == want_out
+        and all(t["dtype"] == "FLOAT" and t["dim_order"] == list(range(len(t["shape"]))) for t in io["inputs"] + io["outputs"])
+    )
+    if not ok:
+        sys.exit(f"{name}: I/O does not match the previous contract: {json.dumps(io)}")
+    return "matches the previous export (shapes, fp32, contiguous dim order, single forward method)"
+
+
+def write_if_changed(path, data):
+    """Leave an identical file alone; otherwise replace it with a new inode, so hard links made to an
+    earlier output elsewhere keep their content."""
+    data = bytes(data)
+    if path.exists():
+        if path.read_bytes() == data:
+            return
+        path.unlink()
+    path.write_bytes(data)
+
+
+def export_programs(enc, dec, enc_std, dec_std, pixels, dec_inputs):
+    with torch.no_grad():
+        ep_enc = torch.export.export(enc, (pixels,), strict=False)
+        ep_dec = torch.export.export(dec, dec_inputs, strict=False)
+        ep_enc_std = torch.export.export(enc_std, (pixels,), strict=False)
+        ep_dec_std = torch.export.export(dec_std, dec_inputs, strict=False)
+        return {
+            "sam2_encoder": (
+                ep_enc,
+                _common.lower_webgpu(ep_enc),
+                _common.lower_portable(ep_enc_std),
+            ),
+            "sam2_decoder": (
+                ep_dec,
+                _common.lower_webgpu(ep_dec),
+                _common.lower_portable(ep_dec_std),
+            ),
+        }
+
+
+def validation_images(args):
+    from huggingface_hub import hf_hub_download
+
+    out = []
+    p = Path(
+        hf_hub_download(
+            IMAGE["repo"],
+            IMAGE["file"],
+            repo_type=IMAGE["repo_type"],
+            revision=IMAGE["revision"],
+        )
+    )
+    digest = _common.sha256_file(p)
+    if digest != IMAGE["sha256"]:
+        sys.exit(f"{IMAGE['file']}: sha256 {digest} != {IMAGE['sha256']}")
+    out.append(
+        (
+            IMAGE["name"],
+            p,
+            digest,
+            IMAGE["prompts"],
+            f"hf://datasets/{IMAGE['repo']}@{IMAGE['revision']}/{IMAGE['file']}",
+        )
+    )
+    if args.assets_dir:
+        for name, spec in ASSET_IMAGES.items():
+            ap = Path(args.assets_dir) / name
+            d = _common.sha256_file(ap)
+            if spec["sha256"] and d != spec["sha256"]:
+                sys.exit(f"{ap}: sha256 {d} != {spec['sha256']}")
+            out.append(
+                (
+                    name,
+                    ap,
+                    d,
+                    spec["prompts"],
+                    f"sample image {name}",
+                )
+            )
+    return out
+
+
+def validate(args, snap, model, tabs, progs, enc, dec, enc_std):
+    from transformers import Sam2Processor
+
+    val = {
+        "cpu_runtime": "executorch pybindings (portable_lib) from the executorch 1.5.1 wheel"
+    }
+    ep_enc, wg_enc, port_enc = progs["sam2_encoder"]
+    ep_dec, wg_dec, port_dec = progs["sam2_decoder"]
+    img_pe = torch.from_numpy(tabs["img_pos"])
+    dense = torch.from_numpy(tabs["dense"])
+
+    ref512 = load_model(snap, args.size)
+    ref1024 = load_model(snap, 1024)
+    processor = Sam2Processor.from_pretrained(snap)
+    other = Encoder(model, add_no_memory=not enc.add_no_memory).eval()
+    # Everything below runs the upstream do_pool; the export used the rewrite (checked to be exact).
+    upstream_pool()
+
+    for p_name, ep, wg in (("encoder", ep_enc, wg_enc), ("decoder", ep_dec, wg_dec)):
+        _, err = _common.try_run_webgpu_pte_on_cpu(wg.buffer, ep.example_inputs[0])
+        val[f"webgpu_{p_name}_pte_on_cpu"] = (
+            f"not runnable on CPU pybindings: {err}" if err else "ran"
+        )
+
+    val["images"] = {}
+    refs = Path(args.refs_dir) if args.refs_dir else None
+    refs_meta = {"prompt_tokens": args.prompt_tokens, "prompts": {}}
+    for name, path, digest, prompts, src in validation_images(args):
+        pixels, img = load_image(path, args.size)
+        W, H = img.size
+        with torch.no_grad():
+            ref = enc(pixels)
+            pool_windows_as_one_image()
+            rewritten = enc(pixels)
+            upstream_pool()
+            exported = ep_enc.module()(pixels)
+            std = enc_std(pixels)
+            emb_other = other(pixels)[0]
+        port = _common.run_pte(port_enc.buffer, (pixels,))
+        rec = {
+            "source": src,
+            "sha256": digest,
+            "size": [W, H],
+            "encoder": {},
+            "prompts": {},
+        }
+        for i, n in enumerate(["image_embeddings", "hr0", "hr1"]):
+            rec["encoder"][n] = {
+                "exported_vs_eager": _common.compare(
+                    exported[i].numpy(), ref[i].numpy()
+                ),
+                "pool_rewrite_vs_upstream_code_max_abs": float((rewritten[i] - ref[i]).abs().max()),
+                "eager_torch_sdpa_vs_eager_et_vk_sdpa": _common.compare(
+                    std[i].numpy(), ref[i].numpy()
+                ),
+                "portable_pte_vs_eager": _common.compare(port[i], ref[i].numpy()),
+            }
+        emb, hr0, hr1 = (t.contiguous() for t in ref)
+        if refs is not None and name == IMAGE["name"]:
+            refs.mkdir(parents=True, exist_ok=True)
+            refs.joinpath("pixels.f32").write_bytes(pixels.numpy().tobytes())
+            for i, t in enumerate((emb, hr0, hr1)):
+                refs.joinpath(f"enc_out{i}.f32").write_bytes(t.numpy().tobytes())
+            refs_meta.update(image=src, image_sha256=digest, size=[W, H])
+        for pname, toks in prompts.items():
+            dtoks = site_tokens(toks, args.prompt_tokens)
+            sparse = encode_tokens(dtoks, W, H, tabs, args.size)
+            ins = (emb, img_pe, sparse, dense, hr0, hr1)
+            with torch.no_grad():
+                eager = [t.numpy() for t in dec(*ins)]
+                exp = [t.numpy() for t in ep_dec.module()(*ins)]
+                alt_out = dec(emb_other, img_pe, sparse, dense, hr0, hr1)
+                alt, alt_iou = alt_out[0].numpy(), alt_out[1].numpy()
+            # End to end through both portable programs (encoder output -> decoder input).
+            pd = _common.run_pte(
+                port_dec.buffer,
+                (
+                    torch.from_numpy(np.ascontiguousarray(port[0])),
+                    img_pe,
+                    sparse,
+                    dense,
+                    torch.from_numpy(np.ascontiguousarray(port[1])),
+                    torch.from_numpy(np.ascontiguousarray(port[2])),
+                ),
+            )
+            m_pte = pd[0][0]
+            u512_m, u512_iou = upstream_512_same_pixels(ref512, pixels, toks, W, H)
+            u_m, u_iou = upstream_reference(ref1024, processor, img, toks)
+            up = upsample(m_pte, W, H)
+            up_alt = upsample(alt[0], W, H)
+            best = int(np.argmax(u_iou))
+            own = int(np.argmax(np.asarray(pd[1]).reshape(-1)))
+            alt_own = int(np.argmax(alt_iou.reshape(-1)))
+            if refs is not None and name == IMAGE["name"]:
+                for i, t in enumerate(eager):
+                    refs.joinpath(f"dec_{pname}_out{i}.f32").write_bytes(np.ascontiguousarray(t, np.float32).tobytes())
+                refs.joinpath(f"sparse_{pname}.f32").write_bytes(sparse.numpy().tobytes())
+                refs.joinpath(f"up512_{pname}_masks.f32").write_bytes(np.ascontiguousarray(u512_m, np.float32).tobytes())
+                refs_meta["prompts"][pname] = {"prompt": [list(t) for t in toks], "tokens": [list(t) for t in dtoks]}
+            rec["prompts"][pname] = {
+                "prompt": [list(t) for t in toks],
+                "decoder_tokens": [t[0] for t in dtoks],
+                "eager_vs_upstream_512_masks": _common.compare(eager[0][0], u512_m),
+                "eager_vs_upstream_512_iou_scores": _common.compare(eager[1], u512_iou),
+                "exported_vs_eager_masks": _common.compare(exp[0], eager[0]),
+                "portable_e2e_vs_eager_masks": _common.compare(pd[0], eager[0]),
+                "portable_e2e_vs_eager_iou_scores": _common.compare(pd[1], eager[1]),
+                "portable_e2e_vs_eager_object_score": _common.compare(pd[2], eager[2]),
+                "mask_iou_portable_vs_eager": [
+                    iou(m_pte[i], eager[0][0][i]) for i in range(3)
+                ],
+                "mask_iou_portable_vs_upstream_512_same_pixels": [
+                    iou(m_pte[i], u512_m[i]) for i in range(3)
+                ],
+                "mask_iou_portable_vs_upstream_1024_pipeline_at_image_res": [
+                    iou(up[i], u_m[i]) for i in range(3)
+                ],
+                "portable_iou_scores": np.asarray(pd[1]).reshape(-1).round(4).tolist(),
+                "upstream_512_iou_scores": u512_iou.round(4).tolist(),
+                "upstream_1024_iou_scores": u_iou.round(4).tolist(),
+                "upstream_1024_best_mask": best,
+                "best_mask_iou_portable_vs_upstream_1024": iou(up[best], u_m[best]),
+                "portable_object_score_logit": round(
+                    float(np.asarray(pd[2]).reshape(-1)[0]), 4
+                ),
+                "mask_area_fraction_portable": [
+                    round(float((m_pte[i] > 0).mean()), 4) for i in range(3)
+                ],
+                "upstream_1024_best_mask_vs_closest_export_mask": max(iou(up[i], u_m[best]) for i in range(3)),
+                "export_best_by_own_iou_head_vs_upstream_1024_best": iou(up[own], u_m[best]),
+                "variant_%s_no_memory_embedding" % ("without" if enc.add_no_memory else "with"): {
+                    "mask_iou_vs_this_export": [iou(alt[0][i], eager[0][0][i]) for i in range(3)],
+                    "mask_iou_vs_upstream_1024_pipeline_at_image_res": [iou(up_alt[i], u_m[i]) for i in range(3)],
+                    "upstream_1024_best_mask_vs_closest_mask": max(iou(up_alt[i], u_m[best]) for i in range(3)),
+                    "best_by_own_iou_head_vs_upstream_1024_best": iou(up_alt[alt_own], u_m[best]),
+                },
+            }
+        val["images"][name] = rec
+        print(json.dumps({name: rec}, indent=1), flush=True)
+    if refs is not None:
+        refs.joinpath("meta.json").write_text(json.dumps(refs_meta, indent=1) + "\n")
+
+    if args.internal_dir:
+        d = Path(args.internal_dir)
+        val["prompt_tables_vs_internal"] = {
+            f"prompt_{k}.bin": "bit-identical"
+            if (d / f"prompt_{k}.bin").read_bytes() == v.astype(np.float32).tobytes()
+            else "max abs %.3g"
+            % float(
+                np.abs(
+                    np.fromfile(d / f"prompt_{k}.bin", dtype=np.float32) - v.reshape(-1)
+                ).max()
+            )
+            for k, v in tabs.items()
+        }
+    return val
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--weights", choices=sorted(WEIGHTS), default="sam2.1")
+    ap.add_argument(
+        "--size", type=int, default=SIZE, help="square input size (multiple of 32)"
+    )
+    ap.add_argument(
+        "--no-memory-embedding",
+        choices=["add", "omit"],
+        default="add",
+        help="add (upstream single-image behaviour, default) or omit (as the previous export)",
+    )
+    ap.add_argument(
+        "--assets-dir",
+        help="optional: directory of sample images for extra validation",
+    )
+    ap.add_argument(
+        "--internal-dir",
+        help="optional: compare prompt tables with an existing prompt_*.bin set",
+    )
+    ap.add_argument(
+        "--prompt-tokens",
+        type=int,
+        default=3,
+        help="sparse prompt tokens per decode (trailing extra pads are masked, so any prompt that fits matches upstream)",
+    )
+    ap.add_argument(
+        "--refs-dir",
+        help="optional: write CPU reference tensors for the browser harness (sam-car.png prompts)",
+    )
+    ap.add_argument(
+        "--skip-validation",
+        action="store_true",
+        help="export only (for the byte-identical re-run check)",
+    )
+    args = ap.parse_args()
+    out = Path(args.out)
+    out.mkdir(parents=True, exist_ok=True)
+    torch.manual_seed(0)
+
+    snap, digest = snapshot(args.weights)
+    register_attention()
+    pool_windows_as_one_image()
+    add_nm = args.no_memory_embedding == "add"
+    model = freeze_hiera_pos_embed(
+        load_model(snap, args.size, attn="et_vk_sdpa"), args.size
+    )
+    enc, dec = Encoder(model, add_nm).eval(), Decoder(model).eval()
+    tabs = prompt_tables(model)
+    # Same weights with torch SDPA: only for the CPU portable .pte (et_vk.sdpa has no portable kernel).
+    model_std = freeze_hiera_pos_embed(
+        load_model(snap, args.size, attn="sdpa"), args.size
+    )
+    model_std_dec = load_model(snap, args.size, attn="torch_sdpa_padmask")
+    enc_std, dec_std = Encoder(model_std, add_nm).eval(), Decoder(model_std_dec).eval()
+
+    # Example inputs for export: a fixed-seed tensor (the graph does not depend on values).
+    gen = torch.Generator().manual_seed(0)
+    pixels = torch.randn(1, 3, args.size, args.size, generator=gen)
+    with torch.no_grad():
+        # Contiguous example inputs: torch.export records each input's dim order.
+        emb, hr0, hr1 = (t.contiguous() for t in enc(pixels))
+    sparse = encode_tokens(
+        site_tokens([("pos", args.size / 2, args.size / 2)], args.prompt_tokens),
+        args.size,
+        args.size,
+        tabs,
+        args.size,
+    )
+    dec_inputs = (
+        emb,
+        torch.from_numpy(tabs["img_pos"]),
+        sparse,
+        torch.from_numpy(tabs["dense"]),
+        hr0,
+        hr1,
+    )
+    progs = export_programs(enc, dec, enc_std, dec_std, pixels, dec_inputs)
+
+    files, graphs = {}, {}
+    for name, (_, wg, port) in progs.items():
+        write_if_changed(out / f"{name}.pte", wg.buffer)
+        write_if_changed(out / f"{name}_portable.pte", port.buffer)
+        graphs[name] = _common.delegate_summary(wg.buffer)
+        graphs[name]["io"] = program_io(wg.buffer)
+        graphs[name]["io_contract"] = check_contract(name, graphs[name]["io"], args.size, args.prompt_tokens)
+        if graphs[name]["non_delegated_kernels"] or len(graphs[name]["delegates"]) != 1:
+            sys.exit(f"{name}: expected one delegate and no CPU kernels: {json.dumps(graphs[name])[:2000]}")
+    for k, v in tabs.items():
+        write_if_changed(out / f"prompt_{k}.bin", v.astype(np.float32).tobytes())
+    for p in sorted(out.iterdir()):
+        if p.suffix in (".pte", ".bin"):
+            files[p.name] = {
+                "bytes": p.stat().st_size,
+                "sha256": _common.sha256_file(p),
+            }
+            if p.name.endswith("_portable.pte"):
+                files[p.name]["note"] = (
+                    "CPU-only validation artifact (portable kernels, no delegate; torch SDPA instead of et_vk.sdpa); not for release"
+                )
+
+    w = WEIGHTS[args.weights]
+    q = args.size // 4
+    recipe = {
+        "model": f"{w['repo'].split('/')[1]} ({args.size} encoder/decoder split)",
+        "upstream": {
+            "model": w["repo"],
+            "revision": w["revision"],
+            "official": True,
+            "license": "Apache-2.0",
+            "files": {"model.safetensors": {"sha256": digest}},
+            "note": "transformers Sam2Model port published by the SAM 2 authors on the Hub",
+        },
+        "versions": _common.versions(),
+        "export": {
+            "api": "torch.export.export(strict=False) -> to_edge_transform_and_lower(partitioner=[WebGPUPartitioner()]) -> to_executorch(), one program per module",
+            "partitioner": "executorch.backends.webgpu.partitioner.WebGPUPartitioner() (default compile options; wraps VulkanPartitioner)",
+            "model_code": "transformers Sam2Model; attention through torch.ops.et_vk.sdpa (transformers attention backend, recipes/sam2-hiera-tiny.py:et_vk_attention); config image_size overridden 1024 -> %d"
+            % args.size,
+            "programs": {
+                "sam2_encoder.pte": {
+                    "method": "forward",
+                    "inputs": [
+                        {
+                            "name": "pixels",
+                            "shape": [1, 3, args.size, args.size],
+                            "dtype": "float32",
+                        }
+                    ],
+                    "outputs": [
+                        {
+                            "name": "image_embeddings",
+                            "shape": list(emb.shape),
+                            "dtype": "float32",
+                            "note": "no_memory_embedding %s"
+                            % ("added" if add_nm else "not added"),
+                        },
+                        {
+                            "name": "hr0",
+                            "shape": list(hr0.shape),
+                            "dtype": "float32",
+                            "note": "before mask_decoder.conv_s0",
+                        },
+                        {
+                            "name": "hr1",
+                            "shape": list(hr1.shape),
+                            "dtype": "float32",
+                            "note": "before mask_decoder.conv_s1",
+                        },
+                    ],
+                },
+                "sam2_decoder.pte": {
+                    "method": "forward",
+                    "prompt_tokens": args.prompt_tokens,
+                    "inputs": [
+                        {"name": n, "shape": list(t.shape), "dtype": "float32"}
+                        for n, t in zip(INPUT_NAMES, dec_inputs)
+                    ],
+                    "outputs": [
+                        {
+                            "name": "masks",
+                            "shape": [1, 3, q, q],
+                            "dtype": "float32",
+                            "note": "logits, mask tokens 1..3",
+                        },
+                        {"name": "iou_scores", "shape": [1, 1, 3], "dtype": "float32"},
+                        {
+                            "name": "object_score_logits",
+                            "shape": [1, 1, 1],
+                            "dtype": "float32",
+                        },
+                    ],
+                },
+            },
+            "prompt_tables": {
+                f"prompt_{k}.bin": {
+                    "shape": list(v.shape),
+                    "dtype": "float32 little-endian",
+                }
+                for k, v in tabs.items()
+            },
+            "prompt_table_sources": {
+                "prompt_gaussian.bin": "shared_image_embedding.positional_embedding",
+                "prompt_point_embed_bg.bin": "prompt_encoder.point_embed.weight[0]",
+                "prompt_point_embed_fg.bin": "prompt_encoder.point_embed.weight[1]",
+                "prompt_box_tl.bin": "prompt_encoder.point_embed.weight[2]",
+                "prompt_box_br.bin": "prompt_encoder.point_embed.weight[3]",
+                "prompt_not_a_point.bin": "prompt_encoder.not_a_point_embed.weight[0]",
+                "prompt_img_pos.bin": "Sam2Model.get_image_wide_positional_embeddings() with image_size=%d"
+                % args.size,
+                "prompt_dense.bin": "prompt_encoder.no_mask_embed.weight broadcast to [1,256,%d,%d]"
+                % (args.size // 16, args.size // 16),
+            },
+            "graph_rewrites": [
+                "Hiera q-pooling of windowed tensors [nW,h,w,C] done as one batch-1 max_pool2d over the windows stacked along H (exact: h is a multiple of the stride); keeps all pools in the delegate",
+                "attention via torch.ops.et_vk.sdpa (same math as torch SDPA; checked in validation.json)",
+                "decoder: a not-a-point token directly after another one gets an additive -1e9 key bias in every attention (computed in-graph from the sparse input), so padding up to the fixed token count does not change the result",
+                "Hiera position embedding (bicubic pos_embed resize + tiled window embedding) for the fixed input size computed in eager and stored as a constant",
+            ],
+            "differences_from_upstream_pipeline": [
+                "input %dx%d instead of 1024x1024 (masks %dx%d instead of 256x256)"
+                % (args.size, args.size, q, q),
+                "fixed %d prompt tokens per decode: points, box corners, one not-a-point pad, then padding pads that the decoder masks out (exactly upstream's prompt for any prompt that fits; larger prompts are composed by the site)" % args.prompt_tokens,
+                "no mask prompt input (dense = no_mask_embed)",
+            ],
+            "preprocessing_outside_graph": "RGB, resize to %dx%d without keeping aspect, /255, ImageNet mean/std"
+            % (args.size, args.size),
+            "dtype": "float32",
+            "quantization": "none (fp32)",
+        },
+        "files": files,
+        "graph": graphs,
+        "command": _common.command_line(out),
+    }
+    (out / "recipe.json").write_text(json.dumps(recipe, indent=2) + "\n")
+    print(json.dumps({"files": files, "graph": graphs}, indent=1), flush=True)
+    if args.skip_validation:
+        return
+    val = validate(args, snap, model, tabs, progs, enc, dec, enc_std)
+    (out / "validation.json").write_text(json.dumps(val, indent=2) + "\n")
+    print(json.dumps(val, indent=1))
+
+
+if __name__ == "__main__":
+    main()

manifest.json

@@ -44,13 +44,13 @@
     },
     {
       "path": "sam2_decoder.pte",
-      "bytes": 16935552,
-      "sha256": "4d6f790ccf3c7ae8f54fff162202703d3f9e78099b25c3e80b70a2c4eb3e1ac8"
+      "bytes": 16949252,
+      "sha256": "57b91034acc71be2ce1835ea546a174295daee1046e40db0743e0238ef11e278"
     },
     {
       "path": "sam2_encoder.pte",
-      "bytes": 137529216,
-      "sha256": "46626cd828cca2f22fb2623ab7b9cafeec8f724355ac8d7265aa1d54245fe5dd"
+      "bytes": 115248512,
+      "sha256": "30f9fb469b779558b6b5c7c3eb995c50e1213ffbe082e31430c929ba2b3ffe20"
     }
   ]
 }

sam2_decoder.pte

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4d6f790ccf3c7ae8f54fff162202703d3f9e78099b25c3e80b70a2c4eb3e1ac8
-size 16935552
+oid sha256:57b91034acc71be2ce1835ea546a174295daee1046e40db0743e0238ef11e278
+size 16949252

sam2_encoder.pte

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:46626cd828cca2f22fb2623ab7b9cafeec8f724355ac8d7265aa1d54245fe5dd
-size 137529216
+oid sha256:30f9fb469b779558b6b5c7c3eb995c50e1213ffbe082e31430c929ba2b3ffe20
+size 115248512