sgl-project--sglang/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py

# Copyright 2025 Black Forest Labs, The HuggingFace Team and The InstantX Team. 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.

from typing import Any, Dict, List, Optional, Tuple

import torch
import torch.nn as nn
from diffusers.models.attention import AttentionModuleMixin
from diffusers.models.embeddings import TimestepEmbedding, Timesteps
from diffusers.models.normalization import AdaLayerNormContinuous

from sglang.multimodal_gen.configs.models.dits.flux import FluxConfig
from sglang.multimodal_gen.runtime.distributed import (
    divide,
    get_tp_world_size,
)
from sglang.multimodal_gen.runtime.distributed.sp_shard_utils import (
    build_shard_plan,
    join_seqs,
    shard_like,
    shard_seq_prefix,
    should_shard_text,
    split_seqs,
    tail_attn_meta,
)
from sglang.multimodal_gen.runtime.layers.attention import USPAttention
from sglang.multimodal_gen.runtime.layers.layernorm import (
    RMSNorm,
    apply_qk_norm_with_optional_rope,
)
from sglang.multimodal_gen.runtime.layers.linear import (
    ColumnParallelLinear,
    MergedColumnParallelLinear,
    RowParallelLinear,
)
from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
    QuantizationConfig,
)
from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
    ModelOptFp4Config,
)
from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
    NDRotaryEmbedding,
    apply_flashinfer_rope_qk_inplace,
)
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
    LayerwiseOffloadableModuleMixin,
)
from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
from sglang.multimodal_gen.runtime.platforms import (
    AttentionBackendEnum,
    current_platform,
)
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger

logger = init_logger(__name__)  # pylint: disable=invalid-name


def _get_qkv_projections(
    attn: "Flux2Attention", hidden_states, encoder_hidden_states=None
):
    if attn.use_fused_qkv:
        qkv, _ = attn.to_qkv(hidden_states)
        query, key, value = [t.contiguous() for t in qkv.chunk(3, dim=-1)]
    else:
        query, _ = attn.to_q(hidden_states)
        key, _ = attn.to_k(hidden_states)
        value, _ = attn.to_v(hidden_states)

    encoder_query = encoder_key = encoder_value = None
    if encoder_hidden_states is not None and attn.added_kv_proj_dim is not None:
        if attn.use_fused_added_qkv:
            added_qkv, _ = attn.to_added_qkv(encoder_hidden_states)
            encoder_query, encoder_key, encoder_value = [
                t.contiguous() for t in added_qkv.chunk(3, dim=-1)
            ]
        else:
            encoder_query, _ = attn.add_q_proj(encoder_hidden_states)
            encoder_key, _ = attn.add_k_proj(encoder_hidden_states)
            encoder_value, _ = attn.add_v_proj(encoder_hidden_states)

    return query, key, value, encoder_query, encoder_key, encoder_value


class Flux2SwiGLU(nn.Module):
    """
    Flux 2 uses a SwiGLU-style activation in the transformer feedforward sub-blocks, but with the linear projection
    layer fused into the first linear layer of the FF sub-block. Thus, this module has no trainable parameters.
    """

    def __init__(self):
        super().__init__()
        self.gate_fn = nn.SiLU()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x1, x2 = x.chunk(2, dim=-1)
        x = self.gate_fn(x1) * x2
        return x


