sgl-project--sglang/python/sglang/multimodal_gen/runtime/models/encoders/gemma_3.py

# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo

# SPDX-License-Identifier: Apache-2.0
# Adapted from sglang: python/sglang/srt/models/gemma3_causal.py

import logging
from functools import partial
from typing import Any, Iterable, Optional, Set, Tuple

import torch
from torch import nn

from sglang.multimodal_gen.configs.models.encoders.base import BaseEncoderOutput
from sglang.multimodal_gen.configs.models.encoders.gemma_3 import Gemma3Config
from sglang.multimodal_gen.runtime.distributed import get_tp_world_size
from sglang.multimodal_gen.runtime.layers.activation import GeluAndMul
from sglang.multimodal_gen.runtime.layers.attention import LocalAttention
from sglang.multimodal_gen.runtime.layers.linear import (
    ColumnParallelLinear,
    MergedColumnParallelLinear,
    QKVParallelLinear,
    RowParallelLinear,
)
from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig
from sglang.multimodal_gen.runtime.layers.rotary_embedding import get_rope
from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
    LayerwiseOffloadableModuleMixin,
)
from sglang.multimodal_gen.runtime.utils.common import add_prefix

logger = logging.getLogger(__name__)


def get_attention_sliding_window_size(config):
    return config.sliding_window - 1


class Gemma3RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.zeros(dim))

    def _norm(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        output = self._norm(x.float())
        output = output * (1.0 + self.weight.float())
        return output.type_as(x)

    def extra_repr(self):
        return f"{tuple(self.weight.shape)}, eps={self.eps}"


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

    def forward(self, x):
        x, _ = self.gate_up_proj(x)
        x = self.act_fn(x)
        x, _ = self.down_proj(x)
        return x


def _rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)


