model/modeling_decoder_mot.py

"""MoT decoder layer with three-path routing (text / vision-input / vision-gen).

Subclasses `HunYuanVLMoTDecoderLayer` from upstream to add the third (`_g`)
path used by Flow Matching. Upstream provides `mlp / mlp_v` and matching
LayerNorms; we add `mlp_g / input_layernorm_g / post_attention_layernorm_g`.

`modality_mask` semantics:
  * 0 = text token            → mlp / input_layernorm / post_attention_layernorm
  * 1 = input vision token    → mlp_v / *_v
  * 2 = generated vision token (FM latent) → mlp_g / *_g

Attention QKVO is **not** extended to a third path — generation tokens reuse
the `_v` projection (matches the existing HY-Unified design where
`q_proj_v / k_proj_v / v_proj_v / o_proj_v` cover both input and generated
vision tokens). We just need to convert the int modality mask to a bool
"vision" mask before calling `self.self_attn`.

mlp_g may optionally have a wider intermediate (Net2Wider scale-up). When
`config.mlp_g_intermediate_size` is set and larger than `intermediate_size`,
we build mlp_g with a cloned config that overrides `intermediate_size`.
"""
from __future__ import annotations

import copy
from typing import Optional

import torch

from transformers.cache_utils import Cache
from transformers.processing_utils import Unpack
from transformers.utils import TransformersKwargs
from transformers.utils.deprecation import deprecate_kwarg

from transformers.models.hunyuan_vl_mot.modeling_hunyuan_vl_mot import (
    HunYuanVLMoTDecoderLayer,
    HunYuanVLMoTMLP,
    HunYuanVLMoTRMSNorm,
)

# Side-effect import: replaces upstream's _flash_attention_forward_mot
from . import attention_mot_packed  # noqa: F401  # pylint: disable=unused-import


# ---------------------------------------------------------------------------
# Three-way mask_apply
# ---------------------------------------------------------------------------

def mask_apply_3way(
    hidden_states: torch.Tensor,
    modality_mask: Optional[torch.Tensor],
    text_funcs,
    vision_funcs,
    gen_funcs,
    out_dims=None,
    padding_mask: Optional[torch.Tensor] = None,
):
    """Routes tokens to text/vision/generation function lists by modality_mask.

    hidden_states: (B, S, D)
    modality_mask: (B, S) int tensor with values in {0, 1, 2}, or None
        0 = text, 1 = vision input, 2 = vision generation
    padding_mask: (B, S) int tensor, 1 = valid token, 0 = padding

    Returns a list of stacked (B, S, out_d) tensors — one per `text_funcs[i]`.
    """
    if modality_mask is None:
        # All-text: skip routing entirely
        return [text_funcs[i](hidden_states) for i in range(len(text_funcs))]

    bsz, seq_len, hidden_dim = hidden_states.size()
    flat = hidden_states.reshape(bsz * seq_len, hidden_dim)
    mask_flat = modality_mask.reshape(bsz * seq_len)
    if padding_mask is not None:
        valid_flat = padding_mask.reshape(bsz * seq_len).bool()
    else:
        valid_flat = None

    placeholder = hidden_states[0:1, 0:1, :]  # (1, 1, D)
    zero_feature = 0

    num_outputs = len(text_funcs)
    if out_dims is None:
        out_dims_resolved = [hidden_dim] * num_outputs
    else:
        out_dims_resolved = list(out_dims)

    # Pre-allocate output buffers (empty, not zeros — we overwrite all valid
    # positions and the rest are masked out by the caller / padding).
    out_flat = [
        torch.empty(bsz * seq_len, od, device=flat.device, dtype=flat.dtype)
        for od in out_dims_resolved
    ]
    # Padding positions need to be zeroed — we won't touch them in the
    # gather/scatter below. Cheaper than zeroing the whole tensor: only
    # zero the rows that won't be hit by any of the three modalities.
    if valid_flat is not None:
        invalid_flat = ~valid_flat
        if invalid_flat.any():
            for buf in out_flat:
                buf[invalid_flat] = 0
    # else: all rows will be hit by exactly one of {text, vision, gen}.

    def _dispatch(idx_mask, funcs):
        """Run `funcs` on rows selected by `idx_mask`, scatter back. If no
        rows are selected, multiply through a placeholder so the params still
        receive grad (avoids "unused parameter" DDP errors)."""
        nonlocal zero_feature
        if idx_mask.any():
            hs_sel = flat[idx_mask]
            for i, fn in enumerate(funcs):
                out_flat[i][idx_mask] = fn(hs_sel)
        else:
            for fn in funcs:
                zero_feature = zero_feature + fn(placeholder).mean() * 0

    # Text: mask == 0
    text_idx = (mask_flat == 0)
    if valid_flat is not None:
        text_idx = text_idx & valid_flat
    _dispatch(text_idx, text_funcs)

    # Vision input: mask == 1
    vis_idx = (mask_flat == 1)
    if valid_flat is not None:
        vis_idx = vis_idx & valid_flat
    _dispatch(vis_idx, vision_funcs)

    # Generation: mask == 2
    gen_idx = (mask_flat == 2)
    if valid_flat is not None:
        gen_idx = gen_idx & valid_flat
    _dispatch(gen_idx, gen_funcs)

    result = [out.view(bsz, seq_len, -1) for out in out_flat]
    result[0] = result[0] + zero_feature
    return result


