diff --git a/CMakeLists.txt b/CMakeLists.txt index 5fd53e108..5fa84cb8f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -419,6 +419,7 @@ add_library(engine_core OBJECT src/framework/runtime/host_ops.cpp src/framework/runtime/kv_cache.cpp src/framework/runtime/bounded_static_kv_decode.cpp + src/framework/runtime/greedy_qwen_decoder.cpp src/framework/runtime/options.cpp src/framework/runtime/session_base.cpp src/framework/runtime/workspace.cpp @@ -500,6 +501,7 @@ add_library(engine_core OBJECT src/framework/modules/optimizations/fast_kv_modules.cpp src/framework/tokenizers/hf_tokenizer_json.cpp src/framework/tokenizers/llama_bpe.cpp + src/framework/tokenizers/qwen_bpe_bundle.cpp external/llama_tokenizer/bpe-core.cpp external/llama_tokenizer/unicode.cpp external/llama_tokenizer/unicode-data.cpp diff --git a/include/engine/community_models/audio8_asr/thinker.h b/include/engine/community_models/audio8_asr/thinker.h index 59c4e0330..66bae0897 100644 --- a/include/engine/community_models/audio8_asr/thinker.h +++ b/include/engine/community_models/audio8_asr/thinker.h @@ -2,6 +2,7 @@ #include "engine/framework/assets/tensor_source.h" #include "engine/framework/core/execution_context.h" +#include "engine/framework/runtime/greedy_qwen_decoder.h" #include "engine/community_models/audio8_asr/types.h" #include @@ -9,9 +10,9 @@ namespace engine::community_models::audio8_asr { -// Greedy causal decoder for the Audio8 8-layer Qwen2-style LM. Audio -// embeddings are injected into the token embedding sequence at the prompt's -// audio placeholder positions before prefill. +// The Audio8 8-layer Qwen2-style decoder, expressed through the framework's +// shared greedy Qwen decoder runtime (prefill with audio-embedding injection +// plus static-cache step decode). Owns only the family-specific spec. class Audio8ThinkerRuntime { public: Audio8ThinkerRuntime( @@ -33,8 +34,8 @@ class Audio8ThinkerRuntime { const Audio8ASRGenerationOptions & options); private: - struct Impl; - std::unique_ptr impl_; + runtime::GreedyQwenDecoderRuntime runtime_; + std::shared_ptr config_; }; } // namespace engine::community_models::audio8_asr diff --git a/include/engine/framework/runtime/greedy_qwen_decoder.h b/include/engine/framework/runtime/greedy_qwen_decoder.h new file mode 100644 index 000000000..d54b6b1d3 --- /dev/null +++ b/include/engine/framework/runtime/greedy_qwen_decoder.h @@ -0,0 +1,72 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/modules/transformers/qwen_causal_decoder.h" + +#include +#include +#include +#include +#include + +namespace engine::runtime { + +// Specification for a greedy Qwen-family causal decoder: how to find its +// tensors in a weight source and how the shared QwenCausalDecoder stack is +// configured. Covers Qwen2-style decoders (attention biases, no Q/K norms, +// as in Audio8-ASR) and Qwen3-style decoders (Q/K norms, no attention +// biases, as in the qwen3_asr thinker), with separate or packed QKV +// projections and tied or separate LM heads. +struct GreedyQwenDecoderSpec { + modules::QwenCausalDecoderConfig decoder; + int64_t vocab_size = 0; + int64_t max_position_embeddings = 0; + bool tie_word_embeddings = false; + bool attention_bias = false; + bool packed_qkv = false; + std::string token_embedding_tensor; + std::string lm_head_tensor; // used when !tie_word_embeddings + std::string final_norm_tensor; + std::string layer_prefix; // e.g. "language_model.model.layers" + std::vector eos_token_ids; +}; + +// Greedy autoregressive decoding over a Qwen-style decoder stack: prefill +// with optional audio-embedding injection (ggml_set_rows at prompt +// positions) and static-cache step decode, hiding the graph lifetime and +// K/V state handoff that model families otherwise duplicate. +class GreedyQwenDecoderRuntime { +public: + struct Injection { + std::vector values; // tokens * hidden, token-major + int64_t tokens = 0; + std::vector positions; // prompt positions to replace + }; + + struct Prompt { + std::vector input_ids; + Injection injection; // optional + }; + + GreedyQwenDecoderRuntime( + std::shared_ptr weights_source, + const GreedyQwenDecoderSpec & spec, + core::ExecutionContext & execution, + size_t prefill_graph_arena_bytes, + size_t decode_graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type); + ~GreedyQwenDecoderRuntime(); + + GreedyQwenDecoderRuntime(const GreedyQwenDecoderRuntime &) = delete; + GreedyQwenDecoderRuntime & operator=(const GreedyQwenDecoderRuntime &) = delete; + + std::vector generate(const Prompt & prompt, int64_t max_new_tokens); + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::runtime diff --git a/include/engine/framework/tokenizers/qwen_bpe_bundle.h b/include/engine/framework/tokenizers/qwen_bpe_bundle.h new file mode 100644 index 000000000..6a5f64265 --- /dev/null +++ b/include/engine/framework/tokenizers/qwen_bpe_bundle.h @@ -0,0 +1,22 @@ +#pragma once + +#include "engine/framework/assets/resource_bundle.h" +#include "engine/framework/tokenizers/llama_bpe.h" + +#include +#include + +namespace engine::tokenizers { + +// Load the Qwen2-pretokenized BPE tokenizer referenced by a model bundle: +// tokenizer_config.json plus vocab.json/merges.txt or tokenizer.json. +std::shared_ptr load_qwen_bpe_tokenizer( + const engine::assets::ResourceBundle & bundle); + +// Look up a special token's id from the bundle's tokenizer.json +// added_tokens list; throws when the token is absent. +int64_t require_added_token_id( + const engine::assets::ResourceBundle & bundle, + std::string_view content); + +} // namespace engine::tokenizers diff --git a/src/community_models/audio8_asr/assets.cpp b/src/community_models/audio8_asr/assets.cpp index 4859bfaf7..003e8a110 100644 --- a/src/community_models/audio8_asr/assets.cpp +++ b/src/community_models/audio8_asr/assets.cpp @@ -3,6 +3,7 @@ #include "engine/framework/assets/tensor_source.h" #include "engine/framework/io/json.h" #include "engine/framework/model_spec/package.h" +#include "engine/framework/tokenizers/qwen_bpe_bundle.h" #include #include @@ -126,19 +127,6 @@ qwen3_asr::Qwen3ASRAudioEncoderConfig parse_audio_encoder_config(const json::Val return config; } -int64_t require_added_token_id(const assets::ResourceBundle & resources, std::string_view content) { - const auto tokenizer = resources.parse_json("tokenizer_json"); - for (const auto & item : tokenizer.require("added_tokens").as_array()) { - const auto * token_content = item.find("content"); - const auto * token_id = item.find("id"); - if (token_content != nullptr && token_content->is_string() && - token_id != nullptr && token_id->is_number() && token_content->as_string() == content) { - return token_id->as_i64(); - } - } - throw std::runtime_error("Audio8 ASR tokenizer.json is missing token: " + std::string(content)); -} - Audio8ASRConfig parse_config(const assets::ResourceBundle & resources) { const auto root = resources.parse_json("config"); @@ -199,11 +187,11 @@ Audio8ASRConfig parse_config(const assets::ResourceBundle & resources) { // Prompt special tokens live in tokenizer.json added_tokens; the audio // token id from the config must match the tokenizer entry. - config.user_token_id = require_added_token_id(resources, "<|user|>"); - config.begin_audio_token_id = require_added_token_id(resources, "<|begin_of_audio|>"); - config.end_audio_token_id = require_added_token_id(resources, "<|end_of_audio|>"); - config.assistant_token_id = require_added_token_id(resources, "<|assistant|>"); - config.text_decoder.audio_token_id = require_added_token_id(resources, "<|audio|>"); + config.user_token_id = engine::tokenizers::require_added_token_id(resources, "<|user|>"); + config.begin_audio_token_id = engine::tokenizers::require_added_token_id(resources, "<|begin_of_audio|>"); + config.end_audio_token_id = engine::tokenizers::require_added_token_id(resources, "<|end_of_audio|>"); + config.assistant_token_id = engine::tokenizers::require_added_token_id(resources, "<|assistant|>"); + config.text_decoder.audio_token_id = engine::tokenizers::require_added_token_id(resources, "<|audio|>"); config.supported_languages = { "Chinese", "English", "Cantonese", "French", "German", "Japanese", "Korean"}; @@ -267,19 +255,7 @@ std::shared_ptr load_audio8_asr_assets(const std::filesys encoder_assets->model_weights = std::move(encoder_source); assets->encoder_assets = std::move(encoder_assets); - engine::tokenizers::LlamaBpeTokenizerSpec tokenizer_spec; - tokenizer_spec.tokenizer_config_path = assets->resources.require_file("tokenizer_config"); - if (const auto * path = assets->resources.find_file("vocab")) { - tokenizer_spec.vocab_path = *path; - } - if (const auto * path = assets->resources.find_file("merges")) { - tokenizer_spec.merges_path = *path; - } - if (const auto * path = assets->resources.find_file("tokenizer_json")) { - tokenizer_spec.tokenizer_json_path = *path; - } - tokenizer_spec.pre_type = engine::tokenizers::LlamaBpePreTokenizer::Qwen2; - assets->tokenizer = engine::tokenizers::load_llama_bpe_tokenizer(tokenizer_spec); + assets->tokenizer = engine::tokenizers::load_qwen_bpe_tokenizer(assets->resources); return assets; } diff --git a/src/community_models/audio8_asr/thinker.cpp b/src/community_models/audio8_asr/thinker.cpp index 707c2fdd5..1913c200a 100644 --- a/src/community_models/audio8_asr/thinker.cpp +++ b/src/community_models/audio8_asr/thinker.cpp @@ -1,617 +1,43 @@ #include "engine/community_models/audio8_asr/thinker.h" -#include "engine/framework/assets/tensor_source.h" -#include "engine/framework/core/backend.h" -#include "engine/framework/core/backend_weight_store.h" -#include "engine/framework/debug/profiler.h" -#include "engine/framework/modules/activation_modules.h" -#include "engine/framework/modules/transformers/qwen_causal_decoder.h" -#include "engine/framework/modules/linear_module.h" -#include "engine/framework/modules/lookup_modules.h" -#include "engine/framework/modules/norm_modules.h" -#include "engine/framework/modules/positional_modules.h" -#include "engine/framework/modules/primitive_modules.h" -#include "engine/framework/modules/structural_modules.h" -#include "engine/framework/runtime/errors.h" -#include "engine/framework/runtime/kv_cache.h" -#include "engine/framework/sampling/decode_modules.h" - -#include -#include - -#include -#include -#include -#include -#include -#include #include -#include #include -#include namespace engine::community_models::audio8_asr { namespace { namespace modules = engine::modules; -using Clock = std::chrono::steady_clock; - -struct GgmlContextDeleter { - void operator()(ggml_context * ctx) const noexcept { - if (ctx != nullptr) { - ggml_free(ctx); - } - } -}; - -struct GgmlGallocrDeleter { - void operator()(ggml_gallocr_t alloc) const noexcept { - if (alloc != nullptr) { - ggml_gallocr_free(alloc); - } - } -}; - -struct TextLayerWeights { - core::TensorValue input_norm; - core::TensorValue q_proj; - core::TensorValue q_bias; - core::TensorValue k_proj; - core::TensorValue k_bias; - core::TensorValue v_proj; - core::TensorValue v_bias; - core::TensorValue o_proj; - core::TensorValue post_norm; - core::TensorValue gate_proj; - core::TensorValue up_proj; - core::TensorValue down_proj; -}; - -struct ThinkerWeights { - std::shared_ptr store; - core::TensorValue token_embedding; - std::vector layers; - core::TensorValue norm; - core::TensorValue lm_head; -}; - -struct PrefillOutput { - std::vector logits; - runtime::TransformerKVState kv_state; -}; - -modules::QwenDecoderLayerWeights to_qwen_layer_weights(const TextLayerWeights & weights) { - modules::QwenDecoderLayerWeights out; - out.input_norm = {weights.input_norm, std::nullopt}; - out.self_attention.q_weight = weights.q_proj; - out.self_attention.q_bias = weights.q_bias; - out.self_attention.k_weight = weights.k_proj; - out.self_attention.k_bias = weights.k_bias; - out.self_attention.v_weight = weights.v_proj; - out.self_attention.v_bias = weights.v_bias; - out.self_attention.out_weight = weights.o_proj; - out.post_norm = {weights.post_norm, std::nullopt}; - out.mlp.gate_proj = {weights.gate_proj, std::nullopt}; - out.mlp.up_proj = {weights.up_proj, std::nullopt}; - out.mlp.down_proj = {weights.down_proj, std::nullopt}; - return out; -} - -// Audio8 decoders are Qwen2-style: no per-head Q/K norms and separate Q/K/V -// projections with attention biases. -modules::QwenCausalDecoderConfig make_qwen_decoder_config(const Audio8ASRDecoderConfig & config) { - modules::QwenCausalDecoderConfig out; - out.stack.hidden_size = config.hidden_size; - out.stack.num_attention_heads = config.num_attention_heads; - out.stack.num_key_value_heads = config.num_key_value_heads; - out.stack.head_dim = config.head_dim; - out.stack.intermediate_size = config.intermediate_size; - out.stack.layers = config.num_hidden_layers; - out.stack.rms_norm_eps = config.rms_norm_eps; - out.stack.rope_theta = config.rope_theta; - out.stack.use_qk_norm = false; - out.stack.runtime.static_cache.update_mode = modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; - out.logits_size = config.vocab_size; - out.logits_mode = modules::QwenCausalDecoderLogitsMode::LastStep; - return out; -} - -modules::QwenCausalDecoderWeights make_qwen_decoder_weights(const ThinkerWeights & weights) { - modules::QwenCausalDecoderWeights out; - out.stack.layers.reserve(weights.layers.size()); - for (const auto & layer : weights.layers) { - out.stack.layers.push_back(to_qwen_layer_weights(layer)); - } - out.final_norm = {weights.norm, std::nullopt}; - out.lm_head = {weights.lm_head, std::nullopt}; - return out; -} - -core::TensorValue prompt_embeddings( - core::ModuleBuildContext & ctx, - const ThinkerWeights & weights, - const Audio8ASRDecoderConfig & config, - ggml_tensor * token_ids, - ggml_tensor * audio_embeddings, - ggml_tensor * audio_positions, - int64_t prompt_steps, - int64_t audio_tokens) { - auto ids = core::wrap_tensor(token_ids, core::TensorShape::from_dims({prompt_steps}), GGML_TYPE_I32); - auto x = modules::EmbeddingModule({config.vocab_size, config.hidden_size}).build(ctx, ids, weights.token_embedding); - if (audio_tokens > 0) { - auto audio = core::wrap_tensor( - audio_embeddings, - core::TensorShape::from_dims({audio_tokens, config.hidden_size}), - GGML_TYPE_F32); - auto positions = core::wrap_tensor( - audio_positions, - core::TensorShape::from_dims({audio_tokens}), - GGML_TYPE_I64); - x = core::wrap_tensor( - ggml_set_rows(ctx.ggml, x.tensor, audio.tensor, positions.tensor), - x.shape, - GGML_TYPE_F32); - } - return core::reshape_tensor(ctx, x, core::TensorShape::from_dims({1, prompt_steps, config.hidden_size})); -} - -ThinkerWeights load_weights( - const assets::TensorSource & source, - const Audio8ASRDecoderConfig & config, - ggml_backend_t backend, - core::BackendType backend_type, - size_t weight_context_bytes, - assets::TensorStorageType storage_type) { - ThinkerWeights weights; - weights.store = std::make_shared( - backend, - backend_type, - "audio8_asr.thinker.weights", - weight_context_bytes); - weights.token_embedding = weights.store->load_tensor( - source, - "language_model.model.embed_tokens.weight", - storage_type, - {config.vocab_size, config.hidden_size}); - weights.layers.reserve(static_cast(config.num_hidden_layers)); - const int64_t dim = config.head_dim; - for (int64_t layer = 0; layer < config.num_hidden_layers; ++layer) { - const std::string prefix = "language_model.model.layers." + std::to_string(layer); - TextLayerWeights w; - w.input_norm = weights.store->load_f32_tensor(source, prefix + ".input_layernorm.weight", {config.hidden_size}); - w.q_proj = weights.store->load_tensor(source, prefix + ".self_attn.q_proj.weight", storage_type, {config.num_attention_heads * dim, config.hidden_size}); - w.q_bias = weights.store->load_f32_tensor(source, prefix + ".self_attn.q_proj.bias", {config.num_attention_heads * dim}); - w.k_proj = weights.store->load_tensor(source, prefix + ".self_attn.k_proj.weight", storage_type, {config.num_key_value_heads * dim, config.hidden_size}); - w.k_bias = weights.store->load_f32_tensor(source, prefix + ".self_attn.k_proj.bias", {config.num_key_value_heads * dim}); - w.v_proj = weights.store->load_tensor(source, prefix + ".self_attn.v_proj.weight", storage_type, {config.num_key_value_heads * dim, config.hidden_size}); - w.v_bias = weights.store->load_f32_tensor(source, prefix + ".self_attn.v_proj.bias", {config.num_key_value_heads * dim}); - w.o_proj = weights.store->load_tensor(source, prefix + ".self_attn.o_proj.weight", storage_type, {config.hidden_size, config.num_attention_heads * dim}); - w.post_norm = weights.store->load_f32_tensor(source, prefix + ".post_attention_layernorm.weight", {config.hidden_size}); - w.gate_proj = weights.store->load_tensor(source, prefix + ".mlp.gate_proj.weight", storage_type, {config.intermediate_size, config.hidden_size}); - w.up_proj = weights.store->load_tensor(source, prefix + ".mlp.up_proj.weight", storage_type, {config.intermediate_size, config.hidden_size}); - w.down_proj = weights.store->load_tensor(source, prefix + ".mlp.down_proj.weight", storage_type, {config.hidden_size, config.intermediate_size}); - weights.layers.push_back(std::move(w)); - } - weights.norm = weights.store->load_f32_tensor(source, "language_model.model.norm.weight", {config.hidden_size}); - if (config.tie_word_embeddings) { - weights.lm_head = weights.token_embedding; - } else { - weights.lm_head = weights.store->load_tensor( - source, - "language_model.lm_head.weight", - storage_type, - {config.vocab_size, config.hidden_size}); - } - weights.store->upload(); - return weights; -} -int32_t argmax_index(const std::vector & values) { - if (values.empty()) { - throw std::runtime_error("Audio8 ASR thinker cannot select from empty logits"); - } - size_t best = 0; - for (size_t i = 1; i < values.size(); ++i) { - if (values[i] > values[best]) { - best = i; - } - } - return static_cast(best); -} - -bool is_eos(const Audio8ASRDecoderConfig & config, int32_t token) { - return std::find(config.eos_token_ids.begin(), config.eos_token_ids.end(), static_cast(token)) != - config.eos_token_ids.end(); +runtime::GreedyQwenDecoderSpec make_decoder_spec(const Audio8ASRDecoderConfig & config) { + runtime::GreedyQwenDecoderSpec spec; + // Qwen2-style decoder: attention biases, no Q/K norms, RoPE theta 1e6. + spec.decoder.stack.hidden_size = config.hidden_size; + spec.decoder.stack.num_attention_heads = config.num_attention_heads; + spec.decoder.stack.num_key_value_heads = config.num_key_value_heads; + spec.decoder.stack.head_dim = config.head_dim; + spec.decoder.stack.intermediate_size = config.intermediate_size; + spec.decoder.stack.layers = config.num_hidden_layers; + spec.decoder.stack.rms_norm_eps = config.rms_norm_eps; + spec.decoder.stack.rope_theta = config.rope_theta; + spec.decoder.stack.use_qk_norm = false; + spec.decoder.stack.runtime.static_cache.update_mode = + modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; + spec.decoder.logits_size = config.vocab_size; + spec.decoder.logits_mode = modules::QwenCausalDecoderLogitsMode::LastStep; + spec.vocab_size = config.vocab_size; + spec.max_position_embeddings = config.max_position_embeddings; + spec.tie_word_embeddings = config.tie_word_embeddings; + spec.attention_bias = true; + spec.token_embedding_tensor = "language_model.model.embed_tokens.weight"; + spec.lm_head_tensor = "language_model.lm_head.weight"; + spec.final_norm_tensor = "language_model.model.norm.weight"; + spec.layer_prefix = "language_model.model.layers"; + spec.eos_token_ids = config.eos_token_ids; + return spec; } -class ThinkerWeightsRuntime { -public: - ThinkerWeightsRuntime( - std::shared_ptr source, - Audio8ASRDecoderConfig config, - core::ExecutionContext & execution, - size_t weight_context_bytes, - assets::TensorStorageType storage_type) - : source_(std::move(source)), - config_(std::make_shared(std::move(config))), - backend_(execution.backend()), - backend_type_(execution.backend_type()), - threads_(std::max(1, execution.config().threads)), - weights_(std::make_shared(load_weights( - *source_, - *config_, - backend_, - backend_type_, - weight_context_bytes, - storage_type))) {} - - const Audio8ASRDecoderConfig & config() const noexcept { - return *config_; - } - - const ThinkerWeights & weights() const noexcept { - return *weights_; - } - - ggml_backend_t backend() const noexcept { - return backend_; - } - - core::BackendType backend_type() const noexcept { - return backend_type_; - } - - int threads() const noexcept { - return threads_; - } - -private: - std::shared_ptr source_; - std::shared_ptr config_; - ggml_backend_t backend_ = nullptr; - core::BackendType backend_type_ = core::BackendType::Cpu; - int threads_ = 1; - std::shared_ptr weights_; -}; - -class PrefillGraph { -public: - PrefillGraph( - std::shared_ptr runtime, - int64_t prompt_steps, - int64_t audio_tokens, - size_t graph_arena_bytes) - : runtime_(std::move(runtime)), - prompt_steps_(prompt_steps), - audio_tokens_(audio_tokens) { - if (prompt_steps_ <= 0) { - throw std::runtime_error("Audio8 ASR thinker prefill requires positive prompt length"); - } - if (audio_tokens_ < 0 || audio_tokens_ > prompt_steps_) { - throw std::runtime_error("Audio8 ASR thinker prefill audio token count is invalid"); - } - const auto build_start = Clock::now(); - ggml_init_params params{graph_arena_bytes, nullptr, true}; - ctx_.reset(ggml_init(params)); - if (ctx_ == nullptr) { - throw std::runtime_error("failed to initialize Audio8 ASR thinker prefill graph context"); - } - const auto & config = runtime_->config(); - const auto & weights = runtime_->weights(); - core::ModuleBuildContext ctx{ctx_.get(), "audio8_asr.thinker.prefill", runtime_->backend_type()}; - token_ids_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, prompt_steps_); - audio_embeddings_ = ggml_new_tensor_2d(ctx_.get(), GGML_TYPE_F32, config.hidden_size, std::max(audio_tokens_, 1)); - audio_positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I64, std::max(audio_tokens_, 1)); - auto x = prompt_embeddings( - ctx, - weights, - config, - token_ids_, - audio_embeddings_, - audio_positions_, - prompt_steps_, - audio_tokens_); - positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, prompt_steps_); - auto positions = core::wrap_tensor(positions_, core::TensorShape::from_dims({prompt_steps_}), GGML_TYPE_I32); - - auto decoder_out = modules::QwenCausalDecoderModule(make_qwen_decoder_config(config)) - .build(ctx, x, positions, make_qwen_decoder_weights(weights)); - for (const auto & layer : decoder_out.state.layers) { - if (!layer.key.has_value() || !layer.value.has_value()) { - throw std::runtime_error("Audio8 ASR thinker prefill decoder did not return K/V state"); - } - // See qwen3_asr PrefillGraph: copy K/V out of the graph-allocated - // intermediates and mark them as outputs so the allocator cannot - // recycle them before run() reads them back. - 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_); - for (auto * key : keys_) { - ggml_build_forward_expand(graph_, key); - } - for (auto * value : values_) { - ggml_build_forward_expand(graph_, value); - } - const auto try_alloc = [&]() { - gallocr_.reset(ggml_gallocr_new(ggml_backend_get_default_buffer_type(runtime_->backend()))); - return gallocr_ != nullptr && - ggml_gallocr_reserve(gallocr_.get(), graph_) && - ggml_gallocr_alloc_graph(gallocr_.get(), graph_); - }; - if (!try_alloc() && - (engine::core::trim_backend_pools(runtime_->backend()), !try_alloc())) { - throw engine::runtime::CapacityError( - "Audio8 ASR prefill graph does not fit in device memory at this size (" - + std::to_string(prompt_steps_) + " prompt steps, of which " - + std::to_string(audio_tokens_) + " are audio tokens)"); - } - position_ids_ = modules::qwen_position_ids(prompt_steps_); - debug::timing_log_scalar("audio8_asr.thinker.prefill.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); - debug::trace_log_scalar("audio8_asr.thinker.prefill_prompt_steps", prompt_steps_); - } - - ~PrefillGraph() { - engine::core::release_backend_graph_resources(runtime_->backend(), graph_, true); - } - - bool matches(const ThinkerWeightsRuntime & runtime, int64_t prompt_steps, int64_t audio_tokens) const { - return runtime_.get() == &runtime && prompt_steps_ == prompt_steps && audio_tokens_ == audio_tokens; - } - - PrefillOutput run( - const std::vector & token_ids, - const std::vector & audio_embeddings, - const std::vector & audio_positions) { - const auto & config = runtime_->config(); - if (static_cast(token_ids.size()) != prompt_steps_) { - throw std::runtime_error("Audio8 ASR thinker prefill token id count mismatch"); - } - if (static_cast(audio_embeddings.size()) != audio_tokens_ * config.hidden_size) { - throw std::runtime_error("Audio8 ASR thinker prefill audio embedding size mismatch"); - } - if (static_cast(audio_positions.size()) != audio_tokens_) { - throw std::runtime_error("Audio8 ASR thinker prefill audio position count mismatch"); - } - auto timing_start = Clock::now(); - // Re-uploaded on every 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_, token_ids.data(), 0, token_ids.size() * sizeof(int32_t)); - if (audio_tokens_ > 0) { - std::vector positions(audio_positions.begin(), audio_positions.end()); - ggml_backend_tensor_set( - audio_embeddings_, - audio_embeddings.data(), - 0, - audio_embeddings.size() * sizeof(float)); - ggml_backend_tensor_set( - audio_positions_, - positions.data(), - 0, - positions.size() * sizeof(int64_t)); - } - core::set_backend_threads(runtime_->backend(), runtime_->threads()); - timing_start = Clock::now(); - const ggml_status status = engine::core::compute_backend_graph(runtime_->backend(), graph_); - ggml_backend_synchronize(runtime_->backend()); - debug::timing_log_scalar("audio8_asr.thinker.prefill.graph.compute_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); - if (status != GGML_STATUS_SUCCESS) { - throw std::runtime_error("Audio8 ASR thinker prefill graph compute failed"); - } - PrefillOutput out; - out.logits.resize(static_cast(config.vocab_size)); - ggml_backend_tensor_get(logits_, out.logits.data(), 0, out.logits.size() * sizeof(float)); - out.kv_state.current_end = prompt_steps_; - out.kv_state.layers.resize(keys_.size()); - const size_t layer_values = static_cast( - prompt_steps_ * config.num_key_value_heads * config.head_dim); - for (size_t layer = 0; layer < keys_.size(); ++layer) { - auto & state = out.kv_state.layers[layer]; - state.valid_steps = prompt_steps_; - state.key.resize(layer_values); - state.value.resize(layer_values); - ggml_backend_tensor_get(keys_[layer], state.key.data(), 0, state.key.size() * sizeof(float)); - ggml_backend_tensor_get(values_[layer], state.value.data(), 0, state.value.size() * sizeof(float)); - } - return out; - } - -private: - std::shared_ptr runtime_; - int64_t prompt_steps_ = 0; - int64_t audio_tokens_ = 0; - std::unique_ptr ctx_; - ggml_tensor * token_ids_ = nullptr; - ggml_tensor * audio_embeddings_ = nullptr; - ggml_tensor * audio_positions_ = nullptr; - ggml_tensor * positions_ = nullptr; - ggml_tensor * logits_ = nullptr; - std::vector keys_; - std::vector values_; - std::vector position_ids_; - ggml_cgraph * graph_ = nullptr; - std::unique_ptr, GgmlGallocrDeleter> gallocr_; -}; - -class DecodeGraph { -public: - DecodeGraph(std::shared_ptr runtime, int64_t cache_steps, size_t graph_arena_bytes) - : runtime_(std::move(runtime)), - cache_steps_(cache_steps) { - if (cache_steps_ <= 0) { - throw std::runtime_error("Audio8 ASR thinker decode requires positive cache length"); - } - const auto build_start = Clock::now(); - ggml_init_params params{graph_arena_bytes, nullptr, true}; - ctx_.reset(ggml_init(params)); - if (ctx_ == nullptr) { - throw std::runtime_error("failed to initialize Audio8 ASR thinker decode graph context"); - } - const auto & config = runtime_->config(); - const auto & weights = runtime_->weights(); - core::ModuleBuildContext ctx{ctx_.get(), "audio8_asr.thinker.decode", runtime_->backend_type()}; - token_id_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, 1); - auto token_id = core::wrap_tensor(token_id_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); - auto x = modules::EmbeddingModule({config.vocab_size, config.hidden_size}) - .build(ctx, token_id, weights.token_embedding); - x = core::reshape_tensor(ctx, x, core::TensorShape::from_dims({1, 1, config.hidden_size})); - positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, 1); - auto positions = core::wrap_tensor(positions_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); - cache_slot_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, 1); - auto cache_slot = core::wrap_tensor(cache_slot_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); - attention_mask_ = ggml_new_tensor_4d(ctx_.get(), GGML_TYPE_F16, cache_steps_, 1, 1, 1); - auto attention_mask = core::wrap_tensor( - attention_mask_, - core::TensorShape::from_dims({1, 1, 1, cache_steps_}), - GGML_TYPE_F16); - graph_ = ggml_new_graph_custom(ctx_.get(), 65536, false); - auto decoder_out = modules::QwenCausalDecoderModule(make_qwen_decoder_config(config)) - .build_static_cache_tail( - ctx, - graph_, - x, - positions, - make_qwen_decoder_weights(weights), - cache_steps_, - attention_mask, - cache_slot); - step_cache_ = std::move(decoder_out.cache); - logits_ = decoder_out.logits.tensor; - ggml_set_output(logits_); - ggml_build_forward_expand(graph_, logits_); - buffer_ = ggml_backend_alloc_ctx_tensors(ctx_.get(), runtime_->backend()); - if (buffer_ == nullptr) { - engine::core::trim_backend_pools(runtime_->backend()); - buffer_ = ggml_backend_alloc_ctx_tensors(ctx_.get(), runtime_->backend()); - } - if (buffer_ == nullptr) { - throw std::runtime_error("failed to allocate Audio8 ASR thinker decode graph"); - } - attention_mask_values_.assign(static_cast(cache_steps_), ggml_fp32_to_fp16(-INFINITY)); - debug::timing_log_scalar("audio8_asr.thinker.decode.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); - debug::trace_log_scalar("audio8_asr.thinker.decode_cache_steps", cache_steps_); - } - - ~DecodeGraph() { - engine::core::release_backend_graph_resources(runtime_->backend(), graph_, true); - if (buffer_ != nullptr) { - ggml_backend_buffer_free(buffer_); - } - } - - bool can_run(const ThinkerWeightsRuntime & runtime, int64_t required_steps) const { - return runtime_.get() == &runtime && cache_steps_ >= required_steps; - } - - void import_state(const runtime::TransformerKVState & state) { - step_cache_.import_state(state); - } - - std::vector run_step(int32_t token) { - const auto & config = runtime_->config(); - if (step_cache_.valid_steps() >= cache_steps_) { - throw std::runtime_error("Audio8 ASR thinker decode cache exhausted"); - } - ggml_backend_tensor_set(token_id_, &token, 0, sizeof(int32_t)); - const int32_t position = static_cast(step_cache_.current_end()); - ggml_backend_tensor_set(positions_, &position, 0, sizeof(int32_t)); - const int32_t cache_slot = static_cast(step_cache_.valid_steps()); - ggml_backend_tensor_set(cache_slot_, &cache_slot, 0, sizeof(int32_t)); - modules::write_qwen_cached_step_mask( - attention_mask_, - attention_mask_values_, - cache_steps_, - step_cache_.valid_steps(), - step_cache_.valid_steps()); - core::set_backend_threads(runtime_->backend(), runtime_->threads()); - const ggml_status status = engine::core::compute_backend_graph(runtime_->backend(), graph_); - ggml_backend_synchronize(runtime_->backend()); - if (status != GGML_STATUS_SUCCESS) { - throw std::runtime_error("Audio8 ASR thinker decode graph compute failed"); - } - logits_buffer_.resize(static_cast(config.vocab_size)); - ggml_backend_tensor_get(logits_, logits_buffer_.data(), 0, logits_buffer_.size() * sizeof(float)); - step_cache_.advance_after_direct_append(1); - // The caller moves out of this buffer before the next step. - return std::move(logits_buffer_); - } - -private: - std::shared_ptr runtime_; - int64_t cache_steps_ = 0; - std::unique_ptr ctx_; - ggml_tensor * token_id_ = nullptr; - ggml_tensor * positions_ = nullptr; - ggml_tensor * cache_slot_ = nullptr; - ggml_tensor * attention_mask_ = nullptr; - ggml_tensor * logits_ = nullptr; - std::vector attention_mask_values_; - std::vector logits_buffer_; - runtime::TransformerKVCache step_cache_; - ggml_cgraph * graph_ = nullptr; - ggml_backend_buffer_t buffer_ = nullptr; -}; - } // namespace -struct Audio8ThinkerRuntime::Impl { - Impl( - std::shared_ptr weights_source, - Audio8ASRDecoderConfig config, - core::ExecutionContext & execution, - size_t prefill_graph_arena_bytes, - size_t decode_graph_arena_bytes, - size_t weight_context_bytes, - assets::TensorStorageType storage_type) - : weights(std::make_shared( - std::move(weights_source), - std::move(config), - execution, - weight_context_bytes, - storage_type)), - prefill_graph_arena_bytes(prefill_graph_arena_bytes), - decode_graph_arena_bytes(decode_graph_arena_bytes) {} - - void validate_prompt_audio( - const Audio8ASRPrompt & prompt, - const Audio8ASRAudioEmbeddings & audio_embeddings) const { - const auto & config = weights->config(); - if (audio_embeddings.hidden_size != config.hidden_size) { - throw std::runtime_error("Audio8 ASR audio embedding hidden size mismatch"); - } - if (audio_embeddings.tokens != static_cast(prompt.audio_token_positions.size())) { - throw std::runtime_error("Audio8 ASR audio embedding token count does not match prompt placeholders"); - } - if (static_cast(audio_embeddings.values.size()) != audio_embeddings.tokens * config.hidden_size) { - throw std::runtime_error("Audio8 ASR audio embedding value count mismatch"); - } - for (const int32_t position : prompt.audio_token_positions) { - if (position < 0 || position >= static_cast(prompt.input_ids.size())) { - throw std::runtime_error("Audio8 ASR audio placeholder position out of range"); - } - } - } - - std::shared_ptr weights; - size_t prefill_graph_arena_bytes = 0; - size_t decode_graph_arena_bytes = 0; - std::unique_ptr prefill_graph; - std::unique_ptr decode_graph; -}; - Audio8ThinkerRuntime::Audio8ThinkerRuntime( std::shared_ptr weights_source, const Audio8ASRDecoderConfig & config, @@ -620,14 +46,15 @@ Audio8ThinkerRuntime::Audio8ThinkerRuntime( size_t decode_graph_arena_bytes, size_t weight_context_bytes, assets::TensorStorageType weight_storage_type) - : impl_(std::make_unique( + : runtime_( std::move(weights_source), - config, + make_decoder_spec(config), execution, prefill_graph_arena_bytes, decode_graph_arena_bytes, weight_context_bytes, - weight_storage_type)) {} + weight_storage_type), + config_(std::make_shared(config)) {} Audio8ThinkerRuntime::~Audio8ThinkerRuntime() = default; @@ -635,54 +62,26 @@ Audio8ASRGeneratedTokens Audio8ThinkerRuntime::generate( const Audio8ASRPrompt & prompt, const Audio8ASRAudioEmbeddings & audio_embeddings, const Audio8ASRGenerationOptions & options) { - const auto & config = impl_->weights->config(); if (prompt.input_ids.empty()) { throw std::runtime_error("Audio8 ASR thinker prompt is empty"); } - if (options.max_new_tokens <= 0) { - throw std::runtime_error("Audio8 ASR max_new_tokens must be positive"); + if (audio_embeddings.hidden_size != config_->hidden_size || + audio_embeddings.tokens != static_cast(prompt.audio_token_positions.size()) || + static_cast(audio_embeddings.values.size()) != audio_embeddings.tokens * config_->hidden_size) { + throw std::runtime_error("Audio8 ASR audio embeddings do not match the prompt placeholders"); } - const int64_t prompt_steps = static_cast(prompt.input_ids.size()); - if (prompt_steps + options.max_new_tokens > config.max_position_embeddings) { - throw std::runtime_error("Audio8 ASR thinker request exceeds max_position_embeddings"); - } - impl_->validate_prompt_audio(prompt, audio_embeddings); - if (impl_->prefill_graph == nullptr || - !impl_->prefill_graph->matches(*impl_->weights, prompt_steps, audio_embeddings.tokens)) { - impl_->prefill_graph.reset(); - impl_->prefill_graph = std::make_unique( - impl_->weights, - prompt_steps, - audio_embeddings.tokens, - impl_->prefill_graph_arena_bytes); - } - auto prefill = impl_->prefill_graph->run( - prompt.input_ids, - audio_embeddings.values, - prompt.audio_token_positions); - const int64_t required_cache_steps = prompt_steps + options.max_new_tokens; - if (impl_->decode_graph == nullptr || - !impl_->decode_graph->can_run(*impl_->weights, required_cache_steps)) { - impl_->decode_graph.reset(); - impl_->decode_graph = std::make_unique( - impl_->weights, - required_cache_steps, - impl_->decode_graph_arena_bytes); - } - impl_->decode_graph->import_state(prefill.kv_state); - - Audio8ASRGeneratedTokens out; - std::vector logits = std::move(prefill.logits); - const auto decode_start = Clock::now(); - for (int64_t step = 0; step < options.max_new_tokens; ++step) { - const int32_t token = argmax_index(logits); - if (is_eos(config, token)) { - break; + for (const int32_t position : prompt.audio_token_positions) { + if (position < 0 || position >= static_cast(prompt.input_ids.size())) { + throw std::runtime_error("Audio8 ASR audio placeholder position out of range"); } - out.token_ids.push_back(token); - logits = impl_->decode_graph->run_step(token); } - debug::timing_log_scalar("audio8_asr.thinker.decode_total_ms", engine::debug::elapsed_ms(decode_start, Clock::now())); + runtime::GreedyQwenDecoderRuntime::Prompt decoder_prompt; + decoder_prompt.input_ids = prompt.input_ids; + decoder_prompt.injection.values = audio_embeddings.values; + decoder_prompt.injection.tokens = audio_embeddings.tokens; + decoder_prompt.injection.positions = prompt.audio_token_positions; + Audio8ASRGeneratedTokens out; + out.token_ids = runtime_.generate(decoder_prompt, options.max_new_tokens); return out; } diff --git a/src/framework/runtime/greedy_qwen_decoder.cpp b/src/framework/runtime/greedy_qwen_decoder.cpp new file mode 100644 index 000000000..e9185c23c --- /dev/null +++ b/src/framework/runtime/greedy_qwen_decoder.cpp @@ -0,0 +1,689 @@ +#include "engine/framework/runtime/greedy_qwen_decoder.h" + +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/modules/lookup_modules.h" +#include "engine/framework/modules/positional_modules.h" +#include "engine/framework/modules/primitive_modules.h" +#include "engine/framework/modules/structural_modules.h" +#include "engine/framework/runtime/errors.h" +#include "engine/framework/runtime/kv_cache.h" +#include "engine/framework/sampling/decode_modules.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::runtime { +namespace { + +namespace modules = engine::modules; +using Clock = std::chrono::steady_clock; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +struct GgmlGallocrDeleter { + void operator()(ggml_gallocr_t alloc) const noexcept { + if (alloc != nullptr) { + ggml_gallocr_free(alloc); + } + } +}; + +struct DecoderLayerWeights { + core::TensorValue input_norm; + core::TensorValue q_proj; + core::TensorValue q_bias; + core::TensorValue k_proj; + core::TensorValue k_bias; + core::TensorValue v_proj; + core::TensorValue v_bias; + core::TensorValue qkv_weight; + core::TensorValue qkv_bias; + core::TensorValue o_proj; + core::TensorValue q_norm; + core::TensorValue k_norm; + core::TensorValue post_norm; + core::TensorValue gate_proj; + core::TensorValue up_proj; + core::TensorValue down_proj; +}; + +struct DecoderWeights { + std::shared_ptr store; + core::TensorValue token_embedding; + std::vector layers; + core::TensorValue norm; + core::TensorValue lm_head; +}; + +struct PrefillOutput { + std::vector logits; + runtime::TransformerKVState kv_state; +}; + +modules::QwenDecoderLayerWeights bind_layer_weights( + const DecoderLayerWeights & weights, + const GreedyQwenDecoderSpec & spec) { + modules::QwenDecoderLayerWeights out; + out.input_norm = {weights.input_norm, std::nullopt}; + out.self_attention.q_weight = weights.q_proj; + if (spec.attention_bias) { + out.self_attention.q_bias = weights.q_bias; + out.self_attention.k_bias = weights.k_bias; + out.self_attention.v_bias = weights.v_bias; + } + out.self_attention.k_weight = weights.k_proj; + out.self_attention.v_weight = weights.v_proj; + if (spec.packed_qkv) { + out.self_attention.qkv_weight = weights.qkv_weight; + if (spec.attention_bias) { + out.self_attention.qkv_bias = weights.qkv_bias; + } + } + out.self_attention.out_weight = weights.o_proj; + if (spec.decoder.stack.use_qk_norm) { + out.q_norm = {weights.q_norm, std::nullopt}; + out.k_norm = {weights.k_norm, std::nullopt}; + } + out.post_norm = {weights.post_norm, std::nullopt}; + out.mlp.gate_proj = {weights.gate_proj, std::nullopt}; + out.mlp.up_proj = {weights.up_proj, std::nullopt}; + out.mlp.down_proj = {weights.down_proj, std::nullopt}; + return out; +} + +modules::QwenCausalDecoderWeights bind_decoder_weights( + const DecoderWeights & weights, + const GreedyQwenDecoderSpec & spec) { + modules::QwenCausalDecoderWeights out; + out.stack.layers.reserve(weights.layers.size()); + for (const auto & layer : weights.layers) { + out.stack.layers.push_back(bind_layer_weights(layer, spec)); + } + out.final_norm = {weights.norm, std::nullopt}; + out.lm_head = {weights.lm_head, std::nullopt}; + return out; +} + +core::TensorValue prompt_embeddings( + core::ModuleBuildContext & ctx, + const DecoderWeights & weights, + const GreedyQwenDecoderSpec & spec, + ggml_tensor * token_ids, + int64_t prompt_steps, + const std::vector & injection_values, + int64_t injection_tokens, + const std::vector & injection_positions, + ggml_tensor * injection_values_tensor, + ggml_tensor * injection_positions_tensor) { + auto ids = core::wrap_tensor(token_ids, core::TensorShape::from_dims({prompt_steps}), GGML_TYPE_I32); + auto x = modules::EmbeddingModule({spec.vocab_size, spec.decoder.stack.hidden_size}) + .build(ctx, ids, weights.token_embedding); + if (injection_tokens > 0) { + auto injection = core::wrap_tensor( + injection_values_tensor, + core::TensorShape::from_dims({injection_tokens, spec.decoder.stack.hidden_size}), + GGML_TYPE_F32); + auto positions = core::wrap_tensor( + injection_positions_tensor, + core::TensorShape::from_dims({injection_tokens}), + GGML_TYPE_I64); + x = core::wrap_tensor( + ggml_set_rows(ctx.ggml, x.tensor, injection.tensor, positions.tensor), + x.shape, + GGML_TYPE_F32); + } + (void)injection_values; + (void)injection_positions; + return core::reshape_tensor( + ctx, x, core::TensorShape::from_dims({1, prompt_steps, spec.decoder.stack.hidden_size})); +} + +DecoderWeights load_weights( + const assets::TensorSource & source, + const GreedyQwenDecoderSpec & spec, + ggml_backend_t backend, + core::BackendType backend_type, + size_t weight_context_bytes, + assets::TensorStorageType storage_type) { + const auto & stack = spec.decoder.stack; + if (stack.hidden_size <= 0 || stack.num_attention_heads <= 0 || stack.head_dim <= 0 || + stack.layers <= 0 || spec.vocab_size <= 0) { + throw std::runtime_error("Greedy Qwen decoder spec is invalid"); + } + DecoderWeights weights; + weights.store = std::make_shared( + backend, + backend_type, + "greedy_qwen_decoder.weights", + weight_context_bytes); + weights.token_embedding = weights.store->load_tensor( + source, + spec.token_embedding_tensor, + storage_type, + {spec.vocab_size, stack.hidden_size}); + const int64_t dim = stack.head_dim; + weights.layers.reserve(static_cast(stack.layers)); + for (int64_t layer = 0; layer < stack.layers; ++layer) { + const std::string prefix = spec.layer_prefix + "." + std::to_string(layer); + DecoderLayerWeights w; + w.input_norm = weights.store->load_f32_tensor(source, prefix + ".input_layernorm.weight", {stack.hidden_size}); + if (spec.packed_qkv) { + const int64_t qkv_rows = + (stack.num_attention_heads + 2 * stack.num_key_value_heads) * dim; + w.qkv_weight = weights.store->load_tensor( + source, prefix + ".self_attn.qkv_proj.weight", storage_type, {qkv_rows, stack.hidden_size}); + if (spec.attention_bias) { + w.qkv_bias = weights.store->load_f32_tensor(source, prefix + ".self_attn.qkv_proj.bias", {qkv_rows}); + } + } else { + w.q_proj = weights.store->load_tensor(source, prefix + ".self_attn.q_proj.weight", storage_type, {stack.num_attention_heads * dim, stack.hidden_size}); + w.k_proj = weights.store->load_tensor(source, prefix + ".self_attn.k_proj.weight", storage_type, {stack.num_key_value_heads * dim, stack.hidden_size}); + w.v_proj = weights.store->load_tensor(source, prefix + ".self_attn.v_proj.weight", storage_type, {stack.num_key_value_heads * dim, stack.hidden_size}); + if (spec.attention_bias) { + w.q_bias = weights.store->load_f32_tensor(source, prefix + ".self_attn.q_proj.bias", {stack.num_attention_heads * dim}); + w.k_bias = weights.store->load_f32_tensor(source, prefix + ".self_attn.k_proj.bias", {stack.num_key_value_heads * dim}); + w.v_bias = weights.store->load_f32_tensor(source, prefix + ".self_attn.v_proj.bias", {stack.num_key_value_heads * dim}); + } + } + w.o_proj = weights.store->load_tensor(source, prefix + ".self_attn.o_proj.weight", storage_type, {stack.hidden_size, stack.num_attention_heads * dim}); + if (stack.use_qk_norm) { + w.q_norm = weights.store->load_f32_tensor(source, prefix + ".self_attn.q_norm.weight", {dim}); + w.k_norm = weights.store->load_f32_tensor(source, prefix + ".self_attn.k_norm.weight", {dim}); + } + w.post_norm = weights.store->load_f32_tensor(source, prefix + ".post_attention_layernorm.weight", {stack.hidden_size}); + w.gate_proj = weights.store->load_tensor(source, prefix + ".mlp.gate_proj.weight", storage_type, {stack.intermediate_size, stack.hidden_size}); + w.up_proj = weights.store->load_tensor(source, prefix + ".mlp.up_proj.weight", storage_type, {stack.intermediate_size, stack.hidden_size}); + w.down_proj = weights.store->load_tensor(source, prefix + ".mlp.down_proj.weight", storage_type, {stack.hidden_size, stack.intermediate_size}); + weights.layers.push_back(std::move(w)); + } + weights.norm = weights.store->load_f32_tensor(source, spec.final_norm_tensor, {stack.hidden_size}); + if (spec.tie_word_embeddings) { + if (spec.decoder.logits_size != 0 && spec.decoder.logits_size != spec.vocab_size) { + throw std::runtime_error("tied output embedding requires logits_size == vocab_size"); + } + weights.lm_head = weights.token_embedding; + } else { + weights.lm_head = weights.store->load_tensor( + source, + spec.lm_head_tensor, + storage_type, + {spec.decoder.logits_size != 0 ? spec.decoder.logits_size : spec.vocab_size, stack.hidden_size}); + } + weights.store->upload(); + return weights; +} + +int32_t argmax_index(const std::vector & values) { + if (values.empty()) { + throw std::runtime_error("Greedy Qwen decoder cannot select from empty logits"); + } + size_t best = 0; + for (size_t i = 1; i < values.size(); ++i) { + if (values[i] > values[best]) { + best = i; + } + } + return static_cast(best); +} + +bool is_eos(const GreedyQwenDecoderSpec & spec, int32_t token) { + return std::find(spec.eos_token_ids.begin(), spec.eos_token_ids.end(), static_cast(token)) != + spec.eos_token_ids.end(); +} + +class ThinkerWeightsRuntime { +public: + ThinkerWeightsRuntime( + std::shared_ptr source, + GreedyQwenDecoderSpec spec, + core::ExecutionContext & execution, + size_t weight_context_bytes, + assets::TensorStorageType storage_type) + : spec_(std::make_shared(std::move(spec))), + backend_(execution.backend()), + backend_type_(execution.backend_type()), + threads_(std::max(1, execution.config().threads)), + weights_(std::make_shared(load_weights( + *source, + *spec_, + backend_, + backend_type_, + weight_context_bytes, + storage_type))) {} + + const GreedyQwenDecoderSpec & spec() const noexcept { + return *spec_; + } + + const DecoderWeights & weights() const noexcept { + return *weights_; + } + + ggml_backend_t backend() const noexcept { + return backend_; + } + + core::BackendType backend_type() const noexcept { + return backend_type_; + } + + int threads() const noexcept { + return threads_; + } + +private: + std::shared_ptr spec_; + ggml_backend_t backend_ = nullptr; + core::BackendType backend_type_ = core::BackendType::Cpu; + int threads_ = 1; + std::shared_ptr weights_; +}; + +class PrefillGraph { +public: + PrefillGraph( + std::shared_ptr runtime, + int64_t prompt_steps, + int64_t injection_tokens, + size_t graph_arena_bytes) + : runtime_(std::move(runtime)), + prompt_steps_(prompt_steps), + injection_tokens_(injection_tokens) { + if (prompt_steps_ <= 0) { + throw std::runtime_error("Greedy Qwen decoder prefill requires positive prompt length"); + } + if (injection_tokens_ < 0 || injection_tokens_ > prompt_steps_) { + throw std::runtime_error("Greedy Qwen decoder prefill injection token count is invalid"); + } + const auto build_start = Clock::now(); + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize greedy Qwen decoder prefill graph context"); + } + const auto & spec = runtime_->spec(); + const auto & weights = runtime_->weights(); + core::ModuleBuildContext ctx{ctx_.get(), "greedy_qwen_decoder.prefill", runtime_->backend_type()}; + token_ids_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, prompt_steps_); + injection_values_ = ggml_new_tensor_2d( + ctx_.get(), GGML_TYPE_F32, spec.decoder.stack.hidden_size, std::max(injection_tokens_, 1)); + injection_positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I64, std::max(injection_tokens_, 1)); + auto x = prompt_embeddings( + ctx, + weights, + spec, + token_ids_, + prompt_steps_, + {}, + injection_tokens_, + {}, + injection_values_, + injection_positions_); + positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, prompt_steps_); + auto positions = core::wrap_tensor(positions_, core::TensorShape::from_dims({prompt_steps_}), GGML_TYPE_I32); + + auto decoder_out = modules::QwenCausalDecoderModule(spec.decoder) + .build(ctx, x, positions, bind_decoder_weights(weights, spec)); + for (const auto & layer : decoder_out.state.layers) { + if (!layer.key.has_value() || !layer.value.has_value()) { + throw std::runtime_error("greedy Qwen decoder prefill did not return K/V state"); + } + // The graph allocator recycles intermediates; copy K/V into their + // own tensors and mark them as outputs so run() can read them back. + 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_); + for (auto * key : keys_) { + ggml_build_forward_expand(graph_, key); + } + for (auto * value : values_) { + ggml_build_forward_expand(graph_, value); + } + const auto try_alloc = [&]() { + gallocr_.reset(ggml_gallocr_new(ggml_backend_get_default_buffer_type(runtime_->backend()))); + return gallocr_ != nullptr && + ggml_gallocr_reserve(gallocr_.get(), graph_) && + ggml_gallocr_alloc_graph(gallocr_.get(), graph_); + }; + if (!try_alloc() && + (engine::core::trim_backend_pools(runtime_->backend()), !try_alloc())) { + throw engine::runtime::CapacityError( + "greedy Qwen decoder prefill graph does not fit in device memory at this size (" + + std::to_string(prompt_steps_) + " prompt steps, of which " + + std::to_string(injection_tokens_) + " are injected tokens)"); + } + position_ids_ = modules::qwen_position_ids(prompt_steps_); + debug::timing_log_scalar("greedy_qwen_decoder.prefill.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); + debug::trace_log_scalar("greedy_qwen_decoder.prefill_prompt_steps", prompt_steps_); + } + + ~PrefillGraph() { + engine::core::release_backend_graph_resources(runtime_->backend(), graph_, true); + } + + bool matches(int64_t prompt_steps, int64_t injection_tokens) const { + return prompt_steps_ == prompt_steps && injection_tokens_ == injection_tokens; + } + + PrefillOutput run( + const std::vector & token_ids, + const std::vector & injection_values, + const std::vector & injection_positions) { + const auto & spec = runtime_->spec(); + if (static_cast(token_ids.size()) != prompt_steps_) { + throw std::runtime_error("greedy Qwen decoder prefill token id count mismatch"); + } + if (static_cast(injection_values.size()) != injection_tokens_ * spec.decoder.stack.hidden_size) { + throw std::runtime_error("greedy Qwen decoder prefill injection value size mismatch"); + } + if (static_cast(injection_positions.size()) != injection_tokens_) { + throw std::runtime_error("greedy Qwen decoder prefill injection position count mismatch"); + } + // Re-uploaded on every 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_, token_ids.data(), 0, token_ids.size() * sizeof(int32_t)); + if (injection_tokens_ > 0) { + std::vector positions(injection_positions.begin(), injection_positions.end()); + ggml_backend_tensor_set( + injection_values_, injection_values.data(), 0, injection_values.size() * sizeof(float)); + ggml_backend_tensor_set( + injection_positions_, positions.data(), 0, positions.size() * sizeof(int64_t)); + } + core::set_backend_threads(runtime_->backend(), runtime_->threads()); + const ggml_status status = engine::core::compute_backend_graph(runtime_->backend(), graph_); + ggml_backend_synchronize(runtime_->backend()); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("greedy Qwen decoder prefill graph compute failed"); + } + PrefillOutput out; + const int64_t logits_size = spec.decoder.logits_size != 0 + ? spec.decoder.logits_size + : spec.vocab_size; + out.logits.resize(static_cast(logits_size)); + ggml_backend_tensor_get(logits_, out.logits.data(), 0, out.logits.size() * sizeof(float)); + out.kv_state.current_end = prompt_steps_; + out.kv_state.layers.resize(keys_.size()); + const size_t layer_values = static_cast( + prompt_steps_ * spec.decoder.stack.num_key_value_heads * spec.decoder.stack.head_dim); + for (size_t layer = 0; layer < keys_.size(); ++layer) { + auto & state = out.kv_state.layers[layer]; + state.valid_steps = prompt_steps_; + state.key.resize(layer_values); + state.value.resize(layer_values); + ggml_backend_tensor_get(keys_[layer], state.key.data(), 0, state.key.size() * sizeof(float)); + ggml_backend_tensor_get(values_[layer], state.value.data(), 0, state.value.size() * sizeof(float)); + } + return out; + } + +private: + std::shared_ptr runtime_; + int64_t prompt_steps_ = 0; + int64_t injection_tokens_ = 0; + std::unique_ptr ctx_; + ggml_tensor * token_ids_ = nullptr; + ggml_tensor * injection_values_ = nullptr; + ggml_tensor * injection_positions_ = nullptr; + ggml_tensor * positions_ = nullptr; + ggml_tensor * logits_ = nullptr; + std::vector keys_; + std::vector values_; + std::vector position_ids_; + ggml_cgraph * graph_ = nullptr; + std::unique_ptr, GgmlGallocrDeleter> gallocr_; +}; + +class DecodeGraph { +public: + DecodeGraph(std::shared_ptr runtime, int64_t cache_steps, size_t graph_arena_bytes) + : runtime_(std::move(runtime)), + cache_steps_(cache_steps) { + if (cache_steps_ <= 0) { + throw std::runtime_error("greedy Qwen decoder decode requires positive cache length"); + } + const auto build_start = Clock::now(); + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize greedy Qwen decoder decode graph context"); + } + const auto & spec = runtime_->spec(); + const auto & weights = runtime_->weights(); + core::ModuleBuildContext ctx{ctx_.get(), "greedy_qwen_decoder.decode", runtime_->backend_type()}; + token_id_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, 1); + auto token_id = core::wrap_tensor(token_id_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); + auto x = modules::EmbeddingModule({spec.vocab_size, spec.decoder.stack.hidden_size}) + .build(ctx, token_id, weights.token_embedding); + x = core::reshape_tensor(ctx, x, core::TensorShape::from_dims({1, 1, spec.decoder.stack.hidden_size})); + positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, 1); + auto positions = core::wrap_tensor(positions_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); + cache_slot_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, 1); + auto cache_slot = core::wrap_tensor(cache_slot_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); + attention_mask_ = ggml_new_tensor_4d(ctx_.get(), GGML_TYPE_F16, cache_steps_, 1, 1, 1); + auto attention_mask = core::wrap_tensor( + attention_mask_, + core::TensorShape::from_dims({1, 1, 1, cache_steps_}), + GGML_TYPE_F16); + graph_ = ggml_new_graph_custom(ctx_.get(), 65536, false); + auto decoder_out = modules::QwenCausalDecoderModule(spec.decoder) + .build_static_cache_tail( + ctx, + graph_, + x, + positions, + bind_decoder_weights(weights, spec), + cache_steps_, + attention_mask, + cache_slot); + step_cache_ = std::move(decoder_out.cache); + logits_ = decoder_out.logits.tensor; + ggml_set_output(logits_); + ggml_build_forward_expand(graph_, logits_); + buffer_ = ggml_backend_alloc_ctx_tensors(ctx_.get(), runtime_->backend()); + if (buffer_ == nullptr) { + engine::core::trim_backend_pools(runtime_->backend()); + buffer_ = ggml_backend_alloc_ctx_tensors(ctx_.get(), runtime_->backend()); + } + if (buffer_ == nullptr) { + throw std::runtime_error("failed to allocate greedy Qwen decoder decode graph"); + } + attention_mask_values_.assign(static_cast(cache_steps_), ggml_fp32_to_fp16(-INFINITY)); + debug::timing_log_scalar("greedy_qwen_decoder.decode.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); + debug::trace_log_scalar("greedy_qwen_decoder.decode_cache_steps", cache_steps_); + } + + ~DecodeGraph() { + engine::core::release_backend_graph_resources(runtime_->backend(), graph_, true); + if (buffer_ != nullptr) { + ggml_backend_buffer_free(buffer_); + } + } + + bool can_run(int64_t required_steps) const { + return cache_steps_ >= required_steps; + } + + void import_state(const runtime::TransformerKVState & state) { + step_cache_.import_state(state); + } + + std::vector run_step(int32_t token) { + const auto & spec = runtime_->spec(); + if (step_cache_.valid_steps() >= cache_steps_) { + throw std::runtime_error("greedy Qwen decoder decode cache exhausted"); + } + ggml_backend_tensor_set(token_id_, &token, 0, sizeof(int32_t)); + const int32_t position = static_cast(step_cache_.current_end()); + ggml_backend_tensor_set(positions_, &position, 0, sizeof(int32_t)); + const int32_t cache_slot = static_cast(step_cache_.valid_steps()); + ggml_backend_tensor_set(cache_slot_, &cache_slot, 0, sizeof(int32_t)); + modules::write_qwen_cached_step_mask( + attention_mask_, + attention_mask_values_, + cache_steps_, + step_cache_.valid_steps(), + step_cache_.valid_steps()); + core::set_backend_threads(runtime_->backend(), runtime_->threads()); + const ggml_status status = engine::core::compute_backend_graph(runtime_->backend(), graph_); + ggml_backend_synchronize(runtime_->backend()); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("greedy Qwen decoder decode graph compute failed"); + } + const int64_t logits_size = spec.decoder.logits_size != 0 + ? spec.decoder.logits_size + : spec.vocab_size; + logits_buffer_.resize(static_cast(logits_size)); + ggml_backend_tensor_get(logits_, logits_buffer_.data(), 0, logits_buffer_.size() * sizeof(float)); + step_cache_.advance_after_direct_append(1); + // The caller moves out of this buffer before the next step. + return std::move(logits_buffer_); + } + +private: + std::shared_ptr runtime_; + int64_t cache_steps_ = 0; + std::unique_ptr ctx_; + ggml_tensor * token_id_ = nullptr; + ggml_tensor * positions_ = nullptr; + ggml_tensor * cache_slot_ = nullptr; + ggml_tensor * attention_mask_ = nullptr; + ggml_tensor * logits_ = nullptr; + std::vector attention_mask_values_; + std::vector logits_buffer_; + runtime::TransformerKVCache step_cache_; + ggml_cgraph * graph_ = nullptr; + ggml_backend_buffer_t buffer_ = nullptr; +}; + +} // namespace + +struct GreedyQwenDecoderRuntime::Impl { + Impl( + std::shared_ptr weights_source, + GreedyQwenDecoderSpec spec, + core::ExecutionContext & execution, + size_t prefill_graph_arena_bytes, + size_t decode_graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType storage_type) + : weights(std::make_shared( + std::move(weights_source), + std::move(spec), + execution, + weight_context_bytes, + storage_type)), + prefill_graph_arena_bytes(prefill_graph_arena_bytes), + decode_graph_arena_bytes(decode_graph_arena_bytes) {} + + std::shared_ptr weights; + size_t prefill_graph_arena_bytes = 0; + size_t decode_graph_arena_bytes = 0; + std::unique_ptr prefill_graph; + std::unique_ptr decode_graph; +}; + +GreedyQwenDecoderRuntime::GreedyQwenDecoderRuntime( + std::shared_ptr weights_source, + const GreedyQwenDecoderSpec & spec, + core::ExecutionContext & execution, + size_t prefill_graph_arena_bytes, + size_t decode_graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type) + : impl_(std::make_unique( + std::move(weights_source), + spec, + execution, + prefill_graph_arena_bytes, + decode_graph_arena_bytes, + weight_context_bytes, + weight_storage_type)) {} + +GreedyQwenDecoderRuntime::~GreedyQwenDecoderRuntime() = default; + +std::vector GreedyQwenDecoderRuntime::generate(const Prompt & prompt, int64_t max_new_tokens) { + const auto & spec = impl_->weights->spec(); + if (prompt.input_ids.empty()) { + throw std::runtime_error("greedy Qwen decoder prompt is empty"); + } + if (max_new_tokens <= 0) { + throw std::runtime_error("greedy Qwen decoder max_new_tokens must be positive"); + } + const int64_t prompt_steps = static_cast(prompt.input_ids.size()); + if (prompt_steps + max_new_tokens > spec.max_position_embeddings) { + throw std::runtime_error("greedy Qwen decoder request exceeds max_position_embeddings"); + } + const auto & injection = prompt.injection; + if (injection.tokens < 0 || injection.tokens > prompt_steps || + static_cast(injection.positions.size()) != injection.tokens || + static_cast(injection.values.size()) != injection.tokens * spec.decoder.stack.hidden_size) { + throw std::runtime_error("greedy Qwen decoder injection shape does not match the prompt"); + } + if (impl_->prefill_graph == nullptr || !impl_->prefill_graph->matches(prompt_steps, injection.tokens)) { + impl_->prefill_graph.reset(); + impl_->prefill_graph = std::make_unique( + impl_->weights, + prompt_steps, + injection.tokens, + impl_->prefill_graph_arena_bytes); + } + auto prefill = impl_->prefill_graph->run( + prompt.input_ids, + injection.values, + injection.positions); + const int64_t required_cache_steps = prompt_steps + max_new_tokens; + if (impl_->decode_graph == nullptr || !impl_->decode_graph->can_run(required_cache_steps)) { + impl_->decode_graph.reset(); + impl_->decode_graph = std::make_unique( + impl_->weights, + required_cache_steps, + impl_->decode_graph_arena_bytes); + } + impl_->decode_graph->import_state(prefill.kv_state); + + std::vector out; + std::vector logits = std::move(prefill.logits); + for (int64_t step = 0; step < max_new_tokens; ++step) { + const int32_t token = argmax_index(logits); + if (is_eos(spec, token)) { + break; + } + out.push_back(token); + logits = impl_->decode_graph->run_step(token); + } + return out; +} + +} // namespace engine::runtime diff --git a/src/framework/tokenizers/qwen_bpe_bundle.cpp b/src/framework/tokenizers/qwen_bpe_bundle.cpp new file mode 100644 index 000000000..74f24316c --- /dev/null +++ b/src/framework/tokenizers/qwen_bpe_bundle.cpp @@ -0,0 +1,41 @@ +#include "engine/framework/tokenizers/qwen_bpe_bundle.h" + +#include +#include + +namespace engine::tokenizers { + +std::shared_ptr load_qwen_bpe_tokenizer( + const engine::assets::ResourceBundle & bundle) { + LlamaBpeTokenizerSpec spec; + spec.tokenizer_config_path = bundle.require_file("tokenizer_config"); + if (const auto * path = bundle.find_file("vocab")) { + spec.vocab_path = *path; + } + if (const auto * path = bundle.find_file("merges")) { + spec.merges_path = *path; + } + if (const auto * path = bundle.find_file("tokenizer_json")) { + spec.tokenizer_json_path = *path; + } + spec.pre_type = LlamaBpePreTokenizer::Qwen2; + return load_llama_bpe_tokenizer(spec); +} + +int64_t require_added_token_id( + const engine::assets::ResourceBundle & bundle, + std::string_view content) { + const auto tokenizer = bundle.parse_json("tokenizer_json"); + for (const auto & item : tokenizer.require("added_tokens").as_array()) { + const auto * token_content = item.find("content"); + const auto * token_id = item.find("id"); + if (token_content != nullptr && token_content->is_string() && + token_id != nullptr && token_id->is_number() && + token_content->as_string() == content) { + return token_id->as_i64(); + } + } + throw std::runtime_error("tokenizer.json is missing token: " + std::string(content)); +} + +} // namespace engine::tokenizers