cache-dit 1.0.3__py3-none-any.whl → 1.0.14__py3-none-any.whl

cache_dit/parallelism/backends/native_diffusers/context_parallelism/cp_plan_ltxvideo.py

@@ -0,0 +1,264 @@
+import torch
+import functools
+from typing import Optional
+import torch.nn.functional as F
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.transformers.transformer_ltx import (
+    LTXVideoTransformer3DModel,
+    LTXAttention,
+    AttentionModuleMixin,
+    LTXVideoAttnProcessor,
+    apply_rotary_emb,
+)
+from diffusers.models.attention_dispatch import dispatch_attention_fn
+
+try:
+    from diffusers.models._modeling_parallel import (
+        ContextParallelInput,
+        ContextParallelOutput,
+        ContextParallelModelPlan,
+    )
+except ImportError:
+    raise ImportError(
+        "Context parallelism requires the 'diffusers>=0.36.dev0'."
+        "Please install latest version of diffusers from source: \n"
+        "pip3 install git+https://github.com/huggingface/diffusers.git"
+    )
+from .cp_plan_registers import (
+    ContextParallelismPlanner,
+    ContextParallelismPlannerRegister,
+)
+
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@ContextParallelismPlannerRegister.register("LTXVideo")
+class LTXVideoContextParallelismPlanner(ContextParallelismPlanner):
+    def apply(
+        self,
+        transformer: Optional[torch.nn.Module | ModelMixin] = None,
+        **kwargs,
+    ) -> ContextParallelModelPlan:
+        assert transformer is not None, "Transformer must be provided."
+        assert isinstance(
+            transformer, LTXVideoTransformer3DModel
+        ), "Transformer must be an instance of LTXVideoTransformer3DModel"
+
+        # NOTE: The atttention_mask preparation in LTXAttention while using
+        # context parallelism is buggy in diffusers v0.36.0.dev0, so we
+        # disable the preference to use native diffusers implementation here.
+        self._cp_planner_preferred_native_diffusers = False
+
+        if (
+            transformer is not None
+            and self._cp_planner_preferred_native_diffusers
+        ):
+            if hasattr(transformer, "_cp_plan"):
+                if transformer._cp_plan is not None:
+                    return transformer._cp_plan
+
+        # Apply monkey patch to fix attention mask preparation at class level
+        assert issubclass(LTXAttention, AttentionModuleMixin)
+        LTXAttention.prepare_attention_mask = (
+            __patch__LTXAttention_prepare_attention_mask__
+        )
+        LTXVideoAttnProcessor.__call__ = __patch__LTXVideoAttnProcessor__call__
+
+        # Otherwise, use the custom CP plan defined here, this maybe
+        # a little different from the native diffusers implementation
+        # for some models.
+
+        _cp_plan = {
+            # Here is a Transformer level CP plan for Flux, which will
+            # only apply the only 1 split hook (pre_forward) on the forward
+            # of Transformer, and gather the output after Transformer forward.
+            # Pattern of transformer forward, split_output=False:
+            #     un-split input -> splited input (inside transformer)
+            # Pattern of the transformer_blocks, single_transformer_blocks:
+            #     splited input (previous splited output) -> to_qkv/...
+            #     -> all2all
+            #     -> attn (local head, full seqlen)
+            #     -> all2all
+            #     -> splited output
+            # The `hidden_states` and `encoder_hidden_states` will still keep
+            # itself splited after block forward, namely, hidden_states will
+            # automatically split by the all2all comm op after attn, and the
+            # encoder_hidden_states will be keep splited after the entrypoint
+            # of transformer forward, for the all blocks.
+            "": {
+                "hidden_states": ContextParallelInput(
+                    split_dim=1, expected_dims=3, split_output=False
+                ),
+                "encoder_hidden_states": ContextParallelInput(
+                    split_dim=1, expected_dims=3, split_output=False
+                ),
+                # NOTE: encoder_attention_mask (namely, attention_mask in cross-attn)
+                # should never be split across seqlen while using context parallelism
+                # for LTXVideoTransformer3DModel. It don't contribute to any computation
+                # in parallel or not. So we comment it out here and handle the head-split
+                # correctly while using context parallel in the patched attention processor.
+                # "encoder_attention_mask": ContextParallelInput(
+                #     split_dim=1, expected_dims=2, split_output=False
+                # ),
+            },
+            # Pattern of rope, split_output=True (split output rather than input):
+            #    un-split input
+            #    -> keep input un-split
+            #    -> rope
+            #    -> splited output
+            "rope": {
+                0: ContextParallelInput(
+                    split_dim=1, expected_dims=3, split_output=True
+                ),
+                1: ContextParallelInput(
+                    split_dim=1, expected_dims=3, split_output=True
+                ),
+            },
+            # Then, the final proj_out will gather the splited output.
+            #     splited input (previous splited output)
+            #     -> all gather
+            #     -> un-split output
+            "proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3),
+        }
+        return _cp_plan
+
+
+@functools.wraps(LTXAttention.prepare_attention_mask)
+def __patch__LTXAttention_prepare_attention_mask__(
+    self: LTXAttention,
+    attention_mask: torch.Tensor,
+    target_length: int,
+    batch_size: int,
+    out_dim: int = 3,
+    # NOTE(DefTruth): Allow specifying head_size for CP
+    head_size: Optional[int] = None,
+) -> torch.Tensor:
+    """
+    Prepare the attention mask for the attention computation.
+
+    Args:
+        attention_mask (`torch.Tensor`): The attention mask to prepare.
+        target_length (`int`): The target length of the attention mask.
+        batch_size (`int`): The batch size for repeating the attention mask.
+        out_dim (`int`, *optional*, defaults to `3`): Output dimension.
+
+    Returns:
+        `torch.Tensor`: The prepared attention mask.
+    """
+    if head_size is None:
+        head_size = self.heads
+    if attention_mask is None:
+        return attention_mask
+
+    current_length: int = attention_mask.shape[-1]
+    if current_length != target_length:
+        if attention_mask.device.type == "mps":
+            # HACK: MPS: Does not support padding by greater than dimension of input tensor.
+            # Instead, we can manually construct the padding tensor.
+            padding_shape = (
+                attention_mask.shape[0],
+                attention_mask.shape[1],
+                target_length,
+            )
+            padding = torch.zeros(
+                padding_shape,
+                dtype=attention_mask.dtype,
+                device=attention_mask.device,
+            )
+            attention_mask = torch.cat([attention_mask, padding], dim=2)
+        else:
+            # TODO: for pipelines such as stable-diffusion, padding cross-attn mask:
+            #       we want to instead pad by (0, remaining_length), where remaining_length is:
+            #       remaining_length: int = target_length - current_length
+            # TODO: re-enable tests/models/test_models_unet_2d_condition.py#test_model_xattn_padding
+            attention_mask = F.pad(
+                attention_mask, (0, target_length), value=0.0
+            )
+
+    if out_dim == 3:
+        if attention_mask.shape[0] < batch_size * head_size:
+            attention_mask = attention_mask.repeat_interleave(head_size, dim=0)
+    elif out_dim == 4:
+        attention_mask = attention_mask.unsqueeze(1)
+        attention_mask = attention_mask.repeat_interleave(head_size, dim=1)
+
+    return attention_mask
+
+
+@functools.wraps(LTXVideoAttnProcessor.__call__)
+def __patch__LTXVideoAttnProcessor__call__(
+    self: LTXVideoAttnProcessor,
+    attn: "LTXAttention",
+    hidden_states: torch.Tensor,
+    encoder_hidden_states: Optional[torch.Tensor] = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    image_rotary_emb: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    batch_size, sequence_length, _ = (
+        hidden_states.shape
+        if encoder_hidden_states is None
+        else encoder_hidden_states.shape
+    )
+
+    if attention_mask is not None:
+        if self._parallel_config is None:
+            attention_mask = attn.prepare_attention_mask(
+                attention_mask, sequence_length, batch_size
+            )
+            attention_mask = attention_mask.view(
+                batch_size, attn.heads, -1, attention_mask.shape[-1]
+            )
+        else:
+            # NOTE(DefTruth): Fix attention mask preparation for context parallelism
+            cp_config = getattr(
+                self._parallel_config, "context_parallel_config", None
+            )
+            if cp_config is not None and cp_config._world_size > 1:
+                head_size = attn.heads // cp_config._world_size
+                attention_mask = attn.prepare_attention_mask(
+                    attention_mask,
+                    sequence_length * cp_config._world_size,
+                    batch_size,
+                    3,
+                    head_size,
+                )
+                attention_mask = attention_mask.view(
+                    batch_size, head_size, -1, attention_mask.shape[-1]
+                )
+
+    if encoder_hidden_states is None:
+        encoder_hidden_states = hidden_states
+
+    query = attn.to_q(hidden_states)
+    key = attn.to_k(encoder_hidden_states)
+    value = attn.to_v(encoder_hidden_states)
+
+    query = attn.norm_q(query)
+    key = attn.norm_k(key)
+
+    if image_rotary_emb is not None:
+        query = apply_rotary_emb(query, image_rotary_emb)
+        key = apply_rotary_emb(key, image_rotary_emb)
+
+    query = query.unflatten(2, (attn.heads, -1))
+    key = key.unflatten(2, (attn.heads, -1))
+    value = value.unflatten(2, (attn.heads, -1))
+
+    hidden_states = dispatch_attention_fn(
+        query,
+        key,
+        value,
+        attn_mask=attention_mask,
+        dropout_p=0.0,
+        is_causal=False,
+        backend=self._attention_backend,
+        parallel_config=self._parallel_config,
+    )
+    hidden_states = hidden_states.flatten(2, 3)
+    hidden_states = hidden_states.to(query.dtype)
+
+    hidden_states = attn.to_out[0](hidden_states)
+    hidden_states = attn.to_out[1](hidden_states)
+    return hidden_states

cache_dit/parallelism/backends/native_diffusers/context_parallelism/cp_plan_nunchaku.py

@@ -0,0 +1,407 @@
+import torch
+import functools
+from typing import Optional, Tuple
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.attention_dispatch import dispatch_attention_fn
+from diffusers.models.transformers.transformer_qwenimage import (
+    apply_rotary_emb_qwen,
+)
+
+try:
+    from nunchaku.models.transformers.transformer_flux_v2 import (
+        NunchakuFluxAttention,
+        NunchakuFluxFA2Processor,
+        NunchakuFluxTransformer2DModelV2,
+    )
+    from nunchaku.ops.fused import fused_qkv_norm_rottary
+    from nunchaku.models.transformers.transformer_qwenimage import (
+        NunchakuQwenAttention,
+        NunchakuQwenImageNaiveFA2Processor,
+        NunchakuQwenImageTransformer2DModel,
+    )
+except ImportError:
+    raise ImportError(
+        "NunchakuFluxTransformer2DModelV2 or NunchakuQwenImageTransformer2DModel "
+        "requires the 'nunchaku' package. Please install nunchaku before using "
+        "the context parallelism for nunchaku 4-bits models."
+    )
+
+try:
+    from diffusers.models._modeling_parallel import (
+        ContextParallelInput,
+        ContextParallelOutput,
+        ContextParallelModelPlan,
+    )
+except ImportError:
+    raise ImportError(
+        "Context parallelism requires the 'diffusers>=0.36.dev0'."
+        "Please install latest version of diffusers from source: \n"
+        "pip3 install git+https://github.com/huggingface/diffusers.git"
+    )
+from .cp_plan_registers import (
+    ContextParallelismPlanner,
+    ContextParallelismPlannerRegister,
+)
+
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@ContextParallelismPlannerRegister.register("NunchakuFlux")
+class NunchakuFluxContextParallelismPlanner(ContextParallelismPlanner):
+    def apply(
+        self,
+        transformer: Optional[torch.nn.Module | ModelMixin] = None,
+        **kwargs,
+    ) -> ContextParallelModelPlan:
+
+        self._cp_planner_preferred_native_diffusers = False
+
+        if (
+            transformer is not None
+            and self._cp_planner_preferred_native_diffusers
+        ):
+
+            assert isinstance(
+                transformer, NunchakuFluxTransformer2DModelV2
+            ), "Transformer must be an instance of NunchakuFluxTransformer2DModelV2"
+            if hasattr(transformer, "_cp_plan"):
+                if transformer._cp_plan is not None:
+                    return transformer._cp_plan
+
+        NunchakuFluxFA2Processor.__call__ = (
+            __patch_NunchakuFluxFA2Processor__call__
+        )
+        # Also need to patch the parallel config and attention backend
+        if not hasattr(NunchakuFluxFA2Processor, "_parallel_config"):
+            NunchakuFluxFA2Processor._parallel_config = None
+        if not hasattr(NunchakuFluxFA2Processor, "_attention_backend"):
+            NunchakuFluxFA2Processor._attention_backend = None
+        if not hasattr(NunchakuFluxAttention, "_parallel_config"):
+            NunchakuFluxAttention._parallel_config = None
+        if not hasattr(NunchakuFluxAttention, "_attention_backend"):
+            NunchakuFluxAttention._attention_backend = None
+
+        # Otherwise, use the custom CP plan defined here, this maybe
+        # a little different from the native diffusers implementation
+        # for some models.
+        _cp_plan = {
+            # Here is a Transformer level CP plan for Flux, which will
+            # only apply the only 1 split hook (pre_forward) on the forward
+            # of Transformer, and gather the output after Transformer forward.
+            # Pattern of transformer forward, split_output=False:
+            #     un-split input -> splited input (inside transformer)
+            # Pattern of the transformer_blocks, single_transformer_blocks:
+            #     splited input (previous splited output) -> to_qkv/...
+            #     -> all2all
+            #     -> attn (local head, full seqlen)
+            #     -> all2all
+            #     -> splited output
+            # The `hidden_states` and `encoder_hidden_states` will still keep
+            # itself splited after block forward (namely, automatic split by
+            # the all2all comm op after attn) for the all blocks.
+            # img_ids and txt_ids will only be splited once at the very beginning,
+            # and keep splited through the whole transformer forward. The all2all
+            # comm op only happens on the `out` tensor after local attn not on
+            # img_ids and txt_ids.
+            "": {
+                "hidden_states": ContextParallelInput(
+                    split_dim=1, expected_dims=3, split_output=False
+                ),
+                "encoder_hidden_states": ContextParallelInput(
+                    split_dim=1, expected_dims=3, split_output=False
+                ),
+                "img_ids": ContextParallelInput(
+                    split_dim=0, expected_dims=2, split_output=False
+                ),
+                "txt_ids": ContextParallelInput(
+                    split_dim=0, expected_dims=2, split_output=False
+                ),
+            },
+            # Then, the final proj_out will gather the splited output.
+            #     splited input (previous splited output)
+            #     -> all gather
+            #     -> un-split output
+            "proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3),
+        }
+        return _cp_plan
+
+
+@functools.wraps(NunchakuFluxFA2Processor.__call__)
+def __patch_NunchakuFluxFA2Processor__call__(
+    self: NunchakuFluxFA2Processor,
+    attn: NunchakuFluxAttention,
+    hidden_states: torch.Tensor,
+    encoder_hidden_states: Optional[torch.Tensor] = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    image_rotary_emb: Tuple[torch.Tensor, torch.Tensor] | torch.Tensor = None,
+    **kwargs,
+) -> torch.Tensor | Tuple[torch.Tensor, torch.Tensor]:
+    # The original implementation of NunchakuFluxFA2Processor.__call__
+    # is not changed here for brevity. In actual implementation, we need to
+    # modify the attention computation to support context parallelism.
+    if attention_mask is not None:
+        raise NotImplementedError("attention_mask is not supported")
+
+    batch_size, _, channels = hidden_states.shape
+    assert channels == attn.heads * attn.head_dim
+    qkv = fused_qkv_norm_rottary(
+        hidden_states,
+        attn.to_qkv,
+        attn.norm_q,
+        attn.norm_k,
+        (
+            image_rotary_emb[0]
+            if isinstance(image_rotary_emb, tuple)
+            else image_rotary_emb
+        ),
+    )
+
+    if attn.added_kv_proj_dim is not None:
+        assert encoder_hidden_states is not None
+        assert isinstance(image_rotary_emb, tuple)
+        qkv_context = fused_qkv_norm_rottary(
+            encoder_hidden_states,
+            attn.add_qkv_proj,
+            attn.norm_added_q,
+            attn.norm_added_k,
+            image_rotary_emb[1],
+        )
+        qkv = torch.cat([qkv_context, qkv], dim=1)
+
+    query, key, value = qkv.chunk(3, dim=-1)
+    # Original implementation:
+    # query = query.view(batch_size, -1, attn.heads, attn.head_dim).transpose(
+    #     1, 2
+    # )
+    # key = key.view(batch_size, -1, attn.heads, attn.head_dim).transpose(1, 2)
+    # value = value.view(batch_size, -1, attn.heads, attn.head_dim).transpose(
+    #     1, 2
+    # )
+    # hidden_states = F.scaled_dot_product_attention(
+    #     query,
+    #     key,
+    #     value,
+    #     attn_mask=attention_mask,
+    #     dropout_p=0.0,
+    #     is_causal=False,
+    # )
+    # hidden_states = hidden_states.transpose(1, 2).reshape(
+    #     batch_size, -1, attn.heads * attn.head_dim
+    # )
+    # hidden_states = hidden_states.to(query.dtype)
+
+    # NOTE(DefTruth): Monkey patch to support context parallelism
+    query = query.view(batch_size, -1, attn.heads, attn.head_dim)
+    key = key.view(batch_size, -1, attn.heads, attn.head_dim)
+    value = value.view(batch_size, -1, attn.heads, attn.head_dim)
+
+    hidden_states = dispatch_attention_fn(
+        query,
+        key,
+        value,
+        attn_mask=attention_mask,
+        backend=getattr(self, "_attention_backend", None),
+        parallel_config=getattr(self, "_parallel_config", None),
+    )
+    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 = (
+            hidden_states[:, : encoder_hidden_states.shape[1]],
+            hidden_states[:, encoder_hidden_states.shape[1] :],
+        )
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+        encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+        return hidden_states, encoder_hidden_states
+    else:
+        # for single transformer block, we split the proj_out into two linear layers
+        hidden_states = attn.to_out(hidden_states)
+        return hidden_states
+
+
+@ContextParallelismPlannerRegister.register("NunchakuQwenImage")
+class NunchakuQwenImageContextParallelismPlanner(ContextParallelismPlanner):
+    def apply(
+        self,
+        transformer: Optional[torch.nn.Module | ModelMixin] = None,
+        **kwargs,
+    ) -> ContextParallelModelPlan:
+
+        self._cp_planner_preferred_native_diffusers = False
+
+        if (
+            transformer is not None
+            and self._cp_planner_preferred_native_diffusers
+        ):
+
+            assert isinstance(
+                transformer, NunchakuQwenImageTransformer2DModel
+            ), "Transformer must be an instance of NunchakuQwenImageTransformer2DModel"
+            if hasattr(transformer, "_cp_plan"):
+                if transformer._cp_plan is not None:
+                    return transformer._cp_plan
+
+        NunchakuQwenImageNaiveFA2Processor.__call__ = (
+            __patch_NunchakuQwenImageNaiveFA2Processor__call__
+        )
+        # Also need to patch the parallel config and attention backend
+        if not hasattr(NunchakuQwenImageNaiveFA2Processor, "_parallel_config"):
+            NunchakuQwenImageNaiveFA2Processor._parallel_config = None
+        if not hasattr(
+            NunchakuQwenImageNaiveFA2Processor, "_attention_backend"
+        ):
+            NunchakuQwenImageNaiveFA2Processor._attention_backend = None
+        if not hasattr(NunchakuQwenAttention, "_parallel_config"):
+            NunchakuQwenAttention._parallel_config = None
+        if not hasattr(NunchakuQwenAttention, "_attention_backend"):
+            NunchakuQwenAttention._attention_backend = None
+
+        # Otherwise, use the custom CP plan defined here, this maybe
+        # a little different from the native diffusers implementation
+        # for some models.
+        _cp_plan = {
+            # Here is a Transformer level CP plan for Flux, which will
+            # only apply the only 1 split hook (pre_forward) on the forward
+            # of Transformer, and gather the output after Transformer forward.
+            # Pattern of transformer forward, split_output=False:
+            #     un-split input -> splited input (inside transformer)
+            # Pattern of the transformer_blocks, single_transformer_blocks:
+            #     splited input (previous splited output) -> to_qkv/...
+            #     -> all2all
+            #     -> attn (local head, full seqlen)
+            #     -> all2all
+            #     -> splited output
+            # The `hidden_states` and `encoder_hidden_states` will still keep
+            # itself splited after block forward (namely, automatic split by
+            # the all2all comm op after attn) for the all blocks.
+            "": {
+                "hidden_states": ContextParallelInput(
+                    split_dim=1, expected_dims=3, split_output=False
+                ),
+                # NOTE: Due to the joint attention implementation of
+                # QwenImageTransformerBlock, we must split the
+                # encoder_hidden_states as well.
+                "encoder_hidden_states": ContextParallelInput(
+                    split_dim=1, expected_dims=3, split_output=False
+                ),
+                # NOTE: But encoder_hidden_states_mask seems never used in
+                # QwenImageTransformerBlock, so we do not split it here.
+                # "encoder_hidden_states_mask": ContextParallelInput(
+                #     split_dim=1, expected_dims=2, split_output=False
+                # ),
+            },
+            # Pattern of pos_embed, split_output=True (split output rather than input):
+            #    un-split input
+            #    -> keep input un-split
+            #    -> rope
+            #    -> splited output
+            "pos_embed": {
+                0: ContextParallelInput(
+                    split_dim=0, expected_dims=2, split_output=True
+                ),
+                1: ContextParallelInput(
+                    split_dim=0, expected_dims=2, split_output=True
+                ),
+            },
+            # Then, the final proj_out will gather the splited output.
+            #     splited input (previous splited output)
+            #     -> all gather
+            #     -> un-split output
+            "proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3),
+        }
+        return _cp_plan
+
+
+@functools.wraps(NunchakuQwenImageNaiveFA2Processor.__call__)
+def __patch_NunchakuQwenImageNaiveFA2Processor__call__(
+    self,
+    attn,
+    hidden_states: torch.FloatTensor,
+    encoder_hidden_states: torch.FloatTensor = None,
+    encoder_hidden_states_mask: torch.FloatTensor = None,
+    attention_mask: Optional[torch.FloatTensor] = None,
+    image_rotary_emb: Optional[torch.Tensor] = None,
+) -> torch.Tensor | Tuple[torch.Tensor, torch.Tensor]:
+    if encoder_hidden_states is None:
+        raise ValueError(
+            "NunchakuQwenImageFA2Processor requires encoder_hidden_states (text stream)"
+        )
+
+    seq_txt = encoder_hidden_states.shape[1]
+
+    # Compute QKV for image stream (sample projections)
+    img_qkv = attn.to_qkv(hidden_states)
+    img_query, img_key, img_value = img_qkv.chunk(3, dim=-1)
+
+    # Compute QKV for text stream (context projections)
+    txt_qkv = attn.add_qkv_proj(encoder_hidden_states)
+    txt_query, txt_key, txt_value = txt_qkv.chunk(3, dim=-1)
+
+    # Reshape for multi-head attention
+    img_query = img_query.unflatten(-1, (attn.heads, -1))  # [B, L, H, D]
+    img_key = img_key.unflatten(-1, (attn.heads, -1))
+    img_value = img_value.unflatten(-1, (attn.heads, -1))
+
+    txt_query = txt_query.unflatten(-1, (attn.heads, -1))
+    txt_key = txt_key.unflatten(-1, (attn.heads, -1))
+    txt_value = txt_value.unflatten(-1, (attn.heads, -1))
+
+    # Apply QK normalization
+    assert attn.norm_q is not None
+    img_query = attn.norm_q(img_query)
+    assert attn.norm_k is not None
+    img_key = attn.norm_k(img_key)
+    assert attn.norm_added_q is not None
+    txt_query = attn.norm_added_q(txt_query)
+    assert attn.norm_added_k is not None
+    txt_key = attn.norm_added_k(txt_key)
+
+    # Apply rotary embeddings
+    if image_rotary_emb is not None:
+        img_freqs, txt_freqs = image_rotary_emb
+        img_query = apply_rotary_emb_qwen(img_query, img_freqs, use_real=False)
+        img_key = apply_rotary_emb_qwen(img_key, img_freqs, use_real=False)
+        txt_query = apply_rotary_emb_qwen(txt_query, txt_freqs, use_real=False)
+        txt_key = apply_rotary_emb_qwen(txt_key, txt_freqs, use_real=False)
+
+    # Concatenate for joint attention: [text, image]
+    joint_query = torch.cat([txt_query, img_query], dim=1)
+    joint_key = torch.cat([txt_key, img_key], dim=1)
+    joint_value = torch.cat([txt_value, img_value], dim=1)
+
+    # Compute joint attention
+    joint_hidden_states = dispatch_attention_fn(
+        joint_query,
+        joint_key,
+        joint_value,
+        attn_mask=attention_mask,
+        dropout_p=0.0,
+        is_causal=False,
+        # NOTE(DefTruth): Use the patched attention backend and
+        # parallel config to make context parallelism work here.
+        backend=getattr(self, "_attention_backend", None),
+        parallel_config=getattr(self, "_parallel_config", None),
+    )
+
+    # Reshape back
+    joint_hidden_states = joint_hidden_states.flatten(2, 3)
+    joint_hidden_states = joint_hidden_states.to(joint_query.dtype)
+
+    # Split attention outputs back
+    txt_attn_output = joint_hidden_states[:, :seq_txt, :]  # Text part
+    img_attn_output = joint_hidden_states[:, seq_txt:, :]  # Image part
+
+    # Apply output projections
+    img_attn_output = attn.to_out[0](img_attn_output)
+    if len(attn.to_out) > 1:
+        img_attn_output = attn.to_out[1](img_attn_output)  # dropout
+
+    txt_attn_output = attn.to_add_out(txt_attn_output)
+
+    return img_attn_output, txt_attn_output