# ---------------------------------------------------------------------------
# Decoder layer subclass
# ---------------------------------------------------------------------------

def _make_g_config(config, mlp_g_intermediate_size: int):
    """Clone config with a wider intermediate_size for mlp_g (Net2Wider).

    MUST be `copy.deepcopy`, NOT `copy.copy`. Reason:
      HunYuanVLMoTConfig.__setattr__ is a proxy that re-routes any attribute
      writes (when the key is in text_config.__dict__) into self.text_config.
      Shallow-copying the outer config keeps `g_cfg.text_config is
      config.text_config` — same object — so the subsequent
      `g_cfg.intermediate_size = ...` setattr ends up mutating the SHARED
      text_config used by every other decoder layer's mlp/mlp_v construction.
      Layer 0 escapes (its mlp/mlp_v are built before _make_g_config runs),
      but layer 1+ then sees text_config.intermediate_size = 12288 and builds
      mlp/mlp_v at the wider size — at that point ckpt weights (which are 6144
      for mlp/mlp_v) no longer match and `from_pretrained` fails with size
      mismatch. Deepcopy the whole config tree to break this aliasing.
    """
    g_cfg = copy.deepcopy(config)
    g_cfg.intermediate_size = int(mlp_g_intermediate_size)
    return g_cfg


class MoTDecoderLayer(HunYuanVLMoTDecoderLayer):
    """Three-path decoder layer: text / vision-input / vision-gen.

    Adds `mlp_g`, `input_layernorm_g`, `post_attention_layernorm_g` on top of
    upstream's two-path layer. `modality_mask` is interpreted as int{0,1,2}
    instead of bool.
    """

    def __init__(self, config, layer_idx: int):
        super().__init__(config, layer_idx)

        # Text-config-aware lookup (HunYuanVLMoTConfig proxies these to text_config)
        intermediate_size = getattr(config, "intermediate_size", None)
        mlp_g_inter = getattr(config, "mlp_g_intermediate_size", None)

        if mlp_g_inter is None or intermediate_size is None or mlp_g_inter == intermediate_size:
            g_cfg = config
        else:
            assert mlp_g_inter >= intermediate_size, (
                f"mlp_g_intermediate_size ({mlp_g_inter}) must be >= "
                f"intermediate_size ({intermediate_size}); shrinking is not supported."
            )
            g_cfg = _make_g_config(config, mlp_g_inter)

        self.mlp_g = HunYuanVLMoTMLP(g_cfg)
        self.input_layernorm_g = HunYuanVLMoTRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm_g = HunYuanVLMoTRMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )

    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Cache] = None,
        use_cache: Optional[bool] = False,
        cache_position: Optional[torch.LongTensor] = None,
        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
        modality_mask: Optional[torch.Tensor] = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> torch.Tensor:
        # Pull padding_mask from attention_mask dict for GEMM exclusion
        padding_mask = None
        if isinstance(attention_mask, dict):
            pm = attention_mask.get("padding_mask", None)
            if pm is not None and hidden_states.shape[1] == pm.shape[1]:
                padding_mask = pm

        # Convert int modality_mask to bool for attention QKVO routing
        # (attention has only 2 paths: text vs vision; gen tokens go through _v)
        attn_modality_mask = None
        if modality_mask is not None:
            attn_modality_mask = (modality_mask > 0)

        residual = hidden_states

        # Pre-attention LayerNorm — three paths
        hidden_states = mask_apply_3way(
            hidden_states, modality_mask,
            [self.input_layernorm],
            [self.input_layernorm_v],
            [self.input_layernorm_g],
            padding_mask=padding_mask,
        )[0]

        # Self-attention (upstream's two-path attention with bool mask)
        hidden_states, _ = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            use_cache=use_cache,
            cache_position=cache_position,
            position_embeddings=position_embeddings,
            modality_mask=attn_modality_mask,
            **kwargs,
        )
        hidden_states = residual + hidden_states

        # MLP — three paths (LayerNorm + MLP fused per path to match upstream layout)
        residual = hidden_states
        hidden_states = mask_apply_3way(
            hidden_states, modality_mask,
            [lambda x: self.mlp(self.post_attention_layernorm(x))],
            [lambda x: self.mlp_v(self.post_attention_layernorm_v(x))],
            [lambda x: self.mlp_g(self.post_attention_layernorm_g(x))],
            padding_mask=padding_mask,
        )[0]
        hidden_states = residual + hidden_states

        # Zero out padding positions to keep residuals clean (matches upstream behavior)
        if padding_mask is not None and hidden_states.shape[1] == padding_mask.shape[1]:
            hidden_states = hidden_states * padding_mask.unsqueeze(-1)

        return hidden_states


__all__ = ["MoTDecoderLayer", "mask_apply_3way"]