class Flux2FeedForward(nn.Module):
    def __init__(
        self,
        dim: int,
        dim_out: Optional[int] = None,
        mult: float = 3.0,
        inner_dim: Optional[int] = None,
        bias: bool = False,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ):
        super().__init__()
        if inner_dim is None:
            inner_dim = int(dim * mult)
        dim_out = dim_out or dim

        # Flux2SwiGLU will reduce the dimension by half
        self.linear_in = MergedColumnParallelLinear(
            dim,
            [inner_dim, inner_dim],
            bias=bias,
            gather_output=False,
            quant_config=quant_config,
            prefix=f"{prefix}.linear_in" if prefix else "linear_in",
        )
        self.act_fn = Flux2SwiGLU()
        self.linear_out = RowParallelLinear(
            inner_dim,
            dim_out,
            bias=bias,
            input_is_parallel=True,
            quant_config=quant_config,
            prefix=f"{prefix}.linear_out" if prefix else "linear_out",
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x, _ = self.linear_in(x)
        x = self.act_fn(x)
        x, _ = self.linear_out(x)
        return x


class Flux2Attention(torch.nn.Module, AttentionModuleMixin):
    def __init__(
        self,
        query_dim: int,
        num_heads: int = 8,
        dim_head: int = 64,
        dropout: float = 0.0,
        bias: bool = False,
        added_kv_proj_dim: Optional[int] = None,
        added_proj_bias: Optional[bool] = True,
        out_bias: bool = True,
        eps: float = 1e-5,
        out_dim: int = None,
        elementwise_affine: bool = True,
        supported_attention_backends: set[AttentionBackendEnum] | None = None,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ):
        super().__init__()

        self.head_dim = dim_head
        self.inner_dim = out_dim if out_dim is not None else dim_head * num_heads
        self.query_dim = query_dim
        self.out_dim = out_dim if out_dim is not None else query_dim
        self.heads = out_dim // dim_head if out_dim is not None else num_heads
        self.tp_size = get_tp_world_size()
        self.local_heads = divide(self.heads, self.tp_size)
        self.local_inner_dim = divide(self.inner_dim, self.tp_size)

        self.use_bias = bias
        self.dropout = dropout

        self.added_kv_proj_dim = added_kv_proj_dim
        self.added_proj_bias = added_proj_bias

        # Some FLUX.2 NVFP4 checkpoints store Q/K/V packed as a single tensor, while
        # ModelOpt's standard diffusers export keeps the original to_q/to_k/to_v layout.
        # Only enable the fused loader path for the packed checkpoint family.
        self.use_fused_qkv = isinstance(quant_config, ModelOptFp4Config) and getattr(
            quant_config, "checkpoint_uses_packed_qkv", False
        )
        self.use_fused_added_qkv = self.use_fused_qkv

        if self.use_fused_qkv:
            self.to_qkv = MergedColumnParallelLinear(
                query_dim,
                [self.inner_dim] * 3,
                bias=bias,
                gather_output=False,
                quant_config=quant_config,
                prefix=f"{prefix}.to_qkv" if prefix else "to_qkv",
            )
        else:
            self.to_q = ColumnParallelLinear(
                query_dim,
                self.inner_dim,
                bias=bias,
                gather_output=False,
                quant_config=quant_config,
                prefix=f"{prefix}.to_q" if prefix else "to_q",
            )
            self.to_k = ColumnParallelLinear(
                query_dim,
                self.inner_dim,
                bias=bias,
                gather_output=False,
                quant_config=quant_config,
                prefix=f"{prefix}.to_k" if prefix else "to_k",
            )
            self.to_v = ColumnParallelLinear(
                query_dim,
                self.inner_dim,
                bias=bias,
                gather_output=False,
                quant_config=quant_config,
                prefix=f"{prefix}.to_v" if prefix else "to_v",
            )

        # QK Norm
        self.norm_q = RMSNorm(dim_head, eps=eps)
        self.norm_k = RMSNorm(dim_head, eps=eps)

        self.to_out = torch.nn.ModuleList([])
        self.to_out.append(
            RowParallelLinear(
                self.inner_dim,
                self.out_dim,
                bias=out_bias,
                input_is_parallel=True,
                quant_config=quant_config,
                prefix=f"{prefix}.to_out.0" if prefix else "to_out.0",
            )
        )
        self.to_out.append(torch.nn.Dropout(dropout))

        if added_kv_proj_dim is not None:
            self.norm_added_q = RMSNorm(dim_head, eps=eps)
            self.norm_added_k = RMSNorm(dim_head, eps=eps)
            if self.use_fused_added_qkv:
                # txt_attn.qkv is always BF16 in the NVFP4 checkpoint — no quant needed
                self.to_added_qkv = MergedColumnParallelLinear(
                    added_kv_proj_dim,
                    [self.inner_dim] * 3,
                    bias=added_proj_bias,
                    gather_output=False,
                    quant_config=None,
                    prefix=f"{prefix}.to_added_qkv" if prefix else "to_added_qkv",
                )
            else:
                self.add_q_proj = ColumnParallelLinear(
                    added_kv_proj_dim,
                    self.inner_dim,
                    bias=added_proj_bias,
                    gather_output=False,
                    quant_config=quant_config,
                    prefix=f"{prefix}.add_q_proj" if prefix else "add_q_proj",
                )
                self.add_k_proj = ColumnParallelLinear(
                    added_kv_proj_dim,
                    self.inner_dim,
                    bias=added_proj_bias,
                    gather_output=False,
                    quant_config=quant_config,
                    prefix=f"{prefix}.add_k_proj" if prefix else "add_k_proj",
                )
                self.add_v_proj = ColumnParallelLinear(
                    added_kv_proj_dim,
                    self.inner_dim,
                    bias=added_proj_bias,
                    gather_output=False,
                    quant_config=quant_config,
                    prefix=f"{prefix}.add_v_proj" if prefix else "add_v_proj",
                )
            self.to_add_out = RowParallelLinear(
                self.inner_dim,
                query_dim,
                bias=out_bias,
                input_is_parallel=True,
                quant_config=quant_config,
                prefix=f"{prefix}.to_add_out" if prefix else "to_add_out",
            )

        self.attn = USPAttention(
            num_heads=self.local_heads,
            head_size=self.head_dim,
            dropout_rate=0,
            softmax_scale=None,
            causal=False,
            supported_attention_backends=supported_attention_backends,
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        num_replicated_prefix: int = 0,
        attn_mask: Optional[torch.Tensor] = None,
        attn_mask_meta: Optional[Dict[str, int]] = None,
    ) -> torch.Tensor:
        (
            query,
            key,
            value,
            encoder_query,
            encoder_key,
            encoder_value,
        ) = _get_qkv_projections(self, hidden_states, encoder_hidden_states)

        query = query.unflatten(-1, (self.local_heads, -1))
        key = key.unflatten(-1, (self.local_heads, -1))
        value = value.unflatten(-1, (self.local_heads, -1))

        cos_sin_cache = None
        if freqs_cis is not None:
            cos, sin = freqs_cis
            cos_sin_cache = torch.cat(
                [
                    cos.to(dtype=torch.float32).contiguous(),
                    sin.to(dtype=torch.float32).contiguous(),
                ],
                dim=-1,
            )

        if self.added_kv_proj_dim is not None:
            encoder_query = encoder_query.unflatten(-1, (self.local_heads, -1))
            encoder_key = encoder_key.unflatten(-1, (self.local_heads, -1))
            encoder_value = encoder_value.unflatten(-1, (self.local_heads, -1))

            text_seq_len = encoder_query.shape[1]
            encoder_query, encoder_key = apply_qk_norm_with_optional_rope(
                q=encoder_query,
                k=encoder_key,
                q_norm=self.norm_added_q,
                k_norm=self.norm_added_k,
                head_dim=self.head_dim,
                cos_sin_cache=cos_sin_cache,
                is_neox=False,
                allow_inplace=True,
            )
            query, key = apply_qk_norm_with_optional_rope(
                q=query,
                k=key,
                q_norm=self.norm_q,
                k_norm=self.norm_k,
                head_dim=self.head_dim,
                cos_sin_cache=cos_sin_cache,
                is_neox=False,
                position_offset=text_seq_len,
                allow_inplace=True,
            )

            # join_seqs relocates any SP text tail-pad behind the image (see
            # sp_shard.join_seqs for why).
            sp_txt_pad = (attn_mask_meta or {}).get("local_pad", 0)
            query = join_seqs(encoder_query, query, sp_txt_pad)
            key = join_seqs(encoder_key, key, sp_txt_pad)
            value = join_seqs(encoder_value, value, sp_txt_pad)
        else:
            query, key = apply_qk_norm_with_optional_rope(
                q=query,
                k=key,
                q_norm=self.norm_q,
                k_norm=self.norm_k,
                head_dim=self.head_dim,
                cos_sin_cache=cos_sin_cache,
                is_neox=False,
                allow_inplace=True,
            )

        hidden_states = self.attn(
            query,
            key,
            value,
            attn_mask=attn_mask,
            attn_mask_meta=attn_mask_meta,
            num_replicated_prefix=num_replicated_prefix,
        )

        hidden_states = hidden_states.flatten(2, 3)
        hidden_states = hidden_states.to(query.dtype)

        if encoder_hidden_states is not None:
            encoder_hidden_states, hidden_states = split_seqs(
                hidden_states, encoder_hidden_states.shape[1], sp_txt_pad
            )
            encoder_hidden_states, _ = self.to_add_out(encoder_hidden_states)

        hidden_states, _ = self.to_out[0](hidden_states)
        hidden_states = self.to_out[1](hidden_states)

        if encoder_hidden_states is not None:
            return hidden_states, encoder_hidden_states
        else:
            return hidden_states


class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin):
    """
    Flux 2 parallel self-attention for the Flux 2 single-stream transformer blocks.

    This implements a parallel transformer block, where the attention QKV projections are fused to the feedforward (FF)
    input projections, and the attention output projections are fused to the FF output projections. See the [ViT-22B
    paper](https://arxiv.org/abs/2302.05442) for a visual depiction of this type of transformer block.
    """

    # Does not support QKV fusion as the QKV projections are always fused
    _supports_qkv_fusion = False

    def __init__(
        self,
        query_dim: int,
        num_heads: int = 8,
        dim_head: int = 64,
        dropout: float = 0.0,
        bias: bool = False,
        out_bias: bool = True,
        eps: float = 1e-5,
        out_dim: int = None,
        elementwise_affine: bool = True,
        mlp_ratio: float = 4.0,
        mlp_mult_factor: int = 2,
        supported_attention_backends: set[AttentionBackendEnum] | None = None,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ):
        super().__init__()

        self.head_dim = dim_head
        self.inner_dim = out_dim if out_dim is not None else dim_head * num_heads
        self.query_dim = query_dim
        self.out_dim = out_dim if out_dim is not None else query_dim
        self.heads = out_dim // dim_head if out_dim is not None else num_heads
        self.tp_size = get_tp_world_size()
        self.local_heads = divide(self.heads, self.tp_size)
        self.local_inner_dim = divide(self.inner_dim, self.tp_size)

        self.use_bias = bias
        self.dropout = dropout

        self.mlp_ratio = mlp_ratio
        self.mlp_hidden_dim = int(query_dim * self.mlp_ratio)
        self.local_mlp_hidden_dim = divide(self.mlp_hidden_dim, self.tp_size)
        self.mlp_mult_factor = mlp_mult_factor

        # Fused QKV projections + MLP input projection
        self.to_qkv_mlp_proj = MergedColumnParallelLinear(
            self.query_dim,
            [self.inner_dim, self.inner_dim, self.inner_dim]
            + [self.mlp_hidden_dim] * self.mlp_mult_factor,
            bias=bias,
            gather_output=False,
            quant_config=quant_config,
            prefix=f"{prefix}.to_qkv_mlp_proj" if prefix else "to_qkv_mlp_proj",
        )
        self.mlp_act_fn = Flux2SwiGLU()

        # QK Norm
        self.norm_q = RMSNorm(dim_head, eps=eps)
        self.norm_k = RMSNorm(dim_head, eps=eps)

        # Fused attention output + MLP output projection.
        # Input is [attn_shard | mlp_shard] (independently sharded by
        # MergedColumnParallelLinear), so patch weight loader to pick the
        # correct non-contiguous columns per rank.
        self.to_out = RowParallelLinear(
            self.inner_dim + self.mlp_hidden_dim,
            self.out_dim,
            bias=out_bias,
            input_is_parallel=True,
            quant_config=quant_config,
            prefix=f"{prefix}.to_out" if prefix else "to_out",
        )
        if self.tp_size > 1:
            self._patch_to_out_weight_loader()

        self.attn = USPAttention(
            num_heads=self.local_heads,
            head_size=self.head_dim,
            dropout_rate=0,
            softmax_scale=None,
            causal=False,
            supported_attention_backends=supported_attention_backends,
        )

    def _patch_to_out_weight_loader(self) -> None:
        inner_dim, mlp_dim = self.inner_dim, self.mlp_hidden_dim
        tp_size, tp_rank = self.tp_size, self.to_out.tp_rank

        def _loader(param, loaded_weight):
            input_dim = getattr(param, "input_dim", None)
            if input_dim is not None:
                a = inner_dim // tp_size
                m = mlp_dim // tp_size
                attn_cols = loaded_weight.narrow(input_dim, tp_rank * a, a)
                mlp_cols = loaded_weight.narrow(input_dim, inner_dim + tp_rank * m, m)
                param.data.copy_(torch.cat([attn_cols, mlp_cols], dim=input_dim))
            else:
                param.data.copy_(loaded_weight)

        self.to_out.weight_loader = _loader
        if hasattr(self.to_out.weight, "_weight_loader"):
            self.to_out.weight._weight_loader = _loader
        else:
            self.to_out.weight.weight_loader = _loader

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        num_replicated_prefix: int = 0,
        **kwargs,
    ) -> torch.Tensor:
        attn_mask = kwargs.get("attn_mask")
        attn_mask_meta = kwargs.get("attn_mask_meta")
        if attn_mask is None:
            attn_mask = attention_mask

        # Parallel in (QKV + MLP in) projection
        hidden_states, _ = self.to_qkv_mlp_proj(hidden_states)
        qkv, mlp_hidden_states = torch.split(
            hidden_states,
            [
                3 * self.local_inner_dim,
                self.local_mlp_hidden_dim * self.mlp_mult_factor,
            ],
            dim=-1,
        )

        # Handle the attention logic
        query, key, value = qkv.chunk(3, dim=-1)

        query = query.unflatten(-1, (self.local_heads, -1))
        key = key.unflatten(-1, (self.local_heads, -1))
        value = value.unflatten(-1, (self.local_heads, -1))

        query = self.norm_q(query)
        key = self.norm_k(key)

        if freqs_cis is not None:
            cos, sin = freqs_cis
            cos_sin_cache = torch.cat(
                [
                    cos.to(dtype=torch.float32).contiguous(),
                    sin.to(dtype=torch.float32).contiguous(),
                ],
                dim=-1,
            )
            query, key = apply_flashinfer_rope_qk_inplace(
                query, key, cos_sin_cache, is_neox=False
            )
        hidden_states = self.attn(
            query,
            key,
            value,
            attn_mask=attn_mask,
            attn_mask_meta=attn_mask_meta,
            num_replicated_prefix=num_replicated_prefix,
        )
        hidden_states = hidden_states.flatten(2, 3)
        hidden_states = hidden_states.to(query.dtype)

        # Handle the feedforward (FF) logic
        mlp_hidden_states = self.mlp_act_fn(mlp_hidden_states)

        # Concatenate and parallel output projection
        hidden_states = torch.cat([hidden_states, mlp_hidden_states], dim=-1)
        hidden_states, _ = self.to_out(hidden_states)

        return hidden_states


