Feat/add gemma3 model

csrc/layers/attention/backends/attention_layer.cpp

@@ -9,16 +9,29 @@ AttentionLayer::AttentionLayer(size_t num_heads,
                                size_t layer_idx,
                                infinicore::Tensor k_scale,
                                infinicore::Tensor v_scale,
-                               ::infinilm::backends::AttentionBackend attn_backend) : k_scale_(k_scale), v_scale_(v_scale), layer_idx_(layer_idx), attn_backend_(attn_backend) {
+                               ::infinilm::backends::AttentionBackend attn_backend,
+                               float softcap) : k_scale_(k_scale), v_scale_(v_scale), layer_idx_(layer_idx), attn_backend_(attn_backend) {
+    if (softcap < 0.0f) {
+        // A negative cap would silently fall through the `> 0` enable checks
+        // and disable the feature; it is never a meaningful configuration.
+        throw std::runtime_error("infinilm::layers::attention::AttentionLayer: softcap must be non-negative");
+    }
     switch (attn_backend) {
     case ::infinilm::backends::AttentionBackend::STATIC_ATTN:
-        attn_backend_impl_ = std::make_shared<backends::StaticAttentionImpl>(num_heads, head_size, scale, num_kv_heads, layer_idx);
+        attn_backend_impl_ = std::make_shared<backends::StaticAttentionImpl>(num_heads, head_size, scale, num_kv_heads, layer_idx, softcap);
         break;
     case ::infinilm::backends::AttentionBackend::PAGED_ATTN:
-        attn_backend_impl_ = std::make_shared<backends::PagedAttentionImpl>(num_heads, head_size, scale, num_kv_heads, layer_idx);
-        break;
     case ::infinilm::backends::AttentionBackend::FLASH_ATTN:
-        attn_backend_impl_ = std::make_shared<backends::FlashAttentionImpl>(num_heads, head_size, scale, num_kv_heads, layer_idx);
+        if (softcap > 0.0f) {
+            // Soft-capping is only wired through the static backend; keep other
+            // backends from silently ignoring it.
+            throw std::runtime_error("infinilm::layers::attention::AttentionLayer: softcap requires the STATIC_ATTN backend");
+        }
+        if (attn_backend == ::infinilm::backends::AttentionBackend::PAGED_ATTN) {
+            attn_backend_impl_ = std::make_shared<backends::PagedAttentionImpl>(num_heads, head_size, scale, num_kv_heads, layer_idx);
+        } else {
+            attn_backend_impl_ = std::make_shared<backends::FlashAttentionImpl>(num_heads, head_size, scale, num_kv_heads, layer_idx);
+        }
         break;
     default:
         throw std::runtime_error("infinilm::layers::attention::AttentionLayer: unsupported attention backend");

csrc/layers/attention/backends/attention_layer.hpp

@@ -31,7 +31,8 @@ class AttentionLayer {
                    size_t layer_idx,
                    infinicore::Tensor k_scale,
                    infinicore::Tensor v_scale,
-                   ::infinilm::backends::AttentionBackend attention_backend);
+                   ::infinilm::backends::AttentionBackend attention_backend,
+                   float softcap = 0.0f);
 
     infinicore::Tensor forward(infinicore::Tensor &query,
                                infinicore::Tensor &key,

csrc/layers/attention/backends/static_attn.cpp

@@ -2,6 +2,7 @@
 #include "../../../utils.hpp"
 #include "attention_layer.hpp"
 #include "infinicore/ops.hpp"
+#include "infinicore/ops/mul_scalar.hpp"
 #include "infinicore/ops/per_tensor_dequant_i8.hpp"
 #include "infinicore/ops/per_tensor_quant_i8.hpp"
 
@@ -11,13 +12,15 @@ StaticAttentionImpl::StaticAttentionImpl(size_t num_heads,
                                          size_t head_size,
                                          float scale,
                                          size_t num_kv_heads,
-                                         size_t layer_idx)
+                                         size_t layer_idx,
+                                         float softcap)
     : num_heads_(num_heads),
       head_size_(head_size),
       scale_(scale),
       num_kv_heads_(num_kv_heads),
       layer_idx_(layer_idx),
-      head_dim_(head_size) {
+      head_dim_(head_size),
+      softcap_(softcap) {
     kv_quant_scheme_ = infinilm::global_state::get_infinilm_config().model_config->get_kv_quant_scheme();
 }
 
@@ -88,6 +91,13 @@ infinicore::Tensor StaticAttentionImpl::forward(const AttentionLayer &layer,
 
         auto attn_weight = infinicore::op::matmul(Q, K_transposed, scale_); // [bs * n_kv_head, ng * seq_len, total_seq_len]
 
+        if (softcap_ > 0.0f) {
+            // Attention logit soft-capping (e.g. Gemma-2): squash the scaled
+            // scores with tanh before the causal mask and softmax.
+            attn_weight = infinicore::op::tanh(infinicore::op::mul_scalar(attn_weight, 1.0f / softcap_));
+            attn_weight = infinicore::op::mul_scalar(attn_weight, softcap_);
+        }
+
         auto attn_weight_softmax = attn_weight->view({batch_size * num_heads_, seq_len, total_seq_len});
         infinicore::op::causal_softmax_(attn_weight_softmax, attn_weight_softmax);
 

csrc/layers/attention/backends/static_attn.hpp

@@ -18,7 +18,8 @@ class StaticAttentionImpl {
                         size_t head_size,
                         float scale,
                         size_t num_kv_heads,
-                        size_t layer_idx);
+                        size_t layer_idx,
+                        float softcap = 0.0f);
 
     infinicore::Tensor forward(const AttentionLayer &layer,
                                infinicore::Tensor &q_reshaped, // query
@@ -40,6 +41,9 @@ class StaticAttentionImpl {
     size_t num_kv_heads_;
     size_t layer_idx_;
     size_t head_dim_; // Note: head_dim equals to head_size
+    // Attention logit soft-capping (e.g. Gemma-2): scores are squashed with
+    // tanh before the causal mask/softmax. 0 disables the feature.
+    float softcap_;
 
     infinilm::quantization::KVQuantAlgo kv_quant_scheme_;
 };

csrc/models/gemma2/gemma2_attention.cpp

@@ -0,0 +1,145 @@
+#include "gemma2_attention.hpp"
+#include "../../utils.hpp"
+
+namespace infinilm::models::gemma2 {
+
+Gemma2Attention::Gemma2Attention(std::shared_ptr<infinilm::config::ModelConfig> model_config,
+                                 size_t layer_idx,
+                                 const infinicore::Device &device) {
+    layer_idx_ = layer_idx;
+    hidden_size_ = model_config->get<size_t>("hidden_size");
+    head_dim_ = model_config->get<size_t>("head_dim");
+
+    const auto &dtype{model_config->get_dtype()};
+    size_t total_num_heads = model_config->get<size_t>("num_attention_heads");
+    size_t total_num_kv_heads = model_config->get<size_t>("num_key_value_heads");
+    bool use_bias = model_config->get_or<bool>("attention_bias", false);
+    bool use_output_bias = model_config->get_or<bool>("attention_output_bias", false);
+
+    attention_backend_ = infinilm::global_state::get_infinilm_config().attention_backend;
+    const engine::distributed::RankInfo &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info();
+    int tp_rank = infinilm::global_state::get_tensor_model_parallel_rank();
+    int tp_size = infinilm::global_state::get_tensor_model_parallel_world_size();
+
+    num_attention_heads_ = total_num_heads / tp_size;
+    num_key_value_heads_ = total_num_kv_heads < tp_size ? 1 : total_num_kv_heads / tp_size;
+
+    auto quantization_method = model_config->get_quantization_method();
+    auto register_fn = [this](const std::string &n, infinicore::nn::Parameter p) { this->register_parameter(n, std::move(p)); };
+    qkv_proj_ = std::make_shared<layers::linear::QKVParallelLinear>(
+        hidden_size_, head_dim_, total_num_heads, total_num_kv_heads,
+        "q_proj", "k_proj", "v_proj", register_fn,
+        quantization_method, use_bias, dtype, device, rank_info);
+    o_proj_ = this->register_module<layers::linear::RowParallelLinear>(
+        "o_proj", total_num_heads * head_dim_, hidden_size_, quantization_method,
+        use_output_bias, dtype, device, tp_rank, tp_size, rank_info.comm);
+
+    rotary_emb_ = infinilm::layers::rotary_embedding::get_rope(model_config, device);
+
+    // Gemma-2 scales the scores by 1 / query_pre_attn_scalar instead of
+    // 1 / sqrt(head_dim), and squashes them with tanh soft-capping.
+    float query_pre_attn_scalar = model_config->get_or<float>("query_pre_attn_scalar", static_cast<float>(head_dim_));
+    float scaling = 1.0f / std::sqrt(query_pre_attn_scalar);
+    softcap_ = model_config->get_or<float>("attn_logit_softcapping", 0.0f);
+    attn_ = std::make_shared<infinilm::layers::attention::AttentionLayer>(
+        num_attention_heads_, head_dim_, scaling, num_key_value_heads_, layer_idx_,
+        kv_cache_k_scale_, kv_cache_v_scale_, attention_backend_, softcap_);
+
+    infinilm::layers::attention::init_kv_cache_quant_params(register_fn, device, kv_cache_k_scale_, kv_cache_v_scale_);
+}
+
+infinicore::Tensor Gemma2Attention::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 Gemma2Attention::forward_static_(const infinicore::Tensor &position_ids,
+                                                    const infinicore::Tensor &hidden_states) const {
+    // hidden_states shape: [batch, seq_len, hidden_size]
+    auto hidden_states_mutable = hidden_states;
+    auto shape = hidden_states->shape();
+    size_t batch_size = shape[0];
+    size_t seq_len = shape[1];
+
+    // 1. Project Q, K, V
+    auto [q, k, v] = qkv_proj_->forward_split(hidden_states_mutable);
+
+    // 2. Reshape for multi-head attention
+    auto q_reshaped = q->view({batch_size, seq_len, num_attention_heads_, head_dim_});
+    auto k_reshaped = k->view({batch_size, seq_len, num_key_value_heads_, head_dim_});
+    auto v_reshaped = v->view({batch_size, seq_len, num_key_value_heads_, head_dim_});
+
+    // 3. Prepare position_ids for RoPE
+    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::gemma2::Gemma2Attention: Unexpected position_ids shape");
+    }
+
+    // 4. Apply RoPE to QK
+    rotary_emb_->forward(q_reshaped, pos_ids_for_rope, true);
+    rotary_emb_->forward(k_reshaped, pos_ids_for_rope, true);
+
+    // 5. Attn backend calculate (soft-capping is applied inside the backend)
+    auto attn_output = attn_->forward(q_reshaped, k_reshaped, v_reshaped);
+
+    // 6. Project output
+    auto output = o_proj_->forward(attn_output);
+    return output;
+}
+
+// Note: forward_paged_ exists for backend parity with the shared Attention
+// interface. With attn_logit_softcapping > 0 the AttentionLayer constructor
+// restricts the model to STATIC_ATTN, so under soft-capping this path is
+// unreachable; it remains valid for softcap-free gemma2 checkpoints.
+infinicore::Tensor Gemma2Attention::forward_paged_(const infinicore::Tensor &position_ids,
+                                                   const infinicore::Tensor &hidden_states) const {
+    // hidden_states shape: [batch, seq_len, hidden_size]
+    auto hidden_states_mutable = hidden_states;
+    auto shape = hidden_states->shape();
+    size_t seq_len = shape[1];
+
+    // Only support batchsize==1, all requests should be flattened along seqlen dimension
+    ASSERT_EQ(shape[0], 1);
+
+    // 1. Project Q, K, V
+    auto [q, k, v] = qkv_proj_->forward_split(hidden_states_mutable);
+
+    // 2. Reshape for multi-head attention
+    auto q_reshaped = q->view({seq_len, num_attention_heads_, head_dim_});
+    auto k_reshaped = k->view({seq_len, num_key_value_heads_, head_dim_});
+    auto v_reshaped = v->view({seq_len, num_key_value_heads_, head_dim_});
+
+    // 3. Prepare position_ids for RoPE
+    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::gemma2::Gemma2Attention: Unexpected position_ids shape");
+    }
+
+    // 4. Apply RoPE to QK
+    rotary_emb_->forward(q_reshaped, pos_ids_for_rope, true);
+    rotary_emb_->forward(k_reshaped, pos_ids_for_rope, true);
+
+    // 5. Attn backend calculate
+    auto attn_output = attn_->forward(q_reshaped, k_reshaped, v_reshaped);
+
+    // 6. Project output
+    auto output = o_proj_->forward(attn_output);
+    return output;
+}
+
+} // namespace infinilm::models::gemma2

csrc/models/gemma2/gemma2_attention.hpp

@@ -0,0 +1,56 @@
+#pragma once
+
+#include "../../layers/common_modules.hpp"
+
+namespace infinilm::models::gemma2 {
+
+class Gemma2Attention : public infinicore::nn::Module {
+public:
+    Gemma2Attention(std::shared_ptr<infinilm::config::ModelConfig> model_config,
+                    size_t layer_idx,
+                    const infinicore::Device &device);
+
+    infinicore::Tensor forward(const infinicore::Tensor &positions,
+                               const infinicore::Tensor &hidden_states) const;
+
+    void process_weights_after_loading() override {
+        qkv_proj_->process_weights_after_loading();
+    }
+
+    void reset_runtime_state() const override {
+        qkv_proj_->reset_runtime_state();
+    }
+
+    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 &positions,
+                                       const infinicore::Tensor &hidden_states) const;
+
+    infinicore::Tensor forward_paged_(const infinicore::Tensor &positions,
+                                      const infinicore::Tensor &hidden_states) const;
+
+protected:
+    std::shared_ptr<infinilm::layers::linear::QKVParallelLinear> qkv_proj_;
+    std::shared_ptr<infinilm::layers::linear::RowParallelLinear> o_proj_;
+    std::shared_ptr<infinicore::nn::RoPE> rotary_emb_;
+
+    std::shared_ptr<infinilm::layers::attention::AttentionLayer> attn_;
+    ::infinilm::backends::AttentionBackend attention_backend_;
+    size_t layer_idx_;
+    size_t num_attention_heads_;
+    size_t num_key_value_heads_;
+    size_t hidden_size_;
+    size_t head_dim_;
+    float softcap_;
+
+    // For off-line kv cache quantization
+    INFINICORE_NN_PARAMETER(kv_cache_k_scale);
+    INFINICORE_NN_PARAMETER(kv_cache_v_scale);
+};
+
+} // namespace infinilm::models::gemma2

