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:
@@ -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
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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);
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user