inference_code/inference/model/dit/dit_module.py

# Copyright (c) 2026 SandAI. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import importlib
import os
from dataclasses import dataclass
from enum import Enum
from typing import Any, Callable, List, Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange, repeat
from inference.common import Modality, VarlenHandler, is_hopper_arch
from inference.infra.parallelism import ulysses_scheduler
from inference.utils import env_is_true
from torch import Tensor
from torch.nn import Parameter

from .stage_profile import stage_profile

try:
    from magi_compiler import magi_compile as _magi_compile
    from magi_compiler.api import magi_register_custom_op as _magi_register_custom_op
    from magi_compiler.config import CompileConfig as _CompileConfig
except Exception:
    _magi_compile = None
    _magi_register_custom_op = None
    _CompileConfig = Any

CompileConfig = _CompileConfig

if _magi_compile is None or env_is_true("DAVINCI_DISABLE_MAGI_COMPILE") or env_is_true("MAGI_DISABLE_COMPILE"):
    def magi_compile(*args, **kwargs):
        def _decorator(obj):
            return obj

        return _decorator

    def magi_register_custom_op(*args, **kwargs):
        def _decorator(obj):
            return obj

        return _decorator
else:
    magi_compile = _magi_compile
    magi_register_custom_op = _magi_register_custom_op


@dataclass
class FFAHandler:
    q_ranges: torch.Tensor
    k_ranges: torch.Tensor
    max_seqlen_q: int
    max_seqlen_k: int
    attn_type_map: torch.Tensor
    softmax_scale: float


# Define the MLP activation type
class MLPActivationType(Enum):
    """Enumeration of supported activation functions for MLP"""

    SWIGLU7 = "swiglu7"
    GELU7 = "gelu7"


def swiglu7(x, alpha: float = 1.702, limit: float = 7.0, out_dtype: Optional[torch.dtype] = None):
    out_dtype = x.dtype if out_dtype is None else out_dtype
    compute_dtype = x.dtype if x.dtype in (torch.bfloat16, torch.float16) else torch.float32
    x = x.to(compute_dtype)
    x_glu, x_linear = x[..., ::2], x[..., 1::2]
    # Clamp the input values
    x_glu = x_glu.clamp(min=None, max=limit)
    x_linear = x_linear.clamp(min=-limit, max=limit)
    out_glu = x_glu * torch.sigmoid(alpha * x_glu)
    # Note we add an extra bias of 1 to the linear layer (from GPT-OSS)
    return (out_glu * (x_linear + 1)).to(out_dtype)


def gelu7(x, alpha: float = 1.702, limit: float = 7.0, out_dtype: Optional[torch.dtype] = None):
    out_dtype = x.dtype if out_dtype is None else out_dtype
    compute_dtype = x.dtype if x.dtype in (torch.bfloat16, torch.float16) else torch.float32
    x = x.to(compute_dtype)
    x_glu = x
    # Clamp the input values
    x_glu = x_glu.clamp(min=None, max=limit)
    out_glu = x_glu * torch.sigmoid(alpha * x_glu)
    # Note we add an extra bias of 1 to the linear layer
    return out_glu.to(out_dtype)


def create_activation_func(activation_type: MLPActivationType) -> Callable:
    match activation_type:
        case MLPActivationType.SWIGLU7:
            return swiglu7
        case MLPActivationType.GELU7:
            return gelu7
        case _:
            raise ValueError(f"Unknown activation type: {activation_type}")


class ModalityDispatcher:
    permuted_modality_mapping: torch.Tensor
    group_size: torch.Tensor
    group_size_cpu: list[int]
    num_modalities: int

    def __init__(self, modality_mapping: torch.Tensor, num_modalities: int):
        """
        Initialize dispatcher.
        This runs once during object construction and precomputes all mappings.
        """
        self.modality_mapping = modality_mapping
        self.num_modalities = num_modalities

        self.permuted_modality_mapping = self._precompute_permute_mapping(modality_mapping)

        self.group_size = torch.bincount(self.permuted_modality_mapping, minlength=num_modalities).to(torch.int32)
        self.group_size_cpu: list[int] = [int(x) for x in self.group_size.to("cpu").tolist()]

    def _precompute_permute_mapping(self, modality_mapping):
        # 1. Compute forward and inverse permutation mappings.
        # argsort is an efficient O(N log N) operation.
        self.permute_mapping = torch.argsort(modality_mapping)
        self.inv_permute_mapping = torch.argsort(self.permute_mapping)

        # 2. Compute group size for each modality.
        # bincount is highly efficient for counting.
        permuted_modality_mapping = modality_mapping[self.permute_mapping]

        return permuted_modality_mapping

    def dispatch(self, x: torch.Tensor) -> list[torch.Tensor]:
        grouped_tensors = torch.split(x, self.group_size_cpu, dim=0)
        return list(grouped_tensors)

    def undispatch(self, *processed_groups: list[torch.Tensor]) -> torch.Tensor:
        return torch.cat(processed_groups, dim=0)

    @staticmethod
    def permute(x: torch.Tensor, permute_mapping: torch.Tensor) -> torch.Tensor:
        """Apply forward permutation to tensor."""
        return x[permute_mapping]

    @staticmethod
    def inv_permute(x: torch.Tensor, inv_permute_mapping: torch.Tensor) -> torch.Tensor:
        """Apply inverse permutation to tensor."""
        return x[inv_permute_mapping]


def freq_bands(
    num_bands: int, temperature: float = 10000.0, step: int = 2, device: Optional[torch.device] = None
) -> torch.Tensor:
    exp = torch.arange(0, num_bands, step, dtype=torch.int64, device=device).to(torch.float32) / num_bands
    bands = 1.0 / (temperature**exp)
    return bands


def rotate_half(x, interleaved=False):
    if not interleaved:
        x1, x2 = x.chunk(2, dim=-1)
        out = x.clone()
        half = out.shape[-1] // 2
        out[..., :half] = -x2
        out[..., half:] = x1
        return out
    else:
        out = x.clone()
        out[..., ::2] = -x[..., 1::2]
        out[..., 1::2] = x[..., ::2]
        return out


