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" diff --git a/tests/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py index 2abd9824b6..9678195942 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,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": [[]], diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 0d1db0b9e1..1dd92fe181 100644 --- a/tests/jax/test_fused_attn.py +++ b/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( diff --git a/tests/pytorch/attention/run_attention_with_cp.py b/tests/pytorch/attention/run_attention_with_cp.py index 82b9df262f..6956be323b 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_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", diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index f2ab4a8495..c1ca876366 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,6 +372,58 @@ 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 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.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) + "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_d256_causal_swa": ModelConfig( + 2, 1024, 16, 256, attn_mask_type="causal", window_size=(128, 0) + ), + # GQA variant (num_gqa_groups < num_heads). + "fused_d256_padding_causal_gqa": ModelConfig( + 2, 1024, 16, 256, num_gqa_groups=4, attn_mask_type="padding_causal" + ), +} + + +@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_d256]) +@pytest.mark.parametrize("model", model_configs_fused_d256.keys()) +@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, 25, 0), + reason="cuDNN 9.25+ is required for THD D=256 fused-attn backward.", + ), + ), + ], +) +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) + + 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), diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index 681ee5e6e0..df4ef1f038 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_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_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) + ), +} + + 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_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_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_D256_CP_CUDNN_MARK = pytest.mark.skipif( + 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 = [ + pytest.param( + "cp_d256_causal", + "bshd", + "p2p", + id="BSHD-P2P-CAUSAL-NO_SWA", + marks=BSHD_FUSED_D256_CP_CUDNN_MARK, + ), + pytest.param( + "cp_d256_causal", + "bshd", + "all_gather", + id="BSHD-ALL_GATHER-CAUSAL-NO_SWA", + marks=BSHD_FUSED_D256_CP_CUDNN_MARK, + ), + pytest.param( + "cp_d256_causal_swa", + "bshd", + "all_gather", + id="BSHD-ALL_GATHER-CAUSAL-SWA", + marks=BSHD_FUSED_D256_CP_CUDNN_MARK, + ), + pytest.param( + "cp_d256_causal", + "thd", + "p2p", + id="THD-P2P-PADDING_CAUSAL-NO_SWA", + marks=THD_FUSED_D256_CP_CUDNN_MARK, + ), + pytest.param( + "cp_d256_causal", + "thd", + "all_gather", + id="THD-ALL_GATHER-PADDING_CAUSAL-NO_SWA", + marks=THD_FUSED_D256_CP_CUDNN_MARK, + ), + pytest.param( + "cp_d256_causal_swa", + "thd", + "all_gather", + id="THD-ALL_GATHER-PADDING_CAUSAL-SWA", + marks=THD_FUSED_D256_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_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_d256[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, + ) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index fc21771297..1ac7a36383 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -240,6 +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; NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); auto cudnn_runtime_version = cudnnGetVersion(); @@ -327,7 +329,23 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( 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)) && + 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.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 ? 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). + 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. + ((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.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 &&