Revert "move kv_chunk_size_min concern to callsite (1c4f107)" because equivalent fast-path for 1 query chunk, 1 kv chunk is already supported inside

Author: Birch-san

src/diffusers/models/cross_attention.py

@@ -17,7 +17,6 @@
 import torch
 import torch.nn.functional as F
 from torch import nn, Tensor
-import math
 
 from ..utils.import_utils import is_xformers_available
 
@@ -319,19 +318,14 @@ def __call__(
         _, k_tokens, _ = key.shape
         qk_matmul_size_bytes = batch_x_heads * bytes_per_token * q_tokens * k_tokens
 
-        kv_chunk_size = min(self.kv_chunk_size or int(math.sqrt(k_tokens)), k_tokens)
-        if self.kv_chunk_size_min is not None:
-            kv_chunk_size = max(kv_chunk_size, self.kv_chunk_size_min)
-        
-        uses_chunking = q_tokens > self.query_chunk_size or k_tokens > kv_chunk_size
-
-        if uses_chunking and (self.chunk_threshold_bytes is None or qk_matmul_size_bytes > self.chunk_threshold_bytes):
+        if self.chunk_threshold_bytes is None or qk_matmul_size_bytes > self.chunk_threshold_bytes:
             hidden_states = efficient_dot_product_attention(
                 query,
                 key,
                 value,
                 query_chunk_size=self.query_chunk_size,
-                kv_chunk_size=kv_chunk_size,
+                kv_chunk_size=self.kv_chunk_size,
+                kv_chunk_size_min=self.kv_chunk_size_min,
                 use_checkpoint=attn.training,
             )
         else:

src/diffusers/models/sub_quadratic_attention.py

@@ -125,6 +125,7 @@ def efficient_dot_product_attention(
     value: Tensor,
     query_chunk_size=1024,
     kv_chunk_size: Optional[int] = None,
+    kv_chunk_size_min: Optional[int] = None,
     use_checkpoint=True,
 ):
     """Computes efficient dot-product attention given query, key, and value.
@@ -139,6 +140,7 @@ def efficient_dot_product_attention(
           `[batch * num_heads, tokens, channels_per_head]`.
         query_chunk_size: int: query chunks size
         kv_chunk_size: Optional[int]: key/value chunks size. if None: defaults to sqrt(key_tokens)
+        kv_chunk_size_min: Optional[int]: key/value minimum chunk size. only considered when kv_chunk_size is None. changes `sqrt(key_tokens)` into `max(sqrt(key_tokens), kv_chunk_size_min)`, to ensure our chunk sizes don't get too small (smaller chunks = more chunks = less concurrent work done).
         use_checkpoint: bool: whether to use checkpointing (recommended True for training, False for inference)
       Returns:
         Output of shape `[batch * num_heads, query_tokens, channels_per_head]`.
@@ -147,6 +149,10 @@ def efficient_dot_product_attention(
     _, k_tokens, _ = key.shape
     scale = q_channels_per_head ** -0.5
 
+    kv_chunk_size = min(kv_chunk_size or int(math.sqrt(k_tokens)), k_tokens)
+    if kv_chunk_size_min is not None:
+        kv_chunk_size = max(kv_chunk_size, kv_chunk_size_min)
+
     def get_query_chunk(chunk_idx: int) -> Tensor:
         return dynamic_slice(
             query,