fix: consolidate metal support

.github/workflows/ci.yml

@@ -3,7 +3,7 @@ name: ci
 on:
   push:
   pull_request:
-  # closed-loop job is also registered here but only runs on manual dispatch.
+  # closed-loop job runs on pull_request (gate) and manual dispatch (see below).
   workflow_dispatch:
 
 jobs:
@@ -35,9 +35,9 @@ jobs:
   # -------------------------------------------------------------------------
   # closed-loop: full end-to-end transcript assertion.
   #
-  # Triggered ONLY via workflow_dispatch (Actions UI or `gh workflow run`).
-  # NOT triggered on push/PR — this job downloads nvidia/parakeet-tdt_ctc-110m
-  # (~440 MB via NeMo/HuggingFace Hub) and needs the nemo_toolkit Python env.
+  # Runs on pull_request (merge gate) and manual workflow_dispatch. Not on every
+  # push, since this job downloads nvidia/parakeet-tdt_ctc-110m (~440 MB via
+  # NeMo/HuggingFace Hub) and installs the nemo_toolkit Python env (~60 min).
   #
   # What it does:
   #   1. Check out the repo (submodules recursive).
@@ -56,7 +56,7 @@ jobs:
   #      Fail the job if it doesn't match.
   # -------------------------------------------------------------------------
   closed-loop:
-    if: github.event_name == 'workflow_dispatch'
+    if: github.event_name == 'pull_request' || github.event_name == 'workflow_dispatch'
     runs-on: ubuntu-latest
     timeout-minutes: 60
     steps:

AGENTS.md

@@ -38,6 +38,23 @@ state) + `<EOU>`/`<EOB>` timed events, exposed via `parakeet_capi_stream_*` and
 `parakeet-cli transcribe --stream`.  The streaming transcript matches NeMo's
 cache-aware streaming byte-for-byte.
 
+## Performance invariants (do not regress)
+
+These are measured wins. An agent "simplifying" them has caused real regressions
+before, so do not change them without an A/B benchmark that proves parity.
+
+- **Keep the persistent `ggml_gallocr`** in `src/backend.cpp`. Reusing one
+  allocator across the many tiny per-utterance graphs (no per-call alloc/free) is
+  the core throughput lever on CPU and GPU. Do NOT replace it with
+  `ggml_backend_sched` on the fast path: sched re-plans the graph split on every
+  call and regressed CUDA by 7-23% when it did. The scheduler is used ONLY as a
+  per-graph fallback, when the active GPU backend lacks a kernel for some op
+  (so the unsupported op can run on CPU); when every op is supported, the fast
+  gallocr path runs. If you think gallocr can go, you are about to reintroduce
+  that regression.
+- **Zero-copy weights.** `clone_weight` returns loader tensors directly so the
+  same device buffer is reused every utterance; do not copy weights per call.
+
 ## Repository layout
 
 ```
@@ -97,7 +114,7 @@ docs/
   quantization.md   , quantization allowlist, policy, measured size + WER per type
   parity.md         , full model coverage matrix + per-stage tensor parity
 .github/workflows/
-  ci.yml            , build job (per-push) + closed-loop job (dispatch-only)
+  ci.yml            , build job (per-push) + closed-loop job (pull_request + dispatch)
 ```
 
 ## Build
@@ -305,17 +322,18 @@ See `models/MANIFEST.md` for the expected set of published GGUFs per checkpoint.
 
 `.github/workflows/ci.yml` has two jobs:
 
-1. **build** (runs on every push + pull_request): cmake configure + build +
-   `ctest -LE model`.  Fast (no model, no Python env needed).  This is the
-   per-push gate.
-
-2. **closed-loop** (runs ONLY on `workflow_dispatch`, manual trigger): sets up
-   the Python venv, installs CPU torch + `scripts/requirements.txt`, builds the
-   project, converts `nvidia/parakeet-tdt_ctc-110m` to GGUF (~440 MB download),
-   runs `parakeet-cli transcribe --decoder tdt` on `tests/fixtures/speech.wav`,
-   and asserts the output is byte-for-byte identical to the committed NeMo TDT
-   reference transcript.  Fails the job on any mismatch.  Not triggered on push
-   because it needs network + model.
+1. **build** (every push + pull_request): cmake build + `ctest -LE model`. Fast.
+2. **closed-loop** (pull_request + `workflow_dispatch`): converts the 110m
+   checkpoint and asserts `parakeet-cli transcribe --decoder tdt` matches the
+   reference transcript below. Heavy (NeMo download, ~60 min); not on every push.
+
+### Reference transcript
+
+`tests/fixtures/speech.wav` on the 110m TDT head decodes (WER 0.0 vs NeMo) to
+exactly the following. This is the closed-loop assertion and the quickest smoke
+test that a build is correct on any backend (CPU, Metal, CUDA):
+
+> Well, I don't wish to see it any more, observed Phoebe, turning away her eyes. It is certainly very like the old portrait.
 
 ## GGUF schema
 

