sub-quadratic attention#1

Open

Birch-san wants to merge 35 commits intomainfrom

Commits on Dec 30, 2022

initial commit of sub-quadratic attention source from https://github.com/AminRezaei0x443/memory-efficient-attention.
Birch-san
committed
invoke efficient_dot_product_attention(). not currently giving correct results.
Birch-san
committed
provide a way to skip checkpointing
Birch-san
committed
MPS fixes; now working
Birch-san
committed
eliminate all einsums. assume 3D tensor [batch * num_heads, tokens, channels_per_head] in order to make use of batched matmuls. fuse multiply into matmul. breaks bias, mask in exchange for massive …
Birch-san
committed
remove the bits that I broke in the pursuit of speed (mask, bias, weights_calc_fn, calc_fn_data) and unused vars
Birch-san
committed
clarify comment; verified that upcast_attention is indeed still helpful for SD 2.1. but remove value float32, having established that it works without.
Birch-san
committed
add TODO about softmax
Birch-san
committed
typings
Birch-san
committed
simplify protocols
Birch-san
committed
remove unused
Birch-san
committed
simplify protocol
Birch-san
committed
fix tensor shape destructuring
Birch-san
committed
simplify dynamic_slice
Birch-san
committed
simplify chunk scanning
Birch-san
committed
inline sole use of map_pt function
Birch-san
committed
simplify
Birch-san
committed
no longer using original utilities from memory-efficient-attention repository
Birch-san
committed
fix query slicing
Birch-san
committed
fix kv chunking
Birch-san
committed
simplify dynamic slicing
Birch-san
committed
removed bias, mask, weights, calc_fn, and the conditions controlling them
Birch-san
committed
device arg fix no longer included
Birch-san
committed
simplify
Birch-san
committed
clarify attributions now that algorithm has been substantially rewritten
Birch-san
committed
add chunk_threshold_bytes to let you specify your safe memory limit, to prefer fast-path whenever unchunked attention would fit into memory. add kv_chunk_size_min to control the kv_chunk_size=None …
Birch-san
committed
fast path for when we're just attention-slicing (i.e. chunking query but not kv)
Birch-san
committed
default kv_chunk_size was meant to be sqrt() of global key size, not of chunk key size. improve separation of concerns.
Birch-san
committed
remove debug notes
Birch-san
committed
explain kv fast-path
Birch-san
committed
add fast-path for "1 query chunk"
Birch-san
committed
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 t…
Birch-san
committed
Revert "move kv_chunk_size_min concern to callsite (1c4f10748e31d18514ff1d5ed9fd9c67a278275b)" because equivalent fast-path for 1 query chunk, 1 kv chunk is already supported inside
Birch-san
committed
de-duplicate fast-path for "matmul < quota". we can just ask for everything in one chunk, to re-use an existing fast-path.
Birch-san
committed
pre-transpose key, rather than transposing it then undoing the transpose during the matmul
Birch-san
committed