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

tts: static decode graphs and device resident hidden bridge

Browse Source 
KV writes go through set_rows with the destination rows carried as data
and the code predictor steps get one arena per sub step, so every
decode graph keeps a fixed topology and fixed tensor addresses step
after step. The talker last hidden stays resident on device in a
persistent bridge tensor, written by the talker graph and concatenated
as a leaf by the predictor prefill. The hot loop uploads sixteen code
ids and one overlay row and reads back the logits alone, host hidden
readbacks survive only under the dump path. Fewer nodes, transfers and
syncs pay on every backend, and the stable topology lets the CUDA
backend replay its captured graphs without an update.
This commit is contained in:
Pascal
2026-07-05 07:43:09 +02:00
parent e2e381f4f0
commit 415ef56330
4 changed files with 209 additions and 107 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/code-predictor-forward.h
+52 -52
View File
@@ -4,8 +4,11 @@
// KV cached. // KV cached.
// //
// Input: // Input:
// talker_hidden_last [hidden] f32 -- last position hidden state from // hidden_bridge [hidden] f32 -- persistent backend tensor holding
// the Talker forward (post final norm) // the talker last position hidden
// (post final norm), written on
// device by the talker graph and
// read here as a graph leaf
// c0 -- semantic code sampled from the // c0 -- semantic code sampled from the
// Talker codec_head (codebook 0) // Talker codec_head (codebook 0)
// Output: // Output:
@@ -86,9 +89,9 @@ static struct ggml_tensor * code_predictor_layer_forward(struct ggml_context *
struct ggml_tensor * x, struct ggml_tensor * x,
struct ggml_tensor * positions, struct ggml_tensor * positions,
struct ggml_tensor * mask, struct ggml_tensor * mask,
struct ggml_tensor * kv_rows,
struct ggml_tensor * k_cache, struct ggml_tensor * k_cache,
struct ggml_tensor * v_cache, struct ggml_tensor * v_cache,
int n_past,
int T, int T,
int n_kv_pad, int n_kv_pad,
bool use_flash_attn, bool use_flash_attn,
@@ -120,20 +123,14 @@ static struct ggml_tensor * code_predictor_layer_forward(struct ggml_context *
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, 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); 0.0f);
// Write the fresh positions into the cache. // Write the fresh positions into the cache via set_rows: positions
// travel as data so every step keeps an identical topology and the
// captured CUDA graph replays without an update.
struct ggml_tensor * k_perm = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3)); // [hd, T, n_kv] 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)); 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]; ggml_build_forward_expand(gf, ggml_set_rows(ctx, k_cache, k_perm, kv_rows));
size_t v_off = (size_t) n_past * v_cache->nb[1]; ggml_build_forward_expand(gf, ggml_set_rows(ctx, v_cache, v_perm, kv_rows));
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);
struct ggml_tensor * k_full = ggml_view_3d(ctx, k_cache, hd, n_kv_pad, n_kv, k_cache->nb[1], k_cache->nb[2], 0); struct ggml_tensor * k_full = ggml_view_3d(ctx, k_cache, hd, n_kv_pad, n_kv, k_cache->nb[1], k_cache->nb[2], 0);
struct ggml_tensor * v_full = ggml_view_3d(ctx, v_cache, hd, n_kv_pad, n_kv, v_cache->nb[1], v_cache->nb[2], 0); struct ggml_tensor * v_full = ggml_view_3d(ctx, v_cache, hd, n_kv_pad, n_kv, v_cache->nb[1], v_cache->nb[2], 0);
@@ -194,11 +191,10 @@ static bool code_predictor_run(const CodePredictorWeights * cw,
ggml_backend_sched_t sched, ggml_backend_sched_t sched,
GraphArena * arena, GraphArena * arena,
struct ggml_tensor * embd_table, struct ggml_tensor * embd_table,
const float * hidden_row, struct ggml_tensor * hidden_bridge,
int32_t code_id, int32_t code_id,
int T, int T,
int n_past, int n_past,
int talker_hidden,
int g_head, int g_head,
bool use_flash_attn, bool use_flash_attn,
bool clamp_fp16, bool clamp_fp16,
@@ -216,25 +212,25 @@ static bool code_predictor_run(const CodePredictorWeights * cw,
struct ggml_context * gctx = graph_arena_begin(arena); struct ggml_context * gctx = graph_arena_begin(arena);
// Inputs: one code id gathered in graph from embd_table, positions, // Inputs: one code id gathered in graph from embd_table, positions,
// attention mask, plus the raw talker hidden row on the prefill // attention mask. The prefill path (T == 2, hidden_bridge non NULL)
// path (T == 2, hidden_row non NULL) where the sequence is // concats the resident talker hidden ahead of embed(c0), both on
// [talker_hidden, embed(c0)]. Steps (T == 1) are pure gathers: the // device: the sequence is [talker_hidden, embed(c0)] with zero row
// only per step upload is 4 bytes of code id. // upload. Steps (T == 1) are pure gathers: the only per step upload
// is 4 bytes of code id.
struct ggml_tensor * ids_in = ggml_new_tensor_1d(gctx, GGML_TYPE_I32, 1); struct ggml_tensor * ids_in = ggml_new_tensor_1d(gctx, GGML_TYPE_I32, 1);
struct ggml_tensor * pos_in = ggml_new_tensor_1d(gctx, GGML_TYPE_I32, T); struct ggml_tensor * pos_in = ggml_new_tensor_1d(gctx, GGML_TYPE_I32, T);
struct ggml_tensor * mask_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F16, n_kv_pad, T); struct ggml_tensor * mask_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F16, n_kv_pad, T);
struct ggml_tensor * rows_in = ggml_new_tensor_1d(gctx, GGML_TYPE_I64, T);
ggml_set_name(ids_in, "sub_code_id"); ggml_set_name(ids_in, "sub_code_id");
ggml_set_name(pos_in, "positions"); ggml_set_name(pos_in, "positions");
ggml_set_name(mask_in, "causal_mask"); ggml_set_name(mask_in, "causal_mask");
ggml_set_name(rows_in, "kv_rows");
ggml_set_input(ids_in); ggml_set_input(ids_in);
ggml_set_input(rows_in);
struct ggml_tensor * x_in = ggml_get_rows(gctx, embd_table, ids_in); struct ggml_tensor * x_in = ggml_get_rows(gctx, embd_table, ids_in);
struct ggml_tensor * hid_in = NULL;
if (T == 2) { if (T == 2) {
hid_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F32, talker_hidden, 1); x_in = ggml_concat(gctx, hidden_bridge, x_in, 1);
ggml_set_name(hid_in, "talker_hidden_row");
ggml_set_input(hid_in);
x_in = ggml_concat(gctx, hid_in, x_in, 1);
} }
ggml_set_name(x_in, "sub_input"); ggml_set_name(x_in, "sub_input");
@@ -252,8 +248,9 @@ static bool code_predictor_run(const CodePredictorWeights * cw,
} }
for (int l = 0; l < n_layers; l++) { 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], h = code_predictor_layer_forward(gctx, cw, cw->layers[(size_t) l], h, pos_in, mask_in, rows_in,
kv->v[(size_t) l], n_past, T, n_kv_pad, use_flash_attn, clamp_fp16, gf); kv->k[(size_t) l], kv->v[(size_t) l], T, n_kv_pad, use_flash_attn, clamp_fp16,
gf);
} }
struct ggml_tensor * h_final = ggml_rms_norm(gctx, h, cw->rms_norm_eps); struct ggml_tensor * h_final = ggml_rms_norm(gctx, h, cw->rms_norm_eps);
@@ -272,9 +269,6 @@ static bool code_predictor_run(const CodePredictorWeights * cw,
} }
ggml_backend_tensor_set(ids_in, &code_id, 0, sizeof(int32_t)); ggml_backend_tensor_set(ids_in, &code_id, 0, sizeof(int32_t));
if (hid_in) {
ggml_backend_tensor_set(hid_in, hidden_row, 0, (size_t) talker_hidden * sizeof(float));
}
{ {
std::vector<int32_t> pos((size_t) T); std::vector<int32_t> pos((size_t) T);
@@ -282,6 +276,12 @@ static bool code_predictor_run(const CodePredictorWeights * cw,
pos[(size_t) i] = n_past + i; pos[(size_t) i] = n_past + i;
} }
ggml_backend_tensor_set(pos_in, pos.data(), 0, (size_t) T * sizeof(int32_t)); ggml_backend_tensor_set(pos_in, pos.data(), 0, (size_t) T * sizeof(int32_t));
std::vector<int64_t> rows((size_t) T);
for (int i = 0; i < T; i++) {
rows[(size_t) i] = (int64_t) (n_past + i);
}
ggml_backend_tensor_set(rows_in, rows.data(), 0, (size_t) T * sizeof(int64_t));
} }
{ {
@@ -314,20 +314,21 @@ static bool code_predictor_run(const CodePredictorWeights * cw,
return true; return true;
} }
// Run the predictor for one audio frame. Caller passes the talker hidden // Run the predictor for one audio frame. Caller passes the persistent
// state for the current frame and the already-sampled c0. Sampling // hidden bridge holding the talker hidden for the current frame and the
// parameters control greedy (temperature <= 0) vs stochastic. subseq_base // already-sampled c0. Sampling parameters control greedy
// is the Philox subsequence of the c0 sample for this step; the 15 // (temperature <= 0) vs stochastic. subseq_base is the Philox
// acoustic samples consume subseq_base + 1 .. subseq_base + 15. // subsequence of the c0 sample for this step; the 15 acoustic samples
// Returns the full vector of 16 codes. dump_dir may be NULL. // consume subseq_base + 1 .. subseq_base + 15. Returns the full vector
// use_flash_attn / clamp_fp16 are forwarded as is to every internal run. // of 16 codes. dump_dir may be NULL. use_flash_attn / clamp_fp16 are
// forwarded as is to every internal run.
static bool code_predictor_step(const TalkerWeights * tw, static bool code_predictor_step(const TalkerWeights * tw,
const CodePredictorWeights * cw, const CodePredictorWeights * cw,
KVCache * kv, KVCache * kv,
ggml_backend_sched_t sched, ggml_backend_sched_t sched,
GraphArena * arena_prefill, GraphArena * arena_prefill,
GraphArena * arena_step, GraphArena * arenas_step,
const float * talker_hidden_last, struct ggml_tensor * hidden_bridge,
int c0, int c0,
float temperature, float temperature,
int top_k, int top_k,
@@ -339,10 +340,9 @@ static bool code_predictor_step(const TalkerWeights * tw,
const char * dump_dir, const char * dump_dir,
CodePredictorOutput * out) { CodePredictorOutput * out) {
// sub_input slots live at the talker hidden dimension: both the // sub_input slots live at the talker hidden dimension: both the
// talker last hidden and the codec_embedding rows feeding the sub // bridge row and the codec_embedding rows feeding the sub network
// network are talker sized in the upstream checkpoint. The graph's // are talker sized in the upstream checkpoint. The graph's mtp_proj
// mtp_proj brings them down to predictor hidden when present. // brings them down to predictor hidden when present.
const int talker_hidden = tw->hidden_size;
const int n_acoustic = cw->num_acoustic_codebooks; const int n_acoustic = cw->num_acoustic_codebooks;
if (n_acoustic + 1 > kv->max_seq_len) { if (n_acoustic + 1 > kv->max_seq_len) {
@@ -354,14 +354,14 @@ static bool code_predictor_step(const TalkerWeights * tw,
out->codes.assign((size_t) (n_acoustic + 1), 0); out->codes.assign((size_t) (n_acoustic + 1), 0);
out->codes[0] = c0; out->codes[0] = c0;
// Prefill: two positions, [talker_hidden_last, embed_talker(c0)]. // Prefill: two positions, [talker_hidden, embed_talker(c0)], both
// The hidden row uploads raw, c0 gathers in graph from the talker // resident on device: the hidden reads from the bridge, c0 gathers
// codec embedding table. // in graph from the talker codec embedding table.
kv_cache_reset(kv); kv_cache_reset(kv);
std::vector<float> logits; std::vector<float> logits;
if (!code_predictor_run(cw, kv, sched, arena_prefill, tw->codec_embedding, talker_hidden_last, c0, 2, 0, if (!code_predictor_run(cw, kv, sched, arena_prefill, tw->codec_embedding, hidden_bridge, c0, 2, 0, 0,
talker_hidden, 0, use_flash_attn, clamp_fp16, &logits)) { use_flash_attn, clamp_fp16, &logits)) {
return false; return false;
} }
{ {
@@ -381,11 +381,11 @@ static bool code_predictor_step(const TalkerWeights * tw,
// Decode loop: 14 single-token steps. At step g (g=1..14) we feed // Decode loop: 14 single-token steps. At step g (g=1..14) we feed
// the id of the code we just sampled, gathered in graph from the // the id of the code we just sampled, gathered in graph from the
// group's private embedding table, and read lm_head[g]. // group's private embedding table, and read lm_head[g]. Each step
// owns its arena so lm_head and the table stay fixed per graph.
for (int g = 1; g < n_acoustic; g++) { for (int g = 1; g < n_acoustic; g++) {
if (!code_predictor_run(cw, kv, sched, arena_step, cw->codec_embedding[(size_t) (g - 1)], NULL, if (!code_predictor_run(cw, kv, sched, &arenas_step[(size_t) (g - 1)], cw->codec_embedding[(size_t) (g - 1)],
out->codes[(size_t) g], 1, kv->cur_len, talker_hidden, g, use_flash_attn, clamp_fp16, NULL, out->codes[(size_t) g], 1, kv->cur_len, g, use_flash_attn, clamp_fp16, &logits)) {
&logits)) {
return false; return false;
} }
float u_g = 0.0f; float u_g = 0.0f;
src/pipeline-tts.cpp
+75 -12
View File
@@ -115,6 +115,9 @@ bool pipeline_tts_load(PipelineTTS * pt,
pt->backend = bp.backend; pt->backend = bp.backend;
pt->sched = NULL; pt->sched = NULL;
pt->has_speaker_encoder = false; pt->has_speaker_encoder = false;
pt->bridge_ctx = NULL;
pt->bridge_buf = NULL;
pt->hidden_bridge = NULL;
// Fused flash attention needs a GPU kernel; CPU only backends fall // Fused flash attention needs a GPU kernel; CPU only backends fall
// back to the F32 manual chain automatically. clamp_fp16 is forwarded // back to the F32 manual chain automatically. clamp_fp16 is forwarded
@@ -224,15 +227,63 @@ bool pipeline_tts_load(PipelineTTS * pt,
return false; return false;
} }
// Hidden bridge: one [talker_hidden] f32 tensor resident on the
// backend, written by the talker graph and read by the code
// predictor prefill graph. Cleared once so debug dumps never see
// stale bytes before the first talker forward.
{
struct ggml_init_params gp = { ggml_tensor_overhead() * 2, NULL, true };
pt->bridge_ctx = ggml_init(gp);
pt->hidden_bridge =
pt->bridge_ctx ? ggml_new_tensor_1d(pt->bridge_ctx, GGML_TYPE_F32, pt->talker.hidden_size) : NULL;
if (pt->hidden_bridge) {
ggml_set_name(pt->hidden_bridge, "talker_hidden_bridge");
pt->bridge_buf = ggml_backend_alloc_ctx_tensors(pt->bridge_ctx, pt->backend);
}
if (!pt->hidden_bridge || !pt->bridge_buf) {
qt_log(QT_LOG_ERROR, "[Pipeline] hidden bridge allocation failed");
if (pt->bridge_ctx) {
ggml_free(pt->bridge_ctx);
pt->bridge_ctx = NULL;
}
pt->hidden_bridge = NULL;
kv_cache_free(&pt->code_predictor_kv);
kv_cache_free(&pt->talker_kv);
ggml_backend_sched_free(pt->sched);
pt->sched = NULL;
pipeline_codec_free(&pt->codec);
code_predictor_weights_free(&pt->code_predictor);
talker_weights_free(&pt->talker);
gf_close(&pt->gguf_talker);
return false;
}
ggml_backend_buffer_clear(pt->bridge_buf, 0);
}
// Persistent graph arenas: one shape class each so the backend CUDA // Persistent graph arenas: one shape class each so the backend CUDA
// graph cache keeps a stable executable per flavor across steps. // graph cache keeps a stable executable per flavor across steps.
if (!graph_arena_init(&pt->talker_arena, talker_graph_max_nodes(pt->talker.num_hidden_layers)) || // The predictor gets one arena per sub step g: each of the 14 step
!graph_arena_init(&pt->cp_prefill_arena, // graphs then keeps a fixed lm_head and embedding table and replays
code_predictor_graph_max_nodes(pt->code_predictor.num_hidden_layers)) || // its captured executable without an update.
!graph_arena_init(&pt->cp_step_arena, code_predictor_graph_max_nodes(pt->code_predictor.num_hidden_layers))) { bool arenas_ok =
graph_arena_init(&pt->talker_arena, talker_graph_max_nodes(pt->talker.num_hidden_layers)) &&
graph_arena_init(&pt->cp_prefill_arena, code_predictor_graph_max_nodes(pt->code_predictor.num_hidden_layers));
pt->cp_step_arenas.resize((size_t) (pt->num_code_groups - 2));
for (size_t g = 0; arenas_ok && g < pt->cp_step_arenas.size(); g++) {
arenas_ok = graph_arena_init(&pt->cp_step_arenas[g],
code_predictor_graph_max_nodes(pt->code_predictor.num_hidden_layers));
}
if (!arenas_ok) {
for (size_t g = 0; g < pt->cp_step_arenas.size(); g++) {
graph_arena_free(&pt->cp_step_arenas[g]);
}
graph_arena_free(&pt->talker_arena); graph_arena_free(&pt->talker_arena);
graph_arena_free(&pt->cp_prefill_arena); graph_arena_free(&pt->cp_prefill_arena);
graph_arena_free(&pt->cp_step_arena); ggml_backend_buffer_free(pt->bridge_buf);
pt->bridge_buf = NULL;
ggml_free(pt->bridge_ctx);
pt->bridge_ctx = NULL;
pt->hidden_bridge = NULL;
kv_cache_free(&pt->code_predictor_kv); kv_cache_free(&pt->code_predictor_kv);
kv_cache_free(&pt->talker_kv); kv_cache_free(&pt->talker_kv);
ggml_backend_sched_free(pt->sched); ggml_backend_sched_free(pt->sched);
@@ -254,9 +305,21 @@ bool pipeline_tts_load(PipelineTTS * pt,
} }
void pipeline_tts_free(PipelineTTS * pt) { void pipeline_tts_free(PipelineTTS * pt) {
graph_arena_free(&pt->cp_step_arena); for (size_t g = 0; g < pt->cp_step_arenas.size(); g++) {
graph_arena_free(&pt->cp_step_arenas[g]);
}
pt->cp_step_arenas.clear();
graph_arena_free(&pt->cp_prefill_arena); graph_arena_free(&pt->cp_prefill_arena);
graph_arena_free(&pt->talker_arena); graph_arena_free(&pt->talker_arena);
if (pt->bridge_buf) {
ggml_backend_buffer_free(pt->bridge_buf);
pt->bridge_buf = NULL;
}
if (pt->bridge_ctx) {
ggml_free(pt->bridge_ctx);
pt->bridge_ctx = NULL;
}
pt->hidden_bridge = NULL;
kv_cache_free(&pt->code_predictor_kv); kv_cache_free(&pt->code_predictor_kv);
kv_cache_free(&pt->talker_kv); kv_cache_free(&pt->talker_kv);
if (pt->sched) { if (pt->sched) {
@@ -604,13 +667,13 @@ qt_status pipeline_tts_synthesize(PipelineTTS * pt,
bool ok; bool ok;
Timer t_talker; Timer t_talker;
if (step == 0) { if (step == 0) {
ok = talker_forward_prefill(&pt->talker, &pt->talker_kv, pt->sched, &pt->talker_arena, ok = talker_forward_prefill(&pt->talker, &pt->talker_kv, pt->sched, &pt->talker_arena, pt->hidden_bridge,
prompt.input_embed.data(), prompt.T_ctx, use_fa, clamp_fp16, step_dump, &fw); prompt.input_embed.data(), prompt.T_ctx, use_fa, clamp_fp16, step_dump, &fw);
} else { } else {
ok = ok = talker_forward_decode(&pt->talker, &pt->talker_kv, pt->sched, &pt->talker_arena, pt->hidden_bridge,
talker_forward_decode(&pt->talker, &pt->talker_kv, pt->sched, &pt->talker_arena, prev_ids.data(), prev_ids.data(), pt->code_predictor.codec_embedding.data(),
pt->code_predictor.codec_embedding.data(), pt->code_predictor.num_acoustic_codebooks, prev_overlay, use_fa, clamp_fp16,
pt->code_predictor.num_acoustic_codebooks, prev_overlay, use_fa, clamp_fp16, &fw); params->dump_dir != NULL, &fw);
} }
if (!ok) { if (!ok) {
return QT_STATUS_GENERATE_FAILED; return QT_STATUS_GENERATE_FAILED;
@@ -664,7 +727,7 @@ qt_status pipeline_tts_synthesize(PipelineTTS * pt,
const char * cp_dump = (params->dump_dir && step == 0) ? params->dump_dir : NULL; const char * cp_dump = (params->dump_dir && step == 0) ? params->dump_dir : NULL;
Timer t_pred; Timer t_pred;
if (!code_predictor_step(&pt->talker, &pt->code_predictor, &pt->code_predictor_kv, pt->sched, if (!code_predictor_step(&pt->talker, &pt->code_predictor, &pt->code_predictor_kv, pt->sched,
&pt->cp_prefill_arena, &pt->cp_step_arena, fw.hidden_last.data(), c0, subtk_T, &pt->cp_prefill_arena, pt->cp_step_arenas.data(), pt->hidden_bridge, c0, subtk_T,
params->subtalker_top_k, params->subtalker_top_p, resolved_seed, subseq_counter - 1, params->subtalker_top_k, params->subtalker_top_p, resolved_seed, subseq_counter - 1,
use_fa, clamp_fp16, cp_dump, &cp)) { use_fa, clamp_fp16, cp_dump, &cp)) {
return QT_STATUS_GENERATE_FAILED; return QT_STATUS_GENERATE_FAILED;
src/pipeline-tts.h
+12 -3
View File
@@ -132,14 +132,23 @@ struct PipelineTTS {
KVCache talker_kv; KVCache talker_kv;
KVCache code_predictor_kv; KVCache code_predictor_kv;
// Hidden bridge: the talker last position hidden stays resident on
// device. The talker graph copies it in, the code predictor prefill
// graph reads it as a leaf, so the AR hot loop never round trips
// the row through the host. [talker_hidden] f32 on `backend`.
struct ggml_context * bridge_ctx;
ggml_backend_buffer_t bridge_buf;
struct ggml_tensor * hidden_bridge;
// Persistent graph arenas, one per graph shape class. Stable node // Persistent graph arenas, one per graph shape class. Stable node
// addresses across rebuilds keep the backend CUDA graph cache hot: // addresses across rebuilds keep the backend CUDA graph cache hot:
// the talker shares one arena for prefill and decode, the predictor // the talker shares one arena for prefill and decode, the predictor
// splits prefill (T=2) and step (T=1) so the two flavors that // splits prefill (T=2) from the steps, with one arena per sub step
// alternate within a frame each keep their own executable. // g so each graph keeps a fixed lm_head and embedding table and the
// captured executable replays without an update.
GraphArena talker_arena; GraphArena talker_arena;
GraphArena cp_prefill_arena; GraphArena cp_prefill_arena;
GraphArena cp_step_arena; std::vector<GraphArena> cp_step_arenas;
}; };
// Open the talker GGUF and the codec GGUF, load every module on the // Open the talker GGUF and the codec GGUF, load every module on the
src/talker-forward.h
+65 -35
View File
@@ -52,7 +52,9 @@
#include <vector> #include <vector>
struct TalkerForwardOutput { struct TalkerForwardOutput {
// Final hidden state for the last position [hidden] f32 (post final norm). // Final hidden state for the last position [hidden] f32 (post final
// norm). Filled only when read_hidden_host is set (dump paths); the
// hot loop consumes the on device hidden bridge instead.
std::vector<float> hidden_last; std::vector<float> hidden_last;
// Codec head logits for the last position [vocab] f32. // Codec head logits for the last position [vocab] f32.
@@ -105,10 +107,11 @@ static struct ggml_tensor * talker_attn_f32(struct ggml_context * ctx,
} }
// Build the per-layer block, KV cached. K and V for the T fresh // Build the per-layer block, KV cached. K and V for the T fresh
// positions get computed normally, then written into the cache at // positions get computed normally, then written into the cache at the
// [n_past, n_past+T) on dim 1. The attention reads the contiguous slice // rows carried by kv_rows via set_rows. The attention reads the fixed
// [0, n_past+T) on the same dim, which covers the full causal context // [0, n_kv_pad) window on the same dim, with the mask killing
// in one tensor view. Returns the layer output [hidden, T]. // everything past the causal context. Returns the layer output
// [hidden, T].
// //
// use_flash_attn picks between the fused ggml_flash_attn_ext kernel // use_flash_attn picks between the fused ggml_flash_attn_ext kernel
// (GPU only, FP16 accumulation guarded with set_prec(F32)) and the // (GPU only, FP16 accumulation guarded with set_prec(F32)) and the
@@ -122,9 +125,9 @@ static struct ggml_tensor * talker_layer_forward(struct ggml_context * ctx,
struct ggml_tensor * x, struct ggml_tensor * x,
struct ggml_tensor * positions, struct ggml_tensor * positions,
struct ggml_tensor * mask, struct ggml_tensor * mask,
struct ggml_tensor * kv_rows,
struct ggml_tensor * k_cache, struct ggml_tensor * k_cache,
struct ggml_tensor * v_cache, struct ggml_tensor * v_cache,
int n_past,
int T, int T,
int n_kv_pad, int n_kv_pad,
bool use_flash_attn, bool use_flash_attn,
@@ -164,23 +167,18 @@ static struct ggml_tensor * talker_layer_forward(struct ggml_context * ctx,
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, 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); 0.0f);
// Write the T fresh positions of K and V into the cache. K and V are // Write the T fresh positions of K and V into the cache via
// [hd, n_kv, T] at this point and the cache lives as [hd, max_T, n_kv] // set_rows: kv_rows carries the destination positions as data, so
// so we permute to [hd, T, n_kv] before the cpy. The destination view // the graph topology stays identical across decode steps and the
// covers [hd, T, n_kv] at dim 1 offset n_past * nb1. // captured CUDA graph replays without an update. 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]; the row ids
// broadcast across the n_kv head dim.
struct ggml_tensor * k_perm = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3)); // [hd, T, n_kv] 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] 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]; ggml_build_forward_expand(gf, ggml_set_rows(ctx, k_cache, k_perm, kv_rows));
size_t v_off = (size_t) n_past * v_cache->nb[1]; ggml_build_forward_expand(gf, ggml_set_rows(ctx, v_cache, v_perm, kv_rows));
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_kv_pad) window for attention. n_kv_pad covers the // Read the [0, n_kv_pad) window for attention. n_kv_pad covers the
// causal context and rounds it up so the view shape stays constant // causal context and rounds it up so the view shape stays constant
@@ -253,18 +251,22 @@ static struct ggml_tensor * talker_layer_forward(struct ggml_context * ctx,
} }
// Shared core that builds the graph, allocates, uploads inputs, runs // Shared core that builds the graph, allocates, uploads inputs, runs
// it and pulls out the last position hidden + logits. T tokens are // it and pulls out the last position logits. T tokens are appended to
// appended to the cache starting at n_past. Input is either a raw // the cache starting at n_past. Input is either a raw embedding upload
// embedding upload (input_embed, prefill) or the previous frame code // (input_embed, prefill) or the previous frame code ids plus overlay
// ids plus overlay row assembled in graph (frame_ids, decode hot // row assembled in graph (frame_ids, decode hot path). The last
// path). When n_past == 0 and dump_dir is set, the bisect taps fire on // position hidden copies in graph into the caller owned persistent
// the prefill path. use_fa / clamp_fp16 are forwarded as is to every // hidden_bridge tensor the code predictor prefill reads on device; the
// host copy in out->hidden_last fills only under read_hidden_host.
// When n_past == 0 and dump_dir is set, the bisect taps fire on the
// prefill path. use_fa / clamp_fp16 are forwarded as is to every
// layer. The graph metadata lives in the caller owned persistent // layer. The graph metadata lives in the caller owned persistent
// arena. // arena.
static bool talker_forward_core(const TalkerWeights * tw, static bool talker_forward_core(const TalkerWeights * tw,
KVCache * kv, KVCache * kv,
ggml_backend_sched_t sched, ggml_backend_sched_t sched,
GraphArena * arena, GraphArena * arena,
struct ggml_tensor * hidden_bridge,
const float * input_embed, const float * input_embed,
const int32_t * frame_ids, const int32_t * frame_ids,
struct ggml_tensor * const * acoustic_embd, struct ggml_tensor * const * acoustic_embd,
@@ -274,6 +276,7 @@ static bool talker_forward_core(const TalkerWeights * tw,
int n_past, int n_past,
bool use_flash_attn, bool use_flash_attn,
bool clamp_fp16, bool clamp_fp16,
bool read_hidden_host,
const char * dump_dir, const char * dump_dir,
TalkerForwardOutput * out) { TalkerForwardOutput * out) {
const int hidden = tw->hidden_size; const int hidden = tw->hidden_size;
@@ -301,6 +304,12 @@ static bool talker_forward_core(const TalkerWeights * tw,
ggml_set_name(pos_in, "positions"); ggml_set_name(pos_in, "positions");
ggml_set_name(mask_in, "causal_mask"); ggml_set_name(mask_in, "causal_mask");
// KV write positions as data: identical topology at every step,
// pure CUDA graph replay across the decode loop.
struct ggml_tensor * rows_in = ggml_new_tensor_1d(gctx, GGML_TYPE_I64, T);
ggml_set_name(rows_in, "kv_rows");
ggml_set_input(rows_in);
// Input: either a raw embedding upload (prefill path) or, on the // Input: either a raw embedding upload (prefill path) or, on the
// decode hot path, the frame codes of the previous step gathered // decode hot path, the frame codes of the previous step gathered
// and summed in graph. x = get_rows(codec_embd, ids[0]) plus the 15 // and summed in graph. x = get_rows(codec_embd, ids[0]) plus the 15
@@ -337,8 +346,8 @@ static bool talker_forward_core(const TalkerWeights * tw,
struct ggml_tensor * h = x_in; struct ggml_tensor * h = x_in;
std::vector<struct ggml_tensor *> taps(TALKER_N_BISECT_LAYERS, NULL); std::vector<struct ggml_tensor *> taps(TALKER_N_BISECT_LAYERS, NULL);
for (int l = 0; l < n_layers; l++) { 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], h = talker_layer_forward(gctx, tw, tw->layers[(size_t) l], h, pos_in, mask_in, rows_in, kv->k[(size_t) l],
kv->v[(size_t) l], n_past, T, n_kv_pad, use_flash_attn, clamp_fp16, gf); kv->v[(size_t) l], T, n_kv_pad, use_flash_attn, clamp_fp16, gf);
if (record_taps && talker_is_bisect_layer(l)) { if (record_taps && talker_is_bisect_layer(l)) {
for (int i = 0; i < TALKER_N_BISECT_LAYERS; i++) { for (int i = 0; i < TALKER_N_BISECT_LAYERS; i++) {
if (TALKER_BISECT_LAYERS[i] == l) { if (TALKER_BISECT_LAYERS[i] == l) {
@@ -356,7 +365,16 @@ static bool talker_forward_core(const TalkerWeights * tw,
struct ggml_tensor * h_final = ggml_rms_norm(gctx, h, tw->rms_norm_eps); struct ggml_tensor * h_final = ggml_rms_norm(gctx, h, tw->rms_norm_eps);
h_final = ggml_mul(gctx, h_final, tw->norm_w); h_final = ggml_mul(gctx, h_final, tw->norm_w);
ggml_set_name(h_final, "hidden_final"); ggml_set_name(h_final, "hidden_final");
if (record_taps || read_hidden_host) {
ggml_set_output(h_final); ggml_set_output(h_final);
}
// Bridge: the last position hidden copies on device into the
// persistent tensor. Constant destination address across steps, so
// the decode graph topology stays replayable.
struct ggml_tensor * h_last =
ggml_view_1d(gctx, h_final, hidden, (size_t) (T - 1) * (size_t) hidden * sizeof(float));
struct ggml_tensor * bridge_cpy = ggml_cpy(gctx, h_last, hidden_bridge);
// codec_head: [hidden, vocab]. ggml_mul_mat returns [vocab, T]. // codec_head: [hidden, vocab]. ggml_mul_mat returns [vocab, T].
struct ggml_tensor * logits = ggml_mul_mat(gctx, tw->codec_head_w, h_final); struct ggml_tensor * logits = ggml_mul_mat(gctx, tw->codec_head_w, h_final);
@@ -371,6 +389,7 @@ static bool talker_forward_core(const TalkerWeights * tw,
ggml_build_forward_expand(gf, h_final); ggml_build_forward_expand(gf, h_final);
} }
ggml_build_forward_expand(gf, logits); ggml_build_forward_expand(gf, logits);
ggml_build_forward_expand(gf, bridge_cpy);
ggml_backend_sched_reset(sched); ggml_backend_sched_reset(sched);
if (!ggml_backend_sched_alloc_graph(sched, gf)) { if (!ggml_backend_sched_alloc_graph(sched, gf)) {
@@ -394,6 +413,12 @@ static bool talker_forward_core(const TalkerWeights * tw,
pos[(size_t) i] = n_past + i; pos[(size_t) i] = n_past + i;
} }
ggml_backend_tensor_set(pos_in, pos.data(), 0, (size_t) T * sizeof(int32_t)); ggml_backend_tensor_set(pos_in, pos.data(), 0, (size_t) T * sizeof(int32_t));
std::vector<int64_t> rows((size_t) T);
for (int i = 0; i < T; i++) {
rows[(size_t) i] = (int64_t) (n_past + i);
}
ggml_backend_tensor_set(rows_in, rows.data(), 0, (size_t) T * sizeof(int64_t));
} }
// Causal mask: 0 where k <= n_past + q, neg inf otherwise. Stored // Causal mask: 0 where k <= n_past + q, neg inf otherwise. Stored
@@ -440,15 +465,17 @@ static bool talker_forward_core(const TalkerWeights * tw,
debug_dump_2d(&d, "talker-hidden-prefill-final", buf.data(), T, hidden); debug_dump_2d(&d, "talker-hidden-prefill-final", buf.data(), T, hidden);
} }
// Pull the last position: final hidden + logits // Pull the last position logits, the only per step readback on the
// hot path. The hidden row comes back to host only for dumps.
out->hidden = hidden; out->hidden = hidden;
out->vocab = vocab; out->vocab = vocab;
out->hidden_last.assign((size_t) hidden, 0.0f);
out->logits_last.assign((size_t) vocab, 0.0f); out->logits_last.assign((size_t) vocab, 0.0f);
{ {
size_t row_bytes = (size_t) vocab * sizeof(float); 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); ggml_backend_tensor_get(logits, out->logits_last.data(), (size_t) (T - 1) * row_bytes, row_bytes);
}
if (read_hidden_host) {
out->hidden_last.assign((size_t) hidden, 0.0f);
size_t hrow_bytes = (size_t) hidden * sizeof(float); 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); ggml_backend_tensor_get(h_final, out->hidden_last.data(), (size_t) (T - 1) * hrow_bytes, hrow_bytes);
} }
@@ -472,6 +499,7 @@ static bool talker_forward_prefill(const TalkerWeights * tw,
KVCache * kv, KVCache * kv,
ggml_backend_sched_t sched, ggml_backend_sched_t sched,
GraphArena * arena, GraphArena * arena,
struct ggml_tensor * hidden_bridge,
const float * input_embed, const float * input_embed,
int T, int T,
bool use_flash_attn, bool use_flash_attn,
@@ -483,8 +511,8 @@ static bool talker_forward_prefill(const TalkerWeights * tw,
fprintf(stderr, "[TalkerForward] FATAL: prefill T=%d exceeds cache max_seq_len=%d\n", T, kv->max_seq_len); fprintf(stderr, "[TalkerForward] FATAL: prefill T=%d exceeds cache max_seq_len=%d\n", T, kv->max_seq_len);
return false; return false;
} }
return talker_forward_core(tw, kv, sched, arena, input_embed, NULL, NULL, 0, NULL, T, 0, use_flash_attn, clamp_fp16, return talker_forward_core(tw, kv, sched, arena, hidden_bridge, input_embed, NULL, NULL, 0, NULL, T, 0,
dump_dir, out); use_flash_attn, clamp_fp16, dump_dir != NULL, dump_dir, out);
} }
// Decode: feed exactly one embedding and append one position to the // Decode: feed exactly one embedding and append one position to the
@@ -498,18 +526,20 @@ static bool talker_forward_decode(const TalkerWeights * tw,
KVCache * kv, KVCache * kv,
ggml_backend_sched_t sched, ggml_backend_sched_t sched,
GraphArena * arena, GraphArena * arena,
struct ggml_tensor * hidden_bridge,
const int32_t * frame_ids, const int32_t * frame_ids,
struct ggml_tensor * const * acoustic_embd, struct ggml_tensor * const * acoustic_embd,
int n_acoustic, int n_acoustic,
const float * overlay, const float * overlay,
bool use_flash_attn, bool use_flash_attn,
bool clamp_fp16, bool clamp_fp16,
bool read_hidden_host,
TalkerForwardOutput * out) { TalkerForwardOutput * out) {
if (kv->cur_len + 1 > kv->max_seq_len) { if (kv->cur_len + 1 > kv->max_seq_len) {
fprintf(stderr, "[TalkerForward] FATAL: decode would overflow cache (%d + 1 > %d)\n", kv->cur_len, fprintf(stderr, "[TalkerForward] FATAL: decode would overflow cache (%d + 1 > %d)\n", kv->cur_len,
kv->max_seq_len); kv->max_seq_len);
return false; return false;
} }
return talker_forward_core(tw, kv, sched, arena, NULL, frame_ids, acoustic_embd, n_acoustic, overlay, 1, return talker_forward_core(tw, kv, sched, arena, hidden_bridge, NULL, frame_ids, acoustic_embd, n_acoustic, overlay,
kv->cur_len, use_flash_attn, clamp_fp16, NULL, out); 1, kv->cur_len, use_flash_attn, clamp_fp16, read_hidden_host, NULL, out);
} }