From 79a99219c6afb5db5177365ec3383726d9e673f7 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Fri, 31 Jul 2026 16:00:33 +0200 Subject: [PATCH] (wip) subgraph --- tools/mtmd/models/models.h | 72 +++++++++++++++++------------- tools/mtmd/models/qwen3tts-gen.cpp | 36 +++++++-------- tools/mtmd/mtmd.cpp | 6 +-- tools/mtmd/mtmd.h | 28 ++++++++---- 4 files changed, 79 insertions(+), 63 deletions(-) diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h index 0cf622fa5f..1331ac8ac3 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -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 & k_cache, - std::vector & 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 & k_cache, - std::vector & 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 & k_cache, + std::vector & 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 & k_cache, + std::vector & 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). diff --git a/tools/mtmd/models/qwen3tts-gen.cpp b/tools/mtmd/models/qwen3tts-gen.cpp index 0014fb6dae..895ee77016 100644 --- a/tools/mtmd/models/qwen3tts-gen.cpp +++ b/tools/mtmd/models/qwen3tts-gen.cpp @@ -3,7 +3,7 @@ #include // 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 & k_cache, std::vector & 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 & k_cache, std::vector & 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 diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 01e9e1fc0f..fb4d7357dc 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -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; diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h index 05520810a4..29fd274860 100644 --- a/tools/mtmd/mtmd.h +++ b/tools/mtmd/mtmd.h @@ -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);