README.md

@@ -82,6 +82,14 @@ cmake --build build-shared -j
 | `PARAKEET_GGML_VULKAN`   | OFF     | Forward GGML_VULKAN to the submodule       |
 | `PARAKEET_GGML_HIP`      | OFF     | Forward GGML_HIP (ROCm) to the submodule   |
 
+To build for a GPU backend, forward its flag, e.g. Apple Metal:
+
+```sh
+cmake -B build -DPARAKEET_GGML_METAL=ON && cmake --build build -j
+```
+
+The CLI auto-selects the first GPU device the ggml registry reports, so no runtime flag is needed (set `PARAKEET_DEVICE=cpu` to force CPU). Ops the chosen backend has no kernel for run on the CPU automatically, so a model always runs even when one op lacks a GPU kernel. On an Apple M4, Metal is up to about 5x faster than CPU on the larger models; see [Apple Metal](benchmarks/BENCHMARK.md#apple-metal-m4).
+
 ---
 
 ## Python environment setup

benchmarks/BENCHMARK.md

@@ -116,6 +116,41 @@ Decode batching coalesces the per-step prediction-LSTM and joint-network GEMMs a
 | tdt_ctc-1.1b | 1.00× | 3.45× | 6.50× | 11.47× | 746.9 |
 | tdt_ctc-110m | 1.01× | 3.29× | 6.33× | 10.89× | 1259.7 |
 
+## Apple Metal (M4)
+
+parakeet.cpp runs on Apple GPUs through ggml's Metal backend
+(`-DPARAKEET_GGML_METAL=ON`). The encoder runs on the GPU; ops without a Metal
+kernel fall back to CPU automatically, so any model runs.
+
+Apple M4, macOS 26.5. q4_k weights, one 7.4 s clip, best-of-6 steady-state runs
+(model load and one-time pipeline compilation excluded), via `parakeet-cli bench`.
+RTFx is audio seconds per second of compute; higher is faster.
+
+| Model | Metal | CPU | Speedup |
+|---|---|---|---|
+| ctc-0.6b | 91.1 | 16.9 | 5.4× |
+| ctc-1.1b | 56.2 | 10.0 | 5.6× |
+| parakeet-tdt_ctc-110m | 99.4 | 65.5 | 1.5× |
+| realtime_eou_120m-v1 | 76.5 | 58.6 | 1.3× |
+| rnnt-0.6b | 53.0 | 20.3 | 2.6× |
+| rnnt-1.1b | 39.4 | 12.4 | 3.2× |
+| tdt-0.6b-v2 | 57.1 | 20.1 | 2.8× |
+| tdt-0.6b-v3 | 54.3 | 19.8 | 2.7× |
+| tdt-1.1b | 43.8 | 12.6 | 3.5× |
+| tdt_ctc-110m | 103.1 | 65.9 | 1.6× |
+| tdt_ctc-1.1b | 41.9 | 12.5 | 3.4× |
+
+Metal helps most on the larger models (about 3× to 5×). The 110m and 120m models
+stay closer to CPU because their short encoders are dominated by per-run fixed
+costs. Two ops the FastConformer encoder needs ship as native Metal kernels here:
+the subsampling depthwise conv (`CONV_2D_DW`) and the conv module's leading-side
+time padding (`PAD`); with these the encoder layers run entirely on the GPU. The
+only work left on the CPU is the one-time log-mel front end: its ops all have
+Metal kernels, but the scheduler keeps that small, once-per-clip subgraph (a
+couple of matmuls plus a log over host-resident inputs) on the CPU rather than
+uploading it, which for its size is the faster choice. Reproduce with
+`scripts/bench_metal_dw.sh <model.gguf> <clip.wav>`.
+
 ## Plots
 
 ### RTFx per model — NeMo vs ours (all dtypes), LibriSpeech

scripts/bench_metal_dw.sh

@@ -0,0 +1,38 @@
+#!/usr/bin/env bash
+# Usage: bench_metal_dw.sh <model.gguf> <input.wav> [reps] [decoder]
+#
+# Measures steady-state inference speed (RTFx = audio_sec / proc_sec, higher is
+# faster) and prints one markdown table row:
+#   Metal (GPU) | CPU-only
+#
+# It drives `parakeet-cli bench`, which loads the model once, WARMS UP once
+# (untimed, so the one-time Metal pipeline JIT + weight upload is excluded), then
+# times ONLY transcription per manifest entry. We list the clip `reps` times and
+# take the best (min proc_ms) to get a stable steady-state number, isolating
+# inference from process-startup noise.
+set -euo pipefail
+MODEL="$1"; WAV="$2"; REPS="${3:-6}"; DEC="${4:-}"
+CLI=./build/examples/cli/parakeet-cli
+
+MAN=$(mktemp)
+trap 'rm -f "$MAN"' EXIT
+for _ in $(seq "$REPS"); do echo "$WAV" >> "$MAN"; done
+decarg=(); [ -n "$DEC" ] && decarg=(--decoder "$DEC")
+
+run() {  # echo best RTFx for the given env prefix; empty string on failure
+  local prefix="$1" json
+  # ${arr[@]+"${arr[@]}"} expands safely for an empty array under `set -u` on
+  # bash 3.2 (macOS default), where a bare "${arr[@]}" is an unbound-variable error.
+  json=$(env $prefix "$CLI" bench --model "$MODEL" --manifest "$MAN" ${decarg[@]+"${decarg[@]}"} 2>/dev/null) || { echo ""; return; }
+  printf '%s' "$json" | python3 -c '
+import sys, json
+d = json.load(sys.stdin)
+ms  = [f["proc_ms"] for f in d["files"]]
+sec = d["files"][0]["audio_sec"]
+print(f"{sec/(min(ms)/1000.0):.1f}")
+' 2>/dev/null || echo ""
+}
+
+metal=$(run "")
+cpu=$(run "PARAKEET_DEVICE=cpu")
+printf "| %s | %s | %s |\n" "$(basename "$MODEL")" "${metal:-ERR}" "${cpu:-ERR}"

src/backend.cpp

@@ -38,8 +38,11 @@ struct PendingCapture {
 } // namespace
 
 struct Backend::Impl {
-    ggml_backend_t    backend  = nullptr;
-    ggml_gallocr_t    galloc   = nullptr;   // created lazily, reused across calls
+    ggml_backend_t       backend     = nullptr;  // primary device (GPU or CPU)
+    ggml_backend_t       cpu_backend = nullptr;  // fallback backend (GPU path only)
+    ggml_gallocr_t       galloc      = nullptr;  // CPU / single-backend path (unchanged)
+    ggml_backend_sched_t sched       = nullptr;  // GPU path: schedules over {backend, cpu_backend}
+    bool                 use_sched   = false;    // true only when `backend` is a GPU device
     // Inputs registered by the build lambda for the IN-FLIGHT compute. Copied
     // into the gallocr-allocated tensors after ggml_gallocr_alloc_graph, then
     // cleared. Never overlaps across calls (compute is not re-entrant).
@@ -71,6 +74,7 @@ Backend::Backend(int n_threads) : impl_(new Impl()) {
                 impl_->backend = ggml_backend_dev_init(dev, nullptr);
                 if (impl_->backend) {
                     device_name_ = ggml_backend_dev_name(dev);
+                    impl_->use_sched = true;   // GPU device: route compute through ggml_backend_sched
                     PK_LOG("pk::Backend using GPU device: %s", device_name_.c_str());
                     break;
                 }
@@ -85,16 +89,27 @@ Backend::Backend(int n_threads) : impl_(new Impl()) {
         PK_LOG("backend init returned null");
         return;
     }
+    // GPU path: create a CPU fallback backend so unsupported ops (e.g. CONV_2D_DW,
+    // which ggml's Metal backend lacks) are offloaded to CPU by the scheduler
+    // instead of aborting. The CPU/single-backend path keeps using the persistent
+    // gallocr below and is untouched.
+    if (impl_->use_sched) {
+        impl_->cpu_backend = ggml_backend_cpu_init();
+        if (!impl_->cpu_backend) {
+            PK_LOG("pk::Backend: CPU fallback init failed; disabling sched");
+            impl_->use_sched = false;
+        }
+    }
     set_n_threads(n_threads);
 }
 
 Backend::~Backend() {
     if (impl_) {
-        // Free the gallocr BEFORE the backend: the gallocr owns the compute
-        // scratch buffer (allocated via the backend's buffer_type). Matches
-        // rt-detr's teardown order.
-        if (impl_->galloc) ggml_gallocr_free(impl_->galloc);
-        if (impl_->backend) ggml_backend_free(impl_->backend);
+        // Free allocators/scheduler BEFORE the backends they reference.
+        if (impl_->sched)       ggml_backend_sched_free(impl_->sched);
+        if (impl_->galloc)      ggml_gallocr_free(impl_->galloc);
+        if (impl_->cpu_backend) ggml_backend_free(impl_->cpu_backend);
+        if (impl_->backend)     ggml_backend_free(impl_->backend);
         delete impl_;
         impl_ = nullptr;
     }
@@ -105,6 +120,9 @@ void Backend::set_n_threads(int n_threads) {
     if (impl_ && impl_->backend && ggml_backend_is_cpu(impl_->backend)) {
         ggml_backend_cpu_set_n_threads(impl_->backend, n_threads_);
     }
+    if (impl_ && impl_->cpu_backend) {
+        ggml_backend_cpu_set_n_threads(impl_->cpu_backend, n_threads_);
+    }
 }
 
 ggml_backend_t Backend::handle() const {
@@ -168,22 +186,64 @@ bool Backend::compute(const std::function<ggml_tensor*(ggml_context*)>& build,
         ggml_build_forward_expand(gf, pc.tensor);
     ggml_build_forward_expand(gf, output);
 
-    // Lazily create the persistent gallocr (reused on every subsequent call; it
-    // only reallocates the underlying buffer when the graph grows beyond the
-    // current high-water mark).
-    if (!impl_->galloc) {
-        impl_->galloc = ggml_gallocr_new(
-            ggml_backend_get_default_buffer_type(impl_->backend));
+    // GPU devices default to the fast persistent-gallocr path (identical to a
+    // single-backend run). Only route THIS graph through the scheduler (which
+    // offloads unsupported ops to CPU) when the GPU backend actually lacks a
+    // kernel for one of its ops. CUDA covers every op parakeet uses, so it stays
+    // on gallocr with zero scheduler overhead; Metal likewise once its kernels
+    // are present; a genuinely missing op still degrades gracefully to CPU. The
+    // per-graph check is a cheap O(nodes) scan, far less than a sched re-plan.
+    bool need_sched = false;
+    if (impl_->use_sched) {
+        const int n_nodes = ggml_graph_n_nodes(gf);
+        for (int i = 0; i < n_nodes; ++i) {
+            if (!ggml_backend_supports_op(impl_->backend, ggml_graph_node(gf, i))) {
+                need_sched = true;
+                break;
+            }
+        }
+    }
+
+    bool alloc_ok = false;
+    if (need_sched) {
+        // GPU path: schedule across {GPU, CPU}. Unsupported ops fall back to CPU.
+        if (!impl_->sched) {
+            ggml_backend_t backs[2] = { impl_->backend, impl_->cpu_backend };
+            impl_->sched = ggml_backend_sched_new(
+                backs, /*bufts=*/nullptr, /*n_backends=*/2,
+                /*graph_size=*/kGraphSize, /*parallel=*/false, /*op_offload=*/true);
+            if (!impl_->sched) {
+                PK_LOG("Backend::compute: ggml_backend_sched_new failed");
+                impl_->pending.clear();
+                impl_->captures.clear();
+                ggml_free(ctx);
+                return false;
+            }
+        }
+        ggml_backend_sched_reset(impl_->sched);
+        alloc_ok = ggml_backend_sched_alloc_graph(impl_->sched, gf);
+        if (!alloc_ok) PK_LOG("Backend::compute: ggml_backend_sched_alloc_graph failed");
+    } else {
+        // Fast path: CPU, or a GPU whose backend supports every op in this graph.
+        // Persistent gallocr over the active backend's buffer type, lazily created
+        // and reused on every subsequent call (it only reallocates the underlying
+        // buffer when the graph grows beyond the current high-water mark). This is
+        // the original single-backend path; weights stay zero-copy on the device.
         if (!impl_->galloc) {
-            PK_LOG("Backend::compute: ggml_gallocr_new failed");
-            impl_->pending.clear();
-            impl_->captures.clear();
-            ggml_free(ctx);
-            return false;
+            impl_->galloc = ggml_gallocr_new(
+                ggml_backend_get_default_buffer_type(impl_->backend));
+            if (!impl_->galloc) {
+                PK_LOG("Backend::compute: ggml_gallocr_new failed");
+                impl_->pending.clear();
+                impl_->captures.clear();
+                ggml_free(ctx);
+                return false;
+            }
         }
+        alloc_ok = ggml_gallocr_alloc_graph(impl_->galloc, gf);
+        if (!alloc_ok) PK_LOG("Backend::compute: ggml_gallocr_alloc_graph failed");
     }
-    if (!ggml_gallocr_alloc_graph(impl_->galloc, gf)) {
-        PK_LOG("Backend::compute: ggml_gallocr_alloc_graph failed");
+    if (!alloc_ok) {
         impl_->pending.clear();
         impl_->captures.clear();
         ggml_free(ctx);
@@ -196,7 +256,9 @@ bool Backend::compute(const std::function<ggml_tensor*(ggml_context*)>& build,
     }
     impl_->pending.clear();
 
-    enum ggml_status status = ggml_backend_graph_compute(impl_->backend, gf);
+    enum ggml_status status = need_sched
+        ? ggml_backend_sched_graph_compute(impl_->sched, gf)
+        : ggml_backend_graph_compute(impl_->backend, gf);
     if (status != GGML_STATUS_SUCCESS) {
         PK_LOG("Backend::compute: ggml_backend_graph_compute failed (status=%d)",
                (int)status);