def _can_use_flash_attention(
    query, key, value, mask=None, is_causal=False, raise_error=False
):
    """Verify the availability of flash attention."""
    try:
        from torch.backends.cuda import SDPAParams
        from torch.backends.cuda import can_use_flash_attention
    except ImportError:
        if raise_error:
            raise ImportError(
                "Flash attention is not supported in your current PyTorch "
                "version. Please update it by following the official guide: "
                "https://pytorch.org/get-started/locally/"
            )
        return False

    try:
        spda_params = SDPAParams(
            query,
            key,
            value,
            mask,
            0.0,  # dropout_p
            is_causal,
            False,  # enable_gqa
        )
    except TypeError:
        # The old function signature for the older version of PyTorch
        spda_params = SDPAParams(
            query,
            key,
            value,
            mask,
            0.0,  # dropout_p
            is_causal,
        )
    if raise_error and can_use_flash_attention(spda_params, True) is False:
        raise RuntimeError(
            "Flash attention is not supported with the provided inputs. "
            "Please check the warnings for more details."
        )
    return can_use_flash_attention(spda_params, False)
Verify the availability of flash attention.
keras/src/backend/torch/nn.py