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

Add fmha dropout and backward support using ck-tiled for ROCM #10

Closed
wants to merge 50 commits into from
Closed

Add fmha dropout and backward support using ck-tiled for ROCM #10

wants to merge 50 commits into from

Conversation

qianfengz
Copy link
Collaborator

With this PR. The following is done:

  1. Fmha forward is completely implemented (added dropout, LSE outputting)
  2. Fmha backward is added and completely implemented (batched mode, grouped mode, flexible mask-type, bias-grad)
  3. Kernel template files removed completely by re-using the kernel files under ck_tile/examples/91_tile_program/fmha
  4. Standalone random-number generating kernel added (located under ck_tile/example/91_tile_program/fmha)

The following scripts are used to verify the implementation:

#> pytest tests/test_mem_eff_attention.py::test_forward
#> pytest tests/test_mem_eff_attention.py::test_mqa_forward
#> pytest tests/test_mem_eff_attention.py::test_backward
#> pytest tests/test_mem_eff_attention.py::test_dropout
#> pytest tests/test_mem_eff_attention.py::test_dropout_backward_ck

qianfengz added 30 commits March 20, 2024 21:50
@qianfengz
Add the c++ extension to the latest change of ck_tile/dev fwd kernel …
2da2927 
…(added droppout)
@qianfengz
Add the c++ extension to use ck_tile/dev/ fmha bwd kernel
9189e45
@qianfengz
Update to add dropout for fmah backward
28e713d
@qianfengz
Update in attention.cpp to align efficient_attention_backward_ck inte…
4ef7eba 
…rface parameters
@qianfengz
Enable BwdOp in ck.py
48a5f3e
@qianfengz
Support grad_out to have different strides as out
2e45012
@qianfengz
Force seqstart_q/seqstart_k to be in device memory in ck.py
566d26f
@qianfengz
Remove duplicated codes in ck_tiled_fmha_grouped_forward.h/infer.h
fc6c4a6
@qianfengz
Use optimized async pipeline where 8x headdim length is assumed
ff0db07
@qianfengz
Fix in batched_infer
0f4a171
@qianfengz
Update to track ck_tile/opt_padding_fa_train_xformers branch
0d6b915
@qianfengz
Update to use the newer FmhaFwdEpilogue
4713576
@qianfengz
Update to use the newer FmhaFwdEpilogue for grouped infer/forward
b1a1009
@qianfengz
Temporarily disable the using of QRKSVSAsync() pipeline
97e4e20
@qianfengz
Implement the ck_rand_uniform interface for generating random number …
6fbd05d 
…tensor
@qianfengz
Add dropout to the infer path (needed by xformers test_dropout)
2ef3b3f
@qianfengz
Update to support test_dropout and test_dropout_backward tests
930bb25
@qianfengz
Update the padding method in batched_backward.h
bdbc956
@qianfengz
Update the OGradDotO kernel padding method
44fff29
@qianfengz
Change the backward padding checking condition
d5c2d88
@qianfengz
Add batch_stride_lse/d parameters to adapt grouped mode forward/backw…
ce9c23c 
…ard to [num_batches, H, MaxSeqlenQ] layout
@qianfengz
Fill the grad_bias in advance
dafea78
@qianfengz
Add support for kHasBiasGrad as instance template
06ad689
@qianfengz
Remove using hdim_stride_do in fmha backward
bdd6291
@qianfengz
Force kPadSeqLenQ/kPadSeqLenK to be true in batched-backward to save …
410f814 
…compiling time
@qianfengz
Fix missing passing of {philox_seed, philox_offset} in inference path
2712dff
@qianfengz
Use SimplifiedGenericAttentionMask to replace GenericAttentionMask
7c27a82
@qianfengz
Shorten the instance file names
46c491e
@qianfengz
Rename the template parameters
4c6c08d
@qianfengz
Simplify the names of the dispatch class and interfaces
411ccd6
qianfengz and others added 17 commits April 16, 2024 06:28
@qianfengz
Changes to reuse the kernel files under ck_tile examples/91_tile_prog…
812a529 
…ram/fmha folder
@qianfengz
Update test_mem_eff_attention.py for test_dropout/test_dropout_backwa…
51b4223 
…rd/test_backward on rocm
@qianfengz
Tiny change to the philox_cuda_state input setting
d10ef79
@qianfengz
Allocate logsumexp to ensure aligned access by each thread-group
25bd720
@qianfengz
Add checking for query/key headdim size attention_backward_generic
abfdc27
@qianfengz
Using ck_tile/opt_padding_fa_train_pr2 and synchronize the backward c…
ff95367 
…odes with the changes
@qianfengz
Enable using async pipeline in the batched inference path for perform…
93469ab 
…ance
@qianfengz
Re-organize cpp instances for calling fmha infer kernel
2c8626b
@qianfengz
Re-organize cpp instances for calling fmha forward kernel
bdd716c
@qianfengz
Re-organize cpp instances for calling fmha backward kernel
44d4592
@qianfengz
Position the composable_kernel_tiled to ck_tile/opt_padding_fa_train …
51ca91b 
…branch
@qianfengz
Update to synchronize with the latest commits in ck_tile/opt_padding_…
1693683 
…fa_train
@carlushuang
update submodule to public
b7aa908
@qianfengz
Merge pull request #7 from ROCm/origin/test_opt_padding_train_public
9a878d9 
update submodule to public
@qianfengz
Update to the criteria for padding seqlen_k in batched infer/forward
b4fa26d
@qianfengz
Keep latest track of ck-tile commits
ee7950f
@qianfengz
Tiny fixing to the decoder including
74dfdfe
@qianfengz qianfengz closed this May 9, 2024
@qianfengz qianfengz deleted the test_opt_padding_train branch May 14, 2024 08:03
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers
No reviews
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@qianfengz @carlushuang