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

codec: adaptive chunk width on the streaming decode

Browse Source 
The stream graphs come in width classes T in {1, 2, 4, 8} sharing the
state and KV ring tensors, built lazily per class. The decoder ramps
1 -> 2 -> 4 -> 8 so the first frame keeps its latency while the steady
state interleaves the codec 8x less often and the batch amortizes the
kernel count; drain flushes the tail at EOS and the reference priming
runs in max width chunks.
This commit is contained in:
Pascal
2026-07-06 22:43:52 +02:00
parent 206ad06fde
commit 0725d2e53b
4 changed files with 261 additions and 149 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/codec-chunked-decode.h
+89 -27
View File
@@ -84,66 +84,128 @@ static inline std::vector<float> codec_chunked_decode(PipelineCodec * pc,
return out;
}
// Stateful streaming decoder over pipeline_codec_decode_stream: every
// pushed frame decodes immediately through the persistent codec state
// and emits its TOKENIZER_HOP_LENGTH samples, no buffering, no left
// context re-decode, no tail to drain at EOS. ICL priming feeds the
// full reference through the same state with the audio discarded,
// which matches the upstream reference plus generated decode exactly.
// Stateful streaming decoder over pipeline_codec_decode_stream with an
// adaptive chunk ramp: the first pushed frame decodes immediately for
// the lowest first byte latency, then chunks grow 2 -> 4 -> 8 frames so
// the steady state interleaves the codec with the talker 8x less often
// and the batch amortizes the kernel count. drain flushes the sub chunk
// tail at EOS with greedy width classes. ICL priming feeds the full
// reference through the same state in max width chunks with the audio
// discarded.
struct codec_stream_decoder {
int K;
// Set true when push_frame returned false because the on_chunk
// Set true when a flush returned false because the on_chunk
// callback requested a cancel. Stays false on decode failures so
// the caller can route to QT_STATUS_CANCELLED vs
// QT_STATUS_GENERATE_FAILED on a negative return.
bool cancelled;
std::vector<float> frame;
static const int MAX_CHUNK = 1 << (CODEC_STREAM_CLASSES - 1);
// Reset the persistent codec state to the zero context. Returns
// false when the state allocation fails.
int target; // current ramp chunk width
int pending_n; // frames accumulated, < target
std::vector<int32_t> pending; // [MAX_CHUNK, K] frame major
std::vector<int32_t> scratch; // [T, K] chunk upload layout
std::vector<float> frame; // [MAX_CHUNK * hop] audio out
// Reset the persistent codec state to the zero context and restart
// the chunk ramp. Returns false when the state allocation fails.
bool init(PipelineCodec * pc, int K_) {
K = K_;
cancelled = false;
frame.assign((size_t) TOKENIZER_HOP_LENGTH, 0.0f);
target = 1;
pending_n = 0;
pending.assign((size_t) MAX_CHUNK * (size_t) K, 0);
scratch.assign((size_t) MAX_CHUNK * (size_t) K, 0);
frame.assign((size_t) MAX_CHUNK * (size_t) TOKENIZER_HOP_LENGTH, 0.0f);
return pipeline_codec_stream_reset(pc);
}
// Decode and emit the first T pending frames, then compact the
// remainder to the front. The scratch reorders frame major pending
// rows into the [T, K] chunk layout (T contiguous per codebook).
bool flush_front(PipelineCodec * pc, int T, qt_audio_chunk_cb cb, void * cb_ud) {
for (int k = 0; k < K; k++) {
for (int t = 0; t < T; t++) {
scratch[(size_t) k * (size_t) T + (size_t) t] = pending[(size_t) t * (size_t) K + (size_t) k];
}
}
if (!pipeline_codec_decode_stream(pc, scratch.data(), T, frame.data())) {
return false;
}
if (cb && !cb(frame.data(), T * TOKENIZER_HOP_LENGTH, cb_ud)) {
cancelled = true;
return false;
}
pending_n -= T;
if (pending_n > 0) {
std::memmove(pending.data(), pending.data() + (size_t) T * (size_t) K,
(size_t) pending_n * (size_t) K * sizeof(int32_t));
}
return true;
}
// Prime the codec state with the full ICL reference: every frame
// runs through the streaming decode with the audio discarded, so
// the first generated frame sees the reference's exact causal
// state. A reference already primed restores its snapshot device
// to device instead of re-decoding; a fresh one saves its primed
// state into the LRU. ref_kt is K major [K, ref_T]. Call once,
// after init and before any push_frame.
// runs through the streaming decode in max width chunks with the
// audio discarded, so the first generated frame sees the
// reference's exact causal state. A reference already primed
// restores its snapshot device to device instead of re-decoding; a
// fresh one saves its primed state into the LRU. ref_kt is K major
// [K, ref_T]. Call once, after init and before any push_frame.
bool seed_reference(PipelineCodec * pc, const int32_t * ref_kt, int ref_T) {
const uint64_t key = pipeline_codec_ref_key(ref_kt, K, ref_T);
if (pipeline_codec_stream_restore(pc, key)) {
return true;
}
std::vector<int32_t> codes((size_t) K);
for (int t = 0; t < ref_T; t++) {
for (int k = 0; k < K; k++) {
codes[(size_t) k] = ref_kt[(size_t) k * (size_t) ref_T + (size_t) t];
int t0 = 0;
while (t0 < ref_T) {
int T = MAX_CHUNK;
while (T > ref_T - t0) {
T >>= 1;
}
if (!pipeline_codec_decode_stream(pc, codes.data(), NULL)) {
for (int k = 0; k < K; k++) {
for (int t = 0; t < T; t++) {
scratch[(size_t) k * (size_t) T + (size_t) t] =
ref_kt[(size_t) k * (size_t) ref_T + (size_t) (t0 + t)];
}
}
if (!pipeline_codec_decode_stream(pc, scratch.data(), T, NULL)) {
return false;
}
t0 += T;
}
return pipeline_codec_stream_snapshot(pc, key);
}
// Decode one frame (K int32 codes, one per codebook) and emit its
// samples through the callback. Returns false on decode failure or
// Accumulate one frame (K int32 codes, one per codebook) and decode
// when the ramp chunk fills. Returns false on decode failure or
// when cb returns false (cancellation).
bool push_frame(PipelineCodec * pc, const int32_t * frame_codes, qt_audio_chunk_cb cb, void * cb_ud) {
if (!pipeline_codec_decode_stream(pc, frame_codes, frame.data())) {
std::memcpy(pending.data() + (size_t) pending_n * (size_t) K, frame_codes, (size_t) K * sizeof(int32_t));
pending_n++;
if (pending_n < target) {
return true;
}
if (!flush_front(pc, target, cb, cb_ud)) {
return false;
}
if (!cb(frame.data(), TOKENIZER_HOP_LENGTH, cb_ud)) {
cancelled = true;
if (target < MAX_CHUNK) {
target <<= 1;
}
return true;
}
// Flush the sub chunk tail at EOS through greedy width classes.
bool drain(PipelineCodec * pc, qt_audio_chunk_cb cb, void * cb_ud) {
while (pending_n > 0) {
int T = MAX_CHUNK;
while (T > pending_n) {
T >>= 1;
}
if (!flush_front(pc, T, cb, cb_ud)) {
return false;
}
}
return true;
}
};
src/pipeline-codec.cpp
+126 -102
View File
@@ -28,9 +28,9 @@ bool pipeline_codec_load(PipelineCodec * pc, const char * gguf_path, BackendPair
pc->stream_ctx = NULL;
pc->stream_buf = NULL;
pc->stream_pos = 0;
pc->stream_graph_ctx = NULL;
pc->stream_gf = NULL;
pc->stream_galloc = NULL;
for (int i = 0; i < CODEC_STREAM_CLASSES; i++) {
pc->stream_graphs[i] = {};
}
for (int i = 0; i < CODEC_SNAP_SLOTS; i++) {
pc->snaps[i] = {};
}
@@ -267,87 +267,6 @@ static bool pipeline_codec_stream_ensure(PipelineCodec * pc) {
return false;
}
// Build the static T=1 frame graph once: fixed topology, fixed
// tensor addresses, computed directly on the backend every frame.
{
const int max_nodes = 4096;
struct ggml_init_params gp2 = {
ggml_tensor_overhead() * (size_t) max_nodes + ggml_graph_overhead_custom(max_nodes, false),
NULL,
true,
};
pc->stream_graph_ctx = ggml_init(gp2);
if (!pc->stream_graph_ctx) {
qt_log(QT_LOG_ERROR, "[Pipeline] stream graph ggml_init failed");
kv_cache_free(&pc->stream_kv);
ggml_backend_buffer_free(pc->stream_buf);
pc->stream_buf = NULL;
ggml_free(pc->stream_ctx);
pc->stream_ctx = NULL;
return false;
}
struct ggml_context * gctx = pc->stream_graph_ctx;
struct ggml_cgraph * gf = ggml_new_graph_custom(gctx, max_nodes, false);
const int ring = CODEC_STREAM_RING;
const int K = TOKENIZER_NUM_CODEBOOKS;
struct ggml_tensor * codes_in = ggml_new_tensor_2d(gctx, GGML_TYPE_I32, 1, K);
struct ggml_tensor * pos_in = ggml_new_tensor_1d(gctx, GGML_TYPE_I32, 1);
struct ggml_tensor * rows_in = ggml_new_tensor_1d(gctx, GGML_TYPE_I64, 1);
struct ggml_tensor * mask_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F32, ring, 1);
ggml_set_name(codes_in, "codes_in");
ggml_set_name(pos_in, "positions");
ggml_set_name(rows_in, "kv_rows");
ggml_set_name(mask_in, "ring_mask");
ggml_set_input(codes_in);
ggml_set_input(pos_in);
ggml_set_input(rows_in);
ggml_set_input(mask_in);
// Same module chain as pipeline_codec_decode with the stateful
// variants threaded through the persistent stream tensors. The
// quantizer uses the alignment safe variant: no scheduler input
// duplication happens on the direct backend compute path.
struct ggml_tensor * h = quant_decode_stream(gctx, &pc->qdec, codes_in); // [512, 1] C-first
h = ggml_cont(gctx, ggml_transpose(gctx, h)); // [1, 512] T-first
h = qwen_causal_conv1d_stream(gctx, gf, pc->pre_conv_w, pc->pre_conv_b, h, 3, 1, pc->stream_pre_conv);
h = ggml_cont(gctx, ggml_transpose(gctx, h)); // [1024, 1] C-first
h = tok_trans_forward_stream(gctx, gf, &pc->transformer, h, pos_in, mask_in, rows_in, &pc->stream_kv);
h = ggml_cont(gctx, ggml_transpose(gctx, h)); // [1, 1024] T-first
h = upsample_stage_forward_stream(gctx, gf, &pc->upsample, h, &pc->stream_up);
h = dac_decoder_forward_stream(gctx, gf, &pc->dac, h, &pc->stream_dac);
h = ggml_clamp(gctx, h, -1.0f, 1.0f);
ggml_set_name(h, "audio_out");
ggml_set_output(h);
ggml_build_forward_expand(gf, h);
pc->stream_galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(pc->backend));
if (!pc->stream_galloc || !ggml_gallocr_alloc_graph(pc->stream_galloc, gf)) {
qt_log(QT_LOG_ERROR, "[Pipeline] stream graph allocation failed");
if (pc->stream_galloc) {
ggml_gallocr_free(pc->stream_galloc);
pc->stream_galloc = NULL;
}
ggml_free(pc->stream_graph_ctx);
pc->stream_graph_ctx = NULL;
kv_cache_free(&pc->stream_kv);
ggml_backend_buffer_free(pc->stream_buf);
pc->stream_buf = NULL;
ggml_free(pc->stream_ctx);
pc->stream_ctx = NULL;
return false;
}
pc->stream_gf = gf;
pc->stream_in_codes = codes_in;
pc->stream_in_pos = pos_in;
pc->stream_in_rows = rows_in;
pc->stream_in_mask = mask_in;
pc->stream_out = h;
}
pc->stream_ready = true;
qt_log(QT_LOG_INFO, "[Pipeline] Codec stream state ready: %d conv contexts, KV ring %d", n_state - 4,
CODEC_STREAM_RING);
@@ -470,35 +389,135 @@ bool pipeline_codec_stream_snapshot(PipelineCodec * pc, uint64_t key) {
return true;
}
bool pipeline_codec_decode_stream(PipelineCodec * pc, const int32_t * codes, float * audio_out) {
// 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
// sized for T.
static bool codec_stream_graph_ensure(PipelineCodec * pc, int cls) {
CodecStreamGraph * sg = &pc->stream_graphs[cls];
if (sg->ctx) {
return true;
}
const int T = 1 << cls;
const int ring = CODEC_STREAM_RING;
const int K = TOKENIZER_NUM_CODEBOOKS;
// The ring must hold the sliding window plus the whole fresh chunk.
if (pc->transformer.sliding_window + T > ring) {
qt_log(QT_LOG_ERROR, "[Pipeline] sliding window %d plus chunk %d exceeds KV ring %d",
pc->transformer.sliding_window, T, ring);
return false;
}
const int max_nodes = 4096;
struct ggml_init_params gp = {
ggml_tensor_overhead() * (size_t) max_nodes + ggml_graph_overhead_custom(max_nodes, false),
NULL,
true,
};
sg->ctx = ggml_init(gp);
if (!sg->ctx) {
qt_log(QT_LOG_ERROR, "[Pipeline] stream graph ggml_init failed (T=%d)", T);
return false;
}
struct ggml_context * gctx = sg->ctx;
struct ggml_cgraph * gf = ggml_new_graph_custom(gctx, max_nodes, false);
struct ggml_tensor * codes_in = ggml_new_tensor_2d(gctx, GGML_TYPE_I32, T, K);
struct ggml_tensor * pos_in = ggml_new_tensor_1d(gctx, GGML_TYPE_I32, T);
struct ggml_tensor * rows_in = ggml_new_tensor_1d(gctx, GGML_TYPE_I64, T);
struct ggml_tensor * mask_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F32, ring, T);
ggml_set_name(codes_in, "codes_in");
ggml_set_name(pos_in, "positions");
ggml_set_name(rows_in, "kv_rows");
ggml_set_name(mask_in, "ring_mask");
ggml_set_input(codes_in);
ggml_set_input(pos_in);
ggml_set_input(rows_in);
ggml_set_input(mask_in);
// Same module chain as pipeline_codec_decode with the stateful
// variants threaded through the persistent stream tensors. The
// quantizer uses the alignment safe variant: no scheduler input
// duplication happens on the direct backend compute path.
struct ggml_tensor * h = quant_decode_stream(gctx, &pc->qdec, codes_in); // [512, T] C-first
h = ggml_cont(gctx, ggml_transpose(gctx, h)); // [T, 512] T-first
h = qwen_causal_conv1d_stream(gctx, gf, pc->pre_conv_w, pc->pre_conv_b, h, 3, 1, pc->stream_pre_conv);
h = ggml_cont(gctx, ggml_transpose(gctx, h)); // [1024, T] C-first
h = tok_trans_forward_stream(gctx, gf, &pc->transformer, h, pos_in, mask_in, rows_in, &pc->stream_kv);
h = ggml_cont(gctx, ggml_transpose(gctx, h)); // [T, 1024] T-first
h = upsample_stage_forward_stream(gctx, gf, &pc->upsample, h, &pc->stream_up);
h = dac_decoder_forward_stream(gctx, gf, &pc->dac, h, &pc->stream_dac);
h = ggml_clamp(gctx, h, -1.0f, 1.0f);
ggml_set_name(h, "audio_out");
ggml_set_output(h);
ggml_build_forward_expand(gf, h);
sg->galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(pc->backend));
if (!sg->galloc || !ggml_gallocr_alloc_graph(sg->galloc, gf)) {
qt_log(QT_LOG_ERROR, "[Pipeline] stream graph allocation failed (T=%d)", T);
if (sg->galloc) {
ggml_gallocr_free(sg->galloc);
}
ggml_free(sg->ctx);
*sg = {};
return false;
}
sg->gf = gf;
sg->codes = codes_in;
sg->pos = pos_in;
sg->rows = rows_in;
sg->mask = mask_in;
sg->out = h;
return true;
}
bool pipeline_codec_decode_stream(PipelineCodec * pc, const int32_t * codes, int T, float * audio_out) {
const int K = TOKENIZER_NUM_CODEBOOKS;
const int ring = CODEC_STREAM_RING;
ggml_backend_tensor_set(pc->stream_in_codes, codes, 0, (size_t) K * sizeof(int32_t));
int cls = 0;
while ((1 << cls) < T) {
cls++;
}
if (cls >= CODEC_STREAM_CLASSES || (1 << cls) != T) {
qt_log(QT_LOG_ERROR, "[Pipeline] invalid stream chunk width %d", T);
return false;
}
if (!codec_stream_graph_ensure(pc, cls)) {
return false;
}
CodecStreamGraph * sg = &pc->stream_graphs[cls];
int32_t pos = pc->stream_pos;
ggml_backend_tensor_set(pc->stream_in_pos, &pos, 0, sizeof(int32_t));
ggml_backend_tensor_set(sg->codes, codes, 0, (size_t) T * (size_t) K * sizeof(int32_t));
int64_t row = (int64_t) (pc->stream_pos % ring);
ggml_backend_tensor_set(pc->stream_in_rows, &row, 0, sizeof(int64_t));
std::vector<int32_t> pos((size_t) T);
std::vector<int64_t> rows((size_t) T);
for (int t = 0; t < T; t++) {
pos[(size_t) t] = pc->stream_pos + t;
rows[(size_t) t] = (int64_t) ((pc->stream_pos + t) % ring);
}
ggml_backend_tensor_set(sg->pos, pos.data(), 0, (size_t) T * sizeof(int32_t));
ggml_backend_tensor_set(sg->rows, rows.data(), 0, (size_t) T * sizeof(int64_t));
std::vector<float> mask_buf;
tok_trans_build_stream_mask(pc->stream_pos, 1, ring, pc->transformer.sliding_window, mask_buf);
ggml_backend_tensor_set(pc->stream_in_mask, mask_buf.data(), 0, mask_buf.size() * sizeof(float));
tok_trans_build_stream_mask(pc->stream_pos, T, ring, pc->transformer.sliding_window, mask_buf);
ggml_backend_tensor_set(sg->mask, mask_buf.data(), 0, mask_buf.size() * sizeof(float));
enum ggml_status st = ggml_backend_graph_compute(pc->backend, pc->stream_gf);
enum ggml_status st = ggml_backend_graph_compute(pc->backend, sg->gf);
if (st != GGML_STATUS_SUCCESS) {
qt_log(QT_LOG_ERROR, "[Pipeline] stream graph_compute status=%d", (int) st);
return false;
}
if (audio_out) {
ggml_backend_tensor_get(pc->stream_out, audio_out, 0, (size_t) TOKENIZER_HOP_LENGTH * sizeof(float));
ggml_backend_tensor_get(sg->out, audio_out, 0, (size_t) T * (size_t) TOKENIZER_HOP_LENGTH * sizeof(float));
}
// The state tensors and the graph allocation persist into the next
// frame.
pc->stream_pos++;
// The state tensors and every class allocation persist into the
// next chunk.
pc->stream_pos += T;
return true;
}
@@ -755,11 +774,16 @@ void pipeline_codec_free(PipelineCodec * pc) {
}
graph_arena_free(&pc->dec_arena);
if (pc->stream_ready) {
ggml_gallocr_free(pc->stream_galloc);
pc->stream_galloc = NULL;
ggml_free(pc->stream_graph_ctx);
pc->stream_graph_ctx = NULL;
pc->stream_gf = NULL;
for (int i = 0; i < CODEC_STREAM_CLASSES; i++) {
CodecStreamGraph * sg = &pc->stream_graphs[i];
if (sg->galloc) {
ggml_gallocr_free(sg->galloc);
}
if (sg->ctx) {
ggml_free(sg->ctx);
}
*sg = {};
}
kv_cache_free(&pc->stream_kv);
ggml_backend_buffer_free(pc->stream_buf);
pc->stream_buf = NULL;
src/pipeline-codec.h
+27 -9
View File
@@ -49,6 +49,23 @@
// Primed stream state snapshots kept per reference, LRU evicted.
static const int CODEC_SNAP_SLOTS = 8;
// Chunk width classes of the streaming decode: T in {1, 2, 4, 8}.
static const int CODEC_STREAM_CLASSES = 4;
// One static stream graph of chunk width T: built once, allocated
// once, replayed with per chunk uploads of codes, positions, ring
// rows, and the sliding window mask.
struct CodecStreamGraph {
struct ggml_context * ctx = nullptr;
struct ggml_cgraph * gf = nullptr;
ggml_gallocr_t galloc = nullptr;
struct ggml_tensor * codes = nullptr;
struct ggml_tensor * pos = nullptr;
struct ggml_tensor * rows = nullptr;
struct ggml_tensor * mask = nullptr;
struct ggml_tensor * out = nullptr;
};
// One primed stream state snapshot: a mirror of every conv context and
// KV ring tensor in a single backend buffer, plus the host position
// cursor. key is the content hash of the reference codes, stamp orders
@@ -112,14 +129,11 @@ struct PipelineCodec {
// constant, so the graph builds and allocates once and every frame
// is input uploads + one backend compute + one readback. The single
// backend runs the whole graph, no scheduler involved.
struct ggml_context * stream_graph_ctx;
struct ggml_cgraph * stream_gf;
ggml_gallocr_t stream_galloc;
struct ggml_tensor * stream_in_codes;
struct ggml_tensor * stream_in_pos;
struct ggml_tensor * stream_in_rows;
struct ggml_tensor * stream_in_mask;
struct ggml_tensor * stream_out;
// Static stream graphs, one per chunk width class (T = 1 << cls).
// Every class shares the stream state and KV ring tensors above;
// only the input shapes and the intermediates differ. Classes
// build lazily on their first decode.
CodecStreamGraph stream_graphs[CODEC_STREAM_CLASSES];
// Snapshot LRU over the stream state: after an ICL reference
// priming the conv contexts, KV ring, and position copy device to
@@ -163,7 +177,11 @@ bool pipeline_codec_stream_reset(PipelineCodec * pc);
// as a side effect. When audio_out is non NULL the frame's
// TOKENIZER_HOP_LENGTH samples copy into it; a NULL audio_out primes
// the state without a readback (ICL reference priming).
bool pipeline_codec_decode_stream(PipelineCodec * pc, const int32_t * codes, float * audio_out);
// Decode one chunk of T frames (T in {1, 2, 4, 8}) through the
// persistent stream state. codes is [T, K] with T contiguous per
// codebook; audio_out receives T * TOKENIZER_HOP_LENGTH samples, NULL
// discards the audio (reference priming).
bool pipeline_codec_decode_stream(PipelineCodec * pc, const int32_t * codes, int T, float * audio_out);
// Content hash of an ICL reference, the snapshot LRU key.
// codes: flat int32, [K, T] row-major.
src/pipeline-tts.cpp
+11 -3
View File
@@ -833,10 +833,18 @@ qt_status pipeline_tts_synthesize(PipelineTTS * pt,
debug_dump_i32_as_f32(&d, "codes-full", flat.data(), shape, 2);
}
// Streaming tail: nothing to drain, every frame already emitted at
// generation time through the stateful decoder. The buffered output
// stays empty in this branch.
// 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;
}
out->samples = NULL;
out->n_samples = 0;
out->sample_rate = TOKENIZER_SAMPLE_RATE;