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

Support AMD ROCm on FlashAttention 2 #1010

Merged
merged 46 commits into from
Jul 23, 2024
Merged

Support AMD ROCm on FlashAttention 2 #1010

merged 46 commits into from
Jul 23, 2024

Commits on May 31, 2024

  1. Support ck in fmha

    @rocking5566
    rocking5566 committed May 31, 2024
    Configuration menu
    Copy the full SHA
    6de9890 View commit details
    Browse the repository at this point in the history

  2. Add ck submodule

    @rocking5566
    rocking5566 committed May 31, 2024
    Configuration menu
    Copy the full SHA
    20d7715 View commit details
    Browse the repository at this point in the history

  3. Do not return lse if return_softmax == false

    @rocking5566
    rocking5566 committed May 31, 2024
    Configuration menu
    Copy the full SHA
    7aeae25 View commit details
    Browse the repository at this point in the history

  4. Use receipt to speed up ck compile time

    @rocking5566
    rocking5566 committed May 31, 2024
    Configuration menu
    Copy the full SHA
    025155d View commit details
    Browse the repository at this point in the history

Commits on Jun 7, 2024

  1. Integrate new version of ck_tile

    @rocking5566
    rocking5566 committed Jun 7, 2024
    Configuration menu
    Copy the full SHA
    5c89d54 View commit details
    Browse the repository at this point in the history

  2. Support dropout for mha_fwd()

    @rocking5566
    rocking5566 committed Jun 7, 2024
    Configuration menu
    Copy the full SHA
    86c2919 View commit details
    Browse the repository at this point in the history

  3. Add dropout to mha_varlen_fwd()

    @rocking5566
    rocking5566 committed Jun 7, 2024
    Configuration menu
    Copy the full SHA
    69715b1 View commit details
    Browse the repository at this point in the history

  4. Update ck to develop

    @rocking5566
    rocking5566 committed Jun 7, 2024
    Configuration menu
    Copy the full SHA
    a3e9e9c View commit details
    Browse the repository at this point in the history

  5. Extract padding function for dropout randval

    @rocking5566
    rocking5566 committed Jun 7, 2024
    Configuration menu
    Copy the full SHA
    6be061e View commit details
    Browse the repository at this point in the history

Commits on Jun 11, 2024

  1. Extract randval transformation function

    @rocking5566
    rocking5566 committed Jun 11, 2024
    Configuration menu
    Copy the full SHA
    161c006 View commit details
    Browse the repository at this point in the history

Commits on Jun 12, 2024

  1. Sync the code structure and coding style with FA

    @rocking5566
    rocking5566 committed Jun 12, 2024
    Configuration menu
    Copy the full SHA
    8569286 View commit details
    Browse the repository at this point in the history

  2. Remove this line, c++ api will handle this. 

    Sync with test_flash_attn.py
    @rocking5566
    rocking5566 committed Jun 12, 2024
    Configuration menu
    Copy the full SHA
    5728fb0 View commit details
    Browse the repository at this point in the history

  3. Merge pull request #61 from rocking5566/ck_tile/flash-attention 

    Ck tile/flash attention
    @carlushuang
    carlushuang authored Jun 12, 2024
    Configuration menu
    Copy the full SHA
    d93ced7 View commit details
    Browse the repository at this point in the history

  4. fix compile error

    @rocking5566
    rocking5566 committed Jun 12, 2024
    Configuration menu
    Copy the full SHA
    fb6a49f View commit details
    Browse the repository at this point in the history

Commits on Jun 20, 2024

  1. Add mha_bwd

    @rocking5566
    rocking5566 committed Jun 20, 2024
    Configuration menu
    Copy the full SHA
    26eab31 View commit details
    Browse the repository at this point in the history

Commits on Jun 21, 2024

  1. Generate dropout seed and offset from user generator

    @rocking5566
    rocking5566 committed Jun 21, 2024
    Configuration menu
    Copy the full SHA
    bf79417 View commit details
    Browse the repository at this point in the history

  2. update CK

    @rocking5566
    rocking5566 committed Jun 21, 2024
    Configuration menu
    Copy the full SHA
    6976d72 View commit details
    Browse the repository at this point in the history

  3. Add mha_varlen_bwd

    @rocking5566
    rocking5566 committed Jun 21, 2024
    Configuration menu
    Copy the full SHA
    09eeacf View commit details
    Browse the repository at this point in the history

Commits on Jun 24, 2024

  1. Use same python as build flash-attn to generate ck kernel

    @oraluben
    oraluben committed Jun 24, 2024
    Configuration menu
    Copy the full SHA
    aaea86c View commit details
    Browse the repository at this point in the history

  2. Fix bug of group mode fwd about returning softmax lse

    @rocking5566
    rocking5566 committed Jun 24, 2024
    Configuration menu
    Copy the full SHA
    52c3cb2 View commit details
    Browse the repository at this point in the history

  3. larger the test tollerance

    @rocking5566
    rocking5566 committed Jun 24, 2024
    Configuration menu
    Copy the full SHA
    08f5ad7 View commit details
    Browse the repository at this point in the history

  4. Add test_flash_attn_output() and test_flash_attn_varlen_output()

    @rocking5566
    rocking5566 committed Jun 24, 2024
    Configuration menu
    Copy the full SHA
    6fe8999 View commit details
    Browse the repository at this point in the history

  5. Always fill softmax_lse

    @rocking5566
    rocking5566 committed Jun 24, 2024
    Configuration menu
    Copy the full SHA
    a1b97f6 View commit details
    Browse the repository at this point in the history

