vllm-omni/vllm_omni/diffusion/models/minimax_h3/encoder.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""MiniMax H3 Qwen3-VL layer-50 text/vision encoder.

The encoder is reimplemented on top of vLLM-style tensor-parallel building
blocks (column/row/vocab-parallel linears bound to the MiniMax H3 encoder
process group) instead of the single-GPU ``transformers`` backbone.  This
makes the large Qwen3-VL text model TP-shardable across ``text_encoder_tp_size``
ranks, which removes the rank-0 memory hotspot that dominates no-offload peak
memory (the retained 50-layer encoder is ~51.5 GB in BF16).

The computation contract is unchanged from the HF reference path:

- plain BF16 weights, only the first ``MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER``
  decoder layers are kept (the checkpoint consumes the unnormalized hidden
  state after decoder layer 49);
- cuDNN SDP is enabled during encode;
- the multimodal backbone runs with an all-ones attention mask and
  ``mm_token_type_ids`` derived from the image/video token ids;
- DeepStack visual features are injected at the first
  ``len(deepstack_visual_indexes)`` decoder layers;
- the output is ``hidden_states[50]`` of shape ``[seq, 5120]``.

Tensor parallelism follows vLLM's TP semantics: vocab-parallel embeddings and
row-parallel projections all-reduce so the hidden state stays fully replicated
on every encoder rank, while column-parallel projections (QKV / MLP up) keep a
local shard.  After the final layer every encoder rank therefore holds the
complete ``[seq, 5120]`` hidden state.
"""

from __future__ import annotations

import itertools
import re
from collections.abc import Iterable
from typing import Any

import torch
import torch.nn as nn
import torch.nn.functional as F
from vllm.logger import init_logger
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig

from vllm_omni.diffusion.layers.norm import RMSNorm as DiffusionRMSNorm
from vllm_omni.diffusion.offloader.module_residency import (
    BoundedAllocatorCache,
    PinnedModuleStager,
)

MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER = 50
MINIMAX_H3_QWEN3VL_HIDDEN_DIM = 5120

logger = init_logger(__name__)


def _validate_encoder_quant_config(quant_config: QuantizationConfig | None) -> None:
    if (
        quant_config is not None
        and quant_config.get_name() == "fp8"
        and bool(getattr(quant_config, "is_checkpoint_fp8_serialized", False))
    ):
        raise ValueError("MiniMax-H3 text encoder supports online FP8 only, not serialized FP8 checkpoints")


def _is_quantized_linear(module: LinearBase) -> bool:
    return not isinstance(module.quant_method, UnquantizedLinearMethod)


def _default_weight_loader(
    param: nn.Parameter,
    loaded_weight: torch.Tensor,
    loaded_shard_id: str | None = None,
) -> None:
    del loaded_shard_id
    param.data.copy_(loaded_weight)


def _tp_range(group: Any) -> tuple[int, int]:
    return group.rank_in_group, group.world_size


# ---------------------------------------------------------------------------
# Tensor-parallel primitives bound to the encoder process group.
#
# These mirror the math of vLLM's ``VocabParallelEmbedding``,
# ``ColumnParallelLinear``, ``MergedColumnParallelLinear``, ``QKVParallelLinear``
# and ``RowParallelLinear`` but route their collectives through the encoder
# group instead of the (DiT) tensor-model-parallel group, so the text encoder
# can be sharded independently of the DiT's TP configuration.
# ---------------------------------------------------------------------------


class MiniMaxH3Qwen3VLVocabParallelEmbedding(nn.Module):
    def __init__(self, group: Any, num_embeddings: int, embedding_dim: int, dtype: torch.dtype) -> None:
        super().__init__()
        self.group = group
        self.num_embeddings = num_embeddings
        self.embedding_dim = embedding_dim
        tp_rank, tp_size = _tp_range(group)
        assert num_embeddings % tp_size == 0, (
            f"vocab_size {num_embeddings} must be divisible by text_encoder_tp_size {tp_size}"
        )
        self.num_embeddings_per_partition = num_embeddings // tp_size
        self.weight = nn.Parameter(torch.empty(self.num_embeddings_per_partition, embedding_dim, dtype=dtype))
        self.weight.weight_loader = self.weight_loader  # type: ignore[attr-defined]
        self._tp_rank = tp_rank
        self._tp_size = tp_size

    def forward(self, input_: torch.Tensor) -> torch.Tensor:
        if self._tp_size == 1:
            return F.embedding(input_, self.weight)
        # Mirror vLLM's `VocabParallelEmbedding`: token ids that fall outside
        # this rank's vocab partition are masked to 0 and in-partition ids are
        # shifted to local indices before the local gather.  The out-of-range
        # rows are zeroed so the subsequent all-reduce reconstructs the full
        # embedding for every token.
        start_idx = self._tp_rank * self.num_embeddings_per_partition
        end_idx = start_idx + self.num_embeddings_per_partition
        input_mask = (input_ < start_idx) | (input_ >= end_idx)
        masked_input = input_.clone() - start_idx
        masked_input[input_mask] = 0
        output = F.embedding(masked_input, self.weight)
        output[input_mask, :] = 0.0
        self.group.all_reduce(output)
        return output

    def weight_loader(
        self, param: nn.Parameter, loaded_weight: torch.Tensor, loaded_shard_id: str | None = None
    ) -> None:
        del loaded_shard_id
        shard_size = self.num_embeddings_per_partition
        start_idx = self._tp_rank * shard_size
        param.data.copy_(loaded_weight.narrow(0, start_idx, shard_size))


class MiniMaxH3Qwen3VLMergedColumnParallelLinear(LinearBase):
    """Packed gate/up projection sharded along the output dimension."""

    def __init__(
        self,
        group: Any,
        input_size: int,
        intermediate_size: int,
        dtype: torch.dtype,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        self.group = group
        self.intermediate_size = intermediate_size
        tp_rank, tp_size = _tp_range(group)
        assert intermediate_size % tp_size == 0, (
            f"intermediate_size {intermediate_size} must be divisible by text_encoder_tp_size {tp_size}"
        )
        self.intermediate_size_per_partition = intermediate_size // tp_size
        self.output_dtype = dtype
        super().__init__(
            input_size=input_size,
            output_size=2 * intermediate_size,
            bias=False,
            params_dtype=dtype,
            quant_config=quant_config,
            prefix=prefix,
            disable_tp=True,
        )
        self.quant_method.create_weights(
            self,
            input_size_per_partition=input_size,
            output_partition_sizes=[
                self.intermediate_size_per_partition,
                self.intermediate_size_per_partition,
            ],
            input_size=input_size,
            output_size=2 * intermediate_size,
            params_dtype=dtype,
            weight_loader=self.weight_loader,
        )
        self._tp_rank = tp_rank
        self._tp_size = tp_size

    def forward(self, input_: torch.Tensor) -> torch.Tensor:
        return self.quant_method.apply(self, input_)

    def weight_loader(
        self, param: nn.Parameter, loaded_weight: torch.Tensor, loaded_shard_id: str | None = None
    ) -> None:
        shard_size = self.intermediate_size_per_partition
        start_idx = self._tp_rank * shard_size
        if loaded_shard_id == 1:  # up_proj
            param.data[shard_size : 2 * shard_size].copy_(loaded_weight.narrow(0, start_idx, shard_size))
        else:  # gate_proj
            param.data[0:shard_size].copy_(loaded_weight.narrow(0, start_idx, shard_size))


class MiniMaxH3Qwen3VLQKVParallelLinear(LinearBase):
    """QKV projection with GQA head sharding along the output dimension."""

    def __init__(
        self,
        group: Any,
        hidden_size: int,
        num_heads: int,
        num_kv_heads: int,
        head_dim: int,
        dtype: torch.dtype,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        self.group = group
        self.hidden_size = hidden_size
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads
        self.head_dim = head_dim
        tp_rank, tp_size = _tp_range(group)
        assert num_heads % tp_size == 0, (
            f"num_attention_heads {num_heads} must be divisible by text_encoder_tp_size {tp_size}"
        )
        assert num_kv_heads % tp_size == 0, (
            f"num_key_value_heads {num_kv_heads} must be divisible by text_encoder_tp_size {tp_size}"
        )
        self.local_num_heads = num_heads // tp_size
        self.local_num_kv_heads = num_kv_heads // tp_size
        self.output_dtype = dtype
        q_local = self.local_num_heads * head_dim
        kv_local = self.local_num_kv_heads * head_dim
        output_size = (num_heads + 2 * num_kv_heads) * head_dim
        super().__init__(
            input_size=hidden_size,
            output_size=output_size,
            bias=False,
            params_dtype=dtype,
            quant_config=quant_config,
            prefix=prefix,
            disable_tp=True,
        )
        self.quant_method.create_weights(
            self,
            input_size_per_partition=hidden_size,
            output_partition_sizes=[q_local, kv_local, kv_local],
            input_size=hidden_size,
            output_size=output_size,
            params_dtype=dtype,
            weight_loader=self.weight_loader,
        )
        self._tp_rank = tp_rank
        self._tp_size = tp_size

    def forward(self, input_: torch.Tensor) -> torch.Tensor:
        return self.quant_method.apply(self, input_)

    def weight_loader(
        self, param: nn.Parameter, loaded_weight: torch.Tensor, loaded_shard_id: str | None = None
    ) -> None:
        head_dim = self.head_dim
        q_local = self.local_num_heads * head_dim
        kv_local = self.local_num_kv_heads * head_dim
        if loaded_shard_id == "q":
            start_idx = self._tp_rank * q_local
            param.data[0:q_local].copy_(loaded_weight.narrow(0, start_idx, q_local))
        elif loaded_shard_id == "k":
            start_idx = self._tp_rank * kv_local
            param.data[q_local : q_local + kv_local].copy_(loaded_weight.narrow(0, start_idx, kv_local))
        elif loaded_shard_id == "v":
            start_idx = self._tp_rank * kv_local
            param.data[q_local + kv_local : q_local + 2 * kv_local].copy_(loaded_weight.narrow(0, start_idx, kv_local))
        else:
            raise ValueError(f"unexpected QKV shard id {loaded_shard_id!r}")


# A fused weight is filled from several checkpoint tensors, so a partial load leaves
# uninitialized `torch.empty` rows that no later check would catch. Keyed by the class
# owning the multi-shard `weight_loader`; each entry is the source name reported to the
# user paired with the shard id that loader expects, in projection order.
_FUSED_SOURCE_SHARDS: dict[type[nn.Module], tuple[tuple[str, str | int], ...]] = {
    MiniMaxH3Qwen3VLQKVParallelLinear: (("q", "q"), ("k", "k"), ("v", "v")),
    MiniMaxH3Qwen3VLMergedColumnParallelLinear: (("gate", 0), ("up", 1)),
}


def _fused_source_shards(module: nn.Module) -> tuple[tuple[str, str | int], ...] | None:
    for cls, sources in _FUSED_SOURCE_SHARDS.items():
        if isinstance(module, cls):
            return sources
    return None


def _verify_fused_shards_complete(
    expected: dict[str, tuple[tuple[str, str | int], ...]],
    loaded: dict[str, set[str | int]],
) -> None:
    missing = {
        name: [source for source, shard_id in sources if shard_id not in loaded[name]]
        for name, sources in expected.items()
    }
    missing = {name: sources for name, sources in missing.items() if sources}
    if not missing:
        return
    details = "; ".join(f"{name}: {sources}" for name, sources in sorted(missing.items()))
    raise RuntimeError(f"MiniMax H3 text encoder fused weights are missing source shards: {details}")


class MiniMaxH3Qwen3VLRowParallelLinear(LinearBase):
    def __init__(
        self,
        group: Any,
        input_size: int,
        output_size: int,
        dtype: torch.dtype,
        input_is_parallel: bool = True,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        self.group = group
        self.input_is_parallel = input_is_parallel
        self.output_dtype = dtype
        tp_rank, tp_size = _tp_range(group)
        assert input_size % tp_size == 0, f"input_size {input_size} must be divisible by text_encoder_tp_size {tp_size}"
        self.input_size_per_partition = input_size // tp_size
        super().__init__(
            input_size=input_size,
            output_size=output_size,
            bias=False,
            params_dtype=dtype,
            quant_config=quant_config,
            prefix=prefix,
            disable_tp=True,
        )
        self.quant_method.create_weights(
            self,
            input_size_per_partition=self.input_size_per_partition,
            output_partition_sizes=[output_size],
            input_size=input_size,
            output_size=output_size,
            params_dtype=dtype,
            weight_loader=self.weight_loader,
        )
        self._tp_rank = tp_rank
        self._tp_size = tp_size

    def forward(self, input_: torch.Tensor) -> torch.Tensor:
        if self.input_is_parallel:
            input_parallel = input_
        else:
            split_input = input_.split(self.input_size_per_partition, dim=-1)
            input_parallel = split_input[self._tp_rank].contiguous()
        output_parallel = self.quant_method.apply(self, input_parallel)
        if self._tp_size > 1:
            # Reduce in fp32: the reference path accumulates the full (K=8192)
            # dot product inside a single cuBLAS GEMM before rounding to bf16.
            # A bf16 all-reduce of the per-rank partial GEMMs would round twice
            # and amplify error on this model's large-magnitude activations.
            output_parallel = output_parallel.float()
            self.group.all_reduce(output_parallel)
            output_parallel = output_parallel.to(self.output_dtype)
        return output_parallel

    def weight_loader(
        self, param: nn.Parameter, loaded_weight: torch.Tensor, loaded_shard_id: str | None = None
    ) -> None:
        del loaded_shard_id
        shard_size = self.input_size_per_partition
        start_idx = self._tp_rank * shard_size
        param.data.copy_(loaded_weight.narrow(1, start_idx, shard_size))


# ---------------------------------------------------------------------------
# Shared math ports (identical to the HF reference path).
# ---------------------------------------------------------------------------


class MiniMaxH3Qwen3VLRMSNorm(DiffusionRMSNorm):
    """Qwen3-VL RMSNorm using the common implementation.

    The model-specific name keeps checkpoint weight keys stable, while the
    common RMSNorm dispatches to torch_npu.npu_rms_norm on Ascend. Gamma uses
    the checkpoint/model dtype, but native fallbacks keep the RMS variance
    reduction and scaling in float32.
    """

    def __init__(
        self,
        hidden_size: int,
        eps: float = 1e-6,
        dtype: torch.dtype = torch.float32,
    ) -> None:
        super().__init__(hidden_size, eps=eps, dtype=dtype)


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_pos_emb(
    q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    cos = cos.unsqueeze(1)
    sin = sin.unsqueeze(1)
    q_embed = (q * cos) + (_rotate_half(q) * sin)
    k_embed = (k * cos) + (_rotate_half(k) * sin)
    return q_embed, k_embed


def _apply_interleaved_mrope(freqs: torch.Tensor, mrope_section: list[int]) -> torch.Tensor:
    """Reorganize chunked [TTT...HHH...WWW] into interleaved [THWTHW...] layout."""
    freqs_t = freqs[0]
    for dim, offset in enumerate((1, 2), start=1):
        length = mrope_section[dim] * 3
        idx = slice(offset, length, 3)
        freqs_t[..., idx] = freqs[dim, ..., idx]
    return freqs_t


class MiniMaxH3Qwen3VLTextRotaryEmbedding(nn.Module):
    """Qwen3-VL text M-RoPE (interleaved sections). Same math as the reference."""

    def __init__(self, config: Any) -> None:
        super().__init__()
        self.config = config
        # ``dim`` and ``base`` are plain attributes so the inverse frequencies
        # are always recomputed in float32 (module-wide dtype casts must not
        # degrade the RoPE cache precision).
        self.dim = int(getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads)
        self.base = float(config.rope_parameters["rope_theta"])
        self.mrope_section = config.rope_parameters.get("mrope_section", [24, 20, 20])

    def _get_inv_freq(self, device: torch.device) -> torch.Tensor:
        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float() / self.dim))
        return inv_freq.to(device)

    def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        # position_ids: (3, 1, seq)
        inv_freq = self._get_inv_freq(position_ids.device)
        inv_freq_expanded = inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1)
        position_ids_expanded = position_ids[:, :, None, :].float()
        freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3)
        freqs = _apply_interleaved_mrope(freqs, self.mrope_section)
        emb = torch.cat((freqs, freqs), dim=-1)
        cos = emb.cos()
        sin = emb.sin()
        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)


# ---------------------------------------------------------------------------
# Qwen3-VL vision tower (replicated on every encoder rank; ~1.2 GB BF16).
# ---------------------------------------------------------------------------


class MiniMaxH3Qwen3VLVisionPatchEmbed(nn.Module):
    def __init__(self, config: Any) -> None:
        super().__init__()
        self.patch_size = config.patch_size
        self.temporal_patch_size = config.temporal_patch_size
        self.in_channels = config.in_channels
        self.embed_dim = config.hidden_size
        kernel_size = [self.temporal_patch_size, self.patch_size, self.patch_size]
        self.proj = nn.Conv3d(self.in_channels, self.embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=True)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        target_dtype = self.proj.weight.dtype
        hidden_states = hidden_states.view(
            -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size
        )
        hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim)
        return hidden_states


class MiniMaxH3Qwen3VLVisionRotaryEmbedding(nn.Module):
    def __init__(self, dim: int, theta: float = 10000.0) -> None:
        super().__init__()
        self.dim = dim
        self.theta = theta
        # Plain attribute (not a registered buffer): the reference keeps this
        # table in float32 and module-wide dtype casts must not touch it.
        self._inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))

    def forward(self, seqlen: int) -> torch.Tensor:
        inv_freq = self._inv_freq.to(device=self._inv_freq.device)
        seq = torch.arange(seqlen, device=inv_freq.device, dtype=torch.float)
        freqs = torch.outer(seq, inv_freq)
        return freqs


class MiniMaxH3Qwen3VLVisionPatchMerger(nn.Module):
    def __init__(self, config: Any, use_postshuffle_norm: bool = False) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size * (config.spatial_merge_size**2)
        self.use_postshuffle_norm = use_postshuffle_norm
        self.norm = nn.LayerNorm(self.hidden_size if use_postshuffle_norm else config.hidden_size, eps=1e-6)
        self.linear_fc1 = nn.Linear(self.hidden_size, self.hidden_size)
        self.act_fn = nn.GELU()
        self.linear_fc2 = nn.Linear(self.hidden_size, config.out_hidden_size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.norm(x.view(-1, self.hidden_size) if self.use_postshuffle_norm else x).view(-1, self.hidden_size)
        x = self.linear_fc2(self.act_fn(self.linear_fc1(x)))
        return x


class MiniMaxH3Qwen3VLVisionMLP(nn.Module):
    def __init__(self, config: Any) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size
        self.linear_fc1 = nn.Linear(self.hidden_size, self.intermediate_size, bias=True)
        self.linear_fc2 = nn.Linear(self.intermediate_size, self.hidden_size, bias=True)
        # The reference vision MLP uses the tanh GELU approximation
        # (``gelu_pytorch_tanh``), which differs from the exact erf GELU.
        self.act_fn = nn.GELU(approximate="tanh")

    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
        return self.linear_fc2(self.act_fn(self.linear_fc1(hidden_state)))


def _apply_rotary_pos_emb_vision(
    q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    orig_q_dtype = q.dtype
    orig_k_dtype = k.dtype
    q, k = q.float(), k.float()
    cos, sin = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float()
    q_embed = (q * cos) + (_rotate_half(q) * sin)
    k_embed = (k * cos) + (_rotate_half(k) * sin)
    q_embed = q_embed.to(orig_q_dtype)
    k_embed = k_embed.to(orig_k_dtype)
    return q_embed, k_embed


class MiniMaxH3Qwen3VLVisionAttention(nn.Module):
    def __init__(self, config: Any) -> None:
        super().__init__()
        self.dim = config.hidden_size
        self.num_heads = config.num_heads
        self.head_dim = self.dim // self.num_heads
        self.num_key_value_groups = 1
        self.qkv = nn.Linear(self.dim, self.dim * 3, bias=True)
        self.proj = nn.Linear(self.dim, self.dim)
        self.scaling = self.head_dim**-0.5

    def forward(
        self,
        hidden_states: torch.Tensor,
        cu_seqlens: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
    ) -> torch.Tensor:
        seq_length = hidden_states.shape[0]
        query_states, key_states, value_states = (
            self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
        )
        cos, sin = position_embeddings
        query_states, key_states = _apply_rotary_pos_emb_vision(query_states, key_states, cos, sin)

        query_states = query_states.transpose(0, 1).unsqueeze(0)
        key_states = key_states.transpose(0, 1).unsqueeze(0)
        value_states = value_states.transpose(0, 1).unsqueeze(0)

        lengths = cu_seqlens[1:] - cu_seqlens[:-1]
        splits = [torch.split(tensor, lengths.tolist(), dim=2) for tensor in (query_states, key_states, value_states)]
        attn_outputs = [
            F.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=False) for q, k, v in zip(*splits)
        ]
        attn_output = torch.cat(attn_outputs, dim=2)
        # The SDPA kernel may return a non-contiguous (1, num_heads, seq_len, head_dim)
        # tensor whose memory layout is seq-major; transpose before flattening exactly
        # like the reference `sdpa_attention_forward` to avoid scrambling heads.
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.reshape(seq_length, -1).contiguous()
        attn_output = self.proj(attn_output)
        return attn_output


class MiniMaxH3Qwen3VLVisionBlock(nn.Module):
    def __init__(self, config: Any) -> None:
        super().__init__()
        self.norm1 = nn.LayerNorm(config.hidden_size, eps=1e-6)
        self.norm2 = nn.LayerNorm(config.hidden_size, eps=1e-6)
        self.attn = MiniMaxH3Qwen3VLVisionAttention(config=config)
        self.mlp = MiniMaxH3Qwen3VLVisionMLP(config=config)

    def forward(
        self,
        hidden_states: torch.Tensor,
        cu_seqlens: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
    ) -> torch.Tensor:
        hidden_states = hidden_states + self.attn(
            self.norm1(hidden_states),
            cu_seqlens=cu_seqlens,
            position_embeddings=position_embeddings,
        )
        hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
        return hidden_states


class MiniMaxH3Qwen3VLVisionModel(nn.Module):
    """Qwen3-VL vision tower (patch embed + blocks + merger + DeepStack)."""

    def __init__(self, config: Any) -> None:
        super().__init__()
        self.config = config
        self.spatial_merge_size = config.spatial_merge_size
        self.patch_size = config.patch_size
        self.spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size

        self.patch_embed = MiniMaxH3Qwen3VLVisionPatchEmbed(config=config)
        self.pos_embed = nn.Embedding(config.num_position_embeddings, config.hidden_size)
        self.num_grid_per_side = int(config.num_position_embeddings**0.5)

        head_dim = config.hidden_size // config.num_heads
        self.rotary_pos_emb = MiniMaxH3Qwen3VLVisionRotaryEmbedding(head_dim // 2)

        self.blocks = nn.ModuleList([MiniMaxH3Qwen3VLVisionBlock(config=config) for _ in range(config.depth)])
        self.merger = MiniMaxH3Qwen3VLVisionPatchMerger(config=config, use_postshuffle_norm=False)
        self.deepstack_visual_indexes = config.deepstack_visual_indexes
        self.deepstack_merger_list = nn.ModuleList(
            [
                MiniMaxH3Qwen3VLVisionPatchMerger(config=config, use_postshuffle_norm=True)
                for _ in range(len(config.deepstack_visual_indexes))
            ]
        )

    def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor:
        merge_size = self.spatial_merge_size
        grid_thw_list = grid_thw.tolist()
        max_hw = max(max(h, w) for _, h, w in grid_thw_list)
        freq_table = self.rotary_pos_emb(max_hw).to(device=grid_thw.device)
        device = freq_table.device
        total_tokens = sum(t * h * w for t, h, w in grid_thw_list)
        pos_ids = torch.empty((total_tokens, 2), dtype=torch.long, device=device)
        offset = 0
        for num_frames, height, width in grid_thw_list:
            merged_h, merged_w = height // merge_size, width // merge_size
            block_rows = torch.arange(merged_h, device=device)
            block_cols = torch.arange(merged_w, device=device)
            intra_row = torch.arange(merge_size, device=device)
            intra_col = torch.arange(merge_size, device=device)
            row_idx = block_rows[:, None, None, None] * merge_size + intra_row[None, None, :, None]
            col_idx = block_cols[None, :, None, None] * merge_size + intra_col[None, None, None, :]
            row_idx = row_idx.expand(merged_h, merged_w, merge_size, merge_size).reshape(-1)
            col_idx = col_idx.expand(merged_h, merged_w, merge_size, merge_size).reshape(-1)
            coords = torch.stack((row_idx, col_idx), dim=-1)
            if num_frames > 1:
                coords = coords.repeat(num_frames, 1)
            num_tokens = coords.shape[0]
            pos_ids[offset : offset + num_tokens] = coords
            offset += num_tokens
        embeddings = freq_table[pos_ids]
        embeddings = embeddings.flatten(1)
        return embeddings

    def fast_pos_embed_interpolate(self, grid_thw: torch.Tensor) -> torch.Tensor:
        grid_thw_list = grid_thw.tolist()
        grid_ts = [row[0] for row in grid_thw_list]
        grid_hs = [row[1] for row in grid_thw_list]
        grid_ws = [row[2] for row in grid_thw_list]
        device = self.pos_embed.weight.device
        idx_list = [[] for _ in range(4)]
        weight_list = [[] for _ in range(4)]
        for _, h, w in grid_thw_list:
            h_idxs = torch.linspace(0, self.num_grid_per_side - 1, h)
            w_idxs = torch.linspace(0, self.num_grid_per_side - 1, w)
            h_idxs_floor = h_idxs.int()
            w_idxs_floor = w_idxs.int()
            h_idxs_ceil = (h_idxs.int() + 1).clip(max=self.num_grid_per_side - 1)
            w_idxs_ceil = (w_idxs.int() + 1).clip(max=self.num_grid_per_side - 1)
            dh = h_idxs - h_idxs_floor
            dw = w_idxs - w_idxs_floor
            base_h = h_idxs_floor * self.num_grid_per_side
            base_h_ceil = h_idxs_ceil * self.num_grid_per_side
            indices = [
                (base_h[None].T + w_idxs_floor[None]).flatten(),
                (base_h[None].T + w_idxs_ceil[None]).flatten(),
                (base_h_ceil[None].T + w_idxs_floor[None]).flatten(),
                (base_h_ceil[None].T + w_idxs_ceil[None]).flatten(),
            ]
            weights = [
                ((1 - dh)[None].T * (1 - dw)[None]).flatten(),
                ((1 - dh)[None].T * dw[None]).flatten(),
                (dh[None].T * (1 - dw)[None]).flatten(),
                (dh[None].T * dw[None]).flatten(),
            ]
            for i in range(4):
                idx_list[i].extend(indices[i].tolist())
                weight_list[i].extend(weights[i].tolist())
        idx_tensor = torch.tensor(idx_list, dtype=torch.long, device=device)
        weight_tensor = torch.tensor(weight_list, dtype=self.pos_embed.weight.dtype, device=device)
        pos_embeds = self.pos_embed(idx_tensor).to(device) * weight_tensor[:, :, None]
        patch_pos_embeds = pos_embeds[0] + pos_embeds[1] + pos_embeds[2] + pos_embeds[3]
        patch_pos_embeds = patch_pos_embeds.split([h * w for h, w in zip(grid_hs, grid_ws)])
        patch_pos_embeds_permute = []
        merge_size = self.spatial_merge_size
        for pos_embed, t, h, w in zip(patch_pos_embeds, grid_ts, grid_hs, grid_ws):
            pos_embed = pos_embed.repeat(t, 1)
            pos_embed = (
                pos_embed.view(t, h // merge_size, merge_size, w // merge_size, merge_size, -1)
                .permute(0, 1, 3, 2, 4, 5)
                .flatten(0, 4)
            )
            patch_pos_embeds_permute.append(pos_embed)
        return torch.cat(patch_pos_embeds_permute)

    def forward(
        self,
        hidden_states: torch.Tensor,
        grid_thw: torch.Tensor,
    ) -> tuple[torch.Tensor, list[torch.Tensor]]:
        hidden_states = self.patch_embed(hidden_states)
        pos_embeds = self.fast_pos_embed_interpolate(grid_thw)
        hidden_states = hidden_states + pos_embeds
        rotary_pos_emb = self.rot_pos_emb(grid_thw)
        seq_len, _ = hidden_states.size()
        hidden_states = hidden_states.reshape(seq_len, -1)
        rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1)
        emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)
        position_embeddings = (emb.cos(), emb.sin())
        cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(
            dim=0, dtype=torch.int32
        )
        cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)
        deepstack_feature_lists = []
        for layer_num, blk in enumerate(self.blocks):
            hidden_states = blk(
                hidden_states,
                cu_seqlens=cu_seqlens,
                position_embeddings=position_embeddings,
            )
            if layer_num in self.deepstack_visual_indexes:
                deepstack_feature = self.deepstack_merger_list[self.deepstack_visual_indexes.index(layer_num)](
                    hidden_states
                )
                deepstack_feature_lists.append(deepstack_feature)
        merged_hidden_states = self.merger(hidden_states)
        return merged_hidden_states, deepstack_feature_lists


# ---------------------------------------------------------------------------
# Qwen3-VL text decoder layers (TP-sharded).
# ---------------------------------------------------------------------------


class MiniMaxH3Qwen3VLTextAttention(nn.Module):
    def __init__(
        self,
        group: Any,
        config: Any,
        dtype: torch.dtype,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.group = group
        self.config = config
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
        self.scaling = self.head_dim**-0.5
        self.is_causal = True

        self.qkv_proj = MiniMaxH3Qwen3VLQKVParallelLinear(
            group,
            hidden_size=self.hidden_size,
            num_heads=self.num_heads,
            num_kv_heads=self.num_kv_heads,
            head_dim=self.head_dim,
            dtype=dtype,
            quant_config=quant_config,
            prefix=f"{prefix}.qkv_proj",
        )
        self.o_proj = MiniMaxH3Qwen3VLRowParallelLinear(
            group,
            input_size=self.num_heads * self.head_dim,
            output_size=self.hidden_size,
            dtype=dtype,
            input_is_parallel=True,
            quant_config=quant_config,
            prefix=f"{prefix}.o_proj",
        )
        self.q_norm = MiniMaxH3Qwen3VLRMSNorm(self.head_dim, eps=config.rms_norm_eps, dtype=dtype)
        self.k_norm = MiniMaxH3Qwen3VLRMSNorm(self.head_dim, eps=config.rms_norm_eps, dtype=dtype)

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
    ) -> torch.Tensor:
        input_shape = hidden_states.shape[:-1]
        hidden_shape = (*input_shape, -1, self.head_dim)
        q_local = self.qkv_proj.local_num_heads * self.head_dim
        kv_local = self.qkv_proj.local_num_kv_heads * self.head_dim
        if _is_quantized_linear(self.qkv_proj):
            query, key, value = self.qkv_proj(hidden_states).split(
                (q_local, kv_local, kv_local),
                dim=-1,
            )
        else:
            q_weight = self.qkv_proj.weight[0:q_local]
            k_weight = self.qkv_proj.weight[q_local : q_local + kv_local]
            v_weight = self.qkv_proj.weight[q_local + kv_local : q_local + 2 * kv_local]
            # Three separate GEMMs keep the numerics identical to the reference
            # path (packed QKV changes cuBLAS tiling and shifts bf16 rounding).
            query = F.linear(hidden_states, q_weight)
            key = F.linear(hidden_states, k_weight)
            value = F.linear(hidden_states, v_weight)
        query_states = self.q_norm(query.view(hidden_shape)).transpose(1, 2)
        key_states = self.k_norm(key.view(hidden_shape)).transpose(1, 2)
        value_states = value.view(hidden_shape).transpose(1, 2)

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

        # Expand GQA keys/values to the local query-head count so SDPA works
        # uniformly across backends (mirrors the reference eager path).
        num_key_value_groups = self.num_heads // self.num_kv_heads
        key_states = key_states.repeat_interleave(num_key_value_groups, dim=1)
        value_states = value_states.repeat_interleave(num_key_value_groups, dim=1)

        attn_output = F.scaled_dot_product_attention(
            query_states,
            key_states,
            value_states,
            dropout_p=0.0,
            is_causal=True,
        )
        attn_output = attn_output.transpose(1, 2).reshape(*input_shape, -1).contiguous()
        attn_output = self.o_proj(attn_output)
        return attn_output


class MiniMaxH3Qwen3VLTextMLP(nn.Module):
    def __init__(
        self,
        group: Any,
        config: Any,
        dtype: torch.dtype,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size
        self.gate_up_proj = MiniMaxH3Qwen3VLMergedColumnParallelLinear(
            group,
            input_size=self.hidden_size,
            intermediate_size=self.intermediate_size,
            dtype=dtype,
            quant_config=quant_config,
            prefix=f"{prefix}.gate_up_proj",
        )
        self.down_proj = MiniMaxH3Qwen3VLRowParallelLinear(
            group,
            input_size=self.intermediate_size,
            output_size=self.hidden_size,
            dtype=dtype,
            input_is_parallel=True,
            quant_config=quant_config,
            prefix=f"{prefix}.down_proj",
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        ip_local = self.gate_up_proj.intermediate_size_per_partition
        if _is_quantized_linear(self.gate_up_proj):
            gate, up = self.gate_up_proj(x).split(ip_local, dim=-1)
        else:
            gate = F.linear(x, self.gate_up_proj.weight[0:ip_local])
            up = F.linear(x, self.gate_up_proj.weight[ip_local : 2 * ip_local])
        x = F.silu(gate) * up
        return self.down_proj(x)


class MiniMaxH3Qwen3VLTextDecoderLayer(nn.Module):
    def __init__(
        self,
        group: Any,
        config: Any,
        dtype: torch.dtype,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size
        self.self_attn = MiniMaxH3Qwen3VLTextAttention(
            group,
            config,
            dtype,
            quant_config=quant_config,
            prefix=f"{prefix}.self_attn",
        )
        self.mlp = MiniMaxH3Qwen3VLTextMLP(
            group,
            config,
            dtype,
            quant_config=quant_config,
            prefix=f"{prefix}.mlp",
        )
        self.input_layernorm = MiniMaxH3Qwen3VLRMSNorm(config.hidden_size, eps=config.rms_norm_eps, dtype=dtype)
        self.post_attention_layernorm = MiniMaxH3Qwen3VLRMSNorm(
            config.hidden_size, eps=config.rms_norm_eps, dtype=dtype
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
    ) -> torch.Tensor:
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        hidden_states = self.self_attn(hidden_states, position_embeddings=position_embeddings)
        hidden_states = residual + hidden_states

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


class MiniMaxH3Qwen3VLTextModel(nn.Module):
    """TP-sharded Qwen3-VL text model returning unnormalized layer-50 states."""

    def __init__(
        self,
        group: Any,
        config: Any,
        selected_layer: int,
        dtype: torch.dtype,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "text_encoder.text_model",
    ) -> None:
        super().__init__()
        self.group = group
        self.config = config
        self.selected_layer = selected_layer
        self.vocab_size = config.vocab_size
        self.hidden_size = config.hidden_size
        self.num_layers = min(int(getattr(config, "num_hidden_layers", 0)), selected_layer)
        self.embed_tokens = MiniMaxH3Qwen3VLVocabParallelEmbedding(
            group, num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, dtype=dtype
        )
        self.rotary_emb = MiniMaxH3Qwen3VLTextRotaryEmbedding(config)
        self.layers = nn.ModuleList(
            [
                MiniMaxH3Qwen3VLTextDecoderLayer(
                    group,
                    config,
                    dtype,
                    quant_config=quant_config,
                    prefix=f"{prefix}.layers.{layer_idx}",
                )
                for layer_idx in range(self.num_layers)
            ]
        )

    def _deepstack_process(
        self,
        hidden_states: torch.Tensor,
        visual_pos_masks: torch.Tensor,
        visual_embeds: torch.Tensor,
    ) -> torch.Tensor:
        visual_pos_masks = visual_pos_masks.to(hidden_states.device)
        visual_embeds = visual_embeds.to(hidden_states.device, hidden_states.dtype)
        hidden_states = hidden_states.clone()
        hidden_states[visual_pos_masks, :] = hidden_states[visual_pos_masks, :] + visual_embeds
        return hidden_states

    def forward(
        self,
        inputs_embeds: torch.Tensor,
        positions: torch.Tensor,
        *,
        visual_pos_masks: torch.Tensor | None = None,
        deepstack_visual_embeds: list[torch.Tensor] | None = None,
    ) -> torch.Tensor:
        hidden_states = inputs_embeds
        position_embeddings = self.rotary_emb(hidden_states, positions)
        for layer_idx, decoder_layer in enumerate(self.layers):
            hidden_states = decoder_layer(hidden_states, position_embeddings=position_embeddings)
            if deepstack_visual_embeds is not None and layer_idx in range(len(deepstack_visual_embeds)):
                hidden_states = self._deepstack_process(
                    hidden_states,
                    visual_pos_masks,
                    deepstack_visual_embeds[layer_idx],
                )
        return hidden_states


# ---------------------------------------------------------------------------
# Multimodal position / placeholder-mask helpers (ported from the reference).
# ---------------------------------------------------------------------------


def _get_vision_position_ids(
    start_position: int,
    grid_thw: list[int] | torch.Tensor,
    temp_merge_size: int,
    spatial_merge_size: int,
    device: str | torch.device,
) -> torch.Tensor:
    llm_grid_t, llm_grid_h, llm_grid_w = (
        grid_thw[0].item() // temp_merge_size,
        grid_thw[1].item() // spatial_merge_size,
        grid_thw[2].item() // spatial_merge_size,
    )
    position_temporal = torch.arange(llm_grid_t, device=device)
    position_width = torch.arange(llm_grid_w, device=device) + start_position
    position_height = torch.arange(llm_grid_h, device=device) + start_position
    position_width = position_width.repeat(llm_grid_h * llm_grid_t)
    position_height = position_height.repeat_interleave(llm_grid_w).repeat(llm_grid_t)
    position_temporal = position_temporal.repeat_interleave(llm_grid_h * llm_grid_w) + start_position
    return torch.stack([position_temporal, position_height, position_width], dim=0)


def _get_rope_index(
    input_ids: torch.Tensor,
    mm_token_type_ids: torch.Tensor,
    image_grid_thw: torch.Tensor | None,
    video_grid_thw: torch.Tensor | None,
    attention_mask: torch.Tensor | None,
    spatial_merge_size: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    if video_grid_thw is not None:
        video_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0)
        video_grid_thw[:, 0] = 1

    mrope_position_deltas = []
    position_ids = torch.zeros(
        3,
        input_ids.shape[0],
        input_ids.shape[1],
        dtype=input_ids.dtype,
        device=input_ids.device,
    )
    grid_iters = {
        1: iter(image_grid_thw) if image_grid_thw is not None else None,
        2: iter(video_grid_thw) if video_grid_thw is not None else None,
    }

    for batch_idx, current_input_ids in enumerate(input_ids):
        input_token_type = mm_token_type_ids[batch_idx]
        if attention_mask is not None:
            current_input_ids = current_input_ids[attention_mask[batch_idx].bool()]
            input_token_type = input_token_type[attention_mask[batch_idx].bool()]

        input_type_group = []
        for key, group in itertools.groupby(enumerate(input_token_type.tolist()), lambda x: x[1]):
            group = list(group)
            start_index = group[0][0]
            end_index = group[-1][0] + 1
            input_type_group.append((key, start_index, end_index))

        current_pos = 0
        llm_pos_ids_list = []
        for modality_type, start_idx, end_idx in input_type_group:
            if modality_type == 0:
                text_len = end_idx - start_idx
                llm_pos_ids_list.append(
                    torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + current_pos
                )
                current_pos += text_len
            else:
                grid_thw = next(grid_iters[modality_type])
                vision_position_ids = _get_vision_position_ids(
                    current_pos, grid_thw, 1, spatial_merge_size, device=input_ids.device
                )
                llm_pos_ids_list.append(vision_position_ids)
                current_pos += max(grid_thw[1], grid_thw[2]) // spatial_merge_size
        llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
        if attention_mask is not None:
            position_ids[:, batch_idx, attention_mask[batch_idx].bool()] = llm_positions.to(position_ids.device)
        else:
            position_ids[:, batch_idx] = llm_positions.to(position_ids.device)
        mrope_position_deltas.append(llm_positions.max() + 1 - len(current_input_ids))
    mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1)
    return position_ids, mrope_position_deltas


# ---------------------------------------------------------------------------
# Encoder entry point.
# ---------------------------------------------------------------------------


class MiniMaxH3Qwen3VLEncoder(nn.Module):
    """Frozen, TP-capable Qwen3-VL backbone returning layer-50 states.

    ``encoder_group`` is a ``GroupCoordinator`` over the encoder tensor-parallel
    ranks (by default the first ``text_encoder_tp_size`` DiT ranks).  Only
    ranks inside the group load weights; other ranks construct a parameter-free
    stub that never runs the forward.
    """

    def __init__(
        self,
        model_path: str,
        *,
        device: torch.device,
        load_model: bool,
        encoder_group: Any | None = None,
        quant_config: QuantizationConfig | None = None,
    ) -> None:
        super().__init__()
        self.device_target = device
        self.encoder_group = encoder_group
        self.quant_config = quant_config
        _validate_encoder_quant_config(quant_config)
        self.image_token_id = 151655
        self.video_token_id = 151656
        self._tp_size = 1
        if not load_model:
            return

        from transformers import Qwen3VLConfig

        config = Qwen3VLConfig.from_pretrained(model_path, trust_remote_code=False)
        self.image_token_id = int(config.image_token_id)
        self.video_token_id = int(config.video_token_id)
        self._tp_size = int(encoder_group.world_size) if encoder_group is not None else 1
        dtype = torch.bfloat16
        self.vision = MiniMaxH3Qwen3VLVisionModel(config.vision_config)
        self.vision.to(dtype=dtype)
        self.text_model = MiniMaxH3Qwen3VLTextModel(
            encoder_group,
            config.text_config,
            MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER,
            dtype,
            quant_config=quant_config,
        )
        self.text_model.to(dtype=dtype)
        logger.info(
            "MiniMax H3 Qwen3-VL encoder: %d retained decoder layers, text_encoder_tp_size=%d, vision replicated",
            MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER,
            self._tp_size,
        )

    @property
    def is_loaded(self) -> bool:
        return hasattr(self, "text_model")

    @property
    def tp_size(self) -> int:
        return self._tp_size

    def _map_weight_name(self, name: str) -> tuple[str, str | int | None] | None:
        if name == "lm_head.weight" or name == "model.language_model.norm.weight":
            return None
        if name.startswith("model.visual."):
            return ("vision." + name[len("model.visual.") :], None)
        if name.startswith("model.language_model."):
            rest = name[len("model.language_model.") :]
            if rest.startswith("embed_tokens."):
                return ("text_model." + rest, None)
            match = re.match(r"layers\.(\d+)\.", rest)
            if match is not None and int(match.group(1)) >= MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER:
                return None
            if ".self_attn.q_proj.weight" in rest:
                return (
                    "text_model." + rest.replace(".self_attn.q_proj.weight", ".self_attn.qkv_proj.weight"),
                    "q",
                )
            if ".self_attn.k_proj.weight" in rest:
                return (
                    "text_model." + rest.replace(".self_attn.k_proj.weight", ".self_attn.qkv_proj.weight"),
                    "k",
                )
            if ".self_attn.v_proj.weight" in rest:
                return (
                    "text_model." + rest.replace(".self_attn.v_proj.weight", ".self_attn.qkv_proj.weight"),
                    "v",
                )
            if ".mlp.gate_proj.weight" in rest:
                return (
                    "text_model." + rest.replace(".mlp.gate_proj.weight", ".mlp.gate_up_proj.weight"),
                    0,
                )
            if ".mlp.up_proj.weight" in rest:
                return (
                    "text_model." + rest.replace(".mlp.up_proj.weight", ".mlp.gate_up_proj.weight"),
                    1,
                )
            return ("text_model." + rest, None)
        return None

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        params = dict(self.named_parameters())
        loaded: set[str] = set()
        # Seeded from modules that own a multi-shard loader, so a fused weight
        # receiving no source shard is caught as well as a partially filled one.
        expected_fused_shards: dict[str, tuple[tuple[str, str | int], ...]] = {
            f"{module_name}.weight": sources
            for module_name, module in self.named_modules()
            if (sources := _fused_source_shards(module)) is not None
        }
        loaded_fused_shards: dict[str, set[str | int]] = {name: set() for name in expected_fused_shards}
        for name, tensor in weights:
            mapped = self._map_weight_name(name)
            if mapped is None:
                continue
            param_name, shard_id = mapped
            param = params.get(param_name)
            if param is None:
                logger.warning("MiniMax H3 text encoder weight %s has no target parameter", name)
                continue
            weight_loader = getattr(param, "weight_loader", _default_weight_loader)
            weight_loader(param, tensor, shard_id)
            loaded.add(param_name)
            if shard_id is not None and param_name in loaded_fused_shards:
                loaded_fused_shards[param_name].add(shard_id)
        _verify_fused_shards_complete(expected_fused_shards, loaded_fused_shards)
        missing = sorted(set(params) - loaded)
        if missing:
            # Listed in full: this aborts startup, so the message is the only diagnostic.
            raise RuntimeError(f"MiniMax H3 text encoder weights not loaded: {len(missing)} params: {missing}")
        return loaded

    def load_to_device(self) -> None:
        if not self.is_loaded:
            return
        if getattr(self, "_omni_layerwise_enabled", False):
            stager = getattr(self, "_omni_non_block_stager", None)
            if stager is None:
                hooks = getattr(self, "_omni_layerwise_hooks", ())
                pin_memory = bool(getattr(hooks[0], "pin_memory", True)) if hooks else True
                stager = PinnedModuleStager(
                    self._omni_non_block_modules(),
                    self.device_target,
                    pin_memory=pin_memory,
                    cache_retention=getattr(self, "_omni_component_cache", None),
                )
                self._omni_non_block_stager = stager
            stager.load()
            return
        self.vision.to(self.device_target)
        self.text_model.to(self.device_target)

    def offload_to_cpu(self) -> None:
        if not self.is_loaded:
            return
        if getattr(self, "_omni_layerwise_enabled", False):
            for hook in self._omni_layerwise_hooks:
                hook.offload_layer()
            stager = getattr(self, "_omni_non_block_stager", None)
            if stager is not None:
                stager.offload()
            else:
                # The initial DLO setup reaches this branch before the first
                # load. These modules were constructed on CPU, so this is not
                # a device-to-host rematerialization.
                for child in self._omni_non_block_modules():
                    child.to("cpu")
                self._release_omni_component_cache()
            return
        self.vision.to("cpu")
        self.text_model.to("cpu")
        torch.accelerator.empty_cache()

    def _omni_non_block_modules(self) -> list[nn.Module]:
        modules = [child for name, child in self.vision.named_children() if name != "blocks"]
        modules.extend(child for name, child in self.text_model.named_children() if name != "layers")
        return modules

    def set_omni_component_cache(self, cache: BoundedAllocatorCache | None) -> None:
        self._omni_component_cache = cache
        stager = getattr(self, "_omni_non_block_stager", None)
        if stager is not None:
            stager.set_cache_retention(cache)

    def _release_omni_component_cache(self) -> None:
        cache = getattr(self, "_omni_component_cache", None)
        if cache is None:
            torch.accelerator.empty_cache()
        else:
            cache.release_if_needed()

    def enable_omni_layerwise_offload(self, *, pin_memory: bool = True) -> None:
        """Stream the TP-local Qwen vision/text blocks for low-HBM serving.

        The encoder has its own TP process group, so these blocks must remain
        rank-local.  Reusing the DiT AllGather group would concatenate
        different TP shards and corrupt the encoder weights.
        """
        if not self.is_loaded or getattr(self, "_omni_layerwise_enabled", False):
            return

        from vllm_omni.diffusion.offloader.layerwise_backend import apply_block_hook
        from vllm_omni.platforms import current_omni_platform

        self._omni_layerwise_hooks = []
        self._omni_layerwise_block_groups = []
        copy_stream = current_omni_platform.Stream()
        for blocks in (self.vision.blocks, self.text_model.layers):
            if len(blocks) <= 1:
                continue
            last_hook = apply_block_hook(
                blocks[-1],
                blocks[0],
                self.device_target,
                copy_stream,
                pin_memory,
            )
            hooks = [last_hook]
            for index, block in enumerate(blocks[:-1]):
                hooks.append(
                    apply_block_hook(
                        block,
                        blocks[index + 1],
                        self.device_target,
                        copy_stream,
                        pin_memory,
                    )
                )
            for index, hook in enumerate(hooks):
                hook._prev_hook = hooks[index - 1]
            self._omni_layerwise_hooks.extend(hooks)
            self._omni_layerwise_block_groups.append(blocks)

        self._omni_layerwise_enabled = bool(self._omni_layerwise_hooks)
        logger.info(
            "MiniMax H3 encoder layerwise offload enabled on %d blocks across %d stacks",
            sum(len(blocks) for blocks in self._omni_layerwise_block_groups),
            len(self._omni_layerwise_block_groups),
        )

    def _encode(
        self,
        input_ids: torch.Tensor,
        *,
        pixel_values: torch.Tensor | None = None,
        image_grid_thw: torch.Tensor | None = None,
        pixel_values_videos: torch.Tensor | None = None,
        video_grid_thw: torch.Tensor | None = None,
    ) -> torch.Tensor:
        device = self.device_target
        ids = input_ids.to(device, torch.long)[None]
        attention_mask = torch.ones_like(ids)
        mm_types = torch.zeros_like(ids, dtype=torch.int32)
        mm_types[ids == self.image_token_id] = 1
        mm_types[ids == self.video_token_id] = 2

        inputs_embeds = self.text_model.embed_tokens(ids)
        image_mask = None
        video_mask = None
        deepstack_image_embeds: list[torch.Tensor] | None = None
        deepstack_video_embeds: list[torch.Tensor] | None = None

        if pixel_values is not None:
            image_embeds, deepstack_image_embeds = self.vision(
                pixel_values.to(device, torch.bfloat16),
                image_grid_thw.to(device, torch.long),
            )
            image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
            image_mask, _ = self._get_placeholder_mask(ids, inputs_embeds=inputs_embeds, image_features=image_embeds)
            inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
        if pixel_values_videos is not None:
            video_embeds, deepstack_video_embeds = self.vision(
                pixel_values_videos.to(device, torch.bfloat16),
                video_grid_thw.to(device, torch.long),
            )
            video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
            _, video_mask = self._get_placeholder_mask(ids, inputs_embeds=inputs_embeds, video_features=video_embeds)
            inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)

        visual_pos_masks, deepstack_visual_embeds = self._aggregate_deepstack(
            image_mask,
            video_mask,
            deepstack_image_embeds,
            deepstack_video_embeds,
        )

        positions, _ = _get_rope_index(
            ids,
            mm_token_type_ids=mm_types,
            image_grid_thw=image_grid_thw.to(device, torch.long) if image_grid_thw is not None else None,
            video_grid_thw=video_grid_thw.to(device, torch.long) if video_grid_thw is not None else None,
            attention_mask=attention_mask,
            spatial_merge_size=int(self.vision.config.spatial_merge_size),
        )
        hidden = self.text_model(
            inputs_embeds,
            positions,
            visual_pos_masks=visual_pos_masks,
            deepstack_visual_embeds=deepstack_visual_embeds,
        )[0]
        expected = (int(ids.shape[1]), MINIMAX_H3_QWEN3VL_HIDDEN_DIM)
        if tuple(hidden.shape) != expected:
            raise ValueError(f"unexpected Qwen3-VL hidden shape {tuple(hidden.shape)}, expected {expected}")
        return hidden.to(torch.bfloat16)

    def _get_placeholder_mask(
        self,
        input_ids: torch.Tensor,
        *,
        inputs_embeds: torch.Tensor,
        image_features: torch.Tensor | None = None,
        video_features: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        special_image_mask = input_ids == self.image_token_id
        special_video_mask = input_ids == self.video_token_id
        special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
        special_video_mask = special_video_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
        if image_features is not None and special_image_mask.sum().item() != image_features.numel():
            raise ValueError(
                "Image features and image tokens do not match: "
                f"tokens: {special_image_mask.sum().item() // inputs_embeds.shape[-1]}, "
                f"features: {image_features.shape[0]}"
            )
        if video_features is not None and special_video_mask.sum().item() != video_features.numel():
            raise ValueError(
                "Video features and video tokens do not match: "
                f"tokens: {special_video_mask.sum().item() // inputs_embeds.shape[-1]}, "
                f"features: {video_features.shape[0]}"
            )
        return special_image_mask, special_video_mask

    def _aggregate_deepstack(
        self,
        image_mask: torch.Tensor | None,
        video_mask: torch.Tensor | None,
        deepstack_image_embeds: list[torch.Tensor] | None,
        deepstack_video_embeds: list[torch.Tensor] | None,
    ) -> tuple[torch.Tensor | None, list[torch.Tensor] | None]:
        if image_mask is not None and video_mask is not None:
            image_mask = image_mask[..., 0]
            video_mask = video_mask[..., 0]
            visual_pos_masks = image_mask | video_mask
            deepstack_visual_embeds = []
            image_mask_joint = image_mask[visual_pos_masks]
            video_mask_joint = video_mask[visual_pos_masks]
            for img_embed, vid_embed in zip(deepstack_image_embeds, deepstack_video_embeds):
                embed_joint = img_embed.new_zeros(visual_pos_masks.sum(), img_embed.shape[-1]).to(img_embed.device)
                embed_joint[image_mask_joint, :] = img_embed
                embed_joint[video_mask_joint, :] = vid_embed
                deepstack_visual_embeds.append(embed_joint)
            return visual_pos_masks, deepstack_visual_embeds
        if image_mask is not None:
            image_mask = image_mask[..., 0]
            return image_mask, deepstack_image_embeds
        if video_mask is not None:
            video_mask = video_mask[..., 0]
            return video_mask, deepstack_video_embeds
        return None, None

    @torch.inference_mode()
    def encode_ids(
        self,
        input_ids: torch.Tensor,
        *,
        pixel_values: torch.Tensor | None = None,
        image_grid_thw: torch.Tensor | None = None,
        pixel_values_videos: torch.Tensor | None = None,
        video_grid_thw: torch.Tensor | None = None,
    ) -> torch.Tensor:
        if not self.is_loaded:
            raise RuntimeError("MiniMax H3 text encoder is only loaded on the encoder TP ranks")
        if input_ids.ndim != 1:
            raise ValueError(f"input_ids must be 1-D, got {list(input_ids.shape)}")
        if (pixel_values is None) != (image_grid_thw is None):
            raise ValueError("pixel_values and image_grid_thw must be provided together")
        if (pixel_values_videos is None) != (video_grid_thw is None):
            raise ValueError("pixel_values_videos and video_grid_thw must be provided together")
        if next(self.parameters()).device.type != self.device_target.type:
            raise RuntimeError("call load_to_device() before encode_ids()")

        torch.backends.cuda.enable_cudnn_sdp(True)
        hidden = self._encode(
            input_ids,
            pixel_values=pixel_values,
            image_grid_thw=image_grid_thw,
            pixel_values_videos=pixel_values_videos,
            video_grid_thw=video_grid_thw,
        )
        return hidden.cpu()

    def forward(
        self,
        input_ids: torch.Tensor,
        *,
        pixel_values: torch.Tensor | None = None,
        image_grid_thw: torch.Tensor | None = None,
        pixel_values_videos: torch.Tensor | None = None,
        video_grid_thw: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Hook-compatible entry point for model-level CPU offloading."""
        kwargs = {
            "pixel_values": pixel_values,
            "image_grid_thw": image_grid_thw,
        }
        if pixel_values_videos is not None or video_grid_thw is not None:
            kwargs.update(
                pixel_values_videos=pixel_values_videos,
                video_grid_thw=video_grid_thw,
            )
        return self.encode_ids(input_ids, **kwargs)


__all__ = ["MiniMaxH3Qwen3VLEncoder"]