diff --git a/src/code-predictor-forward.h b/src/code-predictor-forward.h index 70871ec..ec2ad03 100644 --- a/src/code-predictor-forward.h +++ b/src/code-predictor-forward.h @@ -1,26 +1,29 @@ #pragma once // code-predictor-forward.h: run the 5-layer Qwen3 code predictor over a -// growing context to produce the 15 acoustic codes of one audio frame, -// KV cached. +// growing context to produce the 15 acoustic codes of one audio frame +// for every active slot in one batched pass, KV cached per slot. // // Input: -// hidden_bridge [hidden] f32 -- persistent backend tensor holding -// 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 -// Talker codec_head (codebook 0) +// hidden_bridge [hidden, max_batch] f32 -- persistent backend tensor +// holding the talker last +// position hidden per slot (post +// final norm), written on device +// by the talker graph and read +// here as a graph leaf +// c0[N] -- semantic codes sampled from the +// Talker codec_head (codebook 0), +// one per slot // Output: -// codes[16] = [c0, c1, ..., c15] -- the full set of codes for one -// frame, ready for decode through -// the codec +// codes[N * 16] -- per slot [c0, c1, ..., c15], +// slot major, ready for decode +// through the codec // -// The predictor cache is local to a single frame: we reset it at every -// frame, prefill the first two positions (talker_hidden + embed(c0)), -// 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. +// The predictor cache is local to a single frame and holds one set per +// slot: the prefill writes the first two positions (talker_hidden + +// embed(c0)) of every slot, then 14 batched single-token steps follow. +// Every slot runs the same sub-step sequence every frame, so the batch +// stays in perfect lockstep and the graphs bake positions, kv rows and +// masks at build time. // // Architecture mirrors the Talker block, only differences are: // - 5 layers instead of 28 @@ -28,10 +31,10 @@ // - one private embedding table and one private linear head per // acoustic codebook (1..15) // -// Graph metadata lives in caller owned static graphs, one for the T=2 -// prefill and one per T=1 step: each graph is built and allocated once -// at load, then replayed directly on the backend with a 4 byte code id -// upload per call, the positions, kv rows, and causal mask baked in. +// Graph metadata lives in caller owned static graphs, one per flavor +// and batch width N (prefill plus one per acoustic step), built lazily +// on the first frame at a given N, then replayed directly on the +// backend with an N * 4 byte code id upload per call. #include "code-predictor-graph.h" #include "code-predictor-weights.h" @@ -51,13 +54,15 @@ #include struct CodePredictorOutput { - // Sixteen codes: c0 from the talker plus c1..c15 from the predictor. + // Per slot sixteen codes, slot major: codes[slot * 16 + g] holds c0 + // from the talker plus c1..c15 from the predictor. std::vector codes; }; // Manual F32 attention chain for the code predictor block. Same shape -// contract as talker_attn_f32: q [hd, T, n_q_heads], k/v [hd, T_full, -// n_kv], output [hd, n_q_heads, T]. Used when use_flash_attn is false. +// contract as talker_attn_f32: q [hd, T, n_q_heads, N], k/v +// [hd, T_full, n_kv, N], output [hd, n_q_heads, T, N]; the mul_mat +// broadcasts over dims 2 and 3. Used when use_flash_attn is false. static struct ggml_tensor * code_predictor_attn_f32(struct ggml_context * ctx, struct ggml_tensor * q, struct ggml_tensor * k, @@ -76,10 +81,12 @@ static int code_predictor_graph_max_nodes(int n_layers) { return 48 * n_layers + 64; } -// 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 fixed [0, n_kv_pad) window with the mask carrying -// neg inf beyond n_past+T. Returns the layer output [hidden, T]. +// One batched Qwen3 decoder block, KV cached over sets [0, N). x holds +// the N slots' token columns flattened slot major: column j = n * T + t +// is position t of slot n. K and V for the fresh positions are written +// into each slot's set at the rows carried by kv_rows; the attention +// reads the fixed [0, n_kv_pad) window per set with the mask carrying +// neg inf beyond n_past + T. Returns the layer output [hidden, T * N]. // use_flash_attn and clamp_fp16 follow the same contract as in // talker-forward.h. static struct ggml_tensor * code_predictor_layer_forward(struct ggml_context * ctx, @@ -89,9 +96,10 @@ static struct ggml_tensor * code_predictor_layer_forward(struct ggml_context * struct ggml_tensor * positions, struct ggml_tensor * mask, struct ggml_tensor * kv_rows, - struct ggml_tensor * k_cache, - struct ggml_tensor * v_cache, + struct ggml_tensor * k4, + struct ggml_tensor * v4, int T, + int N, int n_kv_pad, bool use_flash_attn, bool clamp_fp16, @@ -100,6 +108,7 @@ static struct ggml_tensor * code_predictor_layer_forward(struct ggml_context * const int n_kv = cw->num_key_value_heads; const int hd = cw->head_dim; const float eps = cw->rms_norm_eps; + const int TN = T * N; struct ggml_tensor * h = ggml_rms_norm(ctx, x, eps); h = ggml_mul(ctx, h, layer.input_norm_w); @@ -108,34 +117,43 @@ static struct ggml_tensor * code_predictor_layer_forward(struct ggml_context * 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_reshape_3d(ctx, q, hd, n_q_heads, TN); + k = ggml_reshape_3d(ctx, k, hd, n_kv, TN); + v = ggml_reshape_3d(ctx, v, hd, n_kv, TN); 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 over the flattened token axis: positions [T * N] repeat the + // same in-frame offsets for every slot. 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 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 * v_perm = ggml_cont(ctx, ggml_permute(ctx, v, 0, 2, 1, 3)); + // Write the fresh positions into each slot's set via set_rows: + // [hd, heads, T*N] unflattens to [hd, heads, T, N], permutes to + // [hd, T, heads, N] and lands at the kv_rows [T, 1, N] destinations, + // broadcast across the n_kv head dim. + struct ggml_tensor * k_perm = + ggml_cont(ctx, ggml_permute(ctx, ggml_reshape_4d(ctx, k, hd, n_kv, T, N), 0, 2, 1, 3)); + struct ggml_tensor * v_perm = + ggml_cont(ctx, ggml_permute(ctx, ggml_reshape_4d(ctx, v, hd, n_kv, T, N), 0, 2, 1, 3)); - ggml_build_forward_expand(gf, ggml_set_rows(ctx, k_cache, k_perm, kv_rows)); - ggml_build_forward_expand(gf, ggml_set_rows(ctx, v_cache, v_perm, kv_rows)); + struct ggml_tensor * k_sets = ggml_view_4d(ctx, k4, hd, k4->ne[1], n_kv, N, k4->nb[1], k4->nb[2], k4->nb[3], 0); + struct ggml_tensor * v_sets = ggml_view_4d(ctx, v4, hd, v4->ne[1], n_kv, N, v4->nb[1], v4->nb[2], v4->nb[3], 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); + ggml_build_forward_expand(gf, ggml_set_rows(ctx, k_sets, k_perm, kv_rows)); + ggml_build_forward_expand(gf, ggml_set_rows(ctx, v_sets, v_perm, kv_rows)); - // 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); + struct ggml_tensor * k_full = ggml_view_4d(ctx, k4, hd, n_kv_pad, n_kv, N, k4->nb[1], k4->nb[2], k4->nb[3], 0); + struct ggml_tensor * v_full = ggml_view_4d(ctx, v4, hd, n_kv_pad, n_kv, N, v4->nb[1], v4->nb[2], v4->nb[3], 0); + + // Q [hd, n_q_heads, T, N] -> [hd, T, n_q_heads, N] for + // flash_attn_ext, taken as a view like the talker path does. + struct ggml_tensor * q_p = ggml_permute(ctx, ggml_reshape_4d(ctx, q, hd, n_q_heads, T, N), 0, 2, 1, 3); // Clamp V before attention when clamp_fp16 is set, same rationale // as the talker block: sub Ampere CUDA tensor cores accumulate in @@ -144,9 +162,7 @@ static struct ggml_tensor * code_predictor_layer_forward(struct ggml_context * v_full = ggml_clamp(ctx, v_full, -65504.0f, 65504.0f); } - // Attention: fused flash kernel or manual F32 chain. Matches the - // working acestep qw3lm_build_attn pattern on the fused branch and - // the omnivoice qwen3_attn_f32 helper on the manual one. + // Attention: fused flash kernel or manual F32 chain. float scale = 1.0f / sqrtf((float) hd); struct ggml_tensor * attn; if (use_flash_attn) { @@ -156,7 +172,9 @@ static struct ggml_tensor * code_predictor_layer_forward(struct ggml_context * attn = code_predictor_attn_f32(ctx, q_p, k_full, v_full, mask, scale); } - attn = ggml_reshape_2d(ctx, attn, n_q_heads * hd, T); + // [hd, n_q_heads, T, N] -> [n_q_heads*hd, T*N], flatten heads for + // o_proj, token order matching x. + attn = ggml_reshape_2d(ctx, attn, n_q_heads * hd, TN); struct ggml_tensor * o = ggml_mul_mat(ctx, layer.attn.o_proj_w, attn); x = ggml_add(ctx, x, o); @@ -180,18 +198,20 @@ static struct ggml_tensor * code_predictor_layer_forward(struct ggml_context * return x; } -// Build one static predictor graph. A non NULL hidden_bridge selects -// the T=2 prefill flavor reading [talker_hidden, embed(c0)] through -// lm_head[0]; otherwise the graph is the single token step for g_head, -// appending at the fixed cache row g_head + 1. The logits node holds -// the last position only, so every flavor reads back one row at -// offset zero. use_flash_attn / clamp_fp16 apply to every layer. +// Build one static batched predictor graph over sets [0, N). A non +// NULL hidden_bridge selects the T=2 prefill flavor reading +// [talker_hidden, embed(c0)] per slot through lm_head[0]; otherwise +// the graph is the single token step for g_head, appending at the +// fixed cache row g_head + 1. The logits node holds the last position +// of every slot as [Vg, N]. use_flash_attn / clamp_fp16 apply to every +// layer. static bool code_predictor_graph_build(const CodePredictorWeights * cw, KVCache * kv, ggml_backend_t backend, struct ggml_tensor * embd_table, struct ggml_tensor * hidden_bridge, int g_head, + int N, bool use_flash_attn, bool clamp_fp16, CodePredGraph * cp) { @@ -214,17 +234,18 @@ static bool code_predictor_graph_build(const CodePredictorWeights * cw, } struct ggml_context * gctx = cp->ctx; - // Inputs: one code id gathered in graph from embd_table, positions, - // attention mask. The prefill path (T == 2, hidden_bridge non NULL) - // concats the resident talker hidden ahead of embed(c0), both on - // device: the sequence is [talker_hidden, embed(c0)] with zero row - // 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 * 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 * rows_in = ggml_new_tensor_1d(gctx, GGML_TYPE_I64, T); - ggml_set_name(ids_in, "sub_code_id"); + // Inputs: one code id per slot gathered in graph from embd_table, + // positions, kv rows and the attention mask. The prefill path + // (T == 2) concats each slot's resident talker hidden ahead of + // embed(c0), both on device: the per slot sequence is + // [talker_hidden, embed(c0)] with zero row upload. Steps (T == 1) + // are pure gathers: the only per step upload is N * 4 bytes of + // code ids. + struct ggml_tensor * ids_in = ggml_new_tensor_1d(gctx, GGML_TYPE_I32, N); + struct ggml_tensor * pos_in = ggml_new_tensor_1d(gctx, GGML_TYPE_I32, T * N); + struct ggml_tensor * mask_in = ggml_new_tensor_4d(gctx, GGML_TYPE_F16, n_kv_pad, T, 1, N); + struct ggml_tensor * rows_in = ggml_new_tensor_3d(gctx, GGML_TYPE_I64, T, 1, N); + ggml_set_name(ids_in, "sub_code_ids"); ggml_set_name(pos_in, "positions"); ggml_set_name(mask_in, "causal_mask"); ggml_set_name(rows_in, "kv_rows"); @@ -240,9 +261,14 @@ static bool code_predictor_graph_build(const CodePredictorWeights * cw, ggml_set_input(rows_in); ggml_set_output(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); // [hidden, N] if (T == 2) { - x_in = ggml_concat(gctx, hidden_bridge, x_in, 1); + struct ggml_tensor * bridge_cols = + ggml_view_2d(gctx, hidden_bridge, hidden_bridge->ne[0], N, hidden_bridge->nb[1], 0); + struct ggml_tensor * b3 = ggml_reshape_3d(gctx, bridge_cols, hidden_bridge->ne[0], 1, N); + struct ggml_tensor * e3 = ggml_reshape_3d(gctx, x_in, x_in->ne[0], 1, N); + x_in = ggml_concat(gctx, b3, e3, 1); // [in_dim, 2, N] + x_in = ggml_reshape_2d(gctx, x_in, x_in->ne[0], T * N); } ggml_set_name(x_in, "sub_input"); @@ -261,16 +287,17 @@ static bool code_predictor_graph_build(const CodePredictorWeights * cw, 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, rows_in, - kv->k[(size_t) l], kv->v[(size_t) l], T, n_kv_pad, use_flash_attn, clamp_fp16, - gf); + kv->k4[(size_t) l], kv->v4[(size_t) l], T, N, n_kv_pad, use_flash_attn, + clamp_fp16, 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); if (T > 1) { - // Last position only: the prefill pays a single lm_head row. - h_final = ggml_cont( - gctx, ggml_view_2d(gctx, h_final, h_final->ne[0], 1, h_final->nb[1], (size_t) (T - 1) * h_final->nb[1])); + // Last position of every slot: columns (T - 1) + n * T, one + // strided view so the prefill pays N lm_head rows. + h_final = ggml_cont(gctx, ggml_view_2d(gctx, h_final, h_final->ne[0], N, (size_t) T * h_final->nb[1], + (size_t) (T - 1) * h_final->nb[1])); } struct ggml_tensor * logits = ggml_mul_mat(gctx, cw->lm_head[(size_t) g_head], h_final); @@ -286,30 +313,36 @@ static bool code_predictor_graph_build(const CodePredictorWeights * cw, } { - std::vector pos((size_t) T); - for (int i = 0; i < T; i++) { - pos[(size_t) i] = n_past + i; + std::vector pos((size_t) T * (size_t) N); + for (int n = 0; n < N; n++) { + for (int t = 0; t < T; t++) { + pos[(size_t) n * (size_t) T + (size_t) t] = n_past + t; + } } - ggml_backend_tensor_set(pos_in, pos.data(), 0, (size_t) T * sizeof(int32_t)); + ggml_backend_tensor_set(pos_in, pos.data(), 0, pos.size() * sizeof(int32_t)); - std::vector rows((size_t) T); - for (int i = 0; i < T; i++) { - rows[(size_t) i] = (int64_t) (n_past + i); + std::vector rows((size_t) T * (size_t) N); + for (int n = 0; n < N; n++) { + for (int t = 0; t < T; t++) { + rows[(size_t) n * (size_t) T + (size_t) t] = (int64_t) (n_past + t); + } } - ggml_backend_tensor_set(rows_in, rows.data(), 0, (size_t) T * sizeof(int64_t)); + ggml_backend_tensor_set(rows_in, rows.data(), 0, rows.size() * sizeof(int64_t)); } { - std::vector mask((size_t) T * (size_t) n_kv_pad); + std::vector mask((size_t) n_kv_pad * (size_t) T * (size_t) N); 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) n_kv_pad + (size_t) k] = zero; + for (int n = 0; n < N; n++) { + 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) n * (size_t) T + (size_t) q) * (size_t) n_kv_pad + (size_t) k] = zero; + } } } ggml_backend_tensor_set(mask_in, mask.data(), 0, mask.size() * sizeof(ggml_fp16_t)); @@ -318,86 +351,105 @@ static bool code_predictor_graph_build(const CodePredictorWeights * cw, cp->gf = gf; cp->ids_in = ids_in; cp->logits = logits; + cp->N = N; return true; } -// Replay one static predictor graph: upload the code id, run the graph -// directly on the backend, read the single logits row back. +// Replay one static predictor graph: upload the N code ids, run the +// graph directly on the backend, read the [Vg, N] logits back. static bool code_predictor_replay(CodePredGraph * cp, ggml_backend_t backend, - int32_t code_id, + const int32_t * code_ids, + int N, std::vector * logits_out) { - ggml_backend_tensor_set(cp->ids_in, &code_id, 0, sizeof(int32_t)); + ggml_backend_tensor_set(cp->ids_in, code_ids, 0, (size_t) N * sizeof(int32_t)); if (ggml_backend_graph_compute(backend, cp->gf) != GGML_STATUS_SUCCESS) { fprintf(stderr, "[CodePredictor] FATAL: graph compute failed\n"); return false; } - logits_out->resize((size_t) cp->logits->ne[0]); - ggml_backend_tensor_get(cp->logits, logits_out->data(), 0, (size_t) cp->logits->ne[0] * sizeof(float)); + logits_out->resize((size_t) cp->logits->ne[0] * (size_t) N); + ggml_backend_tensor_get(cp->logits, logits_out->data(), 0, logits_out->size() * sizeof(float)); return true; } -// Run the predictor for one audio frame over the static graphs. The -// prefill replay consumes the persistent hidden bridge already written -// by the talker graph and the sampled c0; the 14 step replays each -// feed the code sampled just before. 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. +// Run the predictor for one audio frame over the static graphs, all N +// slots in lockstep. The prefill replay consumes the persistent hidden +// bridge columns already written by the talker graph and the sampled +// c0 of every slot; the 14 step replays each feed the codes sampled +// just before. Per slot sampling parameters control greedy +// (temperature <= 0) vs stochastic; seed and subseq_base index each +// slot's own Philox stream, so per slot outputs match a single +// sequence run bit for bit. subseq_base[i] is the subsequence of slot +// i's c0 sample for this step; its 15 acoustic samples consume +// subseq_base[i] + 1 .. subseq_base[i] + 15. Fills out->codes as +// [N * 16] slot major. dump_dir may be NULL and applies to slot 0. static bool code_predictor_step(const CodePredictorWeights * cw, ggml_backend_t backend, CodePredGraph * prefill_graph, CodePredGraph * step_graphs, - int c0, - float temperature, - int top_k, - float top_p, - int64_t seed, - int64_t subseq_base, + const int32_t * c0, + int N, + const float * temperature, + const int * top_k, + const float * top_p, + const int64_t * seed, + const int64_t * subseq_base, const char * dump_dir, CodePredictorOutput * out) { const int n_acoustic = cw->num_acoustic_codebooks; + const int n_codes = n_acoustic + 1; - out->codes.assign((size_t) (n_acoustic + 1), 0); - out->codes[0] = c0; + out->codes.assign((size_t) N * (size_t) n_codes, 0); + for (int i = 0; i < N; i++) { + out->codes[(size_t) i * (size_t) n_codes] = c0[i]; + } - std::vector logits; - if (!code_predictor_replay(prefill_graph, backend, c0, &logits)) { + std::vector logits; + std::vector ids((size_t) N); + + if (!code_predictor_replay(prefill_graph, backend, c0, N, &logits)) { return false; } - { + const int V0 = (int) (logits.size() / (size_t) N); + for (int i = 0; i < N; i++) { 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); + int cg = sample_top_k_p(logits.data() + (size_t) i * (size_t) V0, V0, temperature[i], top_k[i], top_p[i], 1.0f, + nullptr, 0, seed[i], subseq_base[i] + 1, &u_g); if (cg < 0) { - fprintf(stderr, "[CodePredictor] FATAL: sample returned no candidate at g=0\n"); + fprintf(stderr, "[CodePredictor] FATAL: sample returned no candidate at g=0 (slot %d)\n", i); return false; } - out->codes[1] = cg; + out->codes[(size_t) i * (size_t) n_codes + 1] = cg; } - // Decode loop: 14 single-token replays. At step g (g=1..14) we feed - // the id of the code we just sampled, gathered in graph from the - // group's private embedding table, and read lm_head[g]. + // Decode loop: 14 batched single-token replays. At step g + // (g=1..14) every slot feeds the id of the code it just sampled, + // gathered in graph from the group's private embedding table, and + // reads lm_head[g]. for (int g = 1; g < n_acoustic; g++) { - if (!code_predictor_replay(&step_graphs[(size_t) (g - 1)], backend, out->codes[(size_t) g], &logits)) { + for (int i = 0; i < N; i++) { + ids[(size_t) i] = out->codes[(size_t) i * (size_t) n_codes + (size_t) g]; + } + if (!code_predictor_replay(&step_graphs[(size_t) (g - 1)], backend, ids.data(), N, &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 (cg < 0) { - fprintf(stderr, "[CodePredictor] FATAL: sample returned no candidate at g=%d\n", g); - return false; + const int Vg = (int) (logits.size() / (size_t) N); + for (int i = 0; i < N; i++) { + float u_g = 0.0f; + int cg = sample_top_k_p(logits.data() + (size_t) i * (size_t) Vg, Vg, temperature[i], top_k[i], top_p[i], + 1.0f, nullptr, 0, seed[i], subseq_base[i] + 1 + g, &u_g); + if (cg < 0) { + fprintf(stderr, "[CodePredictor] FATAL: sample returned no candidate at g=%d (slot %d)\n", g, i); + return false; + } + out->codes[(size_t) i * (size_t) n_codes + (size_t) (g + 1)] = cg; } - out->codes[(size_t) (g + 1)] = cg; } if (dump_dir) { DebugDumper d; debug_init(&d, dump_dir); - std::vector codes32(out->codes.begin(), out->codes.end()); + std::vector codes32(out->codes.begin(), out->codes.begin() + n_codes); int n = (int) codes32.size(); debug_dump_i32_as_f32(&d, "codes-step0", codes32.data(), &n, 1); } diff --git a/src/code-predictor-graph.h b/src/code-predictor-graph.h index d1fc1b9..c082814 100644 --- a/src/code-predictor-graph.h +++ b/src/code-predictor-graph.h @@ -1,9 +1,10 @@ #pragma once -// code-predictor-graph.h: one static predictor graph, built once at -// load, allocated once, and replayed with a 4 byte code id upload per -// call. Positions, kv rows, and the causal mask bake in at build time -// because the frame cache layout repeats identically: the prefill -// always writes rows 0..1 and step g always appends at row g + 1. +// code-predictor-graph.h: one static predictor graph per flavor and +// batch width, built lazily, allocated once, and replayed with an N +// code id upload per call. Positions, kv rows, and the causal mask +// bake in at build time because the frame cache layout repeats +// identically for every slot: the prefill always writes rows 0..1 and +// step g always appends at row g + 1. #include "ggml-alloc.h" #include "ggml.h" @@ -12,8 +13,9 @@ struct CodePredGraph { struct ggml_context * ctx = nullptr; struct ggml_cgraph * gf = nullptr; ggml_gallocr_t galloc = nullptr; - struct ggml_tensor * ids_in = nullptr; - struct ggml_tensor * logits = nullptr; + struct ggml_tensor * ids_in = nullptr; // [N] i32, one code id per slot + struct ggml_tensor * logits = nullptr; // [Vg, N] f32 + int N = 0; // batch width this build covers }; static void code_predictor_graph_free(CodePredGraph * cp) { @@ -28,4 +30,5 @@ static void code_predictor_graph_free(CodePredGraph * cp) { cp->gf = nullptr; cp->ids_in = nullptr; cp->logits = nullptr; + cp->N = 0; } diff --git a/src/kv-cache.h b/src/kv-cache.h index ae2b569..f454326 100644 --- a/src/kv-cache.h +++ b/src/kv-cache.h @@ -1,16 +1,19 @@ #pragma once // kv-cache.h: persistent per-layer KV cache for the Talker LM and the -// Code Predictor. Mirrors the standard llama.cpp ring approach but kept +// Code Predictor, batched over n_sets independent sequences. Kept // minimal: the cache is sized at init for a fixed max sequence length -// and never reallocates. Reset just rewinds cur_len to 0. +// and set count and never reallocates. Reset just rewinds cur_len. // // Layout per layer, both K and V : -// ggml_tensor [hd, max_seq_len, n_kv] f32, contiguous on hd -// This matches the layout the attention block already uses for K, so -// the write path is a ggml_cpy of the freshly RoPE'd K into a view of -// the cache, and the read path is just a view spanning [0, cur_len+T) -// on dim 1. V uses the same layout and gets permuted to [T, n_kv, hd] -// at read time for the value matmul, identical to the prefill path. +// 4D ggml_tensor [hd, max_seq_len, n_kv, n_sets] f32, contiguous on hd +// plus one 3D view [hd, max_seq_len, n_kv] per set for the per-sequence +// prefill path and device side set copies. The batched decode graph +// views the 4D tensor directly with ne3 = N over sets [0, N). +// +// The 3D layout matches what the attention block uses for K, so the +// write path is a set_rows of the freshly RoPE'd K into the set view, +// and the read path is a view spanning the padded causal window on +// dim 1. V uses the same layout. #include "ggml-alloc.h" #include "ggml-backend.h" @@ -26,9 +29,16 @@ struct KVCache { int n_kv_heads; int head_dim; int max_seq_len; - int cur_len; + int n_sets; - // One pair per layer, both tensors live in `buffer` allocated below. + // Write head per set. + std::vector cur_len; + + // Per layer 4D tensors, both in `buffer` allocated below. + std::vector k4; + std::vector v4; + + // Per set 3D views into the 4D tensors, indexed [set * n_layers + layer]. std::vector k; std::vector v; @@ -36,25 +46,38 @@ struct KVCache { ggml_backend_buffer_t buffer; }; +static struct ggml_tensor * kv_cache_k(const KVCache * kv, int set, int layer) { + return kv->k[(size_t) set * (size_t) kv->n_layers + (size_t) layer]; +} + +static struct ggml_tensor * kv_cache_v(const KVCache * kv, int set, int layer) { + return kv->v[(size_t) set * (size_t) kv->n_layers + (size_t) layer]; +} + // Allocate a fresh KV cache backed by a dedicated buffer on `backend`. -// Tensors are zero initialised (the attention path never reads past -// cur_len so this is mostly cosmetic, but it keeps debug dumps clean). +// Tensors are zero initialised: the padded attention window reads past +// cur_len and the masked tail must see finite values. static bool kv_cache_init(KVCache * kv, int n_layers, int n_kv_heads, int head_dim, int max_seq_len, + int n_sets, ggml_backend_t backend) { kv->n_layers = n_layers; kv->n_kv_heads = n_kv_heads; kv->head_dim = head_dim; kv->max_seq_len = max_seq_len; - kv->cur_len = 0; - kv->k.assign((size_t) n_layers, NULL); - kv->v.assign((size_t) n_layers, NULL); + kv->n_sets = n_sets; + kv->cur_len.assign((size_t) n_sets, 0); + kv->k4.assign((size_t) n_layers, NULL); + kv->v4.assign((size_t) n_layers, NULL); + kv->k.assign((size_t) n_sets * (size_t) n_layers, NULL); + kv->v.assign((size_t) n_sets * (size_t) n_layers, NULL); - struct ggml_init_params gp = { - ggml_tensor_overhead() * (size_t) (2 * n_layers + 4), + const size_t n_tensors = (size_t) (2 * n_layers) * (size_t) (1 + n_sets) + 4; + struct ggml_init_params gp = { + ggml_tensor_overhead() * n_tensors, NULL, true, }; @@ -65,13 +88,26 @@ static bool kv_cache_init(KVCache * kv, } for (int l = 0; l < n_layers; l++) { - kv->k[(size_t) l] = ggml_new_tensor_3d(kv->ctx, GGML_TYPE_F32, head_dim, max_seq_len, n_kv_heads); - kv->v[(size_t) l] = ggml_new_tensor_3d(kv->ctx, GGML_TYPE_F32, head_dim, max_seq_len, n_kv_heads); + kv->k4[(size_t) l] = ggml_new_tensor_4d(kv->ctx, GGML_TYPE_F32, head_dim, max_seq_len, n_kv_heads, n_sets); + kv->v4[(size_t) l] = ggml_new_tensor_4d(kv->ctx, GGML_TYPE_F32, head_dim, max_seq_len, n_kv_heads, n_sets); char name[64]; snprintf(name, sizeof(name), "kv_k_l%d", l); - ggml_set_name(kv->k[(size_t) l], name); + ggml_set_name(kv->k4[(size_t) l], name); snprintf(name, sizeof(name), "kv_v_l%d", l); - ggml_set_name(kv->v[(size_t) l], name); + ggml_set_name(kv->v4[(size_t) l], name); + + // Per set 3D views, created before the buffer allocation so + // ggml_backend_alloc_ctx_tensors runs its view init on them + // (buffer and data resolve against the owning 4D tensor). + struct ggml_tensor * k4 = kv->k4[(size_t) l]; + struct ggml_tensor * v4 = kv->v4[(size_t) l]; + for (int s = 0; s < n_sets; s++) { + size_t off = (size_t) s * k4->nb[3]; + kv->k[(size_t) s * (size_t) n_layers + (size_t) l] = + ggml_view_3d(kv->ctx, k4, head_dim, max_seq_len, n_kv_heads, k4->nb[1], k4->nb[2], off); + kv->v[(size_t) s * (size_t) n_layers + (size_t) l] = + ggml_view_3d(kv->ctx, v4, head_dim, max_seq_len, n_kv_heads, v4->nb[1], v4->nb[2], off); + } } kv->buffer = ggml_backend_alloc_ctx_tensors(kv->ctx, backend); @@ -85,16 +121,29 @@ static bool kv_cache_init(KVCache * kv, // Zero-init the buffer so any out of bounds read returns a known value. ggml_backend_buffer_clear(kv->buffer, 0); - size_t bytes_per_layer = (size_t) head_dim * (size_t) max_seq_len * (size_t) n_kv_heads * sizeof(float); - size_t total_mb = (size_t) (2 * n_layers) * bytes_per_layer / (1024 * 1024); - fprintf(stderr, "[KVCache] Allocated: %d layers, %d KV heads, head_dim %d, max_seq_len %d -> %zu MB\n", n_layers, - n_kv_heads, head_dim, max_seq_len, total_mb); + size_t bytes_per_layer = + (size_t) head_dim * (size_t) max_seq_len * (size_t) n_kv_heads * (size_t) n_sets * sizeof(float); + size_t total_mb = (size_t) (2 * n_layers) * bytes_per_layer / (1024 * 1024); + fprintf(stderr, "[KVCache] Allocated: %d layers, %d KV heads, head_dim %d, max_seq_len %d, %d sets -> %zu MB\n", + n_layers, n_kv_heads, head_dim, max_seq_len, n_sets, total_mb); return true; } -// Rewind the cache so the next forward starts a fresh sequence. -static void kv_cache_reset(KVCache * kv) { - kv->cur_len = 0; +// Rewind one set so its next forward starts a fresh sequence. +static void kv_cache_reset(KVCache * kv, int set) { + kv->cur_len[(size_t) set] = 0; +} + +// Device side copy of one whole set into another through the +// persistent 3D views. Used by the batch engine to compact the active +// slot range after a retirement: the tail set moves into the freed one +// so the batched decode keeps viewing a consecutive [0, N) span. +static void kv_cache_copy_set(KVCache * kv, int src, int dst) { + for (int l = 0; l < kv->n_layers; l++) { + ggml_backend_tensor_copy(kv_cache_k(kv, src, l), kv_cache_k(kv, dst, l)); + ggml_backend_tensor_copy(kv_cache_v(kv, src, l), kv_cache_v(kv, dst, l)); + } + kv->cur_len[(size_t) dst] = kv->cur_len[(size_t) src]; } static void kv_cache_free(KVCache * kv) { @@ -106,6 +155,9 @@ static void kv_cache_free(KVCache * kv) { ggml_free(kv->ctx); kv->ctx = NULL; } + kv->k4.clear(); + kv->v4.clear(); kv->k.clear(); kv->v.clear(); + kv->cur_len.clear(); } diff --git a/src/pipeline-codec.cpp b/src/pipeline-codec.cpp index 4aff2c4..9e81841 100644 --- a/src/pipeline-codec.cpp +++ b/src/pipeline-codec.cpp @@ -259,7 +259,7 @@ static bool pipeline_codec_stream_ensure(PipelineCodec * pc) { } if (!kv_cache_init(&pc->stream_kv, pc->transformer.num_layers, pc->transformer.num_kv_heads, - pc->transformer.head_dim, CODEC_STREAM_RING, pc->backend)) { + pc->transformer.head_dim, CODEC_STREAM_RING, 1, pc->backend)) { ggml_backend_buffer_free(pc->stream_buf); pc->stream_buf = NULL; ggml_free(pc->stream_ctx); @@ -281,7 +281,7 @@ bool pipeline_codec_stream_reset(PipelineCodec * pc) { // the zeroed KV ring stays hidden behind the sliding window mask. ggml_backend_buffer_clear(pc->stream_buf, 0); ggml_backend_buffer_clear(pc->stream_kv.buffer, 0); - kv_cache_reset(&pc->stream_kv); + kv_cache_reset(&pc->stream_kv, 0); pc->stream_pos = 0; return true; } @@ -305,14 +305,17 @@ static bool codec_snap_ensure(PipelineCodec * pc, CodecStateSnap * s) { if (s->ctx) { return true; } + // Views alias memory the owning tensors already cover, so the + // walker skips them: the mirror holds one duplicate per real + // tensor and the positional pairing in codec_snap_copy holds. int n = 0; for (struct ggml_tensor * t = ggml_get_first_tensor(pc->stream_ctx); t; t = ggml_get_next_tensor(pc->stream_ctx, t)) { - n++; + n += t->view_src ? 0 : 1; } for (struct ggml_tensor * t = ggml_get_first_tensor(pc->stream_kv.ctx); t; t = ggml_get_next_tensor(pc->stream_kv.ctx, t)) { - n++; + n += t->view_src ? 0 : 1; } struct ggml_init_params gp = { ggml_tensor_overhead() * (size_t) n, NULL, true }; s->ctx = ggml_init(gp); @@ -322,11 +325,15 @@ static bool codec_snap_ensure(PipelineCodec * pc, CodecStateSnap * s) { } for (struct ggml_tensor * t = ggml_get_first_tensor(pc->stream_ctx); t; t = ggml_get_next_tensor(pc->stream_ctx, t)) { - ggml_dup_tensor(s->ctx, t); + if (!t->view_src) { + ggml_dup_tensor(s->ctx, t); + } } for (struct ggml_tensor * t = ggml_get_first_tensor(pc->stream_kv.ctx); t; t = ggml_get_next_tensor(pc->stream_kv.ctx, t)) { - ggml_dup_tensor(s->ctx, t); + if (!t->view_src) { + ggml_dup_tensor(s->ctx, t); + } } s->buf = ggml_backend_alloc_ctx_tensors(s->ctx, pc->backend); if (!s->buf) { @@ -343,12 +350,20 @@ static bool codec_snap_ensure(PipelineCodec * pc, CodecStateSnap * s) { static void codec_snap_copy(PipelineCodec * pc, CodecStateSnap * s, bool save) { struct ggml_tensor * m = ggml_get_first_tensor(s->ctx); for (struct ggml_tensor * t = ggml_get_first_tensor(pc->stream_ctx); t; - t = ggml_get_next_tensor(pc->stream_ctx, t), m = ggml_get_next_tensor(s->ctx, m)) { + t = ggml_get_next_tensor(pc->stream_ctx, t)) { + if (t->view_src) { + continue; + } ggml_backend_tensor_copy(save ? t : m, save ? m : t); + m = ggml_get_next_tensor(s->ctx, m); } for (struct ggml_tensor * t = ggml_get_first_tensor(pc->stream_kv.ctx); t; - t = ggml_get_next_tensor(pc->stream_kv.ctx, t), m = ggml_get_next_tensor(s->ctx, m)) { + t = ggml_get_next_tensor(pc->stream_kv.ctx, t)) { + if (t->view_src) { + continue; + } ggml_backend_tensor_copy(save ? t : m, save ? m : t); + m = ggml_get_next_tensor(s->ctx, m); } } @@ -389,6 +404,34 @@ bool pipeline_codec_stream_snapshot(PipelineCodec * pc, uint64_t key) { return true; } +bool pipeline_codec_stream_save(PipelineCodec * pc, CodecStateSnap * s) { + if (!pc->stream_ready || !codec_snap_ensure(pc, s)) { + return false; + } + codec_snap_copy(pc, s, true); + s->pos = pc->stream_pos; + return true; +} + +bool pipeline_codec_stream_load(PipelineCodec * pc, CodecStateSnap * s) { + if (!pc->stream_ready || !s->ctx) { + return false; + } + codec_snap_copy(pc, s, false); + pc->stream_pos = s->pos; + return true; +} + +void pipeline_codec_snap_free(CodecStateSnap * s) { + if (s->buf) { + ggml_backend_buffer_free(s->buf); + } + if (s->ctx) { + ggml_free(s->ctx); + } + *s = {}; +} + // Build the static stream graph of chunk width T = 1 << cls: the // same module chain as the T=1 frame graph, every stream state and // KV ring tensor shared across classes, inputs and intermediates diff --git a/src/pipeline-codec.h b/src/pipeline-codec.h index 8c97266..1eeb09a 100644 --- a/src/pipeline-codec.h +++ b/src/pipeline-codec.h @@ -198,6 +198,15 @@ bool pipeline_codec_stream_restore(PipelineCodec * pc, uint64_t key); // the least recently used slot when all are taken. bool pipeline_codec_stream_snapshot(PipelineCodec * pc, uint64_t key); +// Caller owned stream state mirror, enabling multiple interleaved +// streamed utterances over the single live state: save parks the live +// state (and position cursor) into s, load restores it. The mirror +// allocates lazily on the first save; key and stamp stay unused. Free +// with pipeline_codec_snap_free. +bool pipeline_codec_stream_save(PipelineCodec * pc, CodecStateSnap * s); +bool pipeline_codec_stream_load(PipelineCodec * pc, CodecStateSnap * s); +void pipeline_codec_snap_free(CodecStateSnap * s); + // Encode a 24 kHz mono waveform into RVQ codes. // audio : [n_samples] f32 mono 24 kHz. Must be a multiple of // TOKENIZER_HOP_LENGTH (1920); the caller is expected diff --git a/src/pipeline-tts.cpp b/src/pipeline-tts.cpp index 252dbcb..2df2a60 100644 --- a/src/pipeline-tts.cpp +++ b/src/pipeline-tts.cpp @@ -105,12 +105,41 @@ static void parse_generation_defaults(const GGUFModel & gf, GenerationDefaults & g.max_new_tokens = (int) gf_get_u32(gf, "generation.max_new_tokens"); } +// Ensure the static predictor graph set for batch width N exists: +// prefill (T=2 through lm_head[0]) plus one T=1 step per acoustic +// codebook after the first. Built lazily on the first frame at a given +// width, then replayed for the process lifetime. +static bool pipeline_tts_cp_graphs_ensure(PipelineTTS * pt, int N) { + if ((int) pt->cp_graphs.size() < N) { + pt->cp_graphs.resize((size_t) N); + } + CodePredGraphSet & s = pt->cp_graphs[(size_t) (N - 1)]; + if (s.prefill.ctx) { + return true; + } + if (!code_predictor_graph_build(&pt->code_predictor, &pt->code_predictor_kv, pt->backend, + pt->talker.codec_embedding, pt->hidden_bridge, 0, N, pt->use_flash_attn, + pt->clamp_fp16, &s.prefill)) { + return false; + } + s.steps.resize((size_t) (pt->num_code_groups - 2)); + for (size_t g = 0; g < s.steps.size(); g++) { + if (!code_predictor_graph_build(&pt->code_predictor, &pt->code_predictor_kv, pt->backend, + pt->code_predictor.codec_embedding[g], NULL, (int) g + 1, N, pt->use_flash_attn, + pt->clamp_fp16, &s.steps[g])) { + return false; + } + } + return true; +} + bool pipeline_tts_load(PipelineTTS * pt, const char * talker_gguf_path, const char * codec_gguf_path, BackendPair bp, bool use_fa, - bool clamp_fp16) { + bool clamp_fp16, + int max_batch) { pt->bp = bp; pt->backend = bp.backend; pt->sched = NULL; @@ -118,6 +147,7 @@ bool pipeline_tts_load(PipelineTTS * pt, pt->bridge_ctx = NULL; pt->bridge_buf = NULL; pt->hidden_bridge = NULL; + pt->max_batch = max_batch > 1 ? max_batch : 1; // Fused flash attention needs a GPU kernel; CPU only backends fall // back to the F32 manual chain automatically. clamp_fp16 is forwarded @@ -201,11 +231,12 @@ bool pipeline_tts_load(PipelineTTS * pt, return false; } - // KV caches: talker holds the LM context up to 4096 positions (the - // longest ICL prompt observed is ~250 + max_new_tokens ~ 1500, so - // 4096 has 60% headroom). Predictor holds one frame of 16 sub-steps. + // KV caches, one set per slot: the talker holds the LM context up + // to 4096 positions (the longest ICL prompt observed is ~250 + + // max_new_tokens ~ 1500, so 4096 has 60% headroom). Predictor holds + // one frame of 16 sub-steps per slot. if (!kv_cache_init(&pt->talker_kv, pt->talker.num_hidden_layers, pt->talker.num_key_value_heads, - pt->talker.head_dim, 4096, pt->backend)) { + pt->talker.head_dim, 4096, pt->max_batch, pt->backend)) { ggml_backend_sched_free(pt->sched); pt->sched = NULL; pipeline_codec_free(&pt->codec); @@ -216,7 +247,7 @@ bool pipeline_tts_load(PipelineTTS * pt, } if (!kv_cache_init(&pt->code_predictor_kv, pt->code_predictor.num_hidden_layers, pt->code_predictor.num_key_value_heads, pt->code_predictor.head_dim, pt->num_code_groups, - pt->backend)) { + pt->max_batch, pt->backend)) { kv_cache_free(&pt->talker_kv); ggml_backend_sched_free(pt->sched); pt->sched = NULL; @@ -227,15 +258,16 @@ bool pipeline_tts_load(PipelineTTS * pt, 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. + // Hidden bridge: one [talker_hidden, max_batch] f32 tensor resident + // on the backend, columns written by the talker graphs 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; + pt->bridge_ctx ? ggml_new_tensor_2d(pt->bridge_ctx, GGML_TYPE_F32, pt->talker.hidden_size, pt->max_batch) : + 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); @@ -260,27 +292,24 @@ bool pipeline_tts_load(PipelineTTS * pt, ggml_backend_buffer_clear(pt->bridge_buf, 0); } - // Talker graph arena plus the static predictor graphs: the talker - // keeps one arena per shape class for the CUDA graph cache, the - // predictor builds one static graph per flavor (prefill + one per - // acoustic step), each with its lm_head and embedding table fixed - // and the positions, kv rows, and mask baked in. - bool graphs_ok = - graph_arena_init(&pt->talker_arena, talker_graph_max_nodes(pt->talker.num_hidden_layers)) && - code_predictor_graph_build(&pt->code_predictor, &pt->code_predictor_kv, pt->backend, pt->talker.codec_embedding, - pt->hidden_bridge, 0, pt->use_flash_attn, pt->clamp_fp16, &pt->cp_prefill_graph); + // Talker graph arena plus the batch width 1 predictor graph set: + // the talker keeps one arena per shape class for the CUDA graph + // cache, the predictor builds one static graph per flavor (prefill + // + one per acoustic step) with its lm_head and embedding table + // fixed and the positions, kv rows, and mask baked in. Wider + // predictor sets and the batched talker decode graphs build lazily + // on first use. pt->talker_decode_graphs.resize(((size_t) pt->talker_kv.max_seq_len + 255) / 256); - pt->cp_step_graphs.resize((size_t) (pt->num_code_groups - 2)); - for (size_t g = 0; graphs_ok && g < pt->cp_step_graphs.size(); g++) { - graphs_ok = code_predictor_graph_build(&pt->code_predictor, &pt->code_predictor_kv, pt->backend, - pt->code_predictor.codec_embedding[g], NULL, (int) g + 1, - pt->use_flash_attn, pt->clamp_fp16, &pt->cp_step_graphs[g]); - } + bool graphs_ok = graph_arena_init(&pt->talker_arena, talker_graph_max_nodes(pt->talker.num_hidden_layers)) && + pipeline_tts_cp_graphs_ensure(pt, 1); if (!graphs_ok) { - for (size_t g = 0; g < pt->cp_step_graphs.size(); g++) { - code_predictor_graph_free(&pt->cp_step_graphs[g]); + for (size_t n = 0; n < pt->cp_graphs.size(); n++) { + code_predictor_graph_free(&pt->cp_graphs[n].prefill); + for (size_t g = 0; g < pt->cp_graphs[n].steps.size(); g++) { + code_predictor_graph_free(&pt->cp_graphs[n].steps[g]); + } } - code_predictor_graph_free(&pt->cp_prefill_graph); + pt->cp_graphs.clear(); pt->talker_decode_graphs.clear(); graph_arena_free(&pt->talker_arena); ggml_backend_buffer_free(pt->bridge_buf); @@ -301,19 +330,21 @@ bool pipeline_tts_load(PipelineTTS * pt, qt_log(QT_LOG_INFO, "[Pipeline] Loaded: arch=%s variant=%s tokenizer=%s codebooks=%d speaker_encoder=%s speakers=%zu fa=%s " - "clamp_fp16=%s", + "clamp_fp16=%s max_batch=%d", pt->model_size.c_str(), pt->model_type.c_str(), pt->tokenizer_type.c_str(), pt->num_code_groups, pt->has_speaker_encoder ? "deferred" : "absent", pt->speakers.size(), pt->use_flash_attn ? "on" : "off", - pt->clamp_fp16 ? "on" : "off"); + pt->clamp_fp16 ? "on" : "off", pt->max_batch); return true; } void pipeline_tts_free(PipelineTTS * pt) { - for (size_t g = 0; g < pt->cp_step_graphs.size(); g++) { - code_predictor_graph_free(&pt->cp_step_graphs[g]); + for (size_t n = 0; n < pt->cp_graphs.size(); n++) { + code_predictor_graph_free(&pt->cp_graphs[n].prefill); + for (size_t g = 0; g < pt->cp_graphs[n].steps.size(); g++) { + code_predictor_graph_free(&pt->cp_graphs[n].steps[g]); + } } - pt->cp_step_graphs.clear(); - code_predictor_graph_free(&pt->cp_prefill_graph); + pt->cp_graphs.clear(); for (size_t g = 0; g < pt->talker_decode_graphs.size(); g++) { talker_decode_graph_free(&pt->talker_decode_graphs[g]); } @@ -440,15 +471,136 @@ static void tts_log_perf(const TtsPerf & p) { p.n_frames, per_frame, audio_sec, rtf); } -qt_status pipeline_tts_synthesize(PipelineTTS * pt, - BPETokenizer * tok, - const struct qt_tts_params * params, - int64_t resolved_seed, - struct qt_audio * out) { +// --------------------------------------------------------------------------- +// Batch engine: up to pt->max_batch concurrent synthesis slots in +// lockstep. Slot i owns talker KV set i, predictor KV set i and bridge +// column i; the active range is always [0, N) so the batched decode +// and predictor graphs view consecutive sets. A retirement compacts +// the range by moving the tail slot (host state plus one device side +// talker KV set copy) into the freed index; the bridge and the +// predictor sets rewrite every frame so only the talker cache moves. +// Single threaded: every entry runs on the thread that owns the GPU. +// --------------------------------------------------------------------------- + +struct TtsSlot { + TtsJob * job; + int64_t serial; // stable identity across compaction swaps + PromptBuilderOutput prompt; - const std::string instruct = params->instruct ? params->instruct : ""; - const std::string speaker = params->speaker ? params->speaker : ""; - const std::string ref_text = params->ref_text ? params->ref_text : ""; + + // ICL reference codes kept for the codec: stream seeding at admit, + // buffered decode left context at completion. ref_codes_ptr aims at + // ref_codes_store or at the caller's latent buffer. + std::vector ref_codes_store; + const int32_t * ref_codes_ptr; + int ref_codes_T; + + // Resolved sampling temperatures (0 selects greedy). + float talker_T; + float subtk_T; + + // AR state, one to one with the single sequence loop. + int step; // frames emitted so far + int64_t subseq_counter; // Philox subsequence cursor + std::vector talker_history; // emitted c0, feeds repetition penalty + std::vector prev_ids; // previous frame codes [num_code_groups] + const float * prev_overlay; // trailing text row or tts_pad row + std::vector logits; // pending c0 logits [vocab] + int pending_c0; // c0 of the frame in flight + bool has_frame; // slot emits a frame this engine step + + // Streaming state: the per slot codec stream mirror parks in snap + // whenever another slot takes the live codec state. + bool streaming; + codec_stream_decoder stream; + CodecStateSnap snap; + std::vector> all_codes; + + bool finished; + qt_status fin_status; + TtsPerf perf; + Timer t_total; +}; + +struct TtsEngine { + PipelineTTS * pt; + BPETokenizer * tok; + std::vector slots; + int64_t next_serial; + int64_t stream_owner; // serial of the slot holding the live codec state, -1 none +}; + +TtsEngine * tts_engine_new(PipelineTTS * pt, BPETokenizer * tok) { + TtsEngine * e = new TtsEngine(); + e->pt = pt; + e->tok = tok; + e->next_serial = 0; + e->stream_owner = -1; + e->slots.reserve((size_t) pt->max_batch); + return e; +} + +void tts_engine_free(TtsEngine * e) { + if (!e) { + return; + } + for (TtsSlot & s : e->slots) { + pipeline_codec_snap_free(&s.snap); + } + delete e; +} + +int tts_engine_active(const TtsEngine * e) { + return (int) e->slots.size(); +} + +// Make `s` the owner of the live codec stream state: park the current +// owner's state into its mirror, then restore this slot's. A slot with +// no mirror yet keeps whatever is live (its admit resets the state +// right after). With one streaming slot the ownership never moves and +// no copy is paid. +static bool tts_engine_codec_own(TtsEngine * e, TtsSlot * s) { + if (e->stream_owner == s->serial) { + return true; + } + if (e->stream_owner >= 0) { + for (TtsSlot & o : e->slots) { + if (o.serial == e->stream_owner) { + if (!pipeline_codec_stream_save(&e->pt->codec, &o.snap)) { + return false; + } + break; + } + } + } + e->stream_owner = -1; + if (s->snap.ctx && !pipeline_codec_stream_load(&e->pt->codec, &s->snap)) { + return false; + } + e->stream_owner = s->serial; + return true; +} + +static bool tts_admit_fail(TtsJob * job, qt_status st) { + job->status = st; + job->error = qt_last_error(); + return false; +} + +bool tts_engine_admit(TtsEngine * e, TtsJob * job) { + PipelineTTS * pt = e->pt; + const struct qt_tts_params * params = job->params; + job->status = QT_STATUS_OK; + job->error.clear(); + + if ((int) e->slots.size() >= pt->max_batch) { + qt_set_error("tts_engine_admit: no free slot (%d active, max_batch %d)", (int) e->slots.size(), pt->max_batch); + return tts_admit_fail(job, QT_STATUS_INVALID_PARAMS); + } + + const std::string instruct = params->instruct ? params->instruct : ""; + const std::string speaker = params->speaker ? params->speaker : ""; + const std::string ref_text = params->ref_text ? params->ref_text : ""; // ABI v2 latent reference fields. Callers compiled against ABI 1 // never set them; the abi_version gate keeps their uninitialised @@ -467,14 +619,14 @@ qt_status pipeline_tts_synthesize(PipelineTTS * pt, if (has_ref_audio && (has_lat_spk || has_lat_codes)) { qt_set_error("pipeline_tts_synthesize: ref_audio_24k and ref_spk_emb / ref_codes are mutually exclusive"); qt_log(QT_LOG_ERROR, "[Pipeline] ref_audio_24k and ref_spk_emb / ref_codes are mutually exclusive"); - return QT_STATUS_INVALID_PARAMS; + return tts_admit_fail(job, QT_STATUS_INVALID_PARAMS); } // Latent ICL codes ride on top of the speaker embedding and need the // transcript, mirroring the raw path where mode B implies mode A. if (has_lat_codes && (!has_lat_spk || ref_text.empty())) { qt_set_error("pipeline_tts_synthesize: ref_codes requires ref_spk_emb and ref_text"); qt_log(QT_LOG_ERROR, "[Pipeline] ref_codes requires ref_spk_emb and ref_text"); - return QT_STATUS_INVALID_PARAMS; + return tts_admit_fail(job, QT_STATUS_INVALID_PARAMS); } // Voice clone mode A: a pre-extracted latent embedding feeds the @@ -489,7 +641,7 @@ qt_status pipeline_tts_synthesize(PipelineTTS * pt, pt->talker.hidden_size); qt_log(QT_LOG_ERROR, "[Pipeline] ref_spk_dim %d mismatches talker hidden %d", lat_spk_dim, pt->talker.hidden_size); - return QT_STATUS_INVALID_PARAMS; + return tts_admit_fail(job, QT_STATUS_INVALID_PARAMS); } ref_spk_emb_ptr = lat_spk_emb; qt_log(QT_LOG_INFO, "[Pipeline] Latent speaker embedding: %d values", lat_spk_dim); @@ -497,7 +649,7 @@ qt_status pipeline_tts_synthesize(PipelineTTS * pt, if (!pt->has_speaker_encoder) { qt_set_error("pipeline_tts_synthesize: --ref-wav requires a model with a speaker encoder (Base only)"); qt_log(QT_LOG_ERROR, "[Pipeline] --ref-wav requires a model with a speaker encoder (Base only)"); - return QT_STATUS_GENERATE_FAILED; + return tts_admit_fail(job, QT_STATUS_GENERATE_FAILED); } // Lazy residency: the first reference audio request pays the // weight load once, pre extracted paths never do. @@ -508,86 +660,106 @@ qt_status pipeline_tts_synthesize(PipelineTTS * pt, pt->has_speaker_encoder = false; qt_set_error("pipeline_tts_synthesize: speaker encoder load failed"); qt_log(QT_LOG_ERROR, "[Pipeline] speaker encoder load failed"); - return QT_STATUS_GENERATE_FAILED; + return tts_admit_fail(job, QT_STATUS_GENERATE_FAILED); } pt->spk_enc_loaded = true; qt_log(QT_LOG_INFO, "[Pipeline] Speaker encoder lazy loaded in %.0f ms", t_spk_load.ms()); } if (!speaker_encoder_extract(&pt->speaker_encoder, pt->sched, params->ref_audio_24k, params->ref_n_samples, ref_spk_emb, params->dump_dir)) { - return QT_STATUS_GENERATE_FAILED; + return tts_admit_fail(job, QT_STATUS_GENERATE_FAILED); } if ((int) ref_spk_emb.size() != pt->talker.hidden_size) { qt_set_error("pipeline_tts_synthesize: speaker embedding size %zu mismatches talker hidden %d", ref_spk_emb.size(), pt->talker.hidden_size); qt_log(QT_LOG_ERROR, "[Pipeline] speaker embedding size %zu mismatches talker hidden %d", ref_spk_emb.size(), pt->talker.hidden_size); - return QT_STATUS_GENERATE_FAILED; + return tts_admit_fail(job, QT_STATUS_GENERATE_FAILED); } ref_spk_emb_ptr = ref_spk_emb.data(); } + // Slot construction: everything below fills the tail slot; a + // failure pops it and reports through the job. + e->slots.emplace_back(); + TtsSlot & s = e->slots.back(); + const int slot_idx = (int) e->slots.size() - 1; + s.job = job; + s.serial = e->next_serial++; + s.ref_codes_ptr = NULL; + s.ref_codes_T = 0; + s.step = 0; + s.subseq_counter = 0; + s.prev_overlay = NULL; + s.pending_c0 = -1; + s.has_frame = false; + s.streaming = (params->on_chunk != NULL); + s.snap = {}; + s.finished = false; + s.fin_status = QT_STATUS_OK; + s.perf = {}; + s.t_total.reset(); + // Voice clone mode B: pre-encoded latent codes feed the ICL prompt // directly; otherwise, if ref_text is given, encode the reference // audio into 16 codebook indices via the codec encoder. Layout is // [num_codebooks, T_codec] row major in both cases, matching what // the prompt builder expects for the ICL sum loop. - std::vector ref_codes; - const int32_t * ref_codes_ptr = NULL; - int ref_codes_T = 0; if (has_lat_codes) { - ref_codes_ptr = lat_codes; - ref_codes_T = lat_T; - qt_log(QT_LOG_INFO, "[Pipeline] Latent ICL ref_codes: %d frames at 12.5 Hz", ref_codes_T); + s.ref_codes_ptr = lat_codes; + s.ref_codes_T = lat_T; + qt_log(QT_LOG_INFO, "[Pipeline] Latent ICL ref_codes: %d frames at 12.5 Hz", s.ref_codes_T); } else if (!ref_text.empty()) { if (!has_ref_audio) { qt_set_error("pipeline_tts_synthesize: ref_text requires ref_audio_24k or latent ref_codes"); qt_log(QT_LOG_ERROR, "[Pipeline] ref_text requires ref_audio_24k or latent ref_codes"); - return QT_STATUS_INVALID_PARAMS; + e->slots.pop_back(); + return tts_admit_fail(job, QT_STATUS_INVALID_PARAMS); } // The codec hop is 1920 samples at 24 kHz so n_samples must be // a multiple of 1920. Truncate to the nearest hop boundary. if (params->ref_n_samples < TOKENIZER_HOP_LENGTH) { qt_set_error("pipeline_tts_synthesize: ref_wav too short for ICL (%d samples)", params->ref_n_samples); qt_log(QT_LOG_ERROR, "[Pipeline] ref_wav too short for ICL (%d samples)", params->ref_n_samples); - return QT_STATUS_INVALID_PARAMS; + e->slots.pop_back(); + return tts_admit_fail(job, QT_STATUS_INVALID_PARAMS); } - int aligned_T = (params->ref_n_samples / TOKENIZER_HOP_LENGTH) * TOKENIZER_HOP_LENGTH; - ref_codes = pipeline_codec_encode(&pt->codec, params->ref_audio_24k, aligned_T, params->dump_dir); - if (ref_codes.empty()) { + int aligned_T = (params->ref_n_samples / TOKENIZER_HOP_LENGTH) * TOKENIZER_HOP_LENGTH; + s.ref_codes_store = pipeline_codec_encode(&pt->codec, params->ref_audio_24k, aligned_T, params->dump_dir); + if (s.ref_codes_store.empty()) { qt_set_error("pipeline_tts_synthesize: pipeline_codec_encode returned empty codes"); qt_log(QT_LOG_ERROR, "[Pipeline] pipeline_codec_encode returned empty codes"); - return QT_STATUS_GENERATE_FAILED; + e->slots.pop_back(); + return tts_admit_fail(job, QT_STATUS_GENERATE_FAILED); } - ref_codes_ptr = ref_codes.data(); - ref_codes_T = (int) ref_codes.size() / pt->num_code_groups; - qt_log(QT_LOG_INFO, "[Pipeline] ICL ref_codes: %d frames at 12.5 Hz (%d audio samples)", ref_codes_T, + s.ref_codes_ptr = s.ref_codes_store.data(); + s.ref_codes_T = (int) s.ref_codes_store.size() / pt->num_code_groups; + qt_log(QT_LOG_INFO, "[Pipeline] ICL ref_codes: %d frames at 12.5 Hz (%d audio samples)", s.ref_codes_T, aligned_T); } - TtsPerf perf = {}; - Timer t_total; - // NULL lang selects automatic language: the prompt carries no // language id and the model infers it from the text. const char * lang = params->lang ? params->lang : "auto"; Timer t_build; - if (!prompt_builder_build(pt, tok, params->text, lang, instruct, speaker, ref_spk_emb_ptr, ref_text, ref_codes_ptr, - ref_codes_T, &prompt)) { - return QT_STATUS_GENERATE_FAILED; + if (!prompt_builder_build(pt, e->tok, params->text, lang, instruct, speaker, ref_spk_emb_ptr, ref_text, + s.ref_codes_ptr, s.ref_codes_T, &s.prompt)) { + e->slots.pop_back(); + return tts_admit_fail(job, QT_STATUS_GENERATE_FAILED); } - perf.build_ms = t_build.ms(); + s.perf.build_ms = t_build.ms(); if (params->dump_dir) { DebugDumper d; debug_init(&d, params->dump_dir); - std::vector ids32(prompt.prompt_ids.begin(), prompt.prompt_ids.end()); + std::vector ids32(s.prompt.prompt_ids.begin(), s.prompt.prompt_ids.end()); int n_ids = (int) ids32.size(); debug_dump_i32_as_f32(&d, "prompt-ids", ids32.data(), &n_ids, 1); - debug_dump_2d(&d, "talker-input-embed", prompt.input_embed.data(), prompt.T_ctx, prompt.hidden); - debug_dump_2d(&d, "trailing-text-hidden", prompt.trailing_text_hidden.data(), prompt.T_trailing, prompt.hidden); - debug_dump_1d(&d, "tts-pad-embed", prompt.tts_pad_embed.data(), prompt.hidden); + debug_dump_2d(&d, "talker-input-embed", s.prompt.input_embed.data(), s.prompt.T_ctx, s.prompt.hidden); + debug_dump_2d(&d, "trailing-text-hidden", s.prompt.trailing_text_hidden.data(), s.prompt.T_trailing, + s.prompt.hidden); + debug_dump_1d(&d, "tts-pad-embed", s.prompt.tts_pad_embed.data(), s.prompt.hidden); // Voice clone dumps: spk-emb fires when ref_wav is set // (modes A and B), ref-codes fires only when ref_text is also set @@ -596,329 +768,528 @@ qt_status pipeline_tts_synthesize(PipelineTTS * pt, if (ref_spk_emb_ptr != NULL) { debug_dump_1d(&d, "spk-emb", ref_spk_emb_ptr, pt->talker.hidden_size); } - if (ref_codes_T > 0) { - const int shape[2] = { pt->num_code_groups, ref_codes_T }; - debug_dump_i32_as_f32(&d, "ref-codes", ref_codes_ptr, shape, 2); + if (s.ref_codes_T > 0) { + const int shape[2] = { pt->num_code_groups, s.ref_codes_T }; + debug_dump_i32_as_f32(&d, "ref-codes", s.ref_codes_ptr, shape, 2); } } - // Generation loop: step 0 prefills the talker over the full prompt - // and writes T_ctx positions into the KV cache. Subsequent steps - // feed one next_emb at a time and append one position. The code - // predictor maintains its own per-frame cache that gets reset at - // every step. - const int hidden = prompt.hidden; - const int codec_eos_id = pt->codec_specials.eos_id; - const int num_codebooks = pt->num_code_groups; - const int talker_vocab = pt->talker.vocab_size; - const bool use_fa = pt->use_flash_attn; - const bool clamp_fp16 = pt->clamp_fp16; - const float talker_T = params->do_sample ? params->temperature : 0.0f; - const float subtk_T = params->subtalker_do_sample ? params->subtalker_temperature : 0.0f; - const float talker_rp = params->repetition_penalty; - - // Codec decode framing. The buffered path routes through - // codec_chunked_decode with a rolling left context window mirroring - // the upstream Qwen3-TTS 12 Hz tokenizer chunked_decode rule: every - // chunk re uses up to left_ctx_frames previously decoded frames as - // left context, then the matching audio samples are stripped from - // the head of the decoded chunk. The streaming path decodes frame - // by frame through the stateful codec and ignores both knobs. - const bool streaming = (params->on_chunk != NULL); - const float chunk_sec = params->codec_chunk_sec > 0.0f ? params->codec_chunk_sec : 24.0f; - const float left_ctx_sec = params->codec_left_context_sec >= 0.0f ? params->codec_left_context_sec : 2.0f; - const int chunk_frames = pipeline_tts_duration_sec_to_tokens(pt, chunk_sec); - const int left_ctx_frames = pipeline_tts_duration_sec_to_tokens(pt, left_ctx_sec); - - std::vector> all_codes; - all_codes.reserve((size_t) params->max_new_tokens); - - // c0 codes already emitted, fed to repetition penalty. - std::vector talker_history; - talker_history.reserve((size_t) params->max_new_tokens); - - // Global Philox subsequence counter advances once per primitive - // sample (one for c0 of each step, then 15 for the predictor codes). - int64_t subseq_counter = 0; - - // Decode input state: the codes sampled at the previous frame plus - // the trailing text / pad overlay row for that frame. The talker - // decode graph gathers and sums the 16 embeddings on device. - std::vector prev_ids((size_t) num_codebooks, 0); - const float * prev_overlay = NULL; + s.talker_T = params->do_sample ? params->temperature : 0.0f; + s.subtk_T = params->subtalker_do_sample ? params->subtalker_temperature : 0.0f; + s.prev_ids.assign((size_t) pt->num_code_groups, 0); + s.all_codes.reserve((size_t) params->max_new_tokens); + s.talker_history.reserve((size_t) params->max_new_tokens); // Stateful streaming decoder: every generated frame decodes // immediately through the persistent codec state and emits its // samples, so the first audio callback fires with the first frame. // ICL clone priming runs the full reference through the same state, // matching the upstream reference plus generated decode exactly. - codec_stream_decoder stream; - if (streaming) { - if (!stream.init(&pt->codec, num_codebooks)) { - qt_set_error("pipeline_tts_synthesize: codec stream state init failed"); - return QT_STATUS_GENERATE_FAILED; + // Taking ownership parks the previous streaming slot's state first. + if (s.streaming) { + if (!tts_engine_codec_own(e, &s)) { + qt_set_error("pipeline_tts_synthesize: codec stream state park failed"); + e->slots.pop_back(); + return tts_admit_fail(job, QT_STATUS_GENERATE_FAILED); } - if (ref_codes_ptr != NULL) { + if (!s.stream.init(&pt->codec, pt->num_code_groups)) { + qt_set_error("pipeline_tts_synthesize: codec stream state init failed"); + e->stream_owner = -1; + e->slots.pop_back(); + return tts_admit_fail(job, QT_STATUS_GENERATE_FAILED); + } + if (s.ref_codes_ptr != NULL) { Timer t_seed; - if (!stream.seed_reference(&pt->codec, ref_codes_ptr, ref_codes_T)) { + if (!s.stream.seed_reference(&pt->codec, s.ref_codes_ptr, s.ref_codes_T)) { qt_set_error("pipeline_tts_synthesize: codec stream reference priming failed"); - return QT_STATUS_GENERATE_FAILED; + e->stream_owner = -1; + e->slots.pop_back(); + return tts_admit_fail(job, QT_STATUS_GENERATE_FAILED); } - perf.codec_ms += t_seed.ms(); + s.perf.codec_ms += t_seed.ms(); } } - for (int step = 0; step < params->max_new_tokens; step++) { + // Talker prefill into KV set slot_idx: the joining request stalls + // every already active slot for the duration of one prefill, so the + // measured span is the batch's join cost. + TalkerForwardOutput fw; + Timer t_prefill; + if (!talker_forward_prefill(&pt->talker, &pt->talker_kv, slot_idx, pt->sched, &pt->talker_arena, pt->hidden_bridge, + s.prompt.input_embed.data(), s.prompt.T_ctx, pt->use_flash_attn, pt->clamp_fp16, + params->dump_dir, &fw)) { + if (e->stream_owner == s.serial) { + e->stream_owner = -1; + } + e->slots.pop_back(); + return tts_admit_fail(job, QT_STATUS_GENERATE_FAILED); + } + s.perf.prefill_ms = t_prefill.ms(); + s.logits = std::move(fw.logits_last); + qt_log(QT_LOG_INFO, "[Batch] Admit slot=%d T_ctx=%d prefill=%.1f ms build=%.1f ms (stall for %d active slots)", + slot_idx, s.prompt.T_ctx, s.perf.prefill_ms, s.perf.build_ms, slot_idx); + return true; +} + +// Retire one finished slot: streaming drain or buffered codec decode, +// perf accounting, job status and worker side error capture. The codec +// stream mirror releases here. +static void tts_slot_complete(TtsEngine * e, TtsSlot & s) { + PipelineTTS * pt = e->pt; + TtsJob * job = s.job; + const struct qt_tts_params * params = job->params; + qt_status st = s.fin_status; + + if (st == QT_STATUS_OK) { + qt_log(QT_LOG_INFO, "[Pipeline] Generation done : %zu frames", s.all_codes.size()); + s.perf.n_frames = (int) s.all_codes.size(); + + const int num_codebooks = pt->num_code_groups; + if (params->dump_dir && !s.all_codes.empty()) { + DebugDumper d; + debug_init(&d, params->dump_dir); + int T_frames = (int) s.all_codes.size(); + std::vector flat((size_t) T_frames * (size_t) num_codebooks); + for (int t = 0; t < T_frames; t++) { + for (int k = 0; k < num_codebooks; k++) { + flat[(size_t) t * (size_t) num_codebooks + (size_t) k] = s.all_codes[(size_t) t][(size_t) k]; + } + } + int shape[2] = { T_frames, num_codebooks }; + debug_dump_i32_as_f32(&d, "codes-full", flat.data(), shape, 2); + } + + if (s.streaming) { + // Streaming tail: drain the sub chunk remainder of the + // ramp, then finish with an empty buffered output. + if (!tts_engine_codec_own(e, &s)) { + qt_set_error("pipeline_tts_synthesize: codec stream state park failed"); + st = QT_STATUS_GENERATE_FAILED; + } else if (!s.stream.drain(&pt->codec, params->on_chunk, params->on_chunk_user_data)) { + if (s.stream.cancelled) { + qt_log(QT_LOG_INFO, "[Pipeline] on_chunk callback aborted the synthesis"); + st = QT_STATUS_CANCELLED; + } else { + qt_set_error("pipeline_tts_synthesize: streaming codec drain failed"); + qt_log(QT_LOG_ERROR, "[Pipeline] streaming codec drain failed"); + st = QT_STATUS_GENERATE_FAILED; + } + } else { + if (job->out) { + job->out->samples = NULL; + job->out->n_samples = 0; + job->out->sample_rate = TOKENIZER_SAMPLE_RATE; + job->out->channels = 1; + } + s.perf.total_ms = s.t_total.ms(); + tts_log_perf(s.perf); + } + } else if (s.all_codes.empty()) { + // Buffered path: empty all_codes means EOS at step 0 with + // no audio. Return success and an empty qt_audio struct; + // the facade leaves it to the caller to decide what to do + // with a zero sample synthesis. + job->out->samples = NULL; + job->out->n_samples = 0; + job->out->sample_rate = TOKENIZER_SAMPLE_RATE; + job->out->channels = 1; + s.perf.total_ms = s.t_total.ms(); + tts_log_perf(s.perf); + } else { + // Buffered codec decode through the chunked path : same framing as + // the streaming branch (chunk_frames + left_ctx_frames), bit perfect + // equivalent to a single pipeline_codec_decode call when T_frames + // fits in one chunk, bounded VRAM beyond that. Transpose codes from + // [T_frames, K] to [K, T_frames] because codec_chunked_decode + // expects K major layout. On the ICL clone path the tail of the + // reference codes prepends the buffer so the onset is voiced with + // the reference's causal state, mirroring the upstream pipeline + // which decodes reference plus generated then trims; the seeded + // samples strip from the front afterwards. Raising + // codec_left_context_sec past the reference duration reproduces the + // upstream full reference decode exactly. + const float chunk_sec = params->codec_chunk_sec > 0.0f ? params->codec_chunk_sec : 24.0f; + const float left_ctx_sec = params->codec_left_context_sec >= 0.0f ? params->codec_left_context_sec : 2.0f; + const int chunk_frames = pipeline_tts_duration_sec_to_tokens(pt, chunk_sec); + const int left_ctx_frames = pipeline_tts_duration_sec_to_tokens(pt, left_ctx_sec); + + const int T_frames = (int) s.all_codes.size(); + int seed = 0; + if (s.ref_codes_ptr != NULL) { + seed = s.ref_codes_T < left_ctx_frames ? s.ref_codes_T : left_ctx_frames; + } + const int T_dec = seed + T_frames; + std::vector codes_kt((size_t) num_codebooks * (size_t) T_dec); + for (int k = 0; k < num_codebooks; k++) { + int32_t * row = codes_kt.data() + (size_t) k * (size_t) T_dec; + if (seed > 0) { + std::memcpy(row, + s.ref_codes_ptr + (size_t) k * (size_t) s.ref_codes_T + (size_t) (s.ref_codes_T - seed), + (size_t) seed * sizeof(int32_t)); + } + for (int t = 0; t < T_frames; t++) { + row[(size_t) (seed + t)] = s.all_codes[(size_t) t][(size_t) k]; + } + } + Timer t_codec; + std::vector audio = + codec_chunked_decode(&pt->codec, codes_kt.data(), num_codebooks, T_dec, chunk_frames, left_ctx_frames); + s.perf.codec_ms += t_codec.ms(); + if (audio.empty()) { + qt_set_error("pipeline_tts_synthesize: codec decode returned no audio"); + qt_log(QT_LOG_ERROR, "[Pipeline] codec decode returned no audio"); + st = QT_STATUS_GENERATE_FAILED; + } else { + if (seed > 0) { + audio.erase(audio.begin(), audio.begin() + (size_t) seed * (size_t) TOKENIZER_HOP_LENGTH); + } + if (params->dump_dir) { + DebugDumper d; + debug_init(&d, params->dump_dir); + debug_dump_1d(&d, "output-audio", audio.data(), (int) audio.size()); + } + if (!fill_qt_audio(audio, job->out)) { + st = QT_STATUS_OOM; + } else { + s.perf.total_ms = s.t_total.ms(); + tts_log_perf(s.perf); + } + } + } + } + + if (e->stream_owner == s.serial) { + e->stream_owner = -1; + } + pipeline_codec_snap_free(&s.snap); + + if (st != QT_STATUS_OK) { + job->error = qt_last_error(); + } + job->status = st; +} + +void tts_engine_step(TtsEngine * e, std::vector * retired) { + PipelineTTS * pt = e->pt; + const int N = (int) e->slots.size(); + if (N == 0) { + return; + } + const int hidden = pt->talker.hidden_size; + const int vocab = pt->talker.vocab_size; + const int num_codebooks = pt->num_code_groups; + const int codec_eos_id = pt->codec_specials.eos_id; + const int n_acoustic = pt->code_predictor.num_acoustic_codebooks; + + // 1) Batched talker decode over the slots past their prefill. The + // freshly admitted slots form a contiguous tail (step == 0) and + // consume their prefill logits instead; every retirement happens at + // frame end when all survivors carry step >= 1, so the decode span + // [0, N_dec) stays consecutive by construction. + int N_dec = 0; + while (N_dec < N && e->slots[(size_t) N_dec].step > 0) { + N_dec++; + } + bool any_dump = false; + for (int i = 0; i < N; i++) { + any_dump = any_dump || (e->slots[(size_t) i].job->params->dump_dir != NULL); + } + if (N_dec > 0) { + std::vector ids((size_t) num_codebooks * (size_t) N_dec); + std::vector overlays((size_t) hidden * (size_t) N_dec); + for (int i = 0; i < N_dec; i++) { + TtsSlot & s = e->slots[(size_t) i]; + for (int g = 0; g < num_codebooks; g++) { + ids[(size_t) g * (size_t) N_dec + (size_t) i] = s.prev_ids[(size_t) g]; + } + std::memcpy(overlays.data() + (size_t) i * (size_t) hidden, s.prev_overlay, + (size_t) hidden * sizeof(float)); + } + TalkerForwardOutput fw; + Timer t_talker; + if (!talker_forward_decode(&pt->talker, &pt->talker_kv, pt->backend, pt->talker_decode_graphs, + pt->hidden_bridge, ids.data(), pt->code_predictor.codec_embedding.data(), n_acoustic, + overlays.data(), N_dec, pt->use_flash_attn, pt->clamp_fp16, any_dump, &fw)) { + qt_set_error("pipeline_tts_synthesize: talker decode failed"); + for (TtsSlot & s : e->slots) { + s.finished = true; + s.fin_status = QT_STATUS_GENERATE_FAILED; + } + } else { + const double ms = t_talker.ms(); + for (int i = 0; i < N_dec; i++) { + TtsSlot & s = e->slots[(size_t) i]; + s.perf.talker_ms += ms; + s.logits.assign(fw.logits_last.begin() + (size_t) i * (size_t) vocab, + fw.logits_last.begin() + (size_t) (i + 1) * (size_t) vocab); + + // Bisection dump: the talker hidden_last at step 1 is + // the input the code predictor consumes after consuming + // the next-emb of step 0. Pairing it byte for byte with + // the Python hook tells us whether the next-emb + // composition + talker decode round trip is bit exact + // end to end. + if (s.job->params->dump_dir && s.step == 1) { + DebugDumper d; + debug_init(&d, s.job->params->dump_dir); + debug_dump_1d(&d, "talker-hidden-step1", fw.hidden_last.data() + (size_t) i * (size_t) hidden, + hidden); + } + } + } + } + + // 2) Cancel poll and per slot c0 sampling: suppression, repetition + // penalty over the slot's own history, its own Philox stream. + for (int i = 0; i < N; i++) { + TtsSlot & s = e->slots[(size_t) i]; + s.has_frame = false; + if (s.finished) { + continue; + } + const struct qt_tts_params * p = s.job->params; + // Cooperative cancellation, polled at every step. Granularity is // one AR frame = 1 / 12.5 Hz ~ 83 ms of audio, which is well // below any reasonable UX cancel latency target. - if (params->cancel && params->cancel(params->cancel_user_data)) { - qt_log(QT_LOG_INFO, "[Pipeline] cancelled at step %d", step); - return QT_STATUS_CANCELLED; - } - - TalkerForwardOutput fw; - const char * step_dump = (params->dump_dir && step == 0) ? params->dump_dir : NULL; - bool ok; - Timer t_talker; - if (step == 0) { - 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); - } else { - ok = talker_forward_decode(&pt->talker, &pt->talker_kv, pt->backend, pt->talker_decode_graphs.data(), - pt->hidden_bridge, prev_ids.data(), pt->code_predictor.codec_embedding.data(), - pt->code_predictor.num_acoustic_codebooks, prev_overlay, use_fa, clamp_fp16, - params->dump_dir != NULL, &fw); - } - if (!ok) { - return QT_STATUS_GENERATE_FAILED; - } - if (step == 0) { - perf.prefill_ms = t_talker.ms(); - } else { - perf.talker_ms += t_talker.ms(); - } - - // Bisection dump: the talker hidden_last at step 1 is the input - // the code predictor consumes after consuming the next-emb of - // step 0. Pairing it byte for byte with the Python hook tells us - // whether the next-emb composition + talker decode round trip - // is bit exact end to end. - if (params->dump_dir && step == 1) { - DebugDumper d; - debug_init(&d, params->dump_dir); - debug_dump_1d(&d, "talker-hidden-step1", fw.hidden_last.data(), hidden); + if (p->cancel && p->cancel(p->cancel_user_data)) { + qt_log(QT_LOG_INFO, "[Pipeline] cancelled at step %d (slot %d)", s.step, i); + s.finished = true; + s.fin_status = QT_STATUS_CANCELLED; + continue; } // Apply codec suppression: forbid [vocab - 1024, vocab) except // codec_eos. Then run the upstream sampling chain. Timer t_host; - apply_suppress(fw.logits_last.data(), talker_vocab, talker_vocab - 1024, talker_vocab, codec_eos_id); + apply_suppress(s.logits.data(), vocab, vocab - 1024, vocab, codec_eos_id); float u_c0 = 0.0f; - int c0 = - sample_top_k_p(fw.logits_last.data(), talker_vocab, talker_T, params->top_k, params->top_p, talker_rp, - talker_history.data(), (int) talker_history.size(), resolved_seed, subseq_counter, &u_c0); - perf.host_ms += t_host.ms(); - subseq_counter++; + int c0 = sample_top_k_p(s.logits.data(), vocab, s.talker_T, p->top_k, p->top_p, p->repetition_penalty, + s.talker_history.data(), (int) s.talker_history.size(), s.job->resolved_seed, + s.subseq_counter, &u_c0); + s.perf.host_ms += t_host.ms(); + s.subseq_counter++; if (c0 < 0) { + qt_set_error("pipeline_tts_synthesize: c0 sample returned no candidate"); qt_log(QT_LOG_ERROR, "[Pipeline] c0 sample returned no candidate"); - return QT_STATUS_GENERATE_FAILED; + s.finished = true; + s.fin_status = QT_STATUS_GENERATE_FAILED; + continue; } // Trace the first 32 samples unconditionally so [Sample] lines // up with [Sample-PY] / [Sample-CP] across the 16 codes of step // 0 and step 1 the Python harness emits. - if ((subseq_counter - 1) < 32) { - qt_log(QT_LOG_DEBUG, "[Sample] step=%d c0=%d u=%.10f subseq=%lld", step, c0, (double) u_c0, - (long long) (subseq_counter - 1)); + if ((s.subseq_counter - 1) < 32) { + qt_log(QT_LOG_DEBUG, "[Sample] step=%d c0=%d u=%.10f subseq=%lld", s.step, c0, (double) u_c0, + (long long) (s.subseq_counter - 1)); } if (c0 == codec_eos_id) { - qt_log(QT_LOG_INFO, "[Pipeline] EOS at step %d, stopping", step); - break; + qt_log(QT_LOG_INFO, "[Pipeline] EOS at step %d, stopping (slot %d)", s.step, i); + s.finished = true; + continue; } + s.pending_c0 = c0; + s.has_frame = true; + } + // 3) Batched code predictor over all N lanes in lockstep. Lanes + // whose slot finished this frame ride along with a zero id and get + // discarded; the live lanes each consume their own Philox stream so + // per slot outputs stay identical to a single sequence run. + bool any_live = false; + for (int i = 0; i < N; i++) { + any_live = any_live || e->slots[(size_t) i].has_frame; + } + if (any_live) { CodePredictorOutput cp; - const char * cp_dump = (params->dump_dir && step == 0) ? params->dump_dir : NULL; - Timer t_pred; - if (!code_predictor_step(&pt->code_predictor, pt->backend, &pt->cp_prefill_graph, pt->cp_step_graphs.data(), c0, - subtk_T, params->subtalker_top_k, params->subtalker_top_p, resolved_seed, - subseq_counter - 1, cp_dump, &cp)) { - return QT_STATUS_GENERATE_FAILED; - } - perf.predictor_ms += t_pred.ms(); - if (step == 0) { - perf.ttfa_ms = t_total.ms(); - } - // Predictor consumed (num_codebooks - 1) subsequences after the - // c0 one (subseq_base + 1 .. subseq_base + 15). - subseq_counter += (num_codebooks - 1); - - all_codes.push_back(cp.codes); - talker_history.push_back(c0); - if (streaming) { - Timer t_codec; - bool pushed = stream.push_frame(&pt->codec, cp.codes.data(), params->on_chunk, params->on_chunk_user_data); - perf.codec_ms += t_codec.ms(); - if (!pushed) { - if (stream.cancelled) { - qt_log(QT_LOG_INFO, "[Pipeline] on_chunk callback aborted the synthesis"); - return QT_STATUS_CANCELLED; + if (!pipeline_tts_cp_graphs_ensure(pt, N)) { + qt_set_error("pipeline_tts_synthesize: code predictor graph build failed (N=%d)", N); + for (TtsSlot & s : e->slots) { + s.finished = true; + s.fin_status = QT_STATUS_GENERATE_FAILED; + s.has_frame = false; + } + } else { + CodePredGraphSet & gs = pt->cp_graphs[(size_t) (N - 1)]; + std::vector c0s((size_t) N, 0); + std::vector temps((size_t) N, 0.0f); + std::vector top_ks((size_t) N, 0); + std::vector top_ps((size_t) N, 1.0f); + std::vector seeds((size_t) N, 0); + std::vector subseqs((size_t) N, 0); + const char * cp_dump = NULL; + for (int i = 0; i < N; i++) { + TtsSlot & s = e->slots[(size_t) i]; + if (!s.has_frame) { + continue; } - qt_set_error("pipeline_tts_synthesize: streaming codec decode failed at frame %d", step); - qt_log(QT_LOG_ERROR, "[Pipeline] streaming codec decode failed at frame %d", step); - return QT_STATUS_GENERATE_FAILED; - } - } - - // Next decode input: the 16 frame codes gather and sum in graph - // (codebook 0 from talker.codec_embedding, the 15 acoustic - // groups from the predictor's private tables). The overlay row - // adds the next utterance text hidden while any remains, the - // tts_pad embedding afterwards. - prev_ids[0] = c0; - for (int g = 1; g < num_codebooks; g++) { - prev_ids[(size_t) g] = cp.codes[(size_t) g]; - } - prev_overlay = (step < prompt.T_trailing) ? - prompt.trailing_text_hidden.data() + (size_t) step * (size_t) hidden : - prompt.tts_pad_embed.data(); - - // Bisection dump: reproduce the in graph composition on host so - // the step 0 next embedding stays byte comparable against the - // Python hook (codebook sums plus trailing text overlay). - if (params->dump_dir && step == 0) { - std::vector next_emb((size_t) hidden, 0.0f); - std::vector tmp((size_t) hidden); - embed_row_from_gguf(pt->gguf_talker, "talker.codec_embd.weight", c0, hidden, tmp.data()); - for (int i = 0; i < hidden; i++) { - next_emb[(size_t) i] += tmp[(size_t) i]; - } - for (int g = 0; g < num_codebooks - 1; g++) { - int cg = cp.codes[(size_t) (g + 1)]; - char name[64]; - snprintf(name, sizeof(name), "code_pred.codec_embd.%d.weight", g); - embed_row_from_gguf(pt->gguf_talker, name, cg, hidden, tmp.data()); - for (int i = 0; i < hidden; i++) { - next_emb[(size_t) i] += tmp[(size_t) i]; + const struct qt_tts_params * p = s.job->params; + c0s[(size_t) i] = s.pending_c0; + temps[(size_t) i] = s.subtk_T; + top_ks[(size_t) i] = p->subtalker_top_k; + top_ps[(size_t) i] = p->subtalker_top_p; + seeds[(size_t) i] = s.job->resolved_seed; + subseqs[(size_t) i] = s.subseq_counter - 1; + if (N == 1 && s.step == 0 && p->dump_dir) { + cp_dump = p->dump_dir; } } - for (int i = 0; i < hidden; i++) { - next_emb[(size_t) i] += prev_overlay[(size_t) i]; - } - DebugDumper d; - debug_init(&d, params->dump_dir); - debug_dump_1d(&d, "next-emb-step0", next_emb.data(), hidden); - } + Timer t_pred; + if (!code_predictor_step(&pt->code_predictor, pt->backend, &gs.prefill, gs.steps.data(), c0s.data(), N, + temps.data(), top_ks.data(), top_ps.data(), seeds.data(), subseqs.data(), cp_dump, + &cp)) { + for (TtsSlot & s : e->slots) { + s.finished = true; + s.fin_status = QT_STATUS_GENERATE_FAILED; + s.has_frame = false; + } + } else { + const double ms = t_pred.ms(); - if (((step + 1) % 8) == 0) { - qt_log(QT_LOG_INFO, "[Pipeline] Generated %d frames", step + 1); - } - } + // 4) Per slot frame post: history, codec streaming, + // next decode inputs. + for (int i = 0; i < N; i++) { + TtsSlot & s = e->slots[(size_t) i]; + if (!s.has_frame) { + continue; + } + const struct qt_tts_params * p = s.job->params; + s.perf.predictor_ms += ms; + if (s.step == 0) { + s.perf.ttfa_ms = s.t_total.ms(); + } + // Predictor consumed (num_codebooks - 1) subsequences + // after the c0 one (subseq_base + 1 .. subseq_base + 15). + s.subseq_counter += (num_codebooks - 1); - qt_log(QT_LOG_INFO, "[Pipeline] Generation done : %zu frames", all_codes.size()); - perf.n_frames = (int) all_codes.size(); + std::vector codes(cp.codes.begin() + (size_t) i * (size_t) num_codebooks, + cp.codes.begin() + (size_t) (i + 1) * (size_t) num_codebooks); + s.all_codes.push_back(codes); + s.talker_history.push_back(s.pending_c0); - if (params->dump_dir && !all_codes.empty()) { - DebugDumper d; - debug_init(&d, params->dump_dir); - int T_frames = (int) all_codes.size(); - std::vector flat((size_t) T_frames * (size_t) num_codebooks); - for (int t = 0; t < T_frames; t++) { - for (int k = 0; k < num_codebooks; k++) { - flat[(size_t) t * (size_t) num_codebooks + (size_t) k] = all_codes[(size_t) t][(size_t) k]; + if (s.streaming) { + Timer t_codec; + bool pushed = tts_engine_codec_own(e, &s) && + s.stream.push_frame(&pt->codec, codes.data(), p->on_chunk, p->on_chunk_user_data); + s.perf.codec_ms += t_codec.ms(); + if (!pushed) { + if (s.stream.cancelled) { + qt_log(QT_LOG_INFO, "[Pipeline] on_chunk callback aborted the synthesis (slot %d)", i); + s.finished = true; + s.fin_status = QT_STATUS_CANCELLED; + } else { + qt_set_error("pipeline_tts_synthesize: streaming codec decode failed at frame %d", + s.step); + qt_log(QT_LOG_ERROR, "[Pipeline] streaming codec decode failed at frame %d (slot %d)", + s.step, i); + s.finished = true; + s.fin_status = QT_STATUS_GENERATE_FAILED; + } + continue; + } + } + + // Next decode input: the 16 frame codes gather and sum + // in graph (codebook 0 from talker.codec_embedding, the + // 15 acoustic groups from the predictor's private + // tables). The overlay row adds the next utterance text + // hidden while any remains, the tts_pad embedding + // afterwards. + for (int g = 0; g < num_codebooks; g++) { + s.prev_ids[(size_t) g] = codes[(size_t) g]; + } + s.prev_overlay = (s.step < s.prompt.T_trailing) ? + s.prompt.trailing_text_hidden.data() + (size_t) s.step * (size_t) hidden : + s.prompt.tts_pad_embed.data(); + + // Bisection dump: reproduce the in graph composition on + // host so the step 0 next embedding stays byte + // comparable against the Python hook (codebook sums plus + // trailing text overlay). + if (p->dump_dir && s.step == 0) { + std::vector next_emb((size_t) hidden, 0.0f); + std::vector tmp((size_t) hidden); + embed_row_from_gguf(pt->gguf_talker, "talker.codec_embd.weight", s.pending_c0, hidden, + tmp.data()); + for (int j = 0; j < hidden; j++) { + next_emb[(size_t) j] += tmp[(size_t) j]; + } + for (int g = 0; g < num_codebooks - 1; g++) { + int cg = codes[(size_t) (g + 1)]; + char name[64]; + snprintf(name, sizeof(name), "code_pred.codec_embd.%d.weight", g); + embed_row_from_gguf(pt->gguf_talker, name, cg, hidden, tmp.data()); + for (int j = 0; j < hidden; j++) { + next_emb[(size_t) j] += tmp[(size_t) j]; + } + } + for (int j = 0; j < hidden; j++) { + next_emb[(size_t) j] += s.prev_overlay[(size_t) j]; + } + DebugDumper d; + debug_init(&d, p->dump_dir); + debug_dump_1d(&d, "next-emb-step0", next_emb.data(), hidden); + } + + s.step++; + if ((s.step % 8) == 0) { + qt_log(QT_LOG_INFO, "[Pipeline] Generated %d frames (slot %d)", s.step, i); + } + if (s.step >= p->max_new_tokens) { + s.finished = true; + } + } } } - int shape[2] = { T_frames, num_codebooks }; - debug_dump_i32_as_f32(&d, "codes-full", flat.data(), shape, 2); } - // Streaming tail: drain the sub chunk remainder of the ramp, then - // finish with an empty buffered output. - if (streaming) { - if (!stream.drain(&pt->codec, params->on_chunk, params->on_chunk_user_data)) { - if (stream.cancelled) { - qt_log(QT_LOG_INFO, "[Pipeline] on_chunk callback aborted the synthesis"); - return QT_STATUS_CANCELLED; - } - qt_set_error("pipeline_tts_synthesize: streaming codec drain failed"); - qt_log(QT_LOG_ERROR, "[Pipeline] streaming codec drain failed"); - return QT_STATUS_GENERATE_FAILED; + // Fresh slots that got no frame this step (their very first frame + // ended in EOS or cancel) still advanced past prefill conceptually; + // slots that emitted advanced in the loop above. Slots neither + // finished nor advanced cannot exist: every live slot either emits + // or finishes. + + // 5) Retirement: swap-remove keeps the active range consecutive. + // The tail slot's talker KV set copies device side into the freed + // index; the bridge column and the predictor set rewrite next frame + // before any read, so only the talker cache moves. + for (int i = 0; i < (int) e->slots.size();) { + if (!e->slots[(size_t) i].finished) { + i++; + continue; } - out->samples = NULL; - out->n_samples = 0; - out->sample_rate = TOKENIZER_SAMPLE_RATE; - out->channels = 1; - perf.total_ms = t_total.ms(); - tts_log_perf(perf); - return QT_STATUS_OK; - } - - // Buffered path: empty all_codes means EOS at step 0 with no audio. - // Return success and an empty qt_audio struct; the facade leaves it - // to the caller to decide what to do with a zero sample synthesis. - if (all_codes.empty()) { - out->samples = NULL; - out->n_samples = 0; - out->sample_rate = TOKENIZER_SAMPLE_RATE; - out->channels = 1; - perf.total_ms = t_total.ms(); - tts_log_perf(perf); - return QT_STATUS_OK; - } - - // Buffered codec decode through the chunked path : same framing as - // the streaming branch (chunk_frames + left_ctx_frames), bit perfect - // equivalent to a single pipeline_codec_decode call when T_frames - // fits in one chunk, bounded VRAM beyond that. Transpose codes from - // [T_frames, K] to [K, T_frames] because codec_chunked_decode - // expects K major layout. On the ICL clone path the tail of the - // reference codes prepends the buffer so the onset is voiced with - // the reference's causal state, mirroring the upstream pipeline - // which decodes reference plus generated then trims; the seeded - // samples strip from the front afterwards. Raising - // codec_left_context_sec past the reference duration reproduces the - // upstream full reference decode exactly. - const int T_frames = (int) all_codes.size(); - int seed = 0; - if (ref_codes_ptr != NULL) { - seed = ref_codes_T < left_ctx_frames ? ref_codes_T : left_ctx_frames; - } - const int T_dec = seed + T_frames; - std::vector codes_kt((size_t) num_codebooks * (size_t) T_dec); - for (int k = 0; k < num_codebooks; k++) { - int32_t * row = codes_kt.data() + (size_t) k * (size_t) T_dec; - if (seed > 0) { - std::memcpy(row, ref_codes_ptr + (size_t) k * (size_t) ref_codes_T + (size_t) (ref_codes_T - seed), - (size_t) seed * sizeof(int32_t)); + tts_slot_complete(e, e->slots[(size_t) i]); + if (retired) { + retired->push_back(e->slots[(size_t) i].job); } - for (int t = 0; t < T_frames; t++) { - row[(size_t) (seed + t)] = all_codes[(size_t) t][(size_t) k]; + const int last = (int) e->slots.size() - 1; + if (i != last) { + kv_cache_copy_set(&pt->talker_kv, last, i); + e->slots[(size_t) i] = std::move(e->slots[(size_t) last]); } + e->slots.pop_back(); } - Timer t_codec; - std::vector audio = - codec_chunked_decode(&pt->codec, codes_kt.data(), num_codebooks, T_dec, chunk_frames, left_ctx_frames); - perf.codec_ms += t_codec.ms(); - if (audio.empty()) { - qt_set_error("pipeline_tts_synthesize: codec decode returned no audio"); - qt_log(QT_LOG_ERROR, "[Pipeline] codec decode returned no audio"); - return QT_STATUS_GENERATE_FAILED; - } - if (seed > 0) { - audio.erase(audio.begin(), audio.begin() + (size_t) seed * (size_t) TOKENIZER_HOP_LENGTH); - } - - if (params->dump_dir) { - DebugDumper d; - debug_init(&d, params->dump_dir); - debug_dump_1d(&d, "output-audio", audio.data(), (int) audio.size()); - } +} - if (!fill_qt_audio(audio, out)) { - return QT_STATUS_OOM; +qt_status pipeline_tts_synthesize(PipelineTTS * pt, + BPETokenizer * tok, + const struct qt_tts_params * params, + int64_t resolved_seed, + struct qt_audio * out) { + TtsEngine * e = tts_engine_new(pt, tok); + TtsJob job; + job.params = params; + job.resolved_seed = resolved_seed; + job.out = out; + job.status = QT_STATUS_OK; + job.done = false; + if (tts_engine_admit(e, &job)) { + while (tts_engine_active(e) > 0) { + tts_engine_step(e, NULL); + } } - perf.total_ms = t_total.ms(); - tts_log_perf(perf); - return QT_STATUS_OK; + tts_engine_free(e); + return job.status; } diff --git a/src/pipeline-tts.h b/src/pipeline-tts.h index 8f4941d..96d689d 100644 --- a/src/pipeline-tts.h +++ b/src/pipeline-tts.h @@ -90,6 +90,13 @@ struct PromptCache { size_t max_prefix_entries; }; +// One set of static predictor graphs for a given batch width: the T=2 +// prefill plus one T=1 step per acoustic codebook after the first. +struct CodePredGraphSet { + CodePredGraph prefill; + std::vector steps; +}; + struct PipelineTTS { GGUFModel gguf_talker; TalkerWeights talker; @@ -108,6 +115,11 @@ struct PipelineTTS { std::string model_type; int num_code_groups; + // Batch capacity: number of KV sets, bridge columns and maximum + // concurrent slots the batch engine drives. 1 keeps the exact + // single sequence layout and behavior. + int max_batch; + CodecSpecials codec_specials; TextSpecials text_specials; std::vector languages; @@ -128,43 +140,47 @@ struct PipelineTTS { bool use_flash_attn; bool clamp_fp16; - // Persistent KV caches: the talker holds the LM context, the - // predictor holds one frame's 16 sub-steps and gets reset every - // frame in code_predictor_step. + // Persistent KV caches, one set per slot: the talker holds the LM + // contexts, the predictor holds one frame's 16 sub-steps per slot, + // rewritten every frame at baked rows. KVCache talker_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`. + // Hidden bridge: the talker last position hidden of every slot + // stays resident on device as one column of [talker_hidden, + // max_batch] f32. The talker graphs copy their columns in, the code + // predictor prefill graph reads them as a leaf, so the AR hot loop + // never round trips the rows through the host. struct ggml_context * bridge_ctx; ggml_backend_buffer_t bridge_buf; struct ggml_tensor * hidden_bridge; // Persistent graph arena for the talker prefill (T_ctx varies per // request, rebuilt through the sched). The talker decode and the - // whole predictor run on static graphs instead: the decode keeps - // one graph per attention window class built lazily, the predictor - // one prefill (T=2) plus one per acoustic step built at load, all - // replayed directly on the backend. + // whole predictor run on static batched graphs instead: the decode + // keeps one graph per attention window class built lazily and + // rebuilt when the batch width changes, the predictor one graph + // set (prefill T=2 plus one per acoustic step) per batch width + // built lazily, all replayed directly on the backend. GraphArena talker_arena; std::vector talker_decode_graphs; // one per 256 step window class, lazy - CodePredGraph cp_prefill_graph; - std::vector cp_step_graphs; + std::vector cp_graphs; // index N - 1, lazy per batch width }; // Open the talker GGUF and the codec GGUF, load every module on the // shared backend. Aborts with a logged error on any missing tensor or // invalid metadata. use_fa is gated on bp.has_gpu inside the load: // CPU only runs always use the manual F32 attention chain. clamp_fp16 -// is forwarded as is. Caller frees with pipeline_tts_free. +// is forwarded as is. max_batch sizes the KV sets, the bridge columns +// and the maximum concurrent slots (minimum 1). Caller frees with +// pipeline_tts_free. bool pipeline_tts_load(PipelineTTS * pt, const char * talker_gguf_path, const char * codec_gguf_path, BackendPair bp, bool use_fa, - bool clamp_fp16); + bool clamp_fp16, + int max_batch); void pipeline_tts_free(PipelineTTS * pt); @@ -201,3 +217,49 @@ qt_status pipeline_tts_synthesize(PipelineTTS * pt, // rate (24000 / TOKENIZER_HOP_LENGTH). Clamps to a // minimum of one frame. int pipeline_tts_duration_sec_to_tokens(const PipelineTTS * pt, float duration_sec); + +// One synthesis request driven by the batch engine. The caller owns +// params / out for the whole lifetime of the job; status and error +// fill at retirement. error carries the qt_last_error() text captured +// on the thread that ran the engine, so a scheduler on a worker thread +// can replay it into the caller's thread local slot. done is reserved +// for the owner's completion signaling; the engine never touches it. +struct TtsJob { + const struct qt_tts_params * params; + int64_t resolved_seed; + struct qt_audio * out; + qt_status status; + std::string error; + bool done; +}; + +// Batch engine: drives up to pt->max_batch concurrent synthesis slots +// in lockstep over the batched talker decode and code predictor +// graphs. Slots always occupy KV sets [0, N); a retirement compacts +// the range with one device side set copy so the batched views stay +// consecutive. Single threaded: every call runs on the thread that +// owns the GPU. pipeline_tts_synthesize drives a transient engine +// synchronously for the one request case; the facade scheduler keeps a +// long lived one on its worker thread. +struct TtsEngine; + +TtsEngine * tts_engine_new(PipelineTTS * pt, BPETokenizer * tok); +void tts_engine_free(TtsEngine * e); + +// Admit one job: prompt assembly, reference handling and talker +// prefill into the next free slot (KV set N). Returns false when the +// admit fails; job->status and job->error are then final and the job +// never occupies a slot. Logs the prefill stall the join imposes on +// the already active slots. +bool tts_engine_admit(TtsEngine * e, TtsJob * job); + +// Run one frame for every active slot: batched talker decode for the +// slots past their prefill, per slot c0 sampling, batched code +// predictor, per slot codec streaming, then retirement of finished +// slots (EOS, max_new_tokens, cancel, error) including their buffered +// codec decode or streaming drain. Retired jobs append to *retired +// with status, error and out final. +void tts_engine_step(TtsEngine * e, std::vector * retired); + +// Number of currently occupied slots. +int tts_engine_active(const TtsEngine * e); diff --git a/src/qwen.cpp b/src/qwen.cpp index 1fb45ff..7e34bad 100644 --- a/src/qwen.cpp +++ b/src/qwen.cpp @@ -25,24 +25,49 @@ #include "version.h" #include +#include #include #include #include #include +#include +#include #include #include #include #include +#include #include // Internal definition of the opaque handle. C++ types are fine here // because nothing in this struct ever crosses the public ABI boundary : // callers only ever see `struct qt_context *`. PipelineTTS already // embeds the PipelineCodec, so no separate codec field is needed. +// +// gpu_mu serializes every GPU touching entry: the batch worker holds +// it per engine phase (admit, frame step) so a voice extraction or a +// max_batch == 1 synthesize slips between frames instead of racing the +// backend from another thread. +// +// The scheduler members drive the max_batch > 1 mode: qt_synthesize +// enqueues a TtsJob and blocks on cv_done until the worker retires it. +// The worker owns the long lived TtsEngine, admits queued jobs into +// free slots between frames and steps the batch until both the queue +// and the slots drain. struct qt_context { BackendPair bp; PipelineTTS pt; BPETokenizer tok; + + std::mutex gpu_mu; + + int max_batch = 1; + std::thread worker; + std::mutex mu; + std::condition_variable cv_work; + std::condition_variable cv_done; + std::deque queue; + bool stop = false; }; // Thread-local backing store for qt_last_error(). std::string sized once @@ -193,6 +218,7 @@ void qt_init_default_params(struct qt_init_params * p) { p->codec_path = nullptr; p->use_fa = true; p->clamp_fp16 = false; + p->max_batch = 1; } void qt_tts_default_params(struct qt_tts_params * p) { @@ -236,6 +262,62 @@ int qt_num_codebooks(const struct qt_context * q) { return q->pt.num_code_groups; } +// Batch worker: owns the long lived TtsEngine and the GPU hot loop. +// Sleeps until work arrives, then admits queued jobs into free slots +// between frames and steps the batch until both the queue and the +// slots drain. Every GPU touching phase runs under gpu_mu so a voice +// extraction or another entry can slip between frames from its own +// thread. Retired jobs get their done flag under mu and a cv_done +// broadcast so the blocked qt_synthesize callers wake. +static void qt_batch_worker(qt_context * q) { + TtsEngine * e = tts_engine_new(&q->pt, &q->tok); + std::vector retired; + std::unique_lock lk(q->mu); + for (;;) { + q->cv_work.wait(lk, [&] { return q->stop || !q->queue.empty(); }); + if (q->stop && q->queue.empty()) { + break; + } + while (!q->queue.empty() || tts_engine_active(e) > 0) { + // Admit up to max_batch: joins happen at frame boundaries, + // each one stalls the active slots for one prefill. + while (!q->queue.empty() && tts_engine_active(e) < q->max_batch) { + TtsJob * j = q->queue.front(); + q->queue.pop_front(); + lk.unlock(); + bool admitted; + { + std::lock_guard gpu(q->gpu_mu); + admitted = tts_engine_admit(e, j); + } + lk.lock(); + if (!admitted) { + j->done = true; + q->cv_done.notify_all(); + } + } + if (tts_engine_active(e) == 0) { + break; + } + lk.unlock(); + retired.clear(); + { + std::lock_guard gpu(q->gpu_mu); + tts_engine_step(e, &retired); + } + lk.lock(); + for (TtsJob * j : retired) { + j->done = true; + } + if (!retired.empty()) { + q->cv_done.notify_all(); + } + } + } + lk.unlock(); + tts_engine_free(e); +} + struct qt_context * qt_init(const struct qt_init_params * params) { if (!params || !params->talker_path || !params->codec_path) { qt_set_error("qt_init: params, talker_path or codec_path is NULL"); @@ -252,10 +334,14 @@ struct qt_context * qt_init(const struct qt_init_params * params) { qt_log(QT_LOG_INFO, "[Qwen] qwentts.cpp %s", qt_version()); + // ABI v3 tail field: zero init from older callers means 1. + const int max_batch = (params->abi_version >= 3 && params->max_batch > 1) ? params->max_batch : 1; + // new qt_context() value-initialises every field: POD aggregates // (BackendPair, PipelineTTS) are zero-init, std containers in // BPETokenizer construct empty. qt_context * q = new qt_context(); + q->max_batch = max_batch; // The load chain runs inside a try block. Any failure deep in the // GGUF reader, the codec load or the LM weight load throws via @@ -269,7 +355,7 @@ struct qt_context * qt_init(const struct qt_init_params * params) { } if (!pipeline_tts_load(&q->pt, params->talker_path, params->codec_path, q->bp, params->use_fa, - params->clamp_fp16)) { + params->clamp_fp16, max_batch)) { qt_throw("qt_init: pipeline_tts_load failed for '%s' / '%s'", params->talker_path, params->codec_path); } @@ -292,6 +378,13 @@ struct qt_context * qt_init(const struct qt_init_params * params) { return nullptr; } + // Batch mode: one worker thread owns the long lived engine and the + // GPU hot loop; qt_synthesize enqueues and blocks on completion. + if (q->max_batch > 1) { + q->worker = std::thread(qt_batch_worker, q); + qt_log(QT_LOG_INFO, "[Qwen] Batch scheduler started (max_batch=%d)", q->max_batch); + } + return q; } @@ -299,6 +392,14 @@ void qt_free(struct qt_context * q) { if (!q) { return; } + if (q->worker.joinable()) { + { + std::lock_guard lk(q->mu); + q->stop = true; + } + q->cv_work.notify_all(); + q->worker.join(); + } pipeline_tts_free(&q->pt); backend_release(q->bp.backend, q->bp.cpu_backend); delete q; @@ -353,6 +454,10 @@ enum qt_status qt_extract_voice_ref(struct qt_context * q, } try { + // Serialize against the batch worker and any concurrent + // synthesize: the extraction slips between two engine frames. + std::lock_guard gpu(q->gpu_mu); + // Lazy residency: the first reference audio request pays the // weight load once, mirroring the qt_synthesize ref_audio path. if (!q->pt.spk_enc_loaded) { @@ -534,7 +639,39 @@ enum qt_status qt_synthesize(struct qt_context * q, const struct qt_tts_params * // crosses the extern "C" boundary. try { const int64_t resolved_seed = qt_resolve_seed(params->seed); - return pipeline_tts_synthesize(&q->pt, &q->tok, params, resolved_seed, out); + + if (q->max_batch <= 1) { + // Single sequence mode: run synchronously on the calling + // thread under gpu_mu, so concurrent callers serialize FIFO + // and callbacks fire on their own caller's thread. + std::lock_guard gpu(q->gpu_mu); + return pipeline_tts_synthesize(&q->pt, &q->tok, params, resolved_seed, out); + } + + // Batch mode: enqueue and block until the worker retires the + // job. qt_last_error is thread local, so the engine captured + // the worker side message into job.error at retirement; replay + // it into this caller's slot so the errno style contract holds + // across the thread hop. + TtsJob job; + job.params = params; + job.resolved_seed = resolved_seed; + job.out = out; + job.status = QT_STATUS_OK; + job.done = false; + { + std::lock_guard lk(q->mu); + q->queue.push_back(&job); + } + q->cv_work.notify_all(); + { + std::unique_lock lk(q->mu); + q->cv_done.wait(lk, [&] { return job.done; }); + } + if (job.status != QT_STATUS_OK && !job.error.empty()) { + qt_set_error("%s", job.error.c_str()); + } + return job.status; } catch (const std::exception & e) { qt_set_error("%s", e.what()); qt_log(QT_LOG_ERROR, "[Qwen] %s", e.what()); diff --git a/src/qwen.h b/src/qwen.h index 90f16b6..b3a4ae1 100644 --- a/src/qwen.h +++ b/src/qwen.h @@ -57,7 +57,7 @@ extern "C" { // git short hash + commit date string returned by qt_version(); for // binding compat checks, QT_ABI_VERSION is the only number that // matters. -#define QT_ABI_VERSION 2 +#define QT_ABI_VERSION 3 // Returns a static string of the form " ()" identifying // the exact commit this binary was built from. Safe to call from any @@ -121,6 +121,19 @@ struct qt_init_params { const char * codec_path; bool use_fa; bool clamp_fp16; + + // ABI v3. Maximum number of concurrent synthesis requests batched + // on the GPU. 0 and 1 select the single sequence behavior (zero + // init from older callers keeps the previous semantics); values + // above 1 size the KV cache sets accordingly and start an internal + // worker thread that coalesces concurrent qt_synthesize calls into + // batched decode steps, queueing FIFO beyond max_batch. With + // max_batch > 1 the on_chunk and cancel callbacks of every request + // are invoked from that worker thread, not from the calling + // thread; callbacks must be safe to run there and must not call + // back into the qwen_* API. qt_synthesize itself stays blocking + // and thread safe in both modes. + int max_batch; }; // Initialise to the standard defaults: both paths NULL (caller must set @@ -289,7 +302,12 @@ struct qt_tts_params { // the streaming pipeline: audio chunks emit through on_chunk and // `out` stays empty on success. on_chunk NULL keeps the buffered // path. The last chunk on EOS or max_new flushes whatever frames - // remain. + // remain. With qt_init_params.max_batch > 1 the callback runs on + // the internal batch worker thread, not the qt_synthesize caller + // thread: it must be safe there, must not call back into the + // qwen_* API, and a blocking body stalls every batched request, so + // hand the samples to the consumer thread through a queue instead + // of blocking. qt_audio_chunk_cb on_chunk; void * on_chunk_user_data; diff --git a/src/talker-decode-graph.h b/src/talker-decode-graph.h index 63467d7..bbcce4d 100644 --- a/src/talker-decode-graph.h +++ b/src/talker-decode-graph.h @@ -1,10 +1,11 @@ #pragma once -// talker-decode-graph.h: one static talker decode graph, built once -// per attention window class (the kv window rounds up in 256 step -// spans) and replayed directly on the backend. Every input re-uploads -// before each replay: positions, kv row, and mask carry the moving -// n_past, the frame code ids and the overlay row carry the previous -// frame, so nothing bakes and the allocator contract holds. +// talker-decode-graph.h: one static batched talker decode graph, built +// per (attention window class, batch width N) and replayed directly on +// the backend. The kv window rounds up in 256 step spans; every input +// re-uploads before each replay: positions, kv rows, and masks carry +// the moving per-slot n_past, the frame code ids and the overlay rows +// carry each slot's previous frame, so nothing bakes and the allocator +// contract holds. The graph always covers KV sets [0, N). #include "ggml-alloc.h" #include "ggml.h" @@ -15,14 +16,17 @@ struct TalkerDecodeGraph { struct ggml_context * ctx = nullptr; struct ggml_cgraph * gf = nullptr; ggml_gallocr_t galloc = nullptr; - struct ggml_tensor * ids_in = nullptr; // [1 + n_acoustic] i32 - struct ggml_tensor * overlay = nullptr; // [hidden, 1] f32 - struct ggml_tensor * pos_in = nullptr; // [1] i32 - struct ggml_tensor * rows_in = nullptr; // [1] i64 - struct ggml_tensor * mask_in = nullptr; // [n_kv_pad, 1] f16 - struct ggml_tensor * logits = nullptr; // [vocab, 1] f32 - std::vector mask; // [n_kv_pad] f16 + struct ggml_tensor * ids_in = nullptr; // [(1 + n_acoustic) * N] i32, group major: g * N + slot + struct ggml_tensor * overlay = nullptr; // [hidden, N] f32 + struct ggml_tensor * pos_in = nullptr; // [N] i32 + struct ggml_tensor * rows_in = nullptr; // [1, 1, N] i64 + struct ggml_tensor * mask_in = nullptr; // [n_kv_pad, 1, 1, N] f16 + struct ggml_tensor * logits = nullptr; // [vocab, N] f32 + std::vector mask; // [n_kv_pad * N] f16 + std::vector pos_data; // [N] host staging + std::vector rows_data; // [N] host staging int n_kv_pad = 0; // window class width, 0 marks an empty slot + int N = 0; // batch width this build covers }; static void talker_decode_graph_free(TalkerDecodeGraph * tg) { @@ -42,5 +46,8 @@ static void talker_decode_graph_free(TalkerDecodeGraph * tg) { tg->mask_in = nullptr; tg->logits = nullptr; tg->mask.clear(); + tg->pos_data.clear(); + tg->rows_data.clear(); tg->n_kv_pad = 0; + tg->N = 0; } diff --git a/src/talker-forward.h b/src/talker-forward.h index eb01174..51b44d6 100644 --- a/src/talker-forward.h +++ b/src/talker-forward.h @@ -253,16 +253,17 @@ static struct ggml_tensor * talker_layer_forward(struct ggml_context * ctx, // Prefill core: builds the graph in the caller owned arena, allocates // through the sched, uploads the raw embedding, runs it and pulls out -// the last position logits. T tokens are appended to the cache +// the last position logits. T tokens are appended to KV set `kv_set` // starting at n_past. The last position hidden copies in graph into -// the caller owned persistent 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. use_fa / clamp_fp16 are forwarded as is to -// every layer. The decode hot path runs on the static graphs below -// instead. +// column `slot` of the caller owned persistent hidden_bridge tensor +// [hidden, max_batch] 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. use_fa / +// clamp_fp16 are forwarded as is to every layer. The decode hot path +// runs on the static batched graphs below instead. static bool talker_forward_core(const TalkerWeights * tw, KVCache * kv, + int kv_set, ggml_backend_sched_t sched, GraphArena * arena, struct ggml_tensor * hidden_bridge, @@ -317,8 +318,9 @@ static bool talker_forward_core(const TalkerWeights * tw, struct ggml_tensor * h = x_in; std::vector 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, rows_in, kv->k[(size_t) l], - kv->v[(size_t) l], T, n_kv_pad, use_flash_attn, clamp_fp16, gf); + h = talker_layer_forward(gctx, tw, tw->layers[(size_t) l], h, pos_in, mask_in, rows_in, + kv_cache_k(kv, kv_set, l), kv_cache_v(kv, kv_set, l), T, n_kv_pad, use_flash_attn, + clamp_fp16, 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) { @@ -340,12 +342,14 @@ static bool talker_forward_core(const TalkerWeights * tw, 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. + // Bridge: the last position hidden copies on device into column + // `kv_set` of the persistent [hidden, max_batch] 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); + struct ggml_tensor * bridge_col = ggml_view_1d(gctx, hidden_bridge, hidden, (size_t) kv_set * hidden_bridge->nb[1]); + struct ggml_tensor * bridge_cpy = ggml_cpy(gctx, h_last, bridge_col); // codec_head: [hidden, vocab]. ggml_mul_mat returns [vocab, T]. struct ggml_tensor * logits = ggml_mul_mat(gctx, tw->codec_head_w, h_final); @@ -455,14 +459,15 @@ static bool talker_forward_core(const TalkerWeights * tw, // Advance the cache write head. The graph already executed the cpy // nodes so positions [n_past, n_past + T) are now populated. The // arena and the sched allocation persist into the next forward. - kv->cur_len = T_full; + kv->cur_len[(size_t) kv_set] = T_full; 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. +// Prefill: reset KV set `kv_set` and write T_ctx positions in one +// shot. input_embed is [T, hidden] f32 row-major. dump_dir may be NULL. static bool talker_forward_prefill(const TalkerWeights * tw, KVCache * kv, + int kv_set, ggml_backend_sched_t sched, GraphArena * arena, struct ggml_tensor * hidden_bridge, @@ -472,21 +477,147 @@ static bool talker_forward_prefill(const TalkerWeights * tw, bool clamp_fp16, const char * dump_dir, TalkerForwardOutput * out) { - kv_cache_reset(kv); + kv_cache_reset(kv, kv_set); 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, arena, hidden_bridge, input_embed, T, 0, use_flash_attn, clamp_fp16, - dump_dir != NULL, dump_dir, out); + return talker_forward_core(tw, kv, kv_set, sched, arena, hidden_bridge, input_embed, T, 0, use_flash_attn, + clamp_fp16, dump_dir != NULL, dump_dir, out); } -// Build one static decode graph for the given attention window. The -// previous frame code ids gather and sum in graph with the overlay -// row on top; the last position hidden copies into hidden_bridge and -// the codec head logits are the single output row. Positions, kv row, -// and mask stay plain inputs re-uploaded before every replay since -// n_past moves each step. +// Build the batched per-layer block over KV sets [0, N), one fresh +// token per set. Projections, norms, RoPE and the MLP run on [*, N] +// with per column math identical to the single sequence layer. The +// attention is per set: fresh K,V write into the 4D cache at the rows +// carried by kv_rows via set_rows, the read views the padded causal +// window with ne3 = N, and the mask kills everything past each set's +// own context. Returns the layer output [hidden, N]. +static struct ggml_tensor * talker_layer_forward_batch(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 * kv_rows, + struct ggml_tensor * k4, + struct ggml_tensor * v4, + int N, + int n_kv_pad, + bool use_flash_attn, + bool clamp_fp16, + 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, batched over the N columns + struct ggml_tensor * q = ggml_mul_mat(ctx, layer.attn.q_proj_w, h); // [n_q_heads*hd, N] + struct ggml_tensor * k = ggml_mul_mat(ctx, layer.attn.k_proj_w, h); // [n_kv*hd, N] + struct ggml_tensor * v = ggml_mul_mat(ctx, layer.attn.v_proj_w, h); // [n_kv*hd, N] + + q = ggml_reshape_3d(ctx, q, hd, n_q_heads, N); // [hd, n_q_heads, N] + k = ggml_reshape_3d(ctx, k, hd, n_kv, N); + v = ggml_reshape_3d(ctx, v, hd, n_kv, N); + + // Per-head QK-norm: RMS over hd, then multiply by [hd] gain. + 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: positions [N] map to dim 2 of [hd, heads, N], one + // absolute position per set. + 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 fresh K,V into the 4D cache via set_rows over the N + // consecutive sets: src reshapes [hd, n_kv, N] to [hd, 1, n_kv, N] + // (same layout) and kv_rows [1, 1, N] carries one destination row + // per set, broadcast across the n_kv head dim. The row ids travel + // as data so every step keeps an identical topology and the + // captured CUDA graph replays without an update. + struct ggml_tensor * k_sets = ggml_view_4d(ctx, k4, hd, k4->ne[1], n_kv, N, k4->nb[1], k4->nb[2], k4->nb[3], 0); + struct ggml_tensor * v_sets = ggml_view_4d(ctx, v4, hd, v4->ne[1], n_kv, N, v4->nb[1], v4->nb[2], v4->nb[3], 0); + struct ggml_tensor * k_new = ggml_reshape_4d(ctx, k, hd, 1, n_kv, N); + struct ggml_tensor * v_new = ggml_reshape_4d(ctx, v, hd, 1, n_kv, N); + ggml_build_forward_expand(gf, ggml_set_rows(ctx, k_sets, k_new, kv_rows)); + ggml_build_forward_expand(gf, ggml_set_rows(ctx, v_sets, v_new, kv_rows)); + + // Batched read: [hd, n_kv_pad, n_kv, N] views of the 4D cache. The + // window covers every set's causal context and rounds up so the + // shape stays constant across 256 consecutive decode steps; the + // mask carries neg inf past each set's own context and the cache + // buffer is zero initialized, so the padded tail contributes + // nothing. + struct ggml_tensor * k_batch = ggml_view_4d(ctx, k4, hd, n_kv_pad, n_kv, N, k4->nb[1], k4->nb[2], k4->nb[3], 0); + struct ggml_tensor * v_batch = ggml_view_4d(ctx, v4, hd, n_kv_pad, n_kv, N, v4->nb[1], v4->nb[2], v4->nb[3], 0); + + // Q: [hd, n_q_heads, N] -> [hd, 1, n_q_heads, N] (n_batch=1, ne3=N + // for batched flash_attn). + struct ggml_tensor * q4 = ggml_reshape_4d(ctx, q, hd, 1, n_q_heads, N); + + // Clamp V before attention, same rationale as the prefill block: + // sub Ampere CUDA tensor cores accumulate in FP16 and a V overflow + // corrupts every subsequent attention. + if (clamp_fp16) { + v_batch = ggml_clamp(ctx, v_batch, -65504.0f, 65504.0f); + } + + // Attention: fused flash kernel (set_prec(F32) promotes the + // accumulator) or the manual F32 chain, which broadcasts its + // mul_mat over dims 2 and 3 so the same helper covers 4D. + float scale = 1.0f / sqrtf((float) hd); + struct ggml_tensor * attn; + if (use_flash_attn) { + attn = ggml_flash_attn_ext(ctx, q4, k_batch, v_batch, mask, scale, 0.0f, 0.0f); + ggml_flash_attn_ext_set_prec(attn, GGML_PREC_F32); + } else { + attn = talker_attn_f32(ctx, q4, k_batch, v_batch, mask, scale); + } + + // Output [hd, n_q_heads, 1, N] -> [n_q_heads*hd, N], flatten heads + // for o_proj. + attn = ggml_reshape_2d(ctx, attn, n_q_heads * hd, N); + + struct ggml_tensor * o = ggml_mul_mat(ctx, layer.attn.o_proj_w, attn); + + x = ggml_add(ctx, x, o); + if (clamp_fp16) { + x = ggml_clamp(ctx, x, -65504.0f, 65504.0f); + } + + // 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); + if (clamp_fp16) { + x = ggml_clamp(ctx, x, -65504.0f, 65504.0f); + } + return x; +} + +// Build one static batched decode graph for the given attention window +// and batch width N. Each slot's previous frame code ids gather and +// sum in graph with its overlay row on top; the batch of last position +// hiddens copies into columns [0, N) of hidden_bridge and the codec +// head logits [vocab, N] are the single output. Positions, kv rows, +// and masks stay plain inputs re-uploaded before every replay since +// each slot's n_past moves every step. static bool talker_decode_graph_build(const TalkerWeights * tw, KVCache * kv, ggml_backend_t backend, @@ -494,6 +625,7 @@ static bool talker_decode_graph_build(const TalkerWeights * tw, struct ggml_tensor * const * acoustic_embd, int n_acoustic, int n_kv_pad, + int N, bool use_flash_attn, bool clamp_fp16, TalkerDecodeGraph * tg) { @@ -511,26 +643,29 @@ static bool talker_decode_graph_build(const TalkerWeights * tw, } struct ggml_context * gctx = tg->ctx; - struct ggml_tensor * pos_in = ggml_new_tensor_1d(gctx, GGML_TYPE_I32, 1); - struct ggml_tensor * mask_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F16, n_kv_pad, 1); - struct ggml_tensor * rows_in = ggml_new_tensor_1d(gctx, GGML_TYPE_I64, 1); - struct ggml_tensor * ids_in = ggml_new_tensor_1d(gctx, GGML_TYPE_I32, 1 + n_acoustic); - struct ggml_tensor * overlay = ggml_new_tensor_2d(gctx, GGML_TYPE_F32, hidden, 1); + struct ggml_tensor * pos_in = ggml_new_tensor_1d(gctx, GGML_TYPE_I32, N); + struct ggml_tensor * mask_in = ggml_new_tensor_4d(gctx, GGML_TYPE_F16, n_kv_pad, 1, 1, N); + struct ggml_tensor * rows_in = ggml_new_tensor_3d(gctx, GGML_TYPE_I64, 1, 1, N); + struct ggml_tensor * ids_in = ggml_new_tensor_1d(gctx, GGML_TYPE_I32, (1 + n_acoustic) * N); + struct ggml_tensor * overlay = ggml_new_tensor_2d(gctx, GGML_TYPE_F32, hidden, N); ggml_set_name(pos_in, "positions"); ggml_set_name(mask_in, "causal_mask"); ggml_set_name(rows_in, "kv_rows"); ggml_set_name(ids_in, "frame_code_ids"); - ggml_set_name(overlay, "overlay_row"); + ggml_set_name(overlay, "overlay_rows"); ggml_set_input(pos_in); ggml_set_input(mask_in); ggml_set_input(rows_in); ggml_set_input(ids_in); ggml_set_input(overlay); - struct ggml_tensor * id0 = ggml_view_1d(gctx, ids_in, 1, 0); + // ids_in is group major: entry g * N + slot. Each group slices a + // contiguous [N] view and gathers its own table, summed across the + // 16 codebooks, plus the per slot overlay row. + struct ggml_tensor * id0 = ggml_view_1d(gctx, ids_in, N, 0); struct ggml_tensor * x_in = ggml_get_rows(gctx, tw->codec_embedding, id0); for (int g = 0; g < n_acoustic; g++) { - struct ggml_tensor * idg = ggml_view_1d(gctx, ids_in, 1, (size_t) (g + 1) * sizeof(int32_t)); + struct ggml_tensor * idg = ggml_view_1d(gctx, ids_in, N, (size_t) (g + 1) * (size_t) N * sizeof(int32_t)); x_in = ggml_add(gctx, x_in, ggml_get_rows(gctx, acoustic_embd[g], idg)); } x_in = ggml_add(gctx, x_in, overlay); @@ -540,16 +675,17 @@ static bool talker_decode_graph_build(const TalkerWeights * tw, struct ggml_tensor * h = x_in; for (int l = 0; l < n_layers; 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], 1, n_kv_pad, use_flash_attn, clamp_fp16, gf); + h = talker_layer_forward_batch(gctx, tw, tw->layers[(size_t) l], h, pos_in, mask_in, rows_in, + kv->k4[(size_t) l], kv->v4[(size_t) l], N, n_kv_pad, use_flash_attn, clamp_fp16, + gf); } 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"); - struct ggml_tensor * h_last = ggml_view_1d(gctx, h_final, hidden, 0); - struct ggml_tensor * bridge_cpy = ggml_cpy(gctx, h_last, hidden_bridge); + struct ggml_tensor * bridge_cols = ggml_view_2d(gctx, hidden_bridge, hidden, N, hidden_bridge->nb[1], 0); + struct ggml_tensor * bridge_cpy = ggml_cpy(gctx, h_final, bridge_cols); struct ggml_tensor * logits = ggml_mul_mat(gctx, tw->codec_head_w, h_final); ggml_set_name(logits, "logits"); @@ -571,61 +707,83 @@ static bool talker_decode_graph_build(const TalkerWeights * tw, tg->rows_in = rows_in; tg->mask_in = mask_in; tg->logits = logits; - tg->mask.resize((size_t) n_kv_pad); + tg->mask.resize((size_t) n_kv_pad * (size_t) N); + tg->pos_data.resize((size_t) N); + tg->rows_data.resize((size_t) N); tg->n_kv_pad = n_kv_pad; + tg->N = N; return true; } -// Decode: append one position from the previous frame's codes over the -// static graph of the current window class, built lazily on the first -// step entering the span. frame_ids holds [c0, c1..c15], acoustic_embd -// the 15 group tables owned by the code predictor, overlay the trailing -// text / pad row summed on top. Reads positions [0, kv->cur_len + 1); -// caller ensures kv->cur_len + 1 <= kv->max_seq_len holds by cache -// sizing. read_hidden_host pulls the bridge back for dump paths. -static bool talker_forward_decode(const TalkerWeights * tw, - KVCache * kv, - ggml_backend_t backend, - TalkerDecodeGraph * graphs, - struct ggml_tensor * hidden_bridge, - const int32_t * frame_ids, - struct ggml_tensor * const * acoustic_embd, - int n_acoustic, - const float * overlay, - bool use_flash_attn, - bool clamp_fp16, - bool read_hidden_host, - 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; +// Batched decode: append one position per KV set over the static graph +// of the current window class and batch width, built lazily on the +// first step entering the (span, N) pair. frame_ids holds the previous +// frame codes of every slot, group major: entry g * N + slot. +// acoustic_embd holds the 15 group tables owned by the code predictor, +// overlays the [hidden, N] trailing text / pad rows summed on top. The +// window covers max over sets of (cur_len + 1); caller ensures every +// cur_len + 1 <= max_seq_len holds by cache sizing. Logits fill +// out->logits_last as [vocab, N] column blocks; read_hidden_host pulls +// the bridge back as [hidden, N] for dump paths. +static bool talker_forward_decode(const TalkerWeights * tw, + KVCache * kv, + ggml_backend_t backend, + std::vector & graphs, + struct ggml_tensor * hidden_bridge, + const int32_t * frame_ids, + struct ggml_tensor * const * acoustic_embd, + int n_acoustic, + const float * overlays, + int N, + bool use_flash_attn, + bool clamp_fp16, + bool read_hidden_host, + TalkerForwardOutput * out) { + int max_len = 0; + for (int i = 0; i < N; i++) { + const int len = kv->cur_len[(size_t) i] + 1; + if (len > kv->max_seq_len) { + fprintf(stderr, "[TalkerForward] FATAL: decode would overflow cache (%d > %d, set %d)\n", len, + kv->max_seq_len, i); + return false; + } + if (len > max_len) { + max_len = len; + } } - const int n_past = kv->cur_len; - const int kv_pad_raw = (int) GGML_PAD(n_past + 1, 256); + const int kv_pad_raw = (int) GGML_PAD(max_len, 256); const int n_kv_pad = kv_pad_raw < kv->max_seq_len ? kv_pad_raw : kv->max_seq_len; - TalkerDecodeGraph * tg = &graphs[(n_kv_pad + 255) / 256 - 1]; - if (!tg->ctx && !talker_decode_graph_build(tw, kv, backend, hidden_bridge, acoustic_embd, n_acoustic, n_kv_pad, + TalkerDecodeGraph * tg = &graphs[(size_t) ((n_kv_pad + 255) / 256 - 1)]; + if (tg->ctx && tg->N != N) { + talker_decode_graph_free(tg); + } + if (!tg->ctx && !talker_decode_graph_build(tw, kv, backend, hidden_bridge, acoustic_embd, n_acoustic, n_kv_pad, N, use_flash_attn, clamp_fp16, tg)) { return false; } - ggml_backend_tensor_set(tg->ids_in, frame_ids, 0, (size_t) (1 + n_acoustic) * sizeof(int32_t)); - ggml_backend_tensor_set(tg->overlay, overlay, 0, (size_t) tw->hidden_size * sizeof(float)); + ggml_backend_tensor_set(tg->ids_in, frame_ids, 0, (size_t) (1 + n_acoustic) * (size_t) N * sizeof(int32_t)); + ggml_backend_tensor_set(tg->overlay, overlays, 0, (size_t) tw->hidden_size * (size_t) N * sizeof(float)); - const int32_t pos = n_past; - ggml_backend_tensor_set(tg->pos_in, &pos, 0, sizeof(int32_t)); - const int64_t row = (int64_t) n_past; - ggml_backend_tensor_set(tg->rows_in, &row, 0, sizeof(int64_t)); + for (int i = 0; i < N; i++) { + tg->pos_data[(size_t) i] = kv->cur_len[(size_t) i]; + tg->rows_data[(size_t) i] = (int64_t) kv->cur_len[(size_t) i]; + } + ggml_backend_tensor_set(tg->pos_in, tg->pos_data.data(), 0, (size_t) N * sizeof(int32_t)); + ggml_backend_tensor_set(tg->rows_in, tg->rows_data.data(), 0, (size_t) N * sizeof(int64_t)); - // Causal mask: keys [0, n_past] carry 0, the padded tail neg inf. + // Causal masks: for slot i keys [0, cur_len[i]] carry 0, the rest + // of the padded window neg inf. { std::vector & mask = tg->mask; 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] = (int) i <= n_past ? zero : neg_inf; + for (int i = 0; i < N; i++) { + const int n_past = kv->cur_len[(size_t) i]; + for (int j = 0; j < n_kv_pad; j++) { + mask[(size_t) i * (size_t) n_kv_pad + (size_t) j] = j <= n_past ? zero : neg_inf; + } } ggml_backend_tensor_set(tg->mask_in, mask.data(), 0, mask.size() * sizeof(ggml_fp16_t)); } @@ -637,13 +795,17 @@ static bool talker_forward_decode(const TalkerWeights * tw, out->hidden = tw->hidden_size; out->vocab = tw->vocab_size; - out->logits_last.assign((size_t) tw->vocab_size, 0.0f); - ggml_backend_tensor_get(tg->logits, out->logits_last.data(), 0, (size_t) tw->vocab_size * sizeof(float)); + out->logits_last.assign((size_t) tw->vocab_size * (size_t) N, 0.0f); + ggml_backend_tensor_get(tg->logits, out->logits_last.data(), 0, + (size_t) tw->vocab_size * (size_t) N * sizeof(float)); if (read_hidden_host) { - out->hidden_last.assign((size_t) tw->hidden_size, 0.0f); - ggml_backend_tensor_get(hidden_bridge, out->hidden_last.data(), 0, (size_t) tw->hidden_size * sizeof(float)); + out->hidden_last.assign((size_t) tw->hidden_size * (size_t) N, 0.0f); + ggml_backend_tensor_get(hidden_bridge, out->hidden_last.data(), 0, + (size_t) tw->hidden_size * (size_t) N * sizeof(float)); } - kv->cur_len = n_past + 1; + for (int i = 0; i < N; i++) { + kv->cur_len[(size_t) i]++; + } return true; } diff --git a/src/tokenizer-transformer.h b/src/tokenizer-transformer.h index 060ee78..e9359e7 100644 --- a/src/tokenizer-transformer.h +++ b/src/tokenizer-transformer.h @@ -393,8 +393,8 @@ static struct ggml_tensor * tok_trans_forward_stream(struct ggml_context * h = ggml_add(ctx, h, tr->input_proj_b); for (int l = 0; l < tr->num_layers; l++) { - h = tok_trans_layer_forward_stream(ctx, gf, tr, tr->layers[l], h, positions, mask, kv_rows, kv->k[(size_t) l], - kv->v[(size_t) l], T, ring); + h = tok_trans_layer_forward_stream(ctx, gf, tr, tr->layers[l], h, positions, mask, kv_rows, + kv_cache_k(kv, 0, l), kv_cache_v(kv, 0, l), T, ring); } h = ggml_rms_norm(ctx, h, tr->rms_norm_eps); diff --git a/src/tts-server.h b/src/tts-server.h index 1b0bd5b..e68b941 100644 --- a/src/tts-server.h +++ b/src/tts-server.h @@ -26,14 +26,18 @@ #include "audio-io.h" #include "yyjson.h" +#include #include #include +#include #include #include #include #include +#include #include #include +#include #include // One synthesis request parsed from the OAI JSON body. @@ -95,8 +99,9 @@ struct server_config { int port = 8080; }; -// Single GPU context : synthesis is serialised FIFO across connections. -static std::mutex g_synth_mutex; +// Concurrency lives behind the ABI: qt_synthesize is thread safe and +// batches concurrent requests when the context was initialised with +// max_batch > 1, so connection threads call the backend directly. static httplib::Server * g_svr = nullptr; static void tts_on_signal(int) { @@ -252,18 +257,18 @@ static void tts_handle_speech(const tts_backend & be, const httplib::Request & h } if (req.format == "wav") { - // One-shot : collect the whole utterance, then emit a RIFF file. + // One-shot : collect the whole utterance, then emit a RIFF + // file. The backend call blocks this connection thread; in + // batch mode the chunks arrive from the ABI worker thread while + // this thread waits inside qt_synthesize, so the plain append + // stays single producer and the blocking return synchronizes. std::vector buf; tts_sink sink = [&buf](const float * s, int n) { buf.insert(buf.end(), s, s + n); return true; }; std::string synth_err; - int rc; - { - std::lock_guard lock(g_synth_mutex); - rc = be.synthesize(req, sink, synth_err); - } + int rc = be.synthesize(req, sink, synth_err); if (rc != 0) { tts_json_error(res, tts_status_to_http(rc), "server_error", synth_err.empty() ? "synthesis failed" : synth_err.c_str()); @@ -274,26 +279,71 @@ static void tts_handle_speech(const tts_backend & be, const httplib::Request & h return; } - // Streaming : run synthesis inside the chunked provider on the connection - // thread, pushing s16le frames as the codec produces them. A failed - // sink.write means the client disconnected, which aborts generation and - // frees the GPU instead of finishing a stream nobody reads. - res.set_header("Cache-Control", "no-cache"); - res.set_header("X-Accel-Buffering", "no"); - res.set_chunked_content_provider("audio/pcm", [&be, req](size_t, httplib::DataSink & sink) mutable -> bool { - tts_sink push = [&sink](const float * s, int n) { + // Streaming : the synthesis runs on its own thread pushing s16le + // bytes into a queue; the connection thread drains the queue into + // the chunked sink. The decoupling keeps a slow client from + // stalling the ABI worker (head of line blocking across the whole + // batch), and a client disconnect flips client_gone so the next + // chunk callback aborts generation and frees the GPU instead of + // finishing a stream nobody reads. Backpressure is the utterance + // itself: pending grows at most to the full PCM of one synthesis. + struct stream_state { + std::mutex mu; + std::condition_variable cv; + std::string pending; + bool done = false; + std::atomic client_gone{ false }; + std::thread th; + }; + + auto st = std::make_shared(); + + st->th = std::thread([&be, req, st]() { + tts_sink push = [st](const float * s, int n) { + if (st->client_gone.load(std::memory_order_acquire)) { + return false; + } std::string bytes; tts_append_s16le(bytes, s, n); - return sink.write(bytes.data(), bytes.size()); + std::lock_guard lk(st->mu); + st->pending += bytes; + st->cv.notify_all(); + return true; }; std::string synth_err; - { - std::lock_guard lock(g_synth_mutex); - be.synthesize(req, push, synth_err); - } - sink.done(); - return true; + be.synthesize(req, push, synth_err); + std::lock_guard lk(st->mu); + st->done = true; + st->cv.notify_all(); }); + + res.set_header("Cache-Control", "no-cache"); + res.set_header("X-Accel-Buffering", "no"); + res.set_chunked_content_provider( + "audio/pcm", + [st](size_t, httplib::DataSink & sink) -> bool { + std::string chunk; + { + std::unique_lock lk(st->mu); + st->cv.wait(lk, [&] { return st->done || !st->pending.empty(); }); + chunk.swap(st->pending); + } + if (!chunk.empty()) { + if (!sink.write(chunk.data(), chunk.size())) { + st->client_gone.store(true, std::memory_order_release); + return false; + } + return true; + } + sink.done(); + return true; + }, + [st](bool) { + st->client_gone.store(true, std::memory_order_release); + if (st->th.joinable()) { + st->th.join(); + } + }); } // Decode standard base64 (with optional padding) into out. Returns diff --git a/tools/tts-server.cpp b/tools/tts-server.cpp index a9af221..981297a 100644 --- a/tools/tts-server.cpp +++ b/tools/tts-server.cpp @@ -30,8 +30,11 @@ struct voice_entry { }; // Registered voices, name keyed. Every access happens under -// g_synth_mutex: registration touches the GPU through the extraction -// path and lookups run inside the already serialized synthesize. +// g_voices_mutex; the GPU side of a registration is serialized inside +// the ABI (qt_extract_voice_ref slips between batch frames), and the +// synthesize lookup copies the latents out so a concurrent replace or +// delete never frees buffers a running synthesis still reads. +static std::mutex g_voices_mutex; static std::unordered_map g_voices; static void print_usage(const char * prog) { @@ -46,6 +49,7 @@ static void print_usage(const char * prog) { " --host Listen address (default: 127.0.0.1)\n" " --port Listen port (default: 8080)\n" " --lang Language label (default: auto)\n" + " --max-batch Concurrent requests batched on the GPU (default: 1)\n" " --no-fa Disable flash attention\n" " --clamp-fp16 Clamp hidden states to FP16 range\n", prog); @@ -66,6 +70,7 @@ int main(int argc, char ** argv) { server_config cfg; bool use_fa = true; bool clamp_fp16 = false; + int max_batch = 1; for (int i = 1; i < argc; i++) { const char * arg = argv[i]; @@ -85,6 +90,8 @@ int main(int argc, char ** argv) { use_fa = false; } else if (!std::strcmp(arg, "--clamp-fp16")) { clamp_fp16 = true; + } else if (!std::strcmp(arg, "--max-batch") && i + 1 < argc) { + max_batch = std::atoi(argv[++i]); } else if (!std::strcmp(arg, "--help") || !std::strcmp(arg, "-h")) { print_usage(argv[0]); return 0; @@ -106,6 +113,7 @@ int main(int argc, char ** argv) { iparams.codec_path = codec_path; iparams.use_fa = use_fa; iparams.clamp_fp16 = clamp_fp16; + iparams.max_batch = max_batch; struct qt_context * q = qt_init(&iparams); if (!q) { @@ -136,11 +144,7 @@ int main(int argc, char ** argv) { err = "cannot decode the WAV payload"; return false; } - enum qt_status rc; - { - std::lock_guard lock(g_synth_mutex); - rc = qt_extract_voice_ref(q, pcm, T, &entry.ref); - } + enum qt_status rc = qt_extract_voice_ref(q, pcm, T, &entry.ref); free(pcm); if (rc != QT_STATUS_OK) { err = qt_last_error(); @@ -167,7 +171,7 @@ int main(int argc, char ** argv) { std::memcpy(entry.ref.ref_codes, codes.data(), codes.size() * sizeof(int32_t)); } - std::lock_guard lock(g_synth_mutex); + std::lock_guard lock(g_voices_mutex); auto it = g_voices.find(up.name); if (it != g_voices.end()) { qt_voice_ref_free(&it->second.ref); @@ -180,7 +184,7 @@ int main(int argc, char ** argv) { }; be.remove_voice = [](const std::string & name) -> bool { - std::lock_guard lock(g_synth_mutex); + std::lock_guard lock(g_voices_mutex); auto it = g_voices.find(name); if (it == g_voices.end()) { return false; @@ -191,7 +195,7 @@ int main(int argc, char ** argv) { }; be.registered_voices = []() -> std::vector { - std::lock_guard lock(g_synth_mutex); + std::lock_guard lock(g_voices_mutex); std::vector names; names.reserve(g_voices.size()); for (const auto & kv : g_voices) { @@ -214,15 +218,36 @@ int main(int argc, char ** argv) { p.text = req.input.c_str(); p.lang = lang.c_str(); - auto vit = req.voice.empty() ? g_voices.end() : g_voices.find(req.voice); - if (vit != g_voices.end()) { - const voice_entry & v = vit->second; - p.ref_spk_emb = v.ref.ref_spk_emb; - p.ref_spk_dim = v.ref.ref_spk_dim; - if (!v.ref_text.empty() && v.ref.ref_codes) { - p.ref_codes = v.ref.ref_codes; - p.ref_T = v.ref.ref_T; - p.ref_text = v.ref_text.c_str(); + // Copy the registered voice latents out under the lock: the + // synthesis may run for seconds while another connection + // replaces or deletes the entry. + std::vector voice_spk; + std::vector voice_codes; + std::string voice_ref_text; + int voice_ref_T = 0; + bool have_voice = false; + if (!req.voice.empty()) { + std::lock_guard lock(g_voices_mutex); + auto vit = g_voices.find(req.voice); + if (vit != g_voices.end()) { + const voice_entry & v = vit->second; + have_voice = true; + voice_spk.assign(v.ref.ref_spk_emb, v.ref.ref_spk_emb + v.ref.ref_spk_dim); + if (!v.ref_text.empty() && v.ref.ref_codes) { + voice_codes.assign(v.ref.ref_codes, + v.ref.ref_codes + (size_t) v.ref.num_codebooks * (size_t) v.ref.ref_T); + voice_ref_T = v.ref.ref_T; + voice_ref_text = v.ref_text; + } + } + } + if (have_voice) { + p.ref_spk_emb = voice_spk.data(); + p.ref_spk_dim = (int) voice_spk.size(); + if (!voice_codes.empty()) { + p.ref_codes = voice_codes.data(); + p.ref_T = voice_ref_T; + p.ref_text = voice_ref_text.c_str(); } } else if (!req.voice.empty() && qt_n_speakers(q) > 0) { p.speaker = req.voice.c_str();