issue/294: minicpm-sala model

.gitignore

@@ -30,4 +30,7 @@ __pycache__/
 
 *.http
 
-*.nsys-rep
+**/*.nsys-rep
+**/*.jsonl
+*.jsonl
+**/*.mem

csrc/models/minicpm_sala/minicpm_sala_allocate_kv_cache_tensors.cpp

@@ -32,7 +32,7 @@ std::vector<infinicore::Tensor> minicpm_sala_allocate_kv_cache_tensors(const cac
         const size_t num_key_value_heads = text_config->get<size_t>("num_key_value_heads");
         const size_t max_position_embeddings = text_config->get<size_t>("max_position_embeddings");
 
-        const auto &dtype{text_config->get_dtype()};
+        const auto &dtype{text_config->get_kv_cache_dtype()};
         std::vector<std::string> mixer_types = text_config->get<std::vector<std::string>>("mixer_types");
         size_t current_layer_head_dim, current_layer_num_key_value_heads;
         for (size_t layer_idx = 0; layer_idx < num_hidden_layers; ++layer_idx) {
@@ -70,7 +70,7 @@ std::vector<infinicore::Tensor> minicpm_sala_allocate_kv_cache_tensors(const cac
 
         const size_t head_dim = text_config->get<size_t>("head_dim");
         const size_t num_key_value_heads = text_config->get<size_t>("num_key_value_heads");
-        const auto &dtype{text_config->get_dtype()};
+        const auto &dtype{text_config->get_kv_cache_dtype()};
         std::vector<std::string> mixer_types = text_config->get<std::vector<std::string>>("mixer_types");
         size_t current_layer_head_dim, current_layer_num_key_value_heads;
         for (size_t layer_idx = 0; layer_idx < num_hidden_layers; ++layer_idx) {

csrc/models/minicpm_sala/minicpm_sala_attention.hpp

@@ -1,88 +1,108 @@
 #pragma once
 
-#include "../../layers/common_modules.hpp"
+#include "../../config/model_config.hpp"
+#include "../../layers/rotary_embedding/rotary_embedding.hpp"
 
-namespace infinilm::layers::attention {
-class AttentionLayer;
-}
+#include "infinicore/nn/linear.hpp"
+#include "infinicore/nn/module.hpp"
+#include "infinicore/nn/rmsnorm.hpp"
+#include "infinicore/nn/rope.hpp"
+#include "infinicore/tensor.hpp"
+
+#include <memory>
+#include <string>
 
 namespace infinilm::models::minicpm_sala {
 
-class AttentionBase : public infinicore::nn::Module {
-protected:
-    AttentionBase(std::shared_ptr<infinilm::config::ModelConfig> model_config,
-                  size_t num_attention_heads,
-                  size_t num_key_value_heads,
-                  size_t layer_idx,
-                  const infinicore::Device &device);
+class MiniCPMSALAAttentionBase : public infinicore::nn::Module {
+public:
+    virtual infinicore::Tensor forward(const infinicore::Tensor &position_ids,
+                                       const infinicore::Tensor &hidden_states) const = 0;
+    virtual void reset_state() = 0;
+    virtual ~MiniCPMSALAAttentionBase() = default;
+};
 
+// Lightning attention path (Simple GLA). Parameter names align with HF:
+//   model.layers.N.self_attn.{q_proj,k_proj,v_proj,o_proj,q_norm,k_norm,o_norm,z_proj,...}
+class MiniCPMSALALightningAttention : public MiniCPMSALAAttentionBase {
 public:
-    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_; }
+    MiniCPMSALALightningAttention(std::shared_ptr<infinilm::config::ModelConfig> model_config,
+                                  const infinicore::Device &device,
+                                  size_t layer_idx);
+
+    // Match `infinilm::layers::attention::Attention` API: metadata is pulled from
+    // `global_state::get_forward_context().attn_metadata`.
+    infinicore::Tensor forward(const infinicore::Tensor &position_ids,
+                               const infinicore::Tensor &hidden_states) const override;
+
+    void reset_state() override;
 
 protected:
-    INFINICORE_NN_MODULE(infinilm::layers::linear::ColumnParallelLinear, q_proj);
-    INFINICORE_NN_MODULE(infinilm::layers::linear::ColumnParallelLinear, k_proj);
-    INFINICORE_NN_MODULE(infinilm::layers::linear::ColumnParallelLinear, v_proj);
-    INFINICORE_NN_MODULE(infinilm::layers::linear::RowParallelLinear, o_proj);
+    // Projections (HF-aligned naming)
+    INFINICORE_NN_MODULE(infinicore::nn::Linear, q_proj);
+    INFINICORE_NN_MODULE(infinicore::nn::Linear, k_proj);
+    INFINICORE_NN_MODULE(infinicore::nn::Linear, v_proj);
+    INFINICORE_NN_MODULE(infinicore::nn::Linear, o_proj);
+
+    // Optional (Lightning layers): q_norm/k_norm/o_norm + z_proj
+    INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, q_norm);
+    INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, k_norm);
+    INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, o_norm);
+    INFINICORE_NN_MODULE(infinicore::nn::Linear, z_proj);
 
-    std::shared_ptr<infinilm::layers::attention::AttentionLayer> attn_;
-    ::infinilm::backends::AttentionBackend attention_backend_;
     std::shared_ptr<infinicore::nn::RoPE> rotary_emb_;
 
     size_t layer_idx_;
-    size_t hidden_size_;
     size_t num_attention_heads_;
     size_t num_key_value_heads_;
     size_t head_dim_;
-    bool use_bias_;
-    bool use_output_bias_;
-
-    // For off-line kv cache quantization
-    INFINICORE_NN_PARAMETER(kv_cache_k_scale);
-    INFINICORE_NN_PARAMETER(kv_cache_v_scale);
-};
+    float scaling_;
 
-/**
- * @brief InfLLMv2 attention with optional output gate
- */
-class InfLLMv2Attention : public AttentionBase {
-public:
-    InfLLMv2Attention(std::shared_ptr<infinilm::config::ModelConfig> model_config,
-                      size_t layer_idx,
-                      const infinicore::Device &device);
+    bool use_qk_norm_ = false;
+    bool use_output_gate_ = false;
+    bool use_output_norm_ = false;
+    bool use_rope_ = false;
 
-    infinicore::Tensor forward(const infinicore::Tensor &positions,
-                               const infinicore::Tensor &hidden_states) const;
+    // Lightning layers only: per-head log-decay for Simple GLA (HF _build_slope_tensor * -1).
+    infinicore::Tensor g_gamma_;
 
-protected:
-    bool use_output_gate_;
-    INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, o_gate);
+    // Lightning layers only: recurrent state for fast decode.
+    // Shape: [B, H, D, D] float32. Tracks how many KV tokens are folded into the state.
+    mutable infinicore::Tensor gla_state_;
+    mutable size_t gla_state_cached_len_ = 0;
+    mutable bool gla_state_valid_ = false;
 };
 
