416 lines
19 KiB
C++
416 lines
19 KiB
C++
#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,
|
|
// KV cached.
|
|
//
|
|
// Input:
|
|
// talker_hidden_last [hidden] f32 -- last position hidden state from
|
|
// the Talker forward (post final norm)
|
|
// c0 -- semantic code sampled from the
|
|
// Talker codec_head (codebook 0)
|
|
// Output:
|
|
// codes[16] = [c0, c1, ..., c15] -- the full set of codes for one
|
|
// frame, ready for decode through
|
|
// the codec
|
|
//
|
|
// 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.
|
|
//
|
|
// Architecture mirrors the Talker block, only differences are:
|
|
// - 5 layers instead of 28
|
|
// - plain 1D RoPE (no multimodal sections)
|
|
// - one private embedding table and one private linear head per
|
|
// acoustic codebook (1..15)
|
|
//
|
|
// The single-frame loop here recomputes the full graph at every step g
|
|
// (0..14) over a sequence of length g+2. With 5 layers and at most 16
|
|
// tokens per recompute this is sub-millisecond on modern GPUs.
|
|
|
|
#include "code-predictor-weights.h"
|
|
#include "debug.h"
|
|
#include "ggml-alloc.h"
|
|
#include "ggml-backend.h"
|
|
#include "ggml.h"
|
|
#include "kv-cache.h"
|
|
#include "qt-error.h"
|
|
#include "sampling.h"
|
|
#include "talker-weights.h"
|
|
|
|
#include <cmath>
|
|
#include <cstdint>
|
|
#include <cstdio>
|
|
#include <cstdlib>
|
|
#include <cstring>
|
|
#include <vector>
|
|
|
|
struct CodePredictorOutput {
|
|
// Sixteen codes: c0 from the talker plus c1..c15 from the predictor.
|
|
std::vector<int32_t> codes;
|
|
};
|
|
|
|
// Manual F32 attention chain for the code predictor block. Same shape
|
|
// contract as talker_attn_f32: q [hd, T, n_q_heads], k/v [hd, T_full,
|
|
// n_kv], output [hd, n_q_heads, T]. Used when use_flash_attn is false.
|
|
static struct ggml_tensor * code_predictor_attn_f32(struct ggml_context * ctx,
|
|
struct ggml_tensor * q,
|
|
struct ggml_tensor * k,
|
|
struct ggml_tensor * v,
|
|
struct ggml_tensor * mask,
|
|
float scale) {
|
|
struct ggml_tensor * scores = ggml_mul_mat(ctx, k, q);
|
|
scores = ggml_soft_max_ext(ctx, scores, mask, scale, 0.0f);
|
|
struct ggml_tensor * vt = ggml_cont(ctx, ggml_transpose(ctx, v));
|
|
struct ggml_tensor * out = ggml_mul_mat(ctx, vt, scores);
|
|
return ggml_cont(ctx, ggml_permute(ctx, out, 0, 2, 1, 3));
|
|
}
|
|
|
|
// One Qwen3 decoder block, KV cached. K and V for the T fresh positions
|
|
// are written into the cache at [n_past, n_past+T) on dim 1; the
|
|
// attention reads the contiguous slice [0, n_past+T). Returns the layer
|
|
// output [hidden, T]. use_flash_attn and clamp_fp16 follow the same
|
|
// contract as in talker-forward.h.
|
|
static struct ggml_tensor * code_predictor_layer_forward(struct ggml_context * ctx,
|
|
const CodePredictorWeights * cw,
|
|
const TalkerLayer & layer,
|
|
struct ggml_tensor * x,
|
|
struct ggml_tensor * positions,
|
|
struct ggml_tensor * mask,
|
|
struct ggml_tensor * k_cache,
|
|
struct ggml_tensor * v_cache,
|
|
int n_past,
|
|
int T,
|
|
bool use_flash_attn,
|
|
bool clamp_fp16,
|
|
struct ggml_cgraph * gf) {
|
|
const int n_q_heads = cw->num_attention_heads;
|
|
const int n_kv = cw->num_key_value_heads;
|
|
const int hd = cw->head_dim;
|
|
const float eps = cw->rms_norm_eps;
|
|
|
|
struct ggml_tensor * h = ggml_rms_norm(ctx, x, eps);
|
|
h = ggml_mul(ctx, h, layer.input_norm_w);
|
|
|
|
struct ggml_tensor * q = ggml_mul_mat(ctx, layer.attn.q_proj_w, h);
|
|
struct ggml_tensor * k = ggml_mul_mat(ctx, layer.attn.k_proj_w, h);
|
|
struct ggml_tensor * v = ggml_mul_mat(ctx, layer.attn.v_proj_w, h);
|
|
|
|
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_rms_norm(ctx, q, eps);
|
|
q = ggml_mul(ctx, q, layer.attn.q_norm_w);
|
|
k = ggml_rms_norm(ctx, k, eps);
|
|
k = ggml_mul(ctx, k, layer.attn.k_norm_w);
|
|
|
|
q = ggml_rope_ext(ctx, q, positions, NULL, hd, GGML_ROPE_TYPE_NEOX, 0, cw->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, cw->rope_theta, 1.0f, 0.0f, 1.0f, 0.0f,
|
|
0.0f);
|
|
|
|
// Write the fresh positions into the cache.
|
|
struct ggml_tensor * k_perm = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3)); // [hd, T, n_kv]
|
|
struct ggml_tensor * v_perm = ggml_cont(ctx, ggml_permute(ctx, v, 0, 2, 1, 3));
|
|
|
|
size_t k_off = (size_t) n_past * k_cache->nb[1];
|
|
size_t v_off = (size_t) n_past * v_cache->nb[1];
|
|
|
|
struct ggml_tensor * k_dst = ggml_view_3d(ctx, k_cache, hd, T, n_kv, k_cache->nb[1], k_cache->nb[2], k_off);
|
|
struct ggml_tensor * v_dst = ggml_view_3d(ctx, v_cache, hd, T, n_kv, v_cache->nb[1], v_cache->nb[2], v_off);
|
|
|
|
struct ggml_tensor * k_cpy = ggml_cpy(ctx, k_perm, k_dst);
|
|
struct ggml_tensor * v_cpy = ggml_cpy(ctx, v_perm, v_dst);
|
|
ggml_build_forward_expand(gf, k_cpy);
|
|
ggml_build_forward_expand(gf, v_cpy);
|
|
|
|
const int T_full = n_past + T;
|
|
struct ggml_tensor * k_full = ggml_view_3d(ctx, k_cache, hd, T_full, n_kv, k_cache->nb[1], k_cache->nb[2], 0);
|
|
struct ggml_tensor * v_full = ggml_view_3d(ctx, v_cache, hd, T_full, n_kv, v_cache->nb[1], v_cache->nb[2], 0);
|
|
|
|
// Q permute [hd, n_q_heads, T] -> [hd, T, n_q_heads] for flash_attn_ext.
|
|
struct ggml_tensor * q_p = ggml_permute(ctx, q, 0, 2, 1, 3);
|
|
|
|
// Clamp V before attention when clamp_fp16 is set, same rationale
|
|
// as the talker block: sub Ampere CUDA tensor cores accumulate in
|
|
// FP16 and a V projection overflow corrupts everything downstream.
|
|
if (clamp_fp16) {
|
|
v_full = ggml_clamp(ctx, v_full, -65504.0f, 65504.0f);
|
|
}
|
|
|
|
// Attention: fused flash kernel or manual F32 chain. Matches the
|
|
// working acestep qw3lm_build_attn pattern on the fused branch and
|
|
// the omnivoice qwen3_attn_f32 helper on the manual one.
|
|
float scale = 1.0f / sqrtf((float) hd);
|
|
struct ggml_tensor * attn;
|
|
if (use_flash_attn) {
|
|
attn = ggml_flash_attn_ext(ctx, q_p, k_full, v_full, mask, scale, 0.0f, 0.0f);
|
|
ggml_flash_attn_ext_set_prec(attn, GGML_PREC_F32);
|
|
} else {
|
|
attn = code_predictor_attn_f32(ctx, q_p, k_full, v_full, mask, scale);
|
|
}
|
|
|
|
attn = ggml_reshape_2d(ctx, attn, n_q_heads * hd, T);
|
|
|
|
struct ggml_tensor * o = ggml_mul_mat(ctx, layer.attn.o_proj_w, attn);
|
|
x = ggml_add(ctx, x, o);
|
|
if (clamp_fp16) {
|
|
x = ggml_clamp(ctx, x, -65504.0f, 65504.0f);
|
|
}
|
|
|
|
struct ggml_tensor * h2 = ggml_rms_norm(ctx, x, eps);
|
|
h2 = ggml_mul(ctx, h2, layer.post_attn_norm_w);
|
|
|
|
struct ggml_tensor * gate = ggml_mul_mat(ctx, layer.mlp.gate_proj_w, h2);
|
|
struct ggml_tensor * up = ggml_mul_mat(ctx, layer.mlp.up_proj_w, h2);
|
|
gate = ggml_silu(ctx, gate);
|
|
struct ggml_tensor * gu = ggml_mul(ctx, gate, up);
|
|
struct ggml_tensor * mlp = ggml_mul_mat(ctx, layer.mlp.down_proj_w, gu);
|
|
|
|
x = ggml_add(ctx, x, mlp);
|
|
if (clamp_fp16) {
|
|
x = ggml_clamp(ctx, x, -65504.0f, 65504.0f);
|
|
}
|
|
return x;
|
|
}
|
|
|
|
// Run one predictor pass: feed `T` fresh embeddings starting at cache
|
|
// position `n_past`, run all 5 layers, and pull the logits for the last
|
|
// position through lm_head[g_head]. The cache is written as a side
|
|
// effect so subsequent decode steps can append a single token.
|
|
// use_flash_attn / clamp_fp16 are forwarded as is to every layer.
|
|
static bool code_predictor_run(const CodePredictorWeights * cw,
|
|
KVCache * kv,
|
|
ggml_backend_sched_t sched,
|
|
const float * fresh_input,
|
|
int T,
|
|
int n_past,
|
|
int talker_hidden,
|
|
int g_head,
|
|
bool use_flash_attn,
|
|
bool clamp_fp16,
|
|
std::vector<float> * logits_out) {
|
|
const int vocab = cw->vocab_size;
|
|
const int n_layers = cw->num_hidden_layers;
|
|
const int T_full = n_past + T;
|
|
|
|
const int max_nodes = 48 * n_layers + 64;
|
|
const size_t arena_bytes = ggml_tensor_overhead() * max_nodes + ggml_graph_overhead_custom(max_nodes, false);
|
|
|
|
struct ggml_init_params gp = { arena_bytes, NULL, true };
|
|
struct ggml_context * gctx = ggml_init(gp);
|
|
if (!gctx) {
|
|
fprintf(stderr, "[CodePredictor] FATAL: ggml_init failed\n");
|
|
return false;
|
|
}
|
|
|
|
// Inputs: fresh embeddings (talker_hidden), positions, attention mask
|
|
struct ggml_tensor * x_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F32, talker_hidden, T);
|
|
struct ggml_tensor * pos_in = ggml_new_tensor_1d(gctx, GGML_TYPE_I32, T);
|
|
struct ggml_tensor * mask_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F16, T_full, T);
|
|
ggml_set_name(x_in, "sub_input");
|
|
ggml_set_name(pos_in, "positions");
|
|
ggml_set_name(mask_in, "causal_mask");
|
|
|
|
struct ggml_cgraph * gf = ggml_new_graph_custom(gctx, max_nodes, false);
|
|
|
|
// small_to_mtp projection: Linear(talker_hidden -> hidden) with bias.
|
|
// When absent (Identity case) the input is already at predictor hidden.
|
|
struct ggml_tensor * h = x_in;
|
|
if (cw->mtp_proj_w) {
|
|
h = ggml_mul_mat(gctx, cw->mtp_proj_w, h);
|
|
if (cw->mtp_proj_b) {
|
|
h = ggml_add(gctx, h, cw->mtp_proj_b);
|
|
}
|
|
ggml_set_name(h, "mtp_proj_out");
|
|
}
|
|
|
|
for (int l = 0; l < n_layers; l++) {
|
|
h = code_predictor_layer_forward(gctx, cw, cw->layers[(size_t) l], h, pos_in, mask_in, kv->k[(size_t) l],
|
|
kv->v[(size_t) l], n_past, T, use_flash_attn, clamp_fp16, gf);
|
|
}
|
|
|
|
struct ggml_tensor * h_final = ggml_rms_norm(gctx, h, cw->rms_norm_eps);
|
|
h_final = ggml_mul(gctx, h_final, cw->norm_w);
|
|
|
|
struct ggml_tensor * logits = ggml_mul_mat(gctx, cw->lm_head[(size_t) g_head], h_final);
|
|
ggml_set_name(logits, "logits");
|
|
ggml_set_output(logits);
|
|
ggml_build_forward_expand(gf, logits);
|
|
|
|
if (!ggml_backend_sched_alloc_graph(sched, gf)) {
|
|
fprintf(stderr, "[CodePredictor] FATAL: graph allocation failed\n");
|
|
ggml_backend_sched_reset(sched);
|
|
ggml_free(gctx);
|
|
return false;
|
|
}
|
|
|
|
ggml_backend_tensor_set(x_in, fresh_input, 0, (size_t) T * (size_t) talker_hidden * sizeof(float));
|
|
|
|
{
|
|
std::vector<int32_t> pos((size_t) T);
|
|
for (int i = 0; i < T; i++) {
|
|
pos[(size_t) i] = n_past + i;
|
|
}
|
|
ggml_backend_tensor_set(pos_in, pos.data(), 0, (size_t) T * sizeof(int32_t));
|
|
}
|
|
|
|
{
|
|
std::vector<ggml_fp16_t> mask((size_t) T * (size_t) T_full);
|
|
const ggml_fp16_t zero = ggml_fp32_to_fp16(0.0f);
|
|
const ggml_fp16_t neg_inf = ggml_fp32_to_fp16(-INFINITY);
|
|
for (size_t i = 0; i < mask.size(); i++) {
|
|
mask[i] = neg_inf;
|
|
}
|
|
for (int q = 0; q < T; q++) {
|
|
const int q_pos = n_past + q;
|
|
for (int k = 0; k <= q_pos; k++) {
|
|
mask[(size_t) q * (size_t) T_full + (size_t) k] = zero;
|
|
}
|
|
}
|
|
ggml_backend_tensor_set(mask_in, mask.data(), 0, mask.size() * sizeof(ggml_fp16_t));
|
|
}
|
|
|
|
if (ggml_backend_sched_graph_compute(sched, gf) != GGML_STATUS_SUCCESS) {
|
|
fprintf(stderr, "[CodePredictor] FATAL: graph compute failed\n");
|
|
ggml_backend_sched_reset(sched);
|
|
ggml_free(gctx);
|
|
return false;
|
|
}
|
|
|
|
logits_out->resize((size_t) vocab);
|
|
size_t row_bytes = (size_t) vocab * sizeof(float);
|
|
ggml_backend_tensor_get(logits, logits_out->data(), (size_t) (T - 1) * row_bytes, row_bytes);
|
|
|
|
kv->cur_len = T_full;
|
|
|
|
ggml_backend_sched_reset(sched);
|
|
ggml_free(gctx);
|
|
return true;
|
|
}
|
|
|
|
// Read one row of an embedding table to f32. Reads from the backend
|
|
// (the predictor weights live there) via ggml_backend_tensor_get,
|
|
// dispatched through ggml_get_type_traits so quants are accepted.
|
|
static void embed_row_from_backend(struct ggml_tensor * t, int row_id, int dim, float * dst) {
|
|
if (t->ne[0] != dim) {
|
|
qt_throw("[CodePredictor] embed dim mismatch %lld vs %d", (long long) t->ne[0], dim);
|
|
}
|
|
if (row_id < 0 || row_id >= (int) t->ne[1]) {
|
|
qt_throw("[CodePredictor] row %d out of range (vocab=%lld)", row_id, (long long) t->ne[1]);
|
|
}
|
|
const size_t row_bytes = ggml_row_size(t->type, dim);
|
|
if (t->type == GGML_TYPE_F32) {
|
|
ggml_backend_tensor_get(t, dst, (size_t) row_id * row_bytes, row_bytes);
|
|
return;
|
|
}
|
|
const struct ggml_type_traits * tt = ggml_get_type_traits(t->type);
|
|
if (!tt || !tt->to_float) {
|
|
qt_throw("[CodePredictor] unsupported embed dtype %d", (int) t->type);
|
|
}
|
|
std::vector<uint8_t> tmp(row_bytes);
|
|
ggml_backend_tensor_get(t, tmp.data(), (size_t) row_id * row_bytes, row_bytes);
|
|
tt->to_float(tmp.data(), dst, dim);
|
|
}
|
|
|
|
// Run the predictor for one audio frame. Caller passes the talker hidden
|
|
// state for the current frame and the already-sampled c0. Sampling
|
|
// parameters control greedy (temperature <= 0) vs stochastic. subseq_base
|
|
// is the Philox subsequence of the c0 sample for this step; the 15
|
|
// acoustic samples consume subseq_base + 1 .. subseq_base + 15.
|
|
// Returns the full vector of 16 codes. dump_dir may be NULL.
|
|
// use_flash_attn / clamp_fp16 are forwarded as is to every internal run.
|
|
static bool code_predictor_step(const TalkerWeights * tw,
|
|
const CodePredictorWeights * cw,
|
|
KVCache * kv,
|
|
ggml_backend_sched_t sched,
|
|
const float * talker_hidden_last,
|
|
int c0,
|
|
float temperature,
|
|
int top_k,
|
|
float top_p,
|
|
int64_t seed,
|
|
int64_t subseq_base,
|
|
bool use_flash_attn,
|
|
bool clamp_fp16,
|
|
const char * dump_dir,
|
|
CodePredictorOutput * out) {
|
|
// sub_input slots live at the talker hidden dimension: both the
|
|
// talker last hidden and the codec_embedding rows feeding the sub
|
|
// network are talker sized in the upstream checkpoint. The graph's
|
|
// mtp_proj brings them down to predictor hidden when present.
|
|
const int talker_hidden = tw->hidden_size;
|
|
const int n_acoustic = cw->num_acoustic_codebooks;
|
|
|
|
if (n_acoustic + 1 > kv->max_seq_len) {
|
|
fprintf(stderr, "[CodePredictor] FATAL: frame width %d exceeds cache max_seq_len %d\n", n_acoustic + 1,
|
|
kv->max_seq_len);
|
|
return false;
|
|
}
|
|
|
|
out->codes.assign((size_t) (n_acoustic + 1), 0);
|
|
out->codes[0] = c0;
|
|
|
|
// Prefill: two positions, talker_hidden_last and embed_talker(c0).
|
|
kv_cache_reset(kv);
|
|
std::vector<float> prefill_input((size_t) 2 * (size_t) talker_hidden, 0.0f);
|
|
std::memcpy(prefill_input.data(), talker_hidden_last, (size_t) talker_hidden * sizeof(float));
|
|
embed_row_from_backend(tw->codec_embedding, c0, talker_hidden, prefill_input.data() + (size_t) talker_hidden);
|
|
|
|
std::vector<float> logits;
|
|
if (!code_predictor_run(cw, kv, sched, prefill_input.data(), 2, 0, talker_hidden, 0, use_flash_attn, clamp_fp16,
|
|
&logits)) {
|
|
return false;
|
|
}
|
|
{
|
|
float u_g = 0.0f;
|
|
int cg = sample_top_k_p(logits.data(), (int) logits.size(), temperature, top_k, top_p, 1.0f, nullptr, 0, seed,
|
|
subseq_base + 1, &u_g);
|
|
if (subseq_base + 1 < 32) {
|
|
fprintf(stderr, "[Sample-CP] g=0 c=%d u=%.10f subseq=%lld\n", cg, (double) u_g,
|
|
(long long) (subseq_base + 1));
|
|
}
|
|
if (cg < 0) {
|
|
fprintf(stderr, "[CodePredictor] FATAL: sample returned no candidate at g=0\n");
|
|
return false;
|
|
}
|
|
out->codes[1] = cg;
|
|
}
|
|
|
|
// Decode loop: 14 single-token steps. At step g (g=1..14) we feed
|
|
// the embedding of the code we just sampled and read lm_head[g].
|
|
std::vector<float> step_input((size_t) talker_hidden);
|
|
for (int g = 1; g < n_acoustic; g++) {
|
|
embed_row_from_backend(cw->codec_embedding[(size_t) (g - 1)], out->codes[(size_t) g], talker_hidden,
|
|
step_input.data());
|
|
if (!code_predictor_run(cw, kv, sched, step_input.data(), 1, kv->cur_len, talker_hidden, g, use_flash_attn,
|
|
clamp_fp16, &logits)) {
|
|
return false;
|
|
}
|
|
float u_g = 0.0f;
|
|
int cg = sample_top_k_p(logits.data(), (int) logits.size(), temperature, top_k, top_p, 1.0f, nullptr, 0, seed,
|
|
subseq_base + 1 + g, &u_g);
|
|
if (subseq_base + 1 + g < 32) {
|
|
fprintf(stderr, "[Sample-CP] g=%d c=%d u=%.10f subseq=%lld\n", g, cg, (double) u_g,
|
|
(long long) (subseq_base + 1 + g));
|
|
}
|
|
if (cg < 0) {
|
|
fprintf(stderr, "[CodePredictor] FATAL: sample returned no candidate at g=%d\n", g);
|
|
return false;
|
|
}
|
|
out->codes[(size_t) (g + 1)] = cg;
|
|
}
|
|
|
|
if (dump_dir) {
|
|
DebugDumper d;
|
|
debug_init(&d, dump_dir);
|
|
std::vector<int32_t> codes32(out->codes.begin(), out->codes.end());
|
|
int n = (int) codes32.size();
|
|
debug_dump_i32_as_f32(&d, "codes-step0", codes32.data(), &n, 1);
|
|
}
|
|
|
|
return true;
|
|
}
|