modeling_sfl_encoder_qwen3.py

from torch import nn
import torch
import math
import warnings
from functools import partial
from .configuration_sfl_encoder import SFLVisionEncoderConfigFromQwen3
from transformers.modeling_utils import PreTrainedModel
from transformers.models.qwen3.modeling_qwen3 import Qwen3Model, Qwen3Attention, rotate_half, Qwen3DecoderLayer
from typing import List, Optional, Tuple, Union
from transformers.modeling_outputs import BaseModelOutputWithPast
from transformers.processing_utils import Unpack
from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
from transformers.cache_utils import Cache, DynamicCache
from transformers.utils import logging, is_flash_attn_greater_or_equal_2_10, is_flash_attn_2_available
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
from torch.nn.init import _calculate_fan_in_and_fan_out
import torch.nn.functional as F
if is_flash_attn_2_available():
    from transformers.modeling_flash_attention_utils import _flash_attention_forward
    from flash_attn import flash_attn_varlen_func

logger = logging.get_logger(__name__)

def _trunc_normal_(tensor, mean, std, a, b):
    # Cut & paste from PyTorch official master until it's in a few official releases - RW
    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
    def norm_cdf(x):
        # Computes standard normal cumulative distribution function
        return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0

    if (mean < a - 2 * std) or (mean > b + 2 * std):
        warnings.warn(
            "mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
            "The distribution of values may be incorrect.",
            stacklevel=2,
        )

    # Values are generated by using a truncated uniform distribution and
    # then using the inverse CDF for the normal distribution.
    # Get upper and lower cdf values
    l = norm_cdf((a - mean) / std)
    u = norm_cdf((b - mean) / std)

    # Uniformly fill tensor with values from [l, u], then translate to
    # [2l-1, 2u-1].
    tensor.uniform_(2 * l - 1, 2 * u - 1)

    # Use inverse cdf transform for normal distribution to get truncated
    # standard normal
    tensor.erfinv_()

    # Transform to proper mean, std
    tensor.mul_(std * math.sqrt(2.0))
    tensor.add_(mean)

    # Clamp to ensure it's in the proper range
    tensor.clamp_(min=a, max=b)


def trunc_normal_tf_(
    tensor: torch.Tensor, mean: float = 0.0, std: float = 1.0, a: float = -2.0, b: float = 2.0
) -> torch.Tensor:
    """Fills the input Tensor with values drawn from a truncated
    normal distribution. The values are effectively drawn from the
    normal distribution :math:`\\mathcal{N}(\text{mean}, \text{std}^2)`
    with values outside :math:`[a, b]` redrawn until they are within
    the bounds. The method used for generating the random values works
    best when :math:`a \\leq \text{mean} \\leq b`.

    NOTE: this 'tf' variant behaves closer to Tensorflow / JAX impl where the
    bounds [a, b] are applied when sampling the normal distribution with mean=0, std=1.0
    and the result is subsequently scaled and shifted by the mean and std args.

    Args:
        tensor: an n-dimensional `torch.Tensor`
        mean: the mean of the normal distribution
        std: the standard deviation of the normal distribution
        a: the minimum cutoff value
        b: the maximum cutoff value
    """
    with torch.no_grad():
        _trunc_normal_(tensor, 0, 1.0, a, b)
        tensor.mul_(std).add_(mean)


def variance_scaling_(tensor, scale=1.0, mode="fan_in", distribution="normal"):
    fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
    if mode == "fan_in":
        denom = fan_in
    elif mode == "fan_out":
        denom = fan_out
    elif mode == "fan_avg":
        denom = (fan_in + fan_out) / 2

    variance = scale / denom

    if distribution == "truncated_normal":
        # constant is stddev of standard normal truncated to (-2, 2)
        trunc_normal_tf_(tensor, std=math.sqrt(variance) / 0.87962566103423978)
    elif distribution == "normal":
        with torch.no_grad():
            tensor.normal_(std=math.sqrt(variance))
    elif distribution == "uniform":
        bound = math.sqrt(3 * variance)
        with torch.no_grad():
            tensor.uniform_(-bound, bound)
    else:
        raise ValueError(f"invalid distribution {distribution}")
    