-/**
- * @brief Lightning attention with optional output norm and gate
- */
-class LightningAttention : public AttentionBase {
+// Sparse attention path (`mixer_type=="minicpm4"`) using InfLLM-v2 operators.
+// Parameter names align with HF:
+//   model.layers.N.self_attn.{q_proj,k_proj,v_proj,o_proj,o_gate,...}
+class MiniCPMSALAMinicpm4Attention : public MiniCPMSALAAttentionBase {
 public:
-    LightningAttention(std::shared_ptr<infinilm::config::ModelConfig> model_config,
-                       size_t layer_idx,
-                       const infinicore::Device &device);
+    MiniCPMSALAMinicpm4Attention(std::shared_ptr<infinilm::config::ModelConfig> model_config,
+                                 const infinicore::Device &device,
+                                 size_t layer_idx);
+
+    infinicore::Tensor forward(const infinicore::Tensor &position_ids,
+                               const infinicore::Tensor &hidden_states) const override;
 
-    infinicore::Tensor forward(const infinicore::Tensor &positions,
-                               const infinicore::Tensor &hidden_states) const;
+    void reset_state() override;
 
 protected:
-    bool qk_norm_;
-    bool use_output_norm_;
-    bool use_output_gate_;
-    INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, q_norm);
-    INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, k_norm);
-    INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, o_norm);
-    INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, z_proj);
+    INFINICORE_NN_MODULE(infinicore::nn::Linear, q_proj);
+    INFINICORE_NN_MODULE(infinicore::nn::Linear, k_proj);
+    INFINICORE_NN_MODULE(infinicore::nn::Linear, v_proj);
+    INFINICORE_NN_MODULE(infinicore::nn::Linear, o_proj);
+    INFINICORE_NN_MODULE(infinicore::nn::Linear, o_gate);
+
+    size_t layer_idx_;
+    size_t num_attention_heads_;
+    size_t num_key_value_heads_;
+    size_t head_dim_;
+    float scaling_;
+
+    // InfLLM-v2 local-window masking plumbing.
+    int infllmv2_window_left_ = -1;
+    bool use_local_window_ = false;
 };
 
 } // namespace infinilm::models::minicpm_sala

csrc/models/minicpm_sala/minicpm_sala_decoder_layer.cpp

