[ROCm][Hardware][AMD] Use Triton Kernel for default FA on ROCm #3643

jpvillam-amd · 2024-03-26T19:21:34Z

This PR creates and makes default new triton exclusive backend for attention. Additionally removes some unsupported arguments from AMD's version of the flash_attn_varlen_func function.

Changed selector to allow picking between multiple backends rather than just between two.
Added a new triton FA backend available to ROCm.
Added new VLLM_USE_FLASH_ATTN_TRITON option to be able to swap between Triton and Default FA
Removed unsupported attributes from AMD's flash_attn_varlen_func
Added latest build of AMD's triton to Dockerfile.rocm

Appreciate any and all feedback, and apologies for the long PR

Co-authored-by: Vinayak Gokhale [email protected]

BEFORE SUBMITTING, PLEASE READ THE CHECKLIST BELOW AND FILL IN THE DESCRIPTION ABOVE

PR Checklist (Click to Expand)

Thank you for your contribution to vLLM! Before submitting the pull request, please ensure the PR meets the following criteria. This helps vLLM maintain the code quality and improve the efficiency of the review process.

PR Title and Classification

Only specific types of PRs will be reviewed. The PR title is prefixed appropriately to indicate the type of change. Please use one of the following:

[Bugfix] for bug fixes.
[CI/Build] for build or continuous integration improvements.
[Doc] for documentation fixes and improvements.
[Model] for adding a new model or improving an existing model. Model name should appear in the title.
[Frontend] For changes on the vLLM frontend (e.g., OpenAI API server, LLM class, etc.)
[Kernel] for changes affecting CUDA kernels or other compute kernels.
[Core] for changes in the core vLLM logic (e.g., LLMEngine, AsyncLLMEngine, Scheduler, etc.)
[Hardware][Vendor] for hardware-specific changes. Vendor name should appear in the prefix (e.g., [Hardware][AMD]).
[Misc] for PRs that do not fit the above categories. Please use this sparingly.

Note: If the PR spans more than one category, please include all relevant prefixes.

Code Quality

The PR need to meet the following code quality standards:

We adhere to Google Python style guide and Google C++ style guide.
Pass all linter checks. Please use format.sh to format your code.
The code need to be well-documented to ensure future contributors can easily understand the code.
Include sufficient tests to ensure the project to stay correct and robust. This includes both unit tests and integration tests.
Please add documentation to docs/source/ if the PR modifies the user-facing behaviors of vLLM. It helps vLLM user understand and utilize the new features or changes.

Notes for Large Changes

Please keep the changes as concise as possible. For major architectural changes (>500 LOC excluding kernel/data/config/test), we would expect a GitHub issue (RFC) discussing the technical design and justification. Otherwise, we will tag it with rfc-required and might not go through the PR.

What to Expect for the Reviews

The goal of the vLLM team is to be a transparent reviewing machine. We would like to make the review process transparent and efficient and make sure no contributor feel confused or frustrated. However, the vLLM team is small, so we need to prioritize some PRs over others. Here is what you can expect from the review process:

After the PR is submitted, the PR will be assigned to a reviewer. Every reviewer will pick up the PRs based on their expertise and availability.
After the PR is assigned, the reviewer will provide status update every 2-3 days. If the PR is not reviewed within 7 days, please feel free to ping the reviewer or the vLLM team.
After the review, the reviewer will put an action-required label on the PR if there are changes required. The contributor should address the comments and ping the reviewer to re-review the PR.
Please respond to all comments within a reasonable time frame. If a comment isn't clear or you disagree with a suggestion, feel free to ask for clarification or discuss the suggestion.

Thank You

Finally, thank you for taking the time to read these guidelines and for your interest in contributing to vLLM. Your contributions make vLLM a great tool for everyone!

vllm/attention/selector.py

hongxiayang

I have verified the changes on MI250x and MI300x.
lgtm.

…tream_integration

hongxiayang · 2024-03-28T16:16:03Z

Since latest ray has problem with tp>1 when running on AMD system now, can we also do this in Dockerfile.rocm:
RUN python3 -m pip install --no-cache-dir ray[all]==2.9.3

gshtras · 2024-03-28T16:46:21Z

Since latest ray has problem with tp>1 when running on AMD system now, can we also do this in Dockerfile.rocm: RUN python3 -m pip install --no-cache-dir ray[all]==2.9.3

Also in requirements-rocm.txt

WoosukKwon

Can we get rid of the CK-based FlashAttention, as we discussed offline in today's meeting?

