From 9b5a3f9cffd568428e777fc664817987ae7c81ef Mon Sep 17 00:00:00 2001 From: mirek190 Date: Thu, 3 Sep 2026 10:29:24 +0100 Subject: [PATCH] fix(qwen): preserve compact logits indices across graph reuse --- .../transformers/qwen_causal_decode_runtime.cpp | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp b/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp index 8d0e65ad9..5be454a29 100644 --- a/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp +++ b/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp @@ -291,10 +291,15 @@ ggml_tensor * make_logits_readback_token_ids( if (config.logits_readback_token_ids.empty()) { return nullptr; } - return ggml_new_tensor_1d( + auto * tensor = ggml_new_tensor_1d( ctx, GGML_TYPE_I32, static_cast(config.logits_readback_token_ids.size())); + // This leaf is populated by the host after graph allocation. Marking it as + // an input keeps the graph allocator from aliasing its storage with a + // temporary that executes before ggml_get_rows consumes the token ids. + ggml_set_input(tensor); + return tensor; } core::TensorValue wrap_logits_readback_token_ids( @@ -784,6 +789,9 @@ class QwenCausalDecodeRuntime::Impl { prefill_attention_mask_values_.data(), 0, prefill_attention_mask_values_.size() * sizeof(ggml_fp16_t)); + if (prefill_logits_readback_token_ids_ != nullptr) { + upload_logits_readback_token_ids(prefill_logits_readback_token_ids_, config_); + } core::set_backend_threads(backend_, threads_); const ggml_status status = core::compute_backend_graph(backend_, prefill_graph_); ggml_backend_synchronize(backend_); @@ -970,6 +978,10 @@ class QwenCausalDecodeRuntime::Impl { batched_prefill_attention_mask_values_.data(), 0, batched_prefill_attention_mask_values_.size() * sizeof(ggml_fp16_t)); + if (batched_prefill_logits_readback_token_ids_ != nullptr) { + upload_logits_readback_token_ids( + batched_prefill_logits_readback_token_ids_, config_); + } core::set_backend_threads(backend_, threads_); const ggml_status status = core::compute_backend_graph(backend_, batched_prefill_graph_); ggml_backend_synchronize(backend_);