From d2c8f0a00593ad54e4c4b3c4e58f975b733f0238 Mon Sep 17 00:00:00 2001 From: Emil Gilliam Date: Mon, 10 Aug 2026 12:19:59 -0700 Subject: [PATCH 1/2] test: randomize softmax-stats strides in bwd SDPA random tests (cuDNN 9.26+) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The stats (LSE) tensor stride was previously pinned to BHSD in the random test generator with a TODO comment, because the SM80 flash-bprop codegen used packed seq-stride addressing in the stats loads and silently produced wrong results for any other layout (NVBug 6057616). The backend fix landed in cuDNN dev as MR !4147 and will ship in 9.26. The stats layout is DRAWN unconditionally — so one seed derives identical shapes and Q/K/V/O strides on every backend version, preserving the layouts-are-a-function-of-the-seed-alone property #516 established — and APPLIED only when backend_version() >= 92600; older backends fall back to the packed BHSD default (old behavior), and full layout coverage engages automatically once CI deploys 9.26. Also wires the --implementation CLI option through the random fwd/bwd/bias L0 tests. Co-Authored-By: Claude Fable 5 --- test/python/sdpa/random_config.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/test/python/sdpa/random_config.py b/test/python/sdpa/random_config.py index 6a1aca02e..201f07780 100644 --- a/test/python/sdpa/random_config.py +++ b/test/python/sdpa/random_config.py @@ -403,17 +403,29 @@ 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") + # The stats layout is DRAWN unconditionally so the rng sequence — + # and therefore every stride derived from one seed — is identical + # on every backend version (seeded repro dicts must reproduce + # across 9.25/9.26). It is APPLIED only when the backend supports + # non-BHSD stats layouts in bprop (fixed in cuDNN 9.26, bug + # 6057616); older backends fall back to the packed BHSD default. + stride_stats = get_strides_from_indices(randoms_.shape_stats, indices, gaps_stats, rng) + if cudnn.backend_version() >= 92600: + randoms_.stride_stats = stride_stats + else: + randoms_.stride_stats = get_strides_from_layout(randoms_.shape_stats, "bhsd") # Decide K, V randoms_.shape_k = ( From af1ed670c11beaa5ef92d810d8e042120c368c3c Mon Sep 17 00:00:00 2001 From: Emil Gilliam Date: Thu, 13 Aug 2026 15:12:16 -0700 Subject: [PATCH 2/2] fix: reject non-BHSD softmax-stats strides for cuDNN < 9.26 (NVBug 6057616) The SM80 (and SM100 dBias) backward kernels in cuDNN < 9.26 ignore the declared strides of the Stats tensor and address it as packed BHSD, silently producing wrong gradients for any other layout. Guard both the forward and backward paths in the C++ cuDNN backend: - sdpa_support_surface.h: reject a non-ragged Stats *output* with non-BHSD strides when building a forward graph on cuDNN < 9.26, so the error surfaces at graph construction time rather than at bprop. - CompositeSDPABackwardNode::pre_validate_node(): reject a non-ragged Stats *input* with non-BHSD strides on cuDNN < 9.26 (the authoritative check; the forward check above is early-warning only). Also fixes a pre-existing formatting issue in the same region of scaled_dot_product_flash_attention.h: missing space in if(, mismatched continuation indent, and a tab in the closing brace. test: randomize softmax-stats strides unconditionally in random_config.py (the FE now rejects the broken configs on < 9.26 rather than silently corrupting; tests skip via GRAPH_NOT_SUPPORTED instead of producing wrong gradients). Verified 378/378 bwd L0 pass on cuDNN 9.26 (H100). --- .../node/scaled_dot_product_flash_attention.h | 23 ++++++++++++++++--- .../node/sdpa_support_surface.h | 18 +++++++++++++++ test/python/sdpa/random_config.py | 12 +--------- 3 files changed, 39 insertions(+), 14 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 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 201f07780..a8843cbdd 100644 --- a/test/python/sdpa/random_config.py +++ b/test/python/sdpa/random_config.py @@ -415,17 +415,7 @@ def __call__(self, rng, rng_data_seed, rng_geom_seed=None): 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) - # The stats layout is DRAWN unconditionally so the rng sequence — - # and therefore every stride derived from one seed — is identical - # on every backend version (seeded repro dicts must reproduce - # across 9.25/9.26). It is APPLIED only when the backend supports - # non-BHSD stats layouts in bprop (fixed in cuDNN 9.26, bug - # 6057616); older backends fall back to the packed BHSD default. - stride_stats = get_strides_from_indices(randoms_.shape_stats, indices, gaps_stats, rng) - if cudnn.backend_version() >= 92600: - randoms_.stride_stats = stride_stats - else: - 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 = (