keras-team/keras · error · ValueError

Flash attention is not supported in numpy backend.

Error message

Flash attention is not supported in numpy backend.

What it means

`dot_product_attention` in the NumPy backend rejects `flash_attention=True`. Flash attention is a fused kernel optimization available only on GPU backends (via cuDNN/Triton in the Torch backend); the pure-NumPy implementation computes attention with standard matrix multiplications and has no flash path, so the flag is treated as a hard error instead of being ignored.

Solutions

1. Leave `flash_attention=False` (or unset) when running on the NumPy backend; results are numerically equivalent with the standard attention path
2. Switch to the Torch backend on a CUDA GPU if flash attention kernels are required for performance

When it happens

Trigger: Thrown at keras/src/backend/numpy/nn.py:1426 when the library encounters an invalid state.

Common situations: See trigger scenarios.

AI-assisted analysis of keras-team/keras@7a34a03db6 (2026-08-25). Data as JSON: /api/errors/46e9eba63604d604. Report an issue: GitHub.

Appendix: source

Thrown at keras/src/backend/numpy/nn.py:1426

    encoded = encoded.astype(encoded_dtype)
    return encoded


def dot_product_attention(
    query,
    key,
    value,
    bias=None,
    mask=None,
    scale=None,
    is_causal=False,
    flash_attention=None,
    attn_logits_soft_cap=None,
):
    if flash_attention is None:
        flash_attention = False
    if flash_attention:
        raise ValueError("Flash attention is not supported in numpy backend.")

    # Ref: jax.nn.dot_product_attention
    # https://github.com/jax-ml/jax/blob/jax-v0.4.32/jax/_src/nn/functions.py#L828
    # Not support `query_seq_lengths` and `key_value_seq_lengths` args
    query = convert_to_tensor(query)
    key = convert_to_tensor(key)
    value = convert_to_tensor(value)
    if len(query.shape) != 4:
        raise ValueError(
            "`dot_product_attention` only supports 4D inputs. "
            f"Received: query.shape={query.shape}, key.shape={key.shape}, "
            f"value.shape={value.shape}."
        )
    compute_dtype = backend.result_type(query.dtype, key.dtype, value.dtype)
    query = cast(query, compute_dtype)
    key = cast(key, compute_dtype)
    value = cast(value, compute_dtype)
    if bias is not None:

View on GitHub (pinned to 7a34a03db6)