# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# Copyright 2025 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import copy
import json
from typing import Iterable, List, Optional, Set, Tuple

import einops
import torch
import torch.nn.functional as F
from torch import nn
from transformers import (
    ROPE_INIT_FUNCTIONS,
    Gemma3TextConfig,
    PretrainedConfig,
    PreTrainedModel,
)

from sglang.srt.layers.activation import GeluAndMul
from sglang.srt.layers.layernorm import Gemma3RMSNorm
from sglang.srt.layers.linear import (
    MergedColumnParallelLinear,
    QKVParallelLinear,
    RowParallelLinear,
)
from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.pooler import EmbeddingPoolerOutput, Pooler, PoolingType
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.radix_attention import AttentionType, RadixAttention
from sglang.srt.layers.rotary_embedding import apply_rotary_pos_emb, get_rope
from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_loader.weight_utils import (
    default_weight_loader,
    maybe_remap_kv_scale_name,
)
from sglang.srt.runtime_context import get_parallel
from sglang.srt.utils import add_prefix, cpu_has_amx_support, is_cpu, make_layers
from sglang.srt.utils.hf_transformers.common import _resolve_local_or_cached_file

_is_cpu = is_cpu()
_is_cpu_amx_available = cpu_has_amx_support()


# Aligned with HF's implementation, using sliding window inclusive with the last token
# SGLang assumes exclusive
def get_attention_sliding_window_size(config):
    return config.sliding_window - 1


# Adapted from:
# https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/gemma3.py
def extract_layer_index(prefix: str) -> int:
    """Extract the layer index from a prefix string."""
    parts = prefix.split(".")
    for part in parts:
        if part.startswith("layers."):
            layer_str = part.split(".")[-1]
            try:
                return int(layer_str)
            except ValueError:
                continue
    return -1


class Gemma3MLP(nn.Module):
    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        hidden_activation: str,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.gate_up_proj = MergedColumnParallelLinear(
            hidden_size,
            [intermediate_size] * 2,
            bias=False,
            quant_config=quant_config,
            prefix=add_prefix("gate_up_proj", prefix),
        )
        self.down_proj = RowParallelLinear(
            intermediate_size,
            hidden_size,
            bias=False,
            quant_config=quant_config,
            prefix=add_prefix("down_proj", prefix),
        )
        if hidden_activation != "gelu_pytorch_tanh":
            raise ValueError(
                f"{self.__class__.__name__} uses `gelu_pytorch_tanh` as the hidden activation "
                "function. Please set `hidden_activation` to "
                "`gelu_pytorch_tanh`."
            )
        self.act_fn = GeluAndMul()
        self.prefix = prefix

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate_up, _ = self.gate_up_proj(x)
        x = self.act_fn(gate_up)
        x, _ = self.down_proj(x)
        return x


