codec: primed state snapshot LRU keyed by the reference content

seed_reference hashes the ICL reference codes and restores the conv
contexts, KV ring, and position from a per reference snapshot slot on
a repeat, saving the primed state device to device after a fresh
prime. The reference priming cost amortizes across repeated cloned
voice requests.
This commit is contained in:
Pascal
2026-07-05 14:34:27 +02:00
parent 7a4c799f9c
commit 1a680e8816
4 changed files with 172 additions and 3 deletions
+9
View File
@@ -354,6 +354,15 @@ readback skipped, matching the upstream reference-plus-generated
decode; the transformer receptive field (8 layers x window 72) exceeds decode; the transformer receptive field (8 layers x window 72) exceeds
any reference length, so the full prime is the exact one. any reference length, so the full prime is the exact one.
The primed state feeds a per-reference snapshot LRU
(`CODEC_SNAP_SLOTS`): after a fresh prime the conv contexts, KV ring,
and position copy device to device into a slot keyed by the FNV-1a
hash of the reference codes, and a later request with the same
reference restores the exact state in one pass of tensor copies
instead of re-decoding every reference frame. Slots allocate lazily
and evict least recently used, so the priming cost amortizes across
repeated cloned-voice requests without any registry coupling.
## Inference pipeline ## Inference pipeline
### Prompt assembly ### Prompt assembly
+9 -3
View File
@@ -112,9 +112,15 @@ struct codec_stream_decoder {
// Prime the codec state with the full ICL reference: every frame // Prime the codec state with the full ICL reference: every frame
// runs through the streaming decode with the audio discarded, so // runs through the streaming decode with the audio discarded, so
// the first generated frame sees the reference's exact causal // the first generated frame sees the reference's exact causal
// state. ref_kt is K major [K, ref_T]. Call once, after init and // state. A reference already primed restores its snapshot device
// before any push_frame. // 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) { 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); std::vector<int32_t> codes((size_t) K);
for (int t = 0; t < ref_T; t++) { for (int t = 0; t < ref_T; t++) {
for (int k = 0; k < K; k++) { for (int k = 0; k < K; k++) {
@@ -124,7 +130,7 @@ struct codec_stream_decoder {
return false; return false;
} }
} }
return true; return pipeline_codec_stream_snapshot(pc, key);
} }
// Decode one frame (K int32 codes, one per codebook) and emit its // Decode one frame (K int32 codes, one per codebook) and emit its
+118
View File
@@ -31,6 +31,10 @@ bool pipeline_codec_load(PipelineCodec * pc, const char * gguf_path, BackendPair
pc->stream_graph_ctx = NULL; pc->stream_graph_ctx = NULL;
pc->stream_gf = NULL; pc->stream_gf = NULL;
pc->stream_galloc = NULL; pc->stream_galloc = NULL;
for (int i = 0; i < CODEC_SNAP_SLOTS; i++) {
pc->snaps[i] = {};
}
pc->snap_stamp = 0;
if (!gf_load(&pc->gguf, gguf_path)) { if (!gf_load(&pc->gguf, gguf_path)) {
qt_log(QT_LOG_ERROR, "[Pipeline] failed to load %s", gguf_path); qt_log(QT_LOG_ERROR, "[Pipeline] failed to load %s", gguf_path);
@@ -363,6 +367,109 @@ bool pipeline_codec_stream_reset(PipelineCodec * pc) {
return true; return true;
} }
// FNV-1a 64 over the raw code bytes with T folded in, so a reference
// sharing a prefix with a longer one cannot alias its key.
uint64_t pipeline_codec_ref_key(const int32_t * codes, int K, int T) {
const uint8_t * p = (const uint8_t *) codes;
const size_t n = (size_t) K * (size_t) T * sizeof(int32_t);
uint64_t h = 1469598103934665603ULL;
for (size_t i = 0; i < n; i++) {
h = (h ^ p[i]) * 1099511628211ULL;
}
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.
static bool codec_snap_ensure(PipelineCodec * pc, CodecStateSnap * s) {
if (s->ctx) {
return true;
}
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++;
}
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++;
}
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)) {
ggml_dup_tensor(s->ctx, t);
}
for (struct ggml_tensor * t = ggml_get_first_tensor(pc->stream_kv.ctx); t;
t = ggml_get_next_tensor(pc->stream_kv.ctx, t)) {
ggml_dup_tensor(s->ctx, t);
}
s->buf = ggml_backend_alloc_ctx_tensors(s->ctx, pc->backend);
if (!s->buf) {
qt_log(QT_LOG_ERROR, "[Pipeline] snapshot backend allocation failed");
ggml_free(s->ctx);
s->ctx = NULL;
return false;
}
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) {
struct ggml_tensor * m = ggml_get_first_tensor(s->ctx);
for (struct ggml_tensor * t = ggml_get_first_tensor(pc->stream_ctx); t;
t = ggml_get_next_tensor(pc->stream_ctx, t), m = ggml_get_next_tensor(s->ctx, m)) {
ggml_backend_tensor_copy(save ? t : m, save ? m : t);
}
for (struct ggml_tensor * t = ggml_get_first_tensor(pc->stream_kv.ctx); t;
t = ggml_get_next_tensor(pc->stream_kv.ctx, t), m = ggml_get_next_tensor(s->ctx, m)) {
ggml_backend_tensor_copy(save ? t : m, save ? m : t);
}
}
bool pipeline_codec_stream_restore(PipelineCodec * pc, uint64_t key) {
if (!pc->stream_ready) {
return false;
}
for (int i = 0; i < CODEC_SNAP_SLOTS; i++) {
CodecStateSnap * s = &pc->snaps[i];
if (s->stamp == 0 || s->key != 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);
return true;
}
return false;
}
bool pipeline_codec_stream_snapshot(PipelineCodec * pc, uint64_t key) {
CodecStateSnap * lru = &pc->snaps[0];
for (int i = 1; i < CODEC_SNAP_SLOTS; i++) {
if (pc->snaps[i].stamp < lru->stamp) {
lru = &pc->snaps[i];
}
}
if (!codec_snap_ensure(pc, lru)) {
return false;
}
codec_snap_copy(pc, lru, true);
lru->key = key;
lru->pos = pc->stream_pos;
lru->stamp = ++pc->snap_stamp;
qt_log(QT_LOG_INFO, "[Pipeline] Codec state snapshot saved (%d frames)", lru->pos);
return true;
}
bool pipeline_codec_decode_stream(PipelineCodec * pc, const int32_t * codes, float * audio_out) { bool pipeline_codec_decode_stream(PipelineCodec * pc, const int32_t * codes, float * audio_out) {
const int K = TOKENIZER_NUM_CODEBOOKS; const int K = TOKENIZER_NUM_CODEBOOKS;
const int ring = CODEC_STREAM_RING; const int ring = CODEC_STREAM_RING;
@@ -660,6 +767,17 @@ void pipeline_codec_free(PipelineCodec * pc) {
pc->stream_ctx = NULL; pc->stream_ctx = NULL;
pc->stream_ready = false; pc->stream_ready = false;
} }
for (int i = 0; i < CODEC_SNAP_SLOTS; i++) {
CodecStateSnap * s = &pc->snaps[i];
if (s->buf) {
ggml_backend_buffer_free(s->buf);
}
if (s->ctx) {
ggml_free(s->ctx);
}
*s = {};
}
pc->snap_stamp = 0;
if (pc->enc_loaded) { if (pc->enc_loaded) {
quant_encode_free(&pc->qenc); quant_encode_free(&pc->qenc);
enc_down_free(&pc->enc_downsample); enc_down_free(&pc->enc_downsample);
+36
View File
@@ -46,6 +46,22 @@
#define TOKENIZER_NUM_CODEBOOKS 16 #define TOKENIZER_NUM_CODEBOOKS 16
#define TOKENIZER_CODE_BITS 11 #define TOKENIZER_CODE_BITS 11
// Primed stream state snapshots kept per reference, LRU evicted.
static const int CODEC_SNAP_SLOTS = 8;
// 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.
struct CodecStateSnap {
uint64_t key;
uint64_t stamp;
int pos;
struct ggml_context * ctx;
ggml_backend_buffer_t buf;
};
struct PipelineCodec { struct PipelineCodec {
GGUFModel gguf; GGUFModel gguf;
@@ -105,6 +121,14 @@ struct PipelineCodec {
struct ggml_tensor * stream_in_mask; struct ggml_tensor * stream_in_mask;
struct ggml_tensor * stream_out; struct ggml_tensor * stream_out;
// Snapshot LRU over the stream state: after an ICL reference
// priming the conv contexts, KV ring, and position copy device to
// device into the slot keyed by the reference content hash, so a
// repeated reference restores in one pass of tensor copies instead
// of re-decoding every reference frame.
CodecStateSnap snaps[CODEC_SNAP_SLOTS];
uint64_t snap_stamp;
// CPU mirror of the RVQ encode side, lazy-loaded on first encode call. // CPU mirror of the RVQ encode side, lazy-loaded on first encode call.
QwenQuantizerEncodeHost qenc_sem_host; QwenQuantizerEncodeHost qenc_sem_host;
QwenQuantizerEncodeHost qenc_aco_host; QwenQuantizerEncodeHost qenc_aco_host;
@@ -141,6 +165,18 @@ bool pipeline_codec_stream_reset(PipelineCodec * pc);
// the state without a readback (ICL reference priming). // the state without a readback (ICL reference priming).
bool pipeline_codec_decode_stream(PipelineCodec * pc, const int32_t * codes, float * audio_out); bool pipeline_codec_decode_stream(PipelineCodec * pc, const int32_t * codes, float * audio_out);
// 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);
// 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);
// Encode a 24 kHz mono waveform into RVQ codes. // Encode a 24 kHz mono waveform into RVQ codes.
// audio : [n_samples] f32 mono 24 kHz. Must be a multiple of // audio : [n_samples] f32 mono 24 kHz. Must be a multiple of
// TOKENIZER_HOP_LENGTH (1920); the caller is expected // TOKENIZER_HOP_LENGTH (1920); the caller is expected