export/whisper-base.py

#!/usr/bin/env python3
"""Whisper base (openai/whisper-base) -> six ExecuTorch WebGPU programs + host tables, public ExecuTorch.

Same split and program I/O contract as the playground's previous export (per-program contract in
recipes/_whisper.py):

  whisper_mel.pte            frames[1,3000,400] -> power mel[1,3000,80]                        fp32
  whisper_encoder.pte        log-mel[1,80,3000] -> encoder hidden[1,1500,512]                  fp32
  whisper_cross_kv.pte       encoder hidden -> 12 x [1,8,1500,64] (k0,v0,...,k5,v5)
  whisper_decoder_split.pte  (tok embeds[1,64,512], 12 cross K/V) -> hidden[1,64,512]
  whisper_decode_stream.pte  (embed+pos[1,1,512], pos[1] i64, 12 cross K/V) -> hidden[1,1,512]
  whisper_lmhead.pte         hidden[1,1,512] -> logits[1,1,51872] (rows >= 51865 are -1e9)
  whisper_pos_embed.f32.bin, vocab.json, and the token-embedding table:

  --precision q4   (default; the previous export's numerics) cross_kv / decoder_split / decode_stream
                   / lmhead linears int4 groupwise (et_vk.linear_q4gsw, group 64; quantized weights
                   bit-identical to the previous files) + whisper_embed.i8.bin / .scale.f32.bin.
  --precision fp32 every program fp32 + whisper_embed.f32.bin. Larger, but its greedy transcripts
                   are token-identical to transformers (int4 weights are not; see validation.json).
  --precision hybrid-lmfp32  as q4, but lmhead fp32 (aten.linear).
  --precision hybrid-lmq8    as q4, but lmhead int8 x int8 (static per-tensor input scale calibrated on
                   the 3 LibriSpeech-dummy clips longer than 15 s; et_vk.linear_q8ta_q8csw).
  --precision fp32dec-lmq8   fp32 cross_kv / decoder_split / decode_stream, the int8 x int8 lmhead above
                   (calibrated on the fp32 decoder's hidden states) and the int8 embedding table.

Every program is lowered with executorch.backends.webgpu.partitioner.WebGPUPartitioner() into one
VulkanBackend delegate with no CPU kernels. The host does: framing, the log-mel normalization between
mel and encoder (log10 / max-8 / (x+4)/4 / transpose; the public runtime has no log op), token
embedding lookup (+ positional row for decode_stream), and argmax.

  python recipes/whisper-base.py --out out/whisper-base/q4                     # export + validation
  python recipes/whisper-base.py --out out/whisper-base/fp32 --precision fp32
  python recipes/whisper-base.py --out /tmp/w --no-validate                    # export only

Validation (CPU only; the WebGPU kernels cannot run here): eager vs torch.export program vs portable
.pte (ExecuTorch pybindings) for every program, and greedy transcripts of LibriSpeech clips
(hf-internal-testing/librispeech_asr_dummy) through the full host + program pipeline, compared token
by token with transformers WhisperForConditionalGeneration.
"""

import argparse
import copy
import json
import shutil
import sys
import time
from pathlib import Path

import numpy as np
import torch

sys.path.insert(0, str(Path(__file__).resolve().parent))
import _common  # noqa: E402
import _whisper as W  # noqa: E402

UPSTREAM = {
    "repo": "openai/whisper-base",
    "revision": "e37978b90ca9030d5170a5c07aadb050351a65bb",
    "files": {
        "model.safetensors": "07cadb9f25677c8d50df603e66a98fbd842cce45047139baeb16e6219a1e807b",
        "config.json": "a153c53883a6799b6f056b4a8d1a515c9926d03994682ba88a7616618d7da0c1",
        "generation_config.json": "444b3f636d2fff89dd9ecf549e2a085b61f7ff0fa0246d4628bac6a3b8cc9ba4",
        "preprocessor_config.json": "9b5cd03a36fbb8a627c64d98a5b5b126ead95a77720723944487311f0110b666",
        "vocab.json": "8f680bba319e01a653d2e8a5dbc17a9157179e0576e6ce74ce0c06356c6e24f9",
        "tokenizer.json": "27fc476bfe7f17299480be2273fc0608e4d5a99aba2ab5dec5374b4482d1a566",
        "tokenizer_config.json": "2a4c4281cf9f51ac6ccc406fdc711a087afe6530f671fa7b80953edc498275ce",
        "merges.txt": "2df2990a395e35e8dfbc7511e08c12d56018d8d04691e0133e5d63b21e154dc6",
        "normalizer.json": "bf1c507dc8724ca9cf9903640dacfb69dae2f00edee4f21ceba106a7392f26dd",
        "added_tokens.json": "9715fd2243b6f06a5858b5e32950d2853f73dd5bc201aafcf76f5082a2d8acd1",
        "special_tokens_map.json": "e67ae3a0aaa99abcd9f187138e12db1f65c16a14761c50ef10eef2c174a7a691",
    },
}
SAMPLE = {
    "repo": "hf-internal-testing/librispeech_asr_dummy",
    "repo_type": "dataset",
    "revision": "5be91486e11a2d616f4ec5db8d3fd248585ac07a",
    "file": "clean/validation-00000-of-00001.parquet",
    "sha256": "4e69a06fa5edc90921e5e7e39a7084881f8b3ed9c805c574f4f39c6fde27c603",
}
NAMES = ["mel", "encoder", "cross_kv", "decoder_split", "decode_stream", "lmhead"]
DECODER_SIDE = ["cross_kv", "decoder_split", "decode_stream", "lmhead"]


