diff --git a/tests/pytorch/attention/run_fa4_cp_vs_non_cp.py b/tests/pytorch/attention/run_fa4_cp_vs_non_cp.py new file mode 100644 index 0000000000..b33aa60a69 --- /dev/null +++ b/tests/pytorch/attention/run_fa4_cp_vs_non_cp.py @@ -0,0 +1,185 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Compare FlashAttention v4 context parallelism against full-sequence attention. + +Example: + torchrun --nproc_per_node=2 tests/pytorch/attention/run_fa4_cp_vs_non_cp.py +""" + +import argparse +import os + +import torch +import torch.distributed as dist + + +def _parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--batch-size", type=int, default=2) + parser.add_argument("--seq-len", type=int, default=1024) + parser.add_argument("--num-heads", type=int, default=16) + parser.add_argument("--head-dim", type=int, default=64) + parser.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16") + parser.add_argument("--attn-mask-type", choices=["causal", "no_mask"], default="causal") + parser.add_argument("--cp-comm-type", choices=["p2p"], default="p2p") + parser.add_argument("--atol", type=float, default=5.0e-2) + parser.add_argument("--rtol", type=float, default=5.0e-2) + parser.add_argument("--rmse-tol", type=float, default=2.5e-2) + return parser.parse_args() + + +def _cp_slice(tensor, rank, world_size, seq_dim): + """Select the two load-balanced sequence chunks owned by a CP rank.""" + chunked = tensor.view( + *tensor.shape[:seq_dim], + 2 * world_size, + tensor.shape[seq_dim] // (2 * world_size), + *tensor.shape[(seq_dim + 1) :], + ) + chunk_ids = torch.tensor([rank, 2 * world_size - rank - 1], device=tensor.device) + local = chunked.index_select(seq_dim, chunk_ids) + return local.reshape( + *local.shape[:seq_dim], + -1, + *local.shape[(seq_dim + 2) :], + ).contiguous() + + +def _metric(actual, expected): + diff = (actual.float() - expected.float()).abs() + max_abs = diff.max() + rmse = torch.sqrt(torch.mean(diff * diff)) + return max_abs, rmse + + +def _assert_close(name, actual, expected, atol, rtol, rmse_tol): + max_abs, rmse = _metric(actual, expected) + torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol) + if rmse.item() > rmse_tol: + raise AssertionError(f"{name} RMSE {rmse.item():.6g} exceeds {rmse_tol:.6g}") + return max_abs.detach(), rmse.detach() + + +def _make_attention(num_heads, head_dim, qkv_format, attn_mask_type): + from transformer_engine.pytorch import DotProductAttention + + return DotProductAttention( + num_heads, + head_dim, + num_gqa_groups=num_heads, + attention_dropout=0.0, + qkv_format=qkv_format, + attn_mask_type=attn_mask_type, + ).cuda() + + +def _force_fa4(): + os.environ["NVTE_FLASH_ATTN"] = "1" + os.environ["NVTE_FUSED_ATTN"] = "0" + os.environ["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] = "1" + + from transformer_engine.pytorch.attention.dot_product_attention.backends import ( + _flash_attn_bwd_v4, + _flash_attn_fwd_v4, + ) + from transformer_engine.pytorch.attention.dot_product_attention.utils import FlashAttentionUtils + + if ( + not FlashAttentionUtils.v4_is_installed + or _flash_attn_fwd_v4 is None + or _flash_attn_bwd_v4 is None + ): + raise RuntimeError( + "FlashAttention v4 with raw cute _flash_attn_fwd/_flash_attn_bwd APIs is required." + ) + + # Keep the backend selector on FA4 even when FA2/FA3 are also installed. + FlashAttentionUtils.is_installed = False + FlashAttentionUtils.v3_is_installed = False + + +def main(): + args = _parse_args() + local_rank = int(os.environ["LOCAL_RANK"]) + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + + torch.cuda.set_device(local_rank) + dist.init_process_group("nccl") + _force_fa4() + + if args.seq_len % (2 * world_size) != 0: + raise ValueError("--seq-len must be divisible by 2 * WORLD_SIZE") + dtype = {"bf16": torch.bfloat16, "fp16": torch.float16}[args.dtype] + qkv_format = "bshd" + seq_dim = qkv_format.index("s") + device = torch.device("cuda", local_rank) + + torch.manual_seed(1234) + q_full = torch.randn( + args.batch_size, args.seq_len, args.num_heads, args.head_dim, device=device, dtype=dtype + ).clamp_(-1, 1) + k_full = torch.randn_like(q_full).clamp_(-1, 1) + v_full = torch.randn_like(q_full).clamp_(-1, 1) + dout_full = torch.randn( + args.batch_size, + args.seq_len, + args.num_heads * args.head_dim, + device=device, + dtype=dtype, + ).clamp_(-1, 1) + + q_ref, k_ref, v_ref = [x.detach().clone().requires_grad_() for x in (q_full, k_full, v_full)] + attn_ref = _make_attention(args.num_heads, args.head_dim, qkv_format, args.attn_mask_type) + out_ref = attn_ref(q_ref, k_ref, v_ref) + out_ref.backward(dout_full) + + q_cp, k_cp, v_cp = [ + _cp_slice(x.detach(), rank, world_size, seq_dim).requires_grad_() + for x in (q_full, k_full, v_full) + ] + dout_cp = _cp_slice(dout_full, rank, world_size, seq_dim) + + cp_group = dist.new_group(list(range(world_size)), backend="nccl") + cp_ranks = list(range(world_size)) + attn_cp = _make_attention(args.num_heads, args.head_dim, qkv_format, args.attn_mask_type) + attn_cp.set_context_parallel_group(cp_group, cp_ranks, torch.cuda.Stream(), args.cp_comm_type) + out_cp = attn_cp(q_cp, k_cp, v_cp) + out_cp.backward(dout_cp) + + torch.cuda.synchronize() + + checks = { + "out": (out_cp.detach(), _cp_slice(out_ref.detach(), rank, world_size, seq_dim)), + "dq": (q_cp.grad.detach(), _cp_slice(q_ref.grad.detach(), rank, world_size, seq_dim)), + "dk": (k_cp.grad.detach(), _cp_slice(k_ref.grad.detach(), rank, world_size, seq_dim)), + "dv": (v_cp.grad.detach(), _cp_slice(v_ref.grad.detach(), rank, world_size, seq_dim)), + } + + local_metrics = [] + for name, (actual, expected) in checks.items(): + max_abs, rmse = _assert_close(name, actual, expected, args.atol, args.rtol, args.rmse_tol) + local_metrics.append((name, max_abs, rmse)) + print( + f"[rank {rank}] {name}: max_abs={max_abs.item():.6g}, rmse={rmse.item():.6g}", + flush=True, + ) + + metric_tensor = torch.stack([torch.stack([m[1], m[2]]) for m in local_metrics]) + dist.all_reduce(metric_tensor, op=dist.ReduceOp.MAX) + if rank == 0: + for idx, (name, _, _) in enumerate(local_metrics): + print( + f"[global max] {name}: max_abs={metric_tensor[idx, 0].item():.6g}, " + f"rmse={metric_tensor[idx, 1].item():.6g}", + flush=True, + ) + print("FA4 CP matches non-CP for local sequence chunks.", flush=True) + + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 6a60160f77..717cf473a0 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -166,6 +166,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: from flash_attn.cute.interface import ( # pylint: disable=ungrouped-imports,no-name-in-module flash_attn_func as flash_attn_func_v4, @@ -173,6 +175,15 @@ _validate_head_dims as _fa4_validate_head_dims, ) + try: + from flash_attn.cute.interface import ( # pylint: disable=ungrouped-imports,no-name-in-module + _flash_attn_fwd as _flash_attn_fwd_v4, + _flash_attn_bwd as _flash_attn_bwd_v4, + ) + except ImportError: + _flash_attn_fwd_v4 = None + _flash_attn_bwd_v4 = None + fa_utils.v4_validate_head_dims = _fa4_validate_head_dims fa_utils.set_flash_attention_4_params() @@ -1079,6 +1090,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: 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..c8d6578d94 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,102 @@ def get_fa_args( ] +def _flash_attn_v4_set_window_kwargs(fa_kwargs: dict, window_size: Tuple[int, int]): + """Set window kwargs for the FlashAttention v4 raw API.""" + fa_kwargs["window_size_left"] = window_size[0] + fa_kwargs["window_size_right"] = window_size[1] + + +def _flash_attn_fwd_raw( + qkv_format: str, + fa_forward_kwargs: dict, + flash_attn_fwd, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + max_seqlen_q: int, + max_seqlen_kv: int, + causal: bool, + seqused_q: torch.Tensor = None, + seqused_k: torch.Tensor = None, +): + """Run the FlashAttention v4 raw forward call.""" + if flash_attn_fwd is None: + raise RuntimeError( + "FlashAttention v4 context parallelism requires the raw " + "flash_attn.cute.interface._flash_attn_fwd API." + ) + fa_outputs = flash_attn_fwd( + q, + k, + v, + cu_seqlens_q=cu_seqlens_q if qkv_format == "thd" else None, + cu_seqlens_k=cu_seqlens_kv if qkv_format == "thd" else None, + seqused_q=seqused_q, + seqused_k=seqused_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_kv, + causal=causal, + **fa_forward_kwargs, + ) + return fa_outputs[0], fa_outputs[1], None + + +def _flash_attn_bwd_raw( + qkv_format: str, + fa_backward_kwargs: dict, + flash_attn_bwd, + dout: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + softmax_lse: torch.Tensor, + dq: torch.Tensor, + dk: torch.Tensor, + dv: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + max_seqlen_q: int, + max_seqlen_kv: int, + causal: bool, + seqused_q: torch.Tensor = None, + seqused_k: torch.Tensor = None, +): + """Run the FlashAttention v4 raw backward call.""" + if flash_attn_bwd is None: + raise RuntimeError( + "FlashAttention v4 context parallelism requires the raw " + "flash_attn.cute.interface._flash_attn_bwd API." + ) + grads = flash_attn_bwd( + q, + k, + v, + out, + dout, + softmax_lse, + dq=dq, + dk=dk, + dv=dv, + softmax_scale=fa_backward_kwargs["softmax_scale"], + causal=causal, + softcap=fa_backward_kwargs.get("softcap", 0.0), + window_size_left=fa_backward_kwargs.get("window_size_left", -1), + window_size_right=fa_backward_kwargs.get("window_size_right", 0 if causal else -1), + cu_seqlens_q=cu_seqlens_q if qkv_format == "thd" else None, + cu_seqlens_k=cu_seqlens_kv if qkv_format == "thd" else None, + seqused_q=seqused_q, + seqused_k=seqused_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_kv, + deterministic=fa_backward_kwargs.get("deterministic", False), + ) + return grads[:3] + + def cp_p2p_fwd_prepare_qkv( q_part, k_part, @@ -968,6 +1064,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,7 +1093,9 @@ 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 use_flash_attn_4: + _flash_attn_v4_set_window_kwargs(fa_forward_kwargs, (-1, -1)) + elif 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: fa_forward_kwargs["window_size_left"] = -1 @@ -1016,6 +1115,23 @@ def cp_p2p_fwd_flash_attn( elif section == "upper-triangle": cu_seqlens_q_ = cu_seqlens_q_padded // 2 + if use_flash_attn_4: + return _flash_attn_fwd_raw( + qkv_format, + fa_forward_kwargs, + flash_attn_fwd, + q_part, + k_part, + v_part, + cu_seqlens_q_, + cu_seqlens_kv_, + max_seqlen_q_, + max_seqlen_kv_, + causal_, + seqused_q=seqused_q, + seqused_k=seqused_k, + ) + fa_forward_args_thd = get_fa_args( True, use_flash_attn_3, @@ -1256,6 +1372,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,7 +1400,9 @@ 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 use_flash_attn_4: + _flash_attn_v4_set_window_kwargs(fa_backward_kwargs, (-1, -1)) + elif 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: fa_backward_kwargs["window_size_left"] = -1 @@ -1295,7 +1414,9 @@ def cp_p2p_bwd_flash_attn( softmax_lse__ = softmax_lse causal_ = False if section == "diagonal": - if fa_utils.v2_3_plus and not fa_utils.v2_7_0_plus: + if use_flash_attn_4: + _flash_attn_v4_set_window_kwargs(fa_backward_kwargs, (-1, 0)) + elif 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: fa_backward_kwargs["window_size_left"] = -1 @@ -1321,6 +1442,29 @@ def cp_p2p_bwd_flash_attn( elif section == "upper-triangle": cu_seqlens_q_bwd = cu_seqlens_q_padded // 2 + if use_flash_attn_4: + return _flash_attn_bwd_raw( + qkv_format, + fa_backward_kwargs, + flash_attn_bwd, + dout_part, + q_part, + k_part, + v_part, + out_part, + softmax_lse__, + dq, + dk, + dv, + cu_seqlens_q_bwd, + cu_seqlens_kv_bwd, + max_seqlen_q_, + max_seqlen_kv_, + causal_, + seqused_q=seqused_q, + seqused_k=seqused_k, + ) + fa_backward_args_thd = get_fa_args( False, use_flash_attn_3, @@ -1395,6 +1539,7 @@ def forward( quantizers, pad_between_seqs, use_flash_attn_3, + use_flash_attn_4, fp8_output, layer_number, ): @@ -1622,14 +1767,27 @@ 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 # pylint: disable=possibly-used-before-assignment + ) + fa_forward_kwargs["return_lse"] = True + fa_forward_kwargs["softcap"] = 0.0 + _flash_attn_v4_set_window_kwargs(fa_forward_kwargs, (-1, 0) if causal else (-1, -1)) + elif use_flash_attn_3: from transformer_engine.pytorch.attention.dot_product_attention.backends import ( _flash_attn_fwd_v3, ) @@ -1768,6 +1926,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 +2323,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 +2586,17 @@ 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 # pylint: disable=possibly-used-before-assignment + ) + fa_backward_kwargs["softcap"] = 0.0 + 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 +2727,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 +3141,7 @@ def backward(ctx, dout, *_args): None, None, None, + None, ) @@ -4947,6 +5119,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, @@ -5035,6 +5208,12 @@ def attn_forward_func_with_cp( cp_group, dist_group_type ), f"cp_group must be {dist_group_type} type for {cp_comm_type=}!" + if use_flash_attn_4 and cp_comm_type not in ["p2p", "a2a+p2p"]: + raise ValueError( + "FlashAttention v4 context parallelism is supported for " + "cp_comm_type='p2p' and 'a2a+p2p' only." + ) + assert qkv_format in [ "bshd", "sbhd", @@ -5103,6 +5282,7 @@ def attn_forward_func_with_cp( quantizers, pad_between_seqs, use_flash_attn_3, + use_flash_attn_4, fp8_output, layer_number, ] diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 6c47d0f1cb..5a24dbea5c 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -1144,9 +1144,6 @@ 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_2 or use_flash_attention_3 or use_flash_attention_4 ):