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

src: fuse 3 .cpp into header-only modules to match omnivoice

Browse Source
This commit is contained in:
Pascal
2026-05-14 17:14:30 +02:00
parent e8f1c2a053
commit e55ba50b75
7 changed files with 1291 additions and 1378 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/code-predictor-forward.cpp
-340
View File
@@ -1,340 +0,0 @@
// code-predictor-forward.cpp: eager full-recompute graph for the
// Qwen3-TTS code predictor (5-layer Qwen3 stack with plain 1D NEOX
// RoPE, GQA attention with QK-norm, SwiGLU MLP, head-per-codebook
// output projection).
//
// The predictor architecture mirrors the Talker block, the only
// differences are :
// - 5 layers instead of 28
// - plain 1D RoPE (no multimodal sections)
// - one private embedding table and one private linear head per
// acoustic codebook (1..15)
//
// The single-frame loop here recomputes the full graph at every step g
// (0..14) over a sequence of length g+2. With 5 layers and at most 16
// tokens per recompute this is sub-millisecond on modern GPUs.
#include "code-predictor-forward.h"
#include "debug.h"
#include "ggml-alloc.h"
#include "ggml-backend.h"
#include "ggml.h"
#include "qt-error.h"
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <vector>
// One Qwen3 decoder block, KV cached. K and V for the T fresh positions
// are written into the cache at [n_past, n_past+T) on dim 1; the
// attention reads the contiguous slice [0, n_past+T). Returns the layer
// output [hidden, T].
static struct ggml_tensor * code_predictor_layer_forward(struct ggml_context * ctx,
const CodePredictorWeights * cw,
const TalkerLayer & layer,
struct ggml_tensor * x,
struct ggml_tensor * positions,
struct ggml_tensor * mask,
struct ggml_tensor * k_cache,
struct ggml_tensor * v_cache,
int n_past,
int T,
struct ggml_cgraph * gf) {
const int n_q_heads = cw->num_attention_heads;
const int n_kv = cw->num_key_value_heads;
const int hd = cw->head_dim;
const float eps = cw->rms_norm_eps;
struct ggml_tensor * h = ggml_rms_norm(ctx, x, eps);
h = ggml_mul(ctx, h, layer.input_norm_w);
struct ggml_tensor * q = ggml_mul_mat(ctx, layer.attn.q_proj_w, h);
struct ggml_tensor * k = ggml_mul_mat(ctx, layer.attn.k_proj_w, h);
struct ggml_tensor * v = ggml_mul_mat(ctx, layer.attn.v_proj_w, h);
q = ggml_reshape_3d(ctx, q, hd, n_q_heads, T);
k = ggml_reshape_3d(ctx, k, hd, n_kv, T);
v = ggml_reshape_3d(ctx, v, hd, n_kv, T);
q = ggml_rms_norm(ctx, q, eps);
q = ggml_mul(ctx, q, layer.attn.q_norm_w);
k = ggml_rms_norm(ctx, k, eps);
k = ggml_mul(ctx, k, layer.attn.k_norm_w);
q = ggml_rope_ext(ctx, q, positions, NULL, hd, GGML_ROPE_TYPE_NEOX, 0, cw->rope_theta, 1.0f, 0.0f, 1.0f, 0.0f,
0.0f);
k = ggml_rope_ext(ctx, k, positions, NULL, hd, GGML_ROPE_TYPE_NEOX, 0, cw->rope_theta, 1.0f, 0.0f, 1.0f, 0.0f,
0.0f);
// Write the fresh positions into the cache.
struct ggml_tensor * k_perm = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3)); // [hd, T, n_kv]
struct ggml_tensor * v_perm = ggml_cont(ctx, ggml_permute(ctx, v, 0, 2, 1, 3));
size_t k_off = (size_t) n_past * k_cache->nb[1];
size_t v_off = (size_t) n_past * v_cache->nb[1];
struct ggml_tensor * k_dst = ggml_view_3d(ctx, k_cache, hd, T, n_kv, k_cache->nb[1], k_cache->nb[2], k_off);
struct ggml_tensor * v_dst = ggml_view_3d(ctx, v_cache, hd, T, n_kv, v_cache->nb[1], v_cache->nb[2], v_off);
struct ggml_tensor * k_cpy = ggml_cpy(ctx, k_perm, k_dst);
struct ggml_tensor * v_cpy = ggml_cpy(ctx, v_perm, v_dst);
ggml_build_forward_expand(gf, k_cpy);
ggml_build_forward_expand(gf, v_cpy);
const int T_full = n_past + T;
struct ggml_tensor * k_full = ggml_view_3d(ctx, k_cache, hd, T_full, n_kv, k_cache->nb[1], k_cache->nb[2], 0);
struct ggml_tensor * v_full = ggml_view_3d(ctx, v_cache, hd, T_full, n_kv, v_cache->nb[1], v_cache->nb[2], 0);
// Q permute [hd, n_q_heads, T] -> [hd, T, n_q_heads] for flash_attn_ext.
struct ggml_tensor * q_p = ggml_permute(ctx, q, 0, 2, 1, 3);
// Fused flash attention. Matches the working acestep qw3lm_build_attn
// pattern, fixes the Vulkan autoregressive decode bug.
float scale = 1.0f / sqrtf((float) hd);
struct ggml_tensor * attn = ggml_flash_attn_ext(ctx, q_p, k_full, v_full, mask, scale, 0.0f, 0.0f);
ggml_flash_attn_ext_set_prec(attn, GGML_PREC_F32);
attn = ggml_reshape_2d(ctx, attn, n_q_heads * hd, T);
struct ggml_tensor * o = ggml_mul_mat(ctx, layer.attn.o_proj_w, attn);
x = ggml_add(ctx, x, o);
struct ggml_tensor * h2 = ggml_rms_norm(ctx, x, eps);
h2 = ggml_mul(ctx, h2, layer.post_attn_norm_w);
struct ggml_tensor * gate = ggml_mul_mat(ctx, layer.mlp.gate_proj_w, h2);
struct ggml_tensor * up = ggml_mul_mat(ctx, layer.mlp.up_proj_w, h2);
gate = ggml_silu(ctx, gate);
struct ggml_tensor * gu = ggml_mul(ctx, gate, up);
struct ggml_tensor * mlp = ggml_mul_mat(ctx, layer.mlp.down_proj_w, gu);
x = ggml_add(ctx, x, mlp);
return x;
}
// Run one predictor pass: feed `T` fresh embeddings starting at cache
// position `n_past`, run all 5 layers, and pull the logits for the last
// position through lm_head[g_head]. The cache is written as a side
// effect so subsequent decode steps can append a single token.
static bool code_predictor_run(const CodePredictorWeights * cw,
KVCache * kv,
ggml_backend_sched_t sched,
const float * fresh_input,
int T,
int n_past,
int talker_hidden,
int g_head,
std::vector<float> * logits_out) {
const int vocab = cw->vocab_size;
const int n_layers = cw->num_hidden_layers;
const int T_full = n_past + T;
const int max_nodes = 48 * n_layers + 64;
const size_t arena_bytes = ggml_tensor_overhead() * max_nodes + ggml_graph_overhead_custom(max_nodes, false);
struct ggml_init_params gp = { arena_bytes, NULL, true };
struct ggml_context * gctx = ggml_init(gp);
if (!gctx) {
fprintf(stderr, "[CodePredictor] FATAL: ggml_init failed\n");
return false;
}
// Inputs: fresh embeddings (talker_hidden), positions, attention mask
struct ggml_tensor * x_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F32, talker_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, T_full, T);
ggml_set_name(x_in, "sub_input");
ggml_set_name(pos_in, "positions");
ggml_set_name(mask_in, "causal_mask");
struct ggml_cgraph * gf = ggml_new_graph_custom(gctx, max_nodes, false);
// small_to_mtp projection: Linear(talker_hidden -> hidden) with bias.
// When absent (Identity case) the input is already at predictor hidden.
struct ggml_tensor * h = x_in;
if (cw->mtp_proj_w) {
h = ggml_mul_mat(gctx, cw->mtp_proj_w, h);
if (cw->mtp_proj_b) {
h = ggml_add(gctx, h, cw->mtp_proj_b);
}
ggml_set_name(h, "mtp_proj_out");
}
for (int l = 0; l < n_layers; l++) {
h = code_predictor_layer_forward(gctx, cw, cw->layers[(size_t) l], h, pos_in, mask_in, kv->k[(size_t) l],
kv->v[(size_t) l], n_past, T, gf);
}
struct ggml_tensor * h_final = ggml_rms_norm(gctx, h, cw->rms_norm_eps);
h_final = ggml_mul(gctx, h_final, cw->norm_w);
struct ggml_tensor * logits = ggml_mul_mat(gctx, cw->lm_head[(size_t) g_head], h_final);
ggml_set_name(logits, "logits");
ggml_set_output(logits);
ggml_build_forward_expand(gf, logits);
if (!ggml_backend_sched_alloc_graph(sched, gf)) {
fprintf(stderr, "[CodePredictor] FATAL: graph allocation failed\n");
ggml_backend_sched_reset(sched);
ggml_free(gctx);
return false;
}
ggml_backend_tensor_set(x_in, fresh_input, 0, (size_t) T * (size_t) talker_hidden * sizeof(float));
{
std::vector<int32_t> pos((size_t) T);
for (int i = 0; i < T; i++) {
pos[(size_t) i] = n_past + i;
}
ggml_backend_tensor_set(pos_in, pos.data(), 0, (size_t) T * sizeof(int32_t));
}
{
std::vector<ggml_fp16_t> mask((size_t) T * (size_t) T_full);
const ggml_fp16_t zero = ggml_fp32_to_fp16(0.0f);
const ggml_fp16_t neg_inf = ggml_fp32_to_fp16(-INFINITY);
for (size_t i = 0; i < mask.size(); i++) {
mask[i] = neg_inf;
}
for (int q = 0; q < T; q++) {
const int q_pos = n_past + q;
for (int k = 0; k <= q_pos; k++) {
mask[(size_t) q * (size_t) T_full + (size_t) k] = zero;
}
}
ggml_backend_tensor_set(mask_in, mask.data(), 0, mask.size() * sizeof(ggml_fp16_t));
}
if (ggml_backend_sched_graph_compute(sched, gf) != GGML_STATUS_SUCCESS) {
fprintf(stderr, "[CodePredictor] FATAL: graph compute failed\n");
ggml_backend_sched_reset(sched);
ggml_free(gctx);
return false;
}
logits_out->resize((size_t) vocab);
size_t row_bytes = (size_t) vocab * sizeof(float);
ggml_backend_tensor_get(logits, logits_out->data(), (size_t) (T - 1) * row_bytes, row_bytes);
kv->cur_len = T_full;
ggml_backend_sched_reset(sched);
ggml_free(gctx);
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);
}
bool code_predictor_step(const TalkerWeights * tw,
const CodePredictorWeights * cw,
KVCache * kv,
ggml_backend_sched_t sched,
const float * talker_hidden_last,
int c0,
float temperature,
int top_k,
float top_p,
int64_t seed,
int64_t subseq_base,
const char * dump_dir,
CodePredictorOutput * out) {
// sub_input slots live at the talker hidden dimension: both the
// talker last hidden and the codec_embedding rows feeding the sub
// network are talker sized in the upstream checkpoint. The graph's
// mtp_proj brings them down to predictor hidden when present.
const int talker_hidden = tw->hidden_size;
const int n_acoustic = cw->num_acoustic_codebooks;
if (n_acoustic + 1 > kv->max_seq_len) {
fprintf(stderr, "[CodePredictor] FATAL: frame width %d exceeds cache max_seq_len %d\n", n_acoustic + 1,
kv->max_seq_len);
return false;
}
out->codes.assign((size_t) (n_acoustic + 1), 0);
out->codes[0] = c0;
// Prefill: two positions, talker_hidden_last and embed_talker(c0).
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, prefill_input.data(), 2, 0, talker_hidden, 0, &logits)) {
return false;
}
{
float u_g = 0.0f;
int cg = sample_top_k_p(logits.data(), (int) logits.size(), temperature, top_k, top_p, 1.0f, nullptr, 0, seed,
subseq_base + 1, &u_g);
if (subseq_base + 1 < 32) {
fprintf(stderr, "[Sample-CP] g=0 c=%d u=%.10f subseq=%lld\n", cg, (double) u_g,
(long long) (subseq_base + 1));
}
if (cg < 0) {
fprintf(stderr, "[CodePredictor] FATAL: sample returned no candidate at g=0\n");
return false;
}
out->codes[1] = cg;
}
// 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);
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, step_input.data(), 1, kv->cur_len, talker_hidden, g, &logits)) {
return false;
}
float u_g = 0.0f;
int cg = sample_top_k_p(logits.data(), (int) logits.size(), temperature, top_k, top_p, 1.0f, nullptr, 0, seed,
subseq_base + 1 + g, &u_g);
if (subseq_base + 1 + g < 32) {
fprintf(stderr, "[Sample-CP] g=%d c=%d u=%.10f subseq=%lld\n", g, cg, (double) u_g,
(long long) (subseq_base + 1 + g));
}
if (cg < 0) {
fprintf(stderr, "[CodePredictor] FATAL: sample returned no candidate at g=%d\n", g);
return false;
}
out->codes[(size_t) (g + 1)] = cg;
}
if (dump_dir) {
DebugDumper d;
debug_init(&d, dump_dir);
std::vector<int32_t> codes32(out->codes.begin(), out->codes.end());
int n = (int) codes32.size();
debug_dump_i32_as_f32(&d, "codes-step0", codes32.data(), &n, 1);
}
return true;
}
src/code-predictor-forward.h
+330 -15
View File
@@ -3,12 +3,12 @@
// growing context to produce the 15 acoustic codes of one audio frame,
// KV cached.
//
// Input :
// Input:
// talker_hidden_last [hidden] f32 -- last position hidden state from
// the Talker forward (post final norm)
// c0 -- semantic code sampled from the
// Talker codec_head (codebook 0)
// Output :
// Output:
// codes[16] = [c0, c1, ..., c15] -- the full set of codes for one
// frame, ready for decode through
// the codec
@@ -18,14 +18,32 @@
// then decode 14 single-token steps. Total work drops from
// O(sum_{g=0..14} (g+2)^2) = O(1496 token-steps) to O(16) per frame,
// roughly 90x for the inner loop.
//
// Architecture mirrors the Talker block, only differences are:
// - 5 layers instead of 28
// - plain 1D RoPE (no multimodal sections)
// - one private embedding table and one private linear head per
// acoustic codebook (1..15)
//
// The single-frame loop here recomputes the full graph at every step g
// (0..14) over a sequence of length g+2. With 5 layers and at most 16
// tokens per recompute this is sub-millisecond on modern GPUs.
#include "code-predictor-weights.h"
#include "debug.h"
#include "ggml-alloc.h"
#include "ggml-backend.h"
#include "ggml.h"
#include "kv-cache.h"
#include "qt-error.h"
#include "sampling.h"
#include "talker-weights.h"
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <vector>
struct CodePredictorOutput {
@@ -33,22 +51,319 @@ struct CodePredictorOutput {
std::vector<int32_t> codes;
};
// One Qwen3 decoder block, KV cached. K and V for the T fresh positions
// are written into the cache at [n_past, n_past+T) on dim 1; the
// attention reads the contiguous slice [0, n_past+T). Returns the layer
// output [hidden, T].
static struct ggml_tensor * code_predictor_layer_forward(struct ggml_context * ctx,
const CodePredictorWeights * cw,
const TalkerLayer & layer,
struct ggml_tensor * x,
struct ggml_tensor * positions,
struct ggml_tensor * mask,
struct ggml_tensor * k_cache,
struct ggml_tensor * v_cache,
int n_past,
int T,
struct ggml_cgraph * gf) {
const int n_q_heads = cw->num_attention_heads;
const int n_kv = cw->num_key_value_heads;
const int hd = cw->head_dim;
const float eps = cw->rms_norm_eps;
struct ggml_tensor * h = ggml_rms_norm(ctx, x, eps);
h = ggml_mul(ctx, h, layer.input_norm_w);
struct ggml_tensor * q = ggml_mul_mat(ctx, layer.attn.q_proj_w, h);
struct ggml_tensor * k = ggml_mul_mat(ctx, layer.attn.k_proj_w, h);
struct ggml_tensor * v = ggml_mul_mat(ctx, layer.attn.v_proj_w, h);
q = ggml_reshape_3d(ctx, q, hd, n_q_heads, T);
k = ggml_reshape_3d(ctx, k, hd, n_kv, T);
v = ggml_reshape_3d(ctx, v, hd, n_kv, T);
q = ggml_rms_norm(ctx, q, eps);
q = ggml_mul(ctx, q, layer.attn.q_norm_w);
k = ggml_rms_norm(ctx, k, eps);
k = ggml_mul(ctx, k, layer.attn.k_norm_w);
q = ggml_rope_ext(ctx, q, positions, NULL, hd, GGML_ROPE_TYPE_NEOX, 0, cw->rope_theta, 1.0f, 0.0f, 1.0f, 0.0f,
0.0f);
k = ggml_rope_ext(ctx, k, positions, NULL, hd, GGML_ROPE_TYPE_NEOX, 0, cw->rope_theta, 1.0f, 0.0f, 1.0f, 0.0f,
0.0f);
// Write the fresh positions into the cache.
struct ggml_tensor * k_perm = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3)); // [hd, T, n_kv]
struct ggml_tensor * v_perm = ggml_cont(ctx, ggml_permute(ctx, v, 0, 2, 1, 3));
size_t k_off = (size_t) n_past * k_cache->nb[1];
size_t v_off = (size_t) n_past * v_cache->nb[1];
struct ggml_tensor * k_dst = ggml_view_3d(ctx, k_cache, hd, T, n_kv, k_cache->nb[1], k_cache->nb[2], k_off);
struct ggml_tensor * v_dst = ggml_view_3d(ctx, v_cache, hd, T, n_kv, v_cache->nb[1], v_cache->nb[2], v_off);
struct ggml_tensor * k_cpy = ggml_cpy(ctx, k_perm, k_dst);
struct ggml_tensor * v_cpy = ggml_cpy(ctx, v_perm, v_dst);
ggml_build_forward_expand(gf, k_cpy);
ggml_build_forward_expand(gf, v_cpy);
const int T_full = n_past + T;
struct ggml_tensor * k_full = ggml_view_3d(ctx, k_cache, hd, T_full, n_kv, k_cache->nb[1], k_cache->nb[2], 0);
struct ggml_tensor * v_full = ggml_view_3d(ctx, v_cache, hd, T_full, n_kv, v_cache->nb[1], v_cache->nb[2], 0);
// Q permute [hd, n_q_heads, T] -> [hd, T, n_q_heads] for flash_attn_ext.
struct ggml_tensor * q_p = ggml_permute(ctx, q, 0, 2, 1, 3);
// Fused flash attention. Matches the working acestep qw3lm_build_attn
// pattern, fixes the Vulkan autoregressive decode bug.
float scale = 1.0f / sqrtf((float) hd);
struct ggml_tensor * attn = ggml_flash_attn_ext(ctx, q_p, k_full, v_full, mask, scale, 0.0f, 0.0f);
ggml_flash_attn_ext_set_prec(attn, GGML_PREC_F32);
attn = ggml_reshape_2d(ctx, attn, n_q_heads * hd, T);
struct ggml_tensor * o = ggml_mul_mat(ctx, layer.attn.o_proj_w, attn);
x = ggml_add(ctx, x, o);
struct ggml_tensor * h2 = ggml_rms_norm(ctx, x, eps);
h2 = ggml_mul(ctx, h2, layer.post_attn_norm_w);
struct ggml_tensor * gate = ggml_mul_mat(ctx, layer.mlp.gate_proj_w, h2);
struct ggml_tensor * up = ggml_mul_mat(ctx, layer.mlp.up_proj_w, h2);
gate = ggml_silu(ctx, gate);
struct ggml_tensor * gu = ggml_mul(ctx, gate, up);
struct ggml_tensor * mlp = ggml_mul_mat(ctx, layer.mlp.down_proj_w, gu);
x = ggml_add(ctx, x, mlp);
return x;
}
// Run one predictor pass: feed `T` fresh embeddings starting at cache
// position `n_past`, run all 5 layers, and pull the logits for the last
// position through lm_head[g_head]. The cache is written as a side
// effect so subsequent decode steps can append a single token.
static bool code_predictor_run(const CodePredictorWeights * cw,
KVCache * kv,
ggml_backend_sched_t sched,
const float * fresh_input,
int T,
int n_past,
int talker_hidden,
int g_head,
std::vector<float> * logits_out) {
const int vocab = cw->vocab_size;
const int n_layers = cw->num_hidden_layers;
const int T_full = n_past + T;
const int max_nodes = 48 * n_layers + 64;
const size_t arena_bytes = ggml_tensor_overhead() * max_nodes + ggml_graph_overhead_custom(max_nodes, false);
struct ggml_init_params gp = { arena_bytes, NULL, true };
struct ggml_context * gctx = ggml_init(gp);
if (!gctx) {
fprintf(stderr, "[CodePredictor] FATAL: ggml_init failed\n");
return false;
}
// Inputs: fresh embeddings (talker_hidden), positions, attention mask
struct ggml_tensor * x_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F32, talker_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, T_full, T);
ggml_set_name(x_in, "sub_input");
ggml_set_name(pos_in, "positions");
ggml_set_name(mask_in, "causal_mask");
struct ggml_cgraph * gf = ggml_new_graph_custom(gctx, max_nodes, false);
// small_to_mtp projection: Linear(talker_hidden -> hidden) with bias.
// When absent (Identity case) the input is already at predictor hidden.
struct ggml_tensor * h = x_in;
if (cw->mtp_proj_w) {
h = ggml_mul_mat(gctx, cw->mtp_proj_w, h);
if (cw->mtp_proj_b) {
h = ggml_add(gctx, h, cw->mtp_proj_b);
}
ggml_set_name(h, "mtp_proj_out");
}
for (int l = 0; l < n_layers; l++) {
h = code_predictor_layer_forward(gctx, cw, cw->layers[(size_t) l], h, pos_in, mask_in, kv->k[(size_t) l],
kv->v[(size_t) l], n_past, T, gf);
}
struct ggml_tensor * h_final = ggml_rms_norm(gctx, h, cw->rms_norm_eps);
h_final = ggml_mul(gctx, h_final, cw->norm_w);
struct ggml_tensor * logits = ggml_mul_mat(gctx, cw->lm_head[(size_t) g_head], h_final);
ggml_set_name(logits, "logits");
ggml_set_output(logits);
ggml_build_forward_expand(gf, logits);
if (!ggml_backend_sched_alloc_graph(sched, gf)) {
fprintf(stderr, "[CodePredictor] FATAL: graph allocation failed\n");
ggml_backend_sched_reset(sched);
ggml_free(gctx);
return false;
}
ggml_backend_tensor_set(x_in, fresh_input, 0, (size_t) T * (size_t) talker_hidden * sizeof(float));
{
std::vector<int32_t> pos((size_t) T);
for (int i = 0; i < T; i++) {
pos[(size_t) i] = n_past + i;
}
ggml_backend_tensor_set(pos_in, pos.data(), 0, (size_t) T * sizeof(int32_t));
}
{
std::vector<ggml_fp16_t> mask((size_t) T * (size_t) T_full);
const ggml_fp16_t zero = ggml_fp32_to_fp16(0.0f);
const ggml_fp16_t neg_inf = ggml_fp32_to_fp16(-INFINITY);
for (size_t i = 0; i < mask.size(); i++) {
mask[i] = neg_inf;
}
for (int q = 0; q < T; q++) {
const int q_pos = n_past + q;
for (int k = 0; k <= q_pos; k++) {
mask[(size_t) q * (size_t) T_full + (size_t) k] = zero;
}
}
ggml_backend_tensor_set(mask_in, mask.data(), 0, mask.size() * sizeof(ggml_fp16_t));
}
if (ggml_backend_sched_graph_compute(sched, gf) != GGML_STATUS_SUCCESS) {
fprintf(stderr, "[CodePredictor] FATAL: graph compute failed\n");
ggml_backend_sched_reset(sched);
ggml_free(gctx);
return false;
}
logits_out->resize((size_t) vocab);
size_t row_bytes = (size_t) vocab * sizeof(float);
ggml_backend_tensor_get(logits, logits_out->data(), (size_t) (T - 1) * row_bytes, row_bytes);
kv->cur_len = T_full;
ggml_backend_sched_reset(sched);
ggml_free(gctx);
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
// is the Philox subsequence of the c0 sample for this step; the 15
// acoustic samples consume subseq_base + 1 .. subseq_base + 15.
// Returns the full vector of 16 codes. dump_dir may be NULL.
bool code_predictor_step(const TalkerWeights * tw,
const CodePredictorWeights * cw,
KVCache * kv,
ggml_backend_sched_t sched,
const float * talker_hidden_last,
int c0,
float temperature,
int top_k,
float top_p,
int64_t seed,
int64_t subseq_base,
const char * dump_dir,
CodePredictorOutput * out);
static bool code_predictor_step(const TalkerWeights * tw,
const CodePredictorWeights * cw,
KVCache * kv,
ggml_backend_sched_t sched,
const float * talker_hidden_last,
int c0,
float temperature,
int top_k,
float top_p,
int64_t seed,
int64_t subseq_base,
const char * dump_dir,
CodePredictorOutput * out) {
// sub_input slots live at the talker hidden dimension: both the
// talker last hidden and the codec_embedding rows feeding the sub
// network are talker sized in the upstream checkpoint. The graph's
// mtp_proj brings them down to predictor hidden when present.
const int talker_hidden = tw->hidden_size;
const int n_acoustic = cw->num_acoustic_codebooks;
if (n_acoustic + 1 > kv->max_seq_len) {
fprintf(stderr, "[CodePredictor] FATAL: frame width %d exceeds cache max_seq_len %d\n", n_acoustic + 1,
kv->max_seq_len);
return false;
}
out->codes.assign((size_t) (n_acoustic + 1), 0);
out->codes[0] = c0;
// Prefill: two positions, talker_hidden_last and embed_talker(c0).
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, prefill_input.data(), 2, 0, talker_hidden, 0, &logits)) {
return false;
}
{
float u_g = 0.0f;
int cg = sample_top_k_p(logits.data(), (int) logits.size(), temperature, top_k, top_p, 1.0f, nullptr, 0, seed,
subseq_base + 1, &u_g);
if (subseq_base + 1 < 32) {
fprintf(stderr, "[Sample-CP] g=0 c=%d u=%.10f subseq=%lld\n", cg, (double) u_g,
(long long) (subseq_base + 1));
}
if (cg < 0) {
fprintf(stderr, "[CodePredictor] FATAL: sample returned no candidate at g=0\n");
return false;
}
out->codes[1] = cg;
}
// 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);
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, step_input.data(), 1, kv->cur_len, talker_hidden, g, &logits)) {
return false;
}
float u_g = 0.0f;
int cg = sample_top_k_p(logits.data(), (int) logits.size(), temperature, top_k, top_p, 1.0f, nullptr, 0, seed,
subseq_base + 1 + g, &u_g);
if (subseq_base + 1 + g < 32) {
fprintf(stderr, "[Sample-CP] g=%d c=%d u=%.10f subseq=%lld\n", g, cg, (double) u_g,
(long long) (subseq_base + 1 + g));
}
if (cg < 0) {
fprintf(stderr, "[CodePredictor] FATAL: sample returned no candidate at g=%d\n", g);
return false;
}
out->codes[(size_t) (g + 1)] = cg;
}
if (dump_dir) {
DebugDumper d;
debug_init(&d, dump_dir);
std::vector<int32_t> codes32(out->codes.begin(), out->codes.end());
int n = (int) codes32.size();
debug_dump_i32_as_f32(&d, "codes-step0", codes32.data(), &n, 1);
}
return true;
}
src/prompt-builder.cpp
-620
View File
@@ -1,620 +0,0 @@
// prompt-builder.cpp: CPU-side construction of the talker prefix
// input embedding. Mirrors generate() in qwen_tts/core/models/modeling_qwen3_tts.py
// for the strict subset {non-streaming, no voice clone}.
//
// Two streams are aligned then summed :
// text stream: text_projection(text_embedding(ids)) 151936 -> 2048 -> 1024
// codec stream: codec_embedding(ids) 3072 -> 1024
//
// Layout (lang_id != none, no speaker, no instruct) :
//
// role text(input_id[0:3]) 3 vecs
// prefill_lhs tts_pad x4 + tts_bos 5 vecs
// + codec_emb([think, think_bos, lang_id, think_eos, codec_pad])
// trailing_lhs text(input_id[3:-5]) + tts_eos N_text + 1 vecs
// + codec_emb([codec_pad x (N_text + 1)])
// trailing_rhs tts_pad + codec_emb([codec_bos]) 1 vec
//
// CustomVoice inserts the speaker codec embedding row between think_eos
// and codec_pad in the prefill, growing the prefill by one vector and
// substituting one tts_pad with another in the text stream alignment.
//
// VoiceDesign / CustomVoice may also prepend an instruct segment built
// from text_projection(text_embedding(<|im_start|>user\n{instruct}<|im_end|>\n))
// laid out as N_instruct standalone vectors before the role.
//
// All math is f32. text_embedding and codec_embedding are read from
// the mmapped GGUF in their stored dtype (bf16 by default) and cast
// row by row. The 2-layer ResizeMLP runs as two GEMMs with a SiLU in
// between, with bias on both linear layers.
#include "prompt-builder.h"
#include "ggml.h"
#include "qt-error.h"
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <string>
#include <vector>
// Convert one row of an embedding matrix W [vocab, dim] to f32. Uses the
// ggml type traits to_float dispatch so every dtype shipped by the
// quantizer is supported (F32, BF16, F16, Q8_0, Q4_K_M, etc). The row
// stride is the type block size, computed via ggml_row_size.
static void embed_row_to_f32(const GGUFModel & gf, const char * tensor_name, int row_id, int dim, float * dst) {
struct ggml_tensor * src = ggml_get_tensor(gf.meta, tensor_name);
if (!src) {
qt_throw("[Prompt] tensor '%s' not in meta context", tensor_name);
}
if (src->ne[0] != dim) {
qt_throw("[Prompt] tensor '%s' dim mismatch %lld vs %d", tensor_name, (long long) src->ne[0], dim);
}
if (row_id < 0 || row_id >= (int) src->ne[1]) {
qt_throw("[Prompt] row %d out of range for '%s' (vocab=%lld)", row_id, tensor_name, (long long) src->ne[1]);
}
const uint8_t * base = (const uint8_t *) gf_get_data(gf, tensor_name);
if (!base) {
qt_throw("[Prompt] tensor '%s' has no data", tensor_name);
}
const size_t row_bytes = ggml_row_size(src->type, dim);
const void * row = base + (size_t) row_id * row_bytes;
if (src->type == GGML_TYPE_F32) {
std::memcpy(dst, row, (size_t) dim * sizeof(float));
return;
}
const struct ggml_type_traits * tt = ggml_get_type_traits(src->type);
if (!tt || !tt->to_float) {
qt_throw("[Prompt] unsupported dtype %d for '%s'", (int) src->type, tensor_name);
}
tt->to_float(row, dst, dim);
}
// Read a full small tensor (bias, projection weight) into an f32 buffer.
// Allocates dst.resize internally. Routed through ggml_get_type_traits so
// quants are accepted, same as embed_row_to_f32 above.
static void read_tensor_f32(const GGUFModel & gf, const char * tensor_name, std::vector<float> & dst) {
struct ggml_tensor * src = ggml_get_tensor(gf.meta, tensor_name);
if (!src) {
qt_throw("[Prompt] tensor '%s' not in meta context", tensor_name);
}
int64_t n = ggml_nelements(src);
const uint8_t * base = (const uint8_t *) gf_get_data(gf, tensor_name);
dst.resize((size_t) n);
if (src->type == GGML_TYPE_F32) {
std::memcpy(dst.data(), base, (size_t) n * sizeof(float));
return;
}
const struct ggml_type_traits * tt = ggml_get_type_traits(src->type);
if (!tt || !tt->to_float) {
qt_throw("[Prompt] unsupported dtype %d for '%s'", (int) src->type, tensor_name);
}
tt->to_float(base, dst.data(), (int64_t) n);
}
// y = W @ x + b
// x [in_dim] f32, W [out_dim, in_dim] row-major f32, b [out_dim] f32
// y [out_dim] f32
// Naive dot-product GEMV, fine for small (≤2048) inputs at build time.
static void linear_f32(const float * x, const float * W, const float * b, int in_dim, int out_dim, float * y) {
for (int o = 0; o < out_dim; o++) {
const float * row = W + (size_t) o * (size_t) in_dim;
float acc = b ? b[o] : 0.0f;
for (int i = 0; i < in_dim; i++) {
acc += row[i] * x[i];
}
y[o] = acc;
}
}
static inline float silu(float v) {
return v / (1.0f + std::exp(-v));
}
// Apply text_projection: F1 (text_hidden -> text_hidden) -> SiLU -> F2
// (text_hidden -> hidden), both with bias.
struct TextProjection {
int in_dim; // text_hidden_size
int hid_dim; // intermediate (= text_hidden_size in 0.6B)
int out_dim; // hidden_size
std::vector<float> fc1_w; // [hid_dim, in_dim]
std::vector<float> fc1_b; // [hid_dim]
std::vector<float> fc2_w; // [out_dim, hid_dim]
std::vector<float> fc2_b; // [out_dim]
};
static void text_projection_load(TextProjection * tp, const GGUFModel & gf, int text_hidden_size, int hidden_size) {
tp->in_dim = text_hidden_size;
tp->hid_dim = text_hidden_size;
tp->out_dim = hidden_size;
read_tensor_f32(gf, "talker.text_proj.fc1.weight", tp->fc1_w);
read_tensor_f32(gf, "talker.text_proj.fc1.bias", tp->fc1_b);
read_tensor_f32(gf, "talker.text_proj.fc2.weight", tp->fc2_w);
read_tensor_f32(gf, "talker.text_proj.fc2.bias", tp->fc2_b);
}
static void text_projection_apply(const TextProjection * tp, const float * x, float * y) {
std::vector<float> h((size_t) tp->hid_dim);
linear_f32(x, tp->fc1_w.data(), tp->fc1_b.data(), tp->in_dim, tp->hid_dim, h.data());
for (int i = 0; i < tp->hid_dim; i++) {
h[(size_t) i] = silu(h[(size_t) i]);
}
linear_f32(h.data(), tp->fc2_w.data(), tp->fc2_b.data(), tp->hid_dim, tp->out_dim, y);
}
// Compute text_proj(text_embedding(ids[start:end])) row by row, append
// to dst (which already holds previous rows). Each output row is one
// hidden-dim vector.
static void embed_text_range(const GGUFModel & gf,
const TextProjection * tp,
const int32_t * ids,
int start,
int end,
int text_hidden_size,
int hidden_size,
std::vector<float> & dst) {
std::vector<float> e((size_t) text_hidden_size);
std::vector<float> y((size_t) hidden_size);
for (int i = start; i < end; i++) {
embed_row_to_f32(gf, "talker.text_embd.weight", ids[i], text_hidden_size, e.data());
text_projection_apply(tp, e.data(), y.data());
dst.insert(dst.end(), y.begin(), y.end());
}
}
// Append codec_embedding(id) to dst (one hidden-dim vector).
static void embed_codec(const GGUFModel & gf, int id, int hidden_size, std::vector<float> & dst) {
size_t old = dst.size();
dst.resize(old + (size_t) hidden_size);
embed_row_to_f32(gf, "talker.codec_embd.weight", id, hidden_size, dst.data() + old);
}
// Vector add: a += b, length n.
static void vec_add(float * a, const float * b, int n) {
for (int i = 0; i < n; i++) {
a[i] += b[i];
}
}
bool prompt_builder_build(const PipelineTTS * pt,
const BPETokenizer * tok,
const std::string & utterance_text,
const std::string & language,
const std::string & instruct_text,
const std::string & speaker_name,
const float * ref_spk_emb,
const std::string & ref_text,
const int32_t * ref_codes,
int ref_codes_T,
PromptBuilderOutput * out) {
const int hidden = pt->talker.hidden_size;
const int text_hid = pt->talker.text_hidden_size;
if (!speaker_name.empty() && ref_spk_emb != NULL) {
fprintf(stderr, "[Prompt] FATAL: speaker_name and ref_spk_emb are mutually exclusive\n");
return false;
}
// Voice clone mode B: ref_text and ref_codes drive an ICL prefix.
// Mode B requires ref_spk_emb so the speaker slot is also filled.
const bool icl = !ref_text.empty() && ref_codes != NULL && ref_codes_T > 0;
if (icl && ref_spk_emb == NULL) {
fprintf(stderr, "[Prompt] FATAL: ICL mode requires ref_spk_emb (no --ref-wav?)\n");
return false;
}
// Build the chat-templated prompt fed to the BPE tokenizer.
// Same wrap as the upstream demos: assistant role + utterance +
// im_end + newline + assistant role.
std::string full_text;
full_text.reserve(utterance_text.size() + 64);
full_text = "<|im_start|>assistant\n";
full_text += utterance_text;
full_text += "<|im_end|>\n<|im_start|>assistant\n";
std::vector<int> ids = bpe_encode(tok, full_text, /*add_eos=*/false);
if ((int) ids.size() < 8) {
fprintf(stderr, "[Prompt] FATAL: tokenized prompt too short (%d tokens)\n", (int) ids.size());
return false;
}
out->prompt_ids.assign(ids.begin(), ids.end());
const int N = (int) ids.size();
const int N_text = N - 3 - 5;
if (N_text <= 0) {
fprintf(stderr, "[Prompt] FATAL: no utterance text in prompt (N=%d)\n", N);
return false;
}
// Resolve language: "auto" -> no language id, prefill is 3 codec
// tokens (nothink, think_bos, think_eos). Otherwise insert the
// configured language id between think_bos and think_eos.
int language_id = -1;
{
std::string lang_lc = language;
for (char & c : lang_lc) {
c = (char) std::tolower((unsigned char) c);
}
if (lang_lc != "auto") {
for (const LanguageEntry & e : pt->languages) {
if (e.name == lang_lc) {
language_id = e.id;
break;
}
}
if (language_id < 0) {
fprintf(stderr, "[Prompt] FATAL: unknown language '%s'\n", language.c_str());
return false;
}
}
}
// Resolve speaker: empty name -> no speaker. Otherwise lookup case
// insensitively in pt->speakers and override the language id with the
// dialect entry when the user supplied language is chinese or auto,
// mirroring modeling_qwen3_tts.py lines 2118 to 2122.
int speaker_id = -1;
if (!speaker_name.empty()) {
std::string spk_lc = speaker_name;
for (char & c : spk_lc) {
c = (char) std::tolower((unsigned char) c);
}
const SpeakerEntry * found = NULL;
for (const SpeakerEntry & e : pt->speakers) {
if (e.name == spk_lc) {
found = &e;
break;
}
}
if (!found) {
fprintf(stderr, "[Prompt] FATAL: unknown speaker '%s'\n", speaker_name.c_str());
return false;
}
speaker_id = found->id;
// Dialect override: applied only when the user supplied language
// is chinese or auto, the dialect string is non empty, and the
// dialect resolves to a known language id.
if (!found->dialect.empty()) {
std::string lang_lc = language;
for (char & c : lang_lc) {
c = (char) std::tolower((unsigned char) c);
}
if (lang_lc == "chinese" || lang_lc == "auto") {
int dialect_id = -1;
for (const LanguageEntry & e : pt->languages) {
if (e.name == found->dialect) {
dialect_id = e.id;
break;
}
}
if (dialect_id < 0) {
fprintf(stderr, "[Prompt] FATAL: dialect '%s' not in language table\n", found->dialect.c_str());
return false;
}
language_id = dialect_id;
}
}
}
// Load the small tensors needed for the builder onto the host side.
TextProjection tp;
text_projection_load(&tp, pt->gguf_talker, text_hid, hidden);
// Special embeds (tts_bos, tts_eos, tts_pad, codec_pad, codec_bos)
// computed once.
std::vector<float> tts_bos_emb((size_t) hidden);
std::vector<float> tts_eos_emb((size_t) hidden);
std::vector<float> tts_pad_emb((size_t) hidden);
{
std::vector<float> e((size_t) text_hid);
embed_row_to_f32(pt->gguf_talker, "talker.text_embd.weight", pt->text_specials.tts_bos_id, text_hid, e.data());
text_projection_apply(&tp, e.data(), tts_bos_emb.data());
embed_row_to_f32(pt->gguf_talker, "talker.text_embd.weight", pt->text_specials.tts_eos_id, text_hid, e.data());
text_projection_apply(&tp, e.data(), tts_eos_emb.data());
embed_row_to_f32(pt->gguf_talker, "talker.text_embd.weight", pt->text_specials.tts_pad_id, text_hid, e.data());
text_projection_apply(&tp, e.data(), tts_pad_emb.data());
}
std::vector<float> codec_pad_emb((size_t) hidden);
embed_row_to_f32(pt->gguf_talker, "talker.codec_embd.weight", pt->codec_specials.pad_id, hidden,
codec_pad_emb.data());
// Codec prefill list: 3 ids if auto (no language), 4 otherwise.
// Speaker insertion: if a speaker id is set, the codec embedding row
// for that speaker slips between think_eos and codec_pad in the codec
// stream, mirroring modeling_qwen3_tts.py lines 2167 to 2172.
std::vector<int> codec_prefill;
if (language_id < 0) {
codec_prefill = { pt->codec_specials.nothink_id, pt->codec_specials.think_bos_id,
pt->codec_specials.think_eos_id };
} else {
codec_prefill = { pt->codec_specials.think_id, pt->codec_specials.think_bos_id, language_id,
pt->codec_specials.think_eos_id };
}
if (speaker_id >= 0) {
codec_prefill.push_back(speaker_id);
} else if (ref_spk_emb != NULL) {
// Sentinel: the codec_left builder below copies ref_spk_emb in
// place of an embedding lookup whenever it sees -2.
codec_prefill.push_back(-2);
}
const int n_prefill = (int) codec_prefill.size();
const int T_codec_prefix = n_prefill + 2; // + codec_pad + codec_bos
const int n_pad_pre = T_codec_prefix - 2;
// Tokenize the instruct segment when non empty. The wrapper mirrors
// _build_instruct_text upstream: <|im_start|>user\n{instruct}<|im_end|>\n
// The result is a flat list of text token ids that will be projected
// and placed as standalone vectors at the head of the input embed,
// with no codec stream contribution.
std::vector<int> instruct_ids;
if (!instruct_text.empty()) {
std::string wrapped;
wrapped.reserve(instruct_text.size() + 32);
wrapped = "<|im_start|>user\n";
wrapped += instruct_text;
wrapped += "<|im_end|>\n";
instruct_ids = bpe_encode(tok, wrapped, /*add_eos=*/false);
}
const int N_instruct = (int) instruct_ids.size();
// Tokenize the reference utterance when ICL is active. The wrap is
// identical to the main utterance: assistant role + ref_text +
// im_end + newline + assistant role. We slice [3:-5] later to keep
// only the inner text body, mirroring input_id[:, 3:-5] upstream.
std::vector<int> ref_ids;
int N_ref_text = 0;
if (icl) {
std::string ref_full;
ref_full.reserve(ref_text.size() + 64);
ref_full = "<|im_start|>assistant\n";
ref_full += ref_text;
ref_full += "<|im_end|>\n<|im_start|>assistant\n";
ref_ids = bpe_encode(tok, ref_full, /*add_eos=*/false);
if ((int) ref_ids.size() < 8) {
fprintf(stderr, "[Prompt] FATAL: ref_text tokenized too short (%d tokens)\n", (int) ref_ids.size());
return false;
}
// ref_ids[3: -5] is the inner ref text body without role tokens
N_ref_text = (int) ref_ids.size() - 3 - 5;
if (N_ref_text <= 0) {
fprintf(stderr, "[Prompt] FATAL: empty ref_text body\n");
return false;
}
}
// ICL geometry. text_lens = N_ref_text + N_text + 1 (tts_eos).
// codec_lens = 1 (codec_bos) + ref_codes_T. The non_streaming_mode
// branch upstream pads the shorter stream so they end up the same
// length, except when text > codec where trailing_text_hidden carries
// the leftover text rows.
const int text_lens_icl = icl ? (N_ref_text + N_text + 1) : 0;
const int codec_lens_icl = icl ? (1 + ref_codes_T) : 0;
const int icl_T = icl ? (text_lens_icl > codec_lens_icl ? codec_lens_icl : codec_lens_icl) : 0;
// Allocate the full output buffer.
// Standard layout : N_instruct + 3 (role) + (n_pad_pre + 1) + N_text + 1 (eos) + 1 (final)
// ICL layout : N_instruct + 3 (role) + (n_pad_pre + 1) + icl_T
const int T_ctx =
icl ? (N_instruct + 3 + (n_pad_pre + 1) + icl_T) : (N_instruct + 3 + (n_pad_pre + 1) + N_text + 1 + 1);
out->T_ctx = T_ctx;
out->hidden = hidden;
out->input_embed.assign((size_t) T_ctx * (size_t) hidden, 0.0f);
out->N_text = N_text;
int row = 0;
auto row_ptr = [&](int r) {
return out->input_embed.data() + (size_t) r * (size_t) hidden;
};
// Instruct prefix: text_proj(text_embed(instruct_ids)). Standalone
// vectors with no codec stream (zero pad_id sum, ie nothing added).
if (N_instruct > 0) {
std::vector<float> dst;
embed_text_range(pt->gguf_talker, &tp, instruct_ids.data(), 0, N_instruct, text_hid, hidden, dst);
std::memcpy(row_ptr(row), dst.data(), dst.size() * sizeof(float));
row += N_instruct;
}
// Role: text_proj(text_embed(ids[0:3]))
{
std::vector<float> dst;
dst.reserve((size_t) 3 * (size_t) hidden);
embed_text_range(pt->gguf_talker, &tp, ids.data(), 0, 3, text_hid, hidden, dst);
std::memcpy(row_ptr(row), dst.data(), dst.size() * sizeof(float));
row += 3;
}
// Codec prefix: tts_pad x n_pad_pre + tts_bos, summed with
// codec_emb([codec_prefill_list[:-1]] + codec_pad). The Python code
// takes codec_input_embedding[:, :-1] which drops the codec_bos,
// leaving [codec_prefill_list..., codec_pad].
{
std::vector<int> codec_left = codec_prefill;
codec_left.push_back(pt->codec_specials.pad_id);
for (int i = 0; i < (int) codec_left.size(); i++) {
float * r = row_ptr(row + i);
// text stream: tts_pad * (n - 1) then tts_bos at the end
const float * text_vec = (i == (int) codec_left.size() - 1) ? tts_bos_emb.data() : tts_pad_emb.data();
std::memcpy(r, text_vec, (size_t) hidden * sizeof(float));
// codec stream: either an embedding lookup or, when the
// sentinel -2 marks the speaker slot, a direct copy of the
// user supplied ref_spk_emb (voice clone mode A).
std::vector<float> ce((size_t) hidden);
if (codec_left[(size_t) i] == -2) {
std::memcpy(ce.data(), ref_spk_emb, (size_t) hidden * sizeof(float));
} else {
embed_row_to_f32(pt->gguf_talker, "talker.codec_embd.weight", codec_left[(size_t) i], hidden,
ce.data());
}
vec_add(r, ce.data(), hidden);
}
row += (int) codec_left.size();
}
// From here, two paths: standard (no ICL) builds the trailing
// utterance text + tts_eos + final_pad, ICL builds an aligned
// text/codec block that replaces those rows entirely.
if (!icl) {
// Standard layout: trailing utterance text + tts_eos rows summed
// with codec_pad, then a final tts_pad + codec_bos row.
for (int i = 0; i < N_text; i++) {
std::vector<float> e((size_t) text_hid);
std::vector<float> y((size_t) hidden);
embed_row_to_f32(pt->gguf_talker, "talker.text_embd.weight", ids[3 + i], text_hid, e.data());
text_projection_apply(&tp, e.data(), y.data());
float * r = row_ptr(row);
std::memcpy(r, y.data(), (size_t) hidden * sizeof(float));
vec_add(r, codec_pad_emb.data(), hidden);
row++;
}
{
float * r = row_ptr(row);
std::memcpy(r, tts_eos_emb.data(), (size_t) hidden * sizeof(float));
vec_add(r, codec_pad_emb.data(), hidden);
row++;
}
{
float * r = row_ptr(row);
std::memcpy(r, tts_pad_emb.data(), (size_t) hidden * sizeof(float));
std::vector<float> ce((size_t) hidden);
embed_row_to_f32(pt->gguf_talker, "talker.codec_embd.weight", pt->codec_specials.bos_id, hidden, ce.data());
vec_add(r, ce.data(), hidden);
row++;
}
} else {
// ICL layout: compute the text stream and the codec stream
// separately then add them. The text stream is text_proj of
// [ref_text_ids; utterance_text_ids] followed by tts_eos. The
// codec stream is codec_emb(codec_bos) followed by sum over the
// 16 codebook embeddings of ref_codes[i, t] for each frame t.
// Both streams are aligned to length icl_T per the upstream
// non_streaming_mode=False branch of generate_icl_prompt.
const int T_icl = codec_lens_icl; // text_lens > codec: truncate to codec, else pad text up to codec
// Build the codec stream [T_icl, hidden]. Row 0: codec_emb(codec_bos).
// Row 1..ref_codes_T: sum over k=0..15 of codebook_k_emb(ref_codes[k, t]).
std::vector<float> codec_stream((size_t) T_icl * (size_t) hidden, 0.0f);
{
// Row 0: codec_bos lookup.
embed_row_to_f32(pt->gguf_talker, "talker.codec_embd.weight", pt->codec_specials.bos_id, hidden,
codec_stream.data());
// Row 1..ref_codes_T: sum over codebooks.
std::vector<float> tmp((size_t) hidden);
for (int t = 0; t < ref_codes_T; t++) {
float * dst = codec_stream.data() + (size_t) (1 + t) * (size_t) hidden;
// codebook 0 lives in talker.codec_embd
int code0 = ref_codes[(size_t) 0 * (size_t) ref_codes_T + (size_t) t];
embed_row_to_f32(pt->gguf_talker, "talker.codec_embd.weight", code0, hidden, dst);
// codebooks 1..15 live in code_pred.codec_embd.{i-1}
for (int i = 1; i < pt->num_code_groups; i++) {
int code = ref_codes[(size_t) i * (size_t) ref_codes_T + (size_t) t];
char tname[64];
std::snprintf(tname, sizeof(tname), "code_pred.codec_embd.%d.weight", i - 1);
embed_row_to_f32(pt->gguf_talker, tname, code, hidden, tmp.data());
vec_add(dst, tmp.data(), hidden);
}
}
}
// Build the text stream [text_lens_icl, hidden] = text_proj of
// [ref_text; utterance_text] then tts_eos.
std::vector<float> text_stream((size_t) text_lens_icl * (size_t) hidden, 0.0f);
for (int i = 0; i < N_ref_text; i++) {
std::vector<float> e((size_t) text_hid);
float * r = text_stream.data() + (size_t) i * (size_t) hidden;
embed_row_to_f32(pt->gguf_talker, "talker.text_embd.weight", ref_ids[3 + i], text_hid, e.data());
text_projection_apply(&tp, e.data(), r);
}
for (int i = 0; i < N_text; i++) {
std::vector<float> e((size_t) text_hid);
float * r = text_stream.data() + (size_t) (N_ref_text + i) * (size_t) hidden;
embed_row_to_f32(pt->gguf_talker, "talker.text_embd.weight", ids[3 + i], text_hid, e.data());
text_projection_apply(&tp, e.data(), r);
}
// Append tts_eos at the end of the text stream.
std::memcpy(text_stream.data() + (size_t) (text_lens_icl - 1) * (size_t) hidden, tts_eos_emb.data(),
(size_t) hidden * sizeof(float));
// Align the two streams to T_icl. text_lens > codec: truncate
// text and stash the leftover into trailing_text_hidden. text_lens
// <= codec: pad text with tts_pad up to codec, trailing reduces
// to tts_pad.
std::vector<float> aligned_text((size_t) T_icl * (size_t) hidden, 0.0f);
if (text_lens_icl >= T_icl) {
// truncate text to T_icl rows, leftover goes into trailing
std::memcpy(aligned_text.data(), text_stream.data(), (size_t) T_icl * (size_t) hidden * sizeof(float));
const int trailing_n = text_lens_icl - T_icl;
out->T_trailing = trailing_n > 0 ? trailing_n : 1;
out->trailing_text_hidden.assign((size_t) out->T_trailing * (size_t) hidden, 0.0f);
if (trailing_n > 0) {
std::memcpy(out->trailing_text_hidden.data(), text_stream.data() + (size_t) T_icl * (size_t) hidden,
(size_t) trailing_n * (size_t) hidden * sizeof(float));
} else {
std::memcpy(out->trailing_text_hidden.data(), tts_pad_emb.data(), (size_t) hidden * sizeof(float));
}
} else {
// pad text with tts_pad up to T_icl, trailing = single tts_pad row
std::memcpy(aligned_text.data(), text_stream.data(),
(size_t) text_lens_icl * (size_t) hidden * sizeof(float));
for (int i = text_lens_icl; i < T_icl; i++) {
std::memcpy(aligned_text.data() + (size_t) i * (size_t) hidden, tts_pad_emb.data(),
(size_t) hidden * sizeof(float));
}
out->T_trailing = 1;
out->trailing_text_hidden.assign((size_t) hidden, 0.0f);
std::memcpy(out->trailing_text_hidden.data(), tts_pad_emb.data(), (size_t) hidden * sizeof(float));
}
// Sum aligned_text + codec_stream into the input embed at the
// current row offset.
for (int i = 0; i < T_icl; i++) {
float * r = row_ptr(row + i);
std::memcpy(r, aligned_text.data() + (size_t) i * (size_t) hidden, (size_t) hidden * sizeof(float));
vec_add(r, codec_stream.data() + (size_t) i * (size_t) hidden, hidden);
}
row += T_icl;
}
if (row != T_ctx) {
fprintf(stderr, "[Prompt] FATAL: layout error row=%d expected T_ctx=%d\n", row, T_ctx);
return false;
}
// Trailing text hidden: non streaming mode (no ICL) collapses the
// overlay to a single row equal to tts_pad_embed
// (modeling_qwen3_tts.py line 2227). The full utterance text is
// already integrated into the prefill above as codec_pad summed text
// rows + tts_eos, so the overlay loop only ever needs tts_pad: step
// 0 reads trailing_text_hidden[0] which is tts_pad, every later step
// falls through to the else branch and reads tts_pad_embed. One row,
// bit exact with the Python hook dump.
//
// ICL mode populates out->trailing_text_hidden directly inside the
// ICL branch above, so we only set the default here for non ICL.
if (!icl) {
out->T_trailing = 1;
out->trailing_text_hidden.assign((size_t) hidden, 0.0f);
std::memcpy(out->trailing_text_hidden.data(), tts_pad_emb.data(), (size_t) hidden * sizeof(float));
}
out->tts_pad_embed = tts_pad_emb;
fprintf(stderr,
"[Prompt] Built: %d ids, N_text=%d, N_instruct=%d, T_ctx=%d, hidden=%d, lang=%s (id=%d), speaker=%s "
"(id=%d) ref_spk_emb=%s icl=%s\n",
N, N_text, N_instruct, T_ctx, hidden, language.c_str(), language_id,
speaker_name.empty() ? "none" : speaker_name.c_str(), speaker_id, ref_spk_emb ? "yes" : "no",
icl ? "yes" : "no");
return true;
}
src/prompt-builder.h
+611 -25
View File
@@ -8,7 +8,7 @@
// vectors. The trailing text hidden buffer is also produced for the
// streaming-text overlay used during generation.
//
// Modes :
// Modes:
// base text only, no instruct, no speaker
// voice_design text + instruct (style description), no speaker
// custom_voice text + speaker, optional instruct
@@ -16,11 +16,43 @@
// Empty / NULL strings disable the corresponding stream. The builder runs
// CPU-side using the BF16 weight blocks mmapped from the talker GGUF, no
// backend allocation, no graph compute.
//
// Two streams are aligned then summed:
// text stream: text_projection(text_embedding(ids)) 151936 -> 2048 -> 1024
// codec stream: codec_embedding(ids) 3072 -> 1024
//
// Layout (lang_id != none, no speaker, no instruct):
//
// role text(input_id[0:3]) 3 vecs
// prefill_lhs tts_pad x4 + tts_bos 5 vecs
// + codec_emb([think, think_bos, lang_id, think_eos, codec_pad])
// trailing_lhs text(input_id[3:-5]) + tts_eos N_text + 1 vecs
// + codec_emb([codec_pad x (N_text + 1)])
// trailing_rhs tts_pad + codec_emb([codec_bos]) 1 vec
//
// CustomVoice inserts the speaker codec embedding row between think_eos
// and codec_pad in the prefill, growing the prefill by one vector and
// substituting one tts_pad with another in the text stream alignment.
//
// VoiceDesign / CustomVoice may also prepend an instruct segment built
// from text_projection(text_embedding(<|im_start|>user\n{instruct}<|im_end|>\n))
// laid out as N_instruct standalone vectors before the role.
//
// All math is f32. text_embedding and codec_embedding are read from
// the mmapped GGUF in their stored dtype (bf16 by default) and cast
// row by row. The 2-layer ResizeMLP runs as two GEMMs with a SiLU in
// between, with bias on both linear layers.
#include "bpe.h"
#include "ggml.h"
#include "pipeline-tts.h"
#include "qt-error.h"
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <string>
#include <vector>
@@ -48,27 +80,581 @@ struct PromptBuilderOutput {
int N_text;
};
// Assemble the prefix. instruct_text is the raw user style instruction
// (empty for none). speaker_name is the lowercased speaker key looked up
// in pt->speakers (empty for none). ref_spk_emb is an optional pointer to
// an [hidden] f32 vector that takes the place of the speaker preset row
// for voice clone mode A: when non NULL it is inserted between think_eos
// and codec_pad in the codec stream, mutually exclusive with speaker_name.
// ref_text and ref_codes activate voice clone mode B (ICL): the prompt
// becomes [icl_text + tts_eos] aligned with [codec_bos + ref_codes_summed],
// where ref_codes is a flat [num_code_groups, T_codec] int32 buffer
// produced by pipeline_codec_encode on the resampled reference audio.
// Returns false if BPE encoding produces fewer than the expected
// role/footer tokens, the language is unknown, speaker_name is set but
// not found, or speaker_name and ref_spk_emb are both set.
bool prompt_builder_build(const PipelineTTS * pt,
const BPETokenizer * tok,
const std::string & utterance_text,
const std::string & language,
const std::string & instruct_text,
const std::string & speaker_name,
const float * ref_spk_emb,
const std::string & ref_text,
const int32_t * ref_codes,
int ref_codes_T,
PromptBuilderOutput * out);
// Convert one row of an embedding matrix W [vocab, dim] to f32. Uses the
// ggml type traits to_float dispatch so every dtype shipped by the
// quantizer is supported (F32, BF16, F16, Q8_0, Q4_K_M, etc). The row
// stride is the type block size, computed via ggml_row_size.
static void embed_row_to_f32(const GGUFModel & gf, const char * tensor_name, int row_id, int dim, float * dst) {
struct ggml_tensor * src = ggml_get_tensor(gf.meta, tensor_name);
if (!src) {
qt_throw("[Prompt] tensor '%s' not in meta context", tensor_name);
}
if (src->ne[0] != dim) {
qt_throw("[Prompt] tensor '%s' dim mismatch %lld vs %d", tensor_name, (long long) src->ne[0], dim);
}
if (row_id < 0 || row_id >= (int) src->ne[1]) {
qt_throw("[Prompt] row %d out of range for '%s' (vocab=%lld)", row_id, tensor_name, (long long) src->ne[1]);
}
const uint8_t * base = (const uint8_t *) gf_get_data(gf, tensor_name);
if (!base) {
qt_throw("[Prompt] tensor '%s' has no data", tensor_name);
}
const size_t row_bytes = ggml_row_size(src->type, dim);
const void * row = base + (size_t) row_id * row_bytes;
if (src->type == GGML_TYPE_F32) {
std::memcpy(dst, row, (size_t) dim * sizeof(float));
return;
}
const struct ggml_type_traits * tt = ggml_get_type_traits(src->type);
if (!tt || !tt->to_float) {
qt_throw("[Prompt] unsupported dtype %d for '%s'", (int) src->type, tensor_name);
}
tt->to_float(row, dst, dim);
}
// Read a full small tensor (bias, projection weight) into an f32 buffer.
// Allocates dst.resize internally. Routed through ggml_get_type_traits so
// quants are accepted, same as embed_row_to_f32 above.
static void read_tensor_f32(const GGUFModel & gf, const char * tensor_name, std::vector<float> & dst) {
struct ggml_tensor * src = ggml_get_tensor(gf.meta, tensor_name);
if (!src) {
qt_throw("[Prompt] tensor '%s' not in meta context", tensor_name);
}
int64_t n = ggml_nelements(src);
const uint8_t * base = (const uint8_t *) gf_get_data(gf, tensor_name);
dst.resize((size_t) n);
if (src->type == GGML_TYPE_F32) {
std::memcpy(dst.data(), base, (size_t) n * sizeof(float));
return;
}
const struct ggml_type_traits * tt = ggml_get_type_traits(src->type);
if (!tt || !tt->to_float) {
qt_throw("[Prompt] unsupported dtype %d for '%s'", (int) src->type, tensor_name);
}
tt->to_float(base, dst.data(), (int64_t) n);
}
// y = W @ x + b
// x [in_dim] f32, W [out_dim, in_dim] row-major f32, b [out_dim] f32
// y [out_dim] f32
// Naive dot-product GEMV, fine for small (≤2048) inputs at build time.
static void linear_f32(const float * x, const float * W, const float * b, int in_dim, int out_dim, float * y) {
for (int o = 0; o < out_dim; o++) {
const float * row = W + (size_t) o * (size_t) in_dim;
float acc = b ? b[o] : 0.0f;
for (int i = 0; i < in_dim; i++) {
acc += row[i] * x[i];
}
y[o] = acc;
}
}
static inline float silu(float v) {
return v / (1.0f + std::exp(-v));
}
// Apply text_projection: F1 (text_hidden -> text_hidden) -> SiLU -> F2
// (text_hidden -> hidden), both with bias.
struct TextProjection {
int in_dim; // text_hidden_size
int hid_dim; // intermediate (= text_hidden_size in 0.6B)
int out_dim; // hidden_size
std::vector<float> fc1_w; // [hid_dim, in_dim]
std::vector<float> fc1_b; // [hid_dim]
std::vector<float> fc2_w; // [out_dim, hid_dim]
std::vector<float> fc2_b; // [out_dim]
};
static void text_projection_load(TextProjection * tp, const GGUFModel & gf, int text_hidden_size, int hidden_size) {
tp->in_dim = text_hidden_size;
tp->hid_dim = text_hidden_size;
tp->out_dim = hidden_size;
read_tensor_f32(gf, "talker.text_proj.fc1.weight", tp->fc1_w);
read_tensor_f32(gf, "talker.text_proj.fc1.bias", tp->fc1_b);
read_tensor_f32(gf, "talker.text_proj.fc2.weight", tp->fc2_w);
read_tensor_f32(gf, "talker.text_proj.fc2.bias", tp->fc2_b);
}
static void text_projection_apply(const TextProjection * tp, const float * x, float * y) {
std::vector<float> h((size_t) tp->hid_dim);
linear_f32(x, tp->fc1_w.data(), tp->fc1_b.data(), tp->in_dim, tp->hid_dim, h.data());
for (int i = 0; i < tp->hid_dim; i++) {
h[(size_t) i] = silu(h[(size_t) i]);
}
linear_f32(h.data(), tp->fc2_w.data(), tp->fc2_b.data(), tp->hid_dim, tp->out_dim, y);
}
// Compute text_proj(text_embedding(ids[start:end])) row by row, append
// to dst (which already holds previous rows). Each output row is one
// hidden-dim vector.
static void embed_text_range(const GGUFModel & gf,
const TextProjection * tp,
const int32_t * ids,
int start,
int end,
int text_hidden_size,
int hidden_size,
std::vector<float> & dst) {
std::vector<float> e((size_t) text_hidden_size);
std::vector<float> y((size_t) hidden_size);
for (int i = start; i < end; i++) {
embed_row_to_f32(gf, "talker.text_embd.weight", ids[i], text_hidden_size, e.data());
text_projection_apply(tp, e.data(), y.data());
dst.insert(dst.end(), y.begin(), y.end());
}
}
// Append codec_embedding(id) to dst (one hidden-dim vector).
static void embed_codec(const GGUFModel & gf, int id, int hidden_size, std::vector<float> & dst) {
size_t old = dst.size();
dst.resize(old + (size_t) hidden_size);
embed_row_to_f32(gf, "talker.codec_embd.weight", id, hidden_size, dst.data() + old);
}
// Vector add: a += b, length n.
static void vec_add(float * a, const float * b, int n) {
for (int i = 0; i < n; i++) {
a[i] += b[i];
}
}
static bool prompt_builder_build(const PipelineTTS * pt,
const BPETokenizer * tok,
const std::string & utterance_text,
const std::string & language,
const std::string & instruct_text,
const std::string & speaker_name,
const float * ref_spk_emb,
const std::string & ref_text,
const int32_t * ref_codes,
int ref_codes_T,
PromptBuilderOutput * out) {
const int hidden = pt->talker.hidden_size;
const int text_hid = pt->talker.text_hidden_size;
if (!speaker_name.empty() && ref_spk_emb != NULL) {
fprintf(stderr, "[Prompt] FATAL: speaker_name and ref_spk_emb are mutually exclusive\n");
return false;
}
// Voice clone mode B: ref_text and ref_codes drive an ICL prefix.
// Mode B requires ref_spk_emb so the speaker slot is also filled.
const bool icl = !ref_text.empty() && ref_codes != NULL && ref_codes_T > 0;
if (icl && ref_spk_emb == NULL) {
fprintf(stderr, "[Prompt] FATAL: ICL mode requires ref_spk_emb (no --ref-wav?)\n");
return false;
}
// Build the chat-templated prompt fed to the BPE tokenizer.
// Same wrap as the upstream demos: assistant role + utterance +
// im_end + newline + assistant role.
std::string full_text;
full_text.reserve(utterance_text.size() + 64);
full_text = "<|im_start|>assistant\n";
full_text += utterance_text;
full_text += "<|im_end|>\n<|im_start|>assistant\n";
std::vector<int> ids = bpe_encode(tok, full_text, /*add_eos=*/false);
if ((int) ids.size() < 8) {
fprintf(stderr, "[Prompt] FATAL: tokenized prompt too short (%d tokens)\n", (int) ids.size());
return false;
}
out->prompt_ids.assign(ids.begin(), ids.end());
const int N = (int) ids.size();
const int N_text = N - 3 - 5;
if (N_text <= 0) {
fprintf(stderr, "[Prompt] FATAL: no utterance text in prompt (N=%d)\n", N);
return false;
}
// Resolve language: "auto" -> no language id, prefill is 3 codec
// tokens (nothink, think_bos, think_eos). Otherwise insert the
// configured language id between think_bos and think_eos.
int language_id = -1;
{
std::string lang_lc = language;
for (char & c : lang_lc) {
c = (char) std::tolower((unsigned char) c);
}
if (lang_lc != "auto") {
for (const LanguageEntry & e : pt->languages) {
if (e.name == lang_lc) {
language_id = e.id;
break;
}
}
if (language_id < 0) {
fprintf(stderr, "[Prompt] FATAL: unknown language '%s'\n", language.c_str());
return false;
}
}
}
// Resolve speaker: empty name -> no speaker. Otherwise lookup case
// insensitively in pt->speakers and override the language id with the
// dialect entry when the user supplied language is chinese or auto,
// mirroring modeling_qwen3_tts.py lines 2118 to 2122.
int speaker_id = -1;
if (!speaker_name.empty()) {
std::string spk_lc = speaker_name;
for (char & c : spk_lc) {
c = (char) std::tolower((unsigned char) c);
}
const SpeakerEntry * found = NULL;
for (const SpeakerEntry & e : pt->speakers) {
if (e.name == spk_lc) {
found = &e;
break;
}
}
if (!found) {
fprintf(stderr, "[Prompt] FATAL: unknown speaker '%s'\n", speaker_name.c_str());
return false;
}
speaker_id = found->id;
// Dialect override: applied only when the user supplied language
// is chinese or auto, the dialect string is non empty, and the
// dialect resolves to a known language id.
if (!found->dialect.empty()) {
std::string lang_lc = language;
for (char & c : lang_lc) {
c = (char) std::tolower((unsigned char) c);
}
if (lang_lc == "chinese" || lang_lc == "auto") {
int dialect_id = -1;
for (const LanguageEntry & e : pt->languages) {
if (e.name == found->dialect) {
dialect_id = e.id;
break;
}
}
if (dialect_id < 0) {
fprintf(stderr, "[Prompt] FATAL: dialect '%s' not in language table\n", found->dialect.c_str());
return false;
}
language_id = dialect_id;
}
}
}
// Load the small tensors needed for the builder onto the host side.
TextProjection tp;
text_projection_load(&tp, pt->gguf_talker, text_hid, hidden);
// Special embeds (tts_bos, tts_eos, tts_pad, codec_pad, codec_bos)
// computed once.
std::vector<float> tts_bos_emb((size_t) hidden);
std::vector<float> tts_eos_emb((size_t) hidden);
std::vector<float> tts_pad_emb((size_t) hidden);
{
std::vector<float> e((size_t) text_hid);
embed_row_to_f32(pt->gguf_talker, "talker.text_embd.weight", pt->text_specials.tts_bos_id, text_hid, e.data());
text_projection_apply(&tp, e.data(), tts_bos_emb.data());
embed_row_to_f32(pt->gguf_talker, "talker.text_embd.weight", pt->text_specials.tts_eos_id, text_hid, e.data());
text_projection_apply(&tp, e.data(), tts_eos_emb.data());
embed_row_to_f32(pt->gguf_talker, "talker.text_embd.weight", pt->text_specials.tts_pad_id, text_hid, e.data());
text_projection_apply(&tp, e.data(), tts_pad_emb.data());
}
std::vector<float> codec_pad_emb((size_t) hidden);
embed_row_to_f32(pt->gguf_talker, "talker.codec_embd.weight", pt->codec_specials.pad_id, hidden,
codec_pad_emb.data());
// Codec prefill list: 3 ids if auto (no language), 4 otherwise.
// Speaker insertion: if a speaker id is set, the codec embedding row
// for that speaker slips between think_eos and codec_pad in the codec
// stream, mirroring modeling_qwen3_tts.py lines 2167 to 2172.
std::vector<int> codec_prefill;
if (language_id < 0) {
codec_prefill = { pt->codec_specials.nothink_id, pt->codec_specials.think_bos_id,
pt->codec_specials.think_eos_id };
} else {
codec_prefill = { pt->codec_specials.think_id, pt->codec_specials.think_bos_id, language_id,
pt->codec_specials.think_eos_id };
}
if (speaker_id >= 0) {
codec_prefill.push_back(speaker_id);
} else if (ref_spk_emb != NULL) {
// Sentinel: the codec_left builder below copies ref_spk_emb in
// place of an embedding lookup whenever it sees -2.
codec_prefill.push_back(-2);
}
const int n_prefill = (int) codec_prefill.size();
const int T_codec_prefix = n_prefill + 2; // + codec_pad + codec_bos
const int n_pad_pre = T_codec_prefix - 2;
// Tokenize the instruct segment when non empty. The wrapper mirrors
// _build_instruct_text upstream: <|im_start|>user\n{instruct}<|im_end|>\n
// The result is a flat list of text token ids that will be projected
// and placed as standalone vectors at the head of the input embed,
// with no codec stream contribution.
std::vector<int> instruct_ids;
if (!instruct_text.empty()) {
std::string wrapped;
wrapped.reserve(instruct_text.size() + 32);
wrapped = "<|im_start|>user\n";
wrapped += instruct_text;
wrapped += "<|im_end|>\n";
instruct_ids = bpe_encode(tok, wrapped, /*add_eos=*/false);
}
const int N_instruct = (int) instruct_ids.size();
// Tokenize the reference utterance when ICL is active. The wrap is
// identical to the main utterance: assistant role + ref_text +
// im_end + newline + assistant role. We slice [3:-5] later to keep
// only the inner text body, mirroring input_id[:, 3:-5] upstream.
std::vector<int> ref_ids;
int N_ref_text = 0;
if (icl) {
std::string ref_full;
ref_full.reserve(ref_text.size() + 64);
ref_full = "<|im_start|>assistant\n";
ref_full += ref_text;
ref_full += "<|im_end|>\n<|im_start|>assistant\n";
ref_ids = bpe_encode(tok, ref_full, /*add_eos=*/false);
if ((int) ref_ids.size() < 8) {
fprintf(stderr, "[Prompt] FATAL: ref_text tokenized too short (%d tokens)\n", (int) ref_ids.size());
return false;
}
// ref_ids[3: -5] is the inner ref text body without role tokens
N_ref_text = (int) ref_ids.size() - 3 - 5;
if (N_ref_text <= 0) {
fprintf(stderr, "[Prompt] FATAL: empty ref_text body\n");
return false;
}
}
// ICL geometry. text_lens = N_ref_text + N_text + 1 (tts_eos).
// codec_lens = 1 (codec_bos) + ref_codes_T. The non_streaming_mode
// branch upstream pads the shorter stream so they end up the same
// length, except when text > codec where trailing_text_hidden carries
// the leftover text rows.
const int text_lens_icl = icl ? (N_ref_text + N_text + 1) : 0;
const int codec_lens_icl = icl ? (1 + ref_codes_T) : 0;
const int icl_T = icl ? (text_lens_icl > codec_lens_icl ? codec_lens_icl : codec_lens_icl) : 0;
// Allocate the full output buffer.
// Standard layout : N_instruct + 3 (role) + (n_pad_pre + 1) + N_text + 1 (eos) + 1 (final)
// ICL layout : N_instruct + 3 (role) + (n_pad_pre + 1) + icl_T
const int T_ctx =
icl ? (N_instruct + 3 + (n_pad_pre + 1) + icl_T) : (N_instruct + 3 + (n_pad_pre + 1) + N_text + 1 + 1);
out->T_ctx = T_ctx;
out->hidden = hidden;
out->input_embed.assign((size_t) T_ctx * (size_t) hidden, 0.0f);
out->N_text = N_text;
int row = 0;
auto row_ptr = [&](int r) {
return out->input_embed.data() + (size_t) r * (size_t) hidden;
};
// Instruct prefix: text_proj(text_embed(instruct_ids)). Standalone
// vectors with no codec stream (zero pad_id sum, ie nothing added).
if (N_instruct > 0) {
std::vector<float> dst;
embed_text_range(pt->gguf_talker, &tp, instruct_ids.data(), 0, N_instruct, text_hid, hidden, dst);
std::memcpy(row_ptr(row), dst.data(), dst.size() * sizeof(float));
row += N_instruct;
}
// Role: text_proj(text_embed(ids[0:3]))
{
std::vector<float> dst;
dst.reserve((size_t) 3 * (size_t) hidden);
embed_text_range(pt->gguf_talker, &tp, ids.data(), 0, 3, text_hid, hidden, dst);
std::memcpy(row_ptr(row), dst.data(), dst.size() * sizeof(float));
row += 3;
}
// Codec prefix: tts_pad x n_pad_pre + tts_bos, summed with
// codec_emb([codec_prefill_list[:-1]] + codec_pad). The Python code
// takes codec_input_embedding[:, :-1] which drops the codec_bos,
// leaving [codec_prefill_list..., codec_pad].
{
std::vector<int> codec_left = codec_prefill;
codec_left.push_back(pt->codec_specials.pad_id);
for (int i = 0; i < (int) codec_left.size(); i++) {
float * r = row_ptr(row + i);
// text stream: tts_pad * (n - 1) then tts_bos at the end
const float * text_vec = (i == (int) codec_left.size() - 1) ? tts_bos_emb.data() : tts_pad_emb.data();
std::memcpy(r, text_vec, (size_t) hidden * sizeof(float));
// codec stream: either an embedding lookup or, when the
// sentinel -2 marks the speaker slot, a direct copy of the
// user supplied ref_spk_emb (voice clone mode A).
std::vector<float> ce((size_t) hidden);
if (codec_left[(size_t) i] == -2) {
std::memcpy(ce.data(), ref_spk_emb, (size_t) hidden * sizeof(float));
} else {
embed_row_to_f32(pt->gguf_talker, "talker.codec_embd.weight", codec_left[(size_t) i], hidden,
ce.data());
}
vec_add(r, ce.data(), hidden);
}
row += (int) codec_left.size();
}
// From here, two paths: standard (no ICL) builds the trailing
// utterance text + tts_eos + final_pad, ICL builds an aligned
// text/codec block that replaces those rows entirely.
if (!icl) {
// Standard layout: trailing utterance text + tts_eos rows summed
// with codec_pad, then a final tts_pad + codec_bos row.
for (int i = 0; i < N_text; i++) {
std::vector<float> e((size_t) text_hid);
std::vector<float> y((size_t) hidden);
embed_row_to_f32(pt->gguf_talker, "talker.text_embd.weight", ids[3 + i], text_hid, e.data());
text_projection_apply(&tp, e.data(), y.data());
float * r = row_ptr(row);
std::memcpy(r, y.data(), (size_t) hidden * sizeof(float));
vec_add(r, codec_pad_emb.data(), hidden);
row++;
}
{
float * r = row_ptr(row);
std::memcpy(r, tts_eos_emb.data(), (size_t) hidden * sizeof(float));
vec_add(r, codec_pad_emb.data(), hidden);
row++;
}
{
float * r = row_ptr(row);
std::memcpy(r, tts_pad_emb.data(), (size_t) hidden * sizeof(float));
std::vector<float> ce((size_t) hidden);
embed_row_to_f32(pt->gguf_talker, "talker.codec_embd.weight", pt->codec_specials.bos_id, hidden, ce.data());
vec_add(r, ce.data(), hidden);
row++;
}
} else {
// ICL layout: compute the text stream and the codec stream
// separately then add them. The text stream is text_proj of
// [ref_text_ids; utterance_text_ids] followed by tts_eos. The
// codec stream is codec_emb(codec_bos) followed by sum over the
// 16 codebook embeddings of ref_codes[i, t] for each frame t.
// Both streams are aligned to length icl_T per the upstream
// non_streaming_mode=False branch of generate_icl_prompt.
const int T_icl = codec_lens_icl; // text_lens > codec: truncate to codec, else pad text up to codec
// Build the codec stream [T_icl, hidden]. Row 0: codec_emb(codec_bos).
// Row 1..ref_codes_T: sum over k=0..15 of codebook_k_emb(ref_codes[k, t]).
std::vector<float> codec_stream((size_t) T_icl * (size_t) hidden, 0.0f);
{
// Row 0: codec_bos lookup.
embed_row_to_f32(pt->gguf_talker, "talker.codec_embd.weight", pt->codec_specials.bos_id, hidden,
codec_stream.data());
// Row 1..ref_codes_T: sum over codebooks.
std::vector<float> tmp((size_t) hidden);
for (int t = 0; t < ref_codes_T; t++) {
float * dst = codec_stream.data() + (size_t) (1 + t) * (size_t) hidden;
// codebook 0 lives in talker.codec_embd
int code0 = ref_codes[(size_t) 0 * (size_t) ref_codes_T + (size_t) t];
embed_row_to_f32(pt->gguf_talker, "talker.codec_embd.weight", code0, hidden, dst);
// codebooks 1..15 live in code_pred.codec_embd.{i-1}
for (int i = 1; i < pt->num_code_groups; i++) {
int code = ref_codes[(size_t) i * (size_t) ref_codes_T + (size_t) t];
char tname[64];
std::snprintf(tname, sizeof(tname), "code_pred.codec_embd.%d.weight", i - 1);
embed_row_to_f32(pt->gguf_talker, tname, code, hidden, tmp.data());
vec_add(dst, tmp.data(), hidden);
}
}
}
// Build the text stream [text_lens_icl, hidden] = text_proj of
// [ref_text; utterance_text] then tts_eos.
std::vector<float> text_stream((size_t) text_lens_icl * (size_t) hidden, 0.0f);
for (int i = 0; i < N_ref_text; i++) {
std::vector<float> e((size_t) text_hid);
float * r = text_stream.data() + (size_t) i * (size_t) hidden;
embed_row_to_f32(pt->gguf_talker, "talker.text_embd.weight", ref_ids[3 + i], text_hid, e.data());
text_projection_apply(&tp, e.data(), r);
}
for (int i = 0; i < N_text; i++) {
std::vector<float> e((size_t) text_hid);
float * r = text_stream.data() + (size_t) (N_ref_text + i) * (size_t) hidden;
embed_row_to_f32(pt->gguf_talker, "talker.text_embd.weight", ids[3 + i], text_hid, e.data());
text_projection_apply(&tp, e.data(), r);
}
// Append tts_eos at the end of the text stream.
std::memcpy(text_stream.data() + (size_t) (text_lens_icl - 1) * (size_t) hidden, tts_eos_emb.data(),
(size_t) hidden * sizeof(float));
// Align the two streams to T_icl. text_lens > codec: truncate
// text and stash the leftover into trailing_text_hidden. text_lens
// <= codec: pad text with tts_pad up to codec, trailing reduces
// to tts_pad.
std::vector<float> aligned_text((size_t) T_icl * (size_t) hidden, 0.0f);
if (text_lens_icl >= T_icl) {
// truncate text to T_icl rows, leftover goes into trailing
std::memcpy(aligned_text.data(), text_stream.data(), (size_t) T_icl * (size_t) hidden * sizeof(float));
const int trailing_n = text_lens_icl - T_icl;
out->T_trailing = trailing_n > 0 ? trailing_n : 1;
out->trailing_text_hidden.assign((size_t) out->T_trailing * (size_t) hidden, 0.0f);
if (trailing_n > 0) {
std::memcpy(out->trailing_text_hidden.data(), text_stream.data() + (size_t) T_icl * (size_t) hidden,
(size_t) trailing_n * (size_t) hidden * sizeof(float));
} else {
std::memcpy(out->trailing_text_hidden.data(), tts_pad_emb.data(), (size_t) hidden * sizeof(float));
}
} else {
// pad text with tts_pad up to T_icl, trailing = single tts_pad row
std::memcpy(aligned_text.data(), text_stream.data(),
(size_t) text_lens_icl * (size_t) hidden * sizeof(float));
for (int i = text_lens_icl; i < T_icl; i++) {
std::memcpy(aligned_text.data() + (size_t) i * (size_t) hidden, tts_pad_emb.data(),
(size_t) hidden * sizeof(float));
}
out->T_trailing = 1;
out->trailing_text_hidden.assign((size_t) hidden, 0.0f);
std::memcpy(out->trailing_text_hidden.data(), tts_pad_emb.data(), (size_t) hidden * sizeof(float));
}
// Sum aligned_text + codec_stream into the input embed at the
// current row offset.
for (int i = 0; i < T_icl; i++) {
float * r = row_ptr(row + i);
std::memcpy(r, aligned_text.data() + (size_t) i * (size_t) hidden, (size_t) hidden * sizeof(float));
vec_add(r, codec_stream.data() + (size_t) i * (size_t) hidden, hidden);
}
row += T_icl;
}
if (row != T_ctx) {
fprintf(stderr, "[Prompt] FATAL: layout error row=%d expected T_ctx=%d\n", row, T_ctx);
return false;
}
// Trailing text hidden: non streaming mode (no ICL) collapses the
// overlay to a single row equal to tts_pad_embed
// (modeling_qwen3_tts.py line 2227). The full utterance text is
// already integrated into the prefill above as codec_pad summed text
// rows + tts_eos, so the overlay loop only ever needs tts_pad: step
// 0 reads trailing_text_hidden[0] which is tts_pad, every later step
// falls through to the else branch and reads tts_pad_embed. One row,
// bit exact with the Python hook dump.
//
// ICL mode populates out->trailing_text_hidden directly inside the
// ICL branch above, so we only set the default here for non ICL.
if (!icl) {
out->T_trailing = 1;
out->trailing_text_hidden.assign((size_t) hidden, 0.0f);
std::memcpy(out->trailing_text_hidden.data(), tts_pad_emb.data(), (size_t) hidden * sizeof(float));
}
out->tts_pad_embed = tts_pad_emb;
fprintf(stderr,
"[Prompt] Built: %d ids, N_text=%d, N_instruct=%d, T_ctx=%d, hidden=%d, lang=%s (id=%d), speaker=%s "
"(id=%d) ref_spk_emb=%s icl=%s\n",
N, N_text, N_instruct, T_ctx, hidden, language.c_str(), language_id,
speaker_name.empty() ? "none" : speaker_name.c_str(), speaker_id, ref_spk_emb ? "yes" : "no",
icl ? "yes" : "no");
return true;
}
src/talker-forward.cpp
-359
View File
@@ -1,359 +0,0 @@
// talker-forward.cpp: eager prefill graph for the Talker LM.
//
// Mirrors Qwen3TTSTalkerDecoderLayer for TTS-only operation :
// pre-norm, GQA attention with per-head QK-norm, mrope collapsed to
// 1D NEOX (since the three multimodal axes share position ids in TTS
// mode), SwiGLU MLP, two residuals, repeated 28 times, then final
// RMSNorm and codec_head. Eager softmax in F32, no flash-attention,
// no KV cache.
//
// Tensor shapes follow ggml row-major convention: ne[0] is the fastest
// axis. Our input embedding lives as [hidden, T] inside the graph and
// the loader feeds it from a [T, hidden] f32 row-major host buffer
// (which becomes [hidden, T] in ggml after a 2d view because rows on
// the host are contiguous along the hidden axis).
#include "talker-forward.h"
#include "ggml-alloc.h"
#include "ggml-backend.h"
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <vector>
// Bisect layers dumped when a dump_dir is set. Match the Python hook list
// in tests/debug-tts-cossim.py: 0, 7, 14, 21, 27.
static const int BISECT_LAYERS[] = { 0, 7, 14, 21, 27 };
static const int N_BISECT_LAYERS = (int) (sizeof(BISECT_LAYERS) / sizeof(BISECT_LAYERS[0]));
static bool is_bisect_layer(int l) {
for (int i = 0; i < N_BISECT_LAYERS; i++) {
if (BISECT_LAYERS[i] == l) {
return true;
}
}
return false;
}
// Build the per-layer block, KV cached. K and V for the T fresh
// positions get computed normally, then written into the cache at
// [n_past, n_past+T) on dim 1. The attention reads the contiguous slice
// [0, n_past+T) on the same dim, which covers the full causal context
// in one tensor view. Returns the layer output [hidden, T].
static struct ggml_tensor * talker_layer_forward(struct ggml_context * ctx,
const TalkerWeights * tw,
const TalkerLayer & layer,
struct ggml_tensor * x,
struct ggml_tensor * positions,
struct ggml_tensor * mask,
struct ggml_tensor * k_cache,
struct ggml_tensor * v_cache,
int n_past,
int T,
struct ggml_cgraph * gf) {
const int n_q_heads = tw->num_attention_heads;
const int n_kv = tw->num_key_value_heads;
const int hd = tw->head_dim;
const float eps = tw->rms_norm_eps;
// Pre-norm
struct ggml_tensor * h = ggml_rms_norm(ctx, x, eps);
h = ggml_mul(ctx, h, layer.input_norm_w);
// Q/K/V projections
struct ggml_tensor * q = ggml_mul_mat(ctx, layer.attn.q_proj_w, h); // [n_q_heads*hd, T]
struct ggml_tensor * k = ggml_mul_mat(ctx, layer.attn.k_proj_w, h); // [n_kv*hd, T]
struct ggml_tensor * v = ggml_mul_mat(ctx, layer.attn.v_proj_w, h); // [n_kv*hd, T]
q = ggml_reshape_3d(ctx, q, hd, n_q_heads, T); // [hd, n_q_heads, T]
k = ggml_reshape_3d(ctx, k, hd, n_kv, T);
v = ggml_reshape_3d(ctx, v, hd, n_kv, T);
// Per-head QK-norm: RMS over hd, then multiply by [hd] gain. The
// norm operates on ne[0] = hd, identical layout for q (16 heads) and
// k (8 heads), so the same code path covers both.
q = ggml_rms_norm(ctx, q, eps);
q = ggml_mul(ctx, q, layer.attn.q_norm_w);
k = ggml_rms_norm(ctx, k, eps);
k = ggml_mul(ctx, k, layer.attn.k_norm_w);
// RoPE NEOX (half-split). In TTS-only mode the three mrope axes share
// position ids, so the multimodal interleaved cos/sin collapses to
// plain 1D rotate_half with the same freq base.
q = ggml_rope_ext(ctx, q, positions, NULL, hd, GGML_ROPE_TYPE_NEOX, 0, tw->rope_theta, 1.0f, 0.0f, 1.0f, 0.0f,
0.0f);
k = ggml_rope_ext(ctx, k, positions, NULL, hd, GGML_ROPE_TYPE_NEOX, 0, tw->rope_theta, 1.0f, 0.0f, 1.0f, 0.0f,
0.0f);
// Write the T fresh positions of K and V into the cache. K and V are
// [hd, n_kv, T] at this point and the cache lives as [hd, max_T, n_kv]
// so we permute to [hd, T, n_kv] before the cpy. The destination view
// covers [hd, T, n_kv] at dim 1 offset n_past * nb1.
struct ggml_tensor * k_perm = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3)); // [hd, T, n_kv]
struct ggml_tensor * v_perm = ggml_cont(ctx, ggml_permute(ctx, v, 0, 2, 1, 3)); // [hd, T, n_kv]
size_t k_off = (size_t) n_past * k_cache->nb[1];
size_t v_off = (size_t) n_past * v_cache->nb[1];
struct ggml_tensor * k_dst = ggml_view_3d(ctx, k_cache, hd, T, n_kv, k_cache->nb[1], k_cache->nb[2], k_off);
struct ggml_tensor * v_dst = ggml_view_3d(ctx, v_cache, hd, T, n_kv, v_cache->nb[1], v_cache->nb[2], v_off);
struct ggml_tensor * k_cpy = ggml_cpy(ctx, k_perm, k_dst);
struct ggml_tensor * v_cpy = ggml_cpy(ctx, v_perm, v_dst);
ggml_build_forward_expand(gf, k_cpy);
ggml_build_forward_expand(gf, v_cpy);
// Read the [0, n_past + T) window for attention. The cache slice is
// already in the [hd, T_full, n_kv] layout flash_attn_ext expects for
// K and V (n_embd, n_kv, n_head_kv, ne3), passed directly as views.
const int T_full = n_past + T;
struct ggml_tensor * k_full = ggml_view_3d(ctx, k_cache, hd, T_full, n_kv, k_cache->nb[1], k_cache->nb[2], 0);
struct ggml_tensor * v_full = ggml_view_3d(ctx, v_cache, hd, T_full, n_kv, v_cache->nb[1], v_cache->nb[2], 0);
// Q permute [hd, n_q_heads, T] -> [hd, T, n_q_heads]. flash_attn_ext
// expects ne[1] = n_batch and ne[2] = n_head; the GQA broadcast
// check ggml_can_mul_mat is n_head % n_head_kv == 0, matching K layout.
// No cont: flash_attn_ext takes the view directly, like acestep does.
struct ggml_tensor * q_p = ggml_permute(ctx, q, 0, 2, 1, 3);
// Fused flash attention. The manual mul_mat + soft_max_ext + mul_mat
// path it replaces had a Vulkan bug in autoregressive decode (T=1)
// where the KV cache view stride on dim 2 (= max_T * hd, non
// contiguous with dim 1 of length T_full) caused the second mul_mat
// to diverge silently. Flash attention has a backend-tested kernel
// on every target (CPU, CUDA, Vulkan), matches the working acestep
// qw3lm_build_attn pattern.
float scale = 1.0f / sqrtf((float) hd);
struct ggml_tensor * attn = ggml_flash_attn_ext(ctx, q_p, k_full, v_full, mask, scale, 0.0f, 0.0f);
ggml_flash_attn_ext_set_prec(attn, GGML_PREC_F32);
// Flash attention output is [hd, n_q_heads, T], flatten heads for o_proj.
attn = ggml_reshape_2d(ctx, attn, n_q_heads * hd, T);
struct ggml_tensor * o = ggml_mul_mat(ctx, layer.attn.o_proj_w, attn);
x = ggml_add(ctx, x, o);
// MLP block: pre-norm + SwiGLU + residual
struct ggml_tensor * h2 = ggml_rms_norm(ctx, x, eps);
h2 = ggml_mul(ctx, h2, layer.post_attn_norm_w);
struct ggml_tensor * gate = ggml_mul_mat(ctx, layer.mlp.gate_proj_w, h2);
struct ggml_tensor * up = ggml_mul_mat(ctx, layer.mlp.up_proj_w, h2);
gate = ggml_silu(ctx, gate);
struct ggml_tensor * gu = ggml_mul(ctx, gate, up);
struct ggml_tensor * mlp = ggml_mul_mat(ctx, layer.mlp.down_proj_w, gu);
x = ggml_add(ctx, x, mlp);
return x;
}
// 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.
static bool talker_forward_core(const TalkerWeights * tw,
KVCache * kv,
ggml_backend_sched_t sched,
const float * input_embed,
int T,
int n_past,
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;
const int T_full = n_past + T;
// Dedicated context for graph + IO tensors. Counts approximate :
// per layer: ~38 ops (added 2 cpy + 4 views per layer for KV)
// IO : 4 tensors (input embed, positions, mask, output norm)
// final : norm + codec_head + dump branches
const int max_nodes = 48 * n_layers + 256;
const size_t graph_arena_bytes = ggml_tensor_overhead() * max_nodes + ggml_graph_overhead_custom(max_nodes, false);
struct ggml_init_params gparams = {
graph_arena_bytes,
NULL,
true,
};
struct ggml_context * gctx = ggml_init(gparams);
if (!gctx) {
fprintf(stderr, "[TalkerForward] FATAL: ggml_init failed\n");
return false;
}
// IO tensors: input embedding, positions, causal mask. The mask
// spans [T_full, T]: for each fresh query q in [0, T) we allow
// keys k in [0, n_past + q]. In the decode case (T=1) this is a
// single row of zeros of length 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, T_full, T);
ggml_set_name(x_in, "input_embed");
ggml_set_name(pos_in, "positions");
ggml_set_name(mask_in, "causal_mask");
struct ggml_cgraph * gf = ggml_new_graph_custom(gctx, max_nodes, false);
// Build the layer stack. Bisect taps fire on prefill only.
const bool record_taps = (dump_dir != NULL) && (n_past == 0);
struct ggml_tensor * h = x_in;
std::vector<struct ggml_tensor *> taps(N_BISECT_LAYERS, NULL);
for (int l = 0; l < n_layers; l++) {
h = talker_layer_forward(gctx, tw, tw->layers[(size_t) l], h, pos_in, mask_in, kv->k[(size_t) l],
kv->v[(size_t) l], n_past, T, gf);
if (record_taps && is_bisect_layer(l)) {
for (int i = 0; i < N_BISECT_LAYERS; i++) {
if (BISECT_LAYERS[i] == l) {
char tap_name[64];
snprintf(tap_name, sizeof(tap_name), "tap_l%d", l);
ggml_set_name(h, tap_name);
ggml_set_output(h);
taps[(size_t) i] = h;
break;
}
}
}
}
struct ggml_tensor * h_final = ggml_rms_norm(gctx, h, tw->rms_norm_eps);
h_final = ggml_mul(gctx, h_final, tw->norm_w);
ggml_set_name(h_final, "hidden_final");
ggml_set_output(h_final);
// codec_head: [hidden, vocab]. ggml_mul_mat returns [vocab, T].
struct ggml_tensor * logits = ggml_mul_mat(gctx, tw->codec_head_w, h_final);
ggml_set_name(logits, "logits");
if (record_taps) {
for (int i = 0; i < N_BISECT_LAYERS; i++) {
if (taps[(size_t) i]) {
ggml_build_forward_expand(gf, taps[(size_t) i]);
}
}
ggml_build_forward_expand(gf, h_final);
}
ggml_build_forward_expand(gf, logits);
if (!ggml_backend_sched_alloc_graph(sched, gf)) {
fprintf(stderr, "[TalkerForward] FATAL: graph allocation failed\n");
ggml_backend_sched_reset(sched);
ggml_free(gctx);
return false;
}
// 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));
// Positions: n_past .. n_past + T - 1
{
std::vector<int32_t> pos((size_t) T);
for (int i = 0; i < T; i++) {
pos[(size_t) i] = n_past + i;
}
ggml_backend_tensor_set(pos_in, pos.data(), 0, (size_t) T * sizeof(int32_t));
}
// Causal mask: 0 where k <= n_past + q, -inf otherwise. Stored
// row-major [T_q, T_k] with T_k as the fast axis (ne[0]). F16 dtype
// matches ggml_flash_attn_ext convention used by the attention path.
{
std::vector<ggml_fp16_t> mask((size_t) T * (size_t) T_full);
const ggml_fp16_t zero = ggml_fp32_to_fp16(0.0f);
const ggml_fp16_t neg_inf = ggml_fp32_to_fp16(-INFINITY);
for (size_t i = 0; i < mask.size(); i++) {
mask[i] = neg_inf;
}
for (int q = 0; q < T; q++) {
const int q_pos = n_past + q;
for (int k = 0; k <= q_pos; k++) {
mask[(size_t) q * (size_t) T_full + (size_t) k] = zero;
}
}
ggml_backend_tensor_set(mask_in, mask.data(), 0, mask.size() * sizeof(ggml_fp16_t));
}
if (ggml_backend_sched_graph_compute(sched, gf) != GGML_STATUS_SUCCESS) {
fprintf(stderr, "[TalkerForward] FATAL: graph compute failed\n");
ggml_backend_sched_reset(sched);
ggml_free(gctx);
return false;
}
// Bisect dumps: pull each tap [hidden, T] back to host as [T, hidden].
if (record_taps) {
DebugDumper d;
debug_init(&d, dump_dir);
std::vector<float> buf((size_t) T * (size_t) hidden);
for (int i = 0; i < N_BISECT_LAYERS; i++) {
if (!taps[(size_t) i]) {
continue;
}
ggml_backend_tensor_get(taps[(size_t) i], buf.data(), 0, buf.size() * sizeof(float));
char name[64];
snprintf(name, sizeof(name), "talker-hidden-prefill-l%d", BISECT_LAYERS[i]);
debug_dump_2d(&d, name, buf.data(), T, hidden);
}
ggml_backend_tensor_get(h_final, buf.data(), 0, buf.size() * sizeof(float));
debug_dump_2d(&d, "talker-hidden-prefill-final", buf.data(), T, hidden);
}
// Pull the last position: final hidden + logits
out->hidden = hidden;
out->vocab = vocab;
out->hidden_last.assign((size_t) hidden, 0.0f);
out->logits_last.assign((size_t) vocab, 0.0f);
{
size_t row_bytes = (size_t) vocab * sizeof(float);
ggml_backend_tensor_get(logits, out->logits_last.data(), (size_t) (T - 1) * row_bytes, row_bytes);
size_t hrow_bytes = (size_t) hidden * sizeof(float);
ggml_backend_tensor_get(h_final, out->hidden_last.data(), (size_t) (T - 1) * hrow_bytes, hrow_bytes);
}
if (record_taps) {
DebugDumper d;
debug_init(&d, dump_dir);
debug_dump_1d(&d, "talker-logits-prefill", out->logits_last.data(), vocab);
}
// Advance the cache write head. The graph already executed the cpy
// nodes so positions [n_past, n_past + T) are now populated.
kv->cur_len = T_full;
ggml_backend_sched_reset(sched);
ggml_free(gctx);
return true;
}
bool talker_forward_prefill(const TalkerWeights * tw,
KVCache * kv,
ggml_backend_sched_t sched,
const float * input_embed,
int T,
const char * dump_dir,
TalkerForwardOutput * out) {
kv_cache_reset(kv);
if (T > kv->max_seq_len) {
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, input_embed, T, 0, dump_dir, out);
}
bool talker_forward_decode(const TalkerWeights * tw,
KVCache * kv,
ggml_backend_sched_t sched,
const float * input_embed_1,
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, input_embed_1, 1, kv->cur_len, NULL, out);
}
src/talker-forward.h
+350 -13
View File
@@ -1,6 +1,6 @@
#pragma once
// talker-forward.h: prefill + decode forwards of the Talker LM, KV
// cached.
// cached. Eager prefill graph for the Talker LM.
//
// Both entry points run the same 28-layer Qwen3 decoder stack with
// multimodal RoPE collapsed to 1D NEOX, GQA attention with per-head
@@ -20,17 +20,34 @@
// the predictor has produced its 15 acoustic codes and the loop has
// summed the codec embeddings into next_emb.
//
// Mirrors Qwen3TTSTalkerDecoderLayer for TTS-only operation:
// pre-norm, GQA attention with per-head QK-norm, mrope collapsed to
// 1D NEOX (since the three multimodal axes share position ids in TTS
// mode), SwiGLU MLP, two residuals, repeated 28 times, then final
// RMSNorm and codec_head.
//
// Tensor shapes follow ggml row-major convention: ne[0] is the fastest
// axis. Our input embedding lives as [hidden, T] inside the graph and
// the loader feeds it from a [T, hidden] f32 row-major host buffer
// (which becomes [hidden, T] in ggml after a 2d view because rows on
// the host are contiguous along the hidden axis).
//
// The Python reference uses pure causal attention with no sliding
// window, so the cache is a plain causal ring.
#include "backend.h"
#include "debug.h"
#include "ggml-alloc.h"
#include "ggml-backend.h"
#include "ggml.h"
#include "kv-cache.h"
#include "talker-weights.h"
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <vector>
struct TalkerForwardOutput {
@@ -44,21 +61,341 @@ struct TalkerForwardOutput {
int vocab;
};
// Bisect layers dumped when a dump_dir is set. Match the Python hook list
// in tests/debug-tts-cossim.py: 0, 7, 14, 21, 27.
static const int TALKER_BISECT_LAYERS[] = { 0, 7, 14, 21, 27 };
static const int TALKER_N_BISECT_LAYERS = (int) (sizeof(TALKER_BISECT_LAYERS) / sizeof(TALKER_BISECT_LAYERS[0]));
static bool talker_is_bisect_layer(int l) {
for (int i = 0; i < TALKER_N_BISECT_LAYERS; i++) {
if (TALKER_BISECT_LAYERS[i] == l) {
return true;
}
}
return false;
}
// Build the per-layer block, KV cached. K and V for the T fresh
// positions get computed normally, then written into the cache at
// [n_past, n_past+T) on dim 1. The attention reads the contiguous slice
// [0, n_past+T) on the same dim, which covers the full causal context
// in one tensor view. Returns the layer output [hidden, T].
static struct ggml_tensor * talker_layer_forward(struct ggml_context * ctx,
const TalkerWeights * tw,
const TalkerLayer & layer,
struct ggml_tensor * x,
struct ggml_tensor * positions,
struct ggml_tensor * mask,
struct ggml_tensor * k_cache,
struct ggml_tensor * v_cache,
int n_past,
int T,
struct ggml_cgraph * gf) {
const int n_q_heads = tw->num_attention_heads;
const int n_kv = tw->num_key_value_heads;
const int hd = tw->head_dim;
const float eps = tw->rms_norm_eps;
// Pre-norm
struct ggml_tensor * h = ggml_rms_norm(ctx, x, eps);
h = ggml_mul(ctx, h, layer.input_norm_w);
// Q/K/V projections
struct ggml_tensor * q = ggml_mul_mat(ctx, layer.attn.q_proj_w, h); // [n_q_heads*hd, T]
struct ggml_tensor * k = ggml_mul_mat(ctx, layer.attn.k_proj_w, h); // [n_kv*hd, T]
struct ggml_tensor * v = ggml_mul_mat(ctx, layer.attn.v_proj_w, h); // [n_kv*hd, T]
q = ggml_reshape_3d(ctx, q, hd, n_q_heads, T); // [hd, n_q_heads, T]
k = ggml_reshape_3d(ctx, k, hd, n_kv, T);
v = ggml_reshape_3d(ctx, v, hd, n_kv, T);
// Per-head QK-norm: RMS over hd, then multiply by [hd] gain. The
// norm operates on ne[0] = hd, identical layout for q (16 heads) and
// k (8 heads), so the same code path covers both.
q = ggml_rms_norm(ctx, q, eps);
q = ggml_mul(ctx, q, layer.attn.q_norm_w);
k = ggml_rms_norm(ctx, k, eps);
k = ggml_mul(ctx, k, layer.attn.k_norm_w);
// RoPE NEOX (half-split). In TTS-only mode the three mrope axes share
// position ids, so the multimodal interleaved cos/sin collapses to
// plain 1D rotate_half with the same freq base.
q = ggml_rope_ext(ctx, q, positions, NULL, hd, GGML_ROPE_TYPE_NEOX, 0, tw->rope_theta, 1.0f, 0.0f, 1.0f, 0.0f,
0.0f);
k = ggml_rope_ext(ctx, k, positions, NULL, hd, GGML_ROPE_TYPE_NEOX, 0, tw->rope_theta, 1.0f, 0.0f, 1.0f, 0.0f,
0.0f);
// Write the T fresh positions of K and V into the cache. K and V are
// [hd, n_kv, T] at this point and the cache lives as [hd, max_T, n_kv]
// so we permute to [hd, T, n_kv] before the cpy. The destination view
// covers [hd, T, n_kv] at dim 1 offset n_past * nb1.
struct ggml_tensor * k_perm = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3)); // [hd, T, n_kv]
struct ggml_tensor * v_perm = ggml_cont(ctx, ggml_permute(ctx, v, 0, 2, 1, 3)); // [hd, T, n_kv]
size_t k_off = (size_t) n_past * k_cache->nb[1];
size_t v_off = (size_t) n_past * v_cache->nb[1];
struct ggml_tensor * k_dst = ggml_view_3d(ctx, k_cache, hd, T, n_kv, k_cache->nb[1], k_cache->nb[2], k_off);
struct ggml_tensor * v_dst = ggml_view_3d(ctx, v_cache, hd, T, n_kv, v_cache->nb[1], v_cache->nb[2], v_off);
struct ggml_tensor * k_cpy = ggml_cpy(ctx, k_perm, k_dst);
struct ggml_tensor * v_cpy = ggml_cpy(ctx, v_perm, v_dst);
ggml_build_forward_expand(gf, k_cpy);
ggml_build_forward_expand(gf, v_cpy);
// Read the [0, n_past + T) window for attention. The cache slice is
// already in the [hd, T_full, n_kv] layout flash_attn_ext expects for
// K and V (n_embd, n_kv, n_head_kv, ne3), passed directly as views.
const int T_full = n_past + T;
struct ggml_tensor * k_full = ggml_view_3d(ctx, k_cache, hd, T_full, n_kv, k_cache->nb[1], k_cache->nb[2], 0);
struct ggml_tensor * v_full = ggml_view_3d(ctx, v_cache, hd, T_full, n_kv, v_cache->nb[1], v_cache->nb[2], 0);
// Q permute [hd, n_q_heads, T] -> [hd, T, n_q_heads]. flash_attn_ext
// expects ne[1] = n_batch and ne[2] = n_head; the GQA broadcast
// check ggml_can_mul_mat is n_head % n_head_kv == 0, matching K layout.
// No cont: flash_attn_ext takes the view directly, like acestep does.
struct ggml_tensor * q_p = ggml_permute(ctx, q, 0, 2, 1, 3);
// Fused flash attention. The manual mul_mat + soft_max_ext + mul_mat
// path it replaces had a Vulkan bug in autoregressive decode (T=1)
// where the KV cache view stride on dim 2 (= max_T * hd, non
// contiguous with dim 1 of length T_full) caused the second mul_mat
// to diverge silently. Flash attention has a backend-tested kernel
// on every target (CPU, CUDA, Vulkan), matches the working acestep
// qw3lm_build_attn pattern.
float scale = 1.0f / sqrtf((float) hd);
struct ggml_tensor * attn = ggml_flash_attn_ext(ctx, q_p, k_full, v_full, mask, scale, 0.0f, 0.0f);
ggml_flash_attn_ext_set_prec(attn, GGML_PREC_F32);
// Flash attention output is [hd, n_q_heads, T], flatten heads for o_proj.
attn = ggml_reshape_2d(ctx, attn, n_q_heads * hd, T);
struct ggml_tensor * o = ggml_mul_mat(ctx, layer.attn.o_proj_w, attn);
x = ggml_add(ctx, x, o);
// MLP block: pre-norm + SwiGLU + residual
struct ggml_tensor * h2 = ggml_rms_norm(ctx, x, eps);
h2 = ggml_mul(ctx, h2, layer.post_attn_norm_w);
struct ggml_tensor * gate = ggml_mul_mat(ctx, layer.mlp.gate_proj_w, h2);
struct ggml_tensor * up = ggml_mul_mat(ctx, layer.mlp.up_proj_w, h2);
gate = ggml_silu(ctx, gate);
struct ggml_tensor * gu = ggml_mul(ctx, gate, up);
struct ggml_tensor * mlp = ggml_mul_mat(ctx, layer.mlp.down_proj_w, gu);
x = ggml_add(ctx, x, mlp);
return x;
}
// 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.
static bool talker_forward_core(const TalkerWeights * tw,
KVCache * kv,
ggml_backend_sched_t sched,
const float * input_embed,
int T,
int n_past,
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;
const int T_full = n_past + T;
// Dedicated context for graph + IO tensors. Counts approximate:
// per layer: ~38 ops (added 2 cpy + 4 views per layer for KV)
// IO : 4 tensors (input embed, positions, mask, output norm)
// final : norm + codec_head + dump branches
const int max_nodes = 48 * n_layers + 256;
const size_t graph_arena_bytes = ggml_tensor_overhead() * max_nodes + ggml_graph_overhead_custom(max_nodes, false);
struct ggml_init_params gparams = {
graph_arena_bytes,
NULL,
true,
};
struct ggml_context * gctx = ggml_init(gparams);
if (!gctx) {
fprintf(stderr, "[TalkerForward] FATAL: ggml_init failed\n");
return false;
}
// IO tensors: input embedding, positions, causal mask. The mask
// spans [T_full, T]: for each fresh query q in [0, T) we allow
// keys k in [0, n_past + q]. In the decode case (T=1) this is a
// single row of zeros of length 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, T_full, T);
ggml_set_name(x_in, "input_embed");
ggml_set_name(pos_in, "positions");
ggml_set_name(mask_in, "causal_mask");
struct ggml_cgraph * gf = ggml_new_graph_custom(gctx, max_nodes, false);
// Build the layer stack. Bisect taps fire on prefill only.
const bool record_taps = (dump_dir != NULL) && (n_past == 0);
struct ggml_tensor * h = x_in;
std::vector<struct ggml_tensor *> taps(TALKER_N_BISECT_LAYERS, NULL);
for (int l = 0; l < n_layers; l++) {
h = talker_layer_forward(gctx, tw, tw->layers[(size_t) l], h, pos_in, mask_in, kv->k[(size_t) l],
kv->v[(size_t) l], n_past, T, gf);
if (record_taps && talker_is_bisect_layer(l)) {
for (int i = 0; i < TALKER_N_BISECT_LAYERS; i++) {
if (TALKER_BISECT_LAYERS[i] == l) {
char tap_name[64];
snprintf(tap_name, sizeof(tap_name), "tap_l%d", l);
ggml_set_name(h, tap_name);
ggml_set_output(h);
taps[(size_t) i] = h;
break;
}
}
}
}
struct ggml_tensor * h_final = ggml_rms_norm(gctx, h, tw->rms_norm_eps);
h_final = ggml_mul(gctx, h_final, tw->norm_w);
ggml_set_name(h_final, "hidden_final");
ggml_set_output(h_final);
// codec_head: [hidden, vocab]. ggml_mul_mat returns [vocab, T].
struct ggml_tensor * logits = ggml_mul_mat(gctx, tw->codec_head_w, h_final);
ggml_set_name(logits, "logits");
if (record_taps) {
for (int i = 0; i < TALKER_N_BISECT_LAYERS; i++) {
if (taps[(size_t) i]) {
ggml_build_forward_expand(gf, taps[(size_t) i]);
}
}
ggml_build_forward_expand(gf, h_final);
}
ggml_build_forward_expand(gf, logits);
if (!ggml_backend_sched_alloc_graph(sched, gf)) {
fprintf(stderr, "[TalkerForward] FATAL: graph allocation failed\n");
ggml_backend_sched_reset(sched);
ggml_free(gctx);
return false;
}
// 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));
// Positions: n_past .. n_past + T - 1
{
std::vector<int32_t> pos((size_t) T);
for (int i = 0; i < T; i++) {
pos[(size_t) i] = n_past + i;
}
ggml_backend_tensor_set(pos_in, pos.data(), 0, (size_t) T * sizeof(int32_t));
}
// Causal mask: 0 where k <= n_past + q, -inf otherwise. Stored
// row-major [T_q, T_k] with T_k as the fast axis (ne[0]). F16 dtype
// matches ggml_flash_attn_ext convention used by the attention path.
{
std::vector<ggml_fp16_t> mask((size_t) T * (size_t) T_full);
const ggml_fp16_t zero = ggml_fp32_to_fp16(0.0f);
const ggml_fp16_t neg_inf = ggml_fp32_to_fp16(-INFINITY);
for (size_t i = 0; i < mask.size(); i++) {
mask[i] = neg_inf;
}
for (int q = 0; q < T; q++) {
const int q_pos = n_past + q;
for (int k = 0; k <= q_pos; k++) {
mask[(size_t) q * (size_t) T_full + (size_t) k] = zero;
}
}
ggml_backend_tensor_set(mask_in, mask.data(), 0, mask.size() * sizeof(ggml_fp16_t));
}
if (ggml_backend_sched_graph_compute(sched, gf) != GGML_STATUS_SUCCESS) {
fprintf(stderr, "[TalkerForward] FATAL: graph compute failed\n");
ggml_backend_sched_reset(sched);
ggml_free(gctx);
return false;
}
// Bisect dumps: pull each tap [hidden, T] back to host as [T, hidden].
if (record_taps) {
DebugDumper d;
debug_init(&d, dump_dir);
std::vector<float> buf((size_t) T * (size_t) hidden);
for (int i = 0; i < TALKER_N_BISECT_LAYERS; i++) {
if (!taps[(size_t) i]) {
continue;
}
ggml_backend_tensor_get(taps[(size_t) i], buf.data(), 0, buf.size() * sizeof(float));
char name[64];
snprintf(name, sizeof(name), "talker-hidden-prefill-l%d", TALKER_BISECT_LAYERS[i]);
debug_dump_2d(&d, name, buf.data(), T, hidden);
}
ggml_backend_tensor_get(h_final, buf.data(), 0, buf.size() * sizeof(float));
debug_dump_2d(&d, "talker-hidden-prefill-final", buf.data(), T, hidden);
}
// Pull the last position: final hidden + logits
out->hidden = hidden;
out->vocab = vocab;
out->hidden_last.assign((size_t) hidden, 0.0f);
out->logits_last.assign((size_t) vocab, 0.0f);
{
size_t row_bytes = (size_t) vocab * sizeof(float);
ggml_backend_tensor_get(logits, out->logits_last.data(), (size_t) (T - 1) * row_bytes, row_bytes);
size_t hrow_bytes = (size_t) hidden * sizeof(float);
ggml_backend_tensor_get(h_final, out->hidden_last.data(), (size_t) (T - 1) * hrow_bytes, hrow_bytes);
}
if (record_taps) {
DebugDumper d;
debug_init(&d, dump_dir);
debug_dump_1d(&d, "talker-logits-prefill", out->logits_last.data(), vocab);
}
// Advance the cache write head. The graph already executed the cpy
// nodes so positions [n_past, n_past + T) are now populated.
kv->cur_len = T_full;
ggml_backend_sched_reset(sched);
ggml_free(gctx);
return true;
}
// Prefill: reset the cache and write T_ctx positions in one shot.
// input_embed is [T, hidden] f32 row-major. dump_dir may be NULL.
bool talker_forward_prefill(const TalkerWeights * tw,
KVCache * kv,
ggml_backend_sched_t sched,
const float * input_embed,
int T,
const char * dump_dir,
TalkerForwardOutput * out);
static bool talker_forward_prefill(const TalkerWeights * tw,
KVCache * kv,
ggml_backend_sched_t sched,
const float * input_embed,
int T,
const char * dump_dir,
TalkerForwardOutput * out) {
kv_cache_reset(kv);
if (T > kv->max_seq_len) {
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, input_embed, T, 0, 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.
bool talker_forward_decode(const TalkerWeights * tw,
KVCache * kv,
ggml_backend_sched_t sched,
const float * input_embed_1,
TalkerForwardOutput * out);
static bool talker_forward_decode(const TalkerWeights * tw,
KVCache * kv,
ggml_backend_sched_t sched,
const float * input_embed_1,
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, input_embed_1, 1, kv->cur_len, NULL, out);
}