def fetch_upstream():
    snap, digests = None, {}
    for f, sha in UPSTREAM["files"].items():
        p, d = _common.hf_file(UPSTREAM["repo"], f, UPSTREAM["revision"], sha)
        digests[f] = d
        snap = p.parent
    return snap, digests


def fetch_sample():
    from huggingface_hub import hf_hub_download

    p = hf_hub_download(
        SAMPLE["repo"],
        SAMPLE["file"],
        repo_type=SAMPLE["repo_type"],
        revision=SAMPLE["revision"],
    )
    d = _common.sha256_file(p)
    if d != SAMPLE["sha256"]:
        sys.exit(f"{SAMPLE['repo']}:{SAMPLE['file']} sha256 {d} != {SAMPLE['sha256']}")
    return p


def example_inputs(name):
    g = torch.Generator().manual_seed(0)
    ckv = tuple(torch.randn(1, 8, W.ENC_LEN, 64, generator=g) for _ in range(12))
    return {
        "mel": (torch.randn(1, W.FRAMES, W.N_FFT, generator=g),),
        "encoder": (torch.randn(1, W.N_MELS, W.FRAMES, generator=g),),
        "cross_kv": (torch.randn(1, W.ENC_LEN, W.D_MODEL, generator=g),),
        "decoder_split": (torch.randn(1, W.MAX_DEC, W.D_MODEL, generator=g),) + ckv,
        "decode_stream": (torch.randn(1, 1, W.D_MODEL, generator=g), torch.tensor([0]))
        + ckv,
        "lmhead": (torch.randn(1, 1, W.D_MODEL, generator=g),),
    }[name]


PRECISIONS = {
    # name: (decoder-side linears, lmhead, host embedding table)
    "q4": ("int4", "int4", "int8"),
    "hybrid-lmfp32": ("int4", "fp32", "int8"),
    "hybrid-lmq8": ("int4", "int8 static-act", "int8"),
    "fp32dec-lmq8": ("fp32", "int8 static-act", "int8"),
    "fp32": ("fp32", "fp32", "fp32"),
}


def calibration_hidden(model, fe, dec, tables, max_sec):
    """lmhead inputs for the q8 head's static activation scale: the variant's decoder's hidden states
    for every decoded position of the LibriSpeech dummy clips LONGER than max_sec (3 clips, first
    30 s window), which are disjoint from the evaluation clips. Tokens are chosen with the fp32
    head (greedy, playground prefix)."""
    clips = W.load_librispeech_dummy(fetch_sample(), max_sec, 100, longer=True)
    stream = W.DecodeStreamProgram(dec, "torch")
    ckvp = W.CrossKVProgram(dec)
    mel, enc = W.MelProgram(fe), W.EncoderProgram(model, "eager")
    lm32 = W.LMHeadProgram(model, W.VOCAB_PADDED)
    hs = []
    with torch.no_grad():
        for c in clips:
            lmel = W.host_log_mel(mel(torch.from_numpy(W.build_frames(c["pcm"]))[None]).numpy())
            ckv = ckvp(enc(torch.from_numpy(lmel)))
            stream.reset()
            ids = list(W.PREFIX_EN)
            for p in range(W.MAX_DEC):
                h = stream(torch.from_numpy(tables.embed(ids[p], p)[None, None]), torch.tensor([p]), *ckv)
                if p < len(W.PREFIX_EN) - 1:
                    continue
                hs.append(h)
                nxt = int(lm32(h).reshape(-1)[: W.VOCAB].argmax())
                if nxt == W.TOK["eot"] or len(ids) + 1 >= W.MAX_DEC:
                    break
                ids.append(nxt)
    return torch.cat(hs, 0), [c["id"] for c in clips]


