Enable Paged-Attention in FA3 build

ghstack-source-id: 80565872ee56f53b8604d402e73dc9cc473e6c8d
Pull-Request-resolved: fairinternal/xformers#1392

__original_commit__ = fairinternal/xformers@31dd6f1

Author: danthe3rd

setup.py

@@ -339,7 +339,7 @@ def get_flash_attention3_extensions(cuda_version: int, extra_compile_args):
         # Enabling paged attention causes segfault with some
         # versions of nvcc :(
         # https://github.com/Dao-AILab/flash-attention/issues/1453
-        ("paged", "-DFLASHATTENTION_DISABLE_PAGEDKV"),
+        # ("paged", "-DFLASHATTENTION_DISABLE_PAGEDKV"),
         # We have `CUDA_MINIMUM_COMPUTE_CAPABILITY` set to 9.0
         # ("_sm80.cu", "-DFLASHATTENTION_DISABLE_SM8x"),
     )

xformers/ops/fmha/flash3.py

@@ -101,7 +101,7 @@ def _flash_attention3_incompatible_reason() -> Optional[str]:
     return None
 
 
-FLASH3_HAS_PAGED_ATTENTION = False
+FLASH3_HAS_PAGED_ATTENTION = True
 FLASH3_HAS_FLOAT8 = False
 _C_flashattention3 = None
 if importlib.util.find_spec("...flash_attn_3._C", package=__package__):