SIGN IN SIGN UP

[xcodec2] Fix flex attention and flash dispatch tests (#48244)

#46196 enabled `_supports_sdpa` / `_supports_flex_attn` on
`Wav2Vec2BertPreTrainedModel`. Xcodec2 uses a `Wav2Vec2BertModel` as its
semantic encoder, so two common tests that were silently skipped started
running and failed.

`test_sdpa_can_dispatch_on_flash`: wav2vec2-bert defaults to
`position_embeddings_type="relative_key"`, which feeds a dense additive
`position_bias` mask to SDPA. The flash backend does not accept an
attention mask, hence "No available kernel". Skipped, with the same reason
already used in the wav2vec2-bert tests.

`test_flex_attention_with_grads`: `_prepare_config_headdim` scales
`hidden_size` on the config and, recursively, on every sub-config, but the
tester hardcodes `quantization_dim` to `hidden_size +
semantic_model_config.hidden_size`. That invariant broke, so the quantizer
received 128 features into a `Linear(64, 4)`. `_prepare_config_headdim` is
overridden to restore it, like janus/blip_2/instructblip already do for
their own config invariants.
J
jiqing-feng committed
c7cf04b1e3b1d497dbb1473c2e65e75ee69e12dc
Parent: f440843
Committed by GitHub <noreply@github.com> on 8/24/2026, 7:24:05 AM