def lecun_normal_(tensor):
    variance_scaling_(tensor, mode="fan_in", distribution="truncated_normal")


class SFLVisionEncoderEmbeddings(nn.Module):

    def __init__(self, config: SFLVisionEncoderConfigFromQwen3):
        super().__init__()
        self.config = config
        self.embed_dim = config.hidden_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",
        )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = hidden_states.view(
            -1, self.config.num_channels, self.patch_size, self.patch_size
        )
        patch_embeds = self.patch_embedding(hidden_states)  # shape = [*, width, grid, grid]
        # embeddings = patch_embeds.flatten(2).transpose(1, 2)
        embeddings = patch_embeds.view(-1, self.embed_dim)

        return embeddings
    

class VisualRotaryEmbedding(nn.Module):
    def __init__(
        self,
        dim=None,
        max_position_embeddings=2048,
        base=10000,
        device=None,
        scaling_factor=1.0,
        rope_type="default",
        config = None,
    ):
        super().__init__()
        # TODO (joao): remove the `if` below, only used for BC
        self.rope_kwargs = {}
        if config is None:
            logger.warning_once(
                "`Qwen2VLRotaryEmbedding` can now be fully parameterized by passing the model config through the "
                "`config` argument. All other arguments will be removed in v4.46"
            )
            self.rope_kwargs = {
                "rope_type": rope_type,
                "factor": scaling_factor,
                "dim": dim,
                "base": base,
                "max_position_embeddings": max_position_embeddings,
            }
            self.rope_type = rope_type
            self.max_seq_len_cached = max_position_embeddings
            self.original_max_seq_len = max_position_embeddings
        else:
            # BC: "rope_type" was originally "type"
            if config.rope_scaling is not None:
                self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
            else:
                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
        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]

        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device, **self.rope_kwargs)
        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.rope_kwargs
            )
            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
            self.register_buffer("inv_freq", self.original_inv_freq, persistent=False)
            self.max_seq_len_cached = self.original_max_seq_len

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

        inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(2, position_ids.shape[1], -1, 1)
        position_ids_expanded = position_ids[:, :, None, :].float()  # shape (2, bs, 1, positions)
        # 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() @ position_ids_expanded.float()).transpose(2, 3)
            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)


