Skip to content

mtmd: refactor mtmd gen API, add state object #30902

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ngxson merged 2 commits into from
Oct 21, 2026
Merged

mtmd: refactor mtmd gen API, add state object #30902

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c646c16
mtmd: add state object for mtmd_gen
ngxson Oct 20, 2026
f80943f
Merge branch 'master' into xsn/mtmd_audio_gen_refactor
ngxson Oct 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion tools/mtmd/README-dev.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,7 @@ For example, Qwen3-TTS:
Due to wide variety of audio generation pipelines, the `mtmd_gen_audio` system is designed to be flexible and reusable by new models.

`mtmd_gen_audio` is split into 2 main API:
- Core API `mtmd.h`: handles main inference. Important: the API surface must be stateless; caller must handle state management and audio frame accumulation.
- Core API `mtmd.h`: handles main inference. Important: `mtmd_gen_state` holds only the per-sequence data that the mmproj itself writes and reads back (RNG, decoder KV / conv state, output buffers). Data that goes to or comes from the backbone (hidden states, codes, PCM) must stay explicit in `mtmd_gen_inp` / `mtmd_gen_out`; the caller drives the backbone and handles audio frame accumulation.
- Helper API `mtmd-helper.h`: provides a model-agnostic stateful API. Usage example can be found in the `tools/tts` directory.

### Checklist for porting new audio generation models to mtmd
Expand Down
10 changes: 10 additions & 0 deletions tools/mtmd/clip-impl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
#include <vector>
#include <memory>
#include <fstream>
#include <random>

#ifdef _WIN32
#ifndef NOMINMAX
Expand Down Expand Up @@ -463,6 +464,15 @@ struct clip_enc_state {
std::vector<ggml_tensor *> v;
};

// per-sequence state of audio generation, see clip_gen_state_init()
struct clip_gen_state {
std::mt19937 rng;

ggml_context_ptr ctx_data;
ggml_backend_buffer_ptr buf;
std::vector<ggml_tensor *> slots; // same order as list_gen_state_slots(), F32 [ne0, ne1]
};

