Skip to content

asr: allocate inference graphs with ggml_gallocr, not alloc_ctx_tensors #276

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

Open
derekja wants to merge 1 commit into
base: main
Choose a base branch
from
+52 −21
Open

asr: allocate inference graphs with ggml_gallocr, not alloc_ctx_tensors #276

Show file tree
Hide file tree
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
11 changes: 6 additions & 5 deletions src/models/citrinet_asr/runtime.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -374,12 +374,13 @@ class CitrinetRuntime::Graph {
};
auto input = core::make_tensor(build_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, weights_->config.n_mels, frames_}));
input_ = input.tensor;
ggml_set_input(input_);
output_ = build_citrinet_graph(build_ctx, input, *backend_weights_).tensor;
ggml_set_output(output_);
graph_ = ggml_new_graph_custom(ctx_.get(), 16384, false);
ggml_build_forward_expand(graph_, output_);
buffer_ = ggml_backend_alloc_ctx_tensors(ctx_.get(), backend_);
if (buffer_ == nullptr) {
gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_));
if (gallocr_ == nullptr || !ggml_gallocr_alloc_graph(gallocr_, graph_)) {
throw std::runtime_error("failed to allocate graph");
}
if (engine::core::uses_host_graph_plan(backend_)) {
Expand All @@ -399,8 +400,8 @@ class CitrinetRuntime::Graph {
if (plan_ != nullptr) {
engine::core::free_backend_graph_plan(backend_, plan_);
}
if (buffer_ != nullptr) {
ggml_backend_buffer_free(buffer_);
if (gallocr_ != nullptr) {
ggml_gallocr_free(gallocr_);
}
}

Expand Down Expand Up @@ -466,7 +467,7 @@ class CitrinetRuntime::Graph {
ggml_cgraph * graph_ = nullptr;
ggml_backend_t backend_ = nullptr;
int compute_threads_ = 1;
ggml_backend_buffer_t buffer_ = nullptr;
ggml_gallocr_t gallocr_ = nullptr;
ggml_backend_graph_plan_t plan_ = nullptr;
double plan_create_ms_ = 0.0;
std::vector<float> channels_first_;
Expand Down
62 changes: 46 additions & 16 deletions src/models/qwen3_asr/thinker.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -326,15 +326,37 @@ class PrefillGraph {
if (!layer.key.has_value() || !layer.value.has_value()) {
throw std::runtime_error("Qwen3 ASR thinker prefill decoder did not return K/V state");
}
keys_.push_back(layer.key->tensor);
values_.push_back(layer.value->tensor);
// The graph allocator recycles intermediates, and the decoder K/V is an
// intermediate that run() has to read back afterwards. Copy each one into
// a tensor of its own and mark it as a graph output, which is what keeps
// it off the reuse list. Same shape as QwenCausalDecodeRuntime prefill.
auto * key = ggml_cpy(
ctx_.get(),
layer.key->tensor,
ggml_dup_tensor(ctx_.get(), layer.key->tensor));
auto * value = ggml_cpy(
ctx_.get(),
layer.value->tensor,
ggml_dup_tensor(ctx_.get(), layer.value->tensor));
ggml_set_output(key);
ggml_set_output(value);
keys_.push_back(key);
values_.push_back(value);
}
logits_ = decoder_out.logits.tensor;
ggml_set_output(logits_);
graph_ = ggml_new_graph_custom(ctx_.get(), 65536, false);
ggml_build_forward_expand(graph_, logits_);
buffer_ = ggml_backend_alloc_ctx_tensors(ctx_.get(), runtime_->backend());
if (buffer_ == nullptr) {
for (auto * key : keys_) {
ggml_build_forward_expand(graph_, key);
}
for (auto * value : values_) {
ggml_build_forward_expand(graph_, value);
}
gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(runtime_->backend()));
if (gallocr_ == nullptr ||
!ggml_gallocr_reserve(gallocr_, graph_) ||
!ggml_gallocr_alloc_graph(gallocr_, graph_)) {
// Size, not a fault: the graph scales with prompt_steps_, which the
// caller controls through the transcription prompt and the length of
// the audio. Say which, and by how much, so the remedy is obvious.
Expand All @@ -344,16 +366,15 @@ class PrefillGraph {
+ std::to_string(audio_tokens_) + " are audio tokens); "
"shorten the transcription prompt or the audio");
}
const auto pos = modules::qwen_position_ids(prompt_steps_);
ggml_backend_tensor_set(positions_, pos.data(), 0, pos.size() * sizeof(int32_t));
position_ids_ = modules::qwen_position_ids(prompt_steps_);
debug::timing_log_scalar("qwen3_asr.thinker.prefill.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now()));
debug::trace_log_scalar("qwen3_asr.thinker.prefill_prompt_steps", prompt_steps_);
}

~PrefillGraph() {
engine::core::release_backend_graph_resources(runtime_->backend(), graph_);
if (buffer_ != nullptr) {
ggml_backend_buffer_free(buffer_);
if (gallocr_ != nullptr) {
ggml_gallocr_free(gallocr_);
}
}

Expand All @@ -376,6 +397,10 @@ class PrefillGraph {
throw std::runtime_error("Qwen3 ASR thinker prefill audio position count mismatch");
}
auto timing_start = Clock::now();
// Re-uploaded on every run: the graph allocator may hand this leaf out to a
// later node once the graph has consumed it, so it cannot be written once at
// build time and left alone.
ggml_backend_tensor_set(positions_, position_ids_.data(), 0, position_ids_.size() * sizeof(int32_t));
ggml_backend_tensor_set(token_ids_, token_ids.data(), 0, token_ids.size() * sizeof(int32_t));
if (audio_tokens_ > 0) {
std::vector<int64_t> positions(audio_positions.begin(), audio_positions.end());
Expand Down Expand Up @@ -430,8 +455,9 @@ class PrefillGraph {
ggml_tensor * logits_ = nullptr;
std::vector<ggml_tensor *> keys_;
std::vector<ggml_tensor *> values_;
std::vector<int32_t> position_ids_;
ggml_cgraph * graph_ = nullptr;
ggml_backend_buffer_t buffer_ = nullptr;
ggml_gallocr_t gallocr_ = nullptr;
};

class PromptClassificationGraph {
Expand Down Expand Up @@ -482,20 +508,21 @@ class PromptClassificationGraph {
ggml_set_output(token_ids_);
graph_ = ggml_new_graph_custom(ctx_.get(), 65536, false);
ggml_build_forward_expand(graph_, token_ids_);
buffer_ = ggml_backend_alloc_ctx_tensors(ctx_.get(), runtime_->backend());
if (buffer_ == nullptr) {
gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(runtime_->backend()));
if (gallocr_ == nullptr ||
!ggml_gallocr_reserve(gallocr_, graph_) ||
!ggml_gallocr_alloc_graph(gallocr_, graph_)) {
throw std::runtime_error("failed to allocate Qwen3 ASR thinker classification graph");
}
const auto pos = modules::qwen_position_ids(prompt_steps_);
ggml_backend_tensor_set(positions_, pos.data(), 0, pos.size() * sizeof(int32_t));
position_ids_ = modules::qwen_position_ids(prompt_steps_);
debug::timing_log_scalar("qwen3_asr.thinker.classify.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now()));
debug::trace_log_scalar("qwen3_asr.thinker.classify_prompt_steps", prompt_steps_);
}

~PromptClassificationGraph() {
engine::core::release_backend_graph_resources(runtime_->backend(), graph_);
if (buffer_ != nullptr) {
ggml_backend_buffer_free(buffer_);
if (gallocr_ != nullptr) {
ggml_gallocr_free(gallocr_);
}
}

Expand All @@ -518,6 +545,8 @@ class PromptClassificationGraph {
throw std::runtime_error("Qwen3 ASR thinker classification audio position count mismatch");
}
auto timing_start = Clock::now();
// See PrefillGraph::run: leaves are not pinned by the graph allocator.
ggml_backend_tensor_set(positions_, position_ids_.data(), 0, position_ids_.size() * sizeof(int32_t));
ggml_backend_tensor_set(token_ids_input_, input_ids.data(), 0, input_ids.size() * sizeof(int32_t));
if (audio_tokens_ > 0) {
std::vector<int64_t> positions(audio_positions.begin(), audio_positions.end());
Expand Down Expand Up @@ -558,8 +587,9 @@ class PromptClassificationGraph {
ggml_tensor * audio_positions_ = nullptr;
ggml_tensor * positions_ = nullptr;
ggml_tensor * token_ids_ = nullptr;
std::vector<int32_t> position_ids_;
ggml_cgraph * graph_ = nullptr;
ggml_backend_buffer_t buffer_ = nullptr;
ggml_gallocr_t gallocr_ = nullptr;
};

class DecodeGraph {
Expand Down
Loading