Summary

This PR adds support for ALiBi (Attention with Linear Biases) in TorchInductor’s fused-attention. ALiBi applies a position-based bias to attention scores, improving extrapolation for language modeling tasks. With this addition, ALiBi-based attention can leverage PyTorch’s optimized _scaled_dot_product_attention kernel.

Changes

• New ALiBi Pattern & Replacement
  • _sfdp_pattern_alibi(...): Recognizes [Q @ Kᵀ / √d + alibi_bias] → softmax → dropout → matmul(V).
  • _sfdp_replacement_alibi(...): Fuses into _scaled_dot_product_attention using attn_mask=alibi_bias.
• Test
  • Added _test_sdpa_rewriter_alibi in TestSDPAPatternRewriterTemplate.
  • Confirms forward/backward correctness under dropout.
  • If you get error: torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised: RuntimeError: Duplicate pattern: expand_default = CallFunction(aten.expand.default, KeywordArg('query'), Ignored()),
    -> run export PYTORCH_GEN_PATTERNS=1 in the terminal to generate the attention pattern.

Notes

• If FlashAttention does not support ALiBi directly, PyTorch gracefully falls back to MATH or MEM-EFFICIENT kernels.
• Combining ALiBi with a causal mask can be done by summing the bias and mask if needed.

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang @aakhundov