chore: import upstream snapshot with attribution

python/sglang/srt/models/gemma4_vision.py

@@ -0,0 +1,599 @@
+# 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.
+# ==============================================================================
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from transformers import Gemma4VisionConfig
+
+from sglang.srt.layers.attention.vision import QKV_BACKEND_IMPL
+from sglang.srt.layers.clippable_linear import (
+    ClippableGateUpParallelLinear,
+    ClippableQKVParallelLinear,
+    ClippableRowParallelLinear,
+)
+from sglang.srt.layers.layernorm import Gemma4RMSNorm
+from sglang.srt.layers.quantization.base_config import QuantizationConfig
+from sglang.srt.runtime_context import get_parallel
+from sglang.srt.utils import add_prefix, get_device_capability, is_cuda, is_hip
+
+# ---------------------------------------------------------------------------
+# 2-D Multidimensional RoPE (matches HF Gemma4RotaryEmbedding for vision)
+# ---------------------------------------------------------------------------
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def _apply_rotary(
+    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> torch.Tensor:
+    return (x * cos) + (_rotate_half(x) * sin)
+
+
+class Gemma4VisionRotaryEmbedding(nn.Module):
+    """Compute 2-D multidimensional RoPE cos/sin for patch positions."""
+
+    def __init__(self, config: Gemma4VisionConfig):
+        super().__init__()
+        self.head_dim = config.head_dim
+        self.rope_theta: float = config.rope_parameters["rope_theta"]
+
+    @torch.no_grad()
+    def forward(
+        self, x: torch.Tensor, patch_positions: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            x: [batch, seq, hidden] – only used for device/dtype.
+            patch_positions: [batch, num_patches, 2] – (x, y) coordinates.
+        Returns:
+            (cos, sin) each of shape [batch, num_patches, head_dim].
+        """
+        ndim = patch_positions.shape[-1]  # 2
+        head_dim_per_dim = self.head_dim // ndim
+
+        all_embs = []
+        for d in range(ndim):
+            dim_inv_freq = 1.0 / (
+                self.rope_theta
+                ** (
+                    torch.arange(
+                        0, head_dim_per_dim, 2, device=x.device, dtype=torch.float
+                    )
+                    / head_dim_per_dim
+                )
+            )
+            dim_inv_freq_expanded = dim_inv_freq[None, :, None].expand(
+                patch_positions.shape[0], -1, 1
+            )
+            dim_positions = patch_positions[:, :, d].float()
+            dim_positions_expanded = dim_positions[:, None, :]
+
+            dim_freqs = (dim_inv_freq_expanded @ dim_positions_expanded).transpose(1, 2)
+            dim_emb = torch.cat((dim_freqs, dim_freqs), dim=-1)
+            all_embs.append(dim_emb)
+
+        emb = torch.cat(all_embs, dim=-1)
+        cos = emb.cos().to(dtype=x.dtype)
+        sin = emb.sin().to(dtype=x.dtype)
+        return cos, sin
+
+
+def _apply_multidimensional_rope(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> torch.Tensor:
+    """Apply 2-D RoPE to x of shape [batch*seq, heads, head_dim].
+
+    cos/sin have shape [batch, seq, head_dim]. We split along head_dim into
+    ndim=2 parts and apply standard rotary to each independently.
+    """
+    ndim = 2
+    chunk_size = x.shape[-1] // ndim
+    x_parts = x.split(chunk_size, dim=-1)
+    cos_parts = cos.split(chunk_size, dim=-1)
+    sin_parts = sin.split(chunk_size, dim=-1)
+    y_parts = [
+        _apply_rotary(x_parts[k], cos_parts[k], sin_parts[k]) for k in range(ndim)
+    ]
+    return torch.cat(y_parts, dim=-1)
+
+
+# ---------------------------------------------------------------------------
+# Vision Attention (TP-sharded, fused QKV)
+# ---------------------------------------------------------------------------
+
+
+class Gemma4VisionAttention(nn.Module):
+    """Multi-head attention for the Gemma 4 vision encoder.
+
+    QKV uses a fused ``ClippableQKVParallelLinear`` for efficient matmul with
+    per-projection clip bounds.  Output projection uses ``ClippableLinear``.
+    """
+
+    def __init__(
+        self,
+        config: Gemma4VisionConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.head_dim = config.head_dim
+
+        tp_size = get_parallel().attn_tp_size
+        self.num_heads_per_partition = config.num_attention_heads // tp_size
+        self.num_kv_heads_per_partition = config.num_key_value_heads // tp_size
+
+        self.qkv = ClippableQKVParallelLinear(
+            hidden_size=config.hidden_size,
+            head_size=config.head_dim,
+            total_num_heads=config.num_attention_heads,
+            total_num_kv_heads=config.num_key_value_heads,
+            bias=config.attention_bias,
+            quant_config=quant_config,
+            prefix=prefix,
+        )
+        self.o_proj = ClippableRowParallelLinear(
+            input_size=config.num_attention_heads * config.head_dim,
+            output_size=config.hidden_size,
+            bias=config.attention_bias,
+            quant_config=quant_config,
+            prefix=add_prefix("o_proj", prefix),
+        )
+
+        self.q_norm = Gemma4RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = Gemma4RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.v_norm = Gemma4RMSNorm(
+            self.head_dim, eps=config.rms_norm_eps, scale_shift=0.0, with_scale=False
+        )
+
+        backend = self._select_backend()
+        self.qkv_backend = QKV_BACKEND_IMPL[backend](
+            head_dim=config.head_dim,
+            num_heads=self.num_heads_per_partition,
+            num_kv_heads=self.num_kv_heads_per_partition,
+            dropout=0.0,
+            flatten_batch=True,
+            softmax_in_single_precision=False,
+            softmax_scale=1.0,
+        )
+
+    @staticmethod
+    def _select_backend() -> str:
+        """Mirror VisionAttention._determine_attention_backend for consistency."""
+        from sglang.srt.runtime_context import get_server_args
+
+        override = get_server_args().mm_attention_backend
+        if override is not None:
+            return override
+        if is_cuda():
+            major, _ = get_device_capability()
+            if major == 9:
+                from sglang.srt.utils import is_blackwell_supported
+
+                if is_blackwell_supported():
+                    return "triton_attn"
+                return "fa3"
+            return "triton_attn"
+        if is_hip():
+            # ROCm: use triton_attn to avoid SDPA flatten_batch issues
+            # with multi-image/video inputs
+            return "triton_attn"
+        return "sdpa"
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        bsz, seq_len, _ = hidden_states.shape
+
+        q, k, v = self.qkv(hidden_states)
+
+        q = q.reshape(bsz * seq_len, self.num_heads_per_partition, self.head_dim)
+        k = k.reshape(bsz * seq_len, self.num_kv_heads_per_partition, self.head_dim)
+        v = v.reshape(bsz * seq_len, self.num_kv_heads_per_partition, self.head_dim)
+
+        q = self.q_norm(q.reshape(-1, self.head_dim)).reshape(q.shape)
+        k = self.k_norm(k.reshape(-1, self.head_dim)).reshape(k.shape)
+        v = self.v_norm(v.reshape(-1, self.head_dim)).reshape(v.shape)
+
+        cos_flat = cos.reshape(bsz * seq_len, 1, self.head_dim)
+        sin_flat = sin.reshape(bsz * seq_len, 1, self.head_dim)
+        q = _apply_multidimensional_rope(q, cos_flat, sin_flat)
+        k = _apply_multidimensional_rope(k, cos_flat, sin_flat)
+
+        if attention_mask is not None:
+            attn_mask_4d = (
+                attention_mask.unsqueeze(-1) * attention_mask.unsqueeze(1)
+            ).unsqueeze(1)
+        else:
+            attn_mask_4d = None
+
+        output = self.qkv_backend.forward(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens=None,
+            bsz=bsz,
+            seq_len=seq_len,
+            attention_mask=attn_mask_4d,
+            softmax_scale=1.0,
+        )
+
+        output = rearrange(output, "(b s) h d -> b s (h d)", b=bsz)
+        output = self.o_proj(output)
+        return output
+
+
+# ---------------------------------------------------------------------------
+# Vision MLP (GatedGELU, TP-sharded)
+# ---------------------------------------------------------------------------
+
+
+class Gemma4VisionMLP(nn.Module):
+    def __init__(
+        self,
+        config: Gemma4VisionConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        if config.hidden_activation != "gelu_pytorch_tanh":
+            raise ValueError(
+                f"Gemma4VisionMLP expects hidden_activation='gelu_pytorch_tanh', "
+                f"got {config.hidden_activation!r}"
+            )
+        self.gate_up = ClippableGateUpParallelLinear(
+            input_size=config.hidden_size,
+            intermediate_size=config.intermediate_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=prefix,
+        )
+        self.down_proj = ClippableRowParallelLinear(
+            input_size=config.intermediate_size,
+            output_size=config.hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=add_prefix("down_proj", prefix),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate, up = self.gate_up(x)
+        x = F.gelu(gate, approximate="tanh") * up
+        x = self.down_proj(x)
+        return x
+
+
+# ---------------------------------------------------------------------------
+# Encoder Layer
+# ---------------------------------------------------------------------------
+
+
+class Gemma4VisionEncoderLayer(nn.Module):
+    def __init__(
+        self,
+        config: Gemma4VisionConfig,
+        layer_idx: int,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.self_attn = Gemma4VisionAttention(
+            config,
+            quant_config=quant_config,
+            prefix=add_prefix("self_attn", prefix),
+        )
+        self.mlp = Gemma4VisionMLP(
+            config,
+            quant_config=quant_config,
+            prefix=add_prefix("mlp", prefix),
+        )
+        eps = config.rms_norm_eps
+        hs = config.hidden_size
+        self.input_layernorm = Gemma4RMSNorm(hs, eps=eps)
+        self.post_attention_layernorm = Gemma4RMSNorm(hs, eps=eps)
+        self.pre_feedforward_layernorm = Gemma4RMSNorm(hs, eps=eps)
+        self.post_feedforward_layernorm = Gemma4RMSNorm(hs, eps=eps)
+
+        self.register_buffer("layer_scalar", torch.ones(()))
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.self_attn(hidden_states, cos, sin, attention_mask)
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = residual + hidden_states
+
+        residual = 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 + hidden_states
+
+        hidden_states = hidden_states * self.layer_scalar
+        return hidden_states
+
+
+# ---------------------------------------------------------------------------
+# Vision Transformer (stack of encoder layers + RoPE)
+# ---------------------------------------------------------------------------
+
+
+class Gemma4VisionTransformer(nn.Module):
+    def __init__(
+        self,
+        config: Gemma4VisionConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.config = config
+        self.rotary_emb = Gemma4VisionRotaryEmbedding(config)
+        self.layers = nn.ModuleList(
+            [
+                Gemma4VisionEncoderLayer(
+                    config,
+                    layer_idx=i,
+                    quant_config=quant_config,
+                    prefix=add_prefix(f"layers.{i}", prefix),
+                )
+                for i in range(config.num_hidden_layers)
+            ]
+        )
+
+    def forward(
+        self,
+        inputs_embeds: torch.Tensor,
+        attention_mask: torch.Tensor,
+        patch_positions: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Args:
+            inputs_embeds: [batch, seq, hidden_size]
+            attention_mask: [batch, seq] — True = valid token
+            patch_positions: [batch, seq, 2]
+        Returns:
+            last_hidden_state: [batch, seq, hidden_size]
+        """
+        cos, sin = self.rotary_emb(inputs_embeds, patch_positions)
+        hidden_states = inputs_embeds
+        for layer in self.layers:
+            hidden_states = layer(hidden_states, cos, sin, attention_mask)
+        return hidden_states
+
+
+# ---------------------------------------------------------------------------
+# Patch Embedder
+# ---------------------------------------------------------------------------
+
+
+class Gemma4VisionPatchEmbedder(nn.Module):
+    def __init__(self, config: Gemma4VisionConfig):
+        super().__init__()
+        self.patch_size = config.patch_size
+        self.hidden_size = config.hidden_size
+        self.position_embedding_size = config.position_embedding_size
+
+        self.input_proj = nn.Linear(
+            3 * self.patch_size**2, self.hidden_size, bias=False
+        )
+        self.position_embedding_table = nn.Parameter(
+            torch.ones(2, self.position_embedding_size, self.hidden_size)
+        )
+
+    def _position_embeddings(
+        self, patch_positions: torch.Tensor, padding_positions: torch.Tensor
+    ) -> torch.Tensor:
+        clamped_positions = patch_positions.clamp(min=0)
+        one_hot = F.one_hot(clamped_positions, num_classes=self.position_embedding_size)
+        one_hot = one_hot.permute(0, 2, 1, 3).to(self.position_embedding_table)
+        position_embeddings = one_hot @ self.position_embedding_table
+        position_embeddings = position_embeddings.sum(dim=1)
+        position_embeddings = torch.where(
+            padding_positions.unsqueeze(-1), 0.0, position_embeddings
+        )
+        return position_embeddings
+
+    def _patch_projection(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        """Project pre-patchified pixels into model space.
+
+        Args:
+            pixel_values: [batch, num_patches, patch_pixels] — already patchified
+                          by the image processor, values in [0, 1].
+        """
+        patches = 2 * (pixel_values - 0.5)
+        return self.input_proj(patches.to(self.input_proj.weight.dtype))
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        pixel_position_ids: torch.Tensor,
+        padding_positions: torch.Tensor,
+    ) -> torch.Tensor:
+        """Compute patch embeddings with positional information.
+
+        Args:
+            pixel_values: [batch, num_patches, patch_pixels] — pre-patchified.
+            pixel_position_ids: [batch, num_patches, 2] — (x, y) positions,
+                                -1 for padding patches.
+            padding_positions: [batch, num_patches] — True for padding patches.
+        """
+        hidden_states = self._patch_projection(pixel_values)
+        position_embeddings = self._position_embeddings(
+            pixel_position_ids, padding_positions
+        )
+        return hidden_states + position_embeddings
+
+
+# ---------------------------------------------------------------------------
+# Pooler
+# ---------------------------------------------------------------------------
+
+
+class Gemma4VisionPooler(nn.Module):
+    def __init__(self, config: Gemma4VisionConfig):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.root_hidden_size = self.hidden_size**0.5
+
+    def _avg_pool_by_positions(
+        self, x: torch.Tensor, patch_positions: torch.Tensor, length: int
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        input_seq_len = x.shape[1]
+        k = int((input_seq_len // length) ** 0.5)
+        k_squared = k**2
+        if k_squared * length != input_seq_len:
+            raise ValueError(
+                f"Cannot pool {x.shape} to {length}: {k=}^2 times {length=} must be {input_seq_len}."
+            )
+        clamped_positions = patch_positions.clamp(min=0)
+        max_x = clamped_positions[..., 0].max(dim=-1, keepdim=True)[0] + 1
+        kernel_idxs = torch.div(clamped_positions, k, rounding_mode="floor")
+        kernel_idxs = kernel_idxs[..., 0] + (max_x // k) * kernel_idxs[..., 1]
+
+        weights = F.one_hot(kernel_idxs.long(), length).float() / k_squared
+        output = weights.transpose(1, 2).to(x.dtype) @ x
+        mask = torch.logical_not((weights == 0).all(dim=1))
+        return output, mask
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        patch_positions: torch.Tensor,
+        padding_positions: torch.Tensor,
+        output_length: Optional[int] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Returns:
+            (pooled_hidden_states, mask) where mask is True for valid tokens.
+        """
+        if output_length is None:
+            raise ValueError("output_length is required for Gemma4VisionPooler")
+        if output_length > hidden_states.shape[1]:
+            raise ValueError(
+                f"Cannot output more soft tokens (requested {output_length}) than there are patches"
+                f" ({hidden_states.shape[1]}). Change the value of `num_soft_tokens` when processing."
+            )
+        length = output_length
+        if isinstance(length, (list, tuple)):
+            length = length[0]
+        if hidden_states.shape[1] == length:
+            mask = padding_positions
+        else:
+            hidden_states, mask = self._avg_pool_by_positions(
+                hidden_states, patch_positions, length
+            )
+        hidden_states = hidden_states * self.root_hidden_size
+        return hidden_states, mask
+
+
+# ---------------------------------------------------------------------------
+# Top-level Vision Encoder (patch_embedder → transformer → pooler)
+# ---------------------------------------------------------------------------
+
+
+class Gemma4VisionEncoder(nn.Module):
+    """Drop-in replacement for HF ``Gemma4VisionEncoder`` with TP support."""
+
+    def __init__(
+        self,
+        config: Gemma4VisionConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.config = config
+        self.patch_size = config.patch_size
+        self.pooling_kernel_size = config.pooling_kernel_size
+
+        self.patch_embedder = Gemma4VisionPatchEmbedder(config)
+        self.encoder = Gemma4VisionTransformer(
+            config,
+            quant_config=quant_config,
+            prefix=add_prefix("encoder", prefix),
+        )
+        self.pooler = Gemma4VisionPooler(config)
+
+        # Post-pooling standardization (normalizes vision tokens before projection)
+        self.standardize = getattr(config, "standardize", False)
+        if self.standardize:
+            self.register_buffer("std_bias", torch.zeros(config.hidden_size))
+            self.register_buffer("std_scale", torch.ones(config.hidden_size))
+
+    @property
+    def device(self) -> torch.device:
+        return self.patch_embedder.input_proj.weight.device
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        pixel_position_ids: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Encode pre-patchified pixel_values into soft tokens.
+
+        Args:
+            pixel_values: [batch, num_patches, patch_pixels] — pre-patchified
+                          by the image processor.
+            pixel_position_ids: [batch, num_patches, 2] — (x, y) positions,
+                                -1 for padding patches.
+
+        Returns:
+            (hidden_states, pooler_mask) — hidden_states [batch, output_len, hidden],
+            pooler_mask [batch, output_len] True = valid.
+        """
+        k2 = self.pooling_kernel_size * self.pooling_kernel_size
+        output_length = pixel_values.shape[-2] // k2
+
+        padding_positions = (pixel_position_ids == -1).all(dim=-1)
+
+        inputs_embeds = self.patch_embedder(
+            pixel_values, pixel_position_ids, padding_positions
+        )
+
+        last_hidden = self.encoder(
+            inputs_embeds=inputs_embeds,
+            attention_mask=~padding_positions,
+            patch_positions=pixel_position_ids,
+        )
+
+        pooled, pooler_mask = self.pooler(
+            last_hidden,
+            pixel_position_ids,
+            padding_positions,
+            output_length=output_length,
+        )
+
+        if self.standardize:
+            pooled = (pooled - self.std_bias) * self.std_scale
+
+        return pooled, pooler_mask