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

engine: batched multi lane codec streaming, shared lockstep chunk ramp

Browse Source 
One codec stream state set per lane plus a staging set, static graphs
per (chunk class, lane count) decoding every streaming lane in one
compute. The engine drops the per slot decoders and ownership
save/load for a shared 1 -> 2 -> 4 -> 8 ramp: lanes accumulate and
flush together, an admit drains and restarts at width 1, a retirement
drains then compacts the lane span device side. ICL references prime
through the staging set into per set snapshots. Depthwise and
transposed convs fold lanes into the channel axis; the dense convs
rely on the ggml conv_1d batched layout fix (submodule bump).
This commit is contained in:
Pascal
2026-07-20 19:13:05 +02:00
parent 9d7768e8ee
commit 53c84fe1be
9 changed files with 887 additions and 559 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/causal-trans-conv.h
+78 -44
View File
@@ -79,14 +79,17 @@ static struct ggml_tensor * qwen_load_ctw_f32(WeightCtx * wctx, const GGUFModel
return dst;
}
// Causal ConvTranspose1d forward graph.
// Causal ConvTranspose1d forward graph, batched over the trailing
// lane dim when present.
// w_perm: [IC, K*OC] f32, pre-permuted by qwen_load_ctw_f32
// b: [OC] f32 or NULL
// x: [T_in, IC] f32, T-first
// x: [T_in, IC] or [T_in, IC, N] f32, T-first
// stride: upsample factor
// kernel: kernel size
// oc: output channels (must match the K*OC factorization of w_perm)
// Returns [T_in*stride, OC] f32, T-first.
// Returns [T_in*stride, OC(, N)] f32, T-first. The N > 1 case folds
// the lanes into the col2im channel axis exactly like the streaming
// variant below.
static struct ggml_tensor * qwen_causal_trans_conv1d(struct ggml_context * ctx,
struct ggml_tensor * w_perm,
struct ggml_tensor * b,
@@ -94,25 +97,38 @@ static struct ggml_tensor * qwen_causal_trans_conv1d(struct ggml_context * ctx,
int stride,
int kernel,
int oc) {
int T = (int) x->ne[0];
int N = (int) x->ne[2];
int trim = kernel - stride;
// Transpose x to channels-first [IC, T_in] for the mul_mat contraction
// Transpose x to channels-first for the mul_mat contraction
struct ggml_tensor * xt = ggml_cont(ctx, ggml_transpose(ctx, x));
// mul_mat contracts over IC: col [K*OC, T_in]
// mul_mat contracts over IC: col [K*OC, T_in(, N)]
struct ggml_tensor * col = ggml_mul_mat(ctx, w_perm, xt);
// col2im_1d with padding=0: [T_raw, OC] T-first, T_raw = (T_in-1)*stride + K
struct ggml_tensor * y = ggml_col2im_1d(ctx, col, stride, oc, 0);
struct ggml_tensor * y;
if (N == 1) {
// col2im_1d with padding=0: [T_raw, OC] T-first, T_raw = (T_in-1)*stride + K
y = ggml_col2im_1d(ctx, ggml_reshape_2d(ctx, col, col->ne[0], T), stride, oc, 0);
if (ggml_n_dims(x) > 2) {
y = ggml_reshape_3d(ctx, y, y->ne[0], oc, 1);
}
} else {
struct ggml_tensor * folded = ggml_cont(ctx, ggml_permute(ctx, col, 0, 2, 1, 3)); // [K*OC, N, T]
folded = ggml_reshape_2d(ctx, folded, col->ne[0] * N, T);
y = ggml_col2im_1d(ctx, folded, stride, oc * N, 0); // [T_raw, OC*N]
y = ggml_reshape_3d(ctx, y, y->ne[0], oc, N);
}
// Right-trim K-stride frames -> [T_in*stride, OC] T-first
// Right-trim K-stride frames -> [T_in*stride, OC(, N)] T-first
if (trim > 0) {
int64_t T_keep = y->ne[0] - trim;
y = ggml_view_2d(ctx, y, T_keep, y->ne[1], y->nb[1], 0);
y = ggml_view_3d(ctx, y, T_keep, y->ne[1], y->ne[2], y->nb[1], y->nb[2], 0);
}
if (b) {
// bias [OC] broadcasts as (1, OC) onto (T, OC) via ne[0]=1
// bias [OC] broadcasts as (1, OC) onto (T, OC(, N)) via ne[0]=1
struct ggml_tensor * b2d = ggml_reshape_2d(ctx, b, 1, b->ne[0]);
y = ggml_add(ctx, y, b2d);
}
@@ -201,16 +217,17 @@ static struct ggml_tensor * qwen_causal_conv1d(struct ggml_context * ctx,
return y;
}
// Streaming causal Conv1d, stride 1. The offline zero left pad is
// replaced by a persistent state tensor carrying the last (k-1)*d input
// rows across calls: the graph concats the state ahead of the fresh
// rows, runs a pad free conv, and refreshes the state in graph with the
// tail of the extended input. The state is [(k-1)*d, IC] f32, backend
// resident and zero cleared at stream reset, so the first call matches
// the offline zero pad bit for bit. The state write depends on the
// concat output, so it always executes after the read.
// w: [k, IC, OC] f32, x: [T, IC] f32 T-first
// Returns [T, OC] f32 T-first.
// Streaming causal Conv1d, stride 1, batched over N lanes. The
// offline zero left pad is replaced by a persistent state tensor
// carrying the last (k-1)*d input rows of every lane across calls: the
// graph concats the state ahead of the fresh rows, runs a pad free
// conv over the N lane batch, and refreshes the state in graph with
// the tail of the extended input. The state slice is [(k-1)*d, IC, N]
// f32, backend resident and zero cleared at stream reset, so the first
// call matches the offline zero pad bit for bit. The state write
// depends on the concat output, so it always executes after the read.
// w: [k, IC, OC] f32, x: [T, IC, N] f32 T-first
// Returns [T, OC, N] f32 T-first.
static struct ggml_tensor * qwen_causal_conv1d_stream(struct ggml_context * ctx,
struct ggml_cgraph * gf,
struct ggml_tensor * w,
@@ -222,16 +239,14 @@ static struct ggml_tensor * qwen_causal_conv1d_stream(struct ggml_context * ctx,
int OC = (int) w->ne[2];
int L = (k - 1) * d;
struct ggml_tensor * x_ext = ggml_concat(ctx, state, x, 0); // [L + T, IC]
struct ggml_tensor * x_ext = ggml_concat(ctx, state, x, 0); // [L + T, IC, N]
// State refresh: the last L rows of x_ext feed the next call.
struct ggml_tensor * tail =
ggml_view_2d(ctx, x_ext, L, x_ext->ne[1], x_ext->nb[1], (size_t) (x_ext->ne[0] - L) * x_ext->nb[0]);
struct ggml_tensor * tail = ggml_view_3d(ctx, x_ext, L, x_ext->ne[1], x_ext->ne[2], x_ext->nb[1], x_ext->nb[2],
(size_t) (x_ext->ne[0] - L) * x_ext->nb[0]);
ggml_build_forward_expand(gf, ggml_cpy(ctx, tail, state));
struct ggml_tensor * y = ggml_reshape_3d(ctx, x_ext, x_ext->ne[0], x_ext->ne[1], 1);
y = ggml_conv_1d(ctx, w, y, 1, 0, d);
y = ggml_reshape_2d(ctx, y, y->ne[0], y->ne[1]);
struct ggml_tensor * y = ggml_conv_1d(ctx, w, x_ext, 1, 0, d); // [T, OC, N]
if (b) {
struct ggml_tensor * b2d = ggml_reshape_2d(ctx, b, 1, OC);
@@ -240,15 +255,23 @@ static struct ggml_tensor * qwen_causal_conv1d_stream(struct ggml_context * ctx,
return y;
}
// Streaming causal ConvTranspose1d. The raw col2im output spans
// (T-1)*stride + K rows; the offline path right trims K - stride of
// them, the streaming path instead carries that tail into the next
// call: the persistent carry [K - stride, OC] adds onto the head of the
// raw output and refreshes with the raw tail, bias free (the bias
// applies once, on the emitted rows). The carry consumer expands before
// the carry write so the read always precedes the overwrite.
// w_perm: [IC, K*OC] f32 from qwen_load_ctw_f32, x: [T, IC] f32 T-first
// Returns [T*stride, OC] f32 T-first.
// Streaming causal ConvTranspose1d, batched over N lanes. The raw
// col2im output spans (T-1)*stride + K rows; the offline path right
// trims K - stride of them, the streaming path instead carries that
// tail into the next call: the persistent carry [K - stride, OC, N]
// adds onto the head of the raw output and refreshes with the raw
// tail, bias free (the bias applies once, on the emitted rows). The
// carry consumer expands before the carry write so the read always
// precedes the overwrite.
//
// col2im_1d treats every output channel independently, so the N lane
// batch folds into the channel axis: col [K*OC, T, N] permutes to
// [K*OC, N, T], flattens to [K*(OC*N), T], scatters through the 2D
// col2im with oc' = oc + n*OC, and the [T_raw, OC*N] result reshapes
// straight to [T_raw, OC, N]. The N == 1 branch keeps the fold free
// chain.
// w_perm: [IC, K*OC] f32 from qwen_load_ctw_f32, x: [T, IC, N] f32
// Returns [T*stride, OC, N] f32 T-first.
static struct ggml_tensor * qwen_causal_trans_conv1d_stream(struct ggml_context * ctx,
struct ggml_cgraph * gf,
struct ggml_tensor * w_perm,
@@ -259,28 +282,39 @@ static struct ggml_tensor * qwen_causal_trans_conv1d_stream(struct ggml_context
int oc,
struct ggml_tensor * carry) {
int T = (int) x->ne[0];
int N = (int) x->ne[2];
int trim = kernel - stride;
int emit = T * stride;
struct ggml_tensor * xt = ggml_cont(ctx, ggml_transpose(ctx, x));
struct ggml_tensor * col = ggml_mul_mat(ctx, w_perm, xt);
struct ggml_tensor * xt = ggml_cont(ctx, ggml_transpose(ctx, x)); // [IC, T, N]
struct ggml_tensor * col = ggml_mul_mat(ctx, w_perm, xt); // [K*OC, T, N]
// Raw scatter [(T-1)*stride + K, OC] = [emit + trim, OC]
struct ggml_tensor * raw = ggml_col2im_1d(ctx, col, stride, oc, 0);
// Raw scatter [(T-1)*stride + K, OC, N] = [emit + trim, OC, N]
struct ggml_tensor * raw;
if (N == 1) {
raw = ggml_col2im_1d(ctx, ggml_reshape_2d(ctx, col, col->ne[0], T), stride, oc, 0);
raw = ggml_reshape_3d(ctx, raw, raw->ne[0], oc, 1);
} else {
struct ggml_tensor * folded = ggml_cont(ctx, ggml_permute(ctx, col, 0, 2, 1, 3)); // [K*OC, N, T]
folded = ggml_reshape_2d(ctx, folded, col->ne[0] * N, T);
raw = ggml_col2im_1d(ctx, folded, stride, oc * N, 0); // [emit + trim, OC*N]
raw = ggml_reshape_3d(ctx, raw, raw->ne[0], oc, N);
}
// Head rows [0, trim) receive the previous call's tail.
struct ggml_tensor * head = ggml_view_2d(ctx, raw, trim, raw->ne[1], raw->nb[1], 0);
struct ggml_tensor * head = ggml_view_3d(ctx, raw, trim, raw->ne[1], raw->ne[2], raw->nb[1], raw->nb[2], 0);
struct ggml_tensor * y = ggml_add(ctx, head, carry);
if (emit > trim) {
struct ggml_tensor * mid =
ggml_view_2d(ctx, raw, emit - trim, raw->ne[1], raw->nb[1], (size_t) trim * raw->nb[0]);
y = ggml_concat(ctx, y, mid, 0);
struct ggml_tensor * mid = ggml_view_3d(ctx, raw, emit - trim, raw->ne[1], raw->ne[2], raw->nb[1], raw->nb[2],
(size_t) trim * raw->nb[0]);
y = ggml_concat(ctx, y, mid, 0);
}
ggml_build_forward_expand(gf, y);
// Carry refresh with the raw tail [emit, emit + trim), expanded
// after the head sum so the carry read wins the ordering.
struct ggml_tensor * tail = ggml_view_2d(ctx, raw, trim, raw->ne[1], raw->nb[1], (size_t) emit * raw->nb[0]);
struct ggml_tensor * tail =
ggml_view_3d(ctx, raw, trim, raw->ne[1], raw->ne[2], raw->nb[1], raw->nb[2], (size_t) emit * raw->nb[0]);
ggml_build_forward_expand(gf, ggml_cpy(ctx, tail, carry));
if (b) {
src/codec-chunked-decode.h
+4 -132
View File
@@ -13,18 +13,16 @@
// previously decoded frames and stripping the resulting samples after
// the decode restores continuity.
//
// Two entry points:
// One entry point:
//
// codec_chunked_decode : one shot decode of a full codes buffer.
// Bit perfect equivalent of pipeline_codec_decode when the audio
// fits in a single chunk_frames sized window. Bounds VRAM beyond
// that, mirrors the upstream chunked_decode loop frame for frame.
//
// codec_stream_decoder : stateful frame by frame AR streaming.
// The pipeline pushes one frame at a time as the talker produces
// them ; push_frame decodes and emits a fresh chunk_frames sized
// audio block through the on_chunk callback as soon as enough new
// frames have accumulated. flush drains the tail at EOS.
// Frame by frame AR streaming lives in the batch engine
// (pipeline-tts.cpp), which drives the persistent multi lane stream
// state of pipeline-codec directly.
#include "pipeline-codec.h"
#include "qwen.h"
@@ -83,129 +81,3 @@ static inline std::vector<float> codec_chunked_decode(PipelineCodec * pc,
}
return out;
}
// 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 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;
static const int MAX_CHUNK = 1 << (CODEC_STREAM_CLASSES - 1);
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;
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 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;
}
int t0 = 0;
while (t0 < ref_T) {
int T = MAX_CHUNK;
while (T > ref_T - t0) {
T >>= 1;
}
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);
}
// 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) {
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 (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/convnext-block.h
+27 -10
View File
@@ -198,9 +198,16 @@ struct QwenUpsampleStreamState {
struct ggml_tensor * dw[UPSAMPLE_MAX_BLOCKS];
};
// Streaming ConvNeXt block: the depthwise causal conv reads its left
// context from the persistent state instead of a zero pad. Everything
// else is pointwise and stateless.
// Streaming ConvNeXt block, batched over N lanes: the depthwise causal
// conv reads its left context from the persistent state instead of a
// zero pad. Everything else is pointwise and stateless.
//
// ggml_conv_1d_dw carries no batch dim, but a depthwise conv treats
// every channel independently, so the N lanes fold into the channel
// axis: x_ext [L + T, C, N] reshapes to [L + T, C*N], the kernel tiles
// N times to [k, 1, C*N] (channel' = c + n*C matches the reshape
// order), and the [T, C*N] output reshapes back to [T, C, N]. The
// N == 1 branch keeps the fold free chain.
static struct ggml_tensor * convnext_block_forward_stream(struct ggml_context * ctx,
struct ggml_cgraph * gf,
const QwenConvNeXtBlock & block,
@@ -208,25 +215,35 @@ static struct ggml_tensor * convnext_block_forward_stream(struct ggml_context *
int kernel,
struct ggml_tensor * dw_state) {
int C = (int) x->ne[1];
int N = (int) x->ne[2];
struct ggml_tensor * residual = x;
// dwconv: concat the state ahead of the fresh rows, refresh it with
// the tail, run the depthwise conv pad free.
struct ggml_tensor * x_ext = ggml_concat(ctx, dw_state, x, 0); // [k-1 + T, C]
struct ggml_tensor * tail =
ggml_view_2d(ctx, x_ext, kernel - 1, C, x_ext->nb[1], (size_t) (x_ext->ne[0] - (kernel - 1)) * x_ext->nb[0]);
struct ggml_tensor * x_ext = ggml_concat(ctx, dw_state, x, 0); // [k-1 + T, C, N]
struct ggml_tensor * tail = ggml_view_3d(ctx, x_ext, kernel - 1, C, N, x_ext->nb[1], x_ext->nb[2],
(size_t) (x_ext->ne[0] - (kernel - 1)) * x_ext->nb[0]);
ggml_build_forward_expand(gf, ggml_cpy(ctx, tail, dw_state));
struct ggml_tensor * y = ggml_reshape_3d(ctx, x_ext, x_ext->ne[0], C, 1);
y = ggml_conv_1d_dw(ctx, block.dwconv_w, y, 1, 0, 1); // [T, C, 1]
y = ggml_reshape_2d(ctx, y, y->ne[0], C);
struct ggml_tensor * y;
if (N == 1) {
y = ggml_conv_1d_dw(ctx, block.dwconv_w, x_ext, 1, 0, 1); // [T, C, 1]
} else {
struct ggml_tensor * w_tmpl =
ggml_new_tensor_4d(ctx, block.dwconv_w->type, block.dwconv_w->ne[0], block.dwconv_w->ne[1], C, N);
struct ggml_tensor * w_n = ggml_repeat(ctx, block.dwconv_w, w_tmpl); // [k, 1, C, N]
w_n = ggml_reshape_3d(ctx, w_n, w_n->ne[0], 1, C * N); // [k, 1, C*N]
struct ggml_tensor * folded = ggml_reshape_3d(ctx, x_ext, x_ext->ne[0], C * N, 1);
y = ggml_conv_1d_dw(ctx, w_n, folded, 1, 0, 1); // [T, C*N, 1]
y = ggml_reshape_3d(ctx, y, y->ne[0], C, N);
}
if (block.dwconv_b) {
struct ggml_tensor * b2d = ggml_reshape_2d(ctx, block.dwconv_b, 1, C);
y = ggml_add(ctx, y, b2d);
}
// LayerNorm wants the channel dim on ne[0]: transpose to [C, T].
// LayerNorm wants the channel dim on ne[0]: transpose to [C, T, N].
y = ggml_cont(ctx, ggml_transpose(ctx, y));
y = ggml_norm(ctx, y, 1e-6f);
y = ggml_mul(ctx, y, block.norm_w);
src/dac-decoder-v2.h
+8 -3
View File
@@ -273,8 +273,9 @@ struct QwenDACStreamState {
struct ggml_tensor * post; // conv_post k=7, [6, 96]
};
// Streaming residual unit: conv1 reads its left context from the
// persistent state, conv2 is pointwise.
// Streaming residual unit, batched over N lanes: conv1 reads its left
// context from the persistent state, conv2 is pointwise (k=1, pad
// free) and runs a direct batched conv.
static struct ggml_tensor * dac_res_unit_stream(struct ggml_context * ctx,
struct ggml_cgraph * gf,
const QwenDACResUnit * ru,
@@ -284,7 +285,11 @@ static struct ggml_tensor * dac_res_unit_stream(struct ggml_context * ctx,
x = dac_snake(ctx, x, ru->act1);
x = qwen_causal_conv1d_stream(ctx, gf, ru->c1w, ru->c1b, x, 7, ru->dilation, state);
x = dac_snake(ctx, x, ru->act2);
x = qwen_causal_conv1d(ctx, ru->c2w, ru->c2b, x, 1, 1);
x = ggml_conv_1d(ctx, ru->c2w, x, 1, 0, 1);
if (ru->c2b) {
struct ggml_tensor * b2d = ggml_reshape_2d(ctx, ru->c2b, 1, ru->c2b->ne[0]);
x = ggml_add(ctx, x, b2d);
}
return ggml_add(ctx, skip, x);
}
src/pipeline-codec.cpp
+278 -148
View File
@@ -27,9 +27,13 @@ bool pipeline_codec_load(PipelineCodec * pc, const char * gguf_path, BackendPair
pc->stream_ready = false;
pc->stream_ctx = NULL;
pc->stream_buf = NULL;
pc->stream_pos = 0;
pc->stream_sets = 0;
pc->stream_n_state = 0;
pc->stream_pos.clear();
pc->stream_set_views.clear();
pc->stream_graphs.clear();
for (int i = 0; i < CODEC_STREAM_CLASSES; i++) {
pc->stream_graphs[i] = {};
pc->staging_graphs[i] = {};
}
for (int i = 0; i < CODEC_SNAP_SLOTS; i++) {
pc->snaps[i] = {};
@@ -194,52 +198,87 @@ std::vector<float> pipeline_codec_decode(PipelineCodec * pc, const int32_t * cod
// step graph shape never changes.
static const int CODEC_STREAM_RING = 128;
// Allocate the streaming state on first use: one tensor per causal conv
// left context and per transposed conv carry, plus the transformer KV
// ring and the dedicated graph arena. Idempotent.
// Allocate the streaming state on first use: one [t, c, S] tensor per
// causal conv left context and per transposed conv carry over the
// stream_sets lanes plus staging, the multi-set transformer KV ring,
// and one 2D view per set of every state tensor. The views are
// created before the backend allocation so ggml_backend_alloc_ctx_tensors
// runs its view init on them (buffer and data resolve against the
// owning 3D tensor); they drive the device side set copies and the
// snapshot mirrors. Idempotent.
static bool pipeline_codec_stream_ensure(PipelineCodec * pc) {
if (pc->stream_ready) {
return true;
}
const int S = pc->stream_sets >= 2 ? pc->stream_sets : 2;
pc->stream_sets = S;
const int n_state = 2 + UPSAMPLE_MAX_BLOCKS + DAC_NUM_BLOCKS * (1 + DAC_RES_UNITS) + 4;
struct ggml_init_params gp = { ggml_tensor_overhead() * (size_t) n_state, NULL, true };
struct ggml_init_params gp = { ggml_tensor_overhead() * (size_t) (n_state * (S + 1)), NULL, true };
pc->stream_ctx = ggml_init(gp);
if (!pc->stream_ctx) {
qt_log(QT_LOG_ERROR, "[Pipeline] stream state ggml_init failed");
return false;
}
pc->stream_set_views.clear();
pc->stream_n_state = 0;
char name[64];
auto tensor2d = [&](int t, int c, const char * n) {
struct ggml_tensor * x = ggml_new_tensor_2d(pc->stream_ctx, GGML_TYPE_F32, t, c);
auto tensor3d = [&](int t, int c, const char * n) {
struct ggml_tensor * x = ggml_new_tensor_3d(pc->stream_ctx, GGML_TYPE_F32, t, c, S);
ggml_set_name(x, n);
// Per set 2D views, appended set major after every owner is
// created (see below); record the owner order here.
pc->stream_n_state++;
return x;
};
// pre_conv k=3 over the 512 wide quantizer latents
pc->stream_pre_conv = tensor2d(2, 512, "stream_pre_conv");
pc->stream_pre_conv = tensor3d(2, 512, "stream_pre_conv");
// upsample ConvNeXt depthwise contexts, k=7 over the stage width
for (int i = 0; i < pc->upsample.num_blocks; i++) {
snprintf(name, sizeof(name), "stream_up_dw_%d", i);
pc->stream_up.dw[i] = tensor2d(pc->upsample.dwconv_kernel - 1, pc->upsample.channels, name);
pc->stream_up.dw[i] = tensor3d(pc->upsample.dwconv_kernel - 1, pc->upsample.channels, name);
}
// DAC contexts: conv_pre, per block transconv carry + res unit conv1, conv_post.
// conv_pre reads the 1024 wide upsample output.
pc->stream_dac.pre = tensor2d(6, 1024, "stream_dac_pre");
pc->stream_dac.pre = tensor3d(6, 1024, "stream_dac_pre");
for (int i = 0; i < DAC_NUM_BLOCKS; i++) {
const QwenDACBlock & b = pc->dac.blk[i];
snprintf(name, sizeof(name), "stream_dac_carry_%d", i);
pc->stream_dac.carry[i] = tensor2d(b.kernel - b.stride, b.out_ch, name);
pc->stream_dac.carry[i] = tensor3d(b.kernel - b.stride, b.out_ch, name);
for (int r = 0; r < DAC_RES_UNITS; r++) {
snprintf(name, sizeof(name), "stream_dac_ru_%d_%d", i, r);
pc->stream_dac.ru[i][r] = tensor2d(6 * b.ru[r].dilation, b.out_ch, name);
pc->stream_dac.ru[i][r] = tensor3d(6 * b.ru[r].dilation, b.out_ch, name);
}
}
pc->stream_dac.post = tensor3d(6, pc->dac.channels[DAC_NUM_BLOCKS], "stream_dac_post");
// Per set 2D views, set major: views[s * n_state + i] aliases state
// tensor i of set s. Creation order inside a set matches the owner
// creation order, so the snapshot copy walker pairs positionally.
pc->stream_set_views.assign((size_t) S * (size_t) pc->stream_n_state, nullptr);
{
int i = 0;
for (struct ggml_tensor * t = ggml_get_first_tensor(pc->stream_ctx); t;
t = ggml_get_next_tensor(pc->stream_ctx, t)) {
if (t->view_src) {
continue;
}
for (int s = 0; s < S; s++) {
pc->stream_set_views[(size_t) s * (size_t) pc->stream_n_state + (size_t) i] =
ggml_view_2d(pc->stream_ctx, t, t->ne[0], t->ne[1], t->nb[1], (size_t) s * t->nb[2]);
}
i++;
if (i > pc->stream_n_state) {
break;
}
}
}
pc->stream_dac.post = tensor2d(6, pc->dac.channels[DAC_NUM_BLOCKS], "stream_dac_post");
pc->stream_buf = ggml_backend_alloc_ctx_tensors(pc->stream_ctx, pc->backend);
if (!pc->stream_buf) {
@@ -249,17 +288,17 @@ static bool pipeline_codec_stream_ensure(PipelineCodec * pc) {
return false;
}
// The ring must hold the whole sliding window plus the fresh frame,
// otherwise window slots alias through the modulo and corrupt the
// attention silently.
if (pc->transformer.sliding_window + 1 > CODEC_STREAM_RING) {
// The ring must hold the whole sliding window plus the widest fresh
// chunk, otherwise window slots alias through the modulo and
// corrupt the attention silently.
if (pc->transformer.sliding_window + (1 << (CODEC_STREAM_CLASSES - 1)) > CODEC_STREAM_RING) {
qt_log(QT_LOG_ERROR, "[Pipeline] sliding window %d exceeds KV ring %d", pc->transformer.sliding_window,
CODEC_STREAM_RING);
return false;
}
if (!kv_cache_init(&pc->stream_kv, pc->transformer.num_layers, pc->transformer.num_kv_heads,
pc->transformer.head_dim, CODEC_STREAM_RING, 1, pc->backend)) {
pc->transformer.head_dim, CODEC_STREAM_RING, S, pc->backend)) {
ggml_backend_buffer_free(pc->stream_buf);
pc->stream_buf = NULL;
ggml_free(pc->stream_ctx);
@@ -267,22 +306,45 @@ static bool pipeline_codec_stream_ensure(PipelineCodec * pc) {
return false;
}
pc->stream_graphs.assign((size_t) CODEC_STREAM_CLASSES * (size_t) (S - 1), CodecStreamGraph{});
pc->stream_pos.assign((size_t) S, 0);
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);
qt_log(QT_LOG_INFO, "[Pipeline] Codec stream state ready: %d conv contexts, KV ring %d, %d sets (1 staging)",
pc->stream_n_state, CODEC_STREAM_RING, S);
return true;
}
bool pipeline_codec_stream_reset(PipelineCodec * pc) {
// Zero one set's slice of every state tensor and its KV ring set
// through a host zero upload, so lanes reset independently without
// touching their neighbours.
bool pipeline_codec_stream_reset(PipelineCodec * pc, int set) {
if (!pipeline_codec_stream_ensure(pc)) {
return false;
}
// Zero contexts reproduce the offline zero left pads bit for bit;
// 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, 0);
pc->stream_pos = 0;
// the zeroed KV ring set stays hidden behind the sliding window
// mask.
static thread_local std::vector<float> zeros;
for (int i = 0; i < pc->stream_n_state; i++) {
struct ggml_tensor * v = pc->stream_set_views[(size_t) set * (size_t) pc->stream_n_state + (size_t) i];
size_t n = (size_t) v->ne[0] * (size_t) v->ne[1];
if (zeros.size() < n) {
zeros.assign(n, 0.0f);
}
ggml_backend_tensor_set(v, zeros.data(), 0, n * sizeof(float));
}
for (int l = 0; l < pc->stream_kv.n_layers; l++) {
for (struct ggml_tensor * v : { kv_cache_k(&pc->stream_kv, set, l), kv_cache_v(&pc->stream_kv, set, l) }) {
size_t n = (size_t) ggml_nelements(v);
if (zeros.size() < n) {
zeros.assign(n, 0.0f);
}
ggml_backend_tensor_set(v, zeros.data(), 0, n * sizeof(float));
}
}
kv_cache_reset(&pc->stream_kv, set);
pc->stream_pos[(size_t) set] = 0;
return true;
}
@@ -298,42 +360,27 @@ uint64_t pipeline_codec_ref_key(const int32_t * codes, int K, int T) {
return (h ^ (uint64_t) T) * 1099511628211ULL;
}
// Allocate the mirror tensors of a snapshot slot: one duplicate per
// stream state and KV ring tensor. Creation order is preserved so the
// copy walker pairs source and mirror positionally.
// Allocate the mirror tensors of a snapshot slot: one contiguous
// duplicate per conv state view and per KV ring view of a single set
// (all sets share these shapes). Creation order (conv views, then k
// and v per layer) pairs positionally with the codec_snap_copy walker.
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 += 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 += t->view_src ? 0 : 1;
}
const int n = pc->stream_n_state + 2 * pc->stream_kv.n_layers;
struct ggml_init_params gp = { ggml_tensor_overhead() * (size_t) n, NULL, true };
s->ctx = ggml_init(gp);
if (!s->ctx) {
qt_log(QT_LOG_ERROR, "[Pipeline] snapshot ggml_init failed");
return false;
}
for (struct ggml_tensor * t = ggml_get_first_tensor(pc->stream_ctx); t;
t = ggml_get_next_tensor(pc->stream_ctx, t)) {
if (!t->view_src) {
ggml_dup_tensor(s->ctx, t);
}
for (int i = 0; i < pc->stream_n_state; i++) {
ggml_dup_tensor(s->ctx, pc->stream_set_views[(size_t) i]);
}
for (struct ggml_tensor * t = ggml_get_first_tensor(pc->stream_kv.ctx); t;
t = ggml_get_next_tensor(pc->stream_kv.ctx, t)) {
if (!t->view_src) {
ggml_dup_tensor(s->ctx, t);
}
for (int l = 0; l < pc->stream_kv.n_layers; l++) {
ggml_dup_tensor(s->ctx, kv_cache_k(&pc->stream_kv, 0, l));
ggml_dup_tensor(s->ctx, kv_cache_v(&pc->stream_kv, 0, l));
}
s->buf = ggml_backend_alloc_ctx_tensors(s->ctx, pc->backend);
if (!s->buf) {
@@ -345,29 +392,26 @@ static bool codec_snap_ensure(PipelineCodec * pc, CodecStateSnap * s) {
return true;
}
// Copy every stream state and KV ring tensor to (save) or from its
// slot mirror, device to device on a shared backend.
static void codec_snap_copy(PipelineCodec * pc, CodecStateSnap * s, bool save) {
// Copy one set's conv state views and KV ring views to (save) or from
// its slot mirror, device to device on a shared backend. The set views
// slice the outermost dimension of contiguous owners, so every copy
// moves one dense block.
static void codec_snap_copy(PipelineCodec * pc, CodecStateSnap * s, int set, 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)) {
if (t->view_src) {
continue;
}
ggml_backend_tensor_copy(save ? t : m, save ? m : t);
for (int i = 0; i < pc->stream_n_state; i++) {
struct ggml_tensor * v = pc->stream_set_views[(size_t) set * (size_t) pc->stream_n_state + (size_t) i];
ggml_backend_tensor_copy(save ? v : m, save ? m : v);
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)) {
if (t->view_src) {
continue;
for (int l = 0; l < pc->stream_kv.n_layers; l++) {
for (struct ggml_tensor * v : { kv_cache_k(&pc->stream_kv, set, l), kv_cache_v(&pc->stream_kv, set, l) }) {
ggml_backend_tensor_copy(save ? v : m, save ? m : v);
m = ggml_get_next_tensor(s->ctx, m);
}
ggml_backend_tensor_copy(save ? t : m, save ? m : t);
m = ggml_get_next_tensor(s->ctx, m);
}
}
bool pipeline_codec_stream_restore(PipelineCodec * pc, uint64_t key) {
bool pipeline_codec_stream_restore(PipelineCodec * pc, uint64_t key, int set) {
if (!pc->stream_ready) {
return false;
}
@@ -377,16 +421,19 @@ bool pipeline_codec_stream_restore(PipelineCodec * pc, uint64_t key) {
continue;
}
Timer t;
codec_snap_copy(pc, s, false);
pc->stream_pos = s->pos;
s->stamp = ++pc->snap_stamp;
qt_log(QT_LOG_INFO, "[Pipeline] Codec state restored from snapshot in %.1f ms (%d frames)", t.ms(), s->pos);
codec_snap_copy(pc, s, set, false);
pc->stream_pos[(size_t) set] = s->pos;
s->stamp = ++pc->snap_stamp;
qt_log(QT_LOG_INFO, "[Pipeline] Codec state restored to set %d in %.1f ms (%d frames)", set, t.ms(), s->pos);
return true;
}
return false;
}
bool pipeline_codec_stream_snapshot(PipelineCodec * pc, uint64_t key) {
bool pipeline_codec_stream_snapshot(PipelineCodec * pc, uint64_t key, int set) {
if (!pc->stream_ready) {
return false;
}
CodecStateSnap * lru = &pc->snaps[0];
for (int i = 1; i < CODEC_SNAP_SLOTS; i++) {
if (pc->snaps[i].stamp < lru->stamp) {
@@ -396,29 +443,28 @@ bool pipeline_codec_stream_snapshot(PipelineCodec * pc, uint64_t key) {
if (!codec_snap_ensure(pc, lru)) {
return false;
}
codec_snap_copy(pc, lru, true);
codec_snap_copy(pc, lru, set, true);
lru->key = key;
lru->pos = pc->stream_pos;
lru->pos = pc->stream_pos[(size_t) set];
lru->stamp = ++pc->snap_stamp;
qt_log(QT_LOG_INFO, "[Pipeline] Codec state snapshot saved (%d frames)", lru->pos);
qt_log(QT_LOG_INFO, "[Pipeline] Codec state snapshot saved from set %d (%d frames)", set, lru->pos);
return true;
}
bool pipeline_codec_stream_save(PipelineCodec * pc, CodecStateSnap * s) {
if (!pc->stream_ready || !codec_snap_ensure(pc, s)) {
bool pipeline_codec_stream_copy_set(PipelineCodec * pc, int src, int dst) {
if (!pc->stream_ready) {
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;
// Retiring the last active lane resolves to src == dst: nothing moves.
if (src == dst) {
return true;
}
codec_snap_copy(pc, s, false);
pc->stream_pos = s->pos;
for (int i = 0; i < pc->stream_n_state; i++) {
ggml_backend_tensor_copy(pc->stream_set_views[(size_t) src * (size_t) pc->stream_n_state + (size_t) i],
pc->stream_set_views[(size_t) dst * (size_t) pc->stream_n_state + (size_t) i]);
}
kv_cache_copy_set(&pc->stream_kv, src, dst);
pc->stream_pos[(size_t) dst] = pc->stream_pos[(size_t) src];
return true;
}
@@ -436,14 +482,16 @@ void pipeline_codec_snap_free(CodecStateSnap * s) {
// 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;
}
// Build one static stream graph of chunk width T = 1 << cls over M
// lanes: the same module chain as the offline decode with the stateful
// batched variants threaded through the [0, M) slices of the
// persistent stream tensors, or through the staging set slice when
// staging is set. Inputs and intermediates size for (T, M).
static bool codec_stream_graph_build(PipelineCodec * pc, int cls, int M, bool staging, CodecStreamGraph * sg) {
const int T = 1 << cls;
const int ring = CODEC_STREAM_RING;
const int K = TOKENIZER_NUM_CODEBOOKS;
const int set0 = staging ? pc->stream_sets - 1 : 0;
// The ring must hold the sliding window plus the whole fresh chunk.
if (pc->transformer.sliding_window + T > ring) {
@@ -460,16 +508,16 @@ static bool codec_stream_graph_ensure(PipelineCodec * pc, int cls) {
};
sg->ctx = ggml_init(gp);
if (!sg->ctx) {
qt_log(QT_LOG_ERROR, "[Pipeline] stream graph ggml_init failed (T=%d)", T);
qt_log(QT_LOG_ERROR, "[Pipeline] stream graph ggml_init failed (T=%d, M=%d)", T, M);
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);
struct ggml_tensor * codes_in = ggml_new_tensor_3d(gctx, GGML_TYPE_I32, T, K, M);
struct ggml_tensor * pos_in = ggml_new_tensor_1d(gctx, GGML_TYPE_I32, T * M);
struct ggml_tensor * rows_in = ggml_new_tensor_3d(gctx, GGML_TYPE_I64, T, 1, M);
struct ggml_tensor * mask_in = ggml_new_tensor_4d(gctx, GGML_TYPE_F32, ring, T, 1, M);
ggml_set_name(codes_in, "codes_in");
ggml_set_name(pos_in, "positions");
ggml_set_name(rows_in, "kv_rows");
@@ -479,18 +527,44 @@ static bool codec_stream_graph_ensure(PipelineCodec * pc, int cls) {
ggml_set_input(rows_in);
ggml_set_input(mask_in);
// Lane state slices [t, c, M] at set0: contiguous suffix or prefix
// spans of the [t, c, S] owners, so views reshape freely.
auto slice = [&](struct ggml_tensor * owner) {
return ggml_view_3d(gctx, owner, owner->ne[0], owner->ne[1], M, owner->nb[1], owner->nb[2],
(size_t) set0 * owner->nb[2]);
};
QwenUpsampleStreamState up_sl;
QwenDACStreamState dac_sl;
for (int i = 0; i < pc->upsample.num_blocks; i++) {
up_sl.dw[i] = slice(pc->stream_up.dw[i]);
}
dac_sl.pre = slice(pc->stream_dac.pre);
for (int i = 0; i < DAC_NUM_BLOCKS; i++) {
dac_sl.carry[i] = slice(pc->stream_dac.carry[i]);
for (int r = 0; r < DAC_RES_UNITS; r++) {
dac_sl.ru[i][r] = slice(pc->stream_dac.ru[i][r]);
}
}
dac_sl.post = slice(pc->stream_dac.post);
// KVCache facade over the lane set span: reuse the multi-set ring
// tensors with per graph 4D views built inside tok_trans; staging
// binds the last set through a single set span.
KVCache * kv = &pc->stream_kv;
// 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);
// batched variants threaded through the persistent stream slices.
// The quantizer uses the alignment safe variant: no scheduler input
// duplication happens on the direct backend compute path. Codes
// flatten to [T * M, K] column gathers via the lane major layout.
struct ggml_tensor * h = quant_decode_stream_batch(gctx, &pc->qdec, codes_in); // [512, T, M] C-first
h = ggml_cont(gctx, ggml_transpose(gctx, h)); // [T, 512, M] T-first
h = qwen_causal_conv1d_stream(gctx, gf, pc->pre_conv_w, pc->pre_conv_b, h, 3, 1, slice(pc->stream_pre_conv));
h = ggml_cont(gctx, ggml_transpose(gctx, h)); // [1024, T, M] C-first
h = tok_trans_forward_stream_span(gctx, gf, &pc->transformer, h, pos_in, mask_in, rows_in, kv, set0, M);
h = ggml_cont(gctx, ggml_transpose(gctx, h)); // [T, 1024, M] T-first
h = upsample_stage_forward_stream(gctx, gf, &pc->upsample, h, &up_sl);
h = dac_decoder_forward_stream(gctx, gf, &pc->dac, h, &dac_sl);
h = ggml_clamp(gctx, h, -1.0f, 1.0f);
ggml_set_name(h, "audio_out");
ggml_set_output(h);
@@ -498,7 +572,7 @@ static bool codec_stream_graph_ensure(PipelineCodec * pc, int cls) {
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);
qt_log(QT_LOG_ERROR, "[Pipeline] stream graph allocation failed (T=%d, M=%d)", T, M);
if (sg->galloc) {
ggml_gallocr_free(sg->galloc);
}
@@ -516,35 +590,49 @@ static bool codec_stream_graph_ensure(PipelineCodec * pc, int cls) {
return true;
}
bool pipeline_codec_decode_stream(PipelineCodec * pc, const int32_t * codes, int T, float * audio_out) {
// Lazily build and fetch the (cls, M) lane graph or the cls staging
// graph.
static CodecStreamGraph * codec_stream_graph_ensure(PipelineCodec * pc, int cls, int M, bool staging) {
CodecStreamGraph * sg = staging ?
&pc->staging_graphs[cls] :
&pc->stream_graphs[(size_t) cls * (size_t) (pc->stream_sets - 1) + (size_t) (M - 1)];
if (sg->ctx) {
return sg;
}
return codec_stream_graph_build(pc, cls, M, staging, sg) ? sg : nullptr;
}
// Shared replay: upload codes, per lane positions, ring rows, and
// masks over the sets [set0, set0 + M), compute, read every non NULL
// lane's audio back, advance every lane's position.
static bool codec_stream_run(PipelineCodec * pc,
CodecStreamGraph * sg,
const int32_t * codes,
int T,
int M,
int set0,
float ** audio_out) {
const int K = TOKENIZER_NUM_CODEBOOKS;
const int ring = CODEC_STREAM_RING;
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];
ggml_backend_tensor_set(sg->codes, codes, 0, (size_t) T * (size_t) K * (size_t) M * sizeof(int32_t));
ggml_backend_tensor_set(sg->codes, codes, 0, (size_t) T * (size_t) K * sizeof(int32_t));
sg->pos_buf.resize((size_t) T);
sg->rows_buf.resize((size_t) T);
for (int t = 0; t < T; t++) {
sg->pos_buf[(size_t) t] = pc->stream_pos + t;
sg->rows_buf[(size_t) t] = (int64_t) ((pc->stream_pos + t) % ring);
sg->pos_buf.resize((size_t) T * (size_t) M);
sg->rows_buf.resize((size_t) T * (size_t) M);
static thread_local std::vector<int> pos0;
pos0.assign((size_t) M, 0);
for (int m = 0; m < M; m++) {
int p = pc->stream_pos[(size_t) (set0 + m)];
pos0[m] = p;
for (int t = 0; t < T; t++) {
sg->pos_buf[(size_t) m * (size_t) T + (size_t) t] = p + t;
sg->rows_buf[(size_t) m * (size_t) T + (size_t) t] = (int64_t) ((p + t) % ring);
}
}
ggml_backend_tensor_set(sg->pos, sg->pos_buf.data(), 0, (size_t) T * sizeof(int32_t));
ggml_backend_tensor_set(sg->rows, sg->rows_buf.data(), 0, (size_t) T * sizeof(int64_t));
ggml_backend_tensor_set(sg->pos, sg->pos_buf.data(), 0, sg->pos_buf.size() * sizeof(int32_t));
ggml_backend_tensor_set(sg->rows, sg->rows_buf.data(), 0, sg->rows_buf.size() * sizeof(int64_t));
tok_trans_build_stream_mask(pc->stream_pos, T, ring, pc->transformer.sliding_window, sg->mask_buf);
tok_trans_build_stream_mask(pos0.data(), T, M, ring, pc->transformer.sliding_window, sg->mask_buf);
ggml_backend_tensor_set(sg->mask, sg->mask_buf.data(), 0, sg->mask_buf.size() * sizeof(float));
enum ggml_status st = ggml_backend_graph_compute(pc->backend, sg->gf);
@@ -553,16 +641,50 @@ bool pipeline_codec_decode_stream(PipelineCodec * pc, const int32_t * codes, int
return false;
}
if (audio_out) {
ggml_backend_tensor_get(sg->out, audio_out, 0, (size_t) T * (size_t) TOKENIZER_HOP_LENGTH * sizeof(float));
const size_t lane_samples = (size_t) T * (size_t) TOKENIZER_HOP_LENGTH;
for (int m = 0; m < M; m++) {
if (audio_out && audio_out[m]) {
ggml_backend_tensor_get(sg->out, audio_out[m], (size_t) m * lane_samples * sizeof(float),
lane_samples * sizeof(float));
}
pc->stream_pos[(size_t) (set0 + m)] += T;
}
// The state tensors and every class allocation persist into the
// next chunk.
pc->stream_pos += T;
return true;
}
bool pipeline_codec_decode_stream_batch(PipelineCodec * pc, const int32_t * codes, int T, int M, float ** audio_out) {
int cls = 0;
while ((1 << cls) < T) {
cls++;
}
if (cls >= CODEC_STREAM_CLASSES || (1 << cls) != T || M < 1 || M > pc->stream_sets - 1) {
qt_log(QT_LOG_ERROR, "[Pipeline] invalid stream chunk width %d or lane count %d", T, M);
return false;
}
CodecStreamGraph * sg = codec_stream_graph_ensure(pc, cls, M, false);
if (!sg) {
return false;
}
return codec_stream_run(pc, sg, codes, T, M, 0, audio_out);
}
bool pipeline_codec_decode_stream(PipelineCodec * pc, const int32_t * codes, int T, float * audio_out) {
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;
}
CodecStreamGraph * sg = codec_stream_graph_ensure(pc, cls, 1, true);
if (!sg) {
return false;
}
float * outs[1] = { audio_out };
return codec_stream_run(pc, sg, codes, T, 1, pc->stream_sets - 1, audio_out ? outs : nullptr);
}
bool pipeline_codec_ensure_encoder(PipelineCodec * pc) {
if (pc->enc_loaded) {
return true;
@@ -820,21 +942,29 @@ void pipeline_codec_free(PipelineCodec * pc) {
}
graph_arena_free(&pc->dec_arena);
if (pc->stream_ready) {
auto graph_release = [](CodecStreamGraph & sg) {
if (sg.galloc) {
ggml_gallocr_free(sg.galloc);
}
if (sg.ctx) {
ggml_free(sg.ctx);
}
sg = {};
};
for (CodecStreamGraph & sg : pc->stream_graphs) {
graph_release(sg);
}
pc->stream_graphs.clear();
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 = {};
graph_release(pc->staging_graphs[i]);
}
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;
pc->stream_ctx = NULL;
pc->stream_set_views.clear();
pc->stream_pos.clear();
pc->stream_ready = false;
}
for (int i = 0; i < CODEC_SNAP_SLOTS; i++) {
src/pipeline-codec.h
+76 -63
View File
@@ -52,28 +52,30 @@ 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.
// One static stream graph of chunk width T and lane count M: built
// once, allocated once, replayed with per chunk uploads of codes,
// positions, ring rows, and the sliding window mask of every lane.
// Lane m reads and writes stream state set m; the staging graphs bind
// the single staging set instead (M = 1, set stream_sets - 1).
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;
struct ggml_tensor * codes = nullptr; // [T, K, M] i32
struct ggml_tensor * pos = nullptr; // [T * M] i32
struct ggml_tensor * rows = nullptr; // [T, 1, M] i64
struct ggml_tensor * mask = nullptr; // [ring, T, 1, M] f32
struct ggml_tensor * out = nullptr; // [T * 1920, 1, M] f32
std::vector<int32_t> pos_buf;
std::vector<int64_t> rows_buf;
std::vector<float> mask_buf;
};
// 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
// the slots for LRU eviction, stamp zero marks an empty slot. Mirrors
// allocate lazily on the slot's first save.
// One primed stream state snapshot: a mirror of one set's conv context
// and KV ring tensors in a single backend buffer, plus the host
// position cursor. key is the content hash of the reference codes,
// stamp orders the slots for LRU eviction, stamp zero marks an empty
// slot. Mirrors allocate lazily on the slot's first save.
struct CodecStateSnap {
uint64_t key;
uint64_t stamp;
@@ -113,30 +115,39 @@ struct PipelineCodec {
// so constant size streaming slices replay a captured executable.
GraphArena dec_arena;
// Stateful streaming decoder: every causal conv left context, every
// transposed conv overlap carry, and the transformer sliding window
// KV ring live as backend resident tensors, so a T=1 frame decode
// reproduces the offline full decode exactly with zero re-decoded
// context. Loaded lazily on the first pipeline_codec_stream_reset:
// the buffered chunked path never pays for it.
bool stream_ready;
struct ggml_context * stream_ctx;
ggml_backend_buffer_t stream_buf;
struct ggml_tensor * stream_pre_conv; // pre_conv k=3, [2, 512]
QwenUpsampleStreamState stream_up;
QwenDACStreamState stream_dac;
KVCache stream_kv; // tok transformer ring, [hd, ring, n_kv] per layer
int stream_pos; // absolute frame position, drives RoPE and ring slots
// Stateful streaming decoder, batched over lanes: every causal
// conv left context, every transposed conv overlap carry, and the
// transformer sliding window KV ring live as backend resident
// tensors carrying one set per lane plus a staging set (the last
// one) dedicated to ICL reference priming, so a T=1 frame decode
// of every lane reproduces the offline full decode exactly with
// zero re-decoded context. stream_sets is written by the pipeline
// owner before the first stream use (max_batch + 1, minimum 2);
// state allocates lazily on the first pipeline_codec_stream_reset.
// Per set 2D views of every state tensor are created before the
// backend allocation (the alloc initialises the views) and drive
// the device side set copies of pipeline_codec_stream_copy_set and
// the snapshot mirrors.
int stream_sets;
bool stream_ready;
struct ggml_context * stream_ctx;
ggml_backend_buffer_t stream_buf;
struct ggml_tensor * stream_pre_conv; // pre_conv k=3, [2, 512, S]
QwenUpsampleStreamState stream_up;
QwenDACStreamState stream_dac;
KVCache stream_kv; // tok transformer ring, [hd, ring, n_kv, S] per layer
std::vector<struct ggml_tensor *> stream_set_views; // per set 2D views, set major
int stream_n_state; // state tensors per set
std::vector<int> stream_pos; // absolute frame position per set
// Static frame graph: the T=1 topology and every tensor address are
// 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.
// 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];
// Static stream graphs, one per chunk width class (T = 1 << cls)
// and lane count M in [1, stream_sets - 1], plus one staging class
// set (M = 1 bound to the staging set) for reference priming.
// Every graph shares the stream state and KV ring tensors above;
// only the input shapes, the lane span, and the intermediates
// differ. Graphs build lazily on their first decode.
std::vector<CodecStreamGraph> stream_graphs; // index cls * (stream_sets - 1) + (M - 1)
CodecStreamGraph staging_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
@@ -169,42 +180,44 @@ bool pipeline_codec_ensure_encoder(PipelineCodec * pc);
// Returns audio of length T * TOKENIZER_HOP_LENGTH, empty on failure.
std::vector<float> pipeline_codec_decode(PipelineCodec * pc, const int32_t * codes, int K, int T);
// Reset the stateful streaming decoder to the zero context: allocates
// the state tensors on first call, clears every conv left context,
// transposed conv carry, and the transformer KV ring, and rewinds the
// absolute position. Call once before each streamed utterance.
bool pipeline_codec_stream_reset(PipelineCodec * pc);
// Reset one stream state set to the zero context: allocates the state
// tensors on first call (stream_sets must be written before that),
// clears the set's conv left contexts, transposed conv carries, and
// KV ring set through a host zero upload, and rewinds its position.
// Call once per lane admit and before each staging priming.
bool pipeline_codec_stream_reset(PipelineCodec * pc, int set);
// Decode one frame through the stateful streaming path. codes holds the
// K codebook entries of a single frame; the persistent state advances
// 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).
// 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).
// Decode one chunk of T frames (T in {1, 2, 4, 8}) of M lanes through
// the persistent stream state sets [0, M). codes is [T, K, M] lane
// major (T contiguous per codebook, K codebooks per lane);
// audio_out[m] receives lane m's T * TOKENIZER_HOP_LENGTH samples,
// NULL discards that lane's audio. Every lane's position advances.
bool pipeline_codec_decode_stream_batch(PipelineCodec * pc, const int32_t * codes, int T, int M, float ** audio_out);
// Decode one chunk of T frames through the STAGING set (the last
// one). codes is [T, K]; audio_out receives T * TOKENIZER_HOP_LENGTH
// samples, NULL discards the audio (ICL reference priming). The
// staging position advances.
bool pipeline_codec_decode_stream(PipelineCodec * pc, const int32_t * codes, int T, float * audio_out);
// Copy one stream state set (conv contexts, KV ring, position) device
// to device into another: lane retirement swap-remove and the staging
// to lane handoff after a priming.
bool pipeline_codec_stream_copy_set(PipelineCodec * pc, int src, int dst);
// Content hash of an ICL reference, the snapshot LRU key.
// codes: flat int32, [K, T] row-major.
uint64_t pipeline_codec_ref_key(const int32_t * codes, int K, int T);
// Restore the stream state from the snapshot slot matching key.
// Returns false on a miss; the caller then primes and snapshots.
bool pipeline_codec_stream_restore(PipelineCodec * pc, uint64_t key);
// Restore stream state set `set` from the snapshot slot matching key.
// Returns false on a miss; the caller then primes through the staging
// set and snapshots.
bool pipeline_codec_stream_restore(PipelineCodec * pc, uint64_t key, int set);
// Save the current stream state into the LRU slot for key, evicting
// the least recently used slot when all are taken.
bool pipeline_codec_stream_snapshot(PipelineCodec * pc, uint64_t key);
// Save stream state set `set` into the LRU slot for key, evicting the
// least recently used slot when all are taken.
bool pipeline_codec_stream_snapshot(PipelineCodec * pc, uint64_t key, int set);
// 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.
src/pipeline-tts.cpp
+265 -112
View File
@@ -204,6 +204,10 @@ bool pipeline_tts_load(PipelineTTS * pt,
gf_close(&pt->gguf_talker);
return false;
}
// One codec stream state set per lane plus the staging set for ICL
// reference priming; must land before the first stream call, which
// allocates the [t, c, S] state tensors from it.
pt->codec.stream_sets = pt->max_batch + 1;
// Scheduler shared by talker_forward_* and code_predictor_step.
// Routes ops the GPU backend cannot run (typical case: K-quant
@@ -509,11 +513,11 @@ struct TtsSlot {
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.
// Streaming state: the slot's codec stream lane, an index into the
// compacted [0, codec_M) span of the shared multi set codec state.
// -1 for buffered and wav slots.
bool streaming;
codec_stream_decoder stream;
CodecStateSnap snap;
int codec_set;
std::vector<std::vector<int32_t>> all_codes;
bool finished;
@@ -527,26 +531,45 @@ struct TtsEngine {
BPETokenizer * tok;
std::vector<TtsSlot> slots;
int64_t next_serial;
int64_t stream_owner; // serial of the slot holding the live codec state, -1 none
// Shared codec streaming ramp, lockstep over the compacted lanes
// [0, codec_M): every streaming slot accumulates into the same
// pending rows and flushes through one batched graph compute, so
// the 8x kernel amortization survives multi lane operation. An
// admit resets the ramp to 1 for the new lane's first audio
// latency; a retirement drains the pending rows first so the
// leaving lane's audio is fully dispatched before the swap remove.
int codec_M; // active streaming lanes
int codec_target; // current ramp chunk width
int codec_pending_n; // rows accumulated, < codec_target
std::vector<int32_t> codec_pending; // [row][lane][K] frame rows
std::vector<uint8_t> codec_live; // per lane: rows are real frames
std::vector<int32_t> codec_codes; // [T, K, M] flush upload layout
std::vector<float> codec_audio; // [lane][8 * hop] flush output
std::vector<float *> codec_outs; // per lane out or NULL
};
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);
TtsEngine * e = new TtsEngine();
e->pt = pt;
e->tok = tok;
e->next_serial = 0;
const size_t maxM = (size_t) pt->max_batch;
const size_t K = (size_t) pt->num_code_groups;
const size_t maxT = (size_t) (1 << (CODEC_STREAM_CLASSES - 1));
e->codec_M = 0;
e->codec_target = 1;
e->codec_pending_n = 0;
e->codec_pending.assign(maxT * maxM * K, 0);
e->codec_live.assign(maxM, 0);
e->codec_codes.assign(maxT * maxM * K, 0);
e->codec_audio.assign(maxM * maxT * (size_t) TOKENIZER_HOP_LENGTH, 0.0f);
e->codec_outs.assign(maxM, nullptr);
e->slots.reserve(maxM);
return e;
}
void tts_engine_free(TtsEngine * e) {
if (!e) {
return;
}
for (TtsSlot & s : e->slots) {
pipeline_codec_snap_free(&s.snap);
}
delete e;
}
@@ -554,30 +577,134 @@ 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;
static TtsSlot * tts_engine_codec_lane_slot(TtsEngine * e, int m) {
for (TtsSlot & s : e->slots) {
if (s.codec_set == m) {
return &s;
}
}
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;
return nullptr;
}
// Drain the shared pending rows through greedy width classes: one
// batched graph compute decodes all codec_M lanes per chunk, then each
// lane's samples go to its slot's on_chunk. A cancelled slot or a lane
// riding zero rows gets a NULL out and no callback. The full flush
// wall time lands on every lane's codec_ms, the same convention as the
// batched talker and predictor spans. Returns false only on a decode
// failure; a callback cancel marks the slot and the flush carries on.
static bool tts_engine_codec_flush(TtsEngine * e) {
PipelineTTS * pt = e->pt;
const int M = e->codec_M;
const int K = pt->num_code_groups;
const int hop = TOKENIZER_HOP_LENGTH;
Timer t_codec;
while (e->codec_pending_n > 0) {
int T = 1 << (CODEC_STREAM_CLASSES - 1);
while (T > e->codec_pending_n) {
T >>= 1;
}
for (int m = 0; m < M; m++) {
for (int k = 0; k < K; k++) {
for (int t = 0; t < T; t++) {
e->codec_codes[(size_t) (m * K + k) * (size_t) T + (size_t) t] =
e->codec_pending[((size_t) t * (size_t) M + (size_t) m) * (size_t) K + (size_t) k];
}
break;
}
}
for (int m = 0; m < M; m++) {
TtsSlot * s = tts_engine_codec_lane_slot(e, m);
const bool want = s && e->codec_live[(size_t) m] && s->fin_status != QT_STATUS_CANCELLED &&
s->job->params->on_chunk != NULL;
e->codec_outs[(size_t) m] =
want ? e->codec_audio.data() + (size_t) m * (size_t) (1 << (CODEC_STREAM_CLASSES - 1)) * (size_t) hop :
nullptr;
}
if (!pipeline_codec_decode_stream_batch(&pt->codec, e->codec_codes.data(), T, M, e->codec_outs.data())) {
return false;
}
e->codec_pending_n -= T;
if (e->codec_pending_n > 0) {
std::memmove(e->codec_pending.data(), e->codec_pending.data() + (size_t) T * (size_t) M * (size_t) K,
(size_t) e->codec_pending_n * (size_t) M * (size_t) K * sizeof(int32_t));
}
for (int m = 0; m < M; m++) {
if (!e->codec_outs[(size_t) m]) {
continue;
}
TtsSlot * s = tts_engine_codec_lane_slot(e, m);
const struct qt_tts_params * p = s->job->params;
if (!p->on_chunk(e->codec_outs[(size_t) m], T * hop, p->on_chunk_user_data)) {
qt_log(QT_LOG_INFO, "[Pipeline] on_chunk callback aborted the synthesis (lane %d)", m);
s->finished = true;
s->fin_status = QT_STATUS_CANCELLED;
}
}
}
e->stream_owner = -1;
if (s->snap.ctx && !pipeline_codec_stream_load(&e->pt->codec, &s->snap)) {
const double ms = t_codec.ms();
for (int m = 0; m < M; m++) {
TtsSlot * s = tts_engine_codec_lane_slot(e, m);
if (s) {
s->perf.codec_ms += ms;
}
}
return true;
}
// Attach a fresh streaming lane for an admitted slot: drain the
// pending rows of the current lanes first (their chunks come out
// shorter around the event), zero or ICL prime the new set, restart
// the shared ramp at width 1 so the newcomer's first frame decodes
// immediately. The ICL path restores a previously primed reference
// snapshot device to device, or primes through the staging set in max
// width chunks with the audio discarded and snapshots it for reuse.
static bool tts_engine_codec_admit(TtsEngine * e, TtsSlot * s) {
PipelineTTS * pt = e->pt;
PipelineCodec * pc = &pt->codec;
const int K = pt->num_code_groups;
if (e->codec_pending_n > 0 && !tts_engine_codec_flush(e)) {
return false;
}
e->stream_owner = s->serial;
const int set = e->codec_M;
if (s->ref_codes_ptr != NULL) {
Timer t_seed;
const uint64_t key = pipeline_codec_ref_key(s->ref_codes_ptr, K, s->ref_codes_T);
if (!pipeline_codec_stream_restore(pc, key, set)) {
const int staging = pt->max_batch;
if (!pipeline_codec_stream_reset(pc, staging)) {
return false;
}
std::vector<int32_t> chunk((size_t) (1 << (CODEC_STREAM_CLASSES - 1)) * (size_t) K);
int t0 = 0;
while (t0 < s->ref_codes_T) {
int T = 1 << (CODEC_STREAM_CLASSES - 1);
while (T > s->ref_codes_T - t0) {
T >>= 1;
}
for (int k = 0; k < K; k++) {
for (int t = 0; t < T; t++) {
chunk[(size_t) k * (size_t) T + (size_t) t] =
s->ref_codes_ptr[(size_t) k * (size_t) s->ref_codes_T + (size_t) (t0 + t)];
}
}
if (!pipeline_codec_decode_stream(pc, chunk.data(), T, NULL)) {
return false;
}
t0 += T;
}
if (!pipeline_codec_stream_snapshot(pc, key, staging) ||
!pipeline_codec_stream_copy_set(pc, staging, set)) {
return false;
}
}
s->perf.codec_ms += t_seed.ms();
} else if (!pipeline_codec_stream_reset(pc, set)) {
return false;
}
s->codec_set = set;
e->codec_M = set + 1;
e->codec_target = 1;
e->codec_pending_n = 0;
return true;
}
@@ -694,7 +821,7 @@ bool tts_engine_admit(TtsEngine * e, TtsJob * job) {
s.pending_c0 = -1;
s.has_frame = false;
s.streaming = (params->on_chunk != NULL);
s.snap = {};
s.codec_set = -1;
s.finished = false;
s.fin_status = QT_STATUS_OK;
s.perf = {};
@@ -780,34 +907,17 @@ bool tts_engine_admit(TtsEngine * e, TtsJob * job) {
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.
// Taking ownership parks the previous streaming slot's state first.
// Streaming lane attach: the shared codec ramp restarts at width 1
// so the first generated frame decodes immediately and the audio
// callback fires with it. ICL clone priming runs the full reference
// through the staging set, matching the upstream reference plus
// generated decode exactly.
if (s.streaming) {
if (!tts_engine_codec_own(e, &s)) {
qt_set_error("pipeline_tts_synthesize: codec stream state park failed");
if (!tts_engine_codec_admit(e, &s)) {
qt_set_error("pipeline_tts_synthesize: codec stream lane admit failed");
e->slots.pop_back();
return tts_admit_fail(job, QT_STATUS_GENERATE_FAILED);
}
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 (!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");
e->stream_owner = -1;
e->slots.pop_back();
return tts_admit_fail(job, QT_STATUS_GENERATE_FAILED);
}
s.perf.codec_ms += t_seed.ms();
}
}
// Talker prefill into KV set slot_idx: the joining request stalls
@@ -818,8 +928,10 @@ bool tts_engine_admit(TtsEngine * e, TtsJob * job) {
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;
// The lane just attached is the tail one; dropping it needs no
// compaction.
if (s.codec_set >= 0) {
e->codec_M--;
}
e->slots.pop_back();
return tts_admit_fail(job, QT_STATUS_GENERATE_FAILED);
@@ -860,30 +972,18 @@ static void tts_slot_complete(TtsEngine * e, TtsSlot & s) {
}
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);
// Streaming tail: the engine step drained the shared
// pending rows before entering retirement, so every sample
// is already dispatched; finish with an empty buffered
// output.
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;
@@ -958,11 +1058,6 @@ static void tts_slot_complete(TtsEngine * e, TtsSlot & s) {
}
}
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();
}
@@ -1169,27 +1264,10 @@ void tts_engine_step(TtsEngine * e, std::vector<TtsJob *> * retired) {
s.all_codes.push_back(codes);
s.talker_history.push_back(s.pending_c0);
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;
}
}
// Streaming slots stage this frame through
// all_codes.back() and has_frame; the shared codec
// flush after this loop decodes every lane in one
// batched compute.
// Next decode input: the 16 frame codes gather and sum
// in graph (codebook 0 from talker.codec_embedding, the
@@ -1251,16 +1329,91 @@ void tts_engine_step(TtsEngine * e, std::vector<TtsJob *> * retired) {
// finished nor advanced cannot exist: every live slot either emits
// or finishes.
// 5) Retirement: swap-remove keeps the active range consecutive.
// 5) Shared codec streaming: one lockstep flush cadence over the
// lanes [0, codec_M). A membership change this frame (a streaming
// slot finished) first drains the aligned pending rows, then the
// freshly staged frames ride a single row flush where a lane whose
// slot emitted nothing carries a zero code row and a NULL out, so
// every retiring lane leaves with its audio fully dispatched before
// the swap remove below. Zero rows only ever exist in that single
// row flush, which keeps the per lane audio blocks free of padding.
if (e->codec_M > 0) {
const int num_cg = pt->num_code_groups;
bool any_finish = false;
bool any_stage = false;
for (TtsSlot & s : e->slots) {
if (s.codec_set < 0) {
continue;
}
any_finish = any_finish || s.finished;
any_stage = any_stage || s.has_frame;
}
bool ok = true;
if (any_finish && e->codec_pending_n > 0) {
ok = tts_engine_codec_flush(e);
}
if (ok && any_stage) {
int32_t * row =
e->codec_pending.data() + (size_t) e->codec_pending_n * (size_t) e->codec_M * (size_t) num_cg;
for (int m = 0; m < e->codec_M; m++) {
TtsSlot * s = tts_engine_codec_lane_slot(e, m);
if (s && s->has_frame) {
std::memcpy(row + (size_t) m * (size_t) num_cg, s->all_codes.back().data(),
(size_t) num_cg * sizeof(int32_t));
e->codec_live[(size_t) m] = 1;
} else {
std::memset(row + (size_t) m * (size_t) num_cg, 0, (size_t) num_cg * sizeof(int32_t));
e->codec_live[(size_t) m] = 0;
}
}
e->codec_pending_n++;
if (any_finish) {
ok = tts_engine_codec_flush(e);
} else if (e->codec_pending_n >= e->codec_target) {
ok = tts_engine_codec_flush(e);
if (ok && e->codec_target < (1 << (CODEC_STREAM_CLASSES - 1))) {
e->codec_target <<= 1;
}
}
}
if (!ok) {
qt_set_error("pipeline_tts_synthesize: streaming codec decode failed");
qt_log(QT_LOG_ERROR, "[Pipeline] streaming codec decode failed");
for (TtsSlot & s : e->slots) {
if (s.codec_set >= 0) {
s.finished = true;
s.fin_status = QT_STATUS_GENERATE_FAILED;
}
}
}
}
// 6) 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.
// before any read, so only the talker cache moves. The codec lane
// span compacts the same way: the tail lane's stream state copies
// into the freed lane and its slot reindexes.
for (int i = 0; i < (int) e->slots.size();) {
if (!e->slots[(size_t) i].finished) {
i++;
continue;
}
tts_slot_complete(e, e->slots[(size_t) i]);
if (e->slots[(size_t) i].codec_set >= 0) {
TtsSlot & dead = e->slots[(size_t) i];
const int freed = dead.codec_set;
const int tail = e->codec_M - 1;
pipeline_codec_stream_copy_set(&pt->codec, tail, freed);
for (TtsSlot & o : e->slots) {
if (o.codec_set == tail) {
o.codec_set = freed;
break;
}
}
dead.codec_set = -1;
e->codec_M--;
}
if (retired) {
retired->push_back(e->slots[(size_t) i].job);
}
src/quantizer-decode.h
+36
View File
@@ -197,3 +197,39 @@ static struct ggml_tensor * quant_decode_stream(struct ggml_context * ctx
return ggml_add(ctx, h_sem, h_aco);
}
// Batched variant over M lanes: codes is [T, K, M] lane major. Every
// codebook's T ids of every lane gather flat as [T * M] through a
// strided view + cont (T contiguous inside a lane, lanes strided by
// K * T), keeping the same alignment safe get_rows chain as the single
// lane path. Returns [hidden, T, M] C-first.
static struct ggml_tensor * rvq_group_decode_stream_batch(struct ggml_context * ctx,
const QwenRVQGroup & g,
struct ggml_tensor * codes,
int k0,
int T,
int M) {
struct ggml_tensor * sum = NULL;
for (int k = 0; k < g.num_codebooks; k++) {
// Lane strided [T, M] view of codebook k0 + k, made flat [T * M].
struct ggml_tensor * idx2 = ggml_view_2d(ctx, codes, T, M, codes->nb[2], (size_t) (k0 + k) * codes->nb[1]);
struct ggml_tensor * idx = ggml_reshape_1d(ctx, ggml_cont(ctx, idx2), (int64_t) T * M);
struct ggml_tensor * emb = ggml_get_rows(ctx, g.embed[(size_t) k], idx); // [hidden_in, T * M]
sum = (sum == NULL) ? emb : ggml_add(ctx, sum, emb);
}
return ggml_mul_mat(ctx, g.out_proj_w, sum); // [hidden, T * M]
}
static struct ggml_tensor * quant_decode_stream_batch(struct ggml_context * ctx,
const QwenQuantizerDecoder * dec,
struct ggml_tensor * codes) {
int T = (int) codes->ne[0];
int M = (int) codes->ne[2];
struct ggml_tensor * h_sem = rvq_group_decode_stream_batch(ctx, dec->semantic, codes, 0, T, M);
struct ggml_tensor * h_aco =
rvq_group_decode_stream_batch(ctx, dec->acoustic, codes, dec->num_semantic_quantizers, T, M);
struct ggml_tensor * h = ggml_add(ctx, h_sem, h_aco);
return ggml_reshape_3d(ctx, h, h->ne[0], T, M);
}
src/tokenizer-transformer.h
+115 -47
View File
@@ -294,12 +294,15 @@ static struct ggml_tensor * tok_trans_forward(struct ggml_context * c
return h;
}
// Streaming layer forward: the fresh K and V rows write into a
// persistent ring cache via set_rows at the slots carried by kv_rows,
// and the attention reads the whole ring with the mask killing every
// slot outside the sliding window. Ring slots hold RoPE rotated keys at
// absolute positions, so relative attention falls out as usual. The
// graph topology is constant across steps: pure CUDA graph replay.
// Streaming layer forward, batched over N lanes: the fresh K and V
// rows of every lane write into its own KV ring set via set_rows at
// the slots carried by kv_rows [T, 1, N], and the attention reads the
// whole ring of the [0, N) set span with the mask [ring, T, 1, N]
// killing every slot outside each lane's sliding window. Ring slots
// hold RoPE rotated keys at absolute positions (per lane, through the
// flattened positions [T * N]), so relative attention falls out as
// usual. The graph topology is constant across steps: pure CUDA graph
// replay.
static struct ggml_tensor * tok_trans_layer_forward_stream(struct ggml_context * ctx,
struct ggml_cgraph * gf,
const QwenTokenizerTransformer * tr,
@@ -308,13 +311,15 @@ static struct ggml_tensor * tok_trans_layer_forward_stream(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 ring) {
int n_q_heads = tr->num_attention_heads;
int n_kv = tr->num_kv_heads;
int hd = tr->head_dim;
int TN = T * N;
struct ggml_tensor * ln1 = ggml_rms_norm(ctx, x, tr->rms_norm_eps);
ln1 = ggml_mul(ctx, ln1, layer.input_norm_w);
@@ -323,43 +328,46 @@ static struct ggml_tensor * tok_trans_layer_forward_stream(struct ggml_context *
struct ggml_tensor * k = ggml_mul_mat(ctx, layer.attn.k_proj_w, ln1);
struct ggml_tensor * v = ggml_mul_mat(ctx, layer.attn.v_proj_w, ln1);
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_rope_ext(ctx, q, positions, NULL, hd, GGML_ROPE_TYPE_NEOX, 0, tr->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, tr->rope_theta, 1.0f, 0.0f, 1.0f, 0.0f,
0.0f);
// Ring write: [hd, T, n_kv] rows land at kv_rows, ids broadcast
// across the head dim.
struct ggml_tensor * k_perm = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3));
struct ggml_tensor * v_perm = ggml_cont(ctx, ggml_permute(ctx, v, 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));
// Ring write: [hd, T, n_kv, N] rows land at each lane's set through
// the kv_rows [T, 1, N] destinations, ids broadcast across the
// head dim; k4/v4 are the [0, N) set span views of the ring.
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, k4, k_perm, kv_rows));
ggml_build_forward_expand(gf, ggml_set_rows(ctx, v4, v_perm, kv_rows));
struct ggml_tensor * q_p = ggml_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3));
struct ggml_tensor * q_p =
ggml_cont(ctx, ggml_permute(ctx, ggml_reshape_4d(ctx, q, hd, n_q_heads, T, N), 0, 2, 1, 3));
struct ggml_tensor * k_full = ggml_view_3d(ctx, k_cache, hd, ring, n_kv, k_cache->nb[1], k_cache->nb[2], 0);
struct ggml_tensor * v_full = ggml_view_3d(ctx, v_cache, hd, ring, n_kv, v_cache->nb[1], v_cache->nb[2], 0);
struct ggml_tensor * scores = ggml_mul_mat(ctx, k_full, q_p); // [ring, T, n_q_heads]
struct ggml_tensor * scores = ggml_mul_mat(ctx, k4, q_p); // [ring, T, n_q_heads, N]
float scale = 1.0f / sqrtf((float) hd);
scores = ggml_soft_max_ext(ctx, scores, mask, scale, 0.0f);
struct ggml_tensor * vt = ggml_cont(ctx, ggml_transpose(ctx, v_full)); // [ring, hd, n_kv]
struct ggml_tensor * attn = ggml_mul_mat(ctx, vt, scores); // [hd, T, n_q_heads]
struct ggml_tensor * vt = ggml_cont(ctx, ggml_transpose(ctx, v4)); // [ring, hd, n_kv, N]
struct ggml_tensor * attn = ggml_mul_mat(ctx, vt, scores); // [hd, T, n_q_heads, N]
attn = ggml_cont(ctx, ggml_permute(ctx, attn, 0, 2, 1, 3));
attn = ggml_reshape_2d(ctx, attn, n_q_heads * hd, T);
attn = ggml_cont(ctx, ggml_permute(ctx, attn, 0, 2, 1, 3)); // [hd, n_q_heads, T, N]
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);
o = ggml_mul(ctx, o, layer.attn_scale);
x = ggml_add(ctx, x, o);
(void) ring;
struct ggml_tensor * ln2 = ggml_rms_norm(ctx, x, tr->rms_norm_eps);
ln2 = ggml_mul(ctx, ln2, layer.post_attn_norm_w);
@@ -376,10 +384,15 @@ static struct ggml_tensor * tok_trans_layer_forward_stream(struct ggml_context *
return x;
}
// Streaming transformer forward over a persistent KV ring. kv holds one
// [hd, ring, n_kv] pair per layer; kv_rows carries the ring slots the T
// fresh positions land in; mask is [ring, T] f32 with 0 on the slots
// inside the causal sliding window and neg inf elsewhere.
// Streaming transformer forward over the persistent KV ring, batched
// over N lanes. Lane n reads and writes KV set n; k4/v4 span views
// over sets [0, N) build here from the 4D ring tensors. kv_rows
// [T, 1, N] carries the ring slots the T fresh positions of every lane
// land in; positions is the flattened [T * N] per lane absolute
// positions; mask is [ring, T, 1, N] f32 with 0 on the slots inside
// each lane's causal sliding window and neg inf elsewhere.
// x: [in_dim, T, N] f32 C-first
// Returns [out_dim, T, N] f32 C-first.
static struct ggml_tensor * tok_trans_forward_stream(struct ggml_context * ctx,
struct ggml_cgraph * gf,
const QwenTokenizerTransformer * tr,
@@ -387,16 +400,23 @@ static struct ggml_tensor * tok_trans_forward_stream(struct ggml_context *
struct ggml_tensor * positions,
struct ggml_tensor * mask,
struct ggml_tensor * kv_rows,
KVCache * kv) {
KVCache * kv,
int N) {
int T = (int) x->ne[1];
int ring = kv->max_seq_len;
struct ggml_tensor * h = ggml_mul_mat(ctx, tr->input_proj_w, x);
h = ggml_add(ctx, h, tr->input_proj_b);
h = ggml_reshape_2d(ctx, h, h->ne[0], (int64_t) T * N);
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_cache_k(kv, 0, l), kv_cache_v(kv, 0, l), T, ring);
struct ggml_tensor * k4l = kv->k4[(size_t) l];
struct ggml_tensor * v4l = kv->v4[(size_t) l];
struct ggml_tensor * k4 =
ggml_view_4d(ctx, k4l, k4l->ne[0], ring, k4l->ne[2], N, k4l->nb[1], k4l->nb[2], k4l->nb[3], 0);
struct ggml_tensor * v4 =
ggml_view_4d(ctx, v4l, v4l->ne[0], ring, v4l->ne[2], N, v4l->nb[1], v4l->nb[2], v4l->nb[3], 0);
h = tok_trans_layer_forward_stream(ctx, gf, tr, tr->layers[l], h, positions, mask, kv_rows, k4, v4, T, N, ring);
}
h = ggml_rms_norm(ctx, h, tr->rms_norm_eps);
@@ -405,22 +425,70 @@ static struct ggml_tensor * tok_trans_forward_stream(struct ggml_context *
h = ggml_mul_mat(ctx, tr->output_proj_w, h);
h = ggml_add(ctx, h, tr->output_proj_b);
return h;
return ggml_reshape_3d(ctx, h, h->ne[0], T, N);
}
// Ring mask for the streaming path: [ring, T] f32, row q carries 0 on
// the ring slots holding positions inside [pos_q - window + 1, pos_q]
// and neg inf everywhere else, padded and future slots included.
static void tok_trans_build_stream_mask(int pos0, int T, int ring, int sliding_window, std::vector<float> & dst) {
dst.assign((size_t) ring * (size_t) T, -INFINITY);
for (int q = 0; q < T; q++) {
int pos_q = pos0 + q;
int k_min = pos_q - sliding_window + 1;
if (k_min < 0) {
k_min = 0;
}
for (int p = k_min; p <= pos_q; p++) {
dst[(size_t) q * (size_t) ring + (size_t) (p % ring)] = 0.0f;
// Set span variant: identical to tok_trans_forward_stream with the
// [set0, set0 + N) span of the ring instead of [0, N), so the staging
// graphs bind the last set with the same code path.
static struct ggml_tensor * tok_trans_forward_stream_span(struct ggml_context * ctx,
struct ggml_cgraph * gf,
const QwenTokenizerTransformer * tr,
struct ggml_tensor * x,
struct ggml_tensor * positions,
struct ggml_tensor * mask,
struct ggml_tensor * kv_rows,
KVCache * kv,
int set0,
int N) {
int T = (int) x->ne[1];
int ring = kv->max_seq_len;
struct ggml_tensor * h = ggml_mul_mat(ctx, tr->input_proj_w, x);
h = ggml_add(ctx, h, tr->input_proj_b);
h = ggml_reshape_2d(ctx, h, h->ne[0], (int64_t) T * N);
for (int l = 0; l < tr->num_layers; l++) {
struct ggml_tensor * k4l = kv->k4[(size_t) l];
struct ggml_tensor * v4l = kv->v4[(size_t) l];
struct ggml_tensor * k4 = ggml_view_4d(ctx, k4l, k4l->ne[0], ring, k4l->ne[2], N, k4l->nb[1], k4l->nb[2],
k4l->nb[3], (size_t) set0 * k4l->nb[3]);
struct ggml_tensor * v4 = ggml_view_4d(ctx, v4l, v4l->ne[0], ring, v4l->ne[2], N, v4l->nb[1], v4l->nb[2],
v4l->nb[3], (size_t) set0 * v4l->nb[3]);
h = tok_trans_layer_forward_stream(ctx, gf, tr, tr->layers[l], h, positions, mask, kv_rows, k4, v4, T, N, ring);
}
h = ggml_rms_norm(ctx, h, tr->rms_norm_eps);
h = ggml_mul(ctx, h, tr->norm_w);
h = ggml_mul_mat(ctx, tr->output_proj_w, h);
h = ggml_add(ctx, h, tr->output_proj_b);
return ggml_reshape_3d(ctx, h, h->ne[0], T, N);
}
// Ring mask for the batched streaming path: [ring, T, 1, N] f32, lane
// n row q carries 0 on the ring slots holding positions inside
// [pos_q - window + 1, pos_q] of that lane and neg inf everywhere
// else, padded and future slots included. pos0 holds the absolute
// start position of every lane.
static void tok_trans_build_stream_mask(const int * pos0,
int T,
int N,
int ring,
int sliding_window,
std::vector<float> & dst) {
dst.assign((size_t) ring * (size_t) T * (size_t) N, -INFINITY);
for (int n = 0; n < N; n++) {
for (int q = 0; q < T; q++) {
int pos_q = pos0[n] + q;
int k_min = pos_q - sliding_window + 1;
if (k_min < 0) {
k_min = 0;
}
for (int p = k_min; p <= pos_q; p++) {
dst[((size_t) n * (size_t) T + (size_t) q) * (size_t) ring + (size_t) (p % ring)] = 0.0f;
}
}
}
}