class Gemma3Attention(nn.Module):
    def __init__(
        self,
        layer_id: int,
        config: Gemma3Config,
        hidden_size: int,
        num_heads: int,
        num_kv_heads: int,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.layer_id = layer_id
        self.hidden_size = hidden_size
        tp_size = get_tp_world_size()
        self.total_num_heads = num_heads
        assert self.total_num_heads % tp_size == 0
        self.num_heads = self.total_num_heads // tp_size
        self.total_num_kv_heads = num_kv_heads
        if self.total_num_kv_heads >= tp_size:
            assert self.total_num_kv_heads % tp_size == 0
        else:
            assert tp_size % self.total_num_kv_heads == 0
        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)

        self.head_dim = getattr(
            config.text_config, "head_dim", self.hidden_size // self.total_num_heads
        )

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

        self.qkv_proj = QKVParallelLinear(
            hidden_size=hidden_size,
            head_size=self.head_dim,
            total_num_heads=self.total_num_heads,
            total_num_kv_heads=self.total_num_kv_heads,
            bias=config.text_config.attention_bias,
            quant_config=quant_config,
            prefix=f"{prefix}.qkv_proj",
        )

        self.o_proj = RowParallelLinear(
            input_size=self.total_num_heads * self.head_dim,
            output_size=hidden_size,
            bias=config.text_config.attention_bias,
            quant_config=quant_config,
            prefix=f"{prefix}.o_proj",
        )

        layer_types = getattr(config.text_config, "layer_types", None)
        if layer_types:
            self.layer_type = layer_types[layer_id]
            self.is_sliding = self.layer_type == "sliding_attention"
        else:
            # official Gemma3 uses sliding_window_pattern when layer_types is absent
            sliding_window_pattern = getattr(
                config.text_config, "sliding_window_pattern", None
            )
            self.is_sliding = (
                bool((layer_id + 1) % sliding_window_pattern)
                if sliding_window_pattern
                else False
            )
            self.layer_type = "sliding_attention" if self.is_sliding else None

        rope_parameters = getattr(config.text_config, "rope_parameters", None) or {}
        layer_rope_params = {}
        if self.layer_type is not None and isinstance(rope_parameters, dict):
            layer_rope_params = dict(rope_parameters.get(self.layer_type) or {})

        # Initialize the rotary embedding.
        if self.is_sliding:
            # Local attention.
            self.rope_theta = float(
                layer_rope_params.get(
                    "rope_theta",
                    getattr(
                        config.text_config,
                        "rope_local_base_freq",
                        getattr(
                            getattr(config.text_config, "default_theta", {}),
                            "get",
                            lambda *_: 10_000.0,
                        )("local", 10_000.0),
                    ),
                )
            )
            rope_scaling = layer_rope_params or None
            # sliding window
            self.sliding_window = get_attention_sliding_window_size(config.text_config)
            # (left, right) = (window, 0) effectively for causal
            self.window_size = (self.sliding_window, 0)
        else:
            # Global attention.
            self.rope_theta = float(
                layer_rope_params.get(
                    "rope_theta",
                    getattr(
                        config.text_config,
                        "rope_theta",
                        getattr(
                            getattr(config.text_config, "default_theta", {}),
                            "get",
                            lambda *_: 1_000_000.0,
                        )("global", 1_000_000.0),
                    ),
                )
            )
            rope_scaling = layer_rope_params or getattr(
                config.text_config, "rope_scaling", None
            )
            self.sliding_window = None
            self.window_size = (-1, -1)

        self.rotary_pos_emb = get_rope(
            self.head_dim,
            rotary_dim=self.head_dim,
            max_position=config.text_config.max_position_embeddings,
            base=self.rope_theta,
            rope_scaling=rope_scaling,
            is_neox_style=True,
        )

        # Local Attention not support attention mask, we use global attention instead.
        # self.attn = LocalAttention(
        #     self.num_heads,
        #     self.head_dim,
        #     self.num_kv_heads,
        #     softmax_scale=self.scaling,
        #     causal=True,
        #     supported_attention_backends=config._supported_attention_backends,
        #     window_size=self.window_size,
        # )

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

    def _apply_rotary_pos_emb(self, positions, q, k):
        positions_flat = positions.flatten().to(
            device=self.rotary_pos_emb.cos_sin_cache.device, dtype=torch.long
        )
        cos_sin = self.rotary_pos_emb.cos_sin_cache.index_select(0, positions_flat)
        cos, sin = cos_sin.chunk(2, dim=-1)
        # match HF Gemma3: expand half-dim freqs to full head dim before rotate_half
        cos = torch.cat((cos, cos), dim=-1).to(device=q.device, dtype=q.dtype)
        sin = torch.cat((sin, sin), dim=-1).to(device=q.device, dtype=q.dtype)
        cos = cos.unsqueeze(1)
        sin = sin.unsqueeze(1)
        num_tokens = positions_flat.shape[0]

        q = q.reshape(num_tokens, -1, self.head_dim)
        k = k.reshape(num_tokens, -1, self.head_dim)
        q = q * cos + _rotate_half(q) * sin
        k = k * cos + _rotate_half(k) * sin
        return q, k

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        qkv, _ = self.qkv_proj(hidden_states)
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

        batch_size, seq_len, _ = q.shape
        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim)
        k = k.view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
        v = v.view(batch_size, seq_len, self.num_kv_heads, self.head_dim)

        # Apply QK Norm
        q = self.q_norm(q)
        k = self.k_norm(k)

        # Apply RoPE
        q, k = self._apply_rotary_pos_emb(positions, q, k)
        q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim)
        k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)

        # TODO(FlamingoPg): Support LocalAttention
        query = q.transpose(1, 2)
        key = k.transpose(1, 2)
        value = v.transpose(1, 2)

        attn_mask = torch.ones(
            (seq_len, seq_len),
            device=hidden_states.device,
            dtype=torch.bool,
        )
        causal = torch.triu(
            torch.ones(
                (seq_len, seq_len), device=hidden_states.device, dtype=torch.bool
            ),
            diagonal=1,
        )
        attn_mask = attn_mask.masked_fill(causal, False)
        if self.is_sliding and self.sliding_window is not None:
            idx = torch.arange(seq_len, device=hidden_states.device)
            dist = idx[:, None] - idx[None, :]
            too_far = dist > self.sliding_window
            attn_mask = attn_mask.masked_fill(too_far, False)

        attn_mask = attn_mask[None, None, :, :].expand(batch_size, 1, seq_len, seq_len)
        attn_mask = attn_mask & attention_mask.to(torch.bool)[:, None, None, :]

        if query.shape[1] != key.shape[1]:
            num_key_value_groups = query.shape[1] // key.shape[1]
            key = key[:, :, None, :, :].expand(
                batch_size, key.shape[1], num_key_value_groups, seq_len, self.head_dim
            )
            value = value[:, :, None, :, :].expand(
                batch_size,
                value.shape[1],
                num_key_value_groups,
                seq_len,
                self.head_dim,
            )
            key = key.reshape(batch_size, query.shape[1], seq_len, self.head_dim)
            value = value.reshape(batch_size, query.shape[1], seq_len, self.head_dim)

        attn_output = torch.nn.functional.scaled_dot_product_attention(
            query,
            key,
            value,
            attn_mask=attn_mask,
            dropout_p=0.0,
            is_causal=False,
            scale=self.scaling,
        )
        attn_output = attn_output.transpose(1, 2)

        attn_output = attn_output.reshape(
            batch_size, seq_len, self.num_heads * self.head_dim
        )

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


