From 9a0c63f1647471045d60ba9fd946ed7d541c4306 Mon Sep 17 00:00:00 2001 From: Richard Loyer Date: Fri, 29 May 2026 17:26:54 +0200 Subject: [PATCH] T2.1 CP heads NEON 16-way: CP -33%, pipeline -16%, bit-exact MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Le profile fin (KZTTS_CP_PROFILE=1) a révélé que ~75% du temps CP par frame est dans ggml_graph_compute (4-6 ms/sub-step × 15 = 80 ms) et ~22% dans le delta CPU host (~24 ms/frame). Ce delta était entièrement le sample_head : 2048 dot products de 1024 floats SCALAIRE × 15 heads/frame = 31.5M ops/frame. Implémentation NEON aarch64 16-way unroll (vfmaq_f32 × 4 accumulators) sur dot_neon_1024(). Gated par KZTTS_CP_NO_NEON=1 pour fallback debug. Mesure Pad3 (KZTTS_THREADS=6 GGML_NUM_THREADS=8, seed=42, Bonjour Kazeia) : scalar : CP 106.8 ms/frame, total 8.05s, RTF 2.93 NEON : CP 72.1 ms/frame, total 6.75s, RTF 2.45 (-33% CP, -16% pipe) WAV md5 identiques c3dd71cbc783013842a3dacedfefd758 (BIT-EXACT) Le scalaire compilé par clang -O3 -march=armv8.6-a n'a apparemment pas auto-vectorisé cette boucle malgré l'évidence (peut-être à cause des indices non-trivials avec head_idx). NEON explicite débloque immédiatement. Reste : decoder = 51% du total maintenant (3.34s sur 6.75s). BigVGAN ggml-cpu déjà NEON vectorisé donc moins de marge évidente côté host code. Cibles potentielles : text_projection 50ms one-shot au prefill, mais marginal sur pipeline total. Co-Authored-By: Claude Opus 4.7 (1M context) --- dist/jni/cp_inference.cpp | 63 +++++++++++++++++++++++++++++++++++---- 1 file changed, 58 insertions(+), 5 deletions(-) diff --git a/dist/jni/cp_inference.cpp b/dist/jni/cp_inference.cpp index 7633313..89c5961 100644 --- a/dist/jni/cp_inference.cpp +++ b/dist/jni/cp_inference.cpp @@ -8,6 +8,9 @@ #include #include #include +#if defined(__aarch64__) +# include +#endif // Hyperparams CP (Qwen3-TTS Code Predictor) — fixés par l'architecture du modèle. static const int N_EMBD = 1024; @@ -273,6 +276,27 @@ static bool cp_cache_init(CPState& s) { return true; } +// Dot product NEON 16-way unroll. N=1024 (N_EMBD) ne divisible par 16 garanti. +// Coût : 32M ops / frame partagés sur les 15 sub-steps = scalar 21 ms, NEON ~3 ms. +static inline float dot_neon_1024(const float* a, const float* b) { +#if defined(__aarch64__) + float32x4_t s0 = vdupq_n_f32(0.f), s1 = vdupq_n_f32(0.f); + float32x4_t s2 = vdupq_n_f32(0.f), s3 = vdupq_n_f32(0.f); + for (int i = 0; i < 1024; i += 16) { + s0 = vfmaq_f32(s0, vld1q_f32(a + i + 0), vld1q_f32(b + i + 0)); + s1 = vfmaq_f32(s1, vld1q_f32(a + i + 4), vld1q_f32(b + i + 4)); + s2 = vfmaq_f32(s2, vld1q_f32(a + i + 8), vld1q_f32(b + i + 8)); + s3 = vfmaq_f32(s3, vld1q_f32(a + i + 12), vld1q_f32(b + i + 12)); + } + float32x4_t s = vaddq_f32(vaddq_f32(s0, s1), vaddq_f32(s2, s3)); + return vaddvq_f32(s); +#else + float d = 0.f; + for (int i = 0; i < 1024; i++) d += a[i] * b[i]; + return d; +#endif +} + // Forward un transformer 5L sur X[N_EMBD, L] -> out_hn = hidden après output_norm à la position `last`. // Architecture identique à cp_runner.cpp (la référence bit-exact validée). static void cp_forward_lastpos(CPState& s, const std::vector& embeds_flat, int L, int last, @@ -370,11 +394,18 @@ static void cp_forward_cached_step(CPState& s, const float* embeds_new, int n_ne int pos_off, std::vector& out_hn) { const int T_full = pos_off + n_new; const bool use_sched = (s.sched != nullptr); + const bool prof_global = (getenv("KZTTS_CP_PROFILE") && atoi(getenv("KZTTS_CP_PROFILE")) != 0); + auto nows_fn = []() { + return std::chrono::duration( + std::chrono::steady_clock::now().time_since_epoch()).count(); + }; + double tg0 = prof_global ? nows_fn() : 0; // mem : 32 MB pour le path legacy (no_alloc=false, alloue data activations), // 4 MB en mode sched (juste métadonnées des tensors, allocation backend séparée). size_t mem = use_sched ? (4ULL * 1024 * 1024) : (32ULL * 1024 * 1024); ggml_init_params p = { mem, nullptr, /*no_alloc=*/use_sched }; ggml_context* ctx = ggml_init(p); + double tg_init = prof_global ? nows_fn() : 0; // Pré-calcul host des données d'inputs : pos_ids et mask. Mode legacy on les memcpy // direct dans le tensor data. Mode sched on les uploade via tensor_set après alloc. @@ -490,6 +521,7 @@ static void cp_forward_cached_step(CPState& s, const float* embeds_new, int n_ne // 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); + double tg_build = prof_global ? nows_fn() : 0; out_hn.resize(N_EMBD); if (use_sched) { @@ -534,9 +566,20 @@ static void cp_forward_cached_step(CPState& s, const float* embeds_new, int n_ne (t_c1-t_s1)*1000.0, (t_g1-t_c1)*1000.0); } } else { + double t0 = prof_global ? nows_fn() : 0; ggml_graph_compute_with_ctx(ctx, gf, s.n_threads); + double t1 = prof_global ? nows_fn() : 0; memcpy(out_hn.data(), (float*)x->data + (size_t)N_EMBD * (n_new - 1), N_EMBD * sizeof(float)); + if (prof_global) { + double t2 = nows_fn(); + fprintf(stderr, "cp_sub_cpu n_new=%d pos=%d init=%.2fms build=%.2fms compute=%.2fms memcpy+free=%.2fms\n", + n_new, pos_off, + (tg_init - tg0)*1000.0, + (tg_build - tg_init)*1000.0, + (t1 - t0)*1000.0, + (t2 - t1)*1000.0); + } } ggml_free(ctx); } @@ -586,13 +629,23 @@ void cp_predict_cached(CPState& s, const float* hidden, const float* cb0_emb, in int local_history[N_CB]; // rep_penalty intra-frame, identique à cp_predict int n_local = 0; + // sample_head : 2048 dot products de N_EMBD=1024 floats. KZTTS_CP_NO_NEON=1 force le path + // scalaire (debug). Par défaut NEON 16-way unroll, ~7× plus rapide sur Cortex-X4/A720. + const bool no_neon = (getenv("KZTTS_CP_NO_NEON") && atoi(getenv("KZTTS_CP_NO_NEON")) != 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; + if (no_neon) { + 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.data()[k] * wj[k]; + scores[j] = dot; + } + } else { + const float* hn_ptr = hn.data(); + for (int j = 0; j < N_VOCAB; j++) { + scores[j] = dot_neon_1024(hn_ptr, Wh + (size_t)j * N_EMBD); + } } return sampler_sample_local(s.sampler, scores.data(), N_VOCAB, local_history, n_local); };