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

FA3 kvcache + split kv + gqa parallelization #1236

Merged
merged 106 commits into from
Oct 15, 2024
Merged

FA3 kvcache + split kv + gqa parallelization #1236

merged 106 commits into from
Oct 15, 2024

Commits on Sep 30, 2024

  1. Adding the flash3 kv cache API. Just compiling for now. 

    KV cache functionality not added yet.
    @ganeshcolfax @jayhshah
    ganeshcolfax authored and jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    9dbd114 View commit details
    Browse the repository at this point in the history

  2. start extending seqlen traits for kv cache

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    7ee8ee4 View commit details
    Browse the repository at this point in the history

  3. added cache_batch_idx.

    @ganeshcolfax @jayhshah
    ganeshcolfax authored and jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    3fdf7ee View commit details
    Browse the repository at this point in the history

  4. adding python interface.

    @ganeshcolfax @jayhshah
    ganeshcolfax authored and jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    84e31c2 View commit details
    Browse the repository at this point in the history

  5. add test_kvcache.py.

    @ganeshcolfax @jayhshah
    ganeshcolfax authored and jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    e16053a View commit details
    Browse the repository at this point in the history

  6. enable use of actual seqlen for kv cache

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    be0e36d View commit details
    Browse the repository at this point in the history

  7. add new param to handle cache_batch_size

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    38ad0ac View commit details
    Browse the repository at this point in the history

  8. add semaphore for kv cache causal

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    57de4da View commit details
    Browse the repository at this point in the history

  9. add comparision with fa2.

    @ganeshcolfax @jayhshah
    ganeshcolfax authored and jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    435f86d View commit details
    Browse the repository at this point in the history

  10. change template parameter for SeqLenTraits for ease of further extension

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    74f160b View commit details
    Browse the repository at this point in the history

  11. modify seqlentraits for gqa parallelism

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    13bad55 View commit details
    Browse the repository at this point in the history

  12. modify Ktraits for decoding QO layouts

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    ccf5b9b View commit details
    Browse the repository at this point in the history

  13. decouple types of seqlen traits q and k

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    fc8f704 View commit details
    Browse the repository at this point in the history

  14. change logic of Q loads for gqa parallelization

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    d2f049c View commit details
    Browse the repository at this point in the history

  15. fix o strides

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    c6311e4 View commit details
    Browse the repository at this point in the history

  16. complete gqa parallel changes for non-causal

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    535b827 View commit details
    Browse the repository at this point in the history

  17. fix some errors

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    5704a1f View commit details
    Browse the repository at this point in the history

  18. add causal logic

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    64a9cfb View commit details
    Browse the repository at this point in the history

  19. add to kv cache api

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    a06f1f9 View commit details
    Browse the repository at this point in the history

  20. add in lse writeout and store zero

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    0c4cea9 View commit details
    Browse the repository at this point in the history

  21. refactor for split kv

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    0a1a0c2 View commit details
    Browse the repository at this point in the history

  22. re-enable fp16/bf16 fwd

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    1135dbd View commit details
    Browse the repository at this point in the history

  23. add 1 mma warpgroup option, enable splitkv for hdim 256

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    68ff3f7 View commit details
    Browse the repository at this point in the history

  24. fix bug with finalize for split kv

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    23bf5b0 View commit details
    Browse the repository at this point in the history

  25. delete unused files

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    ac19795 View commit details
    Browse the repository at this point in the history

  26. add hid=64.

    @ganeshcolfax @jayhshah
    ganeshcolfax authored and jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    1a5e40a View commit details
    Browse the repository at this point in the history

  27. change flash api for rebase

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    c75c243 View commit details
    Browse the repository at this point in the history

  28. avoid redundant compilation with combine kernel by only including nee… 

    …ded template params
    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    e9db102 View commit details
    Browse the repository at this point in the history

  29. change Element to OutputType for template param in combine kernel. On… 

    …ly matters for fp8 support
    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    9250969 View commit details
    Browse the repository at this point in the history

  30. fix wrong tile size for hdim 64

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    ec0130f View commit details
    Browse the repository at this point in the history

  31. revert OutputType change

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    68b4bb9 View commit details
    Browse the repository at this point in the history

  32. changes for correct lse write out for splits=1 and splits > 1 case.

    @ganeshcolfax @jayhshah
    ganeshcolfax authored and jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    9c97808 View commit details
    Browse the repository at this point in the history

  33. added num_split_heuristics.

    @ganeshcolfax @jayhshah
    ganeshcolfax authored and jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    ecc5c49 View commit details
    Browse the repository at this point in the history

  34. update parameters

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    f3e5bd4 View commit details
    Browse the repository at this point in the history

  35. remove unused code

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    78736b4 View commit details
    Browse the repository at this point in the history

  36. add num_split_heuristics.

    @ganeshcolfax @jayhshah
    ganeshcolfax authored and jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    e8c7b2e View commit details
    Browse the repository at this point in the history

  37. adding block_n and block_m for different headdim.

    @ganeshcolfax @jayhshah
    ganeshcolfax authored and jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    75a6ce2 View commit details
    Browse the repository at this point in the history

  38. initialize semaphore when num splits != 1

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    cf5bd5c View commit details
    Browse the repository at this point in the history

  39. change combine kernel to condition on lse=-INF

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    ac96c37 View commit details
    Browse the repository at this point in the history

  40. add gqa decoding logic.

    @ganeshcolfax @jayhshah
    ganeshcolfax authored and jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    099ca28 View commit details
    Browse the repository at this point in the history

  41. add split kv heuristic modifications

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    f53703b View commit details
    Browse the repository at this point in the history

  42. recent version.

    @ganeshcolfax @jayhshah
    ganeshcolfax authored and jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    a30863f View commit details
    Browse the repository at this point in the history

  43. more refactoring.

    @ganeshcolfax @jayhshah
    ganeshcolfax authored and jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    ffa48eb View commit details
    Browse the repository at this point in the history

  44. update test script to use heuristic

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    4a77193 View commit details
    Browse the repository at this point in the history

  45. Add some more cases to test script and raise thresholds a bit for max… 

    … error
    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    3615696 View commit details
    Browse the repository at this point in the history

  46. add reference from python.

    @ganeshcolfax @jayhshah
    ganeshcolfax authored and jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    24b4b4f View commit details
    Browse the repository at this point in the history

  47. Adding another test case.

    @ganeshcolfax @jayhshah
    ganeshcolfax authored and jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    c516d63 View commit details
    Browse the repository at this point in the history

  48. add variable seqlen case.

    @ganeshcolfax @jayhshah
    ganeshcolfax authored and jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    9a4941c View commit details
    Browse the repository at this point in the history

  49. all cases passed.

    @ganeshcolfax @jayhshah
    ganeshcolfax authored and jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    70ff847 View commit details
    Browse the repository at this point in the history

  50. change fp8 code path to allow for split kernel and kv cache without p… 

    …erf regression
    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    e36e004 View commit details
    Browse the repository at this point in the history

  51. set correct tolerance limit

    @ganeshcolfax @jayhshah
    ganeshcolfax authored and jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    cd55fb3 View commit details
    Browse the repository at this point in the history

  52. add 'in principle' fp8 kv cache support

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    2472e5e View commit details
    Browse the repository at this point in the history

  53. rebase with Is_local disabled temporarily

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    b5cac6d View commit details
    Browse the repository at this point in the history

  54. consolidate nblock min max methods

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    6111666 View commit details
    Browse the repository at this point in the history

  55. add Is_local back in

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    be481ca View commit details
    Browse the repository at this point in the history

  56. prune unused code

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    81d4024 View commit details
    Browse the repository at this point in the history

  57. enable Is_local with fp8

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    64a0a91 View commit details
    Browse the repository at this point in the history

  58. update composable kernel

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    0f560b7 View commit details
    Browse the repository at this point in the history

  59. separate out fp8 in test_flash_attn

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    cffef15 View commit details
    Browse the repository at this point in the history

  60. fix the test case and re-factor too.

    @ganeshcolfax
    ganeshcolfax committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    2b840ef View commit details
    Browse the repository at this point in the history

  61. Merge branch 'fa3-kvcache-gqa' of github.com:Dao-AILab/flash-attentio… 

    …n into fa3-kvcache-gqa
    @ganeshcolfax
    ganeshcolfax committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    5df67d2 View commit details
    Browse the repository at this point in the history

  62. dont write out zero for split kernel, only lse=-inf

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    aa45d75 View commit details
    Browse the repository at this point in the history

  63. Merge branch 'fa3-kvcache-gqa' of github.com:Dao-AILab/flash-attentio… 

    …n into fa3-kvcache-gqa
    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    33f20a3 View commit details
    Browse the repository at this point in the history

  64. fix composable kernel issue again

    @jayhshah
    jayhshah committed Sep 30, 2024
    Configuration menu
    Copy the full SHA
    f77d9f7 View commit details
    Browse the repository at this point in the history