class Gemma3DecoderLayer(nn.Module):
    def __init__(
        self,
        layer_id: int,
        config: Gemma3Config,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.hidden_size = config.text_config.hidden_size
        self.self_attn = Gemma3Attention(
            layer_id=layer_id,
            config=config,
            hidden_size=self.hidden_size,
            num_heads=config.text_config.num_attention_heads,
            num_kv_heads=getattr(
                config.text_config,
                "num_key_value_heads",
                config.text_config.num_attention_heads,
            ),
            quant_config=quant_config,
            prefix=f"{prefix}.self_attn",
        )
        self.mlp = Gemma3MLP(
            hidden_size=self.hidden_size,
            intermediate_size=config.text_config.intermediate_size,
            hidden_act=config.text_config.hidden_activation,
            quant_config=quant_config,
            prefix=f"{prefix}.mlp",
        )
        self.input_layernorm = Gemma3RMSNorm(
            config.text_config.hidden_size, eps=config.text_config.rms_norm_eps
        )
        self.post_attention_layernorm = Gemma3RMSNorm(
            config.text_config.hidden_size, eps=config.text_config.rms_norm_eps
        )
        self.pre_feedforward_layernorm = Gemma3RMSNorm(
            config.text_config.hidden_size, eps=config.text_config.rms_norm_eps
        )
        self.post_feedforward_layernorm = Gemma3RMSNorm(
            config.text_config.hidden_size, eps=config.text_config.rms_norm_eps
        )

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        residual: torch.Tensor | None,
        attention_mask: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        # Self Attention
        # Gemma3 uses "sandwich norm":
        # x = x + norm(attn(norm(x)))
        # So we treat input hidden_states as the residual base.

        if residual is not None:
            hidden_states = hidden_states + residual
            residual = None

        residual_input = hidden_states
        hidden_states = self.input_layernorm(hidden_states)

        hidden_states = self.self_attn(
            positions=positions,
            hidden_states=hidden_states,
            attention_mask=attention_mask,
        )

        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = residual_input + hidden_states

        # MLP
        residual_mlp = hidden_states
        hidden_states = self.pre_feedforward_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = self.post_feedforward_layernorm(hidden_states)
        hidden_states = residual_mlp + hidden_states

        return hidden_states, None


class Gemma3TextScaledWordEmbedding(nn.Embedding):
    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


# --- Siglip Vision Model Implementation ---


class QuickGELU(nn.Module):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x * torch.sigmoid(1.702 * x)


class SiglipVisionEmbeddings(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.embed_dim = config.hidden_size
        self.image_size = config.image_size
        self.patch_size = config.patch_size

        self.patch_embedding = nn.Conv2d(
            in_channels=config.num_channels,
            out_channels=self.embed_dim,
            kernel_size=self.patch_size,
            stride=self.patch_size,
            padding="valid",
        )

        self.num_patches = (self.image_size // self.patch_size) ** 2
        self.num_positions = self.num_patches
        # Use simple Embedding for position embeddings (usually small enough)
        self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
        self.register_buffer(
            "position_ids",
            torch.arange(self.num_positions).expand((1, -1)),
            persistent=False,
        )

    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
        target_dtype = self.patch_embedding.weight.dtype
        patch_embeds = self.patch_embedding(
            pixel_values.to(dtype=target_dtype)
        )  # shape = [*, width, grid, grid]
        embeddings = patch_embeds.flatten(2).transpose(1, 2)
        embeddings = embeddings + self.position_embedding(self.position_ids)

        return embeddings


class SiglipMLP(nn.Module):
    def __init__(
        self,
        config,
        act_layer: type[nn.Module] = QuickGELU,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ):
        super().__init__()
        self.fc1 = ColumnParallelLinear(
            config.hidden_size,
            config.intermediate_size,
            quant_config=quant_config,
            prefix=add_prefix("fc1", prefix),
        )
        self.act = act_layer()
        self.fc2 = RowParallelLinear(
            config.intermediate_size,
            config.hidden_size,
            quant_config=quant_config,
            prefix=add_prefix("fc2", prefix),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x_parallel, _ = self.fc1(x)
        x_parallel = self.act(x_parallel)
        x, _ = self.fc2(x_parallel)
        return x


class SiglipAttention(nn.Module):
    def __init__(
        self,
        hidden_size: int,
        num_heads: int,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_heads = num_heads
        tp_size = get_tp_world_size()
        self.head_dim = hidden_size // num_heads
        self.num_heads_per_partition = num_heads // tp_size
        # Cache the per-rank projection width so forward() does not re-read the
        # global TP size (which is not patched to the folding group at run time).
        self.embed_dim_per_partition = self.num_heads_per_partition * self.head_dim
        self.scaling = self.head_dim**-0.5

        self.qkv_proj = QKVParallelLinear(
            hidden_size=hidden_size,
            head_size=self.head_dim,
            total_num_heads=num_heads,
            total_num_kv_heads=num_heads,
            bias=True,
            quant_config=quant_config,
            prefix=add_prefix("qkv_proj", prefix),
        )

        self.out_proj = RowParallelLinear(
            input_size=hidden_size,
            output_size=hidden_size,
            bias=True,
            quant_config=quant_config,
            prefix=add_prefix("out_proj", prefix),
        )

        self.attn = LocalAttention(
            num_heads=self.num_heads_per_partition,
            head_size=self.head_dim,
            num_kv_heads=self.num_heads_per_partition,
            softmax_scale=self.scaling,
            causal=False,  # Bidirectional for Vision
        )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        qkv, _ = self.qkv_proj(hidden_states)
        q, k, v = qkv.split([self.embed_dim_per_partition] * 3, dim=-1)

        batch_size, seq_len, _ = q.shape
        q = q.view(batch_size, seq_len, self.num_heads_per_partition, self.head_dim)
        k = k.view(batch_size, seq_len, self.num_heads_per_partition, self.head_dim)
        v = v.view(batch_size, seq_len, self.num_heads_per_partition, self.head_dim)

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

        attn_output = attn_output.reshape(
            batch_size, seq_len, self.embed_dim_per_partition
        )

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


class SiglipEncoderLayer(nn.Module):
    def __init__(
        self,
        config,
        act_layer: type[nn.Module] = QuickGELU,
        norm_layer: type[nn.Module] = None,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        if norm_layer is None:
            norm_layer = partial(nn.LayerNorm, eps=config.layer_norm_eps)
        self.layer_norm1 = norm_layer(config.hidden_size)
        self.layer_norm2 = norm_layer(config.hidden_size)
        self.self_attn = SiglipAttention(
            hidden_size=config.hidden_size,
            num_heads=config.num_attention_heads,
            quant_config=quant_config,
            prefix=add_prefix("self_attn", prefix),
        )
        self.mlp = SiglipMLP(
            config,
            act_layer=act_layer,
            quant_config=quant_config,
            prefix=add_prefix("mlp", prefix),
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        residual = hidden_states
        hidden_states = self.layer_norm1(hidden_states)
        hidden_states = self.self_attn(hidden_states)
        hidden_states = residual + hidden_states

        residual = hidden_states
        hidden_states = self.layer_norm2(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states
        return hidden_states


class SiglipEncoder(nn.Module):
    def __init__(
        self,
        config,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.config = config
        num_hidden_layers = config.num_hidden_layers
        norm_layer = partial(nn.LayerNorm, eps=config.layer_norm_eps)
        self.layers = nn.ModuleList(
            [
                SiglipEncoderLayer(
                    config=config,
                    norm_layer=norm_layer,
                    quant_config=quant_config,
                    prefix=add_prefix(f"layers.{layer_idx}", prefix),
                )
                for layer_idx in range(num_hidden_layers)
            ]
        )

    def forward(
        self,
        inputs_embeds: torch.Tensor,
    ) -> torch.Tensor:
        hidden_states = inputs_embeds
        for encoder_layer in self.layers:
            hidden_states = encoder_layer(hidden_states)
        return hidden_states


class SiglipVisionTransformer(nn.Module):
    def __init__(
        self,
        config,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.config = config
        embed_dim = config.hidden_size
        self.embeddings = SiglipVisionEmbeddings(config)
        self.encoder = SiglipEncoder(
            config=config,
            quant_config=quant_config,
            prefix=add_prefix("encoder", prefix),
        )
        self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)

    @property
    def device(self) -> torch.device:
        return self.encoder.layers[0].layer_norm1.weight.device

    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
        hidden_states = self.embeddings(pixel_values.to(self.device))
        last_hidden_state = self.encoder(inputs_embeds=hidden_states)
        last_hidden_state = self.post_layernorm(last_hidden_state)
        return last_hidden_state


class SiglipVisionModel(nn.Module):
    def __init__(
        self,
        config,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ):
        super().__init__()
        self.vision_model = SiglipVisionTransformer(
            config, quant_config, prefix=add_prefix("vision_model", prefix)
        )

    @property
    def device(self) -> torch.device:
        return self.vision_model.device

    def forward(self, pixel_values: torch.Tensor):
        return self.vision_model(pixel_values)


class Gemma3MultiModalProjector(nn.Module):
    """Projector for Gemma3 multimodal."""

    def __init__(self, config: Gemma3Config):
        super().__init__()

        self.mm_input_projection_weight = nn.Parameter(
            torch.zeros(
                config.vision_config.hidden_size, config.text_config.hidden_size
            )
        )

        self.mm_soft_emb_norm = Gemma3RMSNorm(
            config.vision_config.hidden_size, eps=config.vision_config.layer_norm_eps
        )

        self.patches_per_image = int(
            config.vision_config.image_size // config.vision_config.patch_size
        )
        self.tokens_per_side = int(config.mm_tokens_per_image**0.5)
        self.kernel_size = self.patches_per_image // self.tokens_per_side
        self.avg_pool = nn.AvgPool2d(
            kernel_size=self.kernel_size, stride=self.kernel_size
        )

    def forward(self, vision_outputs: torch.Tensor) -> torch.Tensor:
        batch_size, seq_length, hidden_size = vision_outputs.shape

        # Reshape for pooling
        reshaped_vision_outputs = vision_outputs.transpose(1, 2)
        reshaped_vision_outputs = reshaped_vision_outputs.reshape(
            batch_size, hidden_size, self.patches_per_image, self.patches_per_image
        )
        reshaped_vision_outputs = reshaped_vision_outputs.contiguous()

        # Apply pooling
        pooled_vision_outputs = self.avg_pool(reshaped_vision_outputs)
        pooled_vision_outputs = pooled_vision_outputs.flatten(2)
        pooled_vision_outputs = pooled_vision_outputs.transpose(1, 2)

        # Apply normalization
        normed_vision_outputs = self.mm_soft_emb_norm(pooled_vision_outputs)

        # Project to text embedding space
        projected_vision_outputs = torch.matmul(
            normed_vision_outputs, self.mm_input_projection_weight
        )

        return projected_vision_outputs.type_as(vision_outputs)


class Gemma3TextModel(nn.Module):
    def __init__(self, config: Gemma3Config):
        super().__init__()
        self.config = config
        # TODO(yinfan.1024) support text encoding model quant later
        self.quant_config = None

        # Use VocabParallelEmbedding
        from sglang.multimodal_gen.runtime.layers.vocab_parallel_embedding import (
            VocabParallelEmbedding,
        )

        self.vocab_size = config.text_config.vocab_size
        self.embed_tokens = VocabParallelEmbedding(
            self.vocab_size,
            config.text_config.hidden_size,
            org_num_embeddings=config.text_config.vocab_size,
            quant_config=self.quant_config,
        )
        self.embed_scale = config.text_config.hidden_size**0.5

        self.layers = nn.ModuleList(
            [
                Gemma3DecoderLayer(
                    layer_id=i,
                    config=config,
                    quant_config=self.quant_config,
                    prefix=add_prefix(
                        f"layers.{i}", getattr(config.text_config, "prefix", "")
                    ),
                )
                for i in range(config.text_config.num_hidden_layers)
            ]
        )

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

    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        out = self.embed_tokens(input_ids)
        return out * torch.tensor(self.embed_scale, device=out.device, dtype=out.dtype)

    def forward(
        self,
        input_ids: torch.Tensor | None,
        position_ids: torch.Tensor | None = None,
        attention_mask: torch.Tensor | None = None,
        inputs_embeds: torch.Tensor | None = None,
        output_hidden_states: bool | None = None,
        **kwargs,
    ) -> BaseEncoderOutput:
        output_hidden_states = (
            output_hidden_states
            if output_hidden_states is not None
            else self.config.output_hidden_states
        )

        if inputs_embeds is not None:
            hidden_states = inputs_embeds
        else:
            hidden_states = self.get_input_embeddings(input_ids)

        residual = None

        if position_ids is None:
            position_ids = torch.arange(
                0, hidden_states.shape[1], device=hidden_states.device
            ).unsqueeze(0)

        all_hidden_states: tuple[Any, ...] | None = () if output_hidden_states else None

        for layer in self.layers:
            if all_hidden_states is not None:
                all_hidden_states += (hidden_states,)

            hidden_states, residual = layer(
                position_ids,
                hidden_states,
                residual,
                attention_mask=attention_mask,
            )

        hidden_states = self.norm(hidden_states)

        if all_hidden_states is not None:
            all_hidden_states += (hidden_states,)

        output = BaseEncoderOutput(
            last_hidden_state=hidden_states,
            hidden_states=all_hidden_states,
        )

        return output

    def load_weights(self, weights: Any) -> set[str]:
        # Copied from LlamaModel.load_weights but adapted
        params_dict = dict(self.named_parameters())
        loaded_params: set[str] = set()

        def _load_with_shard_id(
            weight_loader, param, loaded_weight: torch.Tensor, shard_id
        ) -> None:
            """Call param.weight_loader with best-effort shard_id normalization.

            Different fused-QKV implementations expect different shard_id types:
            - Some expect strings: "q"/"k"/"v"
            - Some expect integer indices: 0/1/2
            We try the provided shard_id first, then fall back between str/int forms.
            """
            try:
                weight_loader(param, loaded_weight, shard_id)
                return
            except (AssertionError, TypeError):
                pass

            # Fall back between common representations.
            if isinstance(shard_id, str):
                mapping = {"q": 0, "k": 1, "v": 2}
                if shard_id in mapping:
                    weight_loader(param, loaded_weight, mapping[shard_id])
                    return
                if shard_id.isdigit():
                    weight_loader(param, loaded_weight, int(shard_id))
                    return
            elif isinstance(shard_id, int):
                mapping = {0: "q", 1: "k", 2: "v"}
                if shard_id in mapping:
                    weight_loader(param, loaded_weight, mapping[shard_id])
                    return

            # Re-raise with a clearer message.
            raise TypeError(
                f"Unsupported shard_id={shard_id!r} for weight_loader={weight_loader} "
                f"(param={getattr(param, 'name', '<param>')})."
            )

        stacked_params_mapping = getattr(
            getattr(self.config, "arch_config", object()),
            "stacked_params_mapping",
            None,
        )
        if stacked_params_mapping is None:
            stacked_params_mapping = [
                # Fused QKV shards; downstream loaders may want "q/k/v" or 0/1/2.
                (".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),
            ]

        for name, loaded_weight in weights:
            if "rotary_emb.inv_freq" in name:
                continue

            # The config has stacked_params_mapping
            for (
                param_name,
                weight_name,
                shard_id,
            ) in stacked_params_mapping:
                if weight_name not in name:
                    continue
                name = name.replace(weight_name, param_name)

                if name not in params_dict:
                    continue

                param = params_dict[name]
                weight_loader = param.weight_loader
                _load_with_shard_id(weight_loader, param, loaded_weight, shard_id)
                break
            else:
                if name not in params_dict:
                    continue

                param = params_dict[name]
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
                weight_loader(param, loaded_weight)

            loaded_params.add(name)
        return loaded_params


class Gemma3ForConditionalGeneration(nn.Module, LayerwiseOffloadableModuleMixin):
    # transformers 5.6.0 flattened SiglipVisionModel, dropping the
    # `vision_model` intermediate wrapper. Our reimpl keeps it, so remap
    # HF source keys back into our nested namespace when transferring weights.
    layerwise_offload_dit_group_enabled = False
    layer_names = ["language_model.layers"]

    param_names_mapping = {
        r"^(vision_tower\.)(embeddings|encoder|post_layernorm|head)\.": r"\1vision_model.\2.",
    }
    reverse_param_names_mapping = {
        r"^(vision_tower\.)vision_model\.(embeddings|encoder|post_layernorm|head)\.": r"\1\2.",
    }

    def __init__(
        self,
        config: Gemma3Config,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.config = config
        self.quant_config = quant_config
        self.text_config = config.text_config

        # Vision Tower
        self.vision_tower = SiglipVisionModel(
            config=config.vision_config,
            quant_config=quant_config,
            prefix=add_prefix("vision_tower", prefix),
        )

        # Projector
        self.multi_modal_projector = Gemma3MultiModalProjector(config)

        # Text Model
        self.language_model = Gemma3TextModel(config)

    def get_placeholder_mask(
        self,
        input_ids: torch.LongTensor,
        inputs_embeds: torch.FloatTensor,
        image_features: torch.FloatTensor,
    ) -> torch.Tensor:
        image_token_index = int(getattr(self.config, "image_token_index", -1))
        if image_token_index < 0:
            image_token_index = int(getattr(self.text_config, "image_token_index", -1))
        special_image_mask = input_ids == image_token_index
        n_image_tokens = int(special_image_mask.sum().item())
        special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds)
        n_image_features = int(image_features.shape[0] * image_features.shape[1])
        if inputs_embeds[special_image_mask].numel() != image_features.numel():
            raise ValueError(
                f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
            )
        return special_image_mask

    def forward(
        self,
        input_ids: torch.LongTensor,
        pixel_values: torch.FloatTensor | None = None,
        **kwargs,
    ):
        vocab_size = int(self.language_model.vocab_size)
        image_token_index = int(getattr(self.config, "image_token_index", -1))
        if image_token_index < 0:
            image_token_index = int(getattr(self.text_config, "image_token_index", -1))

        if input_ids is not None and image_token_index >= vocab_size:
            special_image_mask = input_ids == image_token_index
            llm_input_ids = input_ids.clone()
            llm_input_ids[special_image_mask] = 0
        else:
            llm_input_ids = input_ids

        inputs_embeds = self.language_model.get_input_embeddings(llm_input_ids)

        if pixel_values is not None:
            if pixel_values.dim() == 5:
                pixel_values = pixel_values.reshape(
                    -1,
                    pixel_values.shape[2],
                    pixel_values.shape[3],
                    pixel_values.shape[4],
                )
            elif pixel_values.dim() == 3:
                pixel_values = pixel_values.unsqueeze(0)
            elif pixel_values.dim() != 4:
                raise ValueError(f"Unexpected pixel_values shape: {pixel_values.shape}")

            vision_outputs = self.vision_tower(pixel_values)
            image_features = self.multi_modal_projector(vision_outputs)
            image_features = image_features.to(
                device=inputs_embeds.device, dtype=inputs_embeds.dtype
            )
            special_image_mask = self.get_placeholder_mask(
                input_ids, inputs_embeds=inputs_embeds, image_features=image_features
            )
            inputs_embeds = inputs_embeds.masked_scatter(
                special_image_mask, image_features
            )

        return self.language_model.forward(
            llm_input_ids, inputs_embeds=inputs_embeds, **kwargs
        )

    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]:
        loaded_params: Set[str] = set()
        params_dict = dict(self.named_parameters())

        def _load_with_shard_id(
            weight_loader, param, loaded_weight: torch.Tensor, shard_id
        ) -> None:
            """Call param.weight_loader with best-effort shard_id normalization.

            Different fused-QKV implementations expect different shard_id types:
            - Some expect strings: "q"/"k"/"v"
            - Some expect integer indices: 0/1/2
            We try the provided shard_id first, then fall back between str/int forms.
            """
            try:
                weight_loader(param, loaded_weight, shard_id)
                return
            except (AssertionError, TypeError):
                pass

            # Fall back between common representations.
            if isinstance(shard_id, str):
                mapping = {"q": 0, "k": 1, "v": 2}
                if shard_id in mapping:
                    weight_loader(param, loaded_weight, mapping[shard_id])
                    return
                if shard_id.isdigit():
                    weight_loader(param, loaded_weight, int(shard_id))
                    return
            elif isinstance(shard_id, int):
                mapping = {0: "q", 1: "k", 2: "v"}
                if shard_id in mapping:
                    weight_loader(param, loaded_weight, mapping[shard_id])
                    return

            raise TypeError(
                f"Unsupported shard_id={shard_id!r} for weight_loader={weight_loader} "
                f"(param={getattr(param, 'name', '<param>')})."
            )

        # Separate weights
        language_model_weights: list[tuple[str, torch.Tensor]] = []
        other_weights: list[tuple[str, torch.Tensor]] = []

        for name, loaded_weight in weights:
            # Handle prefix mapping if needed
            # HF weights might be "model.vision_tower...", "model.language_model..."

            if "vision_tower" in name or "vision_model" in name:
                # Load vision tower weights
                # Map name to local name
                local_name = name
                if "model.vision_tower" in name:
                    local_name = name.replace("model.vision_tower", "vision_tower")
                elif "vision_tower" in name:
                    pass  # already correct prefix if matching self.vision_tower
                elif local_name.startswith("vision_model."):
                    local_name = (
                        "vision_tower.vision_model."
                        + local_name[len("vision_model.") :]
                    )

                # We need to map HF Siglip names to our Siglip implementation
                # Our Siglip: vision_tower.vision_model.encoder.layers...
                # HF Siglip: vision_model.encoder.layers...

                # If loading from Gemma3 checkpoint, it usually has "model.vision_tower.vision_model..."

                if local_name in params_dict:
                    param = params_dict[local_name]
                    weight_loader = getattr(
                        param, "weight_loader", default_weight_loader
                    )
                    weight_loader(param, loaded_weight)
                    loaded_params.add(local_name)
                else:
                    qkv_shard_id = None
                    fused_name = None
                    if ".self_attn.q_proj." in local_name:
                        fused_name = local_name.replace(
                            ".self_attn.q_proj.", ".self_attn.qkv_proj."
                        )
                        qkv_shard_id = "q"
                    elif ".self_attn.k_proj." in local_name:
                        fused_name = local_name.replace(
                            ".self_attn.k_proj.", ".self_attn.qkv_proj."
                        )
                        qkv_shard_id = "k"
                    elif ".self_attn.v_proj." in local_name:
                        fused_name = local_name.replace(
                            ".self_attn.v_proj.", ".self_attn.qkv_proj."
                        )
                        qkv_shard_id = "v"

                    if fused_name is not None and fused_name in params_dict:
                        param = params_dict[fused_name]
                        weight_loader = getattr(
                            param, "weight_loader", default_weight_loader
                        )
                        _load_with_shard_id(
                            weight_loader, param, loaded_weight, qkv_shard_id
                        )
                        loaded_params.add(fused_name)
                        continue

                    if ".self_attn.proj." in local_name:
                        candidate = local_name.replace(
                            ".self_attn.proj.", ".self_attn.out_proj."
                        )
                        if candidate in params_dict:
                            param = params_dict[candidate]
                            weight_loader = getattr(
                                param, "weight_loader", default_weight_loader
                            )
                            weight_loader(param, loaded_weight)
                            loaded_params.add(candidate)
                            continue
                    if ".self_attn.out_proj." in local_name:
                        candidate = local_name.replace(
                            ".self_attn.out_proj.", ".self_attn.proj."
                        )
                        if candidate in params_dict:
                            param = params_dict[candidate]
                            weight_loader = getattr(
                                param, "weight_loader", default_weight_loader
                            )
                            weight_loader(param, loaded_weight)
                            loaded_params.add(candidate)
                            continue

                    # Try to find match
                    suffix = local_name.split("vision_tower.")[-1]
                    # Try adding vision_model
                    candidate = f"vision_tower.vision_model.{suffix}"
                    if candidate in params_dict:
                        param = params_dict[candidate]
                        weight_loader = getattr(
                            param, "weight_loader", default_weight_loader
                        )
                        weight_loader(param, loaded_weight)
                        loaded_params.add(candidate)

            elif "multi_modal_projector" in name:
                local_name = name
                if "model.multi_modal_projector" in name:
                    local_name = name.replace(
                        "model.multi_modal_projector", "multi_modal_projector"
                    )

                if local_name in params_dict:
                    param = params_dict[local_name]
                    weight_loader = getattr(
                        param, "weight_loader", default_weight_loader
                    )
                    weight_loader(param, loaded_weight)
                    loaded_params.add(local_name)

            elif "language_model" in name or "model.language_model" in name:
                # Strip prefix for language model
                # If name is "model.language_model.model.layers.0...", we want "model.layers.0..." for Gemma3ForCausalLM
                # Gemma3ForCausalLM has .model (Gemma3TextModel) and .lm_head

                # HF: model.language_model.model.layers...
                # Ours: language_model.model.layers...

                # We pass (name, weight) to language_model.load_weights
                # We should strip "model.language_model." or "language_model."

                suffix = name
                if "model.language_model." in name:
                    suffix = name.replace("model.language_model.", "")
                elif "language_model." in name:
                    suffix = name.replace("language_model.", "")
                if suffix.startswith("model."):
                    suffix = suffix[len("model.") :]

                language_model_weights.append((suffix, loaded_weight))

            else:
                # Fallback for other weights (maybe direct lm_head if not nested?)
                other_weights.append((name, loaded_weight))

        if language_model_weights:
            lm_loaded = self.language_model.load_weights(language_model_weights)
            loaded_params.update({f"language_model.{n}" for n in lm_loaded})

        return loaded_params

    def get_attention_sliding_window_size(self):
        if self.text_config is not None and hasattr(
            self.text_config, "get_attention_sliding_window_size"
        ):
            return self.text_config.get_attention_sliding_window_size()
        sliding_window = getattr(self.text_config, "sliding_window", None)
        if sliding_window is None:
            sliding_window = getattr(self.config, "sliding_window", None)
        if sliding_window is None:
            return None
        return int(sliding_window) - 1


EntryClass = Gemma3ForConditionalGeneration