Proposal: add sdpa_memeff attn_implementation for shape combinations no fast backend covers

快速结论：该问题出现在使用 Transformers 的 SDPA 注意力实现（attn_implementation=”sdpa”）时，特定 head_dim 或输入布局导致 Flash/CUDNN 等快速后端拒绝、调度回退到 MATH 并出现 OOM 或 NaN 梯度的场景；优先排查是否需要对 SDPA 后端做显式固定，以及 GQA 分支是否错误传入 enable_gqa=True。

适用环境：Issue 中已确认涉及 CUDA 环境（RTX 5090 sm_120、H100 sm_90 被提及）、bf16、head_dim=320 的 Gemma 4 系列模型（如 google/gemma-4-31B、gemma-4-26B-A4B、gemma-4-E4B、gemma-4-E2B）、Qwen3.5 RLHF 训练场景；完整 Python、PyTorch、CUDA、驱动版本号在 Issue 中未明确给出。

最快修复方案：暂无针对本 Issue 全场景的一步修复方案。Issue 中建议并确认的验证方式是使用 sdpa_kernel 显式固定后端，例如 with torch.nn.attention.sdpa_kernel([torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION]): 包裹 model.generate(…)；针对 Gemma 4（head_dim > 256）的 use_gqa_in_sdpa 问题已在 main 分支修复（v5.13.0 起，via #46960），升级到该版本后可优先尝试。

注意事项：sdpa_kernel 为上下文管理器，只在其作用范围内生效；Issue 指出关闭 MATH 后端（如 torch.backends.cuda.enable_math_sdp(False)）属于进程级开关，存在副作用。此外，未加掩码的 GQA 且 head_dim ≤ 256、缺少可接受 enable_gqa 的融合核（如无 flash、torch < 2.14）的剩余问题被拆分到 #49564，本 Issue 未覆盖，不应期待本页方案解决该场景。

问题场景

用户在使用 Transformers 的 attent_implementation=”sdpa” 运行 Gemma 4 系列模型（head_dim=320）进行推理或训练时，发现 SDPA 在 Flash 与 CUDNN 后端均因库限制（head_dim ≤ 256）拒绝，最终回退到 MATH 后端，导致显存占用为 O(seq²) fp32 softmax（seq=32k 时约 74 GB，在 96 GB 显卡上 OOM）。另一类场景是 Qwen3.5 RLHF 训练中，3D position_ids 导致 4D 密集掩码物化，使 Flash 的 is_causal=True 短路失效，调度回退 MATH 并在 bf16 下出现 NaN 梯度。

Issue 最初提议新增 attn_implementation=”sdpa_memeff” 来固定 Efficient Attention 后端，维护者回复认为无需新增独立实现，直接用 sdpa_kernel 固定后端即可；讨论随后收敛到 GQA 分支判断过宽的问题。

报错原文

Proposal: add sdpa_memeff attn_implementation for shape combinations no fast backend covers

FLASH: REJECT (library cap: head_dim ≤ 256)
CUDNN: REJECT (library cap: head_dim ≤ 256)
EFFICIENT (CUTLASS mem-eff): ACCEPT
MATH: ACCEPT (but O(seq²) fp32 softmax — 74 GB at seq=32k, OOMs on 96 GB card)

原因分析

最可能的原因有两层：

1. 后端选择层面的限制。PyTorch SDPA 调度器按 CUDNN → FA → Memory efficient → MATH 的优先级选择后端。当 head_dim=320 时，Flash 与 CUDNN 均有 head_dim ≤ 256 的库限制而拒绝，只有 Efficient Attention 可用，但标准 sdpa 路径没有简单方式稳定地固定到该后端。

2. Transformers 侧 GQA 判断过宽。sdpa_attention_forward 通过 use_gqa_in_sdpa() 决定使用 repeat_kv 还是 enable_gqa=True；该判断只检查 torch 版本和 attention_mask 是否为 None，没有检查当前激活的 SDPA 后端是否支持 enable_gqa=True（目前仅 Flash 支持）。因此在 sdpa_kernel([EFFICIENT_ATTENTION]) 下运行 GQA 模型时，仍会传入 enable_gqa=True，导致调度继续失败。

维护者评论同时指出，标准 SDPA 本身就会按优先级选择可用后端，若仅想固定后端，调用 sdpa_kernel 装饰器即可，无需新增独立注意力实现。

环境排查

确认 PyTorch 版本以及是否支持 torch.nn.attention.sdpa_kernel 与 torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION。
确认 Transformers 版本：Gemma 4 的 use_gqa_in_sdpa 修复从 v5.13.0 起包含在 main 中（via #46960）；低于该版本可能仍会触发 head_dim > 256 时的 GQA 路径问题。
确认模型 head_dim：Gemma 4 系列为 320，超过 Flash/CUDNN 的 256 上限。
确认是否在 GQA 模型上使用 num_heads_q != num_heads_kv，以及是否传入 attention_mask is None。
确认 CUDA 显卡架构（Issue 提及 RTX 5090 sm_120 与 H100 sm_90），不同架构上后端支持情况可能不同。
确认是否已用 torch.backends.cuda.enable_math_sdp(False) 等全局开关改动后端状态，避免与其他代码产生交互副作用。

解决步骤

先升级 Transformers 到包含 #46960 的版本（v5.13.0 及以后），使 head_dim > 256 的层在 use_gqa_in_sdpa 中返回 False，改走 repeat_kv 路径，从而被 memory-efficient 内核接受。
仅需固定后端时，不要新增独立 attn_implementation，直接在调用处使用上下文管理器包裹生成或前向逻辑：
with torch.nn.attention.sdpa_kernel([torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION]):
model.generate(...)
若处于 GQA 且 enable_gqa=True 导致调度回退，检查 use_gqa_in_sdpa 的实现：Issue 建议其同时检查 torch.backends.cuda.flash_sdp_enabled()（以及未来支持时的 cuDNN），或在不确定时默认选择 repeat_kv；该改动属于上游代码修改，可优先尝试在本地分支验证。
若问题发生在无 flash、PyTorch 版本低于 2.14、且 head_dim ≤ 256 的未掩码 GQA 场景，本 Issue 的修复不覆盖，应跟踪 #49564。

验证方法

在升级或应用上述修改后，对 head_dim > 256 的 Gemma 4 模型在 sdpa_kernel([EFFICIENT_ATTENTION]) 下执行前向或 generate，确认不再回退到 MATH、不再 OOM，并检查显存与输出与可用后端下的结果一致。Issue 作者在本地用 8B Gemma-4-E4B-it（bf16，真实权重）验证过：在标准 sdpa 可正常工作的场景下 top-1 一致率 100%，在标准 sdpa 回退 MATH 并 OOM 的场景下正确执行。对 GQA 路径的修改，可构造 GQA 形状并在固定 EFFICIENT 后端下运行，确认不再传入 enable_gqa=True 导致调度失败。