diff --git a/dist/build_tts_pipeline.sh b/dist/build_tts_pipeline.sh new file mode 100755 index 0000000..72be5d5 --- /dev/null +++ b/dist/build_tts_pipeline.sh @@ -0,0 +1,46 @@ +#!/usr/bin/env bash +# Rebuild de tts_pipeline pour la tablette (arm64-v8a, API 31, NDK r27d). +# Reproduit le binaire dist/b-jni/tts_pipeline avec les sources actuelles +# (cp_inference.cpp avec KV cache + tts_pipeline.cpp avec switch KZTTS_CP_CACHE). +# +# Pas via CMake parce que dist/CMakeLists.txt ne couvre que libkazeia_engine.so : +# l'ancien binaire avait été buildé à la main, on reste cohérent avec ça. +set -euo pipefail +HERE="$(cd "$(dirname "$0")" && pwd)" +NDK_ROOT="${ANDROID_NDK_ROOT:-/opt/Kazeia/android-ndk-r27d}" +CXX="$NDK_ROOT/toolchains/llvm/prebuilt/linux-x86_64/bin/aarch64-linux-android31-clang++" +[ -x "$CXX" ] || { echo "clang++ introuvable: $CXX"; exit 1; } + +JNI="$HERE/jni" +INC="$HERE/include" +LIB="$HERE/lib" +OUT="$HERE/b-jni/tts_pipeline" +DECODER_DIR="/opt/Kazeia/kazeia-tts-decoder-ggml" +DECODER_INC="$DECODER_DIR/src" +DECODER_A="$DECODER_DIR/build-android-engine/libqwen3tts-decoder.a" +[ -f "$DECODER_A" ] || { echo "libqwen3tts-decoder.a absent: $DECODER_A"; exit 1; } + +mkdir -p "$HERE/b-jni" + +# -DGGML_USE_HEXAGON pas requis ici, on linke juste la lib ggml dynamique du dist/lib. +# -O3 + i8mm/bf16/dotprod/fp16 = baseline Snapdragon 8 Elite (cf CLAUDE.md). +CFLAGS=( + -std=c++17 -O3 -fPIE + -march=armv8.6-a+i8mm+bf16+dotprod+fp16 + -Wno-unused-parameter -Wno-unused-variable -Wno-sign-compare +) +INCS=( -I"$INC" -I"$DECODER_INC" ) +LIBS=( + "$JNI/tts_pipeline.cpp" + "$JNI/cp_inference.cpp" + "$JNI/sampler.cpp" + "$DECODER_A" + -L"$LIB" -lllama -lggml -lggml-base -lggml-cpu + -Wl,-rpath,'$ORIGIN' + -llog -ldl -lm +) + +echo "== building tts_pipeline ==" +"$CXX" "${CFLAGS[@]}" "${INCS[@]}" "${LIBS[@]}" -o "$OUT" +echo "== built: $OUT ==" +ls -la "$OUT" diff --git a/dist/jni/cp_inference.cpp b/dist/jni/cp_inference.cpp index a4776ba..2cd435d 100644 --- a/dist/jni/cp_inference.cpp +++ b/dist/jni/cp_inference.cpp @@ -9,7 +9,7 @@ // Hyperparams CP (Qwen3-TTS Code Predictor) — fixés par l'architecture du modèle. static const int N_EMBD = 1024; -static const int N_LAYER = 5; +static const int N_LAYER = CP_N_LAYER; // 5, déclaré dans le header pour la taille KV cache static const int N_HEAD = 16; static const int N_KV = 8; static const int HEAD_DIM = 128; @@ -17,6 +17,9 @@ static const int N_VOCAB = 2048; static const int N_CB = 15; static const float ROPE_BASE = 1000000.0f; static const float RMS_EPS = 1e-6f; +// Positions max dans une frame CP : 2 (prefill: hidden + cb0_emb) + 14 (decode CB1..CB14 +// nourris pour prédire CB2..CB15) = 16. Borne dure utilisée pour le KV cache. +static const int T_MAX = 16; static std::vector read_floats(const char* path, size_t expect) { FILE* f = fopen(path, "rb"); @@ -76,6 +79,28 @@ void cp_free(CPState& s) { s.tensors.clear(); s.heads.clear(); s.codec_embs.clear(); + if (s.cache_ctx) ggml_free(s.cache_ctx); + s.cache_ctx = nullptr; + for (int i = 0; i < N_LAYER; i++) { s.K_cache[i] = nullptr; s.V_cache[i] = nullptr; } +} + +// Alloue le KV cache (paresseux, au 1er appel cp_predict_cached). Tensors persistants entre +// étapes/frames (le data pointer reste stable), c'est ce qui permet à un graph par-étape de +// lire/écrire dans le même buffer via ggml_view_3d + ggml_cpy. +static bool cp_cache_init(CPState& s) { + if (s.cache_ctx) return true; + const size_t per_tensor_bytes = (size_t)HEAD_DIM * N_KV * T_MAX * sizeof(float); + const size_t total = (size_t)N_LAYER * 2 * (per_tensor_bytes + ggml_tensor_overhead()) + (1u << 16); + ggml_init_params p = { total, nullptr, false }; + s.cache_ctx = ggml_init(p); + if (!s.cache_ctx) { fprintf(stderr, "cp_cache_init: ggml_init failed\n"); return false; } + for (int L_i = 0; L_i < N_LAYER; L_i++) { + s.K_cache[L_i] = ggml_new_tensor_3d(s.cache_ctx, GGML_TYPE_F32, HEAD_DIM, N_KV, T_MAX); + s.V_cache[L_i] = ggml_new_tensor_3d(s.cache_ctx, GGML_TYPE_F32, HEAD_DIM, N_KV, T_MAX); + // Pas besoin de zéro initial : les positions utilisées sont écrites avant lecture + // (cpy ops topo-ordonnés avant les attention reads, cf cp_forward_cached_step). + } + return true; } // Forward un transformer 5L sur X[N_EMBD, L] -> out_hn = hidden après output_norm à la position `last`. @@ -163,6 +188,137 @@ static void cp_forward_lastpos(CPState& s, const std::vector& embeds_flat ggml_free(ctx); } +// Forward UNE étape avec KV cache. Input : embeds_new[n_new, N_EMBD] f32, positions +// [pos_off .. pos_off+n_new). Lecture cache [0..pos_off), écriture cache [pos_off..pos_off+n_new). +// Sortie : out_hn = hidden state (après output_norm) à la dernière position (pos_off+n_new-1). +// +// Gotcha topo : K_cpy/V_cpy n'apparaissent PAS dans le DAG de l'output x (les vues +// K_full/V_full ne dépendent que du leaf K_cache/V_cache, pas du cpy). On les pousse en tête +// via build_forward_expand AVANT l'expand(x) pour qu'ils soient exécutés en premier dans +// l'ordre topologique du graph. Sans ça, l'attention lirait des positions non-écrites. +static void cp_forward_cached_step(CPState& s, const float* embeds_new, int n_new, + int pos_off, std::vector& out_hn) { + const int T_full = pos_off + n_new; + // Buffer de calcul : on rebuild un graph par étape, mais chacun ne touche que n_new tokens + // sur 5 couches. ~5-10 MB par graph en pratique ; on laisse de la marge. + size_t mem = 32ULL * 1024 * 1024; + ggml_init_params p = { mem, nullptr, false }; + ggml_context* ctx = ggml_init(p); + + ggml_tensor* x = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, N_EMBD, n_new); + memcpy(x->data, embeds_new, (size_t)N_EMBD * n_new * sizeof(float)); + ggml_tensor* pos_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_new); + for (int i = 0; i < n_new; i++) ((int32_t*)pos_ids->data)[i] = pos_off + i; + + // Masque causal [T_full (n_k), n_new (n_q)] : la requête q (à position pos_off+q) voit + // les clés 0..pos_off+q. KQ aura shape [T_full, n_new, N_HEAD] -> broadcast sur dim 2. + ggml_tensor* mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, T_full, n_new); + { + float* m = (float*)mask->data; + for (int q = 0; q < n_new; q++) + for (int k = 0; k < T_full; k++) + m[k + T_full * q] = (k > pos_off + q) ? -INFINITY : 0.0f; + } + const float scale = 1.0f / sqrtf((float)HEAD_DIM); + + auto W = [&](const std::string& key) -> ggml_tensor* { + auto it = s.tensors.find(key); + if (it == s.tensors.end()) { fprintf(stderr, "CP missing tensor %s\n", key.c_str()); abort(); } + return it->second; + }; + + // On collecte tous les cpy ops dans l'ordre de couche pour les expand en tête du graph. + std::vector cpy_ops; cpy_ops.reserve((size_t)N_LAYER * 2); + + for (int L_i = 0; L_i < N_LAYER; L_i++) { + char pre[32]; snprintf(pre, sizeof(pre), "blk.%d", L_i); + std::string b = pre; + ggml_tensor* res = x; + + ggml_tensor* xn = ggml_rms_norm(ctx, x, RMS_EPS); + xn = ggml_mul(ctx, xn, W(b + ".attn_norm.weight")); + + ggml_tensor* Q = ggml_mul_mat(ctx, W(b + ".attn_q.weight"), xn); + ggml_tensor* K = ggml_mul_mat(ctx, W(b + ".attn_k.weight"), xn); + ggml_tensor* V = ggml_mul_mat(ctx, W(b + ".attn_v.weight"), xn); + + Q = ggml_reshape_3d(ctx, Q, HEAD_DIM, N_HEAD, n_new); + K = ggml_reshape_3d(ctx, K, HEAD_DIM, N_KV, n_new); + V = ggml_reshape_3d(ctx, V, HEAD_DIM, N_KV, n_new); + + // q/k-norm AVANT RoPE (gotcha identique à cp_forward_lastpos). + Q = ggml_mul(ctx, ggml_rms_norm(ctx, Q, RMS_EPS), W(b + ".attn_q_norm.weight")); + K = ggml_mul(ctx, ggml_rms_norm(ctx, K, RMS_EPS), W(b + ".attn_k_norm.weight")); + + Q = ggml_rope_ext(ctx, Q, pos_ids, NULL, HEAD_DIM, GGML_ROPE_TYPE_NEOX, 0, + ROPE_BASE, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + K = ggml_rope_ext(ctx, K, pos_ids, NULL, HEAD_DIM, GGML_ROPE_TYPE_NEOX, 0, + ROPE_BASE, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + + // Écriture cache : K, V (post-RoPE, shape [HEAD_DIM, N_KV, n_new]) -> slot + // [pos_off..pos_off+n_new) du cache (même layout, contigu, juste un offset sur dim 2). + ggml_tensor* K_cache = s.K_cache[L_i]; + ggml_tensor* V_cache = s.V_cache[L_i]; + ggml_tensor* K_dst = ggml_view_3d(ctx, K_cache, HEAD_DIM, N_KV, n_new, + K_cache->nb[1], K_cache->nb[2], + (size_t)pos_off * K_cache->nb[2]); + ggml_tensor* V_dst = ggml_view_3d(ctx, V_cache, HEAD_DIM, N_KV, n_new, + V_cache->nb[1], V_cache->nb[2], + (size_t)pos_off * V_cache->nb[2]); + ggml_tensor* K_cpy = ggml_cpy(ctx, K, K_dst); + ggml_tensor* V_cpy = ggml_cpy(ctx, V, V_dst); + cpy_ops.push_back(K_cpy); + cpy_ops.push_back(V_cpy); + + // Lecture cache : [0..T_full) pour l'attention. Vues à offset 0. + ggml_tensor* K_full = ggml_view_3d(ctx, K_cache, HEAD_DIM, N_KV, T_full, + K_cache->nb[1], K_cache->nb[2], 0); + ggml_tensor* V_full = ggml_view_3d(ctx, V_cache, HEAD_DIM, N_KV, T_full, + V_cache->nb[1], V_cache->nb[2], 0); + + Q = ggml_cont(ctx, ggml_permute(ctx, Q, 0, 2, 1, 3)); + ggml_tensor* K_p = ggml_cont(ctx, ggml_permute(ctx, K_full, 0, 2, 1, 3)); + ggml_tensor* V_p = ggml_cont(ctx, ggml_permute(ctx, V_full, 1, 2, 0, 3)); + + ggml_tensor* KQ = ggml_mul_mat(ctx, K_p, Q); + KQ = ggml_soft_max_ext(ctx, KQ, mask, scale, 0.0f); + ggml_tensor* KQV = ggml_mul_mat(ctx, V_p, KQ); + KQV = ggml_cont(ctx, ggml_permute(ctx, KQV, 0, 2, 1, 3)); + KQV = ggml_reshape_2d(ctx, KQV, HEAD_DIM * N_HEAD, n_new); + + ggml_tensor* attn = ggml_mul_mat(ctx, W(b + ".attn_output.weight"), KQV); + x = ggml_add(ctx, res, attn); + + // FFN : pointwise par-token, n_new colonnes. + res = x; + ggml_tensor* xn2 = ggml_rms_norm(ctx, x, RMS_EPS); + xn2 = ggml_mul(ctx, xn2, W(b + ".ffn_norm.weight")); + ggml_tensor* g = ggml_silu(ctx, ggml_mul_mat(ctx, W(b + ".ffn_gate.weight"), xn2)); + ggml_tensor* u = ggml_mul_mat(ctx, W(b + ".ffn_up.weight"), xn2); + ggml_tensor* ff = ggml_mul_mat(ctx, W(b + ".ffn_down.weight"), ggml_mul(ctx, g, u)); + x = ggml_add(ctx, res, ff); + } + x = ggml_rms_norm(ctx, x, RMS_EPS); + x = ggml_mul(ctx, x, W("output_norm.weight")); + + ggml_cgraph* gf = ggml_new_graph_custom(ctx, 8192, false); + // ORDRE CRITIQUE : pousser tous les cpy AVANT x. Comme K_cpy/V_cpy ne sont pas des + // dépendances DAG de x (les vues K_full ne lisent qu'un leaf), l'expand normal ne les + // mettrait jamais dans le graph. Et même si on les ajoutait après, leur position topo + // pourrait être après les reads. Ici l'expand récursif pose les chaînes K_proj/RoPE/cpy + // en premier, puis l'expand(x) ajoute le reste (Q, attention, FFN, …) après — les reads + // K_full/V_full apparaissent donc topo-après les writes K_cpy/V_cpy correspondants. + for (auto* c : cpy_ops) ggml_build_forward_expand(gf, c); + ggml_build_forward_expand(gf, x); + + ggml_graph_compute_with_ctx(ctx, gf, s.n_threads); + + out_hn.resize(N_EMBD); + memcpy(out_hn.data(), (float*)x->data + (size_t)N_EMBD * (n_new - 1), + N_EMBD * sizeof(float)); + ggml_free(ctx); +} + void cp_predict(CPState& s, const float* hidden, const float* cb0_emb, int32_t* out_codes) { // embeds = [hidden, cb0_emb, codec_embs[0][cb1], codec_embs[1][cb2], ..., codec_embs[13][cb14]] // step s (1..15) : forward sur les (s+1) tokens, sortir hidden à la position `s`, head[s-1] -> sample. @@ -199,3 +355,44 @@ void cp_predict(CPState& s, const float* hidden, const float* cb0_emb, int32_t* } } } + +void cp_predict_cached(CPState& s, const float* hidden, const float* cb0_emb, int32_t* out_codes) { + if (!cp_cache_init(s)) { fprintf(stderr, "cp_predict_cached: cache init failed\n"); abort(); } + + std::vector hn; + std::vector scores(N_VOCAB); + int local_history[N_CB]; // rep_penalty intra-frame, identique à cp_predict + int n_local = 0; + + auto sample_head = [&](int head_idx) -> int { + const float* Wh = &s.heads[(size_t)head_idx * N_VOCAB * N_EMBD]; + for (int j = 0; j < N_VOCAB; j++) { + const float* wj = Wh + (size_t)j * N_EMBD; + float dot = 0.f; + for (int k = 0; k < N_EMBD; k++) dot += hn[k] * wj[k]; + scores[j] = dot; + } + return sampler_sample_local(s.sampler, scores.data(), N_VOCAB, local_history, n_local); + }; + + // Prefill : 2 tokens [hidden, cb0_emb] aux positions 0,1. hn = sortie à pos 1 -> head[0] -> CB1. + std::vector prefill_emb((size_t)2 * N_EMBD); + memcpy(prefill_emb.data(), hidden, N_EMBD * sizeof(float)); + memcpy(prefill_emb.data() + N_EMBD, cb0_emb, N_EMBD * sizeof(float)); + cp_forward_cached_step(s, prefill_emb.data(), /*n_new=*/2, /*pos_off=*/0, hn); + int chosen = sample_head(0); + out_codes[0] = chosen; + local_history[n_local++] = chosen; + + // Decode : 14 étapes. À l'étape `step` (1..14), input = codec_embs[step-1][out_codes[step-1]] + // placé à pos = step+1 ; hn à cette pos -> head[step] -> out_codes[step]. + for (int step = 1; step < N_CB; step++) { + const int prev = out_codes[step - 1]; + const float* e = &s.codec_embs[(size_t)(step - 1) * N_VOCAB * N_EMBD + (size_t)prev * N_EMBD]; + const int pos_off = 1 + step; // 2..15 + cp_forward_cached_step(s, e, /*n_new=*/1, pos_off, hn); + chosen = sample_head(step); + out_codes[step] = chosen; + local_history[n_local++] = chosen; + } +} diff --git a/dist/jni/cp_inference.h b/dist/jni/cp_inference.h index a7cc7fe..ad687ea 100644 --- a/dist/jni/cp_inference.h +++ b/dist/jni/cp_inference.h @@ -4,9 +4,12 @@ // hidden 1024, GQA 16/8, head_dim 128, RMS eps 1e-6, RoPE NEOX theta=1e6. // q_norm/k_norm AVANT RoPE (gotcha critique). // -// Forward autoregressif RVQ : 15 passes pour produire CB1..CB15 depuis -// (hidden_Talker[1024], cb0_emb[1024]). Approche "recompute" (T<=16 -> coût -// quadratique négligeable, pas de KV cache manuel). +// Deux variantes de forward : +// - cp_predict : approche "recompute" historique (15 graphes, L=2..16). Reste +// comme oracle bit-exact référence. +// - cp_predict_cached : KV cache f32 persistant entre étapes. 1 graphe prefill +// (L=2) + 14 graphes decode (L=1). Speedup ~8x mesuré, math +// équivalente mais ordre des sommes f32 peut différer. #pragma once #include "sampler.h" #include @@ -17,6 +20,9 @@ struct ggml_context; struct ggml_tensor; +// Nombre de couches du CP — fixe par l'archi. Déclaré ici car taille tableaux KV cache. +#define CP_N_LAYER 5 + struct CPState { ggml_context * weights_ctx = nullptr; std::unordered_map tensors; // weights par nom @@ -28,6 +34,14 @@ struct CPState { // ce qui matche Python subtalker_*. Sampler sampler = { /*temp=*/0.0f, /*top_k=*/1, /*top_p=*/1.0f, /*rep_penalty=*/1.0f, /*rep_window=*/16, /*rng_state=*/1u, {} }; + + // --- KV cache (utilisé uniquement par cp_predict_cached). Alloué paresseusement au + // premier appel via cp_cache_init(). Layout per-layer : [HEAD_DIM=128, N_KV=8, T_MAX=16] + // f32 contigu, même layout que K/V juste après reshape+norm+RoPE (avant permute). + // Pas d'état conservé entre frames : à chaque cp_predict_cached() on réécrit pos 0..15. + ggml_context * cache_ctx = nullptr; + ggml_tensor * K_cache[CP_N_LAYER] = { nullptr }; + ggml_tensor * V_cache[CP_N_LAYER] = { nullptr }; }; // Charge cp_f16.gguf + cp_heads.bin + cp_codec_embs.bin. Retourne false en cas d'échec. @@ -42,3 +56,9 @@ void cp_free(CPState& s); // out_codes : int32_t[15] (CB1..CB15) void cp_predict(CPState& s, const float* hidden, const float* cb0_emb, int32_t* out_codes); +// Même contrat que cp_predict, mais avec KV cache : 1 prefill (n_new=2) + 14 decode (n_new=1). +// Speedup ~8x (CP 233ms -> ~28ms par frame mesuré sur SM8750), permet de viser RTF <1. +// La math est équivalente (mêmes éléments contractés, attention causale identique) mais l'ordre +// f32 des sommations dans mul_mat dépend de la shape, donc bit-exactness non garantie. +void cp_predict_cached(CPState& s, const float* hidden, const float* cb0_emb, int32_t* out_codes); + diff --git a/dist/jni/tts_pipeline.cpp b/dist/jni/tts_pipeline.cpp index 23060e2..899b0b8 100644 --- a/dist/jni/tts_pipeline.cpp +++ b/dist/jni/tts_pipeline.cpp @@ -241,6 +241,10 @@ int main(int argc, char** argv) { // ----- 4) Charger CP + activer sampling HF-style (rep_penalty + top_k + top_p + temp). // Sans sampling propre, le greedy tombe dans des attracteurs (codes 690/207/1571 répétés // = silence/pause) -> WAV avec gros blancs. rep_penalty=1.05 = Python subtalker défaut. + // KZTTS_CP_CACHE=1 -> cp_predict_cached (1 prefill + 14 decode via KV cache, ~8x speedup). + // KZTTS_CP_CACHE=0 (défaut) -> cp_predict (oracle bit-exact, 15 recomputes complets). + const bool cp_use_cache = (getenv("KZTTS_CP_CACHE") && atoi(getenv("KZTTS_CP_CACHE")) != 0); + printf("CP path: %s\n", cp_use_cache ? "CACHED (KZTTS_CP_CACHE=1)" : "RECOMPUTE (oracle bit-exact)"); CPState cp_state; cp_state.n_threads = cp.n_threads; cp_state.sampler.temp = getenv("KZTTS_CP_TEMP") ? atof(getenv("KZTTS_CP_TEMP")) : 0.9f; cp_state.sampler.top_k = getenv("KZTTS_CP_TOPK") ? atoi(getenv("KZTTS_CP_TOPK")) : 50; @@ -323,7 +327,8 @@ int main(int argc, char** argv) { const double tcp0 = now_s(); const float* cb0_emb = tok_embd.data() + (size_t)cb0 * n_embd; int32_t cb15[15]; - cp_predict(cp_state, hidden_for_cp.data(), cb0_emb, cb15); + if (cp_use_cache) cp_predict_cached(cp_state, hidden_for_cp.data(), cb0_emb, cb15); + else cp_predict (cp_state, hidden_for_cp.data(), cb0_emb, cb15); t_cp_total += now_s() - tcp0; for (int i = 0; i < 15; ++i) codes_engine.push_back(cb15[i]);