class Gemma3Attention(nn.Module):
    def __init__(
        self,
        layer_id: int,
        config: Gemma3TextConfig,
        max_position_embeddings: int,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.layer_id = layer_id
        self.config = config
        tp_size = get_parallel().tp_size

        self.total_num_heads = config.num_attention_heads
        assert self.total_num_heads % tp_size == 0
        self.num_heads = self.total_num_heads // tp_size
        self.total_num_kv_heads = config.num_key_value_heads

        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)

        if self.total_num_kv_heads >= tp_size:
            # Number of KV heads is greater than TP size, so we partition
            # the KV heads across multiple tensor parallel GPUs.
            assert self.total_num_kv_heads % tp_size == 0
        else:
            # Number of KV heads is less than TP size, so we replicate
            # the KV heads across multiple tensor parallel GPUs.
            assert tp_size % self.total_num_kv_heads == 0

        hidden_size = config.hidden_size

        head_dim = getattr(
            config, "head_dim", hidden_size // config.num_attention_heads
        )
        self.head_dim = head_dim
        partial_rotary_factor = getattr(config, "partial_rotary_factor", 1)
        self.rotary_dim = int(partial_rotary_factor * self.head_dim)
        self.q_size = self.num_heads * self.head_dim

        self.kv_size = self.num_kv_heads * self.head_dim
        self.scaling = config.query_pre_attn_scalar**-0.5

        self.qkv_proj = QKVParallelLinear(
            hidden_size,
            self.head_dim,
            self.total_num_heads,
            self.total_num_kv_heads,
            bias=config.attention_bias,
            quant_config=quant_config,
            prefix=add_prefix("qkv_proj", prefix),
        )
        self.o_proj = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            hidden_size,
            bias=config.attention_bias,
            quant_config=quant_config,
            prefix=add_prefix("o_proj", prefix),
        )

        self.is_sliding = config.layer_types[layer_id] == "sliding_attention"

        # In transformers v5, rope_parameters is nested per layer type:
        #   {"sliding_attention": {"rope_theta": 10000}, "full_attention": {"rope_theta": 1000000}}
        # In v4 it was flat: {"rope_type": "default", "rope_theta": ...}
        rope_params = config.rope_parameters
        is_nested = isinstance(rope_params, dict) and "full_attention" in rope_params

        # Initialize the rotary embedding.
        if self.is_sliding:
            # Local attention. Override the values in config.json.
            if is_nested:
                self.rope_theta = rope_params["sliding_attention"].get(
                    "rope_theta", 10000.0
                )
            else:
                self.rope_theta = getattr(config, "rope_local_base_freq", 10000.0)
            self.rope_scaling = {"rope_type": "default"}
            # FIXME(mick): idk why vllm does this
            # self.sliding_window = config.interleaved_sliding_window
            self.sliding_window = get_attention_sliding_window_size(config)
        else:
            # Global attention. Use the values in config.json.
            if is_nested:
                self.rope_theta = rope_params["full_attention"].get(
                    "rope_theta", 1000000.0
                )
            else:
                self.rope_theta = (
                    rope_params.get("rope_theta", 10000.0) if rope_params else 10000.0
                )
            self.rope_scaling = {"rope_type": "default"}
            self.sliding_window = None
        self.rotary_emb = get_rope(
            self.head_dim,
            rotary_dim=self.rotary_dim,
            max_position=max_position_embeddings,
            base=self.rope_theta,
            rope_scaling=self.rope_scaling,
            is_neox_style=getattr(config, "rope_is_neox_style", True),
        )
        self.attn = RadixAttention(
            self.num_heads,
            self.head_dim,
            self.scaling,
            num_kv_heads=self.num_kv_heads,
            layer_id=layer_id,
            logit_cap=0.0,
            # Module must also define `get_attention_sliding_window_size` to correctly initialize
            # attention backend in `ForwardBatch`.
            sliding_window_size=self.sliding_window,
            quant_config=quant_config,
            prefix=add_prefix("attn", prefix),
            # Gemma3 uses this attention implementation for both its causal
            # LMs and EmbeddingGemma.  Only the latter enables bidirectional
            # prompt attention in its upstream config.
            attn_type=(
                AttentionType.DECODER_BIDIRECTIONAL
                if getattr(config, "use_bidirectional_attention", False)
                else AttentionType.DECODER
            ),
        )

        # Gemma3 adds normalization for q and k
        self.q_norm = Gemma3RMSNorm(dim=config.head_dim, eps=config.rms_norm_eps)
        self.k_norm = Gemma3RMSNorm(dim=config.head_dim, eps=config.rms_norm_eps)

    def forward_cpu(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
        forward_batch: ForwardBatch,
        **kwargs,
    ) -> torch.Tensor:
        qkv, _ = self.qkv_proj(hidden_states)
        # [s, h * head_dim]
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

        # [s, h, head_dim]
        q = q.unflatten(-1, (self.num_heads, self.head_dim)).unsqueeze(0)
        q = self.q_norm(q)
        k = k.unflatten(-1, (self.num_kv_heads, self.head_dim)).unsqueeze(0)
        k = self.k_norm(k)
        q, k = self.rotary_emb(positions, q, k)

        attn_output = self.attn(q, k, v, forward_batch=forward_batch)

        # Compatible with triton backend which returns [1, s, h, head_dim]
        if attn_output.dim() == 4 and attn_output.shape[0] == 1:
            attn_output = attn_output.squeeze(0)
            attn_output = attn_output.flatten(-2, -1)
        # [s, h * head_dim]

        output, _ = self.o_proj(attn_output)
        return output

    def forward_native(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
        forward_batch: ForwardBatch,
        **kwargs,
    ) -> torch.Tensor:
        qkv, _ = self.qkv_proj(hidden_states)
        # [s, h * head_dim]
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

        # [s, h, head_dim]
        q = q.unflatten(-1, (self.num_heads, self.head_dim))
        # -> [h, s, head_dim]
        q = q.transpose(0, 1).unsqueeze(0)
        q = self.q_norm(q)
        k = k.unflatten(-1, (self.num_kv_heads, self.head_dim))
        # -> [h, s, head_dim]
        k = k.transpose(0, 1).unsqueeze(0)
        k = self.k_norm(k)

        # q, k = self.rotary_emb(positions, q, k)
        cos, sin = position_embeddings
        q, k = apply_rotary_pos_emb(q, k, cos, sin)

        # [b, h, s, head_dim] ->  [b, s, h, head_dim]
        q = q.permute(0, 2, 1, 3)
        k = k.permute(0, 2, 1, 3)

        attn_output = self.attn(q, k, v, forward_batch=forward_batch)

        # Compatible with triton backend which returns [1, s, h, head_dim]
        if attn_output.dim() == 4 and attn_output.shape[0] == 1:
            attn_output = attn_output.squeeze(0)
            attn_output = attn_output.flatten(-2, -1)
        # [s, h * head_dim]

        output, _ = self.o_proj(attn_output)
        return output

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
        forward_batch: ForwardBatch,
        **kwargs,
    ) -> torch.Tensor:
        if _is_cpu and _is_cpu_amx_available:
            return self.forward_cpu(
                positions, hidden_states, position_embeddings, forward_batch, **kwargs
            )
        return self.forward_native(
            positions, hidden_states, position_embeddings, forward_batch, **kwargs
        )


