diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index ef1ea2d..9537228 100644 --- a/.github/workflows/ci.yml +++ b/.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: diff --git a/AGENTS.md b/AGENTS.md index 0459125..69cf90e 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -38,6 +38,23 @@ state) + ``/`` 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 diff --git a/README.md b/README.md index edc1821..1a9cc3b 100644 --- a/README.md +++ b/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 diff --git a/benchmarks/BENCHMARK.md b/benchmarks/BENCHMARK.md index 3017b5c..cc7187e 100644 --- a/benchmarks/BENCHMARK.md +++ b/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 `. + ## Plots ### RTFx per model — NeMo vs ours (all dtypes), LibriSpeech diff --git a/scripts/bench_metal_dw.sh b/scripts/bench_metal_dw.sh new file mode 100755 index 0000000..9f6f131 --- /dev/null +++ b/scripts/bench_metal_dw.sh @@ -0,0 +1,38 @@ +#!/usr/bin/env bash +# Usage: bench_metal_dw.sh [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}" diff --git a/src/backend.cpp b/src/backend.cpp index 2dada59..3c25409 100644 --- a/src/backend.cpp +++ b/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& 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& 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); diff --git a/third_party/ggml-patches/0002-metal-conv-2d-dw.patch b/third_party/ggml-patches/0002-metal-conv-2d-dw.patch new file mode 100644 index 0000000..236ad09 --- /dev/null +++ b/third_party/ggml-patches/0002-metal-conv-2d-dw.patch @@ -0,0 +1,282 @@ +diff --git a/src/ggml-metal/ggml-metal-device.cpp b/src/ggml-metal/ggml-metal-device.cpp +index ba006d9b..fe8b6a5c 100644 +--- a/src/ggml-metal/ggml-metal-device.cpp ++++ b/src/ggml-metal/ggml-metal-device.cpp +@@ -1818,6 +1818,29 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_2d(ggml_met + return res; + } + ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_2d_dw(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_CONV_2D_DW); ++ ++ GGML_ASSERT(ggml_is_contiguous(op->src[0])); ++ GGML_ASSERT(ggml_is_contiguous(op->src[1])); ++ GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ char base[256]; ++ char name[256]; ++ ++ snprintf(base, 256, "kernel_conv_2d_dw_f32_f32"); ++ snprintf(name, 256, "%s", base); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ return res; ++} ++ + ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_3d(ggml_metal_library_t lib, const ggml_tensor * op) { + assert(op->op == GGML_OP_CONV_3D); + +diff --git a/src/ggml-metal/ggml-metal-device.h b/src/ggml-metal/ggml-metal-device.h +index 1f212a92..c1130431 100644 +--- a/src/ggml-metal/ggml-metal-device.h ++++ b/src/ggml-metal/ggml-metal-device.h +@@ -150,6 +150,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_im2col + struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_transpose_1d (ggml_metal_library_t lib, const struct ggml_tensor * op); + struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_transpose_2d (ggml_metal_library_t lib, const struct ggml_tensor * op); + struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_2d (ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_2d_dw (ggml_metal_library_t lib, const struct ggml_tensor * op); + struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_3d (ggml_metal_library_t lib, const struct ggml_tensor * op); + struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_upscale (ggml_metal_library_t lib, const struct ggml_tensor * op); + struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_pad (ggml_metal_library_t lib, const struct ggml_tensor * op); +diff --git a/src/ggml-metal/ggml-metal-device.m b/src/ggml-metal/ggml-metal-device.m +index 780dfe81..1c9960a8 100644 +--- a/src/ggml-metal/ggml-metal-device.m ++++ b/src/ggml-metal/ggml-metal-device.m +@@ -1122,6 +1122,14 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te + op->src[1]->type == GGML_TYPE_F32 && + op->type == GGML_TYPE_F32 && + (op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_F32); ++ case GGML_OP_CONV_2D_DW: ++ // Depthwise conv2d, WHCN/F32 only (what the Conformer subsampling ++ // uses). CWHN and F16 fall through to the CPU backend. ++ return ggml_is_contiguous(op->src[0]) && ++ ggml_is_contiguous(op->src[1]) && ++ op->src[0]->type == GGML_TYPE_F32 && ++ op->src[1]->type == GGML_TYPE_F32 && ++ op->type == GGML_TYPE_F32; + case GGML_OP_UPSCALE: + return op->src[0]->type == GGML_TYPE_F32; + case GGML_OP_POOL_1D: +diff --git a/src/ggml-metal/ggml-metal-impl.h b/src/ggml-metal/ggml-metal-impl.h +index ff74cafb..d7b4b60a 100644 +--- a/src/ggml-metal/ggml-metal-impl.h ++++ b/src/ggml-metal/ggml-metal-impl.h +@@ -646,6 +646,34 @@ typedef struct { + int32_t d1; + } ggml_metal_kargs_conv_2d; + ++typedef struct { ++ uint64_t nbk0; // kernel nb[0] (KW stride) ++ uint64_t nbk1; // kernel nb[1] (KH stride) ++ uint64_t nbkc; // kernel nb[3] (channel stride; kernel is [KW,KH,1,C]) ++ uint64_t nb10; // src nb[0] ++ uint64_t nb11; // src nb[1] ++ uint64_t nb12; // src nb[2] (channel) ++ uint64_t nb13; // src nb[3] (batch) ++ uint64_t nb0; // dst nb[0] ++ uint64_t nb1; // dst nb[1] ++ uint64_t nb2; // dst nb[2] (channel) ++ uint64_t nb3; // dst nb[3] (batch) ++ int32_t IW; ++ int32_t IH; ++ int32_t KW; ++ int32_t KH; ++ int32_t C; ++ int32_t OW; ++ int32_t OH; ++ int32_t N; ++ int32_t s0; ++ int32_t s1; ++ int32_t p0; ++ int32_t p1; ++ int32_t d0; ++ int32_t d1; ++} ggml_metal_kargs_conv_2d_dw; ++ + typedef struct { + uint64_t ofs0; + uint64_t ofs1; +diff --git a/src/ggml-metal/ggml-metal-ops.cpp b/src/ggml-metal/ggml-metal-ops.cpp +index 206af227..6035ba2a 100644 +--- a/src/ggml-metal/ggml-metal-ops.cpp ++++ b/src/ggml-metal/ggml-metal-ops.cpp +@@ -386,6 +386,10 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) { + { + n_fuse = ggml_metal_op_conv_2d(ctx, idx); + } break; ++ case GGML_OP_CONV_2D_DW: ++ { ++ n_fuse = ggml_metal_op_conv_2d_dw(ctx, idx); ++ } break; + case GGML_OP_CONV_TRANSPOSE_1D: + { + n_fuse = ggml_metal_op_conv_transpose_1d(ctx, idx); +@@ -3727,6 +3731,83 @@ int ggml_metal_op_conv_2d(ggml_metal_op_t ctx, int idx) { + return 1; + } + ++int ggml_metal_op_conv_2d_dw(ggml_metal_op_t ctx, int idx) { ++ ggml_tensor * op = ctx->node(idx); ++ ++ ggml_metal_library_t lib = ctx->lib; ++ ggml_metal_encoder_t enc = ctx->enc; ++ ++ GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); // kernel: KW,KH,1,C ++ GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); // src: IW,IH,C,N ++ GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne, op, ne); // dst: OW,OH,C,N ++ GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); ++ ++ GGML_ASSERT(ggml_is_contiguous(op->src[0])); ++ GGML_ASSERT(ggml_is_contiguous(op->src[1])); ++ GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ const int32_t s0 = ((const int32_t *) op->op_params)[0]; ++ const int32_t s1 = ((const int32_t *) op->op_params)[1]; ++ const int32_t p0 = ((const int32_t *) op->op_params)[2]; ++ const int32_t p1 = ((const int32_t *) op->op_params)[3]; ++ const int32_t d0 = ((const int32_t *) op->op_params)[4]; ++ const int32_t d1 = ((const int32_t *) op->op_params)[5]; ++ ++ ggml_metal_kargs_conv_2d_dw args = { ++ /*.nbk0 =*/ nb00, ++ /*.nbk1 =*/ nb01, ++ /*.nbkc =*/ nb03, ++ /*.nb10 =*/ nb10, ++ /*.nb11 =*/ nb11, ++ /*.nb12 =*/ nb12, ++ /*.nb13 =*/ nb13, ++ /*.nb0 =*/ nb0, ++ /*.nb1 =*/ nb1, ++ /*.nb2 =*/ nb2, ++ /*.nb3 =*/ nb3, ++ /*.IW =*/ ne10, ++ /*.IH =*/ ne11, ++ /*.KW =*/ ne00, ++ /*.KH =*/ ne01, ++ /*.C =*/ ne12, ++ /*.OW =*/ ne0, ++ /*.OH =*/ ne1, ++ /*.N =*/ ne3, ++ /*.s0 =*/ s0, ++ /*.s1 =*/ s1, ++ /*.p0 =*/ p0, ++ /*.p1 =*/ p1, ++ /*.d0 =*/ d0, ++ /*.d1 =*/ d1, ++ }; ++ ++ auto pipeline = ggml_metal_library_get_pipeline_conv_2d_dw(lib, op); ++ ++ int nth = ggml_metal_pipeline_max_theads_per_threadgroup(pipeline); ++ nth = std::min(nth, 256); ++ nth = std::max(nth, 1); ++ ++ const uint64_t n_out = ggml_nelements(op); ++ ++ uint64_t tg = (n_out + nth - 1)/nth; ++ tg = std::max(tg, 1); ++ tg = std::min(tg, (uint64_t) std::numeric_limits::max()); ++ ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); ++ ++ ggml_metal_encoder_dispatch_threadgroups(enc, tg, 1, 1, nth, 1, 1); ++ ++ return 1; ++} ++ + int ggml_metal_op_conv_3d(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + +diff --git a/src/ggml-metal/ggml-metal-ops.h b/src/ggml-metal/ggml-metal-ops.h +index 36c61071..4b6195a1 100644 +--- a/src/ggml-metal/ggml-metal-ops.h ++++ b/src/ggml-metal/ggml-metal-ops.h +@@ -75,6 +75,7 @@ int ggml_metal_op_norm (ggml_metal_op_t ctx, int idx); + int ggml_metal_op_rope (ggml_metal_op_t ctx, int idx); + int ggml_metal_op_im2col (ggml_metal_op_t ctx, int idx); + int ggml_metal_op_conv_2d (ggml_metal_op_t ctx, int idx); ++int ggml_metal_op_conv_2d_dw (ggml_metal_op_t ctx, int idx); + int ggml_metal_op_conv_3d (ggml_metal_op_t ctx, int idx); + int ggml_metal_op_conv_transpose_1d (ggml_metal_op_t ctx, int idx); + int ggml_metal_op_conv_transpose_2d (ggml_metal_op_t ctx, int idx); +diff --git a/src/ggml-metal/ggml-metal.metal b/src/ggml-metal/ggml-metal.metal +index e772664b..7c4fae05 100644 +--- a/src/ggml-metal/ggml-metal.metal ++++ b/src/ggml-metal/ggml-metal.metal +@@ -4864,6 +4864,62 @@ kernel void kernel_conv_2d( + uint3 tpitg[[thread_position_in_threadgroup]], + uint3 ntg[[threads_per_threadgroup]]); + ++kernel void kernel_conv_2d_dw_f32_f32( ++ constant ggml_metal_kargs_conv_2d_dw & args, ++ device const char * weights, ++ device const char * src, ++ device char * dst, ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ uint3 tgpg[[threadgroups_per_grid]], ++ uint3 tpitg[[thread_position_in_threadgroup]], ++ uint3 ntg[[threads_per_threadgroup]]) { ++ ++ const uint threads_per_tg = ntg.x * ntg.y * ntg.z; ++ const uint tg_index = (tgpig.z * tgpg.y + tgpig.y) * tgpg.x + tgpig.x; ++ const uint local_thread = tpitg.z * (ntg.x * ntg.y) + tpitg.y * ntg.x + tpitg.x; ++ const uint thread_index = tg_index * threads_per_tg + local_thread; ++ const uint64_t total_threads = (uint64_t) threads_per_tg * tgpg.x * tgpg.y * tgpg.z; ++ const uint64_t total_outputs = (uint64_t) args.N * args.C * args.OH * args.OW; ++ ++ for (uint64_t index = thread_index; index < total_outputs; index += total_threads) { ++ uint64_t tmp = index; ++ ++ const int32_t ow = tmp % args.OW; tmp /= args.OW; ++ const int32_t oh = tmp % args.OH; tmp /= args.OH; ++ const int32_t c = tmp % args.C; tmp /= args.C; ++ const int32_t n = tmp; ++ ++ const int32_t base_x = ow*args.s0 - args.p0; ++ const int32_t base_y = oh*args.s1 - args.p1; ++ ++ const uint64_t src_base = (uint64_t) n * args.nb13 + (uint64_t) c * args.nb12; ++ const uint64_t w_base = (uint64_t) c * args.nbkc; ++ ++ float acc = 0.0f; ++ for (int32_t ky = 0; ky < args.KH; ++ky) { ++ const int32_t iy = base_y + ky*args.d1; ++ if (iy < 0 || iy >= args.IH) continue; ++ const uint64_t src_row = src_base + (uint64_t) iy * args.nb11; ++ const uint64_t w_row = w_base + (uint64_t) ky * args.nbk1; ++ for (int32_t kx = 0; kx < args.KW; ++kx) { ++ const int32_t ix = base_x + kx*args.d0; ++ if (ix < 0 || ix >= args.IW) continue; ++ const float xv = *(device const float *)(src + src_row + (uint64_t) ix * args.nb10); ++ const float wv = *(device const float *)(weights + w_row + (uint64_t) kx * args.nbk0); ++ acc += xv * wv; ++ } ++ } ++ ++ const uint64_t dst_offs = ++ (uint64_t) n * args.nb3 + ++ (uint64_t) c * args.nb2 + ++ (uint64_t) oh * args.nb1 + ++ (uint64_t) ow * args.nb0; ++ ++ *(device float *)(dst + dst_offs) = acc; ++ } ++} ++ + typedef void (conv_transpose_1d_t)( + constant ggml_metal_kargs_conv_transpose_1d & args, + device const float * src0, diff --git a/third_party/ggml-patches/0003-metal-pad-leading.patch b/third_party/ggml-patches/0003-metal-pad-leading.patch new file mode 100644 index 0000000..629058d --- /dev/null +++ b/third_party/ggml-patches/0003-metal-pad-leading.patch @@ -0,0 +1,105 @@ +diff --git a/src/ggml-metal/ggml-metal-device.m b/src/ggml-metal/ggml-metal-device.m +index 1c9960a8..bcc243ab 100644 +--- a/src/ggml-metal/ggml-metal-device.m ++++ b/src/ggml-metal/ggml-metal-device.m +@@ -1137,13 +1137,10 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te + case GGML_OP_POOL_2D: + return op->src[0]->type == GGML_TYPE_F32; + case GGML_OP_PAD: +- // TODO: add circular padding support for metal, see https://github.com/ggml-org/llama.cpp/pull/16985 +- if (ggml_get_op_params_i32(op, 8) != 0) { +- return false; +- } +- +- return (ggml_get_op_params_i32(op, 0) == 0) && (ggml_get_op_params_i32(op, 2) == 0) && +- (ggml_get_op_params_i32(op, 4) == 0) && (ggml_get_op_params_i32(op, 6) == 0); ++ // Leading and trailing pads on any dim are supported. Circular ++ // padding (op_params[8]) is not yet implemented on Metal. ++ // See https://github.com/ggml-org/llama.cpp/pull/16985 ++ return ggml_get_op_params_i32(op, 8) == 0; + case GGML_OP_PAD_REFLECT_1D: + case GGML_OP_TIMESTEP_EMBEDDING: + case GGML_OP_LEAKY_RELU: +diff --git a/src/ggml-metal/ggml-metal-impl.h b/src/ggml-metal/ggml-metal-impl.h +index d7b4b60a..73d1e220 100644 +--- a/src/ggml-metal/ggml-metal-impl.h ++++ b/src/ggml-metal/ggml-metal-impl.h +@@ -1035,6 +1035,10 @@ typedef struct { + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; ++ int32_t lp0; ++ int32_t lp1; ++ int32_t lp2; ++ int32_t lp3; + } ggml_metal_kargs_pad; + + typedef struct { +diff --git a/src/ggml-metal/ggml-metal-ops.cpp b/src/ggml-metal/ggml-metal-ops.cpp +index 6035ba2a..3039c3b8 100644 +--- a/src/ggml-metal/ggml-metal-ops.cpp ++++ b/src/ggml-metal/ggml-metal-ops.cpp +@@ -4108,6 +4108,13 @@ int ggml_metal_op_pad(ggml_metal_op_t ctx, int idx) { + GGML_TENSOR_LOCALS( int32_t, ne, op, ne); + GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + ++ // Leading pads per dim (op_params even indices); trailing pads are implied ++ // by the dst ne[] and need no explicit value here. ++ const int32_t lp0 = ((const int32_t *) op->op_params)[0]; ++ const int32_t lp1 = ((const int32_t *) op->op_params)[2]; ++ const int32_t lp2 = ((const int32_t *) op->op_params)[4]; ++ const int32_t lp3 = ((const int32_t *) op->op_params)[6]; ++ + ggml_metal_kargs_pad args = { + /*.ne00 =*/ ne00, + /*.ne01 =*/ ne01, +@@ -4124,7 +4131,11 @@ int ggml_metal_op_pad(ggml_metal_op_t ctx, int idx) { + /*.nb0 =*/ nb0, + /*.nb1 =*/ nb1, + /*.nb2 =*/ nb2, +- /*.nb3 =*/ nb3 ++ /*.nb3 =*/ nb3, ++ /*.lp0 =*/ lp0, ++ /*.lp1 =*/ lp1, ++ /*.lp2 =*/ lp2, ++ /*.lp3 =*/ lp3, + }; + + auto pipeline = ggml_metal_library_get_pipeline_pad(lib, op); +diff --git a/src/ggml-metal/ggml-metal.metal b/src/ggml-metal/ggml-metal.metal +index 7c4fae05..1acea264 100644 +--- a/src/ggml-metal/ggml-metal.metal ++++ b/src/ggml-metal/ggml-metal.metal +@@ -5398,9 +5398,17 @@ kernel void kernel_pad_impl( + const int32_t k0 = tgpig.x/args.ne1; + const int32_t i1 = tgpig.x - k0*args.ne1; + +- const int32_t i03 = i3; +- const int32_t i02 = i2; +- const int32_t i01 = i1; ++ // Leading-pad aware: a dst position iN maps to src position (iN - lpN), and ++ // is zero where that falls outside [0, ne0N). (lpN comes from op_params; the ++ // _4/c4 vectorized variant is disabled in the pipeline getter, so lp0 is in ++ // scalar element units here.) ++ const int32_t i03 = i3 - args.lp3; ++ const int32_t i02 = i2 - args.lp2; ++ const int32_t i01 = i1 - args.lp1; ++ ++ const bool row_in = i01 >= 0 && i01 < args.ne01 && ++ i02 >= 0 && i02 < args.ne02 && ++ i03 >= 0 && i03 < args.ne03; + + device const T * src0_ptr = (device const T *) (src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01); + device T * dst_ptr = (device T *) (dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1); +@@ -5411,8 +5419,9 @@ kernel void kernel_pad_impl( + break; + } + +- if (i0 < args.ne00 && i1 < args.ne01 && i2 < args.ne02 && i3 < args.ne03) { +- dst_ptr[i0] = src0_ptr[i0]; ++ const int32_t i00 = i0 - args.lp0; ++ if (row_in && i00 >= 0 && i00 < args.ne00) { ++ dst_ptr[i0] = src0_ptr[i00]; + } else { + dst_ptr[i0] = 0.0f; + }