[2026春季][T2-1-1]xuzheng567

csrc/models/ernie4_5_vl/ernie4_5_attention.hpp

@@ -0,0 +1,66 @@
+#pragma once
+
+#include "../../config/model_config.hpp"
+#include "../../layers/attention/attention.hpp"
+#include "../../layers/common_modules.hpp"
+#include "../../layers/linear/linear.hpp"
+#include "../../layers/rotary_embedding/rotary_embedding.hpp"
+#include "infinicore/device.hpp"
+#include "infinicore/nn/module.hpp"
+#include "infinicore/tensor.hpp"
+
+#include <memory>
+#include <vector>
+
+namespace infinilm::models::ernie4_5_vl {
+
+struct Ernie45MropeCache {
+    std::vector<int> section{22, 22, 20};
+    infinicore::Tensor sin_h;
+    infinicore::Tensor cos_h;
+    infinicore::Tensor sin_w;
+    infinicore::Tensor cos_w;
+    infinicore::Tensor sin_t;
+    infinicore::Tensor cos_t;
+};
+
+std::shared_ptr<const Ernie45MropeCache> build_ernie45_mrope_cache(std::shared_ptr<infinilm::config::ModelConfig> model_config,
+                                                                   const infinicore::Device &device);
+
+class Ernie45Attention : public infinicore::nn::Module {
+public:
+    Ernie45Attention(std::shared_ptr<infinilm::config::ModelConfig> model_config,
+                     size_t layer_idx,
+                     std::shared_ptr<const Ernie45MropeCache> mrope_cache,
+                     const infinicore::Device &device);
+
+    infinicore::Tensor forward(const infinicore::Tensor &positions,
+                               const infinicore::Tensor &hidden_states) const;
+
+private:
+    infinicore::Tensor forward_static_(const infinicore::Tensor &positions,
+                                       const infinicore::Tensor &hidden_states) const;
+    infinicore::Tensor forward_paged_(const infinicore::Tensor &positions,
+                                      const infinicore::Tensor &hidden_states) const;
+
+    size_t layer_idx_{0};
+    size_t hidden_size_{0};
+    size_t head_dim_{0};
+    size_t rotary_dim_{0};
+    size_t num_attention_heads_{0};
+    size_t num_key_value_heads_{0};
+    std::shared_ptr<const Ernie45MropeCache> mrope_cache_;
+    infinilm::backends::AttentionBackend attention_backend_;
+
+    INFINICORE_NN_MODULE(infinilm::layers::linear::ColumnParallelLinear, q_proj);
+    INFINICORE_NN_MODULE(infinilm::layers::linear::ColumnParallelLinear, k_proj);
+    INFINICORE_NN_MODULE(infinilm::layers::linear::ColumnParallelLinear, v_proj);
+    INFINICORE_NN_MODULE(infinilm::layers::linear::RowParallelLinear, o_proj);
+    INFINICORE_NN_MODULE(infinicore::nn::RoPE, rotary_emb);
+
+    infinicore::nn::Parameter kv_cache_k_scale_;
+    infinicore::nn::Parameter kv_cache_v_scale_;
+    std::shared_ptr<infinilm::layers::attention::AttentionLayer> attn_;
+};
+
+} // namespace infinilm::models::ernie4_5_vl

csrc/models/ernie4_5_vl/ernie4_5_decoder_layer.cpp

