enne2/qwentts.cpp
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

tts: lazy encoder loading and in graph embedding gathers

Browse Source 
Load the codec encoder half (seanet, enc_transformer, enc_downsample,
qenc) and the speaker encoder lazily on their first real use: synthesis
from a pre extracted reference (--ref-spk --ref-rvq) now brings up only
the talker and the codec decoder, matching the load profile of a preset
voice server. Assemble the AR inputs on device: the talker decode graph
gathers and sums the 16 frame code embeddings plus the trailing text or
pad overlay via get_rows, and the code predictor gathers c0 and each
sampled sub code from its group table in graph. Per frame host traffic
drops from 16 gguf row reads plus a CPU sum plus 15 synchronous backend
readbacks to 16 code ids and one overlay row uploaded. The next-emb
parity dump reproduces the in graph composition on host under --dump
only, staying byte comparable against the Python hook.
This commit is contained in:
Pascal
2026-07-04 21:35:03 +02:00
parent c9d9fa1c6b
commit 3e3b6c1712
6 changed files with 237 additions and 193 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/code-predictor-forward.h
+35 -41
View File
@@ -193,7 +193,9 @@ static bool code_predictor_run(const CodePredictorWeights * cw,
KVCache * kv,
ggml_backend_sched_t sched,
GraphArena * arena,
const float * fresh_input,
struct ggml_tensor * embd_table,
const float * hidden_row,
int32_t code_id,
int T,
int n_past,
int talker_hidden,
@@ -213,13 +215,28 @@ static bool code_predictor_run(const CodePredictorWeights * cw,
const int max_nodes = code_predictor_graph_max_nodes(n_layers);
struct ggml_context * gctx = graph_arena_begin(arena);
// Inputs: fresh embeddings (talker_hidden), positions, attention mask
struct ggml_tensor * x_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F32, talker_hidden, T);
// Inputs: one code id gathered in graph from embd_table, positions,
// attention mask, plus the raw talker hidden row on the prefill
// path (T == 2, hidden_row non NULL) where the sequence is
// [talker_hidden, embed(c0)]. Steps (T == 1) are pure gathers: the
// only per step upload is 4 bytes of code id.
struct ggml_tensor * ids_in = ggml_new_tensor_1d(gctx, GGML_TYPE_I32, 1);
struct ggml_tensor * pos_in = ggml_new_tensor_1d(gctx, GGML_TYPE_I32, T);
struct ggml_tensor * mask_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F16, n_kv_pad, T);
ggml_set_name(x_in, "sub_input");
ggml_set_name(ids_in, "sub_code_id");
ggml_set_name(pos_in, "positions");
ggml_set_name(mask_in, "causal_mask");
ggml_set_input(ids_in);
struct ggml_tensor * x_in = ggml_get_rows(gctx, embd_table, ids_in);
struct ggml_tensor * hid_in = NULL;
if (T == 2) {
hid_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F32, talker_hidden, 1);
ggml_set_name(hid_in, "talker_hidden_row");
ggml_set_input(hid_in);
x_in = ggml_concat(gctx, hid_in, x_in, 1);
}
ggml_set_name(x_in, "sub_input");
struct ggml_cgraph * gf = ggml_new_graph_custom(gctx, max_nodes, false);
@@ -254,7 +271,10 @@ static bool code_predictor_run(const CodePredictorWeights * cw,
return false;
}
ggml_backend_tensor_set(x_in, fresh_input, 0, (size_t) T * (size_t) talker_hidden * sizeof(float));
ggml_backend_tensor_set(ids_in, &code_id, 0, sizeof(int32_t));
if (hid_in) {
ggml_backend_tensor_set(hid_in, hidden_row, 0, (size_t) talker_hidden * sizeof(float));
}
{
std::vector<int32_t> pos((size_t) T);
@@ -294,30 +314,6 @@ static bool code_predictor_run(const CodePredictorWeights * cw,
return true;
}
// Read one row of an embedding table to f32. Reads from the backend
// (the predictor weights live there) via ggml_backend_tensor_get,
// dispatched through ggml_get_type_traits so quants are accepted.
static void embed_row_from_backend(struct ggml_tensor * t, int row_id, int dim, float * dst) {
if (t->ne[0] != dim) {
qt_throw("[CodePredictor] embed dim mismatch %lld vs %d", (long long) t->ne[0], dim);
}
if (row_id < 0 || row_id >= (int) t->ne[1]) {
qt_throw("[CodePredictor] row %d out of range (vocab=%lld)", row_id, (long long) t->ne[1]);
}
const size_t row_bytes = ggml_row_size(t->type, dim);
if (t->type == GGML_TYPE_F32) {
ggml_backend_tensor_get(t, dst, (size_t) row_id * row_bytes, row_bytes);
return;
}
const struct ggml_type_traits * tt = ggml_get_type_traits(t->type);
if (!tt || !tt->to_float) {
qt_throw("[CodePredictor] unsupported embed dtype %d", (int) t->type);
}
std::vector<uint8_t> tmp(row_bytes);
ggml_backend_tensor_get(t, tmp.data(), (size_t) row_id * row_bytes, row_bytes);
tt->to_float(tmp.data(), dst, dim);
}
// Run the predictor for one audio frame. Caller passes the talker hidden
// state for the current frame and the already-sampled c0. Sampling
// parameters control greedy (temperature <= 0) vs stochastic. subseq_base
@@ -358,15 +354,14 @@ static bool code_predictor_step(const TalkerWeights * tw,
out->codes.assign((size_t) (n_acoustic + 1), 0);
out->codes[0] = c0;
// Prefill: two positions, talker_hidden_last and embed_talker(c0).
// Prefill: two positions, [talker_hidden_last, embed_talker(c0)].
// The hidden row uploads raw, c0 gathers in graph from the talker
// codec embedding table.
kv_cache_reset(kv);
std::vector<float> prefill_input((size_t) 2 * (size_t) talker_hidden, 0.0f);
std::memcpy(prefill_input.data(), talker_hidden_last, (size_t) talker_hidden * sizeof(float));
embed_row_from_backend(tw->codec_embedding, c0, talker_hidden, prefill_input.data() + (size_t) talker_hidden);
std::vector<float> logits;
if (!code_predictor_run(cw, kv, sched, arena_prefill, prefill_input.data(), 2, 0, talker_hidden, 0, use_flash_attn,
clamp_fp16, &logits)) {
if (!code_predictor_run(cw, kv, sched, arena_prefill, tw->codec_embedding, talker_hidden_last, c0, 2, 0,
talker_hidden, 0, use_flash_attn, clamp_fp16, &logits)) {
return false;
}
{
@@ -385,13 +380,12 @@ static bool code_predictor_step(const TalkerWeights * tw,
}
// Decode loop: 14 single-token steps. At step g (g=1..14) we feed
// the embedding of the code we just sampled and read lm_head[g].
std::vector<float> step_input((size_t) talker_hidden);
// the id of the code we just sampled, gathered in graph from the
// group's private embedding table, and read lm_head[g].
for (int g = 1; g < n_acoustic; g++) {
embed_row_from_backend(cw->codec_embedding[(size_t) (g - 1)], out->codes[(size_t) g], talker_hidden,
step_input.data());
if (!code_predictor_run(cw, kv, sched, arena_step, step_input.data(), 1, kv->cur_len, talker_hidden, g,
use_flash_attn, clamp_fp16, &logits)) {
if (!code_predictor_run(cw, kv, sched, arena_step, cw->codec_embedding[(size_t) (g - 1)], NULL,
out->codes[(size_t) g], 1, kv->cur_len, talker_hidden, g, use_flash_attn, clamp_fp16,
&logits)) {
return false;
}
float u_g = 0.0f;
src/pipeline-codec.cpp
+45 -49
View File
@@ -11,6 +11,7 @@
#include "causal-trans-conv.h"
#include "debug.h"
#include "qt-error.h"
#include "timer.h"
#include <cmath>
#include <cstdio>
@@ -73,51 +74,9 @@ bool pipeline_codec_load(PipelineCodec * pc, const char * gguf_path, BackendPair
pc->pre_conv_wctx = std::move(wctx);
}
if (!seanet_encoder_load(&pc->seanet, pc->gguf, pc->backend)) {
wctx_free(&pc->pre_conv_wctx);
dac_decoder_free(&pc->dac);
upsample_stage_free(&pc->upsample);
tok_trans_free(&pc->transformer);
quant_decoder_free(&pc->qdec);
gf_close(&pc->gguf);
return false;
}
if (!enc_trans_load(&pc->enc_transformer, pc->gguf, pc->backend)) {
seanet_encoder_free(&pc->seanet);
wctx_free(&pc->pre_conv_wctx);
dac_decoder_free(&pc->dac);
upsample_stage_free(&pc->upsample);
tok_trans_free(&pc->transformer);
quant_decoder_free(&pc->qdec);
gf_close(&pc->gguf);
return false;
}
if (!enc_down_load(&pc->enc_downsample, pc->gguf, pc->backend)) {
enc_trans_free(&pc->enc_transformer);
seanet_encoder_free(&pc->seanet);
wctx_free(&pc->pre_conv_wctx);
dac_decoder_free(&pc->dac);
upsample_stage_free(&pc->upsample);
tok_trans_free(&pc->transformer);
quant_decoder_free(&pc->qdec);
gf_close(&pc->gguf);
return false;
}
if (!quant_encode_load(&pc->qenc, pc->gguf, pc->backend)) {
enc_down_free(&pc->enc_downsample);
enc_trans_free(&pc->enc_transformer);
seanet_encoder_free(&pc->seanet);
wctx_free(&pc->pre_conv_wctx);
dac_decoder_free(&pc->dac);
upsample_stage_free(&pc->upsample);
tok_trans_free(&pc->transformer);
quant_decoder_free(&pc->qdec);
gf_close(&pc->gguf);
return false;
}
// Encoder half (seanet, enc_transformer, enc_downsample, qenc)
// stays on disk until the first encode request.
pc->enc_loaded = false;
pc->sched = backend_sched_new(bp, 4096);
@@ -217,10 +176,44 @@ std::vector<float> pipeline_codec_decode(PipelineCodec * pc, const int32_t * cod
return audio;
}
bool pipeline_codec_ensure_encoder(PipelineCodec * pc) {
if (pc->enc_loaded) {
return true;
}
Timer t_load;
if (!seanet_encoder_load(&pc->seanet, pc->gguf, pc->backend)) {
return false;
}
if (!enc_trans_load(&pc->enc_transformer, pc->gguf, pc->backend)) {
seanet_encoder_free(&pc->seanet);
return false;
}
if (!enc_down_load(&pc->enc_downsample, pc->gguf, pc->backend)) {
enc_trans_free(&pc->enc_transformer);
seanet_encoder_free(&pc->seanet);
return false;
}
if (!quant_encode_load(&pc->qenc, pc->gguf, pc->backend)) {
enc_down_free(&pc->enc_downsample);
enc_trans_free(&pc->enc_transformer);
seanet_encoder_free(&pc->seanet);
return false;
}
pc->enc_loaded = true;
qt_log(QT_LOG_INFO, "[Pipeline] Codec encoder lazy loaded in %.0f ms", t_load.ms());
return true;
}
std::vector<int32_t> pipeline_codec_encode(PipelineCodec * pc,
const float * audio,
int n_samples,
const char * dump_dir) {
if (!pipeline_codec_ensure_encoder(pc)) {
qt_log(QT_LOG_ERROR, "[Pipeline] codec encoder load failed");
return {};
}
if (n_samples <= 0 || (n_samples % TOKENIZER_HOP_LENGTH) != 0) {
qt_log(QT_LOG_ERROR, "[Pipeline] n_samples must be a positive multiple of %d (got %d)", TOKENIZER_HOP_LENGTH,
n_samples);
@@ -434,10 +427,13 @@ void pipeline_codec_free(PipelineCodec * pc) {
ggml_backend_sched_free(pc->sched);
pc->sched = NULL;
}
quant_encode_free(&pc->qenc);
enc_down_free(&pc->enc_downsample);
enc_trans_free(&pc->enc_transformer);
seanet_encoder_free(&pc->seanet);
if (pc->enc_loaded) {
quant_encode_free(&pc->qenc);
enc_down_free(&pc->enc_downsample);
enc_trans_free(&pc->enc_transformer);
seanet_encoder_free(&pc->seanet);
pc->enc_loaded = false;
}
wctx_free(&pc->pre_conv_wctx);
dac_decoder_free(&pc->dac);
upsample_stage_free(&pc->upsample);
src/pipeline-codec.h
+10 -1
View File
@@ -63,7 +63,12 @@ struct PipelineCodec {
QwenSEANetEncoder seanet;
QwenEncoderTransformer enc_transformer;
QwenEncoderDownsample enc_downsample;
QwenQuantizerEncode qenc;
// Encoder weights (seanet, enc_transformer, enc_downsample, qenc)
// load lazily on the first pipeline_codec_encode call: synthesis
// from pre encoded reference codes never pays for them.
bool enc_loaded;
QwenQuantizerEncode qenc;
// CPU mirror of the RVQ encode side, lazy-loaded on first encode call.
QwenQuantizerEncodeHost qenc_sem_host;
@@ -79,6 +84,10 @@ struct PipelineCodec {
// On failure leaves the struct in a clean state and returns false.
bool pipeline_codec_load(PipelineCodec * pc, const char * gguf_path, BackendPair bp);
// Load the encoder weights on demand. Idempotent, called by
// pipeline_codec_encode; harmless to call when already resident.
bool pipeline_codec_ensure_encoder(PipelineCodec * pc);
// Decode RVQ codes into a 24 kHz mono waveform.
// codes: flat int32 buffer, [K, T] row-major (T fastest).
// Returns audio of length T * TOKENIZER_HOP_LENGTH, empty on failure.
src/pipeline-tts.cpp
+65 -71
View File
@@ -158,23 +158,14 @@ bool pipeline_tts_load(PipelineTTS * pt,
return false;
}
// Speaker encoder is only present in Base checkpoints. Treat absence
// as a soft condition: voice clone path stays disabled, base-direct
// synthesis still works.
if (pt->model_type == "base") {
if (!speaker_encoder_weights_load(&pt->speaker_encoder, pt->gguf_talker, pt->backend)) {
code_predictor_weights_free(&pt->code_predictor);
talker_weights_free(&pt->talker);
gf_close(&pt->gguf_talker);
return false;
}
pt->has_speaker_encoder = (pt->speaker_encoder.weight_buf != NULL);
}
// Speaker encoder tensors are only present in Base checkpoints. The
// weights load lazily on the first --ref-wav request: synthesis from
// pre extracted embeddings never pays for them. has_speaker_encoder
// advertises the capability, spk_enc_loaded tracks residency.
pt->has_speaker_encoder = (pt->model_type == "base");
pt->spk_enc_loaded = false;
if (!pipeline_codec_load(&pt->codec, codec_gguf_path, bp)) {
if (pt->has_speaker_encoder) {
speaker_encoder_weights_free(&pt->speaker_encoder);
}
code_predictor_weights_free(&pt->code_predictor);
talker_weights_free(&pt->talker);
gf_close(&pt->gguf_talker);
@@ -189,9 +180,6 @@ bool pipeline_tts_load(PipelineTTS * pt,
pt->sched = backend_sched_new(bp, 4096);
if (!pt->sched) {
pipeline_codec_free(&pt->codec);
if (pt->has_speaker_encoder) {
speaker_encoder_weights_free(&pt->speaker_encoder);
}
code_predictor_weights_free(&pt->code_predictor);
talker_weights_free(&pt->talker);
gf_close(&pt->gguf_talker);
@@ -204,9 +192,6 @@ bool pipeline_tts_load(PipelineTTS * pt,
ggml_backend_sched_free(pt->sched);
pt->sched = NULL;
pipeline_codec_free(&pt->codec);
if (pt->has_speaker_encoder) {
speaker_encoder_weights_free(&pt->speaker_encoder);
}
code_predictor_weights_free(&pt->code_predictor);
talker_weights_free(&pt->talker);
gf_close(&pt->gguf_talker);
@@ -221,9 +206,6 @@ bool pipeline_tts_load(PipelineTTS * pt,
ggml_backend_sched_free(pt->sched);
pt->sched = NULL;
pipeline_codec_free(&pt->codec);
if (pt->has_speaker_encoder) {
speaker_encoder_weights_free(&pt->speaker_encoder);
}
code_predictor_weights_free(&pt->code_predictor);
talker_weights_free(&pt->talker);
gf_close(&pt->gguf_talker);
@@ -236,9 +218,6 @@ bool pipeline_tts_load(PipelineTTS * pt,
ggml_backend_sched_free(pt->sched);
pt->sched = NULL;
pipeline_codec_free(&pt->codec);
if (pt->has_speaker_encoder) {
speaker_encoder_weights_free(&pt->speaker_encoder);
}
code_predictor_weights_free(&pt->code_predictor);
talker_weights_free(&pt->talker);
gf_close(&pt->gguf_talker);
@@ -259,9 +238,6 @@ bool pipeline_tts_load(PipelineTTS * pt,
ggml_backend_sched_free(pt->sched);
pt->sched = NULL;
pipeline_codec_free(&pt->codec);
if (pt->has_speaker_encoder) {
speaker_encoder_weights_free(&pt->speaker_encoder);
}
code_predictor_weights_free(&pt->code_predictor);
talker_weights_free(&pt->talker);
gf_close(&pt->gguf_talker);
@@ -272,7 +248,7 @@ bool pipeline_tts_load(PipelineTTS * pt,
"[Pipeline] Loaded: arch=%s variant=%s tokenizer=%s codebooks=%d speaker_encoder=%s speakers=%zu fa=%s "
"clamp_fp16=%s",
pt->model_size.c_str(), pt->model_type.c_str(), pt->tokenizer_type.c_str(), pt->num_code_groups,
pt->has_speaker_encoder ? "loaded" : "absent", pt->speakers.size(), pt->use_flash_attn ? "on" : "off",
pt->has_speaker_encoder ? "deferred" : "absent", pt->speakers.size(), pt->use_flash_attn ? "on" : "off",
pt->clamp_fp16 ? "on" : "off");
return true;
}
@@ -288,8 +264,9 @@ void pipeline_tts_free(PipelineTTS * pt) {
pt->sched = NULL;
}
pipeline_codec_free(&pt->codec);
if (pt->has_speaker_encoder) {
if (pt->spk_enc_loaded) {
speaker_encoder_weights_free(&pt->speaker_encoder);
pt->spk_enc_loaded = false;
}
code_predictor_weights_free(&pt->code_predictor);
talker_weights_free(&pt->talker);
@@ -447,11 +424,24 @@ qt_status pipeline_tts_synthesize(PipelineTTS * pt,
qt_log(QT_LOG_INFO, "[Pipeline] Latent speaker embedding: %d values", lat_spk_dim);
} else if (has_ref_audio) {
if (!pt->has_speaker_encoder) {
qt_set_error(
"pipeline_tts_synthesize: --ref-wav requires a model with a loaded speaker encoder (Base only)");
qt_log(QT_LOG_ERROR, "[Pipeline] --ref-wav requires a model with a loaded speaker encoder (Base only)");
qt_set_error("pipeline_tts_synthesize: --ref-wav requires a model with a speaker encoder (Base only)");
qt_log(QT_LOG_ERROR, "[Pipeline] --ref-wav requires a model with a speaker encoder (Base only)");
return QT_STATUS_GENERATE_FAILED;
}
// Lazy residency: the first reference audio request pays the
// weight load once, pre extracted paths never do.
if (!pt->spk_enc_loaded) {
Timer t_spk_load;
if (!speaker_encoder_weights_load(&pt->speaker_encoder, pt->gguf_talker, pt->backend) ||
pt->speaker_encoder.weight_buf == NULL) {
pt->has_speaker_encoder = false;
qt_set_error("pipeline_tts_synthesize: speaker encoder load failed");
qt_log(QT_LOG_ERROR, "[Pipeline] speaker encoder load failed");
return QT_STATUS_GENERATE_FAILED;
}
pt->spk_enc_loaded = true;
qt_log(QT_LOG_INFO, "[Pipeline] Speaker encoder lazy loaded in %.0f ms", t_spk_load.ms());
}
if (!speaker_encoder_extract(&pt->speaker_encoder, pt->sched, params->ref_audio_24k, params->ref_n_samples,
ref_spk_emb, params->dump_dir)) {
return QT_STATUS_GENERATE_FAILED;
@@ -580,7 +570,11 @@ qt_status pipeline_tts_synthesize(PipelineTTS * pt,
// sample (one for c0 of each step, then 15 for the predictor codes).
int64_t subseq_counter = 0;
std::vector<float> next_emb((size_t) hidden, 0.0f);
// Decode input state: the codes sampled at the previous frame plus
// the trailing text / pad overlay row for that frame. The talker
// decode graph gathers and sums the 16 embeddings on device.
std::vector<int32_t> prev_ids((size_t) num_codebooks, 0);
const float * prev_overlay = NULL;
// Streaming rolling decoder. Holds the K major codes buffer, the
// emit cursor and the left context window. push_frame triggers an
@@ -608,8 +602,10 @@ qt_status pipeline_tts_synthesize(PipelineTTS * pt,
ok = talker_forward_prefill(&pt->talker, &pt->talker_kv, pt->sched, &pt->talker_arena,
prompt.input_embed.data(), prompt.T_ctx, use_fa, clamp_fp16, step_dump, &fw);
} else {
ok = talker_forward_decode(&pt->talker, &pt->talker_kv, pt->sched, &pt->talker_arena, next_emb.data(),
use_fa, clamp_fp16, &fw);
ok =
talker_forward_decode(&pt->talker, &pt->talker_kv, pt->sched, &pt->talker_arena, prev_ids.data(),
pt->code_predictor.codec_embedding.data(),
pt->code_predictor.num_acoustic_codebooks, prev_overlay, use_fa, clamp_fp16, &fw);
}
if (!ok) {
return QT_STATUS_GENERATE_FAILED;
@@ -693,43 +689,41 @@ qt_status pipeline_tts_synthesize(PipelineTTS * pt,
}
}
// Build next-token embedding: sum of 16 codebook embeddings.
// codebook 0 uses talker.codec_embedding, the 15 acoustic
// codebooks use the predictor's private embedding tables.
Timer t_emb;
std::fill(next_emb.begin(), next_emb.end(), 0.0f);
std::vector<float> tmp((size_t) hidden);
embed_row_from_gguf(pt->gguf_talker, "talker.codec_embd.weight", c0, hidden, tmp.data());
for (int i = 0; i < hidden; i++) {
next_emb[(size_t) i] += tmp[(size_t) i];
// Next decode input: the 16 frame codes gather and sum in graph
// (codebook 0 from talker.codec_embedding, the 15 acoustic
// groups from the predictor's private tables). The overlay row
// adds the next utterance text hidden while any remains, the
// tts_pad embedding afterwards.
prev_ids[0] = c0;
for (int g = 1; g < num_codebooks; g++) {
prev_ids[(size_t) g] = cp.codes[(size_t) g];
}
for (int g = 0; g < num_codebooks - 1; g++) {
int cg = cp.codes[(size_t) (g + 1)];
char name[64];
snprintf(name, sizeof(name), "code_pred.codec_embd.%d.weight", g);
embed_row_from_gguf(pt->gguf_talker, name, cg, hidden, tmp.data());
prev_overlay = (step < prompt.T_trailing) ?
prompt.trailing_text_hidden.data() + (size_t) step * (size_t) hidden :
prompt.tts_pad_embed.data();
// Bisection dump: reproduce the in graph composition on host so
// the step 0 next embedding stays byte comparable against the
// Python hook (codebook sums plus trailing text overlay).
if (params->dump_dir && step == 0) {
std::vector<float> next_emb((size_t) hidden, 0.0f);
std::vector<float> tmp((size_t) hidden);
embed_row_from_gguf(pt->gguf_talker, "talker.codec_embd.weight", c0, hidden, tmp.data());
for (int i = 0; i < hidden; i++) {
next_emb[(size_t) i] += tmp[(size_t) i];
}
}
// Trailing text overlay: while we still have utterance text
// hiddens to consume, add the next one; otherwise add the
// tts_pad embedding.
const float * overlay = (step < prompt.T_trailing) ?
prompt.trailing_text_hidden.data() + (size_t) step * (size_t) hidden :
prompt.tts_pad_embed.data();
for (int i = 0; i < hidden; i++) {
next_emb[(size_t) i] += overlay[(size_t) i];
}
perf.host_ms += t_emb.ms();
// Bisection dump: the next-token embedding produced at step 0
// is the only thing controlling the talker forward at step 1, so
// matching it bit-exact against Python pinpoints any drift in
// the codebook embedding sums or the trailing text overlay.
if (params->dump_dir && step == 0) {
for (int g = 0; g < num_codebooks - 1; g++) {
int cg = cp.codes[(size_t) (g + 1)];
char name[64];
snprintf(name, sizeof(name), "code_pred.codec_embd.%d.weight", g);
embed_row_from_gguf(pt->gguf_talker, name, cg, hidden, tmp.data());
for (int i = 0; i < hidden; i++) {
next_emb[(size_t) i] += tmp[(size_t) i];
}
}
for (int i = 0; i < hidden; i++) {
next_emb[(size_t) i] += prev_overlay[(size_t) i];
}
DebugDumper d;
debug_init(&d, params->dump_dir);
debug_dump_1d(&d, "next-emb-step0", next_emb.data(), hidden);
src/pipeline-tts.h
+4
View File
@@ -95,6 +95,10 @@ struct PipelineTTS {
SpeakerEncoderWeights speaker_encoder;
bool has_speaker_encoder;
// Speaker encoder weights residency: loaded lazily on the first
// reference audio request, see pipeline-tts.cpp.
bool spk_enc_loaded;
PipelineCodec codec;
std::string tokenizer_type;
src/talker-forward.h
+78 -31
View File
@@ -254,21 +254,28 @@ static struct ggml_tensor * talker_layer_forward(struct ggml_context * ctx,
// Shared core that builds the graph, allocates, uploads inputs, runs
// it and pulls out the last position hidden + logits. T tokens are
// appended to the cache starting at n_past. When n_past == 0 and
// dump_dir is set, the bisect taps fire on the prefill path. use_fa /
// clamp_fp16 are forwarded as is to every layer. The graph metadata
// lives in the caller owned persistent arena.
static bool talker_forward_core(const TalkerWeights * tw,
KVCache * kv,
ggml_backend_sched_t sched,
GraphArena * arena,
const float * input_embed,
int T,
int n_past,
bool use_flash_attn,
bool clamp_fp16,
const char * dump_dir,
TalkerForwardOutput * out) {
// appended to the cache starting at n_past. Input is either a raw
// embedding upload (input_embed, prefill) or the previous frame code
// ids plus overlay row assembled in graph (frame_ids, decode hot
// path). When n_past == 0 and dump_dir is set, the bisect taps fire on
// the prefill path. use_fa / clamp_fp16 are forwarded as is to every
// layer. The graph metadata lives in the caller owned persistent
// arena.
static bool talker_forward_core(const TalkerWeights * tw,
KVCache * kv,
ggml_backend_sched_t sched,
GraphArena * arena,
const float * input_embed,
const int32_t * frame_ids,
struct ggml_tensor * const * acoustic_embd,
int n_acoustic,
const float * overlay,
int T,
int n_past,
bool use_flash_attn,
bool clamp_fp16,
const char * dump_dir,
TalkerForwardOutput * out) {
const int hidden = tw->hidden_size;
const int n_layers = tw->num_hidden_layers;
const int vocab = tw->vocab_size;
@@ -285,17 +292,44 @@ static bool talker_forward_core(const TalkerWeights * tw,
const int max_nodes = talker_graph_max_nodes(n_layers);
struct ggml_context * gctx = graph_arena_begin(arena);
// IO tensors: input embedding, positions, causal mask. The mask
// spans [n_kv_pad, T]: for each fresh query q in [0, T) keys k in
// IO tensors: positions and causal mask. The mask spans
// [n_kv_pad, T]: for each fresh query q in [0, T) keys k in
// [0, n_past + q] carry 0 and every other slot carries neg inf,
// including the padded tail beyond T_full.
struct ggml_tensor * x_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F32, hidden, T);
struct ggml_tensor * pos_in = ggml_new_tensor_1d(gctx, GGML_TYPE_I32, T);
struct ggml_tensor * mask_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F16, n_kv_pad, T);
ggml_set_name(x_in, "input_embed");
ggml_set_name(pos_in, "positions");
ggml_set_name(mask_in, "causal_mask");
// Input: either a raw embedding upload (prefill path) or, on the
// decode hot path, the frame codes of the previous step gathered
// and summed in graph. x = get_rows(codec_embd, ids[0]) plus the 15
// acoustic group gathers plus the trailing text / pad overlay row.
// The only per step uploads are 16 code ids and one overlay row.
struct ggml_tensor * x_in = NULL;
struct ggml_tensor * ids_in = NULL;
struct ggml_tensor * overlay_in = NULL;
if (frame_ids) {
ids_in = ggml_new_tensor_1d(gctx, GGML_TYPE_I32, 1 + n_acoustic);
ggml_set_name(ids_in, "frame_code_ids");
ggml_set_input(ids_in);
overlay_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F32, hidden, 1);
ggml_set_name(overlay_in, "overlay_row");
ggml_set_input(overlay_in);
struct ggml_tensor * id0 = ggml_view_1d(gctx, ids_in, 1, 0);
x_in = ggml_get_rows(gctx, tw->codec_embedding, id0);
for (int g = 0; g < n_acoustic; g++) {
struct ggml_tensor * idg = ggml_view_1d(gctx, ids_in, 1, (size_t) (g + 1) * sizeof(int32_t));
x_in = ggml_add(gctx, x_in, ggml_get_rows(gctx, acoustic_embd[g], idg));
}
x_in = ggml_add(gctx, x_in, overlay_in);
} else {
x_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F32, hidden, T);
ggml_set_input(x_in);
}
ggml_set_name(x_in, "input_embed");
struct ggml_cgraph * gf = ggml_new_graph_custom(gctx, max_nodes, false);
// Build the layer stack. Bisect taps fire on prefill only.
@@ -346,7 +380,12 @@ static bool talker_forward_core(const TalkerWeights * tw,
}
// Upload input embedding (host [T, hidden] -> ggml [hidden, T]).
ggml_backend_tensor_set(x_in, input_embed, 0, (size_t) T * (size_t) hidden * sizeof(float));
if (frame_ids) {
ggml_backend_tensor_set(ids_in, frame_ids, 0, (size_t) (1 + n_acoustic) * sizeof(int32_t));
ggml_backend_tensor_set(overlay_in, overlay, 0, (size_t) hidden * sizeof(float));
} else {
ggml_backend_tensor_set(x_in, input_embed, 0, (size_t) T * (size_t) hidden * sizeof(float));
}
// Positions: n_past .. n_past + T - 1
{
@@ -444,25 +483,33 @@ static bool talker_forward_prefill(const TalkerWeights * tw,
fprintf(stderr, "[TalkerForward] FATAL: prefill T=%d exceeds cache max_seq_len=%d\n", T, kv->max_seq_len);
return false;
}
return talker_forward_core(tw, kv, sched, arena, input_embed, T, 0, use_flash_attn, clamp_fp16, dump_dir, out);
return talker_forward_core(tw, kv, sched, arena, input_embed, NULL, NULL, 0, NULL, T, 0, use_flash_attn, clamp_fp16,
dump_dir, out);
}
// Decode: feed exactly one embedding and append one position to the
// cache. Reads positions [0, kv->cur_len + 1). Caller is responsible
// for ensuring kv->cur_len + 1 <= kv->max_seq_len.
static bool talker_forward_decode(const TalkerWeights * tw,
KVCache * kv,
ggml_backend_sched_t sched,
GraphArena * arena,
const float * input_embed_1,
bool use_flash_attn,
bool clamp_fp16,
TalkerForwardOutput * out) {
// Append one position from the previous frame's codes. frame_ids holds
// [c0, c1..c15], acoustic_embd the 15 group tables owned by the code
// predictor, overlay the trailing text / pad row summed on top. The
// input embedding assembles entirely in graph.
static bool talker_forward_decode(const TalkerWeights * tw,
KVCache * kv,
ggml_backend_sched_t sched,
GraphArena * arena,
const int32_t * frame_ids,
struct ggml_tensor * const * acoustic_embd,
int n_acoustic,
const float * overlay,
bool use_flash_attn,
bool clamp_fp16,
TalkerForwardOutput * out) {
if (kv->cur_len + 1 > kv->max_seq_len) {
fprintf(stderr, "[TalkerForward] FATAL: decode would overflow cache (%d + 1 > %d)\n", kv->cur_len,
kv->max_seq_len);
return false;
}
return talker_forward_core(tw, kv, sched, arena, input_embed_1, 1, kv->cur_len, use_flash_attn, clamp_fp16, NULL,
out);
return talker_forward_core(tw, kv, sched, arena, NULL, frame_ids, acoustic_embd, n_acoustic, overlay, 1,
kv->cur_len, use_flash_attn, clamp_fp16, NULL, out);
}