`koine_machines.inference`: no MPS branch, so on Apple Silicon the published `ink_9um` checkpoints run on the CPU with the GPU idle and nothing is printed

[SurgeFok]

In one sentence: koine_machines.inference has no MPS branch, so Apple Silicon silently runs the published ink_9um checkpoints on the CPU while the GPU sits idle.

I was trying to: Run the published ink_9um checkpoint through the documented First Letters inference workflow on an Apple Silicon Mac, then shorten the turnaround between repeat inference runs.

Using: Current merge-ink-pipelines, torch 2.14.0, an Apple M5 Pro with 24 GB unified memory, a 512x512x20 uint8 zarr input, and python -m koine_machines.inference.infer input.zarr step-075000.pth out.tif --batch-size 1 --no-compile.

What happened: Inference completed with correct output, but selected the CPU without printing a warning; the MPS GPU stayed idle.

What I expected or needed: Automatic MPS selection on Apple Silicon, an explicit device log, and a way to request a device without changing the CUDA default for other users.

Evidence / reproduction: The source path and measured CPU/MPS comparison are below, together with a tested patch.

• I personally encountered or reproduced this using the version and data stated above.

Details

Both ink-detection inference entrypoints resolve the device with CUDA or nothing.
koine_machines/inference/infer.py, in prepare_model_for_inference:

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

infer_full3d_tifxyz.py has the same line in prepare_model. Autocast is gated the same
way, in both files:

autocast_enabled = bool(amp and device.type == "cuda")
...
torch.autocast(device_type="cuda", enabled=True, dtype=amp_dtype)

So on any Mac both entrypoints take the CPU path and disable autocast. The run succeeds and
the output is correct, so there is nothing to notice.

Measured on hardware

Apple M5 Pro, 24 GB unified memory, macOS 26.5.1, torch 2.14.0, checkpoint
hybrid_3d2d-seed42/step-075000.pth from scrollprize/ink_9um, through the documented CLI:

Five timed runs per device after a discarded warm-up, 49 blocks per run, block/s:

device	min	median	max	mean	stdev	log line
mps	33.79	33.88	33.99	33.88	0.08	mps autocast enabled for inference with dtype=float16
cpu	12.65	12.83	13.18	12.88	0.22	Autocast disabled for inference (device=cpu)

The cpu row is current merge-ink-pipelines on this machine. Comparing the two output
TIFFs pixel by pixel: maxdiff=1 on uint8, mean absolute difference 0.084, no pixel
differing by more than one level, which is fp16 accumulation.

Median against median that is 2.64x, on a 512x512 input giving 49 blocks. Model load and
imports dominate wall clock at that size, so a real segment should do better rather than worse.

Why this one is worth fixing

The First Letters workflow post is the front door for new entrants, and it has readers verify
that torch sees CUDA. Its only tuning advice is to lower --batch-size if GPU memory runs
out. Neither statement applies on a Mac, so the people most exposed are the ones with the
least context to spot it: someone following the guide on their own laptop, concluding ink
detection is slow, and running fewer refinement rounds because of it.

The rest of the repo already resolves MPS

vesuvius/src/vesuvius/models/utilities/get_accelerator.py does cuda, then mps, then cpu,
and logs Using MPS device (Apple Silicon). nnU-Net under segmentation/models/arch/nnunet
handles MPS as well. koine_machines is a separate package that predates the helper, so it
never picked it up. Nothing new needs inventing here.

Three things I hit while writing a fix

Recording these in case they save someone the same round trip:

• Under MPS, setting pin_memory makes torch emit 'pin_memory' argument is set as true but not supported on MPS now, device pinned on every loader iteration, while pinning nothing.
  It has to stay gated on CUDA.
• CPU autocast is bfloat16-only, and on the machines that reach the CPU branch I could not
  measure it beating plain float32, so leaving it off seems right.
• --gpus indexes CUDA ordinals and feeds nn.DataParallel, so it cannot mean anything on
  MPS. Raising on --gpus 0 from a Mac is better than quietly resolving it somewhere else.

On the autocast side, torch.autocast(device_type="mps") ran for float16 and for bfloat16,
and the backend default resolves to float16. torch.compile also works on MPS here. I have
not tried it in combination with autocast.

A working patch

I have a fix running, with tests, at
https://github.com/SurgeFok/vesuvius-mps-inference (MIT). It adds
koine_machines/common/accelerator.py mirroring get_accelerator, routes both entrypoints
through it, passes device_type=device.type to autocast, and adds a --device flag with
auto/cuda/mps/cpu. The applied diff is mps-inference.patch in that repo.

One limitation, stated plainly: I have no machine with a usable CUDA GPU, so I could not run
the CUDA branch, which is the default path. I kept its control flow identical and reproduced
both error strings character for character including visible device count is N, and there
are five tests that fake torch.cuda.is_available and device_count to check the decisions.
Somebody with a CUDA box running pytest tests/ -v would close that.

Happy to open a PR if outside PR creation is working again (#1743), otherwise the patch is
there to take.

[SurgeFok]

Full-segment follow-up: I ran the patch on the public PHerc0800 20251028222030-auto_grown_20251028222030940 surface volume. Level 0 is 31x2580x2580 uint8, and the same step-075000.pth checkpoint scheduled 923 blocks per run. This used the same M5 Pro and torch 2.14.0 as the original report, with batch size 1 and --no-compile. I discarded one warm-up and timed five runs per device.

device	min	median	p95	max	mean	stdev	wall median
mps	97.66	97.94	97.97	97.97	97.89	0.13	32.92 s
cpu	16.25	16.40	16.89	16.91	16.55	0.29	79.67 s

That is a 5.97x median inference-loop gain and a 2.42x median end-to-end gain. All six outputs within each backend, including the warm-up, had the same SHA-256. Comparing the retained MPS and CPU TIFFs across 6,656,400 pixels gives a maximum difference of one uint8 level, mean absolute difference 0.0277, and zero pixels differing by more than one level.

As a separate CUDA compatibility check, I ran the same input, checkpoint, batch size, and protocol on a Linux workstation with an 8 GB NVIDIA GTX 1070, driver 570.211.01, and torch 2.2.0+cu118. CUDA measured a median 50.31 blocks/s and a median whole-command time of 22.25 seconds. All six CUDA outputs had the same SHA-256. Compared with the retained MPS output, the CUDA TIFF had a maximum difference of two uint8 levels and a mean absolute difference of 0.0298. Because this used a different machine and torch build, it is not part of the MPS-versus-CPU speedup calculation.

Input:

s3://vesuvius-challenge-open-data/PHerc0800/segments/20251028222030-auto_grown_20251028222030940/surface-volumes/8.64um-1.2m-116keV-volume-20250521135224.zarr/

Per-run results, environment details, and hashes:
https://github.com/SurgeFok/vesuvius-mps-inference/blob/main/benchmarks/full-pherc0800/results.json