keras-team/keras · error · ValueError

Returning attention scores is not supported when flash atten

Error message

Returning attention scores is not supported when flash attention is enabled. Please disable flash attention to access attention scores.

What it means

GroupedQueryAttention can use a fused flash-attention kernel, but flash attention does not materialize attention score matrices, so the layer refuses the combination of flash_attention enabled together with return_attention_scores=True. The check runs at the top of _compute_attention during call(), so the error appears on the first forward pass, not at construction.

Source

Thrown at keras/src/layers/attention/grouped_query_attention.py:466

            ops.arange(q_seq_length, dtype="int32"), (1, q_seq_length, 1)
        )
        col_index = ops.reshape(
            ops.arange(v_seq_length, dtype="int32"), (1, 1, v_seq_length)
        )
        return ops.less(ops.abs(row_index - col_index), self.sliding_window)

    def _compute_attention(
        self,
        query,
        key,
        value,
        attention_mask=None,
        training=None,
        use_causal_mask=False,
    ):
        # Check for flash attention constraints
        if self._flash_attention and self._return_attention_scores:
            raise ValueError(
                "Returning attention scores is not supported when flash "
                "attention is enabled. Please disable flash attention to access"
                " attention scores."
            )

        # Determine whether to use dot-product attention
        use_dot_product_attention = not (
            self.dropout > 0.0
            or self._return_attention_scores
            or (len(query.shape) != 4)
        )

        if use_dot_product_attention:
            if use_causal_mask and attention_mask is None:
                # Skip materializing the [T, S] mask and let the backend
                # use its native causal kernel.
                attention_output = ops.dot_product_attention(
                    query=query,

View on GitHub (pinned to 7a34a03db6)

Solutions

1. Disable flash attention on the layer: GroupedQueryAttention(..., flash_attention=False, return_attention_scores=True).
2. Or keep flash attention and drop return_attention_scores; scores are unavailable by design.
3. For visualization, run a separate cheap forward pass with a non-flash copy of the layer.

Example fix

# before
attn = GroupedQueryAttention(head_dim=64, flash_attention=True, return_attention_scores=True)
out, scores = attn(x)  # -> ValueError

# after
attn = GroupedQueryAttention(head_dim=64, flash_attention=False, return_attention_scores=True)
out, scores = attn(x)

if getattr(layer, '_flash_attention', False) and need_scores:
    layer._flash_attention = False  # or rebuild the layer with flash_attention=False

When it happens

Trigger: GroupedQueryAttention(..., flash_attention=True, return_attention_scores=True) then calling the layer; flash attention enabled by default in some configurations while interpretability code sets return_attention_scores=True.

Common situations: Interpretability code that inspects attention weights on a layer configured for speed; toggling flash_attention on for training while an attention-visualization path still exists.

Related errors

• The last dimension of `query_shape` and `value_shape` must b
• All dimensions of `value` and `key`, except the last one, mu
• Unknown activation function '{activation}' cannot be seriali
• Could not interpret activation function identifier: {identif
• ConvNeXt does not support the `channels_first` image data fo

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