diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index 681ee5e6e0..ceaa9ad382 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -300,7 +300,10 @@ def _submit(pool: PoolWorker, **kwargs) -> None: qkv_formats = ["sbhd", "thd"] -@pytest.mark.skipif(not FlashAttentionUtils.v2_plus, reason="Flash-attn 2.0+ is required.") +@pytest.mark.skipif( + not (FlashAttentionUtils.v2_plus or FlashAttentionUtils.v4_is_installed), + reason="Flash-attn 2.0+ or Flash-attn 4 is required.", +) @pytest.mark.skipif(get_device_compute_capability() < (8, 0), reason="CP tests require sm80+.") @pytest.mark.parametrize("dtype", dtypes) @pytest.mark.parametrize("model", model_configs_flash_attn.keys()) @@ -314,16 +317,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!") @@ -345,9 +342,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 6a60160f77..9b975c3fbf 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -166,11 +166,15 @@ except PackageNotFoundError: flash_attn_func_v4 = None flash_attn_varlen_func_v4 = None + _flash_attn_fwd_v4 = None + _flash_attn_bwd_v4 = None else: from flash_attn.cute.interface import ( # pylint: disable=ungrouped-imports,no-name-in-module 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, ) fa_utils.v4_validate_head_dims = _fa4_validate_head_dims @@ -1079,6 +1083,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: @@ -1138,10 +1143,21 @@ def forward( if inference_params is None: fa_4_optional_forward_kwargs["deterministic"] = self.deterministic if func is flash_attn_varlen_func_v4: - fa_4_optional_forward_kwargs["cu_seqlens_q"] = cu_seqlens_q - fa_4_optional_forward_kwargs["cu_seqlens_k"] = cu_seqlens_kv + fa_4_optional_forward_kwargs["cu_seqlens_q"] = ( + cu_seqlens_q_padded if pad_between_seqs else cu_seqlens_q + ) + fa_4_optional_forward_kwargs["cu_seqlens_k"] = ( + cu_seqlens_kv_padded if pad_between_seqs else cu_seqlens_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 995dd2e90a..f428b42a56 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -688,6 +688,28 @@ def get_fa_args( ] +def get_fa4_thd_kwargs( + qkv_format, + cu_seqlens_q, + cu_seqlens_kv, + max_seqlen_q, + max_seqlen_kv, + seqused_q=None, + seqused_k=None, +): + """Get FA4 varlen kwargs for THD inputs.""" + if qkv_format != "thd": + return {} + return { + "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, + } + + def cp_p2p_fwd_prepare_qkv( q_part, k_part, @@ -968,6 +990,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 +1019,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 +1044,48 @@ 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_fa4_thd_kwargs( + qkv_format, + cu_seqlens_q_, + cu_seqlens_kv_, + max_seqlen_q_, + max_seqlen_kv_, + seqused_q=seqused_q, + seqused_k=seqused_k, + ), + 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 +1305,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 +1333,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 +1371,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 +1381,61 @@ 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_fa4_thd_kwargs( + qkv_format, + cu_seqlens_q_bwd, + cu_seqlens_kv_bwd, + max_seqlen_q_, + max_seqlen_kv_, + seqused_q=seqused_q, + seqused_k=seqused_k, + ) + ) + fa_backward_kwargs["dq"] = dq + fa_backward_kwargs["dk"] = dk + fa_backward_kwargs["dv"] = dv + 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 +1482,7 @@ def forward( quantizers, pad_between_seqs, use_flash_attn_3, + use_flash_attn_4, fp8_output, layer_number, ): @@ -1622,14 +1710,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 +1867,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 +2264,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 +2527,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 +2665,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 +3079,7 @@ def backward(ctx, dout, *_args): None, None, None, + None, ) @@ -3022,7 +3132,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 +3161,7 @@ def forward( cp_group, cp_stream, use_flash_attn_3, + use_flash_attn_4, pad_between_seqs, fp8, fp8_meta, @@ -3087,10 +3198,12 @@ def forward( window_size == (-1, 0) or window_size == (-1, -1) or use_fused_attention + 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=} and {fa_utils.v2_3_plus=}." + f" >= 2.3. Found {use_fused_attention=}, {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='all_gather' requires seq_len % 2 == 0 for Q, K, V. Found seq_len_q =" @@ -3100,7 +3213,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, ) @@ -3339,7 +3459,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 @@ -3458,7 +3578,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] ) @@ -3467,31 +3587,57 @@ 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_fa4_thd_kwargs( + qkv_format, + fa_cu_seqlens_q, + fa_cu_seqlens_kv, + max_seqlen_q, + max_seqlen_kv_, + seqused_q=seqused_q, + seqused_k=seqused_k, + ), + 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] @@ -3644,6 +3790,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": @@ -3812,7 +3959,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, ) @@ -3847,7 +4001,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": @@ -3986,7 +4140,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] ] @@ -4002,7 +4158,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] ) @@ -4011,41 +4169,79 @@ 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_fa4_thd_kwargs( + ctx.qkv_format, + fa_cu_seqlens_q, + fa_cu_seqlens_kv, + ctx.max_seqlen_q, + max_seqlen_kv, + seqused_q=seqused_q, + seqused_k=seqused_k, + ) + ) + fa_backward_kwargs["dq"] = dq_per_step[i] + fa_backward_kwargs["dk"] = dk_per_step[i] + fa_backward_kwargs["dv"] = dv_per_step[i] + 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 @@ -4162,6 +4358,7 @@ def backward(ctx, dout, *_args): None, None, None, + None, ) @@ -4201,6 +4398,7 @@ def forward( quantizers, pad_between_seqs, use_flash_attn_3, + use_flash_attn_4, softmax_type, softmax_offset, fp8_output, @@ -4231,10 +4429,11 @@ def forward( window_size == (-1, 0) or window_size == (-1, -1) or use_fused_attention + 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=} and {fa_utils.v2_3_plus=}." + f" Found {use_fused_attention=}, {use_flash_attn_4=}, 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 =" @@ -4248,7 +4447,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, ) @@ -4439,31 +4647,52 @@ 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_fa4_thd_kwargs( + qkv_format, + fa_cu_seqlens_q, + fa_cu_seqlens_kv, + max_seqlen_q, + max_seqlen_kv, + seqused_q=seqused_q, + seqused_k=seqused_k, + ), + 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: @@ -4590,6 +4819,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 @@ -4687,7 +4917,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, ) @@ -4787,41 +5026,74 @@ 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_fa4_thd_kwargs( + ctx.dqkv_format, + fa_cu_seqlens_q, + fa_cu_seqlens_kv, + ctx.max_seqlen_q, + ctx.max_seqlen_kv, + seqused_q=seqused_q, + seqused_k=seqused_k, + ) + ) + fa_backward_kwargs["dq"] = dq + fa_backward_kwargs["dk"] = dk + fa_backward_kwargs["dv"] = dv + 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 @@ -4913,6 +5185,7 @@ def backward(ctx, dout, *_args): None, None, None, + None, d_softmax_offset, None, ) @@ -4947,6 +5220,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 +5377,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 +5388,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 +5406,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 8a4c6ce286..4bb48ebdc7 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -1046,13 +1046,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") @@ -1144,9 +1143,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 ):