Skip to content

Harden PocketTTS Mimi decoder graph allocation #58

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
0xShug0 merged 1 commit into from
Jul 15, 2026
Merged

Harden PocketTTS Mimi decoder graph allocation #58

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
95 changes: 54 additions & 41 deletions src/models/pocket_tts/mimi_decoder.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,19 @@ class Conv1dRuntime;
class ConvTranspose1dRuntime;
class MimiTransformerRuntime;

void release_partial_graph_runtime(
ggml_gallocr_t & gallocr,
ggml_context *& context) {
if (gallocr != nullptr) {
ggml_gallocr_free(gallocr);
gallocr = nullptr;
}
if (context != nullptr) {
ggml_free(context);
context = nullptr;
}
}

struct TransformerRunResult {
std::vector<float> output_bct;
};
Expand Down Expand Up @@ -207,22 +220,23 @@ class DepthwiseConvTranspose1dRuntime {
false,
}).build(ctx, input_, {weight, std::nullopt});

params_buffer_ = ggml_backend_alloc_ctx_tensors(ggml_ctx_, backend_);
core::write_tensor_f32(input_, std::vector<float>(static_cast<size_t>(channels_ * frames_), 0.0F));
core::set_backend_threads(backend_, threads_);
graph_ = ggml_new_graph_custom(ggml_ctx_, 32768, false);
ggml_build_forward_expand(graph_, output_.tensor);
galloc_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_));
ggml_gallocr_alloc_graph(galloc_, graph_);
if (galloc_ == nullptr ||
!ggml_gallocr_reserve(galloc_, graph_) ||
!ggml_gallocr_alloc_graph(galloc_, graph_)) {
release_partial_graph_runtime(galloc_, ggml_ctx_);
throw std::runtime_error("Mimi depthwise conv transpose graph allocation failed");
}
core::write_tensor_f32(input_, std::vector<float>(static_cast<size_t>(channels_ * frames_), 0.0F));
}

~DepthwiseConvTranspose1dRuntime() {
if (galloc_ != nullptr) {
ggml_gallocr_free(galloc_);
}
if (params_buffer_ != nullptr) {
ggml_backend_buffer_free(params_buffer_);
}
if (ggml_ctx_ != nullptr) {
ggml_free(ggml_ctx_);
}
Expand All @@ -245,7 +259,6 @@ class DepthwiseConvTranspose1dRuntime {
int64_t kernel_size_ = 0;
int stride_ = 1;
ggml_context * ggml_ctx_ = nullptr;
ggml_backend_buffer_t params_buffer_ = nullptr;
ggml_cgraph * graph_ = nullptr;
ggml_gallocr_t galloc_ = nullptr;
core::TensorValue input_;
Expand Down Expand Up @@ -295,22 +308,23 @@ class Conv1dRuntime {
bias.has_value(),
}).build(ctx, input_, modules::Conv1dWeights{weight, bias});

params_buffer_ = ggml_backend_alloc_ctx_tensors(ggml_ctx_, backend_);
core::write_tensor_f32(input_, std::vector<float>(static_cast<size_t>(in_channels_ * frames_), 0.0F));
core::set_backend_threads(backend_, threads_);
graph_ = ggml_new_graph_custom(ggml_ctx_, 32768, false);
ggml_build_forward_expand(graph_, output_.tensor);
galloc_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_));
ggml_gallocr_alloc_graph(galloc_, graph_);
if (galloc_ == nullptr ||
!ggml_gallocr_reserve(galloc_, graph_) ||
!ggml_gallocr_alloc_graph(galloc_, graph_)) {
release_partial_graph_runtime(galloc_, ggml_ctx_);
throw std::runtime_error("Mimi conv1d graph allocation failed");
}
core::write_tensor_f32(input_, std::vector<float>(static_cast<size_t>(in_channels_ * frames_), 0.0F));
}

~Conv1dRuntime() {
if (galloc_ != nullptr) {
ggml_gallocr_free(galloc_);
}
if (params_buffer_ != nullptr) {
ggml_backend_buffer_free(params_buffer_);
}
if (ggml_ctx_ != nullptr) {
ggml_free(ggml_ctx_);
}
Expand All @@ -335,7 +349,6 @@ class Conv1dRuntime {
int stride_ = 1;
int dilation_ = 1;
ggml_context * ggml_ctx_ = nullptr;
ggml_backend_buffer_t params_buffer_ = nullptr;
ggml_cgraph * graph_ = nullptr;
ggml_gallocr_t galloc_ = nullptr;
core::TensorValue input_;
Expand Down Expand Up @@ -383,22 +396,23 @@ class ConvTranspose1dRuntime {
bias.has_value(),
}).build(ctx, input_, modules::ConvTranspose1dWeights{weight, bias});

params_buffer_ = ggml_backend_alloc_ctx_tensors(ggml_ctx_, backend_);
core::write_tensor_f32(input_, std::vector<float>(static_cast<size_t>(in_channels_ * frames_), 0.0F));
core::set_backend_threads(backend_, threads_);
graph_ = ggml_new_graph_custom(ggml_ctx_, 32768, false);
ggml_build_forward_expand(graph_, output_.tensor);
galloc_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_));
ggml_gallocr_alloc_graph(galloc_, graph_);
if (galloc_ == nullptr ||
!ggml_gallocr_reserve(galloc_, graph_) ||
!ggml_gallocr_alloc_graph(galloc_, graph_)) {
release_partial_graph_runtime(galloc_, ggml_ctx_);
throw std::runtime_error("Mimi conv transpose graph allocation failed");
}
core::write_tensor_f32(input_, std::vector<float>(static_cast<size_t>(in_channels_ * frames_), 0.0F));
}

