diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f68b2eb..e6474bf 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -121,9 +121,16 @@ jobs: - os: ubuntu-latest name: linux-x64-cpu backend: cpu - # CPU-only package built with -march=native for the target simd - # (ggml then uses AVX2/AVX-512 etc. via GGML_NATIVE). - cmake_extra: "-DGGML_NATIVE=ON" + # v0.20.2 constraint: GGML_CPU_ALL_VARIANTS requires + # GGML_BACKEND_DL (ISA variants became dlopen MODULE plugins), + # but DL mode does NOT link the CPU backend into the ggml + # umbrella target, so game.cpp's direct references + # (ggml_backend_cpu_init, ggml_threadpool_new, ...) fail to + # link. Non-DL + GGML_NATIVE=OFF builds a portable baseline + # x86-64 CPU backend (no -march=native capture of the runner). + # SIMD dispatch for the CPU package needs a dlopen-path refactor + # of backend.cpp and is tracked separately. + cmake_extra: "-DGGML_NATIVE=OFF" build_jobs: 4 pkg_ext: "" lib_glob: "libggml*.so*" @@ -131,7 +138,13 @@ jobs: - os: ubuntu-22.04 name: linux-x64-cuda backend: cuda - cmake_extra: '-DGAME_GGML_CUDA=ON -DGGML_NATIVE=OFF -DCMAKE_CUDA_ARCHITECTURES=75' + # CUDA 12.9 (newest Jimver/cuda-toolkit ships; 12.8+ required to + # compile Blackwell sm_120). Arch list: native SASS for the + # mainstream generations (Turing→Hopper) plus compute_120 PTX so + # RTX 50-series and future GPUs run native-arch JIT instead of + # falling back to compute_75 PTX (which works but loses all + # newer-architecture kernel optimisations). + cmake_extra: '-DGAME_GGML_CUDA=ON -DGGML_NATIVE=OFF -DCMAKE_CUDA_ARCHITECTURES="75;80;86;89;90;120-virtual"' build_jobs: 2 pkg_ext: "" lib_glob: "libggml*.so*" @@ -163,7 +176,9 @@ jobs: - os: windows-2022 name: windows-x64-cuda backend: cuda - cmake_extra: '-DGAME_GGML_CUDA=ON -DGGML_NATIVE=OFF -DCMAKE_CUDA_ARCHITECTURES=75' + # Same arch policy as linux-x64-cuda (see above): CUDA 12.9 + + # SASS for Turing→Hopper + Blackwell PTX. + cmake_extra: '-DGAME_GGML_CUDA=ON -DGGML_NATIVE=OFF -DCMAKE_CUDA_ARCHITECTURES="75;80;86;89;90;120-virtual"' build_jobs: 2 pkg_ext: ".exe" lib_glob: "ggml*.dll" @@ -189,17 +204,17 @@ jobs: - name: Install CUDA Toolkit (Linux) if: matrix.backend == 'cuda' && runner.os == 'Linux' - uses: Jimver/cuda-toolkit@1a3c14e26833ccf292b268f9a790fb47dea7b2da # v0.2.28 + uses: Jimver/cuda-toolkit@v0.2.29 with: - cuda: "12.6.3" + cuda: "12.9.0" method: network log-file-suffix: "${{ matrix.name }}.txt" - name: Install CUDA Toolkit (Windows) if: matrix.backend == 'cuda' && runner.os == 'Windows' - uses: Jimver/cuda-toolkit@b8bf9c6c28f8a92fbb04dcfcaee872e60c57462d # v0.2.36 + uses: Jimver/cuda-toolkit@v0.2.29 with: - cuda: "12.6.3" + cuda: "12.9.0" method: network sub-packages: '["nvcc", "cudart", "cublas", "cublas_dev", "visual_studio_integration"]' log-file-suffix: "${{ matrix.name }}.txt" diff --git a/.gitignore b/.gitignore index 755b8d8..316afff 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,6 @@ # Build outputs build/ +build-*/ out/ cmake-build-*/ diff --git a/BUILDING.md b/BUILDING.md index a8f3fdc..80fd7f2 100644 --- a/BUILDING.md +++ b/BUILDING.md @@ -22,7 +22,7 @@ sudo apt install libvulkan-dev vulkan-tools # sudo apt install glslc-tools # not available on all distros # Optional: CUDA backend (NVIDIA GPUs) -# Install a CUDA Toolkit supported by your compiler and driver. CI uses CUDA 12.6.3. +# Install a CUDA Toolkit supported by your compiler and driver. CI uses CUDA 12.9.0. # The prebuilt CUDA packages target Turing (CC 7.5) and newer GPUs. # https://developer.nvidia.com/cuda-downloads @@ -57,9 +57,9 @@ brew install cmake ccache # Vulkan SDK (optional, for Vulkan backend) # https://vulkan.lunarg.com/sdk/home # -# CUDA Toolkit 12.6.x (optional, for CUDA backend) +# CUDA Toolkit 12.9.x (optional, for CUDA backend) # https://developer.nvidia.com/cuda-downloads -# CUDA 12.6 supports Visual Studio 2022 / MSVC 193x. +# CUDA 12.9 supports Visual Studio 2019 16.11+ and 2022 (MSVC 192x/193x). ``` > **Windows + CUDA toolchain notes (from local builds):** @@ -73,7 +73,7 @@ brew install cmake ccache > Pin the one you intend to use, e.g. set `CUDA_PATH_V13_0` to the v11.6 path > when building with CUDA 11.6, otherwise nvcc 13 + MSVC 14.29 hits > `__cudaLaunch` macro breakage (`error C4002`). -> - CI reference: `ubuntu-22.04` / `windows-2022` + CUDA 12.6.3. +> - CI reference: `ubuntu-22.04` / `windows-2022` + CUDA 12.9.0. ## Quick start @@ -193,7 +193,7 @@ build/bin/game_ggml_cli serve game_medium.gguf ## CUDA compatibility and CI scope The hosted CI builds Linux x64 and Windows x64 CUDA packages with CUDA Toolkit -12.6.3 and Visual Studio 2022 on Windows. It verifies Toolkit discovery, CUDA +12.9.0 and Visual Studio 2022 on Windows. It verifies Toolkit discovery, CUDA compilation, linking, and packaging. On Windows, `ggml-cuda.dll` imports `nvcuda.dll` (the NVIDIA driver library) at load time, and GitHub-hosted runners have no NVIDIA driver, so the CLI cannot start there even for `--version`; the @@ -203,19 +203,24 @@ dependencies instead, and performs the startup smoke test only when provide an NVIDIA GPU, so actual CUDA inference must still be smoke-tested on an NVIDIA system. -The release architecture is `75`, which emits both native CC 7.5 SASS and CC -7.5 PTX. Turing GPUs (for example, GeForce RTX 20 series) use the native image; -newer Ampere, Ada, and later drivers can JIT the PTX forward-compatible image. -This keeps the hosted build practical: compiling every ggml CUDA translation -unit separately for four real architectures was several times slower and used -substantially more memory. +The release architecture list is `75;80;86;89;90;120-virtual`: +native SASS for Turing (CC 7.5), Ampere data-center (CC 8.0), Ampere consumer +(CC 8.6), Ada (CC 8.9) and Hopper (CC 9.0), plus `compute_120` PTX so RTX +50-series and future GPUs JIT with a Blackwell-targeted image. A plain `75` +build also runs on newer GPUs through its bundled `compute_75` PTX, but the +JIT-ed kernels then miss every newer-architecture optimisation; shipping SASS +for the mainstream generations removes that penalty. Compiling six +architectures roughly multiplies the nvcc workload (the CUDA jobs already run +with `build_jobs: 2` for memory headroom), which is the trade-off for native +performance on each generation. Pascal and Volta are not included in the prebuilt package. Source builds that need these older GPUs can use CUDA 12.x and add `61-real` and/or `70-real`. -Source builds that prefer native images for each newer generation may use +Source builds that want to trim the list back to Turing→Ada may use `75-real;80-real;86-real;89-real`, accepting the longer build and larger binary. CUDA 13.0 removed NVCC offline compilation for architectures older than CC 7.5; -use CUDA 12.9 or earlier when maintaining such builds. +use CUDA 12.9 or earlier when maintaining such builds. sm_100/sm_120 (Blackwell) +compile targets require CUDA 12.8 or newer. CUDA 12.x minor-version compatibility requires at least NVIDIA driver 525.60.13 on Linux or 528.33 on Windows, subject to the limitations documented diff --git a/README.md b/README.md index b38ea1f..10c5dbf 100644 --- a/README.md +++ b/README.md @@ -273,7 +273,7 @@ int main() { params.language = 4; // from lang_map: { "zh": 4 } params.seed = 42; // DBCache (segmenter cross-step reuse; affects nsteps>1 only). - // -1 = auto (CPU 0.25, GPU off); 0 = off; >0 = explicit threshold. + // -1 = auto (0.25 on all backends); 0 = off; >0 = explicit threshold. params.db_cache_threshold = 0.25f; params.db_cache_fn_blocks = 1; params.db_cache_warmup = 1; @@ -346,7 +346,7 @@ trees live under `build/_deps/-src/` after the first configure. | Dependency | Version pin | License | SPDX identifier | |---|---|---|---| -| [ggml](https://github.com/ggerganov/ggml) | `v0.11.0` tag | MIT | MIT | +| [ggml](https://github.com/ggerganov/ggml) | `v0.20.2` tag | MIT | MIT | | [pocketfft](https://gitlab.mpcdf.mpg.de/mtr/pocketfft) | commit `32424d20` on `cpp` branch | BSD-3-Clause | BSD-3-Clause | | [dr_libs](https://github.com/mackron/dr_libs) | commit `243e26ff` on `master` | Public Domain / MIT-0 (dual) | `Unlicense OR MIT-0` | | [GoogleTest](https://github.com/google/googletest) | `v1.14.0` tag (tests only) | BSD-3-Clause | BSD-3-Clause | diff --git a/README_CN.md b/README_CN.md index c839853..0ee5f34 100644 --- a/README_CN.md +++ b/README_CN.md @@ -235,7 +235,7 @@ int main() { params.language = 4; // 来自 lang_map: { "zh": 4 } params.seed = 42; // DBCache（跨步复用，仅 nsteps>1 生效）： - // -1 = 自动（CPU 0.25，GPU 关闭）；0 = 关闭；>0 = 显式阈值。 + // -1 = 自动（所有后端 0.25）；0 = 关闭；>0 = 显式阈值。 params.db_cache_threshold = 0.25f; params.db_cache_fn_blocks = 1; params.db_cache_warmup = 1; @@ -282,7 +282,7 @@ ctest --test-dir ggml_backend/build --output-on-failure | 依赖 | 版本 pin | 许可 | SPDX 标识 | |---|---|---|---| -| [ggml](https://github.com/ggerganov/ggml) | `v0.11.0` tag | MIT | MIT | +| [ggml](https://github.com/ggerganov/ggml) | `v0.20.2` tag | MIT | MIT | | [pocketfft](https://gitlab.mpcdf.mpg.de/mtr/pocketfft) | `cpp` 分支 `32424d20` | BSD-3-Clause | BSD-3-Clause | | [dr_libs](https://github.com/mackron/dr_libs) | `master` 分支 `243e26ff` | Public Domain / MIT-0（双许可） | `Unlicense OR MIT-0` | | [GoogleTest](https://github.com/google/googletest) | `v1.14.0` tag（仅测试） | BSD-3-Clause | BSD-3-Clause | diff --git a/cmake/Dependencies.cmake b/cmake/Dependencies.cmake index 44f56ba..9405e36 100644 --- a/cmake/Dependencies.cmake +++ b/cmake/Dependencies.cmake @@ -62,7 +62,7 @@ endfunction() # --------------------------------------------------------------------------- # SPIRV-Headers shim (Windows Vulkan only) # -# ggml v0.19.0's Vulkan backend hard-requires find_package(SPIRV-Headers CONFIG). +# ggml v0.20.x's Vulkan backend hard-requires find_package(SPIRV-Headers CONFIG). # Windows Vulkan SDKs older than ~1.4.35x ship the headers but not that CMake # config file, so the windows-x64-vulkan CI job fails at configure time. # Generate a minimal config pointing at the SDK headers when the SDK doesn't @@ -87,9 +87,18 @@ endif() # --------------------------------------------------------------------------- FetchContent_Declare( ggml - GIT_REPOSITORY https://github.com/ggerganov/ggml.git - GIT_TAG v0.19.0 - GIT_SHALLOW TRUE + # URL archive instead of git: with FETCHCONTENT_UPDATES_DISCONNECTED=ON + # the populate gitupdate step must resolve the pinned ref locally, and a + # fresh clone cannot — a shallow clone only carries the default-branch + # HEAD, and the v0.20.2 tag/commit is not on it — so populate aborts with + # "requested git ref ... not present locally". A URL archive has no git + # ref semantics: populate is a plain download+extract (network needed on + # first populate / cache miss; the CI _deps cache then makes later runs + # offline). Patches are still applied with `git apply`, which needs no + # .git directory. (No DOWNLOAD_EXTRACT_TIMESTAMP: requires CMake 3.24+, + # project minimum is 3.18.) + URL https://github.com/ggerganov/ggml/archive/refs/tags/v0.20.2.tar.gz + URL_HASH SHA256=55dfd1ea4e6b6b3e25d9411f9525eb4df1c796c03a244e2321388b30f189cd3d ) FetchContent_GetProperties(ggml) diff --git a/cmake/patches/ggml-metal-binary-archive.md b/cmake/patches/ggml-metal-binary-archive.md index 91cb423..af4a2ef 100644 --- a/cmake/patches/ggml-metal-binary-archive.md +++ b/cmake/patches/ggml-metal-binary-archive.md @@ -27,7 +27,7 @@ Controlled by the same `GGML_METAL_ARCHIVE_PATH` env var; disable with ## Baseline -Applied against **ggml v0.19.0** (`ggml-metal-device.m`). Verify with the +Applied against **ggml v0.20.2** (`ggml-metal-device.m`). Verify with the same command used in `cmake/Dependencies.cmake`: ``` diff --git a/cmake/patches/ggml-vulkan-pipeline-cache.md b/cmake/patches/ggml-vulkan-pipeline-cache.md index 1ad3a3b..7c4d789 100644 --- a/cmake/patches/ggml-vulkan-pipeline-cache.md +++ b/cmake/patches/ggml-vulkan-pipeline-cache.md @@ -38,7 +38,7 @@ driver) load precompiled PSO bytes instead of recompiling every shader. ## Baseline & re-apply -- Applies to **ggml v0.19.0** (`ggml-vulkan.cpp`). Pinned by game.cpp +- Applies to **ggml v0.20.2** (`ggml-vulkan.cpp`). Pinned by game.cpp FetchContent; re-apply per ggml upgrade via `cmake/Dependencies.cmake` `game_ggml_apply_patch` (idempotent: skips if already applied). diff --git a/cmake/patches/ggml-vulkan-pipeline-cache.patch b/cmake/patches/ggml-vulkan-pipeline-cache.patch index 8cc2dba..1921e66 100644 --- a/cmake/patches/ggml-vulkan-pipeline-cache.patch +++ b/cmake/patches/ggml-vulkan-pipeline-cache.patch @@ -1,5 +1,5 @@ diff --git a/src/ggml-vulkan/ggml-vulkan.cpp b/src/ggml-vulkan/ggml-vulkan.cpp -index a923755..56304ff 100644 +index 585e10d..4f3686a 100644 --- a/src/ggml-vulkan/ggml-vulkan.cpp +++ b/src/ggml-vulkan/ggml-vulkan.cpp @@ -50,6 +50,7 @@ typedef struct VkPhysicalDeviceCooperativeMatrixDecodeVectorFeaturesNV { @@ -192,7 +192,7 @@ index a923755..56304ff 100644 void vk_command_pool::init(vk_device& device, vk_queue *q_) { cmd_buffers.clear(); q = q_; -@@ -3025,7 +3180,8 @@ static void ggml_vk_create_pipeline_func(vk_device& device, vk_pipeline& pipelin +@@ -3026,7 +3181,8 @@ static void ggml_vk_create_pipeline_func(vk_device& device, vk_pipeline& pipelin #endif try { @@ -202,7 +202,7 @@ index a923755..56304ff 100644 } catch (const vk::SystemError& e) { std::cerr << "ggml_vulkan: Compute pipeline creation failed for " << pipeline->name << std::endl; std::cerr << "ggml_vulkan: " << e.what() << std::endl; -@@ -6910,6 +7066,8 @@ static vk_device ggml_vk_get_device(size_t idx) { +@@ -6934,6 +7090,8 @@ static vk_device ggml_vk_get_device(size_t idx) { .setPEnabledExtensionNames(device_extensions); device_create_info.setPNext(&device_features2); device->device = device->physical_device.createDevice(device_create_info); @@ -211,7 +211,7 @@ index a923755..56304ff 100644 if (device->device_fault) { device->pfn_vkGetDeviceFaultInfoEXT = (PFN_vkGetDeviceFaultInfoEXT) -@@ -17372,6 +17530,11 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg +@@ -17402,6 +17560,11 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ggml_vk_synchronize(ctx); } diff --git a/scripts/verify_backends.py b/scripts/verify_backends.py new file mode 100644 index 0000000..a0f46cc --- /dev/null +++ b/scripts/verify_backends.py @@ -0,0 +1,162 @@ +#!/usr/bin/env python3 +"""Multi-backend verification for game_ggml_cli. + +Runs the same waveform through the CPU / Vulkan / CUDA CLI builds +(nsteps=1 fused path and nsteps=8 DBCache path) and compares outputs: + + * CSV note lists (structure: note count, per-note pitch/offset/duration) + * DBCache hit/miss pattern (via GAME_GGML_DUMP_DBCACHE=1 stderr) + +Usage: + python verify_backends.py --cli-cpu build/bin/game_ggml_cli.exe \ + --cli-vk build-vk/bin/game_ggml_cli.exe \ + --cli-cuda build-cuda/bin/game_ggml_cli.exe \ + --wav test10s.wav -m model.gguf -o outdir + +Exit code 0 = all backend outputs match CPU reference within tolerance. +""" + +from __future__ import annotations + +import argparse +import csv +import pathlib +import re +import subprocess +import sys + +# Note-list tolerance: pitch is quantised (midi cents), offset/duration in +# seconds with 2-decimal CSV precision. GPUs may shift a boundary by one +# frame (11.61 ms) at a note edge. +PITCH_EPS = 5 # cents +TIME_EPS = 0.05 # s + + +def run_cli(cli: pathlib.Path, wav: pathlib.Path, model: pathlib.Path, + out_dir: pathlib.Path, seed: int, nsteps: int, + env_extra: dict | None = None) -> tuple[int, str, str]: + cmd = [str(cli), "extract", str(wav), "-m", str(model), + "--output-dir", str(out_dir), "--seed", str(seed), + "--output-formats", "mid,csv,txt"] + if nsteps: + cmd += ["--nsteps", str(nsteps)] + env = dict(subprocess.os.environ) + env["GAME_GGML_DUMP_DBCACHE"] = "1" + if env_extra: + env.update(env_extra) + r = subprocess.run(cmd, capture_output=True, text=True, env=env, timeout=900) + return r.returncode, r.stdout, r.stderr + + +def read_notes(csv_path: pathlib.Path) -> list[dict]: + if not csv_path.exists(): + return [] + with open(csv_path, newline="") as f: + return list(csv.DictReader(f)) + + +def parse_pitch(p: str) -> int | None: + """'A3-37' -> midi-ish cents value; 'rest' -> None. + + Raises ValueError for any string that is neither 'rest' nor a parsable + pitch — a malformed pitch must not silently compare equal to a rest. + """ + if p == "rest": + return None + m = re.match(r"([A-G])(#?)(\d+)([+-]\d+)?", p) + if not m: + raise ValueError(f"unparsable pitch: {p!r}") + letter, sharp, octave, cents = m.groups() + base = {"C": 0, "D": 2, "E": 4, "F": 5, "G": 7, "A": 9, "B": 11}[letter] + semi = base + (1 if sharp else 0) + (int(octave) + 1) * 12 + return semi * 100 + int(cents or 0) + + +def notes_close(a: list[dict], b: list[dict]) -> tuple[bool, str]: + if len(a) != len(b): + return False, f"note count {len(a)} != {len(b)}" + for i, (ra, rb) in enumerate(zip(a, b, strict=True)): + pa, pb = parse_pitch(ra["pitch"]), parse_pitch(rb["pitch"]) + if (pa is None) != (pb is None): + return False, f"note[{i}] pitch rest-mismatch {ra['pitch']} vs {rb['pitch']}" + if pa is not None and abs(pa - pb) > PITCH_EPS: + return False, f"note[{i}] pitch {ra['pitch']} vs {rb['pitch']}" + for k in ("offset", "duration"): + da, db = float(ra[k]), float(rb[k]) + if abs(da - db) > TIME_EPS: + return False, f"note[{i}] {k} {da} vs {db}" + return True, "" + + +def db_pattern(stderr: str) -> str: + hits = len(re.findall(r"DB.?cache hit|hit", stderr, re.I)) + misses = len(re.findall(r"miss", stderr, re.I)) + m = re.search(r"(?:hit|miss).*?(\d+).*?(\d+)", stderr) + return f"hits~{hits}/misses~{misses}" + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--cli-cpu", required=True, type=pathlib.Path) + ap.add_argument("--cli-vk", type=pathlib.Path, default=None) + ap.add_argument("--cli-cuda", type=pathlib.Path, default=None) + ap.add_argument("--wav", required=True, type=pathlib.Path) + ap.add_argument("-m", "--model", required=True, type=pathlib.Path) + ap.add_argument("-o", "--out-root", required=True, type=pathlib.Path) + ap.add_argument("--seed", type=int, default=42) + args = ap.parse_args() + + backends = [("cpu", args.cli_cpu)] + if args.cli_vk: + backends.append(("vk", args.cli_vk)) + if args.cli_cuda: + backends.append(("cuda", args.cli_cuda)) + + args.out_root.mkdir(parents=True, exist_ok=True) + rc_total = 0 + + for nsteps in (1, 8): + print(f"\n===== nsteps={nsteps} =====") + results = {} + for name, cli in backends: + out = args.out_root / f"{name}_n{nsteps}" + out.mkdir(exist_ok=True) + code, so, se = run_cli(cli, args.wav, args.model, out, args.seed, + nsteps if nsteps > 1 else 0) + csv_path = out / f"{args.wav.stem}.csv" + notes = read_notes(csv_path) + ok = code == 0 + results[name] = (notes, code, db_pattern(se), out) + print(f"[{name}] rc={code} notes={len(notes)} dbc={db_pattern(se)}") + + ref, ref_code = results["cpu"][0], results["cpu"][1] + if ref_code != 0: + print(" !! cpu reference non-zero exit; comparison is meaningless") + rc_total = 1 + continue + if not ref: + print(" !! cpu reference produced no notes; comparison is meaningless") + rc_total = 1 + continue + if len(backends) == 1: + print(" !! no GPU backend supplied; nothing to compare against") + rc_total = 1 + continue + for name in list(results)[1:]: + notes, code, _, out = results[name] + if code != 0: + print(f" !! {name} non-zero exit") + rc_total = 1 + continue + close, why = notes_close(ref, notes) + status = "MATCH" if close else "DIFF" + print(f" {name}: {status} vs cpu" + (f" ({why})" if why else "")) + if not close: + rc_total = 1 + + print(f"\n{'ALL BACKENDS MATCH' if rc_total == 0 else 'MISMATCH DETECTED'}") + return rc_total + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/backend.cpp b/src/backend.cpp index 3f75bca..89e5b1d 100644 --- a/src/backend.cpp +++ b/src/backend.cpp @@ -22,8 +22,10 @@ #include #include #include +#include #include #include +#include // ----------------------------------------------------------------------------- // Public version helpers (declared in version.h) @@ -48,9 +50,10 @@ const char * version_string() noexcept { } const char * ggml_version_string() noexcept { - // ggml v0.11.0 does not export a runtime-queryable version. We log the - // compile-time pin used by this project so the CLI can say "ggml 0.11.0". - return "0.11.0"; + // ggml does not export a runtime-queryable version. Report the tag we + // pin in cmake/Dependencies.cmake (FetchContent GIT_TAG) so --version + // cannot silently drift from the actual dependency. + return "v0.20.2"; } // ----------------------------------------------------------------------------- @@ -93,6 +96,21 @@ int available_backends_count() noexcept { // ----------------------------------------------------------------------------- namespace game_ggml::internal { +namespace { +// CPU threadpool registry: ggml v0.19 creates a *disposable* threadpool on +// every graph compute when none is attached (thread spawn per call). We hold +// one persistent pool per CPU backend and free it together with the backend. +std::mutex g_tp_mutex; +std::unordered_map g_tp; + +ggml_threadpool_t make_cpu_threadpool(int n_threads) { + // Default params: hybrid polling (poll=50) keeps the worker threads warm + // across the many small ops of this model's graphs without pure busy-wait. + struct ggml_threadpool_params tpp = ggml_threadpool_params_default(n_threads); + return ggml_threadpool_new(&tpp); +} +} // namespace + ggml_backend_t init_backend(Backend which) { switch (which) { case Backend::Metal: @@ -136,6 +154,12 @@ ggml_backend_t init_backend(Backend which) { } } ggml_backend_cpu_set_n_threads(b, static_cast(n)); + ggml_threadpool_t tp = make_cpu_threadpool(static_cast(n)); + if (tp) { + ggml_backend_cpu_set_threadpool(b, tp); + std::lock_guard lock(g_tp_mutex); + g_tp[b] = tp; + } } return b; } @@ -159,6 +183,14 @@ ggml_backend_t init_best_backend() { void free_backend(ggml_backend_t backend) { if (backend == nullptr) return; + { + std::lock_guard lock(g_tp_mutex); + auto it = g_tp.find(backend); + if (it != g_tp.end()) { + ggml_threadpool_free(it->second); + g_tp.erase(it); + } + } ggml_backend_free(backend); } diff --git a/src/cli/main.cpp b/src/cli/main.cpp index 05f1461..97f406d 100644 --- a/src/cli/main.cpp +++ b/src/cli/main.cpp @@ -101,7 +101,7 @@ void print_usage(const char * argv0) { " --no-slice Feed the full WAV as one chunk (default: false)\n" " --pitch-format name|number Text output pitch format (default: name)\n" " --round-pitch Round pitch to integer in text output (default: false)\n" - " --cache-threshold DBCache normalized-L1 threshold (default: auto: CPU 0.25, GPU off)\n" + " --cache-threshold DBCache normalized-L1 threshold (default: 0.25 on all backends; 0 = off)\n" " --cache-fn-blocks DBCache front blocks per step (default: 1)\n" " --cache-warmup D3PM steps before caching starts (default: 1)\n" " --cache-window-start only cache from this step fraction on (default: 0)\n" @@ -389,7 +389,7 @@ int cmd_extract(int argc, char ** argv) { std::string pitch_format = "name"; bool round_pitch = false; std::string rng_replay_path; - float db_cache_threshold = -1.0f; // auto: CPU 0.25, GPU 0 (EP-aware) + float db_cache_threshold = -1.0f; // <0: default 0.25 (all backends); 0: off int db_cache_fn_blocks = 1; int db_cache_warmup = 1; float db_cache_window_start = 0.0f; diff --git a/src/mel.cpp b/src/mel.cpp index a2f542d..1bb4777 100644 --- a/src/mel.cpp +++ b/src/mel.cpp @@ -7,6 +7,7 @@ #include #include #include +#include #include namespace game_ggml { @@ -136,32 +137,66 @@ std::vector MelExtractor::forward(const float * wav, std::size_t n) const const int T = num_frames(n); if (T <= 0) return {}; - std::vector frame(n_fft, 0.0f); - std::vector> spec(n_bins); - std::vector mag(n_bins); - std::vector out(static_cast(T) * n_mels); + std::vector out(static_cast(T) * n_mels); pocketfft::shape_t shape = {static_cast(n_fft)}; pocketfft::stride_t stride_in = {sizeof(float)}; pocketfft::stride_t stride_out = {sizeof(std::complex)}; pocketfft::shape_t axes = {0}; - for (int t = 0; t < T; ++t) { - const std::size_t off = static_cast(t) * hop; - std::fill(frame.begin(), frame.end(), 0.0f); - for (int k = 0; k < win; ++k) frame[k] = padded[off + k] * window[k]; - - pocketfft::r2c(shape, stride_in, stride_out, axes, pocketfft::FORWARD, - frame.data(), spec.data(), 1.0f); - for (int k = 0; k < n_bins; ++k) mag[k] = std::hypot(spec[k].real(), spec[k].imag()); - - for (int m = 0; m < n_mels; ++m) { - const float * row = mel_fb.data() + static_cast(m) * n_bins; - float acc = 0.0f; - for (int k = 0; k < n_bins; ++k) acc += row[k] * mag[k]; - out[static_cast(t) * n_mels + m] = std::log(std::max(acc, cfg.clip_val)); + // Frames are independent — split the range into contiguous stripes and + // process them on a small worker pool. At ~1000+ frames per 10 s clip + // this is several ms of single-threaded work that parallelizes cleanly. + // Guard the pool size by frame count so short clips never pay for + // threads they don't use. + unsigned hw = std::thread::hardware_concurrency(); + if (hw == 0) hw = 1; + const int want = static_cast(std::min(hw, 8u)); + const int n_thr = std::max(1, std::min(want, T / 128 + 1)); + + auto worker = [&](int t_begin, int t_end) { + std::vector frame(n_fft, 0.0f); + std::vector> spec(n_bins); + std::vector mag(n_bins); + for (int t = t_begin; t < t_end; ++t) { + const std::size_t off = static_cast(t) * hop; + std::fill(frame.begin(), frame.end(), 0.0f); + for (int k = 0; k < win; ++k) frame[k] = padded[off + k] * window[k]; + + pocketfft::r2c(shape, stride_in, stride_out, axes, pocketfft::FORWARD, + frame.data(), spec.data(), 1.0f); + // Plain sqrt (not std::hypot): no overflow risk at FFT output + // magnitudes, and matches torch/librosa's |·| semantics within + // 1 ulp while running several times faster. + for (int k = 0; k < n_bins; ++k) { + const float re = spec[k].real(), im = spec[k].imag(); + mag[k] = std::sqrt(re * re + im * im); + } + + float * dst = out.data() + static_cast(t) * n_mels; + for (int m = 0; m < n_mels; ++m) { + const float * row = mel_fb.data() + static_cast(m) * n_bins; + float acc = 0.0f; + for (int k = 0; k < n_bins; ++k) acc += row[k] * mag[k]; + dst[m] = std::log(std::max(acc, cfg.clip_val)); + } } + }; + + if (n_thr == 1) { + worker(0, T); + return out; + } + const int per = (T + n_thr - 1) / n_thr; + std::vector pool; + pool.reserve(n_thr); + int t0 = 0; + for (int i = 0; i < n_thr && t0 < T; ++i) { + const int t1 = std::min(T, t0 + per); + pool.emplace_back(worker, t0, t1); + t0 = t1; } + for (auto & th : pool) th.join(); return out; } diff --git a/src/model.cpp b/src/model.cpp index ff7d02a..eab3c3f 100644 --- a/src/model.cpp +++ b/src/model.cpp @@ -10,6 +10,7 @@ #include #include +#include #include #include @@ -19,6 +20,7 @@ #include #include #include +#include #include #include #include @@ -58,6 +60,28 @@ InferResult Model::infer(const float * waveform, std::size_t n_samples, // ============================================================================ Model::Impl::~Impl() { + // Order matters: every backend-owned resource must be released before the + // backend itself. The member destructors would otherwise run AFTER + // ggml_backend_free(backend) and free device buffers that belong to an + // already-destroyed backend (CUDA/Vulkan): PersistentStage ~reset() frees + // the gallocr (whose buffer type came from `backend`) and LoadedWeights + // frees its backend buffers. `dbctx`/`dbbuf` have no owner, so release + // them here explicitly too (leak otherwise, one D×T set per Model). + enc_stage.reset(); + seg_stage.reset(); + seg_front_stage.reset(); + seg_add_stage.reset(); + seg_mid_stage.reset(); + seg_update_stage.reset(); + seg_back_stage.reset(); + seg_head_stage.reset(); + est_stage.reset(); + if (dbctx) { + ggml_backend_buffer_free(dbbuf); + ggml_free(dbctx); + dbctx = nullptr; dbbuf = nullptr; + } + weights.reset(); if (backend) internal::free_backend(backend); } @@ -115,221 +139,255 @@ std::unique_ptr Model::Impl::load(const std::string & path) { namespace { -// Thin guard making a temporary ggml_context + gallocr used by a single stage. -struct StageCtx { - ggml_context * ctx = nullptr; - ggml_cgraph * graph = nullptr; - ggml_gallocr_t alloc = nullptr; - ggml_backend_t backend = nullptr; - - StageCtx(ggml_backend_t b, std::size_t mem_bytes, int graph_nodes) { - backend = b; - ggml_init_params ip{}; - ip.mem_size = mem_bytes; - ip.no_alloc = true; - ctx = ggml_init(ip); - graph = ggml_new_graph_custom(ctx, graph_nodes, /*grads=*/false); - } - - void dump_backend_support(const char * stage) { - const char * env = std::getenv("GAME_GGML_DUMP_OPS"); - if (!env || !*env || env[0] == '0') return; - - std::map unsupported; - const int n_nodes = ggml_graph_n_nodes(graph); - for (int i = 0; i < n_nodes; ++i) { - const ggml_tensor * node = ggml_graph_node(graph, i); - if (!ggml_backend_supports_op(backend, node)) { - unsupported[ggml_op_name(node->op)] += 1; - } - } - - std::fprintf(stderr, - "[GAME_GGML_OPS] stage=%s backend=%s nodes=%d unsupported=%d\n", - stage, game_ggml::internal::backend_name(backend), - n_nodes, - std::accumulate(unsupported.begin(), unsupported.end(), 0, - [](int acc, const auto & kv) { return acc + kv.second; })); - for (const auto & kv : unsupported) { - std::fprintf(stderr, "[GAME_GGML_OPS] unsupported %-24s %d\n", - kv.first.c_str(), kv.second); +// Debug aid gated by GAME_GGML_DUMP_OPS: list ops the backend cannot run. +// (ggml v0.19 graph introspection takes a non-const cgraph.) +void dump_graph_support(const char * stage, ggml_backend_t backend, ggml_cgraph * graph) { + const char * env = std::getenv("GAME_GGML_DUMP_OPS"); + if (!env || !*env || env[0] == '0') return; + + std::map unsupported; + const int n_nodes = ggml_graph_n_nodes(graph); + for (int i = 0; i < n_nodes; ++i) { + const ggml_tensor * node = ggml_graph_node(graph, i); + if (!ggml_backend_supports_op(backend, node)) { + unsupported[ggml_op_name(node->op)] += 1; } } - void finalize(ggml_tensor * out, const char * stage = "stage") { - ggml_set_output(out); - ggml_build_forward_expand(graph, out); - dump_backend_support(stage); - alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!ggml_gallocr_alloc_graph(alloc, graph)) { - throw Error("ggml_gallocr_alloc_graph failed"); - } + std::fprintf(stderr, + "[GAME_GGML_OPS] stage=%s backend=%s nodes=%d unsupported=%d\n", + stage, game_ggml::internal::backend_name(backend), + n_nodes, + std::accumulate(unsupported.begin(), unsupported.end(), 0, + [](int acc, const auto & kv) { return acc + kv.second; })); + for (const auto & kv : unsupported) { + std::fprintf(stderr, "[GAME_GGML_OPS] unsupported %-24s %d\n", + kv.first.c_str(), kv.second); } - - void compute() { - if (ggml_backend_graph_compute(backend, graph) != GGML_STATUS_SUCCESS) { - throw Error("graph compute failed"); - } - } - - ~StageCtx() { - if (alloc) ggml_gallocr_free(alloc); - if (ctx) ggml_free(ctx); - } -}; +} } // namespace // ============================================================================ // Stage 1 — encoder (mel → x_seg, x_est) +// +// The graph is built once per frame count T (PersistentStage) and its two +// outputs — x_seg / x_est — live in the encoder's backend buffer, so they +// stay resident on the device and are referenced directly by the segmenter +// and estimator stages (no host round-trip between stages). // ============================================================================ -void Model::Impl::run_encoder(const float * mel, int T, - std::vector & x_seg_out, - std::vector & x_est_out) +void Model::Impl::run_encoder(const float * mel, int T) { const int D_mel = cfg.in_dim; - const int D_emb = cfg.embedding_dim; - - StageCtx s(backend, 256 * 1024 * 1024, 8192); - - ggml_tensor * mel_in = ggml_new_tensor_3d(s.ctx, GGML_TYPE_F32, D_mel, T, 1); - ggml_set_input(mel_in); - ggml_tensor * pos = ggml_new_tensor_1d(s.ctx, GGML_TYPE_I32, T); - ggml_set_input(pos); - - // spectrogram_projection (mel → D_embed) - ggml_tensor * x_proj = internal::ops::linear(s.ctx, mel_in, w_spec_proj, b_spec_proj); - auto outs = internal::build_encoder_graph(s.ctx, x_proj, encoder_w, pos, cfg.encoder); - - ggml_set_output(outs.x_seg); - ggml_set_output(outs.x_est); - ggml_build_forward_expand(s.graph, outs.x_seg); - ggml_build_forward_expand(s.graph, outs.x_est); - s.dump_backend_support("encoder"); - s.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!ggml_gallocr_alloc_graph(s.alloc, s.graph)) throw Error("alloc failed (encoder)"); - - ggml_backend_tensor_set(mel_in, mel, 0, ggml_nbytes(mel_in)); - std::vector pos_i(T); - for (int i = 0; i < T; ++i) pos_i[i] = i; - ggml_backend_tensor_set(pos, pos_i.data(), 0, pos_i.size() * sizeof(std::int32_t)); - - s.compute(); - - x_seg_out.resize(ggml_nelements(outs.x_seg)); - x_est_out.resize(ggml_nelements(outs.x_est)); - ggml_backend_tensor_get(outs.x_seg, x_seg_out.data(), 0, x_seg_out.size() * sizeof(float)); - ggml_backend_tensor_get(outs.x_est, x_est_out.data(), 0, x_est_out.size() * sizeof(float)); -} -// ============================================================================ -// Stage 2 — segmenter (one D3PM step, DBCache-aware) -// ============================================================================ + ensure_db_tensors(T); -namespace { + if (!enc_stage.matches(T)) { + enc_stage.reset(); -// Normalized L1 residual between the current and previous front output -// (mirrors PyTorch DBCacheSegmenter: mean|x-x_prev| / (mean|x_prev| + eps)). -inline float front_delta(const float * cur, const float * prev, int n) { - float num = 0.0f, den = 0.0f; - for (int i = 0; i < n; ++i) { - num += std::fabs(cur[i] - prev[i]); - den += std::fabs(prev[i]); + ggml_init_params ip{}; + ip.mem_size = 32 * 1024 * 1024; + ip.no_alloc = true; + enc_stage.ctx = ggml_init(ip); + enc_stage.graph = ggml_new_graph_custom(enc_stage.ctx, 8192, /*grads=*/false); + + ggml_tensor * mel_in = ggml_new_tensor_3d(enc_stage.ctx, GGML_TYPE_F32, D_mel, T, 1); + ggml_set_input(mel_in); + ggml_tensor * pos = ggml_new_tensor_1d(enc_stage.ctx, GGML_TYPE_I32, T); + ggml_set_input(pos); + + // spectrogram_projection (mel → D_embed) + ggml_tensor * x_proj = internal::ops::linear(enc_stage.ctx, mel_in, w_spec_proj, b_spec_proj); + auto outs = internal::build_encoder_graph(enc_stage.ctx, x_proj, encoder_w, pos, cfg.encoder); + + // Copy the two outputs into the persistent NONE tensors so downstream + // graphs reference pure leaves (no producer-chain pull-in / recompute). + ggml_tensor * seg_out = ggml_cpy(enc_stage.ctx, outs.x_seg, x_seg_dev); + ggml_tensor * est_out = ggml_cpy(enc_stage.ctx, outs.x_est, x_est_dev); + ggml_set_output(seg_out); + ggml_set_output(est_out); + ggml_build_forward_expand(enc_stage.graph, seg_out); + ggml_build_forward_expand(enc_stage.graph, est_out); + dump_graph_support("encoder", backend, enc_stage.graph); + enc_stage.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!ggml_gallocr_alloc_graph(enc_stage.alloc, enc_stage.graph)) throw Error("alloc failed (encoder)"); + + // Input tensors to refresh on each compute. + enc_stage.inputs = {mel_in, pos}; + enc_stage.T = T; + + // positions are the fixed iota 0..T-1 — set once at build time. + std::vector pos_i(T); + for (int i = 0; i < T; ++i) pos_i[i] = i; + ggml_backend_tensor_set(pos, pos_i.data(), 0, pos_i.size() * sizeof(std::int32_t)); + } + + ggml_backend_tensor_set(enc_stage.inputs[0], mel, 0, ggml_nbytes(enc_stage.inputs[0])); + + if (ggml_backend_graph_compute(backend, enc_stage.graph) != GGML_STATUS_SUCCESS) { + throw Error("graph compute failed"); } - return num / (den + 1e-8f); } -} // namespace +// ============================================================================ +// Stage 2 — segmenter (one D3PM step, DBCache-aware) +// ============================================================================ void Model::Impl::run_segmenter_step( - const float * x_seg_host, int T, + int T, const std::int32_t * noise_mod3, float t_scalar, int language, std::vector & logits_out) { - const int D = cfg.embedding_dim; SegmenterCacheState & cache = seg_cache; // ---------- Fused fast path (DBCache disabled, the default) ---------- // One combined graph + a single backend submit is cheaper than the // 3-stage host-copy split below: fewer gallocr allocations, no extra // host round-trips, and it matches the pre-DBCache behavior exactly. + // The graph is built once per T and reused across all D3PM steps of the + // segment; x_seg is read straight from the encoder's device-resident + // output tensor. if (!cache.enabled) { - StageCtx s(backend, 512 * 1024 * 1024, 16384); - - ggml_tensor * xseg = ggml_new_tensor_3d(s.ctx, GGML_TYPE_F32, D, T, 1); - ggml_tensor * noise = ggml_new_tensor_1d(s.ctx, GGML_TYPE_I32, T); - ggml_tensor * t_tensor = ggml_new_tensor_3d(s.ctx, GGML_TYPE_F32, 1, 1, 1); - ggml_tensor * lang_tensor = ggml_new_tensor_1d(s.ctx, GGML_TYPE_I32, 1); - ggml_tensor * positions = ggml_new_tensor_1d(s.ctx, GGML_TYPE_I32, T); - for (auto * t : {xseg, noise, t_tensor, lang_tensor, positions}) ggml_set_input(t); - - auto outs = internal::build_segmenter_graph( - s.ctx, xseg, noise, t_tensor, lang_tensor, positions, segmenter_w, cfg); - - ggml_set_output(outs.logits); - ggml_build_forward_expand(s.graph, outs.logits); - if (outs.latent) { ggml_set_output(outs.latent); ggml_build_forward_expand(s.graph, outs.latent); } - s.dump_backend_support("segmenter"); - s.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!ggml_gallocr_alloc_graph(s.alloc, s.graph)) throw Error("alloc failed (segmenter/fused)"); - - ggml_backend_tensor_set(xseg, x_seg_host, 0, ggml_nbytes(xseg)); - ggml_backend_tensor_set(noise, noise_mod3, 0, T * sizeof(std::int32_t)); - ggml_backend_tensor_set(t_tensor, &t_scalar, 0, sizeof(float)); + if (!seg_stage.matches(T)) { + seg_stage.reset(); + + ggml_init_params ip{}; + ip.mem_size = 32 * 1024 * 1024; + ip.no_alloc = true; + seg_stage.ctx = ggml_init(ip); + seg_stage.graph = ggml_new_graph_custom(seg_stage.ctx, 16384, /*grads=*/false); + + ggml_tensor * noise = ggml_new_tensor_1d(seg_stage.ctx, GGML_TYPE_I32, T); + ggml_tensor * t_tensor = ggml_new_tensor_3d(seg_stage.ctx, GGML_TYPE_F32, 1, 1, 1); + ggml_tensor * lang_tensor = ggml_new_tensor_1d(seg_stage.ctx, GGML_TYPE_I32, 1); + ggml_tensor * positions = ggml_new_tensor_1d(seg_stage.ctx, GGML_TYPE_I32, T); + for (auto * t : {noise, t_tensor, lang_tensor, positions}) ggml_set_input(t); + + auto outs = internal::build_segmenter_graph( + seg_stage.ctx, x_seg_dev, noise, t_tensor, lang_tensor, positions, segmenter_w, cfg); + + ggml_set_output(outs.logits); + ggml_build_forward_expand(seg_stage.graph, outs.logits); + if (outs.latent) { ggml_set_output(outs.latent); ggml_build_forward_expand(seg_stage.graph, outs.latent); } + dump_graph_support("segmenter", backend, seg_stage.graph); + seg_stage.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!ggml_gallocr_alloc_graph(seg_stage.alloc, seg_stage.graph)) throw Error("alloc failed (segmenter/fused)"); + + seg_stage.inputs = {noise, t_tensor, lang_tensor, positions}; + seg_stage.outs = {outs.logits}; + std::vector pos(T); + for (int i = 0; i < T; ++i) pos[i] = i; + ggml_backend_tensor_set(positions, pos.data(), 0, pos.size() * sizeof(std::int32_t)); + seg_stage.T = T; + } + + ggml_backend_tensor_set(seg_stage.inputs[0], noise_mod3, 0, T * sizeof(std::int32_t)); + ggml_backend_tensor_set(seg_stage.inputs[1], &t_scalar, 0, sizeof(float)); const std::int32_t l = static_cast(language); - ggml_backend_tensor_set(lang_tensor, &l, 0, sizeof(std::int32_t)); - std::vector pos(T); - for (int i = 0; i < T; ++i) pos[i] = i; - ggml_backend_tensor_set(positions, pos.data(), 0, pos.size() * sizeof(std::int32_t)); + ggml_backend_tensor_set(seg_stage.inputs[2], &l, 0, sizeof(std::int32_t)); - s.compute(); + if (ggml_backend_graph_compute(backend, seg_stage.graph) != GGML_STATUS_SUCCESS) { + throw Error("graph compute failed (segmenter/fused)"); + } logits_out.resize(T); - ggml_backend_tensor_get(outs.logits, logits_out.data(), 0, logits_out.size() * sizeof(float)); + ggml_backend_tensor_get(seg_stage.outs[0], logits_out.data(), 0, logits_out.size() * sizeof(float)); + if (std::getenv("GAME_GGML_DUMP_LOGITS")) { + float mn = logits_out[0], mx = logits_out[0], sum = 0.0f; + int nneg = 0; + for (int i = 0; i < T; ++i) { + mn = std::min(mn, logits_out[i]); + mx = std::max(mx, logits_out[i]); + sum += logits_out[i]; + if (logits_out[i] < 0.0f) ++nneg; + } + std::fprintf(stderr, "[LOGITS] fused step t=%.4f n=%d min=%.4f max=%.4f mean=%.4f nneg=%d\n", + t_scalar, T, mn, mx, sum / T, nneg); + } return; } const int nf = std::min(std::max(cache.fn_blocks, 0), cfg.segmenter.num_layers); + const int N_seg = cfg.segmenter.num_layers; + const int nb = std::min(std::max(cache.bn_blocks, 0), N_seg - nf); + const int middle_end = N_seg - nb; // middle = [nf, middle_end) - // ---------- Stage A: front blocks (always executed) ---------- - std::vector x_front(D * T); + // ---------- Stage A: front blocks + cache metric ---------- + // x_front stays on the device; only the 1-float L1 metric is read back. + float fd = std::numeric_limits::infinity(); { - StageCtx s(backend, 256 * 1024 * 1024, 8192); - - ggml_tensor * xseg = ggml_new_tensor_3d(s.ctx, GGML_TYPE_F32, D, T, 1); - ggml_tensor * noise = ggml_new_tensor_1d(s.ctx, GGML_TYPE_I32, T); - ggml_tensor * t_tensor = ggml_new_tensor_3d(s.ctx, GGML_TYPE_F32, 1, 1, 1); - ggml_tensor * lang_tensor = ggml_new_tensor_1d(s.ctx, GGML_TYPE_I32, 1); - ggml_tensor * positions = ggml_new_tensor_1d(s.ctx, GGML_TYPE_I32, T); - for (auto * t : {xseg, noise, t_tensor, lang_tensor, positions}) ggml_set_input(t); - - ggml_tensor * out_front = internal::build_segmenter_front_graph( - s.ctx, xseg, noise, t_tensor, lang_tensor, positions, segmenter_w, cfg, nf); - - s.finalize(out_front, "segmenter/front"); + if (!seg_front_stage.matches(T, nf)) { + seg_front_stage.reset(); + + ggml_init_params ip{}; + ip.mem_size = 32 * 1024 * 1024; + ip.no_alloc = true; + seg_front_stage.ctx = ggml_init(ip); + seg_front_stage.graph = ggml_new_graph_custom(seg_front_stage.ctx, 8192, /*grads=*/false); + + ggml_tensor * noise = ggml_new_tensor_1d(seg_front_stage.ctx, GGML_TYPE_I32, T); + ggml_tensor * t_tensor = ggml_new_tensor_3d(seg_front_stage.ctx, GGML_TYPE_F32, 1, 1, 1); + ggml_tensor * lang_tensor = ggml_new_tensor_1d(seg_front_stage.ctx, GGML_TYPE_I32, 1); + ggml_tensor * positions = ggml_new_tensor_1d(seg_front_stage.ctx, GGML_TYPE_I32, T); + ggml_tensor * eps_t = ggml_new_tensor_1d(seg_front_stage.ctx, GGML_TYPE_F32, 1); + for (auto * t : {noise, t_tensor, lang_tensor, positions, eps_t}) ggml_set_input(t); + + ggml_tensor * out_front = internal::build_segmenter_front_graph( + seg_front_stage.ctx, x_seg_dev, noise, t_tensor, lang_tensor, + positions, segmenter_w, cfg, nf); + + // Copy into the persistent NONE tensor so the metric chain and + // the mid/add/update graphs reference a pure leaf (no producer + // chain pull-in). + ggml_tensor * xf = ggml_cpy(seg_front_stage.ctx, out_front, x_front_dev); + + // Device-side L1 metric — mirrors host front_delta() up to float + // summation order: fd = sum|x - prev| / (sum|prev| + eps). + ggml_tensor * diff = ggml_abs(seg_front_stage.ctx, + ggml_sub(seg_front_stage.ctx, x_front_dev, prev_front_dev)); + ggml_tensor * num = ggml_sum(seg_front_stage.ctx, diff); + ggml_tensor * den = ggml_sum(seg_front_stage.ctx, + ggml_abs(seg_front_stage.ctx, prev_front_dev)); + ggml_tensor * fd_t = ggml_div(seg_front_stage.ctx, num, + ggml_add(seg_front_stage.ctx, den, eps_t)); + + ggml_set_output(xf); + ggml_build_forward_expand(seg_front_stage.graph, xf); + ggml_set_output(fd_t); + ggml_build_forward_expand(seg_front_stage.graph, fd_t); + dump_graph_support("segmenter/front", backend, seg_front_stage.graph); + seg_front_stage.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!ggml_gallocr_alloc_graph(seg_front_stage.alloc, seg_front_stage.graph)) + throw Error("alloc failed (segmenter/front)"); + + seg_front_stage.inputs = {noise, t_tensor, lang_tensor, positions, eps_t}; + seg_front_stage.outs = {xf, fd_t}; + std::vector pos(T); + for (int i = 0; i < T; ++i) pos[i] = i; + ggml_backend_tensor_set(positions, pos.data(), 0, pos.size() * sizeof(std::int32_t)); + const float eps = 1e-8f; + ggml_backend_tensor_set(eps_t, &eps, 0, sizeof(float)); + seg_front_stage.T = T; + seg_front_stage.key = nf; + } - ggml_backend_tensor_set(xseg, x_seg_host, 0, ggml_nbytes(xseg)); - ggml_backend_tensor_set(noise, noise_mod3, 0, T * sizeof(std::int32_t)); - ggml_backend_tensor_set(t_tensor, &t_scalar, 0, sizeof(float)); - const std::int32_t l = static_cast(language); - ggml_backend_tensor_set(lang_tensor, &l, 0, sizeof(std::int32_t)); - std::vector pos(T); - for (int i = 0; i < T; ++i) pos[i] = i; - ggml_backend_tensor_set(positions, pos.data(), 0, pos.size() * sizeof(std::int32_t)); + ggml_backend_tensor_set(seg_front_stage.inputs[0], noise_mod3, 0, T * sizeof(std::int32_t)); + ggml_backend_tensor_set(seg_front_stage.inputs[1], &t_scalar, 0, sizeof(float)); + const std::int32_t l0 = static_cast(language); + ggml_backend_tensor_set(seg_front_stage.inputs[2], &l0, 0, sizeof(std::int32_t)); - s.compute(); - ggml_backend_tensor_get(out_front, x_front.data(), 0, x_front.size() * sizeof(float)); + if (ggml_backend_graph_compute(backend, seg_front_stage.graph) != GGML_STATUS_SUCCESS) { + throw Error("graph compute failed (segmenter/front)"); + } + if (cache.valid) { + ggml_backend_tensor_get(seg_front_stage.outs[1], &fd, 0, sizeof(float)); + } } // ---------- Decide cache hit ---------- bool use_cache = false; - const int N_seg = cfg.segmenter.num_layers; - if (cache.enabled && cache.valid && - cache.step >= cache.warmup && - cache.prev_front.size() == x_front.size()) { - const float fd = front_delta(x_front.data(), cache.prev_front.data(), D * T); - + if (cache.enabled && cache.valid && cache.step >= cache.warmup) { // Reuse window: cache only mid-schedule steps (first/last computed // fully unless the window covers them). bool in_window = true; @@ -351,77 +409,188 @@ void Model::Impl::run_segmenter_step( use_cache = (fd < cache.threshold) && in_window && err_ok && cont_ok; } - // ---------- Stage B: tail blocks (middle reused on hit, back always done) ---------- - const int nb = std::min(std::max(cache.bn_blocks, 0), N_seg - nf); - const int middle_end = N_seg - nb; // middle = [nf, middle_end) - std::vector x_mid(D * T); - std::vector x_out(D * T); + // ---------- Stage B: tail blocks (all device-resident) ---------- + // Small helper: build a persistent graph with a single cpy-to-persistent + // output, run it, and hand the graph handle back. + // * add graph (hit): x_mid_dev = x_front + tail_delta + // * mid graph (miss): x_mid_dev = tail(middle)(x_front) + // * update graph(miss): tail_delta_dev = x_mid - x_front; + // prev_front_dev = x_front + // * back graph: x_out_dev = tail(back)(x_mid_dev) + // All of these are built once per T and reference the device-resident + // x_front / db tensors as leaves, so no D×T host transfer happens on the + // cache path at all. + + if (!seg_add_stage.matches(T, nb ? 1 : 0)) { + seg_add_stage.reset(); + ggml_init_params ip{}; + ip.mem_size = 8 * 1024 * 1024; + ip.no_alloc = true; + seg_add_stage.ctx = ggml_init(ip); + seg_add_stage.graph = ggml_new_graph_custom(seg_add_stage.ctx, 2048, /*grads=*/false); // visited-set must hold the referenced encoder/front chains + ggml_tensor * mid = ggml_add(seg_add_stage.ctx, x_front_dev, tail_delta_dev); + ggml_tensor * md = ggml_cpy(seg_add_stage.ctx, mid, x_mid_dev); + ggml_set_output(md); + ggml_build_forward_expand(seg_add_stage.graph, md); + if (nb == 0) { // no back slice: head reads x_out_dev straight away + ggml_tensor * od = ggml_cpy(seg_add_stage.ctx, x_mid_dev, x_out_dev); + ggml_set_output(od); + ggml_build_forward_expand(seg_add_stage.graph, od); + } + dump_graph_support("segmenter/add", backend, seg_add_stage.graph); + seg_add_stage.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!ggml_gallocr_alloc_graph(seg_add_stage.alloc, seg_add_stage.graph)) + throw Error("alloc failed (segmenter/add)"); + seg_add_stage.outs = {md}; + seg_add_stage.T = T; + seg_add_stage.key = nb ? 1 : 0; + } - auto run_tail_range = [&](const float * in_host, int start, int end, - std::vector & out_host) { - StageCtx s(backend, 512 * 1024 * 1024, 16384); - ggml_tensor * x_in = ggml_new_tensor_3d(s.ctx, GGML_TYPE_F32, D, T, 1); - ggml_tensor * positions = ggml_new_tensor_1d(s.ctx, GGML_TYPE_I32, T); - for (auto * t : {x_in, positions}) ggml_set_input(t); + if (!seg_mid_stage.matches(T, (static_cast(nf) << 20) | middle_end)) { + seg_mid_stage.reset(); + ggml_init_params ip{}; + ip.mem_size = 32 * 1024 * 1024; + ip.no_alloc = true; + seg_mid_stage.ctx = ggml_init(ip); + seg_mid_stage.graph = ggml_new_graph_custom(seg_mid_stage.ctx, 16384, /*grads=*/false); + ggml_tensor * positions = ggml_new_tensor_1d(seg_mid_stage.ctx, GGML_TYPE_I32, T); + ggml_set_input(positions); + auto outs = internal::build_segmenter_tail_graph( + seg_mid_stage.ctx, x_front_dev, positions, segmenter_w, cfg, nf, middle_end); + ggml_tensor * md = ggml_cpy(seg_mid_stage.ctx, outs.x_run, x_mid_dev); + ggml_set_output(md); + ggml_build_forward_expand(seg_mid_stage.graph, md); + dump_graph_support("segmenter/mid", backend, seg_mid_stage.graph); + seg_mid_stage.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!ggml_gallocr_alloc_graph(seg_mid_stage.alloc, seg_mid_stage.graph)) + throw Error("alloc failed (segmenter/mid)"); + seg_mid_stage.inputs = {positions}; + seg_mid_stage.outs = {md}; + if (middle_end > nf) { // empty range => x_run = x_front passthrough, positions unused + std::vector pos(T); + for (int i = 0; i < T; ++i) pos[i] = i; + ggml_backend_tensor_set(positions, pos.data(), 0, pos.size() * sizeof(std::int32_t)); + } + seg_mid_stage.T = T; + seg_mid_stage.key = (static_cast(nf) << 20) | middle_end; + } + if (!seg_update_stage.matches(T, nb ? 1 : 0)) { + seg_update_stage.reset(); + ggml_init_params ip{}; + ip.mem_size = 8 * 1024 * 1024; + ip.no_alloc = true; + seg_update_stage.ctx = ggml_init(ip); + seg_update_stage.graph = ggml_new_graph_custom(seg_update_stage.ctx, 2048, /*grads=*/false); + ggml_tensor * sub = ggml_sub(seg_update_stage.ctx, x_mid_dev, x_front_dev); + ggml_tensor * td = ggml_cpy(seg_update_stage.ctx, sub, tail_delta_dev); + ggml_tensor * pf = ggml_cpy(seg_update_stage.ctx, x_front_dev, prev_front_dev); + ggml_set_output(td); + ggml_build_forward_expand(seg_update_stage.graph, td); + ggml_set_output(pf); + ggml_build_forward_expand(seg_update_stage.graph, pf); + if (nb == 0) { + ggml_tensor * od = ggml_cpy(seg_update_stage.ctx, x_mid_dev, x_out_dev); + ggml_set_output(od); + ggml_build_forward_expand(seg_update_stage.graph, od); + } + dump_graph_support("segmenter/update", backend, seg_update_stage.graph); + seg_update_stage.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!ggml_gallocr_alloc_graph(seg_update_stage.alloc, seg_update_stage.graph)) + throw Error("alloc failed (segmenter/update)"); + seg_update_stage.T = T; + seg_update_stage.key = nb ? 1 : 0; + } + + if (nb > 0 && !seg_back_stage.matches(T, (static_cast(middle_end) << 20) | N_seg)) { + seg_back_stage.reset(); + ggml_init_params ip{}; + ip.mem_size = 32 * 1024 * 1024; + ip.no_alloc = true; + seg_back_stage.ctx = ggml_init(ip); + seg_back_stage.graph = ggml_new_graph_custom(seg_back_stage.ctx, 16384, /*grads=*/false); + ggml_tensor * positions = ggml_new_tensor_1d(seg_back_stage.ctx, GGML_TYPE_I32, T); + ggml_set_input(positions); auto outs = internal::build_segmenter_tail_graph( - s.ctx, x_in, positions, segmenter_w, cfg, start, end); - - // Both outputs must be registered BEFORE gallocr allocation — expanding - // after alloc leaves latent unallocated = UB on compute. - ggml_set_output(outs.x_run); - ggml_build_forward_expand(s.graph, outs.x_run); - if (outs.latent) { ggml_set_output(outs.latent); ggml_build_forward_expand(s.graph, outs.latent); } - s.dump_backend_support("segmenter/tail"); - s.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!ggml_gallocr_alloc_graph(s.alloc, s.graph)) throw Error("alloc failed (segmenter/tail)"); - - ggml_backend_tensor_set(x_in, in_host, 0, x_front.size() * sizeof(float)); + seg_back_stage.ctx, x_mid_dev, positions, segmenter_w, cfg, middle_end, N_seg); + ggml_tensor * od = ggml_cpy(seg_back_stage.ctx, outs.x_run, x_out_dev); + ggml_set_output(od); + ggml_build_forward_expand(seg_back_stage.graph, od); + if (outs.latent) { ggml_set_output(outs.latent); ggml_build_forward_expand(seg_back_stage.graph, outs.latent); } + dump_graph_support("segmenter/back", backend, seg_back_stage.graph); + seg_back_stage.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!ggml_gallocr_alloc_graph(seg_back_stage.alloc, seg_back_stage.graph)) + throw Error("alloc failed (segmenter/back)"); + seg_back_stage.inputs = {positions}; + seg_back_stage.outs = {od}; std::vector pos(T); for (int i = 0; i < T; ++i) pos[i] = i; ggml_backend_tensor_set(positions, pos.data(), 0, pos.size() * sizeof(std::int32_t)); - - s.compute(); - ggml_backend_tensor_get(outs.x_run, out_host.data(), 0, out_host.size() * sizeof(float)); - }; + seg_back_stage.T = T; + seg_back_stage.key = (static_cast(middle_end) << 20) | N_seg; + } if (use_cache) { - // Reconstruct the middle from the cached delta, then always run the - // back slice (if any) so the output-end blocks see accurate input. - for (int i = 0; i < D * T; ++i) x_mid[i] = x_front[i] + cache.tail_delta[i]; - if (nb > 0) run_tail_range(x_mid.data(), middle_end, N_seg, x_out); - else for (int i = 0; i < D * T; ++i) x_out[i] = x_mid[i]; + if (ggml_backend_graph_compute(backend, seg_add_stage.graph) != GGML_STATUS_SUCCESS) { + throw Error("graph compute failed (segmenter/add)"); + } + if (nb > 0) { + if (ggml_backend_graph_compute(backend, seg_back_stage.graph) != GGML_STATUS_SUCCESS) { + throw Error("graph compute failed (segmenter/back)"); + } + } ++cache.hits; ++cache.cont_cnt; } else { - run_tail_range(x_front.data(), nf, middle_end, x_mid); - if (nb > 0) run_tail_range(x_mid.data(), middle_end, N_seg, x_out); - else for (int i = 0; i < D * T; ++i) x_out[i] = x_mid[i]; - cache.tail_delta.resize(D * T); - for (int i = 0; i < D * T; ++i) cache.tail_delta[i] = x_mid[i] - x_front[i]; - cache.prev_front = x_front; + if (ggml_backend_graph_compute(backend, seg_mid_stage.graph) != GGML_STATUS_SUCCESS) { + throw Error("graph compute failed (segmenter/mid)"); + } + if (ggml_backend_graph_compute(backend, seg_update_stage.graph) != GGML_STATUS_SUCCESS) { + throw Error("graph compute failed (segmenter/update)"); + } + if (nb > 0) { + if (ggml_backend_graph_compute(backend, seg_back_stage.graph) != GGML_STATUS_SUCCESS) { + throw Error("graph compute failed (segmenter/back)"); + } + } cache.valid = true; cache.acc_err = 0.0f; cache.cont_cnt = 0; ++cache.misses; } - // ---------- Stage C: head (output norm + proj) ---------- + // ---------- Stage C: head (output norm + proj), input = device x_out_dev ---------- { - StageCtx s(backend, 16 * 1024 * 1024, 256); - - ggml_tensor * x_in = ggml_new_tensor_3d(s.ctx, GGML_TYPE_F32, D, T, 1); - ggml_set_input(x_in); - - ggml_tensor * logits = internal::build_segmenter_head_graph(s.ctx, x_in, segmenter_w, cfg); - - s.finalize(logits, "segmenter/head"); + if (!seg_head_stage.matches(T)) { + seg_head_stage.reset(); + + ggml_init_params ip{}; + ip.mem_size = 16 * 1024 * 1024; + ip.no_alloc = true; + seg_head_stage.ctx = ggml_init(ip); + seg_head_stage.graph = ggml_new_graph_custom(seg_head_stage.ctx, 256, /*grads=*/false); + + ggml_tensor * logits = internal::build_segmenter_head_graph( + seg_head_stage.ctx, x_out_dev, segmenter_w, cfg); + + ggml_set_output(logits); + ggml_build_forward_expand(seg_head_stage.graph, logits); + dump_graph_support("segmenter/head", backend, seg_head_stage.graph); + seg_head_stage.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!ggml_gallocr_alloc_graph(seg_head_stage.alloc, seg_head_stage.graph)) + throw Error("alloc failed (segmenter/head)"); + + seg_head_stage.outs = {logits}; + seg_head_stage.T = T; + } - ggml_backend_tensor_set(x_in, x_out.data(), 0, x_out.size() * sizeof(float)); - s.compute(); + if (ggml_backend_graph_compute(backend, seg_head_stage.graph) != GGML_STATUS_SUCCESS) { + throw Error("graph compute failed (segmenter/head)"); + } logits_out.resize(T); - ggml_backend_tensor_get(logits, logits_out.data(), 0, logits_out.size() * sizeof(float)); + ggml_backend_tensor_get(seg_head_stage.outs[0], logits_out.data(), + 0, logits_out.size() * sizeof(float)); } ++cache.step; @@ -431,48 +600,94 @@ void Model::Impl::run_segmenter_step( } } +// ============================================================================ +// DBCache device-resident tensors +// ============================================================================ + +void Model::Impl::ensure_db_tensors(int T) { + const int D = cfg.embedding_dim; + if (db_T == T && x_seg_dev) return; + + if (dbctx) { + ggml_backend_buffer_free(dbbuf); + ggml_free(dbctx); + dbctx = nullptr; dbbuf = nullptr; + x_seg_dev = x_est_dev = x_front_dev = nullptr; + prev_front_dev = tail_delta_dev = x_mid_dev = x_out_dev = nullptr; + } + ggml_init_params ip{}; + ip.mem_size = 128 * 1024; + ip.no_alloc = true; + dbctx = ggml_init(ip); + x_seg_dev = ggml_new_tensor_3d(dbctx, GGML_TYPE_F32, D, T, 1); + x_est_dev = ggml_new_tensor_3d(dbctx, GGML_TYPE_F32, D, T, 1); + x_front_dev = ggml_new_tensor_3d(dbctx, GGML_TYPE_F32, D, T, 1); + prev_front_dev = ggml_new_tensor_3d(dbctx, GGML_TYPE_F32, D, T, 1); + tail_delta_dev = ggml_new_tensor_3d(dbctx, GGML_TYPE_F32, D, T, 1); + x_mid_dev = ggml_new_tensor_3d(dbctx, GGML_TYPE_F32, D, T, 1); + x_out_dev = ggml_new_tensor_3d(dbctx, GGML_TYPE_F32, D, T, 1); + dbbuf = ggml_backend_alloc_ctx_tensors(dbctx, backend); + if (!dbbuf) throw Error("failed to allocate device-resident tensors"); + // Zero prev_front so the (unread until cache.valid) metric node never + // consumes uninitialised device memory. + std::vector zeros(static_cast(D) * T, 0.0f); + ggml_backend_tensor_set(prev_front_dev, zeros.data(), 0, zeros.size() * sizeof(float)); + db_T = T; +} + // ============================================================================ // Stage 3 — estimator (regions → pool_logits) // ============================================================================ void Model::Impl::run_estimator( - const float * x_est_host, int T, - const std::int32_t * regions, int N, + int T, const std::int32_t * regions, int N, std::vector & pool_logits_out) { - const int D = cfg.embedding_dim; const int S = N + T; - StageCtx s(backend, 512 * 1024 * 1024, 16384); - - ggml_tensor * xest = ggml_new_tensor_3d(s.ctx, GGML_TYPE_F32, D, T, 1); - ggml_tensor * regions_mod = ggml_new_tensor_1d(s.ctx, GGML_TYPE_I32, T); - ggml_tensor * positions = ggml_new_tensor_1d(s.ctx, GGML_TYPE_I32, S); - ggml_tensor * region_ids = ggml_new_tensor_1d(s.ctx, GGML_TYPE_I32, S); - ggml_tensor * mask_fp16 = ggml_new_tensor_4d(s.ctx, GGML_TYPE_F16, S, S, 1, 1); - for (auto * t : {xest, regions_mod, positions, region_ids, mask_fp16}) ggml_set_input(t); - - auto outs = internal::build_estimator_graph( - s.ctx, xest, regions_mod, positions, region_ids, mask_fp16, N, estimator_w, cfg); + if (!est_stage.matches(T, N)) { + est_stage.reset(); - ggml_set_output(outs.pool_logits); - ggml_build_forward_expand(s.graph, outs.pool_logits); - s.dump_backend_support("estimator"); - s.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!ggml_gallocr_alloc_graph(s.alloc, s.graph)) throw Error("alloc failed (estimator)"); + ggml_init_params ip{}; + ip.mem_size = 32 * 1024 * 1024; + ip.no_alloc = true; + est_stage.ctx = ggml_init(ip); + est_stage.graph = ggml_new_graph_custom(est_stage.ctx, 16384, /*grads=*/false); + + ggml_tensor * regions_mod = ggml_new_tensor_1d(est_stage.ctx, GGML_TYPE_I32, T); + ggml_tensor * positions = ggml_new_tensor_1d(est_stage.ctx, GGML_TYPE_I32, S); + ggml_tensor * region_ids = ggml_new_tensor_1d(est_stage.ctx, GGML_TYPE_I32, S); + ggml_tensor * mask_fp16 = ggml_new_tensor_4d(est_stage.ctx, GGML_TYPE_F16, S, S, 1, 1); + for (auto * t : {regions_mod, positions, region_ids, mask_fp16}) ggml_set_input(t); + + // x_est input is the device-resident encoder output — referenced + // cross-context, never copied. + auto outs = internal::build_estimator_graph( + est_stage.ctx, x_est_dev, regions_mod, positions, region_ids, mask_fp16, + N, estimator_w, cfg); + + ggml_set_output(outs.pool_logits); + ggml_build_forward_expand(est_stage.graph, outs.pool_logits); + dump_graph_support("estimator", backend, est_stage.graph); + est_stage.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!ggml_gallocr_alloc_graph(est_stage.alloc, est_stage.graph)) throw Error("alloc failed (estimator)"); + + est_stage.inputs = {regions_mod, positions, region_ids, mask_fp16}; + est_stage.outs = {outs.pool_logits}; + est_stage.T = T; + est_stage.key = N; + } // Host-side prep. - ggml_backend_tensor_set(xest, x_est_host, 0, ggml_nbytes(xest)); - std::vector rmod(T); for (int i = 0; i < T; ++i) rmod[i] = regions[i] % cfg.region_cycle_len; - ggml_backend_tensor_set(regions_mod, rmod.data(), 0, rmod.size() * sizeof(std::int32_t)); + ggml_backend_tensor_set(est_stage.inputs[0], rmod.data(), 0, rmod.size() * sizeof(std::int32_t)); // Global positions: pool 0..N-1, x 0..T-1 std::vector gpos(S); for (int i = 0; i < N; ++i) gpos[i] = i; for (int i = 0; i < T; ++i) gpos[N + i] = i; - ggml_backend_tensor_set(positions, gpos.data(), 0, gpos.size() * sizeof(std::int32_t)); + ggml_backend_tensor_set(est_stage.inputs[1], gpos.data(), 0, gpos.size() * sizeof(std::int32_t)); // Region RoPE indices: pool = 0 (R=1, use_pool_offset=false); // x = local_position_within_region + R (= +1). @@ -487,15 +702,17 @@ void Model::Impl::run_estimator( ++cur_local; } } - ggml_backend_tensor_set(region_ids, ridx.data(), 0, ridx.size() * sizeof(std::int32_t)); + ggml_backend_tensor_set(est_stage.inputs[2], ridx.data(), 0, ridx.size() * sizeof(std::int32_t)); auto mask = internal::ops::build_joint_attn_mask_fp16(regions, T, N); - ggml_backend_tensor_set(mask_fp16, mask.data(), 0, mask.size() * sizeof(std::uint16_t)); + ggml_backend_tensor_set(est_stage.inputs[3], mask.data(), 0, mask.size() * sizeof(std::uint16_t)); - s.compute(); + if (ggml_backend_graph_compute(backend, est_stage.graph) != GGML_STATUS_SUCCESS) { + throw Error("graph compute failed (estimator)"); + } - pool_logits_out.resize(ggml_nelements(outs.pool_logits)); - ggml_backend_tensor_get(outs.pool_logits, pool_logits_out.data(), + pool_logits_out.resize(ggml_nelements(est_stage.outs[0])); + ggml_backend_tensor_get(est_stage.outs[0], pool_logits_out.data(), 0, pool_logits_out.size() * sizeof(float)); } @@ -588,11 +805,11 @@ InferResult Model::Impl::infer_with_rng( // covers: mel extraction, spectrogram_projection, 4× EBF blocks, // output split. The mel sub-stage runs on CPU (pocketfft STFT // + mel filterbank mul + log); everything after is on the backend. - std::vector x_seg_host, x_est_host; + // x_seg / x_est stay resident in the encoder's backend buffer. { auto _ = prof.scope_encoder(); auto mel = mel_extractor->forward(waveform, n_samples); // [T, 80] - run_encoder(mel.data(), T, x_seg_host, x_est_host); + run_encoder(mel.data(), T); } // --- 3) D3PM loop (segmenter) @@ -605,24 +822,14 @@ InferResult Model::Impl::infer_with_rng( // with a single step there is nothing to cache, so the fused single-graph // path stays active (cache.enabled==false) even when a threshold is set // — that avoids paying the 3-stage split cost for --nsteps 1. - // EP-aware default: on CPU the split-cache path wins big (measured -45% - // at nsteps=8); on GPU backends its per-step host round-trips regress - // quantized weights (measured +20% on Vulkan+Q8), so default it off - // unless the user picks a threshold explicitly. + // Default 0.25 on every backend: the DBCache decision metric and the + // middle reconstruction now run fully on-device (1-float readback), so + // GPU backends no longer pay the D×T host round-trip that regressed + // quantized weights (+20% measured on Vulkan+Q8 with the old host-side + // path). Verified on Vulkan/RTX 2070: device-side split path matches + // CPU note output exactly. float thr = params.db_cache_threshold; - if (thr < 0.0f) { - const bool gpu = [this] { - const char * bn = internal::backend_name(backend); - if (!bn) return false; - std::string s(bn); - std::transform(s.begin(), s.end(), s.begin(), - [](unsigned char c) { return static_cast(std::tolower(c)); }); - return s.find("vulkan") != std::string::npos || - s.find("cuda") != std::string::npos || - s.find("metal") != std::string::npos; - }(); - thr = gpu ? 0.0f : 0.25f; - } + if (thr < 0.0f) thr = 0.25f; seg_cache.enabled = thr > 0.0f && ts.size() > 1; seg_cache.threshold = thr; seg_cache.fn_blocks = params.db_cache_fn_blocks; @@ -667,7 +874,7 @@ InferResult Model::Impl::infer_with_rng( { auto _ = prof.scope_segmenter(); - run_segmenter_step(x_seg_host.data(), T, + run_segmenter_step(T, noise_mod.data(), ti, params.language, logits); } @@ -701,7 +908,7 @@ InferResult Model::Impl::infer_with_rng( // --- 5) estimator std::vector pool_logits; { auto _ = prof.scope_estimator(); - run_estimator(x_est_host.data(), T, regions.data(), N, pool_logits); } + run_estimator(T, regions.data(), N, pool_logits); } // --- 6) pitch decode { diff --git a/src/model_encoder.cpp b/src/model_encoder.cpp index 7bf943a..2b2d607 100644 --- a/src/model_encoder.cpp +++ b/src/model_encoder.cpp @@ -31,6 +31,10 @@ static ops::EBFBlockWeights bind_ebf_layer( B.w_norm1 = W.get(p + "norm1.weight"); B.w_ffn1_ln1 = W.get(p + "ffn1.ln1.weight"); B.b_ffn1_ln1 = W.get(p + "ffn1.ln1.bias"); + B.w_ffn1_ln1_a = W.try_get(p + "ffn1.ln1.weight.a"); // F-1 split halves + B.b_ffn1_ln1_a = W.try_get(p + "ffn1.ln1.bias.a"); + B.w_ffn1_ln1_b = W.try_get(p + "ffn1.ln1.weight.b"); + B.b_ffn1_ln1_b = W.try_get(p + "ffn1.ln1.bias.b"); B.w_ffn1_ln2 = W.get(p + "ffn1.ln2.weight"); B.b_ffn1_ln2 = W.get(p + "ffn1.ln2.bias"); if (cfg.use_ls) B.w_lay_scale1 = W.get(p + "lay_scale1.scale"); @@ -39,6 +43,10 @@ static ops::EBFBlockWeights bind_ebf_layer( B.w_norm2 = W.get(p + "norm2.weight"); B.w_ffn2_ln1 = W.get(p + "ffn2.ln1.weight"); B.b_ffn2_ln1 = W.get(p + "ffn2.ln1.bias"); + B.w_ffn2_ln1_a = W.try_get(p + "ffn2.ln1.weight.a"); // F-1 split halves + B.b_ffn2_ln1_a = W.try_get(p + "ffn2.ln1.bias.a"); + B.w_ffn2_ln1_b = W.try_get(p + "ffn2.ln1.weight.b"); + B.b_ffn2_ln1_b = W.try_get(p + "ffn2.ln1.bias.b"); B.w_ffn2_ln2 = W.get(p + "ffn2.ln2.weight"); B.b_ffn2_ln2 = W.get(p + "ffn2.ln2.bias"); if (cfg.use_ls) B.w_lay_scale3 = W.get(p + "lay_scale3.scale"); diff --git a/src/model_estimator.cpp b/src/model_estimator.cpp index add27a7..b525a8a 100644 --- a/src/model_estimator.cpp +++ b/src/model_estimator.cpp @@ -25,10 +25,18 @@ static ops::JEBFBlockWeights bind_jebf_layer( B.w_norm_ffn1_pool = W.get(p + "norm_ffn1_pool.weight"); B.w_ffn1_x_ln1 = W.get(p + "ffn1_x.ln1.weight"); B.b_ffn1_x_ln1 = W.get(p + "ffn1_x.ln1.bias"); + B.w_ffn1_x_ln1_a = W.try_get(p + "ffn1_x.ln1.weight.a"); // F-1 split halves + B.b_ffn1_x_ln1_a = W.try_get(p + "ffn1_x.ln1.bias.a"); + B.w_ffn1_x_ln1_b = W.try_get(p + "ffn1_x.ln1.weight.b"); + B.b_ffn1_x_ln1_b = W.try_get(p + "ffn1_x.ln1.bias.b"); B.w_ffn1_x_ln2 = W.get(p + "ffn1_x.ln2.weight"); B.b_ffn1_x_ln2 = W.get(p + "ffn1_x.ln2.bias"); B.w_ffn1_pool_ln1 = W.get(p + "ffn1_pool.ln1.weight"); B.b_ffn1_pool_ln1 = W.get(p + "ffn1_pool.ln1.bias"); + B.w_ffn1_pool_ln1_a = W.try_get(p + "ffn1_pool.ln1.weight.a"); // F-1 split halves + B.b_ffn1_pool_ln1_a = W.try_get(p + "ffn1_pool.ln1.bias.a"); + B.w_ffn1_pool_ln1_b = W.try_get(p + "ffn1_pool.ln1.weight.b"); + B.b_ffn1_pool_ln1_b = W.try_get(p + "ffn1_pool.ln1.bias.b"); B.w_ffn1_pool_ln2 = W.get(p + "ffn1_pool.ln2.weight"); B.b_ffn1_pool_ln2 = W.get(p + "ffn1_pool.ln2.bias"); if (cfg.use_ls) { @@ -41,10 +49,18 @@ static ops::JEBFBlockWeights bind_jebf_layer( B.w_norm_ffn2_pool = W.get(p + "norm_ffn2_pool.weight"); B.w_ffn2_x_ln1 = W.get(p + "ffn2_x.ln1.weight"); B.b_ffn2_x_ln1 = W.get(p + "ffn2_x.ln1.bias"); + B.w_ffn2_x_ln1_a = W.try_get(p + "ffn2_x.ln1.weight.a"); // F-1 split halves + B.b_ffn2_x_ln1_a = W.try_get(p + "ffn2_x.ln1.bias.a"); + B.w_ffn2_x_ln1_b = W.try_get(p + "ffn2_x.ln1.weight.b"); + B.b_ffn2_x_ln1_b = W.try_get(p + "ffn2_x.ln1.bias.b"); B.w_ffn2_x_ln2 = W.get(p + "ffn2_x.ln2.weight"); B.b_ffn2_x_ln2 = W.get(p + "ffn2_x.ln2.bias"); B.w_ffn2_pool_ln1 = W.get(p + "ffn2_pool.ln1.weight"); B.b_ffn2_pool_ln1 = W.get(p + "ffn2_pool.ln1.bias"); + B.w_ffn2_pool_ln1_a = W.try_get(p + "ffn2_pool.ln1.weight.a"); // F-1 split halves + B.b_ffn2_pool_ln1_a = W.try_get(p + "ffn2_pool.ln1.bias.a"); + B.w_ffn2_pool_ln1_b = W.try_get(p + "ffn2_pool.ln1.weight.b"); + B.b_ffn2_pool_ln1_b = W.try_get(p + "ffn2_pool.ln1.bias.b"); B.w_ffn2_pool_ln2 = W.get(p + "ffn2_pool.ln2.weight"); B.b_ffn2_pool_ln2 = W.get(p + "ffn2_pool.ln2.bias"); if (cfg.use_ls) { diff --git a/src/model_impl.h b/src/model_impl.h index 4c07eb4..34545f6 100644 --- a/src/model_impl.h +++ b/src/model_impl.h @@ -7,6 +7,9 @@ #include "game_ggml/model.h" #include "game_ggml/mel.h" +#include +#include + #include "gguf_io.h" #include "model_encoder.h" #include "model_segmenter.h" @@ -47,8 +50,10 @@ struct SegmenterCacheState { float acc_err = 0.0f; // accumulated skipped delta error int cont_cnt = 0; // consecutive hits - std::vector prev_front; // x_front of the previous step - std::vector tail_delta; // x_out - x_front of the previous full pass + // Device-side cache state (see run_segmenter_step): the decision metric + // and the delta reconstruction live on the backend, so GPU EPs do not pay + // D×T host round-trips per D3PM step. `valid` mirrors "prev_front is + // meaningful" — set after the first full pass. bool valid = false; // tail_delta / prev_front available void reset() { @@ -57,12 +62,44 @@ struct SegmenterCacheState { misses = 0; acc_err = 0.0f; cont_cnt = 0; - prev_front.clear(); - tail_delta.clear(); valid = false; } }; +// Persistent (stage, T) compute context. The graph + gallocr are built once +// per frame count T and reused across the nsteps D3PM iterations of one +// audio segment (and across chunks with identical T), eliminating per-step +// graph construction, gallocr allocation and — on GPU backends — the +// device-buffer alloc/free that goes with it. Rebuilding happens only when +// T (or the caller's stage split parameters) change. +struct PersistentStage { + ggml_context * ctx = nullptr; + ggml_cgraph * graph = nullptr; + ggml_gallocr_t alloc = nullptr; + int T = 0; // frame count this graph was built for + int64_t key = 0; // stage-specific discriminator (e.g. block range) + + // Tensors created in ctx that the caller refreshes via tensor_set before + // each compute (input tensors). The first entry of `outs` is the stage's + // primary output (logits etc.). + std::vector inputs; + std::vector outs; + + bool matches(int t, int64_t k = 0) const { return graph != nullptr && T == t && key == k; } + + void reset() { + if (alloc) ggml_gallocr_free(alloc); + if (ctx) ggml_free(ctx); + alloc = nullptr; ctx = nullptr; graph = nullptr; + inputs.clear(); outs.clear(); + T = 0; key = 0; + } + ~PersistentStage() { reset(); } + PersistentStage() = default; + PersistentStage(const PersistentStage &) = delete; + PersistentStage & operator=(const PersistentStage &) = delete; +}; + struct Model::Impl { GameModelConfig cfg; ggml_backend_t backend = nullptr; @@ -77,6 +114,45 @@ struct Model::Impl { // DBCache state for the segmenter (cross-D3PM-step). SegmenterCacheState seg_cache; + // Persistent stage graphs (rebuilt lazily when T changes). The encoder + // stage also owns the device-resident x_seg / x_est output tensors that + // the segmenter / estimator stages reference directly (no host round-trip + // between stages on GPU backends). + PersistentStage enc_stage; // encoder: mel -> x_seg, x_est (device-resident) + PersistentStage seg_stage; // segmenter fused path (DBCache disabled) + PersistentStage seg_front_stage; // segmenter DBCache: front blocks (+ fd metric) + PersistentStage seg_add_stage; // segmenter DBCache: x_mid = x_front + tail_delta (hit) + PersistentStage seg_mid_stage; // segmenter DBCache: middle tail range (miss) + PersistentStage seg_update_stage; // segmenter DBCache: refresh tail_delta/prev_front + PersistentStage seg_back_stage; // segmenter DBCache: back tail range + PersistentStage seg_head_stage; // segmenter DBCache: output norm + proj + PersistentStage est_stage; // estimator + + // Device-resident cross-stage activations — persistent NONE tensors in + // `dbctx`/`dbbuf`, written by the encoder (x_seg/x_est) or the front + // graph (x_front) via ggml_cpy. Being NONE (not op nodes), downstream + // graphs reference them as pure leaves: ggml never pulls the producer + // chain into a downstream graph, which keeps each stage graph at its own + // size (a CONT-view leaf would drag the entire encoder/front chain into + // every consumer graph and re-compute it). + ggml_tensor * x_seg_dev = nullptr; + ggml_tensor * x_est_dev = nullptr; + + // Device-resident DBCache tensors (D×T each, rebuilt when T changes): + // x_front_dev — front-graph output (metric reference + reuse base) + // prev_front_dev — x_front of the previous full pass (metric reference) + // tail_delta_dev — x_mid - x_front of the previous full pass (hit reuse) + // x_mid_dev — reconstructed/forwarded middle (back-graph input) + // x_out_dev — tail output (head-graph input) + ggml_context * dbctx = nullptr; + ggml_backend_buffer_t dbbuf = nullptr; + ggml_tensor * x_front_dev = nullptr; + ggml_tensor * prev_front_dev = nullptr; + ggml_tensor * tail_delta_dev = nullptr; + ggml_tensor * x_mid_dev = nullptr; + ggml_tensor * x_out_dev = nullptr; + int db_T = 0; + // Top-level weights outside the three sub-models. ggml_tensor * w_spec_proj = nullptr; ggml_tensor * b_spec_proj = nullptr; @@ -99,19 +175,19 @@ struct Model::Impl { private: // Pipeline stages. - void run_encoder(const float * mel, int T, - std::vector & x_seg_out, - std::vector & x_est_out); + void run_encoder(const float * mel, int T); void run_segmenter_step( - const float * x_seg_host, int T, + int T, const std::int32_t * noise_mod3, float t_scalar, int language, std::vector & logits_out); void run_estimator( - const float * x_est_host, int T, - const std::int32_t * regions, int N, + int T, const std::int32_t * regions, int N, std::vector & pool_logits_out); + + // (Re)allocate the device-resident DBCache tensors for frame count T. + void ensure_db_tensors(int T); }; } // namespace game_ggml diff --git a/src/model_segmenter.cpp b/src/model_segmenter.cpp index 31ad24f..13a98d0 100644 --- a/src/model_segmenter.cpp +++ b/src/model_segmenter.cpp @@ -29,6 +29,10 @@ static ops::EBFBlockWeights bind_seg_layer( B.w_norm1 = W.get(p + "norm1.weight"); B.w_ffn1_ln1 = W.get(p + "ffn1.ln1.weight"); B.b_ffn1_ln1 = W.get(p + "ffn1.ln1.bias"); + B.w_ffn1_ln1_a = W.try_get(p + "ffn1.ln1.weight.a"); // F-1 split halves + B.b_ffn1_ln1_a = W.try_get(p + "ffn1.ln1.bias.a"); + B.w_ffn1_ln1_b = W.try_get(p + "ffn1.ln1.weight.b"); + B.b_ffn1_ln1_b = W.try_get(p + "ffn1.ln1.bias.b"); B.w_ffn1_ln2 = W.get(p + "ffn1.ln2.weight"); B.b_ffn1_ln2 = W.get(p + "ffn1.ln2.bias"); if (cfg.use_ls) B.w_lay_scale1 = W.get(p + "lay_scale1.scale"); @@ -37,6 +41,10 @@ static ops::EBFBlockWeights bind_seg_layer( B.w_norm2 = W.get(p + "norm2.weight"); B.w_ffn2_ln1 = W.get(p + "ffn2.ln1.weight"); B.b_ffn2_ln1 = W.get(p + "ffn2.ln1.bias"); + B.w_ffn2_ln1_a = W.try_get(p + "ffn2.ln1.weight.a"); // F-1 split halves + B.b_ffn2_ln1_a = W.try_get(p + "ffn2.ln1.bias.a"); + B.w_ffn2_ln1_b = W.try_get(p + "ffn2.ln1.weight.b"); + B.b_ffn2_ln1_b = W.try_get(p + "ffn2.ln1.bias.b"); B.w_ffn2_ln2 = W.get(p + "ffn2.ln2.weight"); B.b_ffn2_ln2 = W.get(p + "ffn2.ln2.bias"); if (cfg.use_ls) B.w_lay_scale3 = W.get(p + "lay_scale3.scale"); diff --git a/src/ops_attn.cpp b/src/ops_attn.cpp index 20c0e84..4cfc991 100644 --- a/src/ops_attn.cpp +++ b/src/ops_attn.cpp @@ -138,13 +138,12 @@ ggml_tensor * pac( // --------------------------------------------------------------------------- // EBF block // --------------------------------------------------------------------------- - -namespace { -// Elementwise scale by half (used by EBF residuals: x + 0.5 * branch). -ggml_tensor * scale_half(ggml_context * ctx, ggml_tensor * x) { - return ggml_scale(ctx, x, 0.5f); -} -} // namespace +// +// F-2: the branch residuals are x + 0.5·lay_scale(branch). Both factors are +// diagonal and are folded into the producing linear (ffn*.ln2 / merge_linear) +// at load time (tensor_utils.cpp), so the graph no longer emits the lay_scale +// mul and the 0.5 scale node per block. w_lay_scale* are still bound (the +// GGUF keeps the tensors) but are intentionally not referenced here. ggml_tensor * ebf_block( ggml_context * ctx, @@ -157,23 +156,30 @@ ggml_tensor * ebf_block( // FFN 1 (pre-attention) if (W.has_ffn1) { ggml_tensor * h = rms_norm(ctx, x, W.w_norm1); - h = glu_ffn(ctx, h, W.w_ffn1_ln1, W.b_ffn1_ln1, W.w_ffn1_ln2, W.b_ffn1_ln2); - if (W.w_lay_scale1) h = layer_scale(ctx, h, W.w_lay_scale1); - h = scale_half(ctx, h); + if (W.w_ffn1_ln1_a) { + h = glu_ffn_split(ctx, h, + W.w_ffn1_ln1_a, W.b_ffn1_ln1_a, W.w_ffn1_ln1_b, W.b_ffn1_ln1_b, + W.w_ffn1_ln2, W.b_ffn1_ln2); + } else { + h = glu_ffn(ctx, h, W.w_ffn1_ln1, W.b_ffn1_ln1, W.w_ffn1_ln2, W.b_ffn1_ln2); + } x = ggml_add(ctx, x, h); } // PAC ggml_tensor * p = pac(ctx, x, W.pac_w, positions, num_heads, head_dim, theta); - if (W.w_lay_scale2) p = layer_scale(ctx, p, W.w_lay_scale2); x = ggml_add(ctx, x, p); // FFN 2 (post-attention) if (W.has_ffn2) { ggml_tensor * h = rms_norm(ctx, x, W.w_norm2); - h = glu_ffn(ctx, h, W.w_ffn2_ln1, W.b_ffn2_ln1, W.w_ffn2_ln2, W.b_ffn2_ln2); - if (W.w_lay_scale3) h = layer_scale(ctx, h, W.w_lay_scale3); - h = scale_half(ctx, h); + if (W.w_ffn2_ln1_a) { + h = glu_ffn_split(ctx, h, + W.w_ffn2_ln1_a, W.b_ffn2_ln1_a, W.w_ffn2_ln1_b, W.b_ffn2_ln1_b, + W.w_ffn2_ln2, W.b_ffn2_ln2); + } else { + h = glu_ffn(ctx, h, W.w_ffn2_ln1, W.b_ffn2_ln1, W.w_ffn2_ln2, W.b_ffn2_ln2); + } x = ggml_add(ctx, x, h); } diff --git a/src/ops_attn.h b/src/ops_attn.h index f8f1739..be08e73 100644 --- a/src/ops_attn.h +++ b/src/ops_attn.h @@ -98,6 +98,11 @@ struct EBFBlockWeights { ggml_tensor * w_norm1 = nullptr; ggml_tensor * w_ffn1_ln1 = nullptr; // ln1.weight (D -> 2L) ggml_tensor * b_ffn1_ln1 = nullptr; + // F-1: split ln1 halves (load-time; null => use monolithic ln1 above). + ggml_tensor * w_ffn1_ln1_a = nullptr; // .ln1.weight.a (D -> L) + ggml_tensor * b_ffn1_ln1_a = nullptr; + ggml_tensor * w_ffn1_ln1_b = nullptr; // .ln1.weight.b + ggml_tensor * b_ffn1_ln1_b = nullptr; ggml_tensor * w_ffn1_ln2 = nullptr; // ln2.weight (L -> D) ggml_tensor * b_ffn1_ln2 = nullptr; ggml_tensor * w_lay_scale1 = nullptr; // optional @@ -111,6 +116,10 @@ struct EBFBlockWeights { ggml_tensor * w_norm2 = nullptr; ggml_tensor * w_ffn2_ln1 = nullptr; ggml_tensor * b_ffn2_ln1 = nullptr; + ggml_tensor * w_ffn2_ln1_a = nullptr; // F-1 split halves + ggml_tensor * b_ffn2_ln1_a = nullptr; + ggml_tensor * w_ffn2_ln1_b = nullptr; + ggml_tensor * b_ffn2_ln1_b = nullptr; ggml_tensor * w_ffn2_ln2 = nullptr; ggml_tensor * b_ffn2_ln2 = nullptr; ggml_tensor * w_lay_scale3 = nullptr; // optional diff --git a/src/ops_ffn.cpp b/src/ops_ffn.cpp index 0f474be..27e2160 100644 --- a/src/ops_ffn.cpp +++ b/src/ops_ffn.cpp @@ -43,6 +43,21 @@ ggml_tensor * glu_ffn(ggml_context * ctx, return linear(ctx, y, w_ln2, b_ln2); } +// F-1: split-halves variant — the ln1 weight/bias were split into two +// contiguous [in, L] / [L] pieces at load time, so the two mul_mats output +// contiguous [L, T, B] activations and the strided-view + 2x cont of the +// monolithic path disappear. +ggml_tensor * glu_ffn_split(ggml_context * ctx, + ggml_tensor * x, + ggml_tensor * w_ln1_a, ggml_tensor * b_ln1_a, + ggml_tensor * w_ln1_b, ggml_tensor * b_ln1_b, + ggml_tensor * w_ln2, ggml_tensor * b_ln2) { + ggml_tensor * x1 = linear(ctx, x, w_ln1_a, b_ln1_a); // [L, T, B] contiguous + ggml_tensor * x2 = linear(ctx, x, w_ln1_b, b_ln1_b); // [L, T, B] contiguous + ggml_tensor * y = ggml_mul(ctx, ggml_gelu(ctx, x1), x2); + return linear(ctx, y, w_ln2, b_ln2); +} + // --------------------------------------------------------------------------- // CgMLP // --------------------------------------------------------------------------- diff --git a/src/ops_ffn.h b/src/ops_ffn.h index ac4f80f..d0bbf79 100644 --- a/src/ops_ffn.h +++ b/src/ops_ffn.h @@ -29,6 +29,16 @@ ggml_tensor * glu_ffn(ggml_context * ctx, ggml_tensor * w_ln1, ggml_tensor * b_ln1, ggml_tensor * w_ln2, ggml_tensor * b_ln2); +// F-1: same GLU FFN but with the ln1 projection pre-split into two [in, L] +// halves (load-time, see tensor_utils.cpp). Two mul_mats produce contiguous +// [L, T, B] outputs directly, skipping the strided-view + 2x cont of the +// monolithic path. w_ln2/b_ln2 are unchanged (L -> D). +ggml_tensor * glu_ffn_split(ggml_context * ctx, + ggml_tensor * x, + ggml_tensor * w_ln1_a, ggml_tensor * b_ln1_a, + ggml_tensor * w_ln1_b, ggml_tensor * b_ln1_b, + ggml_tensor * w_ln2, ggml_tensor * b_ln2); + // -------------------------------------------------------------------------- // CgMLP (modules.backbones.layers.CgMLP) // diff --git a/src/ops_joint_attn.cpp b/src/ops_joint_attn.cpp index 5a93f19..3760eb2 100644 --- a/src/ops_joint_attn.cpp +++ b/src/ops_joint_attn.cpp @@ -48,35 +48,72 @@ std::vector build_joint_attn_mask_fp16( // - pool[i] region id = i + 1 (1..N) // - x[t] region id = regions[t] (0 = padding, 1..N valid) // - allowed iff (same_stream OR same_region(!=0)) AND both valid + // + // Layout: ggml flash_attn mask is (kv_seq, q_seq, 1, 1). The (i=key, j=query) + // element sits at index j*S + i where i is innermost (ne[0]=S). So mask + // "row" j is query j's allowed key set. + // + // Structure exploited for fast fill (regions is a monotone run-length + // encoding — cumsum of boundaries — so same-region frames are contiguous): + // * pool query rows: keys [0, N) (same stream) + the x frames of the + // matching region (contiguous run) are zero; everything else -inf. + // * x query rows: pool key region-1 (single) + all valid x keys (same + // stream, zero-filled span-by-span) are zero; everything else -inf. + // This replaces a per-element branchy double loop + per-element software + // f16 conversion with bulk fills over contiguous runs. const int S = N + T; - const float kNegInf = -10000.0f; // ggml convention for "blocked" + const std::uint16_t neg_inf_h = f32_to_f16_bits(-10000.0f); // ggml convention for "blocked" + const std::uint16_t zero_h = f32_to_f16_bits(0.0f); + + std::vector mask(static_cast(S) * S, neg_inf_h); + + // Per-region [start, end) intervals of x frames (in x-key index space). + std::vector reg_start(static_cast(N) + 2, -1); + std::vector reg_end(static_cast(N) + 2, -1); + for (int i = 0; i < T; ++i) { + const int r = regions[i]; + if (r <= 0 || r > N) continue; // padding / out-of-range + if (reg_start[r] < 0) reg_start[r] = i; + reg_end[r] = i + 1; + } - auto region = [&](int i) -> int { - if (i < N) return i + 1; - return regions[i - N]; // could be 0 for padding - }; - auto is_pool = [&](int i) { return i < N; }; - auto valid = [&](int i) -> bool { - if (i < N) return true; // B=1 inference: all pool tokens valid - return regions[i - N] != 0; - }; + // Contiguous valid (non-padding) x spans — for the x-x same-stream block. + // The same r in [1, N] predicate as the reg_start/reg_end pass above: + // negative or > N ids are treated as invalid everywhere (a malformed id + // must not be able to attend to valid x keys). + struct Span { int b, e; }; + const auto is_valid_region = [N](int r) { return r >= 1 && r <= N; }; + std::vector x_spans; + for (int i = 0; i < T; ) { + if (!is_valid_region(regions[i])) { ++i; continue; } + int e = i; + while (e < T && is_valid_region(regions[e])) ++e; + x_spans.push_back({i, e}); + i = e; + } - std::vector mask(static_cast(S) * S, - f32_to_f16_bits(kNegInf)); - const std::uint16_t zero_h = f32_to_f16_bits(0.0f); + std::uint16_t * M = mask.data(); - // Layout: ggml flash_attn mask is (kv_seq, q_seq, 1, 1). The (i=key, j=query) - // element sits at index j*S + i where i is innermost (ne[0]=S). - for (int j = 0; j < S; ++j) { - for (int i = 0; i < S; ++i) { - bool allowed = valid(i) && valid(j); - if (allowed) { - const bool same_stream = (is_pool(i) == is_pool(j)); - const bool ri = region(i), rj = region(j); - const bool same_region = (ri != 0 && rj != 0 && region(i) == region(j)); - allowed = same_stream || same_region; - } - mask[static_cast(j) * S + i] = allowed ? zero_h : f32_to_f16_bits(kNegInf); + // Pool query rows (j = 0..N-1): same-stream pool keys + matching region's x. + for (int j = 0; j < N; ++j) { + std::uint16_t * row = M + static_cast(j) * S; + std::fill(row, row + N, zero_h); // pool-pool block (all valid) + const int r = j + 1; + if (reg_start[r] >= 0) { + std::fill(row + N + reg_start[r], row + N + reg_end[r], zero_h); + } + } + + // X query rows (j = 0..T-1, key row N+j): single matching pool key + + // same-stream valid x keys. Padding/out-of-range query ids (incl. + // negative or > N) admit nothing. + for (int j = 0; j < T; ++j) { + const int rj = regions[j]; + if (!is_valid_region(rj)) continue; // padding query: nothing allowed + std::uint16_t * row = M + static_cast(N + j) * S; + row[rj - 1] = zero_h; + for (const Span & sp : x_spans) { + std::fill(row + N + sp.b, row + N + sp.e, zero_h); } } return mask; @@ -274,12 +311,20 @@ namespace { ggml_tensor * apply_ffn_block(ggml_context * ctx, ggml_tensor * x, ggml_tensor * w_norm, ggml_tensor * w_ln1, ggml_tensor * b_ln1, + ggml_tensor * w_ln1_a, ggml_tensor * b_ln1_a, ggml_tensor * w_ln1_b, ggml_tensor * b_ln1_b, ggml_tensor * w_ln2, ggml_tensor * b_ln2, ggml_tensor * w_lay_scale) { + // F-2: lay_scale is folded into the ln2 linear at load time + // (tensor_utils.cpp), so w_lay_scale is intentionally unused here. + (void)w_lay_scale; ggml_tensor * h = rms_norm(ctx, x, w_norm); - h = glu_ffn(ctx, h, w_ln1, b_ln1, w_ln2, b_ln2); - if (w_lay_scale) h = layer_scale(ctx, h, w_lay_scale); + if (w_ln1_a) { + // F-1: pre-split ln1 halves — two mul_mats, no cont copies. + h = glu_ffn_split(ctx, h, w_ln1_a, b_ln1_a, w_ln1_b, b_ln1_b, w_ln2, b_ln2); + } else { + h = glu_ffn(ctx, h, w_ln1, b_ln1, w_ln2, b_ln2); + } // JEBF uses `+ x` (not `* 0.5 + x`) unlike the single-stream EBF. return ggml_add(ctx, x, h); } @@ -300,29 +345,37 @@ JoinResult jebf_block( // --- FFN1 per stream --- if (W.has_ffn1) { x = apply_ffn_block(ctx, x, W.w_norm_ffn1_x, - W.w_ffn1_x_ln1, W.b_ffn1_x_ln1, W.w_ffn1_x_ln2, W.b_ffn1_x_ln2, + W.w_ffn1_x_ln1, W.b_ffn1_x_ln1, + W.w_ffn1_x_ln1_a, W.b_ffn1_x_ln1_a, W.w_ffn1_x_ln1_b, W.b_ffn1_x_ln1_b, + W.w_ffn1_x_ln2, W.b_ffn1_x_ln2, W.w_lay_scale_ffn1_x); pool = apply_ffn_block(ctx, pool, W.w_norm_ffn1_pool, - W.w_ffn1_pool_ln1, W.b_ffn1_pool_ln1, W.w_ffn1_pool_ln2, W.b_ffn1_pool_ln2, + W.w_ffn1_pool_ln1, W.b_ffn1_pool_ln1, + W.w_ffn1_pool_ln1_a, W.b_ffn1_pool_ln1_a, W.w_ffn1_pool_ln1_b, W.b_ffn1_pool_ln1_b, + W.w_ffn1_pool_ln2, W.b_ffn1_pool_ln2, W.w_lay_scale_ffn1_pool); } // --- PJAC --- + // F-2: lay_scale_jpac_{x,pool} folded into merge_linear_{x,pool} at load + // time, so the residual uses att.* directly. auto att = pjac(ctx, pool, x, W.pjac, global_positions, region_indices, attn_mask_fp16, num_heads, head_dim, theta); - ggml_tensor * x_att = W.w_lay_scale_jpac_x ? layer_scale(ctx, att.x, W.w_lay_scale_jpac_x) : att.x; - ggml_tensor * pool_att = W.w_lay_scale_jpac_pool ? layer_scale(ctx, att.pool, W.w_lay_scale_jpac_pool) : att.pool; - x = ggml_add(ctx, x, x_att); - pool = ggml_add(ctx, pool, pool_att); + x = ggml_add(ctx, x, att.x); + pool = ggml_add(ctx, pool, att.pool); // --- FFN2 per stream --- if (W.has_ffn2) { x = apply_ffn_block(ctx, x, W.w_norm_ffn2_x, - W.w_ffn2_x_ln1, W.b_ffn2_x_ln1, W.w_ffn2_x_ln2, W.b_ffn2_x_ln2, + W.w_ffn2_x_ln1, W.b_ffn2_x_ln1, + W.w_ffn2_x_ln1_a, W.b_ffn2_x_ln1_a, W.w_ffn2_x_ln1_b, W.b_ffn2_x_ln1_b, + W.w_ffn2_x_ln2, W.b_ffn2_x_ln2, W.w_lay_scale_ffn2_x); pool = apply_ffn_block(ctx, pool, W.w_norm_ffn2_pool, - W.w_ffn2_pool_ln1, W.b_ffn2_pool_ln1, W.w_ffn2_pool_ln2, W.b_ffn2_pool_ln2, + W.w_ffn2_pool_ln1, W.b_ffn2_pool_ln1, + W.w_ffn2_pool_ln1_a, W.b_ffn2_pool_ln1_a, W.w_ffn2_pool_ln1_b, W.b_ffn2_pool_ln1_b, + W.w_ffn2_pool_ln2, W.b_ffn2_pool_ln2, W.w_lay_scale_ffn2_pool); } diff --git a/src/ops_joint_attn.h b/src/ops_joint_attn.h index 08cf9a6..1ddad61 100644 --- a/src/ops_joint_attn.h +++ b/src/ops_joint_attn.h @@ -98,10 +98,18 @@ struct JEBFBlockWeights { ggml_tensor * w_norm_ffn1_x = nullptr; ggml_tensor * w_norm_ffn1_pool = nullptr; ggml_tensor * w_ffn1_x_ln1 = nullptr; + ggml_tensor * w_ffn1_x_ln1_a = nullptr; // F-1 split halves (.a/.b) + ggml_tensor * b_ffn1_x_ln1_a = nullptr; + ggml_tensor * w_ffn1_x_ln1_b = nullptr; + ggml_tensor * b_ffn1_x_ln1_b = nullptr; ggml_tensor * b_ffn1_x_ln1 = nullptr; ggml_tensor * w_ffn1_x_ln2 = nullptr; ggml_tensor * b_ffn1_x_ln2 = nullptr; ggml_tensor * w_ffn1_pool_ln1 = nullptr; + ggml_tensor * w_ffn1_pool_ln1_a = nullptr; // F-1 split halves + ggml_tensor * b_ffn1_pool_ln1_a = nullptr; + ggml_tensor * w_ffn1_pool_ln1_b = nullptr; + ggml_tensor * b_ffn1_pool_ln1_b = nullptr; ggml_tensor * b_ffn1_pool_ln1 = nullptr; ggml_tensor * w_ffn1_pool_ln2 = nullptr; ggml_tensor * b_ffn1_pool_ln2 = nullptr; @@ -117,10 +125,18 @@ struct JEBFBlockWeights { ggml_tensor * w_norm_ffn2_x = nullptr; ggml_tensor * w_norm_ffn2_pool = nullptr; ggml_tensor * w_ffn2_x_ln1 = nullptr; + ggml_tensor * w_ffn2_x_ln1_a = nullptr; // F-1 split halves + ggml_tensor * b_ffn2_x_ln1_a = nullptr; + ggml_tensor * w_ffn2_x_ln1_b = nullptr; + ggml_tensor * b_ffn2_x_ln1_b = nullptr; ggml_tensor * b_ffn2_x_ln1 = nullptr; ggml_tensor * w_ffn2_x_ln2 = nullptr; ggml_tensor * b_ffn2_x_ln2 = nullptr; ggml_tensor * w_ffn2_pool_ln1 = nullptr; + ggml_tensor * w_ffn2_pool_ln1_a = nullptr; // F-1 split halves + ggml_tensor * b_ffn2_pool_ln1_a = nullptr; + ggml_tensor * w_ffn2_pool_ln1_b = nullptr; + ggml_tensor * b_ffn2_pool_ln1_b = nullptr; ggml_tensor * b_ffn2_pool_ln1 = nullptr; ggml_tensor * w_ffn2_pool_ln2 = nullptr; ggml_tensor * b_ffn2_pool_ln2 = nullptr; diff --git a/src/tensor_utils.cpp b/src/tensor_utils.cpp index 106a671..7d00e0c 100644 --- a/src/tensor_utils.cpp +++ b/src/tensor_utils.cpp @@ -5,8 +5,11 @@ #include #include +#include #include #include +#include +#include #include #include #include @@ -14,37 +17,287 @@ namespace game_ggml::internal { namespace { -constexpr size_t kMetaCtxBytes = 2ull * 1024 * 1024; // 2 MB suffices for 700-tensor metadata + +// Depthwise-conv weights are stored F16 in the GGUF (converter constraint) +// but the dedicated conv_2d_dw kernels want an F32 kernel. Casting inside +// the graph would re-run a conversion node on every stage compute (once per +// D3PM step, per layer). Instead we materialise a persistent F32 copy at +// load time and point the weight table at it; the graph then sees F32 and +// builds zero cast nodes. +bool is_dwconv_weight(const char * name) { + return (std::strstr(name, ".dw.") != nullptr || std::strstr(name, "dw_conv") != nullptr) && + std::strstr(name, ".weight") != nullptr; +} + +// F-1: GLU FFN first linear ("ln1"). The [in, 2L] weight and [2L] bias are +// split at load time into two [in, L] / [L] halves so the graph can run two +// mul_mats that produce contiguous outputs, eliminating the two cont copies +// per GLU FFN (see glu_ffn_split in ops_ffn.cpp). +bool is_glu_ln1_weight(const char * name) { + return std::strstr(name, ".ln1.weight") != nullptr && !is_dwconv_weight(name); +} +bool is_glu_ln1_bias(const char * name) { + return std::strstr(name, ".ln1.bias") != nullptr; +} + +// --------------------------------------------------------------------------- +// F-2: EBF layer-scale folding (load-time, schema-preserving). +// +// The EBF residual is x + 0.5 * lay_scale(branch) where lay_scale is a +// per-channel multiply. Since 0.5 and lay_scale are both diagonal, they can +// be folded into the *producing* linear of the branch at load time: +// +// out' = 0.5 * s ⊙ (W·h + b) == (0.5·s·W)·h + (0.5·s⊙b) +// +// so the graph no longer emits a lay_scale mul + a 0.5 scale node per EBF +// block (two elementwise kernels per FFN, one per PAC branch). The GGUF +// keeps its lay_scale tensors (bind code still finds them) — they are simply +// no longer referenced by the graph. This is idempotent: the file is never +// modified, every load folds the same way. +// +// The fold is lossless for Q8_0 (only the per-block d scalars change) and +// exact for F32/F16 up to float rounding; the graph arithmetic order changes +// (scale applied before the matmul instead of after), so outputs are +// expected to match the unfolded graph to ~1e-7, not bit-exactly. +// --------------------------------------------------------------------------- + +// --------------------------------------------------------------------------- +// F-2: layer-scale folding (load-time, schema-preserving) — covers both the +// single-stream EBF blocks (encoder/segmenter) and the joint EBF (estimator): +// +// * single-stream EBF residual x + 0.5·lay_scale(branch) → mult 0.5 +// * joint EBF residual x + lay_scale(branch) → mult 1.0 +// +// lay_scale is a per-channel multiply, so it folds into the branch's +// producing linear at load time (see F-2 note above). `parse_lay_scale_name` +// maps a GGUF lay_scale tensor name to its producing linear's weight/bias +// tensor names (relative to the block base, with leading '.') and the +// extra multiplier. Returns false for non-lay_scale names. +// --------------------------------------------------------------------------- +struct FoldTarget { + std::string w_suffix; // e.g. ".ffn1.ln2.weight" (base + suffix = full name) + std::string b_suffix; + float mult; // extra factor applied to the lay_scale values +}; + +bool parse_lay_scale_name(const std::string & name, std::string & base_out, + FoldTarget & t_out) { + // ---- Single-stream EBF: {base}.lay_scale{1,2,3}.scale ---- + const std::string ebf = ".lay_scale"; + const std::size_t ep = name.rfind(ebf); + if (ep != std::string::npos) { + const std::string tail = name.substr(ep + ebf.size()); // "N.scale" + if (tail.size() == 7 && tail[1] == '.' && + tail.compare(1, std::string::npos, ".scale") == 0) { + const int which = tail[0] - '0'; + if (which >= 1 && which <= 3) { + base_out = name.substr(0, ep); + if (which == 1) { + t_out.w_suffix = ".ffn1.ln2.weight"; t_out.b_suffix = ".ffn1.ln2.bias"; + } else if (which == 2) { + t_out.w_suffix = ".attn.merge_linear.weight"; t_out.b_suffix = ".attn.merge_linear.bias"; + } else { + t_out.w_suffix = ".ffn2.ln2.weight"; t_out.b_suffix = ".ffn2.ln2.bias"; + } + t_out.mult = (which == 2) ? 1.0f : 0.5f; + return true; + } + } + } + + // ---- Joint EBF (estimator): {base}.lay_scale_{kind}.scale ---- + // kind ∈ {ffn1_x, ffn1_pool, ffn2_x, ffn2_pool, jpac_x, jpac_pool} + const std::string jp = ".lay_scale_"; + const std::size_t jpos = name.rfind(jp); + if (jpos != std::string::npos) { + const std::string kind = name.substr(jpos + jp.size()); // e.g. "ffn1_x.scale" + if (kind.size() >= 7 && kind.compare(kind.size() - 6, 6, ".scale") == 0) { + const std::string k = kind.substr(0, kind.size() - 6); // "ffn1_x" + base_out = name.substr(0, jpos); + if (k.size() >= 6 && k.compare(0, 3, "ffn") == 0 && + (k[3] == '1' || k[3] == '2') && k[4] == '_') { + // k = "ffn{1,2}_{x,pool}" + const std::string stream = k.substr(5); + if (stream == "x" || stream == "pool") { + t_out.w_suffix = ".ffn" + k.substr(3, 1) + "_" + stream + ".ln2.weight"; + t_out.b_suffix = ".ffn" + k.substr(3, 1) + "_" + stream + ".ln2.bias"; + t_out.mult = 1.0f; + return true; + } + } + if (k.size() >= 6 && k.compare(0, 4, "jpac") == 0 && k[4] == '_') { + const std::string stream = k.substr(5); + if (stream == "x" || stream == "pool") { + t_out.w_suffix = ".attn.merge_linear_" + stream + ".weight"; + t_out.b_suffix = ".attn.merge_linear_" + stream + ".bias"; + t_out.mult = 1.0f; + return true; + } + } + } + } + return false; +} + +// Multiply the rows of a weight tensor (ne[0] = in dim, contiguous) by a +// per-output-channel factor vector. In-place on the raw payload bytes. +void fold_linear_weight(ggml_tensor * t, std::vector & data, + const std::vector & factor) { + const int64_t D_in = t->ne[0]; + const int64_t D_out = t->ne[1]; + if (D_out != static_cast(factor.size())) { + throw GgufError(std::string("fold: row count mismatch on '") + t->name + + "' (" + std::to_string(D_out) + " vs " + + std::to_string(factor.size()) + ")"); + } + switch (t->type) { + case GGML_TYPE_F32: { + float * p = reinterpret_cast(data.data()); + for (int64_t o = 0; o < D_out; ++o) { + const float f = factor[static_cast(o)]; + float * row = p + o * D_in; + for (int64_t i = 0; i < D_in; ++i) row[i] *= f; + } + } break; + case GGML_TYPE_F16: { + ggml_fp16_t * p = reinterpret_cast(data.data()); + for (int64_t o = 0; o < D_out; ++o) { + const float f = factor[static_cast(o)]; + ggml_fp16_t * row = p + o * D_in; + for (int64_t i = 0; i < D_in; ++i) + row[i] = ggml_fp32_to_fp16(ggml_fp16_to_fp32(row[i]) * f); + } + } break; + case GGML_TYPE_Q8_0: { + // block_q8_0: { fp16 d; int8 qs[QK8_0] } — d is the row scale of + // the block, so a per-row multiplier becomes a per-block d + // change. Lossless. + constexpr int64_t QK = 32; + if (D_in % QK != 0) + throw GgufError("fold: Q8_0 weight requires ne[0] % 32 == 0"); + const int64_t bpr = D_in / QK; + struct BQ8 { ggml_fp16_t d; int8_t qs[QK]; }; + BQ8 * p = reinterpret_cast(data.data()); + for (int64_t o = 0; o < D_out; ++o) { + const float f = factor[static_cast(o)]; + BQ8 * row = p + o * bpr; + for (int64_t b = 0; b < bpr; ++b) + row[b].d = ggml_fp32_to_fp16(ggml_fp16_to_fp32(row[b].d) * f); + } + } break; + default: + throw GgufError(std::string("fold: unsupported weight type '") + + ggml_type_name(t->type) + "' on '" + t->name + "'"); + } +} + +void fold_linear_bias(ggml_tensor * t, std::vector & data, + const std::vector & factor) { + if (t->ne[0] != static_cast(factor.size())) { + throw GgufError(std::string("fold: bias dim mismatch on '") + t->name + "'"); + } + switch (t->type) { + case GGML_TYPE_F32: { + float * p = reinterpret_cast(data.data()); + for (std::size_t i = 0; i < factor.size(); ++i) p[i] *= factor[i]; + } break; + case GGML_TYPE_F16: { + ggml_fp16_t * p = reinterpret_cast(data.data()); + for (std::size_t i = 0; i < factor.size(); ++i) + p[i] = ggml_fp32_to_fp16(ggml_fp16_to_fp32(p[i]) * factor[i]); + } break; + default: + throw GgufError(std::string("fold: unsupported bias type '") + + ggml_type_name(t->type) + "' on '" + t->name + "'"); + } } +} // namespace + LoadedWeights LoadedWeights::load_all(const GgufFile & gguf, ggml_backend_t backend) { // --- 1. Build a metadata-only ggml_context that mirrors the file's // tensor table. gguf_init_from_file with params.ctx != null will // create ggml_tensor entries in that context but leave them - // un-backed (no_alloc=true). + // un-backed (no_alloc=true). It sizes the context exactly for the + // tensor table, so no extra tensors may be added to it afterwards. ggml_context * ctx = nullptr; - { - ggml_init_params ip{}; - ip.mem_size = kMetaCtxBytes; - ip.mem_buffer = nullptr; - ip.no_alloc = true; - ctx = ggml_init(ip); - if (!ctx) throw GgufError("failed to create weight ggml_context"); - } - gguf_init_params gp{}; gp.no_alloc = true; gp.ctx = &ctx; gguf_context * gctx = gguf_init_from_file(gguf.path().c_str(), gp); if (!gctx) { - ggml_free(ctx); throw GgufError("failed to re-open GGUF for tensor loading: " + gguf.path()); } + // --- 1b. Create persistent F32 copies of F16 depthwise-conv weights + // (see is_dwconv_weight above), plus F-1 GLU ln1 split halves + // (a/b). The gguf-init context is sized exactly for its tensor + // table, so these copies get their own context + backend buffer. + std::map dw_f32; // original name -> F32 copy + std::map> glu_ab; // ln1 weight -> {a, b} + std::map> glu_ab_bias; // ln1 bias -> {a, b} + ggml_context * ctx2 = nullptr; + ggml_backend_buffer_t buf2 = nullptr; + { + ggml_init_params ip{}; + // Metadata budget for the copies: at most 2 new tensors per GGUF + // tensor (GLU ln1 weight a/b halves, plus one dwconv F32 copy), so + // size from the tensor count rather than a fixed 512 KiB (a model + // with many estimator/segmenter layers could exhaust a fixed pool, + // and ggml_new_tensor_2d would then abort instead of throwing). + ip.mem_size = ggml_tensor_overhead() * + (2 * static_cast(gguf_get_n_tensors(gctx)) + 1); + ip.mem_buffer = nullptr; + ip.no_alloc = true; + ctx2 = ggml_init(ip); + if (!ctx2) { + gguf_free(gctx); + ggml_free(ctx); + throw GgufError("failed to create weight-copy context"); + } + + const int64_t n_tensors = gguf_get_n_tensors(gctx); + for (int64_t i = 0; i < n_tensors; ++i) { + const char * name = gguf_get_tensor_name(gctx, i); + ggml_tensor * t = ggml_get_tensor(ctx, name); + if (!t) continue; + if (t->type == GGML_TYPE_F16 && is_dwconv_weight(name)) { + ggml_tensor * w32 = ggml_new_tensor_3d(ctx2, GGML_TYPE_F32, + t->ne[0], t->ne[1], t->ne[2]); + dw_f32.emplace(name, w32); + } else if (is_glu_ln1_weight(name) && + (t->type == GGML_TYPE_F32 || t->type == GGML_TYPE_F16)) { + // w_ln1 ne = [in, 2L] (column-major, out dim contiguous) -> + // a = first L out-columns, b = last L out-columns. + const int64_t L = t->ne[1] / 2; + ggml_tensor * a = ggml_new_tensor_2d(ctx2, GGML_TYPE_F32, t->ne[0], L); + ggml_tensor * b = ggml_new_tensor_2d(ctx2, GGML_TYPE_F32, t->ne[0], L); + glu_ab.emplace(name, std::make_pair(a, b)); + } else if (is_glu_ln1_bias(name) && + (t->type == GGML_TYPE_F32 || t->type == GGML_TYPE_F16)) { + const int64_t L = t->ne[0] / 2; + ggml_tensor * a = ggml_new_tensor_1d(ctx2, GGML_TYPE_F32, L); + ggml_tensor * b = ggml_new_tensor_1d(ctx2, GGML_TYPE_F32, L); + glu_ab_bias.emplace(name, std::make_pair(a, b)); + } + } + if (!dw_f32.empty() || !glu_ab.empty() || !glu_ab_bias.empty()) { + buf2 = ggml_backend_alloc_ctx_tensors(ctx2, backend); + if (!buf2) { + ggml_free(ctx2); + gguf_free(gctx); + ggml_free(ctx); + throw GgufError("failed to allocate weight-copy buffer"); + } + } + } + // --- 2. Allocate backend buffer covering all tensors in ctx. ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); if (!buf) { + if (buf2) ggml_backend_buffer_free(buf2); + ggml_free(ctx2); gguf_free(gctx); ggml_free(ctx); throw GgufError("ggml_backend_alloc_ctx_tensors failed; out of memory?"); @@ -54,6 +307,8 @@ LoadedWeights LoadedWeights::load_all(const GgufFile & gguf, ggml_backend_t back FILE * f = std::fopen(gguf.path().c_str(), "rb"); if (!f) { ggml_backend_buffer_free(buf); + if (buf2) ggml_backend_buffer_free(buf2); + ggml_free(ctx2); gguf_free(gctx); ggml_free(ctx); throw GgufError("failed to open GGUF payload file: " + gguf.path()); @@ -62,9 +317,92 @@ LoadedWeights LoadedWeights::load_all(const GgufFile & gguf, ggml_backend_t back const size_t data_offset = gguf_get_data_offset(gctx); std::vector scratch; + // --- 3a. F-2: EBF layer-scale fold plan. Scan the GGUF tensor table + // for encoder/segmenter EBF lay_scale tensors, read their + // payloads once, and record which producing linear weights and + // biases to fold. The upload loop below applies the fold and + // reuses the cached lay_scale payload (so the file is read + // once per tensor). fold_weights/fold_biases are keyed by + // tensor name; the value is the per-channel factor vector. + std::map> fold_weights; + std::map> fold_biases; + std::map> ls_cache; + { + std::vector ls_scratch; + const int64_t n_tensors = gguf_get_n_tensors(gctx); + for (int64_t i = 0; i < n_tensors; ++i) { + const char * name = gguf_get_tensor_name(gctx, i); + std::string base; + FoldTarget tgt; + if (!parse_lay_scale_name(name, base, tgt)) continue; + ggml_tensor * t = ggml_get_tensor(ctx, name); + if (!t) continue; + + const size_t bytes = ggml_nbytes(t); + const size_t offset = data_offset + gguf_get_tensor_offset(gctx, i); + ls_scratch.resize(bytes); + if (std::fseek(f, static_cast(offset), SEEK_SET) != 0 || + std::fread(ls_scratch.data(), 1, bytes, f) != bytes) { + // Same cleanup as every other failure path: these handles are + // backend buffers / contexts, so leaking them on a retried + // load would exhaust device memory. + std::fclose(f); + ggml_backend_buffer_free(buf); + if (buf2) ggml_backend_buffer_free(buf2); + if (ctx2) ggml_free(ctx2); + gguf_free(gctx); + ggml_free(ctx); + throw GgufError(std::string("short read for lay_scale '") + name + "'"); + } + + // Per-channel scale as f32. + const int64_t D = t->ne[0]; + std::vector s(static_cast(D)); + if (t->type == GGML_TYPE_F32) { + std::memcpy(s.data(), ls_scratch.data(), bytes); + } else if (t->type == GGML_TYPE_F16) { + ggml_fp16_to_fp32_row(reinterpret_cast(ls_scratch.data()), + s.data(), D); + } else { + std::fclose(f); + ggml_backend_buffer_free(buf); + if (buf2) ggml_backend_buffer_free(buf2); + if (ctx2) ggml_free(ctx2); + gguf_free(gctx); + ggml_free(ctx); + throw GgufError(std::string("fold: unsupported lay_scale type '") + + ggml_type_name(t->type) + "' on '" + name + "'"); + } + + // factor = mult · s (0.5 for single-stream FFN branches, 1.0 for + // PAC/PJAC and joint-FFN branches). + std::vector factor = s; + if (tgt.mult != 1.0f) { + for (auto & v : factor) v *= tgt.mult; + } + + const std::string w_name = base + tgt.w_suffix; + if (!ggml_get_tensor(ctx, w_name.c_str())) { + throw GgufError("fold: missing target '" + w_name + "' for '" + name + "'"); + } + fold_weights.emplace(w_name, factor); + const std::string b_name = base + tgt.b_suffix; + if (ggml_get_tensor(ctx, b_name.c_str())) { + fold_biases.emplace(b_name, factor); + } + ls_cache.emplace(name, ls_scratch); + } + if (!fold_weights.empty()) { + std::fprintf(stderr, "[FOLD] folded %zu lay_scale(s) into producing linears\n", + fold_weights.size()); + } + } + LoadedWeights out; - out.ctx_ = ctx; - out.buffer_ = buf; + out.ctx_ = ctx; + out.buffer_ = buf; + out.ctx2_ = ctx2; + out.buffer2_ = buf2; const int64_t n_tensors = gguf_get_n_tensors(gctx); for (int64_t i = 0; i < n_tensors; ++i) { @@ -73,6 +411,8 @@ LoadedWeights LoadedWeights::load_all(const GgufFile & gguf, ggml_backend_t back if (!t) { std::fclose(f); ggml_backend_buffer_free(buf); + if (buf2) ggml_backend_buffer_free(buf2); + if (ctx2) ggml_free(ctx2); gguf_free(gctx); ggml_free(ctx); throw GgufError(std::string("tensor '") + name + "' missing from ggml context"); @@ -84,12 +424,87 @@ LoadedWeights LoadedWeights::load_all(const GgufFile & gguf, ggml_backend_t back std::fread(scratch.data(), 1, bytes, f) != bytes) { std::fclose(f); ggml_backend_buffer_free(buf); + if (buf2) ggml_backend_buffer_free(buf2); + if (ctx2) ggml_free(ctx2); gguf_free(gctx); ggml_free(ctx); throw GgufError(std::string("short read for tensor '") + name + "'"); } - ggml_backend_tensor_set(t, scratch.data(), 0, bytes); - out.tensors_.emplace(name, t); + + // F-2: fold EBF layer-scale into producing linear weights/biases. + { + const auto fw = fold_weights.find(name); + if (fw != fold_weights.end()) fold_linear_weight(t, scratch, fw->second); + const auto fb = fold_biases.find(name); + if (fb != fold_biases.end()) fold_linear_bias(t, scratch, fb->second); + const auto lsc = ls_cache.find(name); + if (lsc != ls_cache.end()) scratch = lsc->second; // reuse cached payload + } + + // F-1: split GLU ln1 weights/biases into a/b F32 halves. The + // original tensor is still uploaded (bind code references it), but + // the graph uses the contiguous .a/.b halves (two mul_mats, no cont). + { + const auto gw = glu_ab.find(name); + if (gw != glu_ab.end()) { + const int64_t D_in = t->ne[0]; + const int64_t L = t->ne[1] / 2; + std::vector half(static_cast(D_in) * L); + auto fill_half = [&](ggml_tensor * dst, int64_t col0) { + if (t->type == GGML_TYPE_F32) { + const float * src = reinterpret_cast(scratch.data()); + for (int64_t o = 0; o < L; ++o) + std::memcpy(half.data() + o * D_in, src + (col0 + o) * D_in, + static_cast(D_in) * sizeof(float)); + } else { // F16 + const ggml_fp16_t * src = reinterpret_cast(scratch.data()); + for (int64_t o = 0; o < L; ++o) + ggml_fp16_to_fp32_row(src + (col0 + o) * D_in, + half.data() + o * D_in, D_in); + } + ggml_backend_tensor_set(dst, half.data(), 0, half.size() * sizeof(float)); + }; + fill_half(gw->second.first, 0); + fill_half(gw->second.second, L); + out.tensors_.emplace(std::string(name) + ".a", gw->second.first); + out.tensors_.emplace(std::string(name) + ".b", gw->second.second); + } + const auto gb = glu_ab_bias.find(name); + if (gb != glu_ab_bias.end()) { + const int64_t L = t->ne[0] / 2; + std::vector half(static_cast(L)); + auto fill_half_b = [&](ggml_tensor * dst, int64_t i0) { + if (t->type == GGML_TYPE_F32) { + const float * src = reinterpret_cast(scratch.data()); + for (int64_t i = 0; i < L; ++i) half[static_cast(i)] = src[i0 + i]; + } else { // F16 + const ggml_fp16_t * src = reinterpret_cast(scratch.data()); + for (int64_t i = 0; i < L; ++i) + half[static_cast(i)] = ggml_fp16_to_fp32(src[i0 + i]); + } + ggml_backend_tensor_set(dst, half.data(), 0, half.size() * sizeof(float)); + }; + fill_half_b(gb->second.first, 0); + fill_half_b(gb->second.second, L); + out.tensors_.emplace(std::string(name) + ".a", gb->second.first); + out.tensors_.emplace(std::string(name) + ".b", gb->second.second); + } + } + + // F16 depthwise-conv weights: store the persistent F32 copy instead. + const auto dup = dw_f32.find(name); + if (dup != dw_f32.end()) { + ggml_tensor * w32 = dup->second; + const int64_t n = ggml_nelements(t); + std::vector f32(static_cast(n)); + ggml_fp16_to_fp32_row(reinterpret_cast(scratch.data()), + f32.data(), n); + ggml_backend_tensor_set(w32, f32.data(), 0, f32.size() * sizeof(float)); + out.tensors_.emplace(name, w32); + } else { + ggml_backend_tensor_set(t, scratch.data(), 0, bytes); + out.tensors_.emplace(name, t); + } } std::fclose(f); @@ -99,25 +514,37 @@ LoadedWeights LoadedWeights::load_all(const GgufFile & gguf, ggml_backend_t back } LoadedWeights::~LoadedWeights() { - if (buffer_) ggml_backend_buffer_free(buffer_); - if (ctx_) ggml_free(ctx_); + if (buffer_) ggml_backend_buffer_free(buffer_); + if (ctx_) ggml_free(ctx_); + if (buffer2_) ggml_backend_buffer_free(buffer2_); + if (ctx2_) ggml_free(ctx2_); } LoadedWeights::LoadedWeights(LoadedWeights && other) noexcept - : ctx_(other.ctx_), buffer_(other.buffer_), tensors_(std::move(other.tensors_)) { + : ctx_(other.ctx_), buffer_(other.buffer_), + ctx2_(other.ctx2_), buffer2_(other.buffer2_), + tensors_(std::move(other.tensors_)) { other.ctx_ = nullptr; other.buffer_ = nullptr; + other.ctx2_ = nullptr; + other.buffer2_ = nullptr; } LoadedWeights & LoadedWeights::operator=(LoadedWeights && other) noexcept { if (this != &other) { - if (buffer_) ggml_backend_buffer_free(buffer_); - if (ctx_) ggml_free(ctx_); + if (buffer_) ggml_backend_buffer_free(buffer_); + if (ctx_) ggml_free(ctx_); + if (buffer2_) ggml_backend_buffer_free(buffer2_); + if (ctx2_) ggml_free(ctx2_); ctx_ = other.ctx_; buffer_ = other.buffer_; + ctx2_ = other.ctx2_; + buffer2_ = other.buffer2_; tensors_ = std::move(other.tensors_); other.ctx_ = nullptr; other.buffer_ = nullptr; + other.ctx2_ = nullptr; + other.buffer2_ = nullptr; } return *this; } diff --git a/src/tensor_utils.h b/src/tensor_utils.h index bbd88c5..a848b9a 100644 --- a/src/tensor_utils.h +++ b/src/tensor_utils.h @@ -53,6 +53,11 @@ class LoadedWeights { ggml_context * ctx_ = nullptr; ggml_backend_buffer_t buffer_ = nullptr; + // Second context/buffer pair holding the load-time F32 copies of the F16 + // depthwise-conv weights (the gguf-init context is sized exactly for its + // tensor table, so the copies cannot live there). + ggml_context * ctx2_ = nullptr; + ggml_backend_buffer_t buffer2_ = nullptr; std::unordered_map tensors_; };