feat(minimax): support MiniMax-Text-01 and MiniMax-M2

.gitignore

@@ -2,6 +2,11 @@
 .xmake/
 build/
 python/infinilm/lib/*.so
+python/infinilm/lib/*.pyd
+python/infinilm/lib/*.dll
+python/infinilm/lib/*.lib
+python/infinilm/lib/*.dylib
+python/infinilm/bin/
 
 # MacOS Cache
 .DS_Store
@@ -34,3 +39,4 @@ __pycache__/
 *.http
 
 *.nsys-rep
+

csrc/engine/rank_worker.cpp

@@ -476,7 +476,12 @@ void RankWorker::thread_loop() {
                             const auto &batch_size{logits_shape[0]};
 
                             auto n_req = local_args.input_offsets.value()->size(0) - 1;
-                            int32_t *input_offsets = (int32_t *)local_args.input_offsets.value()->data();
+                            auto cpu_input_offsets = local_args.input_offsets.value();
+                            if (cpu_input_offsets->device().getType() != infinicore::Device::Type::CPU) {
+                                cpu_input_offsets = cpu_input_offsets->to(infinicore::Device::cpu());
+                                infinicore::context::syncStream();
+                            }
+                            const int32_t *input_offsets = reinterpret_cast<const int32_t *>(cpu_input_offsets->data());
 
                             const bool sample_all_positions = local_args.sample_all_positions;
                             const size_t logits_positions = batch_size * total_len;

csrc/models/minimax/minimax_attention.cpp

@@ -0,0 +1,164 @@
+#include "minimax_attention.hpp"
+
+#include "../../global_state/global_state.hpp"
+#include "../../layers/attention/attention.hpp"
+#include "../../utils.hpp"
+#include <cmath>
+#include <stdexcept>
+
+namespace infinilm::models::minimax {
+
+MiniMaxAttention::MiniMaxAttention(std::shared_ptr<infinilm::config::ModelConfig> model_config,
+                                   size_t layer_idx,
+                                   const infinicore::Device &device)
+    : layer_idx_(layer_idx) {
+    const auto &dtype{model_config->get_dtype()};
+    hidden_size_ = model_config->get<size_t>("hidden_size");
+    head_dim_ = model_config->get_or<size_t>("head_dim", 0);
+    if (head_dim_ == 0) {
+        head_dim_ = hidden_size_ / model_config->get<size_t>("num_attention_heads");
+    }
+    const size_t total_num_heads = model_config->get<size_t>("num_attention_heads");
+    const size_t total_num_kv_heads = model_config->get_or<size_t>("num_key_value_heads", total_num_heads);
+
+    attention_backend_ = infinilm::global_state::get_infinilm_config().attention_backend;
+    const auto &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info();
+    const int tp_rank = rank_info.tp_rank;
+    const int tp_size = rank_info.tp_size;
+
+    num_attention_heads_ = total_num_heads / tp_size;
+    num_key_value_heads_ = total_num_kv_heads < static_cast<size_t>(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)); };
+    qkv_proj_ = std::make_shared<layers::linear::QKVParallelLinear>(
+        hidden_size_, head_dim_, total_num_heads, total_num_kv_heads,
+        "q_proj", "k_proj", "v_proj", register_fn,
+        quantization_method, false, dtype, device, rank_info);
+    o_proj_ = this->register_module<layers::linear::RowParallelLinear>(
+        "o_proj", total_num_heads * head_dim_, hidden_size_, quantization_method,
+        false, dtype, device, tp_rank, tp_size, rank_info.comm);
+
+    rotary_emb_ = infinilm::layers::rotary_embedding::get_rope(model_config, device);
+
+    const float scaling = 1.0f / std::sqrt(static_cast<float>(head_dim_));
+    attn_ = std::make_shared<infinilm::layers::attention::AttentionLayer>(
+        num_attention_heads_, head_dim_, scaling, num_key_value_heads_, layer_idx_,
+        kv_cache_k_scale_, kv_cache_v_scale_, attention_backend_);
+
+    infinilm::layers::attention::init_kv_cache_quant_params(register_fn, device, kv_cache_k_scale_, kv_cache_v_scale_);
+}
+
+infinicore::Tensor MiniMaxAttention::forward(const infinicore::Tensor &positions,
+                                             const infinicore::Tensor &hidden_states) const {
+    if (::infinilm::backends::AttentionBackend::STATIC_ATTN == attention_backend_) {
+        return forward_static_(positions, hidden_states);
+    }
+    return forward_paged_(positions, hidden_states);
+}
+
+infinicore::Tensor MiniMaxAttention::forward_static_(const infinicore::Tensor &position_ids,
+                                                     const infinicore::Tensor &hidden_states) const {
+    auto hidden_states_mutable = hidden_states;
+    auto shape = hidden_states->shape();
+    const size_t batch_size = shape[0];
+    const size_t seq_len = shape[1];
+
+    auto [q_proj_out, k, v] = qkv_proj_->forward_split(hidden_states_mutable);
+    auto q_heads = q_proj_out->as_strided(
+        {batch_size * seq_len, num_attention_heads_, head_dim_},
+        {q_proj_out->stride(1), static_cast<infinicore::Stride>(head_dim_), 1});
+    auto k_heads = k->as_strided(
+        {batch_size * seq_len, num_key_value_heads_, head_dim_},
+        {k->stride(1), static_cast<infinicore::Stride>(head_dim_), 1});
+
+    auto q_reshaped = q_heads->as_strided(
+        {batch_size, seq_len, num_attention_heads_, head_dim_},
+        {static_cast<infinicore::Stride>(seq_len * num_attention_heads_ * head_dim_),
+         static_cast<infinicore::Stride>(num_attention_heads_ * head_dim_),
+         static_cast<infinicore::Stride>(head_dim_),
+         1});
+    auto k_reshaped = k_heads->as_strided(
+        {batch_size, seq_len, num_key_value_heads_, head_dim_},
+        {static_cast<infinicore::Stride>(seq_len * num_key_value_heads_ * head_dim_),
+         static_cast<infinicore::Stride>(num_key_value_heads_ * head_dim_),
+         static_cast<infinicore::Stride>(head_dim_),
+         1});
+    auto v_reshaped = v->as_strided(
+        {batch_size, seq_len, num_key_value_heads_, head_dim_},
+        {v->stride(0), v->stride(1), static_cast<infinicore::Stride>(head_dim_), 1});
+
+    // q/k/v are views into the fused qkv buffer; materialize them so the
+    // in-place RoPE cannot clobber neighbouring q/k/v regions.
+    q_reshaped = q_reshaped->contiguous();
+    k_reshaped = k_reshaped->contiguous();
+    v_reshaped = v_reshaped->contiguous();
+
+    auto pos_shape = position_ids->shape();
+    infinicore::Tensor pos_ids_for_rope = position_ids;
+    if (pos_shape.size() == 2) {
+        auto pos_narrowed = position_ids->narrow({{0, 0, 1}});
+        pos_ids_for_rope = pos_narrowed->contiguous()->view({pos_shape[1]});
+    } else if (pos_shape.size() == 1) {
+        pos_ids_for_rope = position_ids->contiguous();
+    } else {
+        throw std::runtime_error("infinilm::models::minimax::MiniMaxAttention: Unexpected position_ids shape");
+    }
+
+    rotary_emb_->forward(q_reshaped, pos_ids_for_rope, true);
+    rotary_emb_->forward(k_reshaped, pos_ids_for_rope, true);
+
+
+    auto attn_output = attn_->forward(q_reshaped, k_reshaped, v_reshaped);
+    return o_proj_->forward(attn_output);
+}
+
+infinicore::Tensor MiniMaxAttention::forward_paged_(const infinicore::Tensor &position_ids,
+                                                    const infinicore::Tensor &hidden_states) const {
+    auto hidden_states_mutable = hidden_states;
+    auto shape = hidden_states->shape();
+    const size_t batch_size = shape[0];
+    const size_t seq_len = shape[1];
+
+    ASSERT_EQ(batch_size, 1);
+
+    auto [q_proj_out, k, v] = qkv_proj_->forward_split(hidden_states_mutable);
+    auto q_reshaped = q_proj_out->as_strided(
+        {seq_len, num_attention_heads_, head_dim_},
+        {q_proj_out->stride(1), static_cast<infinicore::Stride>(head_dim_), 1});
+    auto k_reshaped = k->as_strided(
+        {seq_len, num_key_value_heads_, head_dim_},
+        {k->stride(1), static_cast<infinicore::Stride>(head_dim_), 1});
+    auto v_reshaped = v->as_strided(
+        {seq_len, num_key_value_heads_, head_dim_},
+        {v->stride(1), static_cast<infinicore::Stride>(head_dim_), 1});
+
+    q_reshaped = q_reshaped->contiguous();
+    k_reshaped = k_reshaped->contiguous();
+    v_reshaped = v_reshaped->contiguous();
+
+    auto pos_shape = position_ids->shape();
+    infinicore::Tensor pos_ids_for_rope = position_ids;
+    if (pos_shape.size() == 2) {
+        auto pos_narrowed = position_ids->narrow({{0, 0, 1}});
+        pos_ids_for_rope = pos_narrowed->view({pos_shape[1]});
+    } else if (pos_shape.size() == 1) {
+        pos_ids_for_rope = position_ids;
+    } else {
+        throw std::runtime_error("infinilm::models::minimax::MiniMaxAttention: Unexpected position_ids shape");
+    }
+
+    rotary_emb_->forward(q_reshaped, pos_ids_for_rope, true);
+    rotary_emb_->forward(k_reshaped, pos_ids_for_rope, true);
+
+
+    auto attn_output = attn_->forward(q_reshaped, k_reshaped, v_reshaped);
+    return o_proj_->forward(attn_output);
+}
+
+} // namespace infinilm::models::minimax
+
+
+

csrc/models/minimax/minimax_attention.hpp

@@ -0,0 +1,46 @@
+#pragma once
+
+#include "../../layers/common_modules.hpp"
+#include <memory>
+
+namespace infinilm::models::minimax {
+
+class MiniMaxAttention : public infinicore::nn::Module {
+public:
+    MiniMaxAttention(std::shared_ptr<infinilm::config::ModelConfig> model_config,
+                     size_t layer_idx,
+                     const infinicore::Device &device);
+
+    infinicore::Tensor forward(const infinicore::Tensor &positions,
+                               const infinicore::Tensor &hidden_states) const;
+
+    size_t layer_idx() const { return layer_idx_; }
+    size_t num_heads() const { return num_attention_heads_; }
+    size_t num_kv_heads() const { return num_key_value_heads_; }
+    size_t head_dim() const { return head_dim_; }
+    size_t hidden_size() const { return hidden_size_; }
+
+private:
+    infinicore::Tensor forward_static_(const infinicore::Tensor &position_ids,
+                                       const infinicore::Tensor &hidden_states) const;
+    infinicore::Tensor forward_paged_(const infinicore::Tensor &position_ids,
+                                      const infinicore::Tensor &hidden_states) const;
+
+    std::shared_ptr<infinilm::layers::linear::QKVParallelLinear> qkv_proj_;
+    std::shared_ptr<infinilm::layers::linear::RowParallelLinear> o_proj_;
+    std::shared_ptr<infinicore::nn::RoPE> rotary_emb_;
+    std::shared_ptr<infinilm::layers::attention::AttentionLayer> attn_;
+
+    ::infinilm::backends::AttentionBackend attention_backend_;
+    size_t layer_idx_;
+    size_t num_attention_heads_;
+    size_t num_key_value_heads_;
+    size_t hidden_size_;
+    size_t head_dim_;
+
+    // For off-line kv cache quantization
+    INFINICORE_NN_PARAMETER(kv_cache_k_scale);
+    INFINICORE_NN_PARAMETER(kv_cache_v_scale);
+};
+
+} // namespace infinilm::models::minimax

csrc/models/minimax/minimax_decoderLayer.cpp

@@ -0,0 +1,76 @@
+#include "minimax_decoderLayer.hpp"
+
+#include <infinicore/ops/add.hpp>
+#include <infinicore/ops/mul_scalar.hpp>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+namespace infinilm::models::minimax {
+
+MiniMaxDecoderLayer::MiniMaxDecoderLayer(std::shared_ptr<infinilm::config::ModelConfig> 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<size_t>("hidden_size");
+    const double rms_norm_eps = model_config->get_or<double>("rms_norm_eps", 1e-5);
+
+    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);
+
+    const std::vector<std::string> layer_types = model_config->get<std::vector<std::string>>("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("infinilm::models::minimax::MiniMaxDecoderLayer: unsupported layer_type '" + layer_type_ + "' for layer " + std::to_string(layer_idx));
+    }
+
+    num_experts_ = model_config->get_or<size_t>("num_experts", 1);
+    if (num_experts_ > 1) {
+        INFINICORE_NN_MODULE_INIT(moe, model_config, layer_idx, device);
+    } else {
+        INFINICORE_NN_MODULE_INIT(mlp, model_config, device);
+    }
+
+    alpha_attn_ = model_config->get_or("layernorm_attention_alpha", model_config->get_or("linear_attn_alpha_factor", 1.0));
+    beta_attn_ = model_config->get_or("layernorm_attention_beta", model_config->get_or("linear_attn_beta_factor", 1.0));
+    alpha_mlp_ = model_config->get_or("layernorm_mlp_alpha", model_config->get_or("mlp_alpha_factor", 1.0));
+    beta_mlp_ = model_config->get_or("layernorm_mlp_beta", model_config->get_or("mlp_beta_factor", 1.0));
+}
+
+infinicore::Tensor MiniMaxDecoderLayer::forward(const infinicore::Tensor &positions,
+                                                infinicore::Tensor &hidden_states) const {
+    // HF transformers `minimax` carries the *normalized* value in both residual
+    // paths: x1 = norm(x0); out_attn = x1 + attn(x1); x2 = norm(out_attn);
+    // out = x2 + mlp(x2).
+    auto x = input_layernorm_->forward(hidden_states);
+    auto residual = x;
+    if ("linear_attention" == layer_type_) {
+        x = linear_attn_->forward(x);
+    } else {
+        x = self_attn_->forward(positions, x);
+    }
+    auto scale_add = [](const infinicore::Tensor &a, double sa,
+                        const infinicore::Tensor &b, double sb) -> infinicore::Tensor {
+        if (sa == 1.0 && sb == 1.0) {
+            return infinicore::op::add(a, b);
+        }
+        auto a_scaled = sa == 1.0 ? a : infinicore::op::mul_scalar(a, sa);
+        auto b_scaled = sb == 1.0 ? b : infinicore::op::mul_scalar(b, sb);
+        return infinicore::op::add(a_scaled, b_scaled);
+    };
+    auto post_input = scale_add(residual, alpha_attn_, x, beta_attn_);
+
+    // Pre-norm MLP sub-block (residual = normalized post-attention value).
+    auto normed = post_attention_layernorm_->forward(post_input);
+    auto mlp_out = num_experts_ > 1 ? moe_->forward(normed) : mlp_->forward(normed);
+    return scale_add(normed, alpha_mlp_, mlp_out, beta_mlp_);
+}
+
+} // namespace infinilm::models::minimax
+
+

csrc/models/minimax/minimax_decoderLayer.hpp

@@ -0,0 +1,40 @@
+#pragma once
+
+#include "../../layers/common_modules.hpp"
+#include "minimax_attention.hpp"
+#include "minimax_lightning_attention.hpp"
+#include "minimax_moe.hpp"
+#include <string>
+
+namespace infinilm::models::minimax {
+
+class MiniMaxDecoderLayer : public infinicore::nn::Module {
+public:
+    MiniMaxDecoderLayer(std::shared_ptr<infinilm::config::ModelConfig> model_config,
+                        size_t layer_idx,
+                        const infinicore::Device &device);
+
+    infinicore::Tensor forward(const infinicore::Tensor &positions,
+                               infinicore::Tensor &hidden_states) const;
+
+    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(MiniMaxLightningAttention, linear_attn);
+    INFINICORE_NN_MODULE(MiniMaxAttention, self_attn);
+    INFINICORE_NN_MODULE(infinilm::layers::mlp::MLP, mlp);
+    INFINICORE_NN_MODULE(MiniMaxMoeBlock, moe);
+
+private:
+    size_t layer_idx_;
+    std::string layer_type_;
+    size_t num_experts_{1};
+    double alpha_attn_{1.0};
+    double beta_attn_{1.0};
+    double alpha_mlp_{1.0};
+    double beta_mlp_{1.0};
+};
+
+} // namespace infinilm::models::minimax