class Flux2SingleTransformerBlock(nn.Module):
    def __init__(
        self,
        dim: int,
        num_attention_heads: int,
        attention_head_dim: int,
        mlp_ratio: float = 3.0,
        eps: float = 1e-6,
        bias: bool = False,
        supported_attention_backends: set[AttentionBackendEnum] | None = None,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ):
        super().__init__()

        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)

        # Note that the MLP in/out linear layers are fused with the attention QKV/out projections, respectively; this
        # is often called a "parallel" transformer block. See the [ViT-22B paper](https://arxiv.org/abs/2302.05442)
        # for a visual depiction of this type of transformer block.
        self.attn = Flux2ParallelSelfAttention(
            query_dim=dim,
            dim_head=attention_head_dim,
            num_heads=num_attention_heads,
            out_dim=dim,
            bias=bias,
            out_bias=bias,
            eps=eps,
            mlp_ratio=mlp_ratio,
            mlp_mult_factor=2,
            supported_attention_backends=supported_attention_backends,
            quant_config=quant_config,
            prefix=f"{prefix}.attn" if prefix else "attn",
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: Optional[torch.Tensor],
        temb_mod_params: Tuple[torch.Tensor, torch.Tensor, torch.Tensor],
        freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
        split_hidden_states: bool = False,
        text_seq_len: Optional[int] = None,
        num_replicated_prefix: int = 0,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        # If encoder_hidden_states is None, hidden_states is assumed to have encoder_hidden_states already
        # concatenated
        if encoder_hidden_states is not None:
            text_seq_len = encoder_hidden_states.shape[1]
            hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)

        mod_shift, mod_scale, mod_gate = temb_mod_params

        norm_hidden_states = self.norm(hidden_states)
        norm_hidden_states = (1 + mod_scale) * norm_hidden_states + mod_shift

        joint_attention_kwargs = joint_attention_kwargs or {}
        attn_output = self.attn(
            hidden_states=norm_hidden_states,
            freqs_cis=freqs_cis,
            num_replicated_prefix=num_replicated_prefix,
            **joint_attention_kwargs,
        )

        hidden_states = hidden_states + mod_gate * attn_output
        if hidden_states.dtype == torch.float16:
            hidden_states = hidden_states.clip(-65504, 65504)

        if split_hidden_states:
            encoder_hidden_states, hidden_states = (
                hidden_states[:, :text_seq_len],
                hidden_states[:, text_seq_len:],
            )
            return encoder_hidden_states, hidden_states
        else:
            return hidden_states


