talker, code predictor: persistent KV cache

Adds a per-layer K/V ring (kv-cache.h) backed by a dedicated backend
buffer, sized at init for max_seq_len positions. The talker holds a
4096 position cache (896 MB f32) for the LM context, the code
predictor a 16 position cache (~80 KB) reset every frame.

talker_forward splits into prefill (resets the cache and writes T
positions in one shot) and decode (appends one position, reads the
[0, cur_len+1) window). code_predictor_step does the same with a T=2
prefill plus 14 single token decodes.

Bit identical audio output, validated by sha256 against the pre KV
cache run on a 64 frame F32 reference seed=42. Walltime drops ~10%
on a single utterance ; the win scales with sequence length and
unlocks frame by frame streaming.
This commit is contained in:
Pascal
2026-05-11 06:44:39 +02:00
parent 05ecde271a
commit 0a7023bd0b
7 changed files with 430 additions and 153 deletions
+9 -6
View File
@@ -1,6 +1,7 @@
#pragma once
// code-predictor-forward.h : run the 5-layer Qwen3 code predictor over a
// growing context to produce the 15 acoustic codes of one audio frame.
// growing context to produce the 15 acoustic codes of one audio frame,
// KV cached.
//
// Input :
// talker_hidden_last [hidden] f32 -- last position hidden state from
@@ -12,14 +13,15 @@
// frame, ready for decode through
// the codec
//
// Phase 4.5 runs the predictor without a KV cache : every step rebuilds
// the full graph over a context of length g+2 (g being the predictor
// step, 0..14). With 5 layers and at most 16 tokens this is well below
// the threshold where caching would matter. A KV-cached variant lands
// in the generation loop phase.
// The predictor cache is local to a single frame : we reset it at every
// frame, prefill the first two positions (talker_hidden + embed(c0)),
// then decode 14 single-token steps. Total work drops from
// O(sum_{g=0..14} (g+2)^2) = O(1496 token-steps) to O(16) per frame,
// roughly 90x for the inner loop.
#include "code-predictor-weights.h"
#include "ggml-backend.h"
#include "kv-cache.h"
#include "sampling.h"
#include "talker-weights.h"
@@ -39,6 +41,7 @@ struct CodePredictorOutput {
// Returns the full vector of 16 codes. dump_dir may be NULL.
bool code_predictor_step(const TalkerWeights * tw,
const CodePredictorWeights * cw,
KVCache * kv,
ggml_backend_sched_t sched,
const float * talker_hidden_last,
int c0,