def build_programs(model, fe, precision, encoder_attn, tables, max_sec):
    """name -> (module lowered for WebGPU, module for the CPU portable .pte, fp32 reference or None).

    The portable module computes the same function with ops the portable runtime has: attention as
    matmul+softmax instead of et_vk.sdpa, the fp32 dequantization of the same int4 weights, and for
    the q8 head an explicit per-tensor fake-quant + dequantized int8 weights."""
    dec_kind, lm_kind, _ = PRECISIONS[precision]
    fdec = model.model.decoder
    lm32 = W.LMHeadProgram(model, W.VOCAB_PADDED)
    info = {}
    if dec_kind == "int4":
        wdec = W.quantize_q4(copy.deepcopy(fdec))
        pdec = W.dequantized_copy(wdec)
        dref = lambda m: m  # noqa: E731
    else:
        wdec = pdec = fdec
        dref = lambda m: None  # noqa: E731
    if lm_kind == "int4":
        wlm = W.quantize_q4(copy.deepcopy(lm32))
        plm, lref = W.dequantized_copy(wlm), lm32
    elif lm_kind == "fp32":
        wlm = plm = lm32
        lref = None
    else:
        calib, calib_ids = calibration_hidden(model, fe, wdec, tables, max_sec)
        wlm, qp, wdq = W.q8_lmhead(lm32, calib)
        plm, lref = W.FakeQ8LMHead(lm32, qp, wdq), lm32
        info = {
            "calibration_clips": calib_ids,
            "calibration_vectors": int(calib.shape[0]),
            "input_scale": qp[0],
            "input_zero_point": qp[1],
        }
    progs = {
        "mel": (W.MelProgram(fe), W.MelProgram(fe), None),
        "encoder": (W.EncoderProgram(model, encoder_attn), W.EncoderProgram(model, "eager"), None),
        "cross_kv": (W.CrossKVProgram(wdec), W.CrossKVProgram(pdec), dref(W.CrossKVProgram(fdec))),
        "decoder_split": (
            W.DecoderSplitProgram(wdec, "et_vk"),
            W.DecoderSplitProgram(pdec, "torch"),
            dref(W.DecoderSplitProgram(fdec, "torch")),
        ),
        "decode_stream": (
            W.DecodeStreamProgram(wdec, "et_vk"),
            W.DecodeStreamProgram(pdec, "torch"),
            dref(W.DecodeStreamProgram(fdec, "torch")),
        ),
        "lmhead": (wlm, plm, lref),
    }
    return progs, info


def export_one(module, inputs):
    if isinstance(module, W.DecodeStreamProgram):
        module.reset()  # the serialized K/V cache must start as zeros
    with torch.no_grad():
        if not isinstance(module, torch.fx.GraphModule):  # pt2e-converted modules have no eval()
            module = module.eval()
        return torch.export.export(module, inputs, strict=False)


# ---------------------------------------------------------------------------------------------
# Runners: run(name, [arrays]) -> [np arrays]; reset() clears decode_stream's K/V cache.
# ---------------------------------------------------------------------------------------------


class ModuleRunner:
    def __init__(self, modules):
        self.m = modules

    def reset(self):
        for m in self.m.values():
            if isinstance(m, W.DecodeStreamProgram):
                m.reset()
            elif hasattr(m, "named_buffers"):
                for n, b in m.named_buffers():
                    if "cache" in n:
                        b.zero_()

    def __call__(self, name, inputs):
        with torch.no_grad():
            out = self.m[name](
                *[torch.from_numpy(np.ascontiguousarray(x)) for x in inputs]
            )
        out = out if isinstance(out, (tuple, list)) else (out,)
        return [o.detach().numpy() for o in out]


class PteRunner:
    def __init__(self, ptes):
        from executorch.extension.pybindings.portable_lib import (
            _load_for_executorch_from_buffer,
        )

        self._load = _load_for_executorch_from_buffer
        self.ptes = ptes
        self.m = {k: self._load(v) for k, v in ptes.items()}

    def reset(self):
        if "decode_stream" in self.ptes:  # fresh method = zeroed mutable K/V buffers
            self.m["decode_stream"] = self._load(self.ptes["decode_stream"])

    def __call__(self, name, inputs):
        out = self.m[name].forward(
            [torch.from_numpy(np.ascontiguousarray(x)) for x in inputs]
        )
        return [
            o.detach().numpy() if isinstance(o, torch.Tensor) else np.asarray(o)
            for o in out
        ]