class Flux2TransformerBlock(nn.Module):
    def __init__(
        self,
        dim: int,
        num_attention_heads: int,
        attention_head_dim: int,
        mlp_ratio: float = 3.0,
        eps: float = 1e-6,
        bias: bool = False,
        supported_attention_backends: set[AttentionBackendEnum] | None = None,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ):
        super().__init__()
        self.mlp_hidden_dim = int(dim * mlp_ratio)

        self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
        self.norm1_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)

        self.attn = Flux2Attention(
            query_dim=dim,
            added_kv_proj_dim=dim,
            dim_head=attention_head_dim,
            num_heads=num_attention_heads,
            out_dim=dim,
            bias=bias,
            added_proj_bias=bias,
            out_bias=bias,
            eps=eps,
            supported_attention_backends=supported_attention_backends,
            quant_config=quant_config,
            prefix=f"{prefix}.attn" if prefix else "attn",
        )

        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
        self.ff = Flux2FeedForward(
            dim=dim,
            dim_out=dim,
            mult=mlp_ratio,
            bias=bias,
            quant_config=quant_config,
            prefix=f"{prefix}.ff" if prefix else "ff",
        )

        self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
        self.ff_context = Flux2FeedForward(
            dim=dim,
            dim_out=dim,
            mult=mlp_ratio,
            bias=bias,
            quant_config=quant_config,
            prefix=f"{prefix}.ff_context" if prefix else "ff_context",
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor,
        temb_mod_params_img: Tuple[
            Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...
        ],
        temb_mod_params_txt: Tuple[
            Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...
        ],
        freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
        num_replicated_prefix: int = 0,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        joint_attention_kwargs = joint_attention_kwargs or {}

        # Modulation parameters shape: [1, 1, self.dim]
        (shift_msa, scale_msa, gate_msa), (
            shift_mlp,
            scale_mlp,
            gate_mlp,
        ) = temb_mod_params_img
        (c_shift_msa, c_scale_msa, c_gate_msa), (
            c_shift_mlp,
            c_scale_mlp,
            c_gate_mlp,
        ) = temb_mod_params_txt

        # Img stream
        norm_hidden_states = self.norm1(hidden_states)
        norm_hidden_states = (1 + scale_msa) * norm_hidden_states + shift_msa

        # Conditioning txt stream
        norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states)
        norm_encoder_hidden_states = (
            1 + c_scale_msa
        ) * norm_encoder_hidden_states + c_shift_msa

        # Attention on concatenated img + txt stream
        attention_outputs = self.attn(
            hidden_states=norm_hidden_states,
            encoder_hidden_states=norm_encoder_hidden_states,
            freqs_cis=freqs_cis,
            num_replicated_prefix=num_replicated_prefix,
            **joint_attention_kwargs,
        )

        attn_output, context_attn_output = attention_outputs

        # Process attention outputs for the image stream (`hidden_states`).
        attn_output = gate_msa * attn_output
        hidden_states = hidden_states + attn_output

        norm_hidden_states = self.norm2(hidden_states)
        norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp

        ff_output = self.ff(norm_hidden_states)
        hidden_states = hidden_states + gate_mlp * ff_output

        # Process attention outputs for the text stream (`encoder_hidden_states`).
        context_attn_output = c_gate_msa * context_attn_output
        encoder_hidden_states = encoder_hidden_states + context_attn_output

        norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
        norm_encoder_hidden_states = (
            norm_encoder_hidden_states * (1 + c_scale_mlp) + c_shift_mlp
        )

        context_ff_output = self.ff_context(norm_encoder_hidden_states)
        encoder_hidden_states = encoder_hidden_states + c_gate_mlp * context_ff_output
        if encoder_hidden_states.dtype == torch.float16:
            encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)

        return encoder_hidden_states, hidden_states


