diff --git a/qa/L3_pytorch_FA_versions_test/test.sh b/qa/L3_pytorch_FA_versions_test/test.sh index 047358b301..e7fd3189f3 100644 --- a/qa/L3_pytorch_FA_versions_test/test.sh +++ b/qa/L3_pytorch_FA_versions_test/test.sh @@ -31,12 +31,13 @@ export NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 # Iterate over Flash Attention versions sm_arch=`python3 -c "import torch; sm = torch.cuda.get_device_capability(0); print(sm[0]*10+sm[1])"` export FLASH_ATTN_CUDA_ARCHS=$sm_arch -# Run one architecture-owned FlashAttention generation. CP remains FA3-only -# until the production selector and runner support FA4 CP end to end. +# Run one architecture-owned FlashAttention generation. sm90 keeps the mature +# FA3 path, while sm100+ exercises the FA4 CP support enabled by this change. CP_FA_VERSION="" if [ $sm_arch -gt 90 ] then FA_versions=(4.0.0b11) + CP_FA_VERSION="4.0.0b11" elif [ $sm_arch -eq 90 ] then FA_versions=(3.0.0b1) diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index d7eb16b862..c8b01e32a2 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -300,8 +300,12 @@ def _submit(pool: PoolWorker, **kwargs) -> None: @pytest.mark.skipif( - not (FlashAttentionUtils.v2_plus or FlashAttentionUtils.v3_is_installed), - reason="Flash-attn v2 or v3 is required.", + not ( + FlashAttentionUtils.v2_plus + or FlashAttentionUtils.v3_is_installed + or FlashAttentionUtils.v4_is_installed + ), + reason="Flash-attn v2, v3, or v4 is required.", ) @pytest.mark.skipif(get_device_compute_capability() < (8, 0), reason="CP tests require sm80+.") @pytest.mark.parametrize("dtype", dtypes) @@ -316,16 +320,10 @@ def test_cp_with_flash_attention(cp_pool, dtype, model, qkv_format, cp_comm_type if pad_between_seqs: if qkv_format != "thd": pytest.skip("pad_between_seqs only applies to THD format!") - if not FlashAttentionUtils.v3_is_installed or get_device_compute_capability() > (9, 0): - pytest.skip("pad_between_seqs with CP requires Flash Attention v3 on Hopper (sm90)!") - if cp_comm_type == "a2a+p2p": - pytest.skip("pad_between_seqs is not yet supported with A2A+P2P CP comm type!") - - if pad_between_seqs: - if qkv_format != "thd": - pytest.skip("pad_between_seqs only applies to THD format!") - if not FlashAttentionUtils.v3_is_installed: - pytest.skip("pad_between_seqs with CP requires Flash Attention v3!") + has_fa3 = FlashAttentionUtils.v3_is_installed and get_device_compute_capability() == (9, 0) + has_fa4 = FlashAttentionUtils.v4_is_installed + if not (has_fa3 or has_fa4): + pytest.skip("pad_between_seqs with CP requires Flash Attention v3 on Hopper or v4!") if cp_comm_type == "a2a+p2p": pytest.skip("pad_between_seqs is not yet supported with A2A+P2P CP comm type!") @@ -342,9 +340,10 @@ def test_cp_with_flash_attention(cp_pool, dtype, model, qkv_format, cp_comm_type qkv_format == "thd" and cp_comm_type == "all_gather" and not FlashAttentionUtils.v3_is_installed + and not FlashAttentionUtils.v4_is_installed ): pytest.skip( - "THD + all_gather requires FA3 (seqused_k) to separate tensor offsets from" + "THD + all_gather requires FA3 or FA4 (seqused_k) to separate tensor offsets from" " visibility limits in the gathered KV buffer." ) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 8a219a6a4d..f340d8d8cd 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -171,6 +171,8 @@ except PackageNotFoundError: flash_attn_func_v4 = None flash_attn_varlen_func_v4 = None + _flash_attn_fwd_v4 = None + _flash_attn_bwd_v4 = None else: try: cutlass_dsl_version = PkgVersion(get_pkg_version("nvidia-cutlass-dsl")) @@ -188,10 +190,14 @@ flash_attn_func as _flash_attn_func_v4, flash_attn_varlen_func as _flash_attn_varlen_func_v4, _validate_head_dims as _fa4_validate_head_dims, + _flash_attn_fwd as _flash_attn_fwd_v4, + _flash_attn_bwd as _flash_attn_bwd_v4, ) except ImportError as exc: flash_attn_func_v4 = None flash_attn_varlen_func_v4 = None + _flash_attn_fwd_v4 = None + _flash_attn_bwd_v4 = None warnings.warn( f"FlashAttention 4 is installed but cannot be loaded: {exc}", RuntimeWarning, @@ -202,6 +208,8 @@ # its kernels through the CUTLASS DSL as it runs. Keep it an eager island. 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) + _flash_attn_bwd_v4 = no_torch_dynamo()(_flash_attn_bwd_v4) fa_utils.v4_validate_head_dims = _fa4_validate_head_dims fa_utils.set_flash_attention_4_params() @@ -1133,6 +1141,7 @@ def forward( quantizers=quantizers, pad_between_seqs=pad_between_seqs, use_flash_attn_3=use_flash_attn_3, + use_flash_attn_4=use_flash_attn_4, fp8_output=fp8_output, ) else: @@ -1192,11 +1201,21 @@ def forward( if inference_params is None: fa_4_optional_forward_kwargs["deterministic"] = self.deterministic if func is flash_attn_varlen_func_v4: - cu_q, cu_kv = _unalias_cu_seqlens(cu_seqlens_q, cu_seqlens_kv) + cu_q, cu_kv = _unalias_cu_seqlens( + cu_seqlens_q_padded if pad_between_seqs else cu_seqlens_q, + cu_seqlens_kv_padded if pad_between_seqs else cu_seqlens_kv, + ) fa_4_optional_forward_kwargs["cu_seqlens_q"] = cu_q fa_4_optional_forward_kwargs["cu_seqlens_k"] = cu_kv fa_4_optional_forward_kwargs["max_seqlen_q"] = max_seqlen_q fa_4_optional_forward_kwargs["max_seqlen_k"] = max_seqlen_kv + if pad_between_seqs: + fa_4_optional_forward_kwargs["seqused_q"] = ( + cu_seqlens_q[1:] - cu_seqlens_q[:-1] + ) + fa_4_optional_forward_kwargs["seqused_k"] = ( + cu_seqlens_kv[1:] - cu_seqlens_kv[:-1] + ) output = func( query_layer, key_layer, 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 ea89ca97eb..eeb99bbaaa 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -614,8 +614,28 @@ def get_fa_args( dv=None, seqused_q=None, seqused_k=None, + use_flash_attn_4: bool = False, ): - """Get forward/backward arguments for flash-attn v2 and v3.""" + """Get positional FA2/FA3 arguments or FA4 keyword arguments. + + FA2/FA3 use version-specific positional layouts, while FA4's raw backward + API orders its optional metadata differently and returns the gradients. + """ + if use_flash_attn_4: + fa_kwargs = {} + if qkv_format == "thd": + fa_kwargs.update( + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_kv, + seqused_q=seqused_q, + seqused_k=seqused_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_kv, + ) + if not forward: + fa_kwargs.update(dq=dq, dk=dk, dv=dv) + return fa_kwargs + if use_flash_attn_3: if forward: if qkv_format == "thd": @@ -968,6 +988,7 @@ def cp_p2p_fwd_fused_attn( def cp_p2p_fwd_flash_attn( use_flash_attn_3, + use_flash_attn_4, qkv_format, fa_forward_kwargs, flash_attn_fwd, @@ -996,15 +1017,20 @@ def cp_p2p_fwd_flash_attn( elif section == "upper-triangle": max_seqlen_q_ = max_seqlen_q // 2 if section in ["lower-triangle", "upper-triangle"]: - if fa_utils.v2_3_plus and not fa_utils.v2_7_0_plus: + if ( + not use_flash_attn_3 + and not use_flash_attn_4 + and fa_utils.v2_3_plus + and not fa_utils.v2_7_0_plus + ): fa_forward_kwargs["window_size"] = (-1, -1) - elif use_flash_attn_3 or fa_utils.v2_7_0_plus: + elif use_flash_attn_3 or use_flash_attn_4 or fa_utils.v2_7_0_plus: fa_forward_kwargs["window_size_left"] = -1 fa_forward_kwargs["window_size_right"] = -1 seqused_q = None seqused_k = None - if pad_between_seqs and use_flash_attn_3 and qkv_format == "thd": + if pad_between_seqs and (use_flash_attn_3 or use_flash_attn_4) and qkv_format == "thd": # Derive actual token counts per batch element from cu_seqlens seqused_q = cu_seqlens_q_per_step[1:] - cu_seqlens_q_per_step[:-1] seqused_k = cu_seqlens_kv_per_step[1:] - cu_seqlens_kv_per_step[:-1] @@ -1016,27 +1042,51 @@ def cp_p2p_fwd_flash_attn( elif section == "upper-triangle": cu_seqlens_q_ = cu_seqlens_q_padded // 2 - fa_forward_args_thd = get_fa_args( - True, - use_flash_attn_3, - qkv_format, - cu_seqlens_q=cu_seqlens_q_, - cu_seqlens_kv=cu_seqlens_kv_, - max_seqlen_q=max_seqlen_q_, - max_seqlen_kv=max_seqlen_kv_, - seqused_q=seqused_q, - seqused_k=seqused_k, - ) - fa_outputs = flash_attn_fwd( - q_part, - k_part, - v_part, - *fa_forward_args_thd, - causal=causal_, - **fa_forward_kwargs, - ) + if use_flash_attn_4: + fa_outputs = flash_attn_fwd( + q_part, + k_part, + v_part, + **get_fa_args( + True, + False, + qkv_format, + cu_seqlens_q=cu_seqlens_q_, + cu_seqlens_kv=cu_seqlens_kv_, + max_seqlen_q=max_seqlen_q_, + max_seqlen_kv=max_seqlen_kv_, + seqused_q=seqused_q, + seqused_k=seqused_k, + use_flash_attn_4=True, + ), + causal=causal_, + **fa_forward_kwargs, + ) + else: + fa_forward_args_thd = get_fa_args( + True, + use_flash_attn_3, + qkv_format, + cu_seqlens_q=cu_seqlens_q_, + cu_seqlens_kv=cu_seqlens_kv_, + max_seqlen_q=max_seqlen_q_, + max_seqlen_kv=max_seqlen_kv_, + seqused_q=seqused_q, + seqused_k=seqused_k, + ) + fa_outputs = flash_attn_fwd( + q_part, + k_part, + v_part, + *fa_forward_args_thd, + causal=causal_, + **fa_forward_kwargs, + ) rng_states = None - if not use_flash_attn_3 and not fa_utils.v2_7_0_plus: + if use_flash_attn_4: + out_per_step = fa_outputs[0] + softmax_lse_per_step = fa_outputs[1] + elif not use_flash_attn_3 and not fa_utils.v2_7_0_plus: out_per_step = fa_outputs[4] softmax_lse_per_step = fa_outputs[5] rng_states = fa_outputs[7] @@ -1256,6 +1306,7 @@ def cp_p2p_bwd_fused_attn( def cp_p2p_bwd_flash_attn( use_flash_attn_3, + use_flash_attn_4, qkv_format, max_seqlen_q, max_seqlen_kv, @@ -1283,21 +1334,31 @@ def cp_p2p_bwd_flash_attn( dq, dk, dv = [torch.zeros_like(x) for x in [q_part, k_part, v_part]] else: dq, dk, dv = [torch.empty_like(x) for x in [q_part, k_part, v_part]] - if fa_utils.v2_3_plus and not fa_utils.v2_7_0_plus: + if ( + not use_flash_attn_3 + and not use_flash_attn_4 + and fa_utils.v2_3_plus + and not fa_utils.v2_7_0_plus + ): fa_backward_kwargs["window_size"] = (-1, -1) - elif use_flash_attn_3 or fa_utils.v2_7_0_plus: + elif use_flash_attn_3 or use_flash_attn_4 or fa_utils.v2_7_0_plus: fa_backward_kwargs["window_size_left"] = -1 fa_backward_kwargs["window_size_right"] = -1 - if not use_flash_attn_3: + if not use_flash_attn_3 and not use_flash_attn_4: fa_backward_kwargs["rng_state"] = rng_states[cp_size - step - 1] max_seqlen_q_ = max_seqlen_q max_seqlen_kv_ = max_seqlen_kv softmax_lse__ = softmax_lse causal_ = False if section == "diagonal": - if fa_utils.v2_3_plus and not fa_utils.v2_7_0_plus: + if ( + not use_flash_attn_3 + and not use_flash_attn_4 + and fa_utils.v2_3_plus + and not fa_utils.v2_7_0_plus + ): fa_backward_kwargs["window_size"] = (-1, 0) - elif use_flash_attn_3 or fa_utils.v2_7_0_plus: + elif use_flash_attn_3 or use_flash_attn_4 or fa_utils.v2_7_0_plus: fa_backward_kwargs["window_size_left"] = -1 fa_backward_kwargs["window_size_right"] = 0 causal_ = True @@ -1311,7 +1372,7 @@ def cp_p2p_bwd_flash_attn( seqused_k = None cu_seqlens_q_bwd = cu_seqlens_q_per_step[cp_size - step - 1] cu_seqlens_kv_bwd = cu_seqlens_kv_per_step[cp_size - step - 1] - if pad_between_seqs and use_flash_attn_3 and qkv_format == "thd": + if pad_between_seqs and (use_flash_attn_3 or use_flash_attn_4) and qkv_format == "thd": seqused_q = cu_seqlens_q_bwd[1:] - cu_seqlens_q_bwd[:-1] seqused_k = cu_seqlens_kv_bwd[1:] - cu_seqlens_kv_bwd[:-1] cu_seqlens_q_bwd = cu_seqlens_q_padded @@ -1321,34 +1382,64 @@ def cp_p2p_bwd_flash_attn( elif section == "upper-triangle": cu_seqlens_q_bwd = cu_seqlens_q_padded // 2 - fa_backward_args_thd = get_fa_args( - False, - use_flash_attn_3, - qkv_format, - cu_seqlens_q=cu_seqlens_q_bwd, - cu_seqlens_kv=cu_seqlens_kv_bwd, - max_seqlen_q=max_seqlen_q_, - max_seqlen_kv=max_seqlen_kv_, - dq=dq, - dk=dk, - dv=dv, - seqused_q=seqused_q, - seqused_k=seqused_k, - ) + if use_flash_attn_4: + fa_backward_kwargs.update( + get_fa_args( + False, + False, + qkv_format, + cu_seqlens_q=cu_seqlens_q_bwd, + cu_seqlens_kv=cu_seqlens_kv_bwd, + max_seqlen_q=max_seqlen_q_, + max_seqlen_kv=max_seqlen_kv_, + dq=dq, + dk=dk, + dv=dv, + seqused_q=seqused_q, + seqused_k=seqused_k, + use_flash_attn_4=True, + ) + ) + else: + fa_backward_args_thd = get_fa_args( + False, + use_flash_attn_3, + qkv_format, + cu_seqlens_q=cu_seqlens_q_bwd, + cu_seqlens_kv=cu_seqlens_kv_bwd, + max_seqlen_q=max_seqlen_q_, + max_seqlen_kv=max_seqlen_kv_, + dq=dq, + dk=dk, + dv=dv, + seqused_q=seqused_q, + seqused_k=seqused_k, + ) if use_flash_attn_3: fa_backward_kwargs["is_causal"] = causal_ else: fa_backward_kwargs["causal"] = causal_ - flash_attn_bwd( - dout_part, - q_part, - k_part, - v_part, - out_part, - softmax_lse__, - *fa_backward_args_thd, - **fa_backward_kwargs, - ) + if use_flash_attn_4: + dq, dk, dv = flash_attn_bwd( + q_part, + k_part, + v_part, + out_part, + dout_part, + softmax_lse__, + **fa_backward_kwargs, + ) + else: + flash_attn_bwd( + dout_part, + q_part, + k_part, + v_part, + out_part, + softmax_lse__, + *fa_backward_args_thd, + **fa_backward_kwargs, + ) return dq, dk, dv @@ -1395,6 +1486,7 @@ def forward( quantizers, pad_between_seqs, use_flash_attn_3, + use_flash_attn_4, fp8_output, layer_number, ): @@ -1622,14 +1714,25 @@ def forward( 0, ) and get_device_compute_capability() != (12, 0) else: - softmax_lse_in_packed_format = fa_utils.v2_6_0_plus or use_flash_attn_3 + softmax_lse_in_packed_format = ( + fa_utils.v2_6_0_plus or use_flash_attn_3 or use_flash_attn_4 + ) # set up args for FlashAttention backend flash_attn_fwd = None fa_forward_kwargs = {} if not use_fused_attention: fa_forward_kwargs = {"softmax_scale": softmax_scale} - if use_flash_attn_3: + if use_flash_attn_4: + from transformer_engine.pytorch.attention.dot_product_attention.backends import ( + _flash_attn_fwd_v4, + ) + + flash_attn_fwd = _flash_attn_fwd_v4 + fa_forward_kwargs["window_size_left"] = -1 + fa_forward_kwargs["window_size_right"] = 0 if causal else -1 + fa_forward_kwargs["return_lse"] = True + elif use_flash_attn_3: from transformer_engine.pytorch.attention.dot_product_attention.backends import ( _flash_attn_fwd_v3, ) @@ -1768,6 +1871,7 @@ def forward( else: flash_attn_inputs = [ use_flash_attn_3, + use_flash_attn_4, qkv_format, fa_forward_kwargs, flash_attn_fwd, @@ -2164,6 +2268,7 @@ def forward( ctx.is_input_fp8 = is_input_fp8 ctx.is_output_fp8 = is_output_fp8 ctx.use_flash_attn_3 = use_flash_attn_3 + ctx.use_flash_attn_4 = use_flash_attn_4 ctx.orig_q_shape = orig_q_shape ctx.orig_k_shape = orig_k_shape @@ -2426,7 +2531,14 @@ def backward(ctx, dout, *_args): flash_attn_bwd = None if not ctx.use_fused_attention: fa_backward_kwargs = {"softmax_scale": ctx.softmax_scale} - if ctx.use_flash_attn_3: + if ctx.use_flash_attn_4: + from transformer_engine.pytorch.attention.dot_product_attention.backends import ( + _flash_attn_bwd_v4, + ) + + flash_attn_bwd = _flash_attn_bwd_v4 + fa_backward_kwargs["deterministic"] = ctx.deterministic + elif ctx.use_flash_attn_3: from transformer_engine.pytorch.attention.dot_product_attention.backends import ( _flash_attn_bwd_v3, ) @@ -2557,6 +2669,7 @@ def backward(ctx, dout, *_args): else: flash_attn_inputs = [ ctx.use_flash_attn_3, + ctx.use_flash_attn_4, ctx.qkv_format, ctx.max_seqlen_q, ctx.max_seqlen_kv, @@ -2970,6 +3083,7 @@ def backward(ctx, dout, *_args): None, None, None, + None, ) @@ -3022,7 +3136,7 @@ class AttnFuncWithCPAndKVAllGather(torch.autograd.Function): FusedAttention carries this split with ``cu_seqlens`` plus ``cu_seqlens_padded``; FlashAttention v3 uses layout ``cu_seqlens`` plus ``seqused_k``. FlashAttention v2 cannot represent both values, so THD - all-gather is restricted to FusedAttention or FlashAttention v3. + all-gather is restricted to FusedAttention, FlashAttention v3, or FlashAttention v4. """ @staticmethod @@ -3051,6 +3165,7 @@ def forward( cp_group, cp_stream, use_flash_attn_3, + use_flash_attn_4, pad_between_seqs, fp8, fp8_meta, @@ -3088,10 +3203,12 @@ def forward( or window_size == (-1, -1) or use_fused_attention or use_flash_attn_3 + or use_flash_attn_4 or fa_utils.v2_3_plus ), ( "cp_comm_type='all_gather' only supports SWA through FusedAttention or FlashAttention" f" >= 2.3. Found {use_fused_attention=}, {use_flash_attn_3=}, " + f"{use_flash_attn_4=}, " f"and {fa_utils.v2_3_plus=}." ) assert q.shape[seq_dim_qkv] % 2 == 0 and k.shape[seq_dim_qkv] % 2 == 0, ( @@ -3102,7 +3219,14 @@ def forward( flash_attn_fwd = None if not use_fused_attention: fa_forward_kwargs = {"softmax_scale": softmax_scale} - if use_flash_attn_3: + if use_flash_attn_4: + from transformer_engine.pytorch.attention.dot_product_attention.backends import ( + _flash_attn_fwd_v4, + ) + + flash_attn_fwd = _flash_attn_fwd_v4 + fa_forward_kwargs["return_lse"] = True + elif use_flash_attn_3: from transformer_engine.pytorch.attention.dot_product_attention.backends import ( _flash_attn_fwd_v3, ) @@ -3341,7 +3465,7 @@ def forward( # FA3 uses internal per-call workspace. Consecutive AG per-step # calls are serialized on GPU streams so that workspace lifetimes # do not overlap. FusedAttention keeps the existing per-step overlap. - if i > 0 and use_flash_attn_3: + if i > 0 and (use_flash_attn_3 or use_flash_attn_4): flash_attn_streams[i].wait_stream(flash_attn_streams[i - 1]) with torch.cuda.stream(flash_attn_streams[i]): new_qkv_layout = qkv_layout @@ -3460,7 +3584,7 @@ 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 and qkv_format == "thd": + if (use_flash_attn_3 or use_flash_attn_4) and qkv_format == "thd": seqused_q = ( thd_cu_seqlens_q_per_step[i][1:] - thd_cu_seqlens_q_per_step[i][:-1] ) @@ -3469,31 +3593,60 @@ def forward( ) fa_cu_seqlens_q = thd_cu_seqlens_q_padded_per_step[i] fa_cu_seqlens_kv = cu_seqlens_kv_padded - fa_forward_args_thd = get_fa_args( - True, - use_flash_attn_3, - qkv_format, - cu_seqlens_q=fa_cu_seqlens_q, - cu_seqlens_kv=fa_cu_seqlens_kv, - max_seqlen_q=max_seqlen_q, - max_seqlen_kv=max_seqlen_kv_, - seqused_q=seqused_q, - seqused_k=seqused_k, - ) - if fa_utils.v2_3_plus and not fa_utils.v2_7_0_plus: + if ( + not use_flash_attn_3 + and not use_flash_attn_4 + and fa_utils.v2_3_plus + and not fa_utils.v2_7_0_plus + ): fa_forward_kwargs["window_size"] = window_size_per_step[i] - elif use_flash_attn_3 or fa_utils.v2_7_0_plus: + elif use_flash_attn_3 or use_flash_attn_4 or fa_utils.v2_7_0_plus: fa_forward_kwargs["window_size_left"] = window_size_per_step[i][0] fa_forward_kwargs["window_size_right"] = window_size_per_step[i][1] - fa_outputs = flash_attn_fwd( - q_part, - k_part, - v_part, - *fa_forward_args_thd, - causal=causal, - **fa_forward_kwargs, - ) - if not use_flash_attn_3 and not fa_utils.v2_7_0_plus: + if use_flash_attn_4: + fa_outputs = flash_attn_fwd( + q_part, + k_part, + v_part, + **get_fa_args( + True, + False, + qkv_format, + cu_seqlens_q=fa_cu_seqlens_q, + cu_seqlens_kv=fa_cu_seqlens_kv, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv_, + seqused_q=seqused_q, + seqused_k=seqused_k, + use_flash_attn_4=True, + ), + causal=causal, + **fa_forward_kwargs, + ) + else: + fa_forward_args_thd = get_fa_args( + True, + use_flash_attn_3, + qkv_format, + cu_seqlens_q=fa_cu_seqlens_q, + cu_seqlens_kv=fa_cu_seqlens_kv, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv_, + seqused_q=seqused_q, + seqused_k=seqused_k, + ) + fa_outputs = flash_attn_fwd( + q_part, + k_part, + v_part, + *fa_forward_args_thd, + causal=causal, + **fa_forward_kwargs, + ) + if use_flash_attn_4: + out_per_step[i] = fa_outputs[0] + softmax_lse_per_step[i] = fa_outputs[1] + elif not use_flash_attn_3 and not fa_utils.v2_7_0_plus: out_per_step[i] = fa_outputs[4] softmax_lse_per_step[i] = fa_outputs[5] rng_states[i] = fa_outputs[7] @@ -3646,6 +3799,7 @@ def forward( ctx.deterministic = deterministic ctx.use_fused_attention = use_fused_attention ctx.use_flash_attn_3 = use_flash_attn_3 + ctx.use_flash_attn_4 = use_flash_attn_4 ctx.pad_between_seqs = pad_between_seqs ctx.window_size = window_size if qkv_format == "thd": @@ -3814,7 +3968,14 @@ def backward(ctx, dout, *_args): flash_attn_bwd = None if not ctx.use_fused_attention: fa_backward_kwargs = {"softmax_scale": ctx.softmax_scale} - if ctx.use_flash_attn_3: + if ctx.use_flash_attn_4: + from transformer_engine.pytorch.attention.dot_product_attention.backends import ( + _flash_attn_bwd_v4, + ) + + flash_attn_bwd = _flash_attn_bwd_v4 + fa_backward_kwargs["deterministic"] = ctx.deterministic + elif ctx.use_flash_attn_3: from transformer_engine.pytorch.attention.dot_product_attention.backends import ( _flash_attn_bwd_v3, ) @@ -3849,7 +4010,7 @@ def backward(ctx, dout, *_args): # backward calls are serialized on GPU streams so that workspace # lifetimes do not overlap. FusedAttention keeps the existing # per-step overlap. - if i > 0 and ctx.use_flash_attn_3: + if i > 0 and (ctx.use_flash_attn_3 or ctx.use_flash_attn_4): flash_attn_streams[i].wait_stream(flash_attn_streams[i - 1]) with torch.cuda.stream(flash_attn_streams[i]): if ctx.qkv_format == "thd": @@ -3988,7 +4149,9 @@ def backward(ctx, dout, *_args): for x in [dq_per_step[i], dk_per_step[i], dv_per_step[i]] ] else: - if ctx.use_flash_attn_3 and ctx.qkv_format == "thd": + if ( + ctx.use_flash_attn_3 or ctx.use_flash_attn_4 + ) and ctx.qkv_format == "thd": dq_per_step[i], dk_per_step[i], dv_per_step[i] = [ torch.zeros_like(x) for x in [q_part, k_part, v_part] ] @@ -4004,7 +4167,9 @@ def backward(ctx, dout, *_args): else cu_seqlens_q ) fa_cu_seqlens_kv = cu_seqlens_kv_per_step[i] - if ctx.use_flash_attn_3 and ctx.qkv_format == "thd": + if ( + ctx.use_flash_attn_3 or ctx.use_flash_attn_4 + ) and ctx.qkv_format == "thd": seqused_q = ( thd_cu_seqlens_q_per_step[i][1:] - thd_cu_seqlens_q_per_step[i][:-1] ) @@ -4013,41 +4178,82 @@ def backward(ctx, dout, *_args): ) fa_cu_seqlens_q = thd_cu_seqlens_q_padded_per_step[i] fa_cu_seqlens_kv = cu_seqlens_kv_padded - fa_backward_args_thd = get_fa_args( - False, - ctx.use_flash_attn_3, - ctx.qkv_format, - cu_seqlens_q=fa_cu_seqlens_q, - cu_seqlens_kv=fa_cu_seqlens_kv, - max_seqlen_q=ctx.max_seqlen_q, - max_seqlen_kv=max_seqlen_kv, - dq=dq_per_step[i], - dk=dk_per_step[i], - dv=dv_per_step[i], - seqused_q=seqused_q, - seqused_k=seqused_k, - ) - if not ctx.use_flash_attn_3: + if ctx.use_flash_attn_4: + fa_backward_kwargs.update( + get_fa_args( + False, + False, + ctx.qkv_format, + cu_seqlens_q=fa_cu_seqlens_q, + cu_seqlens_kv=fa_cu_seqlens_kv, + max_seqlen_q=ctx.max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + dq=dq_per_step[i], + dk=dk_per_step[i], + dv=dv_per_step[i], + seqused_q=seqused_q, + seqused_k=seqused_k, + use_flash_attn_4=True, + ) + ) + else: + fa_backward_args_thd = get_fa_args( + False, + ctx.use_flash_attn_3, + ctx.qkv_format, + cu_seqlens_q=fa_cu_seqlens_q, + cu_seqlens_kv=fa_cu_seqlens_kv, + max_seqlen_q=ctx.max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + dq=dq_per_step[i], + dk=dk_per_step[i], + dv=dv_per_step[i], + seqused_q=seqused_q, + seqused_k=seqused_k, + ) + if ctx.use_flash_attn_4: + fa_backward_kwargs["causal"] = causal + elif not ctx.use_flash_attn_3: fa_backward_kwargs["rng_state"] = rng_states[i] - if fa_utils.v2_3_plus and not fa_utils.v2_7_0_plus: + if ( + not ctx.use_flash_attn_3 + and not ctx.use_flash_attn_4 + and fa_utils.v2_3_plus + and not fa_utils.v2_7_0_plus + ): fa_backward_kwargs["window_size"] = window_size_per_step[i] - elif ctx.use_flash_attn_3 or fa_utils.v2_7_0_plus: + elif ctx.use_flash_attn_3 or ctx.use_flash_attn_4 or fa_utils.v2_7_0_plus: fa_backward_kwargs["window_size_left"] = window_size_per_step[i][0] fa_backward_kwargs["window_size_right"] = window_size_per_step[i][1] if ctx.use_flash_attn_3: fa_backward_kwargs["is_causal"] = causal - else: + elif not ctx.use_flash_attn_4: fa_backward_kwargs["causal"] = causal - flash_attn_bwd( - dout_part, - q_part, - k_part, - v_part, - out_part, - softmax_lse_per_step[i], - *fa_backward_args_thd, - **fa_backward_kwargs, - ) + if ctx.use_flash_attn_4: + ( + dq_per_step[i], + dk_per_step[i], + dv_per_step[i], + ) = flash_attn_bwd( + q_part, + k_part, + v_part, + out_part, + dout_part, + softmax_lse_per_step[i], + **fa_backward_kwargs, + ) + else: + flash_attn_bwd( + dout_part, + q_part, + k_part, + v_part, + out_part, + softmax_lse_per_step[i], + *fa_backward_args_thd, + **fa_backward_kwargs, + ) if i > 0: # dq/dk/dv, dq_per_step/dk_per_step/dv_per_step: ctx.fwd_nominal_dtype @@ -4164,6 +4370,7 @@ def backward(ctx, dout, *_args): None, None, None, + None, ) @@ -4203,6 +4410,7 @@ def forward( quantizers, pad_between_seqs, use_flash_attn_3, + use_flash_attn_4, softmax_type, softmax_offset, fp8_output, @@ -4234,10 +4442,12 @@ def forward( or window_size == (-1, -1) or use_fused_attention or use_flash_attn_3 + or use_flash_attn_4 or fa_utils.v2_3_plus ), ( "cp_comm_type='a2a' only supports SWA through FusedAttention or FlashAttention >= 2.3." - f" Found {use_fused_attention=}, {use_flash_attn_3=}, and {fa_utils.v2_3_plus=}." + f" Found {use_fused_attention=}, {use_flash_attn_3=}, {use_flash_attn_4=}, " + f"and {fa_utils.v2_3_plus=}." ) assert q.shape[seq_dim_qkv] % 2 == 0 and k.shape[seq_dim_qkv] % 2 == 0, ( "cp_comm_type='a2a' requires seq_len % 2 == 0 for Q, K, V. Found seq_len_q =" @@ -4251,7 +4461,16 @@ def forward( flash_attn_fwd = None if not use_fused_attention: fa_forward_kwargs = {"softmax_scale": softmax_scale} - if use_flash_attn_3: + if use_flash_attn_4: + from transformer_engine.pytorch.attention.dot_product_attention.backends import ( + _flash_attn_fwd_v4, + ) + + flash_attn_fwd = _flash_attn_fwd_v4 + fa_forward_kwargs["window_size_left"] = window_size[0] + fa_forward_kwargs["window_size_right"] = window_size[1] + fa_forward_kwargs["return_lse"] = True + elif use_flash_attn_3: from transformer_engine.pytorch.attention.dot_product_attention.backends import ( _flash_attn_fwd_v3, ) @@ -4442,31 +4661,55 @@ def forward( seqused_k = None fa_cu_seqlens_q = cu_seqlens_q fa_cu_seqlens_kv = cu_seqlens_kv - if pad_between_seqs and use_flash_attn_3 and qkv_format == "thd": + if pad_between_seqs and (use_flash_attn_3 or use_flash_attn_4) and qkv_format == "thd": seqused_q = cu_seqlens_q[1:] - cu_seqlens_q[:-1] seqused_k = cu_seqlens_kv[1:] - cu_seqlens_kv[:-1] fa_cu_seqlens_q = cu_seqlens_q_padded fa_cu_seqlens_kv = cu_seqlens_kv_padded - fa_forward_args_thd = get_fa_args( - True, - use_flash_attn_3, - qkv_format, - cu_seqlens_q=fa_cu_seqlens_q, - cu_seqlens_kv=fa_cu_seqlens_kv, - max_seqlen_q=max_seqlen_q, - max_seqlen_kv=max_seqlen_kv, - seqused_q=seqused_q, - seqused_k=seqused_k, - ) - fa_outputs = flash_attn_fwd( - q_part, - k_part, - v_part, - *fa_forward_args_thd, - causal=causal, - **fa_forward_kwargs, - ) - if not use_flash_attn_3 and not fa_utils.v2_7_0_plus: + if use_flash_attn_4: + fa_outputs = flash_attn_fwd( + q_part, + k_part, + v_part, + **get_fa_args( + True, + False, + qkv_format, + cu_seqlens_q=fa_cu_seqlens_q, + cu_seqlens_kv=fa_cu_seqlens_kv, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + seqused_q=seqused_q, + seqused_k=seqused_k, + use_flash_attn_4=True, + ), + causal=causal, + **fa_forward_kwargs, + ) + else: + fa_forward_args_thd = get_fa_args( + True, + use_flash_attn_3, + qkv_format, + cu_seqlens_q=fa_cu_seqlens_q, + cu_seqlens_kv=fa_cu_seqlens_kv, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + seqused_q=seqused_q, + seqused_k=seqused_k, + ) + fa_outputs = flash_attn_fwd( + q_part, + k_part, + v_part, + *fa_forward_args_thd, + causal=causal, + **fa_forward_kwargs, + ) + if use_flash_attn_4: + out_, softmax_lse = fa_outputs[0], fa_outputs[1] + rng_state = None + elif not use_flash_attn_3 and not fa_utils.v2_7_0_plus: out_, softmax_lse = fa_outputs[4], fa_outputs[5] rng_state = fa_outputs[7] else: @@ -4593,6 +4836,7 @@ def forward( ctx.fwd_nominal_dtype = fwd_nominal_dtype ctx.fp8_recipe = fp8_recipe ctx.use_flash_attn_3 = use_flash_attn_3 + ctx.use_flash_attn_4 = use_flash_attn_4 ctx.pad_between_seqs = pad_between_seqs ctx.softmax_type = softmax_type @@ -4690,7 +4934,16 @@ def backward(ctx, dout, *_args): flash_attn_bwd = None if not ctx.use_fused_attention: fa_backward_kwargs = {"softmax_scale": ctx.softmax_scale} - if ctx.use_flash_attn_3: + if ctx.use_flash_attn_4: + from transformer_engine.pytorch.attention.dot_product_attention.backends import ( + _flash_attn_bwd_v4, + ) + + flash_attn_bwd = _flash_attn_bwd_v4 + fa_backward_kwargs["window_size_left"] = ctx.window_size[0] + fa_backward_kwargs["window_size_right"] = ctx.window_size[1] + fa_backward_kwargs["deterministic"] = ctx.deterministic + elif ctx.use_flash_attn_3: from transformer_engine.pytorch.attention.dot_product_attention.backends import ( _flash_attn_bwd_v3, ) @@ -4790,41 +5043,77 @@ def backward(ctx, dout, *_args): seqused_k = None fa_cu_seqlens_q = cu_seqlens_q fa_cu_seqlens_kv = cu_seqlens_kv - if ctx.pad_between_seqs and ctx.use_flash_attn_3 and ctx.dqkv_format == "thd": + if ( + ctx.pad_between_seqs + and (ctx.use_flash_attn_3 or ctx.use_flash_attn_4) + and ctx.dqkv_format == "thd" + ): seqused_q = cu_seqlens_q[1:] - cu_seqlens_q[:-1] seqused_k = cu_seqlens_kv[1:] - cu_seqlens_kv[:-1] fa_cu_seqlens_q = cu_seqlens_q_padded fa_cu_seqlens_kv = cu_seqlens_kv_padded - fa_backward_args_thd = get_fa_args( - False, - ctx.use_flash_attn_3, - ctx.dqkv_format, - cu_seqlens_q=fa_cu_seqlens_q, - cu_seqlens_kv=fa_cu_seqlens_kv, - max_seqlen_q=ctx.max_seqlen_q, - max_seqlen_kv=ctx.max_seqlen_kv, - dq=dq, - dk=dk, - dv=dv, - seqused_q=seqused_q, - seqused_k=seqused_k, - ) - if not ctx.use_flash_attn_3: + if ctx.use_flash_attn_4: + fa_backward_kwargs.update( + get_fa_args( + False, + False, + ctx.dqkv_format, + cu_seqlens_q=fa_cu_seqlens_q, + cu_seqlens_kv=fa_cu_seqlens_kv, + max_seqlen_q=ctx.max_seqlen_q, + max_seqlen_kv=ctx.max_seqlen_kv, + dq=dq, + dk=dk, + dv=dv, + seqused_q=seqused_q, + seqused_k=seqused_k, + use_flash_attn_4=True, + ) + ) + else: + fa_backward_args_thd = get_fa_args( + False, + ctx.use_flash_attn_3, + ctx.dqkv_format, + cu_seqlens_q=fa_cu_seqlens_q, + cu_seqlens_kv=fa_cu_seqlens_kv, + max_seqlen_q=ctx.max_seqlen_q, + max_seqlen_kv=ctx.max_seqlen_kv, + dq=dq, + dk=dk, + dv=dv, + seqused_q=seqused_q, + seqused_k=seqused_k, + ) + if ctx.use_flash_attn_4: + fa_backward_kwargs["causal"] = causal + elif not ctx.use_flash_attn_3: fa_backward_kwargs["rng_state"] = rng_state fa_backward_kwargs["causal"] = causal else: fa_backward_kwargs["is_causal"] = causal - flash_attn_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - *fa_backward_args_thd, - **fa_backward_kwargs, - ) + if ctx.use_flash_attn_4: + dq, dk, dv = flash_attn_bwd( + q, + k, + v, + out, + dout, + softmax_lse, + **fa_backward_kwargs, + ) + else: + flash_attn_bwd( + dout, + q, + k, + v, + out, + softmax_lse, + *fa_backward_args_thd, + **fa_backward_kwargs, + ) # dq, dk, dv: # FP8DS: torch.uint8 @@ -4916,6 +5205,7 @@ def backward(ctx, dout, *_args): None, None, None, + None, d_softmax_offset, None, ) @@ -4950,6 +5240,7 @@ def attn_forward_func_with_cp( quantizers=None, pad_between_seqs=False, use_flash_attn_3=False, + use_flash_attn_4=False, softmax_type="vanilla", softmax_offset=None, fp8_output=False, @@ -5103,6 +5394,7 @@ def attn_forward_func_with_cp( quantizers, pad_between_seqs, use_flash_attn_3, + use_flash_attn_4, fp8_output, layer_number, ] @@ -5113,6 +5405,7 @@ def attn_forward_func_with_cp( cp_group, cp_stream, use_flash_attn_3, + use_flash_attn_4, pad_between_seqs, fp8, fp8_meta, @@ -5130,6 +5423,7 @@ def attn_forward_func_with_cp( quantizers, pad_between_seqs, use_flash_attn_3, + use_flash_attn_4, softmax_type, softmax_offset, fp8_output, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index ba049c9aef..a076430da7 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -1073,13 +1073,12 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt if pad_between_seqs: if ( # pylint: disable=too-many-boolean-expressions use_flash_attention_2 and FlashAttentionUtils.is_installed - ) or (use_flash_attention_4 and FlashAttentionUtils.v4_is_installed): + ): logger.debug( - "Disabling FlashAttention 2 and 4 for qkv_format = thd when there is " + "Disabling FlashAttention 2 for qkv_format = thd when there is " "padding between sequences, i.e. [a, a, PAD, b, b, b, PAD, c, PAD]" ) use_flash_attention_2 = False - use_flash_attention_4 = False # FA3 supports pad_between_seqs via seqused_q/seqused_k if use_unfused_attention: logger.debug("Disabling UnfusedDotProductAttention for pad_between_seqs = True") @@ -1171,9 +1170,14 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt "Disabling UnfusedDotProductAttention as it does not support context parallelism" ) use_unfused_attention = False - if context_parallel and use_flash_attention_4 and FlashAttentionUtils.v4_is_installed: - logger.debug("Disabling FlashAttention 4 as it does not support context parallelism yet") - use_flash_attention_4 = False + if context_parallel and use_flash_attention_4: + fa4_cp_supported = cp_comm_type in ["p2p", "all_gather", "a2a"] + if not fa4_cp_supported: + logger.debug( + "Disabling FlashAttention 4 for context parallelism with cp_comm_type = %s", + cp_comm_type, + ) + use_flash_attention_4 = False if context_parallel and ( use_flash_attention_2 or use_flash_attention_3 or use_flash_attention_4 ):