vllm/attention/backends/flash_attn_triton.py

hongxiayang · 2024-03-28T21:01:42Z

Can we get rid of the CK-based FlashAttention, as we discussed offline in today's meeting?

We can leave the clean up task in a different pull request, and I will re-open the cleanup-xformer PR (#3558) for that purpose.
There is also concern regarding different data type support.

Now users can use the flag to choose not to install FA during building the docker image using Dockerfile.rocm.

Rename backend

WoosukKwon · 2024-03-31T02:25:34Z

vllm/attention/backends/flash_attn.py

+                if is_hip():
+                    output = flash_attn_varlen_func(
+                        q=query,
+                        k=key,
+                        v=value,
+                        cu_seqlens_q=attn_metadata.seq_start_loc,
+                        cu_seqlens_k=attn_metadata.seq_start_loc,
+                        max_seqlen_q=attn_metadata.max_prompt_len,
+                        max_seqlen_k=attn_metadata.max_prompt_len,
+                        softmax_scale=self.scale,
+                        causal=True,
+                    )
+                else:
+                    output = flash_attn_varlen_func(
+                        q=query,
+                        k=key,
+                        v=value,
+                        cu_seqlens_q=attn_metadata.seq_start_loc,
+                        cu_seqlens_k=attn_metadata.seq_start_loc,
+                        max_seqlen_q=attn_metadata.max_prompt_len,
+                        max_seqlen_k=attn_metadata.max_prompt_len,
+                        softmax_scale=self.scale,
+                        causal=True,
+                        window_size=self.sliding_window,
+                        alibi_slopes=self.alibi_slopes,
+                    )
+


QQ: Actually in #3648, we are planning to use FlashAttention's recent APIs that support attention keys and values stored in paged KV cache. I believe the new APIs are incompatible with AMD GPUs.

Can we just use Triton FlashAttention at all times?

WoosukKwon · 2024-04-01T08:52:37Z

@jpvillam-amd Can we have ROCmAttentionBackend that selects one of the four different implementations of prefill Attention (i.e., Triton, CK, xFormers, naive) while using PagedAttention for decoding?

I feel like mixing the code for NVIDIA GPUs and AMD GPUs for attention is not good at the moment, since the APIs are not unified yet. For example, as I mentioned above, FlashAttentionBackend will use flash-attn's paged KV cache APIs, which are not supported for AMD GPUs at the moment.

…tream_integration

…tion' into jpvillam/triton_upstream_integration

jpvillam-amd · 2024-04-02T15:12:27Z

@jpvillam-amd Can we have ROCmAttentionBackend that selects one of the four different implementations of prefill Attention (i.e., Triton, CK, xFormers, naive) while using PagedAttention for decoding?

I feel like mixing the code for NVIDIA GPUs and AMD GPUs for attention is not good at the moment, since the APIs are not unified yet. For example, as I mentioned above, FlashAttentionBackend will use flash-attn's paged KV cache APIs, which are not supported for AMD GPUs at the moment.

I will rename the backend and have it container all the selector logic for ROCm. I will leave xFormers out of the backend and have it selected with the current selector code same as CUDA

@WoosukKwon Does that sounds ok?

…tream_integration

WoosukKwon · 2024-04-08T18:43:30Z

@jpvillam-amd Thanks for updating the PR. Do you mind if I directly edit this PR? I refactored the PR a bit and wanted to directly upstream it to this PR for faster integration. If you don't mind, could you allow editing the PR? Your branch is protected now.

WoosukKwon

@jpvillam-amd LGTM. Thanks for updating the PR. Could you please submit another PR to test the backend? Thanks.

…project#3643) Co-authored-by: jpvillam <[email protected]> Co-authored-by: Gregory Shtrasberg <[email protected]> Co-authored-by: Woosuk Kwon <[email protected]>

jpvillam added 6 commits March 25, 2024 23:33

Add triton version of FA

fb1d862

Re-Enabled the param checks

948684e

Run yapf and ruffwq

906905f

Ununsed variable

80c7cab

Re-ran formater

734fce7

Make variable code more clear and simplify attn selector

c91e5c3

hongxiayang reviewed Mar 26, 2024

View reviewed changes

vllm/attention/selector.py Outdated Show resolved Hide resolved

WoosukKwon added the rocm label Mar 26, 2024

Logic mistake

534e1f9

hongxiayang approved these changes Mar 28, 2024

View reviewed changes

Merge remote-tracking branch 'upstream/main' into jpvillam/triton_ups…

5c4a993

…tream_integration

hongxiayang mentioned this pull request Mar 28, 2024

[ROCm] [Hardware][AMD] Remove xformer patches and ray issue fix #3558

Closed

WoosukKwon reviewed Mar 28, 2024

View reviewed changes

vllm/attention/backends/flash_attn_triton.py Outdated Show resolved Hide resolved

WoosukKwon added the action-required label Mar 28, 2024

jpvillam added 5 commits March 28, 2024 21:46

Review comments

11437bc

Rename backend

File rename

320d7c4

File rename

b8919d1

Run yapf and ruff

ebba198

Extra bracket in dockerfile

f0d1eeb

jpvillam-amd requested a review from WoosukKwon March 29, 2024 20:34

WoosukKwon self-assigned this Mar 29, 2024

WoosukKwon reviewed Mar 31, 2024

View reviewed changes

jpvillam added 2 commits April 1, 2024 15:06

Merge remote-tracking branch 'upstream/main' into jpvillam/triton_ups…

540c20a

…tream_integration

Merge remote-tracking branch 'origin/jpvillam/triton_upstream_integra…

80c936c

…tion' into jpvillam/triton_upstream_integration

File rename to rocm fa

215d15f

Make a ROCM FA backend

c6cfdff

jpvillam-amd requested a review from WoosukKwon April 2, 2024 22:05

Merge remote-tracking branch 'upstream/main' into jpvillam/triton_ups…

f19a592

…tream_integration

WoosukKwon removed the action-required label Apr 4, 2024

WoosukKwon added 8 commits April 8, 2024 17:21

Merge branch 'main' into jpvillam/triton_upstream_integration

a115f65

Add kv_scale to ROCM flash attn

e691488

Remove is_hip from FlashAttentionBackend

26457fa

Move naive attention to rocm flash attn

03679f4

Refactor selector

3d0db08

yapf

0c9b668

Fix use_naive

c939338

yapf

1238bc1

WoosukKwon mentioned this pull request Apr 8, 2024

[Misc]: Question re: ROCm + Triton Backend #3921

Closed

WoosukKwon approved these changes Apr 9, 2024

View reviewed changes

WoosukKwon merged commit 6c0b045 into vllm-project:main Apr 9, 2024
67 checks passed

jamestwhedbee mentioned this pull request Apr 10, 2024

[Model][AMD] ROCm support for 256 head dims for Gemma #3972

Merged

dtrifiro mentioned this pull request May 15, 2024

bump ubi base image tag opendatahub-io/vllm#24

Merged

bringlein mentioned this pull request May 28, 2024

[RFC]: OpenAI Triton-only backend #5083

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[ROCm][Hardware][AMD] Use Triton Kernel for default FA on ROCm #3643

[ROCm][Hardware][AMD] Use Triton Kernel for default FA on ROCm #3643

jpvillam-amd commented Mar 26, 2024 •

edited

Loading

hongxiayang left a comment

hongxiayang commented Mar 28, 2024

gshtras commented Mar 28, 2024

WoosukKwon left a comment

hongxiayang commented Mar 28, 2024 •

edited

Loading

WoosukKwon Mar 31, 2024 •

edited

Loading

WoosukKwon commented Apr 1, 2024

jpvillam-amd commented Apr 2, 2024

WoosukKwon commented Apr 8, 2024

WoosukKwon left a comment

[ROCm][Hardware][AMD] Use Triton Kernel for default FA on ROCm #3643

[ROCm][Hardware][AMD] Use Triton Kernel for default FA on ROCm #3643

Conversation

jpvillam-amd commented Mar 26, 2024 • edited Loading

PR Title and Classification

Code Quality

Notes for Large Changes

What to Expect for the Reviews

Thank You

hongxiayang left a comment

Choose a reason for hiding this comment

hongxiayang commented Mar 28, 2024

gshtras commented Mar 28, 2024

WoosukKwon left a comment

Choose a reason for hiding this comment

hongxiayang commented Mar 28, 2024 • edited Loading

WoosukKwon Mar 31, 2024 • edited Loading

Choose a reason for hiding this comment

WoosukKwon commented Apr 1, 2024

jpvillam-amd commented Apr 2, 2024

WoosukKwon commented Apr 8, 2024

WoosukKwon left a comment

Choose a reason for hiding this comment

jpvillam-amd commented Mar 26, 2024 •

edited

Loading

hongxiayang commented Mar 28, 2024 •

edited

Loading

WoosukKwon Mar 31, 2024 •

edited

Loading