[cuDNN][SDPA] cuDNN SDPA refactor/cleanup, nested tensor backward, test priority bump for sm90, sm100
#149282
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/149282
Note: Links to docs will display an error until the docs builds have been completed. ❌ 2 New Failures, 1 Unrelated FailureAs of commit f1198b0 with merge base 3cd6935 ( NEW FAILURES - The following jobs have failed:
UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
@pytorchmergebot rebase |
|
@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here |
|
@pytorchmergebot rebase |
|
@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here |
|
Successfully rebased |
ac66884 to
f7c76b8
Compare
| } | ||
| auto workspace_size = mha_graph->get_workspace_size(); | ||
| auto workspace_ptr = | ||
| c10::cuda::CUDACachingAllocator::get()->allocate(workspace_size); | ||
| TORCH_CHECK( | ||
| mha_graph->execute(handle, variant_pack, workspace_ptr.get()).is_good()); | ||
| mhagraphcache.update(key, graph_and_tensors_values); | ||
| mhagraphcache.update(key, mha_graph); |
There was a problem hiding this comment.
The update method for mhagraphcache should probably use perfect forward up where the update method is defined instead of an lref. And throughout the file should be to remove extra copies.
| mhagraphcache.update(key, mha_graph); | |
| mhagraphcache.update(key, std::move(mha_graph)); |
sm90, sm100
| @@ -563,7 +581,7 @@ bool check_for_nested_inputs(sdp_params const& params, bool debug) { | |||
|
|
|||
| const auto dprop = at::cuda::getCurrentDeviceProperties(); | |||
| // Check that the input is nested | |||
| if (dprop->major != 9 && has_for_nested_inputs(params)) { | |||
| if (dprop->major < 9 && has_for_nested_inputs(params)) { | |||
There was a problem hiding this comment.
should we guard consumer GPUS? I guess thats handled in the dispatch
| @@ -57,6 +57,8 @@ | |||
| namespace sdp { | |||
| namespace { | |||
|
|
|||
| bool priority_order_init_ = false; | |||
| mhagraphbackwardcache; | ||
|
|
||
| namespace { | ||
|
|
||
| enum UIDS { |
There was a problem hiding this comment.
This was the API we had in the convolution frontend API days, which was not supported when SDPA was first added as a cuDNN composite op
This removal made the use of cuDNN graphs after they had been built very clunky, as we had to explicitly stash each cuDNN frontend tensor (a shared ptr) in order to associate with a data ptr in order to execute the graph. In particular it meant that the function to build a graph had a really awkward return type of a gigantic tuple of "graph and tensors" as we could not drop the frontend object for each tensor.
In reality we know all of the tensors that are needed to execute the graph by unique names, such as 'q', 'k', 'v', 'dO', 'dQ'... etc., so there's no need to hold on to unique objects. The UID is just a lightweight way to encode these names if we give each tensor a UID (currently just an int64_t) and we can just pass a mapping of UID -> data ptr to execute the graph instead of a clunky shared_ptr (to frontend object) -> data_ptr.
| auto q_strides = q.strides(); | ||
| auto k_strides = k.strides(); | ||
| auto v_strides = v.strides(); | ||
| constexpr int strideidx0 = 1; |
There was a problem hiding this comment.
maybe do a NB: sdpa api is transposed vs cudnn
|
SDPA perf. should be the same, no nested tensor benchmarks yet as that doesn't have caching yet 😅 |
cleanup tuple/tensor boilerplate in cuDNN SDPA, preparation for nested/ragged tensor backward