Caching only the raw indice map left two thirds of the runtime on the table. Every
layer still rebuilt the missing->N sentinel substitution over 5.6M elements in numpy
and re-uploaded a 22.6MB index to the GPU. Isolating the kernel showed ~13ms of
actual GPU work behind ~26ms of CPU bookkeeping.
The transposed, sentinel-substituted index depends only on the coordinate set - the
same invariant that justifies caching the indice map - so it is now cached on-device
whole.
m3ultra 128^3/128ch, cumulative:
per-offset loop 81.6ms 2.57 Mvox/s
fused gather+matmul 39.3ms 5.34
cached device index 13.8ms 15.25 <- 5.9x total
A chunk-size sweep confirmed 256MB (11 dispatches) is at the optimum; unchunked is
marginally slower (14.2ms), so the chunking is free insurance for small-memory boxes.
Fleet benchmarking exposed the problem: throughput was ANTI-correlated with GPU
core count. The 80-core M3 Ultra came in slowest at 128^3/128ch (81.6ms) behind a
38-core M2 Max (58.3ms), M1 Ultra (66.7ms) and even a 32-core M1 Max (69.2ms).
That ordering only makes sense if the op is bound by dispatch latency rather than
compute - the per-offset loop issued 2*K^3 = 54 tiny GPU ops per layer, none big
enough to occupy the machine, and the Ultra's fused-die design penalises exactly
that.
Concatenating the K^3 neighbour taps along the channel axis collapses it to a
single [N, K^3*Cin] x [K^3*Cin, Cout] matmul. Chunked over rows so peak memory
stays ~256MB (the unchunked buffer is ~2.9GB at 128^3/128ch - fine on a Studio,
not fine on an 8GB mini).
m3ultra 128^3/128ch: 81.6ms -> 39.3ms (2.08x), 2.57 -> 5.34 Mvox/s
m3ultra 64^3/128ch: 21.9ms -> 6.4ms (3.4x)
7/7 tests still pass against the naive reference.
The blocker for LATO.2 on Apple Silicon is one op, not the whole setup.sh --all
CUDA stack. Measured: 5 of 7 checkpoints are fully dense, and every SparseConv3d
in the model is constructed stride=1/padding=None, which upstream dispatches to
spconv's SubMConv3d. No strided or inverse sparse conv is ever instantiated.
- SubMConv3d in pure MLX via a sorted-key indice map (27 lookups/voxel vectorised,
cached per coordinate set the way spconv uses indice_key)
- SparseTensor container + subdivide upsampling
- Converter handles the 5-D layout collision: spconv KRSC [O,kz,ky,kx,I] vs torch
Conv3d [O,I,kz,ky,kx]. Rank alone is ambiguous; misreading it silently mangles
the voxel encoder.
- 7/7 tests pass vs an independent naive reference, max err 3e-7. spconv has no
Metal build so there is no upstream oracle; the reference shares no indexing code.
Kernel orientation (feats[c+d] vs feats[c-d]) remains unverified and is silent
when wrong; --flip-kernel builds the mirror for an end-to-end A/B.
