[JAX] [PyT] [Common] Enable D=256 BWD cuDNN fused attn for Blackwell CC 10.x

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"
 

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,59 @@ 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
@@ -638,6 +692,119 @@ def test_context_parallel_ring_attn(
             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 = 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"
+                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 = {
     "L0": [[]],

tests/jax/test_fused_attn.py

@@ -469,7 +469,7 @@ 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.")
 
@@ -493,11 +493,62 @@ def _check_configs(self):
             pytest.skip(
                 "seqlen_q > seqlen_kv is not supported with sliding window attention in cuDNN"
             )
+        compute_capability = get_device_compute_capability(0)
+        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.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):
+            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}."
+                )
+            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"
+                    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.
+            # 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:
+                unsupported = (
+                    "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"
+            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."
+                    )
 
-        if get_device_compute_capability(0) >= 100 and self.is_training:
+        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 get_cudnn_version() < 90700
+                or cudnn_version < 90700
             ):
                 pytest.skip(
                     "For sm100+, non-deterministic bprop (cuDNN 9.7+) does not support bias with"
@@ -506,7 +557,7 @@ def _check_configs(self):
             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
+                or cudnn_version < 91801
             ):
                 pytest.skip(
                     "For sm100+, deterministic bprop (cuDNN 9.18.1+) does not support bias or"
@@ -1642,6 +1693,32 @@ 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.
+        # BSHD requires cuDNN FE 1.24 / BE 9.23+; THD requires cuDNN FE 1.26 / BE 9.25+.
+        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.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",
+        ),
     ],
 )
 @pytest.mark.parametrize(

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_d256,
+)
 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_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 [
         "causal",
         "no_mask",