diff --git a/csrc/models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.cpp b/csrc/cache/hybrid_cache.cpp similarity index 51% rename from csrc/models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.cpp rename to csrc/cache/hybrid_cache.cpp index 3aaf898cb..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; @@ -43,8 +43,17 @@ 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) { - auto conv_state = cache::MambaCache::create_layer_conv_state( + if (mamba_state_pool_size == 0) { + mamba_state_pool_size = pool_size; + } else if (mamba_state_pool_size != pool_size) { + throw std::runtime_error( + "allocate_hybrid_cache_tensors: inconsistent mamba state pool size at layer " + + std::to_string(layer_idx)); + } + + auto conv_state = MambaCache::create_layer_conv_state( linear_key_head_dim, linear_value_head_dim, linear_num_key_heads, @@ -52,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, @@ -65,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, @@ -80,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, @@ -96,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) { @@ -108,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; @@ -117,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); @@ -130,18 +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)}; + 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/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/engine/compiler/paged_compiler.cpp b/csrc/engine/compiler/paged_compiler.cpp index df3fd1cb4..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,8 +37,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); - + 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) { + 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(), + 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 +75,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()); + // Models declare their position-id axes explicitly. Single-axis + // models retain the traditional [b] layout. + input.position_ids = infinicore::Tensor::empty( + 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()); set_zeros(input.input_ids.value()); set_zeros(input.position_ids.value()); @@ -126,7 +153,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/csrc/global_state/forward_context.hpp b/csrc/global_state/forward_context.hpp index f395b531a..ed7de761a 100644 --- a/csrc/global_state/forward_context.hpp +++ b/csrc/global_state/forward_context.hpp @@ -62,6 +62,14 @@ 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 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/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/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/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_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/qwen3_5_for_causal_lm.cpp b/csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp index aece19fda..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,46 +1,14 @@ #include "qwen3_5_for_causal_lm.hpp" -#include "../../global_state/global_state.hpp" +#include "../../config/hybrid_model_config.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 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"]; @@ -53,30 +21,36 @@ std::shared_ptr create_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")) { 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; - } + infinilm::config::prepare_hybrid_model_config(model_config); + return model_config; +} - if (!config_json.contains("attention_bias")) { - config_json["attention_bias"] = false; +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 model_config; + return prepare_qwen3_5_model_config(model_config); } } // namespace infinilm::models::qwen3_5 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..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,24 +1,31 @@ #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; + 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); + } +}; - void reset_cache(const cache::CacheConfig *cache_config) override; +using Qwen35ForCausalLM = Qwen35CausalLM; -protected: - INFINICORE_NN_MODULE(Qwen35Model, model); - 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); diff --git a/csrc/models/qwen3_5/qwen3_5_model.cpp b/csrc/models/qwen3_5/qwen3_5_model.cpp index 82c8b0d4d..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 @@ -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,10 @@ 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); } - 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; } @@ -107,31 +106,20 @@ void Qwen35Model::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 Qwen35Model::reset_cache(const cache::CacheConfig *cache_config) { +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(); - 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); + forward_context.mamba_state_pool_size = cache_vectors.mamba_state_pool_size; } } // namespace infinilm::models::qwen3_5 diff --git a/csrc/models/qwen3_5/qwen3_5_model.hpp b/csrc/models/qwen3_5/qwen3_5_model.hpp index bea1b78ee..e9f2080bb 100644 --- a/csrc/models/qwen3_5/qwen3_5_model.hpp +++ b/csrc/models/qwen3_5/qwen3_5_model.hpp @@ -10,24 +10,46 @@ 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_; }; +template +class Qwen35ModelTemplate : public Qwen35ModelBase { +public: + 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(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 new file mode 100644 index 000000000..99dc16d92 --- /dev/null +++ b/csrc/models/qwen3_5_moe/qwen3_5_moe_for_causal_lm.cpp @@ -0,0 +1,34 @@ +#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 { + +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..130eb3543 --- /dev/null +++ b/csrc/models/qwen3_5_moe/qwen3_5_moe_for_causal_lm.hpp @@ -0,0 +1,23 @@ +#pragma once + +#include "../../layers/causal_lm_templates/hybrid_decoder_layer.hpp" +#include "../qwen3_5/qwen3_5_for_causal_lm.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( + std::shared_ptr model_config); + +} // namespace infinilm::models::qwen3_5_moe 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 a4b5190f1..000000000 --- a/csrc/models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.hpp +++ /dev/null @@ -1,26 +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; -}; - -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_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_for_causal_lm.cpp b/csrc/models/qwen3_next/qwen3_next_for_causal_lm.cpp index 1b6354540..e440833db 100644 --- a/csrc/models/qwen3_next/qwen3_next_for_causal_lm.cpp +++ b/csrc/models/qwen3_next/qwen3_next_for_causal_lm.cpp @@ -1,31 +1,14 @@ #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 "qwen3_next_allocate_kv_cache_tensors.hpp" #include #include #include -#include 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); @@ -34,16 +17,15 @@ 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.clear_model_caches(); 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); + 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) { @@ -52,21 +34,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; } 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); diff --git a/csrc/models/qwen3_next/qwen3_next_sparse_moe_block.cpp b/csrc/models/qwen3_next/qwen3_next_sparse_moe_block.cpp index d3548c9fa..dc29fbebb 100644 --- a/csrc/models/qwen3_next/qwen3_next_sparse_moe_block.cpp +++ b/csrc/models/qwen3_next/qwen3_next_sparse_moe_block.cpp @@ -1,72 +1,31 @@ #include "qwen3_next_sparse_moe_block.hpp" -#include "../../global_state/global_state.hpp" +#include "infinicore/ops.hpp" +#include "infinicore/ops/mul.hpp" -#include -#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::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", +Qwen3NextSparseMoeBlock::Qwen3NextSparseMoeBlock( + std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device) + : 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( + 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, @@ -74,33 +33,20 @@ Qwen3NextSparseMoeBlock::Qwen3NextSparseMoeBlock(std::shared_ptrndim() == 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()); +infinicore::Tensor Qwen3NextSparseMoeBlock::forward( + const infinicore::Tensor &hidden_states) const { + auto routed_output = infinilm::layers::moe::SparseMoeBlock::forward(hidden_states); - 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 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); - 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); + return infinicore::op::add(routed_output, shared_output); } } // 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..8cff0848e 100644 --- a/csrc/models/qwen3_next/qwen3_next_sparse_moe_block.hpp +++ b/csrc/models/qwen3_next/qwen3_next_sparse_moe_block.hpp @@ -2,42 +2,30 @@ #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 "../../layers/mlp/mlp.hpp" +#include "../../layers/moe/sparse_moe_block.hpp" +#include "infinicore/nn/module.hpp" +#include #include namespace infinilm::models::qwen3_next { -class Qwen3NextSharedExpert : public infinicore::nn::Module { +class Qwen3NextSparseMoeBlock : public infinilm::layers::moe::SparseMoeBlock { public: - Qwen3NextSharedExpert(std::shared_ptr model_config, - const infinicore::Device &device); + 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_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_; + 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/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/infer_engine.py b/python/infinilm/infer_engine.py index 270af75ec..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() @@ -516,9 +542,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.position_id_axes > 1: + position_ids_list = [position_ids_list] * self.position_id_axes 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: 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): 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()