diff --git a/qa/L3_pytorch_FA_versions_test/test.sh b/qa/L3_pytorch_FA_versions_test/test.sh index e7fd3189f35..13ba92f56b9 100644 --- a/qa/L3_pytorch_FA_versions_test/test.sh +++ b/qa/L3_pytorch_FA_versions_test/test.sh @@ -36,8 +36,8 @@ export FLASH_ATTN_CUDA_ARCHS=$sm_arch CP_FA_VERSION="" if [ $sm_arch -gt 90 ] then - FA_versions=(4.0.0b11) - CP_FA_VERSION="4.0.0b11" + FA_versions=(4.0.0b31) + CP_FA_VERSION="4.0.0b31" elif [ $sm_arch -eq 90 ] then FA_versions=(3.0.0b1) @@ -66,13 +66,13 @@ do elif [[ "${fa_version}" == 4.* ]] then export NVTE_FLASH_ATTN_V4=1 - # FA4 is intentionally last in every version array. Its b11 test pin needs - # CUTLASS DSL 4.4.2, so replace the image-matched stack only for this final + # FA4 is intentionally last in every version array. Its b31 test pin needs + # CUTLASS DSL 4.6.2, so replace the image-matched stack only for this final # iteration; later iterations would otherwise need that stack restored. pip3 uninstall -y nvidia-cutlass-dsl nvidia-cutlass-dsl-libs-base \ nvidia-cutlass-dsl-libs-cu12 nvidia-cutlass-dsl-libs-cu13 \ || error_exit "Failed to isolate CUTLASS DSL for Flash Attention $fa_version" - pip3 install flash-attn-4==${fa_version} nvidia-cutlass-dsl[cu13]==4.4.2 \ + pip3 install flash-attn-4==${fa_version} nvidia-cutlass-dsl[cu13]==4.6.2 \ --no-build-isolation || error_exit "Failed to install Flash Attention $fa_version" else export NVTE_FLASH_ATTN_V3=1 diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 9a339233a42..59dbff646f8 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -6,6 +6,7 @@ from contextlib import nullcontext from dataclasses import dataclass +import functools from importlib.metadata import version as get_pkg_version from importlib.metadata import PackageNotFoundError import inspect @@ -102,6 +103,25 @@ _flash_attn_bwd = None _flash_attn_varlen_fwd = None _flash_attn_varlen_bwd = None + + +def _normalize_fa4_window_kwargs(function): + """Translate TE's negative unbounded-window sentinel to the FA4 API.""" + + @functools.wraps(function) + def wrapper(*args, **kwargs): + if kwargs.get("window_size") is not None: + kwargs["window_size"] = tuple( + None if bound == -1 else bound for bound in kwargs["window_size"] + ) + for name in ("window_size_left", "window_size_right"): + if kwargs.get(name) == -1: + kwargs[name] = None + return function(*args, **kwargs) + + return wrapper + + # Try to import Flash Attention v2 try: fa_utils.version = PkgVersion(get_pkg_version("flash-attn")) @@ -222,6 +242,12 @@ else: # Unlike versions 2 and 3, FlashAttention 4 registers no custom ops: it builds # its kernels through the CUTLASS DSL as it runs. Keep it an eager island. + # FA4 b31 changed its unbounded-window sentinel from -1 to None. + if fa_utils.fa4_version >= fa_utils.v4_0_0_beta31: + _flash_attn_func_v4 = _normalize_fa4_window_kwargs(_flash_attn_func_v4) + _flash_attn_varlen_func_v4 = _normalize_fa4_window_kwargs(_flash_attn_varlen_func_v4) + _flash_attn_fwd_v4 = _normalize_fa4_window_kwargs(_flash_attn_fwd_v4) + _flash_attn_bwd_v4 = _normalize_fa4_window_kwargs(_flash_attn_bwd_v4) flash_attn_func_v4 = no_torch_dynamo()(_flash_attn_func_v4) flash_attn_varlen_func_v4 = no_torch_dynamo()(_flash_attn_varlen_func_v4) _flash_attn_fwd_v4 = no_torch_dynamo()(_flash_attn_fwd_v4) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 510d14ac638..b2bdd495b2b 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -3765,7 +3765,14 @@ def forward( thd_cu_seqlens_q_per_step[i] if qkv_format == "thd" else cu_seqlens_q ) fa_cu_seqlens_kv = cu_seqlens_kv_per_step[i] - if (use_flash_attn_3 or use_flash_attn_4) and qkv_format == "thd": + if ( + (use_flash_attn_3 or use_flash_attn_4) + and qkv_format == "thd" + # NO_LOAD_BALANCE buffers are already compact, so their + # cu_seqlens carry both offsets and sequence lengths. + and load_balancing_strategy + is not CPLoadBalancingStrategy.NO_LOAD_BALANCE + ): seqused_q = ( thd_cu_seqlens_q_per_step[i][1:] - thd_cu_seqlens_q_per_step[i][:-1] ) @@ -4369,8 +4376,13 @@ def backward(ctx, dout, *_args): ) fa_cu_seqlens_kv = cu_seqlens_kv_per_step[i] if ( - ctx.use_flash_attn_3 or ctx.use_flash_attn_4 - ) and ctx.qkv_format == "thd": + (ctx.use_flash_attn_3 or ctx.use_flash_attn_4) + and ctx.qkv_format == "thd" + and ( + ctx.load_balancing_strategy + is not CPLoadBalancingStrategy.NO_LOAD_BALANCE + ) + ): seqused_q = ( thd_cu_seqlens_q_per_step[i][1:] - thd_cu_seqlens_q_per_step[i][:-1] ) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 23b89287dfd..4eea60406e9 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -148,11 +148,12 @@ class FlashAttentionUtils: v4_is_installed = False fa4_version = PkgVersion("0") + v4_0_0_beta31 = PkgVersion("4.0.0b31") use_v4 = False # Set by a signature probe in backends.py; fail-closed default. fa3_supports_softcap = False v4_installation_steps = """\ -pip install flash-attn-4==4.0.0b11 nvidia-cutlass-dsl[cu13]""" +pip install flash-attn-4==4.0.0b31 nvidia-cutlass-dsl[cu13]==4.6.2""" v4_warning_printed = False # Set by backends.py if FA4 is installed; calls flash_attn.cute.interface._validate_head_dims # which raises AssertionError for unsupported (head_dim, head_dim_v) combinations. @@ -1279,6 +1280,22 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt cp_comm_type, ) use_flash_attention_4 = False + elif ( + qkv_format == "thd" + and cp_comm_type == "all_gather" + and (10, 0) <= device_compute_capability < (12, 0) + and head_dim_qk == head_dim_v == 256 + and FlashAttentionUtils.fa4_version < FlashAttentionUtils.v4_0_0_beta31 + ): + # Earlier FA4 releases predate the complete D=256 varlen support + # needed by THD all-gather CP, so compact metadata can be incorrect. + logger.debug( + "Disabling FlashAttention 4 for THD all-gather context parallelism with " + "head_dim=256 on SM100/SM110 with version %s (requires >= %s)", + FlashAttentionUtils.fa4_version, + FlashAttentionUtils.v4_0_0_beta31, + ) + use_flash_attention_4 = False if context_parallel and ( use_flash_attention_2 or use_flash_attention_3 or use_flash_attention_4 ): @@ -1757,6 +1774,19 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt head_dim_v, ) use_flash_attention_3 = False + if use_flash_attention_4 and deterministic and FlashAttentionUtils.v4_is_installed: + if ( + is_training + and (10, 0) <= device_compute_capability < (12, 0) + and head_dim_qk == head_dim_v == 256 + ): + # FA4's dedicated SM100/SM110 D=256 backward kernel rejects + # deterministic execution, so select another backend before launch. + logger.debug( + "Disabling FlashAttention 4 for deterministic backward with " + "head_dim=256 on SM100/SM110." + ) + use_flash_attention_4 = False if use_fused_attention and deterministic: if softmax_type != "vanilla": logger.debug(