@@ -0,0 +1,66 @@
+#include "minicpm_sala_decoder_layer.hpp"
+
+#include "infinicore/ops.hpp"
+#include "infinicore/context/context.hpp"
+#include <cmath>
+#include <cstdio>
+#include <chrono>
+#include <cstdlib>
+#include <fstream>
+#include <vector>
+
+namespace infinilm::models::minicpm_sala {
+
+
+MiniCPMSALADecoderLayer::MiniCPMSALADecoderLayer(std::shared_ptr<infinilm::config::ModelConfig> model_config,
+                                                 const infinicore::Device &device,
+                                                 size_t layer_idx,
+                                                 const std::string &mixer_type) {
+    // Match parameter dtype with checkpoint `torch_dtype` (e.g. BF16 for MiniCPM-SALA).
+    const auto dtype = model_config->get_dtype();
+    const double eps = model_config->get<double>("rms_norm_eps");
+
+    // MuP residual scaling at forward (o_proj/down_proj not scaled in loader for minicpm_sala).
+    const double scale_depth = model_config->get_or<double>("scale_depth", 1.0);
+    const size_t num_layers = model_config->get<size_t>("num_hidden_layers");
+    residual_scale_ = scale_depth / std::sqrt(static_cast<double>(num_layers));
+
+    INFINICORE_NN_MODULE_INIT(input_layernorm, model_config->get<size_t>("hidden_size"), eps, dtype, device);
+    if (mixer_type == "minicpm4") {
+        self_attn_ = this->register_module<MiniCPMSALAMinicpm4Attention>(
+            "self_attn", model_config, device, layer_idx);
+    } else {
+        self_attn_ = this->register_module<MiniCPMSALALightningAttention>(
+            "self_attn", model_config, device, layer_idx);
+    }
+    INFINICORE_NN_MODULE_INIT(post_attention_layernorm, model_config->get<size_t>("hidden_size"), eps, dtype, device);
+    INFINICORE_NN_MODULE_INIT(mlp, model_config, device);
+}
+
+void MiniCPMSALADecoderLayer::reset_attn_state() {
+    self_attn_->reset_state();
+}
+
+infinicore::Tensor MiniCPMSALADecoderLayer::forward(const infinicore::Tensor &hidden_states,
+                                                    const infinicore::Tensor &position_ids) const {
+    // Pre-norm attention
+    auto hs1 = input_layernorm_->forward(hidden_states);
+    auto attn_out = self_attn_->forward(position_ids, hs1);
+
+    // residual + scale_down * attn_out (MuP)
+    auto ones_attn = infinicore::Tensor::empty(attn_out->shape(), attn_out->dtype(), attn_out->device());
+    infinicore::op::ones_(ones_attn);
+    auto out1 = infinicore::op::addcmul(hidden_states, attn_out, ones_attn, static_cast<float>(residual_scale_));
+
+    // Pre-norm MLP
+    auto hs2 = post_attention_layernorm_->forward(out1);
+    auto mlp_out = mlp_->forward(hs2);
+    // residual + scale_down * mlp_out (MuP)
+    auto ones_mlp = infinicore::Tensor::empty(mlp_out->shape(), mlp_out->dtype(), mlp_out->device());
+    infinicore::op::ones_(ones_mlp);
+    auto out2 = infinicore::op::addcmul(out1, mlp_out, ones_mlp, static_cast<float>(residual_scale_));
+
+    return out2;
+}
+
+} // namespace infinilm::models::minicpm_sala

csrc/models/minicpm_sala/minicpm_sala_decoder_layer.hpp

@@ -0,0 +1,40 @@
+#pragma once
+
+#include "minicpm_sala_attention.hpp"
+#include "minicpm_sala_mlp.hpp"
+
+#include "../../config/model_config.hpp"
+
+#include "infinicore/nn/module.hpp"
+#include "infinicore/nn/rmsnorm.hpp"
+#include "infinicore/tensor.hpp"
+
+#include <memory>
+#include <string>
+
+namespace infinilm::models::minicpm_sala {
+
+class MiniCPMSALADecoderLayer : public infinicore::nn::Module {
+public:
+    MiniCPMSALADecoderLayer(std::shared_ptr<infinilm::config::ModelConfig> model_config,
+                            const infinicore::Device &device,
+                            size_t layer_idx,
+                            const std::string &mixer_type);
+
+    infinicore::Tensor forward(const infinicore::Tensor &hidden_states,
+                               const infinicore::Tensor &position_ids) const;
+
+    void reset_attn_state();
+
+private:
+    double residual_scale_ = 1.0;
+
+protected:
+    INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, input_layernorm);
+    // Registered under the HF-compatible name "self_attn" in ctor.
+    std::shared_ptr<MiniCPMSALAAttentionBase> self_attn_;
+    INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, post_attention_layernorm);
+    INFINICORE_NN_MODULE(MiniCPMSALAMLP, mlp);
+};
+
+} // namespace infinilm::models::minicpm_sala