Archive of how to enable Flash Attention-3 on Pytorch2.8.0+cu129+xformers0.0.32.post2

PS D:\userfiles\comfyui> python_embeded\python.exe -m xformers.info
xFormers 0.0.32.post1
memory_efficient_attention.ckF:                    unavailable
memory_efficient_attention.ckB:                    unavailable
memory_efficient_attention.ck_splitKF:             unavailable
memory_efficient_attention.cutlassF-pt:            available
memory_efficient_attention.cutlassB-pt:            available
memory_efficient_attention.fa2F@0.0.0:             unavailable
memory_efficient_attention.fa2B@0.0.0:             unavailable
memory_efficient_attention.fa3F@2.8.0.post2-3-g3ba6f82: available
memory_efficient_attention.fa3B@2.8.0.post2-3-g3ba6f82: available
memory_efficient_attention.fa3F_splitKV@2.8.0.post2-3-g3ba6f82: available
memory_efficient_attention.triton_splitKF:         available
indexing.scaled_index_addF:                        available
indexing.scaled_index_addB:                        available
indexing.index_select:                             available
sp24.sparse24_sparsify_both_ways:                  available
sp24.sparse24_apply:                               available
sp24.sparse24_apply_dense_output:                  available
sp24._sparse24_gemm:                               available
sp24._cslt_sparse_mm_search@0.0.0:                 available
sp24._cslt_sparse_mm@0.0.0:                        available
swiglu.dual_gemm_silu:                             available
swiglu.gemm_fused_operand_sum:                     available
swiglu.fused.p.cpp:                                available
is_triton_available:                               True
pytorch.version:                                   2.8.0+cu129
pytorch.cuda:                                      available
gpu.compute_capability:                            8.9
gpu.name:                                          NVIDIA GeForce RTX 4070
dcgm_profiler:                                     unavailable
build.info:                                        available
build.cuda_version:                                1208
build.hip_version:                                 None
build.python_version:                              3.9.13
build.torch_version:                               2.8.0+cu128
build.env.TORCH_CUDA_ARCH_LIST:                    7.5 8.0+PTX 9.0a
build.env.PYTORCH_ROCM_ARCH:                       None
build.env.XFORMERS_BUILD_TYPE:                     Release
build.env.XFORMERS_ENABLE_DEBUG_ASSERTIONS:        None
build.env.NVCC_FLAGS:                              -allow-unsupported-compiler
build.env.XFORMERS_PACKAGE_FROM:                   wheel-v0.0.32.post1
build.nvcc_version:                                12.8.93
source.privacy:                                    open source
PS D:\userfiles\comfyui>

def fa3_available() -> bool:
    has_cuda = torch.version.cuda is not None
    is_90a = has_cuda and torch.cuda.get_device_capability() >= (9, 0)  # SM90以上
    has_valid_flash3 = flash3._C_flashattention3 is not None
    return is_90a and has_valid_flash3

SUPPORTED_DTYPES: Set[torch.dtype] = {
    torch.half,      # float16
    torch.bfloat16,  # bfloat16
}  # float32は含まれていない

def _dispatch_fw_priority_list(inp: Inputs, needs_gradient: bool):
    if torch.version.cuda:
        flash3_op = [flash3.FwOp] if _get_use_fa3() else []
        priority_list_ops = deque(
            flash3_op  # FA-3が最初に試される
            + [
                flash.FwOp,
                cutlass.FwOp,
            ]
        )

_USE_FLASH_ATTENTION_3 = True  # 定数

def _get_use_fa3() -> bool:
    global _USE_FLASH_ATTENTION_3
    return _USE_FLASH_ATTENTION_3  # 常にTrue

class FwOp(AttentionFwOpBase):
    CUDA_MINIMUM_COMPUTE_CAPABILITY = (8, 0)  # SM80以上

def _run_priority_list(name: str, priority_list: Sequence[T], inp: Inputs):
    for op in priority_list:
        not_supported = op.not_supported_reasons(inp)
        if not not_supported:  # サポートされている場合
            return op

if importlib.util.find_spec("...flash_attn_3._C", package=__package__):
    from ..._cpp_lib import _build_metadata
    from ...flash_attn_3 import _C
    _C_flashattention3 = torch.ops.flash_attn_3

{
  "version": {
    "cuda": 1208,
    "torch": "2.7.1+cu126",
    "flash": "v2.8.0.post2-3-g3ba6f82"
  },
  "env": {
    "TORCH_CUDA_ARCH_LIST": "7.5 8.0+PTX 9.0a"
  }
}

def fa3_available() -> bool:
    has_cuda = torch.version.cuda is not None
    is_90a = has_cuda and torch.cuda.get_device_capability() >= (9, 0)
    has_valid_flash3 = flash3._C_flashattention3 is not None
    # Allow FA-3 on SM89 (Ada Lovelace) for testing
    is_89a = has_cuda and torch.cuda.get_device_capability() >= (8, 9)
    return (is_90a or is_89a) and has_valid_flash3