Commits on Oct 1, 2024

  1. change default output type of fp8 kernel to bf16

    @jayhshah
    jayhshah committed Oct 1, 2024
    Configuration menu
    Copy the full SHA
    5e3864f View commit details
    Browse the repository at this point in the history

  2. remove deprecated fp8 code

    @jayhshah
    jayhshah committed Oct 1, 2024
    Configuration menu
    Copy the full SHA
    16eb1e5 View commit details
    Browse the repository at this point in the history

  3. correct indent

    @jayhshah
    jayhshah committed Oct 1, 2024
    Configuration menu
    Copy the full SHA
    7940377 View commit details
    Browse the repository at this point in the history

  4. add log max splits based on num splits to static switch

    @jayhshah
    jayhshah committed Oct 1, 2024
    Configuration menu
    Copy the full SHA
    31c71e0 View commit details
    Browse the repository at this point in the history

  5. change seq len class per discussion

    @jayhshah
    jayhshah committed Oct 1, 2024
    Configuration menu
    Copy the full SHA
    6bb1092 View commit details
    Browse the repository at this point in the history

Commits on Oct 2, 2024

  1. add fp8 test case.

    @ganeshcolfax
    ganeshcolfax committed Oct 2, 2024
    Configuration menu
    Copy the full SHA
    0085f04 View commit details
    Browse the repository at this point in the history

  2. fix submodule

    @jayhshah
    jayhshah committed Oct 2, 2024
    Configuration menu
    Copy the full SHA
    eaf8898 View commit details
    Browse the repository at this point in the history

  3. re-commiting.

    @ganeshcolfax
    ganeshcolfax committed Oct 2, 2024
    Configuration menu
    Copy the full SHA
    a44596f View commit details
    Browse the repository at this point in the history

  4. Revert "re-commiting." 

    This reverts commit a44596f.
    @ganeshcolfax
    ganeshcolfax committed Oct 2, 2024
    Configuration menu
    Copy the full SHA
    c0c58ee View commit details
    Browse the repository at this point in the history

  5. change fp8 tolerances to be smaller

    @jayhshah
    jayhshah committed Oct 2, 2024
    Configuration menu
    Copy the full SHA
    b8f9dc2 View commit details
    Browse the repository at this point in the history

  6. lower rtol for fp8 a bit

    @jayhshah
    jayhshah committed Oct 2, 2024
    Configuration menu
    Copy the full SHA
    49f1849 View commit details
    Browse the repository at this point in the history

  7. separate gqa compilation

    @jayhshah
    jayhshah committed Oct 2, 2024
    Configuration menu
    Copy the full SHA
    bb230b8 View commit details
    Browse the repository at this point in the history

  8. removed old gqa cu files and unified methods

    @jayhshah
    jayhshah committed Oct 2, 2024
    Configuration menu
    Copy the full SHA
    03200a7 View commit details
    Browse the repository at this point in the history

