完全版 How to fix A1111 for Flash Attention-3＆2

python -m xformers.info

xFormers 0.0.32+8354497d.d20250716
memory_efficient_attention.ckF:                    unavailable
memory_efficient_attention.ckB:                    unavailable
memory_efficient_attention.ck_decoderF:            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.7.1+cu128
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.11.13
build.torch_version:                               2.7.1+cu128
build.env.TORCH_CUDA_ARCH_LIST:                    8.9
build.env.PYTORCH_ROCM_ARCH:                       None
build.env.XFORMERS_BUILD_TYPE:                     None
build.env.XFORMERS_ENABLE_DEBUG_ASSERTIONS:        None
build.env.NVCC_FLAGS:                              None
build.env.XFORMERS_PACKAGE_FROM:                   None
build.nvcc_version:                                12.8.93
source.privacy:                                    open source

def xformers_attention_forward(self, x, context=None, mask=None, **kwargs):
    # ...（q, k, vの生成）...
    # supports()やnot_supported_reasons()で判定
    log_fa_operation_runtime(q, k, v)
    # 判定OKなら
    out = xformers.ops.memory_efficient_attention(q, k, v, ...)
    # NGなら従来のSDP等にフォールバック
    # ...（後略）...

FA3_LOGGED_THIS_GEN = False

def reset_fa3_log():
    global FA3_LOGGED_THIS_GEN
    FA3_LOGGED_THIS_GEN = False

def xformers_attention_forward(self, x, context=None, mask=None):
    global FA3_LOGGED_THIS_GEN
    try:
        # ...（q, k, vの生成）...
        # dtypeをfloat16優先で揃える
        # 必ずmemory_efficient_attentionを呼ぶ
        out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=mask)
        if not FA3_LOGGED_THIS_GEN:
            print("[A1111] xformers.memory_efficient_attention called (FA-3 or best available kernel will be used)")
            FA3_LOGGED_THIS_GEN = True
        # ...（後略）...
    except Exception as e:
        # フォールバック
        if hasattr(self, "_old_attention_forward"):
            return self._old_attention_forward(x, context, mask)
        raise

def xformers_attention_forward(self, x, context=None, mask=None, **kwargs):
    global FA3_LOGGED_THIS_GEN
    # Remove unsupported kwargs for xformers/SDP
    kwargs.pop('additional_tokens', None)
    try:
        import xformers.ops
        h = self.heads
        q_in = self.to_q(x)
        context = context if context is not None else x
        k_in = self.to_k(context)
        v_in = self.to_v(context)
        q, k, v = (t.reshape(t.shape[0], t.shape[1], h, -1) for t in (q_in, k_in, v_in))
        # dtype揃え
        dtypes = [q.dtype, k.dtype, v.dtype]
        if torch.float16 in dtypes:
            target_dtype = torch.float16
        else:
            target_dtype = q.dtype
        q = q.to(target_dtype)
        k = k.to(target_dtype)
        v = v.to(target_dtype)
        out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=mask)
        if not FA3_LOGGED_THIS_GEN:
            print("[A1111] xformers.memory_efficient_attention called (FA-3 or best available kernel will be used)")
            FA3_LOGGED_THIS_GEN = True
        out = out.to(x.dtype)
        b, n, h, d = out.shape
        out = out.reshape(b, n, h * d)
        return self.to_out(out)
    except Exception as e:
        print(f"[A1111] xformers.memory_efficient_attention failed, falling back to SDP. Exception: {e}")
        # Try PyTorch's scaled_dot_product_attention (SDP)
        try:
            h = self.heads
            q_in = self.to_q(x)
            context = context if context is not None else x
            k_in = self.to_k(context)
            v_in = self.to_v(context)
            q, k, v = (t.reshape(t.shape[0], t.shape[1], h, -1) for t in (q_in, k_in, v_in))
            dtype = q.dtype
            q = q.contiguous()
            k = k.contiguous()
            v = v.contiguous()
            out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False)
            out = out.to(dtype)
            b, n, h, d = out.shape
            out = out.reshape(b, n, h * d)
            print("[A1111] Fallback: torch.nn.functional.scaled_dot_product_attention (SDP) used.")
            return self.to_out(out)
        except Exception as e2:
            print(f"[A1111] SDP also failed, falling back to legacy. Exception: {e2}")
            if hasattr(self, "_old_attention_forward"):
                return self._old_attention_forward(x, context, mask)
            raise

