f3177cd807
The flex_gemm_sparse_attn backend is now flash-attention-v2 with simdgroup_matrix_multiply_accumulate for Q@K^T and P@V plus simd-shuffle softmax row reductions, and wins 5–15× over SDPA-padded-CPU-bounce at every measured shape including max_seqlen=2048. Production decoder block on M3 Max: 5.04× wall-clock vs all-SDPA-padded baseline (3.49 ms vs 17.57 ms), entirely from this change being the default. The existing __flex_gemm_works_on_mps() probe covers the same package the attention path lives in, so the gate is identical: if conv probe passes, set both CONV='flex_gemm' and ATTN='flex_gemm_sparse_attn'. SPARSE_ATTN_BACKEND= (or ATTN_BACKEND=) env override is unchanged. Also fix benchmarks/e2e_decoder.py to inject the repo root into sys.path so it runs as `python benchmarks/e2e_decoder.py` from any cwd. |
||
|---|---|---|
| .. | ||
| datasets | ||
| models | ||
| modules | ||
| pipelines | ||
| renderers | ||
| representations | ||
| trainers | ||
| utils | ||
| __init__.py | ||
| backends.py | ||