[CUDNN][SDPA] Experimental cuDNN Flash Attention v2 Inference #115663

eqy · 2023-12-12T19:29:24Z

Going to clean up some of the checks and will remove draft status after.
Can be tested on SM80+ with TORCH_CUDNN_MHA_ENABLED=1.

CC @drisspg @ptrblck

pytorch-bot · 2023-12-12T19:29:28Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/115663

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 3 Unrelated Failures

As of commit 4fc3337 with merge base 3ab0894 ():

NEW FAILURE - The following job has failed:

periodic / linux-focal-rocm5.7-py3.8 / test (distributed, 2, 2, linux.rocm.gpu) (gh)
distributed/test_dynamo_distributed.py::TestMultiProc::test_ddp_baseline_aot_eager_multiprocess

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

periodic / buck-build-test / buck-build-test (default, 1, 1, ubuntu-latest) (gh)
Process completed with exit code 1.
periodic / linux-focal-cuda11.8-py3.9-gcc9 / test (multigpu, 1, 1, linux.g5.12xlarge.nvidia.gpu) (gh)
distributed/test_c10d_nccl.py::NcclErrorHandlingTest::test_nccl_errors_blocking_nonzero_exit
periodic / win-vs2019-cuda11.8-py3 / test (default, 2, 4, windows.g5.4xlarge.nvidia.gpu) (gh)
test_nestedtensor.py::TestNestedTensorSubclassCUDA::test_sdpa_compile_bfloat16_cuda_bfloat16

This comment was automatically generated by Dr. CI and updates every 15 minutes.

torch/backends/cuda/__init__.py

drisspg · 2023-12-18T22:08:13Z

Also feel free to ping me when you think I should do a review

eqy · 2024-01-04T22:55:07Z

@pytorchmergebot rebase

pytorchmergebot · 2024-01-04T22:56:50Z

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

pytorchmergebot · 2024-01-04T22:56:58Z

Successfully rebased cudnnmha3 onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via git checkout cudnnmha3 && git pull --rebase)

eqy · 2024-01-05T18:06:42Z

@drisspg sorry for the delay, I think this should be ready for review now

drisspg · 2024-01-05T23:50:37Z

aten/src/ATen/cuda/Exceptions.h

@@ -21,6 +21,15 @@ class CuDNNError : public c10::Error {

 }  // namespace c10

+#define AT_CUDNN_FRONTEND_CHECK(EXPR, ...)                                                      \


Is this a result of frontend changing to 9.0?

This is kind of just the "new" way that cuDNN-frontend has been doing error reporting and this is done in recent frontend versions that are supported in 8.9.x. It is a bit annoying in that there is now a mechanism for getting an error string and it is different from the previous convention of "return CUDNN_STATUS_SUCCESS."

aten/src/ATen/native/cudnn/MHA.cpp

drisspg · 2024-01-06T00:59:05Z

aten/src/ATen/native/cudnn/MHA.cpp

+    //                         .set_dim({b, 1, s_q, s_kv})
+    //                         .set_stride({s_q * s_kv, s_q * s_kv, s_kv, 1}));
+    auto seed = mha_graph->tensor(fe::graph::Tensor_attributes()
+                                     .set_name("Seed")


Another dumb Q: are these names meaningful? like do they need to match an api for calling the function?

Not a dumb question! The names are potentially meaningful in the future in that we would be able to reference I/O tensors by name rather than by holding onto an explicit reference as we currently do to supply the "variant pack" (the map used to specify the data pointers for each I/O tensor upon invocation). In either case they do not really mean anything to cuDNN and rather are for the user's (us) housekeeping.

drisspg · 2024-01-06T00:59:48Z

aten/src/ATen/native/cudnn/MHA.cpp

+    if (cudnnGetVersion() >= 8904) {
+        //scaled_dot_product_flash_attention_options.set_alibi_mask(true);
+    }
+


Does this only support certain bias types?

I would need check with the cuDNN folks, this doc https://docs.nvidia.com/deeplearning/cudnn/developer-guide/index.html#flash-fused-multi-head-att-fprop doesn't seem to be very clear about this

drisspg · 2024-01-06T01:00:57Z

aten/src/ATen/native/cudnn/MHA.cpp

+    cudnnHandle_t handle = getCudnnHandle();
+    o = at::empty_strided({b, h, s_q, d}, {s_q * h * d, d, h * d, 1}, q.options());
+    if (return_softmaxstats) {
+      // TODO(eqy): fix strides


I assume this stride fix is only needed for backward support right? so is okay for this PR?

I think this comment is out-of-date, but I should verify rather than fix the strides here

aten/src/ATen/native/cudnn/MHA.cpp

aten/src/ATen/native/native_functions.yaml

aten/src/ATen/native/transformers/attention.cpp

drisspg · 2024-01-09T18:32:52Z

One other random question, if I wanted to play around with this early, is there any specific build options I should be using? And any specific version of cudnn I need to have?

eqy · 2024-01-09T21:51:07Z

One other random question, if I wanted to play around with this early, is there any specific build options I should be using? And any specific version of cudnn I need to have?

I don't think you would need any specific build options---cuDNN >= 8.9 would probably work best. The only other requirement should be cuDNN frontend 1.0, but that tag was updated in the cudnn-frontend submodule a while ago.

eqy · 2024-01-10T02:31:05Z

@pytorchmergebot rebase

pytorchmergebot · 2024-02-11T08:28:33Z

Successfully rebased cudnnmha3 onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via git checkout cudnnmha3 && git pull --rebase)

facebook-github-bot · 2024-02-13T18:26:46Z

@drisspg has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

izaitsevfb · 2024-02-14T21:59:36Z

@pytorchbot merge -i

Landed internally as D53716382

pytorchmergebot · 2024-02-14T22:01:33Z

Merge started

Your change will be merged while ignoring the following 4 checks: periodic / buck-build-test / buck-build-test (default, 1, 1, ubuntu-latest), periodic / win-vs2019-cuda11.8-py3 / test (default, 2, 4, windows.g5.4xlarge.nvidia.gpu), periodic / linux-focal-rocm5.7-py3.8 / test (distributed, 2, 2, linux.rocm.gpu), periodic / linux-focal-cuda11.8-py3.9-gcc9 / test (multigpu, 1, 1, linux.g5.12xlarge.nvidia.gpu)

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

Skylion007 · 2024-02-26T20:42:19Z

We should probably update our cudnn binaries as most of the recent changes has been improving this flash attention kernal.

pytorchbot added the open source label Dec 12, 2023

eqy added the topic: not user facing topic category label Dec 12, 2023

eqy marked this pull request as ready for review December 13, 2023 02:16

albanD requested a review from drisspg December 13, 2023 16:28

albanD added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Dec 13, 2023

drisspg reviewed Dec 13, 2023

View reviewed changes

torch/backends/cuda/__init__.py Show resolved Hide resolved

eqy added ciflow/trunk Trigger trunk jobs on your pull request ciflow/periodic Trigger jobs ran periodically on master (periodic.yml) on the PR labels Dec 13, 2023

pytorchmergebot force-pushed the cudnnmha3 branch from fc7cf93 to 7a1e8d9 Compare January 4, 2024 22:57

drisspg reviewed Jan 5, 2024

View reviewed changes