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

Tune paged attention parameters for AMD GPU. #3255

Merged
merged 5 commits into from
Feb 2, 2025
Merged

Tune paged attention parameters for AMD GPU. #3255

merged 5 commits into from
Feb 2, 2025
+13 −2

Conversation

whchung
Copy link
Contributor

@whchung whchung commented Feb 1, 2025 •
edited
Loading

Motivation

Fine tune SGLang page attention kernel performance on AMD MI GPU for LLM.

Modifications

Changes:

  • num_kv_splits : 8 -> 16
  • BLOCK : 64 -> 8
  • num_warps : 2 -> 1 when kv_group_num is more than 1
  • waves_per_cu : 4 -> 1 in grouped paged attention kernel

These knobs have been tested with a couple of workloads on AMD ROCm platform.

@whchung
Tune paged attention parameters for AMD GPU.
c145acb 
Changes:
- num_kv_splits
- BLOCK
- num_warps
@whchung
Additional tuning for grouped page attention kernel.
9c5980e 
Changed:
- waves_per_eu
@zhyncs zhyncs requested a review from HaiShaw February 1, 2025 17:26
whchung and others added 3 commits February 1, 2025 11:30
@whchung
Additional tuning for grouped paged attention kernel
eb34954 
Changed:
- num_stages
@whchung
Segregate AMD-specific tuning with is_hip checks
5c095e5 
Make the tuning on the knobs only applicable on AMD ROCm platform via is_hip checks.
@whchung
Merge branch 'main' into whchung/exp-amd-pa
bf36b5b
@whchung whchung marked this pull request as ready for review February 1, 2025 23:09
@andyluo7 andyluo7 mentioned this pull request Feb 2, 2025
Open
5 tasks
tot0 pushed a commit to tot0/sglang that referenced this pull request Feb 2, 2025
Apply amd paged attention optimization, ref: sgl-project#3255
f2d173a
HaiShaw
HaiShaw reviewed Feb 2, 2025
View reviewed changes
python/sglang/srt/server_args.py
# AMD-specific Triton attention KV splits default number
if is_hip():
self.triton_attention_num_kv_splits = 16

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@whchung we noticed that default 8 works better to most cases? Is 16 slightly better as you see, or better toward long sequences as you interested?

Copy link
Collaborator

@HaiShaw HaiShaw Feb 2, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@this PR boosts grok, etc. as observed.

@HaiShaw HaiShaw merged commit d9eb935 into sgl-project:main Feb 2, 2025
10 checks passed
@HaiShaw HaiShaw added the good first issue Good for newcomers label Feb 2, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@HaiShaw HaiShaw HaiShaw approved these changes

@merrymercy merrymercy Awaiting requested review from merrymercy merrymercy is a code owner

@Ying1123 Ying1123 Awaiting requested review from Ying1123 Ying1123 is a code owner

@zhyncs zhyncs Awaiting requested review from zhyncs zhyncs is a code owner

@ispobock ispobock Awaiting requested review from ispobock ispobock is a code owner

@hnyls2002 hnyls2002 Awaiting requested review from hnyls2002 hnyls2002 is a code owner

@ByronHsu ByronHsu Awaiting requested review from ByronHsu ByronHsu is a code owner

Assignees
No one assigned
Labels
good first issue Good for newcomers
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@whchung @HaiShaw