def encode_window(run, pcm):
    power = run("mel", [W.build_frames(pcm)[None]])[0]
    log_mel = W.host_log_mel(power)
    enc = run("encoder", [log_mel])[0]
    return log_mel, enc, run("cross_kv", [enc])


def greedy_stream(run, ckv, tables, max_dec=W.MAX_DEC):
    """whisper.js #window() with the streaming decoder: one token per call, argmax over 51865."""
    run.reset()
    ids = list(W.PREFIX_EN)
    for p in range(max_dec):
        e = tables.embed(ids[p], p)[None, None]
        h = run("decode_stream", [e, np.array([p], dtype=np.int64), *ckv])[0]
        if p < len(W.PREFIX_EN) - 1:
            continue
        nxt = int(np.argmax(run("lmhead", [h])[0].reshape(-1)[: W.VOCAB]))
        if nxt == W.TOK["eot"]:
            break
        ids.append(nxt)
        if len(ids) >= max_dec:
            break
    return ids[len(W.PREFIX_EN) :]


def greedy_split(run, ckv, tables, max_dec=W.MAX_DEC):
    """whisper.js #window() with the full-recompute decoder_split (no positions on the host)."""
    ids = list(W.PREFIX_EN)
    for p in range(len(W.PREFIX_EN) - 1, max_dec):
        te = np.zeros((1, W.MAX_DEC, W.D_MODEL), dtype=np.float32)
        te[0, : p + 1] = tables.embed(np.array(ids[: p + 1]))
        h = run("decoder_split", [te, *ckv])[0]
        nxt = int(np.argmax(run("lmhead", [h[:, p : p + 1]])[0].reshape(-1)[: W.VOCAB]))
        if nxt == W.TOK["eot"]:
            break
        ids.append(nxt)
        if len(ids) >= max_dec:
            break
    return ids[len(W.PREFIX_EN) :]


def hf_generate(model, feats, max_new):
    with torch.no_grad():
        seq = model.generate(
            feats,
            language="en",
            task="transcribe",
            return_timestamps=False,
            do_sample=False,
            num_beams=1,
            max_new_tokens=max_new,
        )[0].tolist()
    # transformers 5 returns only the generated tokens; older versions also return the prompt.
    out = seq[len(W.PREFIX_EN) :] if seq[: len(W.PREFIX_EN)] == W.PREFIX_EN else seq
    return out[: out.index(W.TOK["eot"])] if W.TOK["eot"] in out else out


def hf_greedy(model, feats, max_dec=W.MAX_DEC):
    """Pure argmax with the forced prefix and no logit processors (the playground's rule)."""
    with torch.no_grad():
        enc = model.model.encoder(feats)
        ids = list(W.PREFIX_EN)
        while len(ids) < max_dec:
            logits = model(
                encoder_outputs=enc, decoder_input_ids=torch.tensor([ids])
            ).logits
            nxt = int(logits[0, -1].argmax())
            if nxt == W.TOK["eot"]:
                break
            ids.append(nxt)
    return ids[len(W.PREFIX_EN) :]


def first_diff(a, b):
    for i, (x, y) in enumerate(zip(a, b)):
        if x != y:
            return i
    return None if len(a) == len(b) else min(len(a), len(b))


def summarize(rs):
    if len(rs) == 1:
        return rs[0]
    return {
        "outputs": len(rs),
        "worst_max_abs": max(d["max_abs"] for d in rs),
        "worst_max_rel_to_absmax": max(d["max_rel_to_absmax"] for d in rs),
        "min_cosine": min(d["cosine"] for d in rs),
    }


def cmp(a, b):
    return summarize([W.isclose_report(x, y) for x, y in zip(a, b)])


# ---------------------------------------------------------------------------------------------


