From 422920689ac0d9a5878cd1b39760ac8f58347b17 Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Fri, 31 Jul 2026 03:07:35 +0000 Subject: [PATCH 1/2] feat(qwen3.6): support dense and MoE models --- .../hybrid_cache.cpp} | 92 ++++++++------ csrc/cache/hybrid_cache.hpp | 26 ++++ csrc/config/hybrid_model_config.cpp | 63 ++++++++++ csrc/config/hybrid_model_config.hpp | 12 ++ csrc/engine/compiler/paged_compiler.cpp | 52 ++++++-- csrc/global_state/forward_context.hpp | 8 ++ .../hybrid_decoder_layer.hpp | 117 ++++++++++++++++++ csrc/models/infinilm_model.cpp | 8 +- csrc/models/infinilm_model.hpp | 3 + csrc/models/qwen3_5/qwen3_5_decoderLayer.cpp | 66 ---------- csrc/models/qwen3_5/qwen3_5_decoderLayer.hpp | 34 +---- csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp | 76 ++++-------- csrc/models/qwen3_5/qwen3_5_for_causal_lm.hpp | 27 ++-- csrc/models/qwen3_5/qwen3_5_model.cpp | 32 ++--- csrc/models/qwen3_5/qwen3_5_model.hpp | 40 ++++-- .../qwen3_5_moe/qwen3_5_moe_for_causal_lm.cpp | 34 +++++ .../qwen3_5_moe/qwen3_5_moe_for_causal_lm.hpp | 23 ++++ .../qwen3_next_allocate_kv_cache_tensors.hpp | 26 ---- .../qwen3_next/qwen3_next_decoderLayer.cpp | 66 ---------- .../qwen3_next/qwen3_next_decoderLayer.hpp | 33 +---- .../qwen3_next/qwen3_next_for_causal_lm.cpp | 44 +------ .../qwen3_next/qwen3_next_for_causal_lm.hpp | 16 +-- .../qwen3_next_sparse_moe_block.cpp | 116 +++++------------ .../qwen3_next_sparse_moe_block.hpp | 40 +++--- python/infinilm/infer_engine.py | 34 ++++- python/infinilm/modeling_utils.py | 55 ++++++++ python/infinilm/processors/__init__.py | 13 +- .../infinilm/processors/qwen3_5_processor.py | 1 + test/models/qwen3_5_moe/test_adaptation.py | 81 ++++++++++++ 29 files changed, 709 insertions(+), 529 deletions(-) rename csrc/{models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.cpp => cache/hybrid_cache.cpp} (51%) create mode 100644 csrc/cache/hybrid_cache.hpp create mode 100644 csrc/config/hybrid_model_config.cpp create mode 100644 csrc/config/hybrid_model_config.hpp create mode 100644 csrc/layers/causal_lm_templates/hybrid_decoder_layer.hpp delete mode 100644 csrc/models/qwen3_5/qwen3_5_decoderLayer.cpp 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 delete mode 100644 csrc/models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.hpp delete mode 100644 csrc/models/qwen3_next/qwen3_next_decoderLayer.cpp create mode 100644 test/models/qwen3_5_moe/test_adaptation.py 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 3aaf898c..f945f1e9 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 00000000..ad078d23 --- /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 00000000..d1bb6c20 --- /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 00000000..745340c9 --- /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 8f7f789b..623574aa 100644 --- a/csrc/engine/compiler/paged_compiler.cpp +++ b/csrc/engine/compiler/paged_compiler.cpp @@ -11,16 +11,7 @@ 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 @@ -61,8 +52,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()); @@ -71,7 +90,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()); @@ -142,7 +168,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 f395b531..ed7de761 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 00000000..345461ee --- /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 5d284a31..6751cefa 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 ac994fd6..27275ee2 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 70964bb6..00000000 --- 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 751586bb..d1223cb0 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 aece19fd..5c0b68be 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 a7c751e3..2aaaf4a1 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 82c8b0d4..754ae3c1 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 bea1b78e..e9f2080b 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 00000000..99dc16d9 --- /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 00000000..130eb354 --- /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 a4b5190f..00000000 --- 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 4c61832c..00000000 --- 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 dd0505bd..df7f4c36 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 1b635454..e440833d 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 0cbe4532..1ee3fba0 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 d3548c9f..dc29fbeb 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 2cec0bd5..8cff0848 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/python/infinilm/infer_engine.py b/python/infinilm/infer_engine.py index 00dae220..11bcdaa0 100644 --- a/python/infinilm/infer_engine.py +++ b/python/infinilm/infer_engine.py @@ -124,6 +124,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, @@ -144,6 +165,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() @@ -531,9 +557,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 9bcefa7f..486df03b 100644 --- a/python/infinilm/modeling_utils.py +++ b/python/infinilm/modeling_utils.py @@ -985,6 +985,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, @@ -994,5 +1048,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 3d5cad4d..d4ab5b3f 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 e550de5f..6b8fee90 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 00000000..e7294fa0 --- /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() From f9ac008461c2b82e33b6b9412c4cec11d2d69009 Mon Sep 17 00:00:00 2001 From: wooway777 Date: Fri, 14 Aug 2026 22:29:26 +0800 Subject: [PATCH 2/2] fix(qwen3.5): support 122B MoE with TP8 --- csrc/cache/hybrid_cache.hpp | 26 ---- csrc/config/hybrid_model_config.cpp | 63 ---------- csrc/config/hybrid_model_config.hpp | 12 -- csrc/engine/compiler/paged_compiler.cpp | 42 +++---- csrc/global_state/forward_context.hpp | 8 -- .../hybrid_decoder_layer.hpp | 117 ------------------ csrc/models/infinilm_model.cpp | 8 +- csrc/models/qwen3_5/qwen3_5_attention.cpp | 7 +- csrc/models/qwen3_5/qwen3_5_decoderLayer.cpp | 66 ++++++++++ csrc/models/qwen3_5/qwen3_5_decoderLayer.hpp | 34 ++++- csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp | 52 +++++++- csrc/models/qwen3_5/qwen3_5_for_causal_lm.hpp | 30 ++--- csrc/models/qwen3_5/qwen3_5_model.cpp | 10 +- .../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.hpp | 26 ++-- .../qwen3_next_allocate_kv_cache_tensors.cpp} | 92 ++++++-------- .../qwen3_next_allocate_kv_cache_tensors.hpp | 26 ++++ .../qwen3_next/qwen3_next_decoderLayer.cpp | 66 ++++++++++ .../qwen3_next/qwen3_next_decoderLayer.hpp | 33 ++++- .../qwen3_next/qwen3_next_for_causal_lm.cpp | 44 ++++++- .../qwen3_next/qwen3_next_for_causal_lm.hpp | 16 ++- .../qwen3_next_sparse_moe_block.cpp | 116 ++++++++++++----- .../qwen3_next_sparse_moe_block.hpp | 40 +++--- python/infinilm/processors/__init__.py | 6 +- test/models/qwen3_5_moe/test_adaptation.py | 11 +- 26 files changed, 640 insertions(+), 422 deletions(-) delete mode 100644 csrc/cache/hybrid_cache.hpp delete mode 100644 csrc/config/hybrid_model_config.cpp delete mode 100644 csrc/config/hybrid_model_config.hpp delete mode 100644 csrc/layers/causal_lm_templates/hybrid_decoder_layer.hpp create mode 100644 csrc/models/qwen3_5/qwen3_5_decoderLayer.cpp 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 rename csrc/{cache/hybrid_cache.cpp => models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.cpp} (51%) create mode 100644 csrc/models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.hpp create mode 100644 csrc/models/qwen3_next/qwen3_next_decoderLayer.cpp diff --git a/csrc/cache/hybrid_cache.hpp b/csrc/cache/hybrid_cache.hpp deleted file mode 100644 index ad078d23..00000000 --- a/csrc/cache/hybrid_cache.hpp +++ /dev/null @@ -1,26 +0,0 @@ -#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 deleted file mode 100644 index d1bb6c20..00000000 --- a/csrc/config/hybrid_model_config.cpp +++ /dev/null @@ -1,63 +0,0 @@ -#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 deleted file mode 100644 index 745340c9..00000000 --- a/csrc/config/hybrid_model_config.hpp +++ /dev/null @@ -1,12 +0,0 @@ -#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 623574aa..dee3123c 100644 --- a/csrc/engine/compiler/paged_compiler.cpp +++ b/csrc/engine/compiler/paged_compiler.cpp @@ -11,7 +11,16 @@ namespace infinilm::engine { namespace { bool has_mamba_cache(const infinilm::global_state::ForwardContext &forward_context) { - return forward_context.mamba_state_pool_size > 0; + 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); } } // namespace @@ -52,6 +61,7 @@ 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 @@ -59,29 +69,8 @@ void PagedCompiler::compile() { 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()); @@ -90,14 +79,11 @@ 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()); - // 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()); + 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()); @@ -168,7 +154,7 @@ void PagedCompiler::compile() { infinicore::context::syncStream(); } - for (size_t b : compile_batch_sizes) { + for (size_t b : decode_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 ed7de761..f395b531 100644 --- a/csrc/global_state/forward_context.hpp +++ b/csrc/global_state/forward_context.hpp @@ -62,14 +62,6 @@ 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 deleted file mode 100644 index 345461ee..00000000 --- a/csrc/layers/causal_lm_templates/hybrid_decoder_layer.hpp +++ /dev/null @@ -1,117 +0,0 @@ -#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 6751cefa..5d284a31 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; - forward_context.kv_cache_vec = std::move( - default_allocate_kv_cache_tensors(cache_config, model_config_, attention_backend)); + 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_attention.cpp b/csrc/models/qwen3_5/qwen3_5_attention.cpp index 3db8aba6..e7a47f4f 100644 --- a/csrc/models/qwen3_5/qwen3_5_attention.cpp +++ b/csrc/models/qwen3_5/qwen3_5_attention.cpp @@ -32,12 +32,15 @@ Qwen35Attention::Qwen35Attention(std::shared_ptr const engine::distributed::RankInfo &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info(); int tp_rank = infinilm::global_state::get_tensor_model_parallel_rank(); int tp_size = infinilm::global_state::get_tensor_model_parallel_world_size(); - if ((total_num_kv_heads < tp_size) || (0 != (total_num_kv_heads % tp_size))) { + if (total_num_kv_heads >= static_cast(tp_size) + && total_num_kv_heads % tp_size != 0) { throw std::runtime_error("infinilm::models::qwen3_5::Qwen35Attention: num_key_value_heads must be divisible by tp_size"); } num_attention_heads_ = total_num_heads / tp_size; - num_key_value_heads_ = total_num_kv_heads / tp_size; + num_key_value_heads_ = total_num_kv_heads < static_cast(tp_size) + ? 1 + : total_num_kv_heads / tp_size; 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)); }; diff --git a/csrc/models/qwen3_5/qwen3_5_decoderLayer.cpp b/csrc/models/qwen3_5/qwen3_5_decoderLayer.cpp new file mode 100644 index 00000000..70964bb6 --- /dev/null +++ b/csrc/models/qwen3_5/qwen3_5_decoderLayer.cpp @@ -0,0 +1,66 @@ +#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 d1223cb0..751586bb 100644 --- a/csrc/models/qwen3_5/qwen3_5_decoderLayer.hpp +++ b/csrc/models/qwen3_5/qwen3_5_decoderLayer.hpp @@ -1,15 +1,37 @@ #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 { -using Qwen35DecoderLayer = infinilm::layers::causal_lm_templates::HybridDecoderLayer< - Qwen35Attention, - qwen3_next::Qwen3NextGatedDeltaNet, - infinilm::layers::MLP>; +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_; +}; } // 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 5c0b68be..72fe1a87 100644 --- a/csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp +++ b/csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp @@ -1,13 +1,41 @@ #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 { +Qwen35ForCausalLM::Qwen35ForCausalLM( + 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 Qwen35ForCausalLM::forward( + const InfinilmModel::Input &input) const { + auto hidden_states = model_->forward(input); + return {lm_head_->forward(hidden_states)}; +} + +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()) { @@ -41,7 +69,25 @@ 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"]; } - infinilm::config::prepare_hybrid_model_config(model_config); + 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("Qwen3.5 full_attention_interval must be positive"); + } + const 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( + (i + 1) % full_attention_interval == 0 + ? "full_attention" + : "linear_attention"); + } + config_json["layer_types"] = std::move(layer_types); + } + if (!config_json.contains("attention_bias")) { + config_json["attention_bias"] = false; + } return model_config; } 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 2aaaf4a1..51211481 100644 --- a/csrc/models/qwen3_5/qwen3_5_for_causal_lm.hpp +++ b/csrc/models/qwen3_5/qwen3_5_for_causal_lm.hpp @@ -1,31 +1,27 @@ #pragma once -#include "../../layers/causal_lm_templates/text_causal_lm.hpp" #include "qwen3_5_model.hpp" #include #include namespace infinilm::models::qwen3_5 { -template -class Qwen35CausalLM : public infinilm::layers::causal_lm_templates::TextCausalLM { +class Qwen35ForCausalLM : public InfinilmModel { public: - 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); - } -}; + Qwen35ForCausalLM(std::shared_ptr model_config, + const infinicore::Device &device); + + Output forward(const Input &input) const override; -using Qwen35ForCausalLM = Qwen35CausalLM; + void reset_cache(const cache::CacheConfig *cache_config) override; + +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 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 754ae3c1..0f66990a 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 @@ -108,18 +108,20 @@ void Qwen35ModelBase::replace_image_embeddings(infinicore::Tensor &inputs_embeds void Qwen35ModelBase::reset_cache(const cache::CacheConfig *cache_config) { auto &forward_context = infinilm::global_state::get_forward_context(); - forward_context.clear_model_caches(); + forward_context.kv_cache_vec.clear(); + forward_context.conv_state_vec.clear(); + forward_context.ssm_state_vec.clear(); if (nullptr == cache_config) { return; } const backends::AttentionBackend attention_backend = infinilm::global_state::get_infinilm_config().attention_backend; - auto cache_vectors = infinilm::cache::allocate_hybrid_cache_tensors(cache_config, model_config_, attention_backend); + auto cache_vectors = infinilm::models::qwen3_next::qwen3_next_allocate_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_moe/qwen3_5_moe_decoder_layer.cpp b/csrc/models/qwen3_5_moe/qwen3_5_moe_decoder_layer.cpp new file mode 100644 index 00000000..aa513284 --- /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 00000000..50d279e0 --- /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_next/qwen3_next_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(qwen3_next::Qwen3NextSparseMoeBlock, 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.hpp b/csrc/models/qwen3_5_moe/qwen3_5_moe_for_causal_lm.hpp index 130eb354..099c544e 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,21 +1,31 @@ #pragma once -#include "../../layers/causal_lm_templates/hybrid_decoder_layer.hpp" +#include "../../layers/causal_lm_templates/text_causal_lm.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 "qwen3_5_moe_decoder_layer.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; + +class Qwen35MoeForConditionalGeneration + : public infinilm::layers::causal_lm_templates::TextCausalLM { +public: + 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); + } +}; std::shared_ptr create_qwen3_5_moe_model_config( std::shared_ptr model_config); diff --git a/csrc/cache/hybrid_cache.cpp b/csrc/models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.cpp similarity index 51% rename from csrc/cache/hybrid_cache.cpp rename to csrc/models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.cpp index f945f1e9..3aaf898c 100644 --- a/csrc/cache/hybrid_cache.cpp +++ b/csrc/models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.cpp @@ -1,40 +1,40 @@ -#include "hybrid_cache.hpp" +#include "qwen3_next_allocate_kv_cache_tensors.hpp" + +#include "../../global_state/global_state.hpp" +#include "../../utils.hpp" +#include "infinicore/context/context.hpp" #include #include #include #include -namespace infinilm::cache { +namespace infinilm::models::qwen3_next { -HybridCacheTensors allocate_hybrid_cache_tensors( - const CacheConfig *cache_config, - const std::shared_ptr &model_config, +AllocatedHybridCache qwen3_next_allocate_cache_tensors( + const cache::CacheConfig *cache_config, + const std::shared_ptr &text_config, const backends::AttentionBackend &attention_backend) { if (nullptr == cache_config) { return {}; } - if (nullptr == model_config) { - throw std::runtime_error("allocate_hybrid_cache_tensors: model_config is null"); + if (nullptr == text_config) { + throw std::runtime_error("infinilm::models::qwen3_next::qwen3_next_allocate_kv_cache_tensors: text_config is null"); } - 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"); - } + 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"); std::vector kv_cache_vec; std::vector conv_state_vec; @@ -43,17 +43,8 @@ HybridCacheTensors allocate_hybrid_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( - "allocate_hybrid_cache_tensors: inconsistent mamba state pool size at layer " - + std::to_string(layer_idx)); - } - - auto conv_state = MambaCache::create_layer_conv_state( + auto conv_state = cache::MambaCache::create_layer_conv_state( linear_key_head_dim, linear_value_head_dim, linear_num_key_heads, @@ -61,7 +52,7 @@ HybridCacheTensors allocate_hybrid_cache_tensors( linear_conv_kernel_dim, dtype, pool_size); - auto ssm_state = MambaCache::create_layer_ssm_state( + auto ssm_state = cache::MambaCache::create_layer_ssm_state( linear_key_head_dim, linear_value_head_dim, linear_num_key_heads, @@ -74,8 +65,8 @@ HybridCacheTensors allocate_hybrid_cache_tensors( ssm_state_vec.push_back(std::move(ssm_state)); }; - auto allocate_static_full_attention_cache = [&](size_t layer_idx, const StaticKVCacheConfig &config) { - auto kv_cache = StaticKVCache::create_layer_kv_cache( + auto allocate_static_full_attention_cache = [&](size_t layer_idx, const cache::StaticKVCacheConfig &config) { + auto kv_cache = cache::StaticKVCache::create_layer_kv_cache( head_dim, head_dim, num_key_value_heads, @@ -89,8 +80,8 @@ HybridCacheTensors allocate_hybrid_cache_tensors( ssm_state_vec.emplace_back(); }; - auto allocate_paged_full_attention_cache = [&](size_t layer_idx, const PagedKVCacheConfig &config) { - auto kv_cache = PagedKVCache::create_layer_kv_cache( + auto allocate_paged_full_attention_cache = [&](size_t layer_idx, const cache::PagedKVCacheConfig &config) { + auto kv_cache = cache::PagedKVCache::create_layer_kv_cache( head_dim, head_dim, num_key_value_heads, @@ -105,9 +96,9 @@ HybridCacheTensors allocate_hybrid_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("allocate_hybrid_cache_tensors: invalid static kv cache config type"); + throw std::runtime_error("infinilm::models::qwen3_next::qwen3_next_allocate_kv_cache_tensors: invalid static kv cache config type"); } for (size_t layer_idx = 0; layer_idx < num_hidden_layers; ++layer_idx) { @@ -117,7 +108,7 @@ HybridCacheTensors allocate_hybrid_cache_tensors( } else if ("full_attention" == layer_type) { allocate_static_full_attention_cache(layer_idx, *static_kv_cache_config); } else { - throw std::runtime_error("allocate_hybrid_cache_tensors: unsupported layer_type '" + layer_type + "' for layer " + std::to_string(layer_idx)); + 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)); } } break; @@ -126,9 +117,9 @@ HybridCacheTensors allocate_hybrid_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("allocate_hybrid_cache_tensors: invalid paged kv cache config type"); + throw std::runtime_error("infinilm::models::qwen3_next::qwen3_next_allocate_kv_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 +130,18 @@ HybridCacheTensors allocate_hybrid_cache_tensors( } else if ("full_attention" == layer_type) { allocate_paged_full_attention_cache(layer_idx, *paged_kv_cache_config); } else { - throw std::runtime_error("allocate_hybrid_cache_tensors: unsupported layer_type '" + layer_type + "' for layer " + std::to_string(layer_idx)); + 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)); } } break; } default: - throw std::runtime_error("allocate_hybrid_cache_tensors: Unsupported attention backend: " + std::to_string(static_cast(attention_backend))); + throw std::runtime_error("infinilm::models::qwen3_next::qwen3_next_allocate_kv_cache_tensors: Unsupported attention backend: " + std::to_string(static_cast(attention_backend))); } - return HybridCacheTensors{ + return AllocatedHybridCache{ std::move(kv_cache_vec), std::move(conv_state_vec), - std::move(ssm_state_vec), - mamba_state_pool_size}; + std::move(ssm_state_vec)}; } -} // namespace infinilm::cache +} // 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 new file mode 100644 index 00000000..a4b5190f --- /dev/null +++ b/csrc/models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.hpp @@ -0,0 +1,26 @@ +#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 new file mode 100644 index 00000000..4c61832c --- /dev/null +++ b/csrc/models/qwen3_next/qwen3_next_decoderLayer.cpp @@ -0,0 +1,66 @@ +#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 df7f4c36..dd0505bd 100644 --- a/csrc/models/qwen3_next/qwen3_next_decoderLayer.hpp +++ b/csrc/models/qwen3_next/qwen3_next_decoderLayer.hpp @@ -1,15 +1,38 @@ #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 { -using Qwen3NextDecoderLayer = infinilm::layers::causal_lm_templates::HybridDecoderLayer< - Qwen3NextAttention, - Qwen3NextGatedDeltaNet, - Qwen3NextSparseMoeBlock>; +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_; +}; } // 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 e440833d..1b635454 100644 --- a/csrc/models/qwen3_next/qwen3_next_for_causal_lm.cpp +++ b/csrc/models/qwen3_next/qwen3_next_for_causal_lm.cpp @@ -1,14 +1,31 @@ #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); @@ -17,15 +34,16 @@ 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.clear_model_caches(); + 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 = cache::allocate_hybrid_cache_tensors(cache_config, model_config_, attention_backend); + auto cache_vectors = qwen3_next_allocate_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) { @@ -34,7 +52,21 @@ 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"); } - infinilm::config::prepare_hybrid_model_config(model_config); + 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; + } 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 1ee3fba0..0cbe4532 100644 --- a/csrc/models/qwen3_next/qwen3_next_for_causal_lm.hpp +++ b/csrc/models/qwen3_next/qwen3_next_for_causal_lm.hpp @@ -1,21 +1,25 @@ #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 infinilm::layers::causal_lm_templates::TextCausalLM { +class Qwen3NextForCausalLM : public InfinilmModel { public: - using Base = infinilm::layers::causal_lm_templates::TextCausalLM; - using Base::Base; + Qwen3NextForCausalLM(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(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 dc29fbeb..d3548c9f 100644 --- a/csrc/models/qwen3_next/qwen3_next_sparse_moe_block.cpp +++ b/csrc/models/qwen3_next/qwen3_next_sparse_moe_block.cpp @@ -1,31 +1,72 @@ #include "qwen3_next_sparse_moe_block.hpp" -#include "infinicore/ops.hpp" -#include "infinicore/ops/mul.hpp" +#include "../../global_state/global_state.hpp" -#include +#include +#include +#include +#include +#include +#include + +#include namespace infinilm::models::qwen3_next { -Qwen3NextSparseMoeBlock::Qwen3NextSparseMoeBlock( - std::shared_ptr model_config, - const infinicore::Device &device) +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) - : 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, +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, @@ -33,20 +74,33 @@ Qwen3NextSparseMoeBlock::Qwen3NextSparseMoeBlock( device); } -infinicore::Tensor Qwen3NextSparseMoeBlock::forward( - const infinicore::Tensor &hidden_states) const { - auto routed_output = infinilm::layers::moe::SparseMoeBlock::forward(hidden_states); +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_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 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); - return infinicore::op::add(routed_output, shared_output); + 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 8cff0848..2cec0bd5 100644 --- a/csrc/models/qwen3_next/qwen3_next_sparse_moe_block.hpp +++ b/csrc/models/qwen3_next/qwen3_next_sparse_moe_block.hpp @@ -2,30 +2,42 @@ #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 "../../layers/moe/experts/fused_moe_experts.hpp" +#include "../../layers/moe/fused_moe.hpp" +#include "../../layers/moe/router/topk_router.hpp" -#include #include namespace infinilm::models::qwen3_next { -class Qwen3NextSparseMoeBlock : public infinilm::layers::moe::SparseMoeBlock { +class Qwen3NextSharedExpert : 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); + Qwen3NextSharedExpert(std::shared_ptr model_config, + 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); + 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_; }; } // namespace infinilm::models::qwen3_next diff --git a/python/infinilm/processors/__init__.py b/python/infinilm/processors/__init__.py index d4ab5b3f..f1a543b6 100644 --- a/python/infinilm/processors/__init__.py +++ b/python/infinilm/processors/__init__.py @@ -39,8 +39,10 @@ def from_pretrained(cls, model_dir_path: str, **kwargs) -> InfinilmProcessor: with raw_config_path.open("r") as f: raw_config = json.load(f) - model_type = str(raw_config.get("model_type", "")).lower() - if not model_type: + raw_model_type = str(raw_config.get("model_type", "")).lower() + if raw_model_type in {"qwen3_5", "qwen3_5_moe"}: + model_type = raw_model_type + else: config = AutoConfig.from_pretrained(model_dir_path, trust_remote_code=True) model_type = config.model_type.lower() diff --git a/test/models/qwen3_5_moe/test_adaptation.py b/test/models/qwen3_5_moe/test_adaptation.py index e7294fa0..3a0a43ba 100644 --- a/test/models/qwen3_5_moe/test_adaptation.py +++ b/test/models/qwen3_5_moe/test_adaptation.py @@ -1,7 +1,6 @@ import unittest import torch - from infinilm.infer_engine import _infer_position_id_axes from infinilm.modeling_utils import _remap_qwen3_5_moe @@ -11,11 +10,7 @@ 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]} - } - } + 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): @@ -61,9 +56,7 @@ def test_splits_packed_expert_weights(self): 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.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)