class Flux2TimestepGuidanceEmbeddings(nn.Module):
    def __init__(
        self,
        in_channels: int = 256,
        embedding_dim: int = 6144,
        bias: bool = False,
        guidance_embeds: bool = True,
    ):
        super().__init__()

        self.time_proj = Timesteps(
            num_channels=in_channels, flip_sin_to_cos=True, downscale_freq_shift=0
        )
        self.timestep_embedder = TimestepEmbedding(
            in_channels=in_channels, time_embed_dim=embedding_dim, sample_proj_bias=bias
        )

        if guidance_embeds:
            self.guidance_embedder = TimestepEmbedding(
                in_channels=in_channels,
                time_embed_dim=embedding_dim,
                sample_proj_bias=bias,
            )
        else:
            self.guidance_embedder = None

    def forward(
        self, timestep: torch.Tensor, guidance: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        timesteps_proj = self.time_proj(timestep)
        timesteps_emb = self.timestep_embedder(
            timesteps_proj.to(timestep.dtype)
        )  # (N, D)

        if guidance is not None and self.guidance_embedder is not None:
            guidance_proj = self.time_proj(guidance)
            guidance_emb = self.guidance_embedder(
                guidance_proj.to(guidance.dtype)
            )  # (N, D)
            time_guidance_emb = timesteps_emb + guidance_emb
            return time_guidance_emb
        else:
            return timesteps_emb


class Flux2Modulation(nn.Module):
    def __init__(self, dim: int, mod_param_sets: int = 2, bias: bool = False):
        super().__init__()
        self.mod_param_sets = mod_param_sets

        self.linear = ColumnParallelLinear(
            dim, dim * 3 * self.mod_param_sets, bias=bias, gather_output=True
        )
        self.act_fn = nn.SiLU()

    def forward(
        self, temb: torch.Tensor
    ) -> Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...]:
        mod = self.act_fn(temb)
        mod, _ = self.linear(mod)

        if mod.ndim == 2:
            mod = mod.unsqueeze(1)
        mod_params = torch.chunk(mod, 3 * self.mod_param_sets, dim=-1)
        # Return tuple of 3-tuples of modulation params shift/scale/gate
        return tuple(
            mod_params[3 * i : 3 * (i + 1)] for i in range(self.mod_param_sets)
        )


