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.

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

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

Merged
Merged
Changes from 1 commit
Commits
Show all changes
26 commits
Select commit Hold shift + click to select a range
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Re-ran formater
jpvillam committed Mar 26, 2024

Unverified

The committer email address is not verified.
commit 734fce7285cff6e665b561096eda9ac366029b82
11 changes: 4 additions & 7 deletions vllm/attention/backends/flash_attn_triton.py
Original file line number Diff line number Diff line change
@@ -8,14 +8,11 @@

import torch

from vllm.attention.backends.abstract import (
AttentionImpl, )
from vllm.attention.backends.flash_attn import (
FlashAttentionBackend,
FlashAttentionMetadata,
)
from vllm.attention.ops.paged_attn import PagedAttention
from vllm.attention.backends.abstract import AttentionImpl
from vllm.attention.backends.flash_attn import (FlashAttentionBackend,
FlashAttentionMetadata)
from vllm.attention.ops.flash_attention_triton import triton_attention
from vllm.attention.ops.paged_attn import PagedAttention


class FlashAttentionTritonBackend(FlashAttentionBackend):
2 changes: 1 addition & 1 deletion vllm/attention/selector.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import os
from functools import lru_cache

import torch
import os

from vllm.attention.backends.abstract import AttentionBackend
from vllm.logger import init_logger