(wip) subgraph

This commit is contained in:
Xuan Son Nguyen
2026-07-31 16:00:33 +02:00
parent 84559b46a0
commit 79a99219c6
4 changed files with 79 additions and 63 deletions
+40 -32
View File
@@ -235,42 +235,50 @@ struct clip_graph_qwen3tts_gen : clip_graph {
int top_k;
float top_p;
ggml_tensor * cache_set(ggml_tensor * cache, int row_idx, ggml_tensor * value) const;
ggml_tensor * do_sampling(ggml_tensor * logits, ggml_tensor * inp_rand, int top_k, float top_p) const;
//
// code_gen: backbone hidden state + sampled code0 -> 16 RVQ codes.
// MTP-style autoregressive code predictor: one token per codebook, causal KV cache.
//
struct code_gen : clip_graph {
code_gen(const clip_graph & parent, int top_k, float top_p)
: clip_graph(parent), top_k(top_k), top_p(top_p) {}
ggml_tensor * const_i32(ggml_tensor * anchor, float value) const;
ggml_tensor * causal_mask_row(int64_t n_kv_pad, int pos) const;
ggml_tensor * project_in(ggml_tensor * cur) const;
int top_k;
float top_p;
ggml_tensor * layer_forward(
ggml_tensor * cur,
const clip_layer & layer,
ggml_tensor * inp_pos,
ggml_tensor * kq_mask,
ggml_tensor *& k_cache_layer,
ggml_tensor *& v_cache_layer,
int64_t n_kv_pad,
int pos,
int il) const;
ggml_tensor * cache_set(ggml_tensor * cache, int row_idx, ggml_tensor * value) const;
ggml_tensor * do_sampling(ggml_tensor * logits, ggml_tensor * inp_rand) const;
void prefill(
std::vector<ggml_tensor *> & k_cache,
std::vector<ggml_tensor *> & v_cache,
ggml_tensor *& out_code_cache,
ggml_tensor * h_state,
ggml_tensor * code0_embd,
ggml_tensor * inp_rand,
int top_k,
float top_p) const;
ggml_tensor * const_i32(ggml_tensor * anchor, float value) const;
ggml_tensor * causal_mask_row(int64_t n_kv_pad, int pos) const;
ggml_tensor * project_in(ggml_tensor * cur) const;
ggml_tensor * step(
std::vector<ggml_tensor *> & k_cache,
std::vector<ggml_tensor *> & v_cache,
ggml_tensor * out_code_cache,
ggml_tensor * inp_rand,
int step_idx,
int top_k,
float top_p) const;
ggml_tensor * layer_forward(
ggml_tensor * cur,
const clip_layer & layer,
ggml_tensor * inp_pos,
ggml_tensor * kq_mask,
ggml_tensor *& k_cache_layer,
ggml_tensor *& v_cache_layer,
int64_t n_kv_pad,
int pos,
int il) const;
void prefill(
std::vector<ggml_tensor *> & k_cache,
std::vector<ggml_tensor *> & v_cache,
ggml_tensor *& out_code_cache,
ggml_tensor * h_state,
ggml_tensor * code0_embd,
ggml_tensor * inp_rand) const;
ggml_tensor * step(
std::vector<ggml_tensor *> & k_cache,
std::vector<ggml_tensor *> & v_cache,
ggml_tensor * out_code_cache,
ggml_tensor * inp_rand,
int step_idx) const;
};
//
// code2wav: RVQ codes -> raw PCM (quantizer + pre_conv + pre_transformer + upsample + DAC).
+17 -19
View File
@@ -3,7 +3,7 @@
#include <string>
// on-device sampling: top-k, top-p, then a random draw
ggml_tensor * clip_graph_qwen3tts_gen::do_sampling(ggml_tensor * logits, ggml_tensor * inp_rand, int top_k, float top_p) const {
ggml_tensor * clip_graph_qwen3tts_gen::code_gen::do_sampling(ggml_tensor * logits, ggml_tensor * inp_rand) const {
logits = ggml_reshape_1d(ctx0, logits, ggml_nelements(logits));
const int64_t n_vocab = logits->ne[0];
@@ -68,7 +68,7 @@ ggml_tensor * clip_graph_qwen3tts_gen::do_sampling(ggml_tensor * logits, ggml_te
}
// returns a new cache with row row_idx set to value
ggml_tensor * clip_graph_qwen3tts_gen::cache_set(ggml_tensor * cache, int row_idx, ggml_tensor * value) const {
ggml_tensor * clip_graph_qwen3tts_gen::code_gen::cache_set(ggml_tensor * cache, int row_idx, ggml_tensor * value) const {
const int64_t n_embd = cache->ne[0];
const int64_t n_cache = cache->ne[1];
GGML_ASSERT(row_idx >= 0 && row_idx < n_cache);
@@ -97,7 +97,7 @@ ggml_tensor * clip_graph_qwen3tts_gen::cache_set(ggml_tensor * cache, int row_id
// builds a const i32 value with no host upload: view any tensor, cast to
// f32 (ggml_scale only supports f32), scale it to 0, add the value, cast to i32
ggml_tensor * clip_graph_qwen3tts_gen::const_i32(ggml_tensor * anchor, float value) const {
ggml_tensor * clip_graph_qwen3tts_gen::code_gen::const_i32(ggml_tensor * anchor, float value) const {
ggml_tensor * v = ggml_view_1d(ctx0, anchor, 1, 0);
if (v->type != GGML_TYPE_F32) {
v = ggml_cast(ctx0, v, GGML_TYPE_F32);
@@ -106,7 +106,7 @@ ggml_tensor * clip_graph_qwen3tts_gen::const_i32(ggml_tensor * anchor, float val
}
// causal keep-mask row for a query at position pos, window size n_kv_pad
ggml_tensor * clip_graph_qwen3tts_gen::causal_mask_row(int64_t n_kv_pad, int pos) const {
ggml_tensor * clip_graph_qwen3tts_gen::code_gen::causal_mask_row(int64_t n_kv_pad, int pos) const {
ggml_tensor * ones = ggml_fill(ctx0, ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_kv_pad, n_kv_pad), 1.0f);
ggml_tensor * keep = ggml_tri(ctx0, ones, GGML_TRI_TYPE_LOWER_DIAG);
ggml_tensor * row = ggml_view_1d(ctx0, keep, n_kv_pad, (size_t) pos * keep->nb[1]);
@@ -115,7 +115,7 @@ ggml_tensor * clip_graph_qwen3tts_gen::causal_mask_row(int64_t n_kv_pad, int pos
}
// talker hidden size -> predictor hidden size (small_to_mtp_projection)
ggml_tensor * clip_graph_qwen3tts_gen::project_in(ggml_tensor * cur) const {
ggml_tensor * clip_graph_qwen3tts_gen::code_gen::project_in(ggml_tensor * cur) const {
if (!model.gen_code_proj_in_w) {
return cur;
}
@@ -128,7 +128,7 @@ ggml_tensor * clip_graph_qwen3tts_gen::project_in(ggml_tensor * cur) const {
// one transformer layer at a single new position pos; writes k/v into
// k_cache_layer/v_cache_layer at row pos
ggml_tensor * clip_graph_qwen3tts_gen::layer_forward(
ggml_tensor * clip_graph_qwen3tts_gen::code_gen::layer_forward(
ggml_tensor * cur,
const clip_layer & layer,
ggml_tensor * inp_pos,
@@ -192,15 +192,13 @@ ggml_tensor * clip_graph_qwen3tts_gen::layer_forward(
// position 0: hidden bridge, no sampling, only seeds the k/v cache.
// position 1: embed(code0) via the talker's out_embd table, sample with
// lm_head[0], write out_code_cache[1].
void clip_graph_qwen3tts_gen::prefill(
void clip_graph_qwen3tts_gen::code_gen::prefill(
std::vector<ggml_tensor *> & k_cache,
std::vector<ggml_tensor *> & v_cache,
ggml_tensor *& out_code_cache,
ggml_tensor * h_state,
ggml_tensor * code0_embd,
ggml_tensor * inp_rand,
int top_k,
float top_p) const {
ggml_tensor * inp_rand) const {
const int64_t n_kv_pad = k_cache[0]->ne[1];
{
@@ -228,7 +226,7 @@ void clip_graph_qwen3tts_gen::prefill(
ggml_tensor * head_g = ggml_view_2d(ctx0, head_w, head_w->ne[0], head_w->ne[1], head_w->nb[1], 0); // lm_head[0]
ggml_tensor * logits = ggml_mul_mat(ctx0, head_g, cur);
ggml_tensor * sampled = do_sampling(logits, inp_rand, top_k, top_p);
ggml_tensor * sampled = do_sampling(logits, inp_rand);
out_code_cache = cache_set(out_code_cache, 1, sampled);
}
}
@@ -242,14 +240,12 @@ void clip_graph_qwen3tts_gen::prefill(
// out_code_cache: [1, n_codes] I32. inp_rand: [1] F32 draw for this step.
// Create all input tensors in build(), not here.
// step_idx range: [1, n_acoustic - 1]. Returns the new out_code_cache.
ggml_tensor * clip_graph_qwen3tts_gen::step(
ggml_tensor * clip_graph_qwen3tts_gen::code_gen::step(
std::vector<ggml_tensor *> & k_cache,
std::vector<ggml_tensor *> & v_cache,
ggml_tensor * out_code_cache,
ggml_tensor * inp_rand,
int step_idx,
int top_k,
float top_p) const {
int step_idx) const {
const int64_t n_acoustic = model.gen_code_head_w->ne[2];
GGML_ASSERT(step_idx >= 1 && step_idx < n_acoustic);
GGML_ASSERT(k_cache.size() == model.layers.size());
@@ -290,7 +286,7 @@ ggml_tensor * clip_graph_qwen3tts_gen::step(
ggml_tensor * logits = ggml_mul_mat(ctx0, head_g, cur);
cb(logits, "step_logits", step_idx);
ggml_tensor * sampled = do_sampling(logits, inp_rand, top_k, top_p);
ggml_tensor * sampled = do_sampling(logits, inp_rand);
cb(sampled, "step_sampled", step_idx);
return cache_set(out_code_cache, pos, sampled);
@@ -568,20 +564,22 @@ ggml_cgraph * clip_graph_qwen3tts_gen::build() {
v_cache[il] = ggml_fill(ctx0, ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, d_head * n_head_kv, n_kv_pad), 0.0f);
}
code_gen cg(*this, top_k, top_p);
ggml_tensor * out_code_cache = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, 1, n_codes);
out_code_cache = cache_set(out_code_cache, 0, code0);
out_code_cache = cg.cache_set(out_code_cache, 0, code0);
ggml_tensor * inp_rand0 = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, 1);
ggml_set_name(inp_rand0, "inp_rand_0");
ggml_set_input(inp_rand0);
prefill(k_cache, v_cache, out_code_cache, h_state, code0_embd, inp_rand0, top_k, top_p);
cg.prefill(k_cache, v_cache, out_code_cache, h_state, code0_embd, inp_rand0);
for (int g = 1; g < n_acoustic; g++) {
ggml_tensor * inp_rand = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, 1);
ggml_set_name(inp_rand, ("inp_rand_" + std::to_string(g)).c_str());
ggml_set_input(inp_rand);
out_code_cache = step(k_cache, v_cache, out_code_cache, inp_rand, g, top_k, top_p);
out_code_cache = cg.step(k_cache, v_cache, out_code_cache, inp_rand, g);
}
// output 1: raw PCM audio for this frame, decoded from the 16 sampled codes
+3 -3
View File
@@ -1579,7 +1579,7 @@ mtmd_gen_audio_type mtmd_gen_audio_get_type(const mtmd_context * ctx) {
}
}
static int32_t mtmd_gen_audio_impl(mtmd_context * ctx, const mtmd_gen_inp * inp, mtmd_gen_out * out) {
static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_inp * inp, mtmd_gen_out * out) {
clip_ctx * ctx_clip = ctx->ctx_gen_a;
if (!ctx_clip) {
LOG_ERR("%s: model does not support audio generation\n", __func__);
@@ -1628,9 +1628,9 @@ static int32_t mtmd_gen_audio_impl(mtmd_context * ctx, const mtmd_gen_inp * inp,
return 0;
}
int32_t mtmd_gen_audio(mtmd_context * ctx, const struct mtmd_gen_inp * inp, struct mtmd_gen_out * out) {
int32_t mtmd_gen_audio_process(mtmd_context * ctx, const struct mtmd_gen_inp * inp, struct mtmd_gen_out * out) {
try {
return mtmd_gen_audio_impl(ctx, inp, out);
return mtmd_gen_audio_process_impl(ctx, inp, out);
} catch (const std::exception & e) {
LOG_ERR("%s: error: %s\n", __func__, e.what());
return 1;
+19 -9
View File
@@ -336,25 +336,35 @@ enum mtmd_gen_audio_type {
};
MTMD_API mtmd_gen_audio_type mtmd_gen_audio_get_type(const mtmd_context * ctx);
enum mtmd_gen_process_type {
MTMD_GEN_PROCESS_TYPE_GEN_CODE, // h_state to codes
MTMD_GEN_PROCESS_TYPE_CODE2WAV, // codes to raw PCM audio
};
struct mtmd_gen_inp {
int32_t code0; // the sampled codebook 0 entry from backbone
float * embd; // the hidden state from backbone, size = n_embd * n_pos
size_t n_embd; // only for validation
mtmd_gen_process_type type;
// sampling params
// for MTMD_GEN_PROCESS_TYPE_GEN_CODE
int32_t code0; // the sampled codebook 0 entry from backbone
float * embd; // the hidden state from backbone, must have n_text_embd elements
int32_t top_k;
float top_p;
// for MTMD_GEN_PROCESS_TYPE_CODE2WAV
int32_t * codes; // the sampled codebook entries, must have n_codes elements
size_t n_codes;
};
struct mtmd_gen_out {
float * embd; // the generated hidden state, to be fed back to backbone
size_t n_embd; // only for validation
// note: output memory is allocated by the context, valid until next process() call
// out: raw PCM samples (F32) decoded for this frame; owned by mtmd_context,
// valid until the next mtmd_gen_audio() call, caller does not allocate this
// for MTMD_GEN_PROCESS_TYPE_GEN_CODE
const float * embd; // the generated hidden state, to be fed back to backbone
// it must have n_text_embd elements
// for MTMD_GEN_PROCESS_TYPE_CODE2WAV
const float * audio;
size_t n_samples;
};
MTMD_API int32_t mtmd_gen_audio(mtmd_context * ctx,
MTMD_API int32_t mtmd_gen_audio_process(mtmd_context * ctx,
const struct mtmd_gen_inp * inp,
struct mtmd_gen_out * out);