def apply_multimodal_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
    rope_section = [cos.shape[-1] // 2, cos.shape[-1] // 2]
    cos = torch.cat([m[i % 2] for i, m in enumerate(cos.split(rope_section, dim=-1))], dim=-1).unsqueeze(unsqueeze_dim)
    sin = torch.cat([m[i % 2] for i, m in enumerate(sin.split(rope_section, dim=-1))], dim=-1).unsqueeze(unsqueeze_dim)

    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed


class SFLQwen3Attention(Qwen3Attention):
    """Multi-headed attention from 'Attention Is All You Need' paper"""

    # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.is_causal = False

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
        attention_mask: Optional[torch.Tensor],
        past_key_value: Optional[Cache] = None,
        cache_position: Optional[torch.LongTensor] = None,
        cu_seqlens: Optional[torch.Tensor] = None,
        **kwargs: Unpack[FlashAttentionKwargs],
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
        input_shape = hidden_states.shape[:-1]
        hidden_shape = (*input_shape, -1, self.head_dim)

        query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
        key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)

        cos, sin = position_embeddings
        query_states, key_states = apply_multimodal_rotary_pos_emb(query_states, key_states, cos, sin)

        if past_key_value is not None:
            # sin and cos are specific to RoPE models; cache_position needed for the static cache
            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)

        # This is before the transpose
        seq_len = query_states.shape[2]

        # In PEFT, usually we cast the layer norms in float32 for training stability reasons
        # therefore the input hidden states gets silently casted in float32. Hence, we need
        # cast them back in the correct dtype just to be sure everything works as expected.
        # This might slowdown training & inference so it is recommended to not cast the LayerNorms
        # in fp32. (usually our RMSNorm modules handle it correctly)
        target_dtype = None
        if query_states.dtype == torch.float32:
            if torch.is_autocast_enabled():
                target_dtype = torch.get_autocast_gpu_dtype()
            # Handle the case where the model is quantized
            elif hasattr(self.config, "_pre_quantization_dtype"):
                target_dtype = self.config._pre_quantization_dtype
            else:
                target_dtype = next(layer for layer in self.modules() if isinstance(layer, torch.nn.Linear)).weight.dtype

        # FA2 always relies on the value set in the module, so remove it if present in kwargs to avoid passing it twice
        kwargs.pop("is_causal", None)

        # Reashape to the expected shape for Flash Attention
        query_states = query_states.transpose(1, 2).squeeze(0)
        key_states = key_states.transpose(1, 2).squeeze(0)
        value_states = value_states.transpose(1, 2).squeeze(0)

        max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
        attn_output = flash_attn_varlen_func(
            query_states,
            key_states,
            value_states,
            cu_seqlens_q=cu_seqlens,
            cu_seqlens_k=cu_seqlens,
            max_seqlen_q=max_seqlen,
            max_seqlen_k=max_seqlen,
            dropout_p=0.0 if not self.training else self.attention_dropout,
            causal=self.is_causal
        )

        # attn_output = _flash_attention_forward(
        #     query_states,
        #     key_states,
        #     value_states,
        #     attention_mask,
        #     q_len,
        #     position_ids=position_ids,
        #     dropout=dropout_rate,
        #     sliding_window=sliding_window,
        #     is_causal=self.is_causal,
        #     use_top_left_mask=self._flash_attn_uses_top_left_mask,
        # )

        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
        attn_output = self.o_proj(attn_output)
        return attn_output, None
    

class SFLQwen3DecoderLayer(Qwen3DecoderLayer):
    def __init__(self, config: SFLVisionEncoderConfigFromQwen3, layer_idx: int):
        super(SFLQwen3DecoderLayer, self).__init__(config, layer_idx)
        self.self_attn = SFLQwen3Attention(config, layer_idx)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_value: Optional[Cache] = None,
        output_attentions: Optional[bool] = False,
        use_cache: Optional[bool] = False,
        cache_position: Optional[torch.LongTensor] = None,
        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC
        cu_seqlens: Optional[torch.Tensor] = None,
        **kwargs: Unpack[FlashAttentionKwargs],
    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
        residual = hidden_states

        hidden_states = self.input_layernorm(hidden_states)

        # Self Attention
        hidden_states, self_attn_weights = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_value=past_key_value,
            output_attentions=output_attentions,
            use_cache=use_cache,
            cache_position=cache_position,
            position_embeddings=position_embeddings,
            cu_seqlens=cu_seqlens,
            **kwargs,
        )
        hidden_states = residual + hidden_states

        # Fully Connected
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states

        outputs = (hidden_states,)
        if output_attentions:
            outputs += (self_attn_weights,)

        return outputs
    
    
