diff --git a/.gitignore b/.gitignore index 149d4bc38..5fd88b8d2 100644 --- a/.gitignore +++ b/.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 + diff --git a/csrc/engine/rank_worker.cpp b/csrc/engine/rank_worker.cpp index 0fa0a84cf..841bd49a2 100644 --- a/csrc/engine/rank_worker.cpp +++ b/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(cpu_input_offsets->data()); const bool sample_all_positions = local_args.sample_all_positions; const size_t logits_positions = batch_size * total_len; diff --git a/csrc/models/minimax/minimax_attention.cpp b/csrc/models/minimax/minimax_attention.cpp new file mode 100644 index 000000000..94c0cd408 --- /dev/null +++ b/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 +#include + +namespace infinilm::models::minimax { + +MiniMaxAttention::MiniMaxAttention(std::shared_ptr 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("hidden_size"); + head_dim_ = model_config->get_or("head_dim", 0); + if (head_dim_ == 0) { + head_dim_ = hidden_size_ / model_config->get("num_attention_heads"); + } + const size_t total_num_heads = model_config->get("num_attention_heads"); + const size_t total_num_kv_heads = model_config->get_or("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(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( + 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( + "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(head_dim_)); + attn_ = std::make_shared( + 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(head_dim_), 1}); + auto k_heads = k->as_strided( + {batch_size * seq_len, num_key_value_heads_, head_dim_}, + {k->stride(1), static_cast(head_dim_), 1}); + + auto q_reshaped = q_heads->as_strided( + {batch_size, seq_len, num_attention_heads_, head_dim_}, + {static_cast(seq_len * num_attention_heads_ * head_dim_), + static_cast(num_attention_heads_ * head_dim_), + static_cast(head_dim_), + 1}); + auto k_reshaped = k_heads->as_strided( + {batch_size, seq_len, num_key_value_heads_, head_dim_}, + {static_cast(seq_len * num_key_value_heads_ * head_dim_), + static_cast(num_key_value_heads_ * head_dim_), + static_cast(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(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(head_dim_), 1}); + auto k_reshaped = k->as_strided( + {seq_len, num_key_value_heads_, head_dim_}, + {k->stride(1), static_cast(head_dim_), 1}); + auto v_reshaped = v->as_strided( + {seq_len, num_key_value_heads_, head_dim_}, + {v->stride(1), static_cast(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 + + + diff --git a/csrc/models/minimax/minimax_attention.hpp b/csrc/models/minimax/minimax_attention.hpp new file mode 100644 index 000000000..6d5f1b354 --- /dev/null +++ b/csrc/models/minimax/minimax_attention.hpp @@ -0,0 +1,46 @@ +#pragma once + +#include "../../layers/common_modules.hpp" +#include + +namespace infinilm::models::minimax { + +class MiniMaxAttention : public infinicore::nn::Module { +public: + MiniMaxAttention(std::shared_ptr 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 qkv_proj_; + std::shared_ptr o_proj_; + std::shared_ptr rotary_emb_; + std::shared_ptr 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 diff --git a/csrc/models/minimax/minimax_decoderLayer.cpp b/csrc/models/minimax/minimax_decoderLayer.cpp new file mode 100644 index 000000000..665334615 --- /dev/null +++ b/csrc/models/minimax/minimax_decoderLayer.cpp @@ -0,0 +1,76 @@ +#include "minimax_decoderLayer.hpp" + +#include +#include +#include +#include +#include + +namespace infinilm::models::minimax { + +MiniMaxDecoderLayer::MiniMaxDecoderLayer(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_or("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 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("infinilm::models::minimax::MiniMaxDecoderLayer: unsupported layer_type '" + layer_type_ + "' for layer " + std::to_string(layer_idx)); + } + + num_experts_ = model_config->get_or("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 + + diff --git a/csrc/models/minimax/minimax_decoderLayer.hpp b/csrc/models/minimax/minimax_decoderLayer.hpp new file mode 100644 index 000000000..107372f88 --- /dev/null +++ b/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 + +namespace infinilm::models::minimax { + +class MiniMaxDecoderLayer : public infinicore::nn::Module { +public: + MiniMaxDecoderLayer(std::shared_ptr 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 diff --git a/csrc/models/minimax/minimax_for_causal_lm.cpp b/csrc/models/minimax/minimax_for_causal_lm.cpp new file mode 100644 index 000000000..70a7a3c9a --- /dev/null +++ b/csrc/models/minimax/minimax_for_causal_lm.cpp @@ -0,0 +1,251 @@ +#include "minimax_for_causal_lm.hpp" + +#include "../../global_state/global_state.hpp" +#include "../models_registry.hpp" + +#include +#include +#include +#include +#include + +namespace infinilm::models::minimax { + +MiniMaxModel::MiniMaxModel(std::shared_ptr model_config, + const infinicore::Device &device) { + const auto &dtype{model_config->get_dtype()}; + const size_t vocab_size = model_config->get("vocab_size"); + const size_t hidden_size = model_config->get("hidden_size"); + const size_t num_hidden_layers = model_config->get("num_hidden_layers"); + const double rms_norm_eps = model_config->get_or("rms_norm_eps", 1e-5); + + INFINICORE_NN_MODULE_INIT(embed_tokens, vocab_size, hidden_size, std::nullopt, dtype, device); + layers_.reserve(num_hidden_layers); + for (size_t i = 0; i < num_hidden_layers; ++i) { + layers_.push_back(this->register_module("layers." + std::to_string(i), model_config, i, device)); + } + INFINICORE_NN_MODULE_INIT(norm, hidden_size, rms_norm_eps, dtype, device); +} + +infinicore::Tensor MiniMaxModel::forward(const infinilm::InfinilmModel::Input &input) const { + auto input_ids = input.input_ids.value(); + if (input_ids->shape().size() == 1) { + input_ids = input_ids->view({1, input_ids->shape()[0]}); + } + auto hidden_states = embed_tokens_->forward(input_ids); + auto positions = input.position_ids.value(); + for (const auto &layer : layers_) { + hidden_states = layer->forward(positions, hidden_states); + } + return norm_->forward(hidden_states); +} + +MiniMaxForCausalLM::MiniMaxForCausalLM(std::shared_ptr model_config, + const infinicore::Device &device) { + model_config_ = model_config; + const auto &dtype{model_config->get_dtype()}; + const size_t hidden_size = model_config->get("hidden_size"); + const size_t vocab_size = model_config->get("vocab_size"); + + INFINICORE_NN_MODULE_INIT(model, model_config, device); + INFINICORE_NN_MODULE_INIT(lm_head, hidden_size, vocab_size, false, dtype, device); +} + +infinilm::InfinilmModel::Output MiniMaxForCausalLM::forward(const infinilm::InfinilmModel::Input &input) const { + auto hidden_states = model_->forward(input); + auto logits = lm_head_->forward(hidden_states); + return {logits}; +} + +void MiniMaxForCausalLM::reset_cache(const cache::CacheConfig *cache_config) { + if (nullptr == cache_config) { + InfinilmModel::reset_cache(nullptr); + 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 size_t num_hidden_layers = model_config_->get("num_hidden_layers"); + const size_t hidden_size = model_config_->get("hidden_size"); + const size_t total_num_heads = model_config_->get("num_attention_heads"); + const size_t head_dim = model_config_->get_or("head_dim", 0); + const size_t resolved_head_dim = head_dim != 0 ? head_dim : hidden_size / total_num_heads; + const size_t num_kv_heads = model_config_->get_or("num_key_value_heads", total_num_heads); + const size_t max_position_embeddings = model_config_->get_or("max_position_embeddings", 8192); + 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"); + const backends::AttentionBackend attention_backend = infinilm::global_state::get_infinilm_config().attention_backend; + + forward_context.kv_cache_vec.reserve(num_hidden_layers); + forward_context.ssm_state_vec.reserve(num_hidden_layers); + + auto allocate_linear_state = [&](size_t layer_idx, size_t pool_size) { + auto state = cache::MambaCache::create_layer_ssm_state( + resolved_head_dim, + resolved_head_dim, + total_num_heads, + total_num_heads, + dtype, + pool_size); + forward_context.kv_cache_vec.emplace_back(); + forward_context.ssm_state_vec.push_back(std::move(state)); + }; + + auto allocate_static_full_attention = [&](size_t layer_idx, const cache::StaticKVCacheConfig &config) { + auto kv_cache = cache::StaticKVCache::create_layer_kv_cache( + resolved_head_dim, + resolved_head_dim, + num_kv_heads, + num_kv_heads, + max_position_embeddings, + kv_cache_dtype, + config); + forward_context.kv_cache_vec.push_back(std::move(kv_cache)); + forward_context.ssm_state_vec.emplace_back(); + }; + + auto allocate_paged_full_attention = [&](size_t layer_idx, const cache::PagedKVCacheConfig &config) { + auto kv_cache = cache::PagedKVCache::create_layer_kv_cache( + resolved_head_dim, + resolved_head_dim, + num_kv_heads, + num_kv_heads, + kv_cache_dtype, + config); + forward_context.kv_cache_vec.push_back(std::move(kv_cache)); + forward_context.ssm_state_vec.emplace_back(); + }; + + switch (attention_backend) { + case backends::AttentionBackend::STATIC_ATTN: { + auto static_kv_cache_config = dynamic_cast(cache_config); + if (nullptr == static_kv_cache_config) { + throw std::runtime_error("infinilm::models::minimax::MiniMaxForCausalLM: invalid static kv cache config type"); + } + for (size_t layer_idx = 0; layer_idx < num_hidden_layers; ++layer_idx) { + if ("linear_attention" == layer_types[layer_idx]) { + allocate_linear_state(layer_idx, static_kv_cache_config->max_batch_size()); + } else { + allocate_static_full_attention(layer_idx, *static_kv_cache_config); + } + } + break; + } + case backends::AttentionBackend::FLASH_ATTN: + case backends::AttentionBackend::PAGED_ATTN: { + auto paged_kv_cache_config = dynamic_cast(cache_config); + if (nullptr == paged_kv_cache_config) { + throw std::runtime_error("infinilm::models::minimax::MiniMaxForCausalLM: invalid paged kv cache config type"); + } + const size_t lightning_pool_size = std::max(2, paged_kv_cache_config->num_blocks() / 4); + for (size_t layer_idx = 0; layer_idx < num_hidden_layers; ++layer_idx) { + if ("linear_attention" == layer_types[layer_idx]) { + allocate_linear_state(layer_idx, lightning_pool_size); + } else { + allocate_paged_full_attention(layer_idx, *paged_kv_cache_config); + } + } + break; + } + default: + throw std::runtime_error("infinilm::models::minimax::MiniMaxForCausalLM: Unsupported attention backend: " + std::to_string(static_cast(attention_backend))); + } +} + +std::shared_ptr create_minimax_model_config(std::shared_ptr model_config) { + const std::string model_type = model_config->get("model_type"); + if ("minimax" != model_type && "minimax_m2" != model_type) { + throw std::runtime_error("infinilm::models::minimax::create_minimax_model_config: model_type is not minimax/minimax_m2"); + } + + nlohmann::json &config_json = model_config->get_config_json(); + + // --- Normalize basic transformer fields --- + config_json["num_key_value_heads"] = config_json.value("num_key_value_heads", config_json.value("num_attention_heads", 0)); + if (!config_json.contains("head_dim")) { + config_json["head_dim"] = config_json["hidden_size"].get() / config_json["num_attention_heads"].get(); + } + config_json["rms_norm_eps"] = config_json.value("rms_norm_eps", config_json.value("layer_norm_epsilon", 1e-5)); + config_json["max_position_embeddings"] = config_json.value("max_position_embeddings", config_json.value("max_model_len", 8192)); + // transformers MiniMax keeps rope theta inside `rope_parameters`; flatten it. + if (!config_json.contains("rope_theta")) { + if (config_json.contains("rope_parameters") && config_json["rope_parameters"].is_object() && + config_json["rope_parameters"].contains("rope_theta")) { + config_json["rope_theta"] = config_json["rope_parameters"]["rope_theta"]; + } else { + config_json["rope_theta"] = 1000000.0; + } + } + + // --- Decode per-layer attention types into `layer_types` --- + if (!config_json.contains("layer_types")) { + const size_t num_hidden_layers = config_json["num_hidden_layers"].get(); + std::vector layer_types; + layer_types.reserve(num_hidden_layers); + if (config_json.contains("attn_type_list") && config_json["attn_type_list"].is_array()) { + for (const auto &v : config_json["attn_type_list"]) { + layer_types.push_back(v.get() == 0 ? "linear_attention" : "full_attention"); + } + } else if (config_json.contains("decoder_attention_types") && config_json["decoder_attention_types"].is_array()) { + for (const auto &v : config_json["decoder_attention_types"]) { + layer_types.push_back(v.get()); + } + } else if (config_json.contains("full_attention_interval")) { + const size_t interval = config_json["full_attention_interval"].get(); + for (size_t i = 0; i < num_hidden_layers; ++i) { + layer_types.push_back(bool((i + 1) % interval) ? "linear_attention" : "full_attention"); + } + } else { + // MiniMax-01 default: one softmax attention layer after every 7 lightning layers. + for (size_t i = 0; i < num_hidden_layers; ++i) { + layer_types.push_back(bool((i + 1) % 8) ? "linear_attention" : "full_attention"); + } + } + if (layer_types.size() != num_hidden_layers) { + throw std::runtime_error("infinilm::models::minimax::create_minimax_model_config: layer_types length mismatch"); + } + config_json["layer_types"] = layer_types; + } + + // --- Lightning attention knobs --- + config_json["block"] = config_json.value("block", 256); + + // --- MoE defaults --- + config_json["hidden_act"] = config_json.value("hidden_act", "silu"); + config_json["num_experts_per_tok"] = config_json.value("num_experts_per_tok", 1); + // HF transformers `minimax` uses `num_local_experts`; normalize the alias. + if (config_json.contains("num_local_experts") && !config_json.contains("num_experts")) { + config_json["num_experts"] = config_json["num_local_experts"]; + } else { + config_json["num_experts"] = config_json.value("num_experts", 1); + } + // Expert FFN dim. transformers `minimax` uses `intermediate_size` for the experts. + config_json["moe_intermediate_size"] = config_json.value("moe_intermediate_size", config_json["intermediate_size"]); + // transformers renormalizes the top-k softmax weights over the selected experts. + config_json["norm_topk_prob"] = config_json.value("norm_topk_prob", true); + config_json["moe_router_backend"] = config_json.value("moe_router_backend", "softmax"); + config_json["shared_intermediate_size"] = config_json.value("shared_intermediate_size", 0); + + return model_config; +} + +} // namespace infinilm::models::minimax + +namespace { +INFINILM_REGISTER_CAUSAL_LM_MODEL( + minimax, + infinilm::models::minimax::MiniMaxForCausalLM, + infinilm::models::minimax::create_minimax_model_config); +INFINILM_REGISTER_CAUSAL_LM_MODEL( + minimax_m2, + infinilm::models::minimax::MiniMaxForCausalLM, + infinilm::models::minimax::create_minimax_model_config); +} // namespace + + + diff --git a/csrc/models/minimax/minimax_for_causal_lm.hpp b/csrc/models/minimax/minimax_for_causal_lm.hpp new file mode 100644 index 000000000..fd733b022 --- /dev/null +++ b/csrc/models/minimax/minimax_for_causal_lm.hpp @@ -0,0 +1,40 @@ +#pragma once + +#include "../../cache/mamba_cache.hpp" +#include "../infinilm_model.hpp" +#include "minimax_decoderLayer.hpp" +#include +#include + +namespace infinilm::models::minimax { + +class MiniMaxModel : public infinicore::nn::Module { +public: + MiniMaxModel(std::shared_ptr model_config, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinilm::InfinilmModel::Input &input) const; + +protected: + INFINICORE_NN_MODULE(infinicore::nn::Embedding, embed_tokens); + std::vector> layers_; + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, norm); +}; + +class MiniMaxForCausalLM : public InfinilmModel { +public: + MiniMaxForCausalLM(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(MiniMaxModel, model); + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, lm_head); +}; + +std::shared_ptr create_minimax_model_config(std::shared_ptr model_config); + +} // namespace infinilm::models::minimax diff --git a/csrc/models/minimax/minimax_lightning_attention.cpp b/csrc/models/minimax/minimax_lightning_attention.cpp new file mode 100644 index 000000000..389483e39 --- /dev/null +++ b/csrc/models/minimax/minimax_lightning_attention.cpp @@ -0,0 +1,180 @@ +#include "minimax_lightning_attention.hpp" + +#include "../../global_state/global_state.hpp" +#include "infinicore/context/context.hpp" +#include "../../utils.hpp" + +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace infinilm::models::minimax { + +std::vector MiniMaxLightningAttention::build_slopes(size_t num_heads) { + // Same construction as HF transformers `minimax` (MiniMaxLightningAttention): + // base = 1 / 2^(8/H); rate[h] = base^(h+1) + const float base = 1.0f / std::pow(2.0f, 8.0f / static_cast(num_heads)); + std::vector slopes(num_heads); + for (size_t h = 0; h < num_heads; ++h) { + slopes[h] = std::pow(base, static_cast(h + 1)); + } + return slopes; +} + +MiniMaxLightningAttention::MiniMaxLightningAttention(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 size_t total_num_heads = model_config->get("num_attention_heads"); + head_dim_ = model_config->get_or("head_dim", 0); + if (head_dim_ == 0) { + head_dim_ = hidden_size / total_num_heads; + } + block_size_ = model_config->get_or("block", 256); + const std::string hidden_act = model_config->get_or("hidden_act", "silu"); + if (hidden_act != "silu") { + throw std::runtime_error("MiniMaxLightningAttention: unsupported hidden_act '" + hidden_act + "'"); + } + silu_act_ = true; + const size_t num_hidden_layers = model_config->get("num_hidden_layers"); + const double rms_norm_eps = model_config->get_or("linear_rms_norm_eps", 1e-6); + + const auto &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info(); + num_heads_ = total_num_heads / rank_info.tp_size; + inner_dim_ = num_heads_ * head_dim_; + + qkv_proj_ = this->register_module( + "qkv_proj", hidden_size, inner_dim_ * 3, false, dtype, device, rank_info.tp_rank, rank_info.tp_size); + output_gate_ = this->register_module( + "output_gate", hidden_size, inner_dim_, false, dtype, device, rank_info.tp_rank, rank_info.tp_size); + out_proj_ = this->register_module( + "out_proj", inner_dim_, hidden_size, model_config->get_quantization_method(), false, dtype, device, + rank_info.tp_rank, rank_info.tp_size, rank_info.comm); + INFINICORE_NN_MODULE_INIT(norm, inner_dim_, rms_norm_eps, dtype, device); + + // Per-head ALiBi slope, decayed by layer index (same as MiniMax-01 / vLLM). + const std::vector slopes = build_slopes(total_num_heads); + float layer_scale = 1.0f; + if (num_hidden_layers > 1) { + // transformers MiniMax: factor = 1 - layer_idx/(num_hidden_layers - 1 + 1e-5) + 1e-5 + layer_scale = 1.0f - static_cast(layer_idx_) / (static_cast(num_hidden_layers - 1) + 1e-5f) + 1e-5f; + } + auto slope_cpu = infinicore::Tensor::empty({total_num_heads}, infinicore::DataType::F32, infinicore::Device::cpu()); + auto *slope_data = reinterpret_cast(slope_cpu->data()); + auto ratio_cpu = infinicore::Tensor::empty({total_num_heads}, infinicore::DataType::F32, infinicore::Device::cpu()); + auto *ratio_data = reinterpret_cast(ratio_cpu->data()); + for (size_t i = 0; i < total_num_heads; ++i) { + slope_data[i] = slopes[i] * layer_scale; + ratio_data[i] = std::exp(-slope_data[i]); + } + slope_ = slope_cpu->to(device); + ratio_ = ratio_cpu->to(device); + if (rank_info.tp_size > 1) { + slope_ = slope_->narrow({{0, static_cast(rank_info.tp_rank) * num_heads_, num_heads_}}); + ratio_ = ratio_->narrow({{0, static_cast(rank_info.tp_rank) * num_heads_, num_heads_}}); + } +} + +namespace { +// Build a one-element int32 index tensor for a single-request op call. +infinicore::Tensor make_index_tensor(size_t value, const infinicore::Device &device) { + auto cpu = infinicore::Tensor::empty({1}, infinicore::DataType::I32, infinicore::Device::cpu()); + reinterpret_cast(cpu->data())[0] = static_cast(value); + return cpu->to(device); +} +} // namespace + +infinicore::Tensor MiniMaxLightningAttention::forward(const infinicore::Tensor &hidden_states) const { + auto shape = hidden_states->shape(); + const size_t batch_size = shape[0]; + const size_t seq_len = shape[1]; + + auto hidden_mutable = hidden_states; + auto qkv = infinicore::op::silu(qkv_proj_->forward(hidden_mutable)); // [B, T, 3 * inner] + auto qkv4 = qkv->view({batch_size, seq_len, num_heads_, 3 * head_dim_}); + auto q = qkv4->narrow({{3, 0, head_dim_}}); // [B, T, H, D] + auto k = qkv4->narrow({{3, head_dim_, head_dim_}}); + auto v = qkv4->narrow({{3, 2 * head_dim_, head_dim_}}); + + auto &forward_context = infinilm::global_state::get_forward_context(); + const auto &mamba_metadata = forward_context.mamba_metadata; + if (!mamba_metadata.input_offsets.has_value() || + !mamba_metadata.init_state_indices.has_value() || + !mamba_metadata.final_state_indices.has_value()) { + throw std::runtime_error("MiniMaxLightningAttention: linear attention requires mamba state indices"); + } + if (forward_context.ssm_state_vec.size() <= layer_idx_ || !forward_context.ssm_state_vec[layer_idx_]) { + throw std::runtime_error("MiniMaxLightningAttention: lightning state cache is not allocated for layer " + std::to_string(layer_idx_)); + } + auto state_pool = forward_context.ssm_state_vec[layer_idx_]; + const auto &init_indices = mamba_metadata.init_state_indices.value(); + const auto &final_indices = mamba_metadata.final_state_indices.value(); + + const bool is_decode = mamba_metadata.input_offsets.value()->shape()[0] - 1 == seq_len; + auto attn_out = infinicore::Tensor::empty( + {batch_size, seq_len, num_heads_, head_dim_}, state_pool->dtype(), state_pool->device()); + if (is_decode) { + // Batched decode: one token per request, one op call. + infinicore::op::lightning_attention_( + attn_out, state_pool, q, k, v, slope_, init_indices, final_indices); + } else { + // Prefill: requests may have different lengths, so run one op call per request. + // The offsets/indices may live on an accelerator; copy them to the host + // first (the per-request loop itself is device-agnostic). + auto cpu_offsets = mamba_metadata.input_offsets.value(); + auto cpu_init = init_indices; + auto cpu_final = final_indices; + if (cpu_offsets->device().getType() != infinicore::Device::Type::CPU) { + cpu_offsets = cpu_offsets->to(infinicore::Device::cpu()); + } + if (cpu_init->device().getType() != infinicore::Device::Type::CPU) { + cpu_init = cpu_init->to(infinicore::Device::cpu()); + } + if (cpu_final->device().getType() != infinicore::Device::Type::CPU) { + cpu_final = cpu_final->to(infinicore::Device::cpu()); + } + infinicore::context::syncStream(); + auto read_index = [](const infinicore::Tensor &t, size_t i) -> size_t { + if (t->dtype() == infinicore::DataType::I32) { + return static_cast(reinterpret_cast(t->data())[i]); + } + return static_cast(reinterpret_cast(t->data())[i]); + }; + const auto *offsets_ptr = reinterpret_cast(cpu_offsets->data()); + const size_t num_requests = cpu_offsets->shape()[0] - 1; + for (size_t r = 0; r < num_requests; ++r) { + const size_t start = static_cast(offsets_ptr[r]); + const size_t end = static_cast(offsets_ptr[r + 1]); + const size_t len = end - start; + if (len == 0) { + continue; + } + auto q_r = q->narrow({{0, 0, 1}})->narrow({{1, start, len}}); + auto k_r = k->narrow({{0, 0, 1}})->narrow({{1, start, len}}); + auto v_r = v->narrow({{0, 0, 1}})->narrow({{1, start, len}}); + auto out_r = attn_out->narrow({{0, 0, 1}})->narrow({{1, start, len}}); + auto init_r = make_index_tensor(read_index(cpu_init, r), state_pool->device()); + auto final_r = make_index_tensor(read_index(cpu_final, r), state_pool->device()); + infinicore::op::lightning_attention_(out_r, state_pool, q_r, k_r, v_r, slope_, init_r, final_r); + } + } + + auto attn_flat = attn_out->view({batch_size, seq_len, inner_dim_}); + auto normed = norm_->forward(attn_flat); + auto gate = infinicore::op::sigmoid(output_gate_->forward(hidden_mutable)); + auto gated = infinicore::op::mul(normed, gate); + return out_proj_->forward(gated); +} + +} // namespace infinilm::models::minimax + + + diff --git a/csrc/models/minimax/minimax_lightning_attention.hpp b/csrc/models/minimax/minimax_lightning_attention.hpp new file mode 100644 index 000000000..20432130b --- /dev/null +++ b/csrc/models/minimax/minimax_lightning_attention.hpp @@ -0,0 +1,42 @@ +#pragma once + +#include "../../layers/common_modules.hpp" +#include +#include +#include + +namespace infinilm::models::minimax { + +class MiniMaxLightningAttention : public infinicore::nn::Module { +public: + MiniMaxLightningAttention(std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const; + + size_t layer_idx() const { return layer_idx_; } + size_t num_heads() const { return num_heads_; } + size_t head_dim() const { return head_dim_; } + +private: + static std::vector build_slopes(size_t num_heads); + + std::shared_ptr qkv_proj_; + std::shared_ptr output_gate_; + std::shared_ptr out_proj_; + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, norm); + + infinicore::Tensor slope_; + infinicore::Tensor ratio_; + + size_t layer_idx_; + size_t num_heads_; + size_t head_dim_; + size_t inner_dim_; + bool silu_act_{true}; + size_t block_size_; +}; + +} // namespace infinilm::models::minimax + diff --git a/csrc/models/minimax/minimax_moe.cpp b/csrc/models/minimax/minimax_moe.cpp new file mode 100644 index 000000000..d84c4a236 --- /dev/null +++ b/csrc/models/minimax/minimax_moe.cpp @@ -0,0 +1,34 @@ +#include "minimax_moe.hpp" + +#include + +namespace infinilm::models::minimax { + +MiniMaxMoeBlock::MiniMaxMoeBlock(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); +} + +infinicore::Tensor MiniMaxMoeBlock::forward(const infinicore::Tensor &hidden_states) const { + auto shape = hidden_states->shape(); + auto hidden_flat = hidden_states->view({shape[0] * shape[1], shape[2]}); + + auto [routing_weights, selected_experts] = gate_->forward(hidden_flat); + infinilm::layers::moe::TopKOutput topk_output{ + routing_weights, + selected_experts, + infinicore::Tensor(), + }; + auto routed_states = fused_moe_->forward(hidden_flat, topk_output, experts_->moe_weights()); + + return routed_states->as_strided( + {shape[0], shape[1], shape[2]}, + {static_cast(shape[1] * shape[2]), + static_cast(shape[2]), + 1}); +} + +} // namespace infinilm::models::minimax diff --git a/csrc/models/minimax/minimax_moe.hpp b/csrc/models/minimax/minimax_moe.hpp new file mode 100644 index 000000000..e108a1a07 --- /dev/null +++ b/csrc/models/minimax/minimax_moe.hpp @@ -0,0 +1,30 @@ +#pragma once + +#include "../../layers/common_modules.hpp" +#include "../../layers/moe/experts/fused_moe_experts.hpp" +#include "../../layers/moe/fused_moe.hpp" +#include "../../layers/moe/router/topk_router.hpp" + +#include + +namespace infinilm::models::minimax { + +// Block-sparse MoE for MiniMax (matching HF transformers `minimax`): +// softmax top-k router + gate_up_proj/down_proj experts. No shared MLP in the +// transformers `minimax` reference (shared_mlp + coefficient is a follow-up for +// MiniMax-M2 style configs). +class MiniMaxMoeBlock : public infinicore::nn::Module { +public: + MiniMaxMoeBlock(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_; +}; + +} // namespace infinilm::models::minimax diff --git a/docs/minimax/HANDOVER.md b/docs/minimax/HANDOVER.md new file mode 100644 index 000000000..bfd9dd5e6 --- /dev/null +++ b/docs/minimax/HANDOVER.md @@ -0,0 +1,210 @@ +# MiniMax 支持：归档与交接说明 + +> 本文记录「为 InfiniCore 重构做归档准备」这一轮（任务 A）实际做了什么、产出了什么、你还需要执行什么。 +> 技术迁移细节见 `docs/minimax/PORTING.md`；算子补丁见 `docs/minimax/lightning-attention-infinicore.patch`。 + +## 1. 现状：重构已经落地 + +撰写本文时已核实上游状态（不是预告，是已发生）： + +| 事实 | 证据 | +|---|---| +| InfiniCore 已完成重构 | `26f7382d refactor!: reduce InfiniCore to unified component architecture (#1406)`，提交时间 **2026-09-11 17:00 +0800** | +| 重构后顶层只剩集成件 | `26f7382d` 的顶层树仅含 `.gitmodules`、`CONTRIBUTING.md`、`LICENSE`、`README.md`、`submodules`；`include/`、`src/`、`xmake.lua`、CI workflow 全部删除 | +| 算子实现归属 InfiniOps | 该提交把 `submodules/InfiniOps` pin 在 `f890afb4b2327f13ccdd3c1b6b0d49567c5fe00d` | +| 我们落后 21 个提交 | 本地/归档基线 `35b46277`，`origin/main` 已是 `26f7382d` | +| **InfiniLM 上游尚未适配新版** | 上游 `InfiniTensor/InfiniLM` main = `270feb3e`，仅比我们的基线 `80bb09e` 多 5 个提交，且全是无关改动（readme、hygon、metax、多模态懒加载），没有任何 InfiniOps 适配 | + +**由此得出的三条结论** + +1. 我们写的 `lightning_attention`（位于 InfiniCore 顶层）**不会在重构后存活**，必须按 InfiniOps 规范重新落地。 +2. 现在**整个 InfiniLM 主干都跑不了新 InfiniCore**（不只是 minimax），所以「适配新版」这件事目前**无法端到端验证**，性价比低、建议等上游跟进。 +3. 我们的成果要能复现与评审，必须**锁定归档基线**：InfiniCore `35b46277` + 本补丁 ↔ InfiniLM `80bb09e` + MiniMax 改动。 + +## 2. 本轮实际执行的动作 + +### 2.1 状态巡检 + +对两个仓库做了完整的未提交状态巡检，把改动分成三类： + +| 类别 | InfiniCore | InfiniLM | +|---|---|---| +| **本次 MiniMax 工作** | 4 个注册文件改动 + 6 个新增路径 | `python/infinilm/modeling_utils.py`（+52 行）、`csrc/models/minimax/`（10 个文件）、`test/models/minimax/`（3 个脚本） | +| **早先的 Windows 本地编译补丁**（与 MiniMax 无关） | `xmake.lua`（export_all） | `csrc/engine/distributed/tcp_rendezvous.cpp`、`csrc/models/kimi_k3/kimi_k3_pipeline_partition.hpp`、`csrc/models/minicpmv/minicpmv_model.cpp`、`xmake.lua` | +| **应排除的产物/垃圾** | `python/infinicore/bin/`、`python/infinicore/lib/` 下的 pyd/dll/lib | `python/infinilm/bin/`、`python/infinilm/lib/*.pyd|*.dll|*.lib`、根目录 `comp2011_1.cpp`（36 字节无关草稿） | + +### 2.2 修复 `.gitignore`（两个仓库） + +原规则只忽略 `*.so`（Linux 产物），Windows 构建产物会长期挂在 `git status` 里、随时可能被 `git add -A` 误提交。已在两个仓库补齐： + +``` +python//lib/*.pyd +python//lib/*.dll +python//lib/*.lib +python//lib/*.dylib +python//bin/ +``` + +效果：两边构建产物已从 `git status` 消失，只剩真正的源码改动。 + +### 2.3 生成算子补丁 + +用 `git add -N`（intent-to-add）把新增文件纳入 `git diff`，再用 **git 自带的 `--output`** 写文件（避免 PowerShell 重定向写出 UTF-16/BOM 污染补丁），最后 `git reset -q` 撤销 intent-to-add，**索引与工作区均未被改动**。 + +补丁覆盖 **17 个文件 / +990 行 / 45,953 字节**： + +``` +include/infiniop.h (注册) +include/infiniop/ops/lightning_attention.h (C API) +include/infinicore/ops.hpp (注册) +include/infinicore/ops/lightning_attention.hpp (C++ 包装) +python/infinicore/__init__.py (注册) +python/infinicore/ops/lightning_attention.py (Python 包装) +src/infiniop/ops/lightning_attention/{info.h,lightning_attention.h,operator.cc} +src/infiniop/ops/lightning_attention/cpu/* (CPU 参考实现) +src/infiniop/ops/lightning_attention/nvidia/* (CUDA kernel + launcher) +src/infinicore/ops/lightning_attention/*.cc (图/调度注册) +src/infinicore/pybind11/ops.hpp + ops/lightning_attention.hpp +``` + +**刻意不包含**：`xmake.lua`（早先的 Windows 本地补丁）、`.gitignore`（仓库卫生修复）——两者都不属于算子贡献，混入会污染将来的 InfiniOps 迁移评审。 + +校验方式： + +- 在 InfiniCore 工作区对补丁执行**反向应用检查** `git apply --check -R`，退出码 0（说明补丁与工作区改动完全一致，因而对基线提交可正向应用）。 +- 字节级检查：补丁仅含 LF（0 个 CR 字节、1134 个 LF），可直接在 Linux 服务器上 `git apply`。 + +### 2.4 CUDA kernel 静态复核（本机无法编译 CUDA，只能逐行审） + +复核了 `nvidia/lightning_attention_nvidia.cu`、`nvidia_*.cuh`、`operator.cc`、`info.h`，并与仓库既有算子（`recurrent_gated_delta_rule`、`causal_softmax`）逐项对照。**发现并修复 6 处问题**，其中第 1 条是真实语义缺陷： + +1. **【语义 bug】状态污染初始行**：原 kernel 让 `S` 直接指向状态池的 `init_row` 并就地累加，最后才把结果拷到 `final_row`。当 `init_row != final_row` 时，`init_row` 被污染；而 CPU 参考实现是把初始状态读进本地缓冲、只写回 `final_row` —— **两者语义不一致，且我们的算子单测正好覆盖该场景**（GPU 上会失败）。 + **修复**：先把 `init_row` 协作式拷贝到 `final_row`（配 `__syncthreads()`），再在 `final_row` 上就地累加。初始行保持只读，CPU/CUDA 语义一致，且不增加额外内存。 +2. **【可移植性】改用 `INFINIOP_CUDA_KERNEL` 宏**：仓库约定用该宏（Hygon 构建下展开为 `__launch_bounds__(1024) __global__ void`），同时补上 `nvidia_kernel_common.cuh` 与 `` 头。 +3. **【健壮性】`slope` 步长**：原 kernel 假设 `slope` 连续（`slope[h]`）。改为传入 `slope_stride` 按步长取值，不再依赖该假设（TP>1 时 `slope` 是 narrow 视图）。 +4. **【潜在死锁】去掉 `if (tid >= D) return;`**：该早退分支位于 `__syncthreads()` 之前，一旦块大小不等于 `D` 就会死锁。改为在 `calculate` 中校验 `D <= maxThreadsPerBlock()`、固定按 `D` 个线程启动，并以注释写明该不变量。 +5. **【编译阻塞项】signed/unsigned 比较**：`_opaque->internal->maxThreadsPerBlock()` 返回 `int`，而 `_info.D` / `_info.B` 是 `size_t`。写成 `size_t > int` 会在 InfiniCore 的 NVIDIA 构建下（`-Xcompiler=-Wall -Werror`）**直接编译失败**。已改为显式 `static_cast(...)` 比较，并把 kernel 的 `slope_stride` 形参改成 `size_t`，避免设备侧混合符号运算。 +6. **【输入校验】** `info.h` 增加索引张量 `stride(0) == 1` 的校验（CPU/CUDA 均按单位步长读索引），避免静默读错状态行；`calculate` 另加 `B <= 65535`（`gridDim.y` 上限）检查与 `(void)workspace` 显式消警。 + +另核实两点（非缺陷）： + +- InfiniLM 的 `mamba_init/final_state_indices` 在各处理器中显式以 **int32** 构造（如 `qwen3_next_processor.py`：`infinicore.from_list(..., dtype=infinicore.int32)`），因此 CUDA 侧 i32-only 与真实调用路径一致；i64 仅 CPU 参考实现支持，CUDA 侧会明确返回 `BAD_TENSOR_DTYPE`。 +- `PagedAttentionInfinilm` 等既有算子确认 `INFINIOP_CUDA_KERNEL`/`nvidia_kernel_common.cuh` 的用法一致，改动符合仓库现状。 + +### 2.5 修改后的回归验证（本机 CPU） + +| 验证项 | 结果 | +|---|---| +| 算子 vs numpy 参考（decode `B=3,T=1`；prefill `B=2,T=6`） | 输出误差 9.54e-7 / 1.91e-6；状态误差 3.58e-7 / 2.38e-7 | +| 模型状态连续性（`prefill(5)` vs `prefill(4)+decode`） | 1.16e-10 | +| vs HF transformers `MiniMaxForCausalLM`（prefill / decode） | 5.59e-4 / 1.59e-4 | + +三个测试全部通过，数值与修复前一致（说明修复未影响已验证的 CPU 语义）。 + +## 3. 产物清单 + +| 文件 | 说明 | +|---|---| +| `docs/minimax/lightning-attention-infinicore.patch` | 45,953 字节，17 文件 / +990 行，基线 InfiniCore `35b46277`，LF-only | +| `docs/minimax/PORTING.md` | 迁移到 InfiniOps 的对照与清单 | +| `docs/minimax/HANDOVER.md` | 本文件 | + +## 4. 你需要执行的提交与推送 + +两个仓库分别操作。**路径列表刻意写全**，避免 `git add -A` 把垃圾文件或无关补丁一起提交。 + +### 4.1 InfiniCore + +```bash +cd <你的 InfiniCore 工作区> +git checkout -b archive/minimax-lightning-attn + +git add .gitignore \ + include/infiniop.h include/infinicore/ops.hpp \ + src/infinicore/pybind11/ops.hpp python/infinicore/__init__.py \ + include/infiniop/ops/lightning_attention.h \ + include/infinicore/ops/lightning_attention.hpp \ + src/infiniop/ops/lightning_attention \ + src/infinicore/ops/lightning_attention \ + src/infinicore/pybind11/ops/lightning_attention.hpp \ + python/infinicore/ops/lightning_attention.py + +git commit -m "feat(infiniop): add lightning_attention op (CPU + NVIDIA) for MiniMax + +Implements indexed-pool lightning attention (MiniMax-01 style): + S <- ratio * S + k^T v ; o = q * S with ratio[h] = exp(-slope[h]) +CPU reference implementation plus a CUDA kernel, wired through the +infiniop C API, the infinicore C++ op layer and the Python bindings." + +git tag minimax-lightning-attn +git push -u origin archive/minimax-lightning-attn +git push origin minimax-lightning-attn +``` + +`xmake.lua` 的本地 Windows 补丁**没有**包含在上面；如需保留本地构建能力，请单独提交（例如 `chore(build): export all symbols for MSVC builds`）。 + +### 4.2 InfiniLM + +```bash +cd <你的 InfiniLM 工作区> +git checkout -b archive/minimax + +git add .gitignore \ + python/infinilm/modeling_utils.py \ + csrc/models/minimax test/models/minimax docs/minimax + +git commit -m "feat(minimax): support MiniMax-Text-01 (lightning attention + MoE) + +- csrc/models/minimax: MiniMax model (hybrid lightning/softmax attention, + block-sparse MoE, dense fallback), registered as minimax/minimax_m2 +- python/infinilm/modeling_utils.py: _remap_minimax weight remapper +- test/models/minimax: op unit test, model smoke test, HF-aligned E2E test +- docs/minimax: porting guide and InfiniCore op patch" + +git tag minimax-support +git push -u origin archive/minimax +git push origin minimax-support +``` + +`csrc/engine/distributed/tcp_rendezvous.cpp`、`csrc/models/kimi_k3/...`、`csrc/models/minicpmv/...`、`xmake.lua` 是早先的 Windows 本地补丁，建议另起提交，不要与 MiniMax 改动混在一起。 + +## 5. 服务器验证：必须锁版本 + +因为 `origin/main` 已经是重构版，服务器上**不能直接用 main**： + +```bash +# InfiniCore：用重构前基线 + 我们的补丁（或直接 checkout 你的 archive 分支） +git clone --recursive https://github.com/<你的账号>/InfiniCore.git +cd InfiniCore +git checkout 35b46277bd666772c11bb417ad4231c5be492822 +git apply /path/to/lightning-attention-infinicore.patch + +# InfiniLM：checkout 你的 archive/minimax 分支（含 minimax 代码与 docs） +``` + +推荐节奏（省租机成本）： + +1. **第一小时只做编译验证**：搭好环境 → 编译 InfiniCore（这一步就能暴露 CUDA kernel 是否可编译，是本机唯一无法验证的部分）。 +2. 编译通过后再续租，跑：`test_lightning_attention_op.py` → `smoke_minimax.py` → `test_minimax_vs_hf.py`。 +3. 若要测多专家 MoE，配 `num_experts=4, num_experts_per_tok=2` 之类的小配置（NVIDIA 上 `FusedMoE` 的 CUDA runner 可用）。 + +**CI 注意**：`.github/ci_config.yaml` 的 nvidia 镜像构建使用 `InfiniCore_BRANCH: __Branch_Name__`。若你的分支名在 InfiniCore 侧不存在，CI 会去拿上游 main（重构版）→ 必然编译失败。要么把 InfiniCore 的 archive 分支推到你的 fork 并使用同名分支，要么显式指定分支。 + +## 6. 三个需要特别注意的坑 + +1. **`comp2011_1.cpp`（InfiniLM 根目录）** 曾是一段 36 字节无关草稿（`for(int i =0; i 配套补丁：`docs/minimax/infiniops-lightning-attention.patch`（CPU + CUDA + pytest，6 文件 / +619 行） +> 旧架构补丁（归档基线用）：`docs/minimax/lightning-attention-infinicore.patch` + +## 1. 为什么要迁到 InfiniOps + +InfiniCore 已完成重构（`26f7382d refactor!: reduce InfiniCore to unified component architecture (#1406)`，2026-09-11）：顶层只剩 `.gitmodules` / `CONTRIBUTING.md` / `LICENSE` / `README.md` / `submodules`，算子实现归属 **InfiniOps**。 + +- 旧的 `infinicore::` C++ API（`op::*`、`nn::*`、`Tensor`、graph）在新版顶层已不存在。 +- **InfiniLM 主干尚未迁移**（上游 `270feb3e`），因此现在无法把 minimax **模型**迁到新版；能且应该迁移的是**算子**。 +- 本补丁交付 `lightning_attention_infinilm` 的 **CPU + NVIDIA CUDA** 后端与 pytest。 + +## 2. 补丁内容（6 个文件） + +| 文件 | 说明 | +|---|---| +| `src/base/lightning_attention_infinilm.h` | 算子类（接口 + 元数据 + assert 校验），所有后端共享 | +| `src/native/cpu/ops/lightning_attention_infinilm/lightning_attention_infinilm.h` | CPU 参考实现（float32 累加，支持 f32/f16/bf16 与任意 stride） | +| `src/native/cuda/ops/lightning_attention_infinilm/kernel.cuh` | CUDA device kernel（每 `(batch, head)` 一块，每线程一列状态） | +| `src/native/cuda/ops/lightning_attention_infinilm/kernel.h` | CUDA launcher（`CudaLightningAttentionInfinilm`，按 dtype/index dtype 分发） | +| `src/native/cuda/nvidia/ops/lightning_attention_infinilm/kernel.h` | NVIDIA vendor 绑定（`Operator<..., kNvidia>`） | +| `tests/test_lightning_attention_infinilm.py` | pytest：4 组形状 × 3 种 dtype，分别断言输出与状态池 | + +## 3. 接口与语义设计 + +### 3.1 对齐目标与命名 + +按 `docs/operator-api-alignment.md` 的对齐顺序（PyTorch → vLLM → SGLang → ONNX → 库级公开封装 → CUDA/vendor → custom），lightning attention 在 PyTorch 中没有对应算子，最接近的公开实现是 Flash-Linear-Attention 的 `fused_recurrent_lightning_attn`。由于本算子带 InfiniLM 专有的"索引状态池"契约，按仓库既有 10+ 个先例（`paged_attention_infinilm`、`causal_softmax_infinilm`…）命名为 **`LightningAttentionInfinilm`**，注释中写明与 `dexp = exp(-slope)` 的对应关系。 + +### 3.2 参数顺序（InfiniOps 规范：输入 → 属性 → 输出） + +```cpp +LightningAttentionInfinilm(q, k, v, slope, + initial_state, initial_state_indices, final_state_indices, + out) +``` + +### 3.3 递推语义 + +``` +ratio[h] = exp(-slope[h]) +S = ratio[h] * S + outer(k_t[h], v_t[h]) # 先更新状态 +out_t[h] = q_t[h] @ S # 再读状态（当前 token 权重为 1） +``` + +请求 `b` 从 `initial_state[initial_state_indices[b]]` 读状态，最终状态写回 `initial_state[final_state_indices[b]]`。 + +**关键契约：初始行不会被修改。** 读/写行相同时表现为就地更新；不同时等价于"先复制到目标行、再累加"。CPU 与 CUDA 实现都遵守，测试特意让读行 ≠ 写行来覆盖该契约（CUDA 侧通过在 kernel 开头把初始行协作拷贝到目标行、再在目标行上累加来保证）。 + +## 4. 服务器上需要下载/安装什么 + +### 4.1 基础依赖（必装） + +| 依赖 | 版本/用途 | 安装方式 | +|---|---|---| +| Linux x86_64 | Ubuntu 22.04 等 | 租的机器自带 | +| **CMake** | ≥ 3.18（两个仓库都是 CMake 工程） | NGC 镜像通常自带，先 `cmake --version`；没有则 `pip install cmake` 或 `apt-get install -y cmake` | +| **C++ 编译器** | gcc-11+ / clang-16+（C++17） | `apt-get install -y build-essential` | +| **OpenMP** | CPU 后端 `find_package(OpenMP REQUIRED)` | `apt-get install -y libgomp1`（多半自带） | +| **Python** | ≥ 3.10 + pip | 镜像自带 | +| Python 包 | `torch`（测试参考实现）、`pytest`、`scikit-build-core` | `pip install torch pytest`；`.[dev]` 会带上构建依赖 | +| **InfiniRT** | InfiniOps 的前置依赖，**必须先装** | `git clone --recursive https://github.com/InfiniTensor/InfiniRT.git` + CMake 安装 | +| **InfiniOps** | 我们的补丁落在这里 | `git clone https://github.com/InfiniTensor/InfiniOps.git` | +| 网络 | configure 阶段访问 GitHub | CUDA 构建会 FetchContent 下载 CUTLASS（固定 commit + SHA256）；CPU 构建不需要 | + +### 4.2 测 NVIDIA GPU 额外需要 + +| 依赖 | 说明 | +|---|---| +| **CUDA Toolkit（nvcc）** | ≥ 12；NGC PyTorch 镜像自带。注意 **InfiniRT 也要用 `-DWITH_NVIDIA=ON` 重新编译安装** | +| **CUTLASS** | 由 CMake `FetchContent` 自动下载，**不需要手动装**；网络受限时要预置或配代理 | +| **PyTorch（CUDA 版）** | 测试参考实现需要；NGC 镜像自带 | + +### 4.3 完整命令（Linux 服务器） + +```bash +# 0) 自检 +cmake --version # >= 3.18 +g++ --version # >= 11 +python3 --version # >= 3.10 +nvcc --version # 只有 GPU 机器需要 + +# 1) 编译安装 InfiniRT（GPU 机器同时打开 WITH_NVIDIA） +git clone --recursive https://github.com/InfiniTensor/InfiniRT.git +cmake -S InfiniRT -B build-rt \ + -DCMAKE_INSTALL_PREFIX=$HOME/infinirt \ + -DWITH_CPU=ON -DWITH_NVIDIA=ON +cmake --build build-rt -j +cmake --install build-rt + +# 2) 取 InfiniOps 并应用补丁 +git clone https://github.com/InfiniTensor/InfiniOps.git +cd InfiniOps +git apply /path/to/docs/minimax/infiniops-lightning-attention.patch + +# 3) 构建安装 InfiniOps（CPU + NVIDIA） +python -m pip install ".[dev]" \ + --config-settings=cmake.define.INFINI_RT_ROOT=$HOME/infinirt \ + --config-settings=cmake.define.WITH_CPU=ON \ + --config-settings=cmake.define.WITH_NVIDIA=ON + +# 4) 跑本算子的 pytest（device fixture 会自动覆盖 cpu 与 cuda） +pytest tests/test_lightning_attention_infinilm.py -v +``` + +- 只想先验 CPU：把第 1、3 步的 `-DWITH_NVIDIA=ON` 去掉即可（更快，也不需要 CUDA/CUTLASS）。 +- 想只编本算子加速 configure/build：追加 `--config-settings=cmake.define.INFINI_OPS_OPS=lightning_attention_infinilm`。 + +## 5. 已知限制与下一步 + +- 本机（Windows，无 CMake）**未编译**本补丁；只做了 `py_compile`、人工复核与补丁反向校验（`git apply --check -R` 通过）。首次在服务器上编译可能有细节需要微调，按报错修即可。 +- CUDA kernel 假定 `head_dim` 能放进一个 block（`head_dim <= Backend::max_block_size`，launcher 里有 assert），典型 MiniMax `head_dim = 128` 满足。 +- **Ascend 后端**下一步：写到 `src/native/ascend/ops/lightning_attention_infinilm/`（InfiniOps 已有 AscendC 自定义 kernel 机制）。 +- **InfiniLM 侧适配**：等上游完成 InfiniLM → 新栈迁移后，把 `MiniMaxLightningAttention` 的调用点换成 `infini::ops::LightningAttentionInfinilm`（一处调用 + 状态池形状对齐），模型其余代码不动。 \ No newline at end of file diff --git a/docs/minimax/PORTING.md b/docs/minimax/PORTING.md new file mode 100644 index 000000000..402493710 --- /dev/null +++ b/docs/minimax/PORTING.md @@ -0,0 +1,149 @@ +# 迁移 `lightning_attention` 到 InfiniOps + +本文件说明如何把 MiniMax 所需的 `lightning_attention` 算子，从**重构前**的 InfiniCore 形态迁移到**重构后**的 `submodules/InfiniOps`。 + +- 配套补丁：`docs/minimax/lightning-attention-infinicore.patch` +- 补丁基线：InfiniCore `35b46277bd666772c11bb417ad4231c5be492822` +- 补丁规模：17 个文件 / +990 行（含 CPU 与 NVIDIA 两套实现） + +## 1. 背景（重构已落地） + +InfiniCore 的重构已经完成并合入上游： + +- 提交：`26f7382d refactor!: reduce InfiniCore to unified component architecture (#1406)`（2026-09-11 17:00 +0800） +- 重构后 InfiniCore 顶层只剩 `.gitmodules` / `CONTRIBUTING.md` / `LICENSE` / `README.md` / `submodules`；`include/`、`src/`、`xmake.lua`、CI workflow 全部移除。 +- 算子实现归属 `submodules/InfiniOps`（该提交 pin 在 `f890afb4b2327f13ccdd3c1b6b0d49567c5fe00d`），要求**接口与主流开源框架一致**、合并要求严格。 +- `InfiniLM` 上游（`270feb3e`）**尚未适配**新版 InfiniCore，因此当前阶段单独移植 minimax 也无法端到端验证。 + +因此当前位于 InfiniCore 顶层的算子实现（`src/infiniop/ops/lightning_attention/` 等）**不会在重构后保留**，需要按 InfiniOps 的规范重新落地。数学推导、CPU 参考实现与 CUDA kernel 主体可以直接复用；需要重写的是接口层与构建/注册部分。 + +## 2. 算子契约（语义不随迁移改变） + +### 2.1 数学语义 + +MiniMax-01 风格的 Lightning Attention（带 ALiBi 式逐头衰减），逐 token 递归： + +``` +ratio[h] = exp(-slope[h]) +S = ratio[h] * S + outer(k_t[h], v_t[h]) # 先更新状态 +o_t[h] = q_t[h] @ S # 再读状态（当前 token 权重为 1） +``` + +即 `S ← ratio ∘ S + kᵀv`，`o = q·S`。注意状态**先更新后读取**，当前 token 对自己的衰减为 0（权重 1），这一点与实现无关，是模型语义的一部分。 + +### 2.2 张量约定（迁移时保持不变） + +| 张量 | 形状 | 说明 | +|---|---|---| +| `out` | `[B, T, H, D]` | 输出，末维连续 | +| `initial_state` | `[pool_size, H, D, D]` | 状态池，F32 累积 | +| `q` / `k` / `v` | `[B, T, H, D]` | 末维连续 | +| `slope` | `[H]` | fp32，逐头衰减系数 | +| `initial_state_indices` | `[B]` | int32/int64，读入状态的行号 | +| `final_state_indices` | `[B]` | int32/int64，写回状态的行号（原位写回池） | + +### 2.3 索引池（indexed-pool）语义 + +每个请求 `b`： + +1. 从 `initial_state[initial_state_indices[b]]` 读取初始状态； +2. 在 `[B, T, H, D]` 上按 token 递归（`T=1` 即 decode）； +3. 把最终状态**原位写回** `initial_state[final_state_indices[b]]`。 + +该语义与 `recurrent_gated_delta_rule` / `chunk_gated_delta_rule` 的 indexed pool 一致，便于 InfiniLM 复用 `mamba_init_state_indices` / `mamba_final_state_indices` 调度机制。 + +## 3. 当前实现（重构前形态） + +| 文件 | 作用 | +|---|---| +| `include/infiniop/ops/lightning_attention.h` | C API：`infiniopCreate/GetWorkspace/Destroy/LightningAttention` | +| `src/infiniop/ops/lightning_attention/info.h` | 描述符校验（形状、dtype、末维连续、索引范围） | +| `src/infiniop/ops/lightning_attention/lightning_attention.h` | `DESCRIPTOR(NAMESPACE)` 宏 | +| `src/infiniop/ops/lightning_attention/operator.cc` | 按设备分发（CPU / NVIDIA） | +| `src/infiniop/ops/lightning_attention/cpu/*` | CPU 参考实现（F32/F16/BF16，支持任意 strides） | +| `src/infiniop/ops/lightning_attention/nvidia/*` | CUDA kernel（**目前仅 F32**）+ launcher | +| `include/infinicore/ops/lightning_attention.hpp` | C++ 包装 `infinicore::op::lightning_attention_` | +| `src/infinicore/ops/lightning_attention/*.cc` | 图/调度注册（`INFINIOP_CACHABLE_DESCRIPTOR`） | +| `src/infinicore/pybind11/ops/lightning_attention.hpp` | Python 绑定 | +| `python/infinicore/ops/lightning_attention.py` | Python 包装 | + +注册点（重构后由代码生成替代）： + +- `include/infiniop.h` 增加一行 `#include "infiniop/ops/lightning_attention.h"` +- `include/infinicore/ops.hpp` 增加 `#include "ops/lightning_attention.hpp"` +- `src/infinicore/pybind11/ops.hpp` 增加 include 与 `bind_lightning_attention(m);` +- `python/infinicore/__init__.py` 增加 `from infinicore.ops.lightning_attention import lightning_attention` + +## 4. 目标形态（InfiniOps）与对照 + +InfiniOps 的算子是一个继承 `Operator` 的类，kernel 与 launcher 分离： + +``` +src/base/.h # 算子类（参数校验 + 元数据） +src/native//[/]ops//kernel.h # launcher +src/native//[/]ops//kernel.cuh # device kernel +``` + +| 当前实现 | InfiniOps 目标 | 需要改什么 | +|---|---|---| +| `include/infiniop/ops/lightning_attention.h`（C API + descriptor） | `src/base/lightning_attention.h`（`class LightningAttention : public Operator`） | **删除 descriptor/创建销毁接口**，改为构造函数做校验 + 记录 strides/元数据 | +| `info.h` 的 `utils::Result` + `CHECK_DTYPE` 校验 | 构造函数中的 `assert` | InfiniOps **禁用异常**；错误信息需含 `__FILE__`/`__LINE__`/`__func__` | +| `operator.cc` 的设备分发 | 由 InfiniOps 的分发机制按 backend 选择实现 | 删除手写分发 | +| `cpu/*.cc` | `src/native/cpu/ops/lightning_attention/` | 复用算法，改为 launcher 形态 | +| `nvidia/*.cu` + `*.cuh` | `src/native/cuda/ops/lightning_attention/{kernel.h, kernel.cuh}` | 复用 kernel 主体；**必须补 fp16/bf16** | +| `include/infinicore/ops.hpp`、pybind、Python 包装 | 生成的 `operator_call_instantiations.h` + `GENERATE_PYTHON_BINDINGS` | 删除手写注册与绑定 | + +参数顺序也要改成 InfiniOps 规范：**输入在前 → 属性居中 → 输出最后**，例如： + +```cpp +LightningAttention(const Tensor q, const Tensor k, const Tensor v, + const Tensor slope, const Tensor initial_state, + const Tensor initial_state_indices, + const Tensor final_state_indices, Tensor out); +``` + +## 5. 可以直接复用 / 必须重写 + +**可直接复用** + +- 递推公式与代数推导（`S ← ratio ∘ S + kᵀv`，`o = q·S`）。 +- CUDA kernel 主体：按 `(batch, head)` 分块的逐列更新逻辑，共享内存缓存 `k`/`q` 行的做法。 +- CPU 参考实现的循环结构与 stride 处理。 +- 边界校验项（末维连续、`D <= 1024`、dtype 一致性、索引 dtype）。 + +**必须重写 / 补齐** + +1. 接口层（类 + assert 校验，见第 4 节）。 +2. **dtype 覆盖**：InfiniOps 的算子普遍只接受 fp16/bf16（例如 `PagedAttentionInfinilm` 明确 `assert` 仅 f16/bf16），当前 kernel 只有 fp32，这是能否被合入的硬门槛。 +3. kernel 文件命名与拆分规范（`kernel.h` + `kernel.cuh`，非模板 kernel 也要求头/源分离）。 +4. 代码风格：Google C++ Style + 仓库自带 `.clang-format`；注释与错误信息用英文完整句。 +5. 构建：改用 InfiniOps 的 CMake（`WITH_NVIDIA` / `WITH_ASCEND` / … 选项），不再依赖 InfiniCore 的 xmake glob。 + +## 6. 接口命名建议（对齐开源） + +InfiniOps 要求接口与主流开源框架一致，命名优先级为 **PyTorch → ONNX → CUDA API**。 + +- Lightning Attention 的开源等价物是 Flash-Linear-Attention（fla）的 `chunk_lightning_attn` / `fused_recurrent_lightning_attn`；vLLM 的 MiniMax 实现也使用同名 kernel。 +- 因此建议参数与语义对齐 fla：`q, k, v, slope, initial_state (+ indices), out`，并在文档中标注「decode 走 recurrent、prefill 走 chunk」。 +- 若维护者认为暂无可对齐的开源算子，可沿用仓库既有先例：`PagedAttentionInfinilm` 采用 `xxx_infinilm` 后缀并标注 + `[[deprecated("Migrate to an open-source-aligned operator when available.")]]`；此时命名为 `LightningAttentionInfinilm` 更符合现状约定。 + +## 7. 昇腾（Ascend）实现位置 + +昇腾实现**不要**写到 InfiniCore 的 `src/infiniop/ops/*/ascend/`（同样会被清理），而应写到： + +``` +src/native/ascend/ops/lightning_attention/ +``` + +InfiniOps 已有 `src/native/ascend/custom/`（AscendC 自定义 kernel）与 CMake 开关 `BUILD_ASCEND_CUSTOM`，可直接复用该机制。InfiniLM 侧 CI 目前 `ascend:` 段是注释状态，需要一并启用。 + +## 8. 迁移后自检清单 + +1. 数值：与补丁中的 CPU 参考实现逐值比对（decode `T=1` 与 prefill `T>1` 两种场景，含「初始行 ≠ 写回行」的索引池场景）。 +2. 状态连续性：`prefill(N)` 与 `prefill(N-1)+decode(1)` 的末位 logits 一致（当前 CPU 实测差异 1e-10 量级）。 +3. 端到端：对照 HF transformers `MiniMaxForCausalLM` 的 prefill/decode logits（当前实测 5.6e-4 / 1.6e-4）。 +4. dtype：至少覆盖 fp16、bf16（数值比对可用 CPU 参考或 torch 参考）。 +5. 形状/边界：`D ∈ {64, 128}`、非连续输入、`B` 多请求、索引池行冲突（多个请求读写同一行）等。 + + diff --git a/docs/minimax/archive-to-fork.ps1 b/docs/minimax/archive-to-fork.ps1 new file mode 100644 index 000000000..84eff52c9 --- /dev/null +++ b/docs/minimax/archive-to-fork.ps1 @@ -0,0 +1,217 @@ +<# +.SYNOPSIS + Archive the MiniMax work into your own GitHub branches and tags (InfiniCore + InfiniLM). + +.DESCRIPTION + - InfiniCore: adds remote `myfork` (https://github.com//InfiniCore.git), + creates the archive branch on base 35b46277, commits, tags and pushes it. + - InfiniLM: creates the archive branch on origin (your own fork), commits, tags and pushes it. + - Only the explicitly listed MiniMax files are staged; the staging area is verified + before committing and the script aborts on any unexpected file. + +.PARAMETER DryRun + Prints the commands instead of running the mutating ones; `git add -n` validates paths for real. + +.PARAMETER SkipPush + Commits and tags locally, but does not push. + +.EXAMPLE + pwsh -File docs/minimax/archive-to-fork.ps1 -DryRun + pwsh -File docs/minimax/archive-to-fork.ps1 +#> +[CmdletBinding()] +param( + [string]$GitHubUser = 'y258dd', + [string]$InfiniCorePath = 'C:\Users\yht13\InfiniCore', + [string]$InfiniLMPath = 'C:\Users\yht13\InfiniLM', + [switch]$DryRun, + [switch]$SkipPush +) + +$ErrorActionPreference = 'Stop' + +$InfiniCoreExpectedHead = '35b46277bd666772c11bb417ad4231c5be492822' +$InfiniCoreBranch = 'archive/minimax-lightning-attn' +$InfiniCoreTag = 'minimax-lightning-attn-35b46277' +$InfiniCoreSubject = 'feat(infiniop): add lightning_attention op (CPU + NVIDIA) for MiniMax' +$InfiniCoreBody = "Implements indexed-pool lightning attention (MiniMax-01 style):`n S <- ratio * S + k^T v ; o = q * S with ratio[h] = exp(-slope[h])`nCPU reference implementation plus a CUDA kernel, wired through the`ninfiniop C API, the infinicore C++ op layer and the Python bindings." +$InfiniCorePaths = @( + '.gitignore', + 'include/infiniop.h', + 'include/infinicore/ops.hpp', + 'src/infinicore/pybind11/ops.hpp', + 'python/infinicore/__init__.py', + 'include/infiniop/ops/lightning_attention.h', + 'include/infinicore/ops/lightning_attention.hpp', + 'src/infiniop/ops/lightning_attention', + 'src/infinicore/ops/lightning_attention', + 'src/infinicore/pybind11/ops/lightning_attention.hpp', + 'python/infinicore/ops/lightning_attention.py' +) + +$InfiniLMBranch = 'archive/minimax' +$InfiniLMTag = 'minimax-support' +$InfiniLMSubject = 'feat(minimax): support MiniMax-Text-01 (lightning attention + MoE)' +$InfiniLMBody = "- csrc/models/minimax: MiniMax model (hybrid lightning/softmax attention,`n block-sparse MoE, dense fallback), registered as minimax/minimax_m2`n- python/infinilm/modeling_utils.py: _remap_minimax weight remapper`n- test/models/minimax: op unit test, model smoke test, HF-aligned E2E test`n- docs/minimax: porting guide, handover notes and the InfiniCore op patch" +$InfiniLMPaths = @( + '.gitignore', + 'python/infinilm/modeling_utils.py', + 'csrc/models/minimax', + 'test/models/minimax', + 'docs/minimax' +) + +function Write-Step { + param([string]$Text) + Write-Host '' + Write-Host "=== $Text ===" -ForegroundColor Cyan +} + +function Invoke-Git { + param( + [Parameter(Mandatory)][string]$Repo, + [Parameter(Mandatory)][string[]]$GitArgs, + [switch]$Mutating + ) + $display = 'git -C "' + $Repo + '" ' + ($GitArgs -join ' ') + if ($DryRun -and $Mutating) { + Write-Host " [dry-run] $display" -ForegroundColor Yellow + return @() + } + Write-Host " > $display" -ForegroundColor DarkGray + $output = & git -C $Repo @GitArgs + if ($LASTEXITCODE -ne 0) { throw "git command failed: $display" } + return $output +} + +function Assert-Repo { + param([string]$Path, [string]$Name) + if (-not (Test-Path (Join-Path $Path '.git'))) { throw "$Name is not a git repository: $Path" } +} + +function Assert-StagedSet { + param([string]$Repo, [string[]]$Expected) + $staged = @(& git -C $Repo diff --cached --name-only | Where-Object { $_ }) + $unexpected = @($staged | Where-Object { + $file = $_ + -not ($Expected | Where-Object { $file -eq $_ -or $file.StartsWith($_ + '/') }) + }) + if ($unexpected.Count -gt 0) { + throw ("Unexpected staged files, aborting:`n " + ($unexpected -join "`n ")) + } + Write-Host (" staged files: {0} (all within the expected list)" -f $staged.Count) -ForegroundColor Green +} + +function Ensure-Branch { + param([string]$Repo, [string]$Branch) + $existing = @(& git -C $Repo branch --list $Branch | Where-Object { $_ }) + if ($existing.Count -gt 0) { + Write-Host " branch already exists, checking it out: $Branch" -ForegroundColor Yellow + Invoke-Git -Repo $Repo -GitArgs @('checkout', $Branch) -Mutating | Out-Null + } else { + Invoke-Git -Repo $Repo -GitArgs @('checkout', '-b', $Branch) -Mutating | Out-Null + } +} + +function Commit-If-Needed { + param( + [string]$Repo, + [string]$Subject, + [string]$Body, + [string[]]$Paths + ) + if ($DryRun) { + Invoke-Git -Repo $Repo -GitArgs (@('add', '-n', '--') + $Paths) | Out-Null + } else { + Invoke-Git -Repo $Repo -GitArgs (@('add', '--') + $Paths) -Mutating | Out-Null + Assert-StagedSet -Repo $Repo -Expected $Paths + $stagedCount = @(& git -C $Repo diff --cached --name-only | Where-Object { $_ }).Count + if ($stagedCount -eq 0) { + Write-Host ' nothing to commit (already committed?), skipping commit.' -ForegroundColor Yellow + return + } + } + Invoke-Git -Repo $Repo -GitArgs @('commit', '-m', $Subject, '-m', $Body) -Mutating | Out-Null +} + +function Tag-If-Missing { + param([string]$Repo, [string]$Tag) + $existing = @(& git -C $Repo tag --list $Tag | Where-Object { $_ }) + if ($existing.Count -gt 0) { + Write-Host " tag already exists, skipping: $Tag" -ForegroundColor Yellow + return + } + Invoke-Git -Repo $Repo -GitArgs @('tag', $Tag) -Mutating | Out-Null +} + +# ---------------------------------------------------------------- InfiniCore +Write-Step "InfiniCore: $InfiniCorePath" +Assert-Repo -Path $InfiniCorePath -Name 'InfiniCore' + +$head = (& git -C $InfiniCorePath rev-parse HEAD).Trim() +if ($head -ne $InfiniCoreExpectedHead) { + Write-Host " WARNING: HEAD = $head, expected base $InfiniCoreExpectedHead" -ForegroundColor Yellow + Write-Host " (ignore if you already archived; otherwise check for stray commits)" -ForegroundColor Yellow +} else { + Write-Host " base commit confirmed: $head" -ForegroundColor Green +} + +$forkUrl = "https://github.com/$GitHubUser/InfiniCore.git" +$remotes = @(& git -C $InfiniCorePath remote | Where-Object { $_ }) +if ($remotes -notcontains 'myfork') { + Invoke-Git -Repo $InfiniCorePath -GitArgs @('remote', 'add', 'myfork', $forkUrl) -Mutating | Out-Null +} else { + $existingUrl = (& git -C $InfiniCorePath remote get-url myfork).Trim() + if ($existingUrl -ne $forkUrl) { + Write-Host " NOTE: myfork exists with url $existingUrl (expected $forkUrl)" -ForegroundColor Yellow + } + Write-Host " myfork configured: $existingUrl" -ForegroundColor Green +} + +Ensure-Branch -Repo $InfiniCorePath -Branch $InfiniCoreBranch +Commit-If-Needed -Repo $InfiniCorePath -Subject $InfiniCoreSubject -Body $InfiniCoreBody -Paths $InfiniCorePaths +Tag-If-Missing -Repo $InfiniCorePath -Tag $InfiniCoreTag + +if (-not $SkipPush) { + Invoke-Git -Repo $InfiniCorePath -GitArgs @('push', '-u', 'myfork', $InfiniCoreBranch) -Mutating | Out-Null + Invoke-Git -Repo $InfiniCorePath -GitArgs @('push', 'myfork', $InfiniCoreTag) -Mutating | Out-Null +} + +# ----------------------------------------------------------------- InfiniLM +Write-Step "InfiniLM: $InfiniLMPath" +Assert-Repo -Path $InfiniLMPath -Name 'InfiniLM' + +$originUrl = (& git -C $InfiniLMPath remote get-url origin).Trim() +Write-Host " origin = $originUrl" +if ($originUrl -notmatch [regex]::Escape($GitHubUser)) { + Write-Host " WARNING: origin does not contain account $GitHubUser; please confirm this is your own fork." -ForegroundColor Yellow +} else { + Write-Host " confirmed as your own fork." -ForegroundColor Green +} + +Ensure-Branch -Repo $InfiniLMPath -Branch $InfiniLMBranch +Commit-If-Needed -Repo $InfiniLMPath -Subject $InfiniLMSubject -Body $InfiniLMBody -Paths $InfiniLMPaths +Tag-If-Missing -Repo $InfiniLMPath -Tag $InfiniLMTag + +if (-not $SkipPush) { + Invoke-Git -Repo $InfiniLMPath -GitArgs @('push', '-u', 'origin', $InfiniLMBranch) -Mutating | Out-Null + Invoke-Git -Repo $InfiniLMPath -GitArgs @('push', 'origin', $InfiniLMTag) -Mutating | Out-Null +} + +# ------------------------------------------------------------------- Summary +Write-Step 'Done' +if ($DryRun) { + Write-Host 'This was a dry run; nothing was changed except read-only checks. Re-run without -DryRun to execute.' -ForegroundColor Yellow +} else { + Write-Host " InfiniCore: branch $InfiniCoreBranch , tag $InfiniCoreTag (remote: myfork)" -ForegroundColor Green + Write-Host " InfiniLM : branch $InfiniLMBranch , tag $InfiniLMTag (remote: origin)" -ForegroundColor Green + if ($SkipPush) { + Write-Host ' Push skipped (-SkipPush). Push manually with:' -ForegroundColor Yellow + Write-Host " git -C `"$InfiniCorePath`" push -u myfork $InfiniCoreBranch" + Write-Host " git -C `"$InfiniCorePath`" push myfork $InfiniCoreTag" + Write-Host " git -C `"$InfiniLMPath`" push -u origin $InfiniLMBranch" + Write-Host " git -C `"$InfiniLMPath`" push origin $InfiniLMTag" + } + Write-Host '' + Write-Host ' Next (server compile check): InfiniCore base 35b46277 + docs/minimax/lightning-attention-infinicore.patch' -ForegroundColor Cyan +} \ No newline at end of file diff --git a/docs/minimax/infiniops-lightning-attention.patch b/docs/minimax/infiniops-lightning-attention.patch new file mode 100644 index 000000000..deb9274f7 --- /dev/null +++ b/docs/minimax/infiniops-lightning-attention.patch @@ -0,0 +1,663 @@ +diff --git a/src/base/lightning_attention_infinilm.h b/src/base/lightning_attention_infinilm.h +new file mode 100644 +index 0000000..3cedb49 +--- /dev/null ++++ b/src/base/lightning_attention_infinilm.h +@@ -0,0 +1,171 @@ ++#ifndef INFINI_OPS_BASE_LIGHTNING_ATTENTION_INFINILM_H_ ++#define INFINI_OPS_BASE_LIGHTNING_ATTENTION_INFINILM_H_ ++ ++#include ++#include ++ ++#include "data_type.h" ++#include "operator.h" ++#include "tensor.h" ++ ++namespace infini::ops { ++ ++/// Lightning attention with an indexed recurrent-state pool. ++/// ++/// The operator evaluates the recurrent form of MiniMax-style linear attention ++/// with a per-head ALiBi-style decay, one token at a time: ++/// ++/// ratio[h] = exp(-slope[h]) ++/// S = ratio[h] * S + outer(k_t[h], v_t[h]) ++/// out_t[h] = q_t[h] @ S ++/// ++/// The recurrent state of request `b` is read from ++/// `initial_state[initial_state_indices[b]]` and the final state is written to ++/// `initial_state[final_state_indices[b]]`. The row referenced by ++/// `initial_state_indices` is left untouched, so a caller may keep using it. ++/// ++/// Requests are independent and may execute concurrently, so a row used as the ++/// destination of one request must not be the source row of another request in ++/// the same call. ++/// ++/// This operator is InfiniLM-specific and is therefore classified as custom ++/// rather than aligned to an open-source operator. The closest public ++/// reference implementation is Flash-Linear-Attention's ++/// `fused_recurrent_lightning_attn`; `dexp` there is the per-head decay factor ++/// `exp(-slope)` used here. ++class LightningAttentionInfinilm ++ : public Operator { ++ public: ++ LightningAttentionInfinilm(const Tensor q, const Tensor k, const Tensor v, ++ const Tensor slope, Tensor initial_state, ++ const Tensor initial_state_indices, ++ const Tensor final_state_indices, Tensor out) ++ : data_type_{q.dtype()}, ++ index_dtype_{initial_state_indices.dtype()}, ++ batch_size_{q.size(0)}, ++ seq_len_{q.size(1)}, ++ num_heads_{q.size(2)}, ++ head_dim_{q.size(3)}, ++ state_pool_size_{initial_state.size(0)}, ++ state_pool_stride_{initial_state.stride(0)}, ++ state_head_stride_{initial_state.stride(1)}, ++ state_row_stride_{initial_state.stride(2)}, ++ state_column_stride_{initial_state.stride(3)}, ++ q_batch_stride_{q.stride(0)}, ++ q_seq_stride_{q.stride(1)}, ++ q_head_stride_{q.stride(2)}, ++ k_batch_stride_{k.stride(0)}, ++ k_seq_stride_{k.stride(1)}, ++ k_head_stride_{k.stride(2)}, ++ v_batch_stride_{v.stride(0)}, ++ v_seq_stride_{v.stride(1)}, ++ v_head_stride_{v.stride(2)}, ++ out_batch_stride_{out.stride(0)}, ++ out_seq_stride_{out.stride(1)}, ++ out_head_stride_{out.stride(2)}, ++ slope_stride_{slope.stride(0)}, ++ initial_index_stride_{initial_state_indices.stride(0)}, ++ final_index_stride_{final_state_indices.stride(0)} { ++ assert(q.ndim() == 4 && k.ndim() == 4 && v.ndim() == 4 && out.ndim() == 4 && ++ "`LightningAttentionInfinilm` expects [batch, seq, heads, head_dim] tensors"); ++ assert(q.dtype() == k.dtype() && k.dtype() == v.dtype() && ++ v.dtype() == out.dtype() && out.dtype() == data_type_ && ++ initial_state.dtype() == data_type_ && ++ "`LightningAttentionInfinilm` requires all data tensors to share one dtype"); ++ assert((data_type_ == DataType::kFloat32 || ++ data_type_ == DataType::kFloat16 || ++ data_type_ == DataType::kBFloat16) && ++ "`LightningAttentionInfinilm` supports float32, float16 and bfloat16"); ++ assert(q.shape() == k.shape() && k.shape() == v.shape() && ++ v.shape() == out.shape() && ++ "`LightningAttentionInfinilm` requires q, k, v and out to share a shape"); ++ assert(slope.dtype() == DataType::kFloat32 && slope.ndim() == 1 && ++ slope.size(0) == num_heads_ && slope_stride_ == 1 && ++ "`LightningAttentionInfinilm` expects `slope` to be a contiguous float32 tensor of size num_heads"); ++ assert(initial_state.ndim() == 4 && initial_state.size(1) == num_heads_ && ++ initial_state.size(2) == head_dim_ && ++ initial_state.size(3) == head_dim_ && ++ "`LightningAttentionInfinilm` expects `initial_state` to be [pool, heads, head_dim, head_dim]"); ++ assert(state_pool_size_ > 0 && state_column_stride_ == 1 && ++ "`LightningAttentionInfinilm` expects a contiguous state pool on the last dimension"); ++ assert(initial_state_indices.ndim() == 1 && ++ final_state_indices.ndim() == 1 && ++ initial_state_indices.size(0) == batch_size_ && ++ final_state_indices.size(0) == batch_size_ && ++ "`LightningAttentionInfinilm` expects one state index per request"); ++ assert(IsIndexDtype(index_dtype_) && ++ final_state_indices.dtype() == index_dtype_ && ++ initial_index_stride_ == 1 && final_index_stride_ == 1 && ++ "`LightningAttentionInfinilm` expects contiguous int32/int64 state indices"); ++ assert(q.stride(3) == 1 && k.stride(3) == 1 && v.stride(3) == 1 && ++ out.stride(3) == 1 && ++ "`LightningAttentionInfinilm` requires a contiguous last dimension"); ++ } ++ ++ virtual void operator()(const Tensor q, const Tensor k, const Tensor v, ++ const Tensor slope, Tensor initial_state, ++ const Tensor initial_state_indices, ++ const Tensor final_state_indices, ++ Tensor out) const = 0; ++ ++ protected: ++ static bool IsIndexDtype(DataType dtype) { ++ return dtype == DataType::kInt32 || dtype == DataType::kInt64; ++ } ++ ++ DataType data_type_; ++ ++ DataType index_dtype_; ++ ++ Tensor::Size batch_size_{0}; ++ ++ Tensor::Size seq_len_{0}; ++ ++ Tensor::Size num_heads_{0}; ++ ++ Tensor::Size head_dim_{0}; ++ ++ Tensor::Size state_pool_size_{0}; ++ ++ Tensor::Stride state_pool_stride_{0}; ++ ++ Tensor::Stride state_head_stride_{0}; ++ ++ Tensor::Stride state_row_stride_{0}; ++ ++ Tensor::Stride state_column_stride_{0}; ++ ++ Tensor::Stride q_batch_stride_{0}; ++ ++ Tensor::Stride q_seq_stride_{0}; ++ ++ Tensor::Stride q_head_stride_{0}; ++ ++ Tensor::Stride k_batch_stride_{0}; ++ ++ Tensor::Stride k_seq_stride_{0}; ++ ++ Tensor::Stride k_head_stride_{0}; ++ ++ Tensor::Stride v_batch_stride_{0}; ++ ++ Tensor::Stride v_seq_stride_{0}; ++ ++ Tensor::Stride v_head_stride_{0}; ++ ++ Tensor::Stride out_batch_stride_{0}; ++ ++ Tensor::Stride out_seq_stride_{0}; ++ ++ Tensor::Stride out_head_stride_{0}; ++ ++ Tensor::Stride slope_stride_{0}; ++ ++ Tensor::Stride initial_index_stride_{0}; ++ ++ Tensor::Stride final_index_stride_{0}; ++}; ++ ++} // namespace infini::ops ++ ++#endif // INFINI_OPS_BASE_LIGHTNING_ATTENTION_INFINILM_H_ +diff --git a/src/native/cpu/ops/lightning_attention_infinilm/lightning_attention_infinilm.h b/src/native/cpu/ops/lightning_attention_infinilm/lightning_attention_infinilm.h +new file mode 100644 +index 0000000..8e8ef1c +--- /dev/null ++++ b/src/native/cpu/ops/lightning_attention_infinilm/lightning_attention_infinilm.h +@@ -0,0 +1,139 @@ ++#ifndef INFINI_OPS_CPU_LIGHTNING_ATTENTION_INFINILM_H_ ++#define INFINI_OPS_CPU_LIGHTNING_ATTENTION_INFINILM_H_ ++ ++#include ++#include ++#include ++ ++#include "base/lightning_attention_infinilm.h" ++#include "common/generic_utils.h" ++#include "data_type.h" ++#include "native/cpu/caster_.h" ++#include "tensor.h" ++ ++namespace infini::ops { ++ ++template <> ++class Operator ++ : public LightningAttentionInfinilm, Caster { ++ public: ++ Operator(const Tensor q, const Tensor k, const Tensor v, const Tensor slope, ++ Tensor initial_state, const Tensor initial_state_indices, ++ const Tensor final_state_indices, Tensor out) ++ : LightningAttentionInfinilm{q, k, v, slope, initial_state, ++ initial_state_indices, final_state_indices, ++ out} {} ++ ++ void operator()(const Tensor q, const Tensor k, const Tensor v, ++ const Tensor slope, Tensor initial_state, ++ const Tensor initial_state_indices, ++ const Tensor final_state_indices, Tensor out) const override { ++ DispatchFunc( ++ out.dtype(), ++ [&](auto tag) { ++ using T = typename decltype(tag)::type; ++ Compute(q, k, v, slope, initial_state, initial_state_indices, ++ final_state_indices, out); ++ }, ++ "`Operator::operator()`"); ++ } ++ ++ private: ++ template ++ void Compute(const Tensor q, const Tensor k, const Tensor v, ++ const Tensor slope, Tensor initial_state, ++ const Tensor initial_state_indices, ++ const Tensor final_state_indices, Tensor out) const { ++ const auto* q_ptr = static_cast(q.data()); ++ const auto* k_ptr = static_cast(k.data()); ++ const auto* v_ptr = static_cast(v.data()); ++ const auto* slope_ptr = static_cast(slope.data()); ++ auto* state_ptr = static_cast(initial_state.data()); ++ auto* out_ptr = static_cast(out.data()); ++ ++ const bool int64_indices = index_dtype_ == DataType::kInt64; ++ const auto* initial_indices = initial_state_indices.data(); ++ const auto* final_indices = final_state_indices.data(); ++ ++ // The recurrent state of one request, accumulated in float32. ++ std::vector state(num_heads_ * head_dim_ * head_dim_); ++ ++ for (Tensor::Size b = 0; b < batch_size_; ++b) { ++ Tensor::Size initial_row; ++ Tensor::Size final_row; ++ if (int64_indices) { ++ initial_row = static_cast( ++ static_cast(initial_indices)[b]); ++ final_row = ++ static_cast(static_cast(final_indices)[b]); ++ } else { ++ initial_row = static_cast( ++ static_cast(initial_indices)[b]); ++ final_row = ++ static_cast(static_cast(final_indices)[b]); ++ } ++ ++ const T* initial_row_ptr = state_ptr + initial_row * state_pool_stride_; ++ for (Tensor::Size h = 0; h < num_heads_; ++h) { ++ const T* head_ptr = initial_row_ptr + h * state_head_stride_; ++ float* state_head = state.data() + h * head_dim_ * head_dim_; ++ for (Tensor::Size i = 0; i < head_dim_; ++i) { ++ for (Tensor::Size j = 0; j < head_dim_; ++j) { ++ state_head[i * head_dim_ + j] = ++ Cast(head_ptr[i * state_row_stride_ + ++ j * state_column_stride_]); ++ } ++ } ++ } ++ ++ for (Tensor::Size t = 0; t < seq_len_; ++t) { ++ for (Tensor::Size h = 0; h < num_heads_; ++h) { ++ const float ratio = std::exp(-slope_ptr[h * slope_stride_]); ++ const T* q_row = q_ptr + b * q_batch_stride_ + t * q_seq_stride_ + ++ h * q_head_stride_; ++ const T* k_row = k_ptr + b * k_batch_stride_ + t * k_seq_stride_ + ++ h * k_head_stride_; ++ const T* v_row = v_ptr + b * v_batch_stride_ + t * v_seq_stride_ + ++ h * v_head_stride_; ++ float* state_head = state.data() + h * head_dim_ * head_dim_; ++ ++ for (Tensor::Size i = 0; i < head_dim_; ++i) { ++ const float k_i = Cast(k_row[i]); ++ for (Tensor::Size j = 0; j < head_dim_; ++j) { ++ state_head[i * head_dim_ + j] = ++ ratio * state_head[i * head_dim_ + j] + ++ k_i * Cast(v_row[j]); ++ } ++ } ++ ++ T* out_row = out_ptr + b * out_batch_stride_ + t * out_seq_stride_ + ++ h * out_head_stride_; ++ for (Tensor::Size j = 0; j < head_dim_; ++j) { ++ float acc = 0.0f; ++ for (Tensor::Size i = 0; i < head_dim_; ++i) { ++ acc += Cast(q_row[i]) * state_head[i * head_dim_ + j]; ++ } ++ out_row[j] = Cast(acc); ++ } ++ } ++ } ++ ++ // Only the destination row is updated; the initial row stays untouched. ++ T* final_row_ptr = state_ptr + final_row * state_pool_stride_; ++ for (Tensor::Size h = 0; h < num_heads_; ++h) { ++ T* head_ptr = final_row_ptr + h * state_head_stride_; ++ const float* state_head = state.data() + h * head_dim_ * head_dim_; ++ for (Tensor::Size i = 0; i < head_dim_; ++i) { ++ for (Tensor::Size j = 0; j < head_dim_; ++j) { ++ head_ptr[i * state_row_stride_ + j * state_column_stride_] = ++ Cast(state_head[i * head_dim_ + j]); ++ } ++ } ++ } ++ } ++ } ++}; ++ ++} // namespace infini::ops ++ ++#endif // INFINI_OPS_CPU_LIGHTNING_ATTENTION_INFINILM_H_ +diff --git a/src/native/cuda/nvidia/ops/lightning_attention_infinilm/kernel.h b/src/native/cuda/nvidia/ops/lightning_attention_infinilm/kernel.h +new file mode 100644 +index 0000000..9a21f93 +--- /dev/null ++++ b/src/native/cuda/nvidia/ops/lightning_attention_infinilm/kernel.h +@@ -0,0 +1,22 @@ ++#ifndef INFINI_OPS_NVIDIA_LIGHTNING_ATTENTION_INFINILM_KERNEL_H_ ++#define INFINI_OPS_NVIDIA_LIGHTNING_ATTENTION_INFINILM_KERNEL_H_ ++ ++#include ++ ++#include "native/cuda/nvidia/caster.cuh" ++#include "native/cuda/nvidia/runtime_.h" ++#include "native/cuda/ops/lightning_attention_infinilm/kernel.h" ++ ++namespace infini::ops { ++ ++template <> ++class Operator ++ : public CudaLightningAttentionInfinilm> { ++ public: ++ using CudaLightningAttentionInfinilm< ++ Runtime>::CudaLightningAttentionInfinilm; ++}; ++ ++} // namespace infini::ops ++ ++#endif // INFINI_OPS_NVIDIA_LIGHTNING_ATTENTION_INFINILM_KERNEL_H_ +diff --git a/src/native/cuda/ops/lightning_attention_infinilm/kernel.cuh b/src/native/cuda/ops/lightning_attention_infinilm/kernel.cuh +new file mode 100644 +index 0000000..35861cf +--- /dev/null ++++ b/src/native/cuda/ops/lightning_attention_infinilm/kernel.cuh +@@ -0,0 +1,98 @@ ++#ifndef INFINI_OPS_CUDA_LIGHTNING_ATTENTION_INFINILM_KERNEL_CUH_ ++#define INFINI_OPS_CUDA_LIGHTNING_ATTENTION_INFINILM_KERNEL_CUH_ ++ ++#include ++ ++#include "native/cuda/caster.cuh" ++#include "native/cuda/kernel_commons.cuh" ++ ++namespace infini::ops { ++ ++/// One block per `(batch, head)` and one thread per state/output column. ++/// ++/// The recurrent state `[head_dim, head_dim]` of the destination pool row is ++/// staged first (the source row is copied into it when the two rows differ, so ++/// the source row stays untouched), then updated in place: ++/// ++/// state[i][j] = ratio * state[i][j] + k[i] * v[j] ++/// out[j] = sum_i q[i] * state[i][j] ++/// ++/// Thread `j` owns column `j`, so the state update needs no cross-thread ++/// synchronization; only the shared `k`/`q` rows do. The kernel must be ++/// launched with exactly `head_dim` threads so that every thread reaches the ++/// barriers. ++template ++__global__ void LightningAttentionInfinilmKernel( ++ Data* out, Data* state_pool, const Data* q, const Data* k, const Data* v, ++ const float* slope, const Index* initial_state_indices, ++ const Index* final_state_indices, size_t seq_len, size_t head_dim, ++ ptrdiff_t state_pool_stride, ptrdiff_t state_head_stride, ++ ptrdiff_t state_row_stride, ptrdiff_t q_batch_stride, ptrdiff_t q_seq_stride, ++ ptrdiff_t q_head_stride, ptrdiff_t k_batch_stride, ptrdiff_t k_seq_stride, ++ ptrdiff_t k_head_stride, ptrdiff_t v_batch_stride, ptrdiff_t v_seq_stride, ++ ptrdiff_t v_head_stride, ptrdiff_t out_batch_stride, ptrdiff_t out_seq_stride, ++ ptrdiff_t out_head_stride, ptrdiff_t slope_stride) { ++ const size_t batch = blockIdx.y; ++ const size_t head = blockIdx.x; ++ const size_t column = threadIdx.x; ++ ++ const size_t initial_row = static_cast(initial_state_indices[batch]); ++ const size_t final_row = static_cast(final_state_indices[batch]); ++ ++ Data* state = state_pool + final_row * state_pool_stride + ++ head * state_head_stride; ++ if (initial_row != final_row) { ++ const Data* source = state_pool + initial_row * state_pool_stride + ++ head * state_head_stride; ++ for (size_t index = column; index < head_dim * head_dim; ++ index += blockDim.x) { ++ const size_t i = index / head_dim; ++ const size_t j = index % head_dim; ++ state[i * state_row_stride + j] = source[i * state_row_stride + j]; ++ } ++ } ++ __syncthreads(); ++ ++ extern __shared__ float shared[]; ++ float* shared_k = shared; ++ float* shared_q = shared + head_dim; ++ ++ const float ratio = expf(-slope[head * slope_stride]); ++ const Data* q_head = q + batch * q_batch_stride + head * q_head_stride; ++ const Data* k_head = k + batch * k_batch_stride + head * k_head_stride; ++ const Data* v_head = v + batch * v_batch_stride + head * v_head_stride; ++ Data* out_head = out + batch * out_batch_stride + head * out_head_stride; ++ ++ for (size_t t = 0; t < seq_len; ++t) { ++ shared_k[column] = Caster::template Cast( ++ k_head[t * k_seq_stride + column]); ++ shared_q[column] = Caster::template Cast( ++ q_head[t * q_seq_stride + column]); ++ __syncthreads(); ++ ++ const float v_column = Caster::template Cast( ++ v_head[t * v_seq_stride + column]); ++ for (size_t i = 0; i < head_dim; ++i) { ++ Data* element = state + i * state_row_stride + column; ++ const float updated = ratio * Caster::template Cast(*element) + ++ shared_k[i] * v_column; ++ *element = Caster::template Cast(updated); ++ } ++ ++ float accumulator = 0.0f; ++ for (size_t i = 0; i < head_dim; ++i) { ++ accumulator += shared_q[i] * Caster::template Cast( ++ state[i * state_row_stride + column]); ++ } ++ out_head[t * out_seq_stride + column] = ++ Caster::template Cast(accumulator); ++ ++ // The next iteration overwrites the shared rows, so all threads must have ++ // finished reading them. ++ __syncthreads(); ++ } ++} ++ ++} // namespace infini::ops ++ ++#endif // INFINI_OPS_CUDA_LIGHTNING_ATTENTION_INFINILM_KERNEL_CUH_ +diff --git a/src/native/cuda/ops/lightning_attention_infinilm/kernel.h b/src/native/cuda/ops/lightning_attention_infinilm/kernel.h +new file mode 100644 +index 0000000..e90c52b +--- /dev/null ++++ b/src/native/cuda/ops/lightning_attention_infinilm/kernel.h +@@ -0,0 +1,78 @@ ++#ifndef INFINI_OPS_CUDA_LIGHTNING_ATTENTION_INFINILM_KERNEL_H_ ++#define INFINI_OPS_CUDA_LIGHTNING_ATTENTION_INFINILM_KERNEL_H_ ++ ++#include ++#include ++#include ++ ++#include "base/lightning_attention_infinilm.h" ++#include "data_type.h" ++#include "dispatcher.h" ++#include "native/cuda/kernel_commons.cuh" ++#include "native/cuda/ops/lightning_attention_infinilm/kernel.cuh" ++#include "native/cuda/runtime_utils.h" ++ ++namespace infini::ops { ++ ++using LightningAttentionInfinilmDataTypes = ++ ConcatType, ReducedFloatTypes>; ++ ++using LightningAttentionInfinilmIndexTypes = ++ List; ++ ++template ++class CudaLightningAttentionInfinilm : public LightningAttentionInfinilm { ++ public: ++ using LightningAttentionInfinilm::LightningAttentionInfinilm; ++ ++ void operator()(const Tensor q, const Tensor k, const Tensor v, ++ const Tensor slope, Tensor initial_state, ++ const Tensor initial_state_indices, ++ const Tensor final_state_indices, Tensor out) const override { ++ auto cuda_stream = ++ static_cast(stream_ ? stream_ : 0); ++ ++ // One thread per state column, so `head_dim` has to fit into one block. ++ assert(head_dim_ > 0 && ++ static_cast(head_dim_) <= BackendMaxBlockSize::value && ++ "`LightningAttentionInfinilm` requires head_dim to fit one block"); ++ assert(batch_size_ <= 65535 && ++ "`LightningAttentionInfinilm` requires batch_size <= 65535"); ++ ++ dim3 grid(static_cast(num_heads_), ++ static_cast(batch_size_)); ++ dim3 block(static_cast(head_dim_)); ++ const size_t shared_bytes = 2 * head_dim_ * sizeof(float); ++ ++ DispatchFunc( ++ {static_cast(out.dtype()), static_cast(index_dtype_)}, ++ [&](auto list_tag) { ++ using T = TypeMapType(list_tag)>; ++ using TIndex = ++ TypeMapType(list_tag)>; ++ ++ LightningAttentionInfinilmKernel ++ <<>>( ++ reinterpret_cast(out.data()), ++ reinterpret_cast(initial_state.data()), ++ reinterpret_cast(q.data()), ++ reinterpret_cast(k.data()), ++ reinterpret_cast(v.data()), ++ static_cast(slope.data()), ++ reinterpret_cast(initial_state_indices.data()), ++ reinterpret_cast(final_state_indices.data()), ++ seq_len_, head_dim_, state_pool_stride_, state_head_stride_, ++ state_row_stride_, q_batch_stride_, q_seq_stride_, ++ q_head_stride_, k_batch_stride_, k_seq_stride_, k_head_stride_, ++ v_batch_stride_, v_seq_stride_, v_head_stride_, ++ out_batch_stride_, out_seq_stride_, out_head_stride_, ++ slope_stride_); ++ }, ++ "CudaLightningAttentionInfinilm::operator()"); ++ } ++}; ++ ++} // namespace infini::ops ++ ++#endif // INFINI_OPS_CUDA_LIGHTNING_ATTENTION_INFINILM_KERNEL_H_ +diff --git a/tests/test_lightning_attention_infinilm.py b/tests/test_lightning_attention_infinilm.py +new file mode 100644 +index 0000000..a31e725 +--- /dev/null ++++ b/tests/test_lightning_attention_infinilm.py +@@ -0,0 +1,119 @@ ++import infini.ops ++import pytest ++import torch ++ ++from tests.utils import Payload, empty_strided, get_stream, randn_strided ++ ++# (batch, seq_len, num_heads, head_dim, pool_size) ++_SHAPES = ( ++ (1, 1, 2, 4, 1), ++ (2, 3, 2, 4, 4), ++ (3, 1, 4, 8, 6), ++ (1, 5, 4, 8, 2), ++) ++ ++_DTYPE_CASES = ( ++ (torch.float32, 1e-5, 1e-6), ++ (torch.float16, 2e-2, 2e-2), ++ (torch.bfloat16, 2e-2, 2e-2), ++) ++ ++ ++def _make_tensors(shape, dtype, device): ++ batch, seq_len, num_heads, head_dim, pool_size = shape ++ ++ q = randn_strided((batch, seq_len, num_heads, head_dim), None, dtype=dtype, device=device) ++ k = randn_strided((batch, seq_len, num_heads, head_dim), None, dtype=dtype, device=device) ++ v = randn_strided((batch, seq_len, num_heads, head_dim), None, dtype=dtype, device=device) ++ # The decay has to be positive, otherwise `exp(-slope)` would grow. ++ slope = randn_strided((num_heads,), None, dtype=torch.float32, device=device).abs() * 0.5 ++ state = randn_strided( ++ (pool_size, num_heads, head_dim, head_dim), None, dtype=dtype, device=device ++ ) ++ ++ # Deliberately read and write different pool rows so that the test also ++ # covers the "initial row must stay untouched" contract. Rows are disjoint ++ # across requests because requests may execute concurrently. ++ initial_indices = torch.arange(batch, dtype=torch.int32, device=device) % pool_size ++ final_indices = (initial_indices + batch) % pool_size ++ ++ out = empty_strided((batch, seq_len, num_heads, head_dim), None, dtype=dtype, device=device) ++ ++ return q, k, v, slope, state, initial_indices, final_indices, out ++ ++ ++def _torch_lightning_attention(q, k, v, slope, state, initial_indices, final_indices): ++ """Recurrent reference: S <- ratio * S + k^T v ; out = q @ S.""" ++ ++ q = q.float() ++ k = k.float() ++ v = v.float() ++ state = state.float().clone() ++ out = torch.empty_like(q) ++ ratio = torch.exp(-slope.float()) ++ ++ for b in range(q.shape[0]): ++ initial_row = int(initial_indices[b].item()) ++ final_row = int(final_indices[b].item()) ++ ++ current = state[initial_row].clone() # [heads, head_dim, head_dim] ++ for t in range(q.shape[1]): ++ current = ratio[:, None, None] * current + k[b, t].unsqueeze(-1) * v[b, t].unsqueeze(-2) ++ # The implementation stores the recurrent state in the tensor dtype, ++ # so the reference has to round it before the next read. ++ current = current.to(state.dtype) ++ out[b, t] = torch.einsum("hd,hde->he", q[b, t], current.float()) ++ ++ state[final_row] = current ++ ++ return out, state ++ ++ ++def _run_lightning_attention(q, k, v, slope, state, initial_indices, final_indices, out): ++ infini.ops.lightning_attention_infinilm( ++ q, ++ k, ++ v, ++ slope, ++ state, ++ initial_indices, ++ final_indices, ++ out, ++ stream=get_stream(q.device), ++ ) ++ return out ++ ++ ++@pytest.mark.auto_act_and_assert ++@pytest.mark.parametrize("shape", _SHAPES) ++@pytest.mark.parametrize(("dtype", "rtol", "atol"), _DTYPE_CASES) ++def test_lightning_attention_infinilm(shape, dtype, device, rtol, atol): ++ tensors = _make_tensors(shape, dtype, device) ++ ++ return Payload(_run_lightning_attention, _reference_out, tensors, {}, rtol=rtol, atol=atol) ++ ++ ++@pytest.mark.auto_act_and_assert ++@pytest.mark.parametrize("shape", _SHAPES) ++@pytest.mark.parametrize(("dtype", "rtol", "atol"), _DTYPE_CASES) ++def test_lightning_attention_infinilm_state_pool(shape, dtype, device, rtol, atol): ++ tensors = _make_tensors(shape, dtype, device) ++ ++ return Payload(_run_lightning_attention_state, _reference_state, tensors, {}, rtol=rtol, atol=atol) ++ ++ ++def _reference_out(q, k, v, slope, state, initial_indices, final_indices, out): ++ reference, _ = _torch_lightning_attention(q, k, v, slope, state, initial_indices, final_indices) ++ out.copy_(reference.to(out.dtype)) ++ return out ++ ++ ++def _reference_state(q, k, v, slope, state, initial_indices, final_indices, out): ++ _, reference = _torch_lightning_attention(q, k, v, slope, state, initial_indices, final_indices) ++ state.copy_(reference.to(state.dtype)) ++ return state ++ ++ ++def _run_lightning_attention_state(q, k, v, slope, state, initial_indices, final_indices, out): ++ _run_lightning_attention(q, k, v, slope, state, initial_indices, final_indices, out) ++ return state diff --git a/docs/minimax/lightning-attention-infinicore.patch b/docs/minimax/lightning-attention-infinicore.patch new file mode 100644 index 000000000..e7ec18626 --- /dev/null +++ b/docs/minimax/lightning-attention-infinicore.patch @@ -0,0 +1,1135 @@ +diff --git a/include/infinicore/ops.hpp b/include/infinicore/ops.hpp +index 5e93e145..da0215b7 100644 +--- a/include/infinicore/ops.hpp ++++ b/include/infinicore/ops.hpp +@@ -45,6 +45,7 @@ + #include "ops/hardswish.hpp" + #include "ops/hardtanh.hpp" + #include "ops/kimi_delta_attention.hpp" ++#include "ops/lightning_attention.hpp" + #include "ops/kv_caching.hpp" + #include "ops/layer_norm.hpp" + #include "ops/linear.hpp" +@@ -114,3 +115,4 @@ + #include "ops/w4a8_group_gemm.hpp" + #include "ops/w8a8_group_gemm.hpp" + #endif ++ +diff --git a/include/infinicore/ops/lightning_attention.hpp b/include/infinicore/ops/lightning_attention.hpp +new file mode 100644 +index 00000000..268d32d1 +--- /dev/null ++++ b/include/infinicore/ops/lightning_attention.hpp +@@ -0,0 +1,41 @@ ++#pragma once ++ ++#include "infinicore.h" ++ ++#include "../device.hpp" ++#include "../graph/graph.hpp" ++#include "common/op.hpp" ++ ++namespace infinicore::op { ++ ++INFINICORE_GRAPH_OP_CLASS(LightningAttention, ++ Tensor, ++ Tensor, ++ const Tensor &, ++ const Tensor &, ++ const Tensor &, ++ const Tensor &, ++ const Tensor &, ++ const Tensor &); ++ ++// Indexed-pool lightning attention (MiniMax-01 style). ++// Returns out [B, T, H, D] and updates `initial_state` in place at ++// `final_state_indices` rows. ++__export Tensor lightning_attention(const Tensor &q, ++ const Tensor &k, ++ const Tensor &v, ++ const Tensor &slope, ++ Tensor initial_state, ++ const Tensor &initial_state_indices, ++ const Tensor &final_state_indices); ++ ++__export void lightning_attention_(Tensor out, ++ Tensor initial_state, ++ const Tensor &q, ++ const Tensor &k, ++ const Tensor &v, ++ const Tensor &slope, ++ const Tensor &initial_state_indices, ++ const Tensor &final_state_indices); ++ ++} // namespace infinicore::op +diff --git a/include/infiniop.h b/include/infiniop.h +index 9f632e27..95b2d6e7 100644 +--- a/include/infiniop.h ++++ b/include/infiniop.h +@@ -85,6 +85,7 @@ + #include "infiniop/ops/layer_norm.h" + #include "infiniop/ops/ldexp.h" + #include "infiniop/ops/lerp.h" ++#include "infiniop/ops/lightning_attention.h" + #include "infiniop/ops/linear_mxfp4.h" + #include "infiniop/ops/log10.h" + #include "infiniop/ops/log1p.h" +@@ -169,3 +170,5 @@ + #include "infiniop/ops/zeros.h" + #include "infiniop/tensor_descriptor.h" + #endif // __INFINIOP_API_H__ ++ ++ +diff --git a/include/infiniop/ops/lightning_attention.h b/include/infiniop/ops/lightning_attention.h +new file mode 100644 +index 00000000..cd9156e3 +--- /dev/null ++++ b/include/infiniop/ops/lightning_attention.h +@@ -0,0 +1,60 @@ ++#ifndef __INFINIOP_LIGHTNING_ATTENTION_API_H__ ++#define __INFINIOP_LIGHTNING_ATTENTION_API_H__ ++ ++#include "../operator_descriptor.h" ++ ++typedef struct InfiniopDescriptor *infiniopLightningAttentionDescriptor_t; ++ ++// Lightning attention (MiniMax-01 style, ALiBi-style per-head decay) with an ++// indexed recurrent-state pool. ++// ++// Recurrence (per head h, per token t): ++// S = ratio[h] * S + k_t^T v_t (ratio[h] = exp(-slope[h])) ++// o_t = q_t @ S ++// i.e. the state is updated *before* the output is read, so each output token ++// attends to itself with weight 1 (no decay within the same position). ++// ++// Tensor layouts: ++// out [B, T, H, D] (last dim contiguous) ++// initial_state (pool) [pool_size, H, D, D] ++// q/k/v [B, T, H, D] (last dim contiguous) ++// slope [H] (fp32) ++// initial/final_state_indices [B] (int32 or int64) ++// ++// Indexed-pool mode only: for each request b the op reads the state row ++// `initial_state[initial_state_indices[b]]` and writes the final state in place ++// to `initial_state[final_state_indices[b]]`. ++__INFINI_C __export infiniStatus_t infiniopCreateLightningAttentionDescriptor( ++ infiniopHandle_t handle, ++ infiniopLightningAttentionDescriptor_t *desc_ptr, ++ infiniopTensorDescriptor_t out_desc, ++ infiniopTensorDescriptor_t initial_state_desc, ++ infiniopTensorDescriptor_t q_desc, ++ infiniopTensorDescriptor_t k_desc, ++ infiniopTensorDescriptor_t v_desc, ++ infiniopTensorDescriptor_t slope_desc, ++ infiniopTensorDescriptor_t initial_state_indices_desc, ++ infiniopTensorDescriptor_t final_state_indices_desc); ++ ++__INFINI_C __export infiniStatus_t infiniopGetLightningAttentionWorkspaceSize( ++ infiniopLightningAttentionDescriptor_t desc, ++ size_t *size); ++ ++__INFINI_C __export infiniStatus_t infiniopLightningAttention( ++ infiniopLightningAttentionDescriptor_t desc, ++ void *workspace, ++ size_t workspace_size, ++ void *out, ++ void *initial_state, ++ const void *q, ++ const void *k, ++ const void *v, ++ const void *slope, ++ const void *initial_state_indices, ++ const void *final_state_indices, ++ void *stream); ++ ++__INFINI_C __export infiniStatus_t infiniopDestroyLightningAttentionDescriptor( ++ infiniopLightningAttentionDescriptor_t desc); ++ ++#endif +diff --git a/python/infinicore/__init__.py b/python/infinicore/__init__.py +index 612db614..ef8a5b72 100644 +--- a/python/infinicore/__init__.py ++++ b/python/infinicore/__init__.py +@@ -121,6 +121,7 @@ from infinicore.ops.kthvalue import kthvalue + from infinicore.ops.kv_caching import kv_caching + from infinicore.ops.ldexp import ldexp + from infinicore.ops.lerp import lerp ++from infinicore.ops.lightning_attention import lightning_attention + from infinicore.ops.logaddexp import logaddexp + from infinicore.ops.logaddexp2 import logaddexp2 + from infinicore.ops.logcumsumexp import logcumsumexp +@@ -388,3 +389,5 @@ __all__ += [ + "w4a8_group_gemm_", + "w8a8_group_gemm_", + ] ++ ++ +diff --git a/python/infinicore/ops/lightning_attention.py b/python/infinicore/ops/lightning_attention.py +new file mode 100644 +index 00000000..ab7a86d7 +--- /dev/null ++++ b/python/infinicore/ops/lightning_attention.py +@@ -0,0 +1,21 @@ ++from infinicore.lib import _infinicore ++from infinicore.tensor import Tensor ++ ++ ++def lightning_attention(q, k, v, slope, initial_state, initial_state_indices, final_state_indices): ++ """Indexed-pool lightning attention (MiniMax-01 style). ++ ++ Returns out [B, T, H, D]; `initial_state` is updated in place at the ++ `final_state_indices` rows. ++ """ ++ return Tensor( ++ _infinicore.lightning_attention( ++ q._underlying, ++ k._underlying, ++ v._underlying, ++ slope._underlying, ++ initial_state._underlying, ++ initial_state_indices._underlying, ++ final_state_indices._underlying, ++ ) ++ ) +diff --git a/src/infinicore/ops/lightning_attention/lightning_attention.cc b/src/infinicore/ops/lightning_attention/lightning_attention.cc +new file mode 100644 +index 00000000..aacb6fb7 +--- /dev/null ++++ b/src/infinicore/ops/lightning_attention/lightning_attention.cc +@@ -0,0 +1,90 @@ ++#include "infinicore/ops/lightning_attention.hpp" ++#include "../../utils.hpp" ++ ++#include ++ ++namespace infinicore::op { ++ ++INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(LightningAttention); ++ ++LightningAttention::LightningAttention(Tensor out, ++ Tensor initial_state, ++ const Tensor &q, ++ const Tensor &k, ++ const Tensor &v, ++ const Tensor &slope, ++ const Tensor &initial_state_indices, ++ const Tensor &final_state_indices) { ++ INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, initial_state, q, k, v, slope); ++ INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, initial_state_indices); ++ INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, final_state_indices); ++ INFINICORE_GRAPH_OP_DISPATCH(out->device().getType(), ++ out, ++ initial_state, ++ q, ++ k, ++ v, ++ slope, ++ initial_state_indices, ++ final_state_indices); ++} ++ ++void LightningAttention::execute(Tensor out, ++ Tensor initial_state, ++ const Tensor &q, ++ const Tensor &k, ++ const Tensor &v, ++ const Tensor &slope, ++ const Tensor &initial_state_indices, ++ const Tensor &final_state_indices) { ++ INFINICORE_GRAPH_OP_RECORD_OR_RUN(LightningAttention, ++ out, ++ initial_state, ++ q, ++ k, ++ v, ++ slope, ++ initial_state_indices, ++ final_state_indices); ++} ++ ++static Tensor ensure_4d_sequence_tensor(const Tensor &x, const char *name) { ++ if (x->shape().size() == 4) { ++ return x; ++ } ++ if (x->shape().size() == 3) { ++ return x->unsqueeze(1); ++ } ++ throw std::runtime_error(std::string("lightning_attention expects ") + name + " with shape [B, T, H, D] or [B, H, D]"); ++} ++ ++Tensor lightning_attention(const Tensor &q, ++ const Tensor &k, ++ const Tensor &v, ++ const Tensor &slope, ++ Tensor initial_state, ++ const Tensor &initial_state_indices, ++ const Tensor &final_state_indices) { ++ Tensor q4 = ensure_4d_sequence_tensor(q, "q"); ++ Tensor k4 = ensure_4d_sequence_tensor(k, "k"); ++ Tensor v4 = ensure_4d_sequence_tensor(v, "v"); ++ auto out = Tensor::empty(v4->shape(), v4->dtype(), v4->device()); ++ lightning_attention_(out, initial_state, q4, k4, v4, slope, initial_state_indices, final_state_indices); ++ return out; ++} ++ ++void lightning_attention_(Tensor out, ++ Tensor initial_state, ++ const Tensor &q, ++ const Tensor &k, ++ const Tensor &v, ++ const Tensor &slope, ++ const Tensor &initial_state_indices, ++ const Tensor &final_state_indices) { ++ Tensor q4 = ensure_4d_sequence_tensor(q, "q"); ++ Tensor k4 = ensure_4d_sequence_tensor(k, "k"); ++ Tensor v4 = ensure_4d_sequence_tensor(v, "v"); ++ LightningAttention::execute(out, initial_state, q4, k4, v4, slope, initial_state_indices, final_state_indices); ++} ++ ++} // namespace infinicore::op +diff --git a/src/infinicore/ops/lightning_attention/lightning_attention_infiniop.cc b/src/infinicore/ops/lightning_attention/lightning_attention_infiniop.cc +new file mode 100644 +index 00000000..c8ec9e3d +--- /dev/null ++++ b/src/infinicore/ops/lightning_attention/lightning_attention_infiniop.cc +@@ -0,0 +1,85 @@ ++#include "infinicore/ops/lightning_attention.hpp" ++ ++#include "../infiniop_impl.hpp" ++ ++namespace infinicore::op::lightning_attention_impl::infiniop { ++ ++INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, LightningAttention, 100); ++ ++struct PlannedMeta { ++ std::shared_ptr descriptor; ++ graph::GraphTensor workspace, out, initial_state, q, k, v, slope; ++ graph::GraphTensor initial_state_indices; ++ graph::GraphTensor final_state_indices; ++}; ++ ++void *plan(Tensor out, ++ Tensor initial_state, ++ const Tensor &q, ++ const Tensor &k, ++ const Tensor &v, ++ const Tensor &slope, ++ const Tensor &initial_state_indices, ++ const Tensor &final_state_indices) { ++ size_t seed = hash_combine(out, ++ initial_state, ++ q, ++ k, ++ v, ++ slope, ++ initial_state_indices, ++ final_state_indices); ++ ++ INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( ++ Descriptor, descriptor, LightningAttention, ++ seed, ++ out->desc(), ++ initial_state->desc(), ++ q->desc(), ++ k->desc(), ++ v->desc(), ++ slope->desc(), ++ initial_state_indices->desc(), ++ final_state_indices->desc()); ++ ++ INFINIOP_WORKSPACE_TENSOR(workspace, LightningAttention, descriptor); ++ ++ return new PlannedMeta{ ++ descriptor, ++ graph::GraphTensor(workspace), ++ graph::GraphTensor(out), ++ graph::GraphTensor(initial_state), ++ graph::GraphTensor(q), ++ graph::GraphTensor(k), ++ graph::GraphTensor(v), ++ graph::GraphTensor(slope), ++ graph::GraphTensor(initial_state_indices), ++ graph::GraphTensor(final_state_indices)}; ++} ++ ++void run(void *planned_meta) { ++ auto planned = reinterpret_cast(planned_meta); ++ ++ INFINICORE_CHECK_ERROR(infiniopLightningAttention( ++ planned->descriptor->desc, ++ planned->workspace->data(), ++ planned->workspace->numel(), ++ planned->out->data(), ++ planned->initial_state->data(), ++ planned->q->data(), ++ planned->k->data(), ++ planned->v->data(), ++ planned->slope->data(), ++ planned->initial_state_indices->data(), ++ planned->final_state_indices->data(), ++ context::getStream())); ++} ++ ++void cleanup(void **planned_meta_ptr) { ++ delete *reinterpret_cast(planned_meta_ptr); ++ *planned_meta_ptr = nullptr; ++} ++ ++INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(LightningAttention, &plan, &run, &cleanup); ++ ++} // namespace infinicore::op::lightning_attention_impl::infiniop +diff --git a/src/infinicore/pybind11/ops.hpp b/src/infinicore/pybind11/ops.hpp +index e78adf27..9b5d9dfa 100644 +--- a/src/infinicore/pybind11/ops.hpp ++++ b/src/infinicore/pybind11/ops.hpp +@@ -77,6 +77,7 @@ + #include "ops/ldexp.hpp" + #include "ops/lerp.hpp" + #include "ops/linear.hpp" ++#include "ops/lightning_attention.hpp" + #include "ops/linear_mxfp4.hpp" + #include "ops/linear_w8a8i8.hpp" + #include "ops/log_softmax.hpp" +@@ -229,6 +230,7 @@ inline void bind(py::module &m) { + bind_hinge_embedding_loss(m); + bind_kv_caching(m); + bind_kimi_delta_attention(m); ++ bind_lightning_attention(m); + bind_fmod(m); + bind_fp8_indexer_logits(m); + bind_fp8_indexer_quant(m); +@@ -345,3 +347,5 @@ inline void bind(py::module &m) { + } + + } // namespace infinicore::ops ++ ++ +diff --git a/src/infinicore/pybind11/ops/lightning_attention.hpp b/src/infinicore/pybind11/ops/lightning_attention.hpp +new file mode 100644 +index 00000000..af819cc0 +--- /dev/null ++++ b/src/infinicore/pybind11/ops/lightning_attention.hpp +@@ -0,0 +1,25 @@ ++#pragma once ++ ++#include ++ ++#include "infinicore/ops/lightning_attention.hpp" ++ ++namespace py = pybind11; ++ ++namespace infinicore::ops { ++ ++inline void bind_lightning_attention(py::module &m) { ++ m.def("lightning_attention", ++ &op::lightning_attention, ++ py::arg("q"), ++ py::arg("k"), ++ py::arg("v"), ++ py::arg("slope"), ++ py::arg("initial_state"), ++ py::arg("initial_state_indices"), ++ py::arg("final_state_indices"), ++ R"doc(Indexed-pool lightning attention (MiniMax-01 style). ++Returns out [B, T, H, D]; updates `initial_state` in place at final_state_indices rows.)doc"); ++} ++ ++} // namespace infinicore::ops +diff --git a/src/infiniop/ops/lightning_attention/cpu/lightning_attention_cpu.cc b/src/infiniop/ops/lightning_attention/cpu/lightning_attention_cpu.cc +new file mode 100644 +index 00000000..61329492 +--- /dev/null ++++ b/src/infiniop/ops/lightning_attention/cpu/lightning_attention_cpu.cc +@@ -0,0 +1,155 @@ ++#include "lightning_attention_cpu.h" ++#include "../../../../infiniop/handle.h" ++#include "../../../../utils.h" ++#include "../../../../utils/custom_types.h" ++#include ++#include ++#include ++ ++namespace op::lightning_attention::cpu { ++ ++Descriptor::~Descriptor() {} ++ ++template ++static infiniStatus_t lightning_attention_cpu_impl(const LightningAttentionInfo &info, ++ T *out, T *initial_state, ++ const T *q, const T *k, const T *v, ++ const float *slope, ++ const void *init_idx_ptr, ++ const void *final_idx_ptr) { ++ const size_t B = info.B; ++ const size_t Tlen = info.T; ++ const size_t H = info.H; ++ const size_t D = info.D; ++ ++ const auto &out_s = info.out_strides; ++ const auto &state_s = info.initial_state_strides; ++ const auto &q_s = info.q_strides; ++ const auto &k_s = info.k_strides; ++ const auto &v_s = info.v_strides; ++ const auto &slope_s = info.slope_strides; ++ ++ const auto read_index = [&](const void *ptr, size_t i) -> size_t { ++ if (info.index_dtype == INFINI_DTYPE_I32) { ++ return static_cast(reinterpret_cast(ptr)[i]); ++ } ++ return static_cast(reinterpret_cast(ptr)[i]); ++ }; ++ ++ // Per-request recurrent state [H, D, D], accumulated in fp32. ++ std::vector S(H * D * D, 0.0f); ++ ++ for (size_t b = 0; b < B; ++b) { ++ const size_t state_row = read_index(init_idx_ptr, b); ++ const size_t final_row = read_index(final_idx_ptr, b); ++ ++ // Load initial state. ++ for (size_t h = 0; h < H; ++h) { ++ for (size_t i = 0; i < D; ++i) { ++ for (size_t j = 0; j < D; ++j) { ++ S[(h * D + i) * D + j] = utils::cast( ++ initial_state[state_row * state_s[0] + h * state_s[1] + i * state_s[2] + j * state_s[3]]); ++ } ++ } ++ } ++ ++ for (size_t t = 0; t < Tlen; ++t) { ++ const size_t q_base = b * q_s[0] + t * q_s[1]; ++ const size_t k_base = b * k_s[0] + t * k_s[1]; ++ const size_t v_base = b * v_s[0] + t * v_s[1]; ++ const size_t out_base = b * out_s[0] + t * out_s[1]; ++ ++ for (size_t h = 0; h < H; ++h) { ++ const float ratio = std::exp(-slope[h * slope_s[0]]); ++ const size_t head_base = (h * D) * D; ++ const size_t q_h = q_base + h * q_s[2]; ++ const size_t k_h = k_base + h * k_s[2]; ++ const size_t v_h = v_base + h * v_s[2]; ++ const size_t out_h = out_base + h * out_s[2]; ++ ++ // S' = ratio * S + k^T v ++ for (size_t i = 0; i < D; ++i) { ++ const float k_i = utils::cast(k[k_h + i * k_s[3]]); ++ float *S_i = S.data() + head_base + i * D; ++ for (size_t j = 0; j < D; ++j) { ++ const float v_j = utils::cast(v[v_h + j * v_s[3]]); ++ S_i[j] = ratio * S_i[j] + k_i * v_j; ++ } ++ } ++ // o[j] = sum_i q[i] * S[i, j] ++ for (size_t j = 0; j < D; ++j) { ++ float acc = 0.0f; ++ for (size_t i = 0; i < D; ++i) { ++ const float q_i = utils::cast(q[q_h + i * q_s[3]]); ++ acc += q_i * S[head_base + i * D + j]; ++ } ++ out[out_h + j * out_s[3]] = utils::cast(acc); ++ } ++ } ++ } ++ ++ // Write the final state back into the pool. ++ for (size_t h = 0; h < H; ++h) { ++ for (size_t i = 0; i < D; ++i) { ++ for (size_t j = 0; j < D; ++j) { ++ initial_state[final_row * state_s[0] + h * state_s[1] + i * state_s[2] + j * state_s[3]] = ++ utils::cast(S[(h * D + i) * D + j]); ++ } ++ } ++ } ++ } ++ return INFINI_STATUS_SUCCESS; ++} ++ ++infiniStatus_t Descriptor::create( ++ infiniopHandle_t handle, ++ Descriptor **desc_ptr, ++ infiniopTensorDescriptor_t out_desc, ++ infiniopTensorDescriptor_t initial_state_desc, ++ infiniopTensorDescriptor_t q_desc, ++ infiniopTensorDescriptor_t k_desc, ++ infiniopTensorDescriptor_t v_desc, ++ infiniopTensorDescriptor_t slope_desc, ++ infiniopTensorDescriptor_t initial_state_indices_desc, ++ infiniopTensorDescriptor_t final_state_indices_desc) { ++ auto result = LightningAttentionInfo::create(out_desc, initial_state_desc, ++ q_desc, k_desc, v_desc, slope_desc, ++ initial_state_indices_desc, ++ final_state_indices_desc); ++ CHECK_RESULT(result); ++ *desc_ptr = new Descriptor(nullptr, result.take(), 0, handle->device, handle->device_id); ++ return INFINI_STATUS_SUCCESS; ++} ++ ++infiniStatus_t Descriptor::calculate( ++ void *workspace, size_t workspace_size, ++ void *out, void *initial_state, ++ const void *q, const void *k, const void *v, ++ const void *slope, ++ const void *initial_state_indices, ++ const void *final_state_indices, ++ void *stream) const { ++ if (_info.data_dtype == INFINI_DTYPE_F32) { ++ return lightning_attention_cpu_impl( ++ _info, (float *)out, (float *)initial_state, ++ (const float *)q, (const float *)k, (const float *)v, ++ (const float *)slope, initial_state_indices, final_state_indices); ++ } ++ if (_info.data_dtype == INFINI_DTYPE_F16) { ++ return lightning_attention_cpu_impl( ++ _info, (fp16_t *)out, (fp16_t *)initial_state, ++ (const fp16_t *)q, (const fp16_t *)k, (const fp16_t *)v, ++ (const float *)slope, initial_state_indices, final_state_indices); ++ } ++ if (_info.data_dtype == INFINI_DTYPE_BF16) { ++ return lightning_attention_cpu_impl( ++ _info, (bf16_t *)out, (bf16_t *)initial_state, ++ (const bf16_t *)q, (const bf16_t *)k, (const bf16_t *)v, ++ (const float *)slope, initial_state_indices, final_state_indices); ++ } ++ return INFINI_STATUS_BAD_TENSOR_DTYPE; ++} ++ ++} // namespace op::lightning_attention::cpu ++ ++ +diff --git a/src/infiniop/ops/lightning_attention/cpu/lightning_attention_cpu.h b/src/infiniop/ops/lightning_attention/cpu/lightning_attention_cpu.h +new file mode 100644 +index 00000000..50c320a0 +--- /dev/null ++++ b/src/infiniop/ops/lightning_attention/cpu/lightning_attention_cpu.h +@@ -0,0 +1,7 @@ ++#ifndef __LIGHTNING_ATTENTION_CPU_H__ ++#define __LIGHTNING_ATTENTION_CPU_H__ ++#include "../lightning_attention.h" ++ ++DESCRIPTOR(cpu) ++ ++#endif +diff --git a/src/infiniop/ops/lightning_attention/info.h b/src/infiniop/ops/lightning_attention/info.h +new file mode 100644 +index 00000000..0d28f7a7 +--- /dev/null ++++ b/src/infiniop/ops/lightning_attention/info.h +@@ -0,0 +1,134 @@ ++// infiniop/ops/lightning_attention/info.h ++ ++#ifndef __LIGHTNING_ATTENTION_INFO_H__ ++#define __LIGHTNING_ATTENTION_INFO_H__ ++ ++#include "../../../utils.h" ++#include "../../tensor.h" ++#include ++ ++namespace op { ++namespace lightning_attention { ++ ++class LightningAttentionInfo { ++ LightningAttentionInfo() = default; ++ ++public: ++ infiniDtype_t data_dtype; ++ infiniDtype_t index_dtype; ++ size_t B, T, H, D, pool_size; ++ ++ std::vector out_strides; ++ std::vector initial_state_strides; ++ std::vector q_strides; ++ std::vector k_strides; ++ std::vector v_strides; ++ std::vector slope_strides; ++ std::vector initial_state_indices_strides; ++ std::vector final_state_indices_strides; ++ ++ static utils::Result ++ create(infiniopTensorDescriptor_t out_desc, ++ infiniopTensorDescriptor_t initial_state_desc, ++ infiniopTensorDescriptor_t q_desc, ++ infiniopTensorDescriptor_t k_desc, ++ infiniopTensorDescriptor_t v_desc, ++ infiniopTensorDescriptor_t slope_desc, ++ infiniopTensorDescriptor_t initial_state_indices_desc, ++ infiniopTensorDescriptor_t final_state_indices_desc) { ++ if (out_desc == nullptr || initial_state_desc == nullptr || q_desc == nullptr || ++ k_desc == nullptr || v_desc == nullptr || slope_desc == nullptr || ++ initial_state_indices_desc == nullptr || final_state_indices_desc == nullptr) { ++ return INFINI_STATUS_NULL_POINTER; ++ } ++ ++ auto data_dtype = q_desc->dtype(); ++ CHECK_DTYPE(data_dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32); ++ if (k_desc->dtype() != data_dtype || v_desc->dtype() != data_dtype || ++ out_desc->dtype() != data_dtype || initial_state_desc->dtype() != data_dtype) { ++ return INFINI_STATUS_BAD_TENSOR_DTYPE; ++ } ++ if (slope_desc->dtype() != INFINI_DTYPE_F32) { ++ return INFINI_STATUS_BAD_TENSOR_DTYPE; ++ } ++ ++ auto index_dtype = initial_state_indices_desc->dtype(); ++ CHECK_DTYPE(index_dtype, INFINI_DTYPE_I32, INFINI_DTYPE_I64); ++ if (final_state_indices_desc->dtype() != index_dtype) { ++ return INFINI_STATUS_BAD_TENSOR_DTYPE; ++ } ++ ++ const auto &q_shape = q_desc->shape(); ++ const auto &k_shape = k_desc->shape(); ++ const auto &v_shape = v_desc->shape(); ++ const auto &out_shape = out_desc->shape(); ++ const auto &state_shape = initial_state_desc->shape(); ++ const auto &slope_shape = slope_desc->shape(); ++ if (q_shape.size() != 4 || k_shape.size() != 4 || v_shape.size() != 4 || out_shape.size() != 4) { ++ return INFINI_STATUS_BAD_TENSOR_SHAPE; ++ } ++ if (k_shape != q_shape || v_shape != q_shape || out_shape != q_shape) { ++ return INFINI_STATUS_BAD_TENSOR_SHAPE; ++ } ++ const size_t B = q_shape[0]; ++ const size_t T = q_shape[1]; ++ const size_t H = q_shape[2]; ++ const size_t D = q_shape[3]; ++ if (B == 0 || T == 0 || H == 0 || D == 0) { ++ return INFINI_STATUS_BAD_TENSOR_SHAPE; ++ } ++ if (state_shape.size() != 4 || state_shape[1] != H || state_shape[2] != D || state_shape[3] != D) { ++ return INFINI_STATUS_BAD_TENSOR_SHAPE; ++ } ++ const size_t pool_size = state_shape[0]; ++ if (pool_size == 0) { ++ return INFINI_STATUS_BAD_TENSOR_SHAPE; ++ } ++ if (slope_shape.size() != 1 || slope_shape[0] != H) { ++ return INFINI_STATUS_BAD_TENSOR_SHAPE; ++ } ++ if (initial_state_indices_desc->shape().size() != 1 || ++ initial_state_indices_desc->shape()[0] != B || ++ final_state_indices_desc->shape().size() != 1 || ++ final_state_indices_desc->shape()[0] != B) { ++ return INFINI_STATUS_BAD_TENSOR_SHAPE; ++ } ++ ++ // last dim must be contiguous for the sequence tensors (as in gated delta rule) ++ if (q_desc->stride(3) != 1 || k_desc->stride(3) != 1 || ++ v_desc->stride(3) != 1 || out_desc->stride(3) != 1) { ++ return INFINI_STATUS_BAD_TENSOR_STRIDES; ++ } ++ ++ // Both implementations read the state indices with unit stride. ++ if (initial_state_indices_desc->stride(0) != 1 || ++ final_state_indices_desc->stride(0) != 1) { ++ return INFINI_STATUS_BAD_TENSOR_STRIDES; ++ } ++ ++ LightningAttentionInfo info; ++ info.data_dtype = data_dtype; ++ info.index_dtype = index_dtype; ++ info.B = B; ++ info.T = T; ++ info.H = H; ++ info.D = D; ++ info.pool_size = pool_size; ++ info.out_strides = out_desc->strides(); ++ info.initial_state_strides = initial_state_desc->strides(); ++ info.q_strides = q_desc->strides(); ++ info.k_strides = k_desc->strides(); ++ info.v_strides = v_desc->strides(); ++ info.slope_strides = slope_desc->strides(); ++ info.initial_state_indices_strides = initial_state_indices_desc->strides(); ++ info.final_state_indices_strides = final_state_indices_desc->strides(); ++ return utils::Result(info); ++ } ++}; ++ ++} // namespace lightning_attention ++} // namespace op ++ ++#endif // __LIGHTNING_ATTENTION_INFO_H__ ++ ++ +diff --git a/src/infiniop/ops/lightning_attention/lightning_attention.h b/src/infiniop/ops/lightning_attention/lightning_attention.h +new file mode 100644 +index 00000000..b7f4f087 +--- /dev/null ++++ b/src/infiniop/ops/lightning_attention/lightning_attention.h +@@ -0,0 +1,57 @@ ++// infiniop/ops/lightning_attention.h ++ ++#ifndef __INFINIOP_LIGHTNING_ATTENTION_H__ ++#define __INFINIOP_LIGHTNING_ATTENTION_H__ ++ ++#include "../../operator.h" ++#include "info.h" ++ ++#define DESCRIPTOR(NAMESPACE) \ ++ \ ++ namespace op::lightning_attention::NAMESPACE { \ ++ class Descriptor final : public InfiniopDescriptor { \ ++ struct Opaque; \ ++ Opaque *_opaque; \ ++ LightningAttentionInfo _info; \ ++ size_t _workspace_size; \ ++ \ ++ Descriptor( \ ++ Opaque *opaque, \ ++ LightningAttentionInfo info, \ ++ size_t workspace_size, \ ++ infiniDevice_t device_type, \ ++ int device_id) \ ++ : InfiniopDescriptor{device_type, device_id}, \ ++ _opaque(opaque), \ ++ _info(info), \ ++ _workspace_size(workspace_size) {} \ ++ \ ++ public: \ ++ ~Descriptor(); \ ++ \ ++ size_t workspaceSize() const { return _workspace_size; } \ ++ \ ++ static infiniStatus_t create( \ ++ infiniopHandle_t handle, \ ++ Descriptor **desc_ptr, \ ++ infiniopTensorDescriptor_t out_desc, \ ++ infiniopTensorDescriptor_t initial_state_desc, \ ++ infiniopTensorDescriptor_t q_desc, \ ++ infiniopTensorDescriptor_t k_desc, \ ++ infiniopTensorDescriptor_t v_desc, \ ++ infiniopTensorDescriptor_t slope_desc, \ ++ infiniopTensorDescriptor_t initial_state_indices_desc, \ ++ infiniopTensorDescriptor_t final_state_indices_desc); \ ++ \ ++ infiniStatus_t calculate( \ ++ void *workspace, size_t workspace_size, \ ++ void *out, void *initial_state, \ ++ const void *q, const void *k, const void *v, \ ++ const void *slope, \ ++ const void *initial_state_indices, \ ++ const void *final_state_indices, \ ++ void *stream) const; \ ++ }; \ ++ } ++ ++#endif // __INFINIOP_LIGHTNING_ATTENTION_H__ +diff --git a/src/infiniop/ops/lightning_attention/nvidia/lightning_attention_nvidia.cu b/src/infiniop/ops/lightning_attention/nvidia/lightning_attention_nvidia.cu +new file mode 100644 +index 00000000..c6164a03 +--- /dev/null ++++ b/src/infiniop/ops/lightning_attention/nvidia/lightning_attention_nvidia.cu +@@ -0,0 +1,160 @@ ++#include "../../../devices/nvidia/nvidia_common.cuh" ++#include "lightning_attention_nvidia.cuh" ++ ++#include "../../../devices/nvidia/nvidia_kernel_common.cuh" ++ ++#include ++#include ++#include ++ ++namespace op::lightning_attention::nvidia { ++ ++struct Descriptor::Opaque { ++ std::shared_ptr internal; ++}; ++ ++Descriptor::~Descriptor() { ++ delete _opaque; ++} ++ ++// One block per (batch, head); one thread per state/output column. The block is ++// launched with exactly `D` threads (`D <= maxThreadsPerBlock()`, validated in ++// `Descriptor::calculate`), so every thread reaches each `__syncthreads()`. ++// ++// The recurrence is staged into the destination row of the state pool: the ++// initial row is copied to the final row first and the accumulation happens in ++// place on the final row. This keeps the initial row untouched, matching the ++// CPU implementation when `initial_state_indices != final_state_indices`. ++// ++// NOTE: fp32 only for now; fp16/bf16 kernels are a follow-up. ++INFINIOP_CUDA_KERNEL lightningAttentionKernel( ++ const float *__restrict__ q, const float *__restrict__ k, const float *__restrict__ v, ++ float *__restrict__ out, float *__restrict__ state_pool, ++ const float *__restrict__ slope, ++ const int32_t *__restrict__ init_idx, const int32_t *__restrict__ final_idx, ++ size_t T, size_t D, ++ ptrdiff_t q_sb, ptrdiff_t q_st, ptrdiff_t q_sh, ptrdiff_t q_sd, ++ ptrdiff_t k_sb, ptrdiff_t k_st, ptrdiff_t k_sh, ptrdiff_t k_sd, ++ ptrdiff_t v_sb, ptrdiff_t v_st, ptrdiff_t v_sh, ptrdiff_t v_sd, ++ ptrdiff_t o_sb, ptrdiff_t o_st, ptrdiff_t o_sh, ptrdiff_t o_sd, ++ ptrdiff_t s_s0, ptrdiff_t s_s1, ptrdiff_t s_s2, ptrdiff_t s_s3, ++ size_t slope_stride) { ++ const size_t b = blockIdx.y; ++ const size_t h = blockIdx.x; ++ const size_t tid = threadIdx.x; ++ ++ extern __shared__ float smem[]; ++ float *s_k = smem; ++ float *s_q = smem + D; ++ ++ const size_t init_row = static_cast(init_idx[b]); ++ const size_t final_row = static_cast(final_idx[b]); ++ const float ratio = __expf(-slope[h * slope_stride]); ++ ++ float *S = state_pool + final_row * s_s0 + h * s_s1; ++ const float *S_init = state_pool + init_row * s_s0 + h * s_s1; ++ if (init_row != final_row) { ++ for (size_t idx = tid; idx < D * D; idx += D) { ++ const size_t i = idx / D; ++ const size_t j = idx % D; ++ S[i * s_s2 + j * s_s3] = S_init[i * s_s2 + j * s_s3]; ++ } ++ } ++ __syncthreads(); ++ ++ for (size_t t = 0; t < T; ++t) { ++ s_k[tid] = k[b * k_sb + t * k_st + h * k_sh + tid * k_sd]; ++ s_q[tid] = q[b * q_sb + t * q_st + h * q_sh + tid * q_sd]; ++ __syncthreads(); ++ ++ // S[i][j] = ratio * S[i][j] + k[i] * v[j] (thread j owns column j) ++ const float v_j = v[b * v_sb + t * v_st + h * v_sh + tid * v_sd]; ++ for (size_t i = 0; i < D; ++i) { ++ float *s_ij = S + i * s_s2 + tid * s_s3; ++ *s_ij = ratio * (*s_ij) + s_k[i] * v_j; ++ } ++ ++ // o[j] = sum_i q[i] * S[i][j] ++ float acc = 0.0f; ++ for (size_t i = 0; i < D; ++i) { ++ acc += s_q[i] * S[i * s_s2 + tid * s_s3]; ++ } ++ out[b * o_sb + t * o_st + h * o_sh + tid * o_sd] = acc; ++ ++ __syncthreads(); // All threads must finish reading s_k/s_q before reloading. ++ } ++} ++ ++infiniStatus_t Descriptor::create( ++ infiniopHandle_t handle, ++ Descriptor **desc_ptr, ++ infiniopTensorDescriptor_t out_desc, ++ infiniopTensorDescriptor_t initial_state_desc, ++ infiniopTensorDescriptor_t q_desc, ++ infiniopTensorDescriptor_t k_desc, ++ infiniopTensorDescriptor_t v_desc, ++ infiniopTensorDescriptor_t slope_desc, ++ infiniopTensorDescriptor_t initial_state_indices_desc, ++ infiniopTensorDescriptor_t final_state_indices_desc) { ++ auto info = LightningAttentionInfo::create(out_desc, initial_state_desc, ++ q_desc, k_desc, v_desc, slope_desc, ++ initial_state_indices_desc, ++ final_state_indices_desc); ++ CHECK_RESULT(info); ++ *desc_ptr = new Descriptor( ++ new Opaque{reinterpret_cast(handle)->internal()}, ++ info.take(), 0, handle->device, handle->device_id); ++ return INFINI_STATUS_SUCCESS; ++} ++ ++infiniStatus_t Descriptor::calculate( ++ void *workspace, size_t workspace_size, ++ void *out, void *initial_state, ++ const void *q, const void *k, const void *v, ++ const void *slope, ++ const void *initial_state_indices, ++ const void *final_state_indices, ++ void *stream_) const { ++ (void)workspace; ++ (void)workspace_size; ++ ++ if (_info.data_dtype != INFINI_DTYPE_F32) { ++ return INFINI_STATUS_BAD_TENSOR_DTYPE; ++ } ++ if (_info.index_dtype != INFINI_DTYPE_I32) { ++ return INFINI_STATUS_BAD_TENSOR_DTYPE; ++ } ++ if (_info.D > static_cast(_opaque->internal->maxThreadsPerBlock())) { ++ return INFINI_STATUS_BAD_TENSOR_SHAPE; ++ } ++ if (_info.B > static_cast(65535)) { ++ return INFINI_STATUS_BAD_TENSOR_SHAPE; ++ } ++ ++ cudaStream_t stream = reinterpret_cast(stream_); ++ const auto &info = _info; ++ dim3 grid(static_cast(info.H), static_cast(info.B)); ++ const size_t smem_bytes = 2 * info.D * sizeof(float); ++ lightningAttentionKernel<<(info.D), smem_bytes, stream>>>( ++ static_cast(q), static_cast(k), static_cast(v), ++ static_cast(out), static_cast(initial_state), ++ static_cast(slope), ++ static_cast(initial_state_indices), ++ static_cast(final_state_indices), ++ info.T, info.D, ++ info.q_strides[0], info.q_strides[1], info.q_strides[2], info.q_strides[3], ++ info.k_strides[0], info.k_strides[1], info.k_strides[2], info.k_strides[3], ++ info.v_strides[0], info.v_strides[1], info.v_strides[2], info.v_strides[3], ++ info.out_strides[0], info.out_strides[1], info.out_strides[2], info.out_strides[3], ++ info.initial_state_strides[0], info.initial_state_strides[1], ++ info.initial_state_strides[2], info.initial_state_strides[3], ++ static_cast(info.slope_strides[0])); ++ ++ if (cudaGetLastError() != cudaSuccess) { ++ return INFINI_STATUS_INTERNAL_ERROR; ++ } ++ return INFINI_STATUS_SUCCESS; ++} ++ ++} // namespace op::lightning_attention::nvidia ++ +diff --git a/src/infiniop/ops/lightning_attention/nvidia/lightning_attention_nvidia.cuh b/src/infiniop/ops/lightning_attention/nvidia/lightning_attention_nvidia.cuh +new file mode 100644 +index 00000000..401ca434 +--- /dev/null ++++ b/src/infiniop/ops/lightning_attention/nvidia/lightning_attention_nvidia.cuh +@@ -0,0 +1,8 @@ ++#ifndef __LIGHTNING_ATTENTION_NVIDIA_CUH__ ++#define __LIGHTNING_ATTENTION_NVIDIA_CUH__ ++ ++#include "../lightning_attention.h" ++ ++DESCRIPTOR(nvidia) ++ ++#endif // __LIGHTNING_ATTENTION_NVIDIA_CUH__ +diff --git a/src/infiniop/ops/lightning_attention/operator.cc b/src/infiniop/ops/lightning_attention/operator.cc +new file mode 100644 +index 00000000..8813b111 +--- /dev/null ++++ b/src/infiniop/ops/lightning_attention/operator.cc +@@ -0,0 +1,135 @@ ++// infiniop/ops/lightning_attention/operator.cc ++ ++#include "../../operator.h" ++#include "../../handle.h" ++#include "infiniop/ops/lightning_attention.h" ++ ++#ifdef ENABLE_NVIDIA_API ++#include "nvidia/lightning_attention_nvidia.cuh" ++#endif ++#ifdef ENABLE_CPU_API ++#include "cpu/lightning_attention_cpu.h" ++#endif ++ ++__INFINI_C infiniStatus_t infiniopCreateLightningAttentionDescriptor( ++ infiniopHandle_t handle, ++ infiniopLightningAttentionDescriptor_t *desc_ptr, ++ infiniopTensorDescriptor_t out_desc, ++ infiniopTensorDescriptor_t initial_state_desc, ++ infiniopTensorDescriptor_t q_desc, ++ infiniopTensorDescriptor_t k_desc, ++ infiniopTensorDescriptor_t v_desc, ++ infiniopTensorDescriptor_t slope_desc, ++ infiniopTensorDescriptor_t initial_state_indices_desc, ++ infiniopTensorDescriptor_t final_state_indices_desc) { ++#define CREATE(CASE, NAMESPACE) \ ++ case CASE: \ ++ return op::lightning_attention::NAMESPACE::Descriptor::create( \ ++ handle, \ ++ reinterpret_cast( \ ++ desc_ptr), \ ++ out_desc, initial_state_desc, q_desc, k_desc, v_desc, slope_desc, \ ++ initial_state_indices_desc, final_state_indices_desc); ++ ++ switch (handle->device) { ++#ifdef ENABLE_CPU_API ++ CREATE(INFINI_DEVICE_CPU, cpu) ++#endif ++#ifdef ENABLE_NVIDIA_API ++ CREATE(INFINI_DEVICE_NVIDIA, nvidia) ++#endif ++#ifdef ENABLE_HYGON_API ++ CREATE(INFINI_DEVICE_HYGON, nvidia) ++#endif ++ ++ default: ++ return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; ++ } ++#undef CREATE ++} ++ ++__INFINI_C infiniStatus_t infiniopGetLightningAttentionWorkspaceSize( ++ infiniopLightningAttentionDescriptor_t desc, size_t *size) { ++#define GET(CASE, NAMESPACE) \ ++ case CASE: \ ++ *size = reinterpret_cast< \ ++ op::lightning_attention::NAMESPACE::Descriptor *>( \ ++ desc) \ ++ ->workspaceSize(); \ ++ return INFINI_STATUS_SUCCESS; ++ ++ switch (desc->device_type) { ++#ifdef ENABLE_CPU_API ++ GET(INFINI_DEVICE_CPU, cpu) ++#endif ++#ifdef ENABLE_NVIDIA_API ++ GET(INFINI_DEVICE_NVIDIA, nvidia) ++#endif ++#ifdef ENABLE_HYGON_API ++ GET(INFINI_DEVICE_HYGON, nvidia) ++#endif ++ ++ default: ++ return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; ++ } ++#undef GET ++} ++ ++__INFINI_C infiniStatus_t infiniopLightningAttention( ++ infiniopLightningAttentionDescriptor_t desc, ++ void *workspace, size_t workspace_size, ++ void *out, void *initial_state, ++ const void *q, const void *k, const void *v, ++ const void *slope, ++ const void *initial_state_indices, ++ const void *final_state_indices, ++ void *stream) { ++#define CALCULATE(CASE, NAMESPACE) \ ++ case CASE: \ ++ return reinterpret_cast< \ ++ op::lightning_attention::NAMESPACE::Descriptor *>(desc) \ ++ ->calculate(workspace, workspace_size, out, initial_state, \ ++ q, k, v, slope, initial_state_indices, \ ++ final_state_indices, stream); ++ ++ switch (desc->device_type) { ++#ifdef ENABLE_CPU_API ++ CALCULATE(INFINI_DEVICE_CPU, cpu) ++#endif ++#ifdef ENABLE_NVIDIA_API ++ CALCULATE(INFINI_DEVICE_NVIDIA, nvidia) ++#endif ++#ifdef ENABLE_HYGON_API ++ CALCULATE(INFINI_DEVICE_HYGON, nvidia) ++#endif ++ ++ default: ++ return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; ++ } ++#undef CALCULATE ++} ++ ++__INFINI_C infiniStatus_t infiniopDestroyLightningAttentionDescriptor( ++ infiniopLightningAttentionDescriptor_t desc) { ++#define DESTROY(CASE, NAMESPACE) \ ++ case CASE: \ ++ delete reinterpret_cast< \ ++ op::lightning_attention::NAMESPACE::Descriptor *>(desc); \ ++ return INFINI_STATUS_SUCCESS; ++ ++ switch (desc->device_type) { ++#ifdef ENABLE_CPU_API ++ DESTROY(INFINI_DEVICE_CPU, cpu) ++#endif ++#ifdef ENABLE_NVIDIA_API ++ DESTROY(INFINI_DEVICE_NVIDIA, nvidia) ++#endif ++#ifdef ENABLE_HYGON_API ++ DESTROY(INFINI_DEVICE_HYGON, nvidia) ++#endif ++ ++ default: ++ return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; ++ } ++#undef DESTROY ++} diff --git a/python/infinilm/modeling_utils.py b/python/infinilm/modeling_utils.py index b210b0c56..c037a80e4 100644 --- a/python/infinilm/modeling_utils.py +++ b/python/infinilm/modeling_utils.py @@ -1071,6 +1071,71 @@ def _remap_kimi_k3(state_dict, config): return state_dict + + +def _remap_minimax(state_dict, config): + """Adapt HF transformers `minimax` (MiniMax-Text-01) keys to InfiniLM keys.""" + # Lightning layers use `self_attn.{qkv_proj,output_gate,out_proj,norm}` in HF + # and `linear_attn.*` in InfiniLM; full-attention `self_attn.q/k/v/o_proj` keep + # their names. + state_dict = rename_keys( + state_dict, + { + "self_attn.qkv_proj": "linear_attn.qkv_proj", + "self_attn.output_gate": "linear_attn.output_gate", + "self_attn.out_proj": "linear_attn.out_proj", + "self_attn.norm": "linear_attn.norm", + }, + ) + + num_experts = config.get("num_local_experts", config.get("num_experts", 8)) + projection_names = { + "w1": "gate_proj", + "w2": "down_proj", + "w3": "up_proj", + } + expert_pattern = re.compile( + r"^(.*\.)block_sparse_moe\.experts\.(\d+)\.(w1|w2|w3)\.weight$" + ) + + remapped = {} + for key, tensor in state_dict.items(): + match = expert_pattern.match(key) + if match: + prefix, expert_idx, projection = match.groups() + target = projection_names[projection] + if num_experts == 1: + remapped[f"{prefix}mlp.{target}.weight"] = tensor + else: + remapped[f"{prefix}moe.experts.{expert_idx}.{target}.weight"] = tensor + continue + + if key.endswith(".block_sparse_moe.gate.weight"): + if num_experts > 1: + prefix = key[: -len("block_sparse_moe.gate.weight")] + remapped[prefix + "moe.gate.weight"] = tensor + continue + + # Compatibility with older HF MiniMax checkpoints that store packed expert weights. + if key.endswith("mlp.experts.gate_up_proj"): + base = key[: -len("mlp.experts.gate_up_proj")] + if num_experts == 1: + gate_up = tensor.squeeze(0) + gate, up = gate_up.chunk(2, dim=0) + remapped[base + "mlp.gate_proj.weight"] = gate + remapped[base + "mlp.up_proj.weight"] = up + else: + remapped[base + "moe.experts.w13_weight"] = tensor + continue + if key.endswith("mlp.experts.down_proj"): + base = key[: -len("mlp.experts.down_proj")] + target = "mlp.down_proj.weight" if num_experts == 1 else "moe.experts.w2_weight" + remapped[base + target] = tensor.squeeze(0) if num_experts == 1 else tensor + continue + + remapped[key] = tensor + return remapped + _WEIGHT_REMAPPER = { "glm4": _remap_glm4, "chatglm": _remap_chatglm, @@ -1082,5 +1147,8 @@ def _remap_kimi_k3(state_dict, config): "ernie4_5_moe_vl": _remap_ernie4_5_moe_vl, "qwen3_5_moe": _remap_qwen3_5_moe, "qwen3_next": _remap_qwen3_next, + "minimax": _remap_minimax, "kimi_k3": _remap_kimi_k3, } + + diff --git a/test/models/minimax/smoke_minimax.py b/test/models/minimax/smoke_minimax.py new file mode 100644 index 000000000..0e20ac01d --- /dev/null +++ b/test/models/minimax/smoke_minimax.py @@ -0,0 +1,215 @@ +#!/usr/bin/env python3 +""" +M1 CPU smoke test for the InfiniLM `minimax` model (MiniMax-Text-01 / MiniMax-M2 family). + +Checks: +1. The C++ `minimax` model type registers and constructs (config normalization runs). +2. Parameter names look sane. +3. Random-weight prefill + decode forward passes produce finite logits. +4. Decode-after-N-tokens logits match a full (N+1)-token prefill logits at the last position + (linear-attention state carry consistency through the lightning layers). +""" +import ctypes +import json +import os +import sys + +import numpy as np +import torch + +import infinicore +from infinicore.lib import _infinicore +from infinilm.lib import _infinilm + +INFINI_DTYPE = { + torch.float32: infinicore.float32, + torch.float16: infinicore.float16, + torch.int32: infinicore.int32, + torch.int64: infinicore.int64, +} + + +def t2i(t: torch.Tensor, dev): + # Return the raw _infinicore.Tensor (the engine bindings expect the pybind object). + t = t.contiguous() + cpu = infinicore.device("cpu", 0) + tensor = infinicore.from_blob( + t.data_ptr(), list(t.shape), dtype=INFINI_DTYPE[t.dtype], device=cpu + ) + if dev != cpu: + tensor = tensor.to(dev) + return tensor._underlying + + +def _np_dtype(infini_dtype): + if infini_dtype == infinicore.float32: + return np.float32 + if infini_dtype == infinicore.float16: + return np.float16 + if infini_dtype == infinicore.int32: + return np.int32 + if infini_dtype == infinicore.int64: + return np.int64 + raise ValueError(f"unsupported dtype {infini_dtype}") + + +def i2t(t) -> torch.Tensor: + if not hasattr(t, "_underlying"): + t = infinicore.Tensor(t) + t = t.to(infinicore.device("cpu", 0)) + t = t.contiguous() + shape = list(t.shape) + np_dtype = _np_dtype(t.dtype) + ctype = {np.float32: ctypes.c_float, np.float16: ctypes.c_uint16, + np.int32: ctypes.c_int32, np.int64: ctypes.c_int64}[np_dtype] + buf = (ctype * int(t.numel())).from_address(t.data_ptr()) + arr = np.frombuffer(buf, dtype=np_dtype).reshape(shape).copy() + return torch.from_numpy(arr) + + +def make_config(): + return { + "model_type": "minimax", + "vocab_size": 128, + "hidden_size": 32, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "head_dim": 8, + "num_hidden_layers": 4, + "intermediate_size": 64, + "num_experts": int(os.environ.get("MINIMAX_NUM_LOCAL_EXPERTS", "1")), + "num_experts_per_tok": int(os.environ.get("MINIMAX_NUM_EXPERTS_PER_TOK", "1")), + "shared_intermediate_size": 0, + "rms_norm_eps": 1e-5, + "max_position_embeddings": 64, + "attn_type_list": [0, 0, 1, 0], # 3 lightning + 1 softmax + "block": 8, + "rope_theta": 10000.0, + "torch_dtype": os.environ.get("MINIMAX_TORCH_DTYPE", "float32"),} + + +def create_engine(cfg): + device_name = os.environ.get("MINIMAX_DEVICE", "cpu").lower() + if device_name == "nvidia": + device_name = "cuda" + device_type = { + "cpu": _infinicore.Device.Type.CPU, + "cuda": _infinicore.Device.Type.NVIDIA, + }[device_name] + dev = infinicore.device(device_name, 0) + engine = _infinilm.InferEngine( + json.dumps(cfg), + _infinilm.DistConfig(1), + device_type, + _infinilm.StaticKVCacheConfig(max_batch_size=1, max_cache_len=64), + False, + "static-attn", + None, + False, + "sync", + False, + ) + return engine, dev + + +def load_random_weights(engine, dev, seed=0): + torch.manual_seed(seed) + state_dict = engine.state_dict()[0] + params = {} + keep = [] + for name, tensor in state_dict.items(): + shape = list(tensor.shape) + w = (torch.randn(shape, dtype=torch.float32) * 0.02).contiguous() + keep.append(w) + params[name] = t2i(w, dev) + engine.load_params(params, strict=True) + engine.process_weights_after_loading() + return list(state_dict.keys()) + + +def make_input(dev, input_ids, position_ids, past, total, offsets, cu, init, final, sample_all=True): + return _infinilm.InferEngine.Input( + input_ids=t2i(input_ids, dev), + position_ids=t2i(position_ids, dev), + past_sequence_lengths=t2i(past, dev), + total_sequence_lengths=t2i(total, dev), + input_offsets=t2i(offsets, dev), + cu_seqlens=t2i(cu, dev), + mamba_init_state_indices=t2i(init, dev), + mamba_final_state_indices=t2i(final, dev), + sample_all_positions=sample_all, + ) + + +def run_prefill(engine, dev, tokens, slot=0): + n = len(tokens) + input_ids = torch.tensor([tokens], dtype=torch.int32) + position_ids = torch.arange(n, dtype=torch.int32).unsqueeze(0) + past = torch.tensor([0], dtype=torch.int32) + total = torch.tensor([n], dtype=torch.int32) + offsets = torch.tensor([0, n], dtype=torch.int32) + cu = torch.tensor([0, n], dtype=torch.int32) + init = torch.tensor([slot], dtype=torch.int32) + final = torch.tensor([slot], dtype=torch.int32) + out = engine.forward(make_input(dev, input_ids, position_ids, past, total, offsets, cu, init, final)) + return i2t(out.logits) + + +def run_decode(engine, dev, token, past_len, slot=0): + input_ids = torch.tensor([[token]], dtype=torch.int32) + position_ids = torch.tensor([[past_len]], dtype=torch.int32) + past = torch.tensor([past_len], dtype=torch.int32) + total = torch.tensor([past_len + 1], dtype=torch.int32) + offsets = torch.tensor([0, 1], dtype=torch.int32) + cu = torch.tensor([0, past_len + 1], dtype=torch.int32) + init = torch.tensor([slot], dtype=torch.int32) + final = torch.tensor([slot], dtype=torch.int32) + out = engine.forward(make_input(dev, input_ids, position_ids, past, total, offsets, cu, init, final)) + return i2t(out.logits) + + +def main(): + cfg = make_config() + print("[1/5] constructing minimax engine (model_type=minimax, 3 lightning + 1 softmax layers) ...") + engine, dev = create_engine(cfg) + + keys = load_random_weights(engine, dev, seed=42) + print(f"[2/5] loaded {len(keys)} random parameters") + for k in keys: + print(" ", k) + + print("[3/5] prefill forward (4 tokens) ...") + logits_prefill4 = run_prefill(engine, dev, [3, 7, 9, 2], slot=0) + assert np.isfinite(logits_prefill4.numpy()).all(), "prefill logits not finite" + print(" prefill4 logits shape:", tuple(logits_prefill4.shape), "finite: True") + + print("[4/5] decode forward (1 token after 4-token context) ...") + logits_decode = run_decode(engine, dev, 11, past_len=4, slot=0) + assert np.isfinite(logits_decode.numpy()).all(), "decode logits not finite" + print(" decode logits shape:", tuple(logits_decode.shape), "finite: True") + + print("[5/5] consistency: fresh 5-token prefill vs prefill4+decode ...") + engine2, _ = create_engine(cfg) + load_random_weights(engine2, dev, seed=42) + logits_prefill5 = run_prefill(engine2, dev, [3, 7, 9, 2, 11], slot=0) + + ref = logits_prefill5[0, -1, :].float() + got = logits_decode[0, 0, :].float() + diff = (ref - got).abs().max().item() + print(f" max |prefill5[-1] - decode| = {diff:.6e}") + assert diff < 1e-2, f"decode/prefill mismatch too large: {diff}" + print("PASS: minimax M1 smoke test") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) + + + + + + + + + diff --git a/test/models/minimax/test_lightning_attention_op.py b/test/models/minimax/test_lightning_attention_op.py new file mode 100644 index 000000000..70e41850d --- /dev/null +++ b/test/models/minimax/test_lightning_attention_op.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python3 +""" +Direct unit test for the InfiniCore `lightning_attention` op (indexed pool mode). + +Validates the op output and the in-place final state against an independent +numpy reference of the MiniMax-01 recurrence: + S = ratio[h] * S + k_t^T v_t (ratio[h] = exp(-slope[h])) + o_t = q_t @ S +""" +import ctypes +import os +import sys + +import numpy as np +import torch + +import infinicore + + +CTYPE = {np.float32: ctypes.c_float, np.int32: ctypes.c_int32} + + +def t2raw(t: torch.Tensor, dev): + # Python Tensor wrapper (the package-level op wrapper expects wrappers). + t = t.contiguous() + cpu = infinicore.device("cpu", 0) + tensor = infinicore.from_blob(t.data_ptr(), list(t.shape), + dtype={torch.float32: infinicore.float32, + torch.int32: infinicore.int32}[t.dtype], + device=cpu) + return tensor.to(dev) if dev != cpu else tensor + + +def read(t, dtype=np.float32): + if hasattr(t, "_underlying"): + t = t.to(infinicore.device("cpu", 0))._underlying + ctype = CTYPE[dtype] + buf = (ctype * int(t.numel())).from_address(t.data_ptr()) + shape = [int(t.size(i)) for i in range(int(t.ndim))] + return np.frombuffer(buf, dtype=dtype).reshape(shape).copy() + + +def reference(q, k, v, slope, pool, init_idx, final_idx): + B, T, H, D = q.shape + out = np.zeros_like(q) + pool = pool.copy() + ratio = np.exp(-slope) # [H] + for b in range(B): + S = pool[init_idx[b]].copy() # [H, D, D] + for t in range(T): + for h in range(H): + kh = k[b, t, h] # [D] + vh = v[b, t, h] # [D] + qh = q[b, t, h] # [D] + S[h] = ratio[h] * S[h] + np.outer(kh, vh) + out[b, t, h] = qh @ S[h] + pool[final_idx[b]] = S + return out, pool + + +def run_case(B, T, H, D, pool_size, seed): + torch.manual_seed(seed) + rng = np.random.default_rng(seed) + q = rng.standard_normal((B, T, H, D)).astype(np.float32) + k = rng.standard_normal((B, T, H, D)).astype(np.float32) + v = rng.standard_normal((B, T, H, D)).astype(np.float32) + slope = rng.uniform(0.0, 1.0, (H,)).astype(np.float32) + pool = rng.standard_normal((pool_size, H, D, D)).astype(np.float32) + assert pool_size >= 2 * B + init_idx = np.arange(B, dtype=np.int32) + final_idx = np.arange(pool_size - 1, pool_size - 1 - B, -1, dtype=np.int32) + + dev = infinicore.device(os.environ.get("MINIMAX_DEVICE", "cpu"), 0) + q_t = t2raw(torch.from_numpy(q), dev) + k_t = t2raw(torch.from_numpy(k), dev) + v_t = t2raw(torch.from_numpy(v), dev) + s_t = t2raw(torch.from_numpy(slope), dev) + p_t = t2raw(torch.from_numpy(pool), dev) + i_t = t2raw(torch.from_numpy(init_idx), dev) + f_t = t2raw(torch.from_numpy(final_idx), dev) + + out_w = infinicore.lightning_attention(q_t, k_t, v_t, s_t, p_t, i_t, f_t) + + out = read(out_w) + pool_after = read(p_t) + ref_out, ref_pool = reference(q, k, v, slope, pool, init_idx, final_idx) + + out_err = float(np.abs(out - ref_out).max()) + pool_err = float(np.abs(pool_after - ref_pool).max()) + print(f" B={B} T={T} H={H} D={D} pool={pool_size}: out_err={out_err:.3e} pool_err={pool_err:.3e}") + assert out_err < 1e-4, f"out mismatch {out_err}" + assert pool_err < 1e-4, f"state mismatch {pool_err}" + return True + + +def main(): + print("[1/2] decode case (T=1, batched requests, in-place state write)") + run_case(B=3, T=1, H=4, D=8, pool_size=6, seed=1) + print("[2/2] prefill case (T=6, per-request state evolution)") + run_case(B=2, T=6, H=4, D=8, pool_size=4, seed=2) + print("PASS: lightning_attention op matches reference") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) + + diff --git a/test/models/minimax/test_minimax_vs_hf.py b/test/models/minimax/test_minimax_vs_hf.py new file mode 100644 index 000000000..54f266e78 --- /dev/null +++ b/test/models/minimax/test_minimax_vs_hf.py @@ -0,0 +1,171 @@ +#!/usr/bin/env python3 +""" +M3 end-to-end comparison: InfiniLM `minimax` vs HF transformers `MiniMaxForCausalLM`. + +Uses a small config with `num_local_experts = 1`, which is mathematically a dense +SwiGLU MLP in both implementations (the fused MoE runner is CUDA-only, so the +multi-expert routing path is validated on NVIDIA in CI). + +Checks: +1. Full-prefill logits match at every position (HF chunked prefill vs our + indexed recurrent op). +2. Decode-after-context logits match (HF MiniMaxCache recurrent path vs ours). +""" +import json +import os +import sys + +import torch +from transformers import MiniMaxConfig, MiniMaxForCausalLM + +import infinicore +from infinicore.lib import _infinicore + +sys.path.insert(0, __file__.rsplit("\\", 1)[0]) +from smoke_minimax import create_engine, i2t, t2i + +from infinilm.modeling_utils import _remap_minimax + +def torch_dtype(): + dtype_name = os.environ.get("MINIMAX_TORCH_DTYPE", "float32") + return { + "float32": torch.float32, + "float16": torch.float16, + "bfloat16": torch.bfloat16, + }[dtype_name] + +def make_hf_config(): + return MiniMaxConfig( + vocab_size=128, + hidden_size=64, + intermediate_size=128, + num_hidden_layers=4, + num_attention_heads=4, + num_key_value_heads=4, + head_dim=16, + hidden_act="silu", + max_position_embeddings=64, + rms_norm_eps=1e-5, + num_experts_per_tok=int(os.environ.get("MINIMAX_NUM_EXPERTS_PER_TOK", "1")), + num_local_experts=int(os.environ.get("MINIMAX_NUM_LOCAL_EXPERTS", "1")), + attention_dropout=0.0, + block_size=16, + layer_types=["linear_attention", "full_attention", "linear_attention", "full_attention"], + rope_parameters={"rope_type": "default", "rope_theta": 1000000.0}, + torch_dtype=torch_dtype(), + ) + + +def hf_config_to_infinilm_dict(hf_config) -> dict: + d = hf_config.to_dict() + d["model_type"] = "minimax" + d["torch_dtype"] = os.environ.get("MINIMAX_TORCH_DTYPE", "float32") + d["block"] = d.pop("block_size", 16) + # Drop null fields (e.g. `torch_dtype: None`) that break ModelConfig::get_dtype. + d = {k: v for k, v in d.items() if v is not None} + return d + + +def load_hf_weights_into_engine(engine, dev, hf_model, hf_config): + hf_sd = hf_model.state_dict() + remapped = _remap_minimax(hf_sd, hf_config.to_dict()) + expected = set(engine.state_dict_keyname()) + params = {} + keep = [] + matched = 0 + for key, tensor in remapped.items(): + if key in expected: + t = tensor.detach().contiguous() + keep.append(t) + params[key] = t2i(t, dev) + matched += 1 + else: + print(f" (skip) {key}") + print(f" matched {matched}/{len(remapped)} weight keys") + engine.load_params(params, strict=False) + engine.process_weights_after_loading() + return matched + + +def run_hf_prefill(hf_model, tokens): + ids = torch.tensor([tokens], dtype=torch.long) + pos = torch.arange(len(tokens), dtype=torch.long).unsqueeze(0) + mask = torch.ones(1, len(tokens), dtype=torch.long) + out = hf_model(input_ids=ids, position_ids=pos, attention_mask=mask, use_cache=False) + return out.logits.detach().float() + + +def run_hf_decode(hf_model, tokens, next_token): + ids4 = torch.tensor([tokens], dtype=torch.long) + pos4 = torch.arange(len(tokens), dtype=torch.long).unsqueeze(0) + mask4 = torch.ones(1, len(tokens), dtype=torch.long) + out4 = hf_model(input_ids=ids4, position_ids=pos4, attention_mask=mask4, use_cache=True) + past = out4.past_key_values + ids1 = torch.tensor([[next_token]], dtype=torch.long) + pos1 = torch.tensor([[len(tokens)]], dtype=torch.long) + mask5 = torch.ones(1, len(tokens) + 1, dtype=torch.long) + out1 = hf_model(input_ids=ids1, position_ids=pos1, attention_mask=mask5, use_cache=True, past_key_values=past) + return out1.logits.detach().float() + + +def run_il_prefill(engine, dev, tokens): + n = len(tokens) + input_ids = torch.tensor([tokens], dtype=torch.int32) + position_ids = torch.arange(n, dtype=torch.int32).unsqueeze(0) + past = torch.tensor([0], dtype=torch.int32) + total = torch.tensor([n], dtype=torch.int32) + offsets = torch.tensor([0, n], dtype=torch.int32) + cu = torch.tensor([0, n], dtype=torch.int32) + idx = torch.tensor([0], dtype=torch.int32) + from smoke_minimax import make_input + out = engine.forward(make_input(dev, input_ids, position_ids, past, total, offsets, cu, idx, idx)) + return i2t(out.logits) + + +def run_il_decode(engine, dev, tokens, next_token): + from smoke_minimax import run_decode + return run_decode(engine, dev, next_token, past_len=len(tokens), slot=0) + + +def main(): + hf_config = make_hf_config() + torch.manual_seed(7) + hf_model = MiniMaxForCausalLM(hf_config).to(torch_dtype()) + hf_model.eval() + + cfg_dict = hf_config_to_infinilm_dict(hf_config) + print("[1/4] constructing InfiniLM minimax engine (dense 1-expert path) ...") + engine, dev = create_engine(cfg_dict) + + print("[2/4] loading HF weights via _remap_minimax ...") + load_hf_weights_into_engine(engine, dev, hf_model, hf_config) + + tokens = [3, 7, 9, 2, 11] + print("[3/4] comparing full-prefill logits ...") + hf_prefill = run_hf_prefill(hf_model, tokens) + il_prefill = run_il_prefill(engine, dev, tokens) + print(" hf_prefill:", tuple(hf_prefill.shape), " il_prefill:", tuple(il_prefill.shape)) + err = (hf_prefill - il_prefill).abs().max().item() + print(f" max |hf - infinilm| (prefill, all positions) = {err:.6e}") + assert err < 1e-2, f"prefill mismatch {err}" + + print("[4/4] comparing decode-after-context logits ...") + # Use a fresh engine: prefill 4 tokens, then decode the 5th. + engine2, dev2 = create_engine(cfg_dict) + load_hf_weights_into_engine(engine2, dev2, hf_model, hf_config) + run_il_prefill(engine2, dev2, tokens[:-1]) + hf_dec = run_hf_decode(hf_model, tokens[:-1], tokens[-1]) + il_dec = run_il_decode(engine2, dev2, tokens[:-1], tokens[-1]) + err2 = (hf_dec - il_dec).abs().max().item() + print(f" max |hf - infinilm| (decode) = {err2:.6e}") + assert err2 < 1e-2, f"decode mismatch {err2}" + + print("PASS: minimax end-to-end matches HF transformers") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) + + +