Commits on Oct 3, 2024

  1. reorg mma code for less redundancy

    @jayhshah
    jayhshah committed Oct 3, 2024
    Configuration menu
    Copy the full SHA
    930c8ca View commit details
    Browse the repository at this point in the history

  2. add crude hdim 64 heuristic

    @jayhshah
    jayhshah committed Oct 3, 2024
    Configuration menu
    Copy the full SHA
    bc4b872 View commit details
    Browse the repository at this point in the history

Commits on Oct 4, 2024

  1. add split kv benchmark script

    @jayhshah
    jayhshah committed Oct 4, 2024
    Configuration menu
    Copy the full SHA
    fff4b5c View commit details
    Browse the repository at this point in the history

  2. move descale tensor declarations outside of conditional

    @jayhshah
    jayhshah committed Oct 4, 2024
    Configuration menu
    Copy the full SHA
    aa0e699 View commit details
    Browse the repository at this point in the history

Commits on Oct 7, 2024

  1. fix bug with fp8 q layout

    @jayhshah
    jayhshah committed Oct 7, 2024
    Configuration menu
    Copy the full SHA
    785d978 View commit details
    Browse the repository at this point in the history

Commits on Oct 8, 2024

  1. adding rmem to gmem. (Not validating yet).

    @ganeshcolfax
    ganeshcolfax committed Oct 8, 2024
    Configuration menu
    Copy the full SHA
    8fbefa8 View commit details
    Browse the repository at this point in the history

  2. changes to use tiledcopy (still not passing).

    @ganeshcolfax
    ganeshcolfax committed Oct 8, 2024
    Configuration menu
    Copy the full SHA
    f0b4946 View commit details
    Browse the repository at this point in the history

Commits on Oct 9, 2024

  1. tests passing now for non-gqa impl

    @jayhshah
    jayhshah committed Oct 9, 2024
    Configuration menu
    Copy the full SHA
    8f45a8c View commit details
    Browse the repository at this point in the history

  2. move IsRegToGmem

    @jayhshah
    jayhshah committed Oct 9, 2024
    Configuration menu
    Copy the full SHA
    4a4dbd2 View commit details
    Browse the repository at this point in the history