def validate(args, out, model, fe, tok, progs, eps, ptes, tables):
    t0 = time.time()
    clips = W.load_librispeech_dummy(fetch_sample(), args.max_sec, args.clips)
    report = {
        "precision": args.precision,
        "cpu_runtime": "executorch pybindings (portable_lib) from the pip wheel, with "
        "executorch.extension.llm.custom_ops loaded (llama::update_cache / custom_sdpa kernels)",
        "webgpu_pte_on_cpu": "not runnable: the pip wheel's pybindings have no WebGPU/Vulkan backend, "
        "so the lowered GPU kernels themselves are not exercised here",
        "portable_pte_note": "portable .pte = the same function exported without a delegate, with "
        "attention as matmul+softmax (et_vk.sdpa has no portable kernel) and, for q4, the fp32 "
        "dequantization of the same int4 weights (torchao ops have no portable kernels). It checks graph "
        "structure, the host contract and the quantized numerics on the ExecuTorch runtime",
        "sample": {
            **SAMPLE,
            "clips_used": [c["id"] for c in clips],
            "max_sec": args.max_sec,
        },
        "programs": {},
    }
    pcm0 = clips[0]["pcm"]
    eager = ModuleRunner({k: v[0] for k, v in progs.items()})
    exported = ModuleRunner({k: e.module() for k, e in eps.items()})
    pte = PteRunner(ptes)
    log_mel0, enc0, ckv0 = encode_window(eager, pcm0)
    dec_in = np.zeros((1, W.MAX_DEC, W.D_MODEL), np.float32)
    dec_in[0, :4] = tables.embed(np.array(W.PREFIX_EN))
    speech = {
        "mel": [W.build_frames(pcm0)[None]],
        "encoder": [log_mel0],
        "cross_kv": [enc0],
        "decoder_split": [dec_in, *ckv0],
    }
    for name in NAMES:
        mod, _, ref32 = progs[name]
        res = {}
        if name != "decode_stream":
            cases = {"random": [x.numpy() for x in example_inputs(name)]}
            if name in speech:
                cases["speech"] = speech[name]
            if name == "lmhead":
                cases["speech"] = [
                    eager("decoder_split", speech["decoder_split"])[0][:, 3:4]
                ]
            for label, inp in cases.items():
                e = eager(name, inp)
                r = {
                    "exported_vs_eager": cmp(exported(name, inp), e),
                    "portable_pte_vs_eager": cmp(pte(name, inp), e),
                }
                if ref32 is not None:
                    r["quantized_eager_vs_fp32_eager"] = cmp(
                        e, ModuleRunner({name: ref32})(name, inp)
                    )
                res[label] = r
        else:
            # 12 steps from a fresh cache: the 4 prefix tokens, then a fixed continuation.
            seq = W.PREFIX_EN + [440, 1029, 291, 4532, 13, 50, 2034, 293]
            runners = {
                "eager": ModuleRunner({name: mod}),
                "exported": exported,
                "portable_pte": PteRunner({name: ptes[name]}),
            }
            if ref32 is not None:
                runners["fp32_eager"] = ModuleRunner({name: ref32})
            outs = {}
            for k, rr in runners.items():
                rr.reset()
                hs = [
                    rr(
                        name,
                        [
                            tables.embed(t_, p_)[None, None],
                            np.array([p_], dtype=np.int64),
                            *ckv0,
                        ],
                    )[0]
                    for p_, t_ in enumerate(seq)
                ]
                outs[k] = np.concatenate(hs, 1)
            r = {
                "exported_vs_eager": W.isclose_report(outs["exported"], outs["eager"]),
                "portable_pte_vs_eager": W.isclose_report(
                    outs["portable_pte"], outs["eager"]
                ),
            }
            if ref32 is not None:
                r["quantized_eager_vs_fp32_eager"] = W.isclose_report(
                    outs["eager"], outs["fp32_eager"]
                )
            te = np.zeros((1, W.MAX_DEC, W.D_MODEL), np.float32)
            te[0, : len(seq)] = tables.embed(np.array(seq))
            hsplit = eager("decoder_split", [te, *ckv0])[0]
            r["stream_vs_decoder_split_eager"] = W.isclose_report(
                outs["eager"], hsplit[:, : len(seq)]
            )
            res["speech_12_steps"] = r
        report["programs"][name] = res
        print(f"[validate] {name}: {json.dumps(res)[:300]}", flush=True)

    e32 = model.model.decoder.embed_tokens.weight.detach().numpy()
    report["host_tables"] = {
        "kind": tables.kind,
        "embed_vs_checkpoint": W.isclose_report(tables.embed(np.arange(W.VOCAB)), e32),
        "pos_embed_equals_checkpoint": bool(
            np.array_equal(
                tables.pos, model.model.decoder.embed_positions.weight.detach().numpy()
            )
        ),
    }

    # End-to-end greedy transcripts.
    keys = ["exported_stream", "exported_split"]
    e2e = []
    for ci, c in enumerate(clips):
        t1 = time.time()
        pcm = c["pcm"]
        feats = fe(pcm, sampling_rate=W.SAMPLE_RATE, return_tensors="pt").input_features
        hf_gen = hf_generate(model, feats, W.MAX_DEC - len(W.PREFIX_EN))
        hf_arg = hf_greedy(model, feats)
        log_mel_x, _, ckv_x = encode_window(exported, pcm)
        row = {
            "id": c["id"],
            "seconds": W.seconds(pcm),
            "reference_text": c["text"],
            "log_mel_vs_hf_features": W.isclose_report(log_mel_x, feats.numpy()),
            "hf_generate": hf_gen,
            "hf_greedy_argmax": hf_arg,
            "exported_stream": greedy_stream(exported, ckv_x, tables),
            "exported_split": greedy_split(exported, ckv_x, tables),
        }
        if ci < args.portable_clips:
            _, _, ckv_p = encode_window(pte, pcm)
            row["portable_pte_stream"] = greedy_stream(pte, ckv_p, tables)
            row["portable_pte_split"] = greedy_split(pte, ckv_p, tables)
        for k in [k for k in row if k.startswith(("exported_", "portable_pte_"))]:
            row[f"{k}_matches_hf_generate"] = row[k] == hf_gen
            row[f"{k}_matches_hf_greedy_argmax"] = row[k] == hf_arg
            if row[k] != hf_gen:
                row[f"{k}_first_diff_vs_hf_generate"] = first_diff(row[k], hf_gen)
        row["text_exported_stream"] = tok.decode(
            row["exported_stream"], skip_special_tokens=True
        ).strip()
        row["text_hf_generate"] = tok.decode(hf_gen, skip_special_tokens=True).strip()
        row["seconds_to_validate"] = round(time.time() - t1, 1)
        e2e.append(row)
        flags = {
            k: row[f"{k}_matches_hf_generate"]
            for k in row
            if f"{k}_matches_hf_generate" in row
        }
        print(
            f"[validate] {c['id']} ({row['seconds']} s) {flags} | {row['text_exported_stream']}",
            flush=True,
        )

    port = [r for r in e2e if "portable_pte_stream" in r]
    summary = {
        "clips": len(e2e),
        "audio_seconds": round(sum(r["seconds"] for r in e2e), 2),
        "hf_tokens": sum(len(r["hf_generate"]) for r in e2e),
        "hf_generate_equals_hf_greedy_argmax": sum(
            r["hf_generate"] == r["hf_greedy_argmax"] for r in e2e
        ),
    }
    for k in keys:
        summary[f"{k}_token_exact_vs_hf_generate"] = sum(
            r[f"{k}_matches_hf_generate"] for r in e2e
        )
        summary[f"{k}_token_exact_vs_hf_greedy_argmax"] = sum(
            r[f"{k}_matches_hf_greedy_argmax"] for r in e2e
        )
    summary["portable_pte_clips"] = len(port)
    for k in ["portable_pte_stream", "portable_pte_split"]:
        summary[f"{k}_token_exact_vs_hf_generate"] = sum(
            r[f"{k}_matches_hf_generate"] for r in port
        )
        summary[f"{k}_equals_exported_stream"] = sum(
            r[k] == r["exported_stream"] for r in port
        )
    summary["exported_stream_equals_exported_split"] = sum(
        r["exported_stream"] == r["exported_split"] for r in e2e
    )
    wer = {"hf_greedy_argmax": [0, 0], "exported_stream": [0, 0]}
    for r in e2e:
        for k in wer:
            e_, n_ = W.word_errors(tok, tok.decode(r[k], skip_special_tokens=True), r["reference_text"])
            wer[k][0] += e_
            wer[k][1] += n_
    summary["wer_vs_librispeech_text"] = {k: round(v[0] / max(1, v[1]), 4) for k, v in wer.items()}
    summary["word_errors"] = {k: v for k, v in wer.items()}
    summary["wer_note"] = "WhisperTokenizer.normalize (English normalizer) on both sides, word-level edit distance"
    summary["worst_log_mel_vs_hf_features_max_abs"] = max(
        r["log_mel_vs_hf_features"]["max_abs"] for r in e2e
    )
    summary["decode_budget"] = (
        f"{W.MAX_DEC} positions including the 4-token prefix (playground maxDec)"
    )
    summary["pipeline"] = (
        "build_frames -> mel -> host_log_mel -> encoder -> cross_kv -> [decode_stream | decoder_split] -> "
        f"lmhead -> argmax over the first {W.VOCAB} logits; token embeddings from the host table "
        "(+ whisper_pos_embed.f32.bin rows for decode_stream)"
    )
    summary["keys"] = {
        "hf_generate": "WhisperForConditionalGeneration.generate(language='en', task='transcribe', "
        "return_timestamps=False, greedy, max_new_tokens=60) with the checkpoint's generation_config "
        "(applies suppress_tokens and begin_suppress_tokens), input_features from WhisperFeatureExtractor",
        "hf_greedy_argmax": "same model and features, forced prefix <|startoftranscript|><|en|>"
        "<|transcribe|><|notimestamps|>, plain argmax, no logit processors (the playground's rule)",
        "exported_*": "the torch.export programs that were lowered for WebGPU (ExportedProgram.module(), "
        "PyTorch CPU kernels), driven by the host pipeline exactly as whisper.js does",
        "portable_pte_*": "the portable .pte programs on the ExecuTorch CPU runtime (first "
        f"{args.portable_clips} clips; slow single-threaded kernels)",
    }
    report["end_to_end_summary"] = summary
    report["end_to_end"] = e2e
    report["validation_seconds"] = round(time.time() - t0, 1)
    (out / "validation.json").write_text(json.dumps(report, indent=2) + "\n")
    print(json.dumps(summary, indent=2))