class Gemma3DecoderLayer(nn.Module):
    def __init__(
        self,
        layer_id: int,
        config: PretrainedConfig,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size
        self.self_attn = Gemma3Attention(
            layer_id=layer_id,
            config=config,
            max_position_embeddings=config.max_position_embeddings,
            quant_config=quant_config,
            prefix=add_prefix("self_attn", prefix),
        )
        self.hidden_size = config.hidden_size
        self.mlp = Gemma3MLP(
            hidden_size=self.hidden_size,
            intermediate_size=config.intermediate_size,
            hidden_activation=config.hidden_activation,
            quant_config=quant_config,
            prefix=add_prefix("mlp", prefix),
        )
        self.input_layernorm = Gemma3RMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )
        self.post_attention_layernorm = Gemma3RMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )
        self.pre_feedforward_layernorm = Gemma3RMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )
        self.post_feedforward_layernorm = Gemma3RMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )
        self.is_sliding = self.self_attn.is_sliding
        self.layer_id = layer_id

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        position_embeddings_global: torch.Tensor,
        position_embeddings_local: torch.Tensor,
        forward_batch: ForwardBatch,
        residual: Optional[torch.Tensor] = None,
        **kwargs,
    ) -> tuple[
        torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]
    ]:
        # Keep the residual live across layers so the add preceding the next
        # RMSNorm is fused by Gemma3RMSNorm. This matches the upstream Gemma3
        # residual layout and is safe to capture in a breakable CUDA graph.
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
            hidden_states, residual = self.input_layernorm(hidden_states, residual)

        # apply global RoPE to non-sliding layer only
        if self.self_attn.is_sliding:
            position_embeddings = position_embeddings_local
        else:
            position_embeddings = position_embeddings_global

        hidden_states = self.self_attn(
            positions=positions,
            hidden_states=hidden_states,
            position_embeddings=position_embeddings,
            forward_batch=forward_batch,
            **kwargs,
        )
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states, residual = self.pre_feedforward_layernorm(
            hidden_states, residual
        )
        hidden_states = self.mlp(hidden_states)
        hidden_states = self.post_feedforward_layernorm(hidden_states)

        outputs = (hidden_states, residual)

        return outputs