@@ -0,0 +1,93 @@
+#include "ernie4_5_decoder_layer.hpp"
+
+#include "infinicore/ops.hpp"
+
+#include <algorithm>
+#include <utility>
+
+namespace infinilm::models::ernie4_5_vl {
+namespace {
+
+size_t min_json_size(const nlohmann::json &value, size_t fallback) {
+    if (value.is_array() && !value.empty()) {
+        size_t result = value.at(0).get<size_t>();
+        for (const auto &item : value) {
+            result = std::min(result, item.get<size_t>());
+        }
+        return result;
+    }
+    return value.is_number_unsigned() ? value.get<size_t>() : fallback;
+}
+
+size_t max_json_size(const nlohmann::json &value, size_t fallback) {
+    if (value.is_array() && !value.empty()) {
+        size_t result = value.at(0).get<size_t>();
+        for (const auto &item : value) {
+            result = std::max(result, item.get<size_t>());
+        }
+        return result;
+    }
+    return value.is_number_unsigned() ? value.get<size_t>() : fallback;
+}
+
+bool is_moe_layer(const nlohmann::json &config, size_t layer_idx) {
+    const size_t interval = config.value("moe_layer_interval", 1);
+    const size_t start = config.contains("moe_layer_start_index") ? min_json_size(config.at("moe_layer_start_index"), 0) : 0;
+    const size_t end = config.contains("moe_layer_end_index") ? max_json_size(config.at("moe_layer_end_index"), config.value("num_hidden_layers", 1) - 1) : config.value("num_hidden_layers", 1) - 1;
+    return config.value("use_moe", false) && interval > 0 && ((layer_idx + 1) % interval == 0) && layer_idx >= start && layer_idx <= end;
+}
+
+} // namespace
+
+Ernie45DecoderLayer::Ernie45DecoderLayer(std::shared_ptr<infinilm::config::ModelConfig> model_config,
+                                         size_t layer_idx,
+                                         std::shared_ptr<const Ernie45MropeCache> mrope_cache,
+                                         const infinicore::Device &device) {
+    const auto &dtype = model_config->get_dtype();
+    const size_t hidden_size = model_config->get<size_t>("hidden_size");
+    const double rms_norm_eps = model_config->get<double>("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(self_attn, model_config, layer_idx, std::move(mrope_cache), device);
+
+    use_moe_ = is_moe_layer(model_config->get_config_json(), layer_idx);
+    if (use_moe_) {
+        mlp_ = this->register_module<Ernie45MoE>("mlp", model_config, device);
+        // ERNIE also defines an unregistered mlp_text path for pure text tokens when multimodal
+        // experts are enabled. This first implementation routes all tokens through mlp.
+    } else {
+        mlp_ = this->register_module<infinilm::layers::mlp::MLP>("mlp", model_config, device);
+    }
+}
+
+std::tuple<infinicore::Tensor, infinicore::Tensor> Ernie45DecoderLayer::forward(const infinicore::Tensor &positions,
+                                                                                infinicore::Tensor &hidden_states,
+                                                                                infinicore::Tensor &residual) {
+    input_layernorm_->forward_inplace(hidden_states, residual);
+    hidden_states = self_attn_->forward(positions, hidden_states);
+
+    post_attention_layernorm_->forward_inplace(hidden_states, residual);
+    hidden_states = use_moe_
+                      ? std::static_pointer_cast<Ernie45MoE>(mlp_)->forward(hidden_states)
+                      : std::static_pointer_cast<infinilm::layers::mlp::MLP>(mlp_)->forward(hidden_states);
+    return {hidden_states, residual};
+}
+
+infinicore::Tensor Ernie45DecoderLayer::forward(const infinicore::Tensor &positions,
+                                                infinicore::Tensor &hidden_states) {
+    auto residual = hidden_states;
+    hidden_states = input_layernorm_->forward(hidden_states);
+    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 = use_moe_
+                      ? std::static_pointer_cast<Ernie45MoE>(mlp_)->forward(hidden_states)
+                      : std::static_pointer_cast<infinilm::layers::mlp::MLP>(mlp_)->forward(hidden_states);
+    hidden_states = infinicore::op::add(residual, hidden_states);
+    return hidden_states;
+}
+
+} // namespace infinilm::models::ernie4_5_vl

csrc/models/ernie4_5_vl/ernie4_5_decoder_layer.hpp

@@ -0,0 +1,38 @@
+#pragma once
+
+#include "../../layers/common_modules.hpp"
+#include "ernie4_5_attention.hpp"
+#include "ernie4_5_moe.hpp"
+#include "infinicore/nn/module.hpp"
+#include "infinicore/tensor.hpp"
+
+#include <memory>
+#include <tuple>
+
+namespace infinilm::models::ernie4_5_vl {
+
+class Ernie45DecoderLayer : public infinicore::nn::Module {
+public:
+    Ernie45DecoderLayer(std::shared_ptr<infinilm::config::ModelConfig> model_config,
+                        size_t layer_idx,
+                        std::shared_ptr<const Ernie45MropeCache> mrope_cache,
+                        const infinicore::Device &device);
+
+    std::tuple<infinicore::Tensor, infinicore::Tensor> forward(const infinicore::Tensor &positions,
+                                                               infinicore::Tensor &hidden_states,
+                                                               infinicore::Tensor &residual);
+
+    infinicore::Tensor forward(const infinicore::Tensor &positions,
+                               infinicore::Tensor &hidden_states);
+
+protected:
+    INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, input_layernorm);
+    INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, post_attention_layernorm);
+    INFINICORE_NN_MODULE(Ernie45Attention, self_attn);
+    std::shared_ptr<infinicore::nn::Module> mlp_;
+
+private:
+    bool use_moe_{false};
+};
+
+} // namespace infinilm::models::ernie4_5_vl

