251 lines
12 KiB
Plaintext
251 lines
12 KiB
Plaintext
// GGUF loader : lit le .gguf produit par convert_to_gguf.py et peuple la
|
||
// table `tensors` du Decoder. Le forward pass dans decoder.cpp accède
|
||
// ensuite aux poids par leur clé string.
|
||
//
|
||
// Phase 1 : chargement seul, pas de forward.
|
||
|
||
#include "decoder.h"
|
||
#include <ggml-backend.h>
|
||
#include <ggml-cpu.h>
|
||
#include <cmath>
|
||
#include <cstdio>
|
||
#include <cstring>
|
||
#include <vector>
|
||
|
||
namespace kazeia::tts {
|
||
|
||
static int gguf_int(gguf_context * g, const char * key, int def = 0) {
|
||
int64_t idx = gguf_find_key(g, key);
|
||
if (idx < 0) return def;
|
||
return (int)gguf_get_val_u32(g, idx);
|
||
}
|
||
static float gguf_float(gguf_context * g, const char * key, float def = 0.0f) {
|
||
int64_t idx = gguf_find_key(g, key);
|
||
if (idx < 0) return def;
|
||
return gguf_get_val_f32(g, idx);
|
||
}
|
||
|
||
bool Decoder::load(const std::string & path) {
|
||
return load_with_buft(path, ggml_backend_cpu_buffer_type());
|
||
}
|
||
|
||
bool Decoder::load_with_buft(const std::string & path, ggml_backend_buffer_type_t buft) {
|
||
// Étape 1 : lire la metadata + créer les tenseurs SANS allouer (no_alloc).
|
||
struct gguf_init_params params;
|
||
params.no_alloc = true;
|
||
params.ctx = &ctx_w;
|
||
gguf = gguf_init_from_file(path.c_str(), params);
|
||
if (!gguf) {
|
||
fprintf(stderr, "Decoder::load failed to open %s\n", path.c_str());
|
||
return false;
|
||
}
|
||
|
||
// Lire metadata
|
||
cfg.latent_dim = gguf_int(gguf, "qwen3-tts-decoder-12hz.latent_dim", 1024);
|
||
cfg.codebook_dim = gguf_int(gguf, "qwen3-tts-decoder-12hz.codebook_dim", 512);
|
||
cfg.codebook_size = gguf_int(gguf, "qwen3-tts-decoder-12hz.codebook_size", 2048);
|
||
cfg.decoder_dim = gguf_int(gguf, "qwen3-tts-decoder-12hz.decoder_dim", 1536);
|
||
cfg.hidden_size = gguf_int(gguf, "qwen3-tts-decoder-12hz.hidden_size", 512);
|
||
cfg.intermediate_size = gguf_int(gguf, "qwen3-tts-decoder-12hz.intermediate_size", 1024);
|
||
cfg.num_attention_heads = gguf_int(gguf, "qwen3-tts-decoder-12hz.num_attention_heads", 16);
|
||
cfg.head_dim = gguf_int(gguf, "qwen3-tts-decoder-12hz.head_dim", 64);
|
||
cfg.num_hidden_layers = gguf_int(gguf, "qwen3-tts-decoder-12hz.num_hidden_layers", 8);
|
||
cfg.num_quantizers = gguf_int(gguf, "qwen3-tts-decoder-12hz.num_quantizers", 16);
|
||
cfg.layer_scale_initial = gguf_float(gguf, "qwen3-tts-decoder-12hz.layer_scale_initial", 0.01f);
|
||
cfg.rope_theta = gguf_float(gguf, "qwen3-tts-decoder-12hz.rope_theta", 10000.0f);
|
||
cfg.rms_norm_eps = gguf_float(gguf, "qwen3-tts-decoder-12hz.rms_norm_eps", 1e-5f);
|
||
cfg.upsample_rates = {8, 5, 4, 3}; // TODO lire array depuis GGUF
|
||
cfg.upsampling_ratios = {2, 2};
|
||
cfg.total_upsample = 1;
|
||
for (int r : cfg.upsample_rates) cfg.total_upsample *= r;
|
||
for (int r : cfg.upsampling_ratios) cfg.total_upsample *= r;
|
||
|
||
// Étape 2 : allouer les data sur le backend buffer demandé (CPU par défaut,
|
||
// ou Hexagon/Vulkan si précisé). Indispensable pour le scheduler
|
||
// ggml_backend_sched (assertion buffer_id >= 0).
|
||
buf_w = ggml_backend_alloc_ctx_tensors_from_buft(ctx_w, buft);
|
||
if (!buf_w) {
|
||
fprintf(stderr, "Decoder::load: failed to allocate backend buffer\n");
|
||
return false;
|
||
}
|
||
// Marquer ce buffer comme "weights" pour que sched sache qu'il peut
|
||
// (et préfère) router les ops weight-side sur le backend qui possède ces
|
||
// données — ou bien accepter le transfert vers un autre backend.
|
||
ggml_backend_buffer_set_usage(buf_w, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
|
||
|
||
// Étape 3 : lire les data depuis le fichier GGUF dans le backend buffer.
|
||
// Ouverture en lecture binaire ; on slurp les data section par section.
|
||
FILE * f = fopen(path.c_str(), "rb");
|
||
if (!f) {
|
||
fprintf(stderr, "Decoder::load: cannot reopen %s for reading data\n", path.c_str());
|
||
return false;
|
||
}
|
||
const size_t data_off = gguf_get_data_offset(gguf);
|
||
|
||
int64_t n_tensors = gguf_get_n_tensors(gguf);
|
||
std::vector<uint8_t> tmp;
|
||
for (int64_t i = 0; i < n_tensors; i++) {
|
||
const char * name = gguf_get_tensor_name(gguf, i);
|
||
ggml_tensor * t = ggml_get_tensor(ctx_w, name);
|
||
if (!t) continue;
|
||
const size_t nbytes = ggml_nbytes(t);
|
||
const size_t offset = data_off + gguf_get_tensor_offset(gguf, i);
|
||
tmp.resize(nbytes);
|
||
if (fseek(f, offset, SEEK_SET) != 0 ||
|
||
fread(tmp.data(), 1, nbytes, f) != nbytes) {
|
||
fprintf(stderr, "Decoder::load: read failed for %s\n", name);
|
||
fclose(f);
|
||
return false;
|
||
}
|
||
ggml_backend_tensor_set(t, tmp.data(), 0, nbytes);
|
||
tensors[name] = t;
|
||
}
|
||
fclose(f);
|
||
|
||
fprintf(stderr, "Decoder loaded : %lld tensors, total_upsample=%d "
|
||
"(buf_w %.1f MB, %s)\n",
|
||
(long long)n_tensors, cfg.total_upsample,
|
||
ggml_backend_buffer_get_size(buf_w) / 1048576.0,
|
||
ggml_backend_buffer_name(buf_w));
|
||
return true;
|
||
}
|
||
|
||
// Charge le modèle avec une liste de backends + crée un sched qui split les ops
|
||
// entre eux. backends[0] est prioritaire ; backends[last] doit être CPU (fallback).
|
||
// Les poids sont alloués sur le buffer du backend[0] (HTP si présent).
|
||
bool Decoder::load_with_backends(const std::string & path,
|
||
const std::vector<ggml_backend_t> & devs) {
|
||
if (devs.empty()) {
|
||
fprintf(stderr, "load_with_backends: empty backend list\n");
|
||
return false;
|
||
}
|
||
backends = devs;
|
||
// Allouer les poids sur le buft du 1er backend (HMX-friendly si HTP).
|
||
ggml_backend_buffer_type_t buft_w = ggml_backend_get_default_buffer_type(backends[0]);
|
||
if (!load_with_buft(path, buft_w)) return false;
|
||
|
||
// --- Précalcul snake constants (Phase 2b) ---
|
||
// Snake activation = x + sin²(α*x)/β. Le code historique faisait host-side
|
||
// exp() à chaque frame dans precompute_snake (incompatible avec sched
|
||
// no_alloc=true). On précalcule a_eff = exp(α), inv_b = 1/(exp(β)+eps) UNE FOIS
|
||
// ici, stockés dans ctx_const + buf_const sur le même backend que buf_w.
|
||
// Le forward stage_bigvgan les utilise comme des poids normaux.
|
||
{
|
||
const float SNAKE_EPS = 1e-9f;
|
||
// 1. Recenser tous les (alpha_key, beta_key) en parcourant les tensors.
|
||
struct SnakePair { std::string prefix; ggml_tensor * alpha; ggml_tensor * beta; };
|
||
std::vector<SnakePair> pairs;
|
||
for (auto & kv : tensors) {
|
||
const std::string & name = kv.first;
|
||
const std::string suf = ".alpha";
|
||
if (name.size() > suf.size() && name.compare(name.size()-suf.size(), suf.size(), suf) == 0) {
|
||
std::string prefix = name.substr(0, name.size() - suf.size());
|
||
std::string beta_name = prefix + ".beta";
|
||
auto bit = tensors.find(beta_name);
|
||
if (bit != tensors.end()) {
|
||
pairs.push_back({prefix, kv.second, bit->second});
|
||
}
|
||
}
|
||
}
|
||
// 2. Allouer ctx_const : 2 tensors par paire (a_eff, inv_b), tous f32 [C].
|
||
size_t mem_meta = (2 * pairs.size() + 4) * ggml_tensor_overhead();
|
||
ggml_init_params pc{mem_meta, nullptr, /*no_alloc=*/true};
|
||
ctx_const = ggml_init(pc);
|
||
std::vector<std::pair<ggml_tensor*, ggml_tensor*>> created;
|
||
created.reserve(pairs.size());
|
||
for (auto & p : pairs) {
|
||
const int C = (int)p.alpha->ne[0];
|
||
ggml_tensor * a_eff = ggml_new_tensor_1d(ctx_const, GGML_TYPE_F32, C);
|
||
ggml_tensor * inv_b = ggml_new_tensor_1d(ctx_const, GGML_TYPE_F32, C);
|
||
ggml_set_name(a_eff, (p.prefix + ".a_eff").c_str());
|
||
ggml_set_name(inv_b, (p.prefix + ".inv_b").c_str());
|
||
created.push_back({a_eff, inv_b});
|
||
}
|
||
// 3. Allouer buf_const sur le même backend que buf_w.
|
||
buf_const = ggml_backend_alloc_ctx_tensors_from_buft(ctx_const, buft_w);
|
||
if (!buf_const) {
|
||
fprintf(stderr, "load_with_backends: snake buf_const alloc FAIL\n");
|
||
return false;
|
||
}
|
||
ggml_backend_buffer_set_usage(buf_const, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
|
||
// 4. Calculer host-side puis tensor_set.
|
||
std::vector<float> tmp_a, tmp_b;
|
||
for (size_t i = 0; i < pairs.size(); ++i) {
|
||
const int C = (int)pairs[i].alpha->ne[0];
|
||
// Lire α et β depuis leur backend buffer (opt_hostbuf=1 -> mappable CPU).
|
||
tmp_a.resize(C); tmp_b.resize(C);
|
||
ggml_backend_tensor_get(pairs[i].alpha, tmp_a.data(), 0, C * sizeof(float));
|
||
ggml_backend_tensor_get(pairs[i].beta, tmp_b.data(), 0, C * sizeof(float));
|
||
std::vector<float> a_eff(C), inv_b(C);
|
||
for (int c = 0; c < C; ++c) {
|
||
a_eff[c] = std::exp(tmp_a[c]);
|
||
inv_b[c] = 1.0f / (std::exp(tmp_b[c]) + SNAKE_EPS);
|
||
}
|
||
ggml_backend_tensor_set(created[i].first, a_eff.data(), 0, C * sizeof(float));
|
||
ggml_backend_tensor_set(created[i].second, inv_b.data(), 0, C * sizeof(float));
|
||
snake_consts[pairs[i].prefix] = created[i];
|
||
}
|
||
fprintf(stderr, "Decoder: %zu snake constants pré-calculées (buf_const %.1f MB sur %s)\n",
|
||
pairs.size(),
|
||
ggml_backend_buffer_get_size(buf_const) / 1048576.0,
|
||
ggml_backend_buffer_name(buf_const));
|
||
}
|
||
|
||
// Mode minimal d'essai : KZTTS_DECODER_SCHED=1 active un vrai sched (refactor stages requis).
|
||
// Par défaut : pas de sched, on profite juste de opt_hostbuf=1 (HTP buffer CPU-mappé)
|
||
// pour que ggml_graph_compute_with_ctx CPU pur lise directement les poids HTP.
|
||
// -> aucun gain HMX dans ce mode, mais valide la chaîne load HTP + accès CPU.
|
||
if (getenv("KZTTS_DECODER_SCHED") && atoi(getenv("KZTTS_DECODER_SCHED")) != 0) {
|
||
std::vector<ggml_backend_buffer_type_t> bufts;
|
||
bufts.reserve(backends.size());
|
||
for (auto b : backends) bufts.push_back(ggml_backend_get_default_buffer_type(b));
|
||
const size_t graph_size = 16384;
|
||
sched = ggml_backend_sched_new(backends.data(), bufts.data(),
|
||
(int)backends.size(), graph_size,
|
||
/*parallel=*/false, /*op_offload=*/true);
|
||
if (!sched) {
|
||
fprintf(stderr, "load_with_backends: ggml_backend_sched_new FAILED\n");
|
||
return false;
|
||
}
|
||
fprintf(stderr, "Decoder: sched créé (refactor stages requis, KZTTS_DECODER_SCHED=1)\n");
|
||
} else {
|
||
fprintf(stderr, "Decoder: poids sur %s, compute CPU pur (sched OFF)\n",
|
||
ggml_backend_name(backends[0]));
|
||
}
|
||
return true;
|
||
}
|
||
|
||
void Decoder::compute_graph(struct ggml_context * ctx, struct ggml_cgraph * gf, int n_threads) {
|
||
if (sched) {
|
||
// Path backend abstrait : split auto CPU/HTP.
|
||
ggml_backend_sched_reset(sched);
|
||
if (!ggml_backend_sched_alloc_graph(sched, gf)) {
|
||
fprintf(stderr, "compute_graph: sched_alloc_graph FAILED\n"); return;
|
||
}
|
||
if (ggml_backend_sched_graph_compute(sched, gf) != GGML_STATUS_SUCCESS) {
|
||
fprintf(stderr, "compute_graph: sched_graph_compute FAILED\n");
|
||
}
|
||
} else {
|
||
// Path legacy CPU pur (load via load() ou load_with_buft sans backend list).
|
||
ggml_graph_compute_with_ctx(ctx, gf, n_threads);
|
||
}
|
||
}
|
||
|
||
void Decoder::unload() {
|
||
if (sched) { ggml_backend_sched_free(sched); sched = nullptr; }
|
||
if (buf_const) { ggml_backend_buffer_free(buf_const); buf_const = nullptr; }
|
||
if (ctx_const) { ggml_free(ctx_const); ctx_const = nullptr; }
|
||
snake_consts.clear();
|
||
if (buf_w) { ggml_backend_buffer_free(buf_w); buf_w = nullptr; }
|
||
if (ctx_w) { ggml_free(ctx_w); ctx_w = nullptr; }
|
||
if (ctx_g) { ggml_free(ctx_g); ctx_g = nullptr; }
|
||
if (gguf) { gguf_free(gguf); gguf = nullptr; }
|
||
if (owns_backends) {
|
||
for (auto b : backends) if (b) ggml_backend_free(b);
|
||
}
|
||
backends.clear();
|
||
tensors.clear();
|
||
}
|
||
|
||
} // namespace kazeia::tts
|