From 5ee99787b100e8de2bb0b64271914507dd783d31 Mon Sep 17 00:00:00 2001 From: derekja Date: Tue, 18 Aug 2026 22:33:58 +0000 Subject: [PATCH] asr: allocate inference graphs with ggml_gallocr, not alloc_ctx_tensors Several ASR graphs are built with no_alloc=true and then allocated with ggml_backend_alloc_ctx_tensors, which allocates every tensor in the context at once. Intermediate activations for all layers therefore stay resident for the whole graph instead of being recycled, so peak memory scales with sequence length at roughly 50x the legitimate cost. Measured on a 12 GB H100 slice, Qwen3-ASR 0.6B, 32 s of audio with a 3000 character recognition prompt: a single cudaMalloc of 9739 MiB for 886 prompt steps, i.e. 11.0 MiB per token, where the KV cache for a 0.6B model is about 224 KB per token. It fails as one monolithic allocation rather than as gradual pressure. The fix is the graph allocator already used elsewhere in this repository, including in qwen3_asr/audio_encoder.cpp and hviske_asr/encoder.cpp -- one file away from two of the sites changed here. citrinet_asr/runtime.cpp 23.96 -> 0.43 MiB per second of audio (56x) qwen3_asr/thinker.cpp PrefillGraph 123.7 -> 31.2 MiB/s (0.6B), 149.3 -> 30.3 (1.7B) 1.657 -> 0.158 MiB per prompt character (10.5x) PromptClassificationGraph same treatment, single output Two cases that previously failed to allocate now run: 0.6B at 32 s + 2996 characters (9739 MiB -> 3227 MiB) and 1.7B at 32 s + 1795 characters (8246 MiB -> 4294 MiB). Two details that are not simply swapping the allocator: K/V readback. PrefillGraph hands run() the decoder's K/V state, which is an intermediate the allocator would recycle. Marking those tensors as outputs is not sufficient: they may be views, and GGML_TENSOR_FLAG_OUTPUT on a view does not protect its view_src. Each is copied into a tensor of its own with ggml_cpy over ggml_dup_tensor and that copy is marked as an output, matching framework/modules/transformers/qwen_causal_decode_runtime.cpp. Constant leaves. Both qwen3 graphs uploaded position ids once at build time. That is safe when every tensor is pinned, but the graph allocator may reuse a leaf once its last consumer has run, so a second run() on a cached graph would read stale positions. The upload moves into run(). Not changed: DecodeGraph in the same file keeps alloc_ctx_tensors. Its context holds a persistent step cache that must retain stable storage across successive decode steps, and its per-step activations are one token wide, so it does not contribute to length scaling. hviske_asr/decoder.cpp is left alone for a similar reason: it builds views onto prefill K/V that separate contexts consume, which needs stable addresses across graph boundaries. Output is unchanged: 22 transcript comparisons between the two builds are byte-identical, covering both models, four audio lengths, four prompt lengths, the CPU backend, and cached-graph batch runs. --- src/models/citrinet_asr/runtime.cpp | 11 ++--- src/models/qwen3_asr/thinker.cpp | 62 +++++++++++++++++++++-------- 2 files changed, 52 insertions(+), 21 deletions(-) diff --git a/src/models/citrinet_asr/runtime.cpp b/src/models/citrinet_asr/runtime.cpp index 39eca33d..0e314dda 100644 --- a/src/models/citrinet_asr/runtime.cpp +++ b/src/models/citrinet_asr/runtime.cpp @@ -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_)) { @@ -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_); } } @@ -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 channels_first_; diff --git a/src/models/qwen3_asr/thinker.cpp b/src/models/qwen3_asr/thinker.cpp index e039cd4a..857b127d 100644 --- a/src/models/qwen3_asr/thinker.cpp +++ b/src/models/qwen3_asr/thinker.cpp @@ -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. @@ -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_); } } @@ -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 positions(audio_positions.begin(), audio_positions.end()); @@ -430,8 +455,9 @@ class PrefillGraph { ggml_tensor * logits_ = nullptr; std::vector keys_; std::vector values_; + std::vector position_ids_; ggml_cgraph * graph_ = nullptr; - ggml_backend_buffer_t buffer_ = nullptr; + ggml_gallocr_t gallocr_ = nullptr; }; class PromptClassificationGraph { @@ -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_); } } @@ -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 positions(audio_positions.begin(), audio_positions.end()); @@ -558,8 +587,9 @@ class PromptClassificationGraph { ggml_tensor * audio_positions_ = nullptr; ggml_tensor * positions_ = nullptr; ggml_tensor * token_ids_ = nullptr; + std::vector position_ids_; ggml_cgraph * graph_ = nullptr; - ggml_backend_buffer_t buffer_ = nullptr; + ggml_gallocr_t gallocr_ = nullptr; }; class DecodeGraph {