diff --git a/dist/jni/kazeia_engine_jni.cpp b/dist/jni/kazeia_engine_jni.cpp index 8abc448..68cfee6 100644 --- a/dist/jni/kazeia_engine_jni.cpp +++ b/dist/jni/kazeia_engine_jni.cpp @@ -215,16 +215,36 @@ Java_com_kazeia_llm_EngineJni_resetEmbeds(JNIEnv*, jobject, jlong h) { // Helper interne : décode un batch d'embeddings, ne demande la sortie que pour la // dernière position (économie KV + logits), récupère hidden[n_embd] dans out_hidden // si fourni. Avance k->pos_embd de T. Renvoie 0 si OK. +// +// ⚠ M-RoPE (talker Qwen3-TTS) : llm_graph_input_pos::set_input attend des positions +// SOIT comme tokens et fait la conversion 1D→4D, SOIT directement T*4 entiers en +// embeds-mode (la branche else copie raw). Si on passe T positions seulement, ggml +// lit out-of-bounds. -> ici on alloue T*4 et duplique pos sur les 3 premiers axes +// (4ème = 0) comme le framework le fait pour les tokens. static int decode_embd_batch(KEngine* k, const float* embd, int T, float* out_hidden) { llama_context* ctx = k->c_h ? k->c_h : k->c_c; llama_set_embeddings(ctx, true); - std::vector pos (T); + auto rt = llama_model_rope_type(k->m_h ? k->m_h : k->m_c); + const int npe = (rt == LLAMA_ROPE_TYPE_MROPE || rt == LLAMA_ROPE_TYPE_IMROPE) ? 4 : 1; + + std::vector pos (T * npe, 0); std::vector nsd (T, 1); std::vector sid0(T, 0); std::vector sids(T); std::vector lg (T, 0); - for (int i = 0; i < T; ++i) { pos[i] = k->pos_embd + i; sids[i] = &sid0[i]; } + for (int i = 0; i < T; ++i) { + const llama_pos p = k->pos_embd + i; + if (npe == 4) { + pos[ i] = p; // axis t + pos[ T + i] = p; // axis h + pos[2 * T + i] = p; // axis w + pos[3 * T + i] = 0; // axis e (zéro pour text-only) + } else { + pos[i] = p; + } + sids[i] = &sid0[i]; + } lg[T-1] = 1; // n'output logits/embeddings que pour la dernière position llama_batch b{}; diff --git a/dist/jni/test_talker_replay.cpp b/dist/jni/test_talker_replay.cpp new file mode 100644 index 0000000..ca882cd --- /dev/null +++ b/dist/jni/test_talker_replay.cpp @@ -0,0 +1,200 @@ +// Test bit-exact du Talker engine vs Python. +// +// Stratégie = teacher-forcing : on rejoue les MÊMES inputs_embeds (prefill + +// step) que Python a effectivement utilisés, et on compare logits/hidden à +// chaque step. Pas de sampling côté engine — on consomme step_inputs[i] tel +// qu'observé côté Python. +// +// Si match (RMSE/cos serrés sur logits ET hidden) → Patch 2 (M-RoPE IMROPE) + +// Patch 1 (embeds API) sont corrects. Sinon : où ça diverge donne la piste. +// +// Fixtures (dans /): +// manifest.txt (T_prefill, N_steps, n_embd, vocab) +// talker_prefill_embeds.bin float32 [T_prefill * n_embd] +// talker_step_inputs.bin float32 [N_steps * n_embd] +// talker_step_logits.bin float32 [N_steps * vocab] +// talker_step_hidden.bin float32 [N_steps * n_embd] +// talker_last_hidden_pf.bin float32 [n_embd] (past_hidden après prefill) +// +// Usage : test_talker_replay [cpu|htp] +#include +#include +#include +#include +#include +#include +#include +#include +#include "llama.h" +#include "ggml-backend.h" + +static ggml_backend_dev_t find_htp() { + for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { + auto d = ggml_backend_dev_get(i); + if (!strcmp(ggml_backend_dev_name(d), "HTP0")) return d; + } + return nullptr; +} + +static std::vector read_f32(const std::string& p, size_t n_expected) { + std::ifstream f(p, std::ios::binary | std::ios::ate); + if (!f) { fprintf(stderr, "open %s\n", p.c_str()); exit(1); } + size_t n = (size_t)f.tellg() / sizeof(float); + if (n != n_expected) { fprintf(stderr, "%s: %zu f32, attendu %zu\n", p.c_str(), n, n_expected); exit(1); } + f.seekg(0); + std::vector v(n); + f.read((char*)v.data(), n * sizeof(float)); + return v; +} + +// RMSE + cos + max_abs entre deux vecteurs de même taille. +struct Cmp { float rmse, cos, mae, max_abs_diff; int argmax_a, argmax_b; }; +static Cmp compare(const float* a, const float* b, int n) { + double se = 0, dot = 0, na = 0, nb = 0, mae = 0; + float mad = 0; + int am = 0, bm = 0; + float av = a[0], bv = b[0]; + for (int i = 0; i < n; ++i) { + double d = a[i] - b[i]; + se += d * d; + mae += std::abs(d); + if (std::abs((float)d) > mad) mad = std::abs((float)d); + dot += (double)a[i] * b[i]; + na += (double)a[i] * a[i]; + nb += (double)b[i] * b[i]; + if (a[i] > av) { av = a[i]; am = i; } + if (b[i] > bv) { bv = b[i]; bm = i; } + } + return Cmp{ + (float)std::sqrt(se / n), + (float)(dot / (std::sqrt(na) * std::sqrt(nb) + 1e-30)), + (float)(mae / n), + mad, + am, bm, + }; +} + +static void print_cmp(const char* tag, int step, const Cmp& c, int extra_argmax = -1) { + printf(" %s[%2d] rmse=%.4e cos=%.6f mae=%.4e max=%.4e argmax(eng,py", tag, step, c.rmse, c.cos, c.mae, c.max_abs_diff); + if (extra_argmax >= 0) printf(",dump"); + printf(")=(%d,%d", c.argmax_a, c.argmax_b); + if (extra_argmax >= 0) printf(",%d", extra_argmax); + printf(")%s\n", c.argmax_a == c.argmax_b ? " ✓" : ""); +} + +int main(int argc, char** argv) { + if (argc < 3) { printf("usage: %s [cpu|htp]\n", argv[0]); return 1; } + const char* gguf = argv[1]; + std::string D = argv[2]; if (D.back() != '/') D += '/'; + bool force_cpu = (argc >= 4 && !strcmp(argv[3], "cpu")); + + // manifest + int T_prefill = 0, N_steps = 0, n_embd = 1024, n_vocab = 3072; + { + std::ifstream f(D + "manifest.txt"); if (!f) { fprintf(stderr, "no manifest\n"); return 1; } + std::string line; + while (std::getline(f, line)) { + if (line.rfind("T_prefill:", 0) == 0) T_prefill = atoi(line.c_str() + 10); + else if (line.rfind("N_steps:", 0) == 0) N_steps = atoi(line.c_str() + 8); + else if (line.rfind("n_embd:", 0) == 0) n_embd = atoi(line.c_str() + 7); + else if (line.rfind("vocab:", 0) == 0) n_vocab = atoi(line.c_str() + 6); + } + printf("manifest: T_prefill=%d N_steps=%d n_embd=%d n_vocab=%d\n", T_prefill, N_steps, n_embd, n_vocab); + } + + // load fixtures + auto prefill_embeds = read_f32(D + "talker_prefill_embeds.bin", (size_t)T_prefill * n_embd); + auto step_inputs = read_f32(D + "talker_step_inputs.bin", (size_t)N_steps * n_embd); + auto step_logits_py = read_f32(D + "talker_step_logits.bin", (size_t)N_steps * n_vocab); + auto step_hidden_py = read_f32(D + "talker_step_hidden.bin", (size_t)N_steps * n_embd); + auto hidden_pf_py = read_f32(D + "talker_last_hidden_pf.bin", (size_t)n_embd); + printf("fixtures loaded (Python)\n"); + + // engine + setenv("GGML_HEXAGON_USE_HMX", "0", 1); + llama_backend_init(); + auto mp = llama_model_default_params(); + ggml_backend_dev_t devs[2] = { force_cpu ? nullptr : find_htp(), nullptr }; + if (devs[0]) { mp.n_gpu_layers = 99; mp.devices = devs; printf("device: HTP0\n"); } + else { mp.n_gpu_layers = 0; printf("device: CPU%s\n", force_cpu ? " (forcé)" : ""); } + auto m = llama_model_load_from_file(gguf, mp); + if (!m) { printf("model load FAILED\n"); return 1; } + + auto cp = llama_context_default_params(); + cp.n_ctx = std::max(512, T_prefill + N_steps + 16); + cp.n_batch = 1024; cp.n_threads = 4; + cp.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED; + cp.embeddings = true; + auto ctx = llama_init_from_model(m, cp); + if (!ctx) { printf("ctx FAILED\n"); return 1; } + printf("ctx ready (n_ctx=%d, n_embd=%d, n_vocab=%d)\n", cp.n_ctx, llama_model_n_embd(m), llama_vocab_n_tokens(llama_model_get_vocab(m))); + + // M-RoPE (talker) : 4 positions par token (t/h/w/e). Embeds-mode = pas de + // conversion auto par le framework, on fournit explicitement T*npe entiers. + auto rt = llama_model_rope_type(m); + const int npe = (rt == LLAMA_ROPE_TYPE_MROPE || rt == LLAMA_ROPE_TYPE_IMROPE) ? 4 : 1; + printf("rope_type=%d, n_pos_per_embd=%d\n", (int)rt, npe); + + // === PREFILL === + { + std::vector pos(T_prefill * npe, 0); + std::vector nsd(T_prefill, 1); + std::vector sid0(T_prefill, 0); + std::vector sids(T_prefill); + std::vector lg(T_prefill, 0); + for (int i = 0; i < T_prefill; ++i) { + if (npe == 4) { pos[i] = i; pos[T_prefill + i] = i; pos[2*T_prefill + i] = i; pos[3*T_prefill + i] = 0; } + else { pos[i] = i; } + sids[i] = &sid0[i]; + } + lg[T_prefill - 1] = 1; + llama_batch b{}; + b.n_tokens = T_prefill; b.embd = prefill_embeds.data(); + b.pos = pos.data(); b.n_seq_id = nsd.data(); b.seq_id = sids.data(); b.logits = lg.data(); + if (llama_decode(ctx, b) != 0) { printf("prefill FAILED\n"); return 1; } + + const float* h = llama_get_embeddings_ith(ctx, -1); + Cmp ch = compare(h, hidden_pf_py.data(), n_embd); + printf("PREFILL OK (T=%d)\n", T_prefill); + printf(" hidden_pf: rmse=%.4e cos=%.6f mae=%.4e max=%.4e\n", ch.rmse, ch.cos, ch.mae, ch.max_abs_diff); + } + + // === STEPS (teacher-forced) === + int n_ok_argmax = 0; + double sum_rmse_h = 0, sum_rmse_l = 0, sum_cos_h = 0, sum_cos_l = 0; + for (int s = 0; s < N_steps; ++s) { + llama_pos pos1[4] = {0,0,0,0}; + const llama_pos p = T_prefill + s; + if (npe == 4) { pos1[0] = p; pos1[1] = p; pos1[2] = p; pos1[3] = 0; } + else { pos1[0] = p; } + int32_t n = 1; + llama_seq_id sd = 0; llama_seq_id* sp = &sd; + int8_t l = 1; + llama_batch b{}; + b.n_tokens = 1; b.embd = step_inputs.data() + (size_t)s * n_embd; + b.pos = pos1; b.n_seq_id = &n; b.seq_id = &sp; b.logits = &l; + if (llama_decode(ctx, b) != 0) { printf("step %d FAILED\n", s); return 1; } + + const float* h_eng = llama_get_embeddings_ith(ctx, -1); + const float* lg_eng = llama_get_logits_ith(ctx, -1); + Cmp ch = compare(h_eng, step_hidden_py.data() + (size_t)s * n_embd, n_embd); + Cmp cl = compare(lg_eng, step_logits_py.data() + (size_t)s * n_vocab, n_vocab); + if (ch.argmax_a == ch.argmax_b) n_ok_argmax++; // matching argmax over hidden (just a curiosity) + sum_rmse_h += ch.rmse; sum_cos_h += ch.cos; + sum_rmse_l += cl.rmse; sum_cos_l += cl.cos; + // n'imprime que les 3 premiers + les 3 derniers steps (diff de pos / cumul KV) + if (s < 3 || s >= N_steps - 3) { + print_cmp("hidden", s, ch); + print_cmp("logits", s, cl); + } + else if (s == 3) printf(" ... (%d steps intermédiaires) ...\n", N_steps - 6); + } + printf("STEPS SUMMARY (N=%d):\n", N_steps); + printf(" hidden: mean_rmse=%.4e mean_cos=%.6f\n", sum_rmse_h/N_steps, sum_cos_h/N_steps); + printf(" logits: mean_rmse=%.4e mean_cos=%.6f (argmax-on-hidden-match=%d/%d, info only)\n", + sum_rmse_l/N_steps, sum_cos_l/N_steps, n_ok_argmax, N_steps); + + llama_free(ctx); + llama_model_free(m); + return 0; +} diff --git a/dist/lib/libkazeia_engine.so b/dist/lib/libkazeia_engine.so index 53d5295..6abfe1a 100755 Binary files a/dist/lib/libkazeia_engine.so and b/dist/lib/libkazeia_engine.so differ