277 lines
13 KiB
C++
277 lines
13 KiB
C++
// Orchestration ggml-side du Talker engine.
|
|
//
|
|
// Boucle decode : engine Talker fait son vrai forward + greedy sur CB0 ;
|
|
// CB1..15 teacher-forcés depuis codes_golden_NxCB.bin (CP non intégré encore).
|
|
// next_embed = Σ talker.tok_embd[CB0] + Σ cp_codec_embs[i-1, CB(i)] + tts_pad_embed.
|
|
// (En x_vector_only_mode, trailing_text_hidden == tts_pad_embed exactement, vérifié.)
|
|
//
|
|
// Sortie : codes_engine.bin [N, 16] int32 = ce que l'engine produit. À décoder
|
|
// ensuite avec qwen3tts-decoder-test pour obtenir un WAV.
|
|
//
|
|
// Validation parallèle : à chaque step, on compare next_embed engine-side vs
|
|
// talker_step_inputs[s] Python (qui a la même formule). RMSE ~0 attendu (lookups
|
|
// dans les mêmes tables, mêmes codes forcés, même sommation).
|
|
//
|
|
// Usage : tts_orchestrate <gguf> <dump_dir> <out_codes.bin> [cpu|htp] [max_steps]
|
|
#include <cstdio>
|
|
#include <cstdlib>
|
|
#include <cstring>
|
|
#include <cmath>
|
|
#include <cstdint>
|
|
#include <vector>
|
|
#include <string>
|
|
#include <fstream>
|
|
#include <chrono>
|
|
#include "llama.h"
|
|
#include "ggml-backend.h"
|
|
#include "cp_inference.h"
|
|
|
|
static double now_s() {
|
|
using clk = std::chrono::steady_clock;
|
|
return std::chrono::duration<double>(clk::now().time_since_epoch()).count();
|
|
}
|
|
|
|
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<float> 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<float> v(n);
|
|
f.read((char*)v.data(), n * sizeof(float));
|
|
return v;
|
|
}
|
|
static std::vector<int32_t> read_i32(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(int32_t);
|
|
if (n != n_expected) { fprintf(stderr, "%s: %zu i32, attendu %zu\n", p.c_str(), n, n_expected); exit(1); }
|
|
f.seekg(0);
|
|
std::vector<int32_t> v(n);
|
|
f.read((char*)v.data(), n * sizeof(int32_t));
|
|
return v;
|
|
}
|
|
|
|
int main(int argc, char** argv) {
|
|
if (argc < 4) { printf("usage: %s <talker_gguf> <dump_dir> <out_codes.bin> [cpu|htp] [max_steps]\n", argv[0]); return 1; }
|
|
const char* gguf = argv[1];
|
|
std::string D = argv[2]; if (D.back() != '/') D += '/';
|
|
const char* out_codes = argv[3];
|
|
bool force_cpu = (argc >= 5 && !strcmp(argv[4], "cpu"));
|
|
int max_steps_arg = (argc >= 6) ? atoi(argv[5]) : 0;
|
|
|
|
// CP runner : charge cp_f16.gguf + cp_heads.bin + cp_codec_embs.bin depuis dump_dir.
|
|
// Si l'un manque -> fallback teacher-forcing (mode legacy de tts_orchestrate).
|
|
CPState cp_state;
|
|
cp_state.n_threads = 4;
|
|
bool cp_ok = cp_load(cp_state,
|
|
(D + "cp_f16.gguf").c_str(),
|
|
(D + "cp_heads.bin").c_str(),
|
|
(D + "cp_codec_embs.bin").c_str());
|
|
printf("CP : %s\n", cp_ok ? "ENABLED" : "DISABLED (teacher-forced CB1..15)");
|
|
|
|
// manifest
|
|
int T_prefill = 0, N_steps_golden = 0, n_embd = 1024, n_vocab = 3072;
|
|
int N_codebooks = 16, cp_vocab = 2048, codec_eos_token_id = 2150;
|
|
{
|
|
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_golden = 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);
|
|
else if (line.rfind("N_codes:", 0) == 0) N_steps_golden = atoi(line.c_str() + 8);
|
|
else if (line.rfind("N_codebooks:", 0) == 0) N_codebooks = atoi(line.c_str() + 12);
|
|
else if (line.rfind("cp_vocab:", 0) == 0) cp_vocab = atoi(line.c_str() + 9);
|
|
else if (line.rfind("codec_eos_token_id:", 0)==0)codec_eos_token_id= atoi(line.c_str() + 19);
|
|
}
|
|
}
|
|
const int N = (max_steps_arg > 0) ? std::min(max_steps_arg, N_steps_golden) : N_steps_golden;
|
|
printf("manifest: T_prefill=%d N_steps_golden=%d N=%d n_embd=%d n_vocab=%d N_cb=%d cp_vocab=%d eos=%d\n",
|
|
T_prefill, N_steps_golden, N, n_embd, n_vocab, N_codebooks, cp_vocab, codec_eos_token_id);
|
|
|
|
// fixtures
|
|
auto prefill_embeds = read_f32(D + "talker_prefill_embeds.bin", (size_t)T_prefill * n_embd);
|
|
auto tts_pad = read_f32(D + "tts_pad_embed.bin", (size_t)n_embd);
|
|
auto tok_embd = read_f32(D + "talker_tok_embd.bin", (size_t)n_vocab * n_embd);
|
|
auto cp_codec_embs = read_f32(D + "cp_codec_embs.bin", (size_t)15 * cp_vocab * n_embd);
|
|
auto codes_golden = read_i32(D + "codes_golden_NxCB.bin", (size_t)N_steps_golden * N_codebooks);
|
|
// step_inputs sera comparé pour validation
|
|
auto step_inputs_py = read_f32(D + "talker_step_inputs.bin", (size_t)N_steps_golden * n_embd);
|
|
printf("fixtures loaded\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\n"); }
|
|
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 + 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; }
|
|
|
|
auto rt = llama_model_rope_type(m);
|
|
const int npe = (rt == LLAMA_ROPE_TYPE_MROPE || rt == LLAMA_ROPE_TYPE_IMROPE) ? 4 : 1;
|
|
|
|
// PREFILL
|
|
std::vector<llama_pos> pos(T_prefill * npe, 0);
|
|
std::vector<int32_t> nsd(T_prefill, 1);
|
|
std::vector<llama_seq_id> sid0(T_prefill, 0);
|
|
std::vector<llama_seq_id*> sids(T_prefill);
|
|
std::vector<int8_t> 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;
|
|
const double t_prefill0 = now_s();
|
|
{
|
|
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 double t_prefill = now_s() - t_prefill0;
|
|
printf("prefill OK (T=%d) en %.3f s\n", T_prefill, t_prefill);
|
|
|
|
// sample CB0 at prefill output (first step's "input_ids")
|
|
auto logits = llama_get_logits_ith(ctx, -1);
|
|
int cb0 = 0; float mx = logits[0];
|
|
for (int i = 1; i < n_vocab; ++i) if (logits[i] > mx) { mx = logits[i]; cb0 = i; }
|
|
printf("prefill argmax CB0 = %d (golden[0,0]=%d)%s\n", cb0, codes_golden[0],
|
|
cb0 == codes_golden[0] ? " ✓" : " (engine greedy ≠ python sample)");
|
|
|
|
// boucle decode
|
|
std::vector<int32_t> codes_engine(N * N_codebooks, 0);
|
|
int n_match_cb0 = 0, n_match_full = 0;
|
|
int n_eos = -1;
|
|
double t_compose = 0, t_decode = 0, t_cp = 0;
|
|
const double t_loop0 = now_s();
|
|
// capture du hidden state du Talker (pour passer au CP) = ce que renvoie
|
|
// llama_get_embeddings_ith(ctx,-1) après chaque decode. On l'a déjà au prefill (h ci-dessous).
|
|
const float* h_last = llama_get_embeddings_ith(ctx, -1);
|
|
std::vector<float> hidden_for_cp(n_embd);
|
|
if (h_last) memcpy(hidden_for_cp.data(), h_last, n_embd * sizeof(float));
|
|
|
|
for (int s = 0; s < N; ++s) {
|
|
// CB0 = engine greedy
|
|
codes_engine[s * N_codebooks + 0] = cb0;
|
|
// CB1..15 : soit CP réel, soit teacher-forced
|
|
const double tcp0 = now_s();
|
|
if (cp_ok) {
|
|
const float* cb0_emb = tok_embd.data() + (size_t)cb0 * n_embd;
|
|
int32_t out_cb15[15];
|
|
cp_predict(cp_state, hidden_for_cp.data(), cb0_emb, out_cb15);
|
|
for (int i = 1; i < N_codebooks; ++i) codes_engine[s * N_codebooks + i] = out_cb15[i - 1];
|
|
} else {
|
|
for (int i = 1; i < N_codebooks; ++i) {
|
|
codes_engine[s * N_codebooks + i] = codes_golden[s * N_codebooks + i];
|
|
}
|
|
}
|
|
t_cp += now_s() - tcp0;
|
|
if (cb0 == codes_golden[s * N_codebooks + 0]) n_match_cb0++;
|
|
{
|
|
int m = 0;
|
|
for (int i = 0; i < N_codebooks; ++i)
|
|
if (codes_engine[s * N_codebooks + i] == codes_golden[s * N_codebooks + i]) m++;
|
|
if (m == N_codebooks) n_match_full++;
|
|
}
|
|
if (cb0 == codec_eos_token_id && n_eos < 0) { n_eos = s; printf(" step %d: EOS atteint\n", s); }
|
|
|
|
// next_embed = tok_embd[cb0] + sum cp_codec_embs[i-1, CB(i)] + tts_pad
|
|
const double tc0 = now_s();
|
|
std::vector<float> next_embed(n_embd, 0.0f);
|
|
const float* e_cb0 = tok_embd.data() + (size_t)cb0 * n_embd;
|
|
for (int d = 0; d < n_embd; ++d) next_embed[d] = e_cb0[d];
|
|
for (int i = 1; i < N_codebooks; ++i) {
|
|
int code = codes_engine[s * N_codebooks + i];
|
|
const float* e = cp_codec_embs.data() + ((size_t)(i-1) * cp_vocab + code) * n_embd;
|
|
for (int d = 0; d < n_embd; ++d) next_embed[d] += e[d];
|
|
}
|
|
for (int d = 0; d < n_embd; ++d) next_embed[d] += tts_pad[d];
|
|
t_compose += now_s() - tc0;
|
|
|
|
// valid : next_embed (= input à injecter au decode step s) vs step_inputs_py[s]
|
|
// (= ce que Python a injecté au decode step s). Si codes_engine[s] == codes_golden[s]
|
|
// alors next_embed doit être strictement = step_inputs_py[s] (sum + pad sont les mêmes
|
|
// ops sur les mêmes valeurs).
|
|
if (s < N_steps_golden) {
|
|
const float* py = step_inputs_py.data() + (size_t)s * n_embd;
|
|
double rmse = 0, dot = 0, na = 0, nb = 0;
|
|
for (int d = 0; d < n_embd; ++d) {
|
|
float a = next_embed[d], b = py[d];
|
|
double diff = a - b; rmse += diff * diff;
|
|
dot += (double)a * b; na += (double)a * a; nb += (double)b * b;
|
|
}
|
|
rmse = std::sqrt(rmse / n_embd);
|
|
double cos_v = dot / (std::sqrt(na) * std::sqrt(nb) + 1e-30);
|
|
if (s < 3 || s == N - 1) {
|
|
printf(" step %d: cb0=%d golden=%d, next vs py: rmse=%.4e cos=%.6f\n",
|
|
s, cb0, codes_golden[s * N_codebooks + 0], rmse, cos_v);
|
|
}
|
|
}
|
|
|
|
// Decode talker avec next_embed (sauf au dernier step où on n'a plus besoin du suivant)
|
|
if (s == N - 1) break;
|
|
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 = next_embed.data();
|
|
b.pos = pos1; b.n_seq_id = &n; b.seq_id = &sp; b.logits = &l;
|
|
const double td0 = now_s();
|
|
if (llama_decode(ctx, b) != 0) { printf("step %d FAILED\n", s); break; }
|
|
t_decode += now_s() - td0;
|
|
|
|
// sample CB0 pour le prochain step + capture hidden pour CP
|
|
logits = llama_get_logits_ith(ctx, -1);
|
|
cb0 = 0; mx = logits[0];
|
|
for (int i = 1; i < n_vocab; ++i) if (logits[i] > mx) { mx = logits[i]; cb0 = i; }
|
|
const float* hh = llama_get_embeddings_ith(ctx, -1);
|
|
if (hh) memcpy(hidden_for_cp.data(), hh, n_embd * sizeof(float));
|
|
}
|
|
const double t_loop = now_s() - t_loop0;
|
|
|
|
const double audio_s = N / 12.0; // 12 Hz codec
|
|
printf("DONE: N=%d (audio %.2f s @ 12Hz)\n", N, audio_s);
|
|
printf(" CB0 match golden : %d/%d\n", n_match_cb0, N);
|
|
printf(" All 16 codes match : %d/%d (info, dépend du sampling Python)\n", n_match_full, N);
|
|
printf("=== TIMING ===\n");
|
|
printf(" Talker prefill (T=%d) : %.3f s\n", T_prefill, t_prefill);
|
|
printf(" Loop (N=%d steps) : %.3f s talker_decode=%.3f, cp=%.3f, compose=%.3f, autres=%.3f\n",
|
|
N, t_loop, t_decode, t_cp, t_compose, t_loop - t_decode - t_cp - t_compose);
|
|
printf(" per-step talker : %.1f ms\n", t_decode * 1000.0 / N);
|
|
printf(" per-step CP (15 pass) : %.1f ms\n", t_cp * 1000.0 / N);
|
|
printf(" Total Talker+CP : %.3f s -> RTF talker+cp side = %.3f\n",
|
|
t_prefill + t_loop, (t_prefill + t_loop) / audio_s);
|
|
cp_free(cp_state);
|
|
|
|
// dump codes_engine
|
|
{
|
|
std::ofstream f(out_codes, std::ios::binary);
|
|
f.write((char*)codes_engine.data(), N * N_codebooks * sizeof(int32_t));
|
|
printf("codes_engine -> %s (%d frames * %d codebooks)\n", out_codes, N, N_codebooks);
|
|
}
|
|
llama_free(ctx);
|
|
llama_model_free(m);
|
|
return 0;
|
|
}
|