move kv_chunk_size_min concern to callsite, since if caller knows final kv_chunk_size: they can notice when no chunking would happen at all, and use fast-path. note: there's a question of whether that concern belongs *inside* the algorithm. but it'd feel weird for chunked attention to have a no-chunking-at-all branch.

Birch-san · Birch-san · commit 1c4f10748e31 · 2022-12-28T02:16:04.000Z
diff --git a/src/diffusers/models/cross_attention.py b/src/diffusers/models/cross_attention.py
@@ -17,6 +17,7 @@
 import torch
 import torch.nn.functional as F
 from torch import nn, Tensor
+import math
 
 from ..utils.import_utils import is_xformers_available
 
@@ -318,14 +319,19 @@ def __call__(
         _, k_tokens, _ = key.shape
         qk_matmul_size_bytes = batch_x_heads * bytes_per_token * q_tokens * k_tokens
 
-        if self.chunk_threshold_bytes is None or qk_matmul_size_bytes > self.chunk_threshold_bytes:
+        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):
             hidden_states = efficient_dot_product_attention(
                 query,
                 key,
                 value,
                 query_chunk_size=self.query_chunk_size,
-                kv_chunk_size=self.kv_chunk_size,
-                kv_chunk_size_min=self.kv_chunk_size_min,
+                kv_chunk_size=kv_chunk_size,
                 use_checkpoint=attn.training,
             )
         else:
diff --git a/src/diffusers/models/sub_quadratic_attention.py b/src/diffusers/models/sub_quadratic_attention.py
@@ -125,7 +125,6 @@ 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.
@@ -140,7 +139,6 @@ 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]`.
@@ -149,10 +147,6 @@ 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,