def process_images_inner(p: StableDiffusionProcessing) -> Processed:
    """this is the main loop ..."""
    # ...（画像生成処理）...

from modules.sd_hijack_optimizations import reset_fa3_log

def process_images_inner(p: StableDiffusionProcessing) -> Processed:
    reset_fa3_log()
    """this is the main loop ..."""
    # ...（画像生成処理）...

def process_images_inner(p: StableDiffusionProcessing) -> Processed:
    # --- Add for xformers kernel log reset and A1111 log ---
    try:
        from xformers.ops.fmha import reset_kernel_log
        reset_kernel_log()
    except ImportError:
        pass
    print("[A1111] xformers.memory_efficient_attention called (FA-3 or best available kernel will be used)")
    # --- End add ---
    reset_fa3_log()
    # ...既存の処理...

def _memory_efficient_attention_forward(
    inp: Inputs, op: Optional[Type[AttentionFwOpBase]]
) -> torch.Tensor:
    inp.validate_inputs()
    output_shape = inp.normalize_bmhk()
    if op is None:
        op = _dispatch_fw(inp, False)
    else:
        _ensure_op_supports_or_raise(ValueError, "memory_efficient_attention", op, inp)

    out, *_ = op.apply(inp, needs_gradient=False)
    return out.reshape(output_shape)

_kernel_log_shown = False

def reset_kernel_log():
    global _kernel_log_shown
    _kernel_log_shown = False

def _memory_efficient_attention_forward(
    inp: Inputs, op: Optional[Type[AttentionFwOpBase]]
) -> torch.Tensor:
    global _kernel_log_shown
    inp.validate_inputs()
    output_shape = inp.normalize_bmhk()
    if op is None:
        op = _dispatch_fw(inp, False)
    else:
        _ensure_op_supports_or_raise(ValueError, "memory_efficient_attention", op, inp)

    # 1生成ごとに1回だけカーネル名をログ出力
    if not _kernel_log_shown:
        print(f"[xformers] memory_efficient_attention: selected kernel = {getattr(op, 'NAME', str(op))}")
        _kernel_log_shown = True

    out, *_ = op.apply(inp, needs_gradient=False)
    return out.reshape(output_shape)

def process_images_inner(p: StableDiffusionProcessing) -> Processed:
    # --- Add for xformers kernel log reset and A1111 log ---
    try:
        from xformers.ops.fmha import reset_kernel_log
        reset_kernel_log()
    except ImportError:
        pass
    print("[A1111] xformers.memory_efficient_attention called (FA-3 or best available kernel will be used)")
    # --- End add ---
    reset_fa3_log()
    # ...既存の処理...

# Always try to import xformers if available
try:
    import xformers.ops
    shared.xformers_available = True
    print("[A1111] xformers successfully imported")
except Exception as e:
    print(f"[A1111] Cannot import xformers: {e}")
    shared.xformers_available = False

def is_available(self):
    # Enable xformers if it's available and CUDA is available (no upper cap on compute capability)
    return shared.xformers_available and torch.cuda.is_available() and (6, 0) <= torch.cuda.get_device_capability(shared.device)

def xformers_attention_forward(self, x, context=None, mask=None, **kwargs):
    global FA3_LOGGED_THIS_GEN
    kwargs.pop('additional_tokens', None)
    try:
        import xformers.ops
        h = self.heads
        q_in = self.to_q(x)
        context = context if context is not None else x
        k_in = self.to_k(context)
        v_in = self.to_v(context)
        
        # シンプルなreshape（head dimension制限は削除済み）
        q, k, v = (t.reshape(t.shape[0], t.shape[1], h, -1) for t in (q_in, k_in, v_in))
        
        del q_in, k_in, v_in

        # 出力dtypeは入力xに合わせる
        dtype = x.dtype
        if shared.opts.upcast_attn:
            q, k, v = q.float(), k.float(), v.float()
        else:
            # Forgeに寄せて、FAカーネルを使うためfloat32なら半精度で実行→出力は元dtypeへ戻す
            if q.dtype == torch.float32:
                q = q.half(); k = k.half(); v = v.half()

        # Forgeのシンプルな実装を参考に、カスタムopを指定
        out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=mask, op=get_xformers_flash_attention_op(q, k, v))
        
        if not FA3_LOGGED_THIS_GEN:
            # 使用されたカーネルを確認
            kernel_info = get_xformers_kernel_info()
            if kernel_info:
                print(f"[A1111] xformers.memory_efficient_attention called - {kernel_info} kernel detected")
            else:
                print("[A1111] xformers.memory_efficient_attention called (FA-3→FA-2→Cutlass priority order)")
            FA3_LOGGED_THIS_GEN = True
        
        out = out.to(dtype)
        b, n, h, d = out.shape
        out = out.reshape(b, n, h * d)
        return self.to_out(out)
    except Exception as e:
        print(f"[A1111] xformers.memory_efficient_attention failed, falling back to SDP. Exception: {e}")
        # ... フォールバック処理