class SFLVisionEncoderFromQwen3Model(Qwen3Model):
    config_class = SFLVisionEncoderConfigFromQwen3
    def __init__(self, config: SFLVisionEncoderConfigFromQwen3):
        super().__init__(config)
        self.layers = nn.ModuleList(
            [SFLQwen3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
        )
        self.rotary_emb = VisualRotaryEmbedding(config=config)
        del self.embed_tokens

    @staticmethod
    def _prepare_4d_causal_attention_mask_with_cache_position(
        attention_mask: torch.Tensor,
        sequence_length: int,
        target_length: int,
        dtype: torch.dtype,
        device: torch.device,
        cache_position: torch.Tensor,
        batch_size: int,
        config: SFLVisionEncoderConfigFromQwen3,
        past_key_values: Cache,
    ):
        """
        Override the original causal mask method to create full attention mask instead.
        Creates a full attention 4D mask of shape `(batch_size, 1, query_length, key_value_length)` 
        from a 2D mask of shape `(batch_size, key_value_length)`.
        
        For vision encoding, we want full attention between all patches, not causal attention.
        """
        if attention_mask is not None and attention_mask.dim() == 4:
            # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.
            full_attention_mask = attention_mask
        else:
            # Create full attention mask (all zeros, meaning attend to all positions)
            # We only mask based on the provided attention_mask for padding
            if attention_mask is not None:
                # Use the provided attention_mask to handle padding
                min_dtype = torch.finfo(dtype).min
                full_attention_mask = torch.zeros(
                    (sequence_length, target_length), dtype=dtype, device=device
                )
                # Expand to 4D
                full_attention_mask = full_attention_mask[None, None, :, :].expand(batch_size, 1, -1, -1)
                
                # Apply padding mask if provided
                full_attention_mask = full_attention_mask.clone()  # copy to contiguous memory for in-place edit
                if attention_mask.shape[-1] > target_length:
                    attention_mask = attention_mask[:, :target_length]
                mask_length = attention_mask.shape[-1]
                padding_mask = attention_mask[:, None, None, :] == 0
                full_attention_mask[:, :, :, :mask_length] = full_attention_mask[:, :, :, :mask_length].masked_fill(
                    padding_mask, min_dtype
                )
            else:
                # No attention mask provided, create all-zeros mask (full attention)
                full_attention_mask = torch.zeros(
                    (batch_size, 1, sequence_length, target_length), dtype=dtype, device=device
                )
        return full_attention_mask

    def get_rope_index(self, grid_sizes, merge_sizes, position_ids):
        position_ids = position_ids.contiguous()
        """
        Generate position indices for RoPE:
        - Vision tokens (vision_mask=True): use 2D position encoding like (0,0), (0,1), (0,2), (1,0), (1,1), (1,2)
        - Text tokens (vision_mask=False): use 1D position encoding like (3,3), (4,4), (5,5)
        """
        batch_size = grid_sizes.shape[0]

        # Vision Part: Generate 2D position indices for vision tokens
        vision_pos_ids = []
        for (t, h, w), merge_size in zip(grid_sizes, merge_sizes):
            # Generate height position indices
            hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w).to(position_ids.device)
            hpos_ids = hpos_ids.reshape(
                h // merge_size,
                merge_size,
                w // merge_size,
                merge_size,
            )
            hpos_ids = hpos_ids.permute(0, 2, 1, 3)
            hpos_ids = hpos_ids.flatten()

            # Generate width position indices  
            wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1).to(position_ids.device)
            wpos_ids = wpos_ids.reshape(
                h // merge_size,
                merge_size,
                w // merge_size,
                merge_size,
            )
            wpos_ids = wpos_ids.permute(0, 2, 1, 3)
            wpos_ids = wpos_ids.flatten()
            
            # Stack height and width to create 2D positions
            vision_pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))

        num_start_idx = 0
        for batch_idx in range(batch_size):
            pos_len = vision_pos_ids[batch_idx].shape[0]
            position_ids[:, 0, num_start_idx: num_start_idx+pos_len] = vision_pos_ids[batch_idx].permute(1, 0)
            num_start_idx += pos_len
        
        return position_ids # shape: (2, batch_size, seq_len)
    
    # def get_rope_index(self, grid_sizes, merge_sizes, position_ids):
    #     position_ids = position_ids.contiguous()
    #     """
    #     Generate polar (r, φ) position indices for RoPE:
    #     - Vision tokens (vision_mask=True): use 2D polar coordinates
    #     r = sqrt((h - c_h)^2 + (w - c_w)^2)
    #     φ = atan2(h - c_h, w - c_w)
    #     - Text tokens (vision_mask=False): keep 1D indices unchanged
    #     """
    #     batch_size = grid_sizes.shape[0]

    #     vision_pos_ids = []
    #     for (t, h, w), merge_size in zip(grid_sizes, merge_sizes):
    #         device = position_ids.device

    #         h_idx = torch.arange(h, device=device)
    #         w_idx = torch.arange(w, device=device)
    #         hh, ww = torch.meshgrid(h_idx, w_idx, indexing='ij')

    #         hh = hh.reshape(h // merge_size, merge_size, w // merge_size, merge_size)
    #         ww = ww.reshape(h // merge_size, merge_size, w // merge_size, merge_size)
    #         hh = hh.permute(0, 2, 1, 3).flatten()
    #         ww = ww.permute(0, 2, 1, 3).flatten()

    #         center_h = (h - 1) / 2
    #         center_w = (w - 1) / 2

    #         rh = hh.float() - center_h
    #         rw = ww.float() - center_w
    #         r = torch.sqrt(rh ** 2 + rw ** 2)
    #         phi = torch.atan2(rh, rw)  # [-pi, pi]

    #         # r_norm = r / r.max()
    #         # phi_norm = (phi + math.pi) / (2 * math.pi)

    #         vision_pos_ids.append(torch.stack([r, phi], dim=-1).repeat(t, 1))

    #     num_start_idx = 0
    #     for batch_idx in range(batch_size):
    #         pos_len = vision_pos_ids[batch_idx].shape[0]
    #         position_ids[:, 0, num_start_idx:num_start_idx + pos_len] = vision_pos_ids[batch_idx].permute(1, 0)
    #         num_start_idx += pos_len

    #     return position_ids  # shape: (2, batch_size, seq_len)


    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Cache] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
        grid_sizes: Optional[torch.Tensor] = None,
        merge_sizes: Optional[torch.Tensor] = None,
        **flash_attn_kwargs: Unpack[FlashAttentionKwargs],
    ) -> BaseModelOutputWithPast:
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        use_cache = use_cache if use_cache is not None else self.config.use_cache

        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

        if self.gradient_checkpointing and self.training and use_cache:
            logger.warning_once(
                "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
            )
            use_cache = False

        # TODO (joao): remove this exception in v4.56 -- it exists for users that try to pass a legacy cache
        if not isinstance(past_key_values, (type(None), Cache)):
            raise ValueError("The `past_key_values` should be either a `Cache` object or `None`.")

        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)

        if use_cache and past_key_values is None:
            past_key_values = DynamicCache()

        if cache_position is None:
            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
            cache_position = torch.arange(
                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
            )

        # the hard coded `2` is for temporal, height and width.
        if position_ids is None:
            position_ids = cache_position.view(1, 1, -1).expand(2, inputs_embeds.shape[0], -1)
        elif position_ids.dim() == 2:
            position_ids = position_ids[None, ...].expand(2, position_ids.shape[0], -1)
        position_ids = self.get_rope_index(grid_sizes, merge_sizes, position_ids)

        causal_mask = None

        hidden_states = inputs_embeds

        # create position embeddings to be shared across the decoder layers
        position_embeddings = self.rotary_emb(hidden_states, position_ids)

        # decoder layers
        all_hidden_states = () if output_hidden_states else None
        all_self_attns = () if output_attentions else None

        # Calculate cumulative sequence lengths for the grid sizes
        cu_seqlens = torch.repeat_interleave(grid_sizes[:, 1] * grid_sizes[:, 2], grid_sizes[:, 0]).cumsum(dim=0, dtype=torch.int32)
        cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)

        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
            if output_hidden_states:
                all_hidden_states += (hidden_states,)

            if self.gradient_checkpointing and self.training:
                layer_outputs = self._gradient_checkpointing_func(
                    partial(decoder_layer.__call__, **flash_attn_kwargs),
                    hidden_states,
                    causal_mask,
                    position_ids,
                    past_key_values,
                    output_attentions,
                    use_cache,
                    cache_position,
                    position_embeddings,
                    cu_seqlens,
                )
            else:
                layer_outputs = decoder_layer(
                    hidden_states,
                    attention_mask=causal_mask,
                    position_ids=position_ids,
                    past_key_value=past_key_values,
                    output_attentions=output_attentions,
                    use_cache=use_cache,
                    cache_position=cache_position,
                    position_embeddings=position_embeddings,
                    cu_seqlens=cu_seqlens,
                    **flash_attn_kwargs,
                )

            hidden_states = layer_outputs[0]

            if output_attentions:
                all_self_attns += (layer_outputs[1],)

        hidden_states = self.norm(hidden_states)

        # add hidden states from the last decoder layer
        if output_hidden_states:
            all_hidden_states += (hidden_states,)

        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=past_key_values if use_cache else None,
            hidden_states=all_hidden_states,
            attentions=all_self_attns,
        )


