diff --git a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py
index b4a588df9..a0f80e8a1 100644
--- a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py
+++ b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py
@@ -573,7 +573,9 @@ class HunyuanImage3Pipeline(
         cond_vit_image_embeds = []
         for batch_idx, image in enumerate(cond_vit_images):
             cur_kwargs = {k: v[batch_idx] for k, v in vit_kwargs.items()}
-            image_embed = self.vision_model(image, **cur_kwargs).last_hidden_state
+            # Siglip2VisionTransformer now returns a plain tensor (B, max_patches,
+            # hidden_size) after the vLLM-layers refactor, not an HF-style output.
+            image_embed = self.vision_model(image, **cur_kwargs)
             image_embed = self.vision_aligner(image_embed)
             n, seq_len, dim = image_embed.shape
             image_embed = image_embed.reshape(n * seq_len, dim)
diff --git a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py
index 03e150409..7770bee95 100644
--- a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py
+++ b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py
@@ -23,7 +23,9 @@ from vllm.config.multimodal import BaseDummyOptions
 from vllm.distributed import (
     get_ep_group,
     get_pp_group,
+    get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
+    tensor_model_parallel_all_gather,
 )
 from vllm.inputs import MultiModalDataDict
 from vllm.logger import init_logger
@@ -1451,6 +1453,12 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
 
     prefer_model_sampler = True
 
+    # Siglip2 ViT supports data-parallel encoding (mm_encoder_tp_mode="data"):
+    # weights are replicated and the image batch is sharded across TP ranks.
+    # Without this flag, vLLM silently falls back to "weights" (TP). See
+    # _vit_encode_dp.
+    supports_encoder_tp_data = True
+
     packed_modules_mapping = {
         "qkv_proj": [
             "q_proj",
@@ -1519,7 +1527,14 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
         self.time_embed = TimestepEmbedder(hidden_size=config.hidden_size)
 
         # vision
-        self.vision_model = Siglip2VisionTransformer(config.vit)
+        multimodal_config = vllm_config.model_config.multimodal_config
+        self.multimodal_config = multimodal_config
+        self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
+        self.vision_model = Siglip2VisionTransformer(
+            config.vit,
+            quant_config=quant_config,
+            prefix="vision_model",
+        )
         self.vision_aligner = LightProjector(config.vit_aligner)
 
         # Used to embed timestep information into the input sequence.
@@ -1825,17 +1840,69 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
         """
         Encode pixel_values through ViT encoder (vision_model and vision_aligner).
         """
-        # Handle empty batch
+        # Handle empty batch. Multimodal inputs are identical across TP ranks,
+        # so this returns on all ranks together without collective-op deadlock.
         if pixel_values.shape[0] == 0:
             return None
 
-        vision_output = self.vision_model(
-            pixel_values, attention_mask=vit_attention_mask, spatial_shapes=vit_spatial_shapes
-        )
-        image_embed = vision_output.last_hidden_state
+        if self.use_data_parallel:
+            # ViT weights are replicated (disable_tp on every linear); shard the
+            # image batch across ranks instead of doing redundant full-batch
+            # work on each one.
+            image_embed = self._vit_encode_dp(pixel_values, vit_attention_mask, vit_spatial_shapes)
+        else:
+            image_embed = self.vision_model(
+                pixel_values,
+                attention_mask=vit_attention_mask,
+                spatial_shapes=vit_spatial_shapes,
+            )
         image_embed = self.vision_aligner(image_embed)
         return image_embed
 
+    def _vit_encode_dp(
+        self,
+        pixel_values: torch.Tensor,
+        vit_attention_mask: torch.Tensor,
+        vit_spatial_shapes: torch.Tensor,
+    ) -> torch.Tensor:
+        """Data-parallel ViT: each rank runs a slice of the image batch, then
+        outputs are all-gathered. Mirrors ``run_dp_sharded_vision_model`` but
+        shards the three batched ViT inputs together (pixel_values /
+        attention_mask / spatial_shapes) instead of a single tensor.
+        """
+        num_images = pixel_values.shape[0]
+        world_size = get_tensor_model_parallel_world_size()
+        rank = get_tensor_model_parallel_rank()
+        num_per_rank = (num_images + world_size - 1) // world_size
+        num_padded = num_per_rank * world_size - num_images
+
+        if num_padded > 0:
+            # Pad the batch so every rank gets an equal slice (required for
+            # all_gather). Dummy images use spatial_shapes (1, 1) with a single
+            # valid patch: Siglip2 interpolates position embeddings per image
+            # and rejects size 0, so (0, 0) padding would crash. Their outputs
+            # are dropped after the gather.
+            pad_pixels = pixel_values.new_zeros((num_padded, *pixel_values.shape[1:]))
+            pixel_values = torch.cat([pixel_values, pad_pixels], dim=0)
+
+            pad_mask = vit_attention_mask.new_zeros((num_padded, *vit_attention_mask.shape[1:]))
+            pad_mask[:, 0] = 1
+            vit_attention_mask = torch.cat([vit_attention_mask, pad_mask], dim=0)
+
+            pad_shapes = vit_spatial_shapes.new_ones((num_padded, *vit_spatial_shapes.shape[1:]))
+            vit_spatial_shapes = torch.cat([vit_spatial_shapes, pad_shapes], dim=0)
+
+        start = rank * num_per_rank
+        end = start + num_per_rank
+        image_embed = self.vision_model(
+            pixel_values[start:end],
+            attention_mask=vit_attention_mask[start:end],
+            spatial_shapes=vit_spatial_shapes[start:end],
+        )
+        # max_patches (dim 1) is identical across ranks, so dim-0 gather is safe.
+        image_embed = tensor_model_parallel_all_gather(image_embed.contiguous(), dim=0)
+        return image_embed[:num_images]
+
     def _timestep_encode(
         self,
         timestep: torch.Tensor,
diff --git a/vllm_omni/model_executor/models/hunyuan_image3/siglip2.py b/vllm_omni/model_executor/models/hunyuan_image3/siglip2.py
index 09b6c1d64..91b68422d 100644
--- a/vllm_omni/model_executor/models/hunyuan_image3/siglip2.py
+++ b/vllm_omni/model_executor/models/hunyuan_image3/siglip2.py
@@ -25,18 +25,37 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+"""SigLIP2 Vision Transformer for HunyuanImage3, rewritten in vLLM style.
 
-import warnings
+Key optimizations over the original HuggingFace-style implementation:
+- QKVParallelLinear: fused QKV projection with tensor parallelism support
+- MMEncoderAttention: FlashAttention / xFormers backend
+- ColumnParallelLinear / RowParallelLinear: TP-aware MLP layers
+- Packed sequence processing: eliminates padding waste via cu_seqlens
+- Data parallel support for multi-GPU ViT inference
+"""
+
+from collections.abc import Iterable
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from transformers.activations import ACT2FN
-from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask
-from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling
+from vllm.distributed import divide, get_tensor_model_parallel_world_size
+from vllm.model_executor.layers.activation import get_act_fn
+from vllm.model_executor.layers.attention import MMEncoderAttention
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.models.vision import is_vit_use_data_parallel
 
 
 class Config:
+    """Convert dict config to object with attribute access."""
+
     def __init__(self, config):
         if config is not None:
             for key, value in config.items():
@@ -65,470 +84,336 @@ class Siglip2VisionEmbeddings(nn.Module):
         self.position_embedding_size = int(self.num_patches**0.5)
         self.position_embedding = nn.Embedding(self.num_patches, self.embed_dim)
 
-    @staticmethod
-    def resize_positional_embeddings(
-        positional_embeddings: torch.Tensor,
+    def forward(
+        self,
+        pixel_values: torch.FloatTensor,
         spatial_shapes: torch.LongTensor,
-        max_length: int,
     ) -> torch.Tensor:
         """
-        Resize positional embeddings to image-specific size and pad to a fixed size.
+        Process packed pixel values with per-image position embeddings.
+
         Args:
-            positional_embeddings (`torch.Tensor`):
-                Position embeddings of shape (height, width, embed_dim)
-            spatial_shapes (`torch.LongTensor`):
-                Spatial shapes of shape (batch_size, 2) to resize the positional embeddings to
-            max_length (`int`):
-                Maximum length of the positional embeddings to pad resized positional embeddings to
+            pixel_values: Packed pixel values
+                (total_real_patches, num_channels * patch_size * patch_size)
+            spatial_shapes: Per-image spatial shapes (B, 2) as [(h, w), ...]
+
         Returns:
-            `torch.Tensor`: Embeddings of shape (batch_size, max_length, embed_dim)
+            Packed embeddings (total_real_patches, embed_dim)
         """
-        batch_size = spatial_shapes.shape[0]
-        embed_dim = positional_embeddings.shape[-1]
-        source_dtype = positional_embeddings.dtype
+        target_dtype = self.patch_embedding.weight.dtype
+        patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype))
 
-        resulted_positional_embeddings = torch.empty(
-            (batch_size, max_length, embed_dim),
-            device=positional_embeddings.device,
-            dtype=source_dtype,
+        # Resize position embeddings per image and concatenate (packed)
+        positional_embeddings = self.position_embedding.weight.reshape(
+            self.position_embedding_size, self.position_embedding_size, -1
         )
+        # (H_pe, W_pe, embed_dim) → (1, embed_dim, H_pe, W_pe) for interpolation
+        pe_for_resize = positional_embeddings.permute(2, 0, 1).unsqueeze(0)
+        # Upcast on CPU since antialias is not supported for bfloat16/float16
+        if pe_for_resize.device.type == "cpu":
+            pe_for_resize = pe_for_resize.to(torch.float32)
 
-        # (height, width, embed_dim) -> (1, embed_dim, height, width) for interpolation
-        positional_embeddings = positional_embeddings.permute(2, 0, 1).unsqueeze(0)
-
-        # Upcast to float32 on CPU because antialias is not supported for bfloat16/float16 on CPU
-        if positional_embeddings.device.type == "cpu":
-            positional_embeddings = positional_embeddings.to(torch.float32)
-
-        for i in range(batch_size):
-            # (1, dim, height, width) -> (1, dim, target_height, target_width)
-            height, width = spatial_shapes[i]
-            resized_embeddings = F.interpolate(
-                positional_embeddings,
+        position_embs: list[torch.Tensor] = []
+        for i in range(spatial_shapes.shape[0]):
+            height, width = int(spatial_shapes[i, 0]), int(spatial_shapes[i, 1])
+            resized = F.interpolate(
+                pe_for_resize,
                 size=(height, width),
                 mode="bilinear",
                 align_corners=False,
                 antialias=True,
             )
+            # (1, embed_dim, h, w) → (h*w, embed_dim)
+            resized = resized.reshape(self.embed_dim, height * width).transpose(0, 1)
+            position_embs.append(resized.to(target_dtype))
 
-            # (1, dim, target_height, target_width) -> (target_height * target_width, dim)
-            resized_embeddings = resized_embeddings.reshape(embed_dim, height * width).transpose(0, 1)
-
-            # Cast to original dtype
-            resized_embeddings = resized_embeddings.to(source_dtype)
-
-            resulted_positional_embeddings[i, : height * width] = resized_embeddings
-            resulted_positional_embeddings[i, height * width :] = resized_embeddings[0]
-
-        return resulted_positional_embeddings
-
-    def forward(self, pixel_values: torch.FloatTensor, spatial_shapes: torch.LongTensor) -> torch.Tensor:
-        """
-        Args:
-            pixel_values (`torch.FloatTensor`):
-                Pixel values of shape (batch_size, max_num_patches, num_channels * patch_size * patch_size)
-            spatial_shapes (`List[Tuple[int, int]]`):
-                Spatial shapes of shape (batch_size, 2) to resize the positional embeddings to
-        """
-
-        # Apply patch embeddings to already patchified pixel values
-        target_dtype = self.patch_embedding.weight.dtype
-        patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype))
-
-        # Get positional resized and padded positional embeddings
-        positional_embeddings = self.position_embedding.weight.reshape(
-            self.position_embedding_size, self.position_embedding_size, -1
-        )
-        resized_positional_embeddings = self.resize_positional_embeddings(
-            positional_embeddings, spatial_shapes, max_length=pixel_values.shape[1]
-        )
-
-        # Add positional embeddings to patch embeddings
-        embeddings = patch_embeds + resized_positional_embeddings
-        return embeddings
+        packed_position_embs = torch.cat(position_embs, dim=0)
+        return patch_embeds + packed_position_embs
 
 
 class Siglip2Attention(nn.Module):
-    """Multi-headed attention from 'Attention Is All You Need' paper"""
+    """Multi-headed attention using QKVParallelLinear and MMEncoderAttention."""
 
-    def __init__(self, config):
+    def __init__(
+        self,
+        config,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
         super().__init__()
-        self.config = config
         self.embed_dim = config.hidden_size
         self.num_heads = config.num_attention_heads
         self.head_dim = self.embed_dim // self.num_heads
-        if self.head_dim * self.num_heads != self.embed_dim:
-            raise ValueError(
-                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
-                f" {self.num_heads})."
-            )
         self.scale = self.head_dim**-0.5
-        self.dropout = config.attention_dropout
 
-        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
-        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)
-        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
-        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: torch.Tensor | None = None,
-        output_attentions: bool | None = False,
-    ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        """Input shape: Batch x Time x Channel"""
-
-        batch_size, q_len, _ = hidden_states.size()
-
-        query_states = self.q_proj(hidden_states)
-        key_states = self.k_proj(hidden_states)
-        value_states = self.v_proj(hidden_states)
-
-        query_states = query_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-        key_states = key_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-        value_states = value_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-
-        k_v_seq_len = key_states.shape[-2]
-        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * self.scale
-
-        if attn_weights.size() != (batch_size, self.num_heads, q_len, k_v_seq_len):
-            raise ValueError(
-                f"Attention weights should be of size {(batch_size, self.num_heads, q_len, k_v_seq_len)}, but is"
-                f" {attn_weights.size()}"
-            )
-
-        if attention_mask is not None:
-            if attention_mask.size() != (batch_size, 1, q_len, k_v_seq_len):
-                raise ValueError(
-                    f"Attention mask should be of size {(batch_size, 1, q_len, k_v_seq_len)}, "
-                    f"but is {attention_mask.size()}"
-                )
-            attn_weights = attn_weights + attention_mask
-
-        # upcast attention to fp32
-        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
-        attn_weights = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
-        attn_output = torch.matmul(attn_weights, value_states)
-
-        if attn_output.size() != (batch_size, self.num_heads, q_len, self.head_dim):
-            raise ValueError(
-                f"`attn_output` should be of size {(batch_size, self.num_heads, q_len, self.head_dim)}, but is"
-                f" {attn_output.size()}"
-            )
-
-        attn_output = attn_output.transpose(1, 2).contiguous()
-        attn_output = attn_output.reshape(batch_size, q_len, self.embed_dim)
-
-        attn_output = self.out_proj(attn_output)
-
-        return attn_output, attn_weights
-
-
-class Siglip2SdpaAttention(Siglip2Attention):
-    """
-    Siglip2 attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
-    `Siglip2Attention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt
-    to SDPA API.
-    """
-
-    is_causal = False
-
-    # Adapted from Siglip2Attention.forward and transformers.models.llama.modeling_llama.LlamaSdpaAttention.forward
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: torch.Tensor | None = None,
-        output_attentions: bool | None = False,
-    ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        if output_attentions:
-            # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"`
-            #  once this is implemented.
-            warnings.warn(
-                "Siglip2Model is using Siglip2SdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` "
-                "does not support `output_attentions=True`. Falling back to the manual attention implementation, "
-                "but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. "
-                'This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
-            )
-            return super().forward(
-                hidden_states=hidden_states,
-                attention_mask=attention_mask,
-                output_attentions=output_attentions,
-            )
-
-        batch_size, q_len, _ = hidden_states.size()
-
-        query_states = self.q_proj(hidden_states)
-        key_states = self.k_proj(hidden_states)
-        value_states = self.v_proj(hidden_states)
-
-        query_states = query_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-        key_states = key_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-        value_states = value_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-
-        # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with
-        # custom attn_mask,
-        # Reference: https://github.com/pytorch/pytorch/issues/112577.
-        if query_states.device.type == "cuda" and attention_mask is not None:
-            query_states = query_states.contiguous()
-            key_states = key_states.contiguous()
-            value_states = value_states.contiguous()
-
-        # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an
-        # inline conditional assignment in SDPA to support both torch.compile's dynamic shapes and full graph options.
-        # An inline conditional prevents dynamic shapes from compiling.
-        is_causal = True if self.is_causal and q_len > 1 else False
-
-        attn_output = torch.nn.functional.scaled_dot_product_attention(
-            query_states,
-            key_states,
-            value_states,
-            attn_mask=attention_mask,
-            dropout_p=self.dropout if self.training else 0.0,
-            is_causal=is_causal,
+        use_data_parallel = is_vit_use_data_parallel()
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size=self.embed_dim,
+            head_size=self.head_dim,
+            total_num_heads=self.num_heads,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+            disable_tp=use_data_parallel,
+        )
+        self.out_proj = RowParallelLinear(
+            input_size=self.embed_dim,
+            output_size=self.embed_dim,
+            quant_config=quant_config,
+            prefix=f"{prefix}.out_proj",
+            disable_tp=use_data_parallel,
         )
 
-        attn_output = attn_output.transpose(1, 2).contiguous()
-        attn_output = attn_output.view(batch_size, q_len, self.embed_dim)
+        self.tp_size = 1 if use_data_parallel else get_tensor_model_parallel_world_size()
+        self.num_heads_per_partition = divide(self.num_heads, self.tp_size)
 
-        attn_output = self.out_proj(attn_output)
+        self.attn = MMEncoderAttention(
+            num_heads=self.num_heads_per_partition,
+            head_size=self.head_dim,
+            scale=self.scale,
+            prefix=f"{prefix}.attn",
+        )
 
-        return attn_output, None
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Args:
+            hidden_states: Packed input (total_tokens, embed_dim)
+            cu_seqlens: Cumulative sequence lengths (B+1,)
+        """
+        seq_length = hidden_states.shape[0]
+
+        qkv, _ = self.qkv_proj(hidden_states)
+        q, k, v = qkv.chunk(3, dim=-1)
+
+        q = q.view(seq_length, self.num_heads_per_partition, self.head_dim)
+        k = k.view(seq_length, self.num_heads_per_partition, self.head_dim)
+        v = v.view(seq_length, self.num_heads_per_partition, self.head_dim)
+
+        max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
+        attn_output = self.attn(
+            query=q.unsqueeze(0),
+            key=k.unsqueeze(0),
+            value=v.unsqueeze(0),
+            cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
+        )
+        attn_output = attn_output.reshape(seq_length, self.num_heads_per_partition * self.head_dim)
+
+        attn_output, _ = self.out_proj(attn_output)
+        return attn_output
 
 
 class Siglip2MLP(nn.Module):
-    def __init__(self, config):
+    def __init__(
+        self,
+        config,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
         super().__init__()
-        self.config = config
-        self.activation_fn = ACT2FN[config.hidden_act]
-        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
-        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
+        use_data_parallel = is_vit_use_data_parallel()
+        self.activation_fn = get_act_fn(config.hidden_act)
+        self.fc1 = ColumnParallelLinear(
+            config.hidden_size,
+            config.intermediate_size,
+            quant_config=quant_config,
+            prefix=f"{prefix}.fc1",
+            disable_tp=use_data_parallel,
+        )
+        self.fc2 = RowParallelLinear(
+            config.intermediate_size,
+            config.hidden_size,
+            quant_config=quant_config,
+            prefix=f"{prefix}.fc2",
+            disable_tp=use_data_parallel,
+        )
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        hidden_states = self.fc1(hidden_states)
+        hidden_states, _ = self.fc1(hidden_states)
         hidden_states = self.activation_fn(hidden_states)
-        hidden_states = self.fc2(hidden_states)
+        hidden_states, _ = self.fc2(hidden_states)
         return hidden_states
 
 
 class Siglip2EncoderLayer(nn.Module):
-    def __init__(self, config):
+    def __init__(
+        self,
+        config,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
         super().__init__()
         self.embed_dim = config.hidden_size
-        self.self_attn = Siglip2Attention(config=config)
+        self.self_attn = Siglip2Attention(
+            config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.self_attn",
+        )
         self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
-        self.mlp = Siglip2MLP(config)
+        self.mlp = Siglip2MLP(
+            config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.mlp",
+        )
         self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
 
-    # Ignore copy
     def forward(
         self,
         hidden_states: torch.Tensor,
-        attention_mask: torch.Tensor,
-        output_attentions: bool | None = False,
-    ) -> tuple[torch.FloatTensor]:
+        cu_seqlens: torch.Tensor,
+    ) -> torch.Tensor:
         """
         Args:
-            hidden_states (`torch.FloatTensor`):
-                Input to the layer of shape `(batch, seq_len, embed_dim)`.
-            attention_mask (`torch.FloatTensor`):
-                Attention mask of shape `(batch, 1, q_len, k_v_seq_len)` where padding elements are indicated by very
-                large negative values.
-            output_attentions (`bool`, *optional*, defaults to `False`):
-                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
-                returned tensors for more detail.
+            hidden_states: Packed input (total_tokens, embed_dim)
+            cu_seqlens: Cumulative sequence lengths (B+1,)
         """
         residual = hidden_states
-
         hidden_states = self.layer_norm1(hidden_states)
-        hidden_states, attn_weights = self.self_attn(
-            hidden_states=hidden_states,
-            attention_mask=attention_mask,
-            output_attentions=output_attentions,
-        )
+        hidden_states = self.self_attn(hidden_states, cu_seqlens=cu_seqlens)
         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
-
-        outputs = (hidden_states,)
-
-        if output_attentions:
-            outputs += (attn_weights,)
-
-        return outputs
+        return hidden_states
 
 
 class Siglip2Encoder(nn.Module):
-    """
-    Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
-    [`Siglip2EncoderLayer`].
-    Args:
-        config: Siglip2Config
-    """
+    """Transformer encoder with packed sequence processing."""
 
-    def __init__(self, config):
+    def __init__(
+        self,
+        config,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
         super().__init__()
-        self.config = config
-        self.layers = nn.ModuleList([Siglip2EncoderLayer(config) for _ in range(config.num_hidden_layers)])
-        self.gradient_checkpointing = True
+        self.layers = nn.ModuleList(
+            [
+                Siglip2EncoderLayer(
+                    config,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.layers.{idx}",
+                )
+                for idx in range(config.num_hidden_layers)
+            ]
+        )
 
-    # Ignore copy
     def forward(
         self,
-        inputs_embeds,
-        attention_mask: torch.Tensor | None = None,
-        output_attentions: bool | None = None,
-        output_hidden_states: bool | None = None,
-        return_dict: bool | None = None,
-    ) -> tuple | BaseModelOutput:
-        r"""
-        Args:
-            inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
-                Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation.
-                This is useful if you want more control over how to convert `input_ids` indices into associated vectors
-                than the model's internal embedding lookup matrix.
-            attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
-                Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-                - 1 for tokens that are **not masked**,
-                - 0 for tokens that are **masked**.
-                What are attention masks? See HuggingFace documentation for details.
-            output_attentions (`bool`, *optional*):
-                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
-                returned tensors for more detail.
-            output_hidden_states (`bool`, *optional*):
-                Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
-                for more detail.
-            return_dict (`bool`, *optional*):
-                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+        hidden_states: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+    ) -> torch.Tensor:
         """
-        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
-        )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        encoder_states = () if output_hidden_states else None
-        all_attentions = () if output_attentions else None
-
-        hidden_states = inputs_embeds
-        for layer_index, encoder_layer in enumerate(self.layers):  # len(self.layers): 27
-            if output_hidden_states:
-                encoder_states = encoder_states + (hidden_states,)
-
-            layer_outputs = encoder_layer(
-                hidden_states,
-                attention_mask,
-                output_attentions=output_attentions,
-            )
-
-            hidden_states = layer_outputs[0]
-
-            if output_attentions:
-                all_attentions = all_attentions + (layer_outputs[1],)
-
-        if output_hidden_states:
-            encoder_states = encoder_states + (hidden_states,)
-
-        if not return_dict:
-            return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
-        return BaseModelOutput(last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions)
-
-
-class Siglip2MultiheadAttentionPoolingHead(nn.Module):
-    """Multihead Attention Pooling."""
-
-    def __init__(self, config):
-        super().__init__()
-
-        self.probe = nn.Parameter(torch.randn(1, 1, config.hidden_size))
-        self.attention = torch.nn.MultiheadAttention(config.hidden_size, config.num_attention_heads, batch_first=True)
-        self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-        self.mlp = Siglip2MLP(config)
-        self.num_heads = config.num_attention_heads
-
-    def forward(self, hidden_state: torch.Tensor, attention_mask: torch.Tensor | None = None):
-        batch_size = hidden_state.shape[0]
-        probe = self.probe.repeat(batch_size, 1, 1)
-
-        if attention_mask is not None:
-            target_len, source_len = probe.shape[1], hidden_state.shape[1]
-            attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_state.dtype, target_len)
-            attention_mask = attention_mask.repeat(1, self.num_heads, target_len, 1)
-            attention_mask = attention_mask.reshape(-1, target_len, source_len)
-
-        hidden_state = self.attention(probe, hidden_state, hidden_state, attn_mask=attention_mask)[0]
-
-        residual = hidden_state
-        hidden_state = self.layernorm(hidden_state)
-        hidden_state = residual + self.mlp(hidden_state)
-
-        return hidden_state[:, 0]
+        Args:
+            hidden_states: Packed input (total_tokens, embed_dim)
+            cu_seqlens: Cumulative sequence lengths (B+1,)
+        """
+        for layer in self.layers:
+            hidden_states = layer(hidden_states, cu_seqlens)
+        return hidden_states
 
 
 class Siglip2VisionTransformer(nn.Module):
-    def __init__(self, config):
+    def __init__(
+        self,
+        config,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
         super().__init__()
         config = Config(config)
         self.config = config
-        embed_dim = config.hidden_size
+        self.embed_dim = config.hidden_size
 
         self.embeddings = Siglip2VisionEmbeddings(config)
-        self.encoder = Siglip2Encoder(config)
-        self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
-        self.use_head = True if not hasattr(config, "vision_use_head") else config.vision_use_head
-        if self.use_head:
-            self.head = Siglip2MultiheadAttentionPoolingHead(config)
-        self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
+        self.encoder = Siglip2Encoder(
+            config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.encoder" if prefix else "encoder",
+        )
+        self.post_layernorm = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
 
     def forward(
         self,
         pixel_values: torch.FloatTensor,
         attention_mask: torch.Tensor,
         spatial_shapes: torch.LongTensor,
-        output_attentions: bool | None = None,
-        output_hidden_states: bool | None = None,
-        return_dict: bool | None = None,
-    ) -> tuple | BaseModelOutputWithPooling:
-        r"""
-        Returns:
+    ) -> torch.Tensor:
         """
-        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
+        Args:
+            pixel_values: Batched pixel values
+                (B, max_num_patches, num_channels * patch_size * patch_size)
+            attention_mask: (B, max_num_patches) with 1 for real, 0 for padding
+            spatial_shapes: (B, 2) with (height, width) per image
+
+        Returns:
+            (B, max_num_patches, hidden_size) with zeros at padding positions
+        """
+        batch_size, max_patches, _ = pixel_values.shape
+
+        # Pack: extract real tokens using attention_mask
+        mask_bool = attention_mask.bool()
+        packed_pixels = pixel_values[mask_bool]  # (total_real_patches, C*P*P)
+
+        # Compute cu_seqlens from spatial_shapes (h * w per image)
+        seq_lens = (spatial_shapes[:, 0] * spatial_shapes[:, 1]).to(torch.int32)
+        cu_seqlens = torch.zeros(batch_size + 1, dtype=torch.int32, device=pixel_values.device)
+        cu_seqlens[1:] = seq_lens.cumsum(0)
+
+        # Embeddings (packed)
+        hidden_states = self.embeddings(packed_pixels, spatial_shapes)
+
+        # Encoder (packed)
+        hidden_states = self.encoder(hidden_states, cu_seqlens)
+
+        # Post layernorm
+        hidden_states = self.post_layernorm(hidden_states)
+
+        # Unpack: scatter back to (B, max_patches, hidden_size)
+        output = torch.zeros(
+            batch_size,
+            max_patches,
+            self.embed_dim,
+            dtype=hidden_states.dtype,
+            device=hidden_states.device,
         )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        output[mask_bool] = hidden_states
 
-        hidden_states = self.embeddings(pixel_values, spatial_shapes)
+        return output
 
-        if attention_mask is not None and not self._use_flash_attention_2:
-            # [batch_size, seq_len] -> [batch_size, 1, tgt_seq_len, src_seq_len]
-            encoder_attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype)
-        else:
-            encoder_attention_mask = attention_mask
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+        ]
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
 
-        encoder_outputs = self.encoder(
-            inputs_embeds=hidden_states,
-            attention_mask=encoder_attention_mask,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-        )
-
-        last_hidden_state = encoder_outputs[0]
-        last_hidden_state = self.post_layernorm(last_hidden_state)
-
-        pooler_output = self.head(last_hidden_state, attention_mask) if self.use_head else None
-        if not return_dict:
-            return (last_hidden_state, pooler_output) + encoder_outputs[1:]
-
-        return BaseModelOutputWithPooling(
-            last_hidden_state=last_hidden_state,
-            pooler_output=pooler_output,
-            hidden_states=encoder_outputs.hidden_states,
-            attentions=encoder_outputs.attentions,
-        )
+        for name, loaded_weight in weights:
+            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)
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                if name not in params_dict:
+                    # Skip weights for removed modules (e.g., pooling head)
+                    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 LightProjector(nn.Module):