def get_xformers_flash_attention_op(q, k, v):
    # Forgeの実装を参考に、より積極的にFA-2を優先
    try:
        # FA-2を優先的に試す
        flash_attention_op = xformers.ops.MemoryEfficientAttentionFlashAttentionOp
        fw, bw = flash_attention_op
        if fw.supports(xformers.ops.fmha.Inputs(query=q, key=k, value=v, attn_bias=None)):
            return flash_attention_op
    except Exception as e:
        # FA-2が使用できない場合はNoneを返してxformersに自動選択させる
        pass
    return None

def get_xformers_kernel_info():
    """Get information about the last used xformers kernel"""
    try:
        from xformers.ops.fmha import get_last_used_kernel
        kernel_name = get_last_used_kernel()
        if kernel_name:
            if 'fa2' in kernel_name.lower() or ('flash' in kernel_name.lower() and '2' in kernel_name):
                return "FA-2 (2nd priority)"
            elif 'fa3' in kernel_name.lower() or ('flash' in kernel_name.lower() and '3' in kernel_name):
                return "FA-3 (1st priority)"
            elif 'cutlass' in kernel_name.lower():
                return "Cutlass (3rd priority)"
            else:
                return f"{kernel_name} (unknown priority)"
    except ImportError:
        pass
    return None

def check_and_configure_fa3():
    """Check FA-3 availability and configure xformers accordingly"""
    try:
        import xformers.ops.fmha.dispatch
        import xformers.ops.fmha.flash3
        
        # Check if FA-3 is actually available
        fa3_available = xformers.ops.fmha.dispatch.fa3_available()
        
        if fa3_available:
            print("[A1111] FA-3 is available - enabling FA-3 priority")
            xformers.ops.fmha.dispatch._set_use_fa3(True)
        else:
            print("[A1111] FA-3 is not available - disabling FA-3 priority (FA-2 will be used)")
            xformers.ops.fmha.dispatch._set_use_fa3(False)
            
    except Exception as e:
        print(f"[A1111] Error checking FA-3 availability: {e}")
        # Default to False if we can't check
        try:
            import xformers.ops.fmha.dispatch
            xformers.ops.fmha.dispatch._set_use_fa3(False)
        except:
            pass

def check_and_configure_fa3():
    """Check FA-3 availability and configure xformers accordingly"""
    try:
        import xformers.ops.fmha.dispatch
        import xformers.ops.fmha.flash3
        
        # Check if FA-3 is actually available
        fa3_available = xformers.ops.fmha.dispatch.fa3_available()
        
        if fa3_available:
            print("[A1111] FA-3 is available - enabling FA-3 priority")
            xformers.ops.fmha.dispatch._set_use_fa3(True)
        else:
            print("[A1111] FA-3 is not available - disabling FA-3 priority (FA-2 will be used)")
            xformers.ops.fmha.dispatch._set_use_fa3(False)
            
    except Exception as e:
        print(f"[A1111] Error checking FA-3 availability: {e}")
        # Default to False if we can't check
        try:
            import xformers.ops.fmha.dispatch
            xformers.ops.fmha.dispatch._set_use_fa3(False)
        except:
            pass

if selection == "Automatic" and len(optimizers) > 0:
    # Prefer xformers when available or explicitly requested
    try:
        xformers_opt = next((x for x in optimizers if isinstance(x, sd_hijack_optimizations.SdOptimizationXformers)), None)
    except Exception:
        xformers_opt = None

    if xformers_opt is not None and xformers_opt.is_available() and (getattr(shared.cmd_opts, "xformers", False) or getattr(shared.cmd_opts, "xformers_flash_attention", False) or getattr(shared, "xformers_available", False)):
        matching_optimizer = xformers_opt
    else:
        matching_optimizer = next(iter([x for x in optimizers if x.cmd_opt and getattr(shared.cmd_opts, x.cmd_opt, False)]), optimizers[0])

