[DO NOT MERGE] Draft of combined cudnn + hipdnn SDPA backend

[rkayaith]

Primary changes are here: 9bb8752

This adds extra include paths to cmake, where shim headers are placed. The shim headers re-route CUDA includes to the HIP equivalent, with the necessary defines etc. to preserve compatibility. Only shims necessary for SDPA were added.

Key files to look at:

• MHA.cpp - the unified backend (look for USE_HIPDNN)
• sdp_utils.cpp - backend selection logic (look for USE_HIPDNN)
• cudnn_frontend.h - graph API shim

The key differences ended up being:

• hipDNN requires scale to be set on the graph as a constant, while cuDNN binds it through a host buffer at execute time. This requires an extra field in the cache key for the scale value + the plumbing to populate it.

  • cuDNN binds via SDPA_attributes::set_attn_scale(attn_scale) + {SCALE, &scaling_factor} in the variant_pack: MHA.cpp:551.
  • hipDNN bakes via SDPA_attributes::set_attn_scale_value(scaling_factor): MHA.cpp:540.
  • extra cache key field: MHA.cpp:239 (MHAParams::scaling_factor under #ifdef USE_HIPDNN).
  • Population of params.scaling_factor: MHA.cpp:291.

• Code dependent on cuDNN frontend version needs hipDNN handled too:

  • SDPA_attributes::set_generate_stats (in hipdnn and newer cudnn) vs SDPA_attributes::set_is_inference (older cudnn): MHA.cpp:530-534.
  • Graph::get_workspace_size: out-param form (hipdnn / newer cudnn) vs by-value (older cudnn): MHA.cpp:1570-1575

• Ragged/nested tensor APIs don't exist in hipDNN yet, calls have to be ifdefd out:

  • SDPA_attributes::set_seq_len_{q,kv} (and the same on SDPA_backward_attributes)
  • SDPA_attributes::set_padding_mask (and the same on SDPA_backward_attributes)
  • Tensor_attributes::set_ragged_offset

• cuDNN constraints are slightly different. cuDNN checks kernel support in PyTorch rather than through API queries.

  • Increases coupling but allows requirements to be checked symbolically without concrete dimension values.
  • The hipDNN side adds at::native::check_cudnn_sdpa_support at the end of the constraint chain: sdp_utils.cpp:833. It builds the graph and runs Graph::is_supported_ext. On cuDNN it short-circuits to return true (MHA.cpp:2051-2053).
  • hipDNN-only constraint helpers in sdp_utils.cpp: check_hipdnn_enabled, check_no_nested_inputs_hipdnn, check_dtypes_hipdnn.

• hipDNN has slight differences in a few graph APIs that the shim handles. All in cudnn_frontend.h:

  • cuDNN's Graph::check_support(handle) and Graph::build_plans(handle) take a handle; hipDNN's don't. Shim adds handle-taking overloads: cudnn_frontend.h:34-35.
  • cuDNN's Graph::query_tensor_attributes_of_uid(uid, attrs) is a per-uid out-param query; hipDNN only offers a one-shot Graph::getTensorsByUid() map. Shim wraps it: cudnn_frontend.h:39-50.
  • cuDNN's HeurMode_t::A (recommended heuristic) maps to hipDNN's FALLBACK. Shim shadows the enum with a struct so call sites can use fe::HeurMode_t::A uniformly: cudnn_frontend.h:55-58.

Besides the conditional logic, there's some code that's technically there on both backends but could be simplified if targeting only one:

• hipDNN doesn't support nested tensors. A bunch of nested-tensor code is left in but is effectively dead on hipDNN and could just be deleted. Functions whose bodies are entirely cuDNN-only (entry-point bails with TORCH_CHECK(false, ...) on hipDNN):

  • build_graph_nestedtensor
  • build_graph_backward_nestedtensor
  • run_cudnn_SDP_fprop_nestedtensor
  • run_cudnn_SDP_bprop_nestedtensor

  Same goes for the ragged-in-dense path. use_ragged is always false on hipDNN, so the supporting logic and cache-key field are dead too:

  • use_ragged_in_dense
  • MHAParams::use_ragged
  • build_graph:553-575, build_graph:630-701
  • build_graph_backward:981-1003, build_graph_backward:1058-1140
  • run_cudnn_SDP_fprop:1438-1466, run_cudnn_SDP_fprop:1559-1569
  • run_cudnn_SDP_bprop:1771-1786, run_cudnn_SDP_bprop:1869-1877

• hipDNN wants to query graph support before actual execution, which complicates logic involving tensors that haven't been allocated yet.

  • On cuDNN, graph building happens at execution time so all output tensors have been allocated and their size/stride info can be queried directly. To support hipDNN, some assumptions about the allocation are baked into the build code, and those assumptions need to be kept in sync with the call sites that actually do the allocation.
  • The assumptions correspond to the arguments dropped from build_graph and build_graph_backward (originally o, softmaxstats, dropoutseed, dropoutoffset for the forward; plus dO, dQ, dK, dV for the backward). Each removed tensor's dim/stride is now derived from b/h/s_q/s_kv/d_qk/d_v + the dtype info inside the build functions, instead of read directly off a tensor:
    • O: shape {b, h, s_q, d_v}, strides via compute_matching_strides against Q's layout: MHA.cpp:619-622 (forward), MHA.cpp:1053-1057 (backward, also reused for dO).
    • softmaxstats / Stats: contiguous {b, h, s_q, 1} (matches at::empty allocation): MHA.cpp:623-628.
    • dropoutseed / dropoutoffset: hardcoded as {1, 1, 1, 1} with dtype INT64 based on the documented allocation contract from attention.cpp / attention.cu: MHA.cpp:577-590 (forward), MHA.cpp:1022-1033 (backward).
    • dQ / dK / dV: reuse q.sizes() / q.strides() etc. (the gradient tensors are at::empty_like(q) so their layout matches): MHA.cpp:1145-1147.

[rocm-repo-management-api]

Jenkins build for 77cc7f93bb76bf91e080646aa085272c7516ae69 commit finished as FAILURE
Links: Pipeline Overview / Build artifacts / Test Results