def main():
    ap = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    ap.add_argument("--out", required=True, help="output directory")
    ap.add_argument("--precision", choices=list(PRECISIONS), default="q4")
    ap.add_argument(
        "--no-validate", action="store_true", help="export only (no CPU validation)"
    )
    ap.add_argument(
        "--clips",
        type=int,
        default=70,
        help="LibriSpeech clips for the transcript check",
    )
    ap.add_argument(
        "--portable-clips",
        type=int,
        default=2,
        help="of those, how many also run as portable .pte",
    )
    ap.add_argument(
        "--max-sec",
        type=float,
        default=15.0,
        help="longest clip (the playground's window is <= 15 s)",
    )
    ap.add_argument(
        "--encoder-attn",
        choices=["eager", "et_vk"],
        default="eager",
        help="encoder attention: matmul+softmax (default; planned activation memory) or the fused "
        "et_vk.sdpa op as in the previous export (public runtime: 2 x 72 MB scratch per layer)",
    )
    ap.add_argument(
        "--keep-portable",
        action="store_true",
        help="also write the CPU-only portable .pte files",
    )
    ap.add_argument(
        "--runtime-ops",
        help="optional text file, one registered WebGPU runtime op per line; every delegated op is checked against it",
    )
    args = ap.parse_args()
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    torch.manual_seed(0)

    from transformers import (
        WhisperFeatureExtractor,
        WhisperForConditionalGeneration,
        WhisperTokenizer,
    )

    snap, digests = fetch_upstream()
    model = WhisperForConditionalGeneration.from_pretrained(
        snap, dtype=torch.float32, attn_implementation="eager"
    ).eval()
    fe = WhisperFeatureExtractor.from_pretrained(snap)
    tables = W.HostTables(model, PRECISIONS[args.precision][2])
    progs, quant_info = build_programs(model, fe, args.precision, args.encoder_attn, tables, args.max_sec)
    runtime_ops = None
    if args.runtime_ops:
        runtime_ops = {
            l.strip()
            for l in Path(args.runtime_ops).read_text().splitlines()
            if l.strip()
        }

    eps, ptes, files, graphs, checks, io_contract = {}, {}, {}, {}, {}, {}
    for name in NAMES:
        t = time.time()
        wg_mod, port_mod, _ = progs[name]
        inputs = example_inputs(name)
        ep = export_one(wg_mod, inputs)
        prog = _common.lower_webgpu(ep)
        path = out / f"whisper_{name}.pte"
        path.unlink(missing_ok=True)  # never write through a hard link shared with another variant
        path.write_bytes(prog.buffer)
        eps[name] = ep
        graphs[name] = _common.delegate_summary(prog.buffer)
        checks[name] = W.vk_checks(prog.buffer, runtime_ops)
        io_contract[name] = W.program_io(prog.buffer)
        files[path.name] = {
            "bytes": path.stat().st_size,
            "sha256": _common.sha256_file(path),
        }
        if not args.no_validate or args.keep_portable:
            ptes[name] = _common.lower_portable(export_one(port_mod, inputs)).buffer
            if args.keep_portable:
                (out / "portable").mkdir(exist_ok=True)
                (out / "portable" / f"whisper_{name}_portable.pte").write_bytes(
                    ptes[name]
                )
        del prog
        print(
            f"[export] {path.name}: {files[path.name]['bytes']} B, {time.time() - t:.1f} s, "
            f"checks ok={checks[name]['ok']}, {json.dumps(graphs[name])[:240]}",
            flush=True,
        )

    for fname, blob in tables.files().items():
        (out / fname).unlink(missing_ok=True)
        (out / fname).write_bytes(blob)
    (out / "vocab.json").unlink(missing_ok=True)
    shutil.copyfile(snap / "vocab.json", out / "vocab.json")
    for fname in [*tables.files(), "vocab.json"]:
        files[fname] = {
            "bytes": (out / fname).stat().st_size,
            "sha256": _common.sha256_file(out / fname),
        }

    dec_kind, lm_kind, emb_kind = PRECISIONS[args.precision]
    lm_desc = {
        "int4": {"config": W.Q4_DESC, "lowered_to": "et_vk.linear_q4gsw (group 64)"},
        "fp32": {"config": "none", "lowered_to": "aten.linear (2-D fp32)"},
        "int8 static-act": {"config": W.Q8_HEAD_DESC, "lowered_to": "quantized_decomposed.quantize_per_tensor + et_vk.linear_q8ta_q8csw", **quant_info},
    }[lm_kind]
    recipe = {
        "model": "whisper-base",
        "precision": args.precision,
        "upstream": {
            "model": UPSTREAM["repo"],
            "revision": UPSTREAM["revision"],
            "official": True,
            "license": "apache-2.0 (model card)",
            "files_sha256": digests,
            "loader": "transformers.WhisperForConditionalGeneration.from_pretrained(dtype=float32, attn_implementation='eager')",
        },
        "versions": _common.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; delegate id VulkanBackend)",
            "model_code": "recipes/_whisper.py wrappers around the transformers modules (weights shared, forward re-expressed): "
            "MelProgram, EncoderProgram, CrossKVProgram, DecoderSplitProgram, DecodeStreamProgram, LMHeadProgram",
            "quantization": {
                "decoder_side_linears": {
                    "programs": ["cross_kv", "decoder_split", "decode_stream"],
                    "config": W.Q4_DESC if dec_kind == "int4" else "none (fp32)",
                    "lowered_to": "et_vk.linear_q4gsw (group 64)" if dec_kind == "int4" else "aten.linear (2-D fp32)",
                },
                "lmhead": lm_desc,
                "fp32_programs": ["mel", "encoder"],
                "embedding_table": "int8 rowwise (whisper_embed.i8.bin + whisper_embed.scale.f32.bin)"
                if emb_kind == "int8"
                else "fp32 (whisper_embed.f32.bin)",
            },
            "encoder_attention": args.encoder_attn,
            "decoder_attention": "et_vk.sdpa (decoder_split self-attention with a [1,8,64,64] additive causal mask "
            "of 0 / -1e30; cross-attention in decoder_split and decode_stream); decode_stream self-attention = "
            "llama.update_cache + llama.custom_sdpa, fused by the Vulkan partitioner into sdpa_with_kv_cache + "
            "et_vk.select_as_symint",
            "lmhead_rows": f"{W.VOCAB_PADDED}: vocab {W.VOCAB} padded with zero weights and bias {W.PAD_LOGIT} "
            "(executorch 1.5.1's q4gsw lowering fails for N % 4 != 0; kept for fp32 too so the contract is identical)",
            "constants": {
                "max_decoder_rows": W.MAX_DEC,
                "self_attention_cache": [1, W.KV_LEN, W.N_HEADS, W.HEAD_DIM],
                "frames": [W.FRAMES, W.N_FFT],
                "hop": W.HOP,
                "prefix_en": W.PREFIX_EN,
            },
            "host_contract": {
                "frames": "whisper.js buildFrames: zero-pad pcm to 480000 samples, reflect-pad 200 each side, 3000 x 400 windows at hop 160",
                "log_mel": "between mel and encoder: x = log10(max(p, 1e-10)); x = max(x, max(x) - 8); x = (x + 4) / 4; "
                "transpose [1,3000,80] -> [1,80,3000]",
                "embeddings": ("i8[id] * scale[id]" if emb_kind == "int8" else "f32[id]")
                + "; decoder_split: rows 0..p = token embeddings (no positions), zero rows after; "
                "decode_stream: token embedding + pos_embed[p]",
                "argmax": f"over logits[..., :{W.VOCAB}] (rows {W.VOCAB}..{W.VOCAB_PADDED - 1} are {W.PAD_LOGIT})",
                "cross_kv": "computed once per 30 s window, then passed as 12 inputs to every decoder call",
            },
            "io": io_contract,
        },
        "files": files,
        "graph": graphs,
        "webgpu_runtime_checks": {
            "op_list": "checked against --runtime-ops"
            if runtime_ops
            else "not checked (no --runtime-ops given)",
            **checks,
        },
        "command": _common.command_line(out),
    }
    (out / "recipe.json").write_text(json.dumps(recipe, indent=2) + "\n")
    print(json.dumps({"files": files}, indent=2), flush=True)

    if not args.no_validate:
        tok = WhisperTokenizer.from_pretrained(snap)
        validate(args, out, model, fe, tok, progs, eps, ptes, tables)


if __name__ == "__main__":
    main()