diffusion: Metal-side Stage-1 sampling + Accelerate host sampler

[JoshuaOliphant]

I've been running this PR on an M3 Max (48GB) this week and ended up writing two patches for the Metal side; this PR contributes them to diffusion-visual-updates.

1. Accelerate-vectorized host sampler (examples/diffusion/diffusion.cpp)

The host worker makes three scalar passes over the 262k vocab per canvas position. On Apple platforms I replaced those with vDSP/vvexpf reductions: vDSP_maxvi for the argmax (matches the scalar loop bit for bit, first-index tie-breaking included), one vvexpf pass with entropy computed as H = log Z - dot(e, z)/Z, and a block-wise scan for the multinomial crossing. Non-Apple builds keep the scalar path unchanged; the diff is purely additive. On my machine this took a 256-token canvas step from 434ms to 406ms (Q4_K_M).

2. Metal Stage-1 device sampling via zero-copy host reduction (ggml-metal)

Metal currently warns "on-device sampling unsupported on this backend" and falls back every run. This adds ggml_backend_metal_diffusion_sample behind the same backend-reg proc-address boundary as the CUDA sampler. The twist: it dispatches no kernel. Shared (and mapped) Metal buffers are unified memory, so the implementation reads sc_dev in place and runs the same Accelerate reduction; private-storage tensors return false and fall back to the host path. llama_diffusion_device_sample resolves CUDA first, then Metal, and only caches a successful resolution (late-registered backends under dynamic loading still get picked up; the cache is a relaxed atomic since this is a public llama.h entry point).

I scoped a real MSL port and concluded it buys little here: the CUDA kernel's main win is skipping the 268MB PCIe D2H, and unified memory never pays that tax. The remaining step time on Apple silicon is the forward pass itself.

Hardening found in review

Before submitting I ran a multi-angle review over the diff; three things got fixed as a result, all worth knowing about:

• Multinomial block-scan continuation: pairwise vDSP_sve block sums vs the sequential in-block scan created a wider mis-fallback window than the CUDA kernel's identical-serial slices. The scan now carries the sequential cum into the next block, so the n_vocab-1 default only remains when the cumulative sum truly never reaches the target.
• -inf logit safety: the closed-form entropy NaNs on -inf logits (the scalar path's p > 0 guard didn't). A vDSP_vthr floor at -80 keeps masked-token rows finite; normal rows are unaffected (tail contribution ~1e-30 against Z >= 1).
• Shared prototype: the extern "C" impl declaration lives once in ggml-metal-device.h so the defining and calling TUs cannot drift silently.

Known and accepted: the reduction is duplicated between the example's __APPLE__ worker and the backend file; ggml's public API has no portable host-visibility test (Metal shared buffers deliberately report is_host=false), so the backend boundary is where the is_shared knowledge lives. Happy to unify if you'd rather take a different shape. ggml_vec_soft_max_f32 in ggml-cpu implements the same core pattern but is an internal header, hence not reused here.

Validation

• DG_DEVSAMPLE_CHECK=1: 0 argmax mismatches, 0 sampled-token diffs, 0 entropy delta on every step, at default and t=[0.2,0.5] temperatures, including multi-block (-n 512) runs (both paths share the reduction code, so agreement is exact rather than within tolerance).
• Same-seed runs are byte-identical (the threaded reduction adds no nondeterminism).
• Full build with zero warnings on touched files; -DGGML_METAL=OFF build compiles and links llama-diffusion-cli (dispatcher degrades cleanly).
• Edge audit: n_tokens=1, n_tokens < thread count, u at both boundaries, private-storage fallback, non-F32/non-contiguous rejection.

Numbers (M3 Max 48GB, 256-token canvas, diffusiongemma-26B-A4B Q4_K_M, -fa on)

config	before	after
per-step time	434ms	~400ms
adaptive steps	49.3 tok/s	53.8 tok/s
6-step cap	96.8 tok/s	105.0 tok/s
3-step cap	179.3 tok/s	196.3 tok/s

🤖 Generated with Claude Code

[JoshuaOliphant]

Closing this — it was opened against a fork and hasn't been picked up. Cleaning up my open PRs. Thanks!