transformers

9d011e83 - Fix silent SDPA math-kernel fallback for GQA when key/value head_dim > 256 or differ (#46960)

Fix silent SDPA math-kernel fallback for GQA when key/value head_dim > 256 or differ (#46960) * Fix silent SDPA math-kernel fallback for GQA with head_dim > 256 GQA layers with head_dim > 256 (e.g. Gemma4's head_dim=512 full-attention layers) passed enable_gqa=True to SDPA whenever attention_mask is None. PyTorch's mem-efficient and flash SDPA kernels do not support the enable_gqa broadcast at head_dim > 256, so SDPA silently fell back to the math kernel, materializing the S x S scores (O(S^2) memory) and OOMing at long sequences. Make use_gqa_in_sdpa return False for head_dim > 256 so we take the repeat_kv (MHA) path, which the mem-efficient kernel supports at head_dim=512. Behavior is unchanged for head_dim <= 256. Signed-off-by: Butterfingrz <2395959141@qq.com> * Fix silent SDPA math-kernel fallback for GQA when key/value head_dim > 256 Signed-off-by: Butterfingrz <2395959141@qq.com> * Fix silent SDPA math-kernel fallback for GQA when key/value head_dim differ At differing key/value head_dim (e.g. MLA models) SDPA's Flash kernel is disqualified (it requires query/key/value to share the last dim) and the mem-efficient kernel does not support the enable_gqa broadcast, so enable_gqa=True is served only by cuDNN and silently falls back to the math kernel (O(S^2)) on builds/GPUs without the cuDNN SDPA backend. Require key head_dim == value head_dim (<= 256) in use_gqa_in_sdpa so GQA-in-SDPA is only used when the Flash kernel can portably serve it; otherwise take the repeat_kv (MHA) path, which the mem-efficient kernel supports at differing head_dim. Add a regression test. Signed-off-by: Butterfingrz <2395959141@qq.com> * Test SDPA GQA stays off the math kernel Force the mem-efficient backend so the silent math fallback becomes a hard error. Signed-off-by: Butterfingrz <2395959141@qq.com> * Consolidate SDPA GQA tests into one forced-backend test Signed-off-by: Butterfingrz <2395959141@qq.com> * Parametrize the SDPA GQA test by expected backend Consolidate the edge cases into one parametrized test that asserts the enable_gqa decision and that each case (matched <= 256, head_dim > 256, MLA mismatch, explicit mask) stays on a fused kernel by restricting sdpa_kernel to the expected backend, per review. Signed-off-by: Butterfingrz <2395959141@qq.com> --------- Signed-off-by: Butterfingrz <2395959141@qq.com> Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com> Co-authored-by: vasqu <antonprogamer@gmail.com>