def _dispatch_fw(inp: Inputs, needs_gradient: bool) -> Type[AttentionFwOpBase]:
    """Computes the best operator for forward
    Raises:
        NotImplementedError: if not operator was found
    Returns:
        AttentionOp: The best operator for the configuration
    """
    op = _run_priority_list(
        "memory_efficient_attention_forward",
        _dispatch_fw_priority_list(inp, needs_gradient),
        inp,
    )

    # Log the selected kernel
    try:
        import logging
        logger = logging.getLogger("xformers_attention_log")
        if not hasattr(_dispatch_fw, "_last_kernel"):
            _dispatch_fw._last_kernel = None
        last_kernel = _dispatch_fw._last_kernel
        if op.NAME != last_kernel:
            logger.info(f"[xformers] memory_efficient_attention: selected kernel = {op.NAME}")
            _dispatch_fw._last_kernel = op.NAME
    except Exception:
        print(f"[xformers] memory_efficient_attention: selected kernel = {getattr(op, 'NAME', 'unknown')}")

    return op

SUPPORTED_DTYPES: Set[torch.dtype] = {
    torch.half,
    torch.bfloat16,
    torch.float,  # Add float32 support for compatibility with 0.0.31
} | ({torch.float8_e4m3fn} if FLASH3_HAS_FLOAT8 else set())

def _flash_attention3_incompatible_reason() -> Optional[str]:
    # PyTorch 2.8.0ではflash_attn_3のopsが利用できない場合があるため、xformersの内部実装を使用
    if not hasattr(torch.ops, 'flash_attn_3'):
        return None  # xformersの内部実装を使用するため、エラーを返さない

    if not hasattr(torch.ops.flash_attn_3, "fwd") or not hasattr(
        torch.ops.flash_attn_3, "bwd"
    ):
        return None  # xformersの内部実装を使用するため、エラーを返さない

    if not torch.ops.flash_attn_3.fwd.default._schema.is_backward_compatible_with(...):
        return None  # xformersの内部実装を使用するため、エラーを返さない

    if not torch.ops.flash_attn_3.bwd.default._schema.is_backward_compatible_with(...):
        return None  # xformersの内部実装を使用するため、エラーを返さない

    return None

# デフォルトでxformers（FA-3）を使用、SageAttentionはKJノードが動的に制御
if model_management.xformers_enabled():
    logging.info("Using xformers attention (FA-3)")
    optimized_attention = attention_xformers
elif model_management.sage_attention_enabled():
    logging.info("Using sage attention")
    optimized_attention = attention_sage
elif model_management.flash_attention_enabled():
    logging.info("Using Flash Attention")
    optimized_attention = attention_flash
elif model_management.pytorch_attention_enabled():
    logging.info("Using pytorch attention")
    optimized_attention = attention_pytorch
else:
    if args.use_split_cross_attention:
        logging.info("Using split optimization for attention")
        optimized_attention = attention_split
    else:
        logging.info("Using sub quadratic optimization for attention, if you have memory or speed issues try using: --use-split-cross-attention")
        optimized_attention = attention_sub_quad

def attention_xformers(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False):
    # ... 既存コード ...

    # Log xformers kernel selection
    try:
        import logging
        logger = logging.getLogger("xformers_attention_log")
        if not logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter('%(message)s')
            handler.setFormatter(formatter)
            logger.addHandler(handler)
            logger.setLevel(logging.INFO)
            logger.propagate = False

        # Force xformers to select kernel and log it
        if not hasattr(attention_xformers, "_last_logged"):
            attention_xformers._last_logged = False
            logger.info("[xformers] attention_xformers called - kernel selection will be logged")
    except Exception:
        pass

    out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=mask)
    # ... 既存コード ...

def xformers_attention(q, k, v):
    # ... 既存コード ...

    # Log xformers kernel selection for VAE
    try:
        import logging
        logger = logging.getLogger("xformers_attention_log")
        if not logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter('%(message)s')
            handler.setFormatter(formatter)
            logger.addHandler(handler)
            logger.setLevel(logging.INFO)
            logger.propagate = False

        # Force xformers to select kernel and log it
        if not hasattr(xformers_attention, "_last_logged"):
            xformers_attention._last_logged = False
            logger.info("[xformers] VAE xformers_attention called - kernel selection will be logged")
    except Exception:
        pass

    try:
        out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None)
        out = out.transpose(1, 2).reshape(orig_shape)
    # ... 既存コード ...

python_embeded\python.exe ComfyUI\main.py --use-sage-attention

python_embeded\python.exe ComfyUI\main.py --use-xformers

Using xformers attention (FA-3)
[xformers] attention_xformers called - kernel selection will be logged
[xformers] memory_efficient_attention: selected kernel = fa3F@2.8.0.post2-3-g3ba6f82

Patching comfy attention to use sageattn
SageAttention kernel is being used for this generation.
[SageAttention][DEBUG] sageattn (auto) called
Restoring initial comfy attention

# xformersを再インストール
python_embeded\python.exe -m pip uninstall xformers
python_embeded\python.exe -m pip install xformers==0.0.32.post1
# 修正を再適用