~ConvTranspose1dRuntime() {
if (galloc_ != nullptr) {
ggml_gallocr_free(galloc_);
}
if (params_buffer_ != nullptr) {
ggml_backend_buffer_free(params_buffer_);
}
if (ggml_ctx_ != nullptr) {
ggml_free(ggml_ctx_);
}
Expand All @@ -422,7 +436,6 @@ class ConvTranspose1dRuntime {
int64_t kernel_size_ = 0;
int stride_ = 1;
ggml_context * ggml_ctx_ = nullptr;
ggml_backend_buffer_t params_buffer_ = nullptr;
ggml_cgraph * graph_ = nullptr;
ggml_gallocr_t galloc_ = nullptr;
core::TensorValue input_;
Expand Down Expand Up @@ -623,7 +636,16 @@ class MimiFullDecoderRuntime {
}
output_ = modules::SliceModule({2, 0, output_full.shape.dims[2] - trim_frames}).build(ctx, output_full);

params_buffer_ = ggml_backend_alloc_ctx_tensors(ggml_ctx_, backend_);
core::set_backend_threads(backend_, threads_);
graph_ = ggml_new_graph_custom(ggml_ctx_, 32768, false);
ggml_build_forward_expand(graph_, output_.tensor);
galloc_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_));
if (galloc_ == nullptr ||
!ggml_gallocr_reserve(galloc_, graph_) ||
!ggml_gallocr_alloc_graph(galloc_, graph_)) {
release_partial_graph_runtime(galloc_, ggml_ctx_);
throw std::runtime_error("Mimi full decoder graph allocation failed");
}
core::write_tensor_f32(latents_bct_, std::vector<float>(static_cast<size_t>(config_.latent_size * steps_), 0.0F));
std::vector<int32_t> positions(static_cast<size_t>(transformed_frames_), 0);
for (int64_t i = 0; i < transformed_frames_; ++i) {
Expand All @@ -640,11 +662,6 @@ class MimiFullDecoderRuntime {
core::write_tensor_f32(zero_prefix_values_[layer], zero_prefix);
}
}
core::set_backend_threads(backend_, threads_);
graph_ = ggml_new_graph_custom(ggml_ctx_, 32768, false);
ggml_build_forward_expand(graph_, output_.tensor);
galloc_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_));
ggml_gallocr_alloc_graph(galloc_, graph_);
if (engine::core::uses_host_graph_plan(backend_)) {
const auto plan_started = std::chrono::steady_clock::now();
plan_ = engine::core::create_backend_graph_plan_if_host(backend_, graph_);
Expand All @@ -662,9 +679,6 @@ class MimiFullDecoderRuntime {
if (galloc_ != nullptr) {
ggml_gallocr_free(galloc_);
}
if (params_buffer_ != nullptr) {
ggml_backend_buffer_free(params_buffer_);
}
if (ggml_ctx_ != nullptr) {
ggml_free(ggml_ctx_);
}
Expand All @@ -689,7 +703,6 @@ class MimiFullDecoderRuntime {
int64_t cache_steps_ = 0;
int64_t head_dim_ = 0;
ggml_context * ggml_ctx_ = nullptr;
ggml_backend_buffer_t params_buffer_ = nullptr;
ggml_cgraph * graph_ = nullptr;
ggml_gallocr_t galloc_ = nullptr;
ggml_backend_graph_plan_t plan_ = nullptr;
Expand Down Expand Up @@ -809,7 +822,16 @@ class MimiTransformerRuntime {
output_bct_ = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, x);
build_transfer_views();

params_buffer_ = ggml_backend_alloc_ctx_tensors(ggml_ctx_, backend_);
core::set_backend_threads(backend_, threads_);
graph_ = ggml_new_graph_custom(ggml_ctx_, 32768, false);
ggml_build_forward_expand(graph_, output_bct_.tensor);
galloc_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_));
if (galloc_ == nullptr ||
!ggml_gallocr_reserve(galloc_, graph_) ||
!ggml_gallocr_alloc_graph(galloc_, graph_)) {
release_partial_graph_runtime(galloc_, ggml_ctx_);
throw std::runtime_error("Mimi transformer graph allocation failed");
}
core::write_tensor_f32(input_bct_, std::vector<float>(static_cast<size_t>(config_.hidden_size * frames_), 0.0F));
core::write_tensor_i32(positions_, std::vector<int32_t>(static_cast<size_t>(frames_), 0));
core::write_tensor_f32(attention_mask_, std::vector<float>(static_cast<size_t>(frames_ * (cache_steps_ + frames_)), -INFINITY));
Expand All @@ -827,20 +849,12 @@ class MimiTransformerRuntime {
core::write_tensor_f32(zero_prefix_values_[layer], zero_prefix);
}
}
core::set_backend_threads(backend_, threads_);
graph_ = ggml_new_graph_custom(ggml_ctx_, 32768, false);
ggml_build_forward_expand(graph_, output_bct_.tensor);
galloc_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_));
ggml_gallocr_alloc_graph(galloc_, graph_);
}

~MimiTransformerRuntime() {
if (galloc_ != nullptr) {
ggml_gallocr_free(galloc_);
}
if (params_buffer_ != nullptr) {
ggml_backend_buffer_free(params_buffer_);
}
if (ggml_ctx_ != nullptr) {
ggml_free(ggml_ctx_);
}
Expand Down Expand Up @@ -982,7 +996,6 @@ class MimiTransformerRuntime {
int64_t cache_steps_ = 0;
int64_t head_dim_ = 0;
ggml_context * ggml_ctx_ = nullptr;
ggml_backend_buffer_t params_buffer_ = nullptr;
ggml_cgraph * graph_ = nullptr;
ggml_gallocr_t galloc_ = nullptr;
core::TensorValue input_bct_;
Expand Down
Loading