class SFLVisionEncoderModelFromQwen3(PreTrainedModel):

    config_class = SFLVisionEncoderConfigFromQwen3
    base_model_prefix = "sfl_vision_encoder_qwen3"
    main_input_name = "pixel_values"
    supports_gradient_checkpointing = True
    _no_split_modules = [
        "SFLVisionEncoderEmbeddings",
    ]
    _supports_flash_attn_2 = True
    _supports_sdpa = True

    def __init__(self, config: SFLVisionEncoderConfigFromQwen3):
        super().__init__(config=config)
        self.embeddings = SFLVisionEncoderEmbeddings(config)
        self.encoder = SFLVisionEncoderFromQwen3Model(config)

        self.post_init()


    def forward(self, pixel_values, grid_sizes, merge_sizes=None) -> torch.Tensor:
        hidden_states = self.embeddings(pixel_values)
        encoder_output = self.encoder(
            inputs_embeds=hidden_states[None, ...],
            grid_sizes=grid_sizes,
            merge_sizes=merge_sizes,
            output_hidden_states=True,
        )
        hidden_states = encoder_output.hidden_states
        # hidden_states = torch.cat([
        #     hidden_states[7],
        #     hidden_states[14],
        #     hidden_states[21],
        #     hidden_states[28],
        # ], dim=-1).squeeze(0)
        hidden_states = hidden_states[-1].squeeze(0)

        hidden_states_chunks = hidden_states.split(grid_sizes.prod(dim=1).tolist(), dim=0)
        outputs = []

        for hidden_states, grid_size, merge_size in zip(hidden_states_chunks, grid_sizes, merge_sizes):
            # NOTE: previous implementation, which supports downsampling with any factor
            c = hidden_states.shape[-1]
            hidden_states = hidden_states.view(
                grid_size[0], grid_size[1] // merge_size, grid_size[2] // merge_size, merge_size, merge_size,  c
            ).permute(0, 1, 3, 2, 4, 5)
            hidden_states = hidden_states.reshape(
                grid_size[0], grid_size[1], grid_size[2], c
            ).permute(0, 3, 1, 2)
            hidden_states = torch.nn.functional.interpolate(
                hidden_states,
                size=(grid_size[1] // merge_size, grid_size[2] // merge_size),
                mode='bilinear'
            )
            hidden_states = hidden_states.permute(0, 2, 3, 1).view(-1, c)

            # NOTE: simplified implementation, which only supports downsampling with integer factor
            # NOTE: this implementation is mathematically equivalent to the previous one when merge_size is 1 or 2 but may cause slightly different results
            # hidden_states = hidden_states.view(-1, merge_size * merge_size, hidden_states.size(-1))
            # hidden_states = hidden_states.mean(dim=1)

            outputs.append(hidden_states)
        return torch.cat(outputs, dim=0)
    

    def _init_weights(self, module):
        """Initialize the weights"""
        if isinstance(module, (nn.Linear, nn.Conv2d)):
            lecun_normal_(module.weight)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)