def apply_rotary_emb_torch(x, cos, sin, interleaved=False):
    """
    x: (batch_size, seqlen, nheads, headdim)
    cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2)
    """
    ro_dim = cos.shape[-1] * 2
    assert ro_dim <= x.shape[-1]

    x_ro = x[..., :ro_dim]
    x_pass = x[..., ro_dim:]

    if not interleaved:
        x1, x2 = x_ro.chunk(2, dim=-1)

        if cos.ndim == x1.ndim - 1:
            cos = cos.unsqueeze(-2)
            sin = sin.unsqueeze(-2)
        elif cos.ndim == x1.ndim - 2:
            cos = cos.unsqueeze(0).unsqueeze(-2)
            sin = sin.unsqueeze(0).unsqueeze(-2)
        elif cos.ndim < x1.ndim:
            for _ in range(x1.ndim - cos.ndim):
                cos = cos.unsqueeze(0)
                sin = sin.unsqueeze(0)

        out = x.clone()
        half = ro_dim // 2
        out1 = out[..., :half]
        out2 = out[..., half:ro_dim]

        out1.copy_(x1)
        out1.mul_(cos)
        out1.addcmul_(x2, sin, value=-1.0)

        out2.copy_(x2)
        out2.mul_(cos)
        out2.addcmul_(x1, sin, value=1.0)

        return out
    else:
        x_ = x_ro.view(*x_ro.shape[:-1], ro_dim // 2, 2)
        x1 = x_[..., 0]
        x2 = x_[..., 1]

        if cos.ndim == x1.ndim - 1:
            cos = cos.unsqueeze(-2)
            sin = sin.unsqueeze(-2)
        elif cos.ndim == x1.ndim - 2:
            cos = cos.unsqueeze(0).unsqueeze(-2)
            sin = sin.unsqueeze(0).unsqueeze(-2)
        elif cos.ndim < x1.ndim:
            for _ in range(x1.ndim - cos.ndim):
                cos = cos.unsqueeze(0)
                sin = sin.unsqueeze(0)

        out = x.clone()
        out_ro_ = out[..., :ro_dim].view(*x_.shape)
        out0 = out_ro_[..., 0]
        out1 = out_ro_[..., 1]

        out0.copy_(x1)
        out0.mul_(cos)
        out0.addcmul_(x2, sin, value=-1.0)

        out1.copy_(x2)
        out1.mul_(cos)
        out1.addcmul_(x1, sin, value=1.0)

        return out


class ElementWiseFourierEmbed(nn.Module):
    def __init__(
        self,
        dim: int,
        max_res: int = 224,
        temperature: float = 10000.0,
        in_pixels: bool = True,
        linear_bands: bool = False,
        learnable: bool = False,
        device: torch.device = torch.device("cpu"),
        dtype: torch.dtype = torch.float32,
    ):
        """
        Args:
            dim: Output feature dimension, total channels, must be divisible by 6
            max_res: Max pixel-frequency resolution for pixel-domain bands
            temperature: Temperature in inverse-frequency mode
            in_pixels: True -> pixel-frequency bands, False -> inverse-frequency bands
            linear_bands: Whether pixel-frequency bands are linearly spaced
            learnable: Whether frequency bands are trainable
        """
        super().__init__()
        self.dim = dim
        self.in_pixels = in_pixels
        self.learnable = learnable
        self.temperature = temperature
        self.max_res = max_res
        self.linear_bands = linear_bands
        self.device = device
        self.dtype = dtype
        # Make frequency bands trainable or register as buffer
        bands = self.get_default_bands()
        if self.learnable:
            self.bands = nn.Parameter(bands)
        else:
            self.register_buffer("bands", bands)

    def forward(self, coords: torch.Tensor) -> torch.Tensor:
        """
        Args:
            coords: [L,9], column order (time, row, col, T, H, W, ref_T, ref_H, ref_W)
        Returns:
            emb: [L, dim] element-wise Fourier embedding
        """
        # Use slicing instead of unbind + stack to reduce intermediates
        coords_xyz = coords[:, :3]  # [L,3] -> (t, h, w)
        sizes = coords[:, 3:6]  # [L,3] -> (T, H, W)
        refs = coords[:, 6:9]  # [L,3] -> (ref_T, ref_H, ref_W)

        # Compute scale factors
        scales = (refs - 1) / (sizes - 1)  # [L,3]

        # NOTE: if both ref and size are 1, scale is fixed to 1; otherwise invalid
        scales[(refs == 1) & (sizes == 1)] = 1
        assert not scales.isnan().any(), "scales has nan"
        assert not scales.isinf().any(), "scales has inf"

        # Center alignment: apply to h,w only (not time)
        centers = (sizes - 1) / 2  # [L,3]
        centers[:, 0] = 0  # Do not center the time dimension
        coords_xyz = coords_xyz - centers  # [L,3]

        # Project to frequency bands in one shot: [L,3,B]
        proj = coords_xyz.unsqueeze(-1) * scales.unsqueeze(-1) * self.bands

        # Compute sin & cos and concatenate
        sin_proj = proj.sin()  # [L,3,B]
        cos_proj = proj.cos()

        return torch.cat((sin_proj, cos_proj), dim=1).flatten(1)

    def reset_parameters(self):
        bands = self.get_default_bands()
        self.bands.copy_(bands)

    def get_default_bands(self):
        if self.in_pixels:
            raise NotImplementedError("in_pixels are not implemented yet")
        else:
            bands = freq_bands(self.dim // 8, temperature=self.temperature, step=1, device=self.device).to(self.dtype)
        return bands


class MultiModalityRMSNorm(nn.Module):
    __constants__ = ["dim", "eps", "num_modality"]
    dim: int
    eps: float
    num_modality: int

    def __init__(self, dim: int, eps: float = 1e-6, device: torch.device | None = None, num_modality: int = 1):
        super().__init__()
        self.dim = dim
        self.eps = eps
        self.num_modality = num_modality

        self.weight = torch.nn.Parameter(torch.zeros(dim * num_modality, device=device, dtype=torch.float32))
        if num_modality > 1:
            self.forward = self.forward_multi_experts
        else:
            self.forward = self.forward_single_expert

        self.reset_parameters()

    def reset_parameters(self):
        nn.init.zeros_(self.weight)

    def rms(self, x: torch.Tensor) -> torch.Tensor:
        t, original_dtype = x.float(), x.dtype
        t = t * torch.rsqrt(torch.mean(t**2, dim=-1, keepdim=True) + self.eps)
        return t

    def forward_multi_experts(self, x: torch.Tensor, modality_dispatcher: ModalityDispatcher) -> torch.Tensor:
        original_dtype = x.dtype
        t = self.rms(x)

        weight_chunked = self.weight.chunk(self.num_modality, dim=0)
        outs: list[torch.Tensor] = []
        start = 0
        for i in range(self.num_modality):
            size = modality_dispatcher.group_size_cpu[i]
            if size <= 0:
                continue
            seg = (t[start : start + size] * (weight_chunked[i] + 1)).to(original_dtype)
            outs.append(seg)
            start += size

        return torch.cat(outs, dim=0)

    def forward_single_expert(self, x: torch.Tensor, modality_dispatcher: Optional[ModalityDispatcher] = None) -> torch.Tensor:
        original_dtype = x.dtype
        chunk_str = os.getenv("MAGI_RMSNORM_CHUNK", "").strip()
        try:
            chunk = int(chunk_str) if chunk_str else 0
        except ValueError:
            chunk = 0

        if chunk <= 0 and x.is_cuda and x.dim() == 2:
            if x.numel() >= 20_000_000:
                chunk = 4096

        if chunk > 0 and x.dim() == 2 and x.shape[0] > chunk:
            w = (self.weight + 1).to(torch.float32)
            out = torch.empty_like(x, dtype=original_dtype)
            n = int(x.shape[0])
            for s in range(0, n, int(chunk)):
                e = min(n, s + int(chunk))
                xs = x[s:e].to(torch.float32)
                xs = xs * torch.rsqrt(torch.mean(xs * xs, dim=-1, keepdim=True) + self.eps)
                out[s:e] = (xs * w).to(original_dtype)
            return out

        t = x.float()
        t = t * torch.rsqrt(torch.mean(t**2, dim=-1, keepdim=True) + self.eps)
        return (t * (self.weight + 1)).to(original_dtype)


class _BF16ComputeLinear(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx,
        input: torch.Tensor,
        weight: torch.Tensor,
        bias: Optional[torch.Tensor],
        output_dtype: Optional[torch.dtype],
        compute_dtype: torch.dtype = torch.bfloat16,
    ):
        ctx.compute_dtype = compute_dtype
        ctx.input_dtype = input.dtype
        ctx.weight_dtype = weight.dtype
        ctx.bias_dtype = bias.dtype if bias is not None else None
        ctx.has_bias = bias is not None
        ctx.input_shape = tuple(input.shape)

        input_cast = input.to(compute_dtype)
        ctx.save_for_backward(input_cast, weight)

        weight_cast = weight.to(compute_dtype)
        output = torch.matmul(input_cast, weight_cast.t())

        if bias is not None:
            output = output + bias.to(compute_dtype)

        return output.to(output_dtype)

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor):
        input_cast, weight = ctx.saved_tensors
        compute_dtype: torch.dtype = ctx.compute_dtype
        input_dtype: torch.dtype = ctx.input_dtype
        input_shape = ctx.input_shape

        grad_output_f32 = grad_output.to(torch.float32)
        input_f32 = input_cast.to(torch.float32)
        weight_f32 = weight.to(torch.float32)

        grad_output_2d = grad_output_f32.reshape(-1, grad_output_f32.shape[-1])
        input_2d = input_f32.reshape(-1, input_f32.shape[-1])

        grad_input = grad_output_2d.matmul(weight_f32).reshape(input_shape).to(input_dtype)
        grad_weight = grad_output_2d.t().matmul(input_2d)
        grad_bias = grad_output_2d.sum(dim=0) if ctx.has_bias else None

        return grad_input, grad_weight, grad_bias, None, None


class BaseLinear(nn.Module):
    __constants__ = ["in_features", "out_features", "num_layers", "num_experts"]
    in_features: int
    out_features: int
    num_layers_for_initialization: int
    num_experts: int
    weight: Tensor

    def __init__(
        self, in_features, out_features, num_layers_for_initialization, num_experts, bias=True, device=None, dtype=None
    ):
        super().__init__()
        factory_kwargs = {"device": device, "dtype": torch.bfloat16}
        self.in_features = in_features
        self.out_features = out_features
        self.num_layers_for_initialization = num_layers_for_initialization
        self.num_experts = num_experts
        self.use_bias = bias
        self.weight = Parameter(torch.empty((out_features * num_experts, in_features), **factory_kwargs))
        if bias:
            self.bias = Parameter(torch.empty(out_features * num_experts, **factory_kwargs))
        else:
            self.register_parameter("bias", None)

    def forward(
        self,
        input: torch.Tensor,
        output_dtype: Optional[torch.dtype] = None,
        modality_dispatcher: Optional[ModalityDispatcher] = None,
    ) -> torch.Tensor:
        output_dtype = input.dtype if output_dtype is None else output_dtype
        return _BF16ComputeLinear.apply(input, self.weight, self.bias, output_dtype, torch.bfloat16)


class NativeMoELinear(BaseLinear):
    def forward(
        self,
        input: torch.Tensor,
        output_dtype: Optional[torch.dtype] = None,
        modality_dispatcher: Optional[ModalityDispatcher] = None,
    ) -> torch.Tensor:
        output_dtype = input.dtype if output_dtype is None else output_dtype

        input_list = modality_dispatcher.dispatch(input)  # type: ignore
        weight_chunked = self.weight.chunk(self.num_experts, dim=0)

        if self.bias is not None:
            bias_chunked = self.bias.chunk(self.num_experts, dim=0)

        for i in range(self.num_experts):
            input_list[i] = _BF16ComputeLinear.apply(
                input_list[i],
                weight_chunked[i],
                bias_chunked[i] if self.bias is not None else None,
                output_dtype,
                torch.bfloat16,
            )
        return modality_dispatcher.undispatch(*input_list)  # type: ignore


def create_linear(
    in_features, out_features, num_layers=1, num_experts=1, bias=True, device=None, dtype=None
) -> BaseLinear | NativeMoELinear:
    if num_experts == 1:
        return BaseLinear(in_features, out_features, num_layers, num_experts, bias, device, dtype)
    else:
        return NativeMoELinear(in_features, out_features, num_layers, num_experts, bias, device, dtype)


HAS_MAGI_ATTENTION = importlib.util.find_spec("magi_attention") is not None
HAS_FA3 = importlib.util.find_spec("flash_attn_interface") is not None


@magi_register_custom_op(name="infra::flash_attn_func", is_subgraph_boundary=True)
def flash_attn_func(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
    if query.dim() != 4 or key.dim() != 4 or value.dim() != 4:
        raise ValueError(
            f"flash_attn_func expects (B, S, H, D). Got q={tuple(query.shape)} k={tuple(key.shape)} v={tuple(value.shape)}"
        )

    if HAS_FA3 and is_hopper_arch():
        try:
            from flash_attn_interface import flash_attn_func as fa3_flash_attn_func

            return fa3_flash_attn_func(query, key, value)
        except ModuleNotFoundError:
            pass

    try:
        from flash_attn.flash_attn_interface import flash_attn_func as fa2_flash_attn_func

        return fa2_flash_attn_func(query, key, value)
    except ModuleNotFoundError:
        q = query.permute(0, 2, 1, 3).contiguous()
        k = key.permute(0, 2, 1, 3).contiguous()
        v = value.permute(0, 2, 1, 3).contiguous()
        if k.shape[1] != q.shape[1]:
            if q.shape[1] % k.shape[1] != 0:
                raise RuntimeError(
                    f"GQA fallback expects num_heads_q % num_heads_kv == 0. Got num_heads_q={q.shape[1]} num_heads_kv={k.shape[1]}"
                )
            repeat_factor = q.shape[1] // k.shape[1]
            k = k.repeat_interleave(repeat_factor, dim=1)
            v = v.repeat_interleave(repeat_factor, dim=1)
        out = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False)
        return out.permute(0, 2, 1, 3).contiguous()


def _split_q_range_with_no_overlap(
    q_ranges: torch.Tensor, k_ranges: torch.Tensor
) -> Tuple[List[List[int]], List[List[List[int]]]]:
    range_boundary = torch.unique(q_ranges, sorted=True).tolist()
    candidates = [[start, end, []] for start, end in zip(range_boundary[:-1], range_boundary[1:])]
    q_ranges = q_ranges.tolist()
    k_ranges = k_ranges.tolist()
    for q_range, k_range in zip(q_ranges, k_ranges):
        q_start, q_end = q_range
        for q_range_cand in candidates:
            if q_start <= q_range_cand[0] and q_range_cand[1] <= q_end:
                q_range_cand[2].append(k_range)
    q_ranges_out = []
    k_ranges_out = []
    for q_range_cand in candidates:
        if len(q_range_cand[2]) > 0:
            q_ranges_out.append(q_range_cand[0:2])
            k_ranges_out.append(q_range_cand[2])
    return q_ranges_out, k_ranges_out


def _flash_attn_with_correction(
    query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, q_ranges: List[List[int]], k_range_list: List[List[List[int]]]
):
    output = torch.zeros_like(query)
    output_lse = torch.zeros((query.shape[0], query.shape[1]), dtype=torch.float32, device=query.device)

    from flash_attn.flash_attn_interface import flash_attn_func

    for q_range, k_ranges in zip(q_ranges, k_range_list):
        q_start, q_end = q_range
        qo_out, qo_lse = None, None
        for k_range in k_ranges:
            k_start, k_end = k_range
            cur_qo_out, cur_qo_lse, _ = flash_attn_func(
                query[q_start:q_end].unsqueeze(0),
                key[k_start:k_end].unsqueeze(0),
                value[k_start:k_end].unsqueeze(0),
                return_attn_probs=True,
            )
            cur_qo_out, cur_qo_lse = cur_qo_out.squeeze(0), cur_qo_lse.squeeze(0)

            if qo_out is None:
                qo_out = cur_qo_out
                qo_lse = cur_qo_lse
            else:
                qo_lse[qo_lse == torch.inf] = -torch.inf
                cur_qo_lse[cur_qo_lse == torch.inf] = -torch.inf
                max_lse = torch.max(qo_lse, cur_qo_lse)
                qo_se, cur_qo_se = torch.exp(qo_lse - max_lse), torch.exp(cur_qo_lse - max_lse)
                sum_se = qo_se + cur_qo_se
                qo_scale, cur_qo_scale = qo_se / sum_se, cur_qo_se / sum_se

                qo_out = qo_out * qo_scale.permute(1, 0).unsqueeze(-1) + cur_qo_out * cur_qo_scale.permute(1, 0).unsqueeze(-1)
                qo_lse = torch.log(sum_se) + max_lse

        output[q_start:q_end] = qo_out
        output_lse[q_start:q_end, :] = qo_lse.permute(1, 0)
    return output, output_lse


def _custom_flex_flash_attn_func(
    query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, q_ranges: torch.Tensor, k_ranges: torch.Tensor, **kwargs
):
    q_ranges, k_range_list = _split_q_range_with_no_overlap(q_ranges, k_ranges)
    return _flash_attn_with_correction(query, key, value, q_ranges, k_range_list)


def _flex_flash_attn_func_infer_output_meta(
    query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, q_ranges: torch.Tensor, k_ranges: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor]:
    output = torch.empty_like(query)
    output_lse = torch.empty((query.shape[0], query.shape[1]), dtype=torch.float32, device=query.device)
    return output, output_lse


@magi_register_custom_op(
    name="infra::flex_flash_attn_func",
    mutates_args=(),
    infer_output_meta_fn=_flex_flash_attn_func_infer_output_meta,
    is_subgraph_boundary=True,
)
def flex_flash_attn_func(
    query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, q_ranges: torch.Tensor, k_ranges: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor]:
    if HAS_MAGI_ATTENTION and is_hopper_arch():
        from magi_attention.api import flex_flash_attn_func as magi_flex_flash_attn_func

        return magi_flex_flash_attn_func(query, key, value, q_ranges, k_ranges)
    else:
        return _custom_flex_flash_attn_func(query, key, value, q_ranges, k_ranges)


def _attention_with_cp_infer_output_meta(q: torch.Tensor, *args, **kwargs) -> torch.Tensor:
    return torch.empty_like(q, dtype=torch.bfloat16).squeeze(0)


@magi_register_custom_op(
    name="infra::flash_attn_with_cp",
    mutates_args=(),
    infer_output_meta_fn=_attention_with_cp_infer_output_meta,
    is_subgraph_boundary=True,
)
def flash_attn_with_cp(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, cp_split_sizes: List[int]) -> torch.Tensor:
    q, k, v = q.to(torch.bfloat16), k.to(torch.bfloat16), v.to(torch.bfloat16)

    from inference.infra.distributed import get_cp_group, get_cp_rank, get_cp_world_size
    from inference.infra.parallelism.all_to_all_primitive import batch_scatter_head_gather_seqlen, scatter_seqlen_gather_head
    from inference.utils import print_rank_0

    # One-shot probe so we can tell all_to_all hang from flash_attn hang.
    _probe = not hasattr(flash_attn_with_cp, "_probed")
    if _probe and get_cp_rank() == 0:
        print_rank_0(
            f"flash_attn_with_cp: enter q={tuple(q.shape)} split={cp_split_sizes} cp={get_cp_world_size()}"
        )

    if get_cp_world_size() > 1:
        q, k, v = batch_scatter_head_gather_seqlen([q.squeeze(0), k.squeeze(0), v.squeeze(0)], cp_split_sizes, get_cp_group())
        q = q.unsqueeze(0)
        k = k.unsqueeze(0)
        v = v.unsqueeze(0)

    if _probe:
        torch.cuda.synchronize()
        if get_cp_rank() == 0:
            print_rank_0(f"flash_attn_with_cp: after all_to_all q={tuple(q.shape)}")

    if env_is_true("DAVINCI_DISABLE_MAGI_COMPILE") or env_is_true("MAGI_DISABLE_COMPILE"):
        self_attn_out = flash_attn_func(q, k, v).squeeze(0)
    else:
        self_attn_out = torch.ops.infra.flash_attn_func(q, k, v).squeeze(0)

    if _probe:
        torch.cuda.synchronize()
        if get_cp_rank() == 0:
            print_rank_0(f"flash_attn_with_cp: after flash_attn out={tuple(self_attn_out.shape)}")

    if get_cp_world_size() > 1:
        self_attn_out = scatter_seqlen_gather_head(self_attn_out, cp_split_sizes, get_cp_group(), async_op=False)
        self_attn_out = rearrange(self_attn_out, "(cp sq) hn hd -> sq (cp hn) hd", cp=get_cp_world_size())

    if _probe:
        torch.cuda.synchronize()
        if get_cp_rank() == 0:
            print_rank_0(f"flash_attn_with_cp: exit out={tuple(self_attn_out.shape)}")
        flash_attn_with_cp._probed = True

    return self_attn_out


@magi_register_custom_op(
    name="infra::flex_flash_attn_with_cp",
    mutates_args=(),
    infer_output_meta_fn=_attention_with_cp_infer_output_meta,
    is_subgraph_boundary=True,
)
def flex_flash_attn_with_cp(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    q_ranges: torch.Tensor,
    k_ranges: torch.Tensor,
    cp_split_sizes: List[int],
) -> torch.Tensor:
    q, k, v = q.to(torch.bfloat16).squeeze(0), k.to(torch.bfloat16).squeeze(0), v.to(torch.bfloat16).squeeze(0)

    from inference.infra.distributed import get_cp_group, get_cp_world_size
    from inference.infra.parallelism.all_to_all_primitive import batch_scatter_head_gather_seqlen, scatter_seqlen_gather_head

    if get_cp_world_size() > 1:
        q, k, v = batch_scatter_head_gather_seqlen([q, k, v], cp_split_sizes, get_cp_group())

    if env_is_true("DAVINCI_DISABLE_MAGI_COMPILE") or env_is_true("MAGI_DISABLE_COMPILE"):
        out, _ = flex_flash_attn_func(q, k, v, q_ranges=q_ranges, k_ranges=k_ranges)
    else:
        out, _ = torch.ops.infra.flex_flash_attn_func(q, k, v, q_ranges=q_ranges, k_ranges=k_ranges)

    if get_cp_world_size() > 1:
        out = scatter_seqlen_gather_head(out, cp_split_sizes, get_cp_group(), async_op=False)
        out = rearrange(out, "(cp sq) hn hd -> sq (cp hn) hd", cp=get_cp_world_size())

    return out


@dataclass
class AttentionConfig:
    hidden_size: int
    num_heads_q: int
    num_heads_kv: int
    head_dim: int
    params_dtype: torch.dtype
    checkpoint_qk_layernorm_rope: bool
    num_modality: int
    num_layers: int
    use_local_attn: bool = False
    enable_attn_gating: bool = False


class Attention(torch.nn.Module):
    config: AttentionConfig

    def __init__(self, config: AttentionConfig):
        super().__init__()
        self.config = config

        self.pre_norm = MultiModalityRMSNorm(config.hidden_size, eps=1e-6, num_modality=config.num_modality)
        self.gating_size = config.num_heads_q if config.enable_attn_gating else 0

        self.linear_qkv = create_linear(
            config.hidden_size,
            config.num_heads_q * config.head_dim + config.num_heads_kv * config.head_dim * 2 + self.gating_size,
            num_experts=config.num_modality,
            bias=False,
            dtype=config.params_dtype,
            num_layers=config.num_layers,
        )
        self.linear_proj = create_linear(
            config.num_heads_q * config.head_dim,
            config.hidden_size,
            bias=False,
            num_experts=config.num_modality,
            dtype=config.params_dtype,
            num_layers=config.num_layers,
        )
        self.q_norm = MultiModalityRMSNorm(config.head_dim, num_modality=config.num_modality)
        self.k_norm = MultiModalityRMSNorm(config.head_dim, num_modality=config.num_modality)

        self.q_size = config.num_heads_q * config.head_dim
        self.kv_size = config.num_heads_kv * config.head_dim

    def reset_parameters(self):
        if hasattr(self.linear_proj, "reset_parameters_output_layer"):
            self.linear_proj.reset_parameters_output_layer()

    def forward(
        self,
        hidden_states: torch.Tensor,
        rope: torch.Tensor,
        permute_mapping: torch.Tensor,
        inv_permute_mapping: torch.Tensor,
        varlen_handler: VarlenHandler,
        local_attn_handler: FFAHandler,
        modality_dispatcher: ModalityDispatcher,
        cp_split_sizes: List[int],
    ) -> torch.Tensor:
        hidden_states = self.pre_norm(hidden_states, modality_dispatcher=modality_dispatcher).to(torch.bfloat16)
        qkv: torch.Tensor = self.linear_qkv(hidden_states, modality_dispatcher=modality_dispatcher).to(hidden_states.dtype)

        q, k, v, g = torch.split(qkv, [self.q_size, self.kv_size, self.kv_size, self.gating_size], dim=1)
        q = q.view(-1, self.config.num_heads_q, self.config.head_dim)
        k = k.view(-1, self.config.num_heads_kv, self.config.head_dim)
        v = v.view(-1, self.config.num_heads_kv, self.config.head_dim)
        g = g.view(k.shape[0], self.config.num_heads_q, -1)

        q = self.q_norm(q, modality_dispatcher=modality_dispatcher)
        k = self.k_norm(k, modality_dispatcher=modality_dispatcher)

        q = ModalityDispatcher.inv_permute(q, inv_permute_mapping).unsqueeze(0)
        k = ModalityDispatcher.inv_permute(k, inv_permute_mapping).unsqueeze(0)
        v = ModalityDispatcher.inv_permute(v, inv_permute_mapping).unsqueeze(0)

        sin_emb, cos_emb = rope.tensor_split(2, -1)
        q = apply_rotary_emb_torch(q, cos_emb, sin_emb)
        k = apply_rotary_emb_torch(k, cos_emb, sin_emb)

        if self.config.use_local_attn:
            self_attn_out = flex_flash_attn_with_cp(
                q, k, v, local_attn_handler.q_ranges, local_attn_handler.k_ranges, cp_split_sizes
            )
        else:
            self_attn_out = flash_attn_with_cp(q, k, v, cp_split_sizes)
        self_attn_out = ModalityDispatcher.permute(self_attn_out, permute_mapping)

        if self.config.enable_attn_gating:
            self_attn_out = self_attn_out * torch.sigmoid(g)

        self_attn_out = self_attn_out.view(-1, self.config.num_heads_q * self.config.head_dim).to(torch.bfloat16)
        out = self.linear_proj(self_attn_out, modality_dispatcher=modality_dispatcher)
        return out


@dataclass
class MLPConfig:
    hidden_size: int
    intermediate_size: int
    activation_type: MLPActivationType
    params_dtype: torch.dtype
    num_modality: int = 1
    num_layers: int = 1
    gated_act: bool = False


class MLP(torch.nn.Module):
    config: MLPConfig

    def __init__(self, config: MLPConfig):
        super().__init__()
        num_experts = config.num_modality
        self.pre_norm = MultiModalityRMSNorm(config.hidden_size, num_modality=config.num_modality)
        intermediate_size_up = config.intermediate_size * 2 if config.gated_act else config.intermediate_size

        self.up_gate_proj = create_linear(
            config.hidden_size,
            intermediate_size_up,
            bias=False,
            dtype=config.params_dtype,
            num_layers=config.num_layers,
            num_experts=num_experts,
        )
        self.down_proj = create_linear(
            config.intermediate_size,
            config.hidden_size,
            bias=False,
            dtype=config.params_dtype,
            num_layers=config.num_layers,
            num_experts=num_experts,
        )
        self.activation_func = create_activation_func(config.activation_type)

    def forward(self, x: torch.Tensor, modality_dispatcher: ModalityDispatcher) -> torch.Tensor:
        x = self.pre_norm(x, modality_dispatcher=modality_dispatcher).to(torch.bfloat16)
        x = self.up_gate_proj(x, modality_dispatcher=modality_dispatcher).to(torch.bfloat16)
        x = self.activation_func(x).to(torch.bfloat16)
        x = self.down_proj(x, modality_dispatcher=modality_dispatcher).to(torch.bfloat16)
        return x

    def extra_repr(self) -> str:
        return f"{self.up_gate_proj.weight.shape=}, {self.down_proj.weight.shape=}"


class ExpertMLP(torch.nn.Module):
    """Single dense expert FFN (SwiGLU7 / GELU7), used inside SparseMoEMLP."""

    def __init__(
        self,
        *,
        hidden_size: int,
        intermediate_size: int,
        activation_type: MLPActivationType,
        gated_act: bool,
        params_dtype: torch.dtype,
    ):
        super().__init__()
        self.hidden_size = int(hidden_size)
        self.intermediate_size = int(intermediate_size)
        self.gated_act = bool(gated_act)
        intermediate_size_up = intermediate_size * 2 if gated_act else intermediate_size
        self.up_gate_proj = nn.Linear(hidden_size, intermediate_size_up, bias=False, dtype=params_dtype)
        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False, dtype=params_dtype)
        self.activation_func = create_activation_func(activation_type)

    def forward_with_intermediate(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        projected = self.up_gate_proj(x.to(torch.bfloat16)).to(torch.bfloat16)
        z = self.activation_func(projected).to(torch.bfloat16)
        y = self.down_proj(z).to(torch.bfloat16)
        return y, z

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        y, _ = self.forward_with_intermediate(x)
        return y


@dataclass
class SparseMoEConfig:
    hidden_size: int
    intermediate_size: int
    activation_type: MLPActivationType
    params_dtype: torch.dtype
    num_experts: int = 3
    top_k: int = 1
    gated_act: bool = True
    router_noise: float = 0.0
    num_modality: int = 1
    layer_idx: int = -1


def token_sparse_moe_mlp_layers(
    sparse_moe_layers: Optional[List[int]],
    mm_layers: Optional[List[int]] = None,
    convert_mm_mlp: bool = False,
) -> List[int]:
    """MLP layers that use token-level SparseMoEMLP (learned top-k router).

    When ``convert_mm_mlp`` is true and ``sparse_moe_layers`` is non-empty,
    mm-layer MLPs are converted too: modality-hard NativeMoELinear is replaced
    by per-token top-k routing. Attention on mm_layers stays per-modality.
    Empty ``sparse_moe_layers`` leaves the MagiHuman MLP unchanged.
    """
    sparse = [int(x) for x in (sparse_moe_layers or [])]
    if not sparse:
        return []
    if not convert_mm_mlp:
        return sorted(set(sparse))
    mm = [int(x) for x in (mm_layers or [])]
    return sorted(set(sparse) | set(mm))


def _env_optional_int(name: str) -> Optional[int]:
    raw = os.environ.get(name, "").strip()
    if not raw:
        return None
    return int(raw)


class SparseMoEMLP(torch.nn.Module):
    """
    Token-level sparse MoE MLP.

    Each token independently scores E experts with Linear(hidden -> E),
    softmax, then top-k (default k=1). Only selected experts run on that
    token; outputs are gate-weighted. This is Mixtral/Switch-style routing,
    not MagiHuman modality-hard NativeMoELinear.

    - Router: Linear(hidden -> E), top-k softmax gates
    - Experts: E independent ExpertMLP
    - Aux loss: Switch-Transformer style load balancing on top-1 assignment
    - pre_norm may keep per-modality RMSNorm affine on mm_layers (not routing)
    """

    def __init__(self, config: SparseMoEConfig):
        super().__init__()
        if config.num_experts < 1:
            raise ValueError(f"num_experts must be >= 1, got {config.num_experts}")
        if config.top_k < 1 or config.top_k > config.num_experts:
            raise ValueError(f"top_k must be in [1, num_experts], got top_k={config.top_k}, E={config.num_experts}")

        self.num_experts = int(config.num_experts)
        self.top_k = int(config.top_k)
        self.router_noise = float(config.router_noise)
        self.layer_idx = int(getattr(config, "layer_idx", -1))
        self.pre_norm = MultiModalityRMSNorm(
            config.hidden_size, num_modality=int(getattr(config, "num_modality", 1) or 1)
        )
        self.router = nn.Linear(config.hidden_size, self.num_experts, bias=False, dtype=torch.float32)
        self.experts = nn.ModuleList(
            [
                ExpertMLP(
                    hidden_size=config.hidden_size,
                    intermediate_size=config.intermediate_size,
                    activation_type=config.activation_type,
                    gated_act=config.gated_act,
                    params_dtype=config.params_dtype,
                )
                for _ in range(self.num_experts)
            ]
        )
        self._last_aux_loss: Optional[torch.Tensor] = None
        self._last_expert_frac: Optional[torch.Tensor] = None
        self._last_router_entropy: Optional[torch.Tensor] = None
        self._last_mod_counts: Optional[torch.Tensor] = None
        self._mid_observer = None
        self.diffusion_timestep = None
        self.diffusion_step_index = None
        self.diffusion_num_steps = None
        self.reset_router_parameters()

    def set_mid_observer(self, observer) -> None:
        """Optional routed-token observer: observer(expert_id, modality, gate, z, y)."""
        self._mid_observer = observer

    def reset_router_parameters(self):
        nn.init.normal_(self.router.weight, mean=0.0, std=0.01)

    def _compute_aux_loss(self, probs: torch.Tensor, top1_idx: torch.Tensor) -> torch.Tensor:
        # Switch Transformer: aux = E * sum_i (f_i * P_i)
        # f_i: fraction of tokens whose top-1 expert is i
        # P_i: mean router probability for expert i
        num_tokens = probs.shape[0]
        if num_tokens == 0:
            zeros = probs.new_zeros(())
            self._last_expert_frac = probs.new_zeros((self.num_experts,))
            self._last_router_entropy = zeros
            return zeros
        one_hot = F.one_hot(top1_idx, num_classes=self.num_experts).to(probs.dtype)
        f = one_hot.mean(dim=0)
        p = probs.mean(dim=0)
        eps = probs.new_tensor(1e-9)
        self._last_expert_frac = f.detach()
        self._last_router_entropy = (-(f.clamp_min(eps) * f.clamp_min(eps).log()).sum()).detach()
        return self.num_experts * torch.sum(f * p)

    def forward(self, x: torch.Tensor, modality_dispatcher: Optional[ModalityDispatcher] = None) -> torch.Tensor:
        _sp = stage_profile()
        _router_tok = _sp.start(f"layer.{self.layer_idx}.moe.router")
        x_norm = self.pre_norm(x, modality_dispatcher=modality_dispatcher).to(torch.bfloat16)
        logits = self.router(x_norm.float())
        if self.training and self.router_noise > 0:
            logits = logits + torch.randn_like(logits) * self.router_noise
        # Infer-only ablations. Training ignores these env vars.
        if not self.training:
            dropped = _env_optional_int("MAGI_MOE_DROP_EXPERT")
            if dropped is not None:
                if dropped < 0 or dropped >= self.num_experts:
                    raise ValueError(
                        f"MAGI_MOE_DROP_EXPERT={dropped} out of range [0, {self.num_experts})"
                    )
                logits = logits.clone()
                logits[:, dropped] = float("-inf")
        probs = F.softmax(logits, dim=-1)
        topk_vals, topk_idx = torch.topk(probs, k=self.top_k, dim=-1)
        topk_vals = topk_vals / topk_vals.sum(dim=-1, keepdim=True).clamp_min(1e-9)
        if not self.training:
            forced = _env_optional_int("MAGI_MOE_FORCE_EXPERT")
            if forced is not None:
                if forced < 0 or forced >= self.num_experts:
                    raise ValueError(
                        f"MAGI_MOE_FORCE_EXPERT={forced} out of range [0, {self.num_experts})"
                    )
                n = probs.shape[0]
                topk_idx = torch.full(
                    (n, self.top_k), int(forced), device=probs.device, dtype=torch.long
                )
                topk_vals = torch.zeros((n, self.top_k), device=probs.device, dtype=probs.dtype)
                topk_vals[:, 0] = 1.0

        if self.training:
            self._last_aux_loss = self._compute_aux_loss(probs, topk_idx[:, 0])
        else:
            self._last_aux_loss = None
            self._last_expert_frac = None
            self._last_router_entropy = None

        # Observe-only: does not change logits, top-k, or expert compute.
        if env_is_true("MAGI_MOE_TRACE") and modality_dispatcher is not None:
            top1 = topk_idx[:, 0]
            mods = modality_dispatcher.permuted_modality_mapping
            if mods is not None and mods.shape[0] == top1.shape[0]:
                idx = top1.long() * 3 + mods.long().clamp(0, 2)
                self._last_mod_counts = torch.bincount(idx, minlength=self.num_experts * 3).reshape(
                    self.num_experts, 3
                ).detach()
            else:
                self._last_mod_counts = None
        else:
            self._last_mod_counts = None

        _sp.stop(_router_tok)
        _experts_tok = _sp.start(f"layer.{self.layer_idx}.moe.experts")
        mode = self._grouped_mode()
        if mode in ("bmm", "concat"):
            out = self._experts_grouped(x_norm, topk_idx, topk_vals, mode)
            _sp.stop(_experts_tok)
            return out

        out = torch.zeros_like(x_norm)
        # Loop experts; gather tokens that selected this expert in any of top-k slots.
        for expert_id, expert in enumerate(self.experts):
            # [N, K] -> which slots match this expert
            match = topk_idx == expert_id
            if not match.any():
                continue
            token_gate = (topk_vals * match.to(topk_vals.dtype)).sum(dim=-1)  # [N]
            token_mask = token_gate > 0
            if not token_mask.any():
                continue
            idx = token_mask.nonzero(as_tuple=False).squeeze(-1)
            x_expert = x_norm.index_select(0, idx)
            if self._mid_observer is None:
                y = expert(x_expert)
            else:
                if modality_dispatcher is None:
                    raise RuntimeError("mid observer requires modality_dispatcher")
                y, z = expert.forward_with_intermediate(x_expert)
                modalities = modality_dispatcher.permuted_modality_mapping.index_select(0, idx)
                gate_obs = token_gate.index_select(0, idx)
                self._mid_observer(
                    expert_id,
                    modalities.detach(),
                    gate_obs.detach(),
                    z.detach(),
                    y.detach(),
                )
            gate = token_gate.index_select(0, idx).unsqueeze(-1).to(y.dtype)
            out.index_add_(0, idx, y * gate)
        _sp.stop(_experts_tok)
        return out

    def extra_repr(self) -> str:
        return (
            f"layer_idx={self.layer_idx}, num_experts={self.num_experts}, top_k={self.top_k}, "
            f"router_noise={self.router_noise}, num_modality={self.pre_norm.num_modality}"
        )

    # --- grouped-GEMM expert execution (probe) ---------------------------
    def _grouped_mode(self) -> str:
        """Return 'off' / 'bmm' / 'concat'.

        Groups the per-expert python-loop GEMMs into one batched GEMM when
        experts share identical shapes and routing is top-1. Controlled by
        env ``MAGI_MOE_GROUPED`` (off|bmm|concat). An instance override
        ``_moe_grouped_override`` wins for tests/benchmarks.

        Default (``auto``) is: use ``bmm`` during inference when it applies,
        and always fall back to the original loop for training / collector
        (``_mid_observer``) / non-uniform experts. Set ``off`` to disable.
        """
        if self.training or self._mid_observer is not None:
            return "off"
        if self.top_k != 1 or len(self.experts) < 2:
            return "off"
        up = [e.up_gate_proj.out_features for e in self.experts]
        dn = [e.down_proj.in_features for e in self.experts]
        if len(set(up)) != 1 or len(set(dn)) != 1:
            return "off"
        override = getattr(self, "_moe_grouped_override", None)
        mode = override if override is not None else os.getenv("MAGI_MOE_GROUPED", "auto").strip().lower()
        if mode in ("0", "off", "false"):
            return "off"
        if mode in ("2", "concat"):
            return "concat"
        if mode in ("1", "bmm", "auto"):
            return "bmm"
        return "bmm"

    def _experts_grouped(self, x_norm: torch.Tensor, topk_idx: torch.Tensor, topk_vals: torch.Tensor, mode: str) -> torch.Tensor:
        """Numerically-equivalent grouped-GEMM replacement of the expert loop.

        Both modes preserve the reference semantics: token t is transformed by
        exactly its top-1 expert and weighted by the routing gate.
        - ``bmm``: batch-GEMM over experts (no FLOP overcompute; pads each
          expert's tokens to the same M for a single cublas batched launch).
        - ``concat``: one giant GEMM over all tokens x all experts (max kernel
          size, E x FLOP overcompute; non-selected slices are zeroed).
        """
        e0 = self.experts[0]
        ids = topk_idx[:, 0]
        g = topk_vals[:, 0]
        n_tok = x_norm.shape[0]
        ecount = self.num_experts
        dtype = x_norm.dtype

        def _bmm():
            order = ids.argsort(stable=True)
            ids_s = ids[order]
            x_s = x_norm[order]
            g_s = g.to(x_norm.dtype)[order]
            counts = torch.bincount(ids_s, minlength=ecount)
            offsets = torch.cumsum(counts, 0) - counts
            col = torch.arange(n_tok, device=x_norm.device) - offsets[ids_s]
            m_max = int(counts.max().item())
            if m_max == 0:
                return torch.zeros_like(x_norm)
            x_pad = x_norm.new_zeros(ecount, m_max, x_norm.shape[1])
            g_pad = x_norm.new_zeros(ecount, m_max)
            x_pad[ids_s, col] = x_s
            g_pad[ids_s, col] = g_s
            w_up = torch.stack([e.up_gate_proj.weight.to(dtype) for e in self.experts], dim=0)
            z = torch.bmm(x_pad, w_up.transpose(-2, -1)).to(dtype)
            z = e0.activation_func(z)
            if z.dtype != dtype:
                z = z.to(dtype)
            w_dn = torch.stack([e.down_proj.weight.to(dtype) for e in self.experts], dim=0)
            y = torch.bmm(z, w_dn.transpose(-2, -1)).to(dtype)
            y = y * g_pad.unsqueeze(-1)
            y_s = y[ids_s, col]
            y_orig = y_s.new_empty(n_tok, y_s.shape[1])
            y_orig[order] = y_s
            return y_orig

        def _concat():
            dd = e0.down_proj.in_features
            w_up = torch.cat([e.up_gate_proj.weight.to(dtype) for e in self.experts], dim=0)  # [E*du, H]
            z = (x_norm @ w_up.t()).to(dtype)                                       # [N, E*du]
            z = e0.activation_func(z)                                               # [N, E*dd]
            if z.dtype != dtype:
                z = z.to(dtype)
            # owned = per-token slice of its expert (block ids*dd .. ids*dd+dd)
            owned = z.view(n_tok, ecount, dd).gather(
                1, ids.view(n_tok, 1, 1).expand(n_tok, 1, dd)
            ).squeeze(1)                                                            # [N, dd]
            col_idx = (ids * dd).unsqueeze(1) + torch.arange(dd, device=x_norm.device).unsqueeze(0)
            z_masked = z.new_zeros(n_tok, ecount * dd)
            z_masked.scatter_(1, col_idx, owned)                                    # zero other experts
            w_dn = torch.cat([e.down_proj.weight.to(dtype).t() for e in self.experts], dim=0)  # [E*dd, H]
            y = (z_masked @ w_dn).to(dtype) * g.to(dtype).unsqueeze(-1)                 # [N, H]
            return y

        if mode == "concat":
            return _concat()
        return _bmm()



def _as_timestep_float(timestep) -> float:
    if torch.is_tensor(timestep):
        return float(timestep.detach().float().reshape(-1)[0].item())
    return float(timestep)


def set_sparse_moe_diffusion_timestep(
    module: torch.nn.Module,
    timestep,
    step_index: Optional[int] = None,
    num_steps: Optional[int] = None,
) -> None:
    """Write the current diffusion step onto every SparseMoEMLP.

    Pruning collectors read these fields; inference is unchanged when unused.
    """
    t = _as_timestep_float(timestep)
    idx = None if step_index is None else int(step_index)
    n = None if num_steps is None else int(num_steps)
    for m in module.modules():
        if isinstance(m, SparseMoEMLP):
            m.diffusion_timestep = t
            m.diffusion_step_index = idx
            m.diffusion_num_steps = n


def collect_moe_aux_loss(module: torch.nn.Module) -> Optional[torch.Tensor]:
    """Average aux loss over all SparseMoEMLP submodules that ran in the last forward."""
    losses: List[torch.Tensor] = []
    for m in module.modules():
        if isinstance(m, SparseMoEMLP) and m._last_aux_loss is not None:
            losses.append(m._last_aux_loss)
            m._last_aux_loss = None
    if not losses:
        return None
    return torch.stack([x.float() for x in losses]).mean()


def collect_moe_routing_stats(module: torch.nn.Module) -> Optional[dict]:
    """
    Aggregate token-routing stats across SparseMoEMLP layers from the last forward.
    Returns mean expert fractions, mean router entropy, and load-imbalance ratio.
    """
    fracs: List[torch.Tensor] = []
    ents: List[torch.Tensor] = []
    for m in module.modules():
        if isinstance(m, SparseMoEMLP) and m._last_expert_frac is not None:
            fracs.append(m._last_expert_frac.float())
            if m._last_router_entropy is not None:
                ents.append(m._last_router_entropy.float())
            m._last_expert_frac = None
            m._last_router_entropy = None
    if not fracs:
        return None
    mean_frac = torch.stack(fracs, dim=0).mean(dim=0)
    mean_ent = torch.stack(ents).mean() if ents else mean_frac.new_zeros(())
    eps = 1e-9
    imbalance = float((mean_frac.max() / mean_frac.clamp_min(eps).min()).item())
    return {
        "expert_frac": [float(x) for x in mean_frac.detach().cpu().tolist()],
        "router_entropy": float(mean_ent.detach().item()),
        "load_imbalance": imbalance,
    }


def remap_dense_mlp_state_to_sparse_moe(
    dense_state: dict,
    *,
    sparse_moe_layers: List[int],
    num_experts: int,
    expert_noise_std: float = 1e-3,
    mm_layers: Optional[List[int]] = None,
    packed_modality_experts: int = 3,
    convert_mm_mlp: bool = False,
) -> dict:
    """
    Remap dense / modality-hard MLP keys into SparseMoEMLP expert[0..E-1].

    Dense middle MLP:
      block.layers.{i}.mlp.up_gate_proj.weight / down_proj.weight
      -> copy into every expert; expert 1..E-1 get Gaussian noise

    mm_layers NativeMoELinear (packed ``out * 3, in``):
      chunk along out dim into 3 modality experts, copy 1:1 into SparseMoEMLP
      experts (video/audio/text weights become the 3 routed experts; router
      is left absent so it random-inits). Extra experts beyond 3 copy expert 0
      plus noise.

    pre_norm.weight is kept (shape already matches SparseMoEMLP when
    mm_layers use num_modality=3). router.weight is left absent.
    """
    target_layers = token_sparse_moe_mlp_layers(
        sparse_moe_layers, mm_layers, convert_mm_mlp=convert_mm_mlp
    )
    mm_set = set(int(x) for x in (mm_layers or []))
    n_packed = int(packed_modality_experts)
    out = dict(dense_state)
    for layer_idx in target_layers:
        prefix = f"block.layers.{layer_idx}.mlp."
        if f"{prefix}experts.0.up_gate_proj.weight" in dense_state:
            continue
        up_key = prefix + "up_gate_proj.weight"
        down_key = prefix + "down_proj.weight"
        if up_key not in dense_state or down_key not in dense_state:
            continue
        up_w = dense_state[up_key]
        down_w = dense_state[down_key]
        out.pop(up_key, None)
        out.pop(down_key, None)
        if layer_idx in mm_set:
            if up_w.shape[0] % n_packed != 0 or down_w.shape[0] % n_packed != 0:
                raise ValueError(
                    f"mm layer {layer_idx} MLP is not packed into {n_packed} experts: "
                    f"up={tuple(up_w.shape)} down={tuple(down_w.shape)}"
                )
            up_chunks = up_w.chunk(n_packed, dim=0)
            down_chunks = down_w.chunk(n_packed, dim=0)
            for e in range(num_experts):
                src = min(e, n_packed - 1)
                up_e = up_chunks[src].clone()
                down_e = down_chunks[src].clone()
                if e >= n_packed and expert_noise_std > 0:
                    up_e = up_e + torch.randn_like(up_e) * expert_noise_std
                    down_e = down_e + torch.randn_like(down_e) * expert_noise_std
                out[f"{prefix}experts.{e}.up_gate_proj.weight"] = up_e
                out[f"{prefix}experts.{e}.down_proj.weight"] = down_e
        else:
            for e in range(num_experts):
                up_e = up_w.clone()
                down_e = down_w.clone()
                if e > 0 and expert_noise_std > 0:
                    up_e = up_e + torch.randn_like(up_e) * expert_noise_std
                    down_e = down_e + torch.randn_like(down_e) * expert_noise_std
                out[f"{prefix}experts.{e}.up_gate_proj.weight"] = up_e
                out[f"{prefix}experts.{e}.down_proj.weight"] = down_e
    return out


def load_dense_weights_into_sparse_moe_model(
    model: torch.nn.Module,
    dense_state: dict,
    *,
    sparse_moe_layers: List[int],
    num_experts: int,
    expert_noise_std: float = 1e-3,
    mm_layers: Optional[List[int]] = None,
    convert_mm_mlp: bool = False,
) -> Tuple[list, list]:
    """Load a dense DiT state_dict into a Sparse-MoE DiT (strict=False after remap).

    Keys whose tensor shapes do not match the target model (e.g. audio Linear
    64→128 when adapting OmniVAE) are dropped before load_state_dict; PyTorch
    treats size mismatches as hard errors even with strict=False.
    """
    remapped = remap_dense_mlp_state_to_sparse_moe(
        dense_state,
        sparse_moe_layers=sparse_moe_layers,
        num_experts=num_experts,
        expert_noise_std=expert_noise_std,
        mm_layers=mm_layers,
        convert_mm_mlp=convert_mm_mlp,
    )
    model_sd = model.state_dict()
    filtered = {}
    shape_skipped = []
    for key, value in remapped.items():
        if key in model_sd and hasattr(value, "shape") and tuple(value.shape) != tuple(model_sd[key].shape):
            shape_skipped.append(key)
            continue
        filtered[key] = value
    if shape_skipped:
        # Attached for callers that want to log OmniVAE / channel-change skips.
        model._omnivae_shape_skipped_keys = list(shape_skipped)  # type: ignore[attr-defined]
    missing, unexpected = model.load_state_dict(filtered, strict=False)
    unexpected = list(unexpected) + [f"{k} (shape_mismatch_skipped)" for k in shape_skipped]
    return list(missing), unexpected


@dataclass
class AdapterConfig:
    hidden_size: int
    num_attention_heads: int
    text_in_channels: int
    video_in_channels: int
    audio_in_channels: int
    params_dtype: torch.dtype


class Adapter(torch.nn.Module):
    config: AdapterConfig

    def __init__(self, config: AdapterConfig):
        super().__init__()
        self.config = config
        self.video_embedder = nn.Linear(config.video_in_channels, config.hidden_size, bias=True, dtype=torch.float32)
        self.text_embedder = nn.Linear(config.text_in_channels, config.hidden_size, bias=True, dtype=torch.float32)
        self.audio_embedder = nn.Linear(config.audio_in_channels, config.hidden_size, bias=True, dtype=torch.float32)
        self.rope = ElementWiseFourierEmbed(config.hidden_size // config.num_attention_heads, in_pixels=False, learnable=False)

    def forward(
        self,
        x: torch.Tensor,
        coords_mapping: torch.Tensor,
        video_mask: torch.Tensor,
        audio_mask: torch.Tensor,
        text_mask: torch.Tensor,
    ):
        rope = self.rope(coords_mapping)
        if torch.is_autocast_enabled() and x.is_cuda:
            out_dtype = torch.get_autocast_gpu_dtype()
        else:
            out_dtype = x.dtype
        output_x = torch.zeros(x.shape[0], self.config.hidden_size, device=x.device, dtype=out_dtype)
        output_x[text_mask] = self.text_embedder(x[text_mask, : self.config.text_in_channels]).to(output_x.dtype)
        output_x[audio_mask] = self.audio_embedder(x[audio_mask, : self.config.audio_in_channels]).to(output_x.dtype)
        output_x[video_mask] = self.video_embedder(x[video_mask, : self.config.video_in_channels]).to(output_x.dtype)
        return output_x, rope


class TransFormerLayer(torch.nn.Module):
    def __init__(self, config: Any, layer_idx: int):
        super().__init__()
        self.layer_idx = int(layer_idx)
        sparse_moe_layers = [int(x) for x in getattr(config, "sparse_moe_layers", []) or []]
        mm_layers = [int(x) for x in getattr(config, "mm_layers", []) or []]
        convert_mm_mlp = bool(getattr(config, "token_sparse_moe_convert_mm_mlp", False))
        mlp_moe_layers = set(
            token_sparse_moe_mlp_layers(sparse_moe_layers, mm_layers, convert_mm_mlp=convert_mm_mlp)
        )
        # Token-level SparseMoEMLP on listed middle MLPs, and mm-layer MLPs when convert_mm_mlp.
        # Attention on mm_layers stays per-modality NativeMoELinear.
        self.use_sparse_moe = layer_idx in mlp_moe_layers

        num_modality = 3 if layer_idx in set(mm_layers) else 1
        use_local_attn = layer_idx in config.local_attn_layers
        self.post_norm = layer_idx in config.post_norm_layers
        attention_config = AttentionConfig(
            hidden_size=config.hidden_size,
            num_heads_q=config.num_heads_q,
            num_heads_kv=config.num_heads_kv,
            head_dim=config.head_dim,
            params_dtype=config.params_dtype,
            checkpoint_qk_layernorm_rope=config.checkpoint_qk_layernorm_rope,
            num_modality=num_modality,
            num_layers=config.num_layers,
            use_local_attn=use_local_attn,
            enable_attn_gating=config.enable_attn_gating,
        )
        self.attention: Attention = Attention(attention_config)

        activation_type = MLPActivationType.GELU7 if layer_idx in config.gelu7_layers else MLPActivationType.SWIGLU7
        if activation_type == MLPActivationType.SWIGLU7:
            gated_act = True
            intermediate_size = int(config.hidden_size * 4 * 2 / 3) // 4 * 4
        else:
            gated_act = False
            intermediate_size = config.hidden_size * 4

        if self.use_sparse_moe:
            # Keep the per-expert MLP structure aligned with the original dense MLP,
            # while routing behavior is controlled by the training config.
            sparse_cfg = SparseMoEConfig(
                hidden_size=config.hidden_size,
                intermediate_size=intermediate_size,
                activation_type=activation_type,
                params_dtype=config.params_dtype,
                num_experts=int(getattr(config, "num_routed_experts", 3)),
                top_k=int(getattr(config, "moe_top_k", 1)),
                gated_act=gated_act,
                router_noise=float(getattr(config, "moe_router_noise", 0.0)),
                num_modality=num_modality,
                layer_idx=int(layer_idx),
            )
            self.mlp = SparseMoEMLP(sparse_cfg)
        else:
            mlp_config = MLPConfig(
                hidden_size=config.hidden_size,
                intermediate_size=intermediate_size,
                activation_type=activation_type,
                params_dtype=config.params_dtype,
                num_modality=num_modality,
                num_layers=config.num_layers,
                gated_act=gated_act,
            )
            self.mlp = MLP(mlp_config)

        if self.post_norm:
            self.attn_post_norm = MultiModalityRMSNorm(config.hidden_size, num_modality=num_modality)
            self.mlp_post_norm = MultiModalityRMSNorm(config.hidden_size, num_modality=num_modality)

    def forward(
        self,
        hidden_states: torch.Tensor,
        rope: torch.Tensor,
        permute_mapping: torch.Tensor,
        inv_permute_mapping: torch.Tensor,
        varlen_handler: VarlenHandler,
        local_attn_handler: FFAHandler,
        modality_dispatcher: ModalityDispatcher,
        cp_split_sizes: List[int],
    ) -> torch.Tensor:
        _sp = stage_profile()
        _attn_tok = _sp.start(f"layer.{self.layer_idx}.attn")
        attn_out = self.attention(
            hidden_states,
            rope,
            permute_mapping,
            inv_permute_mapping,
            varlen_handler,
            local_attn_handler,
            modality_dispatcher,
            cp_split_sizes,
        )
        if self.post_norm:
            attn_out = self.attn_post_norm(attn_out, modality_dispatcher=modality_dispatcher)
        _sp.stop(_attn_tok)
        hidden_states = hidden_states + attn_out

        _mlp_tok = _sp.start(f"layer.{self.layer_idx}.mlp.{'moe' if self.use_sparse_moe else 'dense'}")
        mlp_out = self.mlp(hidden_states, modality_dispatcher)
        if self.post_norm:
            mlp_out = self.mlp_post_norm(mlp_out, modality_dispatcher=modality_dispatcher)
        _sp.stop(_mlp_tok)
        hidden_states = hidden_states + mlp_out
        return hidden_states


is_base_model = True


def config_patch(compile_config: CompileConfig) -> CompileConfig:
    global is_base_model
    if is_base_model:
        is_base_model = False
    else:
        # Fully offload SR model for memory-constrained GPU
        compile_config.offload_config.gpu_resident_weight_ratio = 0.0
    return compile_config


@magi_compile(config_patch=config_patch)
class TransformerBlock(torch.nn.Module):
    def __init__(self, model_config: Any):
        super().__init__()
        self.layers: list[TransFormerLayer] = nn.ModuleList()
        for layer_idx in range(model_config.num_layers):
            self.layers.append(TransFormerLayer(model_config, layer_idx))

    def forward(
        self,
        x: torch.Tensor,
        rope: torch.Tensor,
        permute_mapping: torch.Tensor,
        inv_permute_mapping: torch.Tensor,
        varlen_handler: VarlenHandler,
        local_attn_handler: FFAHandler,
        modality_dispatcher: ModalityDispatcher,
        cp_split_sizes: List[int],
    ) -> torch.Tensor:
        import torch.utils.checkpoint as checkpoint
        _sp = stage_profile()
        _blk_tok = _sp.start("block.total")
        use_checkpointing = torch.is_grad_enabled() and os.getenv("MAGI_GRADIENT_CHECKPOINTING", "1") == "1"

        def create_custom_forward(module):
            def custom_forward(x_in, rope_in, permute_mapping_in, inv_permute_mapping_in):
                return module(
                    x_in,
                    rope_in,
                    permute_mapping_in,
                    inv_permute_mapping_in,
                    varlen_handler,
                    local_attn_handler,
                    modality_dispatcher,
                    cp_split_sizes,
                )
            return custom_forward

        for _, layer in enumerate(self.layers):
            if use_checkpointing:
                x = checkpoint.checkpoint(
                    create_custom_forward(layer),
                    x,
                    rope,
                    permute_mapping,
                    inv_permute_mapping,
                    use_reentrant=False
                )
            else:
                x = layer(
                    x,
                    rope,
                    permute_mapping,
                    inv_permute_mapping,
                    varlen_handler,
                    local_attn_handler,
                    modality_dispatcher,
                    cp_split_sizes,
                )
        _sp.stop(_blk_tok)
        return x


@dataclass
class TransformerConfig:
    hidden_size: int
    video_in_channels: int
    audio_in_channels: int
    text_in_channels: int
    params_dtype: torch.dtype
    post_process_dtype: torch.dtype


class DiTModel(torch.nn.Module):
    config: TransformerConfig

    def __init__(self, model_config: Any):
        super().__init__()
        self.config = TransformerConfig(
            hidden_size=model_config.hidden_size,
            video_in_channels=model_config.video_in_channels,
            audio_in_channels=model_config.audio_in_channels,
            text_in_channels=model_config.text_in_channels,
            params_dtype=model_config.params_dtype,
            post_process_dtype=torch.float32,
        )
        adapter_config = AdapterConfig(
            hidden_size=model_config.hidden_size,
            num_attention_heads=model_config.num_heads_q,
            text_in_channels=model_config.text_in_channels,
            video_in_channels=model_config.video_in_channels,
            audio_in_channels=model_config.audio_in_channels,
            params_dtype=torch.float32,
        )
        self.adapter: Adapter = Adapter(adapter_config)
        self.block: TransformerBlock = TransformerBlock(model_config=model_config)
        self.final_norm_video = MultiModalityRMSNorm(self.config.hidden_size)
        self.final_norm_audio = MultiModalityRMSNorm(self.config.hidden_size)
        self.final_linear_video = nn.Linear(
            self.config.hidden_size, self.config.video_in_channels, bias=False, dtype=torch.float32
        )
        self.final_linear_audio = nn.Linear(
            self.config.hidden_size, self.config.audio_in_channels, bias=False, dtype=torch.float32
        )

    def set_diffusion_timestep(self, timestep, step_index=None, num_steps=None) -> None:
        set_sparse_moe_diffusion_timestep(self, timestep, step_index=step_index, num_steps=num_steps)

    def forward(
        self,
        x: torch.Tensor,
        coords_mapping: torch.Tensor,
        modality_mapping: torch.Tensor,
        varlen_handler: VarlenHandler,
        local_attn_handler: FFAHandler,
    ):
        _sp = stage_profile()
        _model_tok = _sp.start("model.total")
        x = ulysses_scheduler().dispatch(x)
        coords_mapping = ulysses_scheduler().dispatch(coords_mapping)
        modality_mapping = ulysses_scheduler().dispatch(modality_mapping)
        cp_split_sizes = ulysses_scheduler().cp_split_sizes

        modality_dispatcher = ModalityDispatcher(modality_mapping, 3)
        permute_mapping, inv_permute_mapping = modality_dispatcher.permute_mapping, modality_dispatcher.inv_permute_mapping
        video_mask = modality_mapping == Modality.VIDEO
        audio_mask = modality_mapping == Modality.AUDIO
        text_mask = modality_mapping == Modality.TEXT

        x, rope = self.adapter(x, coords_mapping, video_mask, audio_mask, text_mask)
        x = x.to(self.config.params_dtype)
        x = ModalityDispatcher.permute(x, permute_mapping)
        x = self.block(
            x,
            rope,
            permute_mapping=permute_mapping,
            inv_permute_mapping=inv_permute_mapping,
            varlen_handler=varlen_handler,
            local_attn_handler=local_attn_handler,
            modality_dispatcher=modality_dispatcher,
            cp_split_sizes=cp_split_sizes,
        )
        x = ModalityDispatcher.inv_permute(x, inv_permute_mapping)

        x_video = x[video_mask].to(self.final_norm_video.weight.dtype)
        x_video = self.final_norm_video(x_video)
        x_video = self.final_linear_video(x_video)

        x_audio = x[audio_mask].to(self.final_norm_audio.weight.dtype)
        x_audio = self.final_norm_audio(x_audio)
        x_audio = self.final_linear_audio(x_audio)

        x_out = torch.zeros(
            x.shape[0], max(self.config.video_in_channels, self.config.audio_in_channels), device=x.device, dtype=x.dtype
        )
        x_out[video_mask, : self.config.video_in_channels] = x_video
        x_out[audio_mask, : self.config.audio_in_channels] = x_audio
        x_out = ulysses_scheduler().undispatch(x_out)
        _sp.stop(_model_tok)
        return x_out