fix: reject non-BHSD softmax-stats strides on cuDNN < 9.26; randomize in tests (NVBug 6057616)

include/cudnn_frontend/node/scaled_dot_product_flash_attention.h

@@ -1406,15 +1406,32 @@ class CompositeSDPABackwardNode : public NodeCRTP<CompositeSDPABackwardNode> {
             is_deterministic_algorithm_supported_on_blackwell = true;
         }
 
-        if(detail::get_backend_version() >= 91801) {
+        if (detail::get_backend_version() >= 91801) {
             RETURN_CUDNN_FRONTEND_ERROR_IF(is_ragged && (8 == prop_major || 12 == prop_major) && attributes.is_deterministic_algorithm,
                                         error_code_t::GRAPH_NOT_SUPPORTED,
                                         "Deterministic algorithm is not supported for bprop thd on SM8X and SM12X GPUs");
 
-	    RETURN_CUDNN_FRONTEND_ERROR_IF(is_ragged && (8 == prop_major || 12 == prop_major) && attributes.inputs[input_names::Stats]->get_ragged_offset(),
+            RETURN_CUDNN_FRONTEND_ERROR_IF(is_ragged && (8 == prop_major || 12 == prop_major) && attributes.inputs[input_names::Stats]->get_ragged_offset(),
                                         error_code_t::GRAPH_NOT_SUPPORTED,
                                         "Packed/ragged LSE is not supported for bprop thd on SM8X and SM12X GPUs");
-	}
+        }
+
+        // Non-ragged layouts other than BHSD are not correctly supported prior to 9.26.0.
+        // TODO: move to sdpa_support_surface.h (where the forward twin of this check lives)
+        // once the backward path grows a SDPA_backward_attributes support surface there —
+        // today that file serves only the forward attributes.
+        if (detail::get_backend_version() < 92600 && !attributes.inputs.at(input_names::Stats)->get_ragged_offset()) {
+            auto const& stats_dim    = attributes.inputs.at(input_names::Stats)->get_dim();
+            auto const& stats_stride = attributes.inputs.at(input_names::Stats)->get_stride();
+            bool const stats_is_packed_bhsd = stats_stride[3] == 1 &&
+                                              stats_stride[2] == stats_dim[3] &&
+                                              stats_stride[1] == stats_dim[2] * stats_dim[3] &&
+                                              stats_stride[0] == stats_dim[1] * stats_dim[2] * stats_dim[3];
+            RETURN_CUDNN_FRONTEND_ERROR_IF(!stats_is_packed_bhsd,
+                                        error_code_t::GRAPH_NOT_SUPPORTED,
+                                        "For cuDNN version below 9.26.0, a non-ragged Stats input of sdpa_backward must be "
+                                        "a packed BHSD tensor.");
+        }
 
         // version specific validation
         RETURN_CUDNN_FRONTEND_ERROR_IF(detail::get_backend_version() < 90500 && is_dbias && attributes.padding_mask,

include/cudnn_frontend/node/sdpa_support_surface.h

@@ -49,6 +49,9 @@ SDPA_attributes::validate_sdpa_support_surface(const detail::Context& context,
     auto const& rng_tensor = outputs.find(SDPA_attributes::output_names::RNG_DUMP);
     bool const is_rng      = (rng_tensor != outputs.end() && rng_tensor->second != nullptr);
 
+    auto const& stats_out = outputs.find(SDPA_attributes::output_names::Stats);
+    bool const has_stats  = (stats_out != outputs.end()) && (stats_out->second != nullptr);
+
     bool const max_seq_kv_explicit = max_seq_len_kv.has_value();
 
     auto const& attn_scale    = inputs.find(SDPA_attributes::input_names::Attn_scale);
@@ -159,6 +162,21 @@ SDPA_attributes::validate_sdpa_support_surface(const detail::Context& context,
                                    error_code_t::ATTRIBUTE_NOT_SET,
                                    "Intermediate tensor data type needs to be set as internal tensors require it.");
 
+    // Non-ragged layouts other than BHSD are not correctly supported prior to 9.26.0.
+    if (has_stats && !stats_out->second->get_ragged_offset() && detail::get_backend_version() < 92600) {
+        auto const& stats_dim           = stats_out->second->get_dim();
+        auto const& stats_stride        = stats_out->second->get_stride();
+        bool const stats_is_packed_bhsd = stats_dim.size() == 4 && stats_stride.size() == 4 && stats_stride[3] == 1 &&
+                                          stats_stride[2] == stats_dim[3] &&
+                                          stats_stride[1] == stats_dim[2] * stats_dim[3] &&
+                                          stats_stride[0] == stats_dim[1] * stats_dim[2] * stats_dim[3];
+        RETURN_CUDNN_FRONTEND_ERROR_IF(
+            !stats_is_packed_bhsd,
+            error_code_t::GRAPH_NOT_SUPPORTED,
+            "For cuDNN version below 9.26.0, a non-ragged Stats output must be a packed BHSD "
+            "tensor.");
+    }
+
     if (mma_core_mode == DataType_t::FP8_E4M3 || mma_core_mode == DataType_t::FP8_E5M2) {
         // FP8 specific validation
 

test/python/sdpa/random_config.py

@@ -403,17 +403,19 @@ def __call__(self, rng, rng_data_seed, rng_geom_seed=None):
             indices.append(3)
             gaps_q = [0, 0, 0, 0]
             gaps_o = [0, 0, 0, 0]
+            gaps_stats = [0, 0, 0, 0]
 
             if rng.randint(0, 1) == 0:  # 50% chance of gaps
                 gaps_q = [rng.randint(0, 8) for _ in range(3)]
                 gaps_o = [rng.randint(0, 8) for _ in range(3)]
+                gaps_stats = [rng.randint(0, 8) for _ in range(3)]
                 gaps_q.append(elem_align * rng.randint(0, 2))
                 gaps_o.append(elem_align * rng.randint(0, 2))
+                gaps_stats.append(elem_align * rng.randint(0, 2))
 
             randoms_.stride_q = get_strides_from_indices(randoms_.shape_q, indices, gaps_q, rng)
             randoms_.stride_o = get_strides_from_indices(randoms_.shape_o, indices, gaps_o, rng)
-            # TODO: Randomize stride_stats once all layouts are supported correctly.
-            randoms_.stride_stats = get_strides_from_layout(randoms_.shape_stats, "bhsd")
+            randoms_.stride_stats = get_strides_from_indices(randoms_.shape_stats, indices, gaps_stats, rng)
 
         # Decide K, V
         randoms_.shape_k = (