sub-quadratic attention by Birch-san · Pull Request #1 · Birch-san/diffusers

initial commit of sub-quadratic attention source from https://github.com/AminRezaei0x443/memory-efficient-attention.

Birch-san

committed

c810c32

invoke efficient_dot_product_attention(). not currently giving correct results.

Birch-san

committed

c9b3b9f

provide a way to skip checkpointing

Birch-san

committed

70dc50d

MPS fixes; now working

Birch-san

committed

c794f0b

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

04a5cbe

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

b44fa12

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

8694703

add TODO about softmax

Birch-san

committed

5bfe96d

typings

Birch-san

committed

da8901b

simplify protocols

Birch-san

committed

0c4d82f

remove unused

Birch-san

committed

c5e8e31

simplify protocol

Birch-san

committed

b16edc9

fix tensor shape destructuring

Birch-san

committed

b7fc3a8

simplify dynamic_slice

Birch-san

committed

8f003c2

simplify chunk scanning

Birch-san

committed

1334670

inline sole use of map_pt function

Birch-san

committed

0676c13

simplify

Birch-san

committed

264dfb7

no longer using original utilities from memory-efficient-attention repository

Birch-san

committed

205f55b

fix query slicing

Birch-san

committed

1880c0e

fix kv chunking

Birch-san

committed

8603c30

simplify dynamic slicing

Birch-san

committed

96e0d8c

removed bias, mask, weights, calc_fn, and the conditions controlling them

Birch-san

committed

63ca66d

device arg fix no longer included

Birch-san

committed

f4c0bf4

simplify

Birch-san

committed

624123f

clarify attributions now that algorithm has been substantially rewritten

Birch-san

committed

5b92dab

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

60f0a5e

fast path for when we're just attention-slicing (i.e. chunking query but not kv)

Birch-san

committed

48db711

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

ef20fb9

remove debug notes

Birch-san

committed

69a8d2e

explain kv fast-path

Birch-san

committed

db25934

add fast-path for "1 query chunk"

Birch-san

committed

7aa8bac

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

59002c3

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

a3152d8

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

0eafb95

pre-transpose key, rather than transposing it then undoing the transpose during the matmul

Birch-san

committed

9dc6822

Navigation Menu

sub-quadratic attention - #1

sub-quadratic attention#1
Birch-san wants to merge 35 commits into
mainfrom
subquad_attn

Commits on Dec 30, 2022

initial commit of sub-quadratic attention source from https://github.com/AminRezaei0x443/memory-efficient-attention.

invoke efficient_dot_product_attention(). not currently giving correct results.

provide a way to skip checkpointing

MPS fixes; now working

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 …

remove the bits that I broke in the pursuit of speed (mask, bias, weights_calc_fn, calc_fn_data) and unused vars

clarify comment; verified that upcast_attention is indeed still helpful for SD 2.1. but remove value float32, having established that it works without.

add TODO about softmax

typings

simplify protocols

remove unused

simplify protocol

fix tensor shape destructuring

simplify dynamic_slice

simplify chunk scanning

inline sole use of map_pt function

simplify

no longer using original utilities from memory-efficient-attention repository

fix query slicing

fix kv chunking

simplify dynamic slicing

removed bias, mask, weights, calc_fn, and the conditions controlling them

device arg fix no longer included

simplify

clarify attributions now that algorithm has been substantially rewritten

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 …

fast path for when we're just attention-slicing (i.e. chunking query but not kv)

default kv_chunk_size was meant to be sqrt() of global key size, not of chunk key size. improve separation of concerns.

remove debug notes

explain kv fast-path

add fast-path for "1 query chunk"

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…

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

de-duplicate fast-path for "matmul < quota". we can just ask for everything in one chunk, to re-use an existing fast-path.

pre-transpose key, rather than transposing it then undoing the transpose during the matmul