Commits on Oct 10, 2024

  1. handle gqa_parallel with rmem-to-gmem. Not validating yet.

    @ganeshcolfax
    ganeshcolfax committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    a075e76 View commit details
    Browse the repository at this point in the history

  2. compiles and builes. Not validates.

    @ganeshcolfax
    ganeshcolfax committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    dc2c952 View commit details
    Browse the repository at this point in the history

  3. passes except for hdim=256.

    @ganeshcolfax
    ganeshcolfax committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    e49cb5f View commit details
    Browse the repository at this point in the history

Commits on Oct 11, 2024

  1. remove smem usage for when rmem -> gmem epilogue is used

    @jayhshah
    jayhshah committed Oct 11, 2024
    Configuration menu
    Copy the full SHA
    d437d3d View commit details
    Browse the repository at this point in the history

Commits on Oct 12, 2024

  1. better writeout logic with vectorization

    @jayhshah
    jayhshah committed Oct 12, 2024
    Configuration menu
    Copy the full SHA
    ab5d336 View commit details
    Browse the repository at this point in the history

  2. unify rmem -> gmem methods

    @jayhshah
    jayhshah committed Oct 12, 2024
    Configuration menu
    Copy the full SHA
    7169b23 View commit details
    Browse the repository at this point in the history

  3. uniform notation

    @jayhshah
    jayhshah committed Oct 12, 2024
    Configuration menu
    Copy the full SHA
    551b91f View commit details
    Browse the repository at this point in the history

Commits on Oct 14, 2024

  1. add rmem -> gmem for fp8

    @jayhshah
    jayhshah committed Oct 14, 2024
    Configuration menu
    Copy the full SHA
    eb9c0ee View commit details
    Browse the repository at this point in the history

  2. revert epi change for fp8 due to measured perf regression

    @jayhshah
    jayhshah committed Oct 14, 2024
    Configuration menu
    Copy the full SHA
    b0f067e View commit details
    Browse the repository at this point in the history

  3. refactor names

    @jayhshah
    jayhshah committed Oct 14, 2024
    Configuration menu
    Copy the full SHA
    35f3542 View commit details
    Browse the repository at this point in the history

  4. remove test code

    @jayhshah
    jayhshah committed Oct 14, 2024
    Configuration menu
    Copy the full SHA
    8374e1f View commit details
    Browse the repository at this point in the history

  5. remove constexpr checks for actual seqlen in mainloop

    @jayhshah
    jayhshah committed Oct 14, 2024
    Configuration menu
    Copy the full SHA
    1ecf821 View commit details
    Browse the repository at this point in the history

Commits on Oct 15, 2024

  1. remove Is_batch_dynamic from seqlen traits and handle fp8 perf regres… 

    …sion using smem boolean
    @jayhshah
    jayhshah committed Oct 15, 2024
    Configuration menu
    Copy the full SHA
    7c1473e View commit details
    Browse the repository at this point in the history

  2. change cu_seqlens_k to seqused_k for kv cache api

    @jayhshah
    jayhshah committed Oct 15, 2024
    Configuration menu
    Copy the full SHA
    c06cc0b View commit details
    Browse the repository at this point in the history

  3. adjust tolerances in test script for kv cache

    @jayhshah
    jayhshah committed Oct 15, 2024
    Configuration menu
    Copy the full SHA
    a7cce59 View commit details
    Browse the repository at this point in the history

  4. remove commented out code

    @jayhshah
    jayhshah committed Oct 15, 2024
    Configuration menu
    Copy the full SHA
    8efb953 View commit details
    Browse the repository at this point in the history

  5. prune more dead code

    @jayhshah
    jayhshah committed Oct 15, 2024
    Configuration menu
    Copy the full SHA
    b3d60fa View commit details
    Browse the repository at this point in the history

  6. comment out unimplemented kwargs from flash_attn_with_kvcache

    @jayhshah
    jayhshah committed Oct 15, 2024
    Configuration menu
    Copy the full SHA
    50cb90a View commit details
    Browse the repository at this point in the history

  7. fix integer sign compare warning

    @jayhshah
    jayhshah committed Oct 15, 2024
    Configuration menu
    Copy the full SHA
    dec7dee View commit details
    Browse the repository at this point in the history

  8. remove some debug code

    @jayhshah
    jayhshah committed Oct 15, 2024
    Configuration menu
    Copy the full SHA
    9b6cba1 View commit details
    Browse the repository at this point in the history