perf(predict): GPU channel reduction + pinned patch offload; fix CPU device default

[vboussot]

Summary

Run the per-model channel reduction (softmax/argmax over the class dimension) on the GPU when the
assembled chunk fits free VRAM. On multi-class whole-body segmentation this is the dominant reassembly
cost — on TotalSegmentator:total the argmax alone was ~32 s of a ~77 s run (5 models × a strided
argmax over a 122-class whole-body volume, pathologically slow on CPU).

What does not change

The overlap blend stays on CPU and runs first over the full patch-combined volume (Accumulator.assemble
averages over the class channels across overlapping patches; patches are never all held in VRAM). Only the
already-assembled chunk is moved to the device for the reduction, then the small single-channel argmax result
comes back. Order is unchanged: blend, then argmax, at the end. empty_cache() is called first so the
forward's reserved cache is not mistaken for a full GPU; falls back to CPU when there is no CUDA device or the
chunk (plus headroom) does not fit.

Device plumbing this depends on (also fixed here)

The reduction needs to know the GPU, which exposed two latent gaps:

• NeedDevice.device was declared but never initialised — any object off the .to() propagation path
  raised AttributeError. It now defaults to torch.device("cpu").
• The prediction output datasets were never moved to a device. run_process now co-locates them with the
  model, so their reduction (and any device-aware before_reduction transforms) see the GPU. Note the
  NeedDevice convention stores a CUDA ordinal (int) on GPU and a torch.device on CPU, so the reduction
  normalises accordingly.

Effect (measured, TotalSegmentator:total, 5 models, 1.5 mm, RTX PRO 5000)

• End-to-end loop ~74 s → ~42 s (8.16 → 4.70 s/it); the argmax drops from ~32 s to ~2 s.
• Output vs the CPU path: ~15 differing voxels / 15,068,484 (0.0001 %) — argmax tie-breaking between GPU and
  CPU, the same negligible category as fix(patching): keep overlap-blend reassembly in the patch dtype (halve peak memory) #27/perf(patching): separable blend windows (drop distance-map Cosinus, add Gaussian) #29.

Test plan

• test_predictor_reduction_device: GPU when the chunk fits, CPU fallback when oversized / no CUDA / a CPU
  dataset.
• test_predictor_memory and test_runtime_guards still pass.
• Builds on fix(predict): co-locate load-added modules onto the model's device #26 (device crash), fix(patching): keep overlap-blend reassembly in the patch dtype (halve peak memory) #27 (assemble dtype/OOM), perf(patching): separable blend windows (drop distance-map Cosinus, add Gaussian) #29 (separable blend). The four together take
  TotalSegmentator:total from crashing to running in ~47 s without OOM.

---

Folded-in follow-ups

fix(predict): default OutputDataset.device to CPU on CPU-only runs.
Dataset.__init__ does not forward super().__init__(), so the NeedDevice mixin was never
initialised through the MRO. On a --cpu PREDICTION the output-writer device propagation is
CUDA-gated and never calls .to(), so self.device was unset and the reduction path raised
AttributeError. Fixed by initialising NeedDevice explicitly; adds a regression test that
builds the writer via its real constructor (the existing tests used __new__ and missed it).
This is what turned CI red.

perf(predict): stage the per-patch GPU→CPU offload through pinned memory.
The accumulator keeps every patch of a case until assembly, so each large multi-class patch
(5-model TotalSegmentator: 610 ch, ~2.4 GB) was copied to CPU via a slow pageable .cpu().
Now staged through one reusable page-locked buffer (pinned RAM capped at a single patch), then
copied into a fresh pageable tensor. Bit-identical to .cpu(); small/non-CUDA patches keep
the plain path; falls back to .cpu() if the host can't lock memory.
Measured on TotalSegmentator:total, whole-body CT: 29.7 s → 26.0 s, output byte-identical
(0 / 6.4M voxels differ), +0.5 GB peak RAM; isolated 610-ch D2H 3.4× faster (7.5 s → 2.2 s / 9 patches).