diff --git a/csrc/layers/attention/backends/attention_layer.cpp b/csrc/layers/attention/backends/attention_layer.cpp
index fcaefa292..609fac6a7 100644
--- a/csrc/layers/attention/backends/attention_layer.cpp
+++ b/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");
diff --git a/csrc/layers/attention/backends/attention_layer.hpp b/csrc/layers/attention/backends/attention_layer.hpp
index 874110629..eaa447d83 100644
--- a/csrc/layers/attention/backends/attention_layer.hpp
+++ b/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,
diff --git a/csrc/layers/attention/backends/static_attn.cpp b/csrc/layers/attention/backends/static_attn.cpp
index e95bc03b2..cf55b5abb 100644
--- a/csrc/layers/attention/backends/static_attn.cpp
+++ b/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);
 
diff --git a/csrc/layers/attention/backends/static_attn.hpp b/csrc/layers/attention/backends/static_attn.hpp
index 849d87928..5da73ff5e 100644
--- a/csrc/layers/attention/backends/static_attn.hpp
+++ b/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_;
 };
diff --git a/csrc/models/gemma2/gemma2_attention.cpp b/csrc/models/gemma2/gemma2_attention.cpp
new file mode 100644
index 000000000..e02069a62
--- /dev/null
+++ b/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
diff --git a/csrc/models/gemma2/gemma2_attention.hpp b/csrc/models/gemma2/gemma2_attention.hpp
new file mode 100644
index 000000000..9d3245b8b
--- /dev/null
+++ b/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
diff --git a/csrc/models/gemma2/gemma2_decoder_layer.cpp b/csrc/models/gemma2/gemma2_decoder_layer.cpp
new file mode 100644
index 000000000..d4d80cf22
--- /dev/null
+++ b/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
diff --git a/csrc/models/gemma2/gemma2_decoder_layer.hpp b/csrc/models/gemma2/gemma2_decoder_layer.hpp
new file mode 100644
index 000000000..163cba56c
--- /dev/null
+++ b/csrc/models/gemma2/gemma2_decoder_layer.hpp
@@ -0,0 +1,49 @@
+#pragma once
+
+#include "../../layers/common_modules.hpp"
+#include "gemma2_attention.hpp"
+#include "gemma2_mlp.hpp"
+#include "infinicore/nn/rmsnorm.hpp"
+#include <memory>
+
+namespace infinilm::models::gemma2 {
+
+/**
+ * @brief Gemma-2 decoder layer.
+ *
+ * Unlike the llama-style TextDecoderLayer, Gemma-2 applies four RMSNorms and
+ * normalizes each branch *before* adding it back to the residual stream
+ * (`hidden = residual + post_norm(branch)`), so this layer cannot reuse the
+ * fused llama template. The TextModel residual contract is still honored:
+ * on entry `hidden` is the previous branch output (not yet added) and
+ * `residual` is the mainstream; on exit the same holds, deferring the final
+ * addition to the consumer (next layer's input norm or the model's final norm).
+ */
+class Gemma2DecoderLayer : public infinicore::nn::Module {
+public:
+    Gemma2DecoderLayer(std::shared_ptr<infinilm::config::ModelConfig> model_config,
+                       size_t layer_idx,
+                       const infinicore::Device &device);
+
+    std::tuple<infinicore::Tensor, infinicore::Tensor> forward(const infinicore::Tensor &positions,
+                                                               infinicore::Tensor &hidden_states,
+                                                               infinicore::Tensor &residual);
+
+    infinicore::Tensor forward(const infinicore::Tensor &positions,
+                               infinicore::Tensor &hidden_states);
+
+    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(infinicore::nn::RMSNorm, pre_feedforward_layernorm);
+    INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, post_feedforward_layernorm);
+    INFINICORE_NN_MODULE(Gemma2Attention, self_attn);
+    INFINICORE_NN_MODULE(Gemma2MLP, mlp);
+
+    size_t layer_idx_;
+    double rms_norm_eps_;
+};
+
+} // namespace infinilm::models::gemma2
diff --git a/csrc/models/gemma2/gemma2_for_causal_lm.cpp b/csrc/models/gemma2/gemma2_for_causal_lm.cpp
new file mode 100644
index 000000000..ce3859e37
--- /dev/null
+++ b/csrc/models/gemma2/gemma2_for_causal_lm.cpp
@@ -0,0 +1,149 @@
+#include "gemma2_for_causal_lm.hpp"
+#include "../models_registry.hpp"
+#include "infinicore/device.hpp"
+#include "infinicore/ops.hpp"
+#include "infinicore/ops/mul_scalar.hpp"
+#include "infinicore/ops/select_last_token_hidden.hpp"
+#include <spdlog/spdlog.h>
+
+namespace infinilm::models::gemma2 {
+
+Gemma2ForCausalLM::Gemma2ForCausalLM(std::shared_ptr<infinilm::config::ModelConfig> model_config,
+                                     const infinicore::Device &device) {
+    model_config_ = model_config;
+
+    size_t hidden_size = model_config->get<size_t>("hidden_size");
+    size_t vocab_size = model_config->get<size_t>("vocab_size");
+    const auto &dtype{model_config->get_dtype()};
+    const auto &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info();
+    pp_size_ = static_cast<size_t>(rank_info.pp_size);
+    pp_stage_ = static_cast<size_t>(rank_info.pp_stage);
+
+    final_logit_softcapping_ = model_config->get_or<float>("final_logit_softcapping", 0.0f);
+
+    model_ = this->register_module<Gemma2Model>("model", model_config, device);
+    if (is_last_pp_stage()) {
+        lm_head_ = this->register_module<infinilm::layers::linear::ReplicatedLinear>("lm_head", hidden_size, vocab_size, false, dtype, device);
+    }
+}
+
+infinilm::InfinilmModel::Output Gemma2ForCausalLM::forward(const infinilm::InfinilmModel::Input &input) const {
+    auto hidden_states = model_->forward(input);
+    if (!is_last_pp_stage()) {
+        return {infinicore::Tensor(), hidden_states};
+    }
+
+    // Packed prefill only needs last-token logits for sampling; keep the
+    // soft-capping passes off the full [S, vocab] tensor (same as
+    // TextCausalLM).
+    auto lm_head_input = hidden_states;
+    if (!input.sample_all_positions && input.input_offsets.has_value()) {
+        const size_t num_requests = input.input_offsets.value()->numel() - 1;
+        const bool is_packed_prefill = hidden_states->ndim() == 3
+                                    && hidden_states->size(0) == 1
+                                    && hidden_states->size(1) > num_requests;
+        if (is_packed_prefill) {
+            lm_head_input = infinicore::Tensor::empty(
+                {1, num_requests, hidden_states->size(2)},
+                hidden_states->dtype(),
+                hidden_states->device());
+            infinicore::op::select_last_token_hidden_(
+                lm_head_input, hidden_states, input.input_offsets.value());
+        }
+    }
+
+    auto logits = lm_head_->forward(lm_head_input);
+    if (final_logit_softcapping_ > 0.0f) {
+        // Gemma-2 final logit soft-capping: squash logits into [-cap, cap].
+        logits = infinicore::op::tanh(infinicore::op::mul_scalar(logits, 1.0f / final_logit_softcapping_));
+        logits = infinicore::op::mul_scalar(logits, final_logit_softcapping_);
+    }
+    return {logits, hidden_states};
+}
+
+infinicore::Tensor Gemma2ForCausalLM::logits_from_hidden(const infinicore::Tensor &hidden_states) const {
+    if (!lm_head_) {
+        throw std::runtime_error("Gemma2ForCausalLM::logits_from_hidden called on a non-last pipeline stage");
+    }
+    auto logits = lm_head_->forward(const_cast<infinicore::Tensor &>(hidden_states));
+    if (final_logit_softcapping_ > 0.0f) {
+        logits = infinicore::op::tanh(infinicore::op::mul_scalar(logits, 1.0f / final_logit_softcapping_));
+        logits = infinicore::op::mul_scalar(logits, final_logit_softcapping_);
+    }
+    return logits;
+}
+
+std::shared_ptr<infinilm::config::ModelConfig> create_gemma2_model_config(
+    std::shared_ptr<infinilm::config::ModelConfig> model_config) {
+    const std::string &model_type = model_config->get<std::string>("model_type");
+    if ("gemma2" != model_type) {
+        throw std::runtime_error(
+            "infinilm::models::gemma2::create_gemma2_model_config: model_type is not gemma2");
+    }
+
+    nlohmann::json &config_json = model_config->get_config_json();
+
+    // Gemma-2 uses a dedicated head_dim that is NOT hidden_size / num_attention_heads
+    // (e.g. 2b: hidden 2304, 8 heads, head_dim 256). Never fall back to the quotient;
+    // query_pre_attn_scalar is an equal-valued but semantically distinct field.
+    if (!config_json.contains("head_dim")) {
+        if (config_json.contains("query_pre_attn_scalar")) {
+            config_json["head_dim"] = model_config->get<size_t>("query_pre_attn_scalar");
+        } else {
+            throw std::runtime_error(
+                "infinilm::models::gemma2::create_gemma2_model_config: config lacks head_dim and query_pre_attn_scalar");
+        }
+    }
+
+    // The generic attention module defaults attention_bias to true; Gemma-2 has none.
+    if (!config_json.contains("attention_bias")) {
+        config_json["attention_bias"] = false;
+    }
+
+    // Gemma-2 default soft-capping values (HF Gemma2Config defaults).
+    if (!config_json.contains("attn_logit_softcapping")) {
+        config_json["attn_logit_softcapping"] = 50.0;
+    }
+    if (!config_json.contains("final_logit_softcapping")) {
+        config_json["final_logit_softcapping"] = 30.0;
+    }
+    if (model_config->get_or<float>("attn_logit_softcapping", 50.0f) < 0.0f || model_config->get_or<float>("final_logit_softcapping", 30.0f) < 0.0f) {
+        throw std::runtime_error(
+            "infinilm::models::gemma2::create_gemma2_model_config: soft-capping values must be non-negative");
+    }
+
+    // Only the tanh-approximated GELU is implemented (matches the checkpoints).
+    const std::string activation = config_json.value("hidden_activation", "gelu_pytorch_tanh");
+    if (activation != "gelu_pytorch_tanh") {
+        throw std::runtime_error(
+            "infinilm::models::gemma2::create_gemma2_model_config: unsupported hidden_activation: " + activation);
+    }
+
+    // Sliding-window attention (alternating local layers) is intentionally not
+    // implemented; correctness holds for sequences up to sliding_window length,
+    // where the local window covers the full causal context. Warn when the
+    // checkpoint's context can exceed that limit so long prompts fail loudly
+    // in logs instead of silently producing wrong attention.
+    if (config_json.contains("sliding_window") && config_json.contains("max_position_embeddings")) {
+        const size_t sliding_window = config_json["sliding_window"].get<size_t>();
+        const size_t max_position = config_json["max_position_embeddings"].get<size_t>();
+        if (sliding_window < max_position) {
+            spdlog::warn(
+                "infinilm::models::gemma2: sliding-window attention is not implemented; results are only "
+                "correct for sequences up to sliding_window={} (max_position_embeddings={})",
+                sliding_window, max_position);
+        }
+    }
+    return model_config;
+}
+
+} // namespace infinilm::models::gemma2
+
+namespace {
+
+INFINILM_REGISTER_CAUSAL_LM_MODEL(
+    gemma2,
+    infinilm::models::gemma2::Gemma2ForCausalLM,
+    infinilm::models::gemma2::create_gemma2_model_config);
+
+} // namespace
diff --git a/csrc/models/gemma2/gemma2_for_causal_lm.hpp b/csrc/models/gemma2/gemma2_for_causal_lm.hpp
new file mode 100644
index 000000000..583cd54be
--- /dev/null
+++ b/csrc/models/gemma2/gemma2_for_causal_lm.hpp
@@ -0,0 +1,39 @@
+#pragma once
+
+#include "../../layers/common_modules.hpp"
+#include "../../models/infinilm_model.hpp"
+#include "gemma2_decoder_layer.hpp"
+#include <memory>
+
+namespace infinilm::models::gemma2 {
+
+using Gemma2Model = infinilm::layers::causal_lm_templates::TextModel<Gemma2DecoderLayer>;
+
+/**
+ * @brief Gemma-2 causal LM.
+ *
+ * Modeled on TextCausalLM, with one addition: Gemma-2 applies final logit
+ * soft-capping (`logits * (1/cap) * tanh(logits * cap)`) after the LM head.
+ */
+class Gemma2ForCausalLM : public infinilm::InfinilmModel {
+public:
+    Gemma2ForCausalLM(std::shared_ptr<infinilm::config::ModelConfig> model_config,
+                      const infinicore::Device &device);
+
+    infinilm::InfinilmModel::Output forward(const infinilm::InfinilmModel::Input &input) const override;
+
+    infinicore::Tensor logits_from_hidden(const infinicore::Tensor &hidden_states) const;
+
+protected:
+    INFINICORE_NN_MODULE(Gemma2Model, model);
+    INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, lm_head);
+
+private:
+    bool is_last_pp_stage() const { return pp_stage_ + 1 == pp_size_; }
+
+    float final_logit_softcapping_{0.0f};
+    size_t pp_size_{1};
+    size_t pp_stage_{0};
+};
+
+} // namespace infinilm::models::gemma2
diff --git a/csrc/models/gemma2/gemma2_mlp.cpp b/csrc/models/gemma2/gemma2_mlp.cpp
new file mode 100644
index 000000000..9f3f896f9
--- /dev/null
+++ b/csrc/models/gemma2/gemma2_mlp.cpp
@@ -0,0 +1,41 @@
+#include "gemma2_mlp.hpp"
+#include "../../global_state/global_state.hpp"
+#include "infinicore/ops.hpp"
+#include "infinicore/ops/mul.hpp"
+
+namespace infinilm::models::gemma2 {
+
+Gemma2MLP::Gemma2MLP(std::shared_ptr<infinilm::config::ModelConfig> model_config,
+                     const infinicore::Device &device) {
+    const auto &dtype{model_config->get_dtype()};
+    size_t hidden_size = model_config->get<size_t>("hidden_size");
+    size_t intermediate_size = model_config->get<size_t>("intermediate_size");
+    bool use_bias = model_config->get_or<bool>("mlp_bias", false);
+
+    const engine::distributed::RankInfo &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info();
+    int tp_rank = rank_info.tp_rank;
+    int tp_size = rank_info.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)); };
+    gate_up_proj_ = std::make_shared<layers::linear::GateUpParallelLinear>(
+        hidden_size, intermediate_size, "gate_proj", "up_proj", register_fn,
+        quantization_method, use_bias, dtype, device, rank_info);
+    down_proj_ = this->register_module<layers::linear::RowParallelLinear>(
+        "down_proj", intermediate_size, hidden_size, quantization_method,
+        use_bias, dtype, device, tp_rank, tp_size, rank_info.comm);
+}
+
+infinicore::Tensor Gemma2MLP::forward(const infinicore::Tensor &hidden_states) const {
+    // 1. Project to gate and up
+    auto hidden_states_mutable = hidden_states;
+    auto [gate, up] = gate_up_proj_->forward_split(hidden_states_mutable);
+    // 2. Gemma-2 activation: gelu_pytorch_tanh on the gate branch, then element-wise product
+    auto activated = infinicore::op::gelu_tanh(gate);
+    auto intermediate = infinicore::op::mul(activated, up);
+    // 3. Project down
+    auto output = down_proj_->forward(intermediate);
+    return output;
+}
+
+} // namespace infinilm::models::gemma2
diff --git a/csrc/models/gemma2/gemma2_mlp.hpp b/csrc/models/gemma2/gemma2_mlp.hpp
new file mode 100644
index 000000000..0cb92cfbd
--- /dev/null
+++ b/csrc/models/gemma2/gemma2_mlp.hpp
@@ -0,0 +1,19 @@
+#pragma once
+
+#include "../../layers/common_modules.hpp"
+
+namespace infinilm::models::gemma2 {
+
+class Gemma2MLP : public infinicore::nn::Module {
+public:
+    Gemma2MLP(std::shared_ptr<infinilm::config::ModelConfig> model_config,
+              const infinicore::Device &device);
+
+    infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const;
+
+protected:
+    std::shared_ptr<infinilm::layers::linear::GateUpParallelLinear> gate_up_proj_;
+    std::shared_ptr<infinilm::layers::linear::RowParallelLinear> down_proj_;
+};
+
+} // namespace infinilm::models::gemma2
diff --git a/csrc/models/gemma3/gemma3_attention.cpp b/csrc/models/gemma3/gemma3_attention.cpp
new file mode 100644
index 000000000..176b8c9b3
--- /dev/null
+++ b/csrc/models/gemma3/gemma3_attention.cpp
@@ -0,0 +1,302 @@
+#include "gemma3_attention.hpp"
+#include "../../global_state/global_state.hpp"
+#include "../../layers/quantization/quantization_scheme.hpp"
+#include "../../utils.hpp"
+#include "infinicore/ops.hpp"
+#include "infinicore/ops/broadcast_to.hpp"
+#include "infinicore/ops/mul.hpp"
+#include "infinicore/ops/mul_scalar.hpp"
+#include <cstring>
+
+namespace infinilm::models::gemma3 {
+
+namespace {
+constexpr float kMaskValue = -1e9f;
+
+size_t read_first_len_(const infinicore::Tensor &lengths) {
+    return static_cast<size_t>(
+        reinterpret_cast<int32_t *>(lengths->to(infinicore::Device::cpu())->data())[0]);
+}
+} // namespace
+
+Gemma3Attention::Gemma3Attention(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);
+    double rms_norm_eps = model_config->get<double>("rms_norm_eps");
+
+    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;
+
+    // Gemma-3 alternates sliding (local) and full (global) attention layers.
+    const auto &layer_types = model_config->get_ref("layer_types");
+    is_sliding_ = layer_types[layer_idx].get<std::string>() == "sliding_attention";
+    sliding_window_ = model_config->get_or<size_t>("sliding_window", 0);
+    if (is_sliding_ && sliding_window_ == 0) {
+        throw std::runtime_error(
+            "infinilm::models::gemma3::Gemma3Attention: layer_types marks this layer as sliding_attention "
+            "but the config has no positive sliding_window");
+    }
+
+    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);
+
+    INFINICORE_NN_MODULE_INIT(q_norm, head_dim_, rms_norm_eps, dtype, device);
+    INFINICORE_NN_MODULE_INIT(k_norm, head_dim_, rms_norm_eps, dtype, device);
+
+    // Sliding layers use the local RoPE base frequency, full layers the global one.
+    double theta = is_sliding_ ? model_config->get_or<double>("rope_local_base_freq", 10000.0)
+                               : model_config->get<double>("rope_theta");
+    size_t max_position_embeddings = model_config->get<size_t>("max_position_embeddings");
+    rotary_emb_ = infinilm::layers::rotary_embedding::get_rope(
+        head_dim_, head_dim_, max_position_embeddings, theta,
+        model_config->get_rope_algo(), dtype, device, nullptr);
+
+    float query_pre_attn_scalar = model_config->get_or<float>("query_pre_attn_scalar", static_cast<float>(head_dim_));
+    scale_ = 1.0f / std::sqrt(query_pre_attn_scalar);
+
+    if (!is_sliding_) {
+        attn_ = std::make_shared<infinilm::layers::attention::AttentionLayer>(
+            num_attention_heads_, head_dim_, scale_, num_key_value_heads_, layer_idx_,
+            kv_cache_k_scale_, kv_cache_v_scale_, attention_backend_);
+    } else {
+        if (attention_backend_ != ::infinilm::backends::AttentionBackend::STATIC_ATTN) {
+            throw std::runtime_error(
+                "infinilm::models::gemma3::Gemma3Attention: sliding-window attention requires the STATIC_ATTN backend");
+        }
+        // Sliding layers bypass StaticAttentionImpl and write the cache
+        // directly in the activation dtype; with an int8 KV cache those
+        // stores would be silently truncated, so refuse the combination
+        // instead of producing garbage attention.
+        if (model_config->get_kv_quant_scheme() != infinilm::quantization::KVQuantAlgo::NONE) {
+            throw std::runtime_error(
+                "infinilm::models::gemma3::Gemma3Attention: KV-cache quantization is not supported for "
+                "sliding-window attention layers");
+        }
+    }
+
+    infinilm::layers::attention::init_kv_cache_quant_params(register_fn, device, kv_cache_k_scale_, kv_cache_v_scale_);
+}
+
+std::tuple<infinicore::Tensor, infinicore::Tensor, infinicore::Tensor>
+Gemma3Attention::project_and_rotate_(const infinicore::Tensor &positions,
+                                     const infinicore::Tensor &hidden_states) const {
+    auto hidden_states_mutable = hidden_states;
+    auto shape = hidden_states->shape();
+    size_t batch_size = shape[0];
+    size_t seq_len = shape[1];
+
+    auto [q, k, v] = qkv_proj_->forward_split(hidden_states_mutable);
+
+    // QK-norm on per-head vectors (before RoPE), mirroring the HF reference.
+    q = q_norm_->forward(q->view({batch_size * seq_len, num_attention_heads_, head_dim_}));
+    k = k_norm_->forward(k->view({batch_size * seq_len, num_key_value_heads_, head_dim_}));
+
+    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_});
+
+    auto pos_shape = positions->shape();
+    infinicore::Tensor pos_ids_for_rope = positions;
+    if (pos_shape.size() == 2) {
+        auto pos_narrowed = positions->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 = positions->contiguous();
+    } else {
+        throw std::runtime_error("infinilm::models::gemma3::Gemma3Attention: Unexpected position_ids shape");
+    }
+
+    rotary_emb_->forward(q_reshaped, pos_ids_for_rope, true);
+    rotary_emb_->forward(k_reshaped, pos_ids_for_rope, true);
+    return {q_reshaped, k_reshaped, v_reshaped};
+}
+
+infinicore::Tensor Gemma3Attention::forward(const infinicore::Tensor &positions,
+                                            const infinicore::Tensor &hidden_states) const {
+    auto [q, k, v] = project_and_rotate_(positions, hidden_states);
+    if (!is_sliding_) {
+        return forward_full_(q, k, v);
+    }
+    return forward_sliding_(q, k, v);
+}
+
+infinicore::Tensor Gemma3Attention::forward_full_(infinicore::Tensor &q_reshaped,
+                                                  infinicore::Tensor &k_reshaped,
+                                                  infinicore::Tensor &v_reshaped) const {
+    auto attn_output = attn_->forward(q_reshaped, k_reshaped, v_reshaped);
+    return o_proj_->forward(attn_output);
+}
+
+infinicore::Tensor Gemma3Attention::forward_sliding_(infinicore::Tensor &q_reshaped,
+                                                     infinicore::Tensor &k_reshaped,
+                                                     infinicore::Tensor &v_reshaped) const {
+    // q/k/v: [bs, seq, heads, head_dim]
+    auto shape = q_reshaped->shape();
+    size_t batch_size = shape[0];
+    size_t seq_len = shape[1];
+    size_t ngroup = num_attention_heads_ / num_key_value_heads_;
+
+    auto &forward_context = infinilm::global_state::get_forward_context();
+    auto &kv_cache = forward_context.kv_cache_vec[layer_idx_];
+    auto &attn_metadata = forward_context.attn_metadata;
+    size_t past_len = read_first_len_(attn_metadata.past_sequence_lengths.value());
+    size_t total_len = past_len + seq_len;
+
+    // KV cache update, mirroring StaticAttentionImpl::do_kv_cache_update.
+    auto k_perm = k_reshaped->permute({0, 2, 1, 3}); // [bs, nkv, seq, head_dim]
+    auto v_perm = v_reshaped->permute({0, 2, 1, 3});
+    auto k_cache_layer = kv_cache->narrow({{0, 0, 1}})->squeeze(0); // [bs, nkv, max_len, head_dim]
+    auto v_cache_layer = kv_cache->narrow({{0, 1, 1}})->squeeze(0);
+    // Same capacity guard as StaticAttentionImpl::do_kv_cache_update.
+    ASSERT(past_len + seq_len <= k_cache_layer->size(2));
+    k_cache_layer->narrow({{2, past_len, seq_len}})->copy_from(k_perm);
+    v_cache_layer->narrow({{2, past_len, seq_len}})->copy_from(v_perm);
+
+    // q_reshaped is [bs, seq, heads, dim] (seq-major); regroup to
+    // [bs*nkv, ng*seq, dim] head-major so GQA groups align with the cache.
+    auto Q = q_reshaped->permute({0, 2, 1, 3})
+                 ->contiguous()
+                 ->view({batch_size * num_key_value_heads_, ngroup * seq_len, head_dim_});
+
+    infinicore::Tensor attn_output;
+    if (seq_len == 1) {
+        // Decode: only the trailing window is visible, so slice the cache and
+        // skip the mask entirely.
+        size_t start = total_len > sliding_window_ ? total_len - sliding_window_ : 0;
+        size_t window_len = total_len - start;
+        auto K = k_cache_layer->narrow({{2, start, window_len}})->view({batch_size * num_key_value_heads_, window_len, head_dim_});
+        auto V = v_cache_layer->narrow({{2, start, window_len}})->view({batch_size * num_key_value_heads_, window_len, head_dim_});
+        auto K_transposed = K->permute({0, 2, 1});
+        auto scores = infinicore::op::matmul(Q, K_transposed, scale_); // [bs*nkv, ng, win]
+        auto scores_viewed = scores->view({batch_size * num_attention_heads_, seq_len, window_len});
+        infinicore::op::causal_softmax_(scores_viewed, scores_viewed);
+        auto out = infinicore::op::matmul(scores, V); // [bs*nkv, ng, head_dim]
+        attn_output = out->view({batch_size, num_attention_heads_, seq_len, head_dim_})
+                          ->permute({0, 2, 1, 3})
+                          ->contiguous()
+                          ->view({batch_size, seq_len, num_attention_heads_ * head_dim_});
+    } else {
+        // Prefill: trim the key range to the union of visible keys (keys older
+        // than past-window+1 are masked for every query), then apply the
+        // window mask; causal_softmax_ provides the aligned-causal part.
+        size_t k_start = 0; // BISECT: truncation disabled
+        size_t k_len = total_len - k_start;
+        auto K = k_cache_layer->narrow({{2, k_start, k_len}})->view({batch_size * num_key_value_heads_, k_len, head_dim_});
+        auto V = v_cache_layer->narrow({{2, k_start, k_len}})->view({batch_size * num_key_value_heads_, k_len, head_dim_});
+        auto K_transposed = K->permute({0, 2, 1});
+        auto scores = infinicore::op::matmul(Q, K_transposed, scale_); // [bs*nkv, ng*seq, k_len]
+
+        auto scores_viewed = scores->view({batch_size * num_attention_heads_, seq_len, k_len});
+        // causal_softmax_ provides the aligned-causal mask and the softmax; the
+        // additive mask only needs to enforce the window constraint
+        // (key j is invisible to query i when past+i-j >= sliding_window).
+        infinicore::Tensor mask = sliding_mask_(seq_len, past_len, k_start, k_len, scores_viewed->dtype(), scores_viewed->device());
+        mask = infinicore::op::broadcast_to(
+            mask->view({static_cast<int64_t>(1), static_cast<int64_t>(seq_len), static_cast<int64_t>(k_len)}),
+            {static_cast<int64_t>(batch_size * num_attention_heads_),
+             static_cast<int64_t>(seq_len),
+             static_cast<int64_t>(k_len)});
+        auto scores_masked = infinicore::op::add(scores_viewed, mask);
+        infinicore::op::causal_softmax_(scores_masked, scores_masked);
+
+        auto scores_grouped = scores_masked->view({batch_size * num_key_value_heads_, ngroup * seq_len, k_len});
+        auto out = infinicore::op::matmul(scores_grouped, V); // [bs*nkv, ng*seq, head_dim]
+        attn_output = out->view({batch_size, num_attention_heads_, seq_len, head_dim_})
+                          ->permute({0, 2, 1, 3})
+                          ->contiguous()
+                          ->view({batch_size, seq_len, num_attention_heads_ * head_dim_});
+    }
+    return o_proj_->forward(attn_output);
+}
+
+namespace {
+// bf16 is the truncated top half of a float; build the bit pattern directly so
+// the mask can be created in the score dtype without a cast op.
+uint16_t to_bf16_bits(float value) {
+    uint32_t bits = 0;
+    std::memcpy(&bits, &value, sizeof(bits));
+    return static_cast<uint16_t>(bits >> 16);
+}
+} // namespace
+
+infinicore::Tensor Gemma3Attention::sliding_mask_(size_t seq, size_t past, size_t k_start,
+                                                  size_t k_len, const infinicore::DataType &dtype,
+                                                  const infinicore::Device &device) const {
+    const size_t total = k_start + k_len;
+    // Masks depend only on (seq, total, window), so they are cached per shape
+    // to avoid rebuilding the host buffer on every prefill. The cache is keyed
+    // by shape and bounded in practice by the distinct prompt lengths seen;
+    // clear it if it grows beyond a sane bound.
+    if (mask_cache_.size() > 128) {
+        mask_cache_.clear();
+    }
+    size_t key = seq * 1000000000ULL + total;
+    auto it = mask_cache_.find(key);
+    if (it != mask_cache_.end()) {
+        return it->second;
+    }
+
+    auto host = infinicore::Tensor::empty({seq, total}, dtype, infinicore::Device::cpu());
+    if (dtype == infinicore::DataType::F32) {
+        auto *data = reinterpret_cast<float *>(host->data());
+        for (size_t i = 0; i < seq; ++i) {
+            size_t q_pos = past + i;
+            for (size_t j = 0; j < k_len; ++j) {
+                size_t key_pos = k_start + j;
+                bool visible = (key_pos <= q_pos) && (q_pos - key_pos < sliding_window_);
+                data[i * k_len + j] = visible ? 0.0f : kMaskValue;
+            }
+        }
+    } else if (dtype == infinicore::DataType::BF16) {
+        auto *data = reinterpret_cast<uint16_t *>(host->data());
+        const uint16_t masked = to_bf16_bits(kMaskValue);
+        for (size_t i = 0; i < seq; ++i) {
+            size_t q_pos = past + i;
+            for (size_t j = 0; j < k_len; ++j) {
+                size_t key_pos = k_start + j;
+                bool visible = (key_pos <= q_pos) && (q_pos - key_pos < sliding_window_);
+                data[i * k_len + j] = visible ? 0 : masked;
+            }
+        }
+    } else if (dtype == infinicore::DataType::F16) {
+        // -1e9 overflows the fp16 range; use fp16 -inf (0xFC00) instead.
+        auto *data = reinterpret_cast<uint16_t *>(host->data());
+        constexpr uint16_t kFp16NegInf = 0xFC00;
+        for (size_t i = 0; i < seq; ++i) {
+            size_t q_pos = past + i;
+            for (size_t j = 0; j < k_len; ++j) {
+                size_t key_pos = k_start + j;
+                bool visible = (key_pos <= q_pos) && (q_pos - key_pos < sliding_window_);
+                data[i * k_len + j] = visible ? 0 : kFp16NegInf;
+            }
+        }
+    } else {
+        throw std::runtime_error("infinilm::models::gemma3::Gemma3Attention::sliding_mask_: unsupported dtype");
+    }
+    infinicore::Tensor mask = host->to(device);
+    mask_cache_.emplace(key, mask);
+    return mask;
+}
+
+} // namespace infinilm::models::gemma3
diff --git a/csrc/models/gemma3/gemma3_attention.hpp b/csrc/models/gemma3/gemma3_attention.hpp
new file mode 100644
index 000000000..0455a8d92
--- /dev/null
+++ b/csrc/models/gemma3/gemma3_attention.hpp
@@ -0,0 +1,94 @@
+#pragma once
+
+#include "../../layers/common_modules.hpp"
+#include <unordered_map>
+
+namespace infinilm::models::gemma3 {
+
+/**
+ * @brief Gemma-3 attention with QK-norm, per-layer-type RoPE theta and
+ * (for sliding layers) model-level sliding-window attention.
+ *
+ * Gemma-3 alternates "sliding_attention" layers (local RoPE theta, attention
+ * restricted to the last `sliding_window` keys) with "full_attention" layers
+ * (global RoPE theta, plain causal attention) in a
+ * `sliding_window_pattern`-periodic layout (default 6 -> 5:1).
+ *
+ * Full layers delegate to the shared AttentionLayer. Sliding layers manage the
+ * static KV cache directly: the cache layout and update semantics mirror
+ * `StaticAttentionImpl`, the decode path reads only the trailing window, and
+ * the prefill path applies a banded causal mask built on host and cached per
+ * (seq_len, total_len). Requires the STATIC attention backend.
+ */
+class Gemma3Attention : public infinicore::nn::Module {
+public:
+    Gemma3Attention(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:
+    // Shared prologue: project, QK-normalize, apply this layer's RoPE.
+    std::tuple<infinicore::Tensor, infinicore::Tensor, infinicore::Tensor>
+    project_and_rotate_(const infinicore::Tensor &positions,
+                        const infinicore::Tensor &hidden_states) const;
+
+    infinicore::Tensor forward_full_(infinicore::Tensor &q_reshaped,
+                                     infinicore::Tensor &k_reshaped,
+                                     infinicore::Tensor &v_reshaped) const;
+
+    infinicore::Tensor forward_sliding_(infinicore::Tensor &q_reshaped,
+                                        infinicore::Tensor &k_reshaped,
+                                        infinicore::Tensor &v_reshaped) const;
+
+    // Banded causal mask [seq, total] on `device`: entry (i, j) is 0 when key j
+    // is visible to query i (j <= past+i and past+i-j < window), else -1e9.
+    infinicore::Tensor sliding_mask_(size_t seq, size_t past, size_t k_start,
+                                     size_t k_len, const infinicore::DataType &dtype,
+                                     const infinicore::Device &device) const;
+
+protected:
+    std::shared_ptr<infinilm::layers::linear::QKVParallelLinear> qkv_proj_;
+    std::shared_ptr<infinilm::layers::linear::RowParallelLinear> o_proj_;
+    INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, q_norm);
+    INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, k_norm);
+    std::shared_ptr<infinicore::nn::RoPE> rotary_emb_;
+
+    // Only used by full-attention layers.
+    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 scale_;
+    bool is_sliding_{false};
+    size_t sliding_window_{0};
+
+    // Per-(seq, total) cache of built sliding masks (device tensors).
+    mutable std::unordered_map<size_t, infinicore::Tensor> mask_cache_;
+
+    // For off-line kv cache quantization
+    INFINICORE_NN_PARAMETER(kv_cache_k_scale);
+    INFINICORE_NN_PARAMETER(kv_cache_v_scale);
+};
+
+} // namespace infinilm::models::gemma3
diff --git a/csrc/models/gemma3/gemma3_decoder_layer.cpp b/csrc/models/gemma3/gemma3_decoder_layer.cpp
new file mode 100644
index 000000000..4df7def4d
--- /dev/null
+++ b/csrc/models/gemma3/gemma3_decoder_layer.cpp
@@ -0,0 +1,68 @@
+#include "gemma3_decoder_layer.hpp"
+
+namespace infinilm::models::gemma3 {
+
+Gemma3DecoderLayer::Gemma3DecoderLayer(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<Gemma3Attention>("self_attn", model_config, layer_idx, device);
+    mlp_ = this->register_module<Gemma3MLP>("mlp", model_config, device);
+}
+
+std::tuple<infinicore::Tensor, infinicore::Tensor> Gemma3DecoderLayer::forward(const infinicore::Tensor &positions,
+                                                                               infinicore::Tensor &hidden_states,
+                                                                               infinicore::Tensor &residual) {
+    // 1. Normalize the mainstream (fused: residual += incoming branch, hidden = norm(residual)).
+    input_layernorm_->forward_inplace(hidden_states, residual);
+
+    // 2. Attention on the normalized branch.
+    hidden_states = self_attn_->forward(positions, hidden_states);
+
+    // 3. Gemma 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 Gemma3DecoderLayer::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::gemma3
diff --git a/csrc/models/gemma3/gemma3_decoder_layer.hpp b/csrc/models/gemma3/gemma3_decoder_layer.hpp
new file mode 100644
index 000000000..df20218b8
--- /dev/null
+++ b/csrc/models/gemma3/gemma3_decoder_layer.hpp
@@ -0,0 +1,45 @@
+#pragma once
+
+#include "../../layers/common_modules.hpp"
+#include "gemma3_attention.hpp"
+#include "gemma3_mlp.hpp"
+#include "infinicore/nn/rmsnorm.hpp"
+#include <memory>
+
+namespace infinilm::models::gemma3 {
+
+/**
+ * @brief Gemma-3 decoder layer.
+ *
+ * Same four-norm, branch-normalize-then-add structure as Gemma-2 (the residual
+ * contract with TextModel is identical); the attention slot is Gemma3's own
+ * (QK-norm + sliding/global layer types).
+ */
+class Gemma3DecoderLayer : public infinicore::nn::Module {
+public:
+    Gemma3DecoderLayer(std::shared_ptr<infinilm::config::ModelConfig> model_config,
+                       size_t layer_idx,
+                       const infinicore::Device &device);
+
+    std::tuple<infinicore::Tensor, infinicore::Tensor> forward(const infinicore::Tensor &positions,
+                                                               infinicore::Tensor &hidden_states,
+                                                               infinicore::Tensor &residual);
+
+    infinicore::Tensor forward(const infinicore::Tensor &positions,
+                               infinicore::Tensor &hidden_states);
+
+    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(infinicore::nn::RMSNorm, pre_feedforward_layernorm);
+    INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, post_feedforward_layernorm);
+    INFINICORE_NN_MODULE(Gemma3Attention, self_attn);
+    INFINICORE_NN_MODULE(Gemma3MLP, mlp);
+
+    size_t layer_idx_;
+    double rms_norm_eps_;
+};
+
+} // namespace infinilm::models::gemma3
diff --git a/csrc/models/gemma3/gemma3_for_causal_lm.cpp b/csrc/models/gemma3/gemma3_for_causal_lm.cpp
new file mode 100644
index 000000000..76e9c3573
--- /dev/null
+++ b/csrc/models/gemma3/gemma3_for_causal_lm.cpp
@@ -0,0 +1,74 @@
+#include "gemma3_for_causal_lm.hpp"
+#include "../models_registry.hpp"
+
+namespace infinilm::models::gemma3 {
+
+std::shared_ptr<infinilm::config::ModelConfig> create_gemma3_text_model_config(
+    std::shared_ptr<infinilm::config::ModelConfig> model_config) {
+    const std::string &model_type = model_config->get<std::string>("model_type");
+    if ("gemma3_text" != model_type) {
+        throw std::runtime_error(
+            "infinilm::models::gemma3::create_gemma3_text_model_config: model_type is not gemma3_text");
+    }
+
+    nlohmann::json &config_json = model_config->get_config_json();
+
+    // Linear RoPE scaling (gemma-3-4b/12b/27b text configs use factor 8 on the
+    // global layers to reach 128k context) is not implemented. Refuse loudly
+    // instead of silently building global layers with unscaled RoPE, which
+    // corrupts long-context outputs. The text-only gemma-3-1b has no
+    // rope_scaling and is unaffected.
+    if (config_json.contains("rope_scaling") && !config_json["rope_scaling"].is_null()) {
+        throw std::runtime_error(
+            "infinilm::models::gemma3::create_gemma3_text_model_config: rope_scaling is not supported "
+            "(gemma-3-4b/12b/27b use linear scaling on global attention layers); refusing to load a "
+            "checkpoint whose global-layer RoPE cannot be reproduced");
+    }
+
+    // Gemma-3 has a dedicated head_dim that is NOT hidden_size / num_attention_heads.
+    if (!config_json.contains("head_dim")) {
+        if (config_json.contains("query_pre_attn_scalar")) {
+            config_json["head_dim"] = model_config->get<size_t>("query_pre_attn_scalar");
+        } else {
+            throw std::runtime_error(
+                "infinilm::models::gemma3::create_gemma3_text_model_config: config lacks head_dim and query_pre_attn_scalar");
+        }
+    }
+
+    // The generic attention module defaults attention_bias to true; Gemma-3 has none.
+    if (!config_json.contains("attention_bias")) {
+        config_json["attention_bias"] = false;
+    }
+
+    // Only the tanh-approximated GELU is implemented (matches the checkpoints).
+    const std::string activation = config_json.value("hidden_activation", "gelu_pytorch_tanh");
+    if (activation != "gelu_pytorch_tanh") {
+        throw std::runtime_error(
+            "infinilm::models::gemma3::create_gemma3_text_model_config: unsupported hidden_activation: " + activation);
+    }
+
+    // Generate the per-layer attention types exactly like HF Gemma3TextConfig:
+    // every `sliding_window_pattern`-th layer (default 6 -> 5:1) is full attention.
+    if (!config_json.contains("layer_types")) {
+        size_t num_layers = model_config->get<size_t>("num_hidden_layers");
+        size_t pattern = config_json.value("sliding_window_pattern", 6);
+        nlohmann::json layer_types = nlohmann::json::array();
+        for (size_t i = 0; i < num_layers; ++i) {
+            layer_types.push_back(((i + 1) % pattern != 0) ? "sliding_attention" : "full_attention");
+        }
+        config_json["layer_types"] = layer_types;
+    }
+
+    return model_config;
+}
+
+} // namespace infinilm::models::gemma3
+
+namespace {
+
+INFINILM_REGISTER_CAUSAL_LM_MODEL(
+    gemma3_text,
+    infinilm::models::gemma3::Gemma3ForCausalLM,
+    infinilm::models::gemma3::create_gemma3_text_model_config);
+
+} // namespace
diff --git a/csrc/models/gemma3/gemma3_for_causal_lm.hpp b/csrc/models/gemma3/gemma3_for_causal_lm.hpp
new file mode 100644
index 000000000..1329b51b4
--- /dev/null
+++ b/csrc/models/gemma3/gemma3_for_causal_lm.hpp
@@ -0,0 +1,21 @@
+#pragma once
+
+#include "../../layers/common_modules.hpp"
+#include "gemma3_decoder_layer.hpp"
+#include <memory>
+
+namespace infinilm::models::gemma3 {
+
+using Gemma3Model = infinilm::layers::causal_lm_templates::TextModel<Gemma3DecoderLayer>;
+
+// Gemma-3 has no logit soft-capping, so the generic causal-LM template applies
+// as-is (unlike Gemma-2, which needs a custom top level for final soft-cap).
+using Gemma3ForCausalLM = infinilm::layers::causal_lm_templates::TextCausalLM<Gemma3Model>;
+
+} // namespace infinilm::models::gemma3
+
+namespace infinilm::models::gemma3 {
+
+std::shared_ptr<infinilm::config::ModelConfig> create_gemma3_text_model_config(std::shared_ptr<infinilm::config::ModelConfig> model_config);
+
+} // namespace infinilm::models::gemma3
diff --git a/csrc/models/gemma3/gemma3_mlp.hpp b/csrc/models/gemma3/gemma3_mlp.hpp
new file mode 100644
index 000000000..8a972d650
--- /dev/null
+++ b/csrc/models/gemma3/gemma3_mlp.hpp
@@ -0,0 +1,12 @@
+#pragma once
+
+#include "../gemma2/gemma2_mlp.hpp"
+
+namespace infinilm::models::gemma3 {
+
+// Gemma-3's MLP is identical to Gemma-2's (SwiGLU layout with
+// gelu_pytorch_tanh), so reuse the implementation (qwen3_moe aliases qwen3 the
+// same way).
+using Gemma3MLP = infinilm::models::gemma2::Gemma2MLP;
+
+} // namespace infinilm::models::gemma3
diff --git a/python/infinilm/modeling_utils.py b/python/infinilm/modeling_utils.py
index b210b0c56..c1b9be242 100644
--- a/python/infinilm/modeling_utils.py
+++ b/python/infinilm/modeling_utils.py
@@ -18,6 +18,12 @@ def _get_scale_emb(model_path: str) -> float:
         raise FileNotFoundError(f"config.json not found at {config_path}")
     with open(config_path, "r") as f:
         config = json.load(f)
+    # Gemma multiplies the embedding output by sqrt(hidden_size) at forward
+    # time; bake it into the embedding weight at load time instead. The tied
+    # lm_head is filled from the *unscaled* copy (see the load path), which
+    # matches the HF reference semantics.
+    if config.get("model_type") in ("gemma2", "gemma3_text"):
+        return float(config.get("hidden_size", 1.0)) ** 0.5
     if config.get("model_type") not in ("fm9g", "minicpm"):
         return 1.0
     return config.get("scale_emb", 1.0)
@@ -1071,6 +1077,29 @@ def _remap_kimi_k3(state_dict, config):
     return state_dict
 
 
+def _remap_gemma(state_dict, config=None):
+    """Apply Gemma-2 / Gemma-3 load-time weight fixes.
+
+    Gemma checkpoints store each RMSNorm weight as `w - 1` (the module computes
+    `(1 + w) * norm(x)`), so shift them back by +1 before loading. The embedding
+    output scaling by sqrt(hidden_size) is handled by `_get_scale_emb`.
+    Parameter names follow the Llama-style layout and need no renaming.
+    """
+    norm_suffixes = (
+        "input_layernorm.weight",
+        "post_attention_layernorm.weight",
+        "pre_feedforward_layernorm.weight",
+        "post_feedforward_layernorm.weight",
+        "self_attn.q_norm.weight",
+        "self_attn.k_norm.weight",
+        "model.norm.weight",
+    )
+    for key, tensor in state_dict.items():
+        if key.endswith(norm_suffixes):
+            state_dict[key] = tensor + torch.ones_like(tensor)
+    return state_dict
+
+
 _WEIGHT_REMAPPER = {
     "glm4": _remap_glm4,
     "chatglm": _remap_chatglm,
@@ -1083,4 +1112,6 @@ def _remap_kimi_k3(state_dict, config):
     "qwen3_5_moe": _remap_qwen3_5_moe,
     "qwen3_next": _remap_qwen3_next,
     "kimi_k3": _remap_kimi_k3,
+    "gemma2": _remap_gemma,
+    "gemma3_text": _remap_gemma,
 }
diff --git a/test/models/gemma2/test_remap_gemma.py b/test/models/gemma2/test_remap_gemma.py
new file mode 100644
index 000000000..31dabe22d
--- /dev/null
+++ b/test/models/gemma2/test_remap_gemma.py
@@ -0,0 +1,90 @@
+import json
+import math
+import os
+import tempfile
+import unittest
+
+import torch
+from infinilm.modeling_utils import _get_scale_emb, _remap_gemma
+
+
+def _make_state_dict():
+    return {
+        "model.embed_tokens.weight": torch.randn(8, 4),
+        "model.norm.weight": torch.zeros(4),
+        "model.layers.0.input_layernorm.weight": torch.zeros(4),
+        "model.layers.0.post_attention_layernorm.weight": torch.zeros(4),
+        "model.layers.0.pre_feedforward_layernorm.weight": torch.zeros(4),
+        "model.layers.0.post_feedforward_layernorm.weight": torch.zeros(4),
+        "model.layers.0.self_attn.q_norm.weight": torch.zeros(4),
+        "model.layers.0.self_attn.k_norm.weight": torch.zeros(4),
+        "model.layers.0.self_attn.q_proj.weight": torch.randn(6, 4),
+        "model.layers.0.mlp.gate_proj.weight": torch.randn(5, 4),
+        "lm_head.weight": torch.randn(8, 4),
+    }
+
+
+class RemapGemmaTest(unittest.TestCase):
+    def test_norm_weights_shifted_by_one(self):
+        sd = _make_state_dict()
+        out = _remap_gemma(sd, None)
+        for key in (
+            "model.norm.weight",
+            "model.layers.0.input_layernorm.weight",
+            "model.layers.0.post_attention_layernorm.weight",
+            "model.layers.0.pre_feedforward_layernorm.weight",
+            "model.layers.0.post_feedforward_layernorm.weight",
+            "model.layers.0.self_attn.q_norm.weight",
+            "model.layers.0.self_attn.k_norm.weight",
+        ):
+            self.assertTrue(torch.allclose(out[key], torch.ones(4)), key)
+
+    def test_non_norm_weights_untouched(self):
+        sd = _make_state_dict()
+        out = _remap_gemma(sd, None)
+        self.assertTrue(
+            torch.allclose(
+                out["model.layers.0.self_attn.q_proj.weight"],
+                sd["model.layers.0.self_attn.q_proj.weight"],
+            )
+        )
+        self.assertTrue(
+            torch.allclose(
+                out["model.layers.0.mlp.gate_proj.weight"],
+                sd["model.layers.0.mlp.gate_proj.weight"],
+            )
+        )
+        self.assertTrue(
+            torch.allclose(
+                out["model.embed_tokens.weight"], sd["model.embed_tokens.weight"]
+            )
+        )
+
+    def test_tied_lm_head_not_shifted(self):
+        sd = _make_state_dict()
+        out = _remap_gemma(sd, None)
+        self.assertTrue(torch.allclose(out["lm_head.weight"], sd["lm_head.weight"]))
+
+
+class ScaleEmbTest(unittest.TestCase):
+    def _write_config(self, config):
+        d = tempfile.mkdtemp()
+        with open(os.path.join(d, "config.json"), "w") as f:
+            json.dump(config, f)
+        return d
+
+    def test_gemma2_uses_sqrt_hidden(self):
+        path = self._write_config({"model_type": "gemma2", "hidden_size": 2304})
+        self.assertAlmostEqual(_get_scale_emb(path), math.sqrt(2304))
+
+    def test_gemma3_text_uses_sqrt_hidden(self):
+        path = self._write_config({"model_type": "gemma3_text", "hidden_size": 1152})
+        self.assertAlmostEqual(_get_scale_emb(path), math.sqrt(1152))
+
+    def test_other_models_default_to_one(self):
+        path = self._write_config({"model_type": "llama", "hidden_size": 4096})
+        self.assertEqual(_get_scale_emb(path), 1.0)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/test/models/gemma3/test_remap_gemma.py b/test/models/gemma3/test_remap_gemma.py
new file mode 100644
index 000000000..31dabe22d
--- /dev/null
+++ b/test/models/gemma3/test_remap_gemma.py
@@ -0,0 +1,90 @@
+import json
+import math
+import os
+import tempfile
+import unittest
+
+import torch
+from infinilm.modeling_utils import _get_scale_emb, _remap_gemma
+
+
+def _make_state_dict():
+    return {
+        "model.embed_tokens.weight": torch.randn(8, 4),
+        "model.norm.weight": torch.zeros(4),
+        "model.layers.0.input_layernorm.weight": torch.zeros(4),
+        "model.layers.0.post_attention_layernorm.weight": torch.zeros(4),
+        "model.layers.0.pre_feedforward_layernorm.weight": torch.zeros(4),
+        "model.layers.0.post_feedforward_layernorm.weight": torch.zeros(4),
+        "model.layers.0.self_attn.q_norm.weight": torch.zeros(4),
+        "model.layers.0.self_attn.k_norm.weight": torch.zeros(4),
+        "model.layers.0.self_attn.q_proj.weight": torch.randn(6, 4),
+        "model.layers.0.mlp.gate_proj.weight": torch.randn(5, 4),
+        "lm_head.weight": torch.randn(8, 4),
+    }
+
+
+class RemapGemmaTest(unittest.TestCase):
+    def test_norm_weights_shifted_by_one(self):
+        sd = _make_state_dict()
+        out = _remap_gemma(sd, None)
+        for key in (
+            "model.norm.weight",
+            "model.layers.0.input_layernorm.weight",
+            "model.layers.0.post_attention_layernorm.weight",
+            "model.layers.0.pre_feedforward_layernorm.weight",
+            "model.layers.0.post_feedforward_layernorm.weight",
+            "model.layers.0.self_attn.q_norm.weight",
+            "model.layers.0.self_attn.k_norm.weight",
+        ):
+            self.assertTrue(torch.allclose(out[key], torch.ones(4)), key)
+
+    def test_non_norm_weights_untouched(self):
+        sd = _make_state_dict()
+        out = _remap_gemma(sd, None)
+        self.assertTrue(
+            torch.allclose(
+                out["model.layers.0.self_attn.q_proj.weight"],
+                sd["model.layers.0.self_attn.q_proj.weight"],
+            )
+        )
+        self.assertTrue(
+            torch.allclose(
+                out["model.layers.0.mlp.gate_proj.weight"],
+                sd["model.layers.0.mlp.gate_proj.weight"],
+            )
+        )
+        self.assertTrue(
+            torch.allclose(
+                out["model.embed_tokens.weight"], sd["model.embed_tokens.weight"]
+            )
+        )
+
+    def test_tied_lm_head_not_shifted(self):
+        sd = _make_state_dict()
+        out = _remap_gemma(sd, None)
+        self.assertTrue(torch.allclose(out["lm_head.weight"], sd["lm_head.weight"]))
+
+
+class ScaleEmbTest(unittest.TestCase):
+    def _write_config(self, config):
+        d = tempfile.mkdtemp()
+        with open(os.path.join(d, "config.json"), "w") as f:
+            json.dump(config, f)
+        return d
+
+    def test_gemma2_uses_sqrt_hidden(self):
+        path = self._write_config({"model_type": "gemma2", "hidden_size": 2304})
+        self.assertAlmostEqual(_get_scale_emb(path), math.sqrt(2304))
+
+    def test_gemma3_text_uses_sqrt_hidden(self):
+        path = self._write_config({"model_type": "gemma3_text", "hidden_size": 1152})
+        self.assertAlmostEqual(_get_scale_emb(path), math.sqrt(1152))
+
+    def test_other_models_default_to_one(self):
+        path = self._write_config({"model_type": "llama", "hidden_size": 4096})
+        self.assertEqual(_get_scale_emb(path), 1.0)
+
+
+if __name__ == "__main__":
+    unittest.main()