From 143eefb78197175e129d95a0c49e6b928f30ecf3 Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Wed, 29 Jul 2026 14:36:59 -0700 Subject: [PATCH 1/2] Support dense THD stats graphs on SM8x cuDNN accepts THD on SM8x from 9.18.1 but requires dense Stats/LSE and max-sequence graph dimensions there. Use one representation predicate for graph construction, auxiliary allocation, and cache-key normalization so support probing and execution cannot diverge. Broaden the shared Python THD layout filter to mixed Q/KV layouts and add focused predicate, padding-filter, and SM8x support-boundary coverage. Signed-off-by: Sudhakar Singh --- tests/cpp/util/CMakeLists.txt | 1 + tests/cpp/util/test_fused_attn_config.cpp | 25 +++++++ tests/pytorch/attention/test_attention.py | 74 +++++++++++++++++++ .../common/fused_attn/config_and_params.cpp | 2 +- .../common/fused_attn/config_and_params.h | 6 ++ .../fused_attn_f16_arbitrary_seqlen.cu | 34 ++++----- .../attention/dot_product_attention/utils.py | 6 +- 7 files changed, 127 insertions(+), 21 deletions(-) create mode 100644 tests/cpp/util/test_fused_attn_config.cpp diff --git a/tests/cpp/util/CMakeLists.txt b/tests/cpp/util/CMakeLists.txt index 1dfd2fed4e..fb326a6292 100644 --- a/tests/cpp/util/CMakeLists.txt +++ b/tests/cpp/util/CMakeLists.txt @@ -3,6 +3,7 @@ # See LICENSE for license information. add_executable(test_util + test_fused_attn_config.cpp test_nvrtc.cpp test_string.cpp ../test_common.cu) diff --git a/tests/cpp/util/test_fused_attn_config.cpp b/tests/cpp/util/test_fused_attn_config.cpp new file mode 100644 index 0000000000..6bf60bd403 --- /dev/null +++ b/tests/cpp/util/test_fused_attn_config.cpp @@ -0,0 +1,25 @@ +/************************************************************************* + * Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include + +#include "common/fused_attn/config_and_params.h" + +namespace transformer_engine::fused_attn { + +TEST(FusedAttnConfig, PackedRaggedGraphSupport) { + EXPECT_FALSE(supports_packed_ragged_graph(90500, 90)); + EXPECT_TRUE(supports_packed_ragged_graph(90600, 90)); + EXPECT_TRUE(supports_packed_ragged_graph(91801, 100)); + + // SM8x and SM120 require dense Stats/LSE and max-sequence graph dimensions, + // even when the cuDNN runtime supports THD inputs on those architectures. + EXPECT_FALSE(supports_packed_ragged_graph(91801, 80)); + EXPECT_FALSE(supports_packed_ragged_graph(91801, 89)); + EXPECT_FALSE(supports_packed_ragged_graph(91801, 120)); +} + +} // namespace transformer_engine::fused_attn diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 6c876ad7cc..aaca1bf166 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -936,6 +936,23 @@ def test_dpa_qkv_layout(dtype, model_configs, model, qkv_layout): test_dot_product_attention(dtype, model_configs, model, False, qkv_layout, False, False) +def test_mixed_thd_pad_between_seqs_filter(): + """Mixed THD layouts apply the same backend-independent padding filter as pure THD.""" + config = ModelConfig(2, 128, 16, 64, attn_mask_type="padding") + available_backends, _, _ = get_available_attention_backends( + config, + qkv_dtype=torch.bfloat16, + qkv_layout="thd_bshd_bshd", + pad_between_seqs=True, + is_training=False, + skip_fused_attn=True, + ) + assert not available_backends[2], ( + "UnfusedDotProductAttention must be disabled when either Q or KV uses THD " + "with padding between sequences." + ) + + qkv_layouts_packed = [l for l in qkv_layouts if any(c.isdigit() for c in l)] @@ -1012,6 +1029,63 @@ def test_dpa_qkv_layout_declarative(dtype, model_configs, model, qkv_layout): } +model_configs_thd_sm8x = { + "thd_sm8x": ModelConfig(2, 128, 16, 64, attn_mask_type="padding"), +} + + +@pytest.mark.skipif( + device_compute_capability[0] != 8 or get_cudnn_version() >= (9, 18, 1), + reason="Unsupported THD boundary requires SM8x with cuDNN older than 9.18.1.", +) +def test_dpa_thd_sm8x_unsupported_cudnn(): + """The graph probe rejects SM8x THD training before cuDNN 9.18.1.""" + config = model_configs_thd_sm8x["thd_sm8x"] + available_backends, _, _ = get_available_attention_backends( + config, + qkv_dtype=torch.bfloat16, + qkv_layout="thd_thd_thd", + pad_between_seqs=False, + is_training=True, + deterministic=_deterministic, + ) + assert not available_backends[ + 1 + ], "FusedAttention must reject SM8x THD training when cuDNN is older than 9.18.1." + + +@pytest.mark.skipif( + device_compute_capability[0] != 8, + reason="Dense THD Stats/LSE regression is specific to SM8x.", +) +@pytest.mark.skipif( + get_cudnn_version() < (9, 18, 1), + reason="THD training on SM8x requires cuDNN 9.18.1+.", +) +@pytest.mark.parametrize("dtype", param_types_lean) +def test_dpa_thd_sm8x_dense_stats(dtype): + """SM8x THD training uses dense Stats/LSE and executes both forward and backward.""" + config = model_configs_thd_sm8x["thd_sm8x"] + available_backends, _, _ = get_available_attention_backends( + config, + qkv_dtype=dtype, + qkv_layout="thd_thd_thd", + pad_between_seqs=False, + is_training=True, + deterministic=_deterministic, + ) + assert available_backends[1], "FusedAttention should support this SM8x THD training config." + test_dot_product_attention( + dtype, + model_configs_thd_sm8x, + "thd_sm8x", + False, + "thd_thd_thd", + False, + False, + ) + + @pytest.mark.skipif(get_cudnn_version() < (9, 0, 0), reason="cuDNN 9.0.0+ is required.") @pytest.mark.skipif( get_device_compute_capability() < (9, 0), reason="THD is only supported on Hopper+." diff --git a/transformer_engine/common/fused_attn/config_and_params.cpp b/transformer_engine/common/fused_attn/config_and_params.cpp index ca4214dac3..9e944d1039 100644 --- a/transformer_engine/common/fused_attn/config_and_params.cpp +++ b/transformer_engine/common/fused_attn/config_and_params.cpp @@ -112,7 +112,7 @@ FusedAttnConfig FusedAttnConfig::make_cache_key() const { if (cache_cfg.is_ragged_q || cache_cfg.is_ragged_kv) { const auto cudnn_runtime_version = cudnnGetVersion(); const int sm_arch_ = cuda::sm_arch(cuda::current_device()); - if (cudnn_runtime_version >= 90600 && sm_arch_ != 120) { + if (supports_packed_ragged_graph(cudnn_runtime_version, sm_arch_)) { if (cache_cfg.is_ragged_q) { cache_cfg.max_seqlen_q = cache_cfg.bucketed_num_tokens_q; } diff --git a/transformer_engine/common/fused_attn/config_and_params.h b/transformer_engine/common/fused_attn/config_and_params.h index ebc5b3eb07..5cc495c2fa 100644 --- a/transformer_engine/common/fused_attn/config_and_params.h +++ b/transformer_engine/common/fused_attn/config_and_params.h @@ -19,6 +19,12 @@ namespace transformer_engine { namespace fused_attn { +// Packed THD graph dimensions and ragged Stats/LSE are not supported on SM8x or SM120. +// Those architectures require dense, BHSD-like graph dimensions for the auxiliary tensors. +inline constexpr bool supports_packed_ragged_graph(size_t cudnn_runtime_version, int sm_arch) { + return cudnn_runtime_version >= 90600 && sm_arch >= 90 && sm_arch != 120; +} + struct FusedAttnConfig { // basic attention settings bool is_training = true; diff --git a/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu b/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu index b7c7a349af..28d0d34fca 100644 --- a/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu +++ b/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu @@ -81,7 +81,9 @@ void fused_attn_arbitrary_seqlen_fwd_impl( const auto cudnn_runtime_version = cudnnGetVersion(); const int device_id = cuda::current_device(); const int sm_arch_ = cuda::sm_arch(device_id); - bool use_ragged_stats = is_ragged_q && cudnn_runtime_version >= 90600 && sm_arch_ != 120; + const bool use_packed_ragged_graph = + supports_packed_ragged_graph(cudnn_runtime_version, sm_arch_); + const bool use_ragged_stats = is_ragged_q && use_packed_ragged_graph; NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); bool is_paged_kv = cfg.is_paged_kv; @@ -98,10 +100,9 @@ void fused_attn_arbitrary_seqlen_fwd_impl( int64_t actual_b = b; if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600) { NVTE_CHECK(is_padding, "Ragged QKV input requires padding or padding_causal mask!"); - // On SM 120, cuDNN support check treats layouts with stride[0] > dim[1]*dim[2]*dim[3] - // as interleaved and rejects them. Use BHSD-like dimensions/strides with max_seqlen at plan build - // so the check passes; ragged offset still provides variable-length boundaries. - if (sm_arch_ != 120) { + // SM8x and SM120 require BHSD-like dimensions/strides with max_seqlen at plan build. + // Other supported architectures use token-count dimensions for graph reuse. + if (use_packed_ragged_graph) { // replace batch size and maximum sequence lengths with maximum token counts // for query and key/value so the graph is static within each quantization bucket. // When passing cu_seqlens* directly to cuDNN SDPA, keep the true batch size: @@ -660,14 +661,16 @@ void fused_attn_arbitrary_seqlen_bwd_impl( const auto cudnn_runtime_version = cudnnGetVersion(); const int device_id = cuda::current_device(); const int sm_arch_ = cuda::sm_arch(device_id); - bool use_ragged_stats = is_ragged_q && cudnn_runtime_version >= 90600 && sm_arch_ != 120; + const bool use_packed_ragged_graph = + supports_packed_ragged_graph(cudnn_runtime_version, sm_arch_); + const bool use_ragged_stats = is_ragged_q && use_packed_ragged_graph; // keep original batch size because cu_seqlens are created with [b+1] shape int64_t actual_b = b; if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600) { NVTE_CHECK(is_padding, "Ragged QKV input requires padding or padding_causal mask!"); - // On SM 120, cuDNN support check requires BHSD-like strides with max_seqlen (see fwd). - if (sm_arch_ != 120) { + // SM8x and SM120 require BHSD-like strides with max_seqlen (see fwd). + if (use_packed_ragged_graph) { // replace batch size and maximum sequence lengths with maximum token counts // for query and key/value so the graph is static within each quantization bucket b = bucketed_batch_size; @@ -835,7 +838,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl( if (use_ragged_stats) { sdpa_backward_options.set_max_total_seq_len_q(s_q); } - if (is_ragged_kv && cudnn_runtime_version >= 90600 && sm_arch_ != 120) { + if (is_ragged_kv && use_packed_ragged_graph) { sdpa_backward_options.set_max_total_seq_len_kv(s_kv); } @@ -1134,12 +1137,10 @@ void fused_attn_arbitrary_seqlen_fwd(const FusedAttnConfig &cfg, const Tensor *i const size_t max_seqlen_q = cfg.max_seqlen_q; const size_t num_tokens_q = cfg.num_tokens_q; const bool return_max_logit = cfg.return_max_logit; - const NVTE_QKV_Layout qkv_layout = cfg.qkv_layout; const NVTE_Bias_Type bias_type = cfg.bias_type; const NVTE_Softmax_Type softmax_type = cfg.softmax_type; const auto QKV_type = input_Q->data.dtype; - NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); void *devPtrQ = input_Q->data.dptr; void *devPtrK = input_K->data.dptr; void *devPtrV = input_V->data.dptr; @@ -1171,13 +1172,13 @@ void fused_attn_arbitrary_seqlen_fwd(const FusedAttnConfig &cfg, const Tensor *i size_t i = 0; if (Aux_CTX_Tensors->size == 0) { const auto cudnn_runtime_version = cudnnGetVersion(); + const bool use_ragged_stats = + graph_cfg.is_ragged_q && supports_packed_ragged_graph(cudnn_runtime_version, sm_arch_); Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_S->data.dptr = nullptr; - // sm120 does not use ragged stats: the graph declares a dense - // [b, h, s_q, 1] stats tensor, so allocate to match (same as Max below). - if ((q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) && - (sm_arch_ != 120)) { + // Match the packed or dense shape declared by the graph (same as Max below). + if (use_ragged_stats) { output_S->data.shape = {num_tokens_q, num_attn_heads, 1}; } else { output_S->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; @@ -1187,8 +1188,7 @@ void fused_attn_arbitrary_seqlen_fwd(const FusedAttnConfig &cfg, const Tensor *i if (return_max_logit) { Tensor *output_Max = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_Max->data.dptr = nullptr; - if ((q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) && - (sm_arch_ != 120)) { + if (use_ragged_stats) { output_Max->data.shape = {num_tokens_q, num_attn_heads, 1}; } else { output_Max->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 8cbf6342c4..f1a411ee80 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -1111,13 +1111,13 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt use_flash_attention_4 = False # Filter: QKV layout - if qkv_format == "thd": + if "thd" in (q_format, kv_format): if pad_between_seqs: if ( # pylint: disable=too-many-boolean-expressions use_flash_attention_2 and FlashAttentionUtils.is_installed ) or (use_flash_attention_4 and FlashAttentionUtils.v4_is_installed): logger.debug( - "Disabling FlashAttention 2 and 4 for qkv_format = thd when there is " + "Disabling FlashAttention 2 and 4 when Q or KV uses THD and there is " "padding between sequences, i.e. [a, a, PAD, b, b, b, PAD, c, PAD]" ) use_flash_attention_2 = False @@ -1130,7 +1130,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt if cudnn_version < (9, 18, 1): if use_fused_attention: logger.debug( - "Disabling FusedAttention as qkv_format = thd is" + "Disabling FusedAttention when Q or KV uses THD because it is" " not supported for compute capability = sm120 and cuDNN version < 9.18.1" ) use_fused_attention = False From 5a3a7f065d90725519df8d1e8aad8b93e6746fce Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Wed, 29 Jul 2026 14:41:36 -0700 Subject: [PATCH 2/2] Remove SM8x THD test changes Keep the cross-fork follow-up limited to the requested production-code changes. Signed-off-by: Sudhakar Singh --- tests/cpp/util/CMakeLists.txt | 1 - tests/cpp/util/test_fused_attn_config.cpp | 25 -------- tests/pytorch/attention/test_attention.py | 74 ----------------------- 3 files changed, 100 deletions(-) delete mode 100644 tests/cpp/util/test_fused_attn_config.cpp diff --git a/tests/cpp/util/CMakeLists.txt b/tests/cpp/util/CMakeLists.txt index fb326a6292..1dfd2fed4e 100644 --- a/tests/cpp/util/CMakeLists.txt +++ b/tests/cpp/util/CMakeLists.txt @@ -3,7 +3,6 @@ # See LICENSE for license information. add_executable(test_util - test_fused_attn_config.cpp test_nvrtc.cpp test_string.cpp ../test_common.cu) diff --git a/tests/cpp/util/test_fused_attn_config.cpp b/tests/cpp/util/test_fused_attn_config.cpp deleted file mode 100644 index 6bf60bd403..0000000000 --- a/tests/cpp/util/test_fused_attn_config.cpp +++ /dev/null @@ -1,25 +0,0 @@ -/************************************************************************* - * Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -#include - -#include "common/fused_attn/config_and_params.h" - -namespace transformer_engine::fused_attn { - -TEST(FusedAttnConfig, PackedRaggedGraphSupport) { - EXPECT_FALSE(supports_packed_ragged_graph(90500, 90)); - EXPECT_TRUE(supports_packed_ragged_graph(90600, 90)); - EXPECT_TRUE(supports_packed_ragged_graph(91801, 100)); - - // SM8x and SM120 require dense Stats/LSE and max-sequence graph dimensions, - // even when the cuDNN runtime supports THD inputs on those architectures. - EXPECT_FALSE(supports_packed_ragged_graph(91801, 80)); - EXPECT_FALSE(supports_packed_ragged_graph(91801, 89)); - EXPECT_FALSE(supports_packed_ragged_graph(91801, 120)); -} - -} // namespace transformer_engine::fused_attn diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index aaca1bf166..6c876ad7cc 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -936,23 +936,6 @@ def test_dpa_qkv_layout(dtype, model_configs, model, qkv_layout): test_dot_product_attention(dtype, model_configs, model, False, qkv_layout, False, False) -def test_mixed_thd_pad_between_seqs_filter(): - """Mixed THD layouts apply the same backend-independent padding filter as pure THD.""" - config = ModelConfig(2, 128, 16, 64, attn_mask_type="padding") - available_backends, _, _ = get_available_attention_backends( - config, - qkv_dtype=torch.bfloat16, - qkv_layout="thd_bshd_bshd", - pad_between_seqs=True, - is_training=False, - skip_fused_attn=True, - ) - assert not available_backends[2], ( - "UnfusedDotProductAttention must be disabled when either Q or KV uses THD " - "with padding between sequences." - ) - - qkv_layouts_packed = [l for l in qkv_layouts if any(c.isdigit() for c in l)] @@ -1029,63 +1012,6 @@ def test_dpa_qkv_layout_declarative(dtype, model_configs, model, qkv_layout): } -model_configs_thd_sm8x = { - "thd_sm8x": ModelConfig(2, 128, 16, 64, attn_mask_type="padding"), -} - - -@pytest.mark.skipif( - device_compute_capability[0] != 8 or get_cudnn_version() >= (9, 18, 1), - reason="Unsupported THD boundary requires SM8x with cuDNN older than 9.18.1.", -) -def test_dpa_thd_sm8x_unsupported_cudnn(): - """The graph probe rejects SM8x THD training before cuDNN 9.18.1.""" - config = model_configs_thd_sm8x["thd_sm8x"] - available_backends, _, _ = get_available_attention_backends( - config, - qkv_dtype=torch.bfloat16, - qkv_layout="thd_thd_thd", - pad_between_seqs=False, - is_training=True, - deterministic=_deterministic, - ) - assert not available_backends[ - 1 - ], "FusedAttention must reject SM8x THD training when cuDNN is older than 9.18.1." - - -@pytest.mark.skipif( - device_compute_capability[0] != 8, - reason="Dense THD Stats/LSE regression is specific to SM8x.", -) -@pytest.mark.skipif( - get_cudnn_version() < (9, 18, 1), - reason="THD training on SM8x requires cuDNN 9.18.1+.", -) -@pytest.mark.parametrize("dtype", param_types_lean) -def test_dpa_thd_sm8x_dense_stats(dtype): - """SM8x THD training uses dense Stats/LSE and executes both forward and backward.""" - config = model_configs_thd_sm8x["thd_sm8x"] - available_backends, _, _ = get_available_attention_backends( - config, - qkv_dtype=dtype, - qkv_layout="thd_thd_thd", - pad_between_seqs=False, - is_training=True, - deterministic=_deterministic, - ) - assert available_backends[1], "FusedAttention should support this SM8x THD training config." - test_dot_product_attention( - dtype, - model_configs_thd_sm8x, - "thd_sm8x", - False, - "thd_thd_thd", - False, - False, - ) - - @pytest.mark.skipif(get_cudnn_version() < (9, 0, 0), reason="cuDNN 9.0.0+ is required.") @pytest.mark.skipif( get_device_compute_capability() < (9, 0), reason="THD is only supported on Hopper+."