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
@@ -112,9 +112,15 @@ struct codec_stream_decoder {
|
||||
// Prime the codec state with the full ICL reference: every frame
|
||||
// runs through the streaming decode with the audio discarded, so
|
||||
// the first generated frame sees the reference's exact causal
|
||||
// state. ref_kt is K major [K, ref_T]. Call once, after init and
|
||||
// before any push_frame.
|
||||
// state. A reference already primed restores its snapshot device
|
||||
// to device instead of re-decoding; a fresh one saves its primed
|
||||
// state into the LRU. ref_kt is K major [K, ref_T]. Call once,
|
||||
// after init and before any push_frame.
|
||||
bool seed_reference(PipelineCodec * pc, const int32_t * ref_kt, int ref_T) {
|
||||
const uint64_t key = pipeline_codec_ref_key(ref_kt, K, ref_T);
|
||||
if (pipeline_codec_stream_restore(pc, key)) {
|
||||
return true;
|
||||
}
|
||||
std::vector<int32_t> codes((size_t) K);
|
||||
for (int t = 0; t < ref_T; t++) {
|
||||
for (int k = 0; k < K; k++) {
|
||||
@@ -124,7 +130,7 @@ struct codec_stream_decoder {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
return pipeline_codec_stream_snapshot(pc, key);
|
||||
}
|
||||
|
||||
// Decode one frame (K int32 codes, one per codebook) and emit its
|
||||
|
||||
@@ -31,6 +31,10 @@ bool pipeline_codec_load(PipelineCodec * pc, const char * gguf_path, BackendPair
|
||||
pc->stream_graph_ctx = NULL;
|
||||
pc->stream_gf = 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)) {
|
||||
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;
|
||||
}
|
||||
|
||||
// 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) {
|
||||
const int K = TOKENIZER_NUM_CODEBOOKS;
|
||||
const int ring = CODEC_STREAM_RING;
|
||||
@@ -660,6 +767,17 @@ void pipeline_codec_free(PipelineCodec * pc) {
|
||||
pc->stream_ctx = NULL;
|
||||
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) {
|
||||
quant_encode_free(&pc->qenc);
|
||||
enc_down_free(&pc->enc_downsample);
|
||||
|
||||
@@ -46,6 +46,22 @@
|
||||
#define TOKENIZER_NUM_CODEBOOKS 16
|
||||
#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 {
|
||||
GGUFModel gguf;
|
||||
|
||||
@@ -105,6 +121,14 @@ struct PipelineCodec {
|
||||
struct ggml_tensor * stream_in_mask;
|
||||
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.
|
||||
QwenQuantizerEncodeHost qenc_sem_host;
|
||||
QwenQuantizerEncodeHost qenc_aco_host;
|
||||
@@ -141,6 +165,18 @@ bool pipeline_codec_stream_reset(PipelineCodec * pc);
|
||||
// the state without a readback (ICL reference priming).
|
||||
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.
|
||||
// audio : [n_samples] f32 mono 24 kHz. Must be a multiple of
|
||||
// TOKENIZER_HOP_LENGTH (1920); the caller is expected
|
||||
|
||||
Reference in New Issue
Block a user