perf: grouped MoE path and runtime defaults for Apple Silicon (MPS)

[kkilchrist]

Builds on #22, which is included here as the first commit (credited to @Berkkirik) because this work hard-depends on its device resolution and CUDA-only Triton gate. If #22 merges first, this rebases to a clean single-commit diff. #22 makes MPS work; this PR makes it fast and safe by default.

Results

Measured on M4 Pro (20-core GPU, 48 GB), macOS 27.0, Python 3.13.11, torch 2.12.1, released 1.5B checkpoint, ~1.5k-token document:

Configuration	Throughput	Parity vs CPU
CPU	~65 tok/s	—
MPS, #22 alone (per-token gather MoE)	~173 tok/s (2.6×)	identical spans
MPS + this PR (grouped MoE)	~4,300 tok/s (~65×)	identical spans, identical eval metrics

Throughput holds on longer inputs (~4,400 tok/s on a 50k-token document).

Changes

1. Grouped pure-torch MoE path (opf/_model/model.py). Reuses the Triton branch's sort-by-expert packing, but scatters packed rows into a padded per-expert batch and runs one bmm per projection, gathering results back afterward. Each expert's weights are read once per layer instead of ~20 MB of gathered weights per token, which is what makes the per-token fallback path memory-bound. Default-on for MPS via OPF_MOE_GROUPED (Triton cannot target Metal); CPU keeps the existing fallback path, CUDA keeps Triton. ~24× over the gather path on MPS.

2. Default n_ctx capped at 4096 on non-CUDA devices (opf/_core/runtime.py). Windows are non-overlapping (example_to_windows stride == window size) and the attention band is ±128 tokens, so larger windows don't improve throughput — but transient attention memory scales with window length. Measured on MPS with a 12.5k-token document: n_ctx 4096 → 3,969 tok/s; 16384 → 3,835 tok/s at ~2.3× the memory; the checkpoint's 128k default gets the process OOM-killed on a 50k-token input even with 48 GB of unified memory. Same rationale as the existing CPU cap; --n-ctx still overrides.

3. Python API defaults (opf/_api.py). OPF(device=...) now defaults to "auto" and the Literal types accept "mps"/"auto" — feat: auto-detect Apple Silicon (MPS) and keep Triton CUDA-only #22 covers the CLI, but API users on a Mac would still crash on the previous hardcoded "cuda" default.

4. MPS-safe MXFP4 decode (opf/_model/weights.py). torch.ldexp has no MPS kernel (DispatchStub: missing kernel for mps); replaced with an exp2 multiply, which is the same exact power-of-two scaling. Verified bit-identical to ldexp on CPU. (The released checkpoint ships fused bf16 tensors, so this only matters for MXFP4-packed checkpoints loaded directly to MPS.)

5. Eval preprocess workers on MPS (opf/_eval/runner.py). Parallel preprocessing now applies on MPS as on CUDA; only CPU inference competes with preprocessing for cores.

Verification

• opf "Alice was born on 1990-01-02." auto-selects MPS and redacts correctly.
• CPU vs MPS predicted spans identical on multi-thousand-token synthetic PII documents (105 and ~3,430 spans respectively).
• opf eval examples/data/sample_eval_five_examples.jsonl produces identical precision/recall/F1 on --device cpu and --device mps.
• Grouped-path numerics: the padded bmm computes the same per-expert projections as the Triton/gather paths; parity above is exact at the span level.

🤖 Generated with Claude Code