microsoft / onnxruntime

ONNX Runtime: cross-platform, high performance ML inferencing and training accelerator
https://onnxruntime.ai
MIT License
14.77k stars 2.94k forks source link

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

Open xhcao opened 20 hours ago

xhcao commented 20 hours ago

Description

Motivation and Context

xhcao commented 20 hours ago

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 commented 11 hours ago

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

guschmue commented 11 hours ago

/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

azure-pipelines[bot] commented 11 hours ago
Azure Pipelines successfully started running 1 pipeline(s).
guschmue commented 11 hours ago

/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

azure-pipelines[bot] commented 11 hours ago
Azure Pipelines could not run because the pipeline triggers exclude this branch/path.
guschmue commented 11 hours ago

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

azure-pipelines[bot] commented 11 hours ago
Azure Pipelines could not run because the pipeline triggers exclude this branch/path.
azure-pipelines[bot] commented 11 hours ago
Azure Pipelines successfully started running 1 pipeline(s).
xhcao commented 29 minutes ago

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.