chantier B TTS #8: KV cache CP -40% par frame, bit-exact

cp_inference: nouvelle voie cp_predict_cached (1 prefill 2 tokens + 14 decode 1 token via KV cache f32 persistant) à côté de cp_predict (oracle recompute conservé pour A/B). cp_forward_cached_step écrit le K/V post-RoPE dans s.K_cache/V_cache via ggml_view_3d + ggml_cpy, et lit le full cache [0..T_full) pour l'attention. Ordre topo garanti par build_forward_expand(cpy) AVANT build_forward_expand(x). tts_pipeline: switch KZTTS_CP_CACHE=0/1 (défaut 0 = oracle bit-exact). build_tts_pipeline.sh: rebuild reproductible (CMakeLists ne couvrait que libkazeia_engine.so). Mesure tablette Pad3 (talker_f32+tts_dump, CPU 6t, seed=42, 31 frames): oracle : CP 227.9 ms/frame, total 12.90s, RTF 4.99 cached : CP 137.8 ms/frame, total 8.91s, RTF 3.45 codes_ref.bin == codes_cache.bin (cmp), out_ref.wav md5 == out_cache.wav Speedup CP ×1.65, pas le ×8 espéré : le forward est BW-bound (~100 MB de poids f32 streamés par sub-forward, on lit toujours 15 fois). Le gain est sur la part compute (sum L=2..16 -> 16 token-fwd au lieu de 135). Pour aller plus loin : (a) cp_load en f16 conservé (×2 BW), (b) fusion graphe multi-step. Decoder RTF 1.37 reste aussi à attaquer. Bit-exactness : f32 mul_mat stable entre L=k+1 et L=1 à threads=6, les sommations donnent les mêmes bits. Math identique, ordre identique de fait. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-28 19:51:22 +02:00 · 2026-05-28 19:51:22 +02:00 · 7626131d50
parent d20f2b297b
commit 7626131d50
4 changed files with 273 additions and 5 deletions
--- a/dist/build_tts_pipeline.sh
+++ 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"
--- 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<float> 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<float>& 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<float>& 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<ggml_tensor*> 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<float> hn;
+    std::vector<float> 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<float> 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;
+    }
+}
--- 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 <cstdint>
@ -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<std::string, ggml_tensor*> 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);
+
--- 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]);