feat: fuse qkv/norm/rottary for z image (#868)

* fuse qkv/norm/rottary for z image

* add doc string

* style: make linter happy

Author: devgdovg

nunchaku/models/attention_processors/zimage.py

@@ -8,6 +8,8 @@
 from diffusers.models.attention_dispatch import dispatch_attention_fn
 from diffusers.models.transformers.transformer_z_image import ZSingleStreamAttnProcessor
 
+from ...ops.fused import fused_qkv_norm_rottary
+
 
 class NunchakuZSingleStreamAttnProcessor(ZSingleStreamAttnProcessor):
     """
@@ -31,32 +33,18 @@ def __call__(
         """
         Forward pass of the attention module. Adapted from diffusers.models.transformers.transformer_z_image.ZSingleStreamAttnProcessor#__call__.
         """
-
-        qkv = attn.to_qkv(hidden_states)
+        qkv = fused_qkv_norm_rottary(
+            hidden_states,
+            attn.to_qkv,
+            attn.norm_q,
+            attn.norm_k,
+            freqs_cis,
+        )
         query, key, value = qkv.chunk(3, dim=-1)
-
         query = query.unflatten(-1, (attn.heads, -1))
         key = key.unflatten(-1, (attn.heads, -1))
         value = value.unflatten(-1, (attn.heads, -1))
 
-        # Apply Norms
-        if attn.norm_q is not None:
-            query = attn.norm_q(query)
-        if attn.norm_k is not None:
-            key = attn.norm_k(key)
-
-        # Apply RoPE
-        def apply_rotary_emb(x_in: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
-            with torch.amp.autocast("cuda", enabled=False):
-                x = torch.view_as_complex(x_in.float().reshape(*x_in.shape[:-1], -1, 2))
-                freqs_cis = freqs_cis.unsqueeze(2)
-                x_out = torch.view_as_real(x * freqs_cis).flatten(3)
-                return x_out.type_as(x_in)  # todo
-
-        if freqs_cis is not None:
-            query = apply_rotary_emb(query, freqs_cis)
-            key = apply_rotary_emb(key, freqs_cis)
-
         # Cast to correct dtype
         dtype = query.dtype
         query, key = query.to(dtype), key.to(dtype)

nunchaku/models/transformers/transformer_zimage.py

@@ -8,6 +8,7 @@
 from typing import List, Optional
 
 import torch
+import torch.nn as nn
 from diffusers.models.attention import FeedForward
 from diffusers.models.attention_processor import Attention
 from diffusers.models.transformers.transformer_z_image import FeedForward as ZImageFeedForward
@@ -16,14 +17,39 @@
 
 from nunchaku.models.unets.unet_sdxl import NunchakuSDXLFeedForward
 
-from ...utils import get_precision
+from ...utils import get_precision, pad_tensor
 from ..attention import NunchakuBaseAttention
 from ..attention_processors.zimage import NunchakuZSingleStreamAttnProcessor
+from ..embeddings import pack_rotemb
 from ..linear import SVDQW4A4Linear
 from ..utils import fuse_linears
 from .utils import NunchakuModelLoaderMixin, patch_scale_key
 
 
+class NunchakuZImageRopeHook:
+    """
+    Hook class for caching and substition of packed `freqs_cis` tensor.
+    """
+
+    def __init__(self):
+        self.packed_cache = {}
+
+    def __call__(self, module: nn.Module, input_args: tuple, input_kwargs: dict):
+        freqs_cis: torch.Tensor = input_kwargs.get("freqs_cis", None)
+        if freqs_cis is None:
+            return None
+        cache_key = freqs_cis.data_ptr()
+        packed_freqs_cis = self.packed_cache.get(cache_key, None)
+        if packed_freqs_cis is None:
+            packed_freqs_cis = torch.view_as_real(freqs_cis).unsqueeze(3)
+            packed_freqs_cis = torch.flip(packed_freqs_cis, dims=[-1])
+            packed_freqs_cis = pack_rotemb(pad_tensor(packed_freqs_cis, 256, 1))
+            self.packed_cache[cache_key] = packed_freqs_cis
+        new_input_kwargs = input_kwargs.copy()
+        new_input_kwargs["freqs_cis"] = packed_freqs_cis
+        return input_args, new_input_kwargs
+
+
 class NunchakuZImageAttention(NunchakuBaseAttention):
     """
     Nunchaku-optimized Attention module for ZImage with quantized and fused QKV projections.
@@ -198,6 +224,7 @@ def _convert_feed_forward(block_list: List[ZImageTransformerBlock]):
             for _, block in enumerate(block_list):
                 block.feed_forward = _convert_z_image_ff(block.feed_forward)
 
+        self.skip_refiners = skip_refiners
         _patch_transformer_block(self.layers)
         if skip_refiners:
             _convert_feed_forward(self.noise_refiner)
@@ -207,6 +234,43 @@ def _convert_feed_forward(block_list: List[ZImageTransformerBlock]):
             _patch_transformer_block(self.context_refiner)
         return self
 
+    def register_rope_hook(self, rope_hook: NunchakuZImageRopeHook):
+        self.rope_hook_handles = []
+        for _, ly in enumerate(self.layers):
+            self.rope_hook_handles.append(ly.attention.register_forward_pre_hook(rope_hook, with_kwargs=True))
+        if not self.skip_refiners:
+            for _, nr in enumerate(self.noise_refiner):
+                self.rope_hook_handles.append(nr.attention.register_forward_pre_hook(rope_hook, with_kwargs=True))
+            for _, cr in enumerate(self.context_refiner):
+                self.rope_hook_handles.append(cr.attention.register_forward_pre_hook(rope_hook, with_kwargs=True))
+
+    def unregister_rope_hook(self):
+        for h in self.rope_hook_handles:
+            h.remove()
+        self.rope_hook_handles.clear()
+
+    def forward(
+        self,
+        x: List[torch.Tensor],
+        t,
+        cap_feats: List[torch.Tensor],
+        patch_size=2,
+        f_patch_size=1,
+        return_dict: bool = True,
+    ):
+        """
+        Adapted from diffusers.models.transformers.transformer_z_image.ZImageTransformer2DModel#forward
+
+        Register pre-forward hooks for caching and substitution of packed `freqs_cis` tensor for all attention submodules and unregister after forwarding is done.
+        """
+        rope_hook = NunchakuZImageRopeHook()
+        self.register_rope_hook(rope_hook)
+        try:
+            return super().forward(x, t, cap_feats, patch_size, f_patch_size, return_dict)
+        finally:
+            self.unregister_rope_hook()
+            del rope_hook
+
     @classmethod
     @utils.validate_hf_hub_args
     def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike[str], **kwargs):

nunchaku/ops/fused.py

@@ -3,6 +3,7 @@
 """
 
 import torch
+from diffusers.models.normalization import RMSNorm as DiffUsersRMSNorm
 from torch.nn import RMSNorm
 
 from nunchaku.models.linear import SVDQW4A4Linear
@@ -124,8 +125,12 @@ def fused_qkv_norm_rottary(
     - C_in: input features
     - C_out: output features
     """
-    assert norm_q is None or isinstance(norm_q, RMSNorm)
-    assert norm_k is None or isinstance(norm_k, RMSNorm)
+    assert (
+        norm_q is None or isinstance(norm_q, RMSNorm) or (isinstance(norm_q, DiffUsersRMSNorm) and norm_q.bias is None)
+    )
+    assert (
+        norm_k is None or isinstance(norm_k, RMSNorm) or (isinstance(norm_k, DiffUsersRMSNorm) and norm_k.bias is None)
+    )
 
     batch_size, seq_len, channels = x.shape
     x = x.view(batch_size * seq_len, channels)