The 256-cap inside the flex_gemm_sparse_attn branch was a hold-over from
the early naive Metal kernel (round 3). The current backend is flash-
attention-v2 with simdgroup matmul + simd-shuffle softmax row reductions
and wins at every measured shape including max_seqlen=2048. The CUDA
backends above (xformers / flash_attn / flash_attn_3) never fork on
max_seqlen — match that precedent.
Safety-valve preserved as FLEX_GEMM_ATTN_MAX_SEQLEN=N env var: when set,
falls back to SDPA-padded above the cap. Useful only on PyTorch builds
where the Accelerate-SDPA-CPU-bounce happens to win at a specific shape
(measured crossover sits beyond 768 on fp32, higher on fp16).
Wires the new mtlgemm fused sparse attention kernel through the ATTN
backend selector. Dispatches to the fused Metal kernel when
max(max_q, max_kv) <= 256 (where the naive per-thread-serial-KV kernel
beats SDPA-padded on M3 Max), and falls through to an inline SDPA-padded
path for larger max_seqlen. Opt in via ATTN_BACKEND=flex_gemm_sparse_attn
or SPARSE_ATTN_BACKEND=flex_gemm_sparse_attn. Default on Darwin stays
'sdpa' — the threshold-based fallback doesn't yet prove a universal
win across the pipeline's attention shape distribution.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
