GPT2 attention scaling config is ignored when using SDPA / FlashAttention backends

[Qi-Zhan]

System Info

None

Who can help?

@ArthurZucker Hi, I'm new to LLMs and currently learning GPT2 model. I found that

The GPT2 attention configuration options:
• scale_attn_weights
• scale_attn_by_inverse_layer_idx

are respected in eager attention mode but silently ignored when using AttentionInterface backends such as "sdpa" or "flash_attention_2".

In eager mode:
the scaling logic is applied inside eager_attention_forward:
• division by sqrt(head_dim) if scale_attn_weights=True
• division by (layer_idx+1) if scale_attn_by_inverse_layer_idx=True

However, when using sdpa:
torch._C._nn.scaled_dot_product_attention(query_states_3, key_states_3, value_states_3, attn_mask = attention_mask_1, dropout_p = 0.0, scale = None, is_causal = False) which seems to ignore the above config.

I realize that the default configuration (scale_attn_weights =True, scale_attn_by_inverse_layer_idx=False) produces the same results, so I’m not sure whether this is intentional or should be considered a bug.

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

None

Expected behavior

Different attention implementations should produce semantically equivalent results and respect model configuration parameters.

[Saad-Mallebhari]

@Qi-Zhan
You're right that this is a real inconsistency. The scale parameter in scaled_dot_product_attention is being passed as None, which defaults to 1/sqrt(head_dim) , matching scale_attn_weights=True. But scale_attn_by_inverse_layer_idx has no equivalent in the SDPA call, so it's silently ignored when using SDPA or FlashAttention backends.

The fix would be to compute the scale explicitly before calling SDPA and pass it in:

scale = 1.0 / (head_dim ** 0.5) if config.scale_attn_weights else 1.0
if config.scale_attn_by_inverse_layer_idx:
    scale = scale / (layer_idx + 1)

torch._C._nn.scaled_dot_product_attention(..., scale=scale)

This ensures all three backends produce semantically equivalent results regardless of config. The default config (scale_attn_weights=True, scale_attn_by_inverse_layer_idx=False) happens to match SDPA's default behavior, which is why it's not immediately obvious.

[DhruvTilva]

@Saad-Mallebhari @Qi-Zhan will try to resolve this, On it..

[vasqu]

Heya just seeing this now, coming from the PR here #44397

Not to double the effort @DhruvTilva, you can also help out on the review or similar if you want to. Really appreciate digging into this