class Gemma3RotaryEmbedding(nn.Module):
    def __init__(self, config: Gemma3TextConfig, device=None):
        super().__init__()
        # BC: "rope_type" was originally "type"
        rope_scaling = config.rope_parameters
        if rope_scaling is not None:
            self.rope_type = rope_scaling.get(
                "rope_type", rope_scaling.get("type", "default")
            )

        else:
            self.rope_type = "default"

        if self.rope_type is None:
            self.rope_type = "default"

        self.max_seq_len_cached = config.max_position_embeddings
        self.original_max_seq_len = config.max_position_embeddings

        self.config = config

        if self.rope_type == "default":
            self.rope_init_fn = self.compute_default_rope_parameters
        else:
            self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]

        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self.original_inv_freq = self.inv_freq

    def _dynamic_frequency_update(self, position_ids, device):
        """
        dynamic RoPE layers should recompute `inv_freq` in the following situations:
        1 - growing beyond the cached sequence length (allow scaling)
        2 - the current sequence length is in the original scale (avoid losing precision with small sequences)
        """
        seq_len = torch.max(position_ids) + 1
        if seq_len > self.max_seq_len_cached:  # growth
            inv_freq, self.attention_scaling = self.rope_init_fn(
                self.config, device, seq_len=seq_len
            )
            self.register_buffer(
                "inv_freq", inv_freq, persistent=False
            )  # TODO joao: may break with compilation
            self.max_seq_len_cached = seq_len

        if (
            seq_len < self.original_max_seq_len
            and self.max_seq_len_cached > self.original_max_seq_len
        ):  # reset
            # This .to() is needed if the model has been moved to a device after being initialized (because
            # the buffer is automatically moved, but not the original copy)
            self.original_inv_freq = self.original_inv_freq.to(device)
            self.register_buffer("inv_freq", self.original_inv_freq, persistent=False)
            self.max_seq_len_cached = self.original_max_seq_len

    @staticmethod
    def compute_default_rope_parameters(config, device=None, seq_len=None):
        """Standard RoPE: no scaling, just base frequency."""
        rope_params = config.rope_parameters
        if isinstance(rope_params, dict) and "rope_theta" not in rope_params:
            # Nested per-layer-type format; pick the first available theta
            for v in rope_params.values():
                if isinstance(v, dict) and "rope_theta" in v:
                    base = v["rope_theta"]
                    break
            else:
                base = 10000.0
        else:
            base = rope_params.get("rope_theta", 10000.0) if rope_params else 10000.0
        dim = (
            getattr(config, "head_dim", None)
            or config.hidden_size // config.num_attention_heads
        )
        inv_freq = 1.0 / (
            base
            ** (
                torch.arange(0, dim, 2, dtype=torch.int64).to(
                    device=device, dtype=torch.float
                )
                / dim
            )
        )
        return inv_freq, 1.0

    @torch.no_grad()
    def forward(self, x, position_ids):
        if "dynamic" in self.rope_type:
            self._dynamic_frequency_update(position_ids, device=x.device)

        # Core RoPE block
        inv_freq_expanded = (
            self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
        )
        position_ids_expanded = position_ids[:, None, :].float()
        # Force float32 (see https://github.com/huggingface/transformers/pull/29285)
        device_type = x.device.type
        device_type = (
            device_type
            if isinstance(device_type, str) and device_type != "mps"
            else "cpu"
        )
        with torch.autocast(device_type=device_type, enabled=False):
            freqs = (
                inv_freq_expanded.float().to(x.device) @ position_ids_expanded.float()
            ).transpose(1, 2)
            emb = torch.cat((freqs, freqs), dim=-1)
            cos = emb.cos()
            sin = emb.sin()

        # Advanced RoPE types (e.g. yarn) apply a post-processing scaling factor, equivalent to scaling attention
        cos = cos * self.attention_scaling
        sin = sin * self.attention_scaling

        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)


class Gemma3TextScaledWordEmbedding(nn.Embedding):
    """
    This module overrides nn.Embeddings' forward by multiplying with embeddings scale.
    """

    def __init__(
        self,
        num_embeddings: int,
        embedding_dim: int,
        padding_idx: int,
        embed_scale: Optional[float] = 1.0,
    ):
        super().__init__(num_embeddings, embedding_dim, padding_idx)
        self.embed_scale = embed_scale

    def forward(self, input_ids: torch.Tensor):
        return super().forward(input_ids) * self.embed_scale