Commits on Jun 25, 2024

  1. Remove duplicate benchmark script, since we already implement mha_bwd

    @rocking5566
    rocking5566 committed Jun 25, 2024
    Configuration menu
    Copy the full SHA
    41f3c32 View commit details
    Browse the repository at this point in the history

  2. Refine get value from tuple

    @rocking5566
    rocking5566 committed Jun 25, 2024
    Configuration menu
    Copy the full SHA
    e8007c7 View commit details
    Browse the repository at this point in the history

  3. Use default parameter for stream_config

    @rocking5566
    rocking5566 committed Jun 25, 2024
    Configuration menu
    Copy the full SHA
    ab4eca8 View commit details
    Browse the repository at this point in the history

  4. unblock all platform

    @rocking5566
    rocking5566 committed Jun 25, 2024
    Configuration menu
    Copy the full SHA
    89d56d0 View commit details
    Browse the repository at this point in the history

  5. Add comment

    @rocking5566
    rocking5566 committed Jun 25, 2024
    Configuration menu
    Copy the full SHA
    4dbfbd7 View commit details
    Browse the repository at this point in the history

  6. refine the test code

    @rocking5566
    rocking5566 committed Jun 25, 2024
    Configuration menu
    Copy the full SHA
    d218119 View commit details
    Browse the repository at this point in the history

  7. Refine naming

    @rocking5566
    rocking5566 committed Jun 25, 2024
    Configuration menu
    Copy the full SHA
    f41af53 View commit details
    Browse the repository at this point in the history

  8. Add unpack to namespace

    @rocking5566
    rocking5566 committed Jun 25, 2024
    Configuration menu
    Copy the full SHA
    e051724 View commit details
    Browse the repository at this point in the history

  9. Do not hardcode the warp size 64

    @rocking5566
    rocking5566 committed Jun 25, 2024
    Configuration menu
    Copy the full SHA
    7366478 View commit details
    Browse the repository at this point in the history

Commits on Jun 26, 2024

  1. Add more targets

    @rocking5566
    rocking5566 committed Jun 26, 2024
    Configuration menu
    Copy the full SHA
    26025eb View commit details
    Browse the repository at this point in the history

  2. Add README

    @rocking5566
    rocking5566 committed Jun 26, 2024
    Configuration menu
    Copy the full SHA
    2c5764e View commit details
    Browse the repository at this point in the history

  3. Merge pull request #65 from ROCm/ck-tile-bwd 

    Integrate ck tile backward
    @poyenc
    poyenc authored Jun 26, 2024
    Configuration menu
    Copy the full SHA
    64c16de View commit details
    Browse the repository at this point in the history

Commits on Jul 4, 2024

  1. Optimize mha_fwd if seqlen_q == 1

    @rocking5566
    rocking5566 committed Jul 4, 2024
    Configuration menu
    Copy the full SHA
    4d8dea1 View commit details
    Browse the repository at this point in the history

  2. Support get_wheel_url for rocm

    @rocking5566
    rocking5566 committed Jul 4, 2024
    Configuration menu
    Copy the full SHA
    a28673e View commit details
    Browse the repository at this point in the history

  3. Detect rocm environment by pytorch's IS_HIP_EXTENSION

    @rocking5566
    rocking5566 committed Jul 4, 2024
    Configuration menu
    Copy the full SHA
    d118bcc View commit details
    Browse the repository at this point in the history

Commits on Jul 8, 2024

  1. Merge commit 'd118bcc3bcab7b2539bb2e05a48d5a371266c8ef' into ck_tile

    @rocking5566
    rocking5566 committed Jul 8, 2024
    Configuration menu
    Copy the full SHA
    0495f03 View commit details
    Browse the repository at this point in the history

  2. update to lastest ck

    @rocking5566
    rocking5566 committed Jul 8, 2024
    Configuration menu
    Copy the full SHA
    b7ffdbf View commit details
    Browse the repository at this point in the history

Commits on Jul 10, 2024

  1. Merge branch 'ck_tile' into fix-python-location

    @oraluben
    oraluben authored Jul 10, 2024
    Configuration menu
    Copy the full SHA
    43ff3d5 View commit details
    Browse the repository at this point in the history

  2. Add necessary compile flag

    @rocking5566
    rocking5566 committed Jul 10, 2024
    Configuration menu
    Copy the full SHA
    3b62d48 View commit details
    Browse the repository at this point in the history

  3. Merge pull request #66 from oraluben/fix-python-location 

    Use same python as build flash-attn to generate ck kernel
    @rocking5566
    rocking5566 authored Jul 10, 2024
    Configuration menu
    Copy the full SHA
    c0f637a View commit details
    Browse the repository at this point in the history

  4. Merge pull request #67 from ROCm/update-ck 

    [WIP] update to latest ck
    @rocking5566
    rocking5566 authored Jul 10, 2024
    Configuration menu
    Copy the full SHA
    23a2b1c View commit details
    Browse the repository at this point in the history

Commits on Jul 22, 2024

  1. Merge branch 'main' into ck_tile

    @rocking5566
    rocking5566 authored Jul 22, 2024
    Configuration menu
    Copy the full SHA
    a4417c7 View commit details
    Browse the repository at this point in the history

  2. Sync the api with upstream FA

    @rocking5566
    rocking5566 committed Jul 22, 2024
    Configuration menu
    Copy the full SHA
    d83c412 View commit details
    Browse the repository at this point in the history