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 page kvcache in AMD ROCm #1198

Merged
merged 52 commits into from
Sep 16, 2024
Merged

Support page kvcache in AMD ROCm #1198

merged 52 commits into from
Sep 16, 2024

Conversation

rocking5566
Copy link
Contributor

@rocking5566 rocking5566 commented Sep 3, 2024 •
edited
Loading

In this PR,

  1. Update the ROCm backend (CK), so I modify how to call ck due to changing of CK api.
  2. Improve backward performance by updating the CK (1)
  3. Implement mha_fwd_kvcache().
  4. Change of compile flag to support ROCm6.2
  5. Change bf16 rounding to RTN (round to nearest)

pass all the test in MI200 & MI300 with ROCm6.2 (skip some test case in kvcache, follow the original test)
image

rocking5566 and others added 30 commits July 23, 2024 16:21
@rocking5566
Integrate ck branch of ck_tile/fa_bwd_opt
b5a4204
@rocking5566
Assume dq and q share the same stride
1e3416d
@rocking5566
update ck
bc78de1
@rocking5566
Integrate more stride of dq_acc
46d1cff
@rocking5566
Revert fwd dropout
eac0e38
@rocking5566
Fix paremeter order
b14f245
@rocking5566
Integrate ck with more stride
3180632
@rocking5566
update the limit of hdim of bwd
c7ac11f
@rocking5566
Check argument
6af216c
@rocking5566
Add test_flash_attn_causal
05b657e
@rocking5566
Support unpad lse
dbe28cb
@rocking5566
Add test_flash_attn_varlen_causal, test_flash_attn_race_condition, te…
7b712b2 
…st_flash_attn_bwd_overflow, test_flash_attn_bwd_transpose, test_flash_attn_bwd_varlen_overflow, test_flash_attn_deterministic, test_flash_attn_varlen_deterministic
@rocking5566
Fix stride and Kn0
5346b5b
@rocking5566
Fix CK sync issue
381bbdd
@rocking5566
Fix typo
6ac697c
@rocking5566
Merge commit '3669b25206d5938e3cc74a5f7860e31c38af8204' into ck_impro…
bc5dd34 
…ve_v0.1.1
@rocking5566
Update CK for changing of fmha_fwd_args
909d66f
@rocking5566
Add kvcache tmp
6928524
@rocking5566
Add kvcache
23000f7
@rocking5566
Fix comment
2c80b86
@rocking5566
Sync behavior with ck
d2ed413
@rocking5566
Update CK to develop
c0c2f8f
@rocking5566
remove large test case
e037142
@rocking5566
Merge pull request #70 from ROCm/ck_bwd_opt
d38c59b 
Improve FMHA bwd
@rocking5566
Merge remote-tracking branch 'origin/ck_improve_v0.1.1' into ck_tile/…
a37da96 
…kvcache
@rocking5566
Add kvcache test
d8de7a6
@rocking5566
Fix page_block_size in arg
a84d12d
@rocking5566
Minor fix
5b4546c
@rocking5566
Fix stride error
d6aac9e
@rocking5566
Update seqlen of kvcache before splitkv
ae24800
rocking5566 and others added 22 commits August 21, 2024 12:58
@rocking5566
Fix compile error
e2d3f5b
@rocking5566
Fix bug of hdim is not 8x
bb7a439
@rocking5566
Fit ck arg
7b18b87
@rocking5566
support adaptive num_splits
94e054f
@rocking5566
add more tests
9316aa6
@rocking5566
Refine test tolerance
7815e3b
@rocking5566
update CK
27095f2
@rocking5566
Move override_num_splits_if_necessary into cpp
4a25f60
@rocking5566
update ck
22eee22
@rocking5566
Update ck
007ae03
@rocking5566
Merge pull request #74 from ROCm/ck_tile/kvcache
7259227 
Ck tile/kvcache
@rocking5566
Merge branch 'Dao-AILab:main' into ck_improve_v0.1.1
444ab9f
@rocking5566
Support different flag for different version of hip
4c0f9d2
@rocking5566
remove coerce-illegal, becasue this is not required in FA
89ac30b
@rocking5566
Update ck to fix xcratch memory
a381df5
@rocking5566
Add coerce-illegal in some version
6635e24
@rocking5566
Merge pull request #77 from ROCm/ck_tile/rocm6.2-flag
cc01a17 
Sync Ck tile compile flag with rocm6.2
@rocking5566
Add compile flag for rtn rounding
1cb8f8d
@rocking5566
remove redundant init
ba86d74
@rocking5566
Using env var to switch rounding mode
8c4f9cd
@rocking5566
update ck
ece97c7
@rocking5566
Merge pull request #78 from ROCm/ck_tile/bf16_rtn
b40c1a0 
Change rounding of bf16 to rtn
@joellliu
Copy link

joellliu commented Sep 9, 2024 •
edited
Loading

I have tested this PR on MI300. I saw increased accuracy compared to previous versions of ROCm flash attention. My training of OLMo-1B model with flash attention is finally able to converge!

@rocking5566 rocking5566 changed the title Support kvcache in ROCm Support kvcache in AMD ROCm Sep 9, 2024
@rocking5566 rocking5566 changed the title Support kvcache in AMD ROCm Support page + kvcache in AMD ROCm Sep 9, 2024
@rocking5566 rocking5566 changed the title Support page + kvcache in AMD ROCm Support page kvcache in AMD ROCm Sep 9, 2024
@ehartford
Copy link

This change will greatly help us @tridao
Thank you!

@tridao tridao merged commit e2182cc into Dao-AILab:main Sep 16, 2024
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.
4 participants
@rocking5566 @joellliu @ehartford @tridao