Add Metal/MPS GPU backend for the ungapped prefilter on Apple Silicon

[fnachon]

Summary

Adds a Metal-based counterpart to lib/libmarv (the CUDA/HIP GPU backend), so --gpu 1 works on Apple Silicon Macs, behind a new ENABLE_MPS CMake option (mutually exclusive with ENABLE_CUDA).

Scope: accelerates the default ungapped prefilter (AlignmentType::GAPLESS, i.e. --prefilter-mode 0, the same mode --gpu 1 uses on CUDA today). The gapped Smith-Waterman GPU kernel (--prefilter-mode 2) is not ported and remains CUDA/HIP-only; the constructor throws a clear error if that mode is requested on this backend.

Design

• lib/libmarv-mps/: a Metal implementation of the Marv class declared in the existing lib/libmarv/src/marv.h, reused unchanged so GpuUtil.cpp / gpuserver.cpp / ungappedprefilter.cpp need no backend-specific branching beyond HAVE_MPS compile guards mirroring the existing HAVE_CUDA ones.
• Kernel: one Metal simdgroup (32 lanes) cooperates on a single alignment, mirroring libmarv's own warp-cooperative design — each lane owns a slice of the query in threadgroup memory, with a single simd_shuffle_up per subject residue carrying the one cross-lane dependency the recurrence has. The PSSM is uploaded pre-transposed for coalesced access. Work-stealing dispatch (a shared atomic counter) means simdgroups pull the next database entry as they finish, rather than a static one-simdgroup-per-sequence assignment — this matters a lot on real, length-skewed protein databases where a handful of long sequences would otherwise dominate a naive dispatch's wall-clock while short-sequence threads sit idle.
• marv.h gains two new, additive public methods, submitScan()/collectScan() — an optional async 2-phase form of scan() that lets the query loop in ungappedprefilter.cpp overlap GPU execution of query i with building query i+1's profile on the CPU. Only the Metal backend implements them; CUDA/HIP is unaffected since nothing calls them under HAVE_CUDA builds.
• A register-tiled kernel variant (closer to libmarv's actual register-resident design) was also tried and measured slower on Apple Silicon (details/data in the header comment of marv_mps.mm) — kept as documented history rather than shipped, since it regressed performance.

Validation

Tested on an Apple M5 Max:

• Output cross-checked against a scalar CPU oracle and against MMseqs2's own CPU prefilter path (score differences there are explained by the CPU path's intentional 8-bit saturating SIMD arithmetic, not a correctness issue — verified via self-alignment scores, which must be monotonic under BLOSUM62 and match the unsaturated GPU values exactly).
• 25 synthetic boundary-length test cases plus full real-database runs (up to 153.7M sequences, UniRef90) produce byte-identical results across kernel revisions.
• On a real search (one query against UniRef90, 153.7M sequences), GPU prefilter measured ~8.5x faster than the CPU path's k-mer-index prefilter (~1m22s vs ~14m15s) — the CPU path pays a large index-construction cost that scales with database size, which the GPU path (scanning a pre-padded database directly) avoids.

Build

cmake -DENABLE_MPS=1 -DHAVE_ARM8=1 -DCMAKE_BUILD_TYPE=Release ..
make -j

Usage is identical to the existing CUDA GPU workflow (makepaddedseqdb + --gpu 1).

Test plan

• Clean CMake configure + build from scratch with -DENABLE_MPS=1
• Confirm -DENABLE_CUDA=1 / default CPU-only configure is unaffected (byte-identical CMake behavior, no HAVE_MPS-guarded code compiled)
• Correctness: byte-identical output vs. scalar CPU oracle and vs. CPU prefilter (modulo saturation) across boundary-length and real-database tests
• CI/maintainer review on actual CUDA/HIP hardware to confirm the HAVE_CUDA-guarded paths are untouched (I don't have NVIDIA/AMD hardware to test this myself — the diffs there are mechanical macro-condition extensions, but worth a maintainer's eyes)

🤖 Generated with Claude Code

[milot-mirdita]

Interesting, i was playing with this recently too:
https://github.com/milot-mirdita/mmseqs2/tree/metal-gpu-backend

[fnachon]

Thanks for the pointer, @milot-mirdita — A clearly more sophisticated implementation!
I went and actually built and ran your metal-gpu-backend branch side by side with this one. It actually helped me find a bug.

After the fix, I re-ran a real query (a protein I work with, ~290 aa) against the full UniRef90 (153.7M sequences) on both implementations:

• Correctness: 280/281 hits byte-identical (same target, %id, e-value, bitscore). The one difference is a boundary tie-break at the 300-hit prefilter cutoff.
• Performance: your implementation is noticeably faster on this real workload — GPU prefilter stage 14.9s vs my 1m13s (~4.9x), full easy-search 43.0s vs 1m44.6s (~2.4x). Haven't profiled precisely where the time goes on my side yet.

Given that, it makes more sense for the project to converge on your approach!