Add ExecuTorch WebGPU program for phonon-31m-en

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+phonon_flow_lm.pte filter=lfs diff=lfs merge=lfs -text
+phonon_mimi.pte filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

NOTICE

@@ -0,0 +1,6 @@
+This repository contains a modified version of gradium/phonon-31m-en.
+
+Modifications by Meta Platforms, Inc.: conversion to the ExecuTorch program format (.pte) for the
+ExecuTorch WebGPU backend.
+
+The original work and these modifications are distributed under the license in LICENSE.

README.md

@@ -0,0 +1,93 @@
+---
+license: apache-2.0
+base_model:
+- gradium/phonon-31m-en
+library_name: executorch
+pipeline_tag: text-to-speech
+inference: false
+tags:
+- executorch
+- webgpu
+- on-device
+- browser
+---
+# Phonon 31M (English) for ExecuTorch WebGPU
+
+This repository contains an [ExecuTorch](https://github.com/pytorch/executorch) conversion of
+[gradium/phonon-31m-en](https://huggingface.co/gradium/phonon-31m-en) 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](https://executorch.ai/playground/).
+
+- **Base model:** [gradium/phonon-31m-en](https://huggingface.co/gradium/phonon-31m-en)
+- **Format:** ExecuTorch `.pte` programs delegated to the WebGPU backend
+- **Quantization:** fp32 programs; the text embedding table is stored as bf16 (its values are exactly representable in bf16)
+- **Total size:** 150.7 MB
+
+> **Note:** The voices are from gradium/phonon-voices (Apache-2.0); each is stored here as its 125-frame voice prompt projected into the model.
+
+## About the original model
+
+An English text-to-speech model that streams natural 24 kHz speech while it is still generating. A flow-matching transformer predicts an audio latent for each 80 ms frame from the text and a voice prompt, and the Mimi codec's decoder turns the latents into audio.
+
+- **Developed by:** Gradium (builds on Pocket TTS and the Mimi codec by Kyutai)
+- 31M parameters (excluding the audio tokenizer)
+- English
+- Mono 24 kHz output, one 80 ms frame per step
+
+For training data, evaluation results and full details, see the original model card: [gradium/phonon-31m-en](https://huggingface.co/gradium/phonon-31m-en).
+
+## Intended use and limitations
+
+On-device speech for assistants, screen readers and other accessibility tools, games and interactive apps. Do not use it to impersonate a real person, to deceive listeners about who is speaking, or to clone a voice without the speaker's consent; tell listeners when speech is synthetic where that matters to them.
+
+English only: text in other languages is read with the wrong pronunciation. The official runtime normalizes numbers, currency, dates, emails and so on before synthesis; that step is not part of the model and must be done by the application. Long inputs are spoken sentence by sentence, so prosody does not carry across the whole text.
+
+## What changed in this conversion
+
+- Exported as two ExecuTorch programs: the flow-matching transformer with its 1-step sampler and end-of-speech score, and the streaming Mimi decoder
+- A few operations are expressed with equivalent ones the WebGPU backend supports (layer norm, mean, 1-D transposed convolution, ELU, SiLU); each was checked against the original model
+- Weights stay fp32; the text embedding table is stored as bf16, which is exact for these weights
+- Context capped at 512 positions (voice prompt, text and audio frames)
+
+## Inputs and outputs
+
+- `phonon_flow_lm.pte`: embeddings for 1 to 128 positions plus the start position and noise -> the next audio latent and the end-of-speech score; the attention cache lives inside the program
+- `phonon_mimi.pte`: one latent plus 15 streaming state tensors -> 1,920 samples at 24 kHz and the updated state
+- `phonon.json` documents the full contract: text embedding (`phonon_text_embed.bf16.bin`), the first-step embedding, the voices (`voice_*.f32.bin`), noise, the end-of-speech rule and the frame budget; tokenize with `tokenizer.json`
+
+## Files
+
+| File | Size | SHA-256 |
+|---|---|---|
+| `phonon.json` | 8.4 KB | `e8fa6604769aeb3b…` |
+| `phonon_bos_embed.f32.bin` | 2.0 KB | `a4464c50596adea2…` |
+| `phonon_flow_lm.pte` | 106.7 MB | `178698a6c4539ba4…` |
+| `phonon_mimi.pte` | 39.4 MB | `d7acdf2d93ba4897…` |
+| `phonon_text_embed.bf16.bin` | 3.9 MB | `7c6f3394fc86b55e…` |
+| `tokenizer.json` | 271.5 KB | `691a00068fa7b989…` |
+| `voice_default.f32.bin` | 250.0 KB | `d8a8ffc97e46c033…` |
+| `voice_en_us_easton.f32.bin` | 250.0 KB | `1d6ec67cbd3a45b2…` |
+
+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 run in a browser (headless Chromium, WebGPU on SwiftShader) with the WebAssembly runtime build used by the [ExecuTorch WebGPU Playground](https://executorch.ai/playground/), where the full text-to-speech loop matched the CPU programs frame for frame (audio cosine 0.9999997). That runtime build is published in
+[executorch-community/ExecuTorch-WebGPU-Runtime](https://huggingface.co/executorch-community/ExecuTorch-WebGPU-Runtime); the playground is the easiest way to run them. It downloads these files once, caches them in the browser,
+and runs everything locally.
+
+Requirements: a browser with WebGPU enabled.
+
+## Provenance
+
+- **Upstream:** `gradium/phonon-31m-en` at revision `52433c46975bfa13144c56204c9f4ee1c5b148aa` (official)
+- **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/phonon-31m-en.py --out out/phonon-31m-en` (recipe in `export/`)
+- Re-running the recipe produces byte-identical files.
+
+
+## License
+
+This is a derivative of the upstream model and is distributed under the same terms
+(`apache-2.0`); see `LICENSE`.
+Modifications (conversion to the ExecuTorch format) by Meta Platforms, Inc.

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/_phonon.py

@@ -0,0 +1,910 @@
+"""Phonon (gradium/phonon-31m-en) building blocks for recipes/phonon-31m-en.py: weights, eager reference, export programs.
+
+The eager reference follows the official runtime (the `ptts` crate, src/{flow_lm,tts_model,
+transformer,mlp,conditioners,mimi,seanet,conv}.rs) operation by operation:
+
+  voice   speaker_wavs [1,512,125] -> transpose -> speaker_proj (512x512, no bias) -> [1,125,512]
+  prime   transformer(voice) then transformer(text_embed(tokens)), both written to the KV cache
+  step    x = input_linear(prev latent or bos_emb) + cond      (cond = num_speakers("1") LUT
+          + duration_delta(0.0) sinusoid, each through its output_proj)
+          h = out_norm(transformer(x))[:, -1]; eos_logit = out_eos(h)
+          latent = noise + flow_net(h, s=0, t=1, noise)          (LSD, 1 step; noise ~ N(0, temp))
+  decode  (latent * emb_std + emb_mean) -> 1x1 conv 32->512 -> depthwise convtr x16 -> 2-layer
+          transformer (window 250, layer scale) -> SEANet decoder (ratios 6,5,4) -> 1920 samples
+
+Export programs (WebGPU):
+  FlowLMProgram   forward(embeds[1,S,512] S<=128, input_pos[1] i64, noise[1,32])
+                  -> latent[1,32], eos_logit[1,1], next_embeds[1,1,512]  (KV cache in the program)
+  MimiStepProgram forward(latent[1,1,32], *state) -> pcm[1,1,1920], *new_state  (state explicit)
+
+Rewrites for the public WebGPU runtime (same function, checked against the eager reference):
+  * RoPE: q/k rows of in_proj permuted from interleaved pairs to halves, so the runtime's fused
+    HF RoPE (et_vk.apply_rotary_emb_hf) applies; q.k is unchanged.
+  * flow-LM LayerNorm (dynamic S): mean as a 1/512 linear + the rms_norm pattern, because
+    native_layer_norm has no resize hook and mean.dim's resize drops a keepdim'ed dim.
+  * ELU = relu(x) + exp(min(x, 0)) - 1 and SiLU = x * sigmoid(x) (no aten.elu / aten.silu).
+  * ConvTranspose1d as a 4-D (H = 1) conv_transpose2d (3-D convolution has no transposed path).
+  * Mimi streaming state as explicit tensors: mutable buffers other than the llama KV cache are
+    not kept live by the delegate.
+"""
+
+import math
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+D = 512
+H = 8
+HD = 64
+LDIM = 32
+FLOW_LAYERS = 6
+MIMI_LAYERS = 2
+MIMI_CONTEXT = 250
+FRAME = 1920  # samples per latent frame at 24 kHz
+UPS = 16  # mimi transformer tokens per latent frame
+NEG = -1.0e30
+
+
+# ---------------------------------------------------------------------------------------------
+# Weights: a module whose state_dict keys are exactly the checkpoint's keys.
+# ---------------------------------------------------------------------------------------------
+
+
+def _p(*shape):
+    return nn.Parameter(torch.zeros(*shape), requires_grad=False)
+
+
+class _Lin(nn.Module):
+    def __init__(self, i, o, bias):
+        super().__init__()
+        self.weight = _p(o, i)
+        if bias:
+            self.bias = _p(o)
+
+
+class _LN(nn.Module):
+    def __init__(self, n):
+        super().__init__()
+        self.weight = _p(n)
+        self.bias = _p(n)
+
+
+class _Attn(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.in_proj = _Lin(D, 3 * D, False)
+        self.out_proj = _Lin(D, D, False)
+
+
+class _Scale(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.scale = _p(D)
+
+
+class _Layer(nn.Module):
+    def __init__(self, layer_scale):
+        super().__init__()
+        self.self_attn = _Attn()
+        self.norm1 = _LN(D)
+        self.norm2 = _LN(D)
+        self.linear1 = _Lin(D, 2048, False)
+        self.linear2 = _Lin(2048, D, False)
+        if layer_scale:
+            self.layer_scale_1 = _Scale()
+            self.layer_scale_2 = _Scale()
+
+
+class _Transformer(nn.Module):
+    def __init__(self, n, layer_scale):
+        super().__init__()
+        self.layers = nn.ModuleList(_Layer(layer_scale) for _ in range(n))
+
+
+class _Seq(nn.Module):
+    """nn.Sequential-style numbered children with only the indices that hold weights."""
+
+    def __init__(self, items):
+        super().__init__()
+        for k, m in items.items():
+            self.add_module(str(k), m)
+
+    def __getitem__(self, i):
+        return getattr(self, str(i))
+
+
+class _TimeEmbed(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.freqs = _p(128)
+        self.mlp = _Seq({0: _Lin(256, D, True), 2: _Lin(D, D, True), 3: _Alpha()})
+
+
+class _Alpha(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.alpha = _p(D)
+
+
+class _ResBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.in_ln = _LN(D)
+        self.mlp = _Seq({0: _Lin(D, D, True), 2: _Lin(D, D, True)})
+        self.adaLN_modulation = _Seq({1: _Lin(D, 3 * D, True)})
+
+
+class _Final(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear = _Lin(D, LDIM, True)
+        self.adaLN_modulation = _Seq({1: _Lin(D, 2 * D, True)})
+
+
+class _FlowNet(nn.Module):
+    def __init__(self, depth=6):
+        super().__init__()
+        self.time_embed = nn.ModuleList([_TimeEmbed(), _TimeEmbed()])
+        self.cond_embed = _Lin(D, D, True)
+        self.input_proj = _Lin(LDIM, D, True)
+        self.res_blocks = nn.ModuleList(_ResBlock() for _ in range(depth))
+        self.final_layer = _Final()
+
+
+class _Embed(nn.Module):
+    def __init__(self, n, d):
+        super().__init__()
+        self.weight = _p(n, d)
+
+
+class _Cond(nn.Module):
+    def __init__(self, embed=None, proj_in=None):
+        super().__init__()
+        self.learnt_padding = _p(1, 1, D)
+        if embed is not None:
+            self.embed = _Embed(*embed)
+        if proj_in is not None:
+            self.output_proj = _Lin(proj_in, D, False)
+
+
+class _Conditioners(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.duration_delta = _Cond(proj_in=128)
+        self.num_speakers = _Cond(embed=(32, 16), proj_in=16)
+        self.speaker_wavs = _Cond(proj_in=512)
+        self.transcript_in_segment = _Cond(embed=(4001, D))
+
+
+class _CondProvider(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conditioners = _Conditioners()
+
+
+class _WST(nn.Module):
+    """flow.w_s_t: a training-time loss-weighting net; the runtime drops it (loader.rs remap_key)."""
+
+    def __init__(self):
+        super().__init__()
+        self.w_s_t = _Seq(
+            {
+                0: _Lin(2, 32, True),
+                3: _Lin(32, 32, True),
+                6: _Lin(32, 32, True),
+                9: _Lin(32, 1, True),
+            }
+        )
+
+
+class FlowLMWeights(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.bos_emb = _p(LDIM)
+        self.emb_mean = _p(LDIM)
+        self.emb_std = _p(LDIM)
+        self.condition_provider = _CondProvider()
+        self.flow = _WST()
+        self.flow_net = _FlowNet()
+        self.input_linear = _Lin(LDIM, D, False)
+        self.out_eos = _Lin(D, 1, True)
+        self.out_norm = _LN(D)
+        self.transformer = _Transformer(FLOW_LAYERS, False)
+
+
+class _Conv(nn.Module):
+    def __init__(self, i, o, k):
+        super().__init__()
+        self.conv = _ConvW(i, o, k, bias=True)
+
+
+class _ConvW(nn.Module):
+    def __init__(self, i, o, k, bias):
+        super().__init__()
+        self.weight = _p(o, i, k)
+        if bias:
+            self.bias = _p(o)
+
+
+class _ConvTr(nn.Module):
+    def __init__(self, i, o, k):
+        super().__init__()
+        self.convtr = _ConvTrW(i, o, k, True)
+
+
+class _ConvTrW(nn.Module):
+    def __init__(self, i, o_per_group, k, bias):
+        super().__init__()
+        self.weight = _p(i, o_per_group, k)
+        if bias:
+            self.bias = _p(o_per_group)
+
+
+class _Res(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.block = _Seq({1: _Conv(dim, dim // 2, 3), 3: _Conv(dim // 2, dim, 1)})
+
+
+class _SEANetDecoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.model = _Seq(
+            {
+                0: _Conv(512, 512, 7),
+                2: _ConvTr(512, 256, 12),
+                3: _Res(256),
+                5: _ConvTr(256, 128, 10),
+                6: _Res(128),
+                8: _ConvTr(128, 64, 8),
+                9: _Res(64),
+                11: _Conv(64, 1, 3),
+            }
+        )
+
+
+class _ProjT(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.transformer = _Transformer(MIMI_LAYERS, True)
+
+
+class _Quantizer(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.output_proj = _ConvW(LDIM, D, 1, bias=False)
+
+
+class _Upsample(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.convtr = _Seq({})
+        self.convtr.convtr = _ConvTrW(D, 1, 2 * UPS, False)
+
+
+class MimiWeights(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.decoder = _SEANetDecoder()
+        self.decoder_transformer = _ProjT()
+        self.quantizer = _Quantizer()
+        self.upsample = _Upsample()
+
+
+class PhononWeights(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.flow_lm = FlowLMWeights()
+        self.mimi = MimiWeights()
+
+
+def load_weights(safetensors_path):
+    """Every checkpoint tensor, as float32 (bf16 -> f32 is exact), loaded with strict=True."""
+    from safetensors.torch import load_file
+
+    sd = {k: v.float() for k, v in load_file(str(safetensors_path)).items()}
+    w = PhononWeights()
+    w.load_state_dict(sd, strict=True)
+    return w.eval()
+
+
+def load_voice(path):
+    """speaker_wavs [1, 512, T] (stored speaker-Mimi latents) from a voice .safetensors file."""
+    from safetensors.torch import load_file
+
+    t = load_file(str(path))
+    assert list(t) == ["speaker_wavs"], list(t)
+    return t["speaker_wavs"].float()
+
+
+# ---------------------------------------------------------------------------------------------
+# Shared math
+# ---------------------------------------------------------------------------------------------
+
+
+def lin(m, x):
+    """nn.Linear on a 2-D view (the public WebGPU runtime's fp32 aten.linear takes 2-D input)."""
+    shp = x.shape
+    y = F.linear(x.reshape(-1, shp[-1]), m.weight, getattr(m, "bias", None))
+    return y.reshape(*shp[:-1], y.shape[-1])
+
+
+def lin3(m, x):
+    """nn.Linear on the tensor as it is (any rank); the Vulkan lowering adds its own 2-D views."""
+    return F.linear(x, m.weight, getattr(m, "bias", None))
+
+
+def layer_norm_dyn(m, x, eps, mean_w):
+    """LayerNorm for the dynamic-length flow-LM program, from ops that follow a resized sequence
+    length on the public WebGPU runtime (its native_layer_norm has no resize hook, and its mean.dim
+    resize drops a keepdim'ed dim): the mean is a linear with weight 1/512 (exact: a power of two),
+    the rest matches the Vulkan rms_norm pattern (fused into et_vk.rms_norm)."""
+    xc = x - F.linear(x, mean_w)
+    return xc * torch.rsqrt((xc * xc).mean(-1, keepdim=True) + eps) * m.weight + m.bias
+
+
+def layer_norm(m, x, eps):
+    return F.layer_norm(x, (x.shape[-1],), m.weight, m.bias, eps)
+
+
+def silu(x):
+    return x * torch.sigmoid(x)
+
+
+def elu(x):
+    """ELU(alpha=1) from ops the public WebGPU runtime registers (it has no aten.elu)."""
+    return torch.relu(x) + torch.exp(torch.clamp(x, max=0.0)) - 1.0
+
+
+def gelu(x):
+    return F.gelu(x)  # erf
+
+
+def inv_freq():
+    """RoPE frequencies as the runtime computes them (rope.rs): 1 / max_period^(i / half) in f32."""
+    half = HD // 2
+    i = torch.arange(half, dtype=torch.float32)
+    return 1.0 / torch.pow(torch.tensor(10000.0, dtype=torch.float32), i / half)
+
+
+def rope_interleaved(x, pos):
+    """rope_i on [B, T, H, HD] with absolute positions pos [T] (the runtime's convention:
+    pairs (2i, 2i+1) rotate by pos * inv_freq[i])."""
+    ang = pos.to(torch.float32)[:, None] * inv_freq()[None, :]  # [T, half]
+    c, s = torch.cos(ang)[None, :, None, :], torch.sin(ang)[None, :, None, :]
+    xr, xi = x[..., 0::2], x[..., 1::2]
+    return torch.stack([xr * c - xi * s, xr * s + xi * c], dim=-1).flatten(-2)
+
+
+def time_embed_rms(te, t):
+    """TimestepEmbedder with TimeRmsNorm::Rms (phonon config time_rms_norm: rms)."""
+    args = t * te.freqs
+    emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    x = lin(te.mlp[2], silu(lin(te.mlp[0], emb)))
+    ms = (x * x).mean(-1, keepdim=True)
+    return x * torch.rsqrt(ms + 1e-5) * te.mlp[3].alpha
+
+
+def flow_net(fn, c, x, t_comb):
+    """SimpleMLPAdaLN: c [N, 512] backbone output, x [N, 32] noise, t_comb [N or 1, 512]."""
+    x = lin(fn.input_proj, x)
+    y = t_comb + lin(fn.cond_embed, c)
+    sy = silu(y)
+    for b in fn.res_blocks:
+        ada = lin(b.adaLN_modulation[1], sy)
+        shift, scale, gate = ada[..., :D], ada[..., D : 2 * D], ada[..., 2 * D :]
+        h = layer_norm(b.in_ln, x, 1e-6) * (scale + 1.0) + shift
+        h = lin(b.mlp[2], silu(lin(b.mlp[0], h)))
+        x = x + gate * h
+    fl = fn.final_layer
+    ada = lin(fl.adaLN_modulation[1], sy)
+    shift, scale = ada[..., :D], ada[..., D:]
+    x = F.layer_norm(x, (D,), None, None, 1e-6) * (scale + 1.0) + shift
+    return lin(fl.linear, x)
+
+
+def lsd_time(fn):
+    """(time_embed[0](s=0) + time_embed[1](t=1)) / 2, the flow head's constant time condition
+    for lsd_decode_steps = 1."""
+    s = torch.zeros(1, 1)
+    t = torch.ones(1, 1)
+    return (
+        time_embed_rms(fn.time_embed[0], s) + time_embed_rms(fn.time_embed[1], t)
+    ) / 2
+
+
+def sin_embedding(position, dim, max_period):
+    half = dim // 2
+    i = torch.arange(half, dtype=torch.float32)
+    ph = position / torch.pow(
+        torch.tensor(max_period, dtype=torch.float32), i / (half - 1)
+    )
+    return torch.cat([torch.cos(ph), torch.sin(ph)])
+
+
+def summed_conditions(flw, num_speakers="1", duration_delta=0.0):
+    """conditioners.rs Conditioners::sum with the runtime's defaults, [1, 1, 512]."""
+    c = flw.condition_provider.conditioners
+    idx = ["0", "1", "2"].index(num_speakers)
+    ns = F.linear(
+        c.num_speakers.embed.weight[idx : idx + 1], c.num_speakers.output_proj.weight
+    )
+    dd = F.linear(
+        sin_embedding(1000.0 * duration_delta, 128, 10000.0)[None],
+        c.duration_delta.output_proj.weight,
+    )
+    return (ns + dd).view(1, 1, D)
+
+
+def voice_embedding(flw, speaker_wavs):
+    """[1, 512, T] speaker latents -> [1, T, 512] voice prompt (loader.rs project_latents)."""
+    w = flw.condition_provider.conditioners.speaker_wavs.output_proj.weight
+    return F.linear(speaker_wavs.transpose(1, 2), w)
+
+
+def text_embedding(flw, tokens):
+    emb = flw.condition_provider.conditioners.transcript_in_segment.embed.weight
+    return emb[torch.as_tensor(tokens, dtype=torch.long)][None]
+
+
+# ---------------------------------------------------------------------------------------------
+# Eager reference (Python-side state, mirrors the runtime)
+# ---------------------------------------------------------------------------------------------
+
+
+class FlowLMRef:
+    def __init__(self, w: PhononWeights, num_speakers="1", duration_delta=0.0):
+        self.w = w.flow_lm
+        self.cond = summed_conditions(self.w, num_speakers, duration_delta)
+        self.t_comb = lsd_time(self.w.flow_net)
+        self.reset()
+
+    def reset(self):
+        self.k = [None] * FLOW_LAYERS
+        self.v = [None] * FLOW_LAYERS
+        self.pos = 0
+
+    @torch.no_grad()
+    def backbone(self, x):
+        B, T, _ = x.shape
+        pos = torch.arange(self.pos, self.pos + T)
+        for i, L in enumerate(self.w.transformer.layers):
+            h = layer_norm(L.norm1, x, 1e-5)
+            qkv = lin(L.self_attn.in_proj, h).view(B, T, 3, H, HD)
+            q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
+            q, k = rope_interleaved(q, pos), rope_interleaved(k, pos)
+            self.k[i] = k if self.k[i] is None else torch.cat([self.k[i], k], 1)
+            self.v[i] = v if self.v[i] is None else torch.cat([self.v[i], v], 1)
+            kk, vv = self.k[i].transpose(1, 2), self.v[i].transpose(1, 2)
+            qq = q.transpose(1, 2)
+            Tk = kk.shape[2]
+            qpos = torch.arange(Tk - T, Tk)[:, None]
+            mask = torch.where(torch.arange(Tk)[None, :] <= qpos, 0.0, float("-inf"))
+            att = torch.softmax(qq @ kk.transpose(-1, -2) / math.sqrt(HD) + mask, -1)
+            o = (att @ vv).transpose(1, 2).reshape(B, T, D)
+            x = x + lin(L.self_attn.out_proj, o)
+            h = layer_norm(L.norm2, x, 1e-5)
+            x = x + lin(L.linear2, gelu(lin(L.linear1, h)))
+        self.pos += T
+        return x
+
+    def prompt(self, embeds):
+        self.backbone(embeds)
+
+    @torch.no_grad()
+    def step(self, prev_latent, noise):
+        """prev_latent None (BOS) or [1, 1, 32]; noise [1, 32]. -> latent [1,1,32], eos_logit."""
+        seq = self.w.bos_emb.view(1, 1, LDIM) if prev_latent is None else prev_latent
+        x = lin(self.w.input_linear, seq) + self.cond
+        h = layer_norm(self.w.out_norm, self.backbone(x), 1e-5)[:, -1]
+        eos = lin(self.w.out_eos, h)
+        latent = noise + flow_net(self.w.flow_net, h, noise, self.t_comb)
+        return latent.view(1, 1, LDIM), eos
+
+
+def _conv1d_stream(m, x, prev, dilation=1):
+    xx = torch.cat([prev, x], 2) if prev is not None else x
+    y = F.conv1d(xx, m.conv.weight, m.conv.bias, dilation=dilation)
+    tp = prev.shape[2] if prev is not None else 0
+    return y, (xx[:, :, xx.shape[2] - tp :] if tp else None)
+
+
+def _convtr_stream(w, bias, x, partial, stride, groups=1):
+    y = F.conv_transpose1d(x, w, bias, stride=stride, groups=groups)
+    pt = partial.shape[2]
+    y = torch.cat([y[:, :, :pt] + partial, y[:, :, pt:]], 2)
+    new_partial = y[:, :, y.shape[2] - pt :]
+    if bias is not None:
+        new_partial = new_partial - bias.view(1, -1, 1)
+    return y[:, :, : y.shape[2] - pt], new_partial
+
+
+class MimiRef:
+    """Streaming Mimi decoder (MimiDecoder::decode_from_latent_step), one or more frames per call."""
+
+    def __init__(self, w: PhononWeights):
+        self.w = w.mimi
+        self.emb_std = w.flow_lm.emb_std
+        self.emb_mean = w.flow_lm.emb_mean
+        self.reset()
+
+    def reset(self):
+        z = lambda *s: torch.zeros(1, *s)  # noqa: E731
+        self.up = z(D, UPS)
+        self.kv = [[None, None] for _ in range(MIMI_LAYERS)]
+        self.offset = 0
+        self.prev = {
+            "init": z(512, 6),
+            "r1": z(256, 2),
+            "r2": z(128, 2),
+            "r3": z(64, 2),
+            "final": z(64, 2),
+        }
+        self.partial = {"u1": z(256, 6), "u2": z(128, 5), "u3": z(64, 4)}
+
+    @torch.no_grad()
+    def transformer(self, x):  # x [1, T, 512]
+        T = x.shape[1]
+        pos = torch.arange(self.offset, self.offset + T)
+        for i, L in enumerate(self.w.decoder_transformer.transformer.layers):
+            h = layer_norm(L.norm1, x, 1e-5)
+            qkv = lin(L.self_attn.in_proj, h).view(1, T, 3, H, HD)
+            q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
+            q, k = rope_interleaved(q, pos), rope_interleaved(k, pos)
+            q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+            pk, pv = self.kv[i]
+            kk = k if pk is None else torch.cat([pk, k], 2)
+            vv = v if pv is None else torch.cat([pv, v], 2)
+            kvlen = kk.shape[2] - T
+            qi = torch.arange(T)[:, None] + kvlen
+            kj = torch.arange(kk.shape[2])[None, :]
+            ok = (kj <= qi) & (kj >= qi - MIMI_CONTEXT)
+            mask = torch.where(ok, 0.0, float("-inf"))
+            att = torch.softmax(q @ kk.transpose(-1, -2) / math.sqrt(HD) + mask, -1)
+            o = (att @ vv).transpose(1, 2).reshape(1, T, D)
+            self.kv[i] = [kk[:, :, -MIMI_CONTEXT:], vv[:, :, -MIMI_CONTEXT:]]
+            x = x + L.layer_scale_1.scale * lin(L.self_attn.out_proj, o)
+            h = layer_norm(L.norm2, x, 1e-5)
+            x = x + L.layer_scale_2.scale * lin(L.linear2, gelu(lin(L.linear1, h)))
+        self.offset += T
+        return x
+
+    def _res(self, blk, x, key):
+        v, self.prev[key] = _conv1d_stream(blk.block[1], elu(x), self.prev[key])
+        v, _ = _conv1d_stream(blk.block[3], elu(v), None)
+        return x + v
+
+    @torch.no_grad()
+    def decode(self, latent):
+        """latent [1, T, 32] (flow-LM space) -> pcm [1, 1, T * 1920]."""
+        x = (latent * self.emb_std + self.emb_mean).transpose(1, 2)
+        x = F.conv1d(x, self.w.quantizer.output_proj.weight)
+        x, self.up = _convtr_stream(
+            self.w.upsample.convtr.convtr.weight, None, x, self.up, UPS, groups=D
+        )
+        x = self.transformer(x.transpose(1, 2)).transpose(1, 2)
+        m = self.w.decoder.model
+        x, self.prev["init"] = _conv1d_stream(m[0], x, self.prev["init"])
+        for up, res, s, ku, kr in [
+            (m[2], m[3], 6, "u1", "r1"),
+            (m[5], m[6], 5, "u2", "r2"),
+            (m[8], m[9], 4, "u3", "r3"),
+        ]:
+            x, self.partial[ku] = _convtr_stream(
+                up.convtr.weight, up.convtr.bias, elu(x), self.partial[ku], s
+            )
+            x = self._res(res, x, kr)
+        x, self.prev["final"] = _conv1d_stream(m[11], elu(x), self.prev["final"])
+        return x
+
+
+# ---------------------------------------------------------------------------------------------
+# Generation policy (plan.rs, tts_model.rs)
+# ---------------------------------------------------------------------------------------------
+
+
+def prepare_text_prompt(text):
+    text = " ".join(text.strip().replace("\n", " ").replace("\r", " ").split())
+    if not text:
+        return text, 3
+    frames_after_eos = 3 if len(text.split()) <= 4 else 1
+    text = text[0].upper() + text[1:]
+    if text[-1].isalnum():
+        text += "."
+    return text, frames_after_eos
+
+
+def frame_budget(num_tokens, frame_rate=12.5):
+    return math.ceil((num_tokens / 3.0 + 2.0) * frame_rate)
+
+
+class EosPolicy:
+    def __init__(self, frames_after_eos):
+        self.n, self.countdown = frames_after_eos, None
+
+    def should_stop(self, is_eos):
+        if is_eos and self.countdown is None:
+            self.countdown = self.n
+        if self.countdown is None:
+            return False
+        if self.countdown == 0:
+            return True
+        self.countdown -= 1
+        return False
+
+
+def generate(step_fn, tokens, frames_after_eos, noise_fn, eos_threshold=4.0):
+    """The runtime's loop (synth.rs run_inline): BOS first, the EOS frame and its tail are kept.
+    step_fn(prev_latent_or_None, noise) -> (latent [1,1,32], eos_logit). Returns latents, eos logits."""
+    eos = EosPolicy(frames_after_eos)
+    prev, lats, logits = None, [], []
+    for i in range(frame_budget(len(tokens))):
+        lat, e = step_fn(prev, noise_fn(i))
+        lats.append(lat)
+        logits.append(float(e.reshape(-1)[0]))
+        if eos.should_stop(logits[-1] > eos_threshold):
+            break
+        prev = lat
+    return lats, logits
+
+
+def tokenizer(tokenizer_json):
+    from tokenizers import Tokenizer
+
+    return Tokenizer.from_file(str(tokenizer_json))
+
+
+def tokenize(tok, text):
+    prepared, frames_after_eos = prepare_text_prompt(text)
+    return (
+        tok.encode(prepared, add_special_tokens=False).ids,
+        frames_after_eos,
+        prepared,
+    )
+
+
+# ---------------------------------------------------------------------------------------------
+# Export programs
+# ---------------------------------------------------------------------------------------------
+
+MAX_CTX = 512  # flow-LM KV cache positions: voice prompt + text tokens + generated frames
+PREFILL_MAX = 128  # most positions one flow-LM call may write (the 125-frame voice prompt fits)
+
+
+def _half_split_perm():
+    """Within each head, reorder q/k channels from interleaved pairs (2i, 2i+1) to halves
+    (i, i + 32): rope_i on the old layout equals rotate_half RoPE on the new one, and q.k is
+    unchanged because q and k are permuted alike."""
+    per = torch.cat([torch.arange(0, HD, 2), torch.arange(1, HD, 2)])
+    return torch.cat([h * HD + per for h in range(H)])
+
+
+def _permuted_in_proj(w):
+    """in_proj weight [3*512, 512] with the q and k row blocks permuted, v untouched."""
+    perm = _half_split_perm()
+    return torch.cat([w[:D][perm], w[D : 2 * D][perm], w[2 * D :]], 0).contiguous()
+
+
+def rope_tables(n):
+    """cos/sin [n, 64] in the half-split layout, angle = f32(pos) * f32(inv_freq) as rope.rs."""
+    ang = torch.arange(n, dtype=torch.float32)[:, None] * inv_freq()[None, :]
+    c, s = torch.cos(ang), torch.sin(ang)
+    return torch.cat([c, c], -1).contiguous(), torch.cat([s, s], -1).contiguous()
+
+
+def _rope_hf(q, k, cos, sin):
+    from executorch.backends.vulkan.patterns.rope_hf import HfRotaryEmbeddingPattern
+
+    return HfRotaryEmbeddingPattern()(q, k, cos, sin)
+
+
+class _LinP(nn.Module):
+    def __init__(self, weight, bias=None):
+        super().__init__()
+        self.weight = nn.Parameter(weight.detach().clone().contiguous(), requires_grad=False)
+        if bias is not None:
+            self.bias = nn.Parameter(bias.detach().clone().contiguous(), requires_grad=False)
+
+
+class FlowLMProgram(nn.Module):
+    """forward(embeds[1,S,512], input_pos[1] i64, noise[1,32])
+    -> latent[1,32], eos_logit[1,1], next_embeds[1,1,512].
+
+    Writes the S positions [input_pos, input_pos + S) of the in-program KV cache (llama
+    sdpa_with_kv_cache, causal), then runs the head on the last position: eos_logit = out_eos(h),
+    latent = noise + flow_net(h, s=0, t=1, noise) and next_embeds = input_linear(latent) + cond,
+    the embedding of the next step's input. Prompt calls (voice, text) ignore the outputs."""
+
+    def __init__(self, w: PhononWeights, num_speakers="1", duration_delta=0.0, attn="sdpa_kv"):
+        super().__init__()
+        from executorch.extension.llm.custom_ops import custom_ops  # noqa: F401
+
+        f = w.flow_lm
+        self.attn = attn
+        self.layers = nn.ModuleList()
+        for L in f.transformer.layers:
+            m = nn.Module()
+            m.norm1, m.norm2 = L.norm1, L.norm2
+            m.in_proj = _LinP(_permuted_in_proj(L.self_attn.in_proj.weight))
+            m.out_proj, m.linear1, m.linear2 = L.self_attn.out_proj, L.linear1, L.linear2
+            self.layers.append(m)
+        self.out_norm, self.out_eos, self.input_linear = f.out_norm, f.out_eos, f.input_linear
+        self.fn = f.flow_net
+        with torch.no_grad():
+            self.register_buffer("cond", summed_conditions(f, num_speakers, duration_delta).clone(), persistent=False)
+            self.register_buffer("t_comb", lsd_time(f.flow_net).clone(), persistent=False)
+        self.register_buffer("mean_w", torch.full((1, D), 1.0 / D), persistent=False)
+        cos, sin = rope_tables(MAX_CTX)
+        self.register_buffer("rope_cos", cos, persistent=False)
+        self.register_buffer("rope_sin", sin, persistent=False)
+        for i in range(FLOW_LAYERS):
+            self.register_buffer(f"k_cache_{i}", torch.zeros(1, MAX_CTX, H, HD), persistent=False)
+            self.register_buffer(f"v_cache_{i}", torch.zeros(1, MAX_CTX, H, HD), persistent=False)
+
+    def reset(self):
+        for i in range(FLOW_LAYERS):
+            getattr(self, f"k_cache_{i}").zero_()
+            getattr(self, f"v_cache_{i}").zero_()
+
+    def forward(self, embeds, input_pos, noise):
+        S = embeds.shape[1]
+        p = input_pos[0].item()
+        torch._check_is_size(p)
+        torch._check(p + S <= MAX_CTX)
+        cos = self.rope_cos.narrow(0, p, S)
+        sin = self.rope_sin.narrow(0, p, S)
+        x = embeds
+        for i, L in enumerate(self.layers):
+            h = layer_norm_dyn(L.norm1, x, 1e-5, self.mean_w)
+            qkv = lin3(L.in_proj, h)
+            q = qkv[..., :D].reshape(1, S, H, HD)
+            k = qkv[..., D : 2 * D].reshape(1, S, H, HD)
+            v = qkv[..., 2 * D :].reshape(1, S, H, HD)
+            q, k = _rope_hf(q, k, cos, sin)
+            kc, vc = getattr(self, f"k_cache_{i}"), getattr(self, f"v_cache_{i}")
+            o = torch.ops.llama.sdpa_with_kv_cache(q, k, v, kc, vc, p, S, None, 0.0, True, None)
+            x = x + lin3(L.out_proj, o.reshape(1, S, D))
+            h = layer_norm_dyn(L.norm2, x, 1e-5, self.mean_w)
+            x = x + lin3(L.linear2, gelu(lin3(L.linear1, h)))
+        h = layer_norm(self.out_norm, x[:, -1, :], 1e-5)
+        eos = lin3(self.out_eos, h)
+        latent = noise + flow_net(self.fn, h, noise, self.t_comb)
+        nxt = lin3(self.input_linear, latent).reshape(1, 1, D) + self.cond
+        return latent, eos, nxt
+
+
+MIMI_STATE = [
+    ("upsample_partial", (1, D, UPS)),
+    ("k0", (1, H, MIMI_CONTEXT, HD)),
+    ("v0", (1, H, MIMI_CONTEXT, HD)),
+    ("k1", (1, H, MIMI_CONTEXT, HD)),
+    ("v1", (1, H, MIMI_CONTEXT, HD)),
+    ("kv_seen", (1, 1, 1, MIMI_CONTEXT)),
+    ("rope_offset", (1,)),
+    ("init_conv_prev", (1, 512, 6)),
+    ("up1_partial", (1, 256, 6)),
+    ("res1_prev", (1, 256, 2)),
+    ("up2_partial", (1, 128, 5)),
+    ("res2_prev", (1, 128, 2)),
+    ("up3_partial", (1, 64, 4)),
+    ("res3_prev", (1, 64, 2)),
+    ("final_conv_prev", (1, 64, 2)),
+]
+
+
+def mimi_zero_state():
+    return [torch.zeros(s) for _, s in MIMI_STATE]
+
+
+def _convtr(x, w, b, s, groups=1):
+    """ConvTranspose1d as a 4-D (H = 1) conv_transpose2d: the public runtime's 3-D convolution
+    has no transposed path, its 4-D one does (with groups)."""
+    y = F.conv_transpose2d(x.unsqueeze(2), w.unsqueeze(2), b, stride=(1, s), groups=groups)
+    return y.squeeze(2)
+
+
+def _convtr_step(x, w, b, partial, s, groups=1):
+    y = _convtr(x, w, b, s, groups)
+    pt = partial.shape[2]
+    n = y.shape[2]
+    head = y[:, :, :pt] + partial
+    out = torch.cat([head, y[:, :, pt : n - pt]], 2)
+    tail = y[:, :, n - pt :]
+    if b is not None:
+        tail = tail - b.view(1, -1, 1)
+    return out, tail
+
+
+def _conv_step(m, x, prev):
+    xx = torch.cat([prev, x], 2)
+    y = F.conv1d(xx, m.conv.weight, m.conv.bias)
+    return y, xx[:, :, xx.shape[2] - prev.shape[2] :]
+
+
+class MimiStepProgram(nn.Module):
+    """forward(latent[1,1,32], *state) -> pcm[1,1,1920], *new_state (MIMI_STATE order).
+
+    One latent frame per call, exactly the runtime's streaming decoder: every piece of streaming
+    state (conv tails, transposed-conv partial sums, the 2-layer transformer's last 250 keys and
+    values, which of them are filled, and the RoPE offset) is an explicit input and output. The
+    first call takes all-zero state."""
+
+    def __init__(self, w: PhononWeights):
+        super().__init__()
+        m = w.mimi
+        self.emb_std = nn.Parameter(w.flow_lm.emb_std.detach().clone().view(1, 1, LDIM), requires_grad=False)
+        self.emb_mean = nn.Parameter(w.flow_lm.emb_mean.detach().clone().view(1, 1, LDIM), requires_grad=False)
+        self.qproj = _LinP(m.quantizer.output_proj.weight.view(D, LDIM))
+        self.up_w = nn.Parameter(m.upsample.convtr.convtr.weight.detach().clone(), requires_grad=False)
+        self.layers = nn.ModuleList()
+        for L in m.decoder_transformer.transformer.layers:
+            q = nn.Module()
+            q.norm1, q.norm2 = L.norm1, L.norm2
+            q.in_proj = _LinP(_permuted_in_proj(L.self_attn.in_proj.weight))
+            q.out_proj, q.linear1, q.linear2 = L.self_attn.out_proj, L.linear1, L.linear2
+            q.ls1 = nn.Parameter(L.layer_scale_1.scale.detach().clone(), requires_grad=False)
+            q.ls2 = nn.Parameter(L.layer_scale_2.scale.detach().clone(), requires_grad=False)
+            self.layers.append(q)
+        self.dec = m.decoder.model
+        self.register_buffer("inv_freq", inv_freq().view(1, HD // 2), persistent=False)
+        self.register_buffer("steps", torch.arange(UPS, dtype=torch.float32).view(UPS, 1), persistent=False)
+        i = torch.arange(UPS)[:, None]
+        j = torch.arange(MIMI_CONTEXT + UPS)[None, :]
+        band = torch.where((j >= i) & (j <= i + MIMI_CONTEXT), 0.0, NEG)
+        self.register_buffer("band", band.view(1, 1, UPS, MIMI_CONTEXT + UPS), persistent=False)
+        self.register_buffer("ones_seen", torch.ones(1, 1, 1, UPS), persistent=False)
+
+    def _attn(self, L, x, kc, vc, seen, cos, sin):
+        T = UPS
+        h = layer_norm(L.norm1, x, 1e-5)
+        qkv = lin(L.in_proj, h)
+        q = qkv[..., :D].reshape(1, T, H, HD)
+        k = qkv[..., D : 2 * D].reshape(1, T, H, HD)
+        v = qkv[..., 2 * D :].reshape(1, T, H, HD)
+        rh = lambda t: torch.cat([-t[..., HD // 2 :], t[..., : HD // 2]], -1)  # noqa: E731
+        q = q * cos + rh(q) * sin
+        k = k * cos + rh(k) * sin
+        q = q.permute(0, 2, 1, 3)  # [1, H, T, HD]
+        k = k.permute(0, 2, 1, 3)
+        v = v.permute(0, 2, 1, 3)
+        kk = torch.cat([kc, k], 2)  # [1, H, 266, HD]
+        vv = torch.cat([vc, v], 2)
+        seen_all = torch.cat([seen, self.ones_seen], 3)
+        mask = self.band + (seen_all - 1.0) * 1.0e30
+        sc = torch.bmm(q.reshape(H, T, HD), kk.reshape(H, MIMI_CONTEXT + T, HD).transpose(1, 2)) * (HD**-0.5)
+        pr = torch.softmax(sc.reshape(1, H, T, MIMI_CONTEXT + T) + mask, -1)
+        o = torch.bmm(pr.reshape(H, T, MIMI_CONTEXT + T), vv.reshape(H, MIMI_CONTEXT + T, HD))
+        o = o.reshape(1, H, T, HD).permute(0, 2, 1, 3).reshape(1, T, D)
+        x = x + L.ls1 * lin(L.out_proj, o)
+        h = layer_norm(L.norm2, x, 1e-5)
+        x = x + L.ls2 * lin(L.linear2, gelu(lin(L.linear1, h)))
+        return x, kk[:, :, UPS:], vv[:, :, UPS:], seen_all[..., UPS:]
+
+    def _res(self, blk, x, prev):
+        v, prev = _conv_step(blk.block[1], elu(x), prev)
+        v = F.conv1d(elu(v), blk.block[3].conv.weight, blk.block[3].conv.bias)
+        return x + v, prev
+
+    def forward(self, latent, up_p, k0, v0, k1, v1, seen, off, c0, u1, r1, u2, r2, u3, r3, cf):
+        x = latent * self.emb_std + self.emb_mean  # [1, 1, 32]
+        x = lin(self.qproj, x).reshape(1, D, 1)
+        x, up_p = _convtr_step(x, self.up_w, None, up_p, UPS, groups=D)  # [1, 512, 16]
+        pos = self.steps + off  # [16, 1]
+        ang = pos * self.inv_freq  # [16, 32]
+        cos = torch.cat([torch.cos(ang), torch.cos(ang)], -1).reshape(1, UPS, 1, HD)
+        sin = torch.cat([torch.sin(ang), torch.sin(ang)], -1).reshape(1, UPS, 1, HD)
+        x = x.permute(0, 2, 1)  # [1, 16, 512]
+        x, k0, v0, seen0 = self._attn(self.layers[0], x, k0, v0, seen, cos, sin)
+        x, k1, v1, _ = self._attn(self.layers[1], x, k1, v1, seen, cos, sin)
+        x = x.permute(0, 2, 1)  # [1, 512, 16]
+        m = self.dec
+        x, c0 = _conv_step(m[0], x, c0)
+        x, u1 = _convtr_step(elu(x), m[2].convtr.weight, m[2].convtr.bias, u1, 6)
+        x, r1 = self._res(m[3], x, r1)
+        x, u2 = _convtr_step(elu(x), m[5].convtr.weight, m[5].convtr.bias, u2, 5)
+        x, r2 = self._res(m[6], x, r2)
+        x, u3 = _convtr_step(elu(x), m[8].convtr.weight, m[8].convtr.bias, u3, 4)
+        x, r3 = self._res(m[9], x, r3)
+        pcm, cf = _conv_step(m[11], elu(x), cf)
+        return pcm, up_p, k0, v0, k1, v1, seen0, off + float(UPS), c0, u1, r1, u2, r2, u3, r3, cf

export/_whisper.py

@@ -0,0 +1,662 @@
+"""Whisper base building blocks for recipes/whisper-base.py.
+
+Everything here is built from the pinned openai/whisper-base checkpoint loaded through transformers
+(WhisperForConditionalGeneration); the program wrappers reuse its nn.Modules (weights are shared, not
+copied) and only re-express the forward pass so that it lowers to one WebGPU delegate per program.
+
+Program contract (same split and I/O as the playground's previous whisper-base export):
+
+  mel            forward(frames[1,3000,400] f32) -> power_mel[1,3000,80] f32
+                 frames: 30 s of 16 kHz audio, zero-padded to 480000 samples, reflect-padded by 200 on
+                 both ends, 3000 windows of 400 samples at hop 160 (= whisper.js buildFrames).
+                 power_mel = |rfft(hann_periodic * frame)|^2 @ mel_filters (no log; see host_log_mel).
+  encoder        forward(log_mel[1,80,3000] f32) -> encoder_hidden[1,1500,512] f32   (fp32 weights)
+  cross_kv       forward(encoder_hidden[1,1500,512]) -> 12 x [1,8,1500,64] f32, order k0,v0,...,k5,v5
+                 (per decoder layer: encoder_attn.k_proj / v_proj, heads-first [B,H,S,D])
+  decoder_split  forward(tok_embeds[1,64,512], k0, v0, ..., k5, v5) -> hidden[1,64,512]
+                 tok_embeds: token embeddings WITHOUT positions for positions 0..64 (zero rows beyond
+                 the live prefix); the program adds positions 0..63, runs causal self-attention over
+                 the 64 rows and returns final-LayerNorm hidden states (row p predicts token p+1).
+  decode_stream  forward(embed[1,1,512], input_pos[1] i64, k0, v0, ..., k5, v5) -> hidden[1,1,512]
+                 embed = token embedding + positional embedding of `input_pos` (added on the host);
+                 self-attention K/V cache [1,448,8,64] x 6 layers x {k,v} lives inside the program
+                 (mutable buffers owned by the delegate); feed positions 0,1,2,... in order.
+  lmhead         forward(hidden[1,1,512]) -> logits[1,1,51872]  (tied to decoder.embed_tokens; rows
+                 51865..51871 are padding with logit -1e9, see LMHeadProgram)
+
+Precision (whisper-base.py --precision):
+  q4   (previous-export parity) cross_kv / decoder_split / decode_stream / lmhead linears are 4-bit groupwise
+       weight-only (torchao IntxWeightOnlyConfig int4, group 64, symmetric, "affine" qparams ->
+       et_vk.linear_q4gsw); int8 token-embedding table. mel and encoder are fp32.
+  fp32 every program fp32 (aten.linear on 2-D views); fp32 token-embedding table.
+
+Host-side files (reproduced from the checkpoint by the recipe, consumed by the playground JS):
+  whisper_embed.i8.bin         int8 [51865,512]  round(embed_tokens / scale), scale = rowwise absmax/127  (q4)
+  whisper_embed.scale.f32.bin  f32  [51865]      the per-row scales                                       (q4)
+  whisper_embed.f32.bin        f32  [51865,512]  embed_tokens                                             (fp32)
+  whisper_pos_embed.f32.bin    f32  [448,512]    decoder.embed_positions.weight
+  vocab.json                   upstream tokenizer vocab (byte-level BPE, detokenization only)
+"""
+
+import copy
+import io
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+SAMPLE_RATE = 16000
+N_FFT = 400
+HOP = 160
+FRAMES = 3000
+N_SAMPLES = FRAMES * HOP
+PAD = N_FFT // 2
+N_FREQ = N_FFT // 2 + 1
+N_MELS = 80
+ENC_LEN = 1500
+D_MODEL = 512
+N_HEADS = 8
+HEAD_DIM = 64
+N_LAYERS = 6
+MAX_DEC = 64  # decoder_split rows; the playground's decode budget (prefix + generated)
+KV_LEN = 448  # decode_stream self-attention cache length (= max_target_positions)
+VOCAB = 51865
+VOCAB_PADDED = 51872  # lmhead rows (multiple of 8, = the previous export)
+PAD_LOGIT = -1e9
+
+TOK = {
+    "sot": 50258,
+    "eot": 50257,
+    "en": 50259,
+    "transcribe": 50359,
+    "notimestamps": 50363,
+}
+PREFIX_EN = [TOK["sot"], TOK["en"], TOK["transcribe"], TOK["notimestamps"]]
+
+
+# --------------------------------------------------------------------------------------------
+# Host-side math that the playground does in JS (kept here so validation runs the same steps).
+# --------------------------------------------------------------------------------------------
+
+
+def build_frames(pcm):
+    """Mirror of whisper.js buildFrames(): [3000, 400] float32 windows (no window function)."""
+    pcm = np.asarray(pcm, dtype=np.float32)
+    base = np.zeros(N_SAMPLES + 2 * PAD, dtype=np.float32)
+    n = min(pcm.size, N_SAMPLES)
+    base[PAD : PAD + n] = pcm[:n]
+    for i in range(PAD):
+        base[PAD - 1 - i] = base[PAD + 1 + i]
+        base[PAD + N_SAMPLES + i] = base[PAD + N_SAMPLES - 2 - i]
+    idx = np.arange(FRAMES)[:, None] * HOP + np.arange(N_FFT)[None, :]
+    return base[idx]
+
+
+def host_log_mel(power_mel):
+    """Host step between the mel and encoder programs (WhisperFeatureExtractor normalization).
+
+    power_mel [1,3000,80] (or [3000,80]) -> log_mel [1,80,3000]:
+      x = log10(max(p, 1e-10)); x = max(x, max(x) - 8); x = (x + 4) / 4; transpose.
+    float32 throughout (the JS equivalent uses Float32Array + Math.log10).
+    """
+    p = np.asarray(power_mel, dtype=np.float32).reshape(FRAMES, N_MELS)
+    x = np.log10(np.maximum(p, np.float32(1e-10))).astype(np.float32)
+    x = np.maximum(x, x.max() - np.float32(8.0))
+    x = ((x + np.float32(4.0)) / np.float32(4.0)).astype(np.float32)
+    return np.ascontiguousarray(x.T)[None]
+
+
+class HostTables:
+    """Token-embedding lookup as the playground JS does it, plus the files that carry the tables.
+
+    kind="int8": whisper_embed.i8.bin (int8 [51865,512], round(w / s)) + whisper_embed.scale.f32.bin
+                 (s = rowwise absmax / 127): byte-identical to the previous export's files.
+    kind="fp32": whisper_embed.f32.bin (fp32 [51865,512], the checkpoint's embed_tokens).
+    Both: whisper_pos_embed.f32.bin (fp32 [448,512], decoder.embed_positions)."""
+
+    def __init__(self, model, kind):
+        e = model.model.decoder.embed_tokens.weight.detach().float()
+        self.kind = kind
+        self.pos = model.model.decoder.embed_positions.weight.detach().float().numpy()
+        if kind == "int8":
+            scale = e.abs().amax(dim=1) / 127.0
+            self.table = torch.round(e / scale[:, None]).clamp(-127, 127).to(torch.int8).numpy()
+            self.scale = scale.numpy().astype(np.float32)
+        elif kind == "fp32":
+            self.table = e.numpy().astype(np.float32)
+            self.scale = None
+        else:
+            raise ValueError(kind)
+
+    def embed(self, ids, position=None):
+        e = self.table[ids].astype(np.float32)
+        if self.scale is not None:
+            e = e * self.scale[ids][..., None]
+        if position is not None:
+            e = e + self.pos[position]
+        return e.astype(np.float32)
+
+    def files(self):
+        out = {}
+        if self.kind == "int8":
+            out["whisper_embed.i8.bin"] = self.table.tobytes()
+            out["whisper_embed.scale.f32.bin"] = self.scale.astype("<f4").tobytes()
+        else:
+            out["whisper_embed.f32.bin"] = self.table.astype("<f4").tobytes()
+        out["whisper_pos_embed.f32.bin"] = self.pos.astype("<f4").tobytes()
+        return out
+
+
+# --------------------------------------------------------------------------------------------
+# Program wrappers.
+# --------------------------------------------------------------------------------------------
+
+
+def dft_mel_weights(feature_extractor):
+    """cos / -sin DFT rows with the periodic Hann window folded in, plus the HF mel filters."""
+    n = np.arange(N_FFT)
+    k = np.arange(N_FREQ)
+    window = 0.5 - 0.5 * np.cos(
+        2.0 * np.pi * n / N_FFT
+    )  # periodic Hann (torch / HF default)
+    ang = 2.0 * np.pi * ((k[:, None] * n[None, :]) % N_FFT) / N_FFT
+    wc = (np.cos(ang) * window[None, :]).astype(np.float32)
+    ws = (-np.sin(ang) * window[None, :]).astype(np.float32)
+    mel = np.asarray(feature_extractor.mel_filters, dtype=np.float64).T.astype(
+        np.float32
+    )
+    assert mel.shape == (N_MELS, N_FREQ), mel.shape
+    return torch.from_numpy(wc), torch.from_numpy(ws), torch.from_numpy(mel)
+
+
+class MelProgram(torch.nn.Module):
+    def __init__(self, feature_extractor):
+        super().__init__()
+        wc, ws, mel = dft_mel_weights(feature_extractor)
+        self.register_buffer("wc", wc)
+        self.register_buffer("ws", ws)
+        self.register_buffer("mel", mel)
+
+    def forward(self, frames):
+        x = frames.reshape(
+            FRAMES, N_FFT
+        )  # aten.linear on the WebGPU runtime is 2-D only
+        re = F.linear(x, self.wc)
+        im = F.linear(x, self.ws)
+        return F.linear(re * re + im * im, self.mel).reshape(1, FRAMES, N_MELS)
+
+
+def _heads(x, s):
+    """[s, 512] or [1, s, 512] -> [1, 8, s, 64]."""
+    return x.reshape(1, s, N_HEADS, HEAD_DIM).transpose(1, 2)
+
+
+def _merge(o, s):
+    """[1, 8, s, 64] -> [1, s, 512]."""
+    return o.transpose(1, 2).reshape(1, s, D_MODEL)
+
+
+def attention(q, k, v, mask, impl):
+    """q/k/v [1,H,S,D]; scale = 1/sqrt(64) = 0.125 (a power of two, so applying it to q as HF does
+    or to q.k as the kernels do is bit-identical)."""
+    if impl == "et_vk":
+        import executorch.backends.vulkan.custom_ops_lib  # noqa: F401  (torch.ops.et_vk.*)
+
+        return torch.ops.et_vk.sdpa(q, k, v, mask, None)
+    s = torch.matmul(q * 0.125, k.transpose(-2, -1))
+    if mask is not None:
+        s = s + mask
+    return torch.matmul(torch.softmax(s, dim=-1), v)
+
+
+class EncoderProgram(torch.nn.Module):
+    """HF WhisperEncoder forward re-expressed with 2-D fp32 linears.
+
+    attn="eager": matmul + softmax (lowered to aten.bmm/_softmax; activation memory is planned).
+    attn="et_vk": the fused et_vk.sdpa op (what the previous export used); on the public WebGPU
+    runtime each call allocates its own two [8,1500,1500] f32 scratch buffers (72 MB each).
+    """
+
+    def __init__(self, model, attn="eager"):
+        super().__init__()
+        self.enc = model.model.encoder
+        self.attn = attn
+
+    def forward(self, log_mel):
+        e = self.enc
+        x = F.gelu(e.conv1(log_mel))
+        x = F.gelu(e.conv2(x))  # [1, 512, 1500]
+        x = x.permute(0, 2, 1).reshape(ENC_LEN, D_MODEL) + e.embed_positions.weight
+        for layer in e.layers:
+            a = layer.self_attn
+            h = layer.self_attn_layer_norm(x)
+            q = _heads(a.q_proj(h), ENC_LEN)
+            k = _heads(a.k_proj(h), ENC_LEN)
+            v = _heads(a.v_proj(h), ENC_LEN)
+            o = attention(q, k, v, None, "et_vk" if self.attn == "et_vk" else "torch")
+            x = x + a.out_proj(o.transpose(1, 2).reshape(ENC_LEN, D_MODEL))
+            h = layer.final_layer_norm(x)
+            x = x + layer.fc2(F.gelu(layer.fc1(h)))
+        return e.layer_norm(x).reshape(1, ENC_LEN, D_MODEL)
+
+
+def _lin(m, x):
+    """nn.Linear that lowers on the public WebGPU runtime: fp32 aten.linear only accepts 2-D
+    inputs, so fp32 weights get a [-1, K] view; int4 (q4gsw) linears accept any rank and keep it."""
+    if type(m.weight.data) is torch.Tensor and x.dim() != 2:
+        return m(x.reshape(-1, x.shape[-1])).reshape(*x.shape[:-1], m.out_features)
+    return m(x)
+
+
+class CrossKVProgram(torch.nn.Module):
+    def __init__(self, decoder):
+        super().__init__()
+        self.layers = decoder.layers
+
+    def forward(self, enc):
+        out = []
+        for layer in self.layers:
+            a = layer.encoder_attn
+            out.append(_heads(_lin(a.k_proj, enc), ENC_LEN))
+            out.append(_heads(_lin(a.v_proj, enc), ENC_LEN))
+        return tuple(out)
+
+
+def _mlp(layer, x):
+    return x + _lin(layer.fc2, F.gelu(_lin(layer.fc1, layer.final_layer_norm(x))))
+
+
+def _cross(layer, x, k, v, s, impl):
+    a = layer.encoder_attn
+    h = layer.encoder_attn_layer_norm(x)
+    o = attention(_heads(_lin(a.q_proj, h), s), k, v, None, impl)
+    return x + _lin(a.out_proj, _merge(o, s))
+
+
+class DecoderSplitProgram(torch.nn.Module):
+    """Full-recompute decoder over MAX_DEC rows with a causal mask (no KV cache)."""
+
+    def __init__(self, decoder, attn="et_vk"):
+        super().__init__()
+        self.dec = decoder
+        self.attn = attn
+        self.register_buffer(
+            "pos", decoder.embed_positions.weight.detach()[:MAX_DEC].clone()
+        )
+        causal = torch.full((MAX_DEC, MAX_DEC), -1e30).triu(1)  # same fill as the previous export
+        # et_vk.sdpa on the WebGPU runtime requires the full [B, H, S_q, S_kv] mask.
+        self.register_buffer(
+            "mask", causal.expand(1, N_HEADS, MAX_DEC, MAX_DEC).contiguous()
+        )
+
+    def forward(self, tok_embeds, *cross_kv):
+        x = tok_embeds + self.pos
+        for i, layer in enumerate(self.dec.layers):
+            a = layer.self_attn
+            h = layer.self_attn_layer_norm(x)
+            q = _heads(_lin(a.q_proj, h), MAX_DEC)
+            k = _heads(_lin(a.k_proj, h), MAX_DEC)
+            v = _heads(_lin(a.v_proj, h), MAX_DEC)
+            o = attention(q, k, v, self.mask, self.attn)
+            x = x + _lin(a.out_proj, _merge(o, MAX_DEC))
+            x = _cross(
+                layer, x, cross_kv[2 * i], cross_kv[2 * i + 1], MAX_DEC, self.attn
+            )
+            x = _mlp(layer, x)
+        return self.dec.layer_norm(x)
+
+
+class DecodeStreamProgram(torch.nn.Module):
+    """One token per call; self-attention K/V cache [1,448,8,64] per layer kept in the program.
+
+    Self-attention uses llama.update_cache + llama.custom_sdpa (executorch.extension.llm.custom_ops),
+    which the Vulkan partitioner fuses into sdpa_with_kv_cache with an et_vk.select_as_symint
+    position. Cross-attention uses et_vk.sdpa (attn="et_vk") or matmul/softmax (attn="torch").
+    """
+
+    def __init__(self, decoder, attn="et_vk"):
+        super().__init__()
+        from executorch.extension.llm.custom_ops import custom_ops  # noqa: F401
+
+        self.dec = decoder
+        self.attn = attn
+        for i in range(N_LAYERS):
+            self.register_buffer(
+                f"k_cache_{i}",
+                torch.zeros(1, KV_LEN, N_HEADS, HEAD_DIM),
+                persistent=False,
+            )
+            self.register_buffer(
+                f"v_cache_{i}",
+                torch.zeros(1, KV_LEN, N_HEADS, HEAD_DIM),
+                persistent=False,
+            )
+
+    def reset(self):
+        for i in range(N_LAYERS):
+            getattr(self, f"k_cache_{i}").zero_()
+            getattr(self, f"v_cache_{i}").zero_()
+
+    def forward(self, embed, input_pos, *cross_kv):
+        pos = input_pos[0].item()
+        torch._check_is_size(pos)
+        torch._check(pos < KV_LEN)
+        x = embed
+        for i, layer in enumerate(self.dec.layers):
+            a = layer.self_attn
+            h = layer.self_attn_layer_norm(x)
+            q = _lin(a.q_proj, h).view(1, 1, N_HEADS, HEAD_DIM)
+            k = _lin(a.k_proj, h).view(1, 1, N_HEADS, HEAD_DIM)
+            v = _lin(a.v_proj, h).view(1, 1, N_HEADS, HEAD_DIM)
+            kc = getattr(self, f"k_cache_{i}")
+            vc = getattr(self, f"v_cache_{i}")
+            torch.ops.llama.update_cache(k, kc, pos)
+            torch.ops.llama.update_cache(v, vc, pos)
+            o = torch.ops.llama.custom_sdpa(q, kc, vc, pos, None, 0.0, True, None)
+            x = x + _lin(a.out_proj, o.reshape(1, 1, D_MODEL))
+            x = _cross(layer, x, cross_kv[2 * i], cross_kv[2 * i + 1], 1, self.attn)
+            x = _mlp(layer, x)
+        return self.dec.layer_norm(x)
+
+
+class LMHeadProgram(torch.nn.Module):
+    """Tied output projection. The vocab is padded to `pad_to` rows: executorch 1.5.1's q4gsw
+    lowering pads the scales to a multiple of 4 output channels and then fails its own shape
+    check for N = 51865, so N must already be a multiple of 4 (the previous export used 51872).
+    Padded rows have zero weights and a -1e9 bias, so they can never win an argmax."""
+
+    def __init__(self, model, pad_to=0):
+        super().__init__()
+        w = model.model.decoder.embed_tokens.weight.detach().clone()
+        n = w.shape[0]
+        b = torch.zeros(max(n, pad_to))
+        if pad_to and pad_to > n:
+            w = torch.cat([w, w.new_zeros(pad_to - n, w.shape[1])])
+            b[n:] = PAD_LOGIT
+        self.proj = torch.nn.Linear(D_MODEL, w.shape[0], bias=True)
+        self.proj.weight = torch.nn.Parameter(w, requires_grad=False)
+        self.proj.bias = torch.nn.Parameter(b, requires_grad=False)
+
+    def forward(self, hidden):
+        return _lin(self.proj, hidden)
+
+
+# --------------------------------------------------------------------------------------------
+# Quantization.
+# --------------------------------------------------------------------------------------------
+
+Q4_DESC = (
+    "torchao.quantization.quant_api.IntxWeightOnlyConfig(weight_dtype=torch.int4, "
+    "granularity=PerGroup(64), intx_choose_qparams_algorithm='affine') applied with quantize_ to "
+    "every nn.Linear of the program; symmetric, scale = group absmax / 7.5, q in [-8, 7]"
+)
+
+
+def quantize_q4(module, group_size=64):
+    from torchao.quantization.granularity import PerGroup
+    from torchao.quantization.quant_api import IntxWeightOnlyConfig, quantize_
+
+    quantize_(
+        module,
+        IntxWeightOnlyConfig(
+            weight_dtype=torch.int4,
+            granularity=PerGroup(group_size),
+            intx_choose_qparams_algorithm="affine",
+        ),
+    )
+    return module
+
+
+Q8_HEAD_DESC = (
+    "PT2E with executorch.backends.vulkan.quantizer.VulkanQuantizer: weights int8 per-output-channel "
+    "symmetric (PerChannelMinMaxObserver, the quantizer's own weight spec), input activation int8 "
+    "per-tensor affine STATIC (HistogramObserver) calibrated on lmhead inputs; lowers to "
+    "quantized_decomposed.quantize_per_tensor + et_vk.linear_q8ta_q8csw"
+)
+
+
+def q8_lmhead(lm, calib_hidden):
+    """Static-activation int8 x per-channel int8 lmhead (the only 8-bit linear executorch 1.5.1's
+    Vulkan lowering emits that the public WebGPU runtime registers). Returns the converted GraphModule
+    and the (input scale, zero point, dequantized weight) needed for an equivalent portable module."""
+    from executorch.backends.vulkan.quantizer.vulkan_quantizer import (
+        get_symmetric_quantization_config,
+        VulkanQuantizer,
+    )
+    from torchao.quantization.pt2e import HistogramObserver
+    from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
+    from torchao.quantization.pt2e.quantizer import QuantizationConfig, QuantizationSpec
+
+    base = get_symmetric_quantization_config(is_dynamic=False, weight_bits=8)
+    act = QuantizationSpec(
+        dtype=torch.int8,
+        quant_min=-128,
+        quant_max=127,
+        qscheme=torch.per_tensor_affine,
+        is_dynamic=False,
+        observer_or_fake_quant_ctr=HistogramObserver,
+    )
+    cfg = QuantizationConfig(
+        input_activation=act, output_activation=None, weight=base.weight, bias=None, is_qat=False
+    )
+    with torch.no_grad():
+        gm = torch.export.export(lm.eval(), (calib_hidden[:1],), strict=False).module()
+        gm = prepare_pt2e(gm, VulkanQuantizer().set_global(cfg))
+        for i in range(calib_hidden.shape[0]):
+            gm(calib_hidden[i : i + 1])
+        gm = convert_pt2e(gm)
+    qp, wdq = None, None
+    for n in gm.graph.nodes:
+        if n.op != "call_function":
+            continue
+        name = str(n.target)
+        if "quantize_per_tensor" in name and "dequantize" not in name and qp is None:
+            qp = (float(n.args[1]), int(n.args[2]))
+        if "dequantize_per_channel" in name:
+            w = getattr(gm, n.args[0].target)
+            sc = getattr(gm, n.args[1].target)
+            wdq = (w.float() * sc.float()[:, None]).contiguous()
+    if qp is None or wdq is None:
+        raise RuntimeError("q8 lmhead: could not find the quantize/dequantize nodes")
+    return gm, qp, wdq
+
+
+class FakeQ8LMHead(torch.nn.Module):
+    """Portable-runtime equivalent of the q8 lmhead: per-tensor int8 fake-quant of the input, then
+    an fp32 linear on the dequantized int8 weights (same bias)."""
+
+    def __init__(self, lm, qp, wdq):
+        super().__init__()
+        self.scale, self.zp = qp
+        self.register_buffer("w", wdq)
+        self.register_buffer("b", lm.proj.bias.detach().clone())
+
+    def forward(self, hidden):
+        x = hidden.reshape(-1, D_MODEL)
+        q = torch.clamp(torch.round(x / self.scale) + self.zp, -128, 127)
+        x = (q - self.zp) * self.scale
+        return F.linear(x, self.w, self.b).reshape(*hidden.shape[:-1], self.w.shape[0])
+
+
+def word_errors(tokenizer, hyp_text, ref_text):
+    """(edits, reference words) after WhisperTokenizer.normalize (English text normalizer)."""
+    h = tokenizer.normalize(hyp_text).split()
+    r = tokenizer.normalize(ref_text).split()
+    d = list(range(len(r) + 1))
+    for i in range(1, len(h) + 1):
+        prev, d[0] = d[0], i
+        for j in range(1, len(r) + 1):
+            cur = d[j]
+            d[j] = min(d[j] + 1, d[j - 1] + 1, prev + (h[i - 1] != r[j - 1]))
+            prev = cur
+    return d[len(r)], len(r)
+
+
+def dequantized_copy(module):
+    """Same module with every quantized linear weight replaced by its fp32 dequantization.
+
+    Used only for the CPU portable .pte (the portable runtime has no torchao kernels); its numerics
+    are the q4 numerics (dequantize-then-fp32-matmul, as the q4gsw kernels compute)."""
+    m = copy.deepcopy(module)
+    for sub in m.modules():
+        if isinstance(sub, torch.nn.Linear) and hasattr(sub.weight, "dequantize"):
+            sub.weight = torch.nn.Parameter(
+                sub.weight.dequantize().float(), requires_grad=False
+            )
+    return m
+
+
+# --------------------------------------------------------------------------------------------
+# Speech sample.
+# --------------------------------------------------------------------------------------------
+
+
+def load_librispeech_dummy(parquet_path, max_sec=15.0, count=8, longer=False):
+    """First `count` clips (in file order) no longer than max_sec seconds, as 16 kHz float32.
+    longer=True selects the clips *longer* than max_sec instead (used only for calibration, so
+    calibration and evaluation clips never overlap)."""
+    import pyarrow.parquet as pq
+    import soundfile as sf
+
+    table = pq.read_table(parquet_path).to_pylist()
+    clips = []
+    for row in table:
+        audio, sr = sf.read(io.BytesIO(row["audio"]["bytes"]), dtype="float32")
+        if sr != SAMPLE_RATE:
+            raise ValueError(f"unexpected sample rate {sr}")
+        if audio.ndim != 1:
+            audio = audio.mean(axis=1)
+        if (audio.size / SAMPLE_RATE > max_sec) != longer:
+            continue
+        clips.append({"id": row["id"], "text": row["text"], "pcm": audio})
+        if len(clips) >= count:
+            break
+    return clips
+
+
+def seconds(pcm):
+    return round(len(pcm) / SAMPLE_RATE, 3)
+
+
+def isclose_report(a, b):
+    a = np.asarray(a, dtype=np.float64).ravel()
+    b = np.asarray(b, dtype=np.float64).ravel()
+    d = np.abs(a - b)
+    cos = float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-30))
+    return {
+        "max_abs": float(d.max()),
+        "mean_abs": float(d.mean()),
+        "max_rel_to_absmax": float(d.max() / (np.abs(b).max() + 1e-30)),
+        "cosine": cos,
+    }
+
+
+
+# --------------------------------------------------------------------------------------------
+# Lowered-graph inspection (what the WebGPU runtime will see).
+# --------------------------------------------------------------------------------------------
+
+
+def _vk_graphs(pte_bytes):
+    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")
+    out = []
+    for plan in prog.execution_plan:
+        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]
+            out.append(flatbuffer_to_vk_graph(extract_vk_flatbuffer(blob)))
+    return prog, out
+
+
+def program_io(pte_bytes):
+    """Method names and their input/output tensor shapes + dtypes from the serialized 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
+    res = {}
+    for plan in prog.execution_plan:
+
+        def desc(i):
+            v = plan.values[i].val
+            if hasattr(v, "sizes"):
+                return {"shape": list(v.sizes), "dtype": v.scalar_type.name.lower()}
+            return {"kind": type(v).__name__}
+
+        res[plan.name] = {
+            "inputs": [desc(i) for i in plan.inputs],
+            "outputs": [desc(i) for i in plan.outputs],
+        }
+    return res
+
+
+def vk_checks(pte_bytes, runtime_ops=None):
+    """Checks against constraints of the public WebGPU runtime's op implementations
+    (backends/webgpu/runtime/ops): op registered, aten.linear 2-D, et_vk.sdpa mask shape,
+    sdpa_with_kv_cache without mask / with causal, fp32 tensors only."""
+    _, graphs = _vk_graphs(pte_bytes)
+    problems, notes = [], {}
+    for g in graphs:
+        vals = [v.value for v in g.values]
+
+        def t(i):
+            v = vals[i]
+            return v if type(v).__name__ == "VkTensor" else None
+
+        dtypes = sorted({type(v).__name__ if type(v).__name__ != "VkTensor" else v.datatype.name for v in vals})
+        notes["value_kinds"] = dtypes
+        for op in g.chain:
+            name = op.name
+            if runtime_ops is not None and name not in runtime_ops:
+                problems.append(f"{name}: not registered in the WebGPU runtime op list")
+            a = op.args
+            if name == "aten.linear.default":
+                ri, rw = len(t(a[0]).dims), len(t(a[1]).dims)
+                notes.setdefault("aten.linear input ranks", set()).add(ri)
+                if ri != 2 or rw != 2:
+                    problems.append(f"aten.linear with input rank {ri}, weight rank {rw} (runtime needs 2-D)")
+            if name == "et_vk.sdpa.default":
+                q, k = t(a[0]), t(a[1])
+                m = t(a[3])
+                if m is not None:
+                    need = q.dims[0] * q.dims[1] * q.dims[2] * k.dims[2]
+                    got = int(np.prod(m.dims))
+                    if need != got:
+                        problems.append(f"et_vk.sdpa mask numel {got} != B*H*Sq*Skv {need}")
+                notes.setdefault("et_vk.sdpa q/k/mask", set()).add(
+                    (tuple(q.dims), tuple(k.dims), tuple(m.dims) if m is not None else None)
+                )
+            if name == "sdpa_with_kv_cache.default":
+                mask_kind = type(vals[a[7]]).__name__
+                causal = vals[a[9]]
+                if mask_kind != "Null":
+                    problems.append("sdpa_with_kv_cache with attn_mask (runtime rejects)")
+                if not (type(causal).__name__ == "Bool" and causal.bool_val):
+                    problems.append(f"sdpa_with_kv_cache is_causal={causal} (runtime needs true)")
+                notes.setdefault("sdpa_with_kv_cache q/k_cache", set()).add(
+                    (tuple(t(a[0]).dims), tuple(t(a[3]).dims), type(vals[a[5]]).__name__)
+                )
+            if name == "aten.convolution.default":
+                notes.setdefault("aten.convolution input/weight", set()).add(
+                    (tuple(t(a[0]).dims), tuple(t(a[1]).dims))
+                )
+        for v in vals:
+            if type(v).__name__ == "VkTensor" and v.datatype.name not in ("FLOAT32", "INT32", "INT64", "UINT8", "INT8"):
+                problems.append(f"tensor dtype {v.datatype.name}")
+                break
+    for k, v in list(notes.items()):
+        if isinstance(v, set):
+            notes[k] = sorted(str(x) for x in v)
+    return {"problems": problems, "ok": not problems, "notes": notes}

export/phonon-31m-en.py

@@ -0,0 +1,585 @@
+#!/usr/bin/env python3
+"""Phonon 31M English (gradium/phonon-31m-en, Apache-2.0) -> two ExecuTorch WebGPU programs + host
+tables, public ExecuTorch. Model code and per-program contract: recipes/_phonon.py.
+
+  phonon_flow_lm.pte  forward(embeds f32[1,S,512] (1 <= S <= 128), input_pos i64[1], noise f32[1,32])
+                      -> latent f32[1,32], eos_logit f32[1,1], next_embeds f32[1,1,512]
+                      6-layer flow-LM transformer with its KV cache inside the program (512 positions),
+                      the LSD flow head (1 step) and the next step's input embedding.
+  phonon_mimi.pte     forward(latent f32[1,1,32], 15 state tensors) -> pcm f32[1,1,1920], 15 new states
+                      one 80 ms frame of the streaming Mimi decoder; all state is explicit (zeros at start)
+  phonon_text_embed.bf16.bin   the text conditioner table [4001, 512], bf16 exactly as in the checkpoint
+  phonon_bos_embed.f32.bin     [512] the first step's input embedding
+  voice_<name>.f32.bin         [125, 512] voice prompts (speaker latents through the speaker projection)
+  tokenizer.json, phonon.json  tokenizer (upstream file) and the host contract (shapes, policy, state list)
+
+Every program is lowered with executorch.backends.webgpu.partitioner.WebGPUPartitioner() into one
+VulkanBackend delegate with no CPU kernels. The host does: text normalization + sentencepiece
+tokenization, the embedding-table lookup, the noise (N(0, temperature), 32 values per frame), the EOS
+rule (eos_logit > 4.0, then frames_after_eos more frames) and feeding next_embeds / Mimi state back.
+
+  python recipes/phonon-31m-en.py --out out/phonon-31m-en                  # export + CPU validation
+  python recipes/phonon-31m-en.py --out out/phonon-31m-en --no-validate    # export only
+  python recipes/phonon-31m-en.py --out out/phonon-31m-en --ptts           # also compare with the
+                                                                            # official runtime (pip install ptts==0.4.0)
+
+Validation (CPU only; the WebGPU kernels cannot run here): the eager reference vs the torch.export
+programs vs the portable .pte (ExecuTorch pybindings) over the full TTS loop with identical noise,
+and with --ptts the official runtime's audio at temperature 0 (zero noise) for the same sentences.
+"""
+
+import argparse
+import json
+import sys
+import time
+import wave
+from pathlib import Path
+
+import numpy as np
+import torch
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+import _common  # noqa: E402
+import _phonon as P  # noqa: E402
+import _whisper as W  # noqa: E402  (generic lowered-graph checks: vk_checks, program_io, isclose_report)
+
+UPSTREAM = {
+    "repo": "gradium/phonon-31m-en",
+    "revision": "52433c46975bfa13144c56204c9f4ee1c5b148aa",
+    "files": {
+        "model.safetensors": "8a04cf833cfe69b67c2ec968b2d31244655b9accc84ade9c0ebcc176f2074607",
+        "config.json": "7062ea376e6b158543c8abec8e5af67df171c49ad70eb7e2a0fb25c16f265e06",
+        "tokenizer.json": "691a00068fa7b989f2f3e5ae5daa2d1a5bb319f2d5360deafcfb19a8bb6e7936",
+        "tokenizer.model": "82e4572006d64d483317074c22c368e87779640b6f3a9dcd07b7caae32b030ce",
+        "default-voice.safetensors": "f739e8c4c76ee7f92e7c94a82002748c8cd22e99582ac49f7b9b94cd63786e18",
+        "README.md": "00f09ecdcd23ed9904f3a503baecc1f0c77515637ce6172b572ed3f50ce50c29",
+    },
+}
+VOICES = {
+    "repo": "gradium/phonon-voices",
+    "revision": "0ee17205f254b5abc43b266c8a0f215565c352f7",
+    "files": {
+        "voices/LMDIYE0Oa377Qh6Z.safetensors": "45c582673cb20edb04ddcca2b5a8865ddfb0a82599efadb269021e87d20fbdd4",
+        "voices.jsonl": "874f34998ce581b94f2b8912aa015639f52b5fd62eba3ba11e1b64dcaed7d802",
+        "README.md": "0186c7d586a9594db1e0283cb68e84e88bbd9bc3c29047cf0083c7483556f3ee",
+    },
+    # output name -> (repo, file); "default" is the model repo's own voice
+    "export": {
+        "default": ("model", "default-voice.safetensors"),
+        "en_us_easton": ("voices", "voices/LMDIYE0Oa377Qh6Z.safetensors"),
+    },
+}
+SENTENCES = [
+    "The quick brown fox jumps over the lazy dog.",
+    "This voice is generated by a small model running entirely on your device.",
+    "Please call me back when you have a moment, and we can talk about the plan.",
+]
+NOISE_SEED = 1000  # seeded case: noise[k] = randn(1, 32, Generator().manual_seed(NOISE_SEED + k)) * sqrt(temp)
+NOISE_TEMP = 0.3  # the official runtime's default temperature
+EOS_THRESHOLD = 4.0
+SAMPLE_RATE = 24000
+
+
+def fetch():
+    paths, digests = {}, {}
+    for f, sha in UPSTREAM["files"].items():
+        paths[f], digests[f"{UPSTREAM['repo']}:{f}"] = _common.hf_file(
+            UPSTREAM["repo"], f, UPSTREAM["revision"], sha
+        )
+    vpaths = {}
+    for f, sha in VOICES["files"].items():
+        vpaths[f], digests[f"{VOICES['repo']}:{f}"] = _common.hf_file(
+            VOICES["repo"], f, VOICES["revision"], sha
+        )
+    return paths, vpaths, digests
+
+
+def flow_example():
+    return (
+        torch.randn(1, 5, P.D, generator=torch.Generator().manual_seed(0)),
+        torch.tensor([0]),
+        torch.zeros(1, P.LDIM),
+    )
+
+
+def mimi_example():
+    return (
+        torch.randn(1, 1, P.LDIM, generator=torch.Generator().manual_seed(0)),
+        *P.mimi_zero_state(),
+    )
+
+
+def export_flow(w):
+    with torch.no_grad():
+        return torch.export.export(
+            P.FlowLMProgram(w).eval(),
+            flow_example(),
+            dynamic_shapes=(
+                {1: torch.export.Dim("seq", min=1, max=P.PREFILL_MAX)},
+                None,
+                None,
+            ),
+            strict=False,
+        )
+
+
+def export_mimi(w):
+    with torch.no_grad():
+        return torch.export.export(
+            P.MimiStepProgram(w).eval(), mimi_example(), strict=False
+        )
+
+
+def lower_webgpu(ep, dynamic):
+    from executorch.backends.webgpu.partitioner import WebGPUPartitioner
+    from executorch.exir import to_edge_transform_and_lower
+
+    opts = {"require_dynamic_shapes": True, "force_fp16": False} if dynamic else None
+    part = WebGPUPartitioner(compile_options=opts) if opts else WebGPUPartitioner()
+    return to_edge_transform_and_lower(ep, partitioner=[part]).to_executorch()
+
+
+def write_wav(path, pcm):
+    pcm16 = np.clip(
+        np.round(np.asarray(pcm, np.float64) * 32767), -32768, 32767
+    ).astype("<i2")
+    with wave.open(str(path), "wb") as f:
+        f.setnchannels(1)
+        f.setsampwidth(2)
+        f.setframerate(SAMPLE_RATE)
+        f.writeframes(pcm16.tobytes())
+
+
+# ---------------------------------------------------------------------------------------------
+# Runners: flow(embeds, pos, noise) -> [latent, eos, next]; mimi(latent, *state) -> [pcm, *state]
+# ---------------------------------------------------------------------------------------------
+
+
+class ModuleRunner:
+    def __init__(self, flow, mimi):
+        self.f, self.m = flow, mimi
+
+    def flow(self, e, p, n):
+        with torch.no_grad():
+            return [t.clone() for t in self.f(e, torch.tensor([p]), n)]
+
+    def mimi(self, *a):
+        with torch.no_grad():
+            return [t.clone() for t in self.m(*a)]
+
+
+class PteRunner:
+    def __init__(self, flow_bytes, mimi_bytes):
+        import executorch.extension.llm.custom_ops.custom_ops  # noqa: F401  (portable llama::sdpa_with_kv_cache)
+        from executorch.extension.pybindings.portable_lib import (
+            _load_for_executorch_from_buffer,
+        )
+
+        self.buffers = (flow_bytes, mimi_bytes)  # the loaded programs point into these
+        self.f = _load_for_executorch_from_buffer(flow_bytes)
+        self.m = _load_for_executorch_from_buffer(mimi_bytes)
+
+    def flow(self, e, p, n):
+        return [
+            t.clone()
+            for t in self.f.forward([e.contiguous(), torch.tensor([p]), n.contiguous()])
+        ]
+
+    def mimi(self, *a):
+        return [t.clone() for t in self.m.forward([t.contiguous() for t in a])]
+
+
+def tts(run, voice_embeds, text_embeds, bos_embed, tokens, frames_after_eos, noise_fn):
+    """The host loop a browser runs: voice prompt, text prompt, then step -> Mimi per frame."""
+    pos = 0
+    for e in (voice_embeds, text_embeds):
+        run.flow(e, pos, torch.zeros(1, P.LDIM))
+        pos += e.shape[1]
+    st = {"e": bos_embed, "pos": pos, "steps": []}
+
+    def step(prev, noise):
+        lat, eos, nxt = run.flow(st["e"], st["pos"], noise)
+        st["steps"].append(
+            {"in": st["e"], "noise": noise, "latent": lat, "eos": eos, "next": nxt}
+        )
+        st["pos"] += 1
+        st["e"] = nxt
+        return lat.view(1, 1, P.LDIM), eos
+
+    lats, logits = P.generate(step, tokens, frames_after_eos, noise_fn, EOS_THRESHOLD)
+    state, pcm = P.mimi_zero_state(), []
+    for lat in lats:
+        out = run.mimi(lat, *state)
+        pcm.append(out[0])
+        state = out[1:]
+    return (
+        torch.cat(lats, 1).numpy(),
+        np.array(logits),
+        torch.cat(pcm, 2).reshape(-1).numpy(),
+        st["steps"],
+    )
+
+
+def noise_fn(kind):
+    if kind == "zero":
+        return lambda k: torch.zeros(1, P.LDIM)
+    return lambda k: torch.randn(
+        1, P.LDIM, generator=torch.Generator().manual_seed(NOISE_SEED + k)
+    ) * (NOISE_TEMP**0.5)
+
+
+def eager_reference(w, voice_embeds, text_embeds, tokens, frames_after_eos, nfn):
+    lm = P.FlowLMRef(w)
+    lm.prompt(voice_embeds)
+    lm.prompt(text_embeds)
+    lats, logits = P.generate(lm.step, tokens, frames_after_eos, nfn, EOS_THRESHOLD)
+    mimi = P.MimiRef(w)
+    pcm = torch.cat([mimi.decode(lat) for lat in lats], 2).reshape(-1).numpy()
+    return torch.cat(lats, 1).numpy(), np.array(logits), pcm
+
+
+def write_browser_ref(
+    d,
+    sentence,
+    prepared,
+    tokens,
+    fae,
+    voice_embeds,
+    text_embeds,
+    bos_embed,
+    steps,
+    pcm,
+    noise,
+):
+    d.mkdir(parents=True, exist_ok=True)
+    f32 = lambda t: np.ascontiguousarray(np.asarray(t, np.float32))  # noqa: E731
+    f32(voice_embeds).tofile(d / "voice_embeds.f32")
+    f32(text_embeds).tofile(d / "text_embeds.f32")
+    f32(bos_embed).tofile(d / "bos_embed.f32")
+    np.stack([f32(s["noise"]).reshape(P.LDIM) for s in steps]).tofile(d / "noise.f32")
+    np.stack([f32(s["latent"]).reshape(P.LDIM) for s in steps]).tofile(
+        d / "ref_latents.f32"
+    )
+    np.array([float(s["eos"].reshape(-1)[0]) for s in steps], np.float32).tofile(
+        d / "ref_eos.f32"
+    )
+    np.stack([f32(s["next"]).reshape(P.D) for s in steps]).tofile(
+        d / "ref_next_embeds.f32"
+    )
+    f32(pcm).tofile(d / "ref_pcm.f32")
+    meta = {
+        "text": sentence,
+        "prepared": prepared,
+        "tokens": tokens,
+        "frames_after_eos": fae,
+        "frame_budget": P.frame_budget(len(tokens)),
+        "voice_len": int(voice_embeds.shape[1]),
+        "text_len": len(tokens),
+        "steps": len(steps),
+        "eos_threshold": EOS_THRESHOLD,
+        "noise": noise,
+        "pcm_samples": int(pcm.size),
+        "reference": "portable .pte on the ExecuTorch CPU runtime",
+        "mimi_state": [[n, list(s)] for n, s in P.MIMI_STATE],
+    }
+    (d / "meta.json").write_text(json.dumps(meta, indent=1) + "\n")
+
+
+def validate(args, out, w, tok, voices, bos_embed, eps, ptes):
+    t0 = time.time()
+    runners = {
+        "exported": ModuleRunner(eps["flow_lm"].module(), eps["mimi"].module()),
+        "portable_pte": PteRunner(ptes["flow_lm"], ptes["mimi"]),
+    }
+    report = {
+        "cpu_runtime": "executorch pybindings (portable_lib) from the pip wheel, with "
+        "executorch.extension.llm.custom_ops loaded (llama::sdpa_with_kv_cache kernel)",
+        "webgpu_pte_on_cpu": "not runnable: the pip wheel's pybindings have no WebGPU/Vulkan backend",
+        "reference": "_phonon.FlowLMRef / MimiRef: eager PyTorch that follows the official runtime "
+        "op by op (interleaved RoPE, nn.LayerNorm, 3-D conv / conv-transpose, Python-side streaming state)",
+        "cases": [],
+    }
+    oracle = None
+    if args.ptts:
+        import ptts
+
+        oracle = ptts.TTS(
+            f"hf://{UPSTREAM['repo']}@{UPSTREAM['revision']}",
+            quant="f32",
+            lang="en",
+            temperature=0.0,
+        )
+        report["ptts"] = {
+            "build_info": ptts.build_info(),
+            "quant": oracle.quant,
+            "temperature": 0.0,
+            "voice": "default",
+        }
+    venc = voices["default"]
+    cases = [(i, "zero") for i in range(len(SENTENCES))] + [(0, "seeded")]
+    for si, kind in cases:
+        ids, fae, prepared = P.tokenize(tok, SENTENCES[si])
+        temb = P.text_embedding(w.flow_lm, ids)
+        nfn = noise_fn(kind)
+        rl, re_, rpcm = eager_reference(w, venc, temb, ids, fae, nfn)
+        case = {
+            "sentence": si,
+            "text": SENTENCES[si],
+            "prepared": prepared,
+            "tokens": ids,
+            "noise": kind,
+            "reference_frames": int(rl.shape[1]),
+        }
+        if oracle is not None and kind == "zero":
+            opcm = oracle.synth(SENTENCES[si], voice="default")
+            case["ptts_frames"] = int(opcm.size // P.FRAME)
+            case["reference_vs_ptts"] = W.isclose_report(
+                rpcm[: opcm.size], opcm[: rpcm.size]
+            )
+        for rname, run in runners.items():
+            t = time.time()
+            lat, logits, pcm, steps = tts(run, venc, temb, bos_embed, ids, fae, nfn)
+            n = min(lat.shape[1], rl.shape[1])
+            r = {
+                "frames": int(lat.shape[1]),
+                "latents_vs_reference": W.isclose_report(lat[:, :n], rl[:, :n]),
+                "eos_logits_vs_reference": W.isclose_report(logits[:n], re_[:n]),
+                "pcm_vs_reference": W.isclose_report(
+                    pcm[: n * P.FRAME], rpcm[: n * P.FRAME]
+                ),
+                "seconds": round(time.time() - t, 1),
+            }
+            if oracle is not None and kind == "zero":
+                m = min(pcm.size, opcm.size)
+                r["pcm_vs_ptts"] = W.isclose_report(pcm[:m], opcm[:m])
+            case[rname] = r
+            if rname == "portable_pte":
+                if args.wavs:
+                    (out / "samples").mkdir(exist_ok=True)
+                    write_wav(out / "samples" / f"s{si}_{kind}_portable_pte.wav", pcm)
+                if args.browser_ref and (si, kind) in [(0, "seeded"), (2, "zero")]:
+                    noise_desc = (
+                        "zeros"
+                        if kind == "zero"
+                        else f"randn(1,32, Generator().manual_seed({NOISE_SEED}+k)) * sqrt({NOISE_TEMP})"
+                    )
+                    write_browser_ref(
+                        Path(args.browser_ref) / f"s{si}_{kind}",
+                        SENTENCES[si],
+                        prepared,
+                        ids,
+                        fae,
+                        venc.numpy(),
+                        temb.numpy(),
+                        bos_embed.numpy(),
+                        steps,
+                        pcm,
+                        noise_desc,
+                    )
+        report["cases"].append(case)
+        print(
+            f"[validate] s{si} {kind}: "
+            + json.dumps(
+                {k: case[k] for k in case if k in ("reference_frames", "ptts_frames")}
+                | {
+                    k: {
+                        "frames": v["frames"],
+                        "pcm_cos": v["pcm_vs_reference"]["cosine"],
+                    }
+                    for k, v in case.items()
+                    if isinstance(v, dict) and "frames" in v
+                }
+            ),
+            flush=True,
+        )
+    report["seconds"] = round(time.time() - t0, 1)
+    (out / "validation.json").write_text(json.dumps(report, indent=2) + "\n")
+
+
+def main():
+    ap = argparse.ArgumentParser(
+        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
+    )
+    ap.add_argument("--out", required=True, help="output directory")
+    ap.add_argument(
+        "--no-validate", action="store_true", help="export only (no CPU validation)"
+    )
+    ap.add_argument(
+        "--ptts",
+        action="store_true",
+        help="also compare with the official runtime (pip install ptts==0.4.0)",
+    )
+    ap.add_argument(
+        "--wavs",
+        action="store_true",
+        help="write the portable .pte audio to <out>/samples",
+    )
+    ap.add_argument(
+        "--browser-ref", help="directory for the browser harness's reference tensors"
+    )
+    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",
+    )
+    args = ap.parse_args()
+    out = Path(args.out)
+    out.mkdir(parents=True, exist_ok=True)
+    torch.manual_seed(0)
+    torch.set_grad_enabled(False)
+
+    paths, vpaths, digests = fetch()
+    w = P.load_weights(
+        paths["model.safetensors"]
+    )  # strict: every checkpoint tensor has a home
+    tok = P.tokenizer(paths["tokenizer.json"])
+    runtime_ops = None
+    if args.runtime_ops:
+        runtime_ops = {
+            line.strip()
+            for line in Path(args.runtime_ops).read_text().splitlines()
+            if line.strip()
+        }
+
+    files, graphs, checks, io, eps, ptes = {}, {}, {}, {}, {}, {}
+    for name, make, dyn in [
+        ("flow_lm", export_flow, True),
+        ("mimi", export_mimi, False),
+    ]:
+        t = time.time()
+        ep = make(w)
+        prog = lower_webgpu(ep, dyn)
+        path = out / f"phonon_{name}.pte"
+        path.unlink(missing_ok=True)
+        path.write_bytes(prog.buffer)
+        graphs[name] = _common.delegate_summary(prog.buffer)
+        checks[name] = W.vk_checks(prog.buffer, runtime_ops)
+        io[name] = W.program_io(prog.buffer)
+        files[path.name] = {
+            "bytes": path.stat().st_size,
+            "sha256": _common.sha256_file(path),
+        }
+        eps[name] = make(w)  # fresh program: lowering consumes its input
+        if not args.no_validate or args.keep_portable:
+            ptes[name] = _common.lower_portable(make(w)).buffer
+            if args.keep_portable:
+                (out / "portable").mkdir(exist_ok=True)
+                (out / "portable" / f"phonon_{name}_portable.pte").write_bytes(
+                    ptes[name]
+                )
+        del prog
+        print(
+            f"[export] {path.name}: {files[path.name]['bytes']} B, {time.time() - t:.1f} s, checks ok={checks[name]['ok']}, "
+            f"non-delegated={graphs[name]['non_delegated_kernels']}, delegates={len(graphs[name]['delegates'])}",
+            flush=True,
+        )
+
+    # Host tables.
+    from safetensors.torch import load_file
+
+    emb = load_file(str(paths["model.safetensors"]))[
+        "flow_lm.condition_provider.conditioners.transcript_in_segment.embed.weight"
+    ]
+    assert emb.dtype == torch.bfloat16 and tuple(emb.shape) == (4001, P.D)
+    host = {
+        "phonon_text_embed.bf16.bin": emb.contiguous()
+        .view(torch.int16)
+        .numpy()
+        .tobytes()
+    }
+    bos_embed = P.lin(
+        w.flow_lm.input_linear, w.flow_lm.bos_emb.view(1, 1, P.LDIM)
+    ) + P.summed_conditions(w.flow_lm)
+    host["phonon_bos_embed.f32.bin"] = bos_embed.numpy().astype("<f4").tobytes()
+    voices = {}
+    for vname, (src, f) in VOICES["export"].items():
+        sw = P.load_voice(paths[f] if src == "model" else vpaths[f])
+        voices[vname] = P.voice_embedding(w.flow_lm, sw)
+        host[f"voice_{vname}.f32.bin"] = voices[vname].numpy().astype("<f4").tobytes()
+    host["tokenizer.json"] = paths["tokenizer.json"].read_bytes()
+    for fname, blob in host.items():
+        (out / fname).unlink(missing_ok=True)
+        (out / fname).write_bytes(blob)
+        files[fname] = {"bytes": len(blob), "sha256": _common.sha256_file(out / fname)}
+
+    contract = {
+        "sample_rate": SAMPLE_RATE,
+        "frame_samples": P.FRAME,
+        "frame_rate": 12.5,
+        "max_context": P.MAX_CTX,
+        "prefill_max": P.PREFILL_MAX,
+        "programs": io,
+        "flow_lm": "forward(embeds, input_pos, noise): writes KV positions [input_pos, input_pos + S); outputs are for "
+        "the last position. Calls: voice prompt (S = 125) at 0, text prompt (S = #tokens) at 125, then one call per "
+        "frame with S = 1: the first takes phonon_bos_embed, each later one the previous call's next_embeds",
+        "mimi": "forward(latent, *state) -> (pcm, *state): pass the previous call's states back; zeros to start",
+        "mimi_state": [[n, list(s)] for n, s in P.MIMI_STATE],
+        "text": "normalize (the official runtime's English rules; identity for plain sentences), then trim, collapse "
+        "whitespace, capitalize the first letter, append '.' if the text ends with a letter or digit; encode with "
+        "tokenizer.json without special tokens; embeds[i] = bf16 table row tokens[i] (as f32)",
+        "noise": "32 values per frame, N(0, temperature); 0.3 is the official runtime's default, 0 makes it deterministic",
+        "eos": f"stop after frames_after_eos more frames once eos_logit > {EOS_THRESHOLD} (frames_after_eos = 3 for "
+        "texts of <= 4 words, else 1); the EOS frame and its tail are decoded",
+        "frame_budget": "ceil((#tokens / 3 + 2) * 12.5) frames at most",
+        "voices": {
+            k: {"file": f"voice_{k}.f32.bin", "shape": list(v.shape[1:])}
+            for k, v in voices.items()
+        },
+        "conditions": {
+            "num_speakers": "1",
+            "duration_delta": 0.0,
+            "note": "the runtime's defaults, folded into the programs",
+        },
+    }
+    (out / "phonon.json").write_text(json.dumps(contract, indent=2) + "\n")
+    files["phonon.json"] = {
+        "bytes": (out / "phonon.json").stat().st_size,
+        "sha256": _common.sha256_file(out / "phonon.json"),
+    }
+
+    recipe = {
+        "model": "phonon-31m-en",
+        "upstream": {
+            "model": UPSTREAM["repo"],
+            "revision": UPSTREAM["revision"],
+            "voices": {"repo": VOICES["repo"], "revision": VOICES["revision"]},
+            "official": True,
+            "license": "apache-2.0 (model card front matter and License section; phonon-voices card front matter)",
+            "files_sha256": digests,
+            "loader": "recipes/_phonon.py load_weights: model.safetensors (bf16 flow LM, f32 Mimi decoder) -> float32, "
+            "load_state_dict(strict=True) into a module whose keys equal the checkpoint's",
+        },
+        "versions": _common.versions(),
+        "export": {
+            "api": "torch.export.export(strict=False) -> to_edge_transform_and_lower(partitioner=[WebGPUPartitioner(...)]) -> to_executorch()",
+            "partitioner": {
+                "flow_lm": "WebGPUPartitioner(compile_options={'require_dynamic_shapes': True, 'force_fp16': False}), dynamic S in [1, 128]",
+                "mimi": "WebGPUPartitioner() (default compile options, static shapes)",
+            },
+            "dtype": "float32 (bf16 checkpoint weights converted exactly)",
+            "quantization": "none",
+            "model_code": "recipes/_phonon.py FlowLMProgram, MimiStepProgram",
+        },
+        "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, default=str) + "\n")
+    print(json.dumps({"files": files}, indent=2), flush=True)
+
+    if not args.no_validate:
+        validate(args, out, w, tok, voices, bos_embed, eps, ptes)
+
+
+if __name__ == "__main__":
+    main()

export/recipe.json

@@ -0,0 +1,205 @@
+{
+  "model": "phonon-31m-en",
+  "upstream": {
+    "model": "gradium/phonon-31m-en",
+    "revision": "52433c46975bfa13144c56204c9f4ee1c5b148aa",
+    "voices": {
+      "repo": "gradium/phonon-voices",
+      "revision": "0ee17205f254b5abc43b266c8a0f215565c352f7"
+    },
+    "official": true,
+    "license": "apache-2.0 (model card front matter and License section; phonon-voices card front matter)",
+    "files_sha256": {
+      "gradium/phonon-31m-en:model.safetensors": "8a04cf833cfe69b67c2ec968b2d31244655b9accc84ade9c0ebcc176f2074607",
+      "gradium/phonon-31m-en:config.json": "7062ea376e6b158543c8abec8e5af67df171c49ad70eb7e2a0fb25c16f265e06",
+      "gradium/phonon-31m-en:tokenizer.json": "691a00068fa7b989f2f3e5ae5daa2d1a5bb319f2d5360deafcfb19a8bb6e7936",
+      "gradium/phonon-31m-en:tokenizer.model": "82e4572006d64d483317074c22c368e87779640b6f3a9dcd07b7caae32b030ce",
+      "gradium/phonon-31m-en:default-voice.safetensors": "f739e8c4c76ee7f92e7c94a82002748c8cd22e99582ac49f7b9b94cd63786e18",
+      "gradium/phonon-31m-en:README.md": "00f09ecdcd23ed9904f3a503baecc1f0c77515637ce6172b572ed3f50ce50c29",
+      "gradium/phonon-voices:voices/LMDIYE0Oa377Qh6Z.safetensors": "45c582673cb20edb04ddcca2b5a8865ddfb0a82599efadb269021e87d20fbdd4",
+      "gradium/phonon-voices:voices.jsonl": "874f34998ce581b94f2b8912aa015639f52b5fd62eba3ba11e1b64dcaed7d802",
+      "gradium/phonon-voices:README.md": "0186c7d586a9594db1e0283cb68e84e88bbd9bc3c29047cf0083c7483556f3ee"
+    },
+    "loader": "recipes/_phonon.py load_weights: model.safetensors (bf16 flow LM, f32 Mimi decoder) -> float32, load_state_dict(strict=True) into a module whose keys equal the checkpoint's"
+  },
+  "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()",
+    "partitioner": {
+      "flow_lm": "WebGPUPartitioner(compile_options={'require_dynamic_shapes': True, 'force_fp16': False}), dynamic S in [1, 128]",
+      "mimi": "WebGPUPartitioner() (default compile options, static shapes)"
+    },
+    "dtype": "float32 (bf16 checkpoint weights converted exactly)",
+    "quantization": "none",
+    "model_code": "recipes/_phonon.py FlowLMProgram, MimiStepProgram"
+  },
+  "files": {
+    "phonon_flow_lm.pte": {
+      "bytes": 111847044,
+      "sha256": "178698a6c4539ba48c93bd74c4165586c1da4fdd393bd144c27fd09d80021ad7"
+    },
+    "phonon_mimi.pte": {
+      "bytes": 41304196,
+      "sha256": "d7acdf2d93ba4897e63b58d933c7425f49cf855ecbd1a6f51e273c5dd65bb3cc"
+    },
+    "phonon_text_embed.bf16.bin": {
+      "bytes": 4097024,
+      "sha256": "7c6f3394fc86b55ed181ff0b9d2472cf0371d3feeac25a5d0f2eed0a3a2eaae1"
+    },
+    "phonon_bos_embed.f32.bin": {
+      "bytes": 2048,
+      "sha256": "a4464c50596adea27c8f0457124ead461202dc2f67daf3c6dbbc289cf289fd32"
+    },
+    "voice_default.f32.bin": {
+      "bytes": 256000,
+      "sha256": "d8a8ffc97e46c0330ef6a3842130b95cc5ee71759f66bb54bc2ec0b79fc36e57"
+    },
+    "voice_en_us_easton.f32.bin": {
+      "bytes": 256000,
+      "sha256": "1d6ec67cbd3a45b2bacf84ff7ddfb184fd9e8467c2204f2eac0f4975c354b66b"
+    },
+    "tokenizer.json": {
+      "bytes": 277967,
+      "sha256": "691a00068fa7b989f2f3e5ae5daa2d1a5bb319f2d5360deafcfb19a8bb6e7936"
+    },
+    "phonon.json": {
+      "bytes": 8599,
+      "sha256": "e8fa6604769aeb3b63e95c44c912f7ae635a324ca3793ec62fca017346521750"
+    }
+  },
+  "graph": {
+    "flow_lm": {
+      "non_delegated_kernels": {},
+      "delegates": [
+        {
+          "id": "VulkanBackend",
+          "ops": {
+            "et_vk.prepack.default": 23,
+            "sym_size.int": 1,
+            "et_vk.select_as_symint.default": 1,
+            "aten.view_copy.default": 106,
+            "aten.linear.default": 60,
+            "aten.sub.Tensor": 12,
+            "et_vk.rms_norm.default": 12,
+            "aten.add.Tensor": 47,
+            "aten.slice_copy.Tensor": 38,
+            "et_vk.apply_rotary_emb_hf.default": 6,
+            "sdpa_with_kv_cache.default": 6,
+            "aten.gelu.default": 6,
+            "aten.select_copy.int": 1,
+            "aten.native_layer_norm.default": 8,
+            "aten.sigmoid.default": 7,
+            "aten.mul.Tensor": 20,
+            "aten.unsqueeze_copy.default": 1
+          },
+          "inputs": 3,
+          "outputs": 3
+        }
+      ]
+    },
+    "mimi": {
+      "non_delegated_kernels": {},
+      "delegates": [
+        {
+          "id": "VulkanBackend",
+          "ops": {
+            "et_vk.prepack.default": 34,
+            "aten.mul.Tensor": 18,
+            "aten.add.Tensor": 32,
+            "aten.view_copy.default": 30,
+            "aten.linear.default": 9,
+            "aten.clone.default": 14,
+            "aten.unsqueeze_copy.default": 9,
+            "aten.convolution.default": 12,
+            "aten.squeeze_copy.dims": 12,
+            "aten.slice_copy.Tensor": 35,
+            "aten.cat.default": 21,
+            "aten.cos.default": 2,
+            "aten.sin.default": 2,
+            "aten.permute_copy.default": 12,
+            "aten.native_layer_norm.default": 4,
+            "aten.neg.default": 4,
+            "aten.sub.Tensor": 15,
+            "aten.bmm.default": 4,
+            "aten._softmax.default": 2,
+            "aten.gelu.default": 2,
+            "aten.relu.default": 10,
+            "aten.clamp.default": 10,
+            "aten.exp.default": 10
+          },
+          "inputs": 16,
+          "outputs": 16
+        }
+      ]
+    }
+  },
+  "webgpu_runtime_checks": {
+    "op_list": "checked against --runtime-ops",
+    "flow_lm": {
+      "problems": [],
+      "ok": true,
+      "notes": {
+        "value_kinds": [
+          "Bool",
+          "Double",
+          "FLOAT32",
+          "INT32",
+          "Int",
+          "IntList",
+          "Null",
+          "String",
+          "SymInt",
+          "ValueList"
+        ],
+        "aten.linear input ranks": [
+          "2"
+        ],
+        "sdpa_with_kv_cache q/k_cache": [
+          "((1, 128, 8, 64), (1, 512, 8, 64), 'SymInt')"
+        ]
+      }
+    },
+    "mimi": {
+      "problems": [],
+      "ok": true,
+      "notes": {
+        "value_kinds": [
+          "Bool",
+          "Double",
+          "FLOAT32",
+          "Int",
+          "IntList",
+          "Null",
+          "String",
+          "ValueList"
+        ],
+        "aten.linear input ranks": [
+          "2"
+        ],
+        "aten.convolution input/weight": [
+          "((1, 128, 1, 480), (128, 64, 1, 8))",
+          "((1, 128, 482), (64, 128, 3))",
+          "((1, 128, 96), (256, 128, 1))",
+          "((1, 256, 1, 96), (256, 128, 1, 10))",
+          "((1, 256, 98), (128, 256, 3))",
+          "((1, 32, 1920), (64, 32, 1))",
+          "((1, 512, 1, 1), (512, 1, 1, 32))",
+          "((1, 512, 1, 16), (512, 256, 1, 12))",
+          "((1, 512, 22), (512, 512, 7))",
+          "((1, 64, 1922), (1, 64, 3))",
+          "((1, 64, 1922), (32, 64, 3))",
+          "((1, 64, 480), (128, 64, 1))"
+        ]
+      }
+    }
+  },
+  "command": "python recipes/phonon-31m-en.py --out out/phonon-31m-en"
+}

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

manifest.json

@@ -0,0 +1,46 @@
+{
+  "model": "phonon-31m-en",
+  "repo": "executorch-community/Phonon-31M-EN-ExecuTorch-WebGPU",
+  "files": [
+    {
+      "path": "phonon.json",
+      "bytes": 8599,
+      "sha256": "e8fa6604769aeb3b63e95c44c912f7ae635a324ca3793ec62fca017346521750"
+    },
+    {
+      "path": "phonon_bos_embed.f32.bin",
+      "bytes": 2048,
+      "sha256": "a4464c50596adea27c8f0457124ead461202dc2f67daf3c6dbbc289cf289fd32"
+    },
+    {
+      "path": "phonon_flow_lm.pte",
+      "bytes": 111847044,
+      "sha256": "178698a6c4539ba48c93bd74c4165586c1da4fdd393bd144c27fd09d80021ad7"
+    },
+    {
+      "path": "phonon_mimi.pte",
+      "bytes": 41304196,
+      "sha256": "d7acdf2d93ba4897e63b58d933c7425f49cf855ecbd1a6f51e273c5dd65bb3cc"
+    },
+    {
+      "path": "phonon_text_embed.bf16.bin",
+      "bytes": 4097024,
+      "sha256": "7c6f3394fc86b55ed181ff0b9d2472cf0371d3feeac25a5d0f2eed0a3a2eaae1"
+    },
+    {
+      "path": "tokenizer.json",
+      "bytes": 277967,
+      "sha256": "691a00068fa7b989f2f3e5ae5daa2d1a5bb319f2d5360deafcfb19a8bb6e7936"
+    },
+    {
+      "path": "voice_default.f32.bin",
+      "bytes": 256000,
+      "sha256": "d8a8ffc97e46c0330ef6a3842130b95cc5ee71759f66bb54bc2ec0b79fc36e57"
+    },
+    {
+      "path": "voice_en_us_easton.f32.bin",
+      "bytes": 256000,
+      "sha256": "1d6ec67cbd3a45b2bacf84ff7ddfb184fd9e8467c2204f2eac0f4975c354b66b"
+    }
+  ]
+}

phonon.json

@@ -0,0 +1,482 @@
+{
+  "sample_rate": 24000,
+  "frame_samples": 1920,
+  "frame_rate": 12.5,
+  "max_context": 512,
+  "prefill_max": 128,
+  "programs": {
+    "flow_lm": {
+      "forward": {
+        "inputs": [
+          {
+            "shape": [
+              1,
+              128,
+              512
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1
+            ],
+            "dtype": "long"
+          },
+          {
+            "shape": [
+              1,
+              32
+            ],
+            "dtype": "float"
+          }
+        ],
+        "outputs": [
+          {
+            "shape": [
+              1,
+              32
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1,
+              1
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1,
+              1,
+              512
+            ],
+            "dtype": "float"
+          }
+        ]
+      }
+    },
+    "mimi": {
+      "forward": {
+        "inputs": [
+          {
+            "shape": [
+              1,
+              1,
+              32
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1,
+              512,
+              16
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1,
+              8,
+              250,
+              64
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1,
+              8,
+              250,
+              64
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1,
+              8,
+              250,
+              64
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1,
+              8,
+              250,
+              64
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1,
+              1,
+              1,
+              250
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1,
+              512,
+              6
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1,
+              256,
+              6
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1,
+              256,
+              2
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1,
+              128,
+              5
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1,
+              128,
+              2
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1,
+              64,
+              4
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1,
+              64,
+              2
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1,
+              64,
+              2
+            ],
+            "dtype": "float"
+          }
+        ],
+        "outputs": [
+          {
+            "shape": [
+              1,
+              1,
+              1920
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1,
+              512,
+              16
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1,
+              8,
+              250,
+              64
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1,
+              8,
+              250,
+              64
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1,
+              8,
+              250,
+              64
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1,
+              8,
+              250,
+              64
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1,
+              1,
+              1,
+              250
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1,
+              512,
+              6
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1,
+              256,
+              6
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1,
+              256,
+              2
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1,
+              128,
+              5
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1,
+              128,
+              2
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1,
+              64,
+              4
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1,
+              64,
+              2
+            ],
+            "dtype": "float"
+          },
+          {
+            "shape": [
+              1,
+              64,
+              2
+            ],
+            "dtype": "float"
+          }
+        ]
+      }
+    }
+  },
+  "flow_lm": "forward(embeds, input_pos, noise): writes KV positions [input_pos, input_pos + S); outputs are for the last position. Calls: voice prompt (S = 125) at 0, text prompt (S = #tokens) at 125, then one call per frame with S = 1: the first takes phonon_bos_embed, each later one the previous call's next_embeds",
+  "mimi": "forward(latent, *state) -> (pcm, *state): pass the previous call's states back; zeros to start",
+  "mimi_state": [
+    [
+      "upsample_partial",
+      [
+        1,
+        512,
+        16
+      ]
+    ],
+    [
+      "k0",
+      [
+        1,
+        8,
+        250,
+        64
+      ]
+    ],
+    [
+      "v0",
+      [
+        1,
+        8,
+        250,
+        64
+      ]
+    ],
+    [
+      "k1",
+      [
+        1,
+        8,
+        250,
+        64
+      ]
+    ],
+    [
+      "v1",
+      [
+        1,
+        8,
+        250,
+        64
+      ]
+    ],
+    [
+      "kv_seen",
+      [
+        1,
+        1,
+        1,
+        250
+      ]
+    ],
+    [
+      "rope_offset",
+      [
+        1
+      ]
+    ],
+    [
+      "init_conv_prev",
+      [
+        1,
+        512,
+        6
+      ]
+    ],
+    [
+      "up1_partial",
+      [
+        1,
+        256,
+        6
+      ]
+    ],
+    [
+      "res1_prev",
+      [
+        1,
+        256,
+        2
+      ]
+    ],
+    [
+      "up2_partial",
+      [
+        1,
+        128,
+        5
+      ]
+    ],
+    [
+      "res2_prev",
+      [
+        1,
+        128,
+        2
+      ]
+    ],
+    [
+      "up3_partial",
+      [
+        1,
+        64,
+        4
+      ]
+    ],
+    [
+      "res3_prev",
+      [
+        1,
+        64,
+        2
+      ]
+    ],
+    [
+      "final_conv_prev",
+      [
+        1,
+        64,
+        2
+      ]
+    ]
+  ],
+  "text": "normalize (the official runtime's English rules; identity for plain sentences), then trim, collapse whitespace, capitalize the first letter, append '.' if the text ends with a letter or digit; encode with tokenizer.json without special tokens; embeds[i] = bf16 table row tokens[i] (as f32)",
+  "noise": "32 values per frame, N(0, temperature); 0.3 is the official runtime's default, 0 makes it deterministic",
+  "eos": "stop after frames_after_eos more frames once eos_logit > 4.0 (frames_after_eos = 3 for texts of <= 4 words, else 1); the EOS frame and its tail are decoded",
+  "frame_budget": "ceil((#tokens / 3 + 2) * 12.5) frames at most",
+  "voices": {
+    "default": {
+      "file": "voice_default.f32.bin",
+      "shape": [
+        125,
+        512
+      ]
+    },
+    "en_us_easton": {
+      "file": "voice_en_us_easton.f32.bin",
+      "shape": [
+        125,
+        512
+      ]
+    }
+  },
+  "conditions": {
+    "num_speakers": "1",
+    "duration_delta": 0.0,
+    "note": "the runtime's defaults, folded into the programs"
+  }
+}

phonon_bos_embed.f32.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a4464c50596adea27c8f0457124ead461202dc2f67daf3c6dbbc289cf289fd32
+size 2048

phonon_flow_lm.pte

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:178698a6c4539ba48c93bd74c4165586c1da4fdd393bd144c27fd09d80021ad7
+size 111847044

phonon_mimi.pte

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d7acdf2d93ba4897e63b58d933c7425f49cf855ecbd1a6f51e273c5dd65bb3cc
+size 41304196

phonon_text_embed.bf16.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7c6f3394fc86b55ed181ff0b9d2472cf0371d3feeac25a5d0f2eed0a3a2eaae1
+size 4097024

voice_default.f32.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d8a8ffc97e46c0330ef6a3842130b95cc5ee71759f66bb54bc2ec0b79fc36e57
+size 256000

voice_en_us_easton.f32.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1d6ec67cbd3a45b2bacf84ff7ddfb184fd9e8467c2204f2eac0f4975c354b66b
+size 256000