[PyTorch][CP] Support compact metadata with FA4

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

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)

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]
                             )

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(