From b1389d25039f83fee08d6e0a4a7944fbcf3c35c0 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Mon, 17 Aug 2026 21:41:00 -0700 Subject: [PATCH] fix(sdpa): move the pre-9.26 Stats packed-BHSD check to post_validate_node MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The check added in 173c431c (#304) ran in validate_sdpa_support_surface(), which is called from pre_validate_node() — before shape inference. Samples and users that leave the Stats output dim/stride unset (to be inferred) were rejected with GRAPH_NOT_SUPPORTED on every cuDNN < 9.26, breaking the cpp_samples 9.19 CI jobs on develop since 2026-08-15. Move the check to post_validate_node(), which runs after infer_properties_node() has filled an unset Stats with packed BHSD; the check still rejects explicitly-set non-BHSD layouts and still surfaces from validate()/build(). Co-Authored-By: Claude Fable 5 --- .../node/scaled_dot_product_flash_attention.h | 19 ++++++++++++++++++ .../node/sdpa_support_surface.h | 20 +++---------------- 2 files changed, 22 insertions(+), 17 deletions(-) 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 7aaf66e7b..75a69c9a2 100644 --- a/include/cudnn_frontend/node/scaled_dot_product_flash_attention.h +++ b/include/cudnn_frontend/node/scaled_dot_product_flash_attention.h @@ -498,6 +498,25 @@ class SDPANodeBase : public NodeCRTP { CUDNN_FE_VALIDATE_STRIDE(output_names::O, attributes.outputs); + // Non-ragged Stats layouts other than packed BHSD are not correctly supported prior to 9.26.0. + // Runs post shape inference so that an unset Stats layout (always inferred as packed BHSD) + // is not rejected. + auto const& stats_out = attributes.outputs.find(output_names::Stats); + bool const has_stats = (stats_out != attributes.outputs.end()) && (stats_out->second != nullptr); + 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."); + } + #undef CUDNN_FE_VALIDATE_STRIDE return {error_code_t::OK, ""}; diff --git a/include/cudnn_frontend/node/sdpa_support_surface.h b/include/cudnn_frontend/node/sdpa_support_surface.h index 61d149003..d7dc5aba9 100644 --- a/include/cudnn_frontend/node/sdpa_support_surface.h +++ b/include/cudnn_frontend/node/sdpa_support_surface.h @@ -49,9 +49,6 @@ 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); @@ -162,20 +159,9 @@ 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."); - } + // The Stats layout check (packed BHSD required prior to 9.26.0) lives in + // SDPANode::post_validate_node(), as it must run after shape inference has + // filled in the dim/stride of an unset Stats output. if (mma_core_mode == DataType_t::FP8_E4M3 || mma_core_mode == DataType_t::FP8_E5M2) { // FP8 specific validation