Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

sub-quadratic attention #1

Open
wants to merge 35 commits into
base: main
Choose a base branch
from
Open

sub-quadratic attention #1

wants to merge 35 commits into from

Commits on Dec 30, 2022

  1. initial commit of sub-quadratic attention source from https://github.… 

    …com/AminRezaei0x443/memory-efficient-attention.
    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    c810c32 View commit details
    Browse the repository at this point in the history

  2. invoke efficient_dot_product_attention(). not currently giving correc… 

    …t results.
    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    c9b3b9f View commit details
    Browse the repository at this point in the history

  3. provide a way to skip checkpointing

    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    70dc50d View commit details
    Browse the repository at this point in the history

  4. MPS fixes; now working

    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    c794f0b View commit details
    Browse the repository at this point in the history

  5. eliminate all einsums. assume 3D tensor [batch * num_heads, tokens, c… 

    …hannels_per_head] in order to make use of batched matmuls. fuse multiply into matmul. breaks bias, mask in exchange for massive speedup.
    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    04a5cbe View commit details
    Browse the repository at this point in the history

  6. remove the bits that I broke in the pursuit of speed (mask, bias, wei… 

    …ghts_calc_fn, calc_fn_data) and unused vars
    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    b44fa12 View commit details
    Browse the repository at this point in the history

  7. clarify comment; verified that upcast_attention is indeed still helpf… 

    …ul for SD 2.1. but remove value float32, having established that it works without.
    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    8694703 View commit details
    Browse the repository at this point in the history

  8. add TODO about softmax

    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    5bfe96d View commit details
    Browse the repository at this point in the history

  9. typings

    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    da8901b View commit details
    Browse the repository at this point in the history

  10. simplify protocols

    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    0c4d82f View commit details
    Browse the repository at this point in the history

  11. remove unused

    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    c5e8e31 View commit details
    Browse the repository at this point in the history

  12. simplify protocol

    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    b16edc9 View commit details
    Browse the repository at this point in the history

  13. fix tensor shape destructuring

    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    b7fc3a8 View commit details
    Browse the repository at this point in the history

  14. simplify dynamic_slice

    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    8f003c2 View commit details
    Browse the repository at this point in the history

  15. simplify chunk scanning

    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    1334670 View commit details
    Browse the repository at this point in the history

  16. inline sole use of map_pt function

    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    0676c13 View commit details
    Browse the repository at this point in the history

  17. simplify

    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    264dfb7 View commit details
    Browse the repository at this point in the history

  18. no longer using original utilities from memory-efficient-attention re… 

    …pository
    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    205f55b View commit details
    Browse the repository at this point in the history

  19. fix query slicing

    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    1880c0e View commit details
    Browse the repository at this point in the history

  20. fix kv chunking

    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    8603c30 View commit details
    Browse the repository at this point in the history

  21. simplify dynamic slicing

    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    96e0d8c View commit details
    Browse the repository at this point in the history

  22. removed bias, mask, weights, calc_fn, and the conditions controlling … 

    …them
    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    63ca66d View commit details
    Browse the repository at this point in the history

  23. device arg fix no longer included

    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    f4c0bf4 View commit details
    Browse the repository at this point in the history

  24. simplify

    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    624123f View commit details
    Browse the repository at this point in the history

  25. clarify attributions now that algorithm has been substantially rewritten

    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    5b92dab View commit details
    Browse the repository at this point in the history

  26. 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 behaviour, so that sqrt(key_tokens) does not pick too small of a chunk size
    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    60f0a5e View commit details
    Browse the repository at this point in the history

  27. fast path for when we're just attention-slicing (i.e. chunking query … 

    …but not kv)
    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    48db711 View commit details
    Browse the repository at this point in the history

  28. default kv_chunk_size was meant to be sqrt() of global key size, not … 

    …of chunk key size. improve separation of concerns.
    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    ef20fb9 View commit details
    Browse the repository at this point in the history

  29. remove debug notes

    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    69a8d2e View commit details
    Browse the repository at this point in the history

  30. explain kv fast-path

    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    db25934 View commit details
    Browse the repository at this point in the history

  31. add fast-path for "1 query chunk"

    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    7aa8bac View commit details
    Browse the repository at this point in the history

  32. move kv_chunk_size_min concern to callsite, since if caller knows fin… 

    …al kv_chunk_size: they can notice when no chunking would happen at all, and use fast-path. note: there's a question of whether that concern belongs *inside* the algorithm. but it'd feel weird for chunked attention to have a no-chunking-at-all branch.
    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    59002c3 View commit details
    Browse the repository at this point in the history

  33. 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
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    a3152d8 View commit details
    Browse the repository at this point in the history

  34. de-duplicate fast-path for "matmul < quota". we can just ask for ever… 

    …ything in one chunk, to re-use an existing fast-path.
    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    0eafb95 View commit details
    Browse the repository at this point in the history

  35. pre-transpose key, rather than transposing it then undoing the transp… 

    …ose during the matmul
    @Birch-san
    Birch-san committed Dec 30, 2022
    Configuration menu
    Copy the full SHA
    9dc6822 View commit details
    Browse the repository at this point in the history