class Gemma3TextModel(PreTrainedModel):
    def __init__(
        self,
        config: Gemma3TextConfig,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
        super().__init__(config=config)
        self.config = config
        self.quant_config = quant_config

        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size

        # Gemma3 downcasts the below to float16, causing sqrt(3072)=55.4256 to become 55.5. See https://github.com/huggingface/transformers/pull/29402
        self.embed_tokens = Gemma3TextScaledWordEmbedding(
            config.vocab_size,
            config.hidden_size,
            self.padding_idx,
            embed_scale=self.config.hidden_size**0.5,
        )

        self.norm = Gemma3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

        # In transformers v5, rope_parameters is nested per layer type:
        #   {"sliding_attention": {"rope_type": ..., "rope_theta": 10000},
        #    "full_attention":    {"rope_type": ..., "rope_theta": 1000000}}
        # Flatten into the format Gemma3RotaryEmbedding expects.
        rope_params = config.rope_parameters
        if isinstance(rope_params, dict) and "full_attention" in rope_params:
            global_theta = rope_params["full_attention"].get("rope_theta", 1000000.0)
            local_theta = rope_params["sliding_attention"].get("rope_theta", 10000.0)
        else:
            # v4 flat format fallback
            global_theta = (
                rope_params.get("rope_theta", 10000.0) if rope_params else 10000.0
            )
            local_theta = getattr(config, "rope_local_base_freq", 10000.0)

        global_config = copy.deepcopy(config)
        global_config.rope_parameters = {
            **rope_params["full_attention"],
            "rope_theta": global_theta,
        }
        self.rotary_emb = Gemma3RotaryEmbedding(config=global_config)
        self.gradient_checkpointing = False

        local_config = copy.deepcopy(config)
        local_config.rope_parameters = {
            "rope_type": "default",
            "rope_theta": local_theta,
        }
        self.rotary_emb_local = Gemma3RotaryEmbedding(config=local_config)

        self.layers = make_layers(
            config.num_hidden_layers,
            lambda idx, prefix: Gemma3DecoderLayer(
                layer_id=idx,
                config=config,
                quant_config=quant_config,
                prefix=prefix,
            ),
            prefix=add_prefix("layers", prefix),
        )
        self.norm = Gemma3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.layers_to_capture = []
        self.post_init()

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        forward_batch: ForwardBatch,
        input_embeds: torch.Tensor = None,
        **kwargs,
    ) -> torch.Tensor:
        if input_embeds is None:
            hidden_states = self.embed_tokens(input_ids)
        else:
            hidden_states = input_embeds

        aux_hidden_states = []
        residual = None

        num_layers = len(self.layers)
        if _is_cpu and _is_cpu_amx_available:
            for i, layer in enumerate(self.layers):
                if i in self.layers_to_capture:
                    aux_hidden_states.append(hidden_states)
                hidden_states, residual = layer(
                    positions=positions,
                    position_embeddings_global=None,
                    position_embeddings_local=None,
                    hidden_states=hidden_states,
                    forward_batch=forward_batch,
                    residual=residual,
                    **kwargs,
                )
        else:
            if positions.dim() == 1:
                positions = einops.rearrange(positions, "s -> 1 s")

            position_embeddings_global = self.rotary_emb(hidden_states, positions)
            position_embeddings_local = self.rotary_emb_local(hidden_states, positions)
            for i, layer in enumerate(self.layers):
                if i in self.layers_to_capture:
                    aux_hidden_states.append(hidden_states)
                hidden_states, residual = layer(
                    positions=positions,
                    position_embeddings_global=position_embeddings_global,
                    position_embeddings_local=position_embeddings_local,
                    hidden_states=hidden_states,
                    forward_batch=forward_batch,
                    residual=residual,
                    **kwargs,
                )

        # Capture the output of the last layer if requested.
        # layers_to_capture uses +1 offset (captures input of layer i = output of i-1),
        # so index num_layers means the output of the final layer.
        if num_layers in self.layers_to_capture:
            aux_hidden_states.append(hidden_states)

        hidden_states, _ = self.norm(hidden_states, residual)

        if len(aux_hidden_states) == 0:
            return hidden_states

        return hidden_states, aux_hidden_states


