From f108426d003e6f9e6cbe9dda3d0c84a313b97f26 Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani Date: Wed, 27 May 2026 17:41:35 -0700 Subject: [PATCH 01/26] Refine the support for D=256 on Blackwell server type GPUs Signed-off-by: Kshitij Lakhani --- .../common/fused_attn/fused_attn.cpp | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index fc21771297..3e36089252 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -315,6 +315,25 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( (head_dim_qk <= 256 && head_dim_v <= 256 && ((!is_training && sm_arch_ == 90 && cudnn_runtime_version >= 90100) || (is_training && sm_arch_ == 90 && cudnn_runtime_version >= 90500))) || + // 9.23: d_qk = d_v = 256 + SM10.x + bprop + non-paged + // cuDNN's dedicated SM10.x D=256 SDPA backward kernel (cuDNN FE 1.24 / + // BE 9.23+) only supports d_qk == d_v == 256. + (head_dim_qk == 256 && head_dim_v == 256 && is_training && + sm_arch_ >= 100 && sm_arch_ < 110 && cudnn_runtime_version >= 92300 && + layout_group != NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD && + // The FE forces this path onto deterministic bprop, which then rejects alibi and + // dropout, and only supports vanilla softmax. + bias_type != NVTE_Bias_Type::NVTE_ALIBI && dropout == 0.0 && + softmax_type == NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX && + // Non-causal D=256 supports only full-window attention; SWA is allowed only for causal masks. + // NOTE: SWA support for non causal would be available when cuDNN decides to redirect + // D=256 support to a CUDA backend instead of a Python OSS kernel + ((window_size_left == -1 && window_size_right == -1) || + ((attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || + attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK || + attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK || + attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK) && + (window_size_right == -1 || window_size_right == 0)))) || // 9.9: any head_dim + Blackwell + fprop + non_paged + sq > 1 (!is_training && sm_arch_ >= 100 && cudnn_runtime_version >= 90900 && max_seqlen_q > 1 && layout_group != NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD) || From 8ca68e27adf1e4497fa09bb696c8e9ec57018242 Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani Date: Thu, 28 May 2026 11:38:21 -0700 Subject: [PATCH 02/26] Add deterministic tests for D=256 for sm10.x Signed-off-by: Kshitij Lakhani --- tests/jax/test_fused_attn.py | 160 +++++++++++++++++++++++++++++++++++ 1 file changed, 160 insertions(+) diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 0d1db0b9e1..3b6fe10fae 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -520,6 +520,59 @@ def _check_configs(self): "is either BSHD_BSHD_BSHD or THD_THD_THD" ) + # D=256 bprop on SM10.x uses cuDNN's dedicated SDPA bprop kernel + # (cuDNN FE 1.24 / BE 9.23+). FE forces this path onto the deterministic algorithm path, + # which rejects dBias, dropout, and ALiBi. It supports vanilla softmax only and allows SWA + # together with a causal mask only. + compute_capability = get_device_compute_capability(0) + is_sm10x = 100 <= compute_capability < 110 + if self.is_training and is_sm10x and (self.head_dim_qk == 256 or self.head_dim_v == 256): + if self.head_dim_qk != 256 or self.head_dim_v != 256: + pytest.skip( + "D=256 BWD on Blackwell only supports d_qk == d_v == 256;" + f" got d_qk={self.head_dim_qk}, d_v={self.head_dim_v}." + ) + cudnn_version = get_cudnn_version() + if cudnn_version < 92300: + pytest.skip( + "D=256 BWD on Blackwell requires cuDNN 9.23 or newer;" + f" got cuDNN {cudnn_version}." + ) + # Non-learnable bias is fine (bias is allowed as an input); only dBias is + # unsupported. The JAX runner asks for dBias iff the bias shape is [1, h, s, s] + # (see test_backward), so gate on that. + unsupported = None + if self.attn_bias_type == AttnBiasType.PRE_SCALE_BIAS: + unsupported = "pre-scale bias" + elif ( + self.attn_bias_type != AttnBiasType.NO_BIAS + and self.bias_shape == BiasShape._1HSS + ): + unsupported = ( + "bias gradients (dBias); frozen/non-learnable bias inputs" + " (i.e. non-1HSS bias shapes) are supported" + ) + elif self.dropout_prob != 0.0: + unsupported = "dropout" + elif self.softmax_type != AttnSoftmaxType.VANILLA_SOFTMAX: + unsupported = "non-vanilla softmax" + if unsupported is not None: + pytest.skip( + "D=256 BWD on Blackwell uses the deterministic SM100 D=256 SDPA BWD" + f" kernel which does not support {unsupported}." + ) + if self.window_size is not None and self.window_size != (-1, -1): + if not self.attn_mask_type.is_causal(): + pytest.skip( + "D=256 BWD on Blackwell uses the SM100 D=256 SDPA BWD kernel" + " which requires window_size=(-1, -1) for non-causal masks." + ) + if self.window_size[1] not in (-1, 0): + pytest.skip( + "D=256 BWD on Blackwell only supports right window -1 or 0" + " for causal masks." + ) + self.backend = FusedAttnHelper( self.is_training, self.dtype, @@ -1773,3 +1826,110 @@ def test_backward( swa, seq_desc_format, ) + + +@pytest.mark.parametrize( + "attn_mask_type", + [ + pytest.param(AttnMaskType.NO_MASK, id="NO_MASK"), + pytest.param(AttnMaskType.PADDING_MASK, id="PADDING"), + pytest.param(AttnMaskType.CAUSAL_MASK, id="CAUSAL"), + pytest.param(AttnMaskType.PADDING_CAUSAL_MASK, id="PADDING_CAUSAL"), + pytest.param( + AttnMaskType.PADDING_CAUSAL_BOTTOM_RIGHT_MASK, id="PADDING_CAUSAL_BOTTOM_RIGHT" + ), + ], +) +@pytest.mark.parametrize( + "softmax_type", + [ + pytest.param(AttnSoftmaxType.VANILLA_SOFTMAX, id="VANILLA_SOFTMAX"), + ], +) +@pytest.mark.parametrize( + "dropout_prob", + [ + pytest.param(0.0, id="DROP_0.0"), + ], +) +@pytest.mark.parametrize( + "swa", + [ + pytest.param(False, id="NO_SWA"), + ], +) +@pytest.mark.parametrize( + "seq_desc_format", + [ + pytest.param(SeqDescFormat.Seqlens, id="Seqlens"), + ], +) +@pytest.mark.skipif(not _deterministic, reason="Test determinism only") +class TestFusedAttnD256WithDeterminism: + """ + Fused attention D=256 deterministic backward tester. + """ + + @staticmethod + @pytest.mark.parametrize( + "b, s_q, s_kv, h_q, h_kv, d_qk, d_v, dtype, qkv_layout", + [ + pytest.param( + 4, + 128, + 128, + 16, + 16, + 256, + 256, + jnp.float16, + QKVLayout.BSHD_BS2HD, + id="4-128-128-16-16-256-256-FP16-SELF-KV_PACKED", + ), + ], + ) + @pytest.mark.parametrize( + "attn_bias_type, bias_shape", + [ + pytest.param(AttnBiasType.NO_BIAS, None, id="NO_BIAS"), + ], + ) + def test_backward( + b, + s_q, + s_kv, + h_q, + h_kv, + d_qk, + d_v, + attn_bias_type, + attn_mask_type, + softmax_type, + dropout_prob, + dtype, + qkv_layout, + bias_shape, + swa, + seq_desc_format, + ): + """ + Test D=256 backward with the supported deterministic SM100 bprop configuration. + """ + TestFusedAttn.test_backward( + b, + s_q, + s_kv, + h_q, + h_kv, + d_qk, + d_v, + attn_bias_type, + attn_mask_type, + softmax_type, + dropout_prob, + dtype, + qkv_layout, + bias_shape, + swa, + seq_desc_format, + ) From a093127183d9f9859def5bd85c9b9b99dce1a5d1 Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani Date: Thu, 28 May 2026 12:48:10 -0700 Subject: [PATCH 03/26] Test clean up Signed-off-by: Kshitij Lakhani --- tests/jax/test_fused_attn.py | 122 +++++------------------------------ 1 file changed, 15 insertions(+), 107 deletions(-) diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 3b6fe10fae..639caf45b1 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -1695,6 +1695,21 @@ def test_backward( QKVLayout.THD_THD_THD, id="2-1024-2048-12-6-128-64-BF16-CROSS-GQA-RAGGED_SEPARATE", ), + # D=256 deterministic backward on the SM100 dedicated SDPA bprop kernel + # (cuDNN FE 1.24 / BE 9.23+). Unsupported configs (e.g. dBias, non-256 head dims) + # are skipped by FusedAttnRunner._check_configs. + pytest.param( + 4, + 128, + 128, + 16, + 16, + 256, + 256, + jnp.float16, + QKVLayout.BSHD_BS2HD, + id="4-128-128-16-16-256-256-FP16-SELF-KV_PACKED", + ), ], ) @pytest.mark.parametrize( @@ -1826,110 +1841,3 @@ def test_backward( swa, seq_desc_format, ) - - -@pytest.mark.parametrize( - "attn_mask_type", - [ - pytest.param(AttnMaskType.NO_MASK, id="NO_MASK"), - pytest.param(AttnMaskType.PADDING_MASK, id="PADDING"), - pytest.param(AttnMaskType.CAUSAL_MASK, id="CAUSAL"), - pytest.param(AttnMaskType.PADDING_CAUSAL_MASK, id="PADDING_CAUSAL"), - pytest.param( - AttnMaskType.PADDING_CAUSAL_BOTTOM_RIGHT_MASK, id="PADDING_CAUSAL_BOTTOM_RIGHT" - ), - ], -) -@pytest.mark.parametrize( - "softmax_type", - [ - pytest.param(AttnSoftmaxType.VANILLA_SOFTMAX, id="VANILLA_SOFTMAX"), - ], -) -@pytest.mark.parametrize( - "dropout_prob", - [ - pytest.param(0.0, id="DROP_0.0"), - ], -) -@pytest.mark.parametrize( - "swa", - [ - pytest.param(False, id="NO_SWA"), - ], -) -@pytest.mark.parametrize( - "seq_desc_format", - [ - pytest.param(SeqDescFormat.Seqlens, id="Seqlens"), - ], -) -@pytest.mark.skipif(not _deterministic, reason="Test determinism only") -class TestFusedAttnD256WithDeterminism: - """ - Fused attention D=256 deterministic backward tester. - """ - - @staticmethod - @pytest.mark.parametrize( - "b, s_q, s_kv, h_q, h_kv, d_qk, d_v, dtype, qkv_layout", - [ - pytest.param( - 4, - 128, - 128, - 16, - 16, - 256, - 256, - jnp.float16, - QKVLayout.BSHD_BS2HD, - id="4-128-128-16-16-256-256-FP16-SELF-KV_PACKED", - ), - ], - ) - @pytest.mark.parametrize( - "attn_bias_type, bias_shape", - [ - pytest.param(AttnBiasType.NO_BIAS, None, id="NO_BIAS"), - ], - ) - def test_backward( - b, - s_q, - s_kv, - h_q, - h_kv, - d_qk, - d_v, - attn_bias_type, - attn_mask_type, - softmax_type, - dropout_prob, - dtype, - qkv_layout, - bias_shape, - swa, - seq_desc_format, - ): - """ - Test D=256 backward with the supported deterministic SM100 bprop configuration. - """ - TestFusedAttn.test_backward( - b, - s_q, - s_kv, - h_q, - h_kv, - d_qk, - d_v, - attn_bias_type, - attn_mask_type, - softmax_type, - dropout_prob, - dtype, - qkv_layout, - bias_shape, - swa, - seq_desc_format, - ) From b7dc18c7a536b886bc172ee72b76ac15bc067aca Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani Date: Thu, 28 May 2026 13:28:03 -0700 Subject: [PATCH 04/26] Add PyT side tests for D=256 cuDNN fused attn on SM100-110 Signed-off-by: Kshitij Lakhani --- tests/pytorch/attention/test_attention.py | 38 +++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index f2ab4a8495..e30099d226 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -371,6 +371,44 @@ def test_dpa_fa4_hdim256(dtype, model_configs, model): test_dot_product_attention(dtype, model_configs, model, False, None, False, False) +# cuDNN FusedAttention head_dim=256 backward is supported on Blackwell server GPUs +# (SM100/SM103) from cuDNN 9.23 (FE 1.24), via the dedicated deterministic SDPA bprop +# kernel. It requires d_qk == d_v == 256, vanilla softmax, no dropout, no ALiBi, and +# (for non-causal masks) full-window attention. See nvte_get_fused_attn_backend in +# transformer_engine/common/fused_attn/fused_attn.cpp. These configs use d_qk == d_v == 256 +# with s_q == s_kv > 1 so the training (forward + backward) FusedAttention route is exercised +# and compared against the reference backends. +model_configs_fused_hdim256 = { + # test: ModelConfig(b, sq, hq, dqk) -> head_dim_v defaults to head_dim_qk (256) + "fused_hd256_no_mask": ModelConfig(2, 512, 16, 256), + "fused_hd256_causal": ModelConfig(2, 512, 16, 256, attn_mask_type="causal"), + "fused_hd256_padding": ModelConfig(2, 512, 16, 256, attn_mask_type="padding"), + "fused_hd256_padding_causal": ModelConfig(2, 512, 16, 256, attn_mask_type="padding_causal"), + "fused_hd256_padding_causal_br": ModelConfig( + 2, 512, 16, 256, attn_mask_type="padding_causal_bottom_right" + ), + # SWA is allowed only together with a causal mask on the D=256 bprop kernel. + "fused_hd256_causal_swa": ModelConfig( + 2, 512, 16, 256, attn_mask_type="causal", window_size=(128, 0) + ), + # GQA variant (num_gqa_groups < num_heads). + "fused_hd256_gqa": ModelConfig(2, 512, 16, 256, num_gqa_groups=4, attn_mask_type="causal"), +} + + +@pytest.mark.skipif(get_cudnn_version() < (9, 23, 0), reason="cuDNN 9.23+ is required.") +@pytest.mark.skipif( + device_compute_capability not in ((10, 0), (10, 3)), + reason="cuDNN FusedAttention head_dim=256 backward is Blackwell server (SM100/SM103) only.", +) +@pytest.mark.parametrize("dtype", param_types) +@pytest.mark.parametrize("model_configs", [model_configs_fused_hdim256]) +@pytest.mark.parametrize("model", model_configs_fused_hdim256.keys()) +def test_dpa_fused_attn_hdim256(dtype, model_configs, model): + """Test DotProductAttention with cuDNN FusedAttention: head_dim=256 backward on Blackwell""" + test_dot_product_attention(dtype, model_configs, model, False, True, None, False, False) + + model_configs_fa4_mla = { # test: ModelConfig(b, sq, hq, dqk, head_dim_v=dv) "fa4_mla_1": ModelConfig(4, 128, 16, 128, head_dim_v=64), From 28fe7baf19a09ac4d1bc2a05be1ebd22f4c6cc03 Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani Date: Thu, 28 May 2026 16:04:04 -0700 Subject: [PATCH 05/26] Refine the PyT D=256 tests Signed-off-by: Kshitij Lakhani --- tests/pytorch/attention/test_attention.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index e30099d226..ea445f29d5 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -381,18 +381,13 @@ def test_dpa_fa4_hdim256(dtype, model_configs, model): model_configs_fused_hdim256 = { # test: ModelConfig(b, sq, hq, dqk) -> head_dim_v defaults to head_dim_qk (256) "fused_hd256_no_mask": ModelConfig(2, 512, 16, 256), - "fused_hd256_causal": ModelConfig(2, 512, 16, 256, attn_mask_type="causal"), "fused_hd256_padding": ModelConfig(2, 512, 16, 256, attn_mask_type="padding"), - "fused_hd256_padding_causal": ModelConfig(2, 512, 16, 256, attn_mask_type="padding_causal"), - "fused_hd256_padding_causal_br": ModelConfig( - 2, 512, 16, 256, attn_mask_type="padding_causal_bottom_right" - ), # SWA is allowed only together with a causal mask on the D=256 bprop kernel. "fused_hd256_causal_swa": ModelConfig( - 2, 512, 16, 256, attn_mask_type="causal", window_size=(128, 0) + 2, 1024, 16, 256, attn_mask_type="causal", window_size=(128, 0) ), # GQA variant (num_gqa_groups < num_heads). - "fused_hd256_gqa": ModelConfig(2, 512, 16, 256, num_gqa_groups=4, attn_mask_type="causal"), + "fused_hd256_padding_causal_gqa": ModelConfig(2, 1024, 16, 256, num_gqa_groups=4, attn_mask_type="padding_causal"), } From dd6cbf3b0cae9b1a82cb83b380ff6d18525b3890 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 28 May 2026 23:06:28 +0000 Subject: [PATCH 06/26] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/jax/test_fused_attn.py | 5 +---- tests/pytorch/attention/test_attention.py | 4 +++- transformer_engine/common/fused_attn/fused_attn.cpp | 8 ++++---- 3 files changed, 8 insertions(+), 9 deletions(-) diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 639caf45b1..0c69a6004e 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -544,10 +544,7 @@ def _check_configs(self): unsupported = None if self.attn_bias_type == AttnBiasType.PRE_SCALE_BIAS: unsupported = "pre-scale bias" - elif ( - self.attn_bias_type != AttnBiasType.NO_BIAS - and self.bias_shape == BiasShape._1HSS - ): + elif self.attn_bias_type != AttnBiasType.NO_BIAS and self.bias_shape == BiasShape._1HSS: unsupported = ( "bias gradients (dBias); frozen/non-learnable bias inputs" " (i.e. non-1HSS bias shapes) are supported" diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index ea445f29d5..9870eec68c 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -387,7 +387,9 @@ def test_dpa_fa4_hdim256(dtype, model_configs, model): 2, 1024, 16, 256, attn_mask_type="causal", window_size=(128, 0) ), # GQA variant (num_gqa_groups < num_heads). - "fused_hd256_padding_causal_gqa": ModelConfig(2, 1024, 16, 256, num_gqa_groups=4, attn_mask_type="padding_causal"), + "fused_hd256_padding_causal_gqa": ModelConfig( + 2, 1024, 16, 256, num_gqa_groups=4, attn_mask_type="padding_causal" + ), } diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 3e36089252..5b2d4d6b8f 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -318,16 +318,16 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( // 9.23: d_qk = d_v = 256 + SM10.x + bprop + non-paged // cuDNN's dedicated SM10.x D=256 SDPA backward kernel (cuDNN FE 1.24 / // BE 9.23+) only supports d_qk == d_v == 256. - (head_dim_qk == 256 && head_dim_v == 256 && is_training && - sm_arch_ >= 100 && sm_arch_ < 110 && cudnn_runtime_version >= 92300 && + (head_dim_qk == 256 && head_dim_v == 256 && is_training && sm_arch_ >= 100 && + sm_arch_ < 110 && cudnn_runtime_version >= 92300 && layout_group != NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD && // The FE forces this path onto deterministic bprop, which then rejects alibi and // dropout, and only supports vanilla softmax. bias_type != NVTE_Bias_Type::NVTE_ALIBI && dropout == 0.0 && softmax_type == NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX && // Non-causal D=256 supports only full-window attention; SWA is allowed only for causal masks. - // NOTE: SWA support for non causal would be available when cuDNN decides to redirect - // D=256 support to a CUDA backend instead of a Python OSS kernel + // NOTE: SWA support for non causal would be available when cuDNN decides to redirect + // D=256 support to a CUDA backend instead of a Python OSS kernel ((window_size_left == -1 && window_size_right == -1) || ((attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK || From 9920cbeda1d7e1ea7bc6175a41529d5ef1d0a30c Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani Date: Fri, 29 May 2026 15:17:27 -0700 Subject: [PATCH 07/26] Fix the filtering condition for bias type for D=256 on sm10x for cudnn fused attn Signed-off-by: Kshitij Lakhani --- transformer_engine/common/fused_attn/fused_attn.cpp | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 5b2d4d6b8f..74591e28a7 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -321,9 +321,12 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( (head_dim_qk == 256 && head_dim_v == 256 && is_training && sm_arch_ >= 100 && sm_arch_ < 110 && cudnn_runtime_version >= 92300 && layout_group != NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD && - // The FE forces this path onto deterministic bprop, which then rejects alibi and - // dropout, and only supports vanilla softmax. - bias_type != NVTE_Bias_Type::NVTE_ALIBI && dropout == 0.0 && + // The FE forces this path onto the deterministic bprop algorithm, which on + // Blackwell rejects dBias, dropout, and ALiBi (and supports vanilla softmax only). + // Require NO_BIAS: a learnable pre/post-scale bias would request dBias in the bprop + // graph and fail validation. The selector has no visibility into the bias shape, so + // gate on the bias type to avoid over-selecting this backend. + bias_type == NVTE_Bias_Type::NVTE_NO_BIAS && dropout == 0.0 && softmax_type == NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX && // Non-causal D=256 supports only full-window attention; SWA is allowed only for causal masks. // NOTE: SWA support for non causal would be available when cuDNN decides to redirect From 3f0628963c6199b2b5e4a85d73cf9956bcb86521 Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani Date: Fri, 29 May 2026 15:53:16 -0700 Subject: [PATCH 08/26] Code clean up Signed-off-by: Kshitij Lakhani --- tests/pytorch/attention/test_attention.py | 11 ++++------- transformer_engine/common/fused_attn/fused_attn.cpp | 9 +-------- 2 files changed, 5 insertions(+), 15 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 9870eec68c..88182e36ea 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -371,13 +371,10 @@ def test_dpa_fa4_hdim256(dtype, model_configs, model): test_dot_product_attention(dtype, model_configs, model, False, None, False, False) -# cuDNN FusedAttention head_dim=256 backward is supported on Blackwell server GPUs -# (SM100/SM103) from cuDNN 9.23 (FE 1.24), via the dedicated deterministic SDPA bprop -# kernel. It requires d_qk == d_v == 256, vanilla softmax, no dropout, no ALiBi, and -# (for non-causal masks) full-window attention. See nvte_get_fused_attn_backend in -# transformer_engine/common/fused_attn/fused_attn.cpp. These configs use d_qk == d_v == 256 -# with s_q == s_kv > 1 so the training (forward + backward) FusedAttention route is exercised -# and compared against the reference backends. +# cuDNN FusedAttention D=256 bprop is supported on sm10x from cuDNN 9.23 (FE 1.24), +# via the dedicated deterministic SDPA bprop kernel, which supports d_qk == d_v == 256 only, +# vanilla type of softmax only, no dropout, no ALiBi, and (for non-causal masks) full-window attention only. +# (for non-causal masks) full-window attention. model_configs_fused_hdim256 = { # test: ModelConfig(b, sq, hq, dqk) -> head_dim_v defaults to head_dim_qk (256) "fused_hd256_no_mask": ModelConfig(2, 512, 16, 256), diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 74591e28a7..8fab5169cc 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -315,22 +315,15 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( (head_dim_qk <= 256 && head_dim_v <= 256 && ((!is_training && sm_arch_ == 90 && cudnn_runtime_version >= 90100) || (is_training && sm_arch_ == 90 && cudnn_runtime_version >= 90500))) || - // 9.23: d_qk = d_v = 256 + SM10.x + bprop + non-paged - // cuDNN's dedicated SM10.x D=256 SDPA backward kernel (cuDNN FE 1.24 / - // BE 9.23+) only supports d_qk == d_v == 256. + // 9.23: d_qk = d_v = 256 + SM10x (cuDNN FE 1.24 / BE 9.23+) + bprop + non-paged (head_dim_qk == 256 && head_dim_v == 256 && is_training && sm_arch_ >= 100 && sm_arch_ < 110 && cudnn_runtime_version >= 92300 && layout_group != NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD && // The FE forces this path onto the deterministic bprop algorithm, which on // Blackwell rejects dBias, dropout, and ALiBi (and supports vanilla softmax only). - // Require NO_BIAS: a learnable pre/post-scale bias would request dBias in the bprop - // graph and fail validation. The selector has no visibility into the bias shape, so - // gate on the bias type to avoid over-selecting this backend. bias_type == NVTE_Bias_Type::NVTE_NO_BIAS && dropout == 0.0 && softmax_type == NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX && // Non-causal D=256 supports only full-window attention; SWA is allowed only for causal masks. - // NOTE: SWA support for non causal would be available when cuDNN decides to redirect - // D=256 support to a CUDA backend instead of a Python OSS kernel ((window_size_left == -1 && window_size_right == -1) || ((attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK || From 38de70a2236cee0749eed058bbc1afc32c3cba6d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 29 May 2026 22:54:12 +0000 Subject: [PATCH 09/26] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/attention/test_attention.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 88182e36ea..36d078cfa6 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -372,7 +372,7 @@ def test_dpa_fa4_hdim256(dtype, model_configs, model): # cuDNN FusedAttention D=256 bprop is supported on sm10x from cuDNN 9.23 (FE 1.24), -# via the dedicated deterministic SDPA bprop kernel, which supports d_qk == d_v == 256 only, +# via the dedicated deterministic SDPA bprop kernel, which supports d_qk == d_v == 256 only, # vanilla type of softmax only, no dropout, no ALiBi, and (for non-causal masks) full-window attention only. # (for non-causal masks) full-window attention. model_configs_fused_hdim256 = { From 021c45039b22d1ee4d0e581c53256fc523ff15e7 Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani Date: Wed, 3 Jun 2026 10:48:23 -0700 Subject: [PATCH 10/26] Switching test check logic around Signed-off-by: Kshitij Lakhani --- tests/jax/test_fused_attn.py | 66 +++++++++++++++++------------------- 1 file changed, 32 insertions(+), 34 deletions(-) diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 0c69a6004e..5fe50a0c6f 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -470,6 +470,7 @@ def _get_max_segments_per_sequence(self): def _check_configs(self): # TODO(rewang): probably adds this in is_fused_attn_available + # TDOD(KshitijLakhani): probably add/move this to is_fused_attn_available if self.qkv_layout.is_thd() and not self.attn_mask_type.is_padding(): pytest.skip("THD format requires padding masks.") @@ -493,38 +494,11 @@ def _check_configs(self): pytest.skip( "seqlen_q > seqlen_kv is not supported with sliding window attention in cuDNN" ) - - if get_device_compute_capability(0) >= 100 and self.is_training: - if FusedAttnHelper.is_non_deterministic_allowed() and ( - (self.dropout_prob != 0.0 and self.attn_bias_type != AttnBiasType.NO_BIAS) - or get_cudnn_version() < 90700 - ): - pytest.skip( - "For sm100+, non-deterministic bprop (cuDNN 9.7+) does not support bias with" - " dropout" - ) - if not FusedAttnHelper.is_non_deterministic_allowed() and ( - self.dropout_prob != 0.0 - or self.attn_bias_type != AttnBiasType.NO_BIAS - or get_cudnn_version() < 91801 - ): - pytest.skip( - "For sm100+, deterministic bprop (cuDNN 9.18.1+) does not support bias or" - " dropout" - ) - # Test the MLA case where head dims for qk differ from head dims for v, only if the tensors - # are provided in BSHD_BSHD_BSHD or THD_THD_THD formats - if self.head_dim_qk != self.head_dim_v and not self.qkv_layout.is_separate(): - pytest.skip( - "For head_dim_qk != head_dim_v, it is necessary that the QKV layout " - "is either BSHD_BSHD_BSHD or THD_THD_THD" - ) - - # D=256 bprop on SM10.x uses cuDNN's dedicated SDPA bprop kernel - # (cuDNN FE 1.24 / BE 9.23+). FE forces this path onto the deterministic algorithm path, - # which rejects dBias, dropout, and ALiBi. It supports vanilla softmax only and allows SWA - # together with a causal mask only. compute_capability = get_device_compute_capability(0) + cudnn_version = get_cudnn_version() + # D=256 bprop on SM10x (cuDNN FE 1.24 / BE 9.23+) uses the deterministic algorithm path only, + # which rejects dBias, dropout, and ALiBi. It supports vanilla type of softmax only and allows SWA + # together with a causal mask only. is_sm10x = 100 <= compute_capability < 110 if self.is_training and is_sm10x and (self.head_dim_qk == 256 or self.head_dim_v == 256): if self.head_dim_qk != 256 or self.head_dim_v != 256: @@ -532,7 +506,6 @@ def _check_configs(self): "D=256 BWD on Blackwell only supports d_qk == d_v == 256;" f" got d_qk={self.head_dim_qk}, d_v={self.head_dim_v}." ) - cudnn_version = get_cudnn_version() if cudnn_version < 92300: pytest.skip( "D=256 BWD on Blackwell requires cuDNN 9.23 or newer;" @@ -570,6 +543,32 @@ def _check_configs(self): " for causal masks." ) + if compute_capability >= 100 and self.is_training: + if FusedAttnHelper.is_non_deterministic_allowed() and ( + (self.dropout_prob != 0.0 and self.attn_bias_type != AttnBiasType.NO_BIAS) + or cudnn_version < 90700 + ): + pytest.skip( + "For sm100+, non-deterministic bprop (cuDNN 9.7+) does not support bias with" + " dropout" + ) + if not FusedAttnHelper.is_non_deterministic_allowed() and ( + self.dropout_prob != 0.0 + or self.attn_bias_type != AttnBiasType.NO_BIAS + or cudnn_version < 91801 + ): + pytest.skip( + "For sm100+, deterministic bprop (cuDNN 9.18.1+) does not support bias or" + " dropout" + ) + # Test the MLA case where head dims for qk differ from head dims for v, only if the tensors + # are provided in BSHD_BSHD_BSHD or THD_THD_THD formats + if self.head_dim_qk != self.head_dim_v and not self.qkv_layout.is_separate(): + pytest.skip( + "For head_dim_qk != head_dim_v, it is necessary that the QKV layout " + "is either BSHD_BSHD_BSHD or THD_THD_THD" + ) + self.backend = FusedAttnHelper( self.is_training, self.dtype, @@ -1693,8 +1692,7 @@ def test_backward( id="2-1024-2048-12-6-128-64-BF16-CROSS-GQA-RAGGED_SEPARATE", ), # D=256 deterministic backward on the SM100 dedicated SDPA bprop kernel - # (cuDNN FE 1.24 / BE 9.23+). Unsupported configs (e.g. dBias, non-256 head dims) - # are skipped by FusedAttnRunner._check_configs. + # (cuDNN FE 1.24 / BE 9.23+). pytest.param( 4, 128, From b597edaa797ba4302564b0cec789a4c8130f5f98 Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani Date: Wed, 3 Jun 2026 14:46:32 -0700 Subject: [PATCH 11/26] Add THD xfail marked tests for testing Signed-off-by: Kshitij Lakhani --- tests/jax/test_fused_attn.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 5fe50a0c6f..04f71fe8b6 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -1705,6 +1705,22 @@ def test_backward( QKVLayout.BSHD_BS2HD, id="4-128-128-16-16-256-256-FP16-SELF-KV_PACKED", ), + pytest.param( + 4, + 128, + 128, + 16, + 16, + 256, + 256, + jnp.float16, + QKVLayout.THD_T2HD, + id="4-128-128-16-16-256-256-FP16-SELF-RAGGED_KV_PACKED", + marks=pytest.mark.xfail( + reason="cuDNN 9.23 D=256 BWD currently does not build a THD execution plan.", + strict=True, + ), + ), ], ) @pytest.mark.parametrize( From 4882cda79b6033b8f46f20f49237350b946651a6 Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani Date: Thu, 4 Jun 2026 17:23:49 -0700 Subject: [PATCH 12/26] Fix the jax test bias condition fo skipping for D=256 Signed-off-by: Kshitij Lakhani --- tests/jax/test_fused_attn.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 04f71fe8b6..071bedcc77 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -511,16 +511,16 @@ def _check_configs(self): "D=256 BWD on Blackwell requires cuDNN 9.23 or newer;" f" got cuDNN {cudnn_version}." ) - # Non-learnable bias is fine (bias is allowed as an input); only dBias is - # unsupported. The JAX runner asks for dBias iff the bias shape is [1, h, s, s] - # (see test_backward), so gate on that. + # TODO(KshitijLakhani): cuDNN FE can model bias input separately from dBias, + # but TE does not yet plumb whether dBias is requested into the common backend selector. + # Until that distinction is available, the D=256 SM10x gate requires no bias. unsupported = None if self.attn_bias_type == AttnBiasType.PRE_SCALE_BIAS: unsupported = "pre-scale bias" - elif self.attn_bias_type != AttnBiasType.NO_BIAS and self.bias_shape == BiasShape._1HSS: + elif self.attn_bias_type != AttnBiasType.NO_BIAS: unsupported = ( - "bias gradients (dBias); frozen/non-learnable bias inputs" - " (i.e. non-1HSS bias shapes) are supported" + "post-scale bias in TE's D=256 backend gate; bias-input-only" + " support needs TE to distinguish between bias input and dBias" ) elif self.dropout_prob != 0.0: unsupported = "dropout" From 4d5ad9443e2db7384971a59493e8fd0144c23a45 Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani Date: Thu, 4 Jun 2026 17:24:40 -0700 Subject: [PATCH 13/26] nit: Fix comment Signed-off-by: Kshitij Lakhani --- tests/pytorch/attention/test_attention.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 36d078cfa6..273902146d 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -374,7 +374,6 @@ def test_dpa_fa4_hdim256(dtype, model_configs, model): # cuDNN FusedAttention D=256 bprop is supported on sm10x from cuDNN 9.23 (FE 1.24), # via the dedicated deterministic SDPA bprop kernel, which supports d_qk == d_v == 256 only, # vanilla type of softmax only, no dropout, no ALiBi, and (for non-causal masks) full-window attention only. -# (for non-causal masks) full-window attention. model_configs_fused_hdim256 = { # test: ModelConfig(b, sq, hq, dqk) -> head_dim_v defaults to head_dim_qk (256) "fused_hd256_no_mask": ModelConfig(2, 512, 16, 256), From 6184e854710f76874ecdf117c42557a15936bfd2 Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani Date: Thu, 4 Jun 2026 17:40:25 -0700 Subject: [PATCH 14/26] Move the D=256 changes so as to follow an order of increasing cuDNN versions Signed-off-by: Kshitij Lakhani --- .../common/fused_attn/fused_attn.cpp | 28 +++++++++---------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 8fab5169cc..44255146e5 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -315,6 +315,19 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( (head_dim_qk <= 256 && head_dim_v <= 256 && ((!is_training && sm_arch_ == 90 && cudnn_runtime_version >= 90100) || (is_training && sm_arch_ == 90 && cudnn_runtime_version >= 90500))) || + // 9.9: any head_dim + Blackwell + fprop + non_paged + sq > 1 + (!is_training && sm_arch_ >= 100 && cudnn_runtime_version >= 90900 && max_seqlen_q > 1 && + layout_group != NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD) || + // 9.10.2: any head_dim + any arch + fprop + paged + // 9.10.2: any head_dim + any arch + fprop + non_paged + sq > 1 + // 9.10.2: any head_dim + any arch + fprop + non_paged + sq = 1 + {no_mask, padding, BRCM, padding_BRCM} + (!is_training && cudnn_runtime_version >= 91002 && + (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD || max_seqlen_q > 1 || + (max_seqlen_q == 1 && attn_mask_type != NVTE_Mask_Type::NVTE_CAUSAL_MASK && + attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK))) || + // 9.11: d_qk = 192, d_v = 128 + Blackwell + bprop + non-paged + (head_dim_qk == 192 && head_dim_v == 128 && is_training && sm_arch_ >= 100 && + cudnn_runtime_version >= 91100) || // 9.23: d_qk = d_v = 256 + SM10x (cuDNN FE 1.24 / BE 9.23+) + bprop + non-paged (head_dim_qk == 256 && head_dim_v == 256 && is_training && sm_arch_ >= 100 && sm_arch_ < 110 && cudnn_runtime_version >= 92300 && @@ -329,20 +342,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK || attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK || attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK) && - (window_size_right == -1 || window_size_right == 0)))) || - // 9.9: any head_dim + Blackwell + fprop + non_paged + sq > 1 - (!is_training && sm_arch_ >= 100 && cudnn_runtime_version >= 90900 && max_seqlen_q > 1 && - layout_group != NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD) || - // 9.10.2: any head_dim + any arch + fprop + paged - // 9.10.2: any head_dim + any arch + fprop + non_paged + sq > 1 - // 9.10.2: any head_dim + any arch + fprop + non_paged + sq = 1 + {no_mask, padding, BRCM, padding_BRCM} - (!is_training && cudnn_runtime_version >= 91002 && - (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD || max_seqlen_q > 1 || - (max_seqlen_q == 1 && attn_mask_type != NVTE_Mask_Type::NVTE_CAUSAL_MASK && - attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK))) || - // 9.11: d_qk = 192, d_v = 128 + Blackwell + bprop + non-paged - (head_dim_qk == 192 && head_dim_v == 128 && is_training && sm_arch_ >= 100 && - cudnn_runtime_version >= 91100)) && + (window_size_right == -1 || window_size_right == 0))))) && // 9.11+ bug: 128 < d_qk <= 256, 128 < d_v <= 256 + Hopper + bprop + MLA // Conditional to temporarily use blanket cudnn_runtime_version >= 9.11 until fixed (!((cudnn_runtime_version >= 91100) && is_training && sm_arch_ == 90 && From 424b2c49842f05e79ac61a768c03568d26e100f8 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 5 Jun 2026 00:41:54 +0000 Subject: [PATCH 15/26] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/jax/test_fused_attn.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 071bedcc77..264eb6f6b1 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -512,7 +512,7 @@ def _check_configs(self): f" got cuDNN {cudnn_version}." ) # TODO(KshitijLakhani): cuDNN FE can model bias input separately from dBias, - # but TE does not yet plumb whether dBias is requested into the common backend selector. + # but TE does not yet plumb whether dBias is requested into the common backend selector. # Until that distinction is available, the D=256 SM10x gate requires no bias. unsupported = None if self.attn_bias_type == AttnBiasType.PRE_SCALE_BIAS: From fe7c6b23225e1ef8a0149efa2b9aa473f94a43a0 Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani Date: Wed, 1 Jul 2026 16:23:53 -0700 Subject: [PATCH 16/26] Add support for D=256 THD in TE Signed-off-by: Kshitij Lakhani --- tests/jax/test_fused_attn.py | 25 +++++++------ tests/pytorch/attention/test_attention.py | 36 +++++++++++++++---- .../common/fused_attn/fused_attn.cpp | 8 +++-- 3 files changed, 47 insertions(+), 22 deletions(-) diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 264eb6f6b1..b09e641896 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -470,7 +470,7 @@ def _get_max_segments_per_sequence(self): def _check_configs(self): # TODO(rewang): probably adds this in is_fused_attn_available - # TDOD(KshitijLakhani): probably add/move this to is_fused_attn_available + # TODO(KshitijLakhani): probably add/move this to is_fused_attn_available if self.qkv_layout.is_thd() and not self.attn_mask_type.is_padding(): pytest.skip("THD format requires padding masks.") @@ -496,9 +496,10 @@ def _check_configs(self): ) compute_capability = get_device_compute_capability(0) cudnn_version = get_cudnn_version() - # D=256 bprop on SM10x (cuDNN FE 1.24 / BE 9.23+) uses the deterministic algorithm path only, - # which rejects dBias, dropout, and ALiBi. It supports vanilla type of softmax only and allows SWA - # together with a causal mask only. + # D=256 bprop on SM10x uses the deterministic algorithm path only. BSHD support + # starts with cuDNN FE 1.24 / BE 9.23; THD execution-plan support starts with + # cuDNN FE 1.26 / BE 9.30. The kernel rejects dBias, dropout, and ALiBi, supports vanilla + # softmax only, and allows SWA together with a causal mask only. is_sm10x = 100 <= compute_capability < 110 if self.is_training and is_sm10x and (self.head_dim_qk == 256 or self.head_dim_v == 256): if self.head_dim_qk != 256 or self.head_dim_v != 256: @@ -506,10 +507,12 @@ def _check_configs(self): "D=256 BWD on Blackwell only supports d_qk == d_v == 256;" f" got d_qk={self.head_dim_qk}, d_v={self.head_dim_v}." ) - if cudnn_version < 92300: + required_cudnn_version = 93000 if self.qkv_layout.is_thd() else 92300 + required_cudnn_version_label = "9.30" if self.qkv_layout.is_thd() else "9.23" + if cudnn_version < required_cudnn_version: pytest.skip( - "D=256 BWD on Blackwell requires cuDNN 9.23 or newer;" - f" got cuDNN {cudnn_version}." + f"D=256 BWD on Blackwell with {self.qkv_layout} requires cuDNN" + f" {required_cudnn_version_label} or newer; got cuDNN {cudnn_version}." ) # TODO(KshitijLakhani): cuDNN FE can model bias input separately from dBias, # but TE does not yet plumb whether dBias is requested into the common backend selector. @@ -1691,8 +1694,8 @@ def test_backward( QKVLayout.THD_THD_THD, id="2-1024-2048-12-6-128-64-BF16-CROSS-GQA-RAGGED_SEPARATE", ), - # D=256 deterministic backward on the SM100 dedicated SDPA bprop kernel - # (cuDNN FE 1.24 / BE 9.23+). + # D=256 deterministic backward on the SM100 dedicated SDPA bprop kernel. + # BSHD requires cuDNN FE 1.24 / BE 9.23+; THD requires cuDNN FE 1.26 / BE 9.30+. pytest.param( 4, 128, @@ -1716,10 +1719,6 @@ def test_backward( jnp.float16, QKVLayout.THD_T2HD, id="4-128-128-16-16-256-256-FP16-SELF-RAGGED_KV_PACKED", - marks=pytest.mark.xfail( - reason="cuDNN 9.23 D=256 BWD currently does not build a THD execution plan.", - strict=True, - ), ), ], ) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 273902146d..a6c8620a53 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -5,6 +5,7 @@ import os import sys import pathlib +import copy from typing import Any, Dict, Tuple, Union import pytest @@ -139,7 +140,7 @@ def test_dot_product_attention( tols = dict(atol=1e-3, rtol=1e-3) if dtype == torch.bfloat16: tols = dict(atol=1.5e-2, rtol=1.5e-2) - config = model_configs[model] + config = copy.deepcopy(model_configs[model]) is_mla = config.head_dim_qk != config.head_dim_v is_mqa_gqa = config.num_heads != config.num_gqa_groups if qkv_layout is None: @@ -371,9 +372,10 @@ def test_dpa_fa4_hdim256(dtype, model_configs, model): test_dot_product_attention(dtype, model_configs, model, False, None, False, False) -# cuDNN FusedAttention D=256 bprop is supported on sm10x from cuDNN 9.23 (FE 1.24), -# via the dedicated deterministic SDPA bprop kernel, which supports d_qk == d_v == 256 only, -# vanilla type of softmax only, no dropout, no ALiBi, and (for non-causal masks) full-window attention only. +# cuDNN FusedAttention D=256 bprop is supported on sm10x by the dedicated deterministic +# SDPA bprop kernel. BSHD support starts with cuDNN FE 1.24 / BE 9.23; THD support starts +# with cuDNN FE 1.26 / BE 9.30. The kernel supports d_qk == d_v == 256 only, vanilla softmax only, +# no dropout, no ALiBi, and (for non-causal masks) full-window attention only. model_configs_fused_hdim256 = { # test: ModelConfig(b, sq, hq, dqk) -> head_dim_v defaults to head_dim_qk (256) "fused_hd256_no_mask": ModelConfig(2, 512, 16, 256), @@ -389,7 +391,6 @@ def test_dpa_fa4_hdim256(dtype, model_configs, model): } -@pytest.mark.skipif(get_cudnn_version() < (9, 23, 0), reason="cuDNN 9.23+ is required.") @pytest.mark.skipif( device_compute_capability not in ((10, 0), (10, 3)), reason="cuDNN FusedAttention head_dim=256 backward is Blackwell server (SM100/SM103) only.", @@ -397,9 +398,30 @@ def test_dpa_fa4_hdim256(dtype, model_configs, model): @pytest.mark.parametrize("dtype", param_types) @pytest.mark.parametrize("model_configs", [model_configs_fused_hdim256]) @pytest.mark.parametrize("model", model_configs_fused_hdim256.keys()) -def test_dpa_fused_attn_hdim256(dtype, model_configs, model): +@pytest.mark.parametrize( + "qkv_layout", + [ + pytest.param( + "bshd_bs2hd", + id="BSHD_KV_PACKED", + marks=pytest.mark.skipif( + get_cudnn_version() < (9, 23, 0), + reason="cuDNN 9.23+ is required for BSHD D=256 fused-attn backward.", + ), + ), + pytest.param( + "thd_t2hd", + id="THD_KV_PACKED", + marks=pytest.mark.skipif( + get_cudnn_version() < (9, 30, 0), + reason="cuDNN 9.30+ is required for THD D=256 fused-attn backward.", + ), + ), + ], +) +def test_dpa_fused_attn_hdim256(dtype, model_configs, model, qkv_layout): """Test DotProductAttention with cuDNN FusedAttention: head_dim=256 backward on Blackwell""" - test_dot_product_attention(dtype, model_configs, model, False, True, None, False, False) + test_dot_product_attention(dtype, model_configs, model, False, True, qkv_layout, False, False) model_configs_fa4_mla = { diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 44255146e5..834c1805d5 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -240,6 +240,9 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); + const bool is_thd_layout = q_format == NVTE_QKV_Format::NVTE_THD || + kv_format == NVTE_QKV_Format::NVTE_THD || + qkv_format == NVTE_QKV_Format::NVTE_THD; NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); auto cudnn_runtime_version = cudnnGetVersion(); @@ -328,9 +331,10 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( // 9.11: d_qk = 192, d_v = 128 + Blackwell + bprop + non-paged (head_dim_qk == 192 && head_dim_v == 128 && is_training && sm_arch_ >= 100 && cudnn_runtime_version >= 91100) || - // 9.23: d_qk = d_v = 256 + SM10x (cuDNN FE 1.24 / BE 9.23+) + bprop + non-paged + // 9.23: d_qk = d_v = 256 + SM10x (cuDNN FE 1.24 / BE 9.23+) + bprop + non-paged. + // THD layouts require cuDNN FE 1.26 / BE 9.30+ for execution-plan support. (head_dim_qk == 256 && head_dim_v == 256 && is_training && sm_arch_ >= 100 && - sm_arch_ < 110 && cudnn_runtime_version >= 92300 && + sm_arch_ < 110 && cudnn_runtime_version >= (is_thd_layout ? 93000 : 92300) && layout_group != NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD && // The FE forces this path onto the deterministic bprop algorithm, which on // Blackwell rejects dBias, dropout, and ALiBi (and supports vanilla softmax only). From cfed53fdd7f0c3fa207b456cca2b19f16a37f6f3 Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani Date: Tue, 7 Jul 2026 17:11:16 -0700 Subject: [PATCH 17/26] Add TE JAX CP test support to verify D=256 for AG and Ring Signed-off-by: Kshitij Lakhani --- tests/jax/test_distributed_fused_attn.py | 167 +++++++++++++++++++++++ 1 file changed, 167 insertions(+) diff --git a/tests/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py index 2abd9824b6..f2ee2ececc 100644 --- a/tests/jax/test_distributed_fused_attn.py +++ b/tests/jax/test_distributed_fused_attn.py @@ -22,6 +22,7 @@ _has_cudnn_frontend_python, ) from utils import pytest_parametrize_wrapper +from transformer_engine_jax import get_cudnn_version, get_device_compute_capability from transformer_engine.jax.attention import ( is_fused_attn_kernel_available, AttnBiasType, @@ -345,6 +346,60 @@ def test_softcap_score_mod_with_aux_params_backward( pytest.param([4, 256, 16, 64], id="4-256xCPx2-16-64"), ] +DISTRIBUTED_CONTEXT_SELF_ATTN_D256_DATA_SHAPES = { + "L0": [], + "L1": [[2, 128, 16, 256]], + "L2": [], +} + +# Keep these as explicit tuples instead of independent layout/mask/window as: +# BSHD CP uses CAUSAL_MASK, THD CP uses PADDING_CAUSAL_MASK, SWA is +# only valid for THD, and stripe_size behavior is different for +# BSHD vs THD in these tests. Splitting the axes would mostly add +# invalid BSHD+SWA and THD+CAUSAL combinations that fail or skip. +DISTRIBUTED_CONTEXT_SELF_ATTN_D256_LAYOUTS_MASKS_WINDOWS = [ + # BSHD with different layouts, but same causal mask and no sliding window + pytest.param( + QKVLayout.BSHD_BS2HD, + AttnMaskType.CAUSAL_MASK, + (-1, -1), + id="BSHD_KVPACKED-CAUSAL-NO_SWA", + ), + pytest.param( + QKVLayout.BSHD_BSHD_BSHD, + AttnMaskType.CAUSAL_MASK, + (-1, -1), + id="BSHD_SEPARATE-CAUSAL-NO_SWA", + ), + # THD with different sliding window sizes, but same packed layout and padding causal mask + pytest.param( + QKVLayout.THD_T2HD, + AttnMaskType.PADDING_CAUSAL_MASK, + (-1, -1), + id="THD_KVPACKED-PADDING_CAUSAL-NO_SWA", + ), + pytest.param( + QKVLayout.THD_T2HD, + AttnMaskType.PADDING_CAUSAL_MASK, + (20, 0), + id="THD_KVPACKED-PADDING_CAUSAL-SWA", + ), + # THD with different sliding window sizes, but same separate layout and padding causal mask + pytest.param( + QKVLayout.THD_THD_THD, + AttnMaskType.PADDING_CAUSAL_MASK, + (-1, -1), + id="THD_SEPARATE-PADDING_CAUSAL-NO_SWA", + ), + pytest.param( + QKVLayout.THD_THD_THD, + AttnMaskType.PADDING_CAUSAL_MASK, + (20, 0), + id="THD_SEPARATE-PADDING_CAUSAL-SWA", + ), +] + + class TestDistributedContextParallelSelfAttn: # TODO(KshitijLakhani): parametrize num_segments_per_seq for all CP tests @@ -637,6 +692,118 @@ def test_context_parallel_ring_attn( window_size=window_size, stripe_size=stripe_size, ) + # CP ring and all-gather tests for D=256 + # TODO(KshitijLakhani): Replace this with common-provided fused-attn disable reasons once + # they can be surfaced to framework tests. + @staticmethod + def skip_if_d256_cp_unsupported(qkv_layout): + compute_capability = get_device_compute_capability(0) + if not 100 <= compute_capability < 110: + pytest.skip("D=256 CP fused attention is only enabled on Blackwell server GPUs.") + + required_cudnn_version = 93000 if qkv_layout.is_thd() else 92300 + required_cudnn_version_label = "9.30" if qkv_layout.is_thd() else "9.23" + if get_cudnn_version() < required_cudnn_version: + pytest.skip( + f"D=256 CP fused attention with {qkv_layout} requires cuDNN" + f" {required_cudnn_version_label} or newer." + ) + + @pytest_parametrize_wrapper( + "device_count,mesh_shape,mesh_axes,mesh_resource", + generate_context_parallel_configs_for_attn(), + ) + @pytest_parametrize_wrapper( + "data_shape", + DISTRIBUTED_CONTEXT_SELF_ATTN_D256_DATA_SHAPES, + ) + @pytest.mark.parametrize( + "dtype", + [pytest.param(jnp.float16, id="FP16"), pytest.param(jnp.bfloat16, id="BF16")], + ) + @pytest.mark.parametrize( + "qkv_layout, attn_mask_type, window_size", + DISTRIBUTED_CONTEXT_SELF_ATTN_D256_LAYOUTS_MASKS_WINDOWS, + ) + def test_context_parallel_ring_attn_d256( + self, + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + data_shape, + dtype, + qkv_layout, + attn_mask_type, + window_size, + ): + """D=256 CP ring coverage.""" + self.skip_if_d256_cp_unsupported(qkv_layout) + + self.impl_test_context_parallel_attn( + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + data_shape, + 1, + attn_mask_type, + dtype, + qkv_layout, + True, + CPStrategy.RING, + use_scan_ring=False, + window_size=window_size, + stripe_size=1 if qkv_layout.is_thd() else None, + ) + + @pytest_parametrize_wrapper( + "device_count,mesh_shape,mesh_axes,mesh_resource", + generate_context_parallel_configs_for_attn(), + ) + @pytest_parametrize_wrapper( + "data_shape", + DISTRIBUTED_CONTEXT_SELF_ATTN_D256_DATA_SHAPES, + ) + @pytest.mark.parametrize( + "dtype", + [pytest.param(jnp.float16, id="FP16"), pytest.param(jnp.bfloat16, id="BF16")], + ) + @pytest.mark.parametrize( + "qkv_layout, attn_mask_type, window_size", + DISTRIBUTED_CONTEXT_SELF_ATTN_D256_LAYOUTS_MASKS_WINDOWS, + ) + def test_context_parallel_allgather_attn_d256( + self, + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + data_shape, + dtype, + qkv_layout, + attn_mask_type, + window_size, + ): + """D=256 CP all-gather coverage.""" + self.skip_if_d256_cp_unsupported(qkv_layout) + + self.impl_test_context_parallel_attn( + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + data_shape, + 1, + attn_mask_type, + dtype, + qkv_layout, + True, + CPStrategy.ALL_GATHER, + window_size=window_size, + stripe_size=128 if qkv_layout.is_thd() else None, + num_segments_per_seq=5 if qkv_layout.is_thd() else None, + ) REORDER_CAUSAL_LOAD_BALANCING_DATA_SHAPES = { From 5767db0ba8bc0cef2b0a347518a2e304b619fdfd Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani Date: Tue, 7 Jul 2026 18:36:57 -0700 Subject: [PATCH 18/26] Add PyTorch D256 CP tests Signed-off-by: Kshitij Lakhani --- .../attention/run_attention_with_cp.py | 13 +- .../attention/test_attention_with_cp.py | 132 ++++++++++++++++++ 2 files changed, 143 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/attention/run_attention_with_cp.py b/tests/pytorch/attention/run_attention_with_cp.py index 82b9df262f..7cfe2b3212 100644 --- a/tests/pytorch/attention/run_attention_with_cp.py +++ b/tests/pytorch/attention/run_attention_with_cp.py @@ -15,7 +15,11 @@ from transformer_engine.pytorch.attention.dot_product_attention.utils import combine_and_quantize import transformer_engine_torch as tex from transformer_engine.pytorch import DType -from test_attention_with_cp import model_configs_flash_attn, model_configs_fused_attn +from test_attention_with_cp import ( + model_configs_flash_attn, + model_configs_fused_attn, + model_configs_fused_attn_hdim256, +) from transformer_engine.pytorch import ( autocast, DotProductAttention, @@ -231,7 +235,12 @@ def run_dpa_with_cp( config = copy.deepcopy(model_configs_flash_attn[model]) if kernel_backend == "FusedAttention": os.environ["NVTE_FUSED_ATTN"] = "1" - config = copy.deepcopy(model_configs_fused_attn[model]) + if model in model_configs_fused_attn: + config = copy.deepcopy(model_configs_fused_attn[model]) + elif model in model_configs_fused_attn_hdim256: + config = copy.deepcopy(model_configs_fused_attn_hdim256[model]) + else: + assert False, f"{model=} is not a known FusedAttention CP config!" assert config.attn_mask_type in [ "causal", "no_mask", diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index 681ee5e6e0..75b0306037 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -483,6 +483,16 @@ def test_cp_with_flash_attention(cp_pool, dtype, model, qkv_format, cp_comm_type } +model_configs_fused_attn_hdim256 = { + # cuDNN FusedAttention D=256 bprop is Blackwell server only. Keep these + # outside the generic CP sweep so the D=256 signal is focused and inexpensive. + "cp_hd256_causal": ModelConfig(2, 1024, 16, 256, attn_mask_type="causal"), + "cp_hd256_causal_swa": ModelConfig( + 2, 1024, 16, 256, attn_mask_type="causal", window_size=(128, 0) + ), +} + + dtypes = ["bf16", "fp16", "fp8"] qkv_formats = ["bshd", "sbhd", "thd"] cp_comm_types = ["p2p", "all_gather", "a2a", "a2a+p2p"] @@ -696,3 +706,125 @@ def test_cp_with_fused_attention( deterministic=_deterministic, log_level=pytest_logging_level, ) + + +# Keep these as explicit cases because the CP axes are not independent: +# - This runner covers separate-layout CP (bshd_bshd_bshd/thd_thd_thd); KV-packed +# D=256 is covered by test_dpa_fused_attn_hdim256 in test_attention.py. +# - THD starts from a causal config and is rewritten to padding_causal by the runner. +# - SWA is sampled only through all_gather; p2p does not support SWA. +BSHD_FUSED_HD256_CP_CUDNN_MARK = pytest.mark.skipif( + get_cudnn_version() < (9, 23, 0), + reason="cuDNN 9.23+ is required for BSHD D=256 fused-attn CP backward.", +) +THD_FUSED_HD256_CP_CUDNN_MARK = pytest.mark.skipif( + get_cudnn_version() < (9, 30, 0), + reason="cuDNN 9.30+ is required for THD D=256 fused-attn CP backward.", +) + +DPA_CP_FUSED_HD256_CASES = [ + pytest.param( + "cp_hd256_causal", + "bshd", + "p2p", + id="BSHD-P2P-CAUSAL-NO_SWA", + marks=BSHD_FUSED_HD256_CP_CUDNN_MARK, + ), + pytest.param( + "cp_hd256_causal", + "bshd", + "all_gather", + id="BSHD-ALL_GATHER-CAUSAL-NO_SWA", + marks=BSHD_FUSED_HD256_CP_CUDNN_MARK, + ), + pytest.param( + "cp_hd256_causal_swa", + "bshd", + "all_gather", + id="BSHD-ALL_GATHER-CAUSAL-SWA", + marks=BSHD_FUSED_HD256_CP_CUDNN_MARK, + ), + pytest.param( + "cp_hd256_causal", + "thd", + "p2p", + id="THD-P2P-PADDING_CAUSAL-NO_SWA", + marks=THD_FUSED_HD256_CP_CUDNN_MARK, + ), + pytest.param( + "cp_hd256_causal", + "thd", + "all_gather", + id="THD-ALL_GATHER-PADDING_CAUSAL-NO_SWA", + marks=THD_FUSED_HD256_CP_CUDNN_MARK, + ), + pytest.param( + "cp_hd256_causal_swa", + "thd", + "all_gather", + id="THD-ALL_GATHER-PADDING_CAUSAL-SWA", + marks=THD_FUSED_HD256_CP_CUDNN_MARK, + ), +] + + +@pytest.mark.skipif( + get_device_compute_capability() not in ((10, 0), (10, 3)), + reason="cuDNN FusedAttention head_dim=256 backward is SM100/SM103-only.", +) +@pytest.mark.parametrize("dtype", ["bf16", "fp16"]) +@pytest.mark.parametrize("model,qkv_format,cp_comm_type", DPA_CP_FUSED_HD256_CASES) +def test_cp_with_fused_attention_hdim256(cp_pool, dtype, model, qkv_format, cp_comm_type): + """Test cuDNN FusedAttention CP with head_dim=256 backward on Blackwell.""" + pool = cp_pool(2) + config = copy.deepcopy(model_configs_fused_attn_hdim256[model]) + config.context_parallel = True + config.cp_comm_type = cp_comm_type + + if qkv_format == "thd": + config.attn_mask_type = "padding_causal" + + dtype_map = {"fp16": torch.float16, "bf16": torch.bfloat16} + available_backends, _, _ = get_available_attention_backends( + config, + qkv_dtype=dtype_map[dtype], + qkv_layout="_".join([qkv_format] * 3), + is_training=True, + deterministic=_deterministic, + ) + + _, fused_attn_supported, _ = available_backends + if fused_attn_supported and config.attn_mask_type in ["causal", "padding_causal"]: + # CP can internally invoke bottom-right aligned fused-attn subcalls for + # causal masks. Check that inner fused-attn backend directly; the CP wrapper + # itself rejects user-facing bottom-right masks. + config_copy = copy.deepcopy(config) + config_copy.context_parallel = False + config_copy.attn_mask_type = config.attn_mask_type + "_bottom_right" + available_backends, _, _ = get_available_attention_backends( + config_copy, + qkv_dtype=dtype_map[dtype], + qkv_layout="_".join([qkv_format] * 3), + is_training=True, + deterministic=_deterministic, + ) + _, fused_attn_supported, _ = available_backends + if not fused_attn_supported: + pytest.skip("No attention backend available.") + + _submit( + pool, + dtype=dtype, + model=model, + qkv_format=qkv_format, + kernel_backend="FusedAttention", + cp_comm_type=cp_comm_type, + fp8_bwd=False, + fp8_dpa=False, + fp8_mha=False, + scaling_mode=None, + f16_O=True, + is_training=True, + deterministic=_deterministic, + log_level=pytest_logging_level, + ) From 3035301940f498faf253e7b67dea0450d19041db Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 8 Jul 2026 04:57:32 +0000 Subject: [PATCH 19/26] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/jax/test_distributed_fused_attn.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py index f2ee2ececc..4add7342b4 100644 --- a/tests/jax/test_distributed_fused_attn.py +++ b/tests/jax/test_distributed_fused_attn.py @@ -400,7 +400,6 @@ def test_softcap_score_mod_with_aux_params_backward( ] - class TestDistributedContextParallelSelfAttn: # TODO(KshitijLakhani): parametrize num_segments_per_seq for all CP tests def impl_test_context_parallel_attn( @@ -692,6 +691,7 @@ def test_context_parallel_ring_attn( window_size=window_size, stripe_size=stripe_size, ) + # CP ring and all-gather tests for D=256 # TODO(KshitijLakhani): Replace this with common-provided fused-attn disable reasons once # they can be surfaced to framework tests. From 06eee353c41aa2eac0740e98f199c25b3a4efbc0 Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani Date: Wed, 8 Jul 2026 15:05:02 -0700 Subject: [PATCH 20/26] nit: clean up comments Signed-off-by: Kshitij Lakhani --- tests/jax/test_fused_attn.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index b09e641896..faa3da5ed4 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -469,7 +469,6 @@ def _get_max_segments_per_sequence(self): return 1 def _check_configs(self): - # TODO(rewang): probably adds this in is_fused_attn_available # TODO(KshitijLakhani): probably add/move this to is_fused_attn_available if self.qkv_layout.is_thd() and not self.attn_mask_type.is_padding(): pytest.skip("THD format requires padding masks.") From ee79662df4bfb32e83d84a866c492170bbdabbff Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani Date: Tue, 21 Jul 2026 13:55:56 -0700 Subject: [PATCH 21/26] Correct the atgs passed to test dpa Signed-off-by: Kshitij Lakhani --- tests/pytorch/attention/test_attention.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index a6c8620a53..6ec3e406ad 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -421,7 +421,7 @@ def test_dpa_fa4_hdim256(dtype, model_configs, model): ) def test_dpa_fused_attn_hdim256(dtype, model_configs, model, qkv_layout): """Test DotProductAttention with cuDNN FusedAttention: head_dim=256 backward on Blackwell""" - test_dot_product_attention(dtype, model_configs, model, False, True, qkv_layout, False, False) + test_dot_product_attention(dtype, model_configs, model, False, qkv_layout, False, False) model_configs_fa4_mla = { From cbab884a46ec21ce1af533f9f7aee90057e6ac05 Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani Date: Wed, 22 Jul 2026 16:24:43 -0700 Subject: [PATCH 22/26] Simplify THD fused attention layout check Signed-off-by: Kshitij Lakhani --- transformer_engine/common/fused_attn/fused_attn.cpp | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 834c1805d5..7d726762e1 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -241,8 +241,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); const bool is_thd_layout = q_format == NVTE_QKV_Format::NVTE_THD || - kv_format == NVTE_QKV_Format::NVTE_THD || - qkv_format == NVTE_QKV_Format::NVTE_THD; + kv_format == NVTE_QKV_Format::NVTE_THD; NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); auto cudnn_runtime_version = cudnnGetVersion(); From 00d14cd04d4651615f58bed919c18ca242fd039b Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani Date: Wed, 22 Jul 2026 16:25:06 -0700 Subject: [PATCH 23/26] Rename PyTorch D256 fused attention tests Signed-off-by: Kshitij Lakhani --- .../attention/run_attention_with_cp.py | 6 +-- tests/pytorch/attention/test_attention.py | 16 +++---- .../attention/test_attention_with_cp.py | 44 +++++++++---------- 3 files changed, 33 insertions(+), 33 deletions(-) diff --git a/tests/pytorch/attention/run_attention_with_cp.py b/tests/pytorch/attention/run_attention_with_cp.py index 7cfe2b3212..6956be323b 100644 --- a/tests/pytorch/attention/run_attention_with_cp.py +++ b/tests/pytorch/attention/run_attention_with_cp.py @@ -18,7 +18,7 @@ from test_attention_with_cp import ( model_configs_flash_attn, model_configs_fused_attn, - model_configs_fused_attn_hdim256, + model_configs_fused_attn_d256, ) from transformer_engine.pytorch import ( autocast, @@ -237,8 +237,8 @@ def run_dpa_with_cp( os.environ["NVTE_FUSED_ATTN"] = "1" if model in model_configs_fused_attn: config = copy.deepcopy(model_configs_fused_attn[model]) - elif model in model_configs_fused_attn_hdim256: - config = copy.deepcopy(model_configs_fused_attn_hdim256[model]) + elif model in model_configs_fused_attn_d256: + config = copy.deepcopy(model_configs_fused_attn_d256[model]) else: assert False, f"{model=} is not a known FusedAttention CP config!" assert config.attn_mask_type in [ diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 6ec3e406ad..0e9a233f2d 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -376,16 +376,16 @@ def test_dpa_fa4_hdim256(dtype, model_configs, model): # SDPA bprop kernel. BSHD support starts with cuDNN FE 1.24 / BE 9.23; THD support starts # with cuDNN FE 1.26 / BE 9.30. The kernel supports d_qk == d_v == 256 only, vanilla softmax only, # no dropout, no ALiBi, and (for non-causal masks) full-window attention only. -model_configs_fused_hdim256 = { +model_configs_fused_d256 = { # test: ModelConfig(b, sq, hq, dqk) -> head_dim_v defaults to head_dim_qk (256) - "fused_hd256_no_mask": ModelConfig(2, 512, 16, 256), - "fused_hd256_padding": ModelConfig(2, 512, 16, 256, attn_mask_type="padding"), + "fused_d256_no_mask": ModelConfig(2, 512, 16, 256), + "fused_d256_padding": ModelConfig(2, 512, 16, 256, attn_mask_type="padding"), # SWA is allowed only together with a causal mask on the D=256 bprop kernel. - "fused_hd256_causal_swa": ModelConfig( + "fused_d256_causal_swa": ModelConfig( 2, 1024, 16, 256, attn_mask_type="causal", window_size=(128, 0) ), # GQA variant (num_gqa_groups < num_heads). - "fused_hd256_padding_causal_gqa": ModelConfig( + "fused_d256_padding_causal_gqa": ModelConfig( 2, 1024, 16, 256, num_gqa_groups=4, attn_mask_type="padding_causal" ), } @@ -396,8 +396,8 @@ def test_dpa_fa4_hdim256(dtype, model_configs, model): reason="cuDNN FusedAttention head_dim=256 backward is Blackwell server (SM100/SM103) only.", ) @pytest.mark.parametrize("dtype", param_types) -@pytest.mark.parametrize("model_configs", [model_configs_fused_hdim256]) -@pytest.mark.parametrize("model", model_configs_fused_hdim256.keys()) +@pytest.mark.parametrize("model_configs", [model_configs_fused_d256]) +@pytest.mark.parametrize("model", model_configs_fused_d256.keys()) @pytest.mark.parametrize( "qkv_layout", [ @@ -419,7 +419,7 @@ def test_dpa_fa4_hdim256(dtype, model_configs, model): ), ], ) -def test_dpa_fused_attn_hdim256(dtype, model_configs, model, qkv_layout): +def test_dpa_fused_attn_d256(dtype, model_configs, model, qkv_layout): """Test DotProductAttention with cuDNN FusedAttention: head_dim=256 backward on Blackwell""" test_dot_product_attention(dtype, model_configs, model, False, qkv_layout, False, False) diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index 75b0306037..636b1166a3 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -483,11 +483,11 @@ def test_cp_with_flash_attention(cp_pool, dtype, model, qkv_format, cp_comm_type } -model_configs_fused_attn_hdim256 = { +model_configs_fused_attn_d256 = { # cuDNN FusedAttention D=256 bprop is Blackwell server only. Keep these # outside the generic CP sweep so the D=256 signal is focused and inexpensive. - "cp_hd256_causal": ModelConfig(2, 1024, 16, 256, attn_mask_type="causal"), - "cp_hd256_causal_swa": ModelConfig( + "cp_d256_causal": ModelConfig(2, 1024, 16, 256, attn_mask_type="causal"), + "cp_d256_causal_swa": ModelConfig( 2, 1024, 16, 256, attn_mask_type="causal", window_size=(128, 0) ), } @@ -710,60 +710,60 @@ def test_cp_with_fused_attention( # Keep these as explicit cases because the CP axes are not independent: # - This runner covers separate-layout CP (bshd_bshd_bshd/thd_thd_thd); KV-packed -# D=256 is covered by test_dpa_fused_attn_hdim256 in test_attention.py. +# D=256 is covered by test_dpa_fused_attn_d256 in test_attention.py. # - THD starts from a causal config and is rewritten to padding_causal by the runner. # - SWA is sampled only through all_gather; p2p does not support SWA. -BSHD_FUSED_HD256_CP_CUDNN_MARK = pytest.mark.skipif( +BSHD_FUSED_D256_CP_CUDNN_MARK = pytest.mark.skipif( get_cudnn_version() < (9, 23, 0), reason="cuDNN 9.23+ is required for BSHD D=256 fused-attn CP backward.", ) -THD_FUSED_HD256_CP_CUDNN_MARK = pytest.mark.skipif( +THD_FUSED_D256_CP_CUDNN_MARK = pytest.mark.skipif( get_cudnn_version() < (9, 30, 0), reason="cuDNN 9.30+ is required for THD D=256 fused-attn CP backward.", ) -DPA_CP_FUSED_HD256_CASES = [ +DPA_CP_FUSED_D256_CASES = [ pytest.param( - "cp_hd256_causal", + "cp_d256_causal", "bshd", "p2p", id="BSHD-P2P-CAUSAL-NO_SWA", - marks=BSHD_FUSED_HD256_CP_CUDNN_MARK, + marks=BSHD_FUSED_D256_CP_CUDNN_MARK, ), pytest.param( - "cp_hd256_causal", + "cp_d256_causal", "bshd", "all_gather", id="BSHD-ALL_GATHER-CAUSAL-NO_SWA", - marks=BSHD_FUSED_HD256_CP_CUDNN_MARK, + marks=BSHD_FUSED_D256_CP_CUDNN_MARK, ), pytest.param( - "cp_hd256_causal_swa", + "cp_d256_causal_swa", "bshd", "all_gather", id="BSHD-ALL_GATHER-CAUSAL-SWA", - marks=BSHD_FUSED_HD256_CP_CUDNN_MARK, + marks=BSHD_FUSED_D256_CP_CUDNN_MARK, ), pytest.param( - "cp_hd256_causal", + "cp_d256_causal", "thd", "p2p", id="THD-P2P-PADDING_CAUSAL-NO_SWA", - marks=THD_FUSED_HD256_CP_CUDNN_MARK, + marks=THD_FUSED_D256_CP_CUDNN_MARK, ), pytest.param( - "cp_hd256_causal", + "cp_d256_causal", "thd", "all_gather", id="THD-ALL_GATHER-PADDING_CAUSAL-NO_SWA", - marks=THD_FUSED_HD256_CP_CUDNN_MARK, + marks=THD_FUSED_D256_CP_CUDNN_MARK, ), pytest.param( - "cp_hd256_causal_swa", + "cp_d256_causal_swa", "thd", "all_gather", id="THD-ALL_GATHER-PADDING_CAUSAL-SWA", - marks=THD_FUSED_HD256_CP_CUDNN_MARK, + marks=THD_FUSED_D256_CP_CUDNN_MARK, ), ] @@ -773,11 +773,11 @@ def test_cp_with_fused_attention( reason="cuDNN FusedAttention head_dim=256 backward is SM100/SM103-only.", ) @pytest.mark.parametrize("dtype", ["bf16", "fp16"]) -@pytest.mark.parametrize("model,qkv_format,cp_comm_type", DPA_CP_FUSED_HD256_CASES) -def test_cp_with_fused_attention_hdim256(cp_pool, dtype, model, qkv_format, cp_comm_type): +@pytest.mark.parametrize("model,qkv_format,cp_comm_type", DPA_CP_FUSED_D256_CASES) +def test_cp_with_fused_attention_d256(cp_pool, dtype, model, qkv_format, cp_comm_type): """Test cuDNN FusedAttention CP with head_dim=256 backward on Blackwell.""" pool = cp_pool(2) - config = copy.deepcopy(model_configs_fused_attn_hdim256[model]) + config = copy.deepcopy(model_configs_fused_attn_d256[model]) config.context_parallel = True config.cp_comm_type = cp_comm_type From ed4b74526fabe40ea327330bbb683ea2e6b7d540 Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani Date: Wed, 22 Jul 2026 16:25:26 -0700 Subject: [PATCH 24/26] Run JAX fused attention before softmax Signed-off-by: Kshitij Lakhani --- qa/L1_jax_distributed_unittest/test.sh | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/qa/L1_jax_distributed_unittest/test.sh b/qa/L1_jax_distributed_unittest/test.sh index 8e0ef2c267..c734567e1d 100644 --- a/qa/L1_jax_distributed_unittest/test.sh +++ b/qa/L1_jax_distributed_unittest/test.sh @@ -31,12 +31,12 @@ python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/py python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_dist_mlp.xml $TE_PATH/tests/jax/test_distributed_layernorm_mlp.py || test_fail "test_distributed_layernorm_mlp.py" +python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_dist_fused_attn.xml $TE_PATH/tests/jax/test_distributed_fused_attn.py || test_fail "test_distributed_fused_attn.py" + # XLA_FLAGS to WAR for test_distributed_softmax issue with NCCL -# TODO(Kshitij): remove when NCCL issue is fixed +# TODO(KshitijLakhani): remove when NCCL issue is fixed XLA_FLAGS="$XLA_FLAGS --xla_gpu_enable_nccl_comm_splitting=false" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_dist_softmax.xml $TE_PATH/tests/jax/test_distributed_softmax.py || test_fail "test_distributed_softmax.py" -python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_dist_fused_attn.xml $TE_PATH/tests/jax/test_distributed_fused_attn.py || test_fail "test_distributed_fused_attn.py" - # NCCL EP multi-process suite. Self-skips on <4 GPUs. TE_PATH=$TE_PATH bash $TE_PATH/tests/jax/multi_process_launch_ep.sh || test_fail "test_multi_process_ep.py" From e3032d3be364daabbfe93b83cd430b61c1c2aeed Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani Date: Wed, 22 Jul 2026 16:29:25 -0700 Subject: [PATCH 25/26] Relabel THD D256 cuDNN guard to 9.25 Signed-off-by: Kshitij Lakhani --- tests/jax/test_distributed_fused_attn.py | 4 ++-- tests/jax/test_fused_attn.py | 8 ++++---- tests/pytorch/attention/test_attention.py | 6 +++--- tests/pytorch/attention/test_attention_with_cp.py | 4 ++-- transformer_engine/common/fused_attn/fused_attn.cpp | 4 ++-- 5 files changed, 13 insertions(+), 13 deletions(-) diff --git a/tests/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py index 4add7342b4..9678195942 100644 --- a/tests/jax/test_distributed_fused_attn.py +++ b/tests/jax/test_distributed_fused_attn.py @@ -701,8 +701,8 @@ def skip_if_d256_cp_unsupported(qkv_layout): if not 100 <= compute_capability < 110: pytest.skip("D=256 CP fused attention is only enabled on Blackwell server GPUs.") - required_cudnn_version = 93000 if qkv_layout.is_thd() else 92300 - required_cudnn_version_label = "9.30" if qkv_layout.is_thd() else "9.23" + required_cudnn_version = 92500 if qkv_layout.is_thd() else 92300 + required_cudnn_version_label = "9.25" if qkv_layout.is_thd() else "9.23" if get_cudnn_version() < required_cudnn_version: pytest.skip( f"D=256 CP fused attention with {qkv_layout} requires cuDNN" diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index faa3da5ed4..1dd92fe181 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -497,7 +497,7 @@ def _check_configs(self): cudnn_version = get_cudnn_version() # D=256 bprop on SM10x uses the deterministic algorithm path only. BSHD support # starts with cuDNN FE 1.24 / BE 9.23; THD execution-plan support starts with - # cuDNN FE 1.26 / BE 9.30. The kernel rejects dBias, dropout, and ALiBi, supports vanilla + # cuDNN FE 1.26 / BE 9.25. The kernel rejects dBias, dropout, and ALiBi, supports vanilla # softmax only, and allows SWA together with a causal mask only. is_sm10x = 100 <= compute_capability < 110 if self.is_training and is_sm10x and (self.head_dim_qk == 256 or self.head_dim_v == 256): @@ -506,8 +506,8 @@ def _check_configs(self): "D=256 BWD on Blackwell only supports d_qk == d_v == 256;" f" got d_qk={self.head_dim_qk}, d_v={self.head_dim_v}." ) - required_cudnn_version = 93000 if self.qkv_layout.is_thd() else 92300 - required_cudnn_version_label = "9.30" if self.qkv_layout.is_thd() else "9.23" + required_cudnn_version = 92500 if self.qkv_layout.is_thd() else 92300 + required_cudnn_version_label = "9.25" if self.qkv_layout.is_thd() else "9.23" if cudnn_version < required_cudnn_version: pytest.skip( f"D=256 BWD on Blackwell with {self.qkv_layout} requires cuDNN" @@ -1694,7 +1694,7 @@ def test_backward( id="2-1024-2048-12-6-128-64-BF16-CROSS-GQA-RAGGED_SEPARATE", ), # D=256 deterministic backward on the SM100 dedicated SDPA bprop kernel. - # BSHD requires cuDNN FE 1.24 / BE 9.23+; THD requires cuDNN FE 1.26 / BE 9.30+. + # BSHD requires cuDNN FE 1.24 / BE 9.23+; THD requires cuDNN FE 1.26 / BE 9.25+. pytest.param( 4, 128, diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 0e9a233f2d..c1ca876366 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -374,7 +374,7 @@ def test_dpa_fa4_hdim256(dtype, model_configs, model): # cuDNN FusedAttention D=256 bprop is supported on sm10x by the dedicated deterministic # SDPA bprop kernel. BSHD support starts with cuDNN FE 1.24 / BE 9.23; THD support starts -# with cuDNN FE 1.26 / BE 9.30. The kernel supports d_qk == d_v == 256 only, vanilla softmax only, +# with cuDNN FE 1.26 / BE 9.25. The kernel supports d_qk == d_v == 256 only, vanilla softmax only, # no dropout, no ALiBi, and (for non-causal masks) full-window attention only. model_configs_fused_d256 = { # test: ModelConfig(b, sq, hq, dqk) -> head_dim_v defaults to head_dim_qk (256) @@ -413,8 +413,8 @@ def test_dpa_fa4_hdim256(dtype, model_configs, model): "thd_t2hd", id="THD_KV_PACKED", marks=pytest.mark.skipif( - get_cudnn_version() < (9, 30, 0), - reason="cuDNN 9.30+ is required for THD D=256 fused-attn backward.", + get_cudnn_version() < (9, 25, 0), + reason="cuDNN 9.25+ is required for THD D=256 fused-attn backward.", ), ), ], diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index 636b1166a3..df4ef1f038 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -718,8 +718,8 @@ def test_cp_with_fused_attention( reason="cuDNN 9.23+ is required for BSHD D=256 fused-attn CP backward.", ) THD_FUSED_D256_CP_CUDNN_MARK = pytest.mark.skipif( - get_cudnn_version() < (9, 30, 0), - reason="cuDNN 9.30+ is required for THD D=256 fused-attn CP backward.", + get_cudnn_version() < (9, 25, 0), + reason="cuDNN 9.25+ is required for THD D=256 fused-attn CP backward.", ) DPA_CP_FUSED_D256_CASES = [ diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 7d726762e1..fd1a300e93 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -331,9 +331,9 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( (head_dim_qk == 192 && head_dim_v == 128 && is_training && sm_arch_ >= 100 && cudnn_runtime_version >= 91100) || // 9.23: d_qk = d_v = 256 + SM10x (cuDNN FE 1.24 / BE 9.23+) + bprop + non-paged. - // THD layouts require cuDNN FE 1.26 / BE 9.30+ for execution-plan support. + // THD layouts require cuDNN FE 1.26 / BE 9.25+ for execution-plan support. (head_dim_qk == 256 && head_dim_v == 256 && is_training && sm_arch_ >= 100 && - sm_arch_ < 110 && cudnn_runtime_version >= (is_thd_layout ? 93000 : 92300) && + sm_arch_ < 110 && cudnn_runtime_version >= (is_thd_layout ? 92500 : 92300) && layout_group != NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD && // The FE forces this path onto the deterministic bprop algorithm, which on // Blackwell rejects dBias, dropout, and ALiBi (and supports vanilla softmax only). From 6f2b74825aab880c04827796a59e7d3b7747728a Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 23 Jul 2026 00:23:44 +0000 Subject: [PATCH 26/26] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/common/fused_attn/fused_attn.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index fd1a300e93..1ac7a36383 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -240,8 +240,8 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); - const bool is_thd_layout = q_format == NVTE_QKV_Format::NVTE_THD || - kv_format == NVTE_QKV_Format::NVTE_THD; + const bool is_thd_layout = + q_format == NVTE_QKV_Format::NVTE_THD || kv_format == NVTE_QKV_Format::NVTE_THD; NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); auto cudnn_runtime_version = cudnnGetVersion();