Kazeia
/
Kazeia-engine
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Kazeia-engine/dist/jni/cp_inference.cpp

202 lines
8.8 KiB
C++
Raw Blame History

#include "cp_inference.h"
#include <ggml.h>
#include <gguf.h>
#include <ggml-cpu.h>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cmath>
// 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 = 5;
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;
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;
}
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);
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));
bytes += ggml_nelements(mt) * sizeof(float);
}
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;
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* t32 = ggml_new_tensor(s.weights_ctx, GGML_TYPE_F32, ggml_n_dims(mt), mt->ne);
size_t nb = ggml_nbytes(mt);
int64_t ne = ggml_nelements(mt);
tmp.resize(nb);
fseek(f, off + gguf_get_tensor_offset(g, i), SEEK_SET);
if (fread(tmp.data(), 1, nb, f) != nb) { fprintf(stderr, "cp_load read fail %s\n", name); fclose(f); return false; }
if (mt->type == GGML_TYPE_F32) memcpy(t32->data, tmp.data(), nb);
else if (mt->type == GGML_TYPE_F16) ggml_fp16_to_fp32_row((const ggml_fp16_t*)tmp.data(), (float*)t32->data, ne);
else { fprintf(stderr, "cp_load unsupported type %s for %s\n", ggml_type_name(mt->type), name); fclose(f); return false; }
s.tensors[name] = t32;
}
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 + heads(%.0f MB) + codec_embs(%.0f MB) OK\n",
(long long)n, TAB * 4 / 1048576.0, TAB * 4 / 1048576.0);
return true;
}
void cp_free(CPState& s) {
if (s.weights_ctx) ggml_free(s.weights_ctx);
s.weights_ctx = nullptr;
s.tensors.clear();
s.heads.clear();
s.codec_embs.clear();
}
// 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);
}
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);
}
}
}
Reference in New Issue View Git Blame Copy Permalink