Add key pre-transpose to sub-quadratic attention

Author: brkirch

modules/sd_hijack_optimizations.py

@@ -230,10 +230,11 @@ def sub_quad_attention_forward(self, x, context=None, mask=None):
     del context, context_k, context_v, x
 
     q = q.unflatten(-1, (h, -1)).transpose(1,2).flatten(end_dim=1)
-    k = k.unflatten(-1, (h, -1)).transpose(1,2).flatten(end_dim=1)
+    k_t = k.transpose(1,2).unflatten(1, (h, -1)).flatten(end_dim=1)
+    del k
     v = v.unflatten(-1, (h, -1)).transpose(1,2).flatten(end_dim=1)
 
-    x = sub_quad_attention(q, k, v, q_chunk_size=shared.cmd_opts.sub_quad_q_chunk_size, kv_chunk_size=shared.cmd_opts.sub_quad_kv_chunk_size, chunk_threshold=shared.cmd_opts.sub_quad_chunk_threshold, use_checkpoint=self.training)
+    x = sub_quad_attention(q, k_t, v, q_chunk_size=shared.cmd_opts.sub_quad_q_chunk_size, kv_chunk_size=shared.cmd_opts.sub_quad_kv_chunk_size, chunk_threshold=shared.cmd_opts.sub_quad_chunk_threshold, use_checkpoint=self.training, key_needs_transpose=False)
 
     x = x.unflatten(0, (-1, h)).transpose(1,2).flatten(start_dim=2)
 
@@ -243,7 +244,7 @@ def sub_quad_attention_forward(self, x, context=None, mask=None):
 
     return x
 
-def sub_quad_attention(q, k, v, q_chunk_size=1024, kv_chunk_size=None, kv_chunk_size_min=None, chunk_threshold=None, use_checkpoint=True):
+def sub_quad_attention(q, k, v, q_chunk_size=1024, kv_chunk_size=None, kv_chunk_size_min=None, chunk_threshold=None, use_checkpoint=True, key_needs_transpose=True):
     bytes_per_token = torch.finfo(q.dtype).bits//8
     batch_x_heads, q_tokens, _ = q.shape
     _, k_tokens, _ = k.shape
@@ -275,6 +276,7 @@ def sub_quad_attention(q, k, v, q_chunk_size=1024, kv_chunk_size=None, kv_chunk_
         kv_chunk_size=kv_chunk_size,
         kv_chunk_size_min = kv_chunk_size_min,
         use_checkpoint=use_checkpoint,
+        key_needs_transpose=key_needs_transpose,
     )
 
 

modules/sub_quadratic_attention.py

@@ -51,11 +51,12 @@ def _summarize_chunk(
     key: Tensor,
     value: Tensor,
     scale: float,
+    key_needs_transpose: bool,
 ) -> AttnChunk:
     attn_weights = torch.baddbmm(
         torch.empty(1, 1, 1, device=query.device, dtype=query.dtype),
         query,
-        key.transpose(1,2),
+        key.transpose(1,2) if key_needs_transpose else key,
         alpha=scale,
         beta=0,
     )
@@ -72,14 +73,18 @@ def _query_chunk_attention(
     value: Tensor,
     summarize_chunk: SummarizeChunk,
     kv_chunk_size: int,
+    key_needs_transpose: bool,
 ) -> Tensor:
-    batch_x_heads, k_tokens, k_channels_per_head = key.shape
+    if key_needs_transpose:
+        batch_x_heads, k_tokens, k_channels_per_head = key.shape
+    else:
+        batch_x_heads, k_channels_per_head, k_tokens = key.shape
     _, _, v_channels_per_head = value.shape
 
     def chunk_scanner(chunk_idx: int) -> AttnChunk:
         key_chunk = narrow_trunc(
             key,
-            1,
+            1 if key_needs_transpose else 2,
             chunk_idx,
             kv_chunk_size
         )
@@ -112,11 +117,12 @@ def _get_attention_scores_no_kv_chunking(
     key: Tensor,
     value: Tensor,
     scale: float,
+    key_needs_transpose: bool,
 ) -> Tensor:
     attn_scores = torch.baddbmm(
         torch.empty(1, 1, 1, device=query.device, dtype=query.dtype),
         query,
-        key.transpose(1,2),
+        key.transpose(1,2) if key_needs_transpose else key,
         alpha=scale,
         beta=0,
     )
@@ -136,7 +142,8 @@ def efficient_dot_product_attention(
     query_chunk_size=1024,
     kv_chunk_size: Optional[int] = None,
     kv_chunk_size_min: Optional[int] = None,
-    use_checkpoint=True,
+    use_checkpoint: Optional[bool] = True,
+    key_needs_transpose: Optional[bool] = True,
 ):
     """Computes efficient dot-product attention given query, key, and value.
       This is efficient version of attention presented in
@@ -151,12 +158,16 @@ def efficient_dot_product_attention(
         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)
+        use_checkpoint: Optional[bool]: whether to use checkpointing (recommended True for training, False for inference)
+        key_needs_transpose: Optional[bool]: whether key needs a transpose. defaults to True
       Returns:
         Output of shape `[batch * num_heads, query_tokens, channels_per_head]`.
       """
     batch_x_heads, q_tokens, q_channels_per_head = query.shape
-    _, k_tokens, _ = key.shape
+    if key_needs_transpose:
+        _, k_tokens, _ = key.shape
+    else:
+        _, _, 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)
@@ -171,17 +182,19 @@ def get_query_chunk(chunk_idx: int) -> Tensor:
             min(query_chunk_size, q_tokens)
         )
 
-    summarize_chunk: SummarizeChunk = partial(_summarize_chunk, scale=scale)
+    summarize_chunk: SummarizeChunk = partial(_summarize_chunk, scale=scale, key_needs_transpose=key_needs_transpose)
     summarize_chunk: SummarizeChunk = partial(checkpoint, summarize_chunk) if use_checkpoint else summarize_chunk
     compute_query_chunk_attn: ComputeQueryChunkAttn = partial(
         _get_attention_scores_no_kv_chunking,
-        scale=scale
+        scale=scale,
+        key_needs_transpose=key_needs_transpose,
     ) if k_tokens <= kv_chunk_size else (
         # fast-path for when there's just 1 key-value chunk per query chunk (this is just sliced attention btw)
         partial(
             _query_chunk_attention,
             kv_chunk_size=kv_chunk_size,
             summarize_chunk=summarize_chunk,
+            key_needs_transpose=key_needs_transpose,
         )
     )