diff --git a/include/cudnn_frontend/node/scaled_dot_product_flash_attention.h b/include/cudnn_frontend/node/scaled_dot_product_flash_attention.h index 7b2b194e3..7aaf66e7b 100644 --- a/include/cudnn_frontend/node/scaled_dot_product_flash_attention.h +++ b/include/cudnn_frontend/node/scaled_dot_product_flash_attention.h @@ -1406,15 +1406,32 @@ class CompositeSDPABackwardNode : public NodeCRTP { 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, diff --git a/include/cudnn_frontend/node/sdpa_support_surface.h b/include/cudnn_frontend/node/sdpa_support_surface.h index 8a7106e48..61d149003 100644 --- a/include/cudnn_frontend/node/sdpa_support_surface.h +++ b/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 diff --git a/test/python/sdpa/random_config.py b/test/python/sdpa/random_config.py index 6a1aca02e..a8843cbdd 100644 --- a/test/python/sdpa/random_config.py +++ b/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 = (