679 lines
34 KiB
C++
679 lines
34 KiB
C++
#include "cp_inference.h"
|
||
#include <ggml.h>
|
||
#include <gguf.h>
|
||
#include <ggml-backend.h>
|
||
#include <ggml-cpu.h>
|
||
#include <chrono>
|
||
#include <cstdio>
|
||
#include <cstdlib>
|
||
#include <cstring>
|
||
#include <cmath>
|
||
#if defined(__aarch64__)
|
||
# include <arm_neon.h>
|
||
#endif
|
||
|
||
// 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 = 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;
|
||
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");
|
||
if (!f) { fprintf(stderr, "cp_load: open fail %s\n", path); return {}; }
|
||
std::vector<float> v(expect);
|
||
if (fread(v.data(), sizeof(float), expect, f) != expect) {
|
||
fprintf(stderr, "cp_load: short read %s\n", path); fclose(f); return {};
|
||
}
|
||
fclose(f); return v;
|
||
}
|
||
|
||
// Garde-t-on le tenseur en F16 ou on convertit en F32 ?
|
||
// Rule : les poids matmul (attn_q/k/v/output, ffn_gate/up/down) restent F16 -> moitié BW
|
||
// dans les gros mul_mat (~85% du coût CP). Les norms restent en F32 car ggml_mul exige
|
||
// même type des deux opérandes (xn_f32 * w_norm_f32). Économie pratique : embeds ~1024
|
||
// éléments * 5 layers * 6 norms = 30 KB de norms, négligeable. Les matmul valent
|
||
// ~5 * (3*1024*1024 attn + 3*1024*MLP ffn) qui dominent.
|
||
static bool cp_keep_f16(const std::string& name) {
|
||
// RÉGRESSION CORRIGÉE : garder les poids matmul en F16 fait tourner le mul_mat en F16
|
||
// (ggml convertit l'activation f32 -> f16 pour le vec_dot). Les massive activations
|
||
// Qwen3 (hidden ±83, intermédiaires >65504) overflow en f16 -> inf -> NaN -> CP sort
|
||
// des scores NaN -> argmax=0 -> audio inaudible. cp_runner.cpp (référence bit-exact)
|
||
// convertit TOUT en F32 pour cette raison. On fait pareil : F16 désactivé par défaut.
|
||
// KZTTS_CP_KEEP_F16=1 réactive l'ancien comportement (debug perf uniquement).
|
||
if (!(getenv("KZTTS_CP_KEEP_F16") && atoi(getenv("KZTTS_CP_KEEP_F16")) != 0)) return false;
|
||
if (name.find("_norm.weight") != std::string::npos) return false;
|
||
return true;
|
||
}
|
||
|
||
bool cp_load(CPState& s, const char* gguf_path, const char* heads_path, const char* embs_path) {
|
||
ggml_context* meta = nullptr;
|
||
gguf_init_params p; p.no_alloc = true; p.ctx = &meta;
|
||
gguf_context* g = gguf_init_from_file(gguf_path, p);
|
||
if (!g) { fprintf(stderr, "cp_load: gguf open fail %s\n", gguf_path); return false; }
|
||
int64_t n = gguf_get_n_tensors(g);
|
||
|
||
// Budget : on alloue à la taille native (F16 reste F16 si autorisé, sinon F32). Borne
|
||
// sup = somme des max(nb_f32, nb_native) pour être tranquille.
|
||
size_t bytes = 0;
|
||
for (int64_t i = 0; i < n; i++) {
|
||
ggml_tensor* mt = ggml_get_tensor(meta, gguf_get_tensor_name(g, i));
|
||
size_t nb_f32 = (size_t)ggml_nelements(mt) * sizeof(float);
|
||
bytes += nb_f32; // sup
|
||
}
|
||
ggml_init_params ip = { bytes + (size_t)n * ggml_tensor_overhead() + (1u << 20), nullptr, false };
|
||
s.weights_ctx = ggml_init(ip);
|
||
|
||
FILE* f = fopen(gguf_path, "rb");
|
||
const size_t off = gguf_get_data_offset(g);
|
||
std::vector<uint8_t> tmp;
|
||
size_t kept_f16 = 0, conv_f32 = 0;
|
||
for (int64_t i = 0; i < n; i++) {
|
||
const char* name = gguf_get_tensor_name(g, i);
|
||
ggml_tensor* mt = ggml_get_tensor(meta, name);
|
||
size_t nb_src = ggml_nbytes(mt);
|
||
int64_t ne = ggml_nelements(mt);
|
||
tmp.resize(nb_src);
|
||
fseek(f, off + gguf_get_tensor_offset(g, i), SEEK_SET);
|
||
if (fread(tmp.data(), 1, nb_src, f) != nb_src) { fprintf(stderr, "cp_load read fail %s\n", name); fclose(f); return false; }
|
||
|
||
ggml_type dst_type = GGML_TYPE_F32;
|
||
if (mt->type == GGML_TYPE_F16 && cp_keep_f16(name)) dst_type = GGML_TYPE_F16;
|
||
else if (mt->type == GGML_TYPE_F32) dst_type = GGML_TYPE_F32;
|
||
else dst_type = GGML_TYPE_F32; // F16 -> F32 pour les norms
|
||
|
||
ggml_tensor* t = ggml_new_tensor(s.weights_ctx, dst_type, ggml_n_dims(mt), mt->ne);
|
||
if (dst_type == mt->type) {
|
||
memcpy(t->data, tmp.data(), nb_src);
|
||
if (dst_type == GGML_TYPE_F16) kept_f16++;
|
||
} else if (mt->type == GGML_TYPE_F16 && dst_type == GGML_TYPE_F32) {
|
||
ggml_fp16_to_fp32_row((const ggml_fp16_t*)tmp.data(), (float*)t->data, ne);
|
||
conv_f32++;
|
||
} else if (mt->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F32) {
|
||
memcpy(t->data, tmp.data(), nb_src);
|
||
} else {
|
||
fprintf(stderr, "cp_load unsupported cast %s -> %s for %s\n",
|
||
ggml_type_name(mt->type), ggml_type_name(dst_type), name);
|
||
fclose(f); return false;
|
||
}
|
||
s.tensors[name] = t;
|
||
}
|
||
fclose(f);
|
||
gguf_free(g);
|
||
|
||
const size_t TAB = (size_t)N_CB * N_VOCAB * N_EMBD;
|
||
s.heads = read_floats(heads_path, TAB); if (s.heads.empty()) return false;
|
||
s.codec_embs = read_floats(embs_path, TAB); if (s.codec_embs.empty()) return false;
|
||
fprintf(stderr, "cp_load: %lld tensors (f16 kept=%zu, f32=%zu) + heads(%.0f MB) + codec_embs(%.0f MB) OK\n",
|
||
(long long)n, kept_f16, conv_f32 + (size_t)(n - (int64_t)(kept_f16 + conv_f32)),
|
||
TAB * 4 / 1048576.0, TAB * 4 / 1048576.0);
|
||
return true;
|
||
}
|
||
|
||
// Variante backend-aware (Phase F : CP forward sur HMX V79). Charge les poids sur
|
||
// le buft du 1er backend (HTP), crée un sched, et le forward sched-friendly est
|
||
// activé via cp_forward_cached_step_sched.
|
||
bool cp_load_with_backends(CPState& s, const char* gguf_path,
|
||
const char* heads_path, const char* embs_path,
|
||
const std::vector<ggml_backend_t>& devs,
|
||
bool owns) {
|
||
if (devs.empty()) { fprintf(stderr, "cp_load_with_backends: empty backends\n"); return false; }
|
||
s.backends = devs;
|
||
s.owns_backends = owns;
|
||
ggml_backend_buffer_type_t buft_w = ggml_backend_get_default_buffer_type(s.backends[0]);
|
||
|
||
// 1) Lire metadata gguf (no_alloc=true)
|
||
ggml_context* meta = nullptr;
|
||
gguf_init_params p; p.no_alloc = true; p.ctx = &meta;
|
||
gguf_context* g = gguf_init_from_file(gguf_path, p);
|
||
if (!g) { fprintf(stderr, "cp_load_with_backends: gguf open fail %s\n", gguf_path); return false; }
|
||
int64_t n = gguf_get_n_tensors(g);
|
||
|
||
// 2) Créer ctx weights no_alloc=true (juste metadata des tensors)
|
||
size_t mem_meta = (size_t)n * ggml_tensor_overhead() + (1u << 20);
|
||
ggml_init_params ip = { mem_meta, nullptr, /*no_alloc=*/true };
|
||
s.weights_ctx = ggml_init(ip);
|
||
|
||
// 3) Pour chaque tenseur du gguf, créer le tensor avec son type final dans s.weights_ctx
|
||
// Règle conservée : F16 source + nom != "_norm.weight" -> garde F16 (matmul-friendly).
|
||
// sinon -> F32 (norms, etc.).
|
||
size_t kept_f16 = 0, conv_f32 = 0;
|
||
for (int64_t i = 0; i < n; i++) {
|
||
const char* name = gguf_get_tensor_name(g, i);
|
||
ggml_tensor* mt = ggml_get_tensor(meta, name);
|
||
ggml_type dst_type = GGML_TYPE_F32;
|
||
if (mt->type == GGML_TYPE_F16 && cp_keep_f16(name)) dst_type = GGML_TYPE_F16;
|
||
ggml_tensor* t = ggml_new_tensor(s.weights_ctx, dst_type, ggml_n_dims(mt), mt->ne);
|
||
ggml_set_name(t, name);
|
||
s.tensors[name] = t;
|
||
if (dst_type == GGML_TYPE_F16) kept_f16++; else conv_f32++;
|
||
}
|
||
|
||
// 4) Allouer le buffer backend pour tous les tenseurs du ctx d'un coup
|
||
s.weights_buf = ggml_backend_alloc_ctx_tensors_from_buft(s.weights_ctx, buft_w);
|
||
if (!s.weights_buf) {
|
||
fprintf(stderr, "cp_load_with_backends: weights buf alloc FAIL\n");
|
||
gguf_free(g); ggml_free(meta); return false;
|
||
}
|
||
ggml_backend_buffer_set_usage(s.weights_buf, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
|
||
|
||
// 5) Lire les data depuis le gguf et upload via ggml_backend_tensor_set
|
||
FILE* f = fopen(gguf_path, "rb");
|
||
const size_t off = gguf_get_data_offset(g);
|
||
std::vector<uint8_t> tmp;
|
||
std::vector<float> tmp_f32;
|
||
for (int64_t i = 0; i < n; i++) {
|
||
const char* name = gguf_get_tensor_name(g, i);
|
||
ggml_tensor* mt = ggml_get_tensor(meta, name);
|
||
ggml_tensor* t = s.tensors[name];
|
||
size_t nb_src = ggml_nbytes(mt);
|
||
tmp.resize(nb_src);
|
||
fseek(f, off + gguf_get_tensor_offset(g, i), SEEK_SET);
|
||
if (fread(tmp.data(), 1, nb_src, f) != nb_src) {
|
||
fprintf(stderr, "cp_load_with_backends: read fail %s\n", name); fclose(f); return false;
|
||
}
|
||
if (mt->type == t->type) {
|
||
ggml_backend_tensor_set(t, tmp.data(), 0, nb_src);
|
||
} else if (mt->type == GGML_TYPE_F16 && t->type == GGML_TYPE_F32) {
|
||
// Convert F16 -> F32 host-side then upload.
|
||
int64_t ne = ggml_nelements(mt);
|
||
tmp_f32.resize(ne);
|
||
ggml_fp16_to_fp32_row((const ggml_fp16_t*)tmp.data(), tmp_f32.data(), ne);
|
||
ggml_backend_tensor_set(t, tmp_f32.data(), 0, (size_t)ne * sizeof(float));
|
||
} else {
|
||
fprintf(stderr, "cp_load_with_backends: bad cast %s -> %s for %s\n",
|
||
ggml_type_name(mt->type), ggml_type_name(t->type), name);
|
||
fclose(f); return false;
|
||
}
|
||
}
|
||
fclose(f);
|
||
gguf_free(g);
|
||
ggml_free(meta);
|
||
|
||
// 6) Tables host-side (heads, codec_embs) : inchangées, utilisées en dot product CPU.
|
||
const size_t TAB = (size_t)N_CB * N_VOCAB * N_EMBD;
|
||
s.heads = read_floats(heads_path, TAB); if (s.heads.empty()) return false;
|
||
s.codec_embs = read_floats(embs_path, TAB); if (s.codec_embs.empty()) return false;
|
||
|
||
// 7) Sched (gated par KZTTS_CP_SCHED=1 ; off par défaut tant que stage refactor non validé).
|
||
// Sans sched, l'exec retombe sur compute_with_ctx CPU (poids HTP CPU-mappés via opt_hostbuf=1).
|
||
if (getenv("KZTTS_CP_SCHED") && atoi(getenv("KZTTS_CP_SCHED")) != 0) {
|
||
std::vector<ggml_backend_buffer_type_t> bufts;
|
||
bufts.reserve(s.backends.size());
|
||
for (auto b : s.backends) bufts.push_back(ggml_backend_get_default_buffer_type(b));
|
||
// op_offload toggle via KZTTS_CP_OFFLOAD (default 1 = true).
|
||
const bool op_offload = !(getenv("KZTTS_CP_OFFLOAD") && atoi(getenv("KZTTS_CP_OFFLOAD")) == 0);
|
||
s.sched = ggml_backend_sched_new(s.backends.data(), bufts.data(),
|
||
(int)s.backends.size(), /*graph_size=*/8192,
|
||
/*parallel=*/false, op_offload);
|
||
if (!s.sched) { fprintf(stderr, "cp_load_with_backends: sched_new FAIL\n"); return false; }
|
||
fprintf(stderr, "cp_load_with_backends: %lld tensors (f16=%zu f32=%zu) sur %s, sched ON\n",
|
||
(long long)n, kept_f16, conv_f32, ggml_backend_name(s.backends[0]));
|
||
} else {
|
||
fprintf(stderr, "cp_load_with_backends: %lld tensors (f16=%zu f32=%zu) sur %s, sched OFF\n",
|
||
(long long)n, kept_f16, conv_f32, ggml_backend_name(s.backends[0]));
|
||
}
|
||
return true;
|
||
}
|
||
|
||
void cp_free(CPState& s) {
|
||
if (s.sched) { ggml_backend_sched_free(s.sched); s.sched = nullptr; }
|
||
if (s.cache_buf) { ggml_backend_buffer_free(s.cache_buf); s.cache_buf = nullptr; }
|
||
if (s.weights_buf) { ggml_backend_buffer_free(s.weights_buf); s.weights_buf = nullptr; }
|
||
if (s.weights_ctx) { ggml_free(s.weights_ctx); s.weights_ctx = nullptr; }
|
||
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; }
|
||
if (s.owns_backends) {
|
||
for (auto b : s.backends) if (b) ggml_backend_free(b);
|
||
}
|
||
s.backends.clear();
|
||
}
|
||
|
||
// 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.
|
||
//
|
||
// Path A (legacy CPU pur, s.sched=null) : ggml_init no_alloc=false, data sur ctx mem.
|
||
// Path B (sched HTP, s.sched non-null) : ggml_init no_alloc=true, alloué sur backends[0]
|
||
// buffer via ggml_backend_alloc_ctx_tensors_from_buft.
|
||
static bool cp_cache_init(CPState& s) {
|
||
if (s.cache_ctx) return true;
|
||
const bool use_sched = (s.sched != nullptr);
|
||
if (!use_sched) {
|
||
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, /*no_alloc=*/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);
|
||
}
|
||
return true;
|
||
}
|
||
// Path sched : ctx meta-only, buf backend séparé.
|
||
size_t mem_meta = (size_t)N_LAYER * 2 * ggml_tensor_overhead() + (1u << 16);
|
||
ggml_init_params p = { mem_meta, nullptr, /*no_alloc=*/true };
|
||
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);
|
||
char nm[32];
|
||
snprintf(nm, sizeof(nm), "K_cache_%d", L_i); ggml_set_name(s.K_cache[L_i], nm);
|
||
snprintf(nm, sizeof(nm), "V_cache_%d", L_i); ggml_set_name(s.V_cache[L_i], nm);
|
||
}
|
||
ggml_backend_buffer_type_t buft = ggml_backend_get_default_buffer_type(s.backends[0]);
|
||
s.cache_buf = ggml_backend_alloc_ctx_tensors_from_buft(s.cache_ctx, buft);
|
||
if (!s.cache_buf) { fprintf(stderr, "cp_cache_init: backend alloc FAIL\n"); return false; }
|
||
// KV cache n'est PAS un weight (il change entre frames), pas de set_usage WEIGHTS ici.
|
||
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<float>& embeds_flat, int L, int last,
|
||
std::vector<float>& out_hn) {
|
||
size_t mem = 128ULL * 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, L);
|
||
memcpy(x->data, embeds_flat.data(), (size_t)N_EMBD * L * sizeof(float));
|
||
ggml_tensor* pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, L);
|
||
for (int i = 0; i < L; i++) ((int32_t*)pos->data)[i] = i;
|
||
ggml_tensor* mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, L, L);
|
||
{
|
||
float* m = (float*)mask->data;
|
||
for (int q = 0; q < L; q++)
|
||
for (int k = 0; k < L; k++)
|
||
m[k + L * q] = (k > q) ? -INFINITY : 0.0f;
|
||
}
|
||
const float scale = 1.0f / sqrtf((float)HEAD_DIM);
|
||
|
||
auto W = [&](const std::string& k) -> ggml_tensor* {
|
||
auto it = s.tensors.find(k);
|
||
if (it == s.tensors.end()) { fprintf(stderr, "CP missing tensor %s\n", k.c_str()); abort(); }
|
||
return it->second;
|
||
};
|
||
|
||
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, L);
|
||
K = ggml_reshape_3d(ctx, K, HEAD_DIM, N_KV, L);
|
||
V = ggml_reshape_3d(ctx, V, HEAD_DIM, N_KV, L);
|
||
|
||
// q/k-norm AVANT RoPE (gotcha)
|
||
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, 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, NULL, HEAD_DIM, GGML_ROPE_TYPE_NEOX, 0,
|
||
ROPE_BASE, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
|
||
|
||
Q = ggml_cont(ctx, ggml_permute(ctx, Q, 0, 2, 1, 3));
|
||
K = ggml_cont(ctx, ggml_permute(ctx, K, 0, 2, 1, 3));
|
||
V = ggml_cont(ctx, ggml_permute(ctx, V, 1, 2, 0, 3));
|
||
|
||
ggml_tensor* KQ = ggml_mul_mat(ctx, K, Q);
|
||
KQ = ggml_soft_max_ext(ctx, KQ, mask, scale, 0.0f);
|
||
ggml_tensor* KQV = ggml_mul_mat(ctx, V, KQ);
|
||
KQV = ggml_cont(ctx, ggml_permute(ctx, KQV, 0, 2, 1, 3));
|
||
KQV = ggml_reshape_2d(ctx, KQV, HEAD_DIM * N_HEAD, L);
|
||
|
||
ggml_tensor* attn = ggml_mul_mat(ctx, W(b + ".attn_output.weight"), KQV);
|
||
x = ggml_add(ctx, res, attn);
|
||
|
||
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);
|
||
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 * last, N_EMBD * sizeof(float));
|
||
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;
|
||
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<double>(
|
||
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.
|
||
std::vector<int32_t> pos_host(n_new);
|
||
for (int i = 0; i < n_new; i++) pos_host[i] = pos_off + i;
|
||
std::vector<float> mask_host((size_t)T_full * n_new);
|
||
for (int q = 0; q < n_new; q++)
|
||
for (int k = 0; k < T_full; k++)
|
||
mask_host[k + T_full * q] = (k > pos_off + q) ? -INFINITY : 0.0f;
|
||
|
||
ggml_tensor* x = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, N_EMBD, n_new);
|
||
ggml_tensor* pos_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_new);
|
||
ggml_tensor* mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, T_full, n_new);
|
||
if (use_sched) {
|
||
ggml_set_name(x, "x"); ggml_set_input(x);
|
||
ggml_set_name(pos_ids, "pos"); ggml_set_input(pos_ids);
|
||
ggml_set_name(mask, "mask"); ggml_set_input(mask);
|
||
} else {
|
||
// Legacy : data est sur le ctx mem (no_alloc=false), memcpy direct OK.
|
||
memcpy(x->data, embeds_new, (size_t)N_EMBD * n_new * sizeof(float));
|
||
memcpy(pos_ids->data, pos_host.data(), pos_host.size() * sizeof(int32_t));
|
||
memcpy(mask->data, mask_host.data(), mask_host.size() * sizeof(float));
|
||
}
|
||
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"));
|
||
if (use_sched) { ggml_set_name(x, "out_x"); ggml_set_output(x); }
|
||
|
||
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);
|
||
double tg_build = prof_global ? nows_fn() : 0;
|
||
|
||
out_hn.resize(N_EMBD);
|
||
if (use_sched) {
|
||
// Profile fin via KZTTS_CP_PROFILE=1 : décompose alloc / set / compute / get
|
||
// pour identifier où va le temps dans la boucle sched.
|
||
const bool prof = (getenv("KZTTS_CP_PROFILE") && atoi(getenv("KZTTS_CP_PROFILE")) != 0);
|
||
auto nows = []() {
|
||
return std::chrono::duration<double>(
|
||
std::chrono::steady_clock::now().time_since_epoch()).count();
|
||
};
|
||
double t_a0=0, t_a1=0, t_s1=0, t_c1=0, t_g1=0;
|
||
if (prof) t_a0 = nows();
|
||
|
||
ggml_backend_sched_reset(s.sched);
|
||
if (!ggml_backend_sched_alloc_graph(s.sched, gf)) {
|
||
fprintf(stderr, "cp_forward_cached_step: sched_alloc_graph FAIL\n");
|
||
ggml_free(ctx); return;
|
||
}
|
||
if (prof) t_a1 = nows();
|
||
|
||
// Inputs uploadés APRÈS sched_alloc_graph (les tensors ont maintenant leur backend buf).
|
||
ggml_backend_tensor_set(x, embeds_new, 0, (size_t)N_EMBD * n_new * sizeof(float));
|
||
ggml_backend_tensor_set(pos_ids, pos_host.data(), 0, pos_host.size() * sizeof(int32_t));
|
||
ggml_backend_tensor_set(mask, mask_host.data(), 0, mask_host.size() * sizeof(float));
|
||
if (prof) t_s1 = nows();
|
||
|
||
if (ggml_backend_sched_graph_compute(s.sched, gf) != GGML_STATUS_SUCCESS) {
|
||
fprintf(stderr, "cp_forward_cached_step: sched_graph_compute FAIL\n");
|
||
ggml_free(ctx); return;
|
||
}
|
||
if (prof) t_c1 = nows();
|
||
|
||
// hn = colonne (n_new - 1) de x [N_EMBD, n_new]. tensor_get avec offset.
|
||
ggml_backend_tensor_get(x, out_hn.data(),
|
||
(size_t)N_EMBD * (n_new - 1) * sizeof(float),
|
||
(size_t)N_EMBD * sizeof(float));
|
||
if (prof) {
|
||
t_g1 = nows();
|
||
fprintf(stderr, "cp_sub n_new=%d pos=%d reset+alloc=%.2fms set=%.2fms compute=%.2fms get=%.2fms\n",
|
||
n_new, pos_off,
|
||
(t_a1-t_a0)*1000.0, (t_s1-t_a1)*1000.0,
|
||
(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);
|
||
}
|
||
|
||
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.
|
||
std::vector<float> embeds; embeds.reserve((size_t)17 * N_EMBD);
|
||
embeds.insert(embeds.end(), hidden, hidden + N_EMBD);
|
||
embeds.insert(embeds.end(), cb0_emb, cb0_emb + N_EMBD);
|
||
|
||
std::vector<float> hn;
|
||
std::vector<float> scores(N_VOCAB);
|
||
int local_history[N_CB]; // tokens samplés ce step (pour rep_penalty intra-frame)
|
||
int n_local = 0;
|
||
for (int step = 1; step <= N_CB; step++) {
|
||
const int L = step + 1;
|
||
cp_forward_lastpos(s, embeds, L, step, hn);
|
||
|
||
// tête[step-1] : scores[j] = sum_k hn[k] * heads[step-1, j, k]
|
||
const float* Wh = &s.heads[(size_t)(step - 1) * 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;
|
||
}
|
||
// Sample via sampler partagé (HF-style). rep_penalty appliquée sur les tokens déjà
|
||
// émis CE step (intra-frame), pas l'historique long.
|
||
int chosen = sampler_sample_local(s.sampler, scores.data(), N_VOCAB,
|
||
local_history, n_local);
|
||
out_codes[step - 1] = chosen;
|
||
local_history[n_local++] = chosen;
|
||
|
||
if (step < N_CB) {
|
||
const float* e = &s.codec_embs[(size_t)(step - 1) * N_VOCAB * N_EMBD + (size_t)chosen * N_EMBD];
|
||
embeds.insert(embeds.end(), e, e + N_EMBD);
|
||
}
|
||
}
|
||
}
|
||
|
||
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;
|
||
|
||
// 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];
|
||
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);
|
||
};
|
||
|
||
// 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;
|
||
}
|
||
}
|