class Gemma3ForCausalLM(PreTrainedModel):
    config_class = Gemma3TextConfig

    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
    _tp_plan = {"lm_head": "colwise_rep"}
    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
    config_class = Gemma3TextConfig
    base_model_prefix = "language_model"

    # BitandBytes specific attributes
    default_bitsandbytes_target_modules = [
        ".gate_proj.",
        ".down_proj.",
        ".up_proj.",
        ".q_proj.",
        ".k_proj.",
        ".v_proj.",
        ".o_proj.",
    ]
    bitsandbytes_stacked_params_mapping = {
        # shard_name, weight_name, index
        "q_proj": ("qkv_proj", 0),
        "k_proj": ("qkv_proj", 1),
        "v_proj": ("qkv_proj", 2),
        "gate_proj": ("gate_up_proj", 0),
        "up_proj": ("gate_up_proj", 1),
    }

    packed_modules_mapping = {
        "qkv_proj": [
            "q_proj",
            "k_proj",
            "v_proj",
        ],
        "gate_up_proj": [
            "gate_proj",
            "up_proj",
        ],
    }

    # LoRA specific attributes
    supported_lora_modules = [
        "qkv_proj",
        "o_proj",
        "gate_up_proj",
        "down_proj",
    ]
    # Gemma does not apply LoRA to the embedding layer.
    embedding_modules = {}
    embedding_padding_modules = []
    supports_lora = True

    def __init__(
        self,
        config: Gemma3TextConfig,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
        super().__init__(config=config)
        self.config = config
        self.quant_config = quant_config
        self.model = Gemma3TextModel(
            config, quant_config, prefix=add_prefix("model", prefix)
        )
        self.logits_processor = LogitsProcessor(config)

        if self.config.tie_word_embeddings:
            self.lm_head = self.model.embed_tokens
        else:
            self.lm_head = ParallelLMHead(
                config.vocab_size,
                config.hidden_size,
                quant_config=quant_config,
                prefix=add_prefix("lm_head", prefix),
            )
        self.capture_aux_hidden_states = False
        self.post_init()

    def get_input_embeddings(self) -> nn.Embedding:
        return self.model.embed_tokens

    def get_attention_sliding_window_size(self):
        return get_attention_sliding_window_size(self.config)

    def dtype(self) -> torch.dtype:
        return next(self.parameters()).dtype

    @torch.no_grad()
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        forward_batch: ForwardBatch,
        input_embeds: torch.Tensor = None,
        **kwargs,
    ) -> LogitsProcessor:
        hidden_states = self.model(
            input_ids, positions, forward_batch, input_embeds, **kwargs
        )

        aux_hidden_states = None
        if self.capture_aux_hidden_states:
            hidden_states, aux_hidden_states = hidden_states

        return self.logits_processor(
            input_ids,
            hidden_states,
            self.model.embed_tokens,
            forward_batch,
            aux_hidden_states,
        )

    @torch.no_grad()
    def forward_split_prefill(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        forward_batch: ForwardBatch,
        split_interval: Tuple[int, int],  # [start, end) 0-based
        input_embeds: torch.Tensor = None,
    ):
        start, end = split_interval
        # embed
        if start == 0:
            if input_embeds is None:
                hidden_states = self.model.embed_tokens(input_ids)
            else:
                hidden_states = input_embeds

            if positions.dim() == 1:
                positions = einops.rearrange(positions, "s -> 1 s")
            position_embeddings_global = self.model.rotary_emb(hidden_states, positions)
            position_embeddings_local = self.model.rotary_emb_local(
                hidden_states, positions
            )

            forward_batch.hidden_states = hidden_states
            forward_batch.model_specific_states = {
                "positions": positions,
                "position_embeddings_global": position_embeddings_global,
                "position_embeddings_local": position_embeddings_local,
            }

        # decoder layer
        for i in range(start, end):
            layer = self.model.layers[i]
            layer_output = layer(
                positions=forward_batch.model_specific_states["positions"],
                position_embeddings_global=forward_batch.model_specific_states[
                    "position_embeddings_global"
                ],
                position_embeddings_local=forward_batch.model_specific_states[
                    "position_embeddings_local"
                ],
                hidden_states=forward_batch.hidden_states,
                forward_batch=forward_batch,
            )
            forward_batch.hidden_states = layer_output[0]

        if end == self.model.config.num_hidden_layers:
            # norm
            forward_batch.hidden_states = self.model.norm(forward_batch.hidden_states)

            # logits process
            result = self.logits_processor(
                input_ids,
                forward_batch.hidden_states,
                self.model.embed_tokens,
                forward_batch,
            )
        else:
            result = None

        return result

    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            ("qkv_proj", "q_proj", "q"),
            ("qkv_proj", "k_proj", "k"),
            ("qkv_proj", "v_proj", "v"),
            ("gate_up_proj", "gate_proj", 0),
            ("gate_up_proj", "up_proj", 1),
        ]
        params_dict = dict(self.named_parameters())
        loaded_params: Set[str] = set()
        for name, loaded_weight in weights:
            remapped_name = maybe_remap_kv_scale_name(name, params_dict)
            if remapped_name is None:
                continue
            if remapped_name != name:
                param = params_dict[remapped_name]
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
                weight_loader(param, loaded_weight)
                loaded_params.add(remapped_name)
                continue

            for param_name, shard_name, shard_id in stacked_params_mapping:
                # if param_name in name:
                # print(f"{param_name} is already in {name}")
                if shard_name not in name:
                    continue
                name = name.replace(shard_name, param_name)
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
                # lm_head is not used in vllm as it is tied with embed_token.
                # To prevent errors, skip loading lm_head.weight.
                if "lm_head.weight" in name:
                    continue
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
                # Remapping the name of FP8 kv-scale.
                name = maybe_remap_kv_scale_name(name, params_dict)
                if name is None:
                    continue

                param = params_dict[name]
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
                weight_loader(param, loaded_weight)
            loaded_params.add(name)
        # unloaded_params = params_dict.keys() - loaded_params
        # if unloaded_params:
        #     logger.warning(
        #         "Some weights are not initialized from checkpoints: %s", unloaded_params
        #     )
        return loaded_params

    def set_eagle3_layers_to_capture(self, layer_ids: Optional[List[int]] = None):
        if layer_ids is None:
            self.capture_aux_hidden_states = True
            num_layers = self.config.num_hidden_layers
            self.model.layers_to_capture = [2, num_layers // 2, num_layers - 3]
        else:
            self.capture_aux_hidden_states = True
            # we plus 1 here because in sglang, for the ith layer, it takes the output
            # of the (i-1)th layer as aux hidden state
            self.model.layers_to_capture = [val + 1 for val in layer_ids]

    def _shard_weight(self, weight: torch.Tensor) -> torch.Tensor:
        """Shard a full embedding/lm_head weight along vocab dim for the current TP rank.

        Gemma3 uses nn.Embedding (unsharded) but the Eagle3 draft model uses
        VocabParallelEmbedding (sharded). This method extracts the correct
        shard so the weights can be shared.
        """
        tp_size = get_parallel().tp_size
        if tp_size <= 1:
            return weight
        tp_rank = get_parallel().tp_rank
        shard_size = (weight.shape[0] + tp_size - 1) // tp_size
        return weight[tp_rank * shard_size : (tp_rank + 1) * shard_size]

    def get_embed(self):
        return self._shard_weight(self.model.embed_tokens.weight)

    def get_embed_and_head(self):
        embed = self._shard_weight(self.model.embed_tokens.weight)
        head = self._shard_weight(self.lm_head.weight)
        return embed, head


class EmbeddingGemmaModel(Gemma3ForCausalLM):
    """EmbeddingGemma's Gemma3 encoder with normalized mean pooling."""

    def __init__(
        self,
        config: Gemma3TextConfig,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
        # Do not initialize Gemma3ForCausalLM's unused LM head.  Keeping the
        # backbone under ``model`` also lets BCG capture only the transformer
        # body and run this pooler as the eager tail.
        PreTrainedModel.__init__(self, config=config)
        self.config = config
        self.quant_config = quant_config
        self.model = Gemma3TextModel(
            config, quant_config, prefix=add_prefix("model", prefix)
        )
        # SentenceTransformers applies mean pooling, then its optional Dense
        # projector modules, then L2 normalization. Keep normalization outside
        # Pooler so this ordering is preserved.
        self.pooler = Pooler(pooling_type=PoolingType.MEAN, normalize=False)
        self.projector = self._build_sentence_transformer_projector(config)
        self.capture_aux_hidden_states = False

    @staticmethod
    def _build_sentence_transformer_projector(config: Gemma3TextConfig):
        """Create the checkpoint's SentenceTransformers Dense tail, if present."""
        model_path = getattr(config, "_name_or_path", "")
        try:
            modules_path = _resolve_local_or_cached_file(model_path, "modules.json")
            with open(modules_path) as f:
                module_specs = json.load(f)

            layers = []
            for spec in module_specs:
                if spec.get("type") != "sentence_transformers.models.Dense":
                    continue
                dense_config_path = _resolve_local_or_cached_file(
                    model_path, f"{spec['path']}/config.json"
                )
                with open(dense_config_path) as f:
                    dense_config = json.load(f)
                if dense_config.get("activation_function") not in (
                    None,
                    "torch.nn.modules.linear.Identity",
                ):
                    raise ValueError(
                        "EmbeddingGemma only supports identity SentenceTransformers "
                        "Dense activations"
                    )
                layers.append(
                    nn.Linear(
                        dense_config["in_features"],
                        dense_config["out_features"],
                        bias=dense_config.get("bias", True),
                    )
                )
            return nn.Sequential(*layers) if layers else None
        except (FileNotFoundError, OSError, ValueError, KeyError, json.JSONDecodeError):
            return None

    @torch.no_grad()
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        forward_batch: ForwardBatch,
        input_embeds: torch.Tensor = None,
        get_embedding: bool = True,
        **kwargs,
    ) -> EmbeddingPoolerOutput:
        assert get_embedding, "EmbeddingGemmaModel is only used for embeddings"
        hidden_states = self.model(
            input_ids, positions, forward_batch, input_embeds, **kwargs
        )
        pooled = self.pooler(hidden_states, forward_batch).embeddings
        if self.projector is not None:
            pooled = self.projector(pooled)
        return EmbeddingPoolerOutput(embeddings=F.normalize(pooled, p=2, dim=-1))

    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
        """Load both native Gemma3 and Sentence Transformers checkpoints.

        The official-style Gemma3 checkpoints prefix backbone parameters with
        ``model.``, while the Sentence Transformers packaging used by
        EmbeddingGemma stores the same backbone at the checkpoint root (and
        includes unrelated ``*_Dense`` modules).  Normalize the latter form
        before delegating to the Gemma3 loader.
        """

        backbone_prefixes = ("embed_tokens.", "layers.", "norm.")
        remapped_weights = (
            (
                f"model.{name}" if name.startswith(backbone_prefixes) else name,
                weight,
            )
            for name, weight in weights
            if name.startswith("model.") or name.startswith(backbone_prefixes)
        )
        loaded_params = super().load_weights(remapped_weights)
        if self.projector is not None:
            model_path = getattr(self.config, "_name_or_path", "")
            modules_path = _resolve_local_or_cached_file(model_path, "modules.json")
            with open(modules_path) as f:
                module_specs = json.load(f)
            dense_specs = [
                spec
                for spec in module_specs
                if spec.get("type") == "sentence_transformers.models.Dense"
            ]
            from safetensors.torch import load_file

            for layer, spec in zip(self.projector, dense_specs):
                weights_path = _resolve_local_or_cached_file(
                    model_path, f"{spec['path']}/model.safetensors"
                )
                weights = load_file(weights_path, device="cpu")
                weight_key = next(
                    key
                    for key in ("weight", "linear.weight", "dense.weight")
                    if key in weights
                )
                layer.weight.data.copy_(weights[weight_key].to(layer.weight.device))
                if layer.bias is not None:
                    layer.bias.data.copy_(
                        weights[weight_key.replace("weight", "bias")].to(
                            layer.bias.device
                        )
                    )
        return loaded_params


EntryClass = [Gemma3ForCausalLM, EmbeddingGemmaModel]