// forward declaration
// TODO: improve this later
struct clip_ctx;
Expand Down
101 changes: 69 additions & 32 deletions tools/mtmd/clip.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -174,10 +174,6 @@ struct clip_ctx {

bool support_batch = false;

// for audio gen, reseeded only when the caller asks for another seed
std::mt19937 rng{std::random_device{}()};
uint32_t rng_seed = UINT32_MAX;

clip_ctx(clip_context_params & ctx_params) {
flash_attn_type = ctx_params.flash_attn_type;
no_alloc = ctx_params.no_alloc;
Expand Down Expand Up @@ -4717,6 +4713,59 @@ static std::vector<c2w_state_slot> list_gen_state_slots(const clip_hparams & hpa
}
}

// state_out can have another shape than its slot, only the byte size must match
static void copy_gen_state_out(ggml_tensor * src, ggml_tensor * dst) {
GGML_ASSERT(ggml_nbytes(src) == ggml_nbytes(dst));
if (src->type == dst->type && ggml_are_same_shape(src, dst) && ggml_is_contiguous(src)) {
ggml_backend_tensor_copy(src, dst);
return;
}
std::vector<uint8_t> tmp(ggml_nbytes(src));
ggml_backend_tensor_get(src, tmp.data(), 0, tmp.size());
ggml_backend_tensor_set(dst, tmp.data(), 0, tmp.size());
}

clip_gen_state * clip_gen_state_init(clip_ctx * ctx, uint32_t seed) {
auto st = std::make_unique<clip_gen_state>();
st->rng.seed(seed == UINT32_MAX ? std::random_device{}() : seed);

const auto slots = list_gen_state_slots(ctx->model.hparams, ctx->model);
if (slots.empty()) {
return st.release();
}

ggml_init_params params = {
/*.mem_size =*/ ggml_tensor_overhead() * slots.size(),
/*.mem_buffer =*/ nullptr,
/*.no_alloc =*/ true,
};
st->ctx_data.reset(ggml_init(params));
for (const auto & slot : slots) {
ggml_tensor * t = ggml_new_tensor_2d(st->ctx_data.get(), GGML_TYPE_F32, slot.ne0, slot.ne1);
ggml_format_name(t, "gen_state_%s", slot.name.c_str());
st->slots.push_back(t);
}
st->buf.reset(ggml_backend_alloc_ctx_tensors(st->ctx_data.get(), ctx->backend));
if (!st->buf) {
LOG_ERR("%s: failed to allocate the gen state\n", __func__);
return nullptr;
}
// a zero state is the cold start
ggml_backend_buffer_clear(st->buf.get(), 0);

return st.release();
}

void clip_gen_state_free(clip_gen_state * state) {
delete state;
}

void clip_gen_state_reset(clip_gen_state * state) {
if (state->buf) {
ggml_backend_buffer_clear(state->buf.get(), 0);
}
}

clip_enc_state * clip_enc_state_init(clip_ctx * ctx, int32_t n_frames_max) {
const auto & hparams = ctx->model.hparams;
if (ctx->proj_type() != PROJECTOR_TYPE_VOXTRAL_RT || hparams.attn_window_size <= 0) {
Expand Down Expand Up @@ -4782,9 +4831,9 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
clip_model_loader::warmup(*ctx, *params->imgs);
}

if (params->seed != ctx->rng_seed) {
ctx->rng_seed = params->seed;
ctx->rng.seed(params->seed == UINT32_MAX ? std::random_device{}() : params->seed);
if (params->gen_process != CLIP_GEN_PROCESS_GEN_UNKNOWN && params->gen_state == nullptr) {
LOG_ERR("%s: audio gen requires a gen state\n", __func__);
return false;
}

// build the inference graph
Expand Down Expand Up @@ -4848,19 +4897,13 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
ggml_backend_tensor_set(cur, values.data(), 0, ggml_nbytes(cur));
};

// upload the decoder state from the previous call, or zero-fill on a cold start
// state of the previous call, zero on a cold start
auto set_gen_state_in = [&]() {
size_t offset = 0;
const auto & slots = params->gen_state->slots;
size_t i = 0;
for (const auto & slot : list_gen_state_slots(hparams, model)) {
ggml_tensor * t = get_inp_tensor(("state_in_" + slot.name).c_str());
const size_t nb = ggml_nbytes(t);
if (params->state_in && params->state_in->size() >= offset + nb) {
ggml_backend_tensor_set(t, params->state_in->data() + offset, 0, nb);
} else {
std::vector<uint8_t> zeros(nb, 0);
ggml_backend_tensor_set(t, zeros.data(), 0, nb);
}
offset += nb;
GGML_ASSERT(i < slots.size());
ggml_backend_tensor_copy(slots[i++], get_inp_tensor(("state_in_" + slot.name).c_str()));
}
};

Expand Down Expand Up @@ -5584,7 +5627,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
std::normal_distribution<float> dist(0.0f, std::sqrt(temp));
std::vector<float> noise(ggml_nelements(t));
for (auto & v : noise) {
v = dist(ctx->rng);
v = dist(params->gen_state->rng);
}
set_input_f32("inp_noise", noise);
}
Expand Down Expand Up @@ -5732,7 +5775,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
std::uniform_real_distribution<float> dist(0.0f, 1.0f);
const int64_t n_acoustic = model.gen_code_head_w->ne[2];
for (int64_t g = 0; g < n_acoustic; g++) {
std::vector<float> r = { dist(ctx->rng) };
std::vector<float> r = { dist(params->gen_state->rng) };
set_input_f32(("inp_rand_" + std::to_string(g)).c_str(), r);
}
}
Expand Down Expand Up @@ -6262,22 +6305,16 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
st->n_past += n_new;
}