class Flux2PosEmbed(nn.Module):
    def __init__(self, theta: int, axes_dim: List[int]):
        super().__init__()
        self.rope = NDRotaryEmbedding(
            rope_dim_list=axes_dim,
            rope_theta=theta,
            use_real=False,
            repeat_interleave_real=False,
            dtype=(
                torch.float64
                if (
                    current_platform.is_float64_supported()
                    if hasattr(current_platform, "is_float64_supported")
                    else True
                )
                else torch.float32
            ),
        )

    def forward(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        pos = ids.float()
        # TODO: potential error: flux use n_axes = ids.shape[-1]
        # see: https://github.com/huggingface/diffusers/blob/17c0e79dbdf53fb6705e9c09cc1a854b84c39249/src/diffusers/models/transformers/transformer_flux.py#L509
        freqs_cos, freqs_sin = self.rope.forward_uncached(pos=pos)
        return freqs_cos.contiguous().float(), freqs_sin.contiguous().float()


class Flux2Transformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
    """
    The Transformer model introduced in Flux 2.

    Reference: https://blackforestlabs.ai/announcing-black-forest-labs/

    """

    param_names_mapping = FluxConfig().arch_config.param_names_mapping
    scale_shift_swap_params = ("norm_out.linear.weight", "norm_out.linear.bias")
    # FLUX.2 stays closer to the official diffusers output with Torch SDPA.
    # The generic FA path still produces a measurable image-level drift here.
    _supported_attention_backends = {
        AttentionBackendEnum.TORCH_SDPA,
        AttentionBackendEnum.FA,
        AttentionBackendEnum.AITER,
        AttentionBackendEnum.AITER_SAGE,
    }

    def post_load_weights(self) -> None:
        if not isinstance(getattr(self, "quant_config", None), ModelOptFp4Config):
            return

        # BFL/ComfyUI checkpoints store AdaLN modulation params as [scale, shift],
        # while diffusers expects [shift, scale].
        for param_name in self.scale_shift_swap_params:
            parts = param_name.split(".")
            module = self
            for part in parts[:-1]:
                module = getattr(module, part)
            param = getattr(module, parts[-1], None)
            if param is None:
                continue
            half = param.shape[0] // 2
            with torch.no_grad():
                first_half = param[:half].clone()
                param[:half] = param[half:]
                param[half:] = first_half
            logger.info(
                "Swapped scale/shift order for %s (BFL → diffusers)", param_name
            )

    def __init__(
        self,
        config: FluxConfig,
        hf_config: dict[str, Any],
        quant_config: Optional[QuantizationConfig] = None,
    ):
        super().__init__(config=config, hf_config=hf_config)
        patch_size: int = config.patch_size
        in_channels: int = config.in_channels
        out_channels: Optional[int] = config.out_channels
        num_layers: int = config.num_layers
        num_single_layers: int = config.num_single_layers
        attention_head_dim: int = config.attention_head_dim
        num_attention_heads: int = config.num_attention_heads
        joint_attention_dim: int = config.joint_attention_dim
        timestep_guidance_channels: int = config.timestep_guidance_channels
        mlp_ratio: float = config.mlp_ratio
        axes_dims_rope: Tuple[int, ...] = config.axes_dims_rope
        rope_theta: int = config.rope_theta
        eps: float = config.eps
        guidance_embeds: bool = getattr(config, "guidance_embeds", True)
        self.out_channels = out_channels or in_channels
        self.inner_dim = num_attention_heads * attention_head_dim
        self.guidance_embeds = guidance_embeds
        quant_config = quant_config if quant_config is not None else config.quant_config
        self.quant_config = quant_config

        # 1. Sinusoidal positional embedding for RoPE on image and text tokens
        self.rotary_emb = Flux2PosEmbed(theta=rope_theta, axes_dim=axes_dims_rope)

        # 2. Combined timestep + guidance embedding
        self.time_guidance_embed = Flux2TimestepGuidanceEmbeddings(
            in_channels=timestep_guidance_channels,
            embedding_dim=self.inner_dim,
            bias=False,
            guidance_embeds=guidance_embeds,
        )

        # 3. Modulation (double stream and single stream blocks share modulation parameters, resp.)
        # Two sets of shift/scale/gate modulation parameters for the double stream attn and FF sub-blocks
        self.double_stream_modulation_img = Flux2Modulation(
            self.inner_dim, mod_param_sets=2, bias=False
        )
        self.double_stream_modulation_txt = Flux2Modulation(
            self.inner_dim, mod_param_sets=2, bias=False
        )
        # Only one set of modulation parameters as the attn and FF sub-blocks are run in parallel for single stream
        self.single_stream_modulation = Flux2Modulation(
            self.inner_dim, mod_param_sets=1, bias=False
        )

        # 4. Input projections
        self.x_embedder = ColumnParallelLinear(
            in_channels, self.inner_dim, bias=False, gather_output=True
        )
        self.context_embedder = ColumnParallelLinear(
            joint_attention_dim, self.inner_dim, bias=False, gather_output=True
        )

        # 5. Double Stream Transformer Blocks
        self.transformer_blocks = nn.ModuleList(
            [
                Flux2TransformerBlock(
                    dim=self.inner_dim,
                    num_attention_heads=num_attention_heads,
                    attention_head_dim=attention_head_dim,
                    mlp_ratio=mlp_ratio,
                    eps=eps,
                    bias=False,
                    supported_attention_backends=self._supported_attention_backends,
                    quant_config=quant_config,
                    prefix=f"transformer_blocks.{i}",
                )
                for i in range(num_layers)
            ]
        )

        # 6. Single Stream Transformer Blocks
        self.single_transformer_blocks = nn.ModuleList(
            [
                Flux2SingleTransformerBlock(
                    dim=self.inner_dim,
                    num_attention_heads=num_attention_heads,
                    attention_head_dim=attention_head_dim,
                    mlp_ratio=mlp_ratio,
                    eps=eps,
                    bias=False,
                    supported_attention_backends=self._supported_attention_backends,
                    quant_config=quant_config,
                    prefix=f"single_transformer_blocks.{i}",
                )
                for i in range(num_single_layers)
            ]
        )

        # 7. Output layers
        self.norm_out = AdaLayerNormContinuous(
            self.inner_dim,
            self.inner_dim,
            elementwise_affine=False,
            eps=eps,
            bias=False,
        )
        self.proj_out = ColumnParallelLinear(
            self.inner_dim,
            patch_size * patch_size * self.out_channels,
            bias=False,
            gather_output=True,
            quant_config=quant_config,
            prefix="proj_out",
        )

        self.layer_names = ["transformer_blocks", "single_transformer_blocks"]

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor = None,
        timestep: torch.LongTensor = None,
        guidance: torch.Tensor = None,
        freqs_cis: torch.Tensor = None,
        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
    ) -> torch.Tensor:
        """
        The [`FluxTransformer2DModel`] forward method.

        Args:
            hidden_states (`torch.Tensor` of shape `(batch_size, image_sequence_length, in_channels)`):
                Input `hidden_states`.
            encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_sequence_length, joint_attention_dim)`):
                Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
            timestep ( `torch.LongTensor`):
                Used to indicate denoising step.
            joint_attention_kwargs (`dict`, *optional*):
                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
                `self.processor` in
                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).

        """
        # 0. Handle input arguments
        if joint_attention_kwargs is not None:
            joint_attention_kwargs = joint_attention_kwargs.copy()
            joint_attention_kwargs.pop("scale", 1.0)

        num_txt_tokens = encoder_hidden_states.shape[1]

        # 1. Calculate timestep embedding and modulation parameters
        timestep = timestep.to(hidden_states.dtype)
        if guidance is not None:
            guidance = guidance.to(hidden_states.dtype) * 1000

        temb = self.time_guidance_embed(timestep, guidance)

        double_stream_mod_img = self.double_stream_modulation_img(temb)
        double_stream_mod_txt = self.double_stream_modulation_txt(temb)
        single_stream_mod = self.single_stream_modulation(temb)[0]

        # 2. Input projection for image (hidden_states) and conditioning text (encoder_hidden_states)
        hidden_states, _ = self.x_embedder(hidden_states)
        encoder_hidden_states, _ = self.context_embedder(encoder_hidden_states)

        # Shard the replicated text stream across SP ranks (image latents are
        # already sharded); non-divisible lengths tail-pad the last rank and the
        # per-request tail meta lets attention skip the pad for free.
        num_replicated_prefix = num_txt_tokens
        sp_txt_pad = 0
        singles_freqs_cis = freqs_cis
        if should_shard_text(num_txt_tokens):
            txt_shard = build_shard_plan(num_txt_tokens)
            encoder_hidden_states = shard_like(encoder_hidden_states, txt_shard)
            if freqs_cis is not None:
                cos, sin = freqs_cis
                cos = shard_seq_prefix(cos, num_txt_tokens, txt_shard)
                sin = shard_seq_prefix(sin, num_txt_tokens, txt_shard)
                freqs_cis = (cos, sin)
                singles_freqs_cis = freqs_cis
            num_replicated_prefix = 0
            num_txt_tokens = txt_shard.local_len
            tail_meta = tail_attn_meta(
                txt_shard,
                encoder_hidden_states.shape[0],
                hidden_states.device,
                image_seq_len=hidden_states.shape[1],
            )
            if tail_meta is not None:
                joint_attention_kwargs = (
                    joint_attention_kwargs.copy() if joint_attention_kwargs else {}
                )
                joint_attention_kwargs["attn_mask_meta"] = tail_meta
                sp_txt_pad = txt_shard.local_pad
                # The single-stream trunk applies RoPE on the relocated
                # [txt_real, img, pad] layout; reorder its cache to match.
                if freqs_cis is not None:
                    t_loc = txt_shard.local_len
                    singles_freqs_cis = (
                        join_seqs(cos[:t_loc], cos[t_loc:], sp_txt_pad, dim=0),
                        join_seqs(sin[:t_loc], sin[t_loc:], sp_txt_pad, dim=0),
                    )

        # 4. Double Stream Transformer Blocks
        for index_block, block in enumerate(self.transformer_blocks):
            encoder_hidden_states, hidden_states = block(
                hidden_states=hidden_states,
                encoder_hidden_states=encoder_hidden_states,
                temb_mod_params_img=double_stream_mod_img,
                temb_mod_params_txt=double_stream_mod_txt,
                freqs_cis=freqs_cis,
                joint_attention_kwargs=joint_attention_kwargs,
                num_replicated_prefix=num_replicated_prefix,
            )
        # Concatenate text and image streams for single-block inference;
        # join_seqs relocates any SP text tail-pad behind the image once for
        # the whole trunk (see sp_shard.join_seqs for why).
        txt_real = num_txt_tokens - sp_txt_pad
        hidden_states = join_seqs(encoder_hidden_states, hidden_states, sp_txt_pad)

        # 5. Single Stream Transformer Blocks
        for index_block, block in enumerate(self.single_transformer_blocks):
            hidden_states = block(
                hidden_states=hidden_states,
                encoder_hidden_states=None,
                temb_mod_params=single_stream_mod,
                freqs_cis=singles_freqs_cis,
                joint_attention_kwargs=joint_attention_kwargs,
                text_seq_len=txt_real,
                num_replicated_prefix=num_replicated_prefix,
            )
        # Remove text (and any tail pad) from the concatenated stream
        img_end = hidden_states.shape[1] - sp_txt_pad
        hidden_states = hidden_states[:, txt_real:img_end, ...]

        # 6. Output layers
        hidden_states = self.norm_out(hidden_states, temb)
        output, _ = self.proj_out(hidden_states)

        return output


EntryClass = Flux2Transformer2DModel