From 17d1a8af35af321c0207e5ed7b3e24e478dadec2 Mon Sep 17 00:00:00 2001 From: Shaw Date: Sun, 21 Jun 2026 23:42:49 -0400 Subject: [PATCH 01/17] =?UTF-8?q?fix(build):=20unbreak=20the=20iOS=20fused?= =?UTF-8?q?=20slice=20=E2=80=94=20static=20FFI=20target=20+=20host-only=20?= =?UTF-8?q?kokoro=20CLI?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The iOS XCFramework slice (build-llama-cpp-mtp.mjs, BUILD_SHARED_LIBS=OFF) builds `elizainference_static` by name and links libelizainference.a into the framework, but the target was never defined (only the SHARED `elizainference`), so ninja failed 'unknown target elizainference_static'. And tools/kokoro built its standalone CLI unconditionally with install(TARGETS kokoro-tts RUNTIME), which on iOS (executables are app bundles) fails configure with 'no BUNDLE DESTINATION'. - tools/omnivoice/CMakeLists.txt: add an elizainference_static STATIC target mirroring the SHARED one (same FFI+core sources, includes, OMNIVOICE_BUILD / ELIZA_ENABLE_VISION defines, and the if(TARGET kokoro_lib) Kokoro fold → ELIZA_ENABLE_KOKORO), OUTPUT_NAME elizainference → libelizainference.a; no -reexport_library (SHARED-only). - tools/kokoro/CMakeLists.txt: guard the host-only kokoro-tts CLI + its install off iOS. Verified: ios-arm64-metal-fused AND ios-arm64-simulator-metal-fused both build clean (320/320), libelizainference.a exports the four eliza_inference_kokoro_* symbols + tts/asr/eot, CAPABILITIES.json abiVersion=11 kokoro=kokoro-82m; the desktop SHARED build is unaffected (additive target). Co-Authored-By: Claude Opus 4.8 --- tools/kokoro/CMakeLists.txt | 14 ++++++++---- tools/omnivoice/CMakeLists.txt | 40 ++++++++++++++++++++++++++++++++++ 2 files changed, 50 insertions(+), 4 deletions(-) diff --git a/tools/kokoro/CMakeLists.txt b/tools/kokoro/CMakeLists.txt index 7c4be7e4d..bae822186 100644 --- a/tools/kokoro/CMakeLists.txt +++ b/tools/kokoro/CMakeLists.txt @@ -45,10 +45,16 @@ endif() target_compile_features(kokoro_lib PUBLIC cxx_std_17) -# Standalone CLI harness — required by J2 verification (tools/voice-kokoro/). -add_executable(kokoro-tts tools/kokoro-tts.cpp) -target_link_libraries(kokoro-tts PRIVATE kokoro_lib) -install(TARGETS kokoro-tts RUNTIME) +# Standalone CLI harness — host-only (required by J2 verification, +# tools/voice-kokoro/). Skipped on iOS: there an executable is implicitly a +# MACOSX_BUNDLE app, so `install(TARGETS ... RUNTIME)` fails configure with +# "no BUNDLE DESTINATION". The fused mobile slice only needs kokoro_lib folded +# into libelizainference, never the CLI. +if(NOT CMAKE_SYSTEM_NAME STREQUAL "iOS") + add_executable(kokoro-tts tools/kokoro-tts.cpp) + target_link_libraries(kokoro-tts PRIVATE kokoro_lib) + install(TARGETS kokoro-tts RUNTIME) +endif() # Server-mount handler: compiled into kokoro_lib only when the server target # exists. The handler is guarded by `#ifdef LLAMA_BUILD_KOKORO` and pulls in diff --git a/tools/omnivoice/CMakeLists.txt b/tools/omnivoice/CMakeLists.txt index 72088dfa0..6763e099e 100644 --- a/tools/omnivoice/CMakeLists.txt +++ b/tools/omnivoice/CMakeLists.txt @@ -282,6 +282,46 @@ if(TARGET mtmd) "LINKER:-reexport_library,$") endif() add_dependencies(elizainference omnivoice-version) + + # Static variant for the mobile / iOS XCFramework slice. The mtp build + # (build-llama-cpp-mtp.mjs) builds BUILD_SHARED_LIBS=OFF and links the slice + # statically into LlamaCpp.xcframework — it builds `elizainference_static` + # BY NAME and collects libelizainference.a alongside the separately-built + # libllama.a / libmtmd.a / libkokoro_lib.a. A static archive carries only the + # FFI + core OBJECTS; the dependency archives' objects merge at the app's + # final link, so we record usage requirements (include dirs, ELIZA_ENABLE_* + # defines, the `if(TARGET kokoro_lib)` Kokoro fold) but do NOT reexport a + # library (that is a SHARED-only link option). + add_library(elizainference_static STATIC + ${OMNIVOICE_CORE_SOURCES} + ${OMNIVOICE_FFI_SOURCES}) + target_compile_features(elizainference_static PUBLIC cxx_std_17) + target_compile_definitions(elizainference_static + PRIVATE OMNIVOICE_BUILD ${OMNIVOICE_COMPILE_DEFINITIONS}) + if(ELIZA_ENABLE_VISION) + target_compile_definitions(elizainference_static PRIVATE ELIZA_ENABLE_VISION) + endif() + target_include_directories(elizainference_static PUBLIC + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${CMAKE_CURRENT_SOURCE_DIR}/include + ${CMAKE_CURRENT_BINARY_DIR} + ${OMNIVOICE_LLAMA_PRIVATE_INCLUDE_DIRS}) + target_include_directories(elizainference_static PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/../mtmd + ${CMAKE_CURRENT_SOURCE_DIR}/../../include + ${CMAKE_CURRENT_SOURCE_DIR}/../../common) + target_link_libraries(elizainference_static PUBLIC llama mtmd) + target_link_libraries(elizainference_static PRIVATE eliza_voice_classifiers llama-common) + if(TARGET kokoro_lib) + target_link_libraries(elizainference_static PRIVATE kokoro_lib) + target_include_directories(elizainference_static PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/../kokoro/include) + target_compile_definitions(elizainference_static PRIVATE ELIZA_ENABLE_KOKORO) + endif() + set_target_properties(elizainference_static PROPERTIES + OUTPUT_NAME elizainference + POSITION_INDEPENDENT_CODE ON) + add_dependencies(elizainference_static omnivoice-version) endif() # llama-server wiring — mount the OmniVoice POST /v1/audio/speech route From 678473455668e671dd7851852919eecf41b09630 Mon Sep 17 00:00:00 2001 From: Shaw Date: Tue, 23 Jun 2026 21:36:56 -0700 Subject: [PATCH 02/17] =?UTF-8?q?fix(fused-ffi):=20complete=20ABI=20v9=20c?= =?UTF-8?q?ontext=5Fsize=20+=20bump=20stale=20sizeof=20assert=20(80?= =?UTF-8?q?=E2=86=9288)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The streaming-LLM config struct gained `context_size` (int32 @ offset 80) for ABI v9 — `_llm_stream_open` now honors `cfg->context_size` (>0) instead of only the ELIZA_LLM_N_CTX env default. The TS marshaller (ffi-bindings.ts) was already emitting 88 bytes with context_size at offset 80, but the C-side ABI-guard `static_assert(sizeof(eliza_llm_stream_config_t) == 80)` was never bumped, so the fused `elizainference` target failed to compile. Bump it to 88 to match the real layout (8×int32 + 5×ptr, pointer-aligned). Validated: a host Metal build of `libelizainference` (ABI v12) loads + generates on the real google/gemma-4-E2B with context_size=4096 honored (83 tok/s, M4 Max). Co-Authored-By: Claude Opus 4.8 --- tools/omnivoice/include/eliza-inference-ffi.h | 6 +++++- tools/omnivoice/src/eliza-inference-ffi.cpp | 17 +++++++++++------ 2 files changed, 16 insertions(+), 7 deletions(-) diff --git a/tools/omnivoice/include/eliza-inference-ffi.h b/tools/omnivoice/include/eliza-inference-ffi.h index 075660a60..492626a99 100644 --- a/tools/omnivoice/include/eliza-inference-ffi.h +++ b/tools/omnivoice/include/eliza-inference-ffi.h @@ -853,7 +853,10 @@ int eliza_inference_llm_kv_quant_supported(void); * `cache_type_k` / `cache_type_v` (ABI v8): KV-cache quantization type names * (e.g. "f16", "q8_0", "qjl1_256", "q4_polar"). NULL leaves the llama.cpp * default (f16). Mapped to ggml_type and applied to cparams.type_k/type_v. - * Mirrors desktop-llama-adapter.ts's GGML_KV_CACHE_TYPES pass-through. */ + * Mirrors desktop-llama-adapter.ts's GGML_KV_CACHE_TYPES pass-through. + * + * `context_size` (ABI v9): runtime context window in tokens. <=0 falls back + * to ELIZA_LLM_N_CTX or the native default. */ typedef struct { int32_t max_tokens; float temperature; @@ -870,6 +873,7 @@ typedef struct { int32_t n_gpu_layers; /* -1 = default (all), 0 = CPU (ABI v8) */ const char * cache_type_k; /* KV K-cache quant name; NULL = f16 (ABI v8) */ const char * cache_type_v; /* KV V-cache quant name; NULL = f16 (ABI v8) */ + int32_t context_size; /* Runtime context tokens; <=0 = env/default (ABI v9) */ } eliza_llm_stream_config_t; /* Opaque streaming-LLM session. One per active generation. */ diff --git a/tools/omnivoice/src/eliza-inference-ffi.cpp b/tools/omnivoice/src/eliza-inference-ffi.cpp index dc74dd8eb..2821dcd22 100644 --- a/tools/omnivoice/src/eliza-inference-ffi.cpp +++ b/tools/omnivoice/src/eliza-inference-ffi.cpp @@ -22,13 +22,16 @@ // ABI guard: the TS loader (ffi-llm-streaming-abi.ts) marshals // eliza_llm_stream_config_t by hand-written field offsets, so any reorder / // insert / type change on the C side silently corrupts every streaming-LLM -// call. Pin the on-the-wire layout (documented "sizeof config = 80" since v8): -// 6×int32 + 5×ptr + 4-byte fields packed to 80 bytes on a 64-bit ABI. Adding a -// field is an ABI bump — update this assert AND the TS marshaller together. +// call. Pin the on-the-wire layout. ABI v9 appended `context_size` (int32 at +// offset 80), so the packed size on a 64-bit ABI is 88 bytes (8×int32 + 5×ptr, +// pointer-aligned). The TS marshaller (ffi-bindings.ts) already allocs 88 and +// writes context_size at offset 80; this assert had simply not been bumped to +// match. Adding a field is an ABI bump — update this assert AND the TS +// marshaller together. static_assert( - sizeof(eliza_llm_stream_config_t) == 80, + sizeof(eliza_llm_stream_config_t) == 88, "eliza_llm_stream_config_t layout changed — bump ABI + update the TS " - "marshaller in ffi-llm-streaming-abi.ts, then update this assert."); + "marshaller in ffi-bindings.ts, then update this assert."); /* common/ — the same-file MTP speculative-decode engine wired into the * streaming-LLM text path (ABI v8) reuses the DRAFT_MTP implementation in @@ -2922,7 +2925,9 @@ EliLlmStream * eliza_inference_llm_stream_open( * batch / threads / flash-attn / KV-quant). */ llama_context_params cparams = llama_context_default_params(); const int n_ctx_train = llama_model_n_ctx_train(model); - int n_ctx = eliza_int_env_or_default("ELIZA_LLM_N_CTX", 8192); + int n_ctx = cfg->context_size > 0 + ? cfg->context_size + : eliza_int_env_or_default("ELIZA_LLM_N_CTX", 8192); if (n_ctx_train > 0 && n_ctx > n_ctx_train) n_ctx = n_ctx_train; cparams.n_ctx = (uint32_t) n_ctx; cparams.n_batch = (uint32_t) eliza_int_env_or_default("ELIZA_LLM_N_BATCH", 512); From 05baa0f69147294f93e5336d11fb28ac891c5837 Mon Sep 17 00:00:00 2001 From: Shaw Date: Tue, 23 Jun 2026 23:04:33 -0700 Subject: [PATCH 03/17] fix(metal): revert broken TBQ3_0/TBQ4_0 attn-score kernels to verified-correct sources Commit 412b8487b ('fix(metal): correct TBQ3_0/TBQ4_0 attention score kernels') regressed kernel_turbo3_dot/kernel_turbo4_dot in eliza-shipped/: standalone metal_verify drops to 0/8 (outputs 20-170x off) and the built-fork Metal graph-dispatch smoke fails GGML_OP_ATTN_SCORE_TBQ/turbo3+turbo4. The metal-verify gate only tested native/metal/*.metal (the verify-harness copy), not the eliza-shipped/ copies the runtime embeds, so the regression shipped unguarded. Restore eliza-shipped/turbo3.metal + turbo4.metal from the verified native/metal copies (metal_verify 8/8). After the fix: dispatch_smoke PASS 9/9 routes; all shipped kernels pass parity. Co-Authored-By: Claude Opus 4.8 --- .../src/ggml-metal/eliza-shipped/turbo3.metal | 237 ++++++++---------- .../src/ggml-metal/eliza-shipped/turbo4.metal | 98 ++------ 2 files changed, 115 insertions(+), 220 deletions(-) diff --git a/ggml/src/ggml-metal/eliza-shipped/turbo3.metal b/ggml/src/ggml-metal/eliza-shipped/turbo3.metal index 3024402f8..6885fb8d9 100644 --- a/ggml/src/ggml-metal/eliza-shipped/turbo3.metal +++ b/ggml/src/ggml-metal/eliza-shipped/turbo3.metal @@ -1,120 +1,55 @@ -// # ELIZA-KERNEL-PATCH-V1 — copied verbatim from packages/inference/metal/turbo3.metal -// at build time by build-llama-cpp-dflash.mjs. Do not edit in place; -// edit the standalone source and rerun the build. +// HARDWARE VERIFIED on Apple M4 Max (Metal runtime JIT): 8/8 PASS against the +// fixture harness. Source-level verified against fork's dequantize_turbo3_0_t4 +// at ggml/src/ggml-metal/ggml-metal.metal:700 (commit 6575873e9c). // // turbo3 KV cache dequant + Q·K dot product (Metal Shading Language). // -// Block layout (block_tbq3_0 in ggml-common.h, 14 bytes): -// half norm // per-block scale (corrected) -// uchar qs[QK_TBQ*3/8 = 12] // 32 codes × 3 bits, packed straight +// Ports buun-llama-cpp's CUDA dequantize_turbo3_0 from +// ggml/src/ggml-cuda/turbo-quant-cuda.cuh and matches the fork's Metal +// dequantize_turbo3_0_t4 byte-for-byte. // -// Element decode (matches CPU reference dequantize_row_tbq3_0 + -// ggml/src/ggml-quants.c:66 tbq3_get_code): -// elem i in 0..31 of a 32-element block: -// bit = i * 3 -// byte = bit / 8 -// shift = bit % 8 -// bits = qs[byte] >> shift -// if shift > 5 and byte+1 < 12: bits |= qs[byte+1] << (8 - shift) -// code = bits & 0x7 -// rotated = k_tbq3_codebook[code] * norm -// then: y = tbq_uncondition_block(rotated) -// = k_tbq_signs .* H32(rotated) +// Block layout (block_turbo3_0 in ggml-common.h, 14 bytes): +// half norm // [0..1] fp16 corrected group norm +// uchar qs[8] // [2..9] QK_TURBO3/4 = 8 bytes (4 elements per byte, low 2 bits) +// uchar signs[4] // [10..13] QK_TURBO3/8 = 4 bytes (1 sign-bit per element) // -// By the orthogonality of H32 (FWHT normalized by 1/sqrt(32)) and the -// distributivity of pointwise sign multiply, -// = -// so we precompute q_t = H32(q .* sign) once per (q_head, batch) launch in -// threadgroup memory and dot q_t against the raw decoded codebook value. -// Bit-exactly equivalent to the CPU dequant + dot. +// Element decode (matches fork's _t4 path): +// elem 0..31 within a 32-element block: +// qb = qs[elem >> 2] // 4 elements per byte +// low2 = (qb >> ((elem & 3) * 2)) & 0x3 +// sb = signs[elem >> 3] // 1 bit per element +// hi1 = (sb >> (elem & 7)) & 0x1 +// idx = low2 | (hi1 << 2) // full 3-bit index +// k = TURBO_CENTROIDS_3BIT[idx] * norm // -// Dispatch: one threadgroup per group of `blocks_per_threadgroup` consecutive -// KV tokens for a single (q_head, batch). Threadgroup size = 32 (one Apple -// SIMD-group). Each thread handles 4 of the 128 elements per token. +// Four 32-element blocks form one 128-element rotation group. +// +// CORRECTNESS: the standalone verifier fixture uses +// eliza_quantize_turbo3_group() followed by eliza_dot_q_turbo3(). The reference +// dequantizes to the 128-wide turbo-rotated representation and dots the caller's +// Q vector directly against those codebook values. The 32-wide TBQ +// preconditioner is used by turbo4 and other TBQ paths, not by this turbo3 +// fixture contract. +// +// Dispatch: one threadgroup per (n_kv, n_head). Threadgroup size MUST equal +// 32 (one Apple SIMD-group). Each thread handles 4 of the 128 elements and +// the per-threadgroup reduction is a single simd_sum. #include using namespace metal; -// Match block_tbq3_0 layout exactly (14 bytes, alignment 2). +// Match block_turbo3_0 layout exactly (14 bytes, alignment 2). struct block_turbo3_0 { half norm; - uint8_t qs[12]; -}; - -// k_tbq3_codebook from ggml/src/ggml-quants.c:35. Unlike the legacy fork -// centroids, this set has magnitudes ~11.3× larger; the per-block norm -// absorbs the scale during quantization. -constant float TBQ3_CODEBOOK[8] = { - -2.1519457f, -1.3439093f, -0.7560053f, -0.2450942f, - 0.2450942f, 0.7560053f, 1.3439093f, 2.1519457f, + uint8_t qs[8]; + uint8_t signs[4]; }; -// k_tbq_signs[QK_TBQ=32] from ggml/src/ggml-quants.c:59. Reused per block. -constant float TBQ_SIGNS_32[32] = { - 1.0f, -1.0f, 1.0f, 1.0f, -1.0f, 1.0f, -1.0f, -1.0f, - 1.0f, 1.0f, -1.0f, 1.0f, -1.0f, -1.0f, 1.0f, -1.0f, - -1.0f, 1.0f, 1.0f, -1.0f, 1.0f, -1.0f, -1.0f, 1.0f, - 1.0f, -1.0f, 1.0f, -1.0f, -1.0f, 1.0f, -1.0f, 1.0f, +constant float TURBO_CENTROIDS_3BIT[8] = { + -0.190685f, -0.117832f, -0.065717f, -0.021460f, + 0.021460f, 0.065717f, 0.117832f, 0.190685f, }; -// Extract the 3-bit code at element index `i` in 0..31 from the 12-byte -// packed stream stored in device memory. Port of tbq3_get_code in -// ggml-quants.c:66. -static inline uint tbq3_get_code_dev(device const uint8_t * qs, uint i) { - uint bit = i * 3u; - uint byte = bit >> 3u; - uint shift = bit & 7u; - uint bits = uint(qs[byte]) >> shift; - if (shift > 5u && byte + 1u < 12u) { - bits |= uint(qs[byte + 1u]) << (8u - shift); - } - return bits & 0x7u; -} - -// In-place Fast Walsh–Hadamard transform on a 32-element block, with the -// 1/sqrt(32) normalization that makes H32 orthogonal. Mirrors -// tbq_hadamard32 in ggml/src/ggml-quants.c:104. -static inline void tbq_hadamard32_local(thread float * x) { - for (uint len = 1; len < 32u; len <<= 1) { - for (uint i = 0; i < 32u; i += 2u * len) { - for (uint j = 0; j < len; ++j) { - float a = x[i + j]; - float b = x[i + j + len]; - x[i + j] = a + b; - x[i + j + len] = a - b; - } - } - } - const float norm = 0.1767766952966369f; - for (uint i = 0; i < 32u; ++i) { - x[i] *= norm; - } -} - -// Precompute q_t[128] = H32(q .* k_tbq_signs) per 32-element block. -// Called once per threadgroup at launch; q is constant across all KV tokens -// processed by this threadgroup. -// -// Distributed across the 32-thread SIMD-group: threads 0..3 each own one of -// the 4 hadamard-32 blocks. Other threads idle through the barrier. -static inline void eliza_tbq_precompute_qt( - device const float * q_head, - threadgroup float * q_t, - uint tid) { - if (tid < 4u) { - thread float buf[32]; - uint base = tid * 32u; - for (uint i = 0; i < 32u; ++i) { - buf[i] = q_head[base + i] * TBQ_SIGNS_32[i]; - } - tbq_hadamard32_local(buf); - for (uint i = 0; i < 32u; ++i) { - q_t[base + i] = buf[i]; - } - } - threadgroup_barrier(mem_flags::mem_threadgroup); -} - struct turbo_dot_args { uint head_dim; // must be 128 uint n_kv; @@ -130,34 +65,49 @@ kernel void kernel_turbo3_dot( constant turbo_dot_args & args [[buffer(3)]], uint tid [[thread_position_in_threadgroup]], uint kv_idx [[threadgroup_position_in_grid]]) { - if (kv_idx >= args.n_kv) return; + // 32 threads × 4 elements = 128 head_dim entries. Each thread's 4 elements + // (tid*4 + 0..3) lie wholly within ONE 32-element block (since 32 is a + // multiple of 4 and tid*4 ∈ {0,4,...,124}). + uint elem0 = tid * 4u; // 0,4,...,124 + uint blk_idx = elem0 >> 5; // 0..3 + uint within0 = elem0 & 31u; // 0,4,...,28 - threadgroup float q_t[128]; - device const float * q_head = q + args.q_head * args.head_dim; - eliza_tbq_precompute_qt(q_head, q_t, tid); + if (kv_idx >= args.n_kv) return; + // Resolve the 4-block group for this KV index. Cast through uchar* so the + // optional head_offset_bytes can be a non-zero stride (still must be a + // multiple of sizeof(block_turbo3_0) = 14). device const block_turbo3_0 * grp = (device const block_turbo3_0 *)((device const uchar *)k_blocks + args.head_offset_bytes) + kv_idx * args.kv_stride_blocks; - uint elem0 = tid * 4; - uint blk_idx = elem0 >> 5; - uint within0 = elem0 & 31; + device const block_turbo3_0 & blk = grp[blk_idx]; float norm = float(blk.norm); - - float4 qtv = float4(q_t[elem0 + 0], q_t[elem0 + 1], q_t[elem0 + 2], q_t[elem0 + 3]); - - uint c0 = tbq3_get_code_dev(blk.qs, within0 + 0u); - uint c1 = tbq3_get_code_dev(blk.qs, within0 + 1u); - uint c2 = tbq3_get_code_dev(blk.qs, within0 + 2u); - uint c3 = tbq3_get_code_dev(blk.qs, within0 + 3u); + // All four elements of this thread share the same qs[] byte (within>>2 is + // constant for within = within0..within0+3) and the same signs[] byte + // (within>>3 is constant for within = within0..within0+3). + uint qb = blk.qs[within0 >> 2]; + uint sb = blk.signs[within0 >> 3]; + + uint q_base = args.q_head * args.head_dim + elem0; + device const float4 * q4 = (device const float4 *)(q + q_base); + float4 qv = q4[0]; + uint sign_shift = within0 & 7u; + uint idx0 = ((qb >> 0) & 0x3u) | (((sb >> (sign_shift + 0u)) & 0x1u) << 2); + uint idx1 = ((qb >> 2) & 0x3u) | (((sb >> (sign_shift + 1u)) & 0x1u) << 2); + uint idx2 = ((qb >> 4) & 0x3u) | (((sb >> (sign_shift + 2u)) & 0x1u) << 2); + uint idx3 = ((qb >> 6) & 0x3u) | (((sb >> (sign_shift + 3u)) & 0x1u) << 2); float4 kv = float4( - TBQ3_CODEBOOK[c0], - TBQ3_CODEBOOK[c1], - TBQ3_CODEBOOK[c2], - TBQ3_CODEBOOK[c3]) * norm; - float acc = dot(qtv, kv); - + TURBO_CENTROIDS_3BIT[idx0], + TURBO_CENTROIDS_3BIT[idx1], + TURBO_CENTROIDS_3BIT[idx2], + TURBO_CENTROIDS_3BIT[idx3]) * norm; + float acc = dot(qv, kv); + + // Threadgroup reduction. With threadgroup size == SIMD-group size == 32, + // simd_sum returns the full 128-element dot product to every lane and lane + // 0 writes the result. If the dispatch ever uses a larger threadgroup, + // this needs to switch to threadgroup-shared storage + barrier. float sum = simd_sum(acc); if (tid == 0) { scores[args.q_head * args.n_kv + kv_idx] = sum; @@ -166,7 +116,14 @@ kernel void kernel_turbo3_dot( // Multi-block-per-dispatch variant. Identical math; the threadgroup processes // `blocks_per_threadgroup` consecutive KV indices serially in a 32-thread loop, -// trading dispatch grid breadth for amortised launch tax. +// trading dispatch grid breadth for amortised launch tax. Bench shape: +// +// grid_x = ceil(n_kv / blocks_per_threadgroup) +// tg_x = 32 +// +// `args.q_head` / `args.head_offset_bytes` semantics unchanged. The shader +// derives the absolute kv index from `tg_pos.x * blocks_per_threadgroup + b` +// where `b` is the inner loop counter. struct turbo_dot_multi_args { uint head_dim; uint n_kv; @@ -183,14 +140,13 @@ kernel void kernel_turbo3_dot_multi( constant turbo_dot_multi_args & args [[buffer(3)]], uint tid [[thread_position_in_threadgroup]], uint tg_idx [[threadgroup_position_in_grid]]) { - threadgroup float q_t[128]; - device const float * q_head = q + args.q_head * args.head_dim; - eliza_tbq_precompute_qt(q_head, q_t, tid); - - uint elem0 = tid * 4; + uint elem0 = tid * 4u; uint blk_idx = elem0 >> 5; - uint within0 = elem0 & 31; - float4 qtv = float4(q_t[elem0 + 0], q_t[elem0 + 1], q_t[elem0 + 2], q_t[elem0 + 3]); + uint within0 = elem0 & 31u; + + uint q_base = args.q_head * args.head_dim + elem0; + device const float4 * q4 = (device const float4 *)(q + q_base); + float4 qv = q4[0]; uint kv_base = tg_idx * args.blocks_per_threadgroup; for (uint b = 0; b < args.blocks_per_threadgroup; ++b) { @@ -202,17 +158,20 @@ kernel void kernel_turbo3_dot_multi( + kv_idx * args.kv_stride_blocks; device const block_turbo3_0 & blk = grp[blk_idx]; float norm = float(blk.norm); - - uint c0 = tbq3_get_code_dev(blk.qs, within0 + 0u); - uint c1 = tbq3_get_code_dev(blk.qs, within0 + 1u); - uint c2 = tbq3_get_code_dev(blk.qs, within0 + 2u); - uint c3 = tbq3_get_code_dev(blk.qs, within0 + 3u); + uint qb = blk.qs[within0 >> 2]; + uint sb = blk.signs[within0 >> 3]; + + uint sign_shift = within0 & 7u; + uint idx0 = ((qb >> 0) & 0x3u) | (((sb >> (sign_shift + 0u)) & 0x1u) << 2); + uint idx1 = ((qb >> 2) & 0x3u) | (((sb >> (sign_shift + 1u)) & 0x1u) << 2); + uint idx2 = ((qb >> 4) & 0x3u) | (((sb >> (sign_shift + 2u)) & 0x1u) << 2); + uint idx3 = ((qb >> 6) & 0x3u) | (((sb >> (sign_shift + 3u)) & 0x1u) << 2); float4 kv = float4( - TBQ3_CODEBOOK[c0], - TBQ3_CODEBOOK[c1], - TBQ3_CODEBOOK[c2], - TBQ3_CODEBOOK[c3]) * norm; - float acc = dot(qtv, kv); + TURBO_CENTROIDS_3BIT[idx0], + TURBO_CENTROIDS_3BIT[idx1], + TURBO_CENTROIDS_3BIT[idx2], + TURBO_CENTROIDS_3BIT[idx3]) * norm; + float acc = dot(qv, kv); float sum = simd_sum(acc); if (tid == 0) { diff --git a/ggml/src/ggml-metal/eliza-shipped/turbo4.metal b/ggml/src/ggml-metal/eliza-shipped/turbo4.metal index 7554efeeb..c0886d420 100644 --- a/ggml/src/ggml-metal/eliza-shipped/turbo4.metal +++ b/ggml/src/ggml-metal/eliza-shipped/turbo4.metal @@ -1,6 +1,6 @@ -// # ELIZA-KERNEL-PATCH-V1 — copied verbatim from packages/inference/metal/turbo4.metal -// at build time by build-llama-cpp-dflash.mjs. Do not edit in place; -// edit the standalone source and rerun the build. +// HARDWARE VERIFIED on Apple M4 Max (Metal runtime JIT): standalone fixture +// harness plus built-fork GGML_OP_ATTN_SCORE_TBQ graph dispatch. +// Source-level verified against fork block_tbq4_0 in ggml-common.h. // // turbo4 KV cache dequant + Q·K dot product (Metal Shading Language). // @@ -8,24 +8,17 @@ // half norm // block RMS after TurboQuant preconditioning // uchar qs[16] // 4-bit indices packed like q4_0 // -// Element decode (matches CPU reference dequantize_row_tbq4_0): +// Element decode (matches reference / Python ground truth): // elem 0..31 within a 32-element block: // qb = qs[elem & 15] // idx = elem < 16 ? (qb & 0xF) : (qb >> 4) -// rotated = TURBO_CENTROIDS_4BIT[idx] * norm -// then: y = tbq_uncondition_block(rotated) -// = k_tbq_signs .* H32(rotated) +// k = TURBO_CENTROIDS_4BIT[idx] * norm // -// By the orthogonality of H32 (FWHT normalized by 1/sqrt(32)) and the -// distributivity of pointwise sign multiply, -// = -// so we precompute q_t = H32(q .* sign) once per (q_head, batch) launch in -// threadgroup memory and dot q_t against the raw decoded codebook value. -// This is bit-exactly equivalent to the CPU dequant + dot. +// Four 32-element blocks form one 128-element attention row. Graph pre-rotates +// Q, so the shader accumulates directly against the stored rotated codes. // -// Dispatch: one threadgroup per group of `blocks_per_threadgroup` consecutive -// KV tokens for a single (q_head, batch). Threadgroup size = 32 (one Apple -// SIMD-group). Each thread handles 4 of the 128 elements per token. +// Dispatch: one threadgroup per (n_kv, n_head). Threadgroup size MUST equal +// 32 (one Apple SIMD-group). Each thread handles 4 of the 128 elements. #include using namespace metal; @@ -42,58 +35,6 @@ constant float TURBO_CENTROIDS_4BIT[16] = { 1.2557391f, 1.6175243f, 2.0685055f, 2.7321365f, }; -// k_tbq_signs[QK_TBQ=32] from ggml/src/ggml-quants.c:59. Reused per block. -constant float TBQ_SIGNS_32[32] = { - 1.0f, -1.0f, 1.0f, 1.0f, -1.0f, 1.0f, -1.0f, -1.0f, - 1.0f, 1.0f, -1.0f, 1.0f, -1.0f, -1.0f, 1.0f, -1.0f, - -1.0f, 1.0f, 1.0f, -1.0f, 1.0f, -1.0f, -1.0f, 1.0f, - 1.0f, -1.0f, 1.0f, -1.0f, -1.0f, 1.0f, -1.0f, 1.0f, -}; - -// In-place Fast Walsh–Hadamard transform on a 32-element block, with the -// 1/sqrt(32) normalization that makes H32 orthogonal. Mirrors -// tbq_hadamard32 in ggml/src/ggml-quants.c:104. -static inline void tbq_hadamard32_local(thread float * x) { - for (uint len = 1; len < 32u; len <<= 1) { - for (uint i = 0; i < 32u; i += 2u * len) { - for (uint j = 0; j < len; ++j) { - float a = x[i + j]; - float b = x[i + j + len]; - x[i + j] = a + b; - x[i + j + len] = a - b; - } - } - } - const float norm = 0.1767766952966369f; - for (uint i = 0; i < 32u; ++i) { - x[i] *= norm; - } -} - -// Precompute q_t[128] = H32(q .* k_tbq_signs) per 32-element block. -// Called once per threadgroup at launch; q is constant across all KV tokens -// processed by this threadgroup. -// -// Distributed across the 32-thread SIMD-group: threads 0..3 each own one of -// the 4 hadamard-32 blocks. Other threads idle through the barrier. -static inline void eliza_tbq_precompute_qt( - device const float * q_head, - threadgroup float * q_t, - uint tid) { - if (tid < 4u) { - thread float buf[32]; - uint base = tid * 32u; - for (uint i = 0; i < 32u; ++i) { - buf[i] = q_head[base + i] * TBQ_SIGNS_32[i]; - } - tbq_hadamard32_local(buf); - for (uint i = 0; i < 32u; ++i) { - q_t[base + i] = buf[i]; - } - } - threadgroup_barrier(mem_flags::mem_threadgroup); -} - struct turbo_dot_args { uint head_dim; // must be 128 uint n_kv; @@ -111,10 +52,6 @@ kernel void kernel_turbo4_dot( uint kv_idx [[threadgroup_position_in_grid]]) { if (kv_idx >= args.n_kv) return; - threadgroup float q_t[128]; - device const float * q_head = q + args.q_head * args.head_dim; - eliza_tbq_precompute_qt(q_head, q_t, tid); - device const block_turbo4_0 * grp = (device const block_turbo4_0 *)((device const uchar *)k_blocks + args.head_offset_bytes) + kv_idx * args.kv_stride_blocks; @@ -123,9 +60,10 @@ kernel void kernel_turbo4_dot( uint within0 = elem0 & 31; device const block_turbo4_0 & blk = grp[blk_idx]; float norm = float(blk.norm); + uint q_base = args.q_head * args.head_dim + elem0; - float4 qtv = float4(q_t[elem0 + 0], q_t[elem0 + 1], q_t[elem0 + 2], q_t[elem0 + 3]); - + device const float4 * q4 = (device const float4 *)(q + q_base); + float4 qv = q4[0]; uint qb0 = blk.qs[(within0 + 0u) & 15u]; uint qb1 = blk.qs[(within0 + 1u) & 15u]; uint qb2 = blk.qs[(within0 + 2u) & 15u]; @@ -140,7 +78,7 @@ kernel void kernel_turbo4_dot( TURBO_CENTROIDS_4BIT[idx1], TURBO_CENTROIDS_4BIT[idx2], TURBO_CENTROIDS_4BIT[idx3]) * norm; - float acc = dot(qtv, kv); + float acc = dot(qv, kv); float sum = simd_sum(acc); if (tid == 0) { @@ -167,14 +105,12 @@ kernel void kernel_turbo4_dot_multi( constant turbo_dot_multi_args & args [[buffer(3)]], uint tid [[thread_position_in_threadgroup]], uint tg_idx [[threadgroup_position_in_grid]]) { - threadgroup float q_t[128]; - device const float * q_head = q + args.q_head * args.head_dim; - eliza_tbq_precompute_qt(q_head, q_t, tid); - uint elem0 = tid * 4; uint blk_idx = elem0 >> 5; uint within0 = elem0 & 31; - float4 qtv = float4(q_t[elem0 + 0], q_t[elem0 + 1], q_t[elem0 + 2], q_t[elem0 + 3]); + uint q_base = args.q_head * args.head_dim + elem0; + device const float4 * q4 = (device const float4 *)(q + q_base); + float4 qv = q4[0]; uint kv_base = tg_idx * args.blocks_per_threadgroup; for (uint b = 0; b < args.blocks_per_threadgroup; ++b) { @@ -201,7 +137,7 @@ kernel void kernel_turbo4_dot_multi( TURBO_CENTROIDS_4BIT[idx1], TURBO_CENTROIDS_4BIT[idx2], TURBO_CENTROIDS_4BIT[idx3]) * norm; - float acc = dot(qtv, kv); + float acc = dot(qv, kv); float sum = simd_sum(acc); if (tid == 0) { From 50767b8fd0b9bc901e0686510ab5fa975d452cb8 Mon Sep 17 00:00:00 2001 From: Shaw Date: Tue, 23 Jun 2026 23:28:13 -0700 Subject: [PATCH 04/17] feat(gemma4-assistant): add arch, tensors, hparams, model file, ctx_other seam - LLM_ARCH_GEMMA4_ASSISTANT enum + arch name + nextn/masked tensor enums/names/infos - hparams n_embd_inp_impl + n_embd_inp() override - model fields nextn_proj_pre/post + models.h struct decl - llm_graph_result t_h_nextn + get_h_nextn - src/models/gemma4-assistant.cpp (ported from upstream, fork naming) - ctx_other on cparams + llama_context_params + llama_get_ctx_other - llama_memory_params.mem_other + layer_share_cb Co-Authored-By: Claude Opus 4.8 --- include/llama.h | 9 +- src/llama-arch.cpp | 9 ++ src/llama-arch.h | 6 + src/llama-context.cpp | 26 +++- src/llama-cparams.h | 5 + src/llama-graph.h | 2 + src/llama-hparams.cpp | 4 + src/llama-hparams.h | 4 + src/llama-memory.h | 8 ++ src/llama-model.h | 4 + src/models/gemma4-assistant.cpp | 209 ++++++++++++++++++++++++++++++++ src/models/models.h | 13 ++ 12 files changed, 294 insertions(+), 5 deletions(-) create mode 100644 src/models/gemma4-assistant.cpp diff --git a/include/llama.h b/include/llama.h index cd020b09b..b601995c5 100644 --- a/include/llama.h +++ b/include/llama.h @@ -392,6 +392,12 @@ extern "C" { // note: the samplers must be sampler chains (i.e. use llama_sampler_chain_init) struct llama_sampler_seq_config * samplers; size_t n_samplers; + + // sibling context whose model supplies shared inputs / KV cache. + // required for arches that read from another model (gemma4-assistant MTP + // drafter reads the target model's token embeddings and shares its KV). + // null for all other arches. (the caller keeps this context alive.) + struct llama_context * ctx_other; }; struct llama_model_tensor_override { @@ -551,7 +557,8 @@ extern "C" { DEPRECATED(LLAMA_API int32_t llama_n_vocab (const struct llama_vocab * vocab), "use llama_vocab_n_tokens instead"); - LLAMA_API const struct llama_model * llama_get_model (const struct llama_context * ctx); + LLAMA_API const struct llama_model * llama_get_model (const struct llama_context * ctx); + LLAMA_API struct llama_context * llama_get_ctx_other(struct llama_context * ctx); LLAMA_API llama_memory_t llama_get_memory (const struct llama_context * ctx); LLAMA_API enum llama_pooling_type llama_pooling_type(const struct llama_context * ctx); // TODO: rename to llama_get_pooling_type diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 4bd4cbf32..1d4945ec8 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -59,6 +59,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_GEMMA3, "gemma3" }, { LLM_ARCH_GEMMA3N, "gemma3n" }, { LLM_ARCH_GEMMA4, "gemma4" }, + { LLM_ARCH_GEMMA4_ASSISTANT, "gemma4-assistant" }, { LLM_ARCH_GEMMA_EMBEDDING, "gemma-embedding" }, { LLM_ARCH_STARCODER2, "starcoder2" }, { LLM_ARCH_MAMBA, "mamba" }, @@ -461,6 +462,10 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_NEXTN_HNORM, "blk.%d.nextn.hnorm" }, { LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "blk.%d.nextn.shared_head_head" }, { LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "blk.%d.nextn.shared_head_norm" }, + { LLM_TENSOR_NEXTN_PROJ_PRE, "nextn.pre_projection" }, + { LLM_TENSOR_NEXTN_PROJ_POST, "nextn.post_projection" }, + { LLM_TENSOR_MASKED_EMBD_CENTROIDS, "masked_embd_centroids" }, + { LLM_TENSOR_MASKED_EMBD_ORDERING, "masked_embd_ordering" }, { LLM_TENSOR_ATTN_SUB_NORM, "blk.%d.attn_sub_norm" }, { LLM_TENSOR_FFN_SUB_NORM, "blk.%d.ffn_sub_norm" }, { LLM_TENSOR_DEC_OUTPUT_NORM, "dec.output_norm" }, @@ -582,6 +587,10 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_TOKEN_TYPES, {LLM_TENSOR_LAYER_INPUT, GGML_OP_GET_ROWS}}, {LLM_TENSOR_TOKEN_EMBD_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, // do the norms on the first layer (not the input layer) {LLM_TENSOR_OUTPUT, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_NEXTN_PROJ_PRE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_NEXTN_PROJ_POST, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_MASKED_EMBD_CENTROIDS, {LLM_TENSOR_LAYER_INPUT, GGML_OP_NONE}}, + {LLM_TENSOR_MASKED_EMBD_ORDERING, {LLM_TENSOR_LAYER_INPUT, GGML_OP_NONE}}, {LLM_TENSOR_CLS, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, {LLM_TENSOR_CLS_OUT, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, {LLM_TENSOR_CLS_NORM, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL}}, diff --git a/src/llama-arch.h b/src/llama-arch.h index 1dea33689..2990e8c54 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -63,6 +63,7 @@ enum llm_arch { LLM_ARCH_GEMMA3, LLM_ARCH_GEMMA3N, LLM_ARCH_GEMMA4, + LLM_ARCH_GEMMA4_ASSISTANT, LLM_ARCH_GEMMA_EMBEDDING, LLM_ARCH_STARCODER2, LLM_ARCH_MAMBA, @@ -569,6 +570,11 @@ enum llm_tensor { // EAGLE3 draft-model tensors (upstream PR #18039) LLM_TENSOR_EAGLE3_TARGET_FEATURES, LLM_TENSOR_EAGLE3_TARGET_TOK_EMBD, + // gemma4-assistant (MTP next-token drafter) tensors + LLM_TENSOR_NEXTN_PROJ_PRE, + LLM_TENSOR_NEXTN_PROJ_POST, + LLM_TENSOR_MASKED_EMBD_CENTROIDS, + LLM_TENSOR_MASKED_EMBD_ORDERING, }; enum llm_tensor_layer { diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 609e9aa35..a5578c39a 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -106,6 +106,18 @@ llama_context::llama_context( cparams.ctx_type = params.ctx_type; + cparams.ctx_other = nullptr; + + // gemma4-assistant reads the target model's token embeddings and shares its + // KV cache, so it requires a sibling target context to be supplied. + if (model.arch == LLM_ARCH_GEMMA4_ASSISTANT) { + if (params.ctx_other == nullptr) { + throw std::runtime_error("Gemma4Assistant requires ctx_other to be set (this warning is normal during memory fitting)"); + } + + cparams.ctx_other = params.ctx_other; + } + // Initialize backend samplers here so they are part of the sampling graph // before the reserve passes run later in this function. This avoids a later // re-reserve when graph nodes change. @@ -318,10 +330,11 @@ llama_context::llama_context( // init the memory module if (!hparams.vocab_only) { llama_memory_params params_mem = { - /*.type_k =*/ params.type_k, - /*.type_v =*/ params.type_v, - /*.swa_full =*/ params.swa_full, - /*.ctx_type= */ cparams.ctx_type, + /*.type_k =*/ params.type_k, + /*.type_v =*/ params.type_v, + /*.swa_full =*/ params.swa_full, + /*.ctx_type =*/ cparams.ctx_type, + /*.mem_other =*/ cparams.ctx_other ? llama_get_memory(cparams.ctx_other) : nullptr, }; memory.reset(model.create_memory(params_mem, cparams)); @@ -3358,6 +3371,7 @@ llama_context_params llama_context_default_params() { /*.kv_dynamic =*/ false, /*.sampler =*/ nullptr, /*.n_sampler =*/ 0, + /*.ctx_other =*/ nullptr, }; return result; @@ -3497,6 +3511,10 @@ const llama_model * llama_get_model(const llama_context * ctx) { return &ctx->get_model(); } +llama_context * llama_get_ctx_other(llama_context * ctx) { + return ctx->get_cparams().ctx_other; +} + enum llama_pooling_type llama_pooling_type(const llama_context * ctx) { return ctx->pooling_type(); } diff --git a/src/llama-cparams.h b/src/llama-cparams.h index ad1d1d1c8..29af834c9 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -46,6 +46,11 @@ struct llama_cparams { enum llama_context_type ctx_type; enum llama_pooling_type pooling_type; + // sibling context whose model provides shared inputs/KV (gemma4-assistant MTP + // drafter reads the target model's token embeddings + shares its KV cache). + // null for all other arches. + struct llama_context * ctx_other; + ggml_backend_sched_eval_callback cb_eval; void * cb_eval_user_data; }; diff --git a/src/llama-graph.h b/src/llama-graph.h index 9e55d0a67..38c503526 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -646,6 +646,7 @@ class llm_graph_result { ggml_tensor * get_embd() const { return t_embd; } ggml_tensor * get_embd_pooled() const { return t_embd_pooled; } ggml_tensor * get_h_pre_norm() const { return t_h_pre_norm; } + ggml_tensor * get_h_nextn() const { return t_h_nextn; } ggml_cgraph * get_gf() const { return gf; } ggml_context * get_ctx() const { return ctx_compute.get(); } @@ -675,6 +676,7 @@ class llm_graph_result { ggml_tensor * t_embd = nullptr; ggml_tensor * t_embd_pooled = nullptr; ggml_tensor * t_h_pre_norm = nullptr; // [n_embd, n_outputs] hidden state before final output norm + ggml_tensor * t_h_nextn = nullptr; // [n_embd_backbone, n_outputs] next-token hidden state (gemma4-assistant MTP) std::map t_sampled_logits; std::map t_candidates; diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index 2239309c8..81fb53149 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -71,6 +71,10 @@ uint32_t llama_hparams::n_rot(uint32_t il) const { } uint32_t llama_hparams::n_embd_inp() const { + if (n_embd_inp_impl > 0) { + return n_embd_inp_impl; + } + uint32_t n_embd_inp = n_embd; if (n_deepstack_layers > 0) { diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 556cd69d1..cb71a466e 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -188,6 +188,10 @@ struct llama_hparams { // output embedding dimension (0 = use n_embd) uint32_t n_embd_out_impl = 0; + // input/backbone embedding dimension (0 = derive from n_embd; gemma4-assistant + // sets this to the target hidden size so the MTP drafter projects to/from it) + uint32_t n_embd_inp_impl = 0; + // llama4 smallthinker uint32_t n_moe_layer_step = 0; uint32_t n_no_rope_layer_step = 4; diff --git a/src/llama-memory.h b/src/llama-memory.h index 4ad1612e4..527f02de8 100644 --- a/src/llama-memory.h +++ b/src/llama-memory.h @@ -23,6 +23,10 @@ struct llama_memory_params { bool swa_full; llama_context_type ctx_type; + + // sibling memory whose KV cache is shared per-layer (gemma4-assistant MTP + // drafter shares KV with the target). null for all other arches. + llama_memory_t mem_other; }; enum llama_memory_status { @@ -76,6 +80,10 @@ struct llama_memory_i { // return negative value to indicate that the layer il should not reuse memory using layer_reuse_cb = std::function; + // this callback is used to specify which layer of the sibling (mem_other) cache + // a given layer should share its KV with. return negative to indicate no sharing. + using layer_share_cb = std::function; + virtual ~llama_memory_i() = default; // split the input batch into a set of ubatches and verify that they can fit into the cache diff --git a/src/llama-model.h b/src/llama-model.h index de0b1bf85..c7c3b3f2d 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -549,6 +549,10 @@ struct llama_model { struct ggml_tensor * dflash_fc = nullptr; struct ggml_tensor * dflash_hidden_norm = nullptr; + // gemma4-assistant (MTP next-token drafter) projections. + struct ggml_tensor * nextn_proj_pre = nullptr; + struct ggml_tensor * nextn_proj_post = nullptr; + struct ggml_tensor * conv1d = nullptr; struct ggml_tensor * conv1d_b = nullptr; diff --git a/src/models/gemma4-assistant.cpp b/src/models/gemma4-assistant.cpp new file mode 100644 index 000000000..8b147a14a --- /dev/null +++ b/src/models/gemma4-assistant.cpp @@ -0,0 +1,209 @@ +#include "models.h" + +void llama_model_gemma4_assistant::load_arch_hparams(llama_model_loader & ml) { + hparams.n_embd_inp_impl = hparams.n_embd_out(); + + hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; + ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.swa_layers, hparams.n_layer); + + uint32_t n_kv_shared_layers = 0; + ml.get_key(LLM_KV_ATTENTION_SHARED_KV_LAYERS, n_kv_shared_layers, false); + + hparams.f_attention_scale = 1.0f; + + ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.nextn_predict_layers, false); + + ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false); + ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa); + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_SWA, hparams.n_embd_head_k_swa); + ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_SWA, hparams.n_embd_head_v_swa); + + type = LLM_TYPE_E2B; +} + +void llama_model_gemma4_assistant::load_arch_tensors(llama_model_loader &) { + LLAMA_LOAD_LOCALS; + + if (n_embd_head_k != n_embd_head_v) { + throw std::runtime_error("Gemma 4 assistant requires n_embd_head_k == n_embd_head_v"); + } + if (hparams.n_embd_head_k_swa != hparams.n_embd_head_v_swa) { + throw std::runtime_error("Gemma 4 assistant requires n_embd_head_k_swa == n_embd_head_v_swa"); + } + if (hparams.n_embd_out() == n_embd) { + throw std::runtime_error("Gemma 4 assistant requires embedding_length_out to carry the target hidden size"); + } + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0); + output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_DUPLICATED); + + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0); + + create_tensor(tn(LLM_TENSOR_MASKED_EMBD_CENTROIDS, "weight"), {}, TENSOR_NOT_REQUIRED); + create_tensor(tn(LLM_TENSOR_MASKED_EMBD_ORDERING), {}, TENSOR_NOT_REQUIRED); + + const int64_t n_embd_backbone = hparams.n_embd_inp(); + nextn_proj_post = create_tensor(tn(LLM_TENSOR_NEXTN_PROJ_POST, "weight"), { n_embd, n_embd_backbone }, 0); + + int rope_freqs_flag = 0; + + const int n_layer_nextn = (int) hparams.n_layer; + + for (int i = 0; i < n_layer_nextn; ++i) { + auto & layer = layers[i]; + + const int64_t n_head = hparams.n_head(i); + const int64_t n_embd_head = hparams.n_embd_head_k(i); + const int64_t n_ff = hparams.n_ff(i); + + if (i == 0) { + nextn_proj_pre = create_tensor(tn(LLM_TENSOR_NEXTN_PROJ_PRE, "weight", i), { 2*n_embd_backbone, n_embd }, 0); + } + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), { n_embd }, 0); + layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), { n_embd, n_embd_head*n_head }, 0); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head*n_head, n_embd }, 0); + + layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), { n_embd_head }, 0); + layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), { n_embd }, 0); + + layer.out_scale = create_tensor(tn(LLM_TENSOR_LAYER_OUT_SCALE, "weight", i), { 1u }, 0); + + if (!hparams.is_swa(i)) { + layer.rope_freqs = create_tensor(tn(LLM_TENSOR_ROPE_FREQS, "weight", i), { n_embd_head/2 }, rope_freqs_flag); + rope_freqs_flag = TENSOR_DUPLICATED; + } + + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), { n_embd }, 0); + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), { n_embd, n_ff }, 0); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), { n_embd, n_ff }, 0); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd }, 0); + layer.ffn_post_norm = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), { n_embd }, 0); + } +} + +std::unique_ptr llama_model_gemma4_assistant::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*this, params); +} + +llama_model_gemma4_assistant::graph::graph(const llama_model & model, const llm_graph_params & params) : + llm_graph_context(params) { + const int64_t n_embd_backbone = hparams.n_embd_inp(); + + const int n_layer_nextn = (int) hparams.n_layer; + + ggml_tensor * inp_tokens; + ggml_tensor * inp_h; + { + auto inp = std::make_unique(n_embd_backbone); + + inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, ubatch.n_tokens); + cb(inp->tokens, "inp_tokens", -1); + ggml_set_input(inp->tokens); + inp_tokens = inp->tokens; + res->t_inp_tokens = inp->tokens; + + inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_backbone, ubatch.n_tokens); + cb(inp->embd, "inp_h", -1); + ggml_set_input(inp->embd); + inp_h = inp->embd; + res->t_inp_embd = inp->embd; + + res->add_input(std::move(inp)); + } + + GGML_ASSERT(cparams.ctx_other != nullptr); + const auto * model_other = llama_get_model(cparams.ctx_other); + + ggml_tensor * x = ggml_get_rows(ctx0, model_other->tok_embd, inp_tokens); + x = ggml_scale(ctx0, x, sqrtf((float) n_embd_backbone)); + cb(x, "inp_embd_target", -1); + + ggml_tensor * xh = ggml_concat(ctx0, x, inp_h, 0); + cb(xh, "inp_xh", -1); + + ggml_tensor * cur = ggml_mul_mat(ctx0, model.nextn_proj_pre, xh); + cb(cur, "pre_proj", -1); + + auto * inp_attn = build_attn_inp_kv_iswa(); + ggml_tensor * inp_pos = build_inp_pos(); + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + ggml_tensor * inpL = cur; + + for (int il = 0; il < n_layer_nextn; ++il) { + const bool is_swa = hparams.is_swa(il); + + const int64_t n_embd_head = hparams.n_embd_head_k(il); + const int64_t n_head = hparams.n_head(il); + + const float freq_base_l = model.get_rope_freq_base(cparams, il); + const float freq_scale_l = model.get_rope_freq_scale(cparams, il); + const int n_rot_l = hparams.n_rot(il); + + ggml_tensor * cur_norm = build_norm(inpL, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur_norm, "attn_norm", il); + + ggml_tensor * Qcur = build_lora_mm(model.layers[il].wq, cur_norm); + Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens); + Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, nullptr, LLM_NORM_RMS, il); + cb(Qcur, "Qcur_normed", il); + + ggml_tensor * freq_factors = is_swa ? nullptr : model.layers[il].rope_freqs; + Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, freq_factors, n_rot_l, rope_type, n_ctx_orig, + freq_base_l, freq_scale_l, ext_factor, attn_factor, beta_fast, beta_slow); + cb(Qcur, "Qcur_pos", il); + + cur = build_attn(inp_attn, model.layers[il].wo, nullptr, nullptr, + Qcur, nullptr, nullptr, nullptr, nullptr, nullptr, hparams.f_attention_scale, il); + + if (il == n_layer_nextn - 1 && inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + inpL = ggml_get_rows(ctx0, inpL, inp_out_ids); + } + + cur = build_norm(cur, model.layers[il].attn_post_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "attn_post_norm", il); + + ggml_tensor * attn_out = ggml_add(ctx0, cur, inpL); + cb(attn_out, "attn_out", il); + + cur = build_norm(attn_out, model.layers[il].ffn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + cur = build_ffn(cur, + model.layers[il].ffn_up, nullptr, nullptr, + model.layers[il].ffn_gate, nullptr, nullptr, + model.layers[il].ffn_down, nullptr, nullptr, + nullptr, + LLM_FFN_GELU, LLM_FFN_PAR, il); + cb(cur, "ffn_out", il); + + cur = build_norm(cur, model.layers[il].ffn_post_norm, nullptr, LLM_NORM_RMS, -1); + cb(cur, "ffn_post_norm", il); + + cur = ggml_add(ctx0, cur, attn_out); + + cur = ggml_mul(ctx0, cur, model.layers[il].out_scale); + cb(cur, "out_scaled", il); + + inpL = cur; + } + cur = inpL; + + cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1); + cb(cur, "result_norm", -1); + + ggml_tensor * logits = build_lora_mm(model.output, cur); + cb(logits, "result_output", -1); + res->t_logits = logits; + + ggml_tensor * h_next = ggml_mul_mat(ctx0, model.nextn_proj_post, cur); + cb(h_next, "h_nextn", -1); + res->t_h_nextn = h_next; + ggml_set_output(res->t_h_nextn); + + ggml_build_forward_expand(gf, logits); + ggml_build_forward_expand(gf, h_next); +} diff --git a/src/models/models.h b/src/models/models.h index bb6c23a75..955caf427 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -813,6 +813,19 @@ struct llama_model_gemma4 : public llama_model_base { }; +struct llama_model_gemma4_assistant : public llama_model_base { + llama_model_gemma4_assistant(const struct llama_model_params & params) : llama_model_base(params) {} + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + + struct graph : public llm_graph_context { + graph(const llama_model & model, const llm_graph_params & params); + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + + struct llama_model_gemma_embedding : public llama_model_base { llama_model_gemma_embedding(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; From 5b09bc0601358238dbb348b32a1178208bfde8d6 Mon Sep 17 00:00:00 2001 From: Shaw Date: Tue, 23 Jun 2026 23:32:11 -0700 Subject: [PATCH 05/17] feat(gemma4-assistant): KV-cache sharing seam + model dispatch + rope-type - llama_kv_cache / llama_kv_cache_iswa ctors take mem_other + share - per-layer share() maps drafter layers to sibling target KV layers - all existing callers pass nullptr (additive, behavior unchanged) - model factory dispatch + NEOX rope-type for gemma4-assistant - create_memory builds share lambda + mem_other for gemma4-assistant Co-Authored-By: Claude Opus 4.8 --- src/llama-kv-cache-iswa.cpp | 15 ++++++++++++--- src/llama-kv-cache-iswa.h | 4 +++- src/llama-kv-cache.cpp | 24 +++++++++++++++++++++++- src/llama-kv-cache.h | 6 ++++++ src/llama-memory-hybrid-iswa.cpp | 4 +++- src/llama-memory-hybrid.cpp | 4 +++- src/llama-model.cpp | 28 +++++++++++++++++++++++++++- 7 files changed, 77 insertions(+), 8 deletions(-) diff --git a/src/llama-kv-cache-iswa.cpp b/src/llama-kv-cache-iswa.cpp index 26e2cb427..bf0917191 100644 --- a/src/llama-kv-cache-iswa.cpp +++ b/src/llama-kv-cache-iswa.cpp @@ -23,8 +23,10 @@ llama_kv_cache_iswa::llama_kv_cache_iswa( uint32_t n_seq_max, uint32_t n_ubatch, uint32_t n_pad, + llama_memory_t mem_other, const layer_filter_cb & filter, - const layer_reuse_cb & reuse) : hparams(model.hparams), unified(unified) { + const layer_reuse_cb & reuse, + const layer_share_cb & share) : hparams(model.hparams), unified(unified) { // chain filters const layer_filter_cb filter_base = [&](int32_t il) { @@ -59,17 +61,24 @@ llama_kv_cache_iswa::llama_kv_cache_iswa( LLAMA_LOG_INFO("%s: creating non-SWA KV cache, size = %u cells\n", __func__, size_base); + llama_memory_t mem_other_base = nullptr; + llama_memory_t mem_other_swa = nullptr; + if (mem_other) { + mem_other_base = static_cast(mem_other)->get_base(); + mem_other_swa = static_cast(mem_other)->get_swa(); + } + kv_base = std::make_unique( model, type_k, type_v, v_trans, offload, unified, size_base, n_seq_max, n_pad, - 0, LLAMA_SWA_TYPE_NONE, filter_base, reuse); + 0, LLAMA_SWA_TYPE_NONE, mem_other_base, filter_base, reuse, share); LLAMA_LOG_INFO("%s: creating SWA KV cache, size = %u cells\n", __func__, size_swa); kv_swa = std::make_unique( model, type_k, type_v, v_trans, offload, unified, size_swa, n_seq_max, n_pad, - hparams.n_swa, hparams.swa_type, filter_swa, reuse); + hparams.n_swa, hparams.swa_type, mem_other_swa, filter_swa, reuse, share); } void llama_kv_cache_iswa::clear(bool data) { diff --git a/src/llama-kv-cache-iswa.h b/src/llama-kv-cache-iswa.h index 70ab22f0d..dfafc1ef5 100644 --- a/src/llama-kv-cache-iswa.h +++ b/src/llama-kv-cache-iswa.h @@ -25,8 +25,10 @@ class llama_kv_cache_iswa : public llama_memory_i { uint32_t n_seq_max, uint32_t n_ubatch, uint32_t n_pad, + llama_memory_t mem_other, const layer_filter_cb & filter, - const layer_reuse_cb & reuse); + const layer_reuse_cb & reuse, + const layer_share_cb & share); ~llama_kv_cache_iswa() = default; diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index cd9666a21..2d08ab908 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -89,12 +89,16 @@ llama_kv_cache::llama_kv_cache( uint32_t n_pad, uint32_t n_swa, llama_swa_type swa_type, + llama_memory_t mem_other, const layer_filter_cb & filter, const layer_reuse_cb & reuse, + const layer_share_cb & share, uint32_t kv_size_max) : model(model), hparams(model.hparams), v_trans(v_trans), n_seq_max(n_seq_max), n_stream(unified ? 1 : n_seq_max), n_pad(n_pad), n_swa(n_swa), swa_type(swa_type) { + other = static_cast(mem_other); + // save construction parameters for dynamic resize saved_type_k = type_k; saved_type_v = type_v; @@ -194,6 +198,23 @@ llama_kv_cache::llama_kv_cache( continue; } + if (share && other) { + const int32_t il_share = share(il); + + if (il_share >= 0) { + const auto & layer_share = other->layers[other->map_layer_ids[il_share]]; + + LLAMA_LOG_DEBUG("%s: layer %3d: sharing with sibling layer %d\n", __func__, il, il_share); + + map_layer_ids[il] = layers.size(); + + layers.push_back(layer_share); + layers.back().il = il; + + continue; + } + } + if (n_embd_head_k_all == 0) { n_embd_head_k_all = (int32_t) hparams.n_embd_head_k(il); } else if (n_embd_head_k_all > 0 && n_embd_head_k_all != (int32_t) hparams.n_embd_head_k(il)) { @@ -1196,7 +1217,8 @@ bool llama_kv_cache::try_resize() { // NOTE: pass kv_size_max=0 so the constructor does NOT apply // the dynamic start logic (which would shrink back to 256) llama_kv_cache tmp(model, saved_type_k, saved_type_v, saved_v_trans, saved_offload, saved_unified, new_size, - saved_n_seq_max, saved_n_pad, saved_n_swa, saved_swa_type, saved_filter, saved_reuse, + saved_n_seq_max, saved_n_pad, saved_n_swa, saved_swa_type, + /*mem_other=*/nullptr, saved_filter, saved_reuse, /*share=*/nullptr, /*kv_size_max=*/0); // copy existing data diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 60b33b429..23486f322 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -105,8 +105,10 @@ class llama_kv_cache : public llama_memory_i { uint32_t n_pad, uint32_t n_swa, llama_swa_type swa_type, + llama_memory_t mem_other, const layer_filter_cb & filter, const layer_reuse_cb & reuse, + const layer_share_cb & share, uint32_t kv_size_max = 0); ~llama_kv_cache() = default; @@ -276,6 +278,10 @@ class llama_kv_cache : public llama_memory_i { std::vector layers; + // sibling KV cache to share K/V tensors with (gemma4-assistant MTP drafter + // shares KV with the target). null for all other arches. + llama_kv_cache * other = nullptr; + // dynamic resize state uint32_t kv_size_cur = 0; uint32_t kv_size_max_val = 0; diff --git a/src/llama-memory-hybrid-iswa.cpp b/src/llama-memory-hybrid-iswa.cpp index a59561ea5..99c2bb4c7 100644 --- a/src/llama-memory-hybrid-iswa.cpp +++ b/src/llama-memory-hybrid-iswa.cpp @@ -43,10 +43,12 @@ llama_memory_hybrid_iswa::llama_memory_hybrid_iswa( n_seq_max, n_ubatch, n_pad, + nullptr, // mem_other filter_attn == nullptr ? [&](int32_t il) { return !hparams.is_recurrent(il); } : filter_attn, - nullptr + nullptr, // reuse + nullptr // share )), mem_recr(new llama_memory_recurrent( model, diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index e21e42d09..03d23aa26 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -45,10 +45,12 @@ llama_memory_hybrid::llama_memory_hybrid( n_pad, n_swa, swa_type, + nullptr, // mem_other filter_attn == nullptr ? [&](int32_t il) { return !hparams.is_recurrent(il); } : filter_attn, - nullptr, + nullptr, // reuse + nullptr, // share kv_size_max )), mem_recr(new llama_memory_recurrent( diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 864774635..93ed4d7ca 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -136,6 +136,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_gemma3n(params); case LLM_ARCH_GEMMA4: return new llama_model_gemma4(params); + case LLM_ARCH_GEMMA4_ASSISTANT: + return new llama_model_gemma4_assistant(params); case LLM_ARCH_GEMMA_EMBEDDING: return new llama_model_gemma_embedding(params); case LLM_ARCH_STARCODER2: @@ -2046,6 +2048,8 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, } else { llama_memory_i::layer_reuse_cb reuse = nullptr; llama_kv_cache::layer_filter_cb filter = nullptr; + llama_kv_cache::layer_share_cb share = nullptr; + llama_memory_t mem_other = nullptr; if (arch == LLM_ARCH_GEMMA3N || arch == LLM_ARCH_GEMMA4) { reuse = [&](int32_t il) { @@ -2062,6 +2066,23 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, filter = [n_main](int32_t il) { return (uint32_t)il >= n_main; }; } + if (arch == LLM_ARCH_GEMMA4_ASSISTANT) { + // The MTP drafter carries no K/V weights; every drafter layer + // shares the target model's last (dense) or second-to-last (SWA) + // KV cache layer via the sibling ctx_other memory. + mem_other = llama_get_memory(cparams.ctx_other); + + share = [&](int32_t il) { + const llama_model * model_other = llama_get_model(cparams.ctx_other); + + if (hparams.is_swa(il)) { + return (int32_t) llama_model_n_layer(model_other) - 2; + } + + return (int32_t) llama_model_n_layer(model_other) - 1; + }; + } + if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) { GGML_ASSERT(hparams.is_swa_any()); @@ -2077,8 +2098,10 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, cparams.n_seq_max, cparams.n_ubatch, 1, + mem_other, filter, - reuse); + reuse, + share); } else { GGML_ASSERT(!hparams.is_swa_any()); @@ -2094,8 +2117,10 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, 1, hparams.n_swa, hparams.swa_type, + mem_other, filter, nullptr, + share, cparams.kv_dynamic ? cparams.n_ctx_seq : 0); } } @@ -2334,6 +2359,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_GEMMA3: case LLM_ARCH_GEMMA3N: case LLM_ARCH_GEMMA4: + case LLM_ARCH_GEMMA4_ASSISTANT: case LLM_ARCH_GEMMA_EMBEDDING: case LLM_ARCH_STARCODER2: case LLM_ARCH_OPENELM: From eadc1a3e9a2231ae233180840f5f33d53b78ac5b Mon Sep 17 00:00:00 2001 From: Shaw Date: Tue, 23 Jun 2026 23:34:39 -0700 Subject: [PATCH 06/17] feat(gemma4-assistant): wire ctx_other=ctx_tgt for draft-mtp contexts server-context sets cparams.ctx_other = ctx_tgt on both MTP draft-context branches so the gemma4-assistant drafter can read the target's token embeddings and share its KV cache. Harmless for other draft arches (only that arch reads it). Co-Authored-By: Claude Opus 4.8 --- tools/server/server-context.cpp | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index cf875391f..d31480c90 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -883,6 +883,10 @@ struct server_context_impl { COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params_base.speculative.types.end(); if (spec_mtp) { cparams.ctx_type = LLAMA_CONTEXT_TYPE_MTP; + // gemma4-assistant-style MTP drafters read the target model's + // token embeddings and share its KV cache via ctx_other. Setting + // it for other draft arches is harmless (only that arch reads it). + cparams.ctx_other = ctx_tgt; } // note: for small models maybe we can set this to the maximum possible draft from all speculative types @@ -900,8 +904,9 @@ struct server_context_impl { params_base.model.path.c_str()); auto cparams_mtp = common_context_params_to_llama(params_base); - cparams_mtp.ctx_type = LLAMA_CONTEXT_TYPE_MTP; - cparams_mtp.n_rs_seq = 0; + cparams_mtp.ctx_type = LLAMA_CONTEXT_TYPE_MTP; + cparams_mtp.n_rs_seq = 0; + cparams_mtp.ctx_other = ctx_tgt; ctx_dft.reset(llama_init_from_model(model_tgt, cparams_mtp)); if (ctx_dft == nullptr) { From 1e2949607b8662904bf0785b055f70b921cc6fcf Mon Sep 17 00:00:00 2001 From: Shaw Date: Tue, 23 Jun 2026 23:42:54 -0700 Subject: [PATCH 07/17] fix(gemma4-assistant): correct MTP hidden-state feedback for real speedup Three correctness fixes that take draft acceptance from ~2% to ~41%: 1. MTP runtime row width = n_embd_out (backbone 1536), not the drafter's internal n_embd (256); assert it matches the target hidden width. 2. pre-norm extraction (context.cpp) + get_embeddings_pre_norm_ith + embd_pre_norm buffer now sized/strided by n_embd_out, so the n_embd_out-wide nextn hidden survives the round trip. 3. target gemma4 graph now exposes its POST-final-norm hidden via t_h_pre_norm (the LM-head input feature the drafter consumes), matching upstream's t_h_nextn; gemma4-assistant routes its nextn hidden through the same seam. Result on M4 Max Metal: 133.8 t/s with draft-mtp vs 109.5 t/s baseline (1.22x). Co-Authored-By: Claude Opus 4.8 --- common/speculative.cpp | 9 ++++++++- src/llama-context.cpp | 14 +++++++++----- src/models/gemma4-assistant.cpp | 6 ++++++ src/models/gemma4.cpp | 7 +++++++ tools/server/server-context.cpp | 15 +++++++++++---- 5 files changed, 41 insertions(+), 10 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index 8bd76c108..b1cf61c93 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -434,7 +434,14 @@ struct common_speculative_state_draft_mtp : public common_speculative_impl { auto * ctx_dft = this->params.ctx_dft; GGML_ASSERT(ctx_tgt && ctx_dft && "MTP requires ctx_tgt and ctx_dft to be set"); - n_embd = llama_model_n_embd(llama_get_model(ctx_dft)); + // The MTP hidden-state I/O width is the *backbone* width, which can differ + // from the drafter's internal embedding width (gemma4-assistant: internal + // n_embd=256, but it reads/writes target hidden states of width + // n_embd_out=1536). Size every embd row against the output width and assert + // it matches the target's hidden width that we feed in. + n_embd = llama_model_n_embd_out(llama_get_model(ctx_dft)); + GGML_ASSERT(n_embd == llama_model_n_embd(llama_get_model(ctx_tgt)) && + "MTP input row width must match the target hidden width"); const int32_t n_b = (int32_t) llama_n_batch(ctx_dft); batch = llama_batch_init(/*n_tokens=*/ n_b, /*embd=*/ n_embd, /*n_seq_max=*/ 1); diff --git a/src/llama-context.cpp b/src/llama-context.cpp index a5578c39a..6ed600c6b 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -932,7 +932,7 @@ float * llama_context::get_embeddings_pre_norm_ith(int32_t i) { } const int64_t j = output_resolve_row(i); - const uint32_t n_embd = model.hparams.n_embd; + const uint32_t n_embd = model.hparams.n_embd_out(); return embd_pre_norm.data + j*n_embd; } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: invalid pre-norm embeddings id %d, reason: %s\n", __func__, i, err.what()); @@ -1472,12 +1472,15 @@ int llama_context::encode(const llama_batch & batch_inp) { } } - // extract pre-norm embeddings (hidden state before the final output norm) + // extract pre-norm embeddings (hidden state before the final output norm). + // the hidden width is n_embd_out (the backbone width): == n_embd for ordinary + // models, but wider than the drafter's n_embd for gemma4-assistant (where this + // tensor is the n_embd_out-wide nextn hidden state). if (embd_pre_norm.data && t_h_pre_norm && cparams.pooling_type == LLAMA_POOLING_TYPE_NONE) { ggml_backend_t backend_h = ggml_backend_sched_get_tensor_backend(sched.get(), t_h_pre_norm); GGML_ASSERT(backend_h != nullptr); - const uint32_t n_embd = hparams.n_embd; + const uint32_t n_embd = hparams.n_embd_out(); GGML_ASSERT(n_tokens*n_embd <= (int64_t) embd_pre_norm.size); ggml_backend_tensor_get_async(backend_h, t_h_pre_norm, embd_pre_norm.data, 0, n_tokens*n_embd*sizeof(float)); } @@ -1923,11 +1926,12 @@ int llama_context::decode(const llama_batch & batch_inp) { // extract pre-norm embeddings (hidden state before the final output norm) // only meaningful in LLAMA_POOLING_TYPE_NONE (per-token); other pooling modes are ignored. + // width is n_embd_out (backbone width); wider than n_embd for gemma4-assistant. if (embd_pre_norm.data && t_h_pre_norm && n_outputs > 0 && cparams.pooling_type == LLAMA_POOLING_TYPE_NONE) { ggml_backend_t backend_h = ggml_backend_sched_get_tensor_backend(sched.get(), t_h_pre_norm); GGML_ASSERT(backend_h != nullptr); - const uint32_t n_embd = hparams.n_embd; + const uint32_t n_embd = hparams.n_embd_out(); float * embd_pre_norm_out = embd_pre_norm.data + n_outputs_prev*n_embd; GGML_ASSERT( n_outputs_prev + n_outputs <= n_outputs_all); @@ -2038,7 +2042,7 @@ uint32_t llama_context::output_reserve(int32_t n_outputs) { logits.size = has_logits ? n_vocab*n_outputs_max : 0; embd.size = has_embd ? n_embd_out*n_outputs_max : 0; - embd_pre_norm.size = has_embd_pre_norm ? n_embd*n_outputs_max : 0; + embd_pre_norm.size = has_embd_pre_norm ? n_embd_out*n_outputs_max : 0; // Allocate backend sampling output buffers if there are backend samplers configured. const bool has_sampling = !sampling.samplers.empty(); diff --git a/src/models/gemma4-assistant.cpp b/src/models/gemma4-assistant.cpp index 8b147a14a..fc8280ed5 100644 --- a/src/models/gemma4-assistant.cpp +++ b/src/models/gemma4-assistant.cpp @@ -204,6 +204,12 @@ llama_model_gemma4_assistant::graph::graph(const llama_model & model, const llm_ res->t_h_nextn = h_next; ggml_set_output(res->t_h_nextn); + // Route the next-token hidden state through the existing pre-norm extraction + // seam so the MTP runtime (common/speculative.cpp) can chain it back as the + // drafter's input hidden for the next speculative step. The pre-norm/nextn + // hidden width is n_embd_out (the backbone width), not the drafter's n_embd. + res->t_h_pre_norm = h_next; + ggml_build_forward_expand(gf, logits); ggml_build_forward_expand(gf, h_next); } diff --git a/src/models/gemma4.cpp b/src/models/gemma4.cpp index 4f9d8b18b..dc309cf11 100644 --- a/src/models/gemma4.cpp +++ b/src/models/gemma4.cpp @@ -376,6 +376,13 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para model.output_norm, nullptr, LLM_NORM_RMS, -1); + // Expose the post-output-norm hidden state (the LM-head input feature) through + // the pre-norm extraction seam so gemma4-assistant MTP draft contexts can read + // it as the recurrent h input. This matches the reference (transformers/vLLM/ + // SGLang), which feeds the drafter the target's post-final-norm hidden state. + cb(cur, "h_nextn", -1); + res->t_h_pre_norm = cur; + cb(cur, "result_norm", -1); res->t_embd = cur; diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index d31480c90..c8fc91c43 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -866,21 +866,28 @@ struct server_context_impl { return false; } + const bool spec_mtp = std::find(params_base.speculative.types.begin(), + params_base.speculative.types.end(), + COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params_base.speculative.types.end(); + // Upstream PR #22660: for SWA draft models, force swa_full on the // draft context so prefix reuse (seq_rm + seq_add) works beyond // the SWA window during speculation. Without this, the draft has // to re-decode from the window edge on every long-context request // and acceptance length degrades sharply. - if (llama_model_n_swa(model_dft.get()) > 0 && !params_dft.swa_full) { + // + // Exception: an MTP drafter that shares the target's KV cache (e.g. + // gemma4-assistant via ctx_other) must size its SWA cache to match + // the target's exactly. Forcing swa_full here would make the drafter + // expect a full-size SWA cache while the shared target tensor is + // small-SWA-sized, overflowing the view in get_k/get_v. + if (!spec_mtp && llama_model_n_swa(model_dft.get()) > 0 && !params_dft.swa_full) { SRV_INF("%s", "draft model uses SWA - enabling swa_full for the draft context\n"); params_dft.swa_full = true; } auto cparams = common_context_params_to_llama(params_dft); - const bool spec_mtp = std::find(params_base.speculative.types.begin(), - params_base.speculative.types.end(), - COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params_base.speculative.types.end(); if (spec_mtp) { cparams.ctx_type = LLAMA_CONTEXT_TYPE_MTP; // gemma4-assistant-style MTP drafters read the target model's From 2323fa8633580c7aecea39ce77944d4f3f3debda Mon Sep 17 00:00:00 2001 From: Shaw Date: Tue, 23 Jun 2026 23:44:39 -0700 Subject: [PATCH 08/17] chore(gemma4-assistant): drop now-unused n_embd local in output_reserve embd_pre_norm.size now derives from n_embd_out; the n_embd local is dead. Co-Authored-By: Claude Opus 4.8 --- src/llama-context.cpp | 1 - 1 file changed, 1 deletion(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 6ed600c6b..de916a2b4 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2023,7 +2023,6 @@ uint32_t llama_context::output_reserve(int32_t n_outputs) { const auto n_batch = cparams.n_batch; const auto n_vocab = vocab.n_tokens(); - const auto n_embd = hparams.n_embd; const auto n_embd_out = hparams.n_embd_out(); bool has_logits = true; From 0864259de47b11f5382d6dd368219de4bc0bc7b2 Mon Sep 17 00:00:00 2001 From: Shaw Date: Wed, 24 Jun 2026 00:04:40 -0700 Subject: [PATCH 09/17] fix(fused-ffi): set ctx_other on the MTP draft context (gemma4-assistant) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The fused engine builds its MTP draft context (eliza-inference-ffi.cpp) without ctx_other, so a separate-drafter MTP arch that reaches into the target context (gemma4-assistant: target tok_embd + hidden) hard-failed in the llama-context ctor ("Gemma4Assistant requires ctx_other to be set"). server-context.cpp set it; the fused FFI path did not. Set cp.ctx_other = e->ctx_tgt before init (inert for same-file MTP archs that don't consult ctx_other). Validated via the fused FFI on M4 Max Metal: loads the amaranus gemma4-assistant drafter + runs MTP — drafted 28 / accepted 12 (43%), coherent output, 112 tok/s. Co-Authored-By: Claude Opus 4.8 --- tools/omnivoice/src/eliza-inference-ffi.cpp | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tools/omnivoice/src/eliza-inference-ffi.cpp b/tools/omnivoice/src/eliza-inference-ffi.cpp index 2821dcd22..a0390cacc 100644 --- a/tools/omnivoice/src/eliza-inference-ffi.cpp +++ b/tools/omnivoice/src/eliza-inference-ffi.cpp @@ -1407,6 +1407,11 @@ static Engine * create_engine( cp.n_rs_seq = 0; // draft ctx rolls back via PART/checkpoint, not RS cp.n_threads = cparams_tgt.n_threads; cp.n_threads_batch = cparams_tgt.n_threads_batch; + // Separate-drafter MTP archs that reach into the target context for its + // token embeddings + hidden state (e.g. gemma4-assistant) require + // `ctx_other` = the target context; the llama-context ctor hard-fails + // without it. Inert for archs that don't consult ctx_other (same-file MTP). + cp.ctx_other = e->ctx_tgt; e->ctx_dft = llama_init_from_model(model_for_dft, cp); if (!e->ctx_dft) { From 00f1fd5e31717472d089704af3019c081e728f8f Mon Sep 17 00:00:00 2001 From: Shaw Date: Wed, 24 Jun 2026 22:04:38 -0700 Subject: [PATCH 10/17] fix(voice-diarizer): read pyannote-3 LSTM gates in ONNX IOFC order (#9460) The pyannote-segmentation-3.0 GGUF carries ONNX-exported LSTM weights, whose gate blocks are laid out IOFC (input, output, forget, cell). voice_diarizer.c's lstm_run_dir read them in PyTorch IFGO order (gi@0,gf@1,gg@2,go@3), scrambling forget/cell/output -> garbage segmentation -> spurious 3rd speaker + 100+ fragments on clean 2-speaker audio (#9460 over-detection). Reorder to IOFC (gi@0,go@1,gf@2,gg@3). Verified on Apple M-series with the local real-acoustic harness (freeman.wav + Kokoro, 2-speaker GT): IFGO -> 121 segments / 3 speakers / DER ~57%; IOFC -> 2 segments / 2 speakers / DER ~32%. Mirrors develop's landed fix a359942d30 onto the gemma4-assistant fork branch (submodule 0864259de), which still carried the pre-fix order. Co-Authored-By: Claude Opus 4.8 --- .../voice-classifiers/voice_classifier/voice_diarizer.c | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tools/omnivoice/src/voice-classifiers/voice_classifier/voice_diarizer.c b/tools/omnivoice/src/voice-classifiers/voice_classifier/voice_diarizer.c index a10c729c0..4f9540c28 100644 --- a/tools/omnivoice/src/voice-classifiers/voice_classifier/voice_diarizer.c +++ b/tools/omnivoice/src/voice-classifiers/voice_classifier/voice_diarizer.c @@ -223,11 +223,12 @@ static void lstm_run_dir(const float *x_dot_W, gate_buf[g] = acc; } - /* Apply nonlinearities. Gate order I, F, G, O. */ + /* Apply nonlinearities. Gate order I, O, F, C (ONNX IOFC layout — + * pyannote-3 segmentation GGUF, #9460). */ const float *gi = gate_buf + 0 * H; - const float *gf = gate_buf + 1 * H; - const float *gg = gate_buf + 2 * H; - const float *go = gate_buf + 3 * H; + const float *go = gate_buf + 1 * H; + const float *gf = gate_buf + 2 * H; + const float *gg = gate_buf + 3 * H; for (int j = 0; j < H; ++j) { const float i_t = sigmoidf(gi[j]); const float f_t = sigmoidf(gf[j]); From 3d3133b7fa63d6a69d3a1752c22d37455c1d2f48 Mon Sep 17 00:00:00 2001 From: Shaw Date: Thu, 25 Jun 2026 08:48:44 -0700 Subject: [PATCH 11/17] fix(kokoro): accept published GGUF tensor names + emit F16 weights (#9588) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The fused-lib loader required unprefixed dev tensor names (bert.embd.tok.weight, pred.F0_proj.weight, dec.gen.conv_post.weight) while the published elizaOS Kokoro bundles use the kokoro.* namespace and mainline llama.cpp names (kokoro.bert.layer.attn_q.weight). The loader, the shipped GGUF, and the in-tree converter all disagreed, so kokoro_init_from_file failed its weight sanity check for every available GGUF and the engine silently fell through to OmniVoice/stub. Two root causes, two fixes: 1. Tensor-name mismatch. Centralize the accepted name variants in kokoro-tensor-names.h (published + mainline + legacy dev) and look them up via require_tensor_any(). Required tensors are now a HARD load error instead of the old "non-fatal during J2 — treat absent tensors as zero" path, which produced shape-correct but acoustically degraded (noise) output and masked this bug. 2. All-F32 GGUFs load but synthesize noise. The converter now emits weight matrices / conv kernels (ndim >= 2) as F16 and keeps biases/norms F32, matching the dtype layout the fused forward pass expects. The stub emitter also writes kokoro.gen.conv_post.{weight,bias} so it passes the new required-tensor check. Adds test_kokoro_tensor_names.cpp (LLAMA_BUILD_TESTS-gated) asserting the alias picker resolves published, mainline, and legacy schemas and returns null when a tensor is genuinely absent. Closes the loader half of #9588; regenerating + republishing the bundle GGUF with this converter is the remaining ops step. --- tools/kokoro/CMakeLists.txt | 1 + tools/kokoro/convert_kokoro_pth_to_gguf.py | 15 +++- tools/kokoro/include/kokoro-tensor-names.h | 72 +++++++++++++++++ tools/kokoro/src/kokoro.cpp | 70 +++++++++++++++-- tools/kokoro/tests/CMakeLists.txt | 4 + .../kokoro/tests/test_kokoro_tensor_names.cpp | 78 +++++++++++++++++++ 6 files changed, 232 insertions(+), 8 deletions(-) create mode 100644 tools/kokoro/include/kokoro-tensor-names.h create mode 100644 tools/kokoro/tests/test_kokoro_tensor_names.cpp diff --git a/tools/kokoro/CMakeLists.txt b/tools/kokoro/CMakeLists.txt index bae822186..70619f844 100644 --- a/tools/kokoro/CMakeLists.txt +++ b/tools/kokoro/CMakeLists.txt @@ -17,6 +17,7 @@ set(KOKORO_PUBLIC_HEADERS include/kokoro-istft.h include/kokoro-phonemes.h include/kokoro-server-mount.h + include/kokoro-tensor-names.h include/kokoro-layers.h include/kokoro-predictor.h) diff --git a/tools/kokoro/convert_kokoro_pth_to_gguf.py b/tools/kokoro/convert_kokoro_pth_to_gguf.py index 34303fd77..8232c070f 100644 --- a/tools/kokoro/convert_kokoro_pth_to_gguf.py +++ b/tools/kokoro/convert_kokoro_pth_to_gguf.py @@ -95,11 +95,20 @@ def _add_tensor(writer: gguf.GGUFWriter, name: str, data: np.ndarray) -> None: - """Add a tensor as fp32 (cast from fp64/fp16 if needed; ensure c-contig).""" - if data.dtype != np.float32: + """Add tensors with the dtype layout the Kokoro forward pass expects. + + Weight matrices and convolution kernels (ndim >= 2) are emitted as F16; + biases, norms, and other vectors stay F32. All-F32 GGUFs can load but + synthesize noise in the fused runtime path. + """ + if data.dtype not in (np.float32, np.float16): data = data.astype(np.float32) if not data.flags["C_CONTIGUOUS"]: data = np.ascontiguousarray(data) + if data.ndim >= 2: + data = data.astype(np.float16) + elif data.dtype != np.float32: + data = data.astype(np.float32) writer.add_tensor(name, data) @@ -281,6 +290,8 @@ def emit_stub(out_path: str, hp: dict) -> None: _add_tensor(writer, "kokoro.predictor.F0_proj.bias", np.zeros((1,), dtype=np.float32)) _add_tensor(writer, "kokoro.predictor.N_proj.weight", rng.standard_normal((1, hid//2, 1), dtype=np.float32) * scale) _add_tensor(writer, "kokoro.predictor.N_proj.bias", np.zeros((1,), dtype=np.float32)) + _add_tensor(writer, "kokoro.gen.conv_post.weight", rng.standard_normal((22, 128, 7), dtype=np.float32) * scale) + _add_tensor(writer, "kokoro.gen.conv_post.bias", np.zeros((22,), dtype=np.float32)) writer.write_header_to_file() writer.write_kv_data_to_file() diff --git a/tools/kokoro/include/kokoro-tensor-names.h b/tools/kokoro/include/kokoro-tensor-names.h new file mode 100644 index 000000000..50de4c7d5 --- /dev/null +++ b/tools/kokoro/include/kokoro-tensor-names.h @@ -0,0 +1,72 @@ +// SPDX-License-Identifier: MIT +// +// kokoro-tensor-names.h — tensor-name compatibility for Kokoro GGUFs. +// +// The published elizaOS Kokoro bundles use the `kokoro.*` namespace emitted by +// tools/kokoro/convert_kokoro_pth_to_gguf.py. Older local/dev GGUFs used short +// unprefixed names. Keep the aliases centralized so desktop and iOS builds load +// the same schema through the same kokoro_lib code path. + +#pragma once + +namespace eliza_kokoro { + +inline constexpr const char * KOKORO_TENSOR_BERT_TOKEN_EMBD[] = { + "kokoro.bert.token_embd.weight", + "kokoro.bert.embd.tok.weight", + "bert.embd.tok.weight", + "kokoro.token_embd.weight", + nullptr, +}; + +inline constexpr const char * KOKORO_TENSOR_BERT_ATTN_Q[] = { + "kokoro.bert.layer.attn_q.weight", + "kokoro.bert.attn_q.weight", + "bert.layer.attn_q.weight", + "bert.attn_q.weight", + nullptr, +}; + +inline constexpr const char * KOKORO_TENSOR_DURATION_PROJ[] = { + "kokoro.predictor.duration_proj.weight", + "kokoro.predictor.duration.weight", + "predictor.duration_proj.weight", + "pred.duration_proj.weight", + "pred.duration.weight", + nullptr, +}; + +inline constexpr const char * KOKORO_TENSOR_F0_PROJ[] = { + "kokoro.predictor.F0_proj.weight", + "predictor.F0_proj.weight", + "pred.F0_proj.weight", + nullptr, +}; + +inline constexpr const char * KOKORO_TENSOR_N_PROJ[] = { + "kokoro.predictor.N_proj.weight", + "predictor.N_proj.weight", + "pred.N_proj.weight", + nullptr, +}; + +inline constexpr const char * KOKORO_TENSOR_GEN_CONV_POST[] = { + "kokoro.gen.conv_post.weight", + "kokoro.decoder.gen.conv_post.weight", + "decoder.gen.conv_post.weight", + "dec.gen.conv_post.weight", + nullptr, +}; + +inline const char * kokoro_pick_tensor_name( + const char * const * aliases, + bool (*has_tensor)(const char * name, void * user_data), + void * user_data) { + if (!aliases || !has_tensor) return nullptr; + for (const char * const * p = aliases; *p; ++p) { + if (has_tensor(*p, user_data)) return *p; + } + return nullptr; +} + +} // namespace eliza_kokoro diff --git a/tools/kokoro/src/kokoro.cpp b/tools/kokoro/src/kokoro.cpp index 829840231..49741d148 100644 --- a/tools/kokoro/src/kokoro.cpp +++ b/tools/kokoro/src/kokoro.cpp @@ -30,6 +30,7 @@ #include "kokoro.h" #include "kokoro-istft.h" #include "kokoro-phonemes.h" +#include "kokoro-tensor-names.h" #include "ggml.h" #include "ggml-alloc.h" @@ -164,6 +165,39 @@ static ggml_tensor * find_tensor(ggml_context * ctx, const std::string & name) { return ggml_get_tensor(ctx, name.c_str()); } +static bool has_tensor_alias(const char * name, void * user_data) { + return name && ggml_get_tensor((ggml_context *) user_data, name) != nullptr; +} + +static ggml_tensor * find_tensor_any(ggml_context * ctx, const char * const * aliases) { + const char * name = kokoro_pick_tensor_name(aliases, has_tensor_alias, ctx); + return name ? ggml_get_tensor(ctx, name) : nullptr; +} + +static std::string format_aliases(const char * const * aliases) { + std::string out; + for (const char * const * p = aliases; p && *p; ++p) { + if (!out.empty()) out += ", "; + out += "'"; + out += *p; + out += "'"; + } + return out; +} + +static ggml_tensor * require_tensor_any( + ggml_context * ctx, + const char * const * aliases, + const char * label, + std::string & err_out) { + ggml_tensor * t = find_tensor_any(ctx, aliases); + if (!t) { + err_out = std::string("required tensor missing for ") + label + + " (accepted names: " + format_aliases(aliases) + ")"; + } + return t; +} + } // namespace // --------------------------------------------------------------------------- @@ -252,14 +286,38 @@ kokoro_model_ptr kokoro_load_model( } } - // Bind canonical tensors. Missing tensors are non-fatal during the J2 - // ship phase — the synthesis path treats absent tensors as zero, which - // produces shape-correct but acoustically degraded output. See the - // J2-kokoro-port-notes.md gap log. - model->tok_embd = find_tensor(model->ctx, "kokoro.token_embd.weight"); + // Bind the published Kokoro GGUF schema, while accepting the older + // unprefixed dev names from pre-publication GGUFs. Missing required + // tensors are a hard load error: otherwise the synth path can appear to + // work while silently skipping the real model weights. + model->tok_embd = require_tensor_any( + model->ctx, + KOKORO_TENSOR_BERT_TOKEN_EMBD, + "BERT token embedding", + err_out); + if (!model->tok_embd) return {nullptr, kokoro_model_deleter{}}; + + if (!require_tensor_any(model->ctx, KOKORO_TENSOR_BERT_ATTN_Q, "BERT attention Q", err_out)) { + return {nullptr, kokoro_model_deleter{}}; + } + if (!require_tensor_any(model->ctx, KOKORO_TENSOR_F0_PROJ, "F0 projection", err_out)) { + return {nullptr, kokoro_model_deleter{}}; + } + if (!require_tensor_any(model->ctx, KOKORO_TENSOR_N_PROJ, "noise projection", err_out)) { + return {nullptr, kokoro_model_deleter{}}; + } + if (!require_tensor_any(model->ctx, KOKORO_TENSOR_GEN_CONV_POST, "generator post convolution", err_out)) { + return {nullptr, kokoro_model_deleter{}}; + } + model->mel_proj = find_tensor(model->ctx, "kokoro.decoder.mel_proj.weight"); model->phase_proj = find_tensor(model->ctx, "kokoro.decoder.phase_proj.weight"); - model->dur_proj = find_tensor(model->ctx, "kokoro.predictor.duration.weight"); + model->dur_proj = require_tensor_any( + model->ctx, + KOKORO_TENSOR_DURATION_PROJ, + "duration projection", + err_out); + if (!model->dur_proj) return {nullptr, kokoro_model_deleter{}}; model->style_proj = find_tensor(model->ctx, "kokoro.style.proj.weight"); model->out_norm = find_tensor(model->ctx, "kokoro.text.out_norm.weight"); diff --git a/tools/kokoro/tests/CMakeLists.txt b/tools/kokoro/tests/CMakeLists.txt index 81cae041c..1c97a4627 100644 --- a/tools/kokoro/tests/CMakeLists.txt +++ b/tools/kokoro/tests/CMakeLists.txt @@ -10,3 +10,7 @@ add_test(NAME test-kokoro-phonemes COMMAND test-kokoro-phonemes) add_executable(test-kokoro-istft test_kokoro_istft.cpp) target_link_libraries(test-kokoro-istft PRIVATE kokoro_lib) add_test(NAME test-kokoro-istft COMMAND test-kokoro-istft) + +add_executable(test-kokoro-tensor-names test_kokoro_tensor_names.cpp) +target_link_libraries(test-kokoro-tensor-names PRIVATE kokoro_lib) +add_test(NAME test-kokoro-tensor-names COMMAND test-kokoro-tensor-names) diff --git a/tools/kokoro/tests/test_kokoro_tensor_names.cpp b/tools/kokoro/tests/test_kokoro_tensor_names.cpp new file mode 100644 index 000000000..c510db91e --- /dev/null +++ b/tools/kokoro/tests/test_kokoro_tensor_names.cpp @@ -0,0 +1,78 @@ +// SPDX-License-Identifier: MIT +// +// test_kokoro_tensor_names.cpp — regression coverage for issue #9588. +// +// macOS and iOS both link kokoro_lib into the fused libelizainference target. +// If these aliases drift from the published GGUF schema, both platforms can +// export Kokoro symbols yet fail or silently skip the real weights at load time. + +#include "kokoro-tensor-names.h" + +#include +#include +#include +#include +#include + +namespace { + +bool has_name(const char * name, void * user_data) { + const auto * names = static_cast *>(user_data); + return names->find(name) != names->end(); +} + +void expect_pick( + const char * const * aliases, + const std::set & available, + const char * expected) { + const char * actual = eliza_kokoro::kokoro_pick_tensor_name( + aliases, + has_name, + (void *) &available); + assert(actual != nullptr); + assert(std::strcmp(actual, expected) == 0); +} + +} // namespace + +int main() { + using namespace eliza_kokoro; + + const std::set published_schema = { + "kokoro.bert.token_embd.weight", + "kokoro.bert.layer.attn_q.weight", + "kokoro.predictor.duration_proj.weight", + "kokoro.predictor.F0_proj.weight", + "kokoro.predictor.N_proj.weight", + "kokoro.gen.conv_post.weight", + }; + + expect_pick(KOKORO_TENSOR_BERT_TOKEN_EMBD, published_schema, "kokoro.bert.token_embd.weight"); + expect_pick(KOKORO_TENSOR_BERT_ATTN_Q, published_schema, "kokoro.bert.layer.attn_q.weight"); + expect_pick(KOKORO_TENSOR_DURATION_PROJ, published_schema, "kokoro.predictor.duration_proj.weight"); + expect_pick(KOKORO_TENSOR_F0_PROJ, published_schema, "kokoro.predictor.F0_proj.weight"); + expect_pick(KOKORO_TENSOR_N_PROJ, published_schema, "kokoro.predictor.N_proj.weight"); + expect_pick(KOKORO_TENSOR_GEN_CONV_POST, published_schema, "kokoro.gen.conv_post.weight"); + + const std::set legacy_schema = { + "bert.embd.tok.weight", + "bert.layer.attn_q.weight", + "pred.duration_proj.weight", + "pred.F0_proj.weight", + "pred.N_proj.weight", + "dec.gen.conv_post.weight", + }; + + expect_pick(KOKORO_TENSOR_BERT_TOKEN_EMBD, legacy_schema, "bert.embd.tok.weight"); + expect_pick(KOKORO_TENSOR_BERT_ATTN_Q, legacy_schema, "bert.layer.attn_q.weight"); + expect_pick(KOKORO_TENSOR_DURATION_PROJ, legacy_schema, "pred.duration_proj.weight"); + expect_pick(KOKORO_TENSOR_F0_PROJ, legacy_schema, "pred.F0_proj.weight"); + expect_pick(KOKORO_TENSOR_N_PROJ, legacy_schema, "pred.N_proj.weight"); + expect_pick(KOKORO_TENSOR_GEN_CONV_POST, legacy_schema, "dec.gen.conv_post.weight"); + + const std::set empty_schema; + assert(kokoro_pick_tensor_name(KOKORO_TENSOR_BERT_TOKEN_EMBD, has_name, (void *) &empty_schema) == nullptr); + + std::printf("test_kokoro_tensor_names: OK\n"); + return 0; +} From 82925c0e9b1afd2ea0d36608f734a2e0425bc612 Mon Sep 17 00:00:00 2001 From: Shaw Date: Thu, 25 Jun 2026 11:25:29 -0700 Subject: [PATCH 12/17] fix(kokoro): accept published dur_proj tensor alias --- tools/kokoro/include/kokoro-tensor-names.h | 1 + tools/kokoro/tests/test_kokoro_tensor_names.cpp | 16 ++++++++++++++++ 2 files changed, 17 insertions(+) diff --git a/tools/kokoro/include/kokoro-tensor-names.h b/tools/kokoro/include/kokoro-tensor-names.h index 50de4c7d5..f6ff699f0 100644 --- a/tools/kokoro/include/kokoro-tensor-names.h +++ b/tools/kokoro/include/kokoro-tensor-names.h @@ -32,6 +32,7 @@ inline constexpr const char * KOKORO_TENSOR_DURATION_PROJ[] = { "kokoro.predictor.duration.weight", "predictor.duration_proj.weight", "pred.duration_proj.weight", + "pred.dur_proj.weight", "pred.duration.weight", nullptr, }; diff --git a/tools/kokoro/tests/test_kokoro_tensor_names.cpp b/tools/kokoro/tests/test_kokoro_tensor_names.cpp index c510db91e..356df4265 100644 --- a/tools/kokoro/tests/test_kokoro_tensor_names.cpp +++ b/tools/kokoro/tests/test_kokoro_tensor_names.cpp @@ -70,6 +70,22 @@ int main() { expect_pick(KOKORO_TENSOR_N_PROJ, legacy_schema, "pred.N_proj.weight"); expect_pick(KOKORO_TENSOR_GEN_CONV_POST, legacy_schema, "dec.gen.conv_post.weight"); + const std::set published_legacy_schema = { + "bert.embd.tok.weight", + "bert.attn_q.weight", + "pred.dur_proj.weight", + "pred.F0_proj.weight", + "pred.N_proj.weight", + "dec.gen.conv_post.weight", + }; + + expect_pick(KOKORO_TENSOR_BERT_TOKEN_EMBD, published_legacy_schema, "bert.embd.tok.weight"); + expect_pick(KOKORO_TENSOR_BERT_ATTN_Q, published_legacy_schema, "bert.attn_q.weight"); + expect_pick(KOKORO_TENSOR_DURATION_PROJ, published_legacy_schema, "pred.dur_proj.weight"); + expect_pick(KOKORO_TENSOR_F0_PROJ, published_legacy_schema, "pred.F0_proj.weight"); + expect_pick(KOKORO_TENSOR_N_PROJ, published_legacy_schema, "pred.N_proj.weight"); + expect_pick(KOKORO_TENSOR_GEN_CONV_POST, published_legacy_schema, "dec.gen.conv_post.weight"); + const std::set empty_schema; assert(kokoro_pick_tensor_name(KOKORO_TENSOR_BERT_TOKEN_EMBD, has_name, (void *) &empty_schema) == nullptr); From ba598f562576df61b2a3faba69c4c9cf1e1ce40f Mon Sep 17 00:00:00 2001 From: Shaw Date: Thu, 25 Jun 2026 14:17:01 -0700 Subject: [PATCH 13/17] feat(kokoro): real StyleTTS-2/iSTFTNet decoder + dequant-at-load + espeak G2P (#9588) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Kokoro TTS never actually spoke: `kokoro_synthesize` ran a placeholder spectrogram (envelope x 1/f timbre x random phase) and discarded the model, and the scalar predictor read F16/quantized GGUF tensors as raw F32 (garbage), collapsing BERT to a constant -> a fixed-pitch buzz. This implements the real forward pass end-to-end, validated against the PyTorch reference and verified intelligible by whisper on the SHIPPED Q4_K_M bundle (no GGUF republish needed). Root causes fixed: - DTYPE: kokoro_load_model now dequantizes every non-F32 tensor to a parallel all-F32 context (ggml to_float traits: F16 + Q5_0/Q4_K/Q6_K). The predictor/decoder read F32 only. This alone un-breaks the (already-correct) prosody predictor — BERT now matches the reference to 3 decimals. - SYNTHESIS WAS A STUB: implement the iSTFTNet decoder. - kokoro-decoder-front.h: Decoder.forward up to the generator (F0/N conv, encode, asr_res, 4 AdainResBlk1d decode blocks). Validated corr=1.0/stage. - kokoro-generator.{h,cpp}: the Generator — harmonic source (SineGen+Source ModuleHnNSF from F0), center=True STFT, 2 ConvTranspose upsamples, 6 Snake AdaINResBlock1, noise convs/res, conv_post, center=True iSTFT. corr 0.993 vs reference; whisper transcribes the output. - kokoro-decoder.{h,cpp}: binds the model's F32 tensors into the front + generator weight structs and runs them; wired into kokoro_synthesize (predictor -> transpose asr -> decoder), variable-length T. - G2P: kokoro-phonemes now drives libespeak-ng (en-us IPA -> Kokoro vocab ids, reproducing the reference token sequence) when linked; ASCII fallback only when espeak is absent. CMake auto-detects libespeak-ng. Validation (Apple M-series): full text->speech via kokoro-tts on BOTH the all-F32 regen GGUF and the published Q4_K_M bundle, whisper-intelligible across phrases ("Hello, this is a native Kokoro voice test.", "The quick brown fox ...", "Open the pod bay doors please."). Per-component standalone tests (decoder front, Snake resblock, harmonic source+STFT, iSTFT) corr 1.0 vs the PyTorch reference. New unit test test_kokoro_g2p_espeak; kokoro-stage-dump / kokoro-decoder-test dev harnesses. The published bundle GGUF is unchanged. --- tools/kokoro/CMakeLists.txt | 33 +- tools/kokoro/convert_kokoro_pth_to_gguf.py | 12 +- tools/kokoro/include/kokoro-decoder-front.h | 144 ++++++ tools/kokoro/include/kokoro-decoder.h | 31 ++ tools/kokoro/include/kokoro-generator.h | 76 +++ tools/kokoro/include/kokoro-phonemes.h | 70 ++- tools/kokoro/include/kokoro.h | 10 +- tools/kokoro/src/kokoro-decoder.cpp | 141 ++++++ tools/kokoro/src/kokoro-generator.cpp | 435 ++++++++++++++++++ tools/kokoro/src/kokoro-phonemes.cpp | 317 ++++++++++--- tools/kokoro/src/kokoro.cpp | 313 ++++++------- tools/kokoro/tests/CMakeLists.txt | 6 + tools/kokoro/tests/test_kokoro_g2p_espeak.cpp | 87 ++++ tools/kokoro/tests/test_kokoro_phonemes.cpp | 75 ++- tools/kokoro/tools/kokoro-decoder-test.cpp | 57 +++ tools/kokoro/tools/kokoro-stage-dump.cpp | 73 +++ tools/kokoro/tools/kokoro-tts.cpp | 4 +- 17 files changed, 1588 insertions(+), 296 deletions(-) create mode 100644 tools/kokoro/include/kokoro-decoder-front.h create mode 100644 tools/kokoro/include/kokoro-decoder.h create mode 100644 tools/kokoro/include/kokoro-generator.h create mode 100644 tools/kokoro/src/kokoro-decoder.cpp create mode 100644 tools/kokoro/src/kokoro-generator.cpp create mode 100644 tools/kokoro/tests/test_kokoro_g2p_espeak.cpp create mode 100644 tools/kokoro/tools/kokoro-decoder-test.cpp create mode 100644 tools/kokoro/tools/kokoro-stage-dump.cpp diff --git a/tools/kokoro/CMakeLists.txt b/tools/kokoro/CMakeLists.txt index 70619f844..dee882dd1 100644 --- a/tools/kokoro/CMakeLists.txt +++ b/tools/kokoro/CMakeLists.txt @@ -25,7 +25,9 @@ set(KOKORO_CORE_SOURCES src/kokoro.cpp src/kokoro-istft.cpp src/kokoro-phonemes.cpp - src/kokoro-predictor.cpp) + src/kokoro-predictor.cpp + src/kokoro-generator.cpp + src/kokoro-decoder.cpp) add_library(kokoro_lib STATIC ${KOKORO_CORE_SOURCES} @@ -46,6 +48,29 @@ endif() target_compile_features(kokoro_lib PUBLIC cxx_std_17) +# Real G2P via libespeak-ng. When present, kokoro_phonemize() drives +# espeak_TextToPhonemes() (en-us IPA) and maps codepoints to Kokoro vocab ids, +# reproducing the reference token sequence. When absent, the build falls back +# to the degraded ASCII grapheme mapping and the TS layer must supply IPA. +# Override the search with -DKOKORO_ESPEAK_ROOT= (e.g. Homebrew). +option(KOKORO_ENABLE_ESPEAK "Link libespeak-ng for real Kokoro G2P" ON) +if(KOKORO_ENABLE_ESPEAK) + find_path(ESPEAK_NG_INCLUDE_DIR espeak-ng/speak_lib.h + HINTS ${KOKORO_ESPEAK_ROOT} /opt/homebrew /usr/local /usr + PATH_SUFFIXES include) + find_library(ESPEAK_NG_LIBRARY NAMES espeak-ng + HINTS ${KOKORO_ESPEAK_ROOT} /opt/homebrew /usr/local /usr + PATH_SUFFIXES lib lib64) + if(ESPEAK_NG_INCLUDE_DIR AND ESPEAK_NG_LIBRARY) + target_include_directories(kokoro_lib PRIVATE ${ESPEAK_NG_INCLUDE_DIR}) + target_link_libraries(kokoro_lib PRIVATE ${ESPEAK_NG_LIBRARY}) + target_compile_definitions(kokoro_lib PRIVATE KOKORO_USE_ESPEAK) + message(STATUS "Kokoro G2P: libespeak-ng found (${ESPEAK_NG_LIBRARY}) — real IPA path enabled") + else() + message(STATUS "Kokoro G2P: libespeak-ng not found — falling back to ASCII grapheme mapping (TS layer must supply IPA)") + endif() +endif() + # Standalone CLI harness — host-only (required by J2 verification, # tools/voice-kokoro/). Skipped on iOS: there an executable is implicitly a # MACOSX_BUNDLE app, so `install(TARGETS ... RUNTIME)` fails configure with @@ -55,6 +80,12 @@ if(NOT CMAKE_SYSTEM_NAME STREQUAL "iOS") add_executable(kokoro-tts tools/kokoro-tts.cpp) target_link_libraries(kokoro-tts PRIVATE kokoro_lib) install(TARGETS kokoro-tts RUNTIME) + + add_executable(kokoro-stage-dump tools/kokoro-stage-dump.cpp) + target_link_libraries(kokoro-stage-dump PRIVATE kokoro_lib) + + add_executable(kokoro-decoder-test tools/kokoro-decoder-test.cpp) + target_link_libraries(kokoro-decoder-test PRIVATE kokoro_lib) endif() # Server-mount handler: compiled into kokoro_lib only when the server target diff --git a/tools/kokoro/convert_kokoro_pth_to_gguf.py b/tools/kokoro/convert_kokoro_pth_to_gguf.py index 8232c070f..12bd55fff 100644 --- a/tools/kokoro/convert_kokoro_pth_to_gguf.py +++ b/tools/kokoro/convert_kokoro_pth_to_gguf.py @@ -97,9 +97,15 @@ def _add_tensor(writer: gguf.GGUFWriter, name: str, data: np.ndarray) -> None: """Add tensors with the dtype layout the Kokoro forward pass expects. - Weight matrices and convolution kernels (ndim >= 2) are emitted as F16; - biases, norms, and other vectors stay F32. All-F32 GGUFs can load but - synthesize noise in the fused runtime path. + Weight matrices and convolution kernels (ndim >= 2) are emitted as F16 + purely to halve the GGUF download size; biases, norms, and other vectors + stay F32. The GGUF dtype does not affect correctness: the loader + dequantizes every tensor to F32 at load time, so an all-F32 and an + F16-weights GGUF produce identical synthesis. (An earlier note here + claimed all-F32 GGUFs synthesized noise — that was a misdiagnosis: the + fused path was a stub that ignored the weights, and the real defect was + the loader reading non-F32 tensors as raw F32. Both are fixed; F16 is + kept only for bundle size.) """ if data.dtype not in (np.float32, np.float16): data = data.astype(np.float32) diff --git a/tools/kokoro/include/kokoro-decoder-front.h b/tools/kokoro/include/kokoro-decoder-front.h new file mode 100644 index 000000000..4a89770c5 --- /dev/null +++ b/tools/kokoro/include/kokoro-decoder-front.h @@ -0,0 +1,144 @@ +// SPDX-License-Identifier: MIT +// kokoro-decoder-front.h — Decoder.forward up to the generator (validated port, #9588). +#pragma once +#include +#include +#include +#include "kokoro-layers.h" // conv1d_forward, adain1d_forward, convtranspose1d_depthwise_forward, convtranspose1d_out_len + +namespace eliza_kokoro { + +struct DecAdainResBlk { + int Cin = 0, Cout = 0, Sdim = 128; + bool upsample = false; + bool learned_sc = false; // dim_in != dim_out + const float * norm1_fc_w = nullptr; // [2*Cin, Sdim] + const float * norm1_fc_b = nullptr; // [2*Cin] + const float * norm2_fc_w = nullptr; // [2*Cout, Sdim] + const float * norm2_fc_b = nullptr; // [2*Cout] + const float * conv1_w = nullptr; // [Cout, Cin, 3] + const float * conv1_b = nullptr; // [Cout] + const float * conv2_w = nullptr; // [Cout, Cout, 3] + const float * conv2_b = nullptr; // [Cout] + const float * conv1x1_w = nullptr; // [Cout, Cin, 1] (learned_sc only) + const float * conv1x1_b = nullptr; // [Cout] (null — conv1x1 bias=False) + const float * pool_w = nullptr; // [Cin, 1, 3] (upsample only) + const float * pool_b = nullptr; // [Cin] (upsample only) +}; + +// AdainResBlk1d (decode-block flavor: leaky_relu 0.2; pool/shortcut; +// out = (residual + shortcut)/sqrt(2)). Output y [Cout, T_out]. +inline void dec_adainresblk1d_forward( + const DecAdainResBlk & w, const float * x, int T_in, const float * s, + std::vector & y, int & T_out) { + const int Cin = w.Cin, Cout = w.Cout, Sdim = w.Sdim; + + // residual branch: norm1 -> leaky_relu(0.2) -> [pool] -> conv1 -> norm2 -> leaky_relu -> conv2 + std::vector r(x, x + (size_t)Cin * T_in); + adain1d_forward(r.data(), Cin, T_in, s, Sdim, w.norm1_fc_w, w.norm1_fc_b); + for (size_t i = 0; i < r.size(); ++i) if (r[i] < 0) r[i] *= 0.2f; + + int T_pool = T_in; + if (w.upsample) { + T_pool = convtranspose1d_out_len(T_in, 3, 2, 1, 1); + std::vector r2((size_t)Cin * T_pool); + convtranspose1d_depthwise_forward(r.data(), Cin, T_in, w.pool_w, w.pool_b, 3, 2, 1, 1, r2.data(), T_pool); + r.swap(r2); + } + std::vector r3((size_t)Cout * T_pool); + conv1d_forward(r.data(), Cin, T_pool, w.conv1_w, w.conv1_b, Cout, 3, 1, 1, 1, r3.data(), T_pool); + adain1d_forward(r3.data(), Cout, T_pool, s, Sdim, w.norm2_fc_w, w.norm2_fc_b); + for (size_t i = 0; i < r3.size(); ++i) if (r3[i] < 0) r3[i] *= 0.2f; + std::vector r4((size_t)Cout * T_pool); + conv1d_forward(r3.data(), Cout, T_pool, w.conv2_w, w.conv2_b, Cout, 3, 1, 1, 1, r4.data(), T_pool); + + // shortcut branch: [nearest-upsample x2] -> [conv1x1 if learned_sc] + T_out = T_pool; + std::vector sc; + if (w.upsample) { + const int T_up = T_in * 2; // == T_pool + std::vector up((size_t)Cin * T_up); + for (int c = 0; c < Cin; ++c) + for (int t = 0; t < T_up; ++t) + up[(size_t)c * T_up + t] = x[(size_t)c * T_in + (t / 2)]; + if (w.learned_sc) { + sc.assign((size_t)Cout * T_up, 0.0f); + conv1d_forward(up.data(), Cin, T_up, w.conv1x1_w, w.conv1x1_b, Cout, 1, 1, 0, 1, sc.data(), T_up); + } else sc.swap(up); + } else { + if (w.learned_sc) { + sc.assign((size_t)Cout * T_in, 0.0f); + conv1d_forward(x, Cin, T_in, w.conv1x1_w, w.conv1x1_b, Cout, 1, 1, 0, 1, sc.data(), T_in); + } else sc.assign(x, x + (size_t)Cin * T_in); + } + + y.assign((size_t)Cout * T_out, 0.0f); + const float rsqrt2 = 1.0f / std::sqrt(2.0f); + for (size_t i = 0; i < y.size(); ++i) y[i] = (r4[i] + sc[i]) * rsqrt2; +} + +struct DecoderFrontWeights { + const float * F0_conv_w = nullptr; // [1,1,3] + const float * F0_conv_b = nullptr; // [1] + const float * N_conv_w = nullptr; // [1,1,3] + const float * N_conv_b = nullptr; // [1] + const float * asr_res_w = nullptr; // [64,512,1] + const float * asr_res_b = nullptr; // [64] + DecAdainResBlk encode; // 514 -> 1024, learned_sc + DecAdainResBlk decode[4]; // 1090->1024 (x3), 1090->512 upsample +}; + +// Decoder.forward up to (not including) the generator. +// asr[512,T_asr] (T_asr=132), F0_curve[2*T_asr], N[2*T_asr], s[128] +// Output: x_out [512, 2*T_asr] (== generator_in_0); also returns the +// stride-2 conv outputs F0_down[T_asr], N_down[T_asr] (caller passes them, +// together with the ORIGINAL F0_curve, into the generator). +inline void decoder_front( + const DecoderFrontWeights & W, + const float * asr, int Cin_asr, int T_asr, + const float * F0_curve, const float * N_in, const float * s, + std::vector & x_out, + std::vector & F0_down, + std::vector & N_down) { + const int Tc = 2 * T_asr; // 264 + + F0_down.assign(T_asr, 0.0f); + conv1d_forward(F0_curve, 1, Tc, W.F0_conv_w, W.F0_conv_b, 1, 3, 2, 1, 1, F0_down.data(), T_asr); + N_down.assign(T_asr, 0.0f); + conv1d_forward(N_in, 1, Tc, W.N_conv_w, W.N_conv_b, 1, 3, 2, 1, 1, N_down.data(), T_asr); + + // x = cat([asr, F0, N], dim=channels) -> [514, T_asr] + std::vector xcat((size_t)(Cin_asr + 2) * T_asr); + std::memcpy(xcat.data(), asr, sizeof(float) * (size_t)Cin_asr * T_asr); + std::memcpy(xcat.data() + (size_t)Cin_asr * T_asr, F0_down.data(), sizeof(float) * T_asr); + std::memcpy(xcat.data() + (size_t)(Cin_asr + 1) * T_asr, N_down.data(), sizeof(float) * T_asr); + + std::vector x; int T_x; + dec_adainresblk1d_forward(W.encode, xcat.data(), T_asr, s, x, T_x); // encode 514->1024 + + std::vector asr_res((size_t)64 * T_asr); // asr_res Conv1d k1 512->64 + conv1d_forward(asr, Cin_asr, T_asr, W.asr_res_w, W.asr_res_b, 64, 1, 1, 0, 1, asr_res.data(), T_asr); + + bool res = true; + for (int b = 0; b < 4; ++b) { + std::vector blk_in; + if (res) { // cat([x, asr_res, F0, N]) -> 1024+64+1+1 = 1090 + const int Cx = (int)(x.size() / T_x); + const int Cin_blk = Cx + 64 + 1 + 1; + blk_in.assign((size_t)Cin_blk * T_x, 0.0f); + std::memcpy(blk_in.data(), x.data(), sizeof(float) * (size_t)Cx * T_x); + std::memcpy(blk_in.data() + (size_t)Cx * T_x, asr_res.data(), sizeof(float) * 64 * T_x); + std::memcpy(blk_in.data() + (size_t)(Cx + 64) * T_x, F0_down.data(), sizeof(float) * T_x); + std::memcpy(blk_in.data() + (size_t)(Cx + 65) * T_x, N_down.data(), sizeof(float) * T_x); + } else { + blk_in.assign(x.begin(), x.end()); + } + std::vector y; int T_y; + dec_adainresblk1d_forward(W.decode[b], blk_in.data(), T_x, s, y, T_y); + x.swap(y); T_x = T_y; + if (W.decode[b].upsample) res = false; // decode3 upsamples -> res stops + } + x_out.swap(x); // [512, 264] +} + +} // namespace eliza_kokoro \ No newline at end of file diff --git a/tools/kokoro/include/kokoro-decoder.h b/tools/kokoro/include/kokoro-decoder.h new file mode 100644 index 000000000..3de5d2562 --- /dev/null +++ b/tools/kokoro/include/kokoro-decoder.h @@ -0,0 +1,31 @@ +// SPDX-License-Identifier: MIT +// +// kokoro-decoder.h — StyleTTS-2 / iSTFTNet decoder: predictor outputs -> 24 kHz audio. +// +// Wires the validated decoder_front (kokoro-decoder-front.h) + Generator +// (kokoro-generator.h) against the model's all-F32 ggml context. Replaces the +// J2-ship placeholder spectrogram in kokoro_synthesize (#9588). + +#pragma once + +#include +#include + +namespace eliza_kokoro { + +struct kokoro_model; + +// Run the full decoder. Inputs come from kokoro_predictor_forward: +// asr_ct : [512, T_frame] channel-major (transpose of PredictorOut.asr [T,512]) +// F0, N : [2*T_frame] (PredictorOut.F0_pred / N_pred — the up-2x curves) +// ref_s_dec: [128] decoder-half style (ref_s[:128]) +// Output: audio (24 kHz mono), resized to (2*T_frame)*300. +bool kokoro_decoder_forward( + const kokoro_model * model, + const float * asr_ct, int T_frame, + const float * F0, const float * N, + const float * ref_s_dec, + std::vector & audio, + std::string & err); + +} // namespace eliza_kokoro diff --git a/tools/kokoro/include/kokoro-generator.h b/tools/kokoro/include/kokoro-generator.h new file mode 100644 index 000000000..d1d88f696 --- /dev/null +++ b/tools/kokoro/include/kokoro-generator.h @@ -0,0 +1,76 @@ +// SPDX-License-Identifier: MIT +// +// kokoro-generator.h — iSTFTNet Generator.forward (StyleTTS-2 decoder back-end). +// +// The generator turns the decoder body output `x` [512, 264], the style +// vector `s` [128], and the (un-downsampled) F0 curve `f0_curve` [264] into +// `audio` [79200] (24 kHz). +// +// Weights are raw float pointers (PyTorch row-major, weight_norm-fused): +// Conv1d weight [Cout, Cin, K] +// ConvTranspose1d wt [Cin, Cout, K] +// Linear weight [out, in] +// AdaIN1d fc.weight [2C, style_dim] (style_dim = 128) +// Snake alpha [C] +// The caller supplies them via GeneratorWeights so the function composes with +// any weight-loading boundary (GGUF tensor lookup, raw .f32 fixtures, …). + +#pragma once + +#include + +namespace eliza_kokoro { + +// One AdaINResBlock1 sub-block (the block has three, sharing the same channel +// count). convs use [Cout=Cin=C, Cin=C, K]. +struct GenSubBlockWeights { + const float * conv1_w = nullptr; // [C, C, K] + const float * conv1_b = nullptr; // [C] + const float * conv2_w = nullptr; // [C, C, K] + const float * conv2_b = nullptr; // [C] + const float * adain1_fc_w = nullptr; // [2C, 128] + const float * adain1_fc_b = nullptr; // [2C] + const float * adain2_fc_w = nullptr; // [2C, 128] + const float * adain2_fc_b = nullptr; // [2C] + const float * alpha1 = nullptr; // [C] + const float * alpha2 = nullptr; // [C] +}; + +struct GenAdaResBlockWeights { + GenSubBlockWeights sub[3]; +}; + +struct GeneratorWeights { + // m_source.l_linear: Linear(9 -> 1). + const float * l_linear_w = nullptr; // [1, 9] + const float * l_linear_b = nullptr; // [1] + + // ups[0], ups[1]: ConvTranspose1d. weight [Cin, Cout, K], bias [Cout]. + const float * ups_w[2] = { nullptr, nullptr }; + const float * ups_b[2] = { nullptr, nullptr }; + + // noise_convs[0], noise_convs[1]: Conv1d. weight [Cout, 22, K], bias [Cout]. + const float * noise_convs_w[2] = { nullptr, nullptr }; + const float * noise_convs_b[2] = { nullptr, nullptr }; + + // noise_res[0] (k=7), noise_res[1] (k=11): AdaINResBlock1. + GenAdaResBlockWeights noise_res[2]; + + // resblocks[0..5]: AdaINResBlock1 (stage0: k=3,7,11 ch=256; stage1: ch=128). + GenAdaResBlockWeights resblocks[6]; + + // conv_post: Conv1d(128 -> 22, k=7, pad=3). weight [22, 128, 7], bias [22]. + const float * conv_post_w = nullptr; + const float * conv_post_b = nullptr; +}; + +// Generator.forward. audio is resized to T0 * 300 (== 79200 for T0=264). +void kokoro_generator_forward( + const float * x_in, // [512, T0] channel-major + int T0, // input time (== 2 * predictor T_frame) + const float * s, // [128] + const float * f0_curve, // [T0] + const GeneratorWeights & w, + std::vector & audio); + +} // namespace eliza_kokoro diff --git a/tools/kokoro/include/kokoro-phonemes.h b/tools/kokoro/include/kokoro-phonemes.h index cc9bb3fd6..068dc0f79 100644 --- a/tools/kokoro/include/kokoro-phonemes.h +++ b/tools/kokoro/include/kokoro-phonemes.h @@ -1,22 +1,26 @@ // SPDX-License-Identifier: MIT // -// kokoro-phonemes.h — minimal ASCII text → Kokoro phoneme-id mapping. +// kokoro-phonemes.h — text → Kokoro phoneme-id mapping. // -// Kokoro v1.0 uses espeak-ng's phoneme inventory + a small set of control -// tokens (BOS, EOS, PAD, blanks). The training-time path passes text through -// `phonemize` (Python wrapper around espeak-ng) before tokenizing. +// Kokoro v1.0 tokenizes espeak-ng IPA against a small fixed vocab +// (`tts/kokoro/tokenizer.json`, `model.vocab`). The reference Python path is: // -// Adding an espeak-ng dependency to the fork is overkill for a TTS that -// is being ported as a one-release deprecation runway. This header -// implements a deterministic grapheme→phoneme mapping that: +// text → espeak-ng (en-us, --ipa) → IPA string → per-codepoint vocab lookup +// → ids → model input_ids = [0, *ids, 0] (0 = the "$" pad symbol) // -// 1. covers the basic Latin alphabet + common digraphs (sh, ch, th, ng); -// 2. maps every other ASCII printable to PAD; -// 3. returns ids in the same value range as kokoro-onnx's tokenizer -// (PAD=0, BOS=1, EOS=2, then phonemes from offset 3). +// Every vocab key is a single Unicode codepoint, so the mapping is a pure +// codepoint→id table lookup over the IPA string (no multi-char digraph +// handling is needed — espeak already emits the canonical IPA codepoints, +// e.g. eɪ is two codepoints 'e'+'ɪ', each with its own id). // -// The synthesis quality this produces is noticeably worse than the -// espeak-ng path — that is the documented gap in J2-kokoro-port-notes.md. +// Two build modes: +// * KOKORO_USE_ESPEAK (default when libespeak-ng is linked) — the real G2P +// path: `phonemize_ipa()` drives espeak_TextToPhonemes() to get IPA, then +// maps to ids. This reproduces the kokoro reference ids exactly. +// * fallback — when espeak is unavailable the caller may pass pre-computed +// IPA from the TS layer (which already runs espeak) into +// `ipa_to_token_ids()`. `phonemize_ipa()` then returns an empty vector and +// the caller must supply IPA. #pragma once @@ -26,12 +30,42 @@ namespace eliza_kokoro { -// Tokenize a UTF-8 / ASCII text string into a phoneme-id vector. -// Always returns a sequence of length <= 510 (the BERT encoder cap in -// Kokoro v1.0 — anything longer is split at the caller). -std::vector phonemize_ascii(const std::string & text); +// Kokoro pad/boundary token. model.vocab maps '$' → 0; the reference wraps the +// phoneme ids as [PAD, *ids, PAD] to form the model input_ids. +inline constexpr int32_t KOKORO_PAD_ID = 0; + +// Map a single Unicode codepoint (an espeak IPA symbol) to its Kokoro vocab id. +// Returns -1 if the codepoint is not in the vocab (caller drops it, matching +// the reference which silently skips unmapped codepoints). +int32_t kokoro_codepoint_to_id(char32_t cp) noexcept; + +// Map an espeak-ng IPA string (UTF-8) to the bare Kokoro phoneme-id sequence +// (no pad wrapping). Codepoints absent from the vocab are dropped. This is the +// `ids` array in reference-ids.json — its length is the style-row index. +std::vector ipa_to_token_ids(const std::string & ipa); + +// Phonemize text to bare Kokoro phoneme ids via espeak-ng (en-us IPA). +// Returns the same sequence as `ipa_to_token_ids(espeak_ipa(text))`. +// When KOKORO_USE_ESPEAK is not compiled in, returns an empty vector — the +// caller must supply IPA from the TS layer and call `ipa_to_token_ids()`. +std::vector phonemize_ipa(const std::string & text); + +// Wrap a bare phoneme-id sequence as the model input_ids: [PAD, *ids, PAD]. +std::vector wrap_input_ids(const std::vector & ids); -// Diagnostic — total phoneme vocab size (for hparams cross-check). +// Convenience: text → model input_ids [PAD, *ipa_ids, PAD] via espeak. +// Equivalent to `wrap_input_ids(phonemize_ipa(text))`. +std::vector phonemize_to_input_ids(const std::string & text); + +// True when this build links libespeak-ng (the real G2P path is available). +bool espeak_available() noexcept; + +// Total Kokoro vocab size (highest id + 1 = 178 for v1.0). int phoneme_vocab_size() noexcept; +// --- Legacy ASCII fallback (retained only for callers not yet migrated) --- +// Deprecated: returns the degraded ASCII grapheme mapping. New code uses +// phonemize_to_input_ids(). +std::vector phonemize_ascii(const std::string & text); + } // namespace eliza_kokoro diff --git a/tools/kokoro/include/kokoro.h b/tools/kokoro/include/kokoro.h index 809362962..b218677cf 100644 --- a/tools/kokoro/include/kokoro.h +++ b/tools/kokoro/include/kokoro.h @@ -116,10 +116,12 @@ kokoro_status kokoro_load_voice_preset( kokoro_voice_preset & out, std::string & err_out) noexcept; -// Phonemize an input text into Kokoro's int phoneme ids. The implementation -// uses a deterministic ASCII grapheme→phoneme mapping (no espeak-ng -// dependency). This is intentionally lossy vs the upstream phonemizer — -// quality recovery is part of the gap documented in J2-kokoro-port-notes.md. +// Phonemize an input text into Kokoro's int phoneme ids (the model input_ids, +// wrapped as [PAD, *ids, PAD]). When the build links libespeak-ng this is the +// real G2P path (text → en-us IPA → Kokoro vocab ids), reproducing the +// upstream phonemizer's token sequence. Without libespeak-ng it falls back to +// a deterministic (lossy) ASCII grapheme mapping; in that case the TS voice +// layer should phonemize and pass IPA (see kokoro-phonemes.h ipa_to_token_ids). std::vector kokoro_phonemize(const std::string & text); // Synthesize a single utterance. `text` is the natural-language input, diff --git a/tools/kokoro/src/kokoro-decoder.cpp b/tools/kokoro/src/kokoro-decoder.cpp new file mode 100644 index 000000000..388f124b9 --- /dev/null +++ b/tools/kokoro/src/kokoro-decoder.cpp @@ -0,0 +1,141 @@ +// SPDX-License-Identifier: MIT +// +// kokoro-decoder.cpp — assemble decoder_front + Generator into the full +// StyleTTS-2 / iSTFTNet decoder, reading weights from the model's all-F32 +// ggml context (dequantized at load). Validated against the PyTorch reference +// stage-by-stage (#9588). + +#include "kokoro-decoder.h" +#include "kokoro-decoder-front.h" // DecoderFrontWeights, DecAdainResBlk, decoder_front +#include "kokoro-generator.h" // GeneratorWeights, kokoro_generator_forward + +#include "ggml.h" + +#include + +namespace eliza_kokoro { + +// Defined in kokoro.cpp — the all-F32 working context the predictor reads. +ggml_context * kokoro_model_ggml_ctx(const kokoro_model * model); + +namespace { + +struct Lk { + ggml_context * ctx; + const float * get(const std::string & name) const { + ggml_tensor * t = ggml_get_tensor(ctx, name.c_str()); + return t ? (const float *) t->data : nullptr; + } +}; + +// Fill an AdainResBlk1d (decode flavor) from a tensor-name prefix. +void fill_dec_block(const Lk & L, DecAdainResBlk & b, const std::string & pfx, + int Cin, int Cout, bool upsample) { + b.Cin = Cin; b.Cout = Cout; b.Sdim = 128; + b.upsample = upsample; + b.learned_sc = (Cin != Cout); + b.norm1_fc_w = L.get(pfx + ".norm1.fc.weight"); + b.norm1_fc_b = L.get(pfx + ".norm1.fc.bias"); + b.norm2_fc_w = L.get(pfx + ".norm2.fc.weight"); + b.norm2_fc_b = L.get(pfx + ".norm2.fc.bias"); + b.conv1_w = L.get(pfx + ".conv1.weight"); + b.conv1_b = L.get(pfx + ".conv1.bias"); + b.conv2_w = L.get(pfx + ".conv2.weight"); + b.conv2_b = L.get(pfx + ".conv2.bias"); + b.conv1x1_w = b.learned_sc ? L.get(pfx + ".conv1x1.weight") : nullptr; + b.conv1x1_b = nullptr; // conv1x1 bias=False upstream + b.pool_w = upsample ? L.get(pfx + ".pool.weight") : nullptr; + b.pool_b = upsample ? L.get(pfx + ".pool.bias") : nullptr; +} + +// Fill an AdaINResBlock1 (generator flavor: 3 sub-blocks, Snake1D) from a prefix. +void fill_gen_block(const Lk & L, GenAdaResBlockWeights & b, const std::string & pfx) { + for (int j = 0; j < 3; ++j) { + const std::string js = std::to_string(j); + GenSubBlockWeights & s = b.sub[j]; + s.conv1_w = L.get(pfx + ".convs1." + js + ".weight"); + s.conv1_b = L.get(pfx + ".convs1." + js + ".bias"); + s.conv2_w = L.get(pfx + ".convs2." + js + ".weight"); + s.conv2_b = L.get(pfx + ".convs2." + js + ".bias"); + s.adain1_fc_w = L.get(pfx + ".adain1." + js + ".fc.weight"); + s.adain1_fc_b = L.get(pfx + ".adain1." + js + ".fc.bias"); + s.adain2_fc_w = L.get(pfx + ".adain2." + js + ".fc.weight"); + s.adain2_fc_b = L.get(pfx + ".adain2." + js + ".fc.bias"); + s.alpha1 = L.get(pfx + ".alpha1." + js); + s.alpha2 = L.get(pfx + ".alpha2." + js); + } +} + +} // namespace + +bool kokoro_decoder_forward( + const kokoro_model * model, + const float * asr_ct, int T_frame, + const float * F0, const float * N, + const float * ref_s_dec, + std::vector & audio, + std::string & err) { + audio.clear(); + if (!model) { err = "null model"; return false; } + if (T_frame <= 0) { err = "non-positive T_frame"; return false; } + + ggml_context * ctx = kokoro_model_ggml_ctx(model); + if (!ctx) { err = "null model context"; return false; } + Lk L{ctx}; + + // --- decoder_front weights --- + DecoderFrontWeights W; + W.F0_conv_w = L.get("kokoro.decoder.F0_conv.weight"); + W.F0_conv_b = L.get("kokoro.decoder.F0_conv.bias"); + W.N_conv_w = L.get("kokoro.decoder.N_conv.weight"); + W.N_conv_b = L.get("kokoro.decoder.N_conv.bias"); + W.asr_res_w = L.get("kokoro.decoder.asr_res.weight"); + W.asr_res_b = L.get("kokoro.decoder.asr_res.bias"); + fill_dec_block(L, W.encode, "kokoro.decoder.encode", 514, 1024, /*upsample*/false); + fill_dec_block(L, W.decode[0], "kokoro.decoder.decode.0", 1090, 1024, false); + fill_dec_block(L, W.decode[1], "kokoro.decoder.decode.1", 1090, 1024, false); + fill_dec_block(L, W.decode[2], "kokoro.decoder.decode.2", 1090, 1024, false); + fill_dec_block(L, W.decode[3], "kokoro.decoder.decode.3", 1090, 512, /*upsample*/true); + + if (!W.F0_conv_w || !W.asr_res_w || !W.encode.conv1_w || !W.decode[3].pool_w) { + err = "missing decoder weights (is the GGUF a full Kokoro model?)"; + return false; + } + + // --- generator weights --- + GeneratorWeights G; + G.l_linear_w = L.get("kokoro.gen.m_source.l_linear.weight"); + G.l_linear_b = L.get("kokoro.gen.m_source.l_linear.bias"); + for (int i = 0; i < 2; ++i) { + const std::string is = std::to_string(i); + G.ups_w[i] = L.get("kokoro.gen.ups." + is + ".weight"); + G.ups_b[i] = L.get("kokoro.gen.ups." + is + ".bias"); + G.noise_convs_w[i] = L.get("kokoro.gen.noise_convs." + is + ".weight"); + G.noise_convs_b[i] = L.get("kokoro.gen.noise_convs." + is + ".bias"); + fill_gen_block(L, G.noise_res[i], "kokoro.gen.noise_res." + is); + } + for (int i = 0; i < 6; ++i) { + fill_gen_block(L, G.resblocks[i], "kokoro.gen.resblocks." + std::to_string(i)); + } + G.conv_post_w = L.get("kokoro.gen.conv_post.weight"); + G.conv_post_b = L.get("kokoro.gen.conv_post.bias"); + + if (!G.l_linear_w || !G.ups_w[0] || !G.conv_post_w || !G.resblocks[5].sub[2].conv2_w) { + err = "missing generator weights (is the GGUF a full Kokoro model?)"; + return false; + } + + // --- run: decoder_front -> generator --- + std::vector x, F0_down, N_down; + decoder_front(W, asr_ct, /*Cin_asr*/512, T_frame, F0, N, ref_s_dec, x, F0_down, N_down); + + const int T0 = 2 * T_frame; // decoder_front upsamples T_frame -> 2*T_frame + if ((int) (x.size() / 512) != T0) { + err = "decoder_front output width mismatch"; + return false; + } + kokoro_generator_forward(x.data(), T0, ref_s_dec, F0, G, audio); + return true; +} + +} // namespace eliza_kokoro diff --git a/tools/kokoro/src/kokoro-generator.cpp b/tools/kokoro/src/kokoro-generator.cpp new file mode 100644 index 000000000..6a4423b16 --- /dev/null +++ b/tools/kokoro/src/kokoro-generator.cpp @@ -0,0 +1,435 @@ +// SPDX-License-Identifier: MIT +// +// kokoro-generator.cpp — iSTFTNet Generator.forward, CPU scalar. +// +// Direct port of kokoro/istftnet.py `Generator.forward` (StyleTTS-2 + +// iSTFTNet decoder back-end). Implements: +// - SourceModuleHnNSF / SineGen harmonic source (deterministic: zero +// initial phase noise, zero additive noise — see note below), +// - forward STFT (center=True) of the harmonic source, +// - the two upsample stages (ConvTranspose ups[i] + noise_convs[i] + +// noise_res[i] (AdaINResBlock1) + 3 resblocks[i*3+j] (AdaINResBlock1, +// Snake1D activation), with leaky_relu(0.1) and a reflection_pad(1,0) +// on the final stage), +// - conv_post (Conv1d 128->22, k7, pad3), +// - spec = exp(x[:11]), phase = sin(x[11:22]), +// - inverse STFT (center=True) -> audio. +// +// Config (kokoro v1.0 istftnet): +// style_dim=128, upsample_initial_channel=512, +// upsample_rates=[10,6], upsample_kernel_sizes=[20,12], +// resblock_kernel_sizes=[3,7,11], resblock_dilation_sizes=[[1,3,5]]x3, +// gen_istft_n_fft=20, gen_istft_hop_size=5, +// m_source: harmonic_num=8 (dim=9), upsample_scale=300, voiced_threshold=10. +// +// All conv/linear weights are PyTorch row-major and weight_norm-fused: +// Conv1d weight [Cout, Cin, K] +// ConvTranspose1d wt [Cin, Cout, K] +// Linear weight [out, in] +// AdaIN1d fc.weight [2C, style_dim] +// alpha (Snake) [C] +// +// Determinism note: PyTorch's SineGen seeds an initial random phase +// (`rand_ini`, zeroed for the fundamental) and SourceModuleHnNSF adds +// Gaussian noise. For a reproducible / deterministic on-device renderer we +// drop both (zero phase noise on every harmonic, zero additive noise), so the +// harmonic source and the final audio will not bit-match a seeded PyTorch run +// — they are validated by correlation + structure instead. Every other stage +// (convs, resblocks, ups, conv_post, STFT/iSTFT) is exactly reproduced. +// +// Validated against per-stage reference activations regenerated from the +// real hexgrad/Kokoro-82M v1.0 decoder weights — see kokoro-generator.h and +// the integration notes returned with this work. + +#include "kokoro-generator.h" +#include "kokoro-layers.h" + +#include +#include +#include +#include +#include + +#ifdef KOKORO_GEN_DEBUG +#include +static void dbg_dump(const char * name, const float * d, size_t n) { + std::ofstream f(std::string("/tmp/gendbg_") + name + ".f32", std::ios::binary); + f.write((const char *) d, n * 4); +} +#define DBG(name, d, n) dbg_dump(name, d, n) +#else +#define DBG(name, d, n) do {} while (0) +#endif + +namespace eliza_kokoro { + +namespace { + +static constexpr double K_PI = 3.14159265358979323846; + +// ---------------------------------------------------------------------------- +// Periodic Hann window of length N: w[i] = 0.5 - 0.5*cos(2*pi*i/N). +// Matches scipy.get_window('hann', N, fftbins=True) / torch.hann_window(N, +// periodic=True). +// ---------------------------------------------------------------------------- +static std::vector hann_periodic(int n) { + std::vector w((size_t) n); + const double scale = 2.0 * K_PI / (double) n; + for (int i = 0; i < n; ++i) { + w[(size_t) i] = (float) (0.5 - 0.5 * std::cos(scale * (double) i)); + } + return w; +} + +// ---------------------------------------------------------------------------- +// PyTorch F.interpolate(mode='linear', align_corners=False), 1D. +// in: [C, T_in] (channel-major) -> out: [C, T_out]. +// Matches the half-pixel coordinate transform PyTorch uses by default. +// ---------------------------------------------------------------------------- +static void interp_linear(const float * x, int C, int T_in, int T_out, float * y) { + const double scale = (double) T_in / (double) T_out; + for (int o = 0; o < T_out; ++o) { + // src = (o + 0.5) * scale - 0.5 (align_corners=False half-pixel) + double src = ((double) o + 0.5) * scale - 0.5; + if (src < 0.0) src = 0.0; + int s0 = (int) std::floor(src); + int s1 = s0 + 1; + double frac = src - (double) s0; + if (s1 > T_in - 1) { s1 = T_in - 1; } + if (s0 > T_in - 1) { s0 = T_in - 1; } + for (int c = 0; c < C; ++c) { + const float a = x[(size_t) c * T_in + s0]; + const float b = x[(size_t) c * T_in + s1]; + y[(size_t) c * T_out + o] = (float) (a + (b - a) * frac); + } + } +} + +// ---------------------------------------------------------------------------- +// Forward STFT, center=True, return polar (magnitude, phase). +// input: signal[N]; n_fft, hop, win (win==n_fft here). Reflect-pad n_fft/2 +// each side (PyTorch center=True default pad_mode='reflect'). Periodic Hann +// window. Output: mag[F, n_frames], phase[F, n_frames], F=n_fft/2+1, +// n_frames = N/hop + 1. +// ---------------------------------------------------------------------------- +static void stft_center(const float * sig, int N, int n_fft, int hop, int win, + std::vector & mag, std::vector & phase, + int & F, int & n_frames) { + F = n_fft / 2 + 1; + const int pad = n_fft / 2; + const int Np = N + 2 * pad; + std::vector padded((size_t) Np); + // reflect pad: padded[pad + i] = sig[i]; left/right reflect (no edge dup). + for (int i = 0; i < N; ++i) padded[(size_t) (pad + i)] = sig[i]; + for (int i = 0; i < pad; ++i) { + padded[(size_t) (pad - 1 - i)] = sig[std::min(i + 1, N - 1)]; + padded[(size_t) (pad + N + i)] = sig[std::max(N - 2 - i, 0)]; + } + n_frames = N / hop + 1; + const std::vector window = hann_periodic(win); + mag.assign((size_t) F * n_frames, 0.0f); + phase.assign((size_t) F * n_frames, 0.0f); + for (int t = 0; t < n_frames; ++t) { + const int off = t * hop; + for (int f = 0; f < F; ++f) { + double re = 0.0, im = 0.0; + const double w0 = -2.0 * K_PI * (double) f / (double) n_fft; + for (int k = 0; k < win; ++k) { + const double v = (double) padded[(size_t) (off + k)] * (double) window[(size_t) k]; + const double ang = w0 * (double) k; + re += v * std::cos(ang); + im += v * std::sin(ang); + } + mag[(size_t) f * n_frames + t] = (float) std::sqrt(re * re + im * im); + phase[(size_t) f * n_frames + t] = (float) std::atan2(im, re); + } + } +} + +// ---------------------------------------------------------------------------- +// Inverse STFT, center=True. magnitude/phase: [F, n_frames] (polar). Periodic +// Hann window, win==n_fft. Overlap-add with window^2 normalization, then trim +// n_fft/2 from each end (the center=True crop). Output length = (n_frames-1)*hop. +// ---------------------------------------------------------------------------- +static void istft_center(const float * mag, const float * phase, + int n_fft, int hop, int win, int n_frames, + std::vector & out) { + const int F = n_fft / 2 + 1; + const int pad = n_fft / 2; + const int n_full = (n_frames - 1) * hop + win; // before crop + std::vector acc((size_t) n_full, 0.0); + std::vector wsum((size_t) n_full, 0.0); + const std::vector window = hann_periodic(win); + + std::vector frame((size_t) n_fft); + for (int t = 0; t < n_frames; ++t) { + // irfft of the hermitian spectrum -> n_fft real samples. + for (int n = 0; n < n_fft; ++n) { + double s = 0.0; + for (int f = 0; f < F; ++f) { + const double m = mag[(size_t) f * n_frames + t]; + const double p = phase[(size_t) f * n_frames + t]; + const double re = m * std::cos(p); + const double im = m * std::sin(p); + const double ang = 2.0 * K_PI * (double) f * (double) n / (double) n_fft; + double term = re * std::cos(ang) - im * std::sin(ang); + if (f != 0 && !(n_fft % 2 == 0 && f == F - 1)) term *= 2.0; + s += term; + } + frame[(size_t) n] = s / (double) n_fft; + } + const int off = t * hop; + for (int k = 0; k < win; ++k) { + const double w = (double) window[(size_t) k]; + acc[(size_t) (off + k)] += frame[(size_t) k] * w; + wsum[(size_t) (off + k)] += w * w; + } + } + for (int i = 0; i < n_full; ++i) { + if (wsum[(size_t) i] > 1e-11) acc[(size_t) i] /= wsum[(size_t) i]; + } + const int n_out = n_full - 2 * pad; + out.assign((size_t) n_out, 0.0f); + for (int i = 0; i < n_out; ++i) out[(size_t) i] = (float) acc[(size_t) (i + pad)]; +} + +// ---------------------------------------------------------------------------- +// AdaINResBlock1.forward (Snake activation). x[C,T] in place. +// Three sub-blocks: adain1 -> snake(alpha1) -> conv1[dilated] -> +// adain2 -> snake(alpha2) -> conv2[dilation 1] -> residual add. +// convs1 dilations = dil[0..2]; convs2 dilation = 1. +// padding = get_padding(K, d) = (K*d - d)/2. +// ---------------------------------------------------------------------------- +static int get_padding(int k, int d) { return (k * d - d) / 2; } + +static void adain_resblock1(const GenAdaResBlockWeights & w, + float * x, int C, int T, const float * s, int Sdim, + int K, const int dil[3]) { + std::vector xt((size_t) C * T), y1((size_t) C * T), y2((size_t) C * T); + for (int i = 0; i < 3; ++i) { + const GenSubBlockWeights & sb = w.sub[i]; + std::memcpy(xt.data(), x, sizeof(float) * (size_t) C * T); + adain1d_forward(xt.data(), C, T, s, Sdim, sb.adain1_fc_w, sb.adain1_fc_b); + snake1d_forward(xt.data(), C, T, sb.alpha1); + const int d1 = dil[i]; + conv1d_forward(xt.data(), C, T, sb.conv1_w, sb.conv1_b, C, K, 1, + get_padding(K, d1), d1, y1.data(), T); + adain1d_forward(y1.data(), C, T, s, Sdim, sb.adain2_fc_w, sb.adain2_fc_b); + snake1d_forward(y1.data(), C, T, sb.alpha2); + conv1d_forward(y1.data(), C, T, sb.conv2_w, sb.conv2_b, C, K, 1, + get_padding(K, 1), 1, y2.data(), T); + for (size_t j = 0; j < (size_t) C * T; ++j) x[j] = y2[j] + x[j]; + } +} + +} // namespace + +// ============================================================================ +// kokoro_generator_forward +// ============================================================================ +void kokoro_generator_forward( + const float * x_in, // [512, T0] + int T0, // input time (== 2 * predictor T_frame) + const float * s, // [128] + const float * f0_curve, // [T0] + const GeneratorWeights & w, + std::vector & audio /* out [T0 * 300] */) { + + const int Sdim = 128; + const int up0 = 10, up1 = 6; // upsample rates + const int hop = 5; + const int n_fft = 20, win = 20; + const int upsample_scale = up0 * up1 * hop; // 300 + const int dim = 9; // harmonic_num (8) + fundamental + const float sine_amp = 0.1f; + const float voiced_threshold = 10.0f; + const float sr = 24000.0f; + + // ------------------------------------------------------------------ + // 1. Harmonic source. f0_up = nearest-upsample(f0_curve) x300. + // ------------------------------------------------------------------ + const int L = T0 * upsample_scale; // 79200 + std::vector f0_up((size_t) L); + for (int t = 0; t < T0; ++t) + for (int r = 0; r < upsample_scale; ++r) + f0_up[(size_t) t * upsample_scale + r] = f0_curve[t]; + + // fn = f0 * [1..9]; rad = (fn / sr) % 1 ; [dim, L] channel-major. + std::vector rad((size_t) dim * L); + for (int h = 0; h < dim; ++h) { + const float mul = (float) (h + 1); + for (int t = 0; t < L; ++t) { + float v = f0_up[(size_t) t] * mul / sr; + v = v - std::floor(v); // % 1 + rad[(size_t) h * L + t] = v; + } + } + // Deterministic: rad[:,0,:] += 0 (no random initial phase). + // _f02sine: downsample rad by 1/upsample_scale (linear), cumsum*2pi, + // upsample by *upsample_scale (with phase scaled by upsample_scale), sin. + const int Lds = L / upsample_scale; // 264 + std::vector rad_ds((size_t) dim * Lds); + interp_linear(rad.data(), dim, L, Lds, rad_ds.data()); + // cumsum over time per channel, * 2*pi. + std::vector phase_ds((size_t) dim * Lds); + for (int h = 0; h < dim; ++h) { + double cum = 0.0; + for (int t = 0; t < Lds; ++t) { + cum += (double) rad_ds[(size_t) h * Lds + t]; + phase_ds[(size_t) h * Lds + t] = (float) (cum * 2.0 * K_PI); + } + } + // F.interpolate(phase * upsample_scale, scale=upsample_scale, linear). + std::vector phase_scaled((size_t) dim * Lds); + for (size_t i = 0; i < (size_t) dim * Lds; ++i) + phase_scaled[i] = phase_ds[i] * (float) upsample_scale; + std::vector phase_up((size_t) dim * L); + interp_linear(phase_scaled.data(), dim, Lds, L, phase_up.data()); + // sines = sin(phase) * sine_amp ; voiced mask uv = (f0 > thr). + // sine_waves = sines * uv (+ noise = 0). sine_merge = tanh(l_linear(sine_waves)). + std::vector sine_waves((size_t) dim * L); + for (int h = 0; h < dim; ++h) { + for (int t = 0; t < L; ++t) { + const float uv = (f0_up[(size_t) t] > voiced_threshold) ? 1.0f : 0.0f; + sine_waves[(size_t) h * L + t] = + std::sin(phase_up[(size_t) h * L + t]) * sine_amp * uv; + } + } + // l_linear: Linear(9 -> 1). weight [1, 9], bias [1]. har_source[L]. + std::vector har_source((size_t) L); + for (int t = 0; t < L; ++t) { + double acc = w.l_linear_b[0]; + for (int h = 0; h < dim; ++h) + acc += (double) w.l_linear_w[h] * (double) sine_waves[(size_t) h * L + t]; + har_source[(size_t) t] = std::tanh((float) acc); + } + DBG("har_source", har_source.data(), har_source.size()); + + // ------------------------------------------------------------------ + // 2. STFT of har_source -> har = cat(mag, phase) [22, n_frames]. + // ------------------------------------------------------------------ + std::vector hmag, hphase; + int F = 0, n_frames = 0; + stft_center(har_source.data(), L, n_fft, hop, win, hmag, hphase, F, n_frames); + const int Hc = 2 * F; // 22 + std::vector har((size_t) Hc * n_frames); + std::memcpy(har.data(), hmag.data(), sizeof(float) * (size_t) F * n_frames); + std::memcpy(har.data() + (size_t) F * n_frames, hphase.data(), + sizeof(float) * (size_t) F * n_frames); + DBG("har", har.data(), har.size()); +#ifdef KOKORO_GEN_INJECT_HAR + { + std::ifstream hf(KOKORO_GEN_INJECT_HAR, std::ios::binary); + if (hf) { hf.read((char *) har.data(), sizeof(float) * har.size()); } + } +#endif + + // ------------------------------------------------------------------ + // 3. Upsample stages. + // ------------------------------------------------------------------ + const int ups_rate[2] = { up0, up1 }; + const int ups_k[2] = { 20, 12 }; + const int ch_after[2] = { 256, 128 }; // upsample_initial_channel >> (i+1) + const int dil135[3] = { 1, 3, 5 }; + + // current x: [512, 264] (copy, mutable). + int curC = 512, curT = T0; + std::vector x(x_in, x_in + (size_t) curC * curT); + + for (int i = 0; i < 2; ++i) { + const int u = ups_rate[i], k = ups_k[i]; + const int outC = ch_after[i]; + + // leaky_relu(x, 0.1) (in place). + leaky_relu(x.data(), curC * curT, 0.1f); + + // x_source = noise_convs[i](har); then noise_res[i](x_source, s). + // noise_convs[0]: Conv1d(22->256, k=stride_f0*2=12, stride=6, pad=(6+1)/2=3). + // noise_convs[1]: Conv1d(22->128, k=1, stride=1, pad=0). + std::vector x_source; + int xsT = 0; + if (i == 0) { + const int stride_f0 = up1; // prod(upsample_rates[1:]) = 6 + const int nk = stride_f0 * 2; // 12 + const int npad = (stride_f0 + 1) / 2; // 3 + xsT = (n_frames + 2 * npad - (nk - 1) - 1) / stride_f0 + 1; + x_source.assign((size_t) outC * xsT, 0.0f); + conv1d_forward(har.data(), Hc, n_frames, w.noise_convs_w[i], w.noise_convs_b[i], + outC, nk, stride_f0, npad, 1, x_source.data(), xsT); + } else { + xsT = n_frames; // k=1 stride=1 pad=0 + x_source.assign((size_t) outC * xsT, 0.0f); + conv1d_forward(har.data(), Hc, n_frames, w.noise_convs_w[i], w.noise_convs_b[i], + outC, 1, 1, 0, 1, x_source.data(), xsT); + } + if (i == 0) DBG("noise_convs0", x_source.data(), x_source.size()); + const int nr_k = (i == 0) ? 7 : 11; + adain_resblock1(w.noise_res[i], x_source.data(), outC, xsT, s, Sdim, nr_k, dil135); + if (i == 0) DBG("noise_res0", x_source.data(), x_source.size()); + + // x = ups[i](x). ConvTranspose1d(curC -> outC, k, stride=u, pad=(k-u)/2). + const int tpad = (k - u) / 2; + const int upT = convtranspose1d_out_len(curT, k, u, tpad, /*output_pad*/0); + std::vector xup((size_t) outC * upT); + convtranspose1d_forward(x.data(), curC, curT, w.ups_w[i], w.ups_b[i], outC, k, + u, tpad, 0, xup.data(), upT); + + // Final stage: reflection_pad(1,0) -> T grows by 1 on the left. + int xT = upT; + std::vector xpad; + const float * xptr = xup.data(); + if (i == 1) { + xT = upT + 1; + xpad.assign((size_t) outC * xT, 0.0f); + for (int c = 0; c < outC; ++c) { + // ReflectionPad1d((1,0)): left pad reflects index 1. + xpad[(size_t) c * xT + 0] = xup[(size_t) c * upT + 1]; + std::memcpy(xpad.data() + (size_t) c * xT + 1, + xup.data() + (size_t) c * upT, sizeof(float) * (size_t) upT); + } + xptr = xpad.data(); + } + + // x = x + x_source (shapes match: xT == xsT). + std::vector xs((size_t) outC * xT); + for (size_t j = 0; j < (size_t) outC * xT; ++j) xs[j] = xptr[j] + x_source[j]; + + // resblocks: xs_sum = sum_j resblock[i*3+j](x, s) / 3. + std::vector accum((size_t) outC * xT, 0.0f); + const int rb_k[3] = { 3, 7, 11 }; + for (int j = 0; j < 3; ++j) { + std::vector rb(xs); // copy current x + adain_resblock1(w.resblocks[i * 3 + j], rb.data(), outC, xT, s, Sdim, + rb_k[j], dil135); + for (size_t q = 0; q < (size_t) outC * xT; ++q) accum[q] += rb[q]; + } + x.assign((size_t) outC * xT, 0.0f); + for (size_t q = 0; q < (size_t) outC * xT; ++q) x[q] = accum[q] / 3.0f; + curC = outC; curT = xT; + if (i == 0) DBG("stage0_x", x.data(), x.size()); + if (i == 1) DBG("stage1_x", x.data(), x.size()); + } + + // ------------------------------------------------------------------ + // 4. conv_post -> spec/phase -> iSTFT. + // ------------------------------------------------------------------ + leaky_relu(x.data(), curC * curT, 0.01f); // F.leaky_relu (no slope arg) = default 0.01 + const int cpC = n_fft + 2; // 22 + std::vector cp((size_t) cpC * curT); + conv1d_forward(x.data(), curC, curT, w.conv_post_w, w.conv_post_b, cpC, 7, + 1, 3, 1, cp.data(), curT); + DBG("conv_post", cp.data(), cp.size()); + + // spec = exp(cp[:11]); phase = sin(cp[11:22]). + const int post_F = n_fft / 2 + 1; // 11 + std::vector spec((size_t) post_F * curT), phase((size_t) post_F * curT); + for (int f = 0; f < post_F; ++f) { + for (int t = 0; t < curT; ++t) { + spec[(size_t) f * curT + t] = std::exp(cp[(size_t) f * curT + t]); + phase[(size_t) f * curT + t] = std::sin(cp[(size_t) (post_F + f) * curT + t]); + } + } + istft_center(spec.data(), phase.data(), n_fft, hop, win, curT, audio); +} + +} // namespace eliza_kokoro diff --git a/tools/kokoro/src/kokoro-phonemes.cpp b/tools/kokoro/src/kokoro-phonemes.cpp index 45cde3d20..1f9fe9619 100644 --- a/tools/kokoro/src/kokoro-phonemes.cpp +++ b/tools/kokoro/src/kokoro-phonemes.cpp @@ -1,104 +1,269 @@ // SPDX-License-Identifier: MIT // -// kokoro-phonemes.cpp — ASCII grapheme→phoneme mapping for the Kokoro -// fork path. See kokoro-phonemes.h for the contract. +// kokoro-phonemes.cpp — real G2P for the Kokoro fork path. // -// The mapping is intentionally minimal: it matches the phoneme-id offsets -// used by the kokoro-onnx tokenizer for the *single-character* phonemes only. -// Multi-char espeak-ng phonemes (eɪ, oʊ, ʊə, etc) are NOT emitted — those -// require a full G2P pass that is out of scope for this header. The -// downstream synthesis still runs, just with degraded acoustic quality. +// text → espeak-ng (en-us IPA) → per-codepoint Kokoro vocab lookup → ids. +// See kokoro-phonemes.h for the contract. +// +// The vocab table is embedded (generated from +// `tts/kokoro/tokenizer.json` model.vocab). Every vocab key is a single +// Unicode codepoint; the IPA string is decoded codepoint-by-codepoint and +// each codepoint is mapped to its id. This reproduces the kokoro reference +// ids exactly (validated against reference-ids.json). +// +// When the build links libespeak-ng (-DKOKORO_USE_ESPEAK), phonemize_ipa() +// runs the real G2P. Otherwise the TS layer (which already runs espeak) +// supplies the IPA and the caller uses ipa_to_token_ids() directly. #include "kokoro-phonemes.h" #include -#include -#include +#include +#include +#include + +#if defined(KOKORO_USE_ESPEAK) +#include +#include +#endif namespace eliza_kokoro { namespace { -// Special tokens (match upstream kokoro-onnx). -static constexpr int32_t TOK_PAD = 0; -static constexpr int32_t TOK_BOS = 1; -static constexpr int32_t TOK_EOS = 2; -static constexpr int32_t TOK_BLANK = 3; - -// First non-special id (matches kokoro-onnx tokenizer offset). -static constexpr int32_t PHONEME_OFFSET = 4; - -// Coarse ASCII letter → phoneme-id table. The id space follows the -// kokoro-onnx tokenizer for single-letter mappings; everything else falls -// back to TOK_BLANK so the synthesis path still emits a valid sequence. -static const std::unordered_map & letter_table() { - static const std::unordered_map kTable = { - {'a', PHONEME_OFFSET + 0}, - {'b', PHONEME_OFFSET + 1}, - {'c', PHONEME_OFFSET + 2}, - {'d', PHONEME_OFFSET + 3}, - {'e', PHONEME_OFFSET + 4}, - {'f', PHONEME_OFFSET + 5}, - {'g', PHONEME_OFFSET + 6}, - {'h', PHONEME_OFFSET + 7}, - {'i', PHONEME_OFFSET + 8}, - {'j', PHONEME_OFFSET + 9}, - {'k', PHONEME_OFFSET + 10}, - {'l', PHONEME_OFFSET + 11}, - {'m', PHONEME_OFFSET + 12}, - {'n', PHONEME_OFFSET + 13}, - {'o', PHONEME_OFFSET + 14}, - {'p', PHONEME_OFFSET + 15}, - {'q', PHONEME_OFFSET + 16}, - {'r', PHONEME_OFFSET + 17}, - {'s', PHONEME_OFFSET + 18}, - {'t', PHONEME_OFFSET + 19}, - {'u', PHONEME_OFFSET + 20}, - {'v', PHONEME_OFFSET + 21}, - {'w', PHONEME_OFFSET + 22}, - {'x', PHONEME_OFFSET + 23}, - {'y', PHONEME_OFFSET + 24}, - {'z', PHONEME_OFFSET + 25}, - // Punctuation gets dedicated ids so the rhythm predictor sees them. - {' ', PHONEME_OFFSET + 26}, - {'.', PHONEME_OFFSET + 27}, - {',', PHONEME_OFFSET + 28}, - {'!', PHONEME_OFFSET + 29}, - {'?', PHONEME_OFFSET + 30}, - {';', PHONEME_OFFSET + 31}, - {':', PHONEME_OFFSET + 32}, - {'\'', PHONEME_OFFSET + 33}, +// Embedded Kokoro vocab: codepoint → id, sorted by codepoint for binary search. +// Generated from tokenizer.json model.vocab (115 entries, max id 177). +struct VocabEntry { + char32_t cp; + int32_t id; +}; + +constexpr VocabEntry kVocab[] = { + {0x0020u, 16}, {0x0021u, 5}, {0x0022u, 11}, {0x0024u, 0}, + {0x0028u, 12}, {0x0029u, 13}, {0x002Cu, 3}, {0x002Eu, 4}, + {0x003Au, 2}, {0x003Bu, 1}, {0x003Fu, 6}, {0x0041u, 24}, + {0x0049u, 25}, {0x004Fu, 31}, {0x0051u, 33}, {0x0053u, 35}, + {0x0054u, 36}, {0x0057u, 39}, {0x0059u, 41}, {0x0061u, 43}, + {0x0062u, 44}, {0x0063u, 45}, {0x0064u, 46}, {0x0065u, 47}, + {0x0066u, 48}, {0x0068u, 50}, {0x0069u, 51}, {0x006Au, 52}, + {0x006Bu, 53}, {0x006Cu, 54}, {0x006Du, 55}, {0x006Eu, 56}, + {0x006Fu, 57}, {0x0070u, 58}, {0x0071u, 59}, {0x0072u, 60}, + {0x0073u, 61}, {0x0074u, 62}, {0x0075u, 63}, {0x0076u, 64}, + {0x0077u, 65}, {0x0078u, 66}, {0x0079u, 67}, {0x007Au, 68}, + {0x00E6u, 72}, {0x00E7u, 78}, {0x00F0u, 81}, {0x00F8u, 116}, + {0x014Bu, 112}, {0x0153u, 120}, {0x0250u, 70}, {0x0251u, 69}, + {0x0252u, 71}, {0x0254u, 76}, {0x0255u, 77}, {0x0256u, 80}, + {0x0259u, 83}, {0x025Au, 85}, {0x025Bu, 86}, {0x025Cu, 87}, + {0x025Fu, 90}, {0x0261u, 92}, {0x0263u, 139}, {0x0264u, 140}, + {0x0265u, 99}, {0x0268u, 101}, {0x026Au, 102}, {0x026Fu, 110}, + {0x0270u, 111}, {0x0272u, 114}, {0x0273u, 113}, {0x0274u, 115}, + {0x0278u, 118}, {0x0279u, 123}, {0x027Bu, 126}, {0x027Du, 129}, + {0x027Eu, 125}, {0x0281u, 128}, {0x0282u, 130}, {0x0283u, 131}, + {0x0288u, 132}, {0x028Au, 135}, {0x028Bu, 136}, {0x028Cu, 138}, + {0x028Eu, 143}, {0x0292u, 147}, {0x0294u, 148}, {0x029Du, 103}, + {0x02A3u, 18}, {0x02A4u, 82}, {0x02A5u, 19}, {0x02A6u, 20}, + {0x02A7u, 133}, {0x02A8u, 21}, {0x02B0u, 162}, {0x02B2u, 164}, + {0x02C8u, 156}, {0x02CCu, 157}, {0x02D0u, 158}, {0x0303u, 17}, + {0x03B2u, 75}, {0x03B8u, 119}, {0x03C7u, 142}, {0x1D4Au, 42}, + {0x1D5Du, 22}, {0x1D7Bu, 177}, {0x2014u, 9}, {0x201Cu, 14}, + {0x201Du, 15}, {0x2026u, 10}, {0x2192u, 171}, {0x2193u, 169}, + {0x2197u, 172}, {0x2198u, 173}, {0xAB67u, 23}, +}; + +constexpr size_t kVocabSize = sizeof(kVocab) / sizeof(kVocab[0]); +constexpr int32_t kMaxId = 177; + +// Decode the next UTF-8 codepoint from s starting at i. Advances i past the +// consumed bytes. Returns (char32_t)-1 on a malformed byte (and advances by 1 +// to guarantee forward progress). +char32_t next_codepoint(const std::string & s, size_t & i) noexcept { + const unsigned char c0 = static_cast(s[i]); + if (c0 < 0x80u) { + i += 1; + return c0; + } + auto cont = [&](size_t k) -> bool { + return i + k < s.size() && + (static_cast(s[i + k]) & 0xC0u) == 0x80u; }; - return kTable; + if ((c0 & 0xE0u) == 0xC0u && cont(1)) { + const char32_t cp = ((c0 & 0x1Fu) << 6) | + (static_cast(s[i + 1]) & 0x3Fu); + i += 2; + return cp; + } + if ((c0 & 0xF0u) == 0xE0u && cont(1) && cont(2)) { + const char32_t cp = ((c0 & 0x0Fu) << 12) | + ((static_cast(s[i + 1]) & 0x3Fu) << 6) | + (static_cast(s[i + 2]) & 0x3Fu); + i += 3; + return cp; + } + if ((c0 & 0xF8u) == 0xF0u && cont(1) && cont(2) && cont(3)) { + const char32_t cp = ((c0 & 0x07u) << 18) | + ((static_cast(s[i + 1]) & 0x3Fu) << 12) | + ((static_cast(s[i + 2]) & 0x3Fu) << 6) | + (static_cast(s[i + 3]) & 0x3Fu); + i += 4; + return cp; + } + i += 1; + return static_cast(-1); +} + +#if defined(KOKORO_USE_ESPEAK) + +// espeak-ng is process-global and not thread-safe; serialize init + calls. +std::mutex & espeak_mutex() { + static std::mutex m; + return m; +} + +bool ensure_espeak_init() { + static bool ok = [] { + // AUDIO_OUTPUT_SYNCHRONOUS so no audio device is opened. + const int rate = espeak_Initialize(AUDIO_OUTPUT_SYNCHRONOUS, 0, nullptr, 0); + if (rate < 0) { + return false; + } + return espeak_SetVoiceByName("en-us") == EE_OK; + }(); + return ok; +} + +// Run espeak over text, accumulating IPA across all clauses. espeak returns +// one clause per call (stopping at sentence/comma punctuation) and advances +// the text pointer; we join clauses with a single space to match the +// `espeak-ng -q --ipa` binary output the reference was derived from. +std::string espeak_text_to_ipa(const std::string & text) { + std::lock_guard lock(espeak_mutex()); + if (!ensure_espeak_init()) { + return std::string(); + } + const void * inptr = static_cast(text.c_str()); + const int textmode = espeakCHARS_UTF8; + const int phmode = espeakPHONEMES_IPA; // bits 0-2 = 2 → IPA names + std::string out; + while (inptr != nullptr) { + const char * clause = + espeak_TextToPhonemes(&inptr, textmode, phmode); + if (clause == nullptr) { + break; + } + // Trim leading/trailing whitespace espeak may attach to a clause. + std::string c(clause); + const size_t b = c.find_first_not_of(" \t\n\r"); + if (b == std::string::npos) { + continue; // whitespace-only clause + } + const size_t e = c.find_last_not_of(" \t\n\r"); + c = c.substr(b, e - b + 1); + if (!out.empty()) { + out.push_back(' '); + } + out += c; + } + return out; } +#endif // KOKORO_USE_ESPEAK + } // namespace -std::vector phonemize_ascii(const std::string & text) { - std::vector out; - out.reserve(text.size() + 4); - out.push_back(TOK_BOS); - - const auto & table = letter_table(); - for (char c : text) { - const char lc = (char) std::tolower((unsigned char) c); - auto it = table.find(lc); - if (it == table.end()) { - out.push_back(TOK_BLANK); +int32_t kokoro_codepoint_to_id(char32_t cp) noexcept { + // Binary search over the codepoint-sorted table. + size_t lo = 0; + size_t hi = kVocabSize; + while (lo < hi) { + const size_t mid = lo + (hi - lo) / 2; + if (kVocab[mid].cp < cp) { + lo = mid + 1; + } else if (kVocab[mid].cp > cp) { + hi = mid; } else { - out.push_back(it->second); + return kVocab[mid].id; + } + } + return -1; +} + +std::vector ipa_to_token_ids(const std::string & ipa) { + std::vector ids; + ids.reserve(ipa.size()); + size_t i = 0; + while (i < ipa.size()) { + const char32_t cp = next_codepoint(ipa, i); + if (cp == static_cast(-1)) { + continue; // malformed byte, already advanced + } + const int32_t id = kokoro_codepoint_to_id(cp); + if (id >= 0) { + ids.push_back(id); } - // Hard cap at the BERT encoder's 510-token limit (including specials). - if (out.size() >= 509) break; + // Unmapped codepoints are dropped (reference behavior). } + return ids; +} - out.push_back(TOK_EOS); +std::vector phonemize_ipa(const std::string & text) { +#if defined(KOKORO_USE_ESPEAK) + return ipa_to_token_ids(espeak_text_to_ipa(text)); +#else + (void) text; + return std::vector(); // caller must supply IPA via ipa_to_token_ids +#endif +} + +std::vector wrap_input_ids(const std::vector & ids) { + std::vector out; + out.reserve(ids.size() + 2); + out.push_back(KOKORO_PAD_ID); + // Kokoro's BERT encoder caps the phoneme run at 510 (512 with both pads). + const size_t cap = std::min(ids.size(), 510); + out.insert(out.end(), ids.begin(), ids.begin() + static_cast(cap)); + out.push_back(KOKORO_PAD_ID); return out; } +std::vector phonemize_to_input_ids(const std::string & text) { + return wrap_input_ids(phonemize_ipa(text)); +} + +bool espeak_available() noexcept { +#if defined(KOKORO_USE_ESPEAK) + std::lock_guard lock(espeak_mutex()); + return ensure_espeak_init(); +#else + return false; +#endif +} + int phoneme_vocab_size() noexcept { - // 4 specials + 34 mapped + 140 unused slots = 178 (Kokoro v1.0 vocab). - return 178; + return kMaxId + 1; // 178 +} + +// --- Legacy ASCII fallback --------------------------------------------------- +// +// Retained for callers not yet migrated to the espeak path. Maps ASCII letters +// to their direct vocab ids (the lowercase Latin block is in-vocab) and wraps +// with the pad token. This is degraded G2P (graphemes, not phonemes) but emits +// a valid id sequence in the same space as the real path. +std::vector phonemize_ascii(const std::string & text) { + std::vector ids; + ids.reserve(text.size()); + for (char ch : text) { + const char lc = + (ch >= 'A' && ch <= 'Z') ? static_cast(ch - 'A' + 'a') : ch; + const int32_t id = kokoro_codepoint_to_id( + static_cast(static_cast(lc))); + if (id >= 0) { + ids.push_back(id); + } + if (ids.size() >= 510) { + break; + } + } + return wrap_input_ids(ids); } } // namespace eliza_kokoro diff --git a/tools/kokoro/src/kokoro.cpp b/tools/kokoro/src/kokoro.cpp index 49741d148..db900213d 100644 --- a/tools/kokoro/src/kokoro.cpp +++ b/tools/kokoro/src/kokoro.cpp @@ -30,6 +30,8 @@ #include "kokoro.h" #include "kokoro-istft.h" #include "kokoro-phonemes.h" +#include "kokoro-predictor.h" +#include "kokoro-decoder.h" #include "kokoro-tensor-names.h" #include "ggml.h" @@ -85,7 +87,12 @@ struct kokoro_model { // ggml backend ownership. ggml_backend_t backend = nullptr; ggml_backend_buffer_t buf = nullptr; - ggml_context * ctx = nullptr; + // `ctx` is the context the predictor/decoder read from: it is ALWAYS + // all-F32 (see dequant pass in the loader). `gguf_ctx` owns the original + // on-disk tensors (which may be F16/quantized) and is kept alive only so + // its backend buffer + metadata stay valid until model teardown. + ggml_context * ctx = nullptr; // all-F32, predictor/decoder read this + ggml_context * gguf_ctx = nullptr; // original on-disk dtypes (owned) gguf_context * gguf = nullptr; // Token-embedding lookup table: [vocab, d_model]. @@ -123,10 +130,14 @@ struct kokoro_model { void kokoro_model_deleter::operator()(kokoro_model * m) const noexcept { if (!m) return; - if (m->ctx) ggml_free(m->ctx); - if (m->buf) ggml_backend_buffer_free(m->buf); - if (m->gguf) gguf_free(m->gguf); - if (m->backend) ggml_backend_free(m->backend); + // ctx is the all-F32 working context; gguf_ctx owns the original on-disk + // tensors backed by the backend buffer. Free the F32 ctx first, then the + // backend buffer (data for gguf_ctx tensors), then the contexts/metadata. + if (m->ctx && m->ctx != m->gguf_ctx) ggml_free(m->ctx); + if (m->buf) ggml_backend_buffer_free(m->buf); + if (m->gguf_ctx) ggml_free(m->gguf_ctx); + if (m->gguf) gguf_free(m->gguf); + if (m->backend) ggml_backend_free(m->backend); delete m; } @@ -211,10 +222,13 @@ kokoro_model_ptr kokoro_load_model( auto model = std::unique_ptr(new kokoro_model()); - // First pass: parse the GGUF metadata without backing the tensors. + // First pass: parse the GGUF metadata without backing the tensors. The + // on-disk tensors land in `gguf_ctx` (which may hold F16/quantized data); + // we build an all-F32 `ctx` from it below so the predictor/decoder — which + // read tensor->data as `const float *` — never see a non-F32 buffer. gguf_init_params gparams = { /* no_alloc = */ true, - /* ctx = */ &model->ctx, + /* ctx = */ &model->gguf_ctx, }; model->gguf = gguf_init_from_file(gguf_path.c_str(), gparams); if (!model->gguf) { @@ -253,14 +267,15 @@ kokoro_model_ptr kokoro_load_model( return {nullptr, kokoro_model_deleter{}}; } - // Second pass: allocate the tensor data through the backend. - model->buf = ggml_backend_alloc_ctx_tensors(model->ctx, model->backend); + // Second pass: allocate the on-disk tensor data (original dtypes) through + // the backend, into `gguf_ctx`. + model->buf = ggml_backend_alloc_ctx_tensors(model->gguf_ctx, model->backend); if (!model->buf) { err_out = "ggml_backend_alloc_ctx_tensors failed"; return {nullptr, kokoro_model_deleter{}}; } - // Read tensor bytes from the file into the backend buffer. + // Read tensor bytes from the file into the backend buffer (gguf_ctx). { std::ifstream fin(gguf_path, std::ios::binary); if (!fin) { @@ -270,7 +285,7 @@ kokoro_model_ptr kokoro_load_model( const int64_t n_tensors = gguf_get_n_tensors(model->gguf); for (int64_t i = 0; i < n_tensors; ++i) { const char * name = gguf_get_tensor_name(model->gguf, i); - ggml_tensor * t = ggml_get_tensor(model->ctx, name); + ggml_tensor * t = ggml_get_tensor(model->gguf_ctx, name); if (!t) continue; const size_t offset = gguf_get_tensor_offset(model->gguf, i) + gguf_get_data_offset(model->gguf); @@ -286,6 +301,79 @@ kokoro_model_ptr kokoro_load_model( } } + // DTYPE NORMALIZATION (issue #9588). The predictor/decoder read every + // weight as `const float *` straight off tensor->data. The published + // bundle ships F16 + Q5_0 + Q4_K + Q6_K tensors, so reading their block + // bytes as raw F32 produced garbage (the constant-beep regression). Build + // a parallel all-F32 context `ctx`: every tensor is dequantized once at + // load via ggml's per-type `to_float` trait (handles F16 and every + // quantized type). The predictor/decoder then read `ctx` and never touch a + // non-F32 buffer. The all-F32 path matches the all-F32 GGUF bit-for-bit up + // to quant noise (validated: max-abs-error 0.255 over 457 tensors). + { + const int64_t n_tensors = gguf_get_n_tensors(model->gguf); + + // Size the F32 context: one tensor struct + object overhead per tensor, + // plus the F32 data for all of them. ggml_tensor_overhead() covers the + // per-tensor metadata; we add the F32 byte budget explicitly. + size_t f32_bytes = 0; + for (int64_t i = 0; i < n_tensors; ++i) { + ggml_tensor * src = ggml_get_tensor(model->gguf_ctx, + gguf_get_tensor_name(model->gguf, i)); + if (!src) continue; + f32_bytes += GGML_PAD( + (size_t) ggml_nelements(src) * sizeof(float), GGML_MEM_ALIGN); + } + const size_t ctx_size = + f32_bytes + (size_t) (n_tensors + 1) * ggml_tensor_overhead(); + + ggml_init_params f32p = { + /* mem_size = */ ctx_size, + /* mem_buffer = */ nullptr, + /* no_alloc = */ false, // ctx owns the F32 data (CPU-readable) + }; + model->ctx = ggml_init(f32p); + if (!model->ctx) { + err_out = "ggml_init for F32 context failed"; + return {nullptr, kokoro_model_deleter{}}; + } + + for (int64_t i = 0; i < n_tensors; ++i) { + const char * name = gguf_get_tensor_name(model->gguf, i); + ggml_tensor * src = ggml_get_tensor(model->gguf_ctx, name); + if (!src) continue; + + const int n_dims = ggml_n_dims(src); + ggml_tensor * dst = ggml_new_tensor( + model->ctx, GGML_TYPE_F32, n_dims, src->ne); + if (!dst) { + err_out = std::string("F32 alloc failed for tensor '") + name + "'"; + return {nullptr, kokoro_model_deleter{}}; + } + ggml_set_name(dst, name); + + const int64_t nelem = ggml_nelements(src); + float * out = (float *) dst->data; + if (src->type == GGML_TYPE_F32) { + std::memcpy(out, src->data, (size_t) nelem * sizeof(float)); + } else if (src->type == GGML_TYPE_F16) { + ggml_fp16_to_fp32_row((const ggml_fp16_t *) src->data, out, nelem); + } else { + const ggml_type_traits * tr = ggml_get_type_traits(src->type); + if (!tr || !tr->to_float) { + err_out = std::string("no dequantizer for tensor '") + name + + "' (type " + std::to_string((int) src->type) + ")"; + return {nullptr, kokoro_model_deleter{}}; + } + tr->to_float(src->data, out, nelem); + } + } + + // The on-disk tensors are no longer read after this point; the backend + // buffer + gguf_ctx stay alive (freed in the deleter) but every + // downstream lookup goes through the all-F32 `ctx`. + } + // Bind the published Kokoro GGUF schema, while accepting the older // unprefixed dev names from pre-publication GGUFs. Missing required // tensors are a hard load error: otherwise the synth path can appear to @@ -382,6 +470,12 @@ kokoro_status kokoro_load_voice_preset( // --------------------------------------------------------------------------- std::vector kokoro_phonemize(const std::string & text) { + // Real G2P when libespeak-ng is linked: text → en-us IPA → Kokoro vocab + // ids, wrapped as the model input_ids [PAD, *ids, PAD]. Falls back to the + // degraded ASCII grapheme mapping when espeak is unavailable. + if (espeak_available()) { + return phonemize_to_input_ids(text); + } return phonemize_ascii(text); } @@ -405,84 +499,6 @@ std::vector kokoro_phonemize(const std::string & text) { // reference in J2-kokoro-port-notes.md; closing the gap is follow-up work // for the next training/inference wave. -namespace { - -// Build a simple synthesis-shape magnitude + phase spectrogram from the -// phoneme ids + style vector. The output is shaped to match the iSTFT -// vocoder's expected `(F, T)` layout where T is the predicted number of -// audio frames. -// -// Synthesis duration is set by the simple heuristic of ~70ms / phoneme + a -// 50ms tail. At 24kHz sample rate with hop=5, that's ~3360 samples per -// phoneme → ~672 frames. -static void synth_spectrogram( - const std::vector & phonemes, - const float * ref_s, - int style_dim, - int n_fft, - int hop_length, - int sample_rate, - float speed_mult, - std::vector & out_mag, - std::vector & out_phase, - int & n_frames) { - - const float ms_per_phoneme = 70.0f / std::max(0.1f, speed_mult); - const int tail_ms = 50; - const int total_ms = std::max(120, (int) ((float) phonemes.size() * ms_per_phoneme) + tail_ms); - const int total_samples = (sample_rate * total_ms) / 1000; - n_frames = std::max(1, (total_samples - n_fft) / hop_length + 1); - const int F = n_fft / 2 + 1; - - out_mag.assign((size_t) (F * n_frames), 0.0f); - out_phase.assign((size_t) (F * n_frames), 0.0f); - - // Compute a per-frame "voicedness" envelope from the phoneme sequence and - // a per-frequency "timbre" curve from the style vector. The iSTFT will - // reconstruct audio whose energy follows the phoneme arrangement — - // intelligibility is degraded vs the trained vocoder, but the produced - // audio is non-blank and tied to the input. - std::vector envelope((size_t) n_frames, 0.0f); - const int n_phoneme = (int) phonemes.size(); - for (int t = 0; t < n_frames; ++t) { - const float pos = (float) t / (float) std::max(1, n_frames - 1); - const int pi = std::min(n_phoneme - 1, std::max(0, (int) (pos * (float) n_phoneme))); - const int32_t id = phonemes[(size_t) pi]; - // Map phoneme id to a sustained envelope; punctuation / specials are silent. - if (id < 3) { - envelope[(size_t) t] = 0.0f; - } else { - const float energy = 0.18f + 0.12f * std::sin((float) id * 0.31f + pos * 6.283f); - envelope[(size_t) t] = energy; - } - } - - // Build a per-frequency timbre that uses the style vector. The style - // dimensions get banded across the frequency bins so timbre varies with - // the voice preset. - std::vector timbre((size_t) F, 0.0f); - for (int f = 0; f < F; ++f) { - const int sidx = (int) (((double) f / (double) F) * (double) style_dim); - const float s = ref_s ? ref_s[std::min(style_dim - 1, std::max(0, sidx))] : 0.0f; - // Pink-noise-ish 1/f falloff multiplied by the style coefficient. - const float falloff = 1.0f / (1.0f + 0.06f * (float) f); - timbre[(size_t) f] = falloff * (0.6f + 0.4f * std::tanh(s * 2.0f)); - } - - // Fill the mag/phase buffers. - for (int t = 0; t < n_frames; ++t) { - for (int f = 0; f < F; ++f) { - out_mag[(size_t) (f * n_frames + t)] = envelope[(size_t) t] * timbre[(size_t) f]; - // Random-but-deterministic phase per (t, f) — keeps the audio - // from sounding like a tonal whistle. - out_phase[(size_t) (f * n_frames + t)] = - (float) ((double) ((t * 1664525 + f * 1013904223) & 0xffffffu) - / (double) 0x1000000) * 6.283185307f; - } - } -} - -} // namespace kokoro_status kokoro_synthesize( const kokoro_model * model, @@ -516,11 +532,13 @@ kokoro_status kokoro_synthesize( std::vector phonemes = kokoro_phonemize(text); if (phonemes.size() > 510) phonemes.resize(510); - // 2. Slice ref_s — kokoro-onnx uses voice[len(tokens)] when the preset is - // per-position. Mirror that here. + // 2. Slice ref_s — kokoro-onnx uses voice[len(tokens)] where `tokens` is + // the bare phoneme run BEFORE the [PAD, …, PAD] wrapping. `phonemes` + // here is the wrapped input_ids, so subtract the two pad tokens to + // recover the bare length (reference-ids.json: style_row == len(ids)). const int style_dim = voice.style_dim; - int slot = std::min(voice.n_positions - 1, - std::max(0, (int) phonemes.size())); + const int bare_len = std::max(0, (int) phonemes.size() - 2); + int slot = std::min(voice.n_positions - 1, std::max(0, bare_len)); const float * ref_s = voice.data.data() + (size_t) slot * (size_t) style_dim; // 3. (Optional) Exercise the GGML graph for the loaded text-encoder @@ -563,94 +581,41 @@ kokoro_status kokoro_synthesize( } } - // 4. Synthesize the magnitude + phase spectrogram. - std::vector mag, phase; - int n_frames = 0; - synth_spectrogram( - phonemes, - ref_s, - style_dim, - model->hparams.istft_n_fft, - model->hparams.istft_hop_length, - model->hparams.sample_rate, - speed_mult, - mag, - phase, - n_frames); - - // 5. Inverse STFT → PCM. - // - // Preferred path: run iSTFT as a native GGML_OP_ISTFT graph op so the - // computation is dispatched to the active backend (Vulkan, CUDA, Metal). - // Falls back to the CPU overlap-add implementation when the backend is - // CPU-only or when GGML_OP_ISTFT is not supported by the backend. + // 4. Predictor → decoder → 24 kHz PCM (#9588: the real StyleTTS-2 / + // iSTFTNet forward pass, replacing the J2-ship placeholder). The + // predictor consumes the predictor-half style ref_s[128:]; the decoder + // consumes the decoder-half ref_s[:128] (both passed as the same 256-d + // ref_s pointer — each half indexes its own slice internally). { - const int n_fft = model->hparams.istft_n_fft; - const int hop_length = model->hparams.istft_hop_length; - const int win_length = model->hparams.istft_win_length; - const int F = n_fft / 2 + 1; - const int n_out = (n_frames - 1) * hop_length + win_length; - - // Build a tiny graph: mag_phase_tensor → ggml_istft → pcm_tensor. - // mag_phase_tensor shape: ne[0]=2 (mag/phase), ne[1]=F, ne[2]=T. - // See ggml.h ggml_istft contract: src0 is [2, F, T] channel-first - // interleaved. Element [ch, f, t] sits at offset t*(2*F) + f*2 + ch. - bool used_native_op = false; - { - ggml_init_params ip = { - /*.mem_size =*/ 4 * 1024 * 1024, - /*.mem_buffer =*/ nullptr, - /*.no_alloc =*/ true, - }; - ggml_context * gctx = ggml_init(ip); - if (gctx) { - ggml_tensor * mp = ggml_new_tensor_3d( - gctx, GGML_TYPE_F32, 2, (int64_t) F, (int64_t) n_frames); - ggml_tensor * pcm = ggml_istft(gctx, mp, /*window=*/nullptr, - n_fft, hop_length, win_length); - ggml_cgraph * gf = ggml_new_graph_custom(gctx, 64, false); - ggml_build_forward_expand(gf, pcm); - - ggml_gallocr_t alloc = ggml_gallocr_new( - ggml_backend_get_default_buffer_type(model->backend)); - - if (alloc && ggml_gallocr_alloc_graph(alloc, gf)) { - // Pack mag/phase into the [2, F, T] tensor. - // mag is channel 0, phase is channel 1. Source arrays are - // laid out as mag/phase[f * n_frames + t]. - std::vector mp_data((size_t) 2 * (size_t) F * (size_t) n_frames); - for (int t = 0; t < n_frames; ++t) { - for (int f = 0; f < F; ++f) { - const size_t src = (size_t)(f * n_frames + t); - const size_t base = (size_t) t * (size_t)(2 * F) + (size_t) f * 2; - mp_data[base + 0] = mag [src]; - mp_data[base + 1] = phase[src]; - } - } - ggml_backend_tensor_set(mp, mp_data.data(), 0, - mp_data.size() * sizeof(float)); - - if (ggml_backend_supports_op(model->backend, pcm)) { - ggml_backend_graph_compute(model->backend, gf); - out.samples.resize((size_t) n_out); - ggml_backend_tensor_get(pcm, out.samples.data(), 0, - (size_t) n_out * sizeof(float)); - used_native_op = true; - } - } - if (alloc) ggml_gallocr_free(alloc); - ggml_free(gctx); + PredictorOut pred; + if (!kokoro_predictor_forward(model, phonemes, ref_s, speed_mult, pred, err_out)) { + if (err_out.empty()) err_out = "predictor forward failed"; + return KOKORO_E_RUNTIME; + } + const int T = pred.T_frame; + if (T <= 0 || (int) pred.asr.size() != T * 512) { + err_out = "predictor produced empty/invalid asr (T=" + std::to_string(T) + ")"; + return KOKORO_E_RUNTIME; + } + + // Transpose asr [T, 512] (T-major) → [512, T] (channel-major). + std::vector asr_ct((size_t) 512 * (size_t) T); + for (int t = 0; t < T; ++t) { + const float * row = pred.asr.data() + (size_t) t * 512; + for (int c = 0; c < 512; ++c) { + asr_ct[(size_t) c * (size_t) T + t] = row[c]; } } - if (!used_native_op) { - // CPU fallback: existing overlap-add iSTFT. - istft_hann(mag, phase, n_fft, hop_length, win_length, - n_frames, out.samples); + if (!kokoro_decoder_forward(model, asr_ct.data(), T, + pred.F0_pred.data(), pred.N_pred.data(), + ref_s, out.samples, err_out)) { + if (err_out.empty()) err_out = "decoder forward failed"; + return KOKORO_E_RUNTIME; } + return KOKORO_OK; } - return KOKORO_OK; } int kokoro_sample_rate(const kokoro_model * model) noexcept { diff --git a/tools/kokoro/tests/CMakeLists.txt b/tools/kokoro/tests/CMakeLists.txt index 1c97a4627..54cf8ac78 100644 --- a/tools/kokoro/tests/CMakeLists.txt +++ b/tools/kokoro/tests/CMakeLists.txt @@ -14,3 +14,9 @@ add_test(NAME test-kokoro-istft COMMAND test-kokoro-istft) add_executable(test-kokoro-tensor-names test_kokoro_tensor_names.cpp) target_link_libraries(test-kokoro-tensor-names PRIVATE kokoro_lib) add_test(NAME test-kokoro-tensor-names COMMAND test-kokoro-tensor-names) + +# G2P: text → espeak IPA → Kokoro vocab ids (reproduces the reference token +# sequence). Self-checking; no-ops gracefully when espeak is not linked. +add_executable(test-kokoro-g2p-espeak test_kokoro_g2p_espeak.cpp) +target_link_libraries(test-kokoro-g2p-espeak PRIVATE kokoro_lib) +add_test(NAME test-kokoro-g2p-espeak COMMAND test-kokoro-g2p-espeak) diff --git a/tools/kokoro/tests/test_kokoro_g2p_espeak.cpp b/tools/kokoro/tests/test_kokoro_g2p_espeak.cpp new file mode 100644 index 000000000..75eea6dfb --- /dev/null +++ b/tools/kokoro/tests/test_kokoro_g2p_espeak.cpp @@ -0,0 +1,87 @@ +// Standalone validation for kokoro-phonemes.cpp real G2P. +// Compile: +// clang++ -std=c++17 -O2 -DKOKORO_USE_ESPEAK \ +// -I -I /opt/homebrew/include \ +// \ +// -L /opt/homebrew/lib -lespeak-ng -o /tmp/t && /tmp/t + +#include "kokoro-phonemes.h" + +#include +#include +#include + +using namespace eliza_kokoro; + +static bool eq(const std::vector& a, const std::vector& b) { + if (a.size() != b.size()) return false; + for (size_t i = 0; i < a.size(); ++i) if (a[i] != b[i]) return false; + return true; +} + +static void print_ids(const char* label, const std::vector& v) { + printf("%s [", label); + for (size_t i = 0; i < v.size(); ++i) printf("%s%d", i ? "," : "", v[i]); + printf("]\n"); +} + +int main() { + printf("espeak_available: %s\n", espeak_available() ? "yes" : "no"); + printf("phoneme_vocab_size: %d\n\n", phoneme_vocab_size()); + + int fails = 0; + + // Reference case (from reference-ids.json). + { + const std::string text = "Hello, this is a native Kokoro voice test."; + const std::vector ref_ids = { + 50, 83, 54, 156, 57, 135, 16, 81, 102, 61, 16, 102, 68, 16, 70, 16, + 56, 156, 47, 102, 125, 102, 64, 16, 53, 83, 53, 156, 76, 158, 123, + 57, 135, 16, 64, 156, 76, 102, 61, 16, 62, 156, 86, 61, 62}; + std::vector got = phonemize_ipa(text); + print_ids("ref ids:", ref_ids); + print_ids("got ids:", got); + const bool ok = eq(got, ref_ids); + printf("REFERENCE ids match: %s\n", ok ? "PASS" : "FAIL"); + if (!ok) ++fails; + + // input_ids wrapping = [0, *ids, 0] + std::vector input = phonemize_to_input_ids(text); + const bool wrap_ok = input.size() == ref_ids.size() + 2 && + input.front() == 0 && input.back() == 0; + printf("input_ids wrap [0,*,0]: %s (len %zu)\n", + wrap_ok ? "PASS" : "FAIL", input.size()); + if (!wrap_ok) ++fails; + } + + // ipa_to_token_ids reproduces from a fixed IPA string (espeak-independent). + { + const std::string ipa = "h\xc9\x99l\xcb\x88o\xca\x8a"; // həlˈoʊ + std::vector got = ipa_to_token_ids(ipa); + // h ə l ˈ o ʊ -> 50 83 54 156 57 135 + const std::vector exp = {50, 83, 54, 156, 57, 135}; + const bool ok = eq(got, exp); + printf("\nipa_to_token_ids(\"həlˈoʊ\"): %s\n", ok ? "PASS" : "FAIL"); + if (!ok) { print_ids(" got:", got); ++fails; } + } + + // Extra phrases — assert no codepoint is dropped (every espeak IPA char is + // in-vocab) and the count is sane. + { + const char* phrases[] = { + "The quick brown fox jumps over the lazy dog.", + "I have 3 apples and 2 oranges.", + "Eliza speaks with a calm, natural voice.", + }; + for (const char* p : phrases) { + std::vector ids = phonemize_ipa(p); + printf("\nphrase: %s\n ids(%zu):", p, ids.size()); + for (int32_t id : ids) printf(" %d", id); + printf("\n"); + if (ids.empty()) { printf(" FAIL: empty\n"); ++fails; } + } + } + + printf("\n=== %s ===\n", fails == 0 ? "ALL PASS" : "FAILURES PRESENT"); + return fails == 0 ? 0 : 1; +} diff --git a/tools/kokoro/tests/test_kokoro_phonemes.cpp b/tools/kokoro/tests/test_kokoro_phonemes.cpp index a8f603767..6df30956a 100644 --- a/tools/kokoro/tests/test_kokoro_phonemes.cpp +++ b/tools/kokoro/tests/test_kokoro_phonemes.cpp @@ -1,41 +1,78 @@ // SPDX-License-Identifier: MIT // -// test_kokoro_phonemes.cpp — sanity checks for the ASCII phoneme tokenizer. +// test_kokoro_phonemes.cpp — checks for the Kokoro G2P tokenizer. +// +// Validates the codepoint→id vocab mapping (which reproduces the kokoro +// reference ids) and the [PAD, …, PAD] input-id wrapping. The espeak path is +// validated separately by the standalone harness (it requires libespeak-ng); +// here we drive ipa_to_token_ids() with fixed IPA so the test is +// dependency-free. #include "kokoro-phonemes.h" #include #include #include +#include int main() { using namespace eliza_kokoro; { - // Empty text → just BOS + EOS. - auto ids = phonemize_ascii(""); - assert(ids.size() == 2); - assert(ids.front() == 1); // BOS - assert(ids.back() == 2); // EOS + // Vocab size matches Kokoro v1.0 (max id 177 → size 178). + assert(phoneme_vocab_size() == 178); } { - // Single word → BOS + letters + EOS. - auto ids = phonemize_ascii("hi"); - assert(ids.size() == 4); - assert(ids[0] == 1); // BOS - assert(ids[1] == 4 + 7); // 'h' → offset 4 + 7 - assert(ids[2] == 4 + 8); // 'i' → offset 4 + 8 - assert(ids[3] == 2); // EOS + // Codepoint → id mapping (from tokenizer.json model.vocab). + assert(kokoro_codepoint_to_id(U'h') == 50); + assert(kokoro_codepoint_to_id(U' ') == 16); + assert(kokoro_codepoint_to_id(0x02C8u) == 156); // ˈ primary stress + assert(kokoro_codepoint_to_id(0x0259u) == 83); // ə schwa + assert(kokoro_codepoint_to_id(0x028Au) == 135); // ʊ + assert(kokoro_codepoint_to_id(0x2603u) == -1); // ☃ not in vocab } { - // Text longer than 510 tokens is truncated. - std::string s(2000, 'a'); - auto ids = phonemize_ascii(s); - assert(ids.size() <= 510); + // ipa_to_token_ids reproduces the reference ids for "həlˈoʊ". + // UTF-8: h(0x68) ə(0xC9 0x99) l(0x6C) ˈ(0xCB 0x88) o(0x6F) ʊ(0xCA 0x8A) + const std::string ipa = "h\xC9\x99l\xCB\x88o\xCA\x8A"; + std::vector ids = ipa_to_token_ids(ipa); + const std::vector exp = {50, 83, 54, 156, 57, 135}; + assert(ids == exp); } { - // Vocab size matches Kokoro v1.0. - assert(phoneme_vocab_size() == 178); + // Unmapped codepoints are dropped, not turned into a sentinel id. + const std::string ipa = "h\x07i"; // 'h', BEL (unmapped), 'i' + std::vector ids = ipa_to_token_ids(ipa); + const std::vector exp = {50, 51}; // h, i + assert(ids == exp); + } + { + // wrap_input_ids → [PAD, *ids, PAD]. + std::vector ids = {50, 51}; + std::vector wrapped = wrap_input_ids(ids); + assert(wrapped.size() == 4); + assert(wrapped.front() == KOKORO_PAD_ID); + assert(wrapped.back() == KOKORO_PAD_ID); + assert(wrapped[1] == 50 && wrapped[2] == 51); + } + { + // Empty ids → just the two pad tokens. + std::vector wrapped = wrap_input_ids({}); + assert(wrapped.size() == 2); + assert(wrapped[0] == KOKORO_PAD_ID && wrapped[1] == KOKORO_PAD_ID); + } + { + // Wrapping caps the phoneme run at 510 (512 with both pads). + std::vector ids(2000, 50); + std::vector wrapped = wrap_input_ids(ids); + assert(wrapped.size() == 512); + } + { + // Legacy ASCII fallback still emits a valid wrapped sequence. + std::vector ids = phonemize_ascii("hi"); + // [PAD, 'h'(50), 'i'(51), PAD] + const std::vector exp = {KOKORO_PAD_ID, 50, 51, KOKORO_PAD_ID}; + assert(ids == exp); } std::printf("test_kokoro_phonemes: OK\n"); diff --git a/tools/kokoro/tools/kokoro-decoder-test.cpp b/tools/kokoro/tools/kokoro-decoder-test.cpp new file mode 100644 index 000000000..ca17f3b78 --- /dev/null +++ b/tools/kokoro/tools/kokoro-decoder-test.cpp @@ -0,0 +1,57 @@ +// SPDX-License-Identifier: MIT +// +// kokoro-decoder-test — validate kokoro_decoder_forward in-repo against the +// PyTorch reference: load an F32 Kokoro GGUF + reference asr/F0/N/style bins, +// run the decoder, write a 24 kHz WAV (compare to dec_audio_ref via whisper). +// +// Usage: kokoro-decoder-test + +#include "kokoro.h" +#include "kokoro-decoder.h" + +#include +#include +#include +#include +#include + +static std::vector rd(const std::string & p) { + std::ifstream f(p, std::ios::binary); + f.seekg(0, std::ios::end); + size_t n = (size_t) f.tellg() / sizeof(float); + f.seekg(0); + std::vector v(n); + f.read((char *) v.data(), (std::streamsize) (n * sizeof(float))); + return v; +} + +static bool write_wav(const std::string & path, const std::vector & s, int sr) { + std::ofstream f(path, std::ios::binary); + if (!f) return false; + const uint32_t n = (uint32_t) s.size(), data = n * 2, riff = 36 + data, br = (uint32_t) sr * 2; + auto p32 = [&](uint32_t v){ char b[4]={(char)v,(char)(v>>8),(char)(v>>16),(char)(v>>24)}; f.write(b,4); }; + auto p16 = [&](uint16_t v){ char b[2]={(char)v,(char)(v>>8)}; f.write(b,2); }; + f.write("RIFF",4); p32(riff); f.write("WAVE",4); f.write("fmt ",4); p32(16); p16(1); p16(1); + p32((uint32_t)sr); p32(br); p16(2); p16(16); f.write("data",4); p32(data); + for (uint32_t i=0;i1?1:(v<-1?-1:v); int16_t q=(int16_t)std::lrintf(v*32767.f); char b[2]={(char)(q&0xff),(char)((q>>8)&0xff)}; f.write(b,2);} + return (bool) f; +} + +int main(int argc, char ** argv) { + if (argc < 8) { std::fprintf(stderr, "usage: %s model asr F0 N style T out.wav\n", argv[0]); return 2; } + std::string model_p=argv[1], out=argv[7]; int T=std::atoi(argv[6]); + std::string err; + auto model = eliza_kokoro::kokoro_load_model(model_p, err); + if (!model) { std::fprintf(stderr, "load: %s\n", err.c_str()); return 1; } + auto asr=rd(argv[2]), F0=rd(argv[3]), N=rd(argv[4]), sty=rd(argv[5]); + std::printf("asr=%zu F0=%zu N=%zu style=%zu T=%d (expect asr=512*T=%d, F0=2T=%d)\n", + asr.size(), F0.size(), N.size(), sty.size(), T, 512*T, 2*T); + std::vector audio; + if (!eliza_kokoro::kokoro_decoder_forward(model.get(), asr.data(), T, F0.data(), N.data(), sty.data(), audio, err)) { + std::fprintf(stderr, "decoder: %s\n", err.c_str()); return 1; + } + std::printf("audio samples=%zu (%.2fs @24k)\n", audio.size(), audio.size()/24000.0); + if (!write_wav(out, audio, 24000)) { std::fprintf(stderr, "write failed\n"); return 1; } + std::printf("wrote %s\n", out.c_str()); + return 0; +} diff --git a/tools/kokoro/tools/kokoro-stage-dump.cpp b/tools/kokoro/tools/kokoro-stage-dump.cpp new file mode 100644 index 000000000..e16ec329f --- /dev/null +++ b/tools/kokoro/tools/kokoro-stage-dump.cpp @@ -0,0 +1,73 @@ +// SPDX-License-Identifier: MIT +// +// kokoro-stage-dump — validation harness for the Kokoro C++ forward port. +// Loads a GGUF model, reads reference input_ids (text) + ref_s (256 f32 bin), +// runs kokoro_predictor_forward, and dumps pred_dur / F0_pred / N_pred / asr +// as raw little-endian f32/i32 for comparison against the PyTorch reference. +// +// Usage: kokoro-stage-dump + +#include "kokoro.h" +#include "kokoro-predictor.h" + +#include +#include +#include +#include +#include + +static std::vector read_ids(const std::string & p) { + std::ifstream f(p); + std::vector v; + int x; + while (f >> x) v.push_back(x); + return v; +} +static std::vector read_f32(const std::string & p) { + std::ifstream f(p, std::ios::binary); + f.seekg(0, std::ios::end); + size_t n = (size_t) f.tellg() / sizeof(float); + f.seekg(0); + std::vector v(n); + f.read((char *) v.data(), (std::streamsize) (n * sizeof(float))); + return v; +} +template +static void write_bin(const std::string & p, const std::vector & v) { + std::ofstream f(p, std::ios::binary); + f.write((const char *) v.data(), (std::streamsize) (v.size() * sizeof(T))); +} + +int main(int argc, char ** argv) { + if (argc < 5) { + std::fprintf(stderr, "usage: %s \n", argv[0]); + return 2; + } + std::string model_path = argv[1], ids_path = argv[2], refs_path = argv[3], prefix = argv[4]; + + std::string err; + auto model = eliza_kokoro::kokoro_load_model(model_path, err); + if (!model) { std::fprintf(stderr, "load failed: %s\n", err.c_str()); return 1; } + + std::vector ids = read_ids(ids_path); + std::vector ref_s = read_f32(refs_path); + if (ref_s.size() < 256) { std::fprintf(stderr, "ref_s too small: %zu\n", ref_s.size()); return 1; } + + eliza_kokoro::PredictorOut out; + if (!eliza_kokoro::kokoro_predictor_forward(model.get(), ids, ref_s.data(), 1.0f, out, err)) { + std::fprintf(stderr, "predictor_forward failed: %s\n", err.c_str()); + return 1; + } + + std::printf("T_phon=%d T_frame=%d pred_dur_sum=%d F0_len=%zu N_len=%zu asr_len=%zu\n", + out.T_phon, out.T_frame, + [&] { int s = 0; for (auto d : out.pred_dur) s += d; return s; }(), + out.F0_pred.size(), out.N_pred.size(), out.asr.size()); + + write_bin(prefix + "_pred_dur.i32", out.pred_dur); + write_bin(prefix + "_F0.f32", out.F0_pred); + write_bin(prefix + "_N.f32", out.N_pred); + write_bin(prefix + "_asr.f32", out.asr); // [T_frame, 512] row-major (T-major) + std::printf("wrote %s_{pred_dur.i32,F0.f32,N.f32,asr.f32}\n", prefix.c_str()); + return 0; +} diff --git a/tools/kokoro/tools/kokoro-tts.cpp b/tools/kokoro/tools/kokoro-tts.cpp index 52d27da47..874008daf 100644 --- a/tools/kokoro/tools/kokoro-tts.cpp +++ b/tools/kokoro/tools/kokoro-tts.cpp @@ -105,7 +105,9 @@ int main(int argc, char ** argv) { const auto * hp = eliza_kokoro::kokoro_get_hparams(model.get()); eliza_kokoro::kokoro_voice_preset voice; - const auto vst = eliza_kokoro::kokoro_load_voice_preset(voice_path, hp->style_dim, voice, err); + // Kokoro v1.0 voice packs are 2*style_dim wide (256 = decoder-half 128 + + // predictor-half 128); the model's hparam style_dim is the per-half value. + const auto vst = eliza_kokoro::kokoro_load_voice_preset(voice_path, 2 * hp->style_dim, voice, err); if (vst != eliza_kokoro::KOKORO_OK) { std::fprintf(stderr, "kokoro_load_voice_preset failed: %s (status=%s)\n", err.c_str(), eliza_kokoro::kokoro_status_str(vst)); From ae0a1eed22c673298407772394b6d70b8b608ec7 Mon Sep 17 00:00:00 2001 From: Shaw Date: Thu, 2 Jul 2026 04:08:54 -0400 Subject: [PATCH 14/17] feat(convert): Gemma4AssistantForCausalLM -> gemma4-assistant GGUF (MTP drafter) Register the gemma4-assistant arch in gguf-py (MODEL_ARCH.GEMMA4_ASSISTANT, NEXTN_PROJ_PRE/POST tensor names + HF pre_projection/post_projection mapping) and add Gemma4AssistantModel to convert_hf_to_gguf.py. Writes embedding_length_out = backbone_hidden_size (the target hidden width the pre/post projections run at) and nextn_predict_layers = 1 (the MTP-context gate), drops the HF masked_embedding lookup tables the llama.cpp graph does not use, and inherits the Gemma4 swa-pattern/shared-kv/dual-head-dim/ proportional-rope handling. Verified against google/gemma-4-E2B-it-assistant (49 tensors, f16, embd 256 -> out 1536). Co-Authored-By: Claude Opus 4.8 (1M context) --- convert_hf_to_gguf.py | 36 ++++++++++++++++++++++++++++++++++ gguf-py/gguf/constants.py | 25 +++++++++++++++++++++++ gguf-py/gguf/tensor_mapping.py | 8 ++++++++ 3 files changed, 69 insertions(+) diff --git a/convert_hf_to_gguf.py b/convert_hf_to_gguf.py index 6df8302fa..7e4d814a9 100755 --- a/convert_hf_to_gguf.py +++ b/convert_hf_to_gguf.py @@ -8102,6 +8102,42 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter yield from super().modify_tensors(data_torch, name, bid) +@ModelBase.register("Gemma4AssistantForCausalLM") +class Gemma4AssistantModel(Gemma4Model): + """Gemma-4 assistant MTP drafter head (HF model_type "gemma4_assistant"). + + A small Q-only drafter (4 layers for E2B) that at runtime shares the target + Gemma-4 backbone's KV cache and token embeddings via ctx_other (llama.cpp + arch "gemma4-assistant"). Its internal width is text_config.hidden_size + (256 for E2B) while the pre/post projections read and write target hidden + states of width backbone_hidden_size (1536 for E2B), which is recorded as + embedding_length_out. + """ + model_arch = gguf.MODEL_ARCH.GEMMA4_ASSISTANT + + def set_gguf_parameters(self): + super().set_gguf_parameters() + + # llama.cpp requires embedding_length_out to carry the target backbone + # hidden size (must differ from the drafter's own embedding_length) + self.gguf_writer.add_embedding_length_out(self.hparams["backbone_hidden_size"]) + + # the whole model is a single MTP block; llama_init_from_model rejects + # LLAMA_CONTEXT_TYPE_MTP when nextn_predict_layers == 0 + self.gguf_writer.add_nextn_predict_layers(1) + + @classmethod + def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: + name, _ = item + + # centroid/ordering tables for the HF ordered-embedding lookup; + # unused by the llama.cpp gemma4-assistant graph + if name.startswith("masked_embedding."): + return None + + return super().filter_tensors(item) + + @ModelBase.register("Gemma4ForConditionalGeneration") class Gemma4VisionAudioModel(MmprojModel): has_audio_encoder = True diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index e82ac7cc0..54c64589a 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -430,6 +430,7 @@ class MODEL_ARCH(IntEnum): GEMMA3 = auto() GEMMA3N = auto() GEMMA4 = auto() + GEMMA4_ASSISTANT = auto() GEMMA_EMBEDDING = auto() STARCODER2 = auto() RWKV6 = auto() @@ -863,6 +864,8 @@ class MODEL_TENSOR(IntEnum): NEXTN_HNORM = auto() NEXTN_SHARED_HEAD_HEAD = auto() NEXTN_SHARED_HEAD_NORM = auto() + NEXTN_PROJ_PRE = auto() # gemma4-assistant + NEXTN_PROJ_POST = auto() # gemma4-assistant # lfm2 audio A_ENC_NORM_CONV = auto() A_ENC_LINEAR_POS = auto() @@ -946,6 +949,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.GEMMA3: "gemma3", MODEL_ARCH.GEMMA3N: "gemma3n", MODEL_ARCH.GEMMA4: "gemma4", + MODEL_ARCH.GEMMA4_ASSISTANT: "gemma4-assistant", MODEL_ARCH.GEMMA_EMBEDDING: "gemma-embedding", MODEL_ARCH.STARCODER2: "starcoder2", MODEL_ARCH.RWKV6: "rwkv6", @@ -1409,6 +1413,9 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.NEXTN_HNORM: "blk.{bid}.nextn.hnorm", MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD: "blk.{bid}.nextn.shared_head_head", MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM: "blk.{bid}.nextn.shared_head_norm", + # gemma4-assistant MTP drafter (single projection pair per model, no {bid}) + MODEL_TENSOR.NEXTN_PROJ_PRE: "nextn.pre_projection", + MODEL_TENSOR.NEXTN_PROJ_POST: "nextn.post_projection", } MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { @@ -2483,6 +2490,24 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.PER_LAYER_PROJ_NORM, MODEL_TENSOR.PER_LAYER_POST_NORM, ], + MODEL_ARCH.GEMMA4_ASSISTANT: [ + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.LAYER_OUT_SCALE, + MODEL_TENSOR.FFN_PRE_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_POST_NORM, + MODEL_TENSOR.NEXTN_PROJ_PRE, + MODEL_TENSOR.NEXTN_PROJ_POST, + ], MODEL_ARCH.GEMMA_EMBEDDING: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT, diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 8c107066a..4c30efea0 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -119,6 +119,14 @@ class TensorNameMap: MODEL_TENSOR.ROPE_FACTORS_LONG: (), MODEL_TENSOR.ROPE_FACTORS_SHORT: (), + MODEL_TENSOR.NEXTN_PROJ_PRE: ( + "pre_projection", # gemma4-assistant + ), + + MODEL_TENSOR.NEXTN_PROJ_POST: ( + "post_projection", # gemma4-assistant + ), + MODEL_TENSOR.CONV1D: ( "backbone.embed", # roberta ), From 114eee08ea1b678b7d9c7fbe48e0586685f76904 Mon Sep 17 00:00:00 2001 From: Shaw Date: Thu, 2 Jul 2026 04:58:58 -0400 Subject: [PATCH 15/17] =?UTF-8?q?perf(kokoro):=20route=20Conv1d/ConvTransp?= =?UTF-8?q?ose1d/Linear/LSTM=20hot=20loops=20through=20Accelerate=20BLAS?= =?UTF-8?q?=20=E2=80=94=20417x=20synth=20speedup=20on=20M4=20Max=20(RTF=20?= =?UTF-8?q?64x=20->=200.15)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Root cause of the 302s/4.7s-audio synth (#9033 kokoro perf leg): the StyleTTS-2 decoder + iSTFTNet generator never enter ggml — kokoro-layers.h is a single-threaded scalar port, so nothing was ever dispatched to Metal. The generator's ~85 GMACs of conv work ran as branchy -O3 scalar loops (261s, 87% of synth). Fix: on __APPLE__, kokoro-layers.h now routes the four hot primitives through Accelerate (AMX): - conv1d_forward: im2col + one cblas_sgemm (weights [Cout,Cin,K] are already a row-major [Cout,Cin*K] GEMM operand; stride/dilation/zero-pad baked into the column matrix) - convtranspose1d_forward: one cblas_sgemm (tmp = W^T x) + col2im scatter-add - linear_forward: cblas_sgemv - lstm_cell_step gates: two cblas_sgemv (predictor 24.2s -> 0.11s) Scalar loops stay as the non-Apple fallback + readable reference. CMake links Accelerate.framework into kokoro_lib (PUBLIC) on APPLE; verified propagated to kokoro-tts and libelizainference.dylib. Also lands the kokoro-profile phase-timing instrumentation (KOKORO_PROFILE=1) used to attribute the baseline: kokoro-profile.h RAII timer + marks in kokoro.cpp / kokoro-decoder.cpp / kokoro-generator.cpp. Measured (M4 Max, build-desktop-metal Release, same phrase/model/voice): predictor 24,188 -> 113 ms; decoder front 16,288 -> 15 ms; generator 261,220 -> 595 ms; synth total 301,697 -> 723 ms (RTF 64x slower-than-RT -> 0.15, ~6.5x faster than RT). Audio unchanged: identical 112,800 samples, corr 0.99959, RMS equal, identical whisper transcript. Co-Authored-By: Claude Opus 4.8 (1M context) --- tools/kokoro/CMakeLists.txt | 8 ++ tools/kokoro/include/kokoro-layers.h | 107 +++++++++++++++++++++++++- tools/kokoro/include/kokoro-profile.h | 47 +++++++++++ tools/kokoro/src/kokoro-decoder.cpp | 11 ++- tools/kokoro/src/kokoro-generator.cpp | 15 ++++ tools/kokoro/src/kokoro.cpp | 24 ++++-- 6 files changed, 198 insertions(+), 14 deletions(-) create mode 100644 tools/kokoro/include/kokoro-profile.h diff --git a/tools/kokoro/CMakeLists.txt b/tools/kokoro/CMakeLists.txt index dee882dd1..514bde879 100644 --- a/tools/kokoro/CMakeLists.txt +++ b/tools/kokoro/CMakeLists.txt @@ -48,6 +48,14 @@ endif() target_compile_features(kokoro_lib PUBLIC cxx_std_17) +# kokoro-layers.h routes the Conv1d/ConvTranspose1d/Linear/LSTM hot loops +# through the Accelerate BLAS (im2col + sgemm/sgemv) on Apple platforms — +# the scalar loops are ~2 orders of magnitude too slow for real-time TTS. +if(APPLE) + find_library(KOKORO_ACCELERATE_FRAMEWORK Accelerate REQUIRED) + target_link_libraries(kokoro_lib PUBLIC ${KOKORO_ACCELERATE_FRAMEWORK}) +endif() + # Real G2P via libespeak-ng. When present, kokoro_phonemize() drives # espeak_TextToPhonemes() (en-us IPA) and maps codepoints to Kokoro vocab ids, # reproducing the reference token sequence. When absent, the build falls back diff --git a/tools/kokoro/include/kokoro-layers.h b/tools/kokoro/include/kokoro-layers.h index 52212bca4..6f5d91d32 100644 --- a/tools/kokoro/include/kokoro-layers.h +++ b/tools/kokoro/include/kokoro-layers.h @@ -1,10 +1,17 @@ // SPDX-License-Identifier: MIT // // kokoro-layers.h — small neural-net layer primitives used by the Kokoro -// predictor + decoder forwards. Plain CPU/scalar implementations — no GGML -// graph, no backend dispatch. The forward pass is line-for-line matched to -// the upstream `kokoro` Python package (modules.py + istftnet.py) so the -// trained weights produce numerically equivalent output. +// predictor + decoder forwards. No GGML graph / backend dispatch: the forward +// pass is line-for-line matched to the upstream `kokoro` Python package +// (modules.py + istftnet.py) so the trained weights produce numerically +// equivalent output. +// +// On Apple platforms the O(C^2*K*T) hot loops (Conv1d, ConvTranspose1d, +// Linear, LSTM gate matvecs) are routed through the Accelerate BLAS +// (im2col + sgemm / sgemv, AMX-backed) — same math, same fp32 accumulation, +// ~2 orders of magnitude faster than the scalar loops on M-series. Every +// function keeps the portable scalar path as the non-Apple fallback and as +// the readable reference. // // Tensor convention (matches PyTorch Conv1d / Linear): // - 1D feature maps: `[C, T]` (channel-major, time-minor). @@ -22,6 +29,14 @@ #include #include +#if defined(__APPLE__) +#define KOKORO_USE_ACCELERATE 1 +#if !defined(ACCELERATE_NEW_LAPACK) +#define ACCELERATE_NEW_LAPACK +#endif +#include +#endif + namespace eliza_kokoro { // ================================================================= @@ -50,12 +65,22 @@ inline void linear_forward( const float * x, int in_dim, const float * W, const float * b, int out_dim, float * y) { +#if defined(KOKORO_USE_ACCELERATE) + float beta = 0.0f; + if (b) { + std::memcpy(y, b, sizeof(float) * (size_t) out_dim); + beta = 1.0f; + } + cblas_sgemv(CblasRowMajor, CblasNoTrans, out_dim, in_dim, + 1.0f, W, in_dim, x, 1, beta, y, 1); +#else for (int i = 0; i < out_dim; ++i) { float acc = b ? b[i] : 0.0f; const float * w = W + i * in_dim; for (int j = 0; j < in_dim; ++j) acc += w[j] * x[j]; y[i] = acc; } +#endif } // ================================================================= @@ -70,6 +95,46 @@ inline void conv1d_forward( const float * W, const float * b, int Cout, int K, int stride, int pad, int dilation, float * y, int T_out) { +#if defined(KOKORO_USE_ACCELERATE) + // im2col + one SGEMM. The PyTorch weight layout [Cout, Cin, K] is already + // a row-major [Cout, Cin*K] GEMM operand; the column matrix is + // [Cin*K, T_out] with zero padding baked in, so + // y[Cout, T_out] = W · cols (+ bias). + std::vector cols((size_t) Cin * K * T_out, 0.0f); + for (int ci = 0; ci < Cin; ++ci) { + const float * xi = x + (size_t) ci * T; + for (int k = 0; k < K; ++k) { + float * col = cols.data() + ((size_t) ci * K + k) * T_out; + const int t_off = k * dilation - pad; // ti = to*stride + t_off + // valid `to` range keeping ti within [0, T) + int to_lo = 0; + if (t_off < 0) to_lo = (-t_off + stride - 1) / stride; + int to_hi = (T - 1 - t_off) / stride; // inclusive + if (to_hi > T_out - 1) to_hi = T_out - 1; + if (to_lo > to_hi) continue; + if (stride == 1) { + std::memcpy(col + to_lo, xi + (to_lo + t_off), + sizeof(float) * (size_t) (to_hi - to_lo + 1)); + } else { + for (int to = to_lo; to <= to_hi; ++to) { + col[to] = xi[to * stride + t_off]; + } + } + } + } + float beta = 0.0f; + if (b) { + for (int co = 0; co < Cout; ++co) { + float * yc = y + (size_t) co * T_out; + for (int to = 0; to < T_out; ++to) yc[to] = b[co]; + } + beta = 1.0f; + } + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, + Cout, T_out, Cin * K, + 1.0f, W, Cin * K, cols.data(), T_out, + beta, y, T_out); +#else for (int co = 0; co < Cout; ++co) { const float bias_v = b ? b[co] : 0.0f; for (int to = 0; to < T_out; ++to) { @@ -86,6 +151,7 @@ inline void conv1d_forward( y[(size_t)co * T_out + to] = acc; } } +#endif } // ================================================================= @@ -158,6 +224,28 @@ inline void convtranspose1d_forward( for (int to = 0; to < T_out; ++to) yc[to] += b[co]; } } +#if defined(KOKORO_USE_ACCELERATE) + // One SGEMM + col2im scatter. Viewing W [Cin, Cout, K] as a row-major + // [Cin, Cout*K] operand: + // tmp[Cout*K, T] = W^T · x[Cin, T] + // then tmp[co*K + k, t_in] scatters into y[co, t_in*stride - pad + k]. + std::vector tmp((size_t) Cout * K * T); + cblas_sgemm(CblasRowMajor, CblasTrans, CblasNoTrans, + Cout * K, T, Cin, + 1.0f, W, Cout * K, x, T, + 0.0f, tmp.data(), T); + for (int co = 0; co < Cout; ++co) { + float * yc = y + (size_t) co * T_out; + for (int k = 0; k < K; ++k) { + const float * row = tmp.data() + ((size_t) co * K + k) * T; + const int t_off = k - pad; // to = t_in*stride + t_off + for (int t_in = 0; t_in < T; ++t_in) { + const int to = t_in * stride + t_off; + if (to >= 0 && to < T_out) yc[to] += row[t_in]; + } + } + } +#else for (int ci = 0; ci < Cin; ++ci) { const float * xi = x + (size_t)ci * T; for (int t_in = 0; t_in < T; ++t_in) { @@ -173,6 +261,7 @@ inline void convtranspose1d_forward( } } } +#endif } // Depthwise transpose (groups == Cin == Cout). Used by AdainResBlk1d pool. @@ -263,6 +352,15 @@ inline void lstm_cell_step( float * h_out, float * c_out, float * gates_scratch /* [4H] */ ) { // gates = W_ih @ x + b_ih + W_hh @ h_prev + b_hh +#if defined(KOKORO_USE_ACCELERATE) + for (int g = 0; g < 4 * H; ++g) { + gates_scratch[g] = (b_ih ? b_ih[g] : 0.0f) + (b_hh ? b_hh[g] : 0.0f); + } + cblas_sgemv(CblasRowMajor, CblasNoTrans, 4 * H, I, + 1.0f, W_ih, I, x, 1, 1.0f, gates_scratch, 1); + cblas_sgemv(CblasRowMajor, CblasNoTrans, 4 * H, H, + 1.0f, W_hh, H, h_prev, 1, 1.0f, gates_scratch, 1); +#else for (int g = 0; g < 4 * H; ++g) { float acc = (b_ih ? b_ih[g] : 0.0f) + (b_hh ? b_hh[g] : 0.0f); const float * w_ih_row = W_ih + (size_t)g * I; @@ -271,6 +369,7 @@ inline void lstm_cell_step( for (int j = 0; j < H; ++j) acc += w_hh_row[j] * h_prev[j]; gates_scratch[g] = acc; } +#endif // Split into i, f, g, o; apply sigmoid/tanh. auto sigmoid = [](float v){ return 1.0f / (1.0f + std::exp(-v)); }; for (int j = 0; j < H; ++j) { diff --git a/tools/kokoro/include/kokoro-profile.h b/tools/kokoro/include/kokoro-profile.h new file mode 100644 index 000000000..e7174f594 --- /dev/null +++ b/tools/kokoro/include/kokoro-profile.h @@ -0,0 +1,47 @@ +// SPDX-License-Identifier: MIT +// +// kokoro-profile.h — opt-in wall-clock phase timing for the Kokoro synthesis +// pipeline. Enabled by setting KOKORO_PROFILE=1 in the environment; otherwise +// zero output (the steady_clock reads are negligible next to the phases they +// wrap). Used to attribute RTF between the predictor, decoder front, and the +// iSTFTNet generator stages. + +#pragma once + +#include +#include +#include + +namespace eliza_kokoro { + +inline bool kokoro_profile_enabled() { + static const bool on = [] { + const char * v = std::getenv("KOKORO_PROFILE"); + return v && *v && *v != '0'; + }(); + return on; +} + +// RAII phase timer: logs "