From 8b6f1d22d6b4fb3429a3239d172728a281a0ba32 Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Fri, 31 Jul 2026 03:07:35 +0000 Subject: [PATCH 1/9] feat(qwen3.6): support dense and MoE models --- csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp | 15 ++-- csrc/models/qwen3_5/qwen3_5_for_causal_lm.hpp | 2 + csrc/models/qwen3_5/qwen3_5_model.cpp | 15 ++-- csrc/models/qwen3_5/qwen3_5_model.hpp | 24 +++--- .../qwen3_5_moe/qwen3_5_moe_decoder_layer.cpp | 71 +++++++++++++++++ .../qwen3_5_moe/qwen3_5_moe_decoder_layer.hpp | 40 ++++++++++ .../qwen3_5_moe/qwen3_5_moe_for_causal_lm.cpp | 62 +++++++++++++++ .../qwen3_5_moe/qwen3_5_moe_for_causal_lm.hpp | 27 +++++++ csrc/models/qwen3_5_moe/qwen3_5_moe_model.cpp | 21 +++++ csrc/models/qwen3_5_moe/qwen3_5_moe_model.hpp | 21 +++++ .../qwen3_5_moe_sparse_moe_block.cpp | 78 +++++++++++++++++++ .../qwen3_5_moe_sparse_moe_block.hpp | 34 ++++++++ examples/bench.py | 3 + python/infinilm/modeling_utils.py | 55 +++++++++++++ python/infinilm/processors/__init__.py | 13 +++- .../infinilm/processors/qwen3_5_processor.py | 1 + 16 files changed, 458 insertions(+), 24 deletions(-) create mode 100644 csrc/models/qwen3_5_moe/qwen3_5_moe_decoder_layer.cpp create mode 100644 csrc/models/qwen3_5_moe/qwen3_5_moe_decoder_layer.hpp create mode 100644 csrc/models/qwen3_5_moe/qwen3_5_moe_for_causal_lm.cpp create mode 100644 csrc/models/qwen3_5_moe/qwen3_5_moe_for_causal_lm.hpp create mode 100644 csrc/models/qwen3_5_moe/qwen3_5_moe_model.cpp create mode 100644 csrc/models/qwen3_5_moe/qwen3_5_moe_model.hpp create mode 100644 csrc/models/qwen3_5_moe/qwen3_5_moe_sparse_moe_block.cpp create mode 100644 csrc/models/qwen3_5_moe/qwen3_5_moe_sparse_moe_block.hpp diff --git a/csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp b/csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp index aece19fda..1151551fb 100644 --- a/csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp +++ b/csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp @@ -35,12 +35,7 @@ void Qwen35ForCausalLM::reset_cache(const cache::CacheConfig *cache_config) { model_->reset_cache(cache_config); } -std::shared_ptr create_qwen3_5_model_config(std::shared_ptr model_config) { - const std::string model_type = model_config->get("model_type"); - if ("qwen3_5" != model_type) { - throw std::runtime_error("infinilm::models::qwen3_5::create_qwen3_next_model_config: model_type is not qwen3_5"); - } - +std::shared_ptr prepare_qwen3_5_model_config(std::shared_ptr model_config) { nlohmann::json &config_json = model_config->get_config_json(); if (config_json.contains("text_config") && config_json["text_config"].is_object()) { const nlohmann::json &text_config_json = config_json["text_config"]; @@ -79,6 +74,14 @@ std::shared_ptr create_qwen3_5_model_config(std:: return model_config; } +std::shared_ptr create_qwen3_5_model_config(std::shared_ptr model_config) { + const std::string model_type = model_config->get("model_type"); + if ("qwen3_5" != model_type) { + throw std::runtime_error("infinilm::models::qwen3_5::create_qwen3_5_model_config: model_type is not qwen3_5"); + } + return prepare_qwen3_5_model_config(model_config); +} + } // namespace infinilm::models::qwen3_5 namespace { diff --git a/csrc/models/qwen3_5/qwen3_5_for_causal_lm.hpp b/csrc/models/qwen3_5/qwen3_5_for_causal_lm.hpp index a7c751e30..787a345fd 100644 --- a/csrc/models/qwen3_5/qwen3_5_for_causal_lm.hpp +++ b/csrc/models/qwen3_5/qwen3_5_for_causal_lm.hpp @@ -20,6 +20,8 @@ class Qwen35ForCausalLM : public InfinilmModel { INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, lm_head); }; +std::shared_ptr prepare_qwen3_5_model_config(std::shared_ptr model_config); + std::shared_ptr create_qwen3_5_model_config(std::shared_ptr model_config); } // namespace infinilm::models::qwen3_5 diff --git a/csrc/models/qwen3_5/qwen3_5_model.cpp b/csrc/models/qwen3_5/qwen3_5_model.cpp index 82c8b0d4d..3dd3e61d1 100644 --- a/csrc/models/qwen3_5/qwen3_5_model.cpp +++ b/csrc/models/qwen3_5/qwen3_5_model.cpp @@ -31,8 +31,8 @@ std::vector tensor_to_i32_vector(const infinicore::Tensor &tensor) { } // namespace -Qwen35Model::Qwen35Model(std::shared_ptr model_config, - const infinicore::Device &device) +Qwen35ModelBase::Qwen35ModelBase(std::shared_ptr model_config, + const infinicore::Device &device) : model_config_(model_config) { const auto &dtype{model_config->get_dtype()}; nlohmann::json &config_json = model_config->get_config_json(); @@ -40,11 +40,16 @@ Qwen35Model::Qwen35Model(std::shared_ptr model_co if (config_json.contains("vision_config") && !config_json["vision_config"].is_null()) { INFINICORE_NN_MODULE_INIT(visual, config_json["vision_config"], dtype, device); } +} + +Qwen35Model::Qwen35Model(std::shared_ptr model_config, + const infinicore::Device &device) + : Qwen35ModelBase(model_config, device) { INFINICORE_NN_MODULE_INIT(language_model, model_config, device); } -void Qwen35Model::replace_image_embeddings(infinicore::Tensor &inputs_embeds, - const InfinilmModel::Input &input) const { +void Qwen35ModelBase::replace_image_embeddings(infinicore::Tensor &inputs_embeds, + const InfinilmModel::Input &input) const { if (!input.pixel_values.has_value() || input.pixel_values->empty()) { return; } @@ -116,7 +121,7 @@ infinicore::Tensor Qwen35Model::forward(const InfinilmModel::Input &input) const return language_model_->forward(input); } -void Qwen35Model::reset_cache(const cache::CacheConfig *cache_config) { +void Qwen35ModelBase::reset_cache(const cache::CacheConfig *cache_config) { if (nullptr == cache_config) { return; } diff --git a/csrc/models/qwen3_5/qwen3_5_model.hpp b/csrc/models/qwen3_5/qwen3_5_model.hpp index bea1b78ee..c82763c22 100644 --- a/csrc/models/qwen3_5/qwen3_5_model.hpp +++ b/csrc/models/qwen3_5/qwen3_5_model.hpp @@ -10,24 +10,30 @@ namespace infinilm::models::qwen3_5 { using Qwen35LanguageModel = infinilm::layers::causal_lm_templates::TextModel; -class Qwen35Model : public infinicore::nn::Module { +class Qwen35ModelBase : public infinicore::nn::Module { public: - Qwen35Model(std::shared_ptr model_config, - const infinicore::Device &device); - - infinicore::Tensor forward(const InfinilmModel::Input &input) const; + Qwen35ModelBase(std::shared_ptr model_config, + const infinicore::Device &device); void reset_cache(const cache::CacheConfig *cache_config); -private: +protected: void replace_image_embeddings(infinicore::Tensor &inputs_embeds, const infinilm::InfinilmModel::Input &input) const; -protected: INFINICORE_NN_MODULE(Qwen35VisionModel, visual); - INFINICORE_NN_MODULE(Qwen35LanguageModel, language_model); - std::shared_ptr model_config_; }; +class Qwen35Model : public Qwen35ModelBase { +public: + Qwen35Model(std::shared_ptr model_config, + const infinicore::Device &device); + + infinicore::Tensor forward(const InfinilmModel::Input &input) const; + +protected: + INFINICORE_NN_MODULE(Qwen35LanguageModel, language_model); +}; + } // namespace infinilm::models::qwen3_5 diff --git a/csrc/models/qwen3_5_moe/qwen3_5_moe_decoder_layer.cpp b/csrc/models/qwen3_5_moe/qwen3_5_moe_decoder_layer.cpp new file mode 100644 index 000000000..aa513284e --- /dev/null +++ b/csrc/models/qwen3_5_moe/qwen3_5_moe_decoder_layer.cpp @@ -0,0 +1,71 @@ +#include "qwen3_5_moe_decoder_layer.hpp" + +#include "infinicore/ops.hpp" + +#include +#include +#include + +namespace infinilm::models::qwen3_5_moe { + +Qwen35MoeDecoderLayer::Qwen35MoeDecoderLayer( + std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device) + : layer_idx_(layer_idx) { + const auto &dtype = model_config->get_dtype(); + const size_t hidden_size = model_config->get("hidden_size"); + const double rms_norm_eps = model_config->get("rms_norm_eps"); + + INFINICORE_NN_MODULE_INIT(input_layernorm, hidden_size, rms_norm_eps, dtype, device); + INFINICORE_NN_MODULE_INIT(post_attention_layernorm, hidden_size, rms_norm_eps, dtype, device); + INFINICORE_NN_MODULE_INIT(mlp, model_config, layer_idx, device); + + const auto layer_types = model_config->get>("layer_types"); + layer_type_ = layer_types.at(layer_idx); + if ("linear_attention" == layer_type_) { + INFINICORE_NN_MODULE_INIT(linear_attn, model_config, layer_idx, device); + } else if ("full_attention" == layer_type_) { + INFINICORE_NN_MODULE_INIT(self_attn, model_config, layer_idx, device); + } else { + throw std::runtime_error( + "Qwen35MoeDecoderLayer: unsupported layer_type '" + layer_type_ + + "' for layer " + std::to_string(layer_idx)); + } +} + +std::tuple Qwen35MoeDecoderLayer::forward( + const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) { + input_layernorm_->forward_inplace(hidden_states, residual); + if ("linear_attention" == layer_type_) { + hidden_states = linear_attn_->forward(hidden_states); + } else { + hidden_states = self_attn_->forward(positions, hidden_states); + } + + post_attention_layernorm_->forward_inplace(hidden_states, residual); + hidden_states = mlp_->forward(hidden_states); + return std::make_tuple(hidden_states, residual); +} + +infinicore::Tensor Qwen35MoeDecoderLayer::forward( + const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states) { + auto residual = hidden_states; + hidden_states = input_layernorm_->forward(hidden_states); + if ("linear_attention" == layer_type_) { + hidden_states = linear_attn_->forward(hidden_states); + } else { + hidden_states = self_attn_->forward(positions, hidden_states); + } + hidden_states = infinicore::op::add(residual, hidden_states); + + residual = hidden_states; + hidden_states = post_attention_layernorm_->forward(hidden_states); + hidden_states = mlp_->forward(hidden_states); + return infinicore::op::add(residual, hidden_states); +} + +} // namespace infinilm::models::qwen3_5_moe diff --git a/csrc/models/qwen3_5_moe/qwen3_5_moe_decoder_layer.hpp b/csrc/models/qwen3_5_moe/qwen3_5_moe_decoder_layer.hpp new file mode 100644 index 000000000..27b5f932f --- /dev/null +++ b/csrc/models/qwen3_5_moe/qwen3_5_moe_decoder_layer.hpp @@ -0,0 +1,40 @@ +#pragma once + +#include "../qwen3_5/qwen3_5_attention.hpp" +#include "../qwen3_next/qwen3_next_gated_deltanet.hpp" +#include "qwen3_5_moe_sparse_moe_block.hpp" + +#include +#include + +namespace infinilm::models::qwen3_5_moe { + +class Qwen35MoeDecoderLayer : public infinicore::nn::Module { +public: + Qwen35MoeDecoderLayer(std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device); + + std::tuple forward( + const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual); + + infinicore::Tensor forward(const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states); + + size_t layer_idx() const { return layer_idx_; } + +protected: + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, input_layernorm); + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, post_attention_layernorm); + INFINICORE_NN_MODULE(qwen3_5::Qwen35Attention, self_attn); + INFINICORE_NN_MODULE(qwen3_next::Qwen3NextGatedDeltaNet, linear_attn); + INFINICORE_NN_MODULE(Qwen35MoeSparseMoeBlock, mlp); + +private: + size_t layer_idx_; + std::string layer_type_; +}; + +} // namespace infinilm::models::qwen3_5_moe diff --git a/csrc/models/qwen3_5_moe/qwen3_5_moe_for_causal_lm.cpp b/csrc/models/qwen3_5_moe/qwen3_5_moe_for_causal_lm.cpp new file mode 100644 index 000000000..7797c17ae --- /dev/null +++ b/csrc/models/qwen3_5_moe/qwen3_5_moe_for_causal_lm.cpp @@ -0,0 +1,62 @@ +#include "qwen3_5_moe_for_causal_lm.hpp" + +#include "../models_registry.hpp" +#include "../qwen3_5/qwen3_5_for_causal_lm.hpp" + +#include +#include + +namespace infinilm::models::qwen3_5_moe { + +Qwen35MoeForConditionalGeneration::Qwen35MoeForConditionalGeneration( + std::shared_ptr model_config, + const infinicore::Device &device) { + model_config_ = model_config; + const size_t hidden_size = model_config->get("hidden_size"); + const size_t vocab_size = model_config->get("vocab_size"); + const auto &dtype = model_config->get_dtype(); + + INFINICORE_NN_MODULE_INIT(model, model_config, device); + INFINICORE_NN_MODULE_INIT(lm_head, hidden_size, vocab_size, false, dtype, device); +} + +InfinilmModel::Output Qwen35MoeForConditionalGeneration::forward( + const InfinilmModel::Input &input) const { + auto hidden_states = model_->forward(input); + return {lm_head_->forward(hidden_states)}; +} + +void Qwen35MoeForConditionalGeneration::reset_cache( + const cache::CacheConfig *cache_config) { + if (cache_config == nullptr) { + cache_config_.reset(); + } else { + cache_config_ = cache_config->unique_copy(); + } + model_->reset_cache(cache_config); +} + +std::shared_ptr create_qwen3_5_moe_model_config( + std::shared_ptr model_config) { + const std::string model_type = model_config->get("model_type"); + if ("qwen3_5_moe" != model_type) { + throw std::runtime_error( + "create_qwen3_5_moe_model_config: model_type is not qwen3_5_moe"); + } + + model_config = qwen3_5::prepare_qwen3_5_model_config(model_config); + auto &config_json = model_config->get_config_json(); + if (!config_json.contains("norm_topk_prob")) { + config_json["norm_topk_prob"] = true; + } + return model_config; +} + +} // namespace infinilm::models::qwen3_5_moe + +namespace { +INFINILM_REGISTER_CAUSAL_LM_MODEL( + qwen3_5_moe, + infinilm::models::qwen3_5_moe::Qwen35MoeForConditionalGeneration, + infinilm::models::qwen3_5_moe::create_qwen3_5_moe_model_config); +} // namespace diff --git a/csrc/models/qwen3_5_moe/qwen3_5_moe_for_causal_lm.hpp b/csrc/models/qwen3_5_moe/qwen3_5_moe_for_causal_lm.hpp new file mode 100644 index 000000000..2ac0abe30 --- /dev/null +++ b/csrc/models/qwen3_5_moe/qwen3_5_moe_for_causal_lm.hpp @@ -0,0 +1,27 @@ +#pragma once + +#include "qwen3_5_moe_model.hpp" + +#include + +namespace infinilm::models::qwen3_5_moe { + +class Qwen35MoeForConditionalGeneration : public InfinilmModel { +public: + Qwen35MoeForConditionalGeneration( + std::shared_ptr model_config, + const infinicore::Device &device); + + Output forward(const Input &input) const override; + + void reset_cache(const cache::CacheConfig *cache_config) override; + +protected: + INFINICORE_NN_MODULE(Qwen35MoeModel, model); + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, lm_head); +}; + +std::shared_ptr create_qwen3_5_moe_model_config( + std::shared_ptr model_config); + +} // namespace infinilm::models::qwen3_5_moe diff --git a/csrc/models/qwen3_5_moe/qwen3_5_moe_model.cpp b/csrc/models/qwen3_5_moe/qwen3_5_moe_model.cpp new file mode 100644 index 000000000..ab54a1f18 --- /dev/null +++ b/csrc/models/qwen3_5_moe/qwen3_5_moe_model.cpp @@ -0,0 +1,21 @@ +#include "qwen3_5_moe_model.hpp" + +namespace infinilm::models::qwen3_5_moe { + +Qwen35MoeModel::Qwen35MoeModel( + std::shared_ptr model_config, + const infinicore::Device &device) + : Qwen35ModelBase(model_config, device) { + INFINICORE_NN_MODULE_INIT(language_model, model_config, device); +} + +infinicore::Tensor Qwen35MoeModel::forward(const InfinilmModel::Input &input) const { + if (input.pixel_values.has_value() && !input.pixel_values->empty()) { + auto inputs_embeds = language_model_->embed_tokens(input.input_ids.value()); + replace_image_embeddings(inputs_embeds, input); + return language_model_->forward_embeds(inputs_embeds, input.position_ids.value()); + } + return language_model_->forward(input); +} + +} // namespace infinilm::models::qwen3_5_moe diff --git a/csrc/models/qwen3_5_moe/qwen3_5_moe_model.hpp b/csrc/models/qwen3_5_moe/qwen3_5_moe_model.hpp new file mode 100644 index 000000000..7d944513f --- /dev/null +++ b/csrc/models/qwen3_5_moe/qwen3_5_moe_model.hpp @@ -0,0 +1,21 @@ +#pragma once + +#include "../qwen3_5/qwen3_5_model.hpp" +#include "qwen3_5_moe_decoder_layer.hpp" + +namespace infinilm::models::qwen3_5_moe { + +using Qwen35MoeLanguageModel = infinilm::layers::causal_lm_templates::TextModel; + +class Qwen35MoeModel : public qwen3_5::Qwen35ModelBase { +public: + Qwen35MoeModel(std::shared_ptr model_config, + const infinicore::Device &device); + + infinicore::Tensor forward(const InfinilmModel::Input &input) const; + +protected: + INFINICORE_NN_MODULE(Qwen35MoeLanguageModel, language_model); +}; + +} // namespace infinilm::models::qwen3_5_moe diff --git a/csrc/models/qwen3_5_moe/qwen3_5_moe_sparse_moe_block.cpp b/csrc/models/qwen3_5_moe/qwen3_5_moe_sparse_moe_block.cpp new file mode 100644 index 000000000..293f6578a --- /dev/null +++ b/csrc/models/qwen3_5_moe/qwen3_5_moe_sparse_moe_block.cpp @@ -0,0 +1,78 @@ +#include "qwen3_5_moe_sparse_moe_block.hpp" + +#include "infinicore/ops.hpp" +#include "infinicore/ops/broadcast_to.hpp" +#include "infinicore/ops/mul.hpp" + +namespace infinilm::models::qwen3_5_moe { + +Qwen35MoeSparseMoeBlock::Qwen35MoeSparseMoeBlock( + std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device) { + use_legacy_moe_ = model_config->get_or("use_legacy_moe", false); + if (use_legacy_moe_) { + legacy_gate_ = this->register_module( + "gate", model_config, device); + legacy_experts_ = this->register_module( + "experts", model_config, device); + } else { + gate_ = this->register_module("gate", model_config, device); + experts_ = this->register_module("experts", model_config, device); + fused_moe_ = this->register_module("fused_moe", model_config, device, layer_idx); + } + + auto shared_config_json = model_config->get_config_json(); + shared_config_json["intermediate_size"] = model_config->get("shared_expert_intermediate_size"); + auto shared_config = std::make_shared(shared_config_json); + INFINICORE_NN_MODULE_INIT(shared_expert, shared_config, device); + + const auto &dtype = model_config->get_dtype(); + const size_t hidden_size = model_config->get("hidden_size"); + INFINICORE_NN_MODULE_INIT(shared_expert_gate, hidden_size, 1, false, dtype, device); +} + +infinicore::Tensor Qwen35MoeSparseMoeBlock::forward( + const infinicore::Tensor &hidden_states) const { + ASSERT(hidden_states->ndim() == 3); + + const auto shape = hidden_states->shape(); + auto hidden_states_reshaped = hidden_states->view({shape[0] * shape[1], shape[2]}); + + infinicore::Tensor expert_output; + if (use_legacy_moe_) { + auto [routing_weights, selected_experts] = legacy_gate_->forward(hidden_states_reshaped); + expert_output = legacy_experts_->forward( + hidden_states_reshaped, + selected_experts, + routing_weights); + } else { + auto [routing_weights, selected_experts] = gate_->forward(hidden_states_reshaped); + infinilm::layers::moe::TopKOutput topk_output{ + routing_weights, + selected_experts, + infinicore::Tensor(), + }; + expert_output = fused_moe_->forward( + hidden_states_reshaped, + topk_output, + experts_->moe_weights()); + } + + auto shared_output = shared_expert_->forward(hidden_states); + auto shared_gate_input = hidden_states; + auto shared_gate = infinicore::op::sigmoid(shared_expert_gate_->forward(shared_gate_input)); + const auto &shared_shape = shared_output->shape(); + auto shared_gate_broadcast = infinicore::op::broadcast_to( + shared_gate, + { + static_cast(shared_shape[0]), + static_cast(shared_shape[1]), + static_cast(shared_shape[2]), + }); + shared_output = infinicore::op::mul(shared_output, shared_gate_broadcast); + + return infinicore::op::add(expert_output->view(shape), shared_output); +} + +} // namespace infinilm::models::qwen3_5_moe diff --git a/csrc/models/qwen3_5_moe/qwen3_5_moe_sparse_moe_block.hpp b/csrc/models/qwen3_5_moe/qwen3_5_moe_sparse_moe_block.hpp new file mode 100644 index 000000000..a6bf9e308 --- /dev/null +++ b/csrc/models/qwen3_5_moe/qwen3_5_moe_sparse_moe_block.hpp @@ -0,0 +1,34 @@ +#pragma once + +#include "../../layers/common_modules.hpp" +#include "../../layers/moe/experts/fused_moe_experts.hpp" +#include "../../layers/moe/fused_moe.hpp" +#include "../../layers/moe/router/topk_router.hpp" +#include "../qwen3_moe/qwen3_moe_experts.hpp" +#include "../qwen3_moe/qwen3_moe_topk_router.hpp" + +#include +#include + +namespace infinilm::models::qwen3_5_moe { + +class Qwen35MoeSparseMoeBlock final : public infinicore::nn::Module { +public: + Qwen35MoeSparseMoeBlock(std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const; + +protected: + std::shared_ptr legacy_gate_; + std::shared_ptr legacy_experts_; + std::shared_ptr gate_; + std::shared_ptr experts_; + std::shared_ptr fused_moe_; + INFINICORE_NN_MODULE(infinilm::layers::MLP, shared_expert); + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, shared_expert_gate); + bool use_legacy_moe_{false}; +}; + +} // namespace infinilm::models::qwen3_5_moe diff --git a/examples/bench.py b/examples/bench.py index bd424e836..6866c1c8d 100644 --- a/examples/bench.py +++ b/examples/bench.py @@ -197,6 +197,7 @@ def __init__( weight_load_mode="async", moe_ep_backend="disabled", moe_ep_size=1, + use_legacy_moe=False, ) -> None: model_path = os.path.expanduser(model_path) self.draft_model_path = draft_model_path @@ -240,6 +241,7 @@ def __init__( kv_cache_dtype=cfg.kv_cache_dtype, use_mla=use_mla, weight_load_mode=weight_load_mode, + use_legacy_moe=use_legacy_moe, ) # ---------------------------------------------------------------------------- # @@ -442,6 +444,7 @@ def run( weight_load_mode=cfg.weight_load_mode, moe_ep_backend=moe_ep_backend, moe_ep_size=ep, + use_legacy_moe=cfg.use_legacy_moe, ) # ---------------------------------------------------------------------------- # diff --git a/python/infinilm/modeling_utils.py b/python/infinilm/modeling_utils.py index 88802500d..7cfe54d52 100644 --- a/python/infinilm/modeling_utils.py +++ b/python/infinilm/modeling_utils.py @@ -960,6 +960,60 @@ def _remap_qwen3_next(state_dict, config): return state_dict +def _remap_qwen3_5_moe(state_dict, config): + """Adapt packed Qwen3.5-MoE experts to InfiniLM expert parameter names.""" + state_dict = _remap_qwen3_5(state_dict, config) + text_config = config.get("text_config", config) + expected_num_experts = text_config["num_experts"] + expected_intermediate_size = text_config["moe_intermediate_size"] + + remapped = {} + for key, tensor in state_dict.items(): + if key.endswith(".mlp.experts.gate_up_proj"): + if tensor.ndim != 3: + raise ValueError( + f"Expected packed gate_up_proj to be 3D, got {tensor.shape} for {key}" + ) + if tensor.shape[0] != expected_num_experts: + raise ValueError( + f"Expected {expected_num_experts} experts, got {tensor.shape[0]} for {key}" + ) + if tensor.shape[1] != expected_intermediate_size * 2: + raise ValueError( + f"Expected packed gate/up size {expected_intermediate_size * 2}, " + f"got {tensor.shape[1]} for {key}" + ) + + prefix = key[: -len("gate_up_proj")] + for expert_idx, expert_gate_up in enumerate(tensor.unbind(0)): + gate, up = expert_gate_up.chunk(2, dim=0) + expert_prefix = f"{prefix}{expert_idx}." + remapped[f"{expert_prefix}gate_proj.weight"] = gate + remapped[f"{expert_prefix}up_proj.weight"] = up + elif key.endswith(".mlp.experts.down_proj"): + if tensor.ndim != 3: + raise ValueError( + f"Expected packed down_proj to be 3D, got {tensor.shape} for {key}" + ) + if tensor.shape[0] != expected_num_experts: + raise ValueError( + f"Expected {expected_num_experts} experts, got {tensor.shape[0]} for {key}" + ) + if tensor.shape[2] != expected_intermediate_size: + raise ValueError( + f"Expected down projection input size {expected_intermediate_size}, " + f"got {tensor.shape[2]} for {key}" + ) + + prefix = key[: -len("down_proj")] + for expert_idx, expert_down in enumerate(tensor.unbind(0)): + remapped[f"{prefix}{expert_idx}.down_proj.weight"] = expert_down + else: + remapped[key] = tensor + + return remapped + + _WEIGHT_REMAPPER = { "glm4": _remap_glm4, "chatglm": _remap_chatglm, @@ -969,5 +1023,6 @@ def _remap_qwen3_next(state_dict, config): "videonsa": _remap_videonsa, "qwen3_5": _remap_qwen3_5, "ernie4_5_moe_vl": _remap_ernie4_5_moe_vl, + "qwen3_5_moe": _remap_qwen3_5_moe, "qwen3_next": _remap_qwen3_next, } diff --git a/python/infinilm/processors/__init__.py b/python/infinilm/processors/__init__.py index 3d5cad4dc..d4ab5b3fc 100644 --- a/python/infinilm/processors/__init__.py +++ b/python/infinilm/processors/__init__.py @@ -33,13 +33,18 @@ def from_pretrained(cls, model_dir_path: str, **kwargs) -> InfinilmProcessor: registered Processor. Falls back to the registered default processor for unregistered or standard architectures. """ - config = AutoConfig.from_pretrained(model_dir_path, trust_remote_code=True) - model_type = config.model_type.lower() raw_config_path = Path(model_dir_path) / "config.json" - architectures = [] + raw_config = {} if raw_config_path.exists(): with raw_config_path.open("r") as f: - architectures = json.load(f).get("architectures", []) or [] + raw_config = json.load(f) + + model_type = str(raw_config.get("model_type", "")).lower() + if not model_type: + config = AutoConfig.from_pretrained(model_dir_path, trust_remote_code=True) + model_type = config.model_type.lower() + + architectures = raw_config.get("architectures", []) or [] if ( model_type == "qwen2_5_vl" and "VideoNSAForConditionalGeneration" in architectures diff --git a/python/infinilm/processors/qwen3_5_processor.py b/python/infinilm/processors/qwen3_5_processor.py index e550de5fd..6b8fee906 100644 --- a/python/infinilm/processors/qwen3_5_processor.py +++ b/python/infinilm/processors/qwen3_5_processor.py @@ -10,6 +10,7 @@ from .processor import register_processor +@register_processor("qwen3_5_moe") @register_processor("qwen3_5") class Qwen35Processor(BasicLLMProcessor): def __init__(self, model_dir_path: str): From 09664dc659a1ab605fa3e8e6a8c4344f4557d2d3 Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Fri, 31 Jul 2026 03:07:35 +0000 Subject: [PATCH 2/9] fix(graph): support Qwen3.5 hybrid position metadata --- csrc/engine/compiler/paged_compiler.cpp | 41 +++++++++++++++++++++++-- python/infinilm/infer_engine.py | 8 +++-- 2 files changed, 44 insertions(+), 5 deletions(-) diff --git a/csrc/engine/compiler/paged_compiler.cpp b/csrc/engine/compiler/paged_compiler.cpp index df3fd1cb4..7a5c57fd0 100644 --- a/csrc/engine/compiler/paged_compiler.cpp +++ b/csrc/engine/compiler/paged_compiler.cpp @@ -45,8 +45,36 @@ void PagedCompiler::compile() { size_t nblocks = dynamic_cast(model_->get_cache_config())->num_blocks(); auto &forward_context = infinilm::global_state::get_forward_context(); const bool has_mamba_state = has_mamba_cache(forward_context); - + auto compile_batch_sizes = decode_batch_sizes_; size_t max_batch_size = *std::max_element(decode_batch_sizes_.begin(), decode_batch_sizes_.end()); + if (has_mamba_state) { + size_t max_mamba_batch_size = max_batch_size; + auto clamp_to_state_pool = [&max_mamba_batch_size](const std::vector &state_vec) { + for (const auto &state : state_vec) { + if (!state) { + continue; + } + const size_t pool_size = state->size(0); + const size_t usable_rows = pool_size > 0 ? pool_size - 1 : 0; + max_mamba_batch_size = std::min(max_mamba_batch_size, usable_rows); + } + }; + clamp_to_state_pool(forward_context.conv_state_vec); + clamp_to_state_pool(forward_context.ssm_state_vec); + compile_batch_sizes.erase( + std::remove_if( + compile_batch_sizes.begin(), + compile_batch_sizes.end(), + [max_mamba_batch_size](size_t b) { + return b > max_mamba_batch_size; + }), + compile_batch_sizes.end()); + if (compile_batch_sizes.empty()) { + return; + } + max_batch_size = *std::max_element( + compile_batch_sizes.begin(), compile_batch_sizes.end()); + } compiled_map_decode_.clear(); block_tables_holder_ = infinicore::Tensor::empty( {nblocks * max_batch_size}, infinicore::DataType::I32, infinicore::context::getDevice()); @@ -55,7 +83,14 @@ void PagedCompiler::compile() { auto make_decode_input = [&](size_t b) { InfinilmModel::Input input; input.input_ids = infinicore::Tensor::empty({1, b}, infinicore::DataType::I64, infinicore::context::getDevice()); - input.position_ids = infinicore::Tensor::empty({b}, infinicore::DataType::I64, infinicore::context::getDevice()); + // Qwen3.5 hybrid models use three MRoPE axes even for text-only + // decode. Their scheduler supplies position_ids with shape [3, b]. + // Other paged models continue to use the traditional [b] layout. + input.position_ids = infinicore::Tensor::empty( + has_mamba_state ? std::vector{3, b} + : std::vector{b}, + infinicore::DataType::I64, + infinicore::context::getDevice()); input.total_sequence_lengths = infinicore::Tensor::empty({b}, infinicore::DataType::I32, infinicore::context::getDevice()); set_zeros(input.input_ids.value()); set_zeros(input.position_ids.value()); @@ -126,7 +161,7 @@ void PagedCompiler::compile() { infinicore::context::syncStream(); } - for (size_t b : decode_batch_sizes_) { + for (size_t b : compile_batch_sizes) { auto input = make_decode_input(b); barrier_->wait(); diff --git a/python/infinilm/infer_engine.py b/python/infinilm/infer_engine.py index 270af75ec..6b1499fb5 100644 --- a/python/infinilm/infer_engine.py +++ b/python/infinilm/infer_engine.py @@ -516,9 +516,13 @@ def generate( if self.enable_paged_attn: input_ids = input_ids.view([1, batch_size * seq_len]) + position_ids_list = ( + list(range(past_seq_len, past_seq_len + seq_len)) * batch_size + ) + if self.model_type in ("qwen3_5", "qwen3_5_moe"): + position_ids_list = [position_ids_list] * 3 position_ids = infinicore.from_list( - list(range(past_seq_len, past_seq_len + seq_len)) * batch_size, - dtype=infinicore.int64, + position_ids_list, dtype=infinicore.int64 ) if iter == 0: From ebcf0d189642a84348d515dda5f7b152b35ff868 Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Mon, 3 Aug 2026 03:53:12 +0000 Subject: [PATCH 3/9] refactor: share hybrid decoder and MoE blocks --- .../hybrid_decoder_layer.hpp | 117 ++++++++++++++++++ .../moe/shared_expert_sparse_moe_block.cpp | 70 +++++++++++ .../moe/shared_expert_sparse_moe_block.hpp | 37 ++++++ csrc/models/qwen3_5/qwen3_5_decoderLayer.cpp | 66 ---------- csrc/models/qwen3_5/qwen3_5_decoderLayer.hpp | 34 +---- .../qwen3_5_moe/qwen3_5_moe_decoder_layer.cpp | 71 ----------- .../qwen3_5_moe/qwen3_5_moe_decoder_layer.hpp | 35 +----- .../qwen3_5_moe_sparse_moe_block.cpp | 78 ------------ .../qwen3_5_moe_sparse_moe_block.hpp | 29 +---- .../qwen3_next/qwen3_next_decoderLayer.cpp | 66 ---------- .../qwen3_next/qwen3_next_decoderLayer.hpp | 33 +---- .../qwen3_next_sparse_moe_block.cpp | 106 ---------------- .../qwen3_next_sparse_moe_block.hpp | 38 +----- 13 files changed, 244 insertions(+), 536 deletions(-) create mode 100644 csrc/layers/causal_lm_templates/hybrid_decoder_layer.hpp create mode 100644 csrc/layers/moe/shared_expert_sparse_moe_block.cpp create mode 100644 csrc/layers/moe/shared_expert_sparse_moe_block.hpp delete mode 100644 csrc/models/qwen3_5/qwen3_5_decoderLayer.cpp delete mode 100644 csrc/models/qwen3_5_moe/qwen3_5_moe_decoder_layer.cpp delete mode 100644 csrc/models/qwen3_5_moe/qwen3_5_moe_sparse_moe_block.cpp delete mode 100644 csrc/models/qwen3_next/qwen3_next_decoderLayer.cpp delete mode 100644 csrc/models/qwen3_next/qwen3_next_sparse_moe_block.cpp diff --git a/csrc/layers/causal_lm_templates/hybrid_decoder_layer.hpp b/csrc/layers/causal_lm_templates/hybrid_decoder_layer.hpp new file mode 100644 index 000000000..345461ee5 --- /dev/null +++ b/csrc/layers/causal_lm_templates/hybrid_decoder_layer.hpp @@ -0,0 +1,117 @@ +#pragma once + +#include "../../config/model_config.hpp" +#include "infinicore/device.hpp" +#include "infinicore/nn/module.hpp" +#include "infinicore/nn/rmsnorm.hpp" +#include "infinicore/ops.hpp" +#include "infinicore/tensor.hpp" + +#include +#include +#include +#include +#include +#include + +namespace infinilm::layers::causal_lm_templates { + +template +class HybridDecoderLayer : public infinicore::nn::Module { +public: + HybridDecoderLayer( + std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device) + : layer_idx_(layer_idx) { + const auto &dtype = model_config->get_dtype(); + const size_t hidden_size = model_config->get("hidden_size"); + const double rms_norm_eps = model_config->get("rms_norm_eps"); + + input_layernorm_ = this->register_module( + "input_layernorm", hidden_size, rms_norm_eps, dtype, device); + post_attention_layernorm_ = this->register_module( + "post_attention_layernorm", hidden_size, rms_norm_eps, dtype, device); + mlp_ = register_mlp(model_config, layer_idx, device); + + const auto layer_types = model_config->get>("layer_types"); + const std::string &layer_type = layer_types.at(layer_idx); + if (layer_type == "linear_attention") { + is_linear_attention_ = true; + linear_attn_ = this->register_module( + "linear_attn", model_config, layer_idx, device); + } else if (layer_type == "full_attention") { + self_attn_ = this->register_module( + "self_attn", model_config, layer_idx, device); + } else { + throw std::runtime_error( + "HybridDecoderLayer: unsupported layer_type '" + layer_type + + "' for layer " + std::to_string(layer_idx)); + } + } + + std::tuple forward( + const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) { + input_layernorm_->forward_inplace(hidden_states, residual); + hidden_states = forward_mixer(positions, hidden_states); + post_attention_layernorm_->forward_inplace(hidden_states, residual); + hidden_states = mlp_->forward(hidden_states); + return std::make_tuple(hidden_states, residual); + } + + infinicore::Tensor forward( + const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states) { + auto residual = hidden_states; + hidden_states = input_layernorm_->forward(hidden_states); + hidden_states = forward_mixer(positions, hidden_states); + hidden_states = infinicore::op::add(residual, hidden_states); + + residual = hidden_states; + hidden_states = post_attention_layernorm_->forward(hidden_states); + hidden_states = mlp_->forward(hidden_states); + return infinicore::op::add(residual, hidden_states); + } + + size_t layer_idx() const { return layer_idx_; } + +protected: + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, input_layernorm); + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, post_attention_layernorm); + INFINICORE_NN_MODULE(Attention, self_attn); + INFINICORE_NN_MODULE(LinearAttention, linear_attn); + INFINICORE_NN_MODULE(MLP, mlp); + +private: + infinicore::Tensor forward_mixer( + const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states) const { + if (is_linear_attention_) { + return linear_attn_->forward(hidden_states); + } + return self_attn_->forward(positions, hidden_states); + } + + std::shared_ptr register_mlp( + std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device) { + if constexpr (std::is_constructible_v< + MLP, + std::shared_ptr, + size_t, + const infinicore::Device &>) { + return this->register_module( + "mlp", model_config, layer_idx, device); + } else { + return this->register_module("mlp", model_config, device); + } + } + + size_t layer_idx_; + bool is_linear_attention_{false}; +}; + +} // namespace infinilm::layers::causal_lm_templates diff --git a/csrc/layers/moe/shared_expert_sparse_moe_block.cpp b/csrc/layers/moe/shared_expert_sparse_moe_block.cpp new file mode 100644 index 000000000..dc80e9e66 --- /dev/null +++ b/csrc/layers/moe/shared_expert_sparse_moe_block.cpp @@ -0,0 +1,70 @@ +#include "shared_expert_sparse_moe_block.hpp" + +#include "infinicore/ops.hpp" +#include "infinicore/ops/mul.hpp" + +#include + +namespace infinilm::layers::moe { + +SharedExpertSparseMoeBlock::SharedExpertSparseMoeBlock( + std::shared_ptr model_config, + const infinicore::Device &device) + : SharedExpertSparseMoeBlock(model_config, 0, device) { +} + +SharedExpertSparseMoeBlock::SharedExpertSparseMoeBlock( + std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device) { + INFINICORE_NN_MODULE_INIT(gate, model_config, device); + INFINICORE_NN_MODULE_INIT(experts, model_config, device); + INFINICORE_NN_MODULE_INIT(fused_moe, model_config, device, layer_idx); + + auto shared_config_json = model_config->get_config_json(); + shared_config_json["intermediate_size"] = model_config->get("shared_expert_intermediate_size"); + auto shared_config = std::make_shared( + std::move(shared_config_json)); + INFINICORE_NN_MODULE_INIT(shared_expert, shared_config, device); + + INFINICORE_NN_MODULE_INIT( + shared_expert_gate, + model_config->get("hidden_size"), + 1, + false, + model_config->get_dtype(), + device); +} + +infinicore::Tensor SharedExpertSparseMoeBlock::forward( + const infinicore::Tensor &hidden_states) const { + ASSERT(hidden_states->ndim() == 3); + + const auto shape = hidden_states->shape(); + auto hidden_states_reshaped = hidden_states->view( + {shape[0] * shape[1], shape[2]}); + + auto [routing_weights, selected_experts] = gate_->forward(hidden_states_reshaped); + TopKOutput topk_output{ + routing_weights, + selected_experts, + infinicore::Tensor(), + }; + auto routed_output = fused_moe_->forward( + hidden_states_reshaped, + topk_output, + experts_->moe_weights()); + + auto shared_output = shared_expert_->forward(hidden_states); + auto shared_gate_input = hidden_states; + auto shared_gate = infinicore::op::sigmoid( + shared_expert_gate_->forward(shared_gate_input)); + shared_gate = shared_gate->as_strided( + shared_output->shape(), + {shared_gate->stride(0), shared_gate->stride(1), 0}); + shared_output = infinicore::op::mul(shared_output, shared_gate); + + return infinicore::op::add(routed_output->view(shape), shared_output); +} + +} // namespace infinilm::layers::moe diff --git a/csrc/layers/moe/shared_expert_sparse_moe_block.hpp b/csrc/layers/moe/shared_expert_sparse_moe_block.hpp new file mode 100644 index 000000000..249115164 --- /dev/null +++ b/csrc/layers/moe/shared_expert_sparse_moe_block.hpp @@ -0,0 +1,37 @@ +#pragma once + +#include "experts/fused_moe_experts.hpp" +#include "fused_moe.hpp" +#include "router/topk_router.hpp" + +#include "../../config/model_config.hpp" +#include "../linear/linear.hpp" +#include "../mlp/mlp.hpp" +#include "infinicore/nn/module.hpp" + +#include +#include + +namespace infinilm::layers::moe { + +class SharedExpertSparseMoeBlock : public infinicore::nn::Module { +public: + SharedExpertSparseMoeBlock( + std::shared_ptr model_config, + const infinicore::Device &device); + SharedExpertSparseMoeBlock( + std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const; + +protected: + INFINICORE_NN_MODULE(TopKRouter, gate); + INFINICORE_NN_MODULE(FusedMoeExperts, experts); + INFINICORE_NN_MODULE(FusedMoE, fused_moe); + INFINICORE_NN_MODULE(infinilm::layers::mlp::MLP, shared_expert); + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, shared_expert_gate); +}; + +} // namespace infinilm::layers::moe diff --git a/csrc/models/qwen3_5/qwen3_5_decoderLayer.cpp b/csrc/models/qwen3_5/qwen3_5_decoderLayer.cpp deleted file mode 100644 index 70964bb69..000000000 --- a/csrc/models/qwen3_5/qwen3_5_decoderLayer.cpp +++ /dev/null @@ -1,66 +0,0 @@ -#include "qwen3_5_decoderLayer.hpp" -#include "infinicore/ops.hpp" -#include -#include -#include - -namespace infinilm::models::qwen3_5 { - -Qwen35DecoderLayer::Qwen35DecoderLayer(std::shared_ptr model_config, - size_t layer_idx, - const infinicore::Device &device) - : layer_idx_(layer_idx) { - - const auto &dtype{model_config->get_dtype()}; - size_t hidden_size = model_config->get("hidden_size"); - double rms_norm_eps = model_config->get("rms_norm_eps"); - - INFINICORE_NN_MODULE_INIT(input_layernorm, hidden_size, rms_norm_eps, dtype, device); - INFINICORE_NN_MODULE_INIT(post_attention_layernorm, hidden_size, rms_norm_eps, dtype, device); - INFINICORE_NN_MODULE_INIT(mlp, model_config, device); - - const std::vector layer_types = model_config->get>("layer_types"); - layer_type_ = layer_types[layer_idx]; - if ("linear_attention" == layer_type_) { - INFINICORE_NN_MODULE_INIT(linear_attn, model_config, layer_idx, device); - } else if ("full_attention" == layer_type_) { - INFINICORE_NN_MODULE_INIT(self_attn, model_config, layer_idx, device); - } else { - throw std::runtime_error("infinilm::models::qwen3_5::Qwen35DecoderLayer: unsupported layer_type '" + layer_type_ + "' for layer " + std::to_string(layer_idx)); - } -} - -std::tuple Qwen35DecoderLayer::forward(const infinicore::Tensor &positions, - infinicore::Tensor &hidden_states, - infinicore::Tensor &residual) { - input_layernorm_->forward_inplace(hidden_states, residual); - if ("linear_attention" == layer_type_) { - hidden_states = linear_attn_->forward(hidden_states); - } else if ("full_attention" == layer_type_) { - hidden_states = self_attn_->forward(positions, hidden_states); - } - - post_attention_layernorm_->forward_inplace(hidden_states, residual); - hidden_states = mlp_->forward(hidden_states); - return std::make_tuple(hidden_states, residual); -} - -infinicore::Tensor Qwen35DecoderLayer::forward(const infinicore::Tensor &positions, - infinicore::Tensor &hidden_states) { - auto residual = hidden_states; - hidden_states = input_layernorm_->forward(hidden_states); - if ("linear_attention" == layer_type_) { - hidden_states = linear_attn_->forward(hidden_states); - } else if ("full_attention" == layer_type_) { - hidden_states = self_attn_->forward(positions, hidden_states); - } - hidden_states = infinicore::op::add(residual, hidden_states); - - residual = hidden_states; - hidden_states = post_attention_layernorm_->forward(hidden_states); - hidden_states = mlp_->forward(hidden_states); - hidden_states = infinicore::op::add(residual, hidden_states); - return hidden_states; -} - -} // namespace infinilm::models::qwen3_5 diff --git a/csrc/models/qwen3_5/qwen3_5_decoderLayer.hpp b/csrc/models/qwen3_5/qwen3_5_decoderLayer.hpp index 751586bbe..d1223cb03 100644 --- a/csrc/models/qwen3_5/qwen3_5_decoderLayer.hpp +++ b/csrc/models/qwen3_5/qwen3_5_decoderLayer.hpp @@ -1,37 +1,15 @@ #pragma once +#include "../../layers/causal_lm_templates/hybrid_decoder_layer.hpp" +#include "../../layers/common_modules.hpp" #include "../qwen3_next/qwen3_next_gated_deltanet.hpp" #include "qwen3_5_attention.hpp" -#include -#include namespace infinilm::models::qwen3_5 { -class Qwen35DecoderLayer : public infinicore::nn::Module { -public: - Qwen35DecoderLayer(std::shared_ptr model_config, - size_t layer_idx, - const infinicore::Device &device); - - std::tuple forward(const infinicore::Tensor &positions, - infinicore::Tensor &hidden_states, - infinicore::Tensor &residual); - - infinicore::Tensor forward(const infinicore::Tensor &positions, - infinicore::Tensor &hidden_states); - - size_t layer_idx() const { return layer_idx_; } - -protected: - INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, input_layernorm); - INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, post_attention_layernorm); - INFINICORE_NN_MODULE(Qwen35Attention, self_attn); - INFINICORE_NN_MODULE(qwen3_next::Qwen3NextGatedDeltaNet, linear_attn); - INFINICORE_NN_MODULE(infinilm::layers::MLP, mlp); - -private: - size_t layer_idx_; - std::string layer_type_; -}; +using Qwen35DecoderLayer = infinilm::layers::causal_lm_templates::HybridDecoderLayer< + Qwen35Attention, + qwen3_next::Qwen3NextGatedDeltaNet, + infinilm::layers::MLP>; } // namespace infinilm::models::qwen3_5 diff --git a/csrc/models/qwen3_5_moe/qwen3_5_moe_decoder_layer.cpp b/csrc/models/qwen3_5_moe/qwen3_5_moe_decoder_layer.cpp deleted file mode 100644 index aa513284e..000000000 --- a/csrc/models/qwen3_5_moe/qwen3_5_moe_decoder_layer.cpp +++ /dev/null @@ -1,71 +0,0 @@ -#include "qwen3_5_moe_decoder_layer.hpp" - -#include "infinicore/ops.hpp" - -#include -#include -#include - -namespace infinilm::models::qwen3_5_moe { - -Qwen35MoeDecoderLayer::Qwen35MoeDecoderLayer( - std::shared_ptr model_config, - size_t layer_idx, - const infinicore::Device &device) - : layer_idx_(layer_idx) { - const auto &dtype = model_config->get_dtype(); - const size_t hidden_size = model_config->get("hidden_size"); - const double rms_norm_eps = model_config->get("rms_norm_eps"); - - INFINICORE_NN_MODULE_INIT(input_layernorm, hidden_size, rms_norm_eps, dtype, device); - INFINICORE_NN_MODULE_INIT(post_attention_layernorm, hidden_size, rms_norm_eps, dtype, device); - INFINICORE_NN_MODULE_INIT(mlp, model_config, layer_idx, device); - - const auto layer_types = model_config->get>("layer_types"); - layer_type_ = layer_types.at(layer_idx); - if ("linear_attention" == layer_type_) { - INFINICORE_NN_MODULE_INIT(linear_attn, model_config, layer_idx, device); - } else if ("full_attention" == layer_type_) { - INFINICORE_NN_MODULE_INIT(self_attn, model_config, layer_idx, device); - } else { - throw std::runtime_error( - "Qwen35MoeDecoderLayer: unsupported layer_type '" + layer_type_ - + "' for layer " + std::to_string(layer_idx)); - } -} - -std::tuple Qwen35MoeDecoderLayer::forward( - const infinicore::Tensor &positions, - infinicore::Tensor &hidden_states, - infinicore::Tensor &residual) { - input_layernorm_->forward_inplace(hidden_states, residual); - if ("linear_attention" == layer_type_) { - hidden_states = linear_attn_->forward(hidden_states); - } else { - hidden_states = self_attn_->forward(positions, hidden_states); - } - - post_attention_layernorm_->forward_inplace(hidden_states, residual); - hidden_states = mlp_->forward(hidden_states); - return std::make_tuple(hidden_states, residual); -} - -infinicore::Tensor Qwen35MoeDecoderLayer::forward( - const infinicore::Tensor &positions, - infinicore::Tensor &hidden_states) { - auto residual = hidden_states; - hidden_states = input_layernorm_->forward(hidden_states); - if ("linear_attention" == layer_type_) { - hidden_states = linear_attn_->forward(hidden_states); - } else { - hidden_states = self_attn_->forward(positions, hidden_states); - } - hidden_states = infinicore::op::add(residual, hidden_states); - - residual = hidden_states; - hidden_states = post_attention_layernorm_->forward(hidden_states); - hidden_states = mlp_->forward(hidden_states); - return infinicore::op::add(residual, hidden_states); -} - -} // namespace infinilm::models::qwen3_5_moe diff --git a/csrc/models/qwen3_5_moe/qwen3_5_moe_decoder_layer.hpp b/csrc/models/qwen3_5_moe/qwen3_5_moe_decoder_layer.hpp index 27b5f932f..eee84a0b5 100644 --- a/csrc/models/qwen3_5_moe/qwen3_5_moe_decoder_layer.hpp +++ b/csrc/models/qwen3_5_moe/qwen3_5_moe_decoder_layer.hpp @@ -1,40 +1,15 @@ #pragma once +#include "../../layers/causal_lm_templates/hybrid_decoder_layer.hpp" #include "../qwen3_5/qwen3_5_attention.hpp" #include "../qwen3_next/qwen3_next_gated_deltanet.hpp" #include "qwen3_5_moe_sparse_moe_block.hpp" -#include -#include - namespace infinilm::models::qwen3_5_moe { -class Qwen35MoeDecoderLayer : public infinicore::nn::Module { -public: - Qwen35MoeDecoderLayer(std::shared_ptr model_config, - size_t layer_idx, - const infinicore::Device &device); - - std::tuple forward( - const infinicore::Tensor &positions, - infinicore::Tensor &hidden_states, - infinicore::Tensor &residual); - - infinicore::Tensor forward(const infinicore::Tensor &positions, - infinicore::Tensor &hidden_states); - - size_t layer_idx() const { return layer_idx_; } - -protected: - INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, input_layernorm); - INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, post_attention_layernorm); - INFINICORE_NN_MODULE(qwen3_5::Qwen35Attention, self_attn); - INFINICORE_NN_MODULE(qwen3_next::Qwen3NextGatedDeltaNet, linear_attn); - INFINICORE_NN_MODULE(Qwen35MoeSparseMoeBlock, mlp); - -private: - size_t layer_idx_; - std::string layer_type_; -}; +using Qwen35MoeDecoderLayer = infinilm::layers::causal_lm_templates::HybridDecoderLayer< + qwen3_5::Qwen35Attention, + qwen3_next::Qwen3NextGatedDeltaNet, + Qwen35MoeSparseMoeBlock>; } // namespace infinilm::models::qwen3_5_moe diff --git a/csrc/models/qwen3_5_moe/qwen3_5_moe_sparse_moe_block.cpp b/csrc/models/qwen3_5_moe/qwen3_5_moe_sparse_moe_block.cpp deleted file mode 100644 index 293f6578a..000000000 --- a/csrc/models/qwen3_5_moe/qwen3_5_moe_sparse_moe_block.cpp +++ /dev/null @@ -1,78 +0,0 @@ -#include "qwen3_5_moe_sparse_moe_block.hpp" - -#include "infinicore/ops.hpp" -#include "infinicore/ops/broadcast_to.hpp" -#include "infinicore/ops/mul.hpp" - -namespace infinilm::models::qwen3_5_moe { - -Qwen35MoeSparseMoeBlock::Qwen35MoeSparseMoeBlock( - std::shared_ptr model_config, - size_t layer_idx, - const infinicore::Device &device) { - use_legacy_moe_ = model_config->get_or("use_legacy_moe", false); - if (use_legacy_moe_) { - legacy_gate_ = this->register_module( - "gate", model_config, device); - legacy_experts_ = this->register_module( - "experts", model_config, device); - } else { - gate_ = this->register_module("gate", model_config, device); - experts_ = this->register_module("experts", model_config, device); - fused_moe_ = this->register_module("fused_moe", model_config, device, layer_idx); - } - - auto shared_config_json = model_config->get_config_json(); - shared_config_json["intermediate_size"] = model_config->get("shared_expert_intermediate_size"); - auto shared_config = std::make_shared(shared_config_json); - INFINICORE_NN_MODULE_INIT(shared_expert, shared_config, device); - - const auto &dtype = model_config->get_dtype(); - const size_t hidden_size = model_config->get("hidden_size"); - INFINICORE_NN_MODULE_INIT(shared_expert_gate, hidden_size, 1, false, dtype, device); -} - -infinicore::Tensor Qwen35MoeSparseMoeBlock::forward( - const infinicore::Tensor &hidden_states) const { - ASSERT(hidden_states->ndim() == 3); - - const auto shape = hidden_states->shape(); - auto hidden_states_reshaped = hidden_states->view({shape[0] * shape[1], shape[2]}); - - infinicore::Tensor expert_output; - if (use_legacy_moe_) { - auto [routing_weights, selected_experts] = legacy_gate_->forward(hidden_states_reshaped); - expert_output = legacy_experts_->forward( - hidden_states_reshaped, - selected_experts, - routing_weights); - } else { - auto [routing_weights, selected_experts] = gate_->forward(hidden_states_reshaped); - infinilm::layers::moe::TopKOutput topk_output{ - routing_weights, - selected_experts, - infinicore::Tensor(), - }; - expert_output = fused_moe_->forward( - hidden_states_reshaped, - topk_output, - experts_->moe_weights()); - } - - auto shared_output = shared_expert_->forward(hidden_states); - auto shared_gate_input = hidden_states; - auto shared_gate = infinicore::op::sigmoid(shared_expert_gate_->forward(shared_gate_input)); - const auto &shared_shape = shared_output->shape(); - auto shared_gate_broadcast = infinicore::op::broadcast_to( - shared_gate, - { - static_cast(shared_shape[0]), - static_cast(shared_shape[1]), - static_cast(shared_shape[2]), - }); - shared_output = infinicore::op::mul(shared_output, shared_gate_broadcast); - - return infinicore::op::add(expert_output->view(shape), shared_output); -} - -} // namespace infinilm::models::qwen3_5_moe diff --git a/csrc/models/qwen3_5_moe/qwen3_5_moe_sparse_moe_block.hpp b/csrc/models/qwen3_5_moe/qwen3_5_moe_sparse_moe_block.hpp index a6bf9e308..42bad48b3 100644 --- a/csrc/models/qwen3_5_moe/qwen3_5_moe_sparse_moe_block.hpp +++ b/csrc/models/qwen3_5_moe/qwen3_5_moe_sparse_moe_block.hpp @@ -1,34 +1,9 @@ #pragma once -#include "../../layers/common_modules.hpp" -#include "../../layers/moe/experts/fused_moe_experts.hpp" -#include "../../layers/moe/fused_moe.hpp" -#include "../../layers/moe/router/topk_router.hpp" -#include "../qwen3_moe/qwen3_moe_experts.hpp" -#include "../qwen3_moe/qwen3_moe_topk_router.hpp" - -#include -#include +#include "../../layers/moe/shared_expert_sparse_moe_block.hpp" namespace infinilm::models::qwen3_5_moe { -class Qwen35MoeSparseMoeBlock final : public infinicore::nn::Module { -public: - Qwen35MoeSparseMoeBlock(std::shared_ptr model_config, - size_t layer_idx, - const infinicore::Device &device); - - infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const; - -protected: - std::shared_ptr legacy_gate_; - std::shared_ptr legacy_experts_; - std::shared_ptr gate_; - std::shared_ptr experts_; - std::shared_ptr fused_moe_; - INFINICORE_NN_MODULE(infinilm::layers::MLP, shared_expert); - INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, shared_expert_gate); - bool use_legacy_moe_{false}; -}; +using Qwen35MoeSparseMoeBlock = infinilm::layers::moe::SharedExpertSparseMoeBlock; } // namespace infinilm::models::qwen3_5_moe diff --git a/csrc/models/qwen3_next/qwen3_next_decoderLayer.cpp b/csrc/models/qwen3_next/qwen3_next_decoderLayer.cpp deleted file mode 100644 index 4c61832c8..000000000 --- a/csrc/models/qwen3_next/qwen3_next_decoderLayer.cpp +++ /dev/null @@ -1,66 +0,0 @@ -#include "qwen3_next_decoderLayer.hpp" -#include "infinicore/ops.hpp" -#include -#include -#include - -namespace infinilm::models::qwen3_next { - -Qwen3NextDecoderLayer::Qwen3NextDecoderLayer(std::shared_ptr model_config, - size_t layer_idx, - const infinicore::Device &device) - : layer_idx_(layer_idx) { - - const auto &dtype{model_config->get_dtype()}; - size_t hidden_size = model_config->get("hidden_size"); - double rms_norm_eps = model_config->get("rms_norm_eps"); - - INFINICORE_NN_MODULE_INIT(input_layernorm, hidden_size, rms_norm_eps, dtype, device); - INFINICORE_NN_MODULE_INIT(post_attention_layernorm, hidden_size, rms_norm_eps, dtype, device); - INFINICORE_NN_MODULE_INIT(mlp, model_config, device); - - const std::vector layer_types = model_config->get>("layer_types"); - layer_type_ = layer_types[layer_idx]; - if ("linear_attention" == layer_type_) { - INFINICORE_NN_MODULE_INIT(linear_attn, model_config, layer_idx, device); - } else if ("full_attention" == layer_type_) { - INFINICORE_NN_MODULE_INIT(self_attn, model_config, layer_idx, device); - } else { - throw std::runtime_error("infinilm::models::qwen3_next::Qwen3NextDecoderLayer: unsupported layer_type '" + layer_type_ + "' for layer " + std::to_string(layer_idx)); - } -} - -std::tuple Qwen3NextDecoderLayer::forward(const infinicore::Tensor &positions, - infinicore::Tensor &hidden_states, - infinicore::Tensor &residual) { - input_layernorm_->forward_inplace(hidden_states, residual); - if ("linear_attention" == layer_type_) { - hidden_states = linear_attn_->forward(hidden_states); - } else if ("full_attention" == layer_type_) { - hidden_states = self_attn_->forward(positions, hidden_states); - } - - post_attention_layernorm_->forward_inplace(hidden_states, residual); - hidden_states = mlp_->forward(hidden_states); - return std::make_tuple(hidden_states, residual); -} - -infinicore::Tensor Qwen3NextDecoderLayer::forward(const infinicore::Tensor &positions, - infinicore::Tensor &hidden_states) { - auto residual = hidden_states; - hidden_states = input_layernorm_->forward(hidden_states); - if ("linear_attention" == layer_type_) { - hidden_states = linear_attn_->forward(hidden_states); - } else if ("full_attention" == layer_type_) { - hidden_states = self_attn_->forward(positions, hidden_states); - } - hidden_states = infinicore::op::add(residual, hidden_states); - - residual = hidden_states; - hidden_states = post_attention_layernorm_->forward(hidden_states); - hidden_states = mlp_->forward(hidden_states); - hidden_states = infinicore::op::add(residual, hidden_states); - return hidden_states; -} - -} // namespace infinilm::models::qwen3_next diff --git a/csrc/models/qwen3_next/qwen3_next_decoderLayer.hpp b/csrc/models/qwen3_next/qwen3_next_decoderLayer.hpp index dd0505bd5..df7f4c36e 100644 --- a/csrc/models/qwen3_next/qwen3_next_decoderLayer.hpp +++ b/csrc/models/qwen3_next/qwen3_next_decoderLayer.hpp @@ -1,38 +1,15 @@ #pragma once +#include "../../layers/causal_lm_templates/hybrid_decoder_layer.hpp" #include "qwen3_next_attention.hpp" #include "qwen3_next_gated_deltanet.hpp" #include "qwen3_next_sparse_moe_block.hpp" -#include -#include namespace infinilm::models::qwen3_next { -class Qwen3NextDecoderLayer : public infinicore::nn::Module { -public: - Qwen3NextDecoderLayer(std::shared_ptr model_config, - size_t layer_idx, - const infinicore::Device &device); - - std::tuple forward(const infinicore::Tensor &positions, - infinicore::Tensor &hidden_states, - infinicore::Tensor &residual); - - infinicore::Tensor forward(const infinicore::Tensor &positions, - infinicore::Tensor &hidden_states); - - size_t layer_idx() const { return layer_idx_; } - -protected: - INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, input_layernorm); - INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, post_attention_layernorm); - INFINICORE_NN_MODULE(Qwen3NextAttention, self_attn); - INFINICORE_NN_MODULE(Qwen3NextGatedDeltaNet, linear_attn); - INFINICORE_NN_MODULE(Qwen3NextSparseMoeBlock, mlp); - -private: - size_t layer_idx_; - std::string layer_type_; -}; +using Qwen3NextDecoderLayer = infinilm::layers::causal_lm_templates::HybridDecoderLayer< + Qwen3NextAttention, + Qwen3NextGatedDeltaNet, + Qwen3NextSparseMoeBlock>; } // namespace infinilm::models::qwen3_next diff --git a/csrc/models/qwen3_next/qwen3_next_sparse_moe_block.cpp b/csrc/models/qwen3_next/qwen3_next_sparse_moe_block.cpp deleted file mode 100644 index d3548c9fa..000000000 --- a/csrc/models/qwen3_next/qwen3_next_sparse_moe_block.cpp +++ /dev/null @@ -1,106 +0,0 @@ -#include "qwen3_next_sparse_moe_block.hpp" - -#include "../../global_state/global_state.hpp" - -#include -#include -#include -#include -#include -#include - -#include - -namespace infinilm::models::qwen3_next { - -Qwen3NextSharedExpert::Qwen3NextSharedExpert(std::shared_ptr model_config, - const infinicore::Device &device) { - const auto &dtype{model_config->get_dtype()}; - const size_t hidden_size = model_config->get("hidden_size"); - const size_t intermediate_size = model_config->get("shared_expert_intermediate_size"); - - const engine::distributed::RankInfo &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info(); - auto quantization_method = model_config->get_quantization_method(); - auto register_fn = [this](const std::string &n, infinicore::nn::Parameter p) { this->register_parameter(n, std::move(p)); }; - gate_up_proj_ = std::make_shared( - hidden_size, - intermediate_size, - "gate_proj", - "up_proj", - register_fn, - quantization_method, - false, - dtype, - device, - rank_info); - down_proj_ = this->register_module( - "down_proj", - intermediate_size, - hidden_size, - quantization_method, - false, - dtype, - device, - rank_info.tp_rank, - rank_info.tp_size, - rank_info.comm); -} - -infinicore::Tensor Qwen3NextSharedExpert::forward(const infinicore::Tensor &hidden_states) const { - auto hidden_states_mutable = hidden_states; - auto [gate, up] = gate_up_proj_->forward_split(hidden_states_mutable); - auto intermediate = infinicore::op::swiglu(up, gate); - return down_proj_->forward(intermediate); -} - -Qwen3NextSparseMoeBlock::Qwen3NextSparseMoeBlock(std::shared_ptr model_config, - const infinicore::Device &device) - : Qwen3NextSparseMoeBlock(model_config, 0, device) { -} - -Qwen3NextSparseMoeBlock::Qwen3NextSparseMoeBlock(std::shared_ptr model_config, - size_t layer_idx, - const infinicore::Device &device) { - gate_ = this->register_module("gate", model_config, device); - experts_ = this->register_module("experts", model_config, device); - fused_moe_ = this->register_module("fused_moe", model_config, device, layer_idx); - shared_expert_ = this->register_module("shared_expert", model_config, device); - shared_expert_gate_ = this->register_module( - "shared_expert_gate", - model_config->get("hidden_size"), - 1, - false, - model_config->get_dtype(), - device); -} - -infinicore::Tensor Qwen3NextSparseMoeBlock::forward(const infinicore::Tensor &hidden_states) const { - ASSERT(hidden_states->ndim() == 3); - - auto shape = hidden_states->shape(); - auto hidden_states_reshaped = hidden_states->view({shape[0] * shape[1], shape[2]}); - - auto [routing_weights, selected_experts] = gate_->forward(hidden_states_reshaped); - infinilm::layers::moe::TopKOutput topk_output{ - routing_weights, - selected_experts, - infinicore::Tensor(), - }; - auto routed_states = fused_moe_->forward( - hidden_states_reshaped, - topk_output, - experts_->moe_weights()); - - auto shared_states = shared_expert_->forward(hidden_states); - auto hidden_states_for_gate = hidden_states; - auto shared_gate = infinicore::op::sigmoid(shared_expert_gate_->forward(hidden_states_for_gate)); - shared_gate = shared_gate->as_strided(shared_states->shape(), {shared_gate->stride(0), shared_gate->stride(1), 0}); - shared_states = infinicore::op::mul(shared_states, shared_gate); - - auto routed_states_3d = routed_states->as_strided( - {shape[0], shape[1], shape[2]}, - {static_cast(shape[1] * shape[2]), static_cast(shape[2]), 1}); - return infinicore::op::add(routed_states_3d, shared_states); -} - -} // namespace infinilm::models::qwen3_next diff --git a/csrc/models/qwen3_next/qwen3_next_sparse_moe_block.hpp b/csrc/models/qwen3_next/qwen3_next_sparse_moe_block.hpp index 2cec0bd50..dd302355a 100644 --- a/csrc/models/qwen3_next/qwen3_next_sparse_moe_block.hpp +++ b/csrc/models/qwen3_next/qwen3_next_sparse_moe_block.hpp @@ -1,43 +1,9 @@ #pragma once -#include "../../config/model_config.hpp" -#include "../../layers/linear/linear.hpp" -#include "../../layers/moe/experts/fused_moe_experts.hpp" -#include "../../layers/moe/fused_moe.hpp" -#include "../../layers/moe/router/topk_router.hpp" - -#include +#include "../../layers/moe/shared_expert_sparse_moe_block.hpp" namespace infinilm::models::qwen3_next { -class Qwen3NextSharedExpert : public infinicore::nn::Module { -public: - Qwen3NextSharedExpert(std::shared_ptr model_config, - const infinicore::Device &device); - - infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const; - -protected: - std::shared_ptr gate_up_proj_; - std::shared_ptr down_proj_; -}; - -class Qwen3NextSparseMoeBlock : public infinicore::nn::Module { -public: - Qwen3NextSparseMoeBlock(std::shared_ptr model_config, - const infinicore::Device &device); - Qwen3NextSparseMoeBlock(std::shared_ptr model_config, - size_t layer_idx, - const infinicore::Device &device); - - infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const; - -protected: - std::shared_ptr gate_; - std::shared_ptr experts_; - std::shared_ptr fused_moe_; - std::shared_ptr shared_expert_; - std::shared_ptr shared_expert_gate_; -}; +using Qwen3NextSparseMoeBlock = infinilm::layers::moe::SharedExpertSparseMoeBlock; } // namespace infinilm::models::qwen3_next From 082245f42ce3c574c69374db4e1c4fd59b284699 Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Mon, 3 Aug 2026 03:54:42 +0000 Subject: [PATCH 4/9] fix(graph): make hybrid metadata config-driven --- csrc/engine/compiler/paged_compiler.cpp | 48 ++++++++----------- csrc/global_state/forward_context.hpp | 1 + csrc/models/infinilm_model.hpp | 3 ++ csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp | 17 +++++-- csrc/models/qwen3_5/qwen3_5_model.cpp | 2 + .../qwen3_next_allocate_kv_cache_tensors.cpp | 12 ++++- .../qwen3_next_allocate_kv_cache_tensors.hpp | 1 + .../qwen3_next/qwen3_next_for_causal_lm.cpp | 2 + python/infinilm/infer_engine.py | 30 +++++++++++- 9 files changed, 82 insertions(+), 34 deletions(-) diff --git a/csrc/engine/compiler/paged_compiler.cpp b/csrc/engine/compiler/paged_compiler.cpp index 7a5c57fd0..1ec2ebde3 100644 --- a/csrc/engine/compiler/paged_compiler.cpp +++ b/csrc/engine/compiler/paged_compiler.cpp @@ -4,22 +4,14 @@ #include #include +#include #include namespace infinilm::engine { namespace { bool has_mamba_cache(const infinilm::global_state::ForwardContext &forward_context) { - auto has_state = [](const std::vector &state_vec) { - for (const auto &state : state_vec) { - if (state) { - return true; - } - } - return false; - }; - - return has_state(forward_context.conv_state_vec) || has_state(forward_context.ssm_state_vec); + return forward_context.mamba_state_pool_size > 0; } } // namespace @@ -45,22 +37,22 @@ void PagedCompiler::compile() { size_t nblocks = dynamic_cast(model_->get_cache_config())->num_blocks(); auto &forward_context = infinilm::global_state::get_forward_context(); const bool has_mamba_state = has_mamba_cache(forward_context); + const auto &model_config = model_->get_model_config(); + const size_t position_id_axes = model_config == nullptr + ? 1 + : model_config->get_or("position_id_axes", 1); + if (position_id_axes == 0) { + throw std::runtime_error("PagedCompiler: position_id_axes must be positive"); + } auto compile_batch_sizes = decode_batch_sizes_; size_t max_batch_size = *std::max_element(decode_batch_sizes_.begin(), decode_batch_sizes_.end()); if (has_mamba_state) { - size_t max_mamba_batch_size = max_batch_size; - auto clamp_to_state_pool = [&max_mamba_batch_size](const std::vector &state_vec) { - for (const auto &state : state_vec) { - if (!state) { - continue; - } - const size_t pool_size = state->size(0); - const size_t usable_rows = pool_size > 0 ? pool_size - 1 : 0; - max_mamba_batch_size = std::min(max_mamba_batch_size, usable_rows); - } - }; - clamp_to_state_pool(forward_context.conv_state_vec); - clamp_to_state_pool(forward_context.ssm_state_vec); + if (forward_context.mamba_state_pool_size < 2) { + throw std::runtime_error( + "PagedCompiler: mamba state pool must reserve row 0 and at least one request row"); + } + const size_t max_mamba_batch_size = std::min( + max_batch_size, forward_context.mamba_state_pool_size - 1); compile_batch_sizes.erase( std::remove_if( compile_batch_sizes.begin(), @@ -83,12 +75,12 @@ void PagedCompiler::compile() { auto make_decode_input = [&](size_t b) { InfinilmModel::Input input; input.input_ids = infinicore::Tensor::empty({1, b}, infinicore::DataType::I64, infinicore::context::getDevice()); - // Qwen3.5 hybrid models use three MRoPE axes even for text-only - // decode. Their scheduler supplies position_ids with shape [3, b]. - // Other paged models continue to use the traditional [b] layout. + // Models declare their position-id axes explicitly. Single-axis + // models retain the traditional [b] layout. input.position_ids = infinicore::Tensor::empty( - has_mamba_state ? std::vector{3, b} - : std::vector{b}, + position_id_axes > 1 + ? std::vector{position_id_axes, b} + : std::vector{b}, infinicore::DataType::I64, infinicore::context::getDevice()); input.total_sequence_lengths = infinicore::Tensor::empty({b}, infinicore::DataType::I32, infinicore::context::getDevice()); diff --git a/csrc/global_state/forward_context.hpp b/csrc/global_state/forward_context.hpp index f395b531a..732bf0f5b 100644 --- a/csrc/global_state/forward_context.hpp +++ b/csrc/global_state/forward_context.hpp @@ -62,6 +62,7 @@ struct ForwardContext { std::vector kv_cache_vec; std::vector conv_state_vec; std::vector ssm_state_vec; + size_t mamba_state_pool_size{0}; }; void initialize_forward_context(ForwardContext &forward_context); diff --git a/csrc/models/infinilm_model.hpp b/csrc/models/infinilm_model.hpp index ac994fd6d..27275ee26 100644 --- a/csrc/models/infinilm_model.hpp +++ b/csrc/models/infinilm_model.hpp @@ -70,6 +70,9 @@ class InfinilmModel : public infinicore::nn::Module { virtual const cache::CacheConfig *get_cache_config() const { return cache_config_.get(); } + const std::shared_ptr &get_model_config() const { + return model_config_; + } void process_weights_after_loading(); void reset_runtime_state() const; diff --git a/csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp b/csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp index 1151551fb..6aa7129ef 100644 --- a/csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp +++ b/csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp @@ -48,10 +48,21 @@ std::shared_ptr prepare_qwen3_5_model_config(std: config_json["dtype"] = config_json["torch_dtype"]; } } + if (!config_json.contains("position_id_axes")) { + size_t position_id_axes = 1; + if (config_json.contains("rope_parameters") + && config_json["rope_parameters"].is_object()) { + const auto &rope_parameters = config_json["rope_parameters"]; + if (rope_parameters.contains("mrope_section") + && rope_parameters["mrope_section"].is_array() + && !rope_parameters["mrope_section"].empty()) { + position_id_axes = rope_parameters["mrope_section"].size(); + } + } + config_json["position_id_axes"] = position_id_axes; + } if (!config_json.contains("rope_theta") && config_json.contains("rope_parameters") && config_json["rope_parameters"].is_object() && config_json["rope_parameters"].contains("rope_theta")) { - // TODO: This is only a temporary loader shim. Qwen3.6 uses mRoPE, - // which needs proper support in InfiniCore instead of treating it as - // plain RoPE through a top-level rope_theta. + // Normalize the nested HuggingFace field for the Qwen3.5 attention module. config_json["rope_theta"] = config_json["rope_parameters"]["rope_theta"]; } if (!config_json.contains("partial_rotary_factor") && config_json.contains("rope_parameters") && config_json["rope_parameters"].is_object() && config_json["rope_parameters"].contains("partial_rotary_factor")) { diff --git a/csrc/models/qwen3_5/qwen3_5_model.cpp b/csrc/models/qwen3_5/qwen3_5_model.cpp index 3dd3e61d1..a30d94cbf 100644 --- a/csrc/models/qwen3_5/qwen3_5_model.cpp +++ b/csrc/models/qwen3_5/qwen3_5_model.cpp @@ -130,6 +130,7 @@ void Qwen35ModelBase::reset_cache(const cache::CacheConfig *cache_config) { forward_context.kv_cache_vec.clear(); forward_context.conv_state_vec.clear(); forward_context.ssm_state_vec.clear(); + forward_context.mamba_state_pool_size = 0; const backends::AttentionBackend attention_backend = infinilm::global_state::get_infinilm_config().attention_backend; @@ -137,6 +138,7 @@ void Qwen35ModelBase::reset_cache(const cache::CacheConfig *cache_config) { forward_context.kv_cache_vec = std::move(cache_vectors.kv_cache_tensors); forward_context.conv_state_vec = std::move(cache_vectors.conv_state_tensors); forward_context.ssm_state_vec = std::move(cache_vectors.ssm_state_tensors); + forward_context.mamba_state_pool_size = cache_vectors.mamba_state_pool_size; } } // namespace infinilm::models::qwen3_5 diff --git a/csrc/models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.cpp b/csrc/models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.cpp index 3aaf898cb..3ba3fdbdc 100644 --- a/csrc/models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.cpp +++ b/csrc/models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.cpp @@ -43,7 +43,16 @@ AllocatedHybridCache qwen3_next_allocate_cache_tensors( conv_state_vec.reserve(num_hidden_layers); ssm_state_vec.reserve(num_hidden_layers); + size_t mamba_state_pool_size = 0; auto allocate_linear_attention_cache = [&](size_t layer_idx, size_t pool_size) { + if (mamba_state_pool_size == 0) { + mamba_state_pool_size = pool_size; + } else if (mamba_state_pool_size != pool_size) { + throw std::runtime_error( + "qwen3_next_allocate_cache_tensors: inconsistent mamba state pool size at layer " + + std::to_string(layer_idx)); + } + auto conv_state = cache::MambaCache::create_layer_conv_state( linear_key_head_dim, linear_value_head_dim, @@ -141,7 +150,8 @@ AllocatedHybridCache qwen3_next_allocate_cache_tensors( return AllocatedHybridCache{ std::move(kv_cache_vec), std::move(conv_state_vec), - std::move(ssm_state_vec)}; + std::move(ssm_state_vec), + mamba_state_pool_size}; } } // namespace infinilm::models::qwen3_next diff --git a/csrc/models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.hpp b/csrc/models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.hpp index a4b5190f1..1d5a45dea 100644 --- a/csrc/models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.hpp +++ b/csrc/models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.hpp @@ -16,6 +16,7 @@ struct AllocatedHybridCache { std::vector kv_cache_tensors; std::vector conv_state_tensors; std::vector ssm_state_tensors; + size_t mamba_state_pool_size{0}; }; AllocatedHybridCache qwen3_next_allocate_cache_tensors( diff --git a/csrc/models/qwen3_next/qwen3_next_for_causal_lm.cpp b/csrc/models/qwen3_next/qwen3_next_for_causal_lm.cpp index 1b6354540..1a9eaeb3b 100644 --- a/csrc/models/qwen3_next/qwen3_next_for_causal_lm.cpp +++ b/csrc/models/qwen3_next/qwen3_next_for_causal_lm.cpp @@ -37,6 +37,7 @@ void Qwen3NextForCausalLM::reset_cache(const cache::CacheConfig *cache_config) { forward_context.kv_cache_vec.clear(); forward_context.conv_state_vec.clear(); forward_context.ssm_state_vec.clear(); + forward_context.mamba_state_pool_size = 0; const backends::AttentionBackend attention_backend = infinilm::global_state::get_infinilm_config().attention_backend; @@ -44,6 +45,7 @@ void Qwen3NextForCausalLM::reset_cache(const cache::CacheConfig *cache_config) { forward_context.kv_cache_vec = std::move(cache_vectors.kv_cache_tensors); forward_context.conv_state_vec = std::move(cache_vectors.conv_state_tensors); forward_context.ssm_state_vec = std::move(cache_vectors.ssm_state_tensors); + forward_context.mamba_state_pool_size = cache_vectors.mamba_state_pool_size; } std::shared_ptr create_qwen3_next_model_config(std::shared_ptr model_config) { diff --git a/python/infinilm/infer_engine.py b/python/infinilm/infer_engine.py index 6b1499fb5..a015c2eb7 100644 --- a/python/infinilm/infer_engine.py +++ b/python/infinilm/infer_engine.py @@ -106,6 +106,27 @@ class GenerationConfig: stop_on_eos: bool = True +def _infer_position_id_axes(hf_config: dict) -> int: + text_config = hf_config.get("text_config", hf_config) + if not isinstance(text_config, dict): + return 1 + + explicit_axes = text_config.get( + "position_id_axes", hf_config.get("position_id_axes") + ) + if explicit_axes is not None: + axes = int(explicit_axes) + if axes < 1: + raise ValueError("position_id_axes must be positive") + return axes + + rope_parameters = text_config.get("rope_parameters") or {} + mrope_section = rope_parameters.get("mrope_section") + if isinstance(mrope_section, (list, tuple)) and mrope_section: + return len(mrope_section) + return 1 + + class InferEngine(_infinilm.InferEngine): def __init__( self, @@ -125,6 +146,11 @@ def __init__( self.hf_config = read_hf_config(model_path) self.hf_generation_config = read_hf_generation_config(model_path) self.hf_config["use_legacy_moe"] = bool(use_legacy_moe) + self.position_id_axes = _infer_position_id_axes(self.hf_config) + self.hf_config["position_id_axes"] = self.position_id_axes + text_config = self.hf_config.get("text_config") + if isinstance(text_config, dict): + text_config.setdefault("position_id_axes", self.position_id_axes) if device is None: device = infinicore.device() @@ -519,8 +545,8 @@ def generate( position_ids_list = ( list(range(past_seq_len, past_seq_len + seq_len)) * batch_size ) - if self.model_type in ("qwen3_5", "qwen3_5_moe"): - position_ids_list = [position_ids_list] * 3 + if self.position_id_axes > 1: + position_ids_list = [position_ids_list] * self.position_id_axes position_ids = infinicore.from_list( position_ids_list, dtype=infinicore.int64 ) From 34f76b4592c57e66609ce0392f842fc60d165add Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Mon, 3 Aug 2026 05:08:19 +0000 Subject: [PATCH 5/9] refactor(qwen3.5): share model wrappers and reset caches --- csrc/global_state/forward_context.hpp | 7 +++++ .../moe/shared_expert_sparse_moe_block.cpp | 26 +++-------------- .../moe/shared_expert_sparse_moe_block.hpp | 9 ++---- csrc/models/infinilm_model.cpp | 8 +++--- csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp | 28 ------------------- csrc/models/qwen3_5/qwen3_5_for_causal_lm.hpp | 27 ++++++++++-------- csrc/models/qwen3_5/qwen3_5_model.cpp | 23 ++------------- csrc/models/qwen3_5/qwen3_5_model.hpp | 28 +++++++++++++++---- .../qwen3_5_moe/qwen3_5_moe_for_causal_lm.cpp | 28 ------------------- .../qwen3_5_moe/qwen3_5_moe_for_causal_lm.hpp | 16 ++--------- csrc/models/qwen3_5_moe/qwen3_5_moe_model.cpp | 21 -------------- csrc/models/qwen3_5_moe/qwen3_5_moe_model.hpp | 12 +------- .../qwen3_next/qwen3_next_for_causal_lm.cpp | 5 +--- 13 files changed, 61 insertions(+), 177 deletions(-) delete mode 100644 csrc/models/qwen3_5_moe/qwen3_5_moe_model.cpp diff --git a/csrc/global_state/forward_context.hpp b/csrc/global_state/forward_context.hpp index 732bf0f5b..ed7de761a 100644 --- a/csrc/global_state/forward_context.hpp +++ b/csrc/global_state/forward_context.hpp @@ -63,6 +63,13 @@ struct ForwardContext { std::vector conv_state_vec; std::vector ssm_state_vec; size_t mamba_state_pool_size{0}; + + void clear_model_caches() { + kv_cache_vec.clear(); + conv_state_vec.clear(); + ssm_state_vec.clear(); + mamba_state_pool_size = 0; + } }; void initialize_forward_context(ForwardContext &forward_context); diff --git a/csrc/layers/moe/shared_expert_sparse_moe_block.cpp b/csrc/layers/moe/shared_expert_sparse_moe_block.cpp index dc80e9e66..696e57c99 100644 --- a/csrc/layers/moe/shared_expert_sparse_moe_block.cpp +++ b/csrc/layers/moe/shared_expert_sparse_moe_block.cpp @@ -16,11 +16,8 @@ SharedExpertSparseMoeBlock::SharedExpertSparseMoeBlock( SharedExpertSparseMoeBlock::SharedExpertSparseMoeBlock( std::shared_ptr model_config, size_t layer_idx, - const infinicore::Device &device) { - INFINICORE_NN_MODULE_INIT(gate, model_config, device); - INFINICORE_NN_MODULE_INIT(experts, model_config, device); - INFINICORE_NN_MODULE_INIT(fused_moe, model_config, device, layer_idx); - + const infinicore::Device &device) + : SparseMoeBlock(model_config, device, layer_idx) { auto shared_config_json = model_config->get_config_json(); shared_config_json["intermediate_size"] = model_config->get("shared_expert_intermediate_size"); auto shared_config = std::make_shared( @@ -38,22 +35,7 @@ SharedExpertSparseMoeBlock::SharedExpertSparseMoeBlock( infinicore::Tensor SharedExpertSparseMoeBlock::forward( const infinicore::Tensor &hidden_states) const { - ASSERT(hidden_states->ndim() == 3); - - const auto shape = hidden_states->shape(); - auto hidden_states_reshaped = hidden_states->view( - {shape[0] * shape[1], shape[2]}); - - auto [routing_weights, selected_experts] = gate_->forward(hidden_states_reshaped); - TopKOutput topk_output{ - routing_weights, - selected_experts, - infinicore::Tensor(), - }; - auto routed_output = fused_moe_->forward( - hidden_states_reshaped, - topk_output, - experts_->moe_weights()); + auto routed_output = SparseMoeBlock::forward(hidden_states); auto shared_output = shared_expert_->forward(hidden_states); auto shared_gate_input = hidden_states; @@ -64,7 +46,7 @@ infinicore::Tensor SharedExpertSparseMoeBlock::forward( {shared_gate->stride(0), shared_gate->stride(1), 0}); shared_output = infinicore::op::mul(shared_output, shared_gate); - return infinicore::op::add(routed_output->view(shape), shared_output); + return infinicore::op::add(routed_output, shared_output); } } // namespace infinilm::layers::moe diff --git a/csrc/layers/moe/shared_expert_sparse_moe_block.hpp b/csrc/layers/moe/shared_expert_sparse_moe_block.hpp index 249115164..8ad23b3de 100644 --- a/csrc/layers/moe/shared_expert_sparse_moe_block.hpp +++ b/csrc/layers/moe/shared_expert_sparse_moe_block.hpp @@ -1,8 +1,6 @@ #pragma once -#include "experts/fused_moe_experts.hpp" -#include "fused_moe.hpp" -#include "router/topk_router.hpp" +#include "sparse_moe_block.hpp" #include "../../config/model_config.hpp" #include "../linear/linear.hpp" @@ -14,7 +12,7 @@ namespace infinilm::layers::moe { -class SharedExpertSparseMoeBlock : public infinicore::nn::Module { +class SharedExpertSparseMoeBlock : public SparseMoeBlock { public: SharedExpertSparseMoeBlock( std::shared_ptr model_config, @@ -27,9 +25,6 @@ class SharedExpertSparseMoeBlock : public infinicore::nn::Module { infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const; protected: - INFINICORE_NN_MODULE(TopKRouter, gate); - INFINICORE_NN_MODULE(FusedMoeExperts, experts); - INFINICORE_NN_MODULE(FusedMoE, fused_moe); INFINICORE_NN_MODULE(infinilm::layers::mlp::MLP, shared_expert); INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, shared_expert_gate); }; diff --git a/csrc/models/infinilm_model.cpp b/csrc/models/infinilm_model.cpp index c2cc68762..2e362259e 100644 --- a/csrc/models/infinilm_model.cpp +++ b/csrc/models/infinilm_model.cpp @@ -7,16 +7,16 @@ namespace infinilm { void InfinilmModel::reset_cache(const cache::CacheConfig *cache_config) { + auto &forward_context = global_state::get_forward_context(); + forward_context.clear_model_caches(); if (cache_config == nullptr) { cache_config_.reset(); - global_state::get_forward_context().kv_cache_vec.clear(); return; } cache_config_ = cache_config->unique_copy(); - auto &kv_cache_vec = global_state::get_forward_context().kv_cache_vec; - kv_cache_vec.clear(); const backends::AttentionBackend attention_backend = infinilm::global_state::get_infinilm_config().attention_backend; - kv_cache_vec = std::move(default_allocate_kv_cache_tensors(cache_config, model_config_, attention_backend)); + forward_context.kv_cache_vec = std::move( + default_allocate_kv_cache_tensors(cache_config, model_config_, attention_backend)); } std::vector InfinilmModel::default_allocate_kv_cache_tensors( diff --git a/csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp b/csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp index 6aa7129ef..07836eab7 100644 --- a/csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp +++ b/csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp @@ -1,40 +1,12 @@ #include "qwen3_5_for_causal_lm.hpp" -#include "../../global_state/global_state.hpp" #include "../models_registry.hpp" -#include "../qwen3_next/qwen3_next_for_causal_lm.hpp" #include #include #include namespace infinilm::models::qwen3_5 { -Qwen35ForCausalLM::Qwen35ForCausalLM(std::shared_ptr model_config, - const infinicore::Device &device) { - model_config_ = model_config; - size_t hidden_size = model_config->get("hidden_size"); - size_t vocab_size = model_config->get("vocab_size"); - const auto &dtype{model_config->get_dtype()}; - - INFINICORE_NN_MODULE_INIT(model, model_config, device); - INFINICORE_NN_MODULE_INIT(lm_head, hidden_size, vocab_size, false, dtype, device); -} - -infinilm::InfinilmModel::Output Qwen35ForCausalLM::forward(const infinilm::InfinilmModel::Input &input) const { - auto hidden_states = model_->forward(input); - auto logits = lm_head_->forward(hidden_states); - return {logits}; -} - -void Qwen35ForCausalLM::reset_cache(const cache::CacheConfig *cache_config) { - if (cache_config == nullptr) { - cache_config_.reset(); - } else { - cache_config_ = cache_config->unique_copy(); - } - model_->reset_cache(cache_config); -} - std::shared_ptr prepare_qwen3_5_model_config(std::shared_ptr model_config) { nlohmann::json &config_json = model_config->get_config_json(); if (config_json.contains("text_config") && config_json["text_config"].is_object()) { diff --git a/csrc/models/qwen3_5/qwen3_5_for_causal_lm.hpp b/csrc/models/qwen3_5/qwen3_5_for_causal_lm.hpp index 787a345fd..2aaaf4a12 100644 --- a/csrc/models/qwen3_5/qwen3_5_for_causal_lm.hpp +++ b/csrc/models/qwen3_5/qwen3_5_for_causal_lm.hpp @@ -1,25 +1,30 @@ #pragma once +#include "../../layers/causal_lm_templates/text_causal_lm.hpp" #include "qwen3_5_model.hpp" #include #include namespace infinilm::models::qwen3_5 { -class Qwen35ForCausalLM : public InfinilmModel { +template +class Qwen35CausalLM : public infinilm::layers::causal_lm_templates::TextCausalLM { public: - Qwen35ForCausalLM(std::shared_ptr model_config, - const infinicore::Device &device); - - Output forward(const Input &input) const override; - - void reset_cache(const cache::CacheConfig *cache_config) override; - -protected: - INFINICORE_NN_MODULE(Qwen35Model, model); - INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, lm_head); + using Base = infinilm::layers::causal_lm_templates::TextCausalLM; + using Base::Base; + + void reset_cache(const cache::CacheConfig *cache_config) override { + if (cache_config == nullptr) { + this->cache_config_.reset(); + } else { + this->cache_config_ = cache_config->unique_copy(); + } + this->model().reset_cache(cache_config); + } }; +using Qwen35ForCausalLM = Qwen35CausalLM; + std::shared_ptr prepare_qwen3_5_model_config(std::shared_ptr model_config); std::shared_ptr create_qwen3_5_model_config(std::shared_ptr model_config); diff --git a/csrc/models/qwen3_5/qwen3_5_model.cpp b/csrc/models/qwen3_5/qwen3_5_model.cpp index a30d94cbf..969ea5485 100644 --- a/csrc/models/qwen3_5/qwen3_5_model.cpp +++ b/csrc/models/qwen3_5/qwen3_5_model.cpp @@ -42,12 +42,6 @@ Qwen35ModelBase::Qwen35ModelBase(std::shared_ptr } } -Qwen35Model::Qwen35Model(std::shared_ptr model_config, - const infinicore::Device &device) - : Qwen35ModelBase(model_config, device) { - INFINICORE_NN_MODULE_INIT(language_model, model_config, device); -} - void Qwen35ModelBase::replace_image_embeddings(infinicore::Tensor &inputs_embeds, const InfinilmModel::Input &input) const { if (!input.pixel_values.has_value() || input.pixel_values->empty()) { @@ -112,26 +106,13 @@ void Qwen35ModelBase::replace_image_embeddings(infinicore::Tensor &inputs_embeds } } -infinicore::Tensor Qwen35Model::forward(const InfinilmModel::Input &input) const { - if (input.pixel_values.has_value() && !input.pixel_values->empty()) { - auto inputs_embeds = language_model_->embed_tokens(input.input_ids.value()); - replace_image_embeddings(inputs_embeds, input); - return language_model_->forward_embeds(inputs_embeds, input.position_ids.value()); - } - return language_model_->forward(input); -} - void Qwen35ModelBase::reset_cache(const cache::CacheConfig *cache_config) { + auto &forward_context = infinilm::global_state::get_forward_context(); + forward_context.clear_model_caches(); if (nullptr == cache_config) { return; } - auto &forward_context = infinilm::global_state::get_forward_context(); - forward_context.kv_cache_vec.clear(); - forward_context.conv_state_vec.clear(); - forward_context.ssm_state_vec.clear(); - forward_context.mamba_state_pool_size = 0; - const backends::AttentionBackend attention_backend = infinilm::global_state::get_infinilm_config().attention_backend; auto cache_vectors = infinilm::models::qwen3_next::qwen3_next_allocate_cache_tensors(cache_config, model_config_, attention_backend); diff --git a/csrc/models/qwen3_5/qwen3_5_model.hpp b/csrc/models/qwen3_5/qwen3_5_model.hpp index c82763c22..e9f2080bb 100644 --- a/csrc/models/qwen3_5/qwen3_5_model.hpp +++ b/csrc/models/qwen3_5/qwen3_5_model.hpp @@ -25,15 +25,31 @@ class Qwen35ModelBase : public infinicore::nn::Module { std::shared_ptr model_config_; }; -class Qwen35Model : public Qwen35ModelBase { +template +class Qwen35ModelTemplate : public Qwen35ModelBase { public: - Qwen35Model(std::shared_ptr model_config, - const infinicore::Device &device); - - infinicore::Tensor forward(const InfinilmModel::Input &input) const; + Qwen35ModelTemplate( + std::shared_ptr model_config, + const infinicore::Device &device) + : Qwen35ModelBase(model_config, device) { + language_model_ = this->register_module( + "language_model", model_config, device); + } + + infinicore::Tensor forward(const InfinilmModel::Input &input) const { + if (input.pixel_values.has_value() && !input.pixel_values->empty()) { + auto inputs_embeds = language_model_->embed_tokens(input.input_ids.value()); + replace_image_embeddings(inputs_embeds, input); + return language_model_->forward_embeds( + inputs_embeds, input.position_ids.value()); + } + return language_model_->forward(input); + } protected: - INFINICORE_NN_MODULE(Qwen35LanguageModel, language_model); + INFINICORE_NN_MODULE(LanguageModel, language_model); }; +using Qwen35Model = Qwen35ModelTemplate; + } // namespace infinilm::models::qwen3_5 diff --git a/csrc/models/qwen3_5_moe/qwen3_5_moe_for_causal_lm.cpp b/csrc/models/qwen3_5_moe/qwen3_5_moe_for_causal_lm.cpp index 7797c17ae..99dc16d92 100644 --- a/csrc/models/qwen3_5_moe/qwen3_5_moe_for_causal_lm.cpp +++ b/csrc/models/qwen3_5_moe/qwen3_5_moe_for_causal_lm.cpp @@ -8,34 +8,6 @@ namespace infinilm::models::qwen3_5_moe { -Qwen35MoeForConditionalGeneration::Qwen35MoeForConditionalGeneration( - std::shared_ptr model_config, - const infinicore::Device &device) { - model_config_ = model_config; - const size_t hidden_size = model_config->get("hidden_size"); - const size_t vocab_size = model_config->get("vocab_size"); - const auto &dtype = model_config->get_dtype(); - - INFINICORE_NN_MODULE_INIT(model, model_config, device); - INFINICORE_NN_MODULE_INIT(lm_head, hidden_size, vocab_size, false, dtype, device); -} - -InfinilmModel::Output Qwen35MoeForConditionalGeneration::forward( - const InfinilmModel::Input &input) const { - auto hidden_states = model_->forward(input); - return {lm_head_->forward(hidden_states)}; -} - -void Qwen35MoeForConditionalGeneration::reset_cache( - const cache::CacheConfig *cache_config) { - if (cache_config == nullptr) { - cache_config_.reset(); - } else { - cache_config_ = cache_config->unique_copy(); - } - model_->reset_cache(cache_config); -} - std::shared_ptr create_qwen3_5_moe_model_config( std::shared_ptr model_config) { const std::string model_type = model_config->get("model_type"); diff --git a/csrc/models/qwen3_5_moe/qwen3_5_moe_for_causal_lm.hpp b/csrc/models/qwen3_5_moe/qwen3_5_moe_for_causal_lm.hpp index 2ac0abe30..19babff67 100644 --- a/csrc/models/qwen3_5_moe/qwen3_5_moe_for_causal_lm.hpp +++ b/csrc/models/qwen3_5_moe/qwen3_5_moe_for_causal_lm.hpp @@ -1,25 +1,13 @@ #pragma once +#include "../qwen3_5/qwen3_5_for_causal_lm.hpp" #include "qwen3_5_moe_model.hpp" #include namespace infinilm::models::qwen3_5_moe { -class Qwen35MoeForConditionalGeneration : public InfinilmModel { -public: - Qwen35MoeForConditionalGeneration( - std::shared_ptr model_config, - const infinicore::Device &device); - - Output forward(const Input &input) const override; - - void reset_cache(const cache::CacheConfig *cache_config) override; - -protected: - INFINICORE_NN_MODULE(Qwen35MoeModel, model); - INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, lm_head); -}; +using Qwen35MoeForConditionalGeneration = qwen3_5::Qwen35CausalLM; std::shared_ptr create_qwen3_5_moe_model_config( std::shared_ptr model_config); diff --git a/csrc/models/qwen3_5_moe/qwen3_5_moe_model.cpp b/csrc/models/qwen3_5_moe/qwen3_5_moe_model.cpp deleted file mode 100644 index ab54a1f18..000000000 --- a/csrc/models/qwen3_5_moe/qwen3_5_moe_model.cpp +++ /dev/null @@ -1,21 +0,0 @@ -#include "qwen3_5_moe_model.hpp" - -namespace infinilm::models::qwen3_5_moe { - -Qwen35MoeModel::Qwen35MoeModel( - std::shared_ptr model_config, - const infinicore::Device &device) - : Qwen35ModelBase(model_config, device) { - INFINICORE_NN_MODULE_INIT(language_model, model_config, device); -} - -infinicore::Tensor Qwen35MoeModel::forward(const InfinilmModel::Input &input) const { - if (input.pixel_values.has_value() && !input.pixel_values->empty()) { - auto inputs_embeds = language_model_->embed_tokens(input.input_ids.value()); - replace_image_embeddings(inputs_embeds, input); - return language_model_->forward_embeds(inputs_embeds, input.position_ids.value()); - } - return language_model_->forward(input); -} - -} // namespace infinilm::models::qwen3_5_moe diff --git a/csrc/models/qwen3_5_moe/qwen3_5_moe_model.hpp b/csrc/models/qwen3_5_moe/qwen3_5_moe_model.hpp index 7d944513f..8e10cc349 100644 --- a/csrc/models/qwen3_5_moe/qwen3_5_moe_model.hpp +++ b/csrc/models/qwen3_5_moe/qwen3_5_moe_model.hpp @@ -6,16 +6,6 @@ namespace infinilm::models::qwen3_5_moe { using Qwen35MoeLanguageModel = infinilm::layers::causal_lm_templates::TextModel; - -class Qwen35MoeModel : public qwen3_5::Qwen35ModelBase { -public: - Qwen35MoeModel(std::shared_ptr model_config, - const infinicore::Device &device); - - infinicore::Tensor forward(const InfinilmModel::Input &input) const; - -protected: - INFINICORE_NN_MODULE(Qwen35MoeLanguageModel, language_model); -}; +using Qwen35MoeModel = qwen3_5::Qwen35ModelTemplate; } // namespace infinilm::models::qwen3_5_moe diff --git a/csrc/models/qwen3_next/qwen3_next_for_causal_lm.cpp b/csrc/models/qwen3_next/qwen3_next_for_causal_lm.cpp index 1a9eaeb3b..c1e8e2bfb 100644 --- a/csrc/models/qwen3_next/qwen3_next_for_causal_lm.cpp +++ b/csrc/models/qwen3_next/qwen3_next_for_causal_lm.cpp @@ -34,10 +34,7 @@ void Qwen3NextForCausalLM::reset_cache(const cache::CacheConfig *cache_config) { cache_config_ = cache_config->unique_copy(); auto &forward_context = infinilm::global_state::get_forward_context(); - forward_context.kv_cache_vec.clear(); - forward_context.conv_state_vec.clear(); - forward_context.ssm_state_vec.clear(); - forward_context.mamba_state_pool_size = 0; + forward_context.clear_model_caches(); const backends::AttentionBackend attention_backend = infinilm::global_state::get_infinilm_config().attention_backend; From 2c91d30f06b9c675341c659d379b1893ae5df4c9 Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Mon, 3 Aug 2026 05:15:15 +0000 Subject: [PATCH 6/9] refactor(cache): centralize hybrid cache allocation --- .../hybrid_cache.cpp} | 82 +++++++++---------- csrc/cache/hybrid_cache.hpp | 26 ++++++ csrc/models/qwen3_5/qwen3_5_model.cpp | 4 +- .../qwen3_next_allocate_kv_cache_tensors.hpp | 27 ------ .../qwen3_next/qwen3_next_for_causal_lm.cpp | 4 +- 5 files changed, 71 insertions(+), 72 deletions(-) rename csrc/{models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.cpp => cache/hybrid_cache.cpp} (54%) create mode 100644 csrc/cache/hybrid_cache.hpp delete mode 100644 csrc/models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.hpp diff --git a/csrc/models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.cpp b/csrc/cache/hybrid_cache.cpp similarity index 54% rename from csrc/models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.cpp rename to csrc/cache/hybrid_cache.cpp index 3ba3fdbdc..f945f1e94 100644 --- a/csrc/models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.cpp +++ b/csrc/cache/hybrid_cache.cpp @@ -1,40 +1,40 @@ -#include "qwen3_next_allocate_kv_cache_tensors.hpp" - -#include "../../global_state/global_state.hpp" -#include "../../utils.hpp" -#include "infinicore/context/context.hpp" +#include "hybrid_cache.hpp" #include #include #include #include -namespace infinilm::models::qwen3_next { +namespace infinilm::cache { -AllocatedHybridCache qwen3_next_allocate_cache_tensors( - const cache::CacheConfig *cache_config, - const std::shared_ptr &text_config, +HybridCacheTensors allocate_hybrid_cache_tensors( + const CacheConfig *cache_config, + const std::shared_ptr &model_config, const backends::AttentionBackend &attention_backend) { if (nullptr == cache_config) { return {}; } - if (nullptr == text_config) { - throw std::runtime_error("infinilm::models::qwen3_next::qwen3_next_allocate_kv_cache_tensors: text_config is null"); + if (nullptr == model_config) { + throw std::runtime_error("allocate_hybrid_cache_tensors: model_config is null"); } - const size_t num_hidden_layers = text_config->get("num_hidden_layers"); - const size_t head_dim = text_config->get("head_dim"); - const size_t num_key_value_heads = text_config->get("num_key_value_heads"); - const size_t max_position_embeddings = text_config->get("max_position_embeddings"); - - const size_t linear_conv_kernel_dim = text_config->get("linear_conv_kernel_dim"); - const size_t linear_key_head_dim = text_config->get("linear_key_head_dim"); - const size_t linear_num_key_heads = text_config->get("linear_num_key_heads"); - const size_t linear_num_value_heads = text_config->get("linear_num_value_heads"); - const size_t linear_value_head_dim = text_config->get("linear_value_head_dim"); - - const auto &dtype{text_config->get_dtype()}; - const auto &kv_cache_dtype{text_config->get_kv_cache_dtype()}; - const std::vector layer_types = text_config->get>("layer_types"); + const size_t num_hidden_layers = model_config->get("num_hidden_layers"); + const size_t head_dim = model_config->get("head_dim"); + const size_t num_key_value_heads = model_config->get("num_key_value_heads"); + const size_t max_position_embeddings = model_config->get("max_position_embeddings"); + + const size_t linear_conv_kernel_dim = model_config->get("linear_conv_kernel_dim"); + const size_t linear_key_head_dim = model_config->get("linear_key_head_dim"); + const size_t linear_num_key_heads = model_config->get("linear_num_key_heads"); + const size_t linear_num_value_heads = model_config->get("linear_num_value_heads"); + const size_t linear_value_head_dim = model_config->get("linear_value_head_dim"); + + const auto &dtype{model_config->get_dtype()}; + const auto &kv_cache_dtype{model_config->get_kv_cache_dtype()}; + const std::vector layer_types = model_config->get>("layer_types"); + if (layer_types.size() != num_hidden_layers) { + throw std::runtime_error( + "allocate_hybrid_cache_tensors: layer_types size must match num_hidden_layers"); + } std::vector kv_cache_vec; std::vector conv_state_vec; @@ -49,11 +49,11 @@ AllocatedHybridCache qwen3_next_allocate_cache_tensors( mamba_state_pool_size = pool_size; } else if (mamba_state_pool_size != pool_size) { throw std::runtime_error( - "qwen3_next_allocate_cache_tensors: inconsistent mamba state pool size at layer " + "allocate_hybrid_cache_tensors: inconsistent mamba state pool size at layer " + std::to_string(layer_idx)); } - auto conv_state = cache::MambaCache::create_layer_conv_state( + auto conv_state = MambaCache::create_layer_conv_state( linear_key_head_dim, linear_value_head_dim, linear_num_key_heads, @@ -61,7 +61,7 @@ AllocatedHybridCache qwen3_next_allocate_cache_tensors( linear_conv_kernel_dim, dtype, pool_size); - auto ssm_state = cache::MambaCache::create_layer_ssm_state( + auto ssm_state = MambaCache::create_layer_ssm_state( linear_key_head_dim, linear_value_head_dim, linear_num_key_heads, @@ -74,8 +74,8 @@ AllocatedHybridCache qwen3_next_allocate_cache_tensors( ssm_state_vec.push_back(std::move(ssm_state)); }; - auto allocate_static_full_attention_cache = [&](size_t layer_idx, const cache::StaticKVCacheConfig &config) { - auto kv_cache = cache::StaticKVCache::create_layer_kv_cache( + auto allocate_static_full_attention_cache = [&](size_t layer_idx, const StaticKVCacheConfig &config) { + auto kv_cache = StaticKVCache::create_layer_kv_cache( head_dim, head_dim, num_key_value_heads, @@ -89,8 +89,8 @@ AllocatedHybridCache qwen3_next_allocate_cache_tensors( ssm_state_vec.emplace_back(); }; - auto allocate_paged_full_attention_cache = [&](size_t layer_idx, const cache::PagedKVCacheConfig &config) { - auto kv_cache = cache::PagedKVCache::create_layer_kv_cache( + auto allocate_paged_full_attention_cache = [&](size_t layer_idx, const PagedKVCacheConfig &config) { + auto kv_cache = PagedKVCache::create_layer_kv_cache( head_dim, head_dim, num_key_value_heads, @@ -105,9 +105,9 @@ AllocatedHybridCache qwen3_next_allocate_cache_tensors( switch (attention_backend) { case backends::AttentionBackend::STATIC_ATTN: { - auto static_kv_cache_config = dynamic_cast(cache_config); + auto static_kv_cache_config = dynamic_cast(cache_config); if (nullptr == static_kv_cache_config) { - throw std::runtime_error("infinilm::models::qwen3_next::qwen3_next_allocate_kv_cache_tensors: invalid static kv cache config type"); + throw std::runtime_error("allocate_hybrid_cache_tensors: invalid static kv cache config type"); } for (size_t layer_idx = 0; layer_idx < num_hidden_layers; ++layer_idx) { @@ -117,7 +117,7 @@ AllocatedHybridCache qwen3_next_allocate_cache_tensors( } else if ("full_attention" == layer_type) { allocate_static_full_attention_cache(layer_idx, *static_kv_cache_config); } else { - throw std::runtime_error("infinilm::models::qwen3_next::qwen3_next_allocate_kv_cache_tensors: unsupported layer_type '" + layer_type + "' for layer " + std::to_string(layer_idx)); + throw std::runtime_error("allocate_hybrid_cache_tensors: unsupported layer_type '" + layer_type + "' for layer " + std::to_string(layer_idx)); } } break; @@ -126,9 +126,9 @@ AllocatedHybridCache qwen3_next_allocate_cache_tensors( ; } case backends::AttentionBackend::PAGED_ATTN: { - auto paged_kv_cache_config = dynamic_cast(cache_config); + auto paged_kv_cache_config = dynamic_cast(cache_config); if (nullptr == paged_kv_cache_config) { - throw std::runtime_error("infinilm::models::qwen3_next::qwen3_next_allocate_kv_cache_tensors: invalid paged kv cache config type"); + throw std::runtime_error("allocate_hybrid_cache_tensors: invalid paged kv cache config type"); } const size_t mamba_pool_size = std::max(2, paged_kv_cache_config->num_blocks() / 4); @@ -139,19 +139,19 @@ AllocatedHybridCache qwen3_next_allocate_cache_tensors( } else if ("full_attention" == layer_type) { allocate_paged_full_attention_cache(layer_idx, *paged_kv_cache_config); } else { - throw std::runtime_error("infinilm::models::qwen3_next::qwen3_next_allocate_kv_cache_tensors: unsupported layer_type '" + layer_type + "' for layer " + std::to_string(layer_idx)); + throw std::runtime_error("allocate_hybrid_cache_tensors: unsupported layer_type '" + layer_type + "' for layer " + std::to_string(layer_idx)); } } break; } default: - throw std::runtime_error("infinilm::models::qwen3_next::qwen3_next_allocate_kv_cache_tensors: Unsupported attention backend: " + std::to_string(static_cast(attention_backend))); + throw std::runtime_error("allocate_hybrid_cache_tensors: Unsupported attention backend: " + std::to_string(static_cast(attention_backend))); } - return AllocatedHybridCache{ + return HybridCacheTensors{ std::move(kv_cache_vec), std::move(conv_state_vec), std::move(ssm_state_vec), mamba_state_pool_size}; } -} // namespace infinilm::models::qwen3_next +} // namespace infinilm::cache diff --git a/csrc/cache/hybrid_cache.hpp b/csrc/cache/hybrid_cache.hpp new file mode 100644 index 000000000..ad078d230 --- /dev/null +++ b/csrc/cache/hybrid_cache.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include "../backends/attention_backends.hpp" +#include "../config/model_config.hpp" +#include "kv_cache.hpp" +#include "mamba_cache.hpp" + +#include +#include +#include + +namespace infinilm::cache { + +struct HybridCacheTensors { + std::vector kv_cache_tensors; + std::vector conv_state_tensors; + std::vector ssm_state_tensors; + size_t mamba_state_pool_size{0}; +}; + +HybridCacheTensors allocate_hybrid_cache_tensors( + const CacheConfig *cache_config, + const std::shared_ptr &model_config, + const backends::AttentionBackend &attention_backend); + +} // namespace infinilm::cache diff --git a/csrc/models/qwen3_5/qwen3_5_model.cpp b/csrc/models/qwen3_5/qwen3_5_model.cpp index 969ea5485..754ae3c13 100644 --- a/csrc/models/qwen3_5/qwen3_5_model.cpp +++ b/csrc/models/qwen3_5/qwen3_5_model.cpp @@ -1,7 +1,7 @@ #include "qwen3_5_model.hpp" +#include "../../cache/hybrid_cache.hpp" #include "../../global_state/global_state.hpp" -#include "../qwen3_next/qwen3_next_allocate_kv_cache_tensors.hpp" #include #include @@ -115,7 +115,7 @@ void Qwen35ModelBase::reset_cache(const cache::CacheConfig *cache_config) { const backends::AttentionBackend attention_backend = infinilm::global_state::get_infinilm_config().attention_backend; - auto cache_vectors = infinilm::models::qwen3_next::qwen3_next_allocate_cache_tensors(cache_config, model_config_, attention_backend); + auto cache_vectors = infinilm::cache::allocate_hybrid_cache_tensors(cache_config, model_config_, attention_backend); forward_context.kv_cache_vec = std::move(cache_vectors.kv_cache_tensors); forward_context.conv_state_vec = std::move(cache_vectors.conv_state_tensors); forward_context.ssm_state_vec = std::move(cache_vectors.ssm_state_tensors); diff --git a/csrc/models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.hpp b/csrc/models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.hpp deleted file mode 100644 index 1d5a45dea..000000000 --- a/csrc/models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.hpp +++ /dev/null @@ -1,27 +0,0 @@ -#pragma once - -#include "../../backends/attention_backends.hpp" -#include "../../cache/kv_cache.hpp" -#include "../../cache/mamba_cache.hpp" -#include "../../config/model_config.hpp" - -#include -#include -#include -#include - -namespace infinilm::models::qwen3_next { - -struct AllocatedHybridCache { - std::vector kv_cache_tensors; - std::vector conv_state_tensors; - std::vector ssm_state_tensors; - size_t mamba_state_pool_size{0}; -}; - -AllocatedHybridCache qwen3_next_allocate_cache_tensors( - const cache::CacheConfig *cache_config, - const std::shared_ptr &text_config, - const backends::AttentionBackend &attention_backend); - -} // namespace infinilm::models::qwen3_next diff --git a/csrc/models/qwen3_next/qwen3_next_for_causal_lm.cpp b/csrc/models/qwen3_next/qwen3_next_for_causal_lm.cpp index c1e8e2bfb..9a3525c1d 100644 --- a/csrc/models/qwen3_next/qwen3_next_for_causal_lm.cpp +++ b/csrc/models/qwen3_next/qwen3_next_for_causal_lm.cpp @@ -1,7 +1,7 @@ #include "qwen3_next_for_causal_lm.hpp" +#include "../../cache/hybrid_cache.hpp" #include "../../global_state/global_state.hpp" #include "../models_registry.hpp" -#include "qwen3_next_allocate_kv_cache_tensors.hpp" #include #include #include @@ -38,7 +38,7 @@ void Qwen3NextForCausalLM::reset_cache(const cache::CacheConfig *cache_config) { const backends::AttentionBackend attention_backend = infinilm::global_state::get_infinilm_config().attention_backend; - auto cache_vectors = qwen3_next_allocate_cache_tensors(cache_config, model_config_, attention_backend); + auto cache_vectors = cache::allocate_hybrid_cache_tensors(cache_config, model_config_, attention_backend); forward_context.kv_cache_vec = std::move(cache_vectors.kv_cache_tensors); forward_context.conv_state_vec = std::move(cache_vectors.conv_state_tensors); forward_context.ssm_state_vec = std::move(cache_vectors.ssm_state_tensors); From cce29c0cd05cc02dd0487e80cd6814ea3701af70 Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Mon, 3 Aug 2026 05:19:28 +0000 Subject: [PATCH 7/9] refactor(config): share hybrid model defaults --- csrc/config/hybrid_model_config.cpp | 63 +++++++++++++++++++ csrc/config/hybrid_model_config.hpp | 12 ++++ csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp | 18 +----- .../qwen3_next/qwen3_next_for_causal_lm.cpp | 18 +----- 4 files changed, 80 insertions(+), 31 deletions(-) create mode 100644 csrc/config/hybrid_model_config.cpp create mode 100644 csrc/config/hybrid_model_config.hpp diff --git a/csrc/config/hybrid_model_config.cpp b/csrc/config/hybrid_model_config.cpp new file mode 100644 index 000000000..d1bb6c209 --- /dev/null +++ b/csrc/config/hybrid_model_config.cpp @@ -0,0 +1,63 @@ +#include "hybrid_model_config.hpp" + +#include +#include +#include +#include + +namespace infinilm::config { + +void prepare_hybrid_model_config( + const std::shared_ptr &model_config) { + if (model_config == nullptr) { + throw std::runtime_error( + "prepare_hybrid_model_config: model_config is null"); + } + + auto &config_json = model_config->get_config_json(); + const size_t num_hidden_layers = model_config->get("num_hidden_layers"); + + if (!config_json.contains("layer_types")) { + const size_t full_attention_interval = model_config->get("full_attention_interval"); + if (full_attention_interval == 0) { + throw std::runtime_error( + "prepare_hybrid_model_config: full_attention_interval must be positive"); + } + + std::vector layer_types; + layer_types.reserve(num_hidden_layers); + for (size_t layer_idx = 0; layer_idx < num_hidden_layers; ++layer_idx) { + layer_types.push_back( + (layer_idx + 1) % full_attention_interval == 0 + ? "full_attention" + : "linear_attention"); + } + config_json["layer_types"] = std::move(layer_types); + } + + const auto &layer_types = config_json["layer_types"]; + if (!layer_types.is_array() + || layer_types.size() != num_hidden_layers) { + throw std::runtime_error( + "prepare_hybrid_model_config: layer_types size must match num_hidden_layers"); + } + for (size_t layer_idx = 0; layer_idx < num_hidden_layers; ++layer_idx) { + if (!layer_types[layer_idx].is_string()) { + throw std::runtime_error( + "prepare_hybrid_model_config: layer_types entries must be strings"); + } + const auto &layer_type = layer_types[layer_idx].get_ref(); + if (layer_type != "full_attention" + && layer_type != "linear_attention") { + throw std::runtime_error( + "prepare_hybrid_model_config: unsupported layer_type '" + + layer_type + "' at layer " + std::to_string(layer_idx)); + } + } + + if (!config_json.contains("attention_bias")) { + config_json["attention_bias"] = false; + } +} + +} // namespace infinilm::config diff --git a/csrc/config/hybrid_model_config.hpp b/csrc/config/hybrid_model_config.hpp new file mode 100644 index 000000000..745340c9c --- /dev/null +++ b/csrc/config/hybrid_model_config.hpp @@ -0,0 +1,12 @@ +#pragma once + +#include "model_config.hpp" + +#include + +namespace infinilm::config { + +void prepare_hybrid_model_config( + const std::shared_ptr &model_config); + +} // namespace infinilm::config diff --git a/csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp b/csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp index 07836eab7..5c0b68beb 100644 --- a/csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp +++ b/csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp @@ -1,9 +1,10 @@ #include "qwen3_5_for_causal_lm.hpp" +#include "../../config/hybrid_model_config.hpp" + #include "../models_registry.hpp" #include #include -#include namespace infinilm::models::qwen3_5 { @@ -40,20 +41,7 @@ std::shared_ptr prepare_qwen3_5_model_config(std: if (!config_json.contains("partial_rotary_factor") && config_json.contains("rope_parameters") && config_json["rope_parameters"].is_object() && config_json["rope_parameters"].contains("partial_rotary_factor")) { config_json["partial_rotary_factor"] = config_json["rope_parameters"]["partial_rotary_factor"]; } - if (!config_json.contains("layer_types")) { - size_t full_attention_interval = model_config->get("full_attention_interval"); - size_t num_hidden_layers = model_config->get("num_hidden_layers"); - std::vector layer_types; - layer_types.reserve(num_hidden_layers); - for (size_t i = 0; i < num_hidden_layers; i++) { - layer_types.push_back(bool((i + 1) % full_attention_interval) ? "linear_attention" : "full_attention"); - } - config_json["layer_types"] = layer_types; - } - - if (!config_json.contains("attention_bias")) { - config_json["attention_bias"] = false; - } + infinilm::config::prepare_hybrid_model_config(model_config); return model_config; } diff --git a/csrc/models/qwen3_next/qwen3_next_for_causal_lm.cpp b/csrc/models/qwen3_next/qwen3_next_for_causal_lm.cpp index 9a3525c1d..de23d3ccd 100644 --- a/csrc/models/qwen3_next/qwen3_next_for_causal_lm.cpp +++ b/csrc/models/qwen3_next/qwen3_next_for_causal_lm.cpp @@ -1,11 +1,11 @@ #include "qwen3_next_for_causal_lm.hpp" #include "../../cache/hybrid_cache.hpp" +#include "../../config/hybrid_model_config.hpp" #include "../../global_state/global_state.hpp" #include "../models_registry.hpp" #include #include #include -#include namespace infinilm::models::qwen3_next { @@ -51,21 +51,7 @@ std::shared_ptr create_qwen3_next_model_config(st throw std::runtime_error("infinilm::models::qwen3_next::create_qwen3_next_model_config: model_type is not qwen3_next"); } - nlohmann::json &config_json = model_config->get_config_json(); - if (!config_json.contains("layer_types")) { - size_t full_attention_interval = model_config->get("full_attention_interval"); - size_t num_hidden_layers = model_config->get("num_hidden_layers"); - std::vector layer_types; - layer_types.reserve(num_hidden_layers); - for (size_t i = 0; i < num_hidden_layers; i++) { - layer_types.push_back(bool((i + 1) % full_attention_interval) ? "linear_attention" : "full_attention"); - } - config_json["layer_types"] = layer_types; - } - - if (!config_json.contains("attention_bias")) { - config_json["attention_bias"] = false; - } + infinilm::config::prepare_hybrid_model_config(model_config); return model_config; } From ce8761d4dd084ff40a74cbf6e8844b8d8a1e89c5 Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Mon, 3 Aug 2026 05:22:31 +0000 Subject: [PATCH 8/9] refactor(qwen3-next): reuse text causal LM template --- .../qwen3_next/qwen3_next_for_causal_lm.cpp | 17 ----------------- .../qwen3_next/qwen3_next_for_causal_lm.hpp | 16 ++++++---------- 2 files changed, 6 insertions(+), 27 deletions(-) diff --git a/csrc/models/qwen3_next/qwen3_next_for_causal_lm.cpp b/csrc/models/qwen3_next/qwen3_next_for_causal_lm.cpp index de23d3ccd..e440833db 100644 --- a/csrc/models/qwen3_next/qwen3_next_for_causal_lm.cpp +++ b/csrc/models/qwen3_next/qwen3_next_for_causal_lm.cpp @@ -9,23 +9,6 @@ namespace infinilm::models::qwen3_next { -Qwen3NextForCausalLM::Qwen3NextForCausalLM(std::shared_ptr model_config, - const infinicore::Device &device) { - model_config_ = model_config; - size_t hidden_size = model_config->get("hidden_size"); - size_t vocab_size = model_config->get("vocab_size"); - const auto &dtype{model_config->get_dtype()}; - - INFINICORE_NN_MODULE_INIT(model, model_config, device); - INFINICORE_NN_MODULE_INIT(lm_head, hidden_size, vocab_size, false, dtype, device); -} - -infinilm::InfinilmModel::Output Qwen3NextForCausalLM::forward(const infinilm::InfinilmModel::Input &input) const { - auto hidden_states = model_->forward(input); - auto logits = lm_head_->forward(hidden_states); - return {logits}; -} - void Qwen3NextForCausalLM::reset_cache(const cache::CacheConfig *cache_config) { if (nullptr == cache_config) { InfinilmModel::reset_cache(nullptr); diff --git a/csrc/models/qwen3_next/qwen3_next_for_causal_lm.hpp b/csrc/models/qwen3_next/qwen3_next_for_causal_lm.hpp index 0cbe45320..1ee3fba00 100644 --- a/csrc/models/qwen3_next/qwen3_next_for_causal_lm.hpp +++ b/csrc/models/qwen3_next/qwen3_next_for_causal_lm.hpp @@ -1,25 +1,21 @@ #pragma once +#include "../../layers/causal_lm_templates/text_causal_lm.hpp" #include "qwen3_next_decoderLayer.hpp" + #include -#include namespace infinilm::models::qwen3_next { using Qwen3NextModel = infinilm::layers::causal_lm_templates::TextModel; -class Qwen3NextForCausalLM : public InfinilmModel { +class Qwen3NextForCausalLM + : public infinilm::layers::causal_lm_templates::TextCausalLM { public: - Qwen3NextForCausalLM(std::shared_ptr model_config, - const infinicore::Device &device); - - Output forward(const Input &input) const override; + using Base = infinilm::layers::causal_lm_templates::TextCausalLM; + using Base::Base; void reset_cache(const cache::CacheConfig *cache_config) override; - -protected: - INFINICORE_NN_MODULE(Qwen3NextModel, model); - INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, lm_head); }; std::shared_ptr create_qwen3_next_model_config(std::shared_ptr model_config); From a33beaf6ebc19be085885970d3ac2c929c967f17 Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Tue, 4 Aug 2026 02:34:26 +0000 Subject: [PATCH 9/9] refactor(qwen3.5): align shared MoE with model family --- .../moe/shared_expert_sparse_moe_block.hpp | 32 -------- .../qwen3_5_moe/qwen3_5_moe_decoder_layer.hpp | 15 ---- .../qwen3_5_moe/qwen3_5_moe_for_causal_lm.hpp | 10 ++- csrc/models/qwen3_5_moe/qwen3_5_moe_model.hpp | 11 --- .../qwen3_5_moe_sparse_moe_block.hpp | 9 --- .../qwen3_next_sparse_moe_block.cpp} | 18 ++--- .../qwen3_next_sparse_moe_block.hpp | 26 +++++- test/models/qwen3_5_moe/test_adaptation.py | 81 +++++++++++++++++++ 8 files changed, 123 insertions(+), 79 deletions(-) delete mode 100644 csrc/layers/moe/shared_expert_sparse_moe_block.hpp delete mode 100644 csrc/models/qwen3_5_moe/qwen3_5_moe_decoder_layer.hpp delete mode 100644 csrc/models/qwen3_5_moe/qwen3_5_moe_model.hpp delete mode 100644 csrc/models/qwen3_5_moe/qwen3_5_moe_sparse_moe_block.hpp rename csrc/{layers/moe/shared_expert_sparse_moe_block.cpp => models/qwen3_next/qwen3_next_sparse_moe_block.cpp} (73%) create mode 100644 test/models/qwen3_5_moe/test_adaptation.py diff --git a/csrc/layers/moe/shared_expert_sparse_moe_block.hpp b/csrc/layers/moe/shared_expert_sparse_moe_block.hpp deleted file mode 100644 index 8ad23b3de..000000000 --- a/csrc/layers/moe/shared_expert_sparse_moe_block.hpp +++ /dev/null @@ -1,32 +0,0 @@ -#pragma once - -#include "sparse_moe_block.hpp" - -#include "../../config/model_config.hpp" -#include "../linear/linear.hpp" -#include "../mlp/mlp.hpp" -#include "infinicore/nn/module.hpp" - -#include -#include - -namespace infinilm::layers::moe { - -class SharedExpertSparseMoeBlock : public SparseMoeBlock { -public: - SharedExpertSparseMoeBlock( - std::shared_ptr model_config, - const infinicore::Device &device); - SharedExpertSparseMoeBlock( - std::shared_ptr model_config, - size_t layer_idx, - const infinicore::Device &device); - - infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const; - -protected: - INFINICORE_NN_MODULE(infinilm::layers::mlp::MLP, shared_expert); - INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, shared_expert_gate); -}; - -} // namespace infinilm::layers::moe diff --git a/csrc/models/qwen3_5_moe/qwen3_5_moe_decoder_layer.hpp b/csrc/models/qwen3_5_moe/qwen3_5_moe_decoder_layer.hpp deleted file mode 100644 index eee84a0b5..000000000 --- a/csrc/models/qwen3_5_moe/qwen3_5_moe_decoder_layer.hpp +++ /dev/null @@ -1,15 +0,0 @@ -#pragma once - -#include "../../layers/causal_lm_templates/hybrid_decoder_layer.hpp" -#include "../qwen3_5/qwen3_5_attention.hpp" -#include "../qwen3_next/qwen3_next_gated_deltanet.hpp" -#include "qwen3_5_moe_sparse_moe_block.hpp" - -namespace infinilm::models::qwen3_5_moe { - -using Qwen35MoeDecoderLayer = infinilm::layers::causal_lm_templates::HybridDecoderLayer< - qwen3_5::Qwen35Attention, - qwen3_next::Qwen3NextGatedDeltaNet, - Qwen35MoeSparseMoeBlock>; - -} // namespace infinilm::models::qwen3_5_moe diff --git a/csrc/models/qwen3_5_moe/qwen3_5_moe_for_causal_lm.hpp b/csrc/models/qwen3_5_moe/qwen3_5_moe_for_causal_lm.hpp index 19babff67..130eb3543 100644 --- a/csrc/models/qwen3_5_moe/qwen3_5_moe_for_causal_lm.hpp +++ b/csrc/models/qwen3_5_moe/qwen3_5_moe_for_causal_lm.hpp @@ -1,12 +1,20 @@ #pragma once +#include "../../layers/causal_lm_templates/hybrid_decoder_layer.hpp" #include "../qwen3_5/qwen3_5_for_causal_lm.hpp" -#include "qwen3_5_moe_model.hpp" +#include "../qwen3_next/qwen3_next_gated_deltanet.hpp" +#include "../qwen3_next/qwen3_next_sparse_moe_block.hpp" #include namespace infinilm::models::qwen3_5_moe { +using Qwen35MoeDecoderLayer = infinilm::layers::causal_lm_templates::HybridDecoderLayer< + qwen3_5::Qwen35Attention, + qwen3_next::Qwen3NextGatedDeltaNet, + qwen3_next::Qwen3NextSparseMoeBlock>; +using Qwen35MoeLanguageModel = infinilm::layers::causal_lm_templates::TextModel; +using Qwen35MoeModel = qwen3_5::Qwen35ModelTemplate; using Qwen35MoeForConditionalGeneration = qwen3_5::Qwen35CausalLM; std::shared_ptr create_qwen3_5_moe_model_config( diff --git a/csrc/models/qwen3_5_moe/qwen3_5_moe_model.hpp b/csrc/models/qwen3_5_moe/qwen3_5_moe_model.hpp deleted file mode 100644 index 8e10cc349..000000000 --- a/csrc/models/qwen3_5_moe/qwen3_5_moe_model.hpp +++ /dev/null @@ -1,11 +0,0 @@ -#pragma once - -#include "../qwen3_5/qwen3_5_model.hpp" -#include "qwen3_5_moe_decoder_layer.hpp" - -namespace infinilm::models::qwen3_5_moe { - -using Qwen35MoeLanguageModel = infinilm::layers::causal_lm_templates::TextModel; -using Qwen35MoeModel = qwen3_5::Qwen35ModelTemplate; - -} // namespace infinilm::models::qwen3_5_moe diff --git a/csrc/models/qwen3_5_moe/qwen3_5_moe_sparse_moe_block.hpp b/csrc/models/qwen3_5_moe/qwen3_5_moe_sparse_moe_block.hpp deleted file mode 100644 index 42bad48b3..000000000 --- a/csrc/models/qwen3_5_moe/qwen3_5_moe_sparse_moe_block.hpp +++ /dev/null @@ -1,9 +0,0 @@ -#pragma once - -#include "../../layers/moe/shared_expert_sparse_moe_block.hpp" - -namespace infinilm::models::qwen3_5_moe { - -using Qwen35MoeSparseMoeBlock = infinilm::layers::moe::SharedExpertSparseMoeBlock; - -} // namespace infinilm::models::qwen3_5_moe diff --git a/csrc/layers/moe/shared_expert_sparse_moe_block.cpp b/csrc/models/qwen3_next/qwen3_next_sparse_moe_block.cpp similarity index 73% rename from csrc/layers/moe/shared_expert_sparse_moe_block.cpp rename to csrc/models/qwen3_next/qwen3_next_sparse_moe_block.cpp index 696e57c99..dc29fbebb 100644 --- a/csrc/layers/moe/shared_expert_sparse_moe_block.cpp +++ b/csrc/models/qwen3_next/qwen3_next_sparse_moe_block.cpp @@ -1,23 +1,23 @@ -#include "shared_expert_sparse_moe_block.hpp" +#include "qwen3_next_sparse_moe_block.hpp" #include "infinicore/ops.hpp" #include "infinicore/ops/mul.hpp" #include -namespace infinilm::layers::moe { +namespace infinilm::models::qwen3_next { -SharedExpertSparseMoeBlock::SharedExpertSparseMoeBlock( +Qwen3NextSparseMoeBlock::Qwen3NextSparseMoeBlock( std::shared_ptr model_config, const infinicore::Device &device) - : SharedExpertSparseMoeBlock(model_config, 0, device) { + : Qwen3NextSparseMoeBlock(model_config, 0, device) { } -SharedExpertSparseMoeBlock::SharedExpertSparseMoeBlock( +Qwen3NextSparseMoeBlock::Qwen3NextSparseMoeBlock( std::shared_ptr model_config, size_t layer_idx, const infinicore::Device &device) - : SparseMoeBlock(model_config, device, layer_idx) { + : infinilm::layers::moe::SparseMoeBlock(model_config, device, layer_idx) { auto shared_config_json = model_config->get_config_json(); shared_config_json["intermediate_size"] = model_config->get("shared_expert_intermediate_size"); auto shared_config = std::make_shared( @@ -33,9 +33,9 @@ SharedExpertSparseMoeBlock::SharedExpertSparseMoeBlock( device); } -infinicore::Tensor SharedExpertSparseMoeBlock::forward( +infinicore::Tensor Qwen3NextSparseMoeBlock::forward( const infinicore::Tensor &hidden_states) const { - auto routed_output = SparseMoeBlock::forward(hidden_states); + auto routed_output = infinilm::layers::moe::SparseMoeBlock::forward(hidden_states); auto shared_output = shared_expert_->forward(hidden_states); auto shared_gate_input = hidden_states; @@ -49,4 +49,4 @@ infinicore::Tensor SharedExpertSparseMoeBlock::forward( return infinicore::op::add(routed_output, shared_output); } -} // namespace infinilm::layers::moe +} // namespace infinilm::models::qwen3_next diff --git a/csrc/models/qwen3_next/qwen3_next_sparse_moe_block.hpp b/csrc/models/qwen3_next/qwen3_next_sparse_moe_block.hpp index dd302355a..8cff0848e 100644 --- a/csrc/models/qwen3_next/qwen3_next_sparse_moe_block.hpp +++ b/csrc/models/qwen3_next/qwen3_next_sparse_moe_block.hpp @@ -1,9 +1,31 @@ #pragma once -#include "../../layers/moe/shared_expert_sparse_moe_block.hpp" +#include "../../config/model_config.hpp" +#include "../../layers/linear/linear.hpp" +#include "../../layers/mlp/mlp.hpp" +#include "../../layers/moe/sparse_moe_block.hpp" +#include "infinicore/nn/module.hpp" + +#include +#include namespace infinilm::models::qwen3_next { -using Qwen3NextSparseMoeBlock = infinilm::layers::moe::SharedExpertSparseMoeBlock; +class Qwen3NextSparseMoeBlock : public infinilm::layers::moe::SparseMoeBlock { +public: + Qwen3NextSparseMoeBlock( + std::shared_ptr model_config, + const infinicore::Device &device); + Qwen3NextSparseMoeBlock( + std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const; + +protected: + INFINICORE_NN_MODULE(infinilm::layers::mlp::MLP, shared_expert); + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, shared_expert_gate); +}; } // namespace infinilm::models::qwen3_next diff --git a/test/models/qwen3_5_moe/test_adaptation.py b/test/models/qwen3_5_moe/test_adaptation.py new file mode 100644 index 000000000..e7294fa02 --- /dev/null +++ b/test/models/qwen3_5_moe/test_adaptation.py @@ -0,0 +1,81 @@ +import unittest + +import torch + +from infinilm.infer_engine import _infer_position_id_axes +from infinilm.modeling_utils import _remap_qwen3_5_moe + + +class PositionIdAxesTest(unittest.TestCase): + def test_defaults_to_one_axis(self): + self.assertEqual(_infer_position_id_axes({"text_config": {}}), 1) + + def test_infers_axes_from_mrope_section(self): + config = { + "text_config": { + "rope_parameters": {"mrope_section": [11, 11, 10]} + } + } + self.assertEqual(_infer_position_id_axes(config), 3) + + def test_explicit_axes_take_precedence(self): + config = { + "position_id_axes": 2, + "text_config": { + "position_id_axes": 4, + "rope_parameters": {"mrope_section": [11, 11, 10]}, + }, + } + self.assertEqual(_infer_position_id_axes(config), 4) + + def test_rejects_non_positive_axes(self): + with self.assertRaisesRegex(ValueError, "must be positive"): + _infer_position_id_axes({"text_config": {"position_id_axes": 0}}) + + +class Qwen35MoeWeightRemapTest(unittest.TestCase): + def setUp(self): + self.config = { + "text_config": { + "linear_key_head_dim": 2, + "linear_num_key_heads": 1, + "num_experts": 2, + "moe_intermediate_size": 3, + } + } + + def test_splits_packed_expert_weights(self): + gate_up = torch.arange(2 * 6 * 4).reshape(2, 6, 4) + down = torch.arange(2 * 4 * 3).reshape(2, 4, 3) + state_dict = { + "model.language_model.layers.0.mlp.experts.gate_up_proj": gate_up, + "model.language_model.layers.0.mlp.experts.down_proj": down, + } + + remapped = _remap_qwen3_5_moe(state_dict, self.config) + + prefix = "model.language_model.layers.0.mlp.experts." + self.assertTrue( + torch.equal(remapped[f"{prefix}0.gate_proj.weight"], gate_up[0, :3]) + ) + self.assertTrue( + torch.equal(remapped[f"{prefix}0.up_proj.weight"], gate_up[0, 3:]) + ) + self.assertTrue( + torch.equal(remapped[f"{prefix}1.down_proj.weight"], down[1]) + ) + self.assertNotIn(f"{prefix}gate_up_proj", remapped) + self.assertNotIn(f"{prefix}down_proj", remapped) + + def test_rejects_wrong_expert_count(self): + state_dict = { + "model.language_model.layers.0.mlp.experts.gate_up_proj": torch.zeros( + 1, 6, 4 + ) + } + with self.assertRaisesRegex(ValueError, "Expected 2 experts"): + _remap_qwen3_5_moe(state_dict, self.config) + + +if __name__ == "__main__": + unittest.main()