Code

sub-quadratic attention#1

Code

Open

Birch-san wants to merge 35 commits intomainBirch-san/diffusers:mainfrom subquad_attnBirch-san/diffusers:subquad_attnCopy head branch name to clipboard

+317-1

Commits

Commits on Dec 30, 2022

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

  Birch-san

  committed
  c810c32

  Copy full SHA for c810c32

• invoke efficient_dot_product_attention(). not currently giving correct results.

  Birch-san

  committed
  c9b3b9f

  Copy full SHA for c9b3b9f

• provide a way to skip checkpointing

  Birch-san

  committed
  70dc50d

  Copy full SHA for 70dc50d

• MPS fixes; now working

  Birch-san

  committed
  c794f0b

  Copy full SHA for 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 …

  Show description for 04a5cbe
  Birch-san

  committed
  04a5cbe

  Copy full SHA for 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

  Copy full SHA for 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

  Copy full SHA for 8694703

• add TODO about softmax

  Birch-san

  committed
  5bfe96d

  Copy full SHA for 5bfe96d

• typings

  Birch-san

  committed
  da8901b

  Copy full SHA for da8901b

• simplify protocols

  Birch-san

  committed
  0c4d82f

  Copy full SHA for 0c4d82f

• remove unused

  Birch-san

  committed
  c5e8e31

  Copy full SHA for c5e8e31

• simplify protocol

  Birch-san

  committed
  b16edc9

  Copy full SHA for b16edc9

• fix tensor shape destructuring

  Birch-san

  committed
  b7fc3a8

  Copy full SHA for b7fc3a8

• simplify dynamic_slice

  Birch-san

  committed
  8f003c2

  Copy full SHA for 8f003c2

• simplify chunk scanning

  Birch-san

  committed
  1334670

  Copy full SHA for 1334670

• inline sole use of map_pt function

  Birch-san

  committed
  0676c13

  Copy full SHA for 0676c13

• simplify

  Birch-san

  committed
  264dfb7

  Copy full SHA for 264dfb7

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

  Birch-san

  committed
  205f55b

  Copy full SHA for 205f55b

• fix query slicing

  Birch-san

  committed
  1880c0e

  Copy full SHA for 1880c0e

• fix kv chunking

  Birch-san

  committed
  8603c30

  Copy full SHA for 8603c30

• simplify dynamic slicing

  Birch-san

  committed
  96e0d8c

  Copy full SHA for 96e0d8c

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

  Birch-san

  committed
  63ca66d

  Copy full SHA for 63ca66d

• device arg fix no longer included

  Birch-san

  committed
  f4c0bf4

  Copy full SHA for f4c0bf4

• simplify

  Birch-san

  committed
  624123f

  Copy full SHA for 624123f

• clarify attributions now that algorithm has been substantially rewritten

  Birch-san

  committed
  5b92dab

  Copy full SHA for 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 …

  Show description for 60f0a5e
  Birch-san

  committed
  60f0a5e

  Copy full SHA for 60f0a5e

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

  Birch-san

  committed
  48db711

  Copy full SHA for 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

  Copy full SHA for ef20fb9

• remove debug notes

  Birch-san

  committed
  69a8d2e

  Copy full SHA for 69a8d2e

• explain kv fast-path

  Birch-san

  committed
  db25934

  Copy full SHA for db25934

• add fast-path for "1 query chunk"

  Birch-san

  committed
  7aa8bac

  Copy full SHA for 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…

  Show description for 59002c3
  Birch-san

  committed
  59002c3

  Copy full SHA for 59002c3

• Revert "move kv_chunk_size_min concern to callsite (1c4f107)" because equivalent fast-path for 1 query chunk, 1 kv chunk is already supported inside

  Birch-san

  committed
  a3152d8

  Copy full SHA for 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

  Copy full SHA for 0eafb95

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

  Birch-san

  committed
  9dc6822

  Copy full SHA for 9dc6822