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

[js/webgpu] support FlashAttention-2 for attention operator #22915

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

xhcao
Copy link
Contributor

@xhcao xhcao commented Nov 21, 2024

Description

Motivation and Context

@xhcao
Copy link
Contributor Author

xhcao commented Nov 21, 2024

The FlashAttention-2 algorithm is based on the paper https://tridao.me/publications/flash2/flash2.pdf.
Memory increase quadratically in the sequence length in original algorithm, and need a lot of global memory read and write accesses. If sequence is large, for example running stable diffusion 2.1 with attention nodes, original algorithm will lead chrome crash for out of memory and inefficient.
Please take a look, @axinging @hujiajie @jiechen0826

@guschmue
Copy link
Contributor

/azp run ONNX Runtime Web CI Pipeline,Windows GPU CI Pipeline,Linux Android Emulator QNN CI Pipeline

@guschmue
Copy link
Contributor

/azp run Linux CPU CI Pipeline,Linux CPU Minimal Build E2E CI Pipeline,Linux GPU CI Pipeline,Linux GPU TensorRT CI Pipeline,Linux OpenVINO CI Pipeline,Linux QNN CI Pipeline,MacOS CI Pipeline,Windows ARM64 QNN CI Pipeline,Windows CPU CI Pipeline

Copy link

Azure Pipelines successfully started running 1 pipeline(s).

@guschmue
Copy link
Contributor

/azp run Windows GPU TensorRT CI Pipeline,onnxruntime-binary-size-checks-ci-pipeline,orttraining-linux-ci-pipeline,orttraining-linux-gpu-ci-pipeline,orttraining-ortmodule-distributed,Windows x64 QNN CI Pipeline,Big Models

Copy link

Azure Pipelines could not run because the pipeline triggers exclude this branch/path.

@guschmue
Copy link
Contributor

/azp run Windows GPU CUDA CI Pipeline,Windows GPU DML CI Pipeline,Windows GPU Doc Gen CI Pipeline

Copy link

Azure Pipelines could not run because the pipeline triggers exclude this branch/path.

Copy link

Azure Pipelines successfully started running 1 pipeline(s).

@guschmue guschmue added the ep:WebGPU ort-web webgpu provider label Nov 21, 2024
@xhcao
Copy link
Contributor Author

xhcao commented Nov 22, 2024

Sorry for the error, I wanted to keep the names as the paper, but the variables‘ names mismatched rules. I had already modified the names. Thanks.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
ep:WebGPU ort-web webgpu provider
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants