Skip to content

[Feature] Faster Custom Paged Attention kernels #385

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
Closed
wants to merge 18 commits into from
Closed

[Feature] Faster Custom Paged Attention kernels #385

wants to merge 18 commits into from

Conversation

tjtanaa
Copy link

@tjtanaa tjtanaa commented Jan 24, 2025 •
edited by github-actions bot
Loading

Description

This PR implements a faster Custom Paged Attention (CPA) kernel based on mfma16x16x16 instructions.
This feature is from ROCm/vllm (#372).

maleksan85 and others added 18 commits November 13, 2024 10:22
@maleksan85
corrected types for strides in triton FA (ROCm#274) (ROCm#276)
efb0432 
Co-authored-by: Aleksandr Malyshev <[email protected]>
(cherry picked from commit 9a46e97)
@dhonnappa-amd
Update test-template.j2 (ROCm#283)
d291770 
Adding build only k8s node and queue names update
@hyoon1
Fix max_seqlens_q/k initialization for Navi GPUs (ROCm#310)
679a15c 
- max_seqlens_q/k variables were not correctly initialized for Navi GPUs leading to incorrect outputs.
- ensure that the correct values are passed to the attn_fwd kernel based on the GPU type.
@gshtras
Merge remote-tracking branch 'origin/develop'
fb82bf1
@gshtras
Setting the value for the scpecilative decoding worker class on rocm …
22f9066 
…platform (ROCm#313)

Signed-off-by: Gregory Shtrasberg <[email protected]>
@sanyalington
fp8 cpa
f57aa62
@sanyalington
load nt optional for kv
4145bae
@sanyalington
enable alibi; fix gfx90a compile
7bb4a48
@sanyalington
checkpoint with head size 64 supported
bca5077
@sanyalington
block size 32 fix
e2e7d19
@vllmellm
Merge remote-tracking branch 'origin/shsanyal_develop_cpa_fp8' into p…
3c5b26e 
…g_attn_to_llama_fp8
@sanyalington
clean up
2a669e4
@sanyalington
further clean up and comments
51623be
@vllmellm
kernel bug fixes and code cleaning.
e52eb1f 
adjusted attention kernel unit test.
@vllmellm
Merge remote-tracking branch 'origin/shsanyal_develop_cpa_fp8' into p…
3a6f752 
…g_attn_to_llama_fp8
@vllmellm
fix unit test for rocm custom attention kernel
7267062
@vllmellm
fix benchmark paged attention
b8e66a9
@vllmellm
[Bugfix]: fix v1/v2 paged attention kernel unit test.
0f6ff75
This was referenced Jan 24, 2025
Closed
Closed
@tjtanaa tjtanaa closed this Feb 10, 2025
@tjtanaa tjtanaa deleted the pg_attn_to_llama_fp8 branch February 10, 2025 06:51
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@Alexei-V-Ivanov-AMD Alexei-V-Ivanov-AMD Awaiting requested review from Alexei-V-Ivanov-AMD Alexei-V-Ivanov-AMD will be requested when the pull request is marked ready for review Alexei-V-Ivanov-AMD is a code owner

@shajrawi shajrawi Awaiting requested review from shajrawi shajrawi will be requested when the pull request is marked ready for review shajrawi is a code owner

@gshtras gshtras Awaiting requested review from gshtras gshtras will be requested when the pull request is marked ready for review gshtras is a code owner

@maleksan85 maleksan85 Awaiting requested review from maleksan85 maleksan85 will be requested when the pull request is marked ready for review maleksan85 is a code owner

@sunway513 sunway513 Awaiting requested review from sunway513 sunway513 will be requested when the pull request is marked ready for review sunway513 is a code owner

@charlifu charlifu Awaiting requested review from charlifu charlifu will be requested when the pull request is marked ready for review charlifu is a code owner

@mawong-amd mawong-amd Awaiting requested review from mawong-amd mawong-amd will be requested when the pull request is marked ready for review mawong-amd is a code owner

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
7 participants
@tjtanaa @maleksan85 @dhonnappa-amd @hyoon1 @gshtras @sanyalington @vllmellm