if (params->state_out != nullptr) {
auto & state_out = *params->state_out;
size_t total = 0;
for (const auto & slot : list_gen_state_slots(hparams, model)) {
total += (size_t) (slot.ne0 * slot.ne1) * sizeof(float);
}
state_out.resize(total);
size_t offset = 0;
if (params->gen_process == CLIP_GEN_PROCESS_GEN_WAV) {
const auto & slots = params->gen_state->slots;
size_t i = 0;
for (const auto & slot : list_gen_state_slots(hparams, model)) {
ggml_tensor * t = ggml_graph_get_tensor(gf, ("state_out_" + slot.name).c_str());
if (t == nullptr) {
GGML_ABORT("state_out requested but graph has no \"state_out_%s\" tensor", slot.name.c_str());
GGML_ABORT("graph has no \"state_out_%s\" tensor", slot.name.c_str());
}
const size_t nb = ggml_nbytes(t);
ggml_backend_tensor_get(t, state_out.data() + offset, 0, nb);
offset += nb;
GGML_ASSERT(i < slots.size());
copy_gen_state_out(t, slots[i++]);
}
}

Expand Down
13 changes: 10 additions & 3 deletions tools/mtmd/clip.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@

struct clip_ctx;
struct clip_enc_state;
struct clip_gen_state;

struct clip_image_size {
int width;
Expand Down Expand Up @@ -110,16 +111,16 @@ struct clip_encode_params {
float top_p = 1.0f;
std::vector<int32_t> * out_codes = nullptr; // this frame's 16 sampled codes
std::vector<float> * out_feats = nullptr; // continuous counterpart of out_codes
uint32_t seed = UINT32_MAX; // UINT32_MAX for random
float temp = 0.0f; // sampling temperature, noise scale for flow-matching decoders
bool * out_is_eos = nullptr;

// GEN_WAV
const std::vector<int32_t> * codes = nullptr; // this frame's 16 RVQ codes
const std::vector<float> * feats = nullptr; // continuous counterpart of codes
std::vector<float> * out_audio = nullptr; // decoded PCM samples, F32
const std::vector<uint8_t> * state_in = nullptr; // state from previous call, null or wrong size means cold start
std::vector<uint8_t> * state_out = nullptr; // state for the next call

// audio gen: state of one sequence, clip_encode() updates it in place
struct clip_gen_state * gen_state = nullptr;

// streaming audio input, updated in place by clip_encode()
struct clip_enc_state * enc_state = nullptr;
Expand All @@ -132,6 +133,12 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params);
struct clip_enc_state * clip_enc_state_init(struct clip_ctx * ctx, int32_t n_frames_max);
void clip_enc_state_free(struct clip_enc_state * state);

// per-sequence state of audio generation: RNG and the slots kept between GEN_WAV calls
// seed UINT32_MAX means random; reset() clears the slots and keeps the RNG stream
struct clip_gen_state * clip_gen_state_init(struct clip_ctx * ctx, uint32_t seed);
void clip_gen_state_free(struct clip_gen_state * state);
void clip_gen_state_reset(struct clip_gen_state * state);

bool clip_is_llava(const struct clip_ctx * ctx);
// note for contributor: this clip_is_(model) pattern is deprecated
// do NOT add new functions like this
Expand Down
50 changes: 26 additions & 24 deletions tools/mtmd/mtmd-helper-gen.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -105,12 +105,25 @@ class mtmd_gen_audio_pipeline {
virtual int32_t flush() = 0;

protected:
// one state per generation, so a fixed seed always gives the same output
bool init_gen_state(uint32_t seed) {
mtmd_gen_state_params params = mtmd_gen_state_params_default();
params.seed = seed;
gen_state.reset(mtmd_gen_state_init(mctx, params));
if (!gen_state) {
LOG_ERR("mtmd_helper_gen_audio: failed to create the gen state\n");
return false;
}
return true;
}

llama_context * lctx;
mtmd_context * mctx;
const llama_model * model;
const llama_vocab * vocab;
int n_embd;
mtmd_gen_audio_info info;
mtmd::gen_state_ptr gen_state;
};

// Qwen3-TTS: backbone samples codec_0, code_predictor gives the other 15 codebooks,
Expand All @@ -123,7 +136,7 @@ class qwen3tts_gen_audio_pipeline : public mtmd_gen_audio_pipeline {
seq_id = 0;
pos = 0;
codes_buf.clear();
c2w_state.clear();
gen_state.reset();
audio_pcm.clear();
overlay.clear();
h_state_buf.clear();
Expand Down Expand Up @@ -220,9 +233,11 @@ class qwen3tts_gen_audio_pipeline : public mtmd_gen_audio_pipeline {
const mtmd_gen_inp def = mtmd_gen_inp_default(mctx);
top_k = inp->top_k > 0 ? inp->top_k : def.top_k;
top_p = inp->top_p > 0 ? inp->top_p : def.top_p;
seed = inp->seed;
out_type = inp->out_type;
stream = inp->stream;
if (!init_gen_state(inp->seed)) {
return 1;
}

// the prompt above holds the whole text stream up to tts_eos, so every generated
// frame adds tts_pad on top of the codes embedding
Expand Down Expand Up @@ -279,9 +294,8 @@ class qwen3tts_gen_audio_pipeline : public mtmd_gen_audio_pipeline {
inp.embd = const_cast<float *>(h_state_in);
inp.top_k = top_k;
inp.top_p = top_p;
inp.seed = seed;
mtmd_gen_out out{};
if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) {
if (mtmd_gen_audio_process(mctx, gen_state.get(), &inp, &out) != 0) {
LOG_ERR("mtmd_helper_gen_audio: gen_code process failed\n");
return 1;
}
Expand Down Expand Up @@ -449,16 +463,12 @@ class qwen3tts_gen_audio_pipeline : public mtmd_gen_audio_pipeline {
inp.type = MTMD_GEN_PROCESS_TYPE_GEN_WAV;
inp.codes = codes_buf.data();
inp.n_codes = codes_buf.size();
inp.seed = seed; // same seed as gen_code, else clip reseeds mid-generation
inp.state_data = c2w_state.empty() ? nullptr : (const char *) c2w_state.data();
inp.state_size = c2w_state.size();
mtmd_gen_out out{};
if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) {
if (mtmd_gen_audio_process(mctx, gen_state.get(), &inp, &out) != 0) {
LOG_ERR("mtmd_helper_gen_audio: gen_wav process failed\n");
return false;
}
audio_pcm.insert(audio_pcm.end(), out.audio, out.audio + out.n_samples);
c2w_state.assign(out.state_data, out.state_data + out.state_size);
codes_buf.clear();
return true;
}
Expand Down Expand Up @@ -491,9 +501,7 @@ class qwen3tts_gen_audio_pipeline : public mtmd_gen_audio_pipeline {
int prompt_pos = 0;
int32_t top_k = 50;
float top_p = 1.0f;
uint32_t seed = UINT32_MAX;
std::vector<int32_t> codes_buf;
std::vector<uint8_t> c2w_state;
std::vector<float> audio_pcm;
std::vector<float> overlay;
std::vector<float> h_state_buf;
Expand Down Expand Up @@ -539,7 +547,7 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline {
seq_id = 0;
pos = 0;
feats_buf.clear();
dec_state.clear();
gen_state.reset();
audio_pcm.clear();
h_state_buf.clear();
out_buf.clear();
Expand Down Expand Up @@ -623,9 +631,11 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline {
prompt_batch->set_position_normal(0, seq_id);
prompt_pos = 0;

seed = inp->seed;
out_type = inp->out_type;
stream = inp->stream;
if (!init_gen_state(inp->seed)) {
return 1;
}

return 0;
}
Expand Down Expand Up @@ -663,13 +673,11 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline {
mtmd_gen_inp inp = mtmd_gen_inp_default(mctx);
inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE;
inp.embd = const_cast<float *>(h_state_in);
// clip only reseeds when the seed changes, so pass the same one on every step
inp.seed = seed;
if (pack.temp > 0.0f) {
inp.temp = pack.temp;
}
mtmd_gen_out out{};
if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) {
if (mtmd_gen_audio_process(mctx, gen_state.get(), &inp, &out) != 0) {
LOG_ERR("mtmd_helper_gen_audio: flow decode failed\n");
return 1;
}
Expand Down Expand Up @@ -895,7 +903,7 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline {
return 1;
}
// the decoder restarts too, the next chunk's audio is not continuous with this one
dec_state.clear();
mtmd_gen_state_reset(gen_state.get());

if (chunk_idx + 1 >= chunks.size()) {
*out_stop = true;
Expand Down Expand Up @@ -1023,16 +1031,12 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline {
inp.type = MTMD_GEN_PROCESS_TYPE_GEN_WAV;
inp.feats = feats_buf.data();
inp.n_feats = feats_buf.size();
inp.seed = seed;
inp.state_data = dec_state.empty() ? nullptr : (const char *) dec_state.data();
inp.state_size = dec_state.size();
mtmd_gen_out out{};
if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) {
if (mtmd_gen_audio_process(mctx, gen_state.get(), &inp, &out) != 0) {
LOG_ERR("mtmd_helper_gen_audio: mimi decode failed\n");
return false;
}
audio_pcm.insert(audio_pcm.end(), out.audio, out.audio + out.n_samples);
dec_state.assign(out.state_data, out.state_data + out.state_size);
feats_buf.clear();
return true;
}
Expand All @@ -1049,7 +1053,6 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline {
std::unique_ptr<decode_embd_batch> prompt_batch;
int n_prompt = 0;
int prompt_pos = 0;
uint32_t seed = UINT32_MAX;
// end-of-speech is latched, then a few more frames are generated as tail padding
int step_idx = 0;
int eos_step = -1;
Expand All @@ -1064,7 +1067,6 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline {
// latents are decoded a window at a time, the decoder state bridges the windows
size_t window_frames = 8;
std::vector<float> feats_buf;
std::vector<uint8_t> dec_state;
std::vector<float> audio_pcm;
std::vector<float> h_state_buf;
mtmd_helper_gen_audio_outtype out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV;
Expand Down
Loading
Loading