# 1生成ごとに1回だけカーネル名をログ出力
if not _kernel_log_shown:
    global _last_used_kernel
    kernel_name = getattr(op, 'NAME', str(op))
    _last_used_kernel = kernel_name
    
    # カーネル名をより詳細に解析
    if 'flash' in kernel_name.lower():
        if '3' in kernel_name or 'fa3' in kernel_name.lower():
            kernel_type = "Flash Attention 3 (FA-3)"
            priority_info = "✓ Highest priority (FA3→FA2→Cutlass)"
        elif '2' in kernel_name or 'fa2' in kernel_name.lower():
            kernel_type = "Flash Attention 2 (FA-2)"
            priority_info = "✓ Second priority (FA3→FA2→Cutlass)"
        else:
            kernel_type = "Flash Attention"
            priority_info = "? Unknown Flash Attention version"
    elif 'cutlass' in kernel_name.lower():
        kernel_type = "Cutlass"
        priority_info = "⚠ Third priority (FA3→FA2→Cutlass) - FA3/FA2 unavailable"
    elif 'ck' in kernel_name.lower():
        kernel_type = "CK"
        priority_info = "? CK kernel selected"
    elif 'triton' in kernel_name.lower():
        kernel_type = "Triton"
        priority_info = "? Triton kernel selected"
    else:
        kernel_type = kernel_name
        priority_info = "? Unknown kernel type"
    
    print(f"[xformers] memory_efficient_attention: selected kernel = {kernel_type} ({kernel_name}) - {priority_info}")
    _kernel_log_shown = True

def get_last_used_kernel():
    """Get the name of the last used kernel"""
    global _last_used_kernel
    return _last_used_kernel

def reset_kernel_log():
    """Reset the kernel log for the next generation"""
    global _kernel_log_shown, _last_used_kernel
    _kernel_log_shown = False
    _last_used_kernel = None

_USE_FLASH_ATTENTION_3 = True

def _dispatch_fw_priority_list(
    inp: Inputs, needs_gradient: bool
) -> Sequence[Type[AttentionFwOpBase]]:
    if torch.version.cuda:
        flash3_op = [flash3.FwOp] if _get_use_fa3() else []
        
        # Check if FA-2 is available
        fa2_available = False
        fa2_reasons = []
        try:
            from xformers.ops.fmha import flash
            fa2_available = flash.FwOp.supports(inp)
            if not fa2_available:
                fa2_reasons = flash.FwOp.not_supported_reasons(inp)
        except Exception as e:
            fa2_reasons = [f"Exception: {e}"]
        
        # Debug output for first few calls
        import os
        if not hasattr(_dispatch_fw_priority_list, '_debug_count'):
            _dispatch_fw_priority_list._debug_count = 0
        _dispatch_fw_priority_list._debug_count += 1
        
        if True:  # デバッグ情報を有効化
            print(f"[xformers] Dispatch debug #{_dispatch_fw_priority_list._debug_count}:")
            print(f"  Input shapes: q={inp.query.shape}, k={inp.key.shape}, v={inp.value.shape}")
            print(f"  Input dtypes: q={inp.query.dtype}, k={inp.key.dtype}, v={inp.value.dtype}")
            print(f"  FA-3 enabled: {_get_use_fa3()}")
            print(f"  FA-2 available: {fa2_available}")
            if not fa2_available:
                print(f"  FA-2 not supported reasons: {fa2_reasons}")
        
        # Build priority list based on availability
        if fa2_available:
            priority_list_ops = deque(
                flash3_op
                + [
                    flash.FwOp,  # FA-2
                    cutlass.FwOp,
                ]
            )
        else:
            priority_list_ops = deque(
                flash3_op
                + [
                    cutlass.FwOp,
                ]
            )

from modules import script_callbacks, sd_hijack_optimizations, sd_hijack
script_callbacks.on_list_optimizers(sd_hijack_optimizations.list_optimizers)
sd_hijack.list_optimizers()

# Check and configure FA-3 availability
sd_hijack_optimizations.check_and_configure_fa3()

startup_timer.record("scripts list_optimizers")

print("[A1111] xformers.memory_efficient_attention called (FA-3→FA-2→Cutlass priority order)")