csrc/models/ernie4_5_vl/ernie4_5_for_causal_lm.cpp

@@ -0,0 +1,136 @@
+#include "ernie4_5_for_causal_lm.hpp"
+
+#include "../../global_state/global_state.hpp"
+#include "../models_registry.hpp"
+
+#include <stdexcept>
+#include <string>
+
+namespace infinilm::models::ernie4_5_vl {
+namespace {
+
+size_t get_first_size(const nlohmann::json &config, const char *key, size_t default_value) {
+    if (!config.contains(key) || config.at(key).is_null()) {
+        return default_value;
+    }
+    const auto &value = config.at(key);
+    if (value.is_array()) {
+        return value.empty() ? default_value : value.at(0).get<size_t>();
+    }
+    return value.get<size_t>();
+}
+
+void normalize_ernie_config(nlohmann::json &config_json) {
+    if (!config_json.contains("dtype") && config_json.contains("torch_dtype")) {
+        config_json["dtype"] = config_json["torch_dtype"];
+    }
+    if (!config_json.contains("head_dim")) {
+        config_json["head_dim"] = config_json["hidden_size"].get<size_t>() / config_json["num_attention_heads"].get<size_t>();
+    }
+    if (!config_json.contains("partial_rotary_factor")) {
+        config_json["partial_rotary_factor"] = 1.0;
+    }
+    if (!config_json.contains("compression_ratio")) {
+        config_json["compression_ratio"] = 1.0;
+    }
+    if (!config_json.contains("use_flash_attention")) {
+        config_json["use_flash_attention"] = true;
+    }
+    if (!config_json.contains("attention_probs_dropout_prob")) {
+        config_json["attention_probs_dropout_prob"] = 0.0;
+    }
+    if (!config_json.contains("hidden_dropout_prob")) {
+        config_json["hidden_dropout_prob"] = 0.0;
+    }
+    if (!config_json.contains("mlp_bias")) {
+        config_json["mlp_bias"] = config_json.value("use_bias", false);
+    }
+    if (!config_json.contains("attention_bias")) {
+        config_json["attention_bias"] = config_json.value("use_bias", false);
+    }
+    if (!config_json.contains("norm_topk_prob")) {
+        config_json["norm_topk_prob"] = config_json.value("moe_norm_gate_logits", true);
+    }
+    if (!config_json.contains("num_experts")) {
+        config_json["num_experts"] = get_first_size(config_json, "moe_num_experts", 0);
+    }
+    if (!config_json.contains("num_experts_per_tok")) {
+        config_json["num_experts_per_tok"] = config_json.value("moe_k", 1);
+    }
+    if (!config_json.contains("use_moe")) {
+        const size_t num_experts = config_json.value("num_experts", 0);
+        config_json["use_moe"] = num_experts > 0;
+    }
+    if (!config_json.contains("moe_dropout_prob")) {
+        config_json["moe_dropout_prob"] = 0.0;
+    }
+    if (!config_json.contains("moe_reverse_token_drop")) {
+        config_json["moe_reverse_token_drop"] = false;
+    }
+    if (!config_json.contains("moe_group")) {
+        config_json["moe_group"] = "world";
+    }
+    if (!config_json.contains("moe_all_to_all_dropout")) {
+        config_json["moe_all_to_all_dropout"] = 0.0;
+    }
+}
+
+} // namespace
+
+Ernie45ForConditionalGeneration::Ernie45ForConditionalGeneration(std::shared_ptr<infinilm::config::ModelConfig> model_config,
+                                                                 const infinicore::Device &device) {
+    model_config_ = model_config;
+    const size_t hidden_size = model_config->get<size_t>("hidden_size");
+    const size_t vocab_size = model_config->get<size_t>("vocab_size");
+    const auto &dtype = model_config->get_dtype();
+    auto &config_json = model_config->get_config_json();
+    if (config_json.contains("vision_config") && config_json["vision_config"].is_object()) {
+        INFINICORE_NN_MODULE_INIT(vision_model, config_json["vision_config"], dtype, device);
+    }
+
+    INFINICORE_NN_MODULE_INIT(model, model_config, device);
+    INFINICORE_NN_MODULE_INIT(lm_head, hidden_size, vocab_size, false, dtype, device);
+}
+
+infinilm::InfinilmModel::Output Ernie45ForConditionalGeneration::forward(const infinilm::InfinilmModel::Input &input) const {
+    auto hidden_states = (input.pixel_values.has_value() && !input.pixel_values.value().empty())
+                           ? model_->forward(input, vision_model_.get())
+                           : model_->forward(input);
+    auto logits = lm_head_->forward(hidden_states);
+    return {logits};
+}
+
+void Ernie45ForConditionalGeneration::reset_cache(const cache::CacheConfig *cache_config) {
+    if (cache_config == nullptr) {
+        cache_config_.reset();
+        return;
+    }
+    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();
+
+    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));
+}
+
+std::shared_ptr<infinilm::config::ModelConfig> create_ernie4_5_moe_vl_model_config(std::shared_ptr<infinilm::config::ModelConfig> model_config) {
+    const std::string model_type = model_config->get<std::string>("model_type");
+    if ("ernie4_5_moe_vl" != model_type) {
+        throw std::runtime_error("infinilm::models::ernie4_5_vl::create_ernie4_5_moe_vl_model_config: model_type is not ernie4_5_moe_vl");
+    }
+    normalize_ernie_config(model_config->get_config_json());
+    model_config->set_rope_algo(infinicore::nn::RoPE::Algo::GPT_J);
+    return model_config;
+}
+
+} // namespace infinilm::models::ernie4_5_vl
+
+namespace {
+INFINILM_REGISTER_CAUSAL_LM_MODEL(
+    ernie4_5_moe_vl,
+    infinilm::models::ernie4_5_vl::Ernie45ForConditionalGeneration,
+    infinilm::models::ernie4_5_vl::create_ernie4_5_moe_vl_model_config);
+} // namespace

csrc/models/ernie4_5_vl/ernie4_5_for_causal_lm.hpp

@@ -0,0 +1,24 @@
+#pragma once
+
+#include "ernie4_5_model.hpp"
+
+namespace infinilm::models::ernie4_5_vl {
+
+class Ernie45ForConditionalGeneration : public InfinilmModel {
+public:
+    Ernie45ForConditionalGeneration(std::shared_ptr<infinilm::config::ModelConfig> 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(Ernie45VisionModel, vision_model);
+    INFINICORE_NN_MODULE(Ernie45Model, model);
+    INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, lm_head);
+};
+
+std::shared_ptr<infinilm::config::ModelConfig> create_ernie4_5_moe_vl_model_config(std::shared_ptr<infinilm::config::ModelConfig> model_config);
+
+} // namespace infinilm::models::ernie4_5_vl