Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit fbd3ac7

Browse files
committedDec 28, 2022
add fast-path for "1 query chunk"
1 parent dcc96c9 commit fbd3ac7

File tree

1 file changed

+8
-0
lines changed

1 file changed

+8
-0
lines changed
 

‎src/diffusers/models/sub_quadratic_attention.py

+8
Original file line numberDiff line numberDiff line change
@@ -173,6 +173,14 @@ def get_query_chunk(chunk_idx: int) -> Tensor:
173173
summarize_chunk=summarize_chunk,
174174
)
175175
)
176+
177+
if q_tokens <= query_chunk_size:
178+
# fast-path for when there's just 1 query chunk
179+
return compute_query_chunk_attn(
180+
query=query,
181+
key=key,
182+
value=value,
183+
)
176184

177185
# TODO: maybe we should use torch.empty_like(query) to allocate storage in-advance,
178186
# and pass slices to be mutated, instead of torch.cat()ing the returned slices

0 commit comments

Comments
 (0)
Please sign in to comment.