csrc/models/gemma2/gemma2_decoder_layer.cpp

@@ -0,0 +1,69 @@
+#include "gemma2_decoder_layer.hpp"
+
+namespace infinilm::models::gemma2 {
+
+Gemma2DecoderLayer::Gemma2DecoderLayer(std::shared_ptr<infinilm::config::ModelConfig> model_config,
+                                       size_t layer_idx,
+                                       const infinicore::Device &device)
+    : layer_idx_(layer_idx),
+      rms_norm_eps_(model_config->get<double>("rms_norm_eps")) {
+    const auto &dtype{model_config->get_dtype()};
+    size_t hidden_size = model_config->get<size_t>("hidden_size");
+    double rms_norm_eps = model_config->get<double>("rms_norm_eps");
+
+    input_layernorm_ = this->register_module<infinicore::nn::RMSNorm>("input_layernorm", hidden_size, rms_norm_eps, dtype, device);
+    post_attention_layernorm_ = this->register_module<infinicore::nn::RMSNorm>("post_attention_layernorm", hidden_size, rms_norm_eps, dtype, device);
+    pre_feedforward_layernorm_ = this->register_module<infinicore::nn::RMSNorm>("pre_feedforward_layernorm", hidden_size, rms_norm_eps, dtype, device);
+    post_feedforward_layernorm_ = this->register_module<infinicore::nn::RMSNorm>("post_feedforward_layernorm", hidden_size, rms_norm_eps, dtype, device);
+    self_attn_ = this->register_module<Gemma2Attention>("self_attn", model_config, layer_idx, device);
+    mlp_ = this->register_module<Gemma2MLP>("mlp", model_config, device);
+}
+
+std::tuple<infinicore::Tensor, infinicore::Tensor> Gemma2DecoderLayer::forward(const infinicore::Tensor &positions,
+                                                                               infinicore::Tensor &hidden_states,
+                                                                               infinicore::Tensor &residual) {
+    // 1. Normalize the mainstream (fused: residual += incoming branch, hidden = norm(residual)).
+    //    This matches Gemma-2's `residual = hidden; hidden = input_layernorm(hidden)`.
+    input_layernorm_->forward_inplace(hidden_states, residual);
+
+    // 2. Attention on the normalized branch.
+    hidden_states = self_attn_->forward(positions, hidden_states);
+
+    // 3. Gemma-2 order: normalize the branch FIRST, then add it to the residual stream.
+    hidden_states = post_attention_layernorm_->forward(hidden_states);
+
+    // 4. Fuse the branch addition with the pre-feedforward norm:
+    //    add_rms_norm(residual, branch, w) returns (norm(residual+branch, w),
+    //    residual+branch), replacing a separate add + norm pair.
+    auto fused = infinicore::op::add_rms_norm(residual, hidden_states,
+                                              pre_feedforward_layernorm_->weight(),
+                                              static_cast<float>(rms_norm_eps_));
+    residual = std::move(fused.second);
+    hidden_states = std::move(fused.first);
+    hidden_states = mlp_->forward(hidden_states);
+    hidden_states = post_feedforward_layernorm_->forward(hidden_states);
+
+    // 5. Contract: leave the branch un-added; the consumer (next layer's input
+    //    norm or the model's final norm) performs `residual + hidden`.
+    return std::make_tuple(hidden_states, residual);
+}
+
+infinicore::Tensor Gemma2DecoderLayer::forward(const infinicore::Tensor &positions,
+                                               infinicore::Tensor &hidden_states) {
+    // Naive (debug) path mirroring the HF reference exactly.
+    infinicore::Tensor residual = hidden_states;
+
+    hidden_states = input_layernorm_->forward(hidden_states);
+    hidden_states = self_attn_->forward(positions, hidden_states);
+    hidden_states = post_attention_layernorm_->forward(hidden_states);
+    hidden_states = infinicore::op::add(residual, hidden_states);
+
+    residual = hidden_states;
+    hidden_states = pre_feedforward_layernorm_->forward(hidden_states);
+    hidden_states = mlp_->forward(hidden_states);
+    hidden_states = post_feedforward_layernorm_->forward(hidden_states);
+    hidden_states = infinicore::op::add(residual, hidden_states);
+    return hidden_states;
+}
+
+} // namespace infinilm::models::gemma2