c869eb995a
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> |
||
|---|---|---|
| .. | ||
| datasets | ||
| models | ||
| modules | ||
| pipelines | ||
| renderers | ||
| representations | ||
| trainers